Skip to content

Commit

Permalink
refactor(rust): Ensure fallback node gets correct length df even if n…
Browse files Browse the repository at this point in the history
…o columns selected (#18772)
  • Loading branch information
orlp authored Sep 16, 2024
1 parent f7b6d86 commit 6561eba
Showing 1 changed file with 8 additions and 1 deletion.
9 changes: 8 additions & 1 deletion crates/polars-stream/src/physical_plan/lower_expr.rs
Original file line number Diff line number Diff line change
Expand Up @@ -299,10 +299,17 @@ fn build_fallback_node_with_ctx(
) -> PolarsResult<PhysNodeKey> {
// Pre-select only the columns that are needed for this fallback expression.
let input_schema = &ctx.phys_sm[input].output_schema;
let select_names: PlHashSet<_> = exprs
let mut select_names: PlHashSet<_> = exprs
.iter()
.flat_map(|expr| polars_plan::utils::aexpr_to_leaf_names_iter(expr.node(), ctx.expr_arena))
.collect();
// To keep the length correct we have to ensure we select at least one
// column.
if select_names.is_empty() {
if let Some(name) = input_schema.iter_names().next() {
select_names.insert(name.clone());
}
}
let input_node = if input_schema
.iter_names()
.any(|name| !select_names.contains(name.as_str()))
Expand Down

0 comments on commit 6561eba

Please sign in to comment.