From 6561ebab214695e76695ff77435169b91b29079b Mon Sep 17 00:00:00 2001 From: Orson Peters Date: Mon, 16 Sep 2024 16:29:37 +0200 Subject: [PATCH] refactor(rust): Ensure fallback node gets correct length df even if no columns selected (#18772) --- crates/polars-stream/src/physical_plan/lower_expr.rs | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/crates/polars-stream/src/physical_plan/lower_expr.rs b/crates/polars-stream/src/physical_plan/lower_expr.rs index 62ccdef807c4..8e891bd408af 100644 --- a/crates/polars-stream/src/physical_plan/lower_expr.rs +++ b/crates/polars-stream/src/physical_plan/lower_expr.rs @@ -299,10 +299,17 @@ fn build_fallback_node_with_ctx( ) -> PolarsResult { // 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()))