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
  • Loading branch information
orlp committed Sep 16, 2024
1 parent a53bf03 commit 45dac76
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 @@ -313,10 +313,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 45dac76

Please sign in to comment.