feat(sql): stddev, variance, skewness, kurtosis, corr and covariance aggregates

crates/krishiv-common/src/config/registry.rs

@@ -61,6 +61,7 @@ pub static REGISTRY: &[&ConfigEntry] = &[
     &RUNTIME_FILTER_ADAPTIVE_MAX_PASS_PERCENT,
     &RUNTIME_FILTER_ADAPTIVE_MIN_ROWS,
     &WINDOW_GROUP_LIMIT_THRESHOLD,
+    &LEGACY_STATISTICAL_AGGREGATE,
 ];
 
 /// Every key and Spark alias with its entry, sorted by (length, name) for
@@ -165,6 +166,10 @@ pub static INDEX: &[(&str, &ConfigEntry)] = &[
         "spark.sql.leafNodeDefaultParallelism",
         &LEAF_NODE_DEFAULT_PARALLELISM,
     ),
+    (
+        "spark.sql.legacy.statisticalAggregate",
+        &LEGACY_STATISTICAL_AGGREGATE,
+    ),
     (
         "krishiv.runtimeFilter.adaptive.minRows",
         &RUNTIME_FILTER_ADAPTIVE_MIN_ROWS,
@@ -181,6 +186,10 @@ pub static INDEX: &[(&str, &ConfigEntry)] = &[
         "spark.sql.sources.bucketing.maxBuckets",
         &BUCKETING_MAX_BUCKETS,
     ),
+    (
+        "krishiv.sql.legacy.statisticalAggregate",
+        &LEGACY_STATISTICAL_AGGREGATE,
+    ),
     (
         "krishiv.runtimeFilter.exactSet.maxValues",
         &RUNTIME_FILTER_EXACT_SET_MAX_VALUES,
@@ -859,3 +868,16 @@ pub static WINDOW_GROUP_LIMIT_THRESHOLD: ConfigEntry = ConfigEntry {
     version_added: "0.1.0",
     spark_alias: Some("spark.sql.optimizer.windowGroupLimitThreshold"),
 };
+
+/// `spark.sql.legacy.statisticalAggregate`.
+pub static LEGACY_STATISTICAL_AGGREGATE: ConfigEntry = ConfigEntry {
+    key: "krishiv.sql.legacy.statisticalAggregate",
+    value_type: ValueType::Bool,
+    default: Some(ConfigValue::Bool(false)),
+    doc: "When true, `stddev_samp`, `var_samp`, `covar_samp` and `corr` of a single row, and \
+          `skewness` and `kurtosis` of values without variance, return NaN; when false they \
+          return NULL, as Spark 4 does.",
+    scope: Scope::Session,
+    version_added: "0.1.0",
+    spark_alias: Some("spark.sql.legacy.statisticalAggregate"),
+};

crates/krishiv-exec/benches/aggregate.rs

@@ -94,6 +94,7 @@ fn decimal_plan(batches: &[Batch], f: AggregateFunction) -> Arc<dyn ExecNode> {
         ansi: true,
         spill: None,
         require_unit_weights: false,
+        legacy_statistical: false,
     };
     let scan = Arc::new(Scan(Batches::new(s, batches.to_vec())));
     AggregateNode::new(scan, HashAggregate::new(spec).expect("spec"))
@@ -132,6 +133,7 @@ fn plan(batches: &[Batch], spill: Option<SpillScope>) -> Arc<dyn ExecNode> {
         ansi: true,
         spill,
         require_unit_weights: false,
+        legacy_statistical: false,
     };
     let scan = Arc::new(Scan(Batches::new(s, batches.to_vec())));
     AggregateNode::new(scan, HashAggregate::new(spec).expect("spec"))

crates/krishiv-exec/examples/agg_table_matrix.rs

@@ -1359,6 +1359,7 @@ fn run_op(x: &Ctx<'_>) -> Timing {
         ansi: false,
         spill: None,
         require_unit_weights: false,
+        legacy_statistical: false,
     };
     let batches: Vec<Batch> = x.data.batches.iter().cloned().map(Batch::new).collect();
     let scan = Arc::new(Scan(Batches::new(s, batches)));

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

@@ -37,6 +37,7 @@ use arrow_schema::{DataType, TimeUnit};
 use krishiv_kernels::decimal::DecimalNative;
 use krishiv_types::batch::Weights;
 
+use super::moments::Moments;
 use super::{
     AggInput, AggregateExpr, AggregateFunction, decimal_sum_overflow, long_overflow, out_of_range,
     unsupported,
@@ -252,6 +253,8 @@ enum Kind {
         null: Vec<bool>,
         heap: usize,
     },
+    /// The moment aggregates: Spark's buffers of doubles (see `moments`).
+    Moments(Moments),
 }
 
 /// One aggregate's state for every group of a table.
@@ -547,6 +550,18 @@ fn counts_of(a: &ArrayRef) -> Result<&[i64]> {
         .ok_or_else(|| internal(format!("count state of type {}", a.data_type())))
 }
 
+/// A moment aggregate's state column: `double`s, none NULL (a buffer is
+/// never NULL; a NULL would silently skip a partial's rows).
+fn state_doubles(a: &ArrayRef) -> Result<&[f64]> {
+    match a.as_primitive_opt::<Float64Type>() {
+        Some(c) if c.null_count() == 0 => Ok(c.values().as_ref()),
+        _ => Err(internal(format!(
+            "moment state of type {} with NULLs or not double",
+            a.data_type()
+        ))),
+    }
+}
+
 /// A `boolean` state column.
 fn booleans_of(a: &ArrayRef) -> Result<&BooleanArray> {
     a.as_boolean_opt()
@@ -917,6 +932,18 @@ impl Accumulator {
                 null: Vec::new(),
                 heap: 0,
             },
+            (
+                f @ (F::StddevSamp
+                | F::StddevPop
+                | F::VarSamp
+                | F::VarPop
+                | F::Skewness
+                | F::Kurtosis
+                | F::Corr
+                | F::CovarSamp
+                | F::CovarPop),
+                _,
+            ) => Kind::Moments(Moments::new(f, ansi)?),
         };
         let sum_bound = match (expr.function, expr.avg_sum_type()) {
             (F::Sum | F::Avg, DataType::Decimal64(p, _) | DataType::Decimal128(p, _)) => {
@@ -935,6 +962,17 @@ impl Accumulator {
         })
     }
 
+    /// Whether the sample moment aggregates give NaN instead of NULL where
+    /// Spark divides by zero (`spark.sql.legacy.statisticalAggregate`); the
+    /// other aggregates ignore it.
+    #[must_use]
+    pub(crate) fn legacy_statistical(mut self, on: bool) -> Self {
+        if let Kind::Moments(m) = &mut self.kind {
+            m.legacy(on);
+        }
+        self
+    }
+
     /// Bytes one group's state takes (heap bytes of string values aside).
     pub(crate) fn group_bytes(&self) -> usize {
         self.vecs().iter().map(|&(_, _, e)| e).sum()
@@ -967,6 +1005,7 @@ impl Accumulator {
                 null,
                 ..
             } => vec![values.vec(), keys.vec(), v(set), v(null)],
+            Kind::Moments(m) => m.vecs(),
         }
     }
 
@@ -1068,6 +1107,7 @@ impl Accumulator {
                 null.resize(n, false);
                 values.resize(n);
             }
+            Kind::Moments(m) => m.resize(n),
             Kind::By {
                 values,
                 keys,
@@ -1123,6 +1163,7 @@ impl Accumulator {
                 null.reserve(new);
                 values.reserve(new);
             }
+            Kind::Moments(m) => m.reserve(new),
             Kind::By {
                 values,
                 keys,
@@ -1215,7 +1256,7 @@ impl Accumulator {
         if (state
             || matches!(
                 self.kind,
-                Kind::MinMax { .. } | Kind::Pick { .. } | Kind::By { .. }
+                Kind::MinMax { .. } | Kind::Pick { .. } | Kind::By { .. } | Kind::Moments(_)
             ))
             && !is_unit(w)
         {
@@ -1251,6 +1292,7 @@ impl Accumulator {
                 self.kind,
                 Kind::SumDec { .. }
                     | Kind::AvgDec { .. }
+                    | Kind::Moments(_)
                     | Kind::MinMax {
                         values: Values::Bytes(_),
                         ..
@@ -1526,6 +1568,23 @@ impl Accumulator {
                     }
                 }
             }
+            (Kind::Moments(m), false) => {
+                let cols = args.iter().map(|a| floats(a)).collect::<Result<Vec<_>>>()?;
+                let cols: Vec<&[f64]> = cols.iter().map(AsRef::as_ref).collect();
+                let all_valid = |i: usize| {
+                    nulls
+                        .iter()
+                        .all(|n| n.as_ref().is_none_or(|n| n.is_valid(i)))
+                };
+                m.update(groups, &cols, all_valid)?;
+            }
+            (Kind::Moments(m), true) => {
+                let cols = args
+                    .iter()
+                    .map(|a| state_doubles(a))
+                    .collect::<Result<Vec<_>>>()?;
+                m.merge(groups, &cols)?;
+            }
             (
                 Kind::By {
                     max,
@@ -1592,6 +1651,7 @@ impl Accumulator {
             | Kind::AvgDec { .. }
             | Kind::Pick { .. }
             | Kind::By { .. } => 2,
+            Kind::Moments(ref m) => m.width(),
         }
     }
 
@@ -1651,6 +1711,7 @@ impl Accumulator {
                     ids.iter().map(|i| Some(set[pick(i)])),
                 )),
             ],
+            Kind::Moments(m) => m.state(ids)?,
             Kind::By {
                 values,
                 keys,
@@ -1796,6 +1857,7 @@ impl Accumulator {
                 }
                 self.decimal(out.into_iter(), &self.result)?
             }
+            Kind::Moments(m) => m.evaluate(ids)?,
             // Only the value: the ordering key is state, not a result.
             Kind::By {
                 values, set, null, ..

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

@@ -331,6 +331,48 @@ fn min_max_on_decimal64() {
     }
 }
 
+/// The moment aggregates: `(group, x, y, x is valid, y is valid)`.
+type MomentRow = (u32, f64, f64, bool, bool);
+
+fn moment_acc(f: AggregateFunction, groups: u32) -> Accumulator {
+    let args = if crate::aggregate::arity(f) == 2 {
+        vec![0, 1]
+    } else {
+        vec![0]
+    };
+    let e = AggregateExpr {
+        function: f,
+        args,
+        input: AggInput::Raw,
+        arg_type: DataType::Float64,
+        name: "m".to_owned(),
+    };
+    // Without ANSI: `corr` of a constant column is NULL, not an error.
+    let mut a = Accumulator::new(&e, false).unwrap();
+    a.resize(groups as usize);
+    a
+}
+
+/// The columns of `rows` as the arguments of a moment aggregate (garbage
+/// under the NULL slots).
+fn moment_columns(rows: &[MomentRow]) -> Vec<ArrayRef> {
+    let col = |value: fn(&MomentRow) -> f64, valid: fn(&MomentRow) -> bool| -> ArrayRef {
+        let nulls = NullBuffer::from(rows.iter().map(valid).collect::<Vec<_>>());
+        let v: Vec<f64> = rows
+            .iter()
+            .map(|r| if valid(r) { value(r) } else { f64::NAN })
+            .collect();
+        Arc::new(Float64Array::new(v.into(), Some(nulls)))
+    };
+    vec![col(|r| r.1, |r| r.3), col(|r| r.2, |r| r.4)]
+}
+
+fn moment_values(a: &Accumulator, groups: u32) -> Vec<Option<f64>> {
+    let ids: Vec<u32> = (0..groups).collect();
+    let out = a.evaluate(&ids).unwrap();
+    out.as_primitive::<Float64Type>().iter().collect()
+}
+
 /// 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] = [
@@ -387,6 +429,74 @@ fn by_expr(f: AggregateFunction, strings: (bool, bool)) -> AggregateExpr {
     }
 }
 
+proptest! {
+    #![proptest_config(ProptestConfig {
+        cases: 96,
+        failure_persistence: None,
+        ..ProptestConfig::default()
+    })]
+
+    /// Partial states of any split of the input merge, in any order and
+    /// through a second level, to the single pass's answer (up to the
+    /// reassociation of the floating-point operations).
+    #[test]
+    fn moments_merge_differential(
+        groups in 1..6_u32,
+        rows in proptest::collection::vec(
+            (0..6_u32, -1000.0..1000.0_f64, -1000.0..1000.0_f64,
+             proptest::bool::weighted(0.85), proptest::bool::weighted(0.85)),
+            0..160,
+        ),
+        parts in 1..7_usize,
+        rotate in 0..7_usize,
+        reverse in proptest::bool::ANY,
+    ) {
+        let rows: Vec<MomentRow> = rows.into_iter().map(|r| (r.0 % groups, r.1, r.2, r.3, r.4)).collect();
+        let ids: Vec<u32> = (0..groups).collect();
+        for f in MOMENT_FUNCTIONS {
+            let mut once = moment_acc(f, groups);
+            let g: Vec<u32> = rows.iter().map(|r| r.0).collect();
+            once.update(&g, &moment_columns(&rows), &Weights::Const(1)).unwrap();
+            let mut states: Vec<Vec<ArrayRef>> = rows
+                .chunks(rows.len().div_ceil(parts).max(1))
+                .map(|chunk| {
+                    let mut a = moment_acc(f, groups);
+                    let g: Vec<u32> = chunk.iter().map(|r| r.0).collect();
+                    a.update(&g, &moment_columns(chunk), &Weights::Const(1)).unwrap();
+                    a.state(&ids).unwrap()
+                })
+                .collect();
+            let len = states.len().max(1);
+            states.rotate_left(rotate % len);
+            if reverse {
+                states.reverse();
+            }
+            // Two levels: halves merged apart, then together.
+            let (lo, hi) = states.split_at(states.len() / 2);
+            let mut top = moment_acc(f, groups);
+            for half in [lo, hi] {
+                let mut mid = moment_acc(f, groups);
+                for st in half {
+                    mid.merge(&ids, st).unwrap();
+                }
+                top.merge(&ids, &mid.state(&ids).unwrap()).unwrap();
+            }
+            for (got, want) in moment_values(&top, groups).iter().zip(moment_values(&once, groups)) {
+                match (got, want) {
+                    (None, None) => {}
+                    (Some(g), Some(w)) if g.is_nan() || w.is_nan() => {
+                        prop_assert!(g.is_nan() && w.is_nan(), "{:?}: {} vs {}", f, g, w);
+                    }
+                    (Some(g), Some(w)) => {
+                        prop_assert!((g - w).abs() <= 1e-8 * g.abs().max(w.abs()).max(1.0), "{:?}: {} vs {}", f, g, w);
+                    }
+                    (g, w) => prop_assert!(false, "{:?}: {:?} vs {:?}", f, g, w),
+                }
+            }
+        }
+    }
+}
+
 proptest! {
     #![proptest_config(ProptestConfig {
         cases: 256,
@@ -464,3 +574,15 @@ proptest! {
         }
     }
 }
+
+const MOMENT_FUNCTIONS: [AggregateFunction; 9] = [
+    AggregateFunction::StddevSamp,
+    AggregateFunction::StddevPop,
+    AggregateFunction::VarSamp,
+    AggregateFunction::VarPop,
+    AggregateFunction::Skewness,
+    AggregateFunction::Kurtosis,
+    AggregateFunction::Corr,
+    AggregateFunction::CovarSamp,
+    AggregateFunction::CovarPop,
+];