Fix shared-subplan merges across outer joins and limits, and cuDF string MIN/MAX beside NUNIQUE

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::<i32>().unwrap();
+    let nunique_data = result.series("value_nunique").unwrap().to_vec_strict::<i32>().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::<Vec<_>>();
+    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),

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"]);

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,15 +176,18 @@ 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],
             };
             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(_)

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

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 {

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