feat(sql): count_if, boolean aggregates, any_value, max_by and min_by

crates/krishiv-exec/src/aggregate/accumulator/differential.rs

@@ -330,3 +330,137 @@ fn min_max_on_decimal64() {
         assert_eq!(v.as_primitive::<Decimal64Type>().value(0), want);
     }
 }
+
+/// Keys of the `max_by`/`min_by` merge test: NaN (the largest), the
+/// infinities, and `-0.0 = 0.0` (a tie).
+const KEYS: [f64; 7] = [
+    f64::NAN,
+    f64::NEG_INFINITY,
+    -1.0,
+    -0.0,
+    0.0,
+    1.0,
+    f64::INFINITY,
+];
+
+/// `(group, value, key)`: a NULL value is legal, a NULL key skips the row.
+type ByRow = (u32, Option<i64>, Option<u8>);
+
+/// The value and key columns of `rows`, as strings (`strings`) or as
+/// `bigint`/`double` (keys by their index in strings, so ties stay ties).
+fn by_columns(rows: &[ByRow], strings: (bool, bool)) -> Vec<ArrayRef> {
+    let x: ArrayRef = if strings.0 {
+        Arc::new(StringViewArray::from_iter(
+            rows.iter().map(|r| r.1.map(|v| format!("v{v:03}"))),
+        ))
+    } else {
+        Arc::new(Int64Array::from_iter(rows.iter().map(|r| r.1)))
+    };
+    let y: ArrayRef = if strings.1 {
+        Arc::new(StringViewArray::from_iter(
+            rows.iter().map(|r| r.2.map(|k| format!("k{k}"))),
+        ))
+    } else {
+        Arc::new(Float64Array::from_iter(
+            rows.iter().map(|r| r.2.map(|k| KEYS[usize::from(k)])),
+        ))
+    };
+    vec![x, y]
+}
+
+fn by_expr(f: AggregateFunction, strings: (bool, bool)) -> AggregateExpr {
+    let x = if strings.0 {
+        DataType::Utf8View
+    } else {
+        DataType::Int64
+    };
+    let y = if strings.1 {
+        DataType::Utf8View
+    } else {
+        DataType::Float64
+    };
+    match f {
+        AggregateFunction::MaxBy | AggregateFunction::MinBy => {
+            AggregateExpr::by(f, (0, 1), (x, y), "b")
+        }
+        _ => AggregateExpr::new(f, 0, x, "a"),
+    }
+}
+
+proptest! {
+    #![proptest_config(ProptestConfig {
+        cases: 256,
+        failure_persistence: None,
+        ..ProptestConfig::default()
+    })]
+
+    /// `max_by`, `min_by` and `any_value`: the input cut anywhere into
+    /// consecutive parts, each part's partial states merged in order (also
+    /// through a second level), give the single pass's states and result,
+    /// ties and NULLs included.
+    #[test]
+    fn by_merge_differential(
+        rows in proptest::collection::vec(
+            (0..4_u32, proptest::option::of(0..40_i64), proptest::option::of(0..7_u8)),
+            0..60,
+        ),
+        cuts in proptest::collection::vec(0..61_usize, 0..4),
+        strings in (any::<bool>(), any::<bool>()),
+    ) {
+        let groups = 4_u32;
+        let ids: Vec<u32> = (0..groups).collect();
+        let mut cuts: Vec<usize> = cuts.into_iter().map(|c| c.min(rows.len())).collect();
+        cuts.push(0);
+        cuts.push(rows.len());
+        cuts.sort_unstable();
+        let fold = |range: &[ByRow]| {
+            (range.iter().map(|r| r.0).collect::<Vec<u32>>(), by_columns(range, strings))
+        };
+        for f in [
+            AggregateFunction::MaxBy,
+            AggregateFunction::MinBy,
+            AggregateFunction::AnyValue { ignore_nulls: false },
+            AggregateFunction::AnyValue { ignore_nulls: true },
+        ] {
+            let expr = by_expr(f, strings);
+            let new = || {
+                let mut a = Accumulator::new(&expr, true).unwrap();
+                a.resize(groups as usize);
+                a
+            };
+            let mut single = new();
+            let (g, cols) = fold(&rows);
+            single.update(&g, &cols, &Weights::Const(1)).unwrap();
+            let partials: Vec<Vec<ArrayRef>> = cuts
+                .windows(2)
+                .map(|w| {
+                    let mut a = new();
+                    let (g, cols) = fold(&rows[w[0]..w[1]]);
+                    a.update(&g, &cols, &Weights::Const(1)).unwrap();
+                    a.state(&ids).unwrap()
+                })
+                .collect();
+            let mut flat = new();
+            for st in &partials {
+                flat.merge(&ids, st).unwrap();
+            }
+            // Two levels: the halves merged apart, then together.
+            let half = partials.len() / 2;
+            let mut top = new();
+            for part in [&partials[..half], &partials[half..]] {
+                let mut mid = new();
+                for st in part {
+                    mid.merge(&ids, st).unwrap();
+                }
+                top.merge(&ids, &mid.state(&ids).unwrap()).unwrap();
+            }
+            let want = (single.state(&ids).unwrap(), single.evaluate(&ids).unwrap());
+            for merged in [&flat, &top] {
+                prop_assert_eq!(
+                    &(merged.state(&ids).unwrap(), merged.evaluate(&ids).unwrap()),
+                    &want
+                );
+            }
+        }
+    }
+}

crates/krishiv-exec/src/aggregate/mod.rs

@@ -1,6 +1,9 @@
 // SPDX-License-Identifier: Apache-2.0
-//! Hash aggregation (ADR 0011): `count`, `sum`, `avg`, `min`, `max` with
-//! Spark 4.2 result types, in two phases, with spill and `DISTINCT`.
+//! Hash aggregation (ADR 0011): `count`, `sum`, `avg`, `min`, `max`,
+//! `first`, `last`, `any_value`, `max_by`, `min_by` with Spark 4.2 result
+//! types, in two phases, with spill and `DISTINCT`. `count_if` and the
+//! boolean aggregates are `count` and `min`/`max` over a rewritten argument
+//! (the physical planner).
 //!
 //! # Operators
 //! - [`HashAggregate`] is a pipeline breaker. Each driver aggregates into
@@ -49,9 +52,10 @@
 //! # Weights
 //! Rows carry Z-set weights ([`Weights`](krishiv_types::batch::Weights)):
 //! `count` adds the weight, `sum` adds `value × weight`, `avg` divides the
-//! weighted sum by the weighted count. `min`/`max` need retractions to be
-//! maintained incrementally (P4 IVM), so they and `DISTINCT` reject any weight
-//! other than 1 with `KRISHIV_UNSUPPORTED_FEATURE`. Output rows have weight 1;
+//! weighted sum by the weighted count. `min`/`max`, `first`, `last`,
+//! `any_value`, `max_by` and `min_by` need retractions to be maintained
+//! incrementally (P4 IVM), so they and `DISTINCT` reject any weight other
+//! than 1 with `KRISHIV_UNSUPPORTED_FEATURE`. Output rows have weight 1;
 //! a group exists once it received a row.
 
 mod accumulator;
@@ -66,7 +70,7 @@ mod tests;
 
 use std::sync::Arc;
 
-use arrow_schema::{DataType, Field, Schema, SchemaRef};
+use arrow_schema::{DataType, Field, Fields, Schema, SchemaRef};
 use krishiv_common::error::{
     ArithmeticOverflow, KrishivError, KrishivUnsupportedFeature,
     NumericValueOutOfRangeWithSuggestion,
@@ -112,6 +116,23 @@ pub enum AggregateFunction {
         /// `last(x, true)` / `IGNORE NULLS`: the last non-NULL value.
         ignore_nulls: bool,
     },
+    /// `any_value(x)`: some row's value of the group; here the first row's,
+    /// as `first` (a legal choice); NULL for no value. The input type.
+    /// In-order merges keep the first row in input order; a `DISTINCT`
+    /// first phase groups in hash order.
+    AnyValue {
+        /// `any_value(x, true)` / `IGNORE NULLS`: the first non-NULL value.
+        ignore_nulls: bool,
+    },
+    /// `max_by(x, y)`: `x` of the row with the largest non-NULL `y` (Spark
+    /// ordering), the first such row on a tie for in-order merges; rows
+    /// with a NULL `y` are ignored and a NULL `x` is a result. NULL for no
+    /// row. The type of `x`. Built with [`AggregateExpr::by`]. A
+    /// `DISTINCT` first phase groups in hash order, so a tie there keeps
+    /// the first row in hash order, not in input order.
+    MaxBy,
+    /// `min_by(x, y)`, as `max_by` for the smallest `y` (same tie caveat).
+    MinBy,
 }
 
 /// Where an aggregate's arguments come from.
@@ -142,7 +163,9 @@ pub struct AggregateExpr {
     /// Where the arguments come from.
     pub input: AggInput,
     /// The type of the raw argument (`count` ignores it). Fixed by the
-    /// planner; state columns derive from it.
+    /// planner; state columns derive from it. For `max_by`/`min_by`, a
+    /// struct of the value's type then the ordering key's (see
+    /// [`AggregateExpr::by`]).
     pub arg_type: DataType,
     /// The output column name.
     pub name: String,
@@ -165,6 +188,27 @@ impl AggregateExpr {
         }
     }
 
+    /// `max_by(x, y)` / `min_by(x, y)` of columns `value` and `key`, of
+    /// types `value_type` and `key_type`. Both must be types `min`/`max`
+    /// take.
+    pub fn by(
+        function: AggregateFunction,
+        (value, key): (usize, usize),
+        (value_type, key_type): (DataType, DataType),
+        name: impl Into<String>,
+    ) -> Self {
+        Self {
+            function,
+            args: vec![value, key],
+            input: AggInput::Raw,
+            arg_type: DataType::Struct(Fields::from(vec![
+                Field::new("value", value_type, true),
+                Field::new("key", key_type, true),
+            ])),
+            name: name.into(),
+        }
+    }
+
     /// `count(*)`.
     pub fn count_star(name: impl Into<String>) -> Self {
         Self {
@@ -207,11 +251,24 @@ impl AggregateExpr {
             (F::Avg, DataType::Decimal64(p, s) | DataType::Decimal128(p, s)) if *s >= 0 => {
                 decimal_on(t, (*p + 4).min(38), s.saturating_add(4).min(38))
             }
-            (F::Min | F::Max | F::First { .. } | F::Last { .. }, t)
+            (F::Min | F::Max | F::First { .. } | F::Last { .. } | F::AnyValue { .. }, t)
                 if accumulator::min_max_supported(t) =>
             {
                 t.clone()
             }
+            (F::MaxBy | F::MinBy, _) => match self.by_types() {
+                Some((v, k))
+                    if accumulator::min_max_supported(v) && accumulator::min_max_supported(k) =>
+                {
+                    v.clone()
+                }
+                _ => {
+                    return Err(unsupported(format!(
+                        "{:?} of {}",
+                        self.function, self.arg_type
+                    )));
+                }
+            },
             (f, t) => return Err(unsupported(format!("{f:?} of {t}"))),
         })
     }
@@ -237,13 +294,30 @@ impl AggregateExpr {
             ],
             (F::Min | F::Max, _) => vec![Field::new(n("value"), out, true)],
             // Whether a row was taken: a NULL value is a value.
-            (F::First { .. } | F::Last { .. }, _) => vec![
+            (F::First { .. } | F::Last { .. } | F::AnyValue { .. }, _) => vec![
                 Field::new(n("value"), out, true),
                 Field::new(n("valueSet"), DataType::Boolean, false),
             ],
+            // The key is NULL until a row with a key was taken: it says
+            // whether the value (which may be NULL) counts.
+            (F::MaxBy | F::MinBy, _) => {
+                let key = self.by_types().map_or(DataType::Null, |(_, k)| k.clone());
+                vec![
+                    Field::new(n("value"), out, true),
+                    Field::new(n("key"), key, true),
+                ]
+            }
         })
     }
 
+    /// The value and ordering key types of a `max_by`/`min_by`.
+    fn by_types(&self) -> Option<(&DataType, &DataType)> {
+        match &self.arg_type {
+            DataType::Struct(f) if f.len() == 2 => Some((f[0].data_type(), f[1].data_type())),
+            _ => None,
+        }
+    }
+
     /// The type `avg` sums in: `double`, or `sum`'s decimal type.
     fn avg_sum_type(&self) -> DataType {
         match self.arg_type {

crates/krishiv-exec/src/aggregate/rowtable.rs

@@ -719,7 +719,7 @@ impl RowTable {
             .map(|(part, n)| {
                 part.accs
                     .iter()
-                    .map(|a| a.projected(n, None))
+                    .map(|a| a.projected(n, None, None))
                     .sum::<usize>()
             })
             .sum();
@@ -728,12 +728,16 @@ impl RowTable {
                 .iter()
                 .zip(&self.core.ranges)
                 .map(|(a, r)| {
-                    let at = match fold {
-                        Fold::States { offset } => r.start.checked_sub(offset),
-                        Fold::Raw => a.args().first().copied(),
+                    // The value column, then (`max_by`/`min_by`) the key.
+                    let at = |n: usize| {
+                        let i = match fold {
+                            Fold::States { offset } => r.start.checked_sub(offset).map(|i| i + n),
+                            Fold::Raw => a.args().get(n).copied(),
+                        };
+                        i.and_then(|i| cols.get(i))
                     };
-                    let col = at.and_then(|i| cols.get(i));
-                    a.projected(0, col).saturating_sub(a.projected(0, None))
+                    a.projected(0, at(0), at(1))
+                        .saturating_sub(a.projected(0, None, None))
                 })
                 .sum()
         });

crates/krishiv-exec/src/aggregate/table.rs

@@ -245,12 +245,16 @@ impl AggTable {
                 .iter()
                 .zip(&self.core.ranges)
                 .map(|(a, r)| {
-                    let at = if p.states {
-                        Some(r.start)
-                    } else {
-                        a.args().first().copied()
+                    // The value column, then (`max_by`/`min_by`) the key.
+                    let at = |n: usize| {
+                        let i = if p.states {
+                            Some(r.start + n)
+                        } else {
+                            a.args().get(n).copied()
+                        };
+                        i.and_then(|i| p.cols.get(i))
                     };
-                    a.projected(new, at.and_then(|i| p.cols.get(i)))
+                    a.projected(new, at(0), at(1))
                 })
                 .sum::<usize>()
     }

crates/krishiv-exec/src/aggregate/tests.rs

@@ -1461,6 +1461,8 @@ async fn first_and_last_pick_rows_in_input_order() {
     );
 }
 
+mod simple;
+
 #[test]
 fn first_and_last_types_and_states() {
     use AggregateFunction::{First, Last};