Skip to content

Commit

Permalink
fix(rust): Window function had incorrect output name on ExprIR (#18970)
Browse files Browse the repository at this point in the history
  • Loading branch information
orlp authored Sep 27, 2024
1 parent 8bb5a02 commit 538fd6c
Show file tree
Hide file tree
Showing 2 changed files with 13 additions and 7 deletions.
4 changes: 3 additions & 1 deletion crates/polars-plan/src/plans/conversion/expr_to_ir.rs
Original file line number Diff line number Diff line change
Expand Up @@ -304,14 +304,16 @@ pub(super) fn to_aexpr_impl(
order_by,
options,
} => {
// Process function first so name is correct.
let function = to_aexpr_impl(owned(function), arena, state)?;
let order_by = if let Some((e, options)) = order_by {
Some((to_aexpr_impl(owned(e.clone()), arena, state)?, options))
} else {
None
};

AExpr::Window {
function: to_aexpr_impl(owned(function), arena, state)?,
function,
partition_by: to_aexprs(partition_by, arena, state)?,
order_by,
options,
Expand Down
16 changes: 10 additions & 6 deletions crates/polars-plan/src/plans/expr_ir.rs
Original file line number Diff line number Diff line change
Expand Up @@ -28,17 +28,21 @@ pub enum OutputName {
}

impl OutputName {
pub fn unwrap(&self) -> &PlSmallStr {
pub fn get(&self) -> Option<&PlSmallStr> {
match self {
OutputName::Alias(name) => name,
OutputName::ColumnLhs(name) => name,
OutputName::LiteralLhs(name) => name,
OutputName::Alias(name) => Some(name),
OutputName::ColumnLhs(name) => Some(name),
OutputName::LiteralLhs(name) => Some(name),
#[cfg(feature = "dtype-struct")]
OutputName::Field(name) => name,
OutputName::None => panic!("no output name set"),
OutputName::Field(name) => Some(name),
OutputName::None => None,
}
}

pub fn unwrap(&self) -> &PlSmallStr {
self.get().expect("no output name set")
}

pub(crate) fn is_none(&self) -> bool {
matches!(self, OutputName::None)
}
Expand Down

0 comments on commit 538fd6c

Please sign in to comment.