From 97ea333d6a47ecaa3a2165417146365dab21ad59 Mon Sep 17 00:00:00 2001 From: Morris Tai Date: Thu, 22 Oct 2026 04:21:56 +0000 Subject: [PATCH 1/2] fix: stop shared-subplan lineage at outer-join null sides and limits The shared_subplan_materialization rule treats a grouped branch's reader predicate as post-aggregate filterable when the group column is a source column carried unchanged to the aggregate. materialized_column_lineage followed a join output to its input regardless of join type, so a predicate on the null-supplying input of an outer join was re-applied after the shared aggregate and dropped the NULL group each branch keeps. It also followed a Limit, so two LIMIT branches over different predicates were merged into one limit over the widened rows, which keeps a different prefix. JoinType::passes_through_rows_of names the inputs a join passes through unchanged (both for inner and cross, the left for left outer, semi, anti and mark, neither for full). Lineage now stops at a join output from any other input and at a Limit. TPC-H and TPC-DS applied merges are unchanged. Co-Authored-By: Claude Fable 5.1 --- .../shared_subplan_materialization.rs | 210 ++++++++++++++++++ .../shared_subplan_materialization/lineage.rs | 13 +- .../shared_subplan_materialization/mod.rs | 4 +- crates/algeon-query-planning/src/plan/join.rs | 17 ++ 4 files changed, 241 insertions(+), 3 deletions(-) diff --git a/crates/algeon-query-engine/tests/native_contract_tests_cpu/optimizer_contracts/shared_subplan_materialization.rs b/crates/algeon-query-engine/tests/native_contract_tests_cpu/optimizer_contracts/shared_subplan_materialization.rs index b64c817f1..dfcd13ba7 100644 --- a/crates/algeon-query-engine/tests/native_contract_tests_cpu/optimizer_contracts/shared_subplan_materialization.rs +++ b/crates/algeon-query-engine/tests/native_contract_tests_cpu/optimizer_contracts/shared_subplan_materialization.rs @@ -696,6 +696,14 @@ fn t_shared_subplan_rejects_cases() { "optimizer_contracts::shared_subplan_materialization::native_optimizer_rejects_grouped_materialization_when_predicate_column_has_different_group_lineage", native_optimizer_rejects_grouped_materialization_when_predicate_column_has_different_group_lineage, ), + ( + "optimizer_contracts::shared_subplan_materialization::native_optimizer_rejects_grouped_materialization_when_predicate_column_is_null_supplied_by_outer_join", + native_optimizer_rejects_grouped_materialization_when_predicate_column_is_null_supplied_by_outer_join, + ), + ( + "optimizer_contracts::shared_subplan_materialization::native_optimizer_rejects_grouped_materialization_when_predicate_is_read_below_a_limit", + native_optimizer_rejects_grouped_materialization_when_predicate_is_read_below_a_limit, + ), ( "optimizer_contracts::shared_subplan_materialization::native_optimizer_rejects_grouped_materialization_with_duplicate_source_occurrences", native_optimizer_rejects_grouped_materialization_with_duplicate_source_occurrences, @@ -886,6 +894,96 @@ fn native_optimizer_rejects_grouped_materialization_when_predicate_column_has_di ); } +/// A reader predicate on the null-supplying input of an outer join also +/// decides which preserved rows are padded with nulls, so re-applying it to +/// the shared aggregate would drop the NULL group that each branch keeps. +fn native_optimizer_rejects_grouped_materialization_when_predicate_column_is_null_supplied_by_outer_join() { + for (join_type, predicate_side) in [(JoinType::Left, JoinSide::Right), (JoinType::Full, JoinSide::Left)] { + let case = format!("{join_type} join, predicate on {predicate_side}"); + let plan = repeated_grouped_sum_over_join_with_side_predicate_plan(join_type, predicate_side); + let options = NativeOptimizerOptions { + shared_subplan_materialization: SharedSubplanMaterializationOptions::default(), + ..NativeOptimizerOptions::disabled() + }; + + let result = optimize_query_plan( + plan.clone(), + &options, + rapids_analysis_context(PlanningPolicy::default()), + local_dimension_pruning_resolver(), + ) + .unwrap(); + + assert_eq!(result.plan, plan, "{case}"); + let diagnostics = &result.diagnostics.shared_subplan_materialization; + assert_eq!(diagnostics.candidate_count, 1, "{case}"); + assert_eq!(diagnostics.applied_count, 0, "{case}"); + assert_eq!( + diagnostics.candidates[0].rejected_reason, + Some(SharedSubplanMaterializationRejectedReason::PredicateNotPostAggregateFilterable), + "{case}" + ); + } +} + +/// `LIMIT 10` over rows with `g = 1` keeps ten such rows; a limit over the +/// widened `g = 1 OR g = 2` rows filtered to `g = 1` afterwards may keep fewer. +fn native_optimizer_rejects_grouped_materialization_when_predicate_is_read_below_a_limit() { + let plan = repeated_grouped_count_over_limited_predicate_sources_plan(); + let options = NativeOptimizerOptions { + shared_subplan_materialization: SharedSubplanMaterializationOptions::default(), + ..NativeOptimizerOptions::disabled() + }; + + let result = optimize_query_plan( + plan.clone(), + &options, + rapids_analysis_context(PlanningPolicy::default()), + local_dimension_pruning_resolver(), + ) + .unwrap(); + + assert_eq!(result.plan, plan); + let diagnostics = &result.diagnostics.shared_subplan_materialization; + assert_eq!(diagnostics.candidate_count, 1); + assert_eq!(diagnostics.applied_count, 0); + assert_eq!( + diagnostics.candidates[0].rejected_reason, + Some(SharedSubplanMaterializationRejectedReason::PredicateNotPostAggregateFilterable) + ); +} + +/// The same branches combine when the predicate input's rows reach the join +/// output unchanged: either input of an inner join, or the preserved input of +/// a left join. +#[test] +fn native_optimizer_combines_grouped_join_branches_with_predicate_on_preserved_input() { + for (join_type, predicate_side) in [(JoinType::Inner, JoinSide::Right), (JoinType::Left, JoinSide::Left)] { + let case = format!("{join_type} join, predicate on {predicate_side}"); + let plan = repeated_grouped_sum_over_join_with_side_predicate_plan(join_type, predicate_side); + let options = NativeOptimizerOptions { + shared_subplan_materialization: SharedSubplanMaterializationOptions::default(), + ..NativeOptimizerOptions::disabled() + }; + + let result = optimize_query_plan( + plan, + &options, + rapids_analysis_context(PlanningPolicy::default()), + local_dimension_pruning_resolver(), + ) + .unwrap(); + + let diagnostics = &result.diagnostics.shared_subplan_materialization; + assert_eq!(diagnostics.applied_count, 1, "{case}"); + assert_eq!( + result.plan.nodes().iter().filter(|node| matches!(node.kind, PlanNode::Aggregate(_))).count(), + 1, + "{case}" + ); + } +} + fn native_optimizer_rejects_grouped_materialization_with_duplicate_source_occurrences() { let plan = repeated_grouped_aggregate_self_join_duplicate_source_plan(); let options = NativeOptimizerOptions { @@ -1635,6 +1733,118 @@ fn repeated_grouped_aggregate_predicate_group_name_lineage_collision_plan() -> Q builder.finish_gpu_resident(union).unwrap() } +/// Two `SELECT g, SUM(amount) ... GROUP BY g` branches over one join whose +/// `predicate_side` input reads `g` with the reader predicate `g = 1` or +/// `g = 2`; the other input supplies `amount`. +fn repeated_grouped_sum_over_join_with_side_predicate_plan(join_type: JoinType, predicate_side: JoinSide) -> QueryPlan { + let nullable_outputs = match join_type { + JoinType::Left => [predicate_side == JoinSide::Right, predicate_side == JoinSide::Left], + JoinType::Full => [true, true], + _ => [false, false], + }; + let join_schema = Arc::new(Schema::new(vec![ + Field::new("g", DataType::Int32, nullable_outputs[0]), + Field::new("amount", DataType::Int32, nullable_outputs[1]), + ])); + let aggregate_schema = Arc::new(Schema::new(vec![ + Field::new("g", DataType::Int32, nullable_outputs[0]), + Field::new("total", DataType::Int64, true), + ])); + let mut aggregates = Vec::new(); + let mut builder = QueryPlanBuilder::new(); + for value in 1..=2 { + let filtered = builder.add_source(estimated_local_source_with_stable_id_and_predicate( + "join_group_input", + "join_group_input", + &["k", "g"], + 512 * 1024 * 1024, + 10_000_000, + Expr::binary(Expr::column("g"), BinaryOp::Eq, Expr::literal(LiteralValue::Int32(value))), + )); + let other = builder.add_source(estimated_local_source_with_stable_id( + "join_amount_input", + "join_amount_input", + &["k", "amount"], + 512 * 1024 * 1024, + 10_000_000, + )); + let (left, right) = match predicate_side { + JoinSide::Left => (filtered, other), + JoinSide::Right => (other, filtered), + }; + let predicate_index = usize::from(predicate_side == JoinSide::Right); + let other_side = match predicate_side { + JoinSide::Left => JoinSide::Right, + JoinSide::Right => JoinSide::Left, + }; + let joined = builder + .add_join( + left, + right, + JoinSpec::new( + join_type.clone(), + vec![JoinKey::new(ColumnRef::input(0, "k", "k"), ColumnRef::input(1, "k", "k"))], + ) + .with_output_projection(vec![ + JoinOutputColumn::new(predicate_side, ColumnRef::input(predicate_index, "g", "g"), "g"), + JoinOutputColumn::new( + other_side, + ColumnRef::input(1 - predicate_index, "amount", "amount"), + "amount", + ), + ]), + Arc::clone(&join_schema), + ) + .unwrap(); + aggregates.push( + builder + .add_aggregate( + joined, + AggregateSpec::new( + vec![Expr::column("g")], + vec![AggregateExpr::new(AggregateFunction::Sum, Expr::column("amount"), "total")], + ), + Arc::clone(&aggregate_schema), + ) + .unwrap(), + ); + } + let union = builder.add_union(aggregates, aggregate_schema).unwrap(); + builder.finish_gpu_resident(union).unwrap() +} + +fn repeated_grouped_count_over_limited_predicate_sources_plan() -> QueryPlan { + let mut builder = QueryPlanBuilder::new(); + let aggregate_schema = + Arc::new(Schema::new(vec![Field::new("g", DataType::Int32, false), Field::new("n", DataType::Int64, true)])); + let mut aggregates = Vec::new(); + for value in 1..=2 { + let source = builder.add_source(estimated_local_source_with_stable_id_and_predicate( + "limited", + "limited", + &["g", "a"], + 512 * 1024 * 1024, + 10_000_000, + Expr::binary(Expr::column("g"), BinaryOp::Eq, Expr::literal(LiteralValue::Int32(value))), + )); + let limited = builder.add_limit(source, LimitSpec::fetch(10)).unwrap(); + aggregates.push( + builder + .add_aggregate( + limited, + AggregateSpec::new( + vec![Expr::column("g")], + vec![AggregateExpr::new(AggregateFunction::Count, Expr::literal(LiteralValue::Int64(1)), "n")], + ), + Arc::clone(&aggregate_schema), + ) + .unwrap(), + ); + } + let union = builder.add_union(aggregates, aggregate_schema).unwrap(); + builder.finish_gpu_resident(union).unwrap() +} + fn repeated_grouped_aggregate_self_join_duplicate_source_plan() -> QueryPlan { let mut builder = QueryPlanBuilder::new(); let join_schema = schema(&["customer_id", "sales"]); diff --git a/crates/algeon-query-planning/src/optimizer/shared_subplan_materialization/lineage.rs b/crates/algeon-query-planning/src/optimizer/shared_subplan_materialization/lineage.rs index 5f12eb9c3..229651a54 100644 --- a/crates/algeon-query-planning/src/optimizer/shared_subplan_materialization/lineage.rs +++ b/crates/algeon-query-planning/src/optimizer/shared_subplan_materialization/lineage.rs @@ -145,6 +145,12 @@ pub(super) fn scalar_input_expr_lineage_signature(plan: &QueryPlan, node_id: Nod } } +/// Source column that `column` at `node_id` carries unchanged, through +/// operators across which a reader predicate over that column selects the same +/// rows before or after them. A limit stops the lineage, since it keeps a +/// different prefix of the widened rows. A join carries a column only from an +/// input whose rows it passes through +/// ([`JoinType::passes_through_rows_of`](crate::plan::JoinType::passes_through_rows_of)). pub(super) fn materialized_column_lineage( plan: &QueryPlan, node_id: NodeId, @@ -170,6 +176,9 @@ pub(super) fn materialized_column_lineage( }, PlanNode::Join(join) => { let output = join.spec.output_projection.iter().find(|output| output.output_name == column)?; + if !join.spec.join_type.passes_through_rows_of(output.side) { + return None; + } let input = match output.side { JoinSide::Left => join.inputs[0], JoinSide::Right => join.inputs[1], @@ -177,8 +186,8 @@ pub(super) fn materialized_column_lineage( materialized_column_lineage(plan, input, output.column.name()) }, PlanNode::Sort(sort) => materialized_column_lineage(plan, sort.input, column), - PlanNode::Limit(limit) => materialized_column_lineage(plan, limit.input, column), - PlanNode::Aggregate(_) + PlanNode::Limit(_) + | PlanNode::Aggregate(_) | PlanNode::Window(_) | PlanNode::Union(_) | PlanNode::EdgeNormalize(_) diff --git a/crates/algeon-query-planning/src/optimizer/shared_subplan_materialization/mod.rs b/crates/algeon-query-planning/src/optimizer/shared_subplan_materialization/mod.rs index 913d51baf..f028225c4 100644 --- a/crates/algeon-query-planning/src/optimizer/shared_subplan_materialization/mod.rs +++ b/crates/algeon-query-planning/src/optimizer/shared_subplan_materialization/mod.rs @@ -11,7 +11,9 @@ //! aggregate (grouped) or a projection that renames its branch's `FILTER`ed //! aggregate outputs (scalar). As a //! precondition, a grouped branch predicate must be evaluable after the -//! aggregate, over group columns the source emits unchanged; a scalar branch +//! aggregate, over group columns the source emits unchanged through filters, +//! sorts, column projections and joins that pass its rows through (no limit, +//! and not the null-supplying input of an outer join); a scalar branch //! predicate must reference source columns the widened subtree can carry to the //! aggregate input. A family that fails either is reported with a //! [`SharedSubplanMaterializationRejectedReason`] and left as it is. Facts: a diff --git a/crates/algeon-query-planning/src/plan/join.rs b/crates/algeon-query-planning/src/plan/join.rs index e3543159f..85659c07e 100644 --- a/crates/algeon-query-planning/src/plan/join.rs +++ b/crates/algeon-query-planning/src/plan/join.rs @@ -296,6 +296,23 @@ impl JoinType { matches!(self, Self::LeftAnti { null_aware: true }) } + /// Whether the join passes rows of its `side` input through unchanged, so + /// that a predicate over `side` columns selects the same output rows when + /// it filters that input before the join as when it filters the join + /// output. Inner and cross joins pass both inputs through; left outer, + /// semi, anti and mark joins pass the left input through. The + /// null-supplying input of an outer join does not qualify, because which of + /// its rows survive also decides which rows of the other input are padded + /// with nulls. The right input of semi, anti and mark joins does not + /// qualify either, because its columns never reach the output. + pub(crate) fn passes_through_rows_of(&self, side: JoinSide) -> bool { + match self { + Self::Inner | Self::Cross => true, + Self::Left | Self::LeftSemi | Self::LeftAnti { .. } | Self::LeftMark { .. } => side == JoinSide::Left, + Self::Full => false, + } + } + /// The mark column of a `LeftMark` join. pub fn mark_column(&self) -> Option<&str> { match self { From e85a3fe0ec765632db306e6e006eacbf25c79328 Mon Sep 17 00:00:00 2001 From: Morris Tai Date: Thu, 22 Oct 2026 05:12:51 +0000 Subject: [PATCH 2/2] fix(groupby): keep grouped values in key sort order beside sorted aggregations store_result_functor::get_grouped_values() returned the cached within-group sorted values once NUNIQUE, MEDIAN or a quantile had built them. ARGMIN and ARGMAX map a position in the grouped values back to a row through the key sort order, and NTH_ELEMENT and COLLECT_LIST read the rows of a group in that order, so they received the wrong rows: MIN and MAX of a string column returned the group's first and last row in key order. Grouped values are now always the values gathered by the key sort order; the scan override that existed to avoid the substitution is removed. Found by a DuckDB cross-check of an Algeon query combining COUNT(DISTINCT x), MIN(x) and median(y) over an Open Targets table. --- .../ops/groupby/high_level_groupby_tests.rs | 40 ++++++++++++++++++- components/cudf/cudf | 2 +- website/docs/reference/forks.md | 5 +++ 3 files changed, 45 insertions(+), 2 deletions(-) diff --git a/components/cudf/algeon-cudf/tests/semantic_contracts_gpu/ops/groupby/high_level_groupby_tests.rs b/components/cudf/algeon-cudf/tests/semantic_contracts_gpu/ops/groupby/high_level_groupby_tests.rs index 237247cb4..d15fe1f23 100644 --- a/components/cudf/algeon-cudf/tests/semantic_contracts_gpu/ops/groupby/high_level_groupby_tests.rs +++ b/components/cudf/algeon-cudf/tests/semantic_contracts_gpu/ops/groupby/high_level_groupby_tests.rs @@ -1,6 +1,8 @@ use cudf::{AggSpec, DataFrame, ErrorCode, ErrorKind, Exec, GroupByAgg}; -use crate::common::{assert_missing_column_error, column_from_slice, dataframe_from_columns, error_context_value}; +use crate::common::{ + assert_missing_column_error, column_from_slice, column_from_strs, dataframe_from_columns, error_context_value, +}; fn assert_groupby_validation_error(err: &cudf::Error, operation: &str, reason: &str, remedy: &str) { assert_eq!(err.kind(), ErrorKind::Unsupported); @@ -393,6 +395,38 @@ fn groupby_nunique_shortcut() { assert_eq!(pairs, vec![(1, 2), (2, 1)]); } +fn groupby_string_min_max_beside_nunique_on_same_column() { + // NUNIQUE sorts each group's values. MIN and MAX of a string column read + // row positions, so they must not take those within-group sorted values. + // No group's first or last row holds its minimum or maximum. + let keys = column_from_slice(&[1i32, 2, 1, 2, 1, 2]).unwrap(); + let vals = column_from_strs(&["m", "q", "z", "x", "a", "c"]).unwrap(); + let df = dataframe_from_columns(vec!["key".into(), "value".into()], vec![keys, vals]).unwrap(); + + let result = df + .groupby(["key"]) + .on(Exec::RETAINED) + .agg(&[ + AggSpec::new("value", GroupByAgg::NUnique), + AggSpec::new("value", GroupByAgg::Min), + AggSpec::new("value", GroupByAgg::Max), + ]) + .unwrap(); + + let key_data = result.series("key").unwrap().to_vec_strict::().unwrap(); + let nunique_data = result.series("value_nunique").unwrap().to_vec_strict::().unwrap(); + let min_data = result.series("value_min").unwrap().to_strings_strict().unwrap(); + let max_data = result.series("value_max").unwrap().to_strings_strict().unwrap(); + let mut rows = key_data + .into_iter() + .zip(nunique_data) + .zip(min_data.into_iter().zip(max_data)) + .map(|((key, nunique), (min, max))| (key, nunique, min, max)) + .collect::>(); + rows.sort_by_key(|row| row.0); + assert_eq!(rows, vec![(1, 3, "a".to_string(), "z".to_string()), (2, 3, "c".to_string(), "x".to_string())]); +} + fn groupby_std_ddof1() { // keys: [1, 1, 2, 2], values: [10.0, 20.0, 100.0, 200.0] // key=1: std(ddof=1) = sqrt(((10-15)^2 + (20-15)^2) / 1) = sqrt(50) ≈ 7.0711 @@ -614,6 +648,10 @@ pub(super) const CASES: &[(&str, fn())] = &[ ("high_level_groupby_tests::groupby_median_shortcut", groupby_median_shortcut), ("high_level_groupby_tests::groupby_nunique", groupby_nunique), ("high_level_groupby_tests::groupby_nunique_shortcut", groupby_nunique_shortcut), + ( + "high_level_groupby_tests::groupby_string_min_max_beside_nunique_on_same_column", + groupby_string_min_max_beside_nunique_on_same_column, + ), ("high_level_groupby_tests::groupby_std_ddof1", groupby_std_ddof1), ("high_level_groupby_tests::groupby_std_ddof0", groupby_std_ddof0), ("high_level_groupby_tests::groupby_std_shortcut", groupby_std_shortcut), diff --git a/components/cudf/cudf b/components/cudf/cudf index 00e395cf1..496993804 160000 --- a/components/cudf/cudf +++ b/components/cudf/cudf @@ -1 +1 @@ -Subproject commit 00e395cf1aa09f68a97eb2a7547edfc1cc231c16 +Subproject commit 4969938047d5294ffa422c2ca0a9533130c5eb17 diff --git a/website/docs/reference/forks.md b/website/docs/reference/forks.md index d156a7632..1e30abde0 100644 --- a/website/docs/reference/forks.md +++ b/website/docs/reference/forks.md @@ -54,6 +54,11 @@ Bug fixes: reader filter no longer crashes, the stats survivor count is reported only when the stats stage ran, the decimal literal type check honors scale, and decimal fixed-length dictionary values are decoded big-endian. +- Sort-based groupby: after NUNIQUE, MEDIAN or a quantile on a column, + ARGMIN, ARGMAX, NTH_ELEMENT and COLLECT_LIST on the same column read the + within-group sorted values in place of the grouped values. ARGMIN and + ARGMAX, and with them MIN and MAX of a string column, returned another row + of the group. Grouped values now always keep the key sort order. - cuco exception safety, applied as a build-time patch: temporary storage is neither leaked nor double-freed when an allocation throws. - A GCC 14.2 internal compiler error in `make_datasources` is avoided by