From c92340b6f93bebde3e0f23812cedb574da84ae1e Mon Sep 17 00:00:00 2001 From: techysneha14 <343708768+shehatathe@users.noreply.github.com> Date: Sun, 25 Oct 2026 01:34:00 +0200 Subject: [PATCH 1/5] Add Spark statistical aggregates for moments, correlation, and covariance Adds nine aggregate functions (stddev_samp, stddev_pop, var_samp, var_pop, skewness, kurtosis, corr, covar_samp, and covar_pop, with their Spark aliases) using the Spark pairwise-merge moment formulas for numerically stable partial merging. Integer and decimal inputs are cast to double by the analyzer, and a legacy statistical aggregate config flag is registered to match Spark behavior. Correlation over a constant column raises DIVIDE_BY_ZERO under ANSI and returns NULL otherwise. These aggregates are refused in window definitions. Plan-codec encode and decode arms are added for the new aggregates. TPC-DS q17 and q39 now plan but are not yet encodable. Diff coverage for this change was not measured. Signed-off-by: techysneha14 <343708768+shehatathe@users.noreply.github.com> --- crates/krishiv-common/src/config/registry.rs | 22 + crates/krishiv-exec/benches/aggregate.rs | 2 + .../krishiv-exec/examples/agg_table_matrix.rs | 1 + .../krishiv-exec/src/aggregate/accumulator.rs | 70 +- .../src/aggregate/accumulator/differential.rs | 122 +++ crates/krishiv-exec/src/aggregate/mod.rs | 63 +- crates/krishiv-exec/src/aggregate/moments.rs | 427 ++++++++++ .../src/aggregate/moments/tests.rs | 766 ++++++++++++++++++ crates/krishiv-exec/src/aggregate/node.rs | 1 + crates/krishiv-exec/src/aggregate/rowtable.rs | 5 +- crates/krishiv-exec/src/aggregate/table.rs | 11 +- crates/krishiv-exec/src/aggregate/tests.rs | 464 ++++++++++- .../krishiv-exec/src/physical/instantiate.rs | 8 +- crates/krishiv-exec/src/physical/mod.rs | 5 +- crates/krishiv-exec/src/planner/aggregate.rs | 17 +- crates/krishiv-exec/src/planner/eval/mod.rs | 13 +- crates/krishiv-exec/src/profile/tests.rs | 1 + crates/krishiv-exec/tests/aggregate_tpch.rs | 1 + crates/krishiv-exec/tests/planner/physical.rs | 1 + crates/krishiv-exec/tests/planner/write.rs | 1 + crates/krishiv-functions/src/registry/defs.rs | 19 + crates/krishiv-functions/src/registry/mod.rs | 18 + .../krishiv-functions/src/registry/tests.rs | 63 ++ crates/krishiv-plan-codec/src/decode.rs | 3 + crates/krishiv-plan-codec/src/encode.rs | 28 +- crates/krishiv-plan-codec/src/features.rs | 13 +- .../tests/common/inventory.rs | 1 + crates/krishiv-plan-codec/tests/ratchet.rs | 21 +- .../tests/aggregate_moments.rs | 407 ++++++++++ crates/krishiv-testkit/tests/type_coercion.rs | 4 +- .../analysis/registry-matrix-manifest.json | 390 ++++++++- .../analysis/type-coercion-manifest.json | 30 - 32 files changed, 2934 insertions(+), 64 deletions(-) create mode 100644 crates/krishiv-exec/src/aggregate/moments.rs create mode 100644 crates/krishiv-exec/src/aggregate/moments/tests.rs create mode 100644 crates/krishiv-session/tests/aggregate_moments.rs diff --git a/crates/krishiv-common/src/config/registry.rs b/crates/krishiv-common/src/config/registry.rs index 06882316f..448599955 100644 --- a/crates/krishiv-common/src/config/registry.rs +++ b/crates/krishiv-common/src/config/registry.rs @@ -57,6 +57,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 @@ -158,6 +159,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, @@ -170,6 +175,10 @@ pub static INDEX: &[(&str, &ConfigEntry)] = &[ "krishiv.sql.leafNodeDefaultParallelism", &LEAF_NODE_DEFAULT_PARALLELISM, ), + ( + "krishiv.sql.legacy.statisticalAggregate", + &LEGACY_STATISTICAL_AGGREGATE, + ), ( "krishiv.runtimeFilter.exactSet.maxValues", &RUNTIME_FILTER_EXACT_SET_MAX_VALUES, @@ -787,3 +796,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"), +}; diff --git a/crates/krishiv-exec/benches/aggregate.rs b/crates/krishiv-exec/benches/aggregate.rs index f7b2d14ba..e073909f9 100644 --- a/crates/krishiv-exec/benches/aggregate.rs +++ b/crates/krishiv-exec/benches/aggregate.rs @@ -94,6 +94,7 @@ fn decimal_plan(batches: &[Batch], f: AggregateFunction) -> Arc { 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) -> Arc { 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")) diff --git a/crates/krishiv-exec/examples/agg_table_matrix.rs b/crates/krishiv-exec/examples/agg_table_matrix.rs index 92421b797..f1afbe9a6 100644 --- a/crates/krishiv-exec/examples/agg_table_matrix.rs +++ b/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 = x.data.batches.iter().cloned().map(Batch::new).collect(); let scan = Arc::new(Scan(Batches::new(s, batches))); diff --git a/crates/krishiv-exec/src/aggregate/accumulator.rs b/crates/krishiv-exec/src/aggregate/accumulator.rs index 958b2ac25..bdbd2403c 100644 --- a/crates/krishiv-exec/src/aggregate/accumulator.rs +++ b/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, @@ -155,6 +156,8 @@ enum Kind { null: Vec, heap: usize, }, + /// The moment aggregates: Spark's buffers of doubles (see `moments`). + Moments(Moments), } /// One aggregate's state for every group of a table. @@ -450,6 +453,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::() { + 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() @@ -802,6 +817,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, _)) => { @@ -820,6 +847,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() @@ -845,6 +883,7 @@ impl Accumulator { Kind::Pick { values, set, null, .. } => vec![values.vec(), v(set), v(null)], + Kind::Moments(m) => m.vecs(), } } @@ -931,6 +970,7 @@ impl Accumulator { null.resize(n, false); values.resize(n); } + Kind::Moments(m) => m.resize(n), } } @@ -974,6 +1014,7 @@ impl Accumulator { null.reserve(new); values.reserve(new); } + Kind::Moments(m) => m.reserve(new), } } @@ -1051,7 +1092,13 @@ impl Accumulator { args.len() ))); } - if (state || matches!(self.kind, Kind::MinMax { .. } | Kind::Pick { .. })) && !is_unit(w) { + if (state + || matches!( + self.kind, + Kind::MinMax { .. } | Kind::Pick { .. } | Kind::Moments(_) + )) + && !is_unit(w) + { return Err(unsupported(format!( "{:?} over rows with weights other than 1 (needs incremental view maintenance)", self.expr.function @@ -1084,6 +1131,7 @@ impl Accumulator { self.kind, Kind::SumDec { .. } | Kind::AvgDec { .. } + | Kind::Moments(_) | Kind::MinMax { values: Values::Bytes(_), .. @@ -1340,6 +1388,23 @@ impl Accumulator { } } } + (Kind::Moments(m), false) => { + let cols = args.iter().map(|a| floats(a)).collect::>>()?; + 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::>>()?; + m.merge(groups, &cols)?; + } } if let Kind::SumDec { sum, .. } | Kind::AvgDec { sum, .. } = &self.kind { self.max_sum = refresh_max(self.max_sum, sum, groups); @@ -1360,6 +1425,7 @@ impl Accumulator { | Kind::AvgFloat { .. } | Kind::AvgDec { .. } | Kind::Pick { .. } => 2, + Kind::Moments(ref m) => m.width(), } } @@ -1419,6 +1485,7 @@ impl Accumulator { ids.iter().map(|i| Some(set[pick(i)])), )), ], + Kind::Moments(m) => m.state(ids)?, }) } @@ -1548,6 +1615,7 @@ impl Accumulator { } self.decimal(out.into_iter(), &self.result)? } + Kind::Moments(m) => m.evaluate(ids)?, Kind::Count(_) | Kind::SumInt { .. } | Kind::SumFloat { .. } diff --git a/crates/krishiv-exec/src/aggregate/accumulator/differential.rs b/crates/krishiv-exec/src/aggregate/accumulator/differential.rs index 5f9ecfb9f..a27ea6c56 100644 --- a/crates/krishiv-exec/src/aggregate/accumulator/differential.rs +++ b/crates/krishiv-exec/src/aggregate/accumulator/differential.rs @@ -330,3 +330,125 @@ fn min_max_on_decimal64() { assert_eq!(v.as_primitive::().value(0), want); } } + +/// 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 { + let col = |value: fn(&MomentRow) -> f64, valid: fn(&MomentRow) -> bool| -> ArrayRef { + let nulls = NullBuffer::from(rows.iter().map(valid).collect::>()); + let v: Vec = 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> { + let ids: Vec = (0..groups).collect(); + let out = a.evaluate(&ids).unwrap(); + out.as_primitive::().iter().collect() +} + +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 = rows.into_iter().map(|r| (r.0 % groups, r.1, r.2, r.3, r.4)).collect(); + let ids: Vec = (0..groups).collect(); + for f in MOMENT_FUNCTIONS { + let mut once = moment_acc(f, groups); + let g: Vec = rows.iter().map(|r| r.0).collect(); + once.update(&g, &moment_columns(&rows), &Weights::Const(1)).unwrap(); + let mut states: Vec> = rows + .chunks(rows.len().div_ceil(parts).max(1)) + .map(|chunk| { + let mut a = moment_acc(f, groups); + let g: Vec = 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), + } + } + } + } +} + +const MOMENT_FUNCTIONS: [AggregateFunction; 9] = [ + AggregateFunction::StddevSamp, + AggregateFunction::StddevPop, + AggregateFunction::VarSamp, + AggregateFunction::VarPop, + AggregateFunction::Skewness, + AggregateFunction::Kurtosis, + AggregateFunction::Corr, + AggregateFunction::CovarSamp, + AggregateFunction::CovarPop, +]; diff --git a/crates/krishiv-exec/src/aggregate/mod.rs b/crates/krishiv-exec/src/aggregate/mod.rs index ffc87e155..d0a080821 100644 --- a/crates/krishiv-exec/src/aggregate/mod.rs +++ b/crates/krishiv-exec/src/aggregate/mod.rs @@ -1,6 +1,8 @@ // 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` and the moment aggregates (`stddev`, `variance`, +//! `skewness`, `kurtosis`, `covar_*`, `corr`) with Spark 4.2 result types, in +//! two phases, with spill and `DISTINCT`. //! //! # Operators //! - [`HashAggregate`] is a pipeline breaker. Each driver aggregates into @@ -50,11 +52,13 @@ //! 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; -//! a group exists once it received a row. +//! maintained incrementally (P4 IVM), so they, the moment aggregates 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; +mod moments; mod node; mod partial; mod rowtable; @@ -76,6 +80,7 @@ use krishiv_memory::SpillScope; #[cfg(all(test, not(loom)))] pub(crate) use accumulator::plain; pub(crate) use accumulator::{Accumulator, Running, avg_decimal_wide, wide_sum_value}; +pub(crate) use moments::arity; pub use node::{AggregateNode, PartialNode, distinct_plan, union_distinct}; pub use partial::{BypassPolicy, PartialAggregate}; #[cfg(test)] @@ -112,6 +117,33 @@ pub enum AggregateFunction { /// `last(x, true)` / `IGNORE NULLS`: the last non-NULL value. ignore_nulls: bool, }, + /// `stddev(x)` / `std` / `stddev_samp`: `double`; NULL for no rows, and + /// for one row unless `legacy_statistical` (then NaN). Rows with a NULL + /// are skipped. Likewise for the moment aggregates below. + StddevSamp, + /// `stddev_pop(x)`: `double`; NULL for no rows. + StddevPop, + /// `variance(x)` / `var_samp`, as `StddevSamp` without the root. + VarSamp, + /// `var_pop(x)`, as `StddevPop` without the root. + VarPop, + /// `skewness(x)`: `double`; NULL for no rows, and for no variance + /// unless `legacy_statistical` (then NaN). + Skewness, + /// `kurtosis(x)`, as `Skewness`. + Kurtosis, + /// `corr(x, y)`: Pearson correlation, `double`; both arguments are read + /// as `double` and a row with a NULL in either is skipped. NULL for no + /// rows, and for one row unless `legacy_statistical` (then NaN). Over + /// several rows of which one argument is constant it is + /// `DIVIDE_BY_ZERO` with `ansi` and NULL without, whatever + /// `legacy_statistical` says. + Corr, + /// `covar_samp(x, y)`: as `Corr` for its arguments and its single-row + /// rule. + CovarSamp, + /// `covar_pop(x, y)`: as `Corr` for its arguments; NULL for no rows. + CovarPop, } /// Where an aggregate's arguments come from. @@ -212,6 +244,11 @@ impl AggregateExpr { { t.clone() } + // The analyzer casts every argument to `double` (decimals + // included); the executor also reads integers and `float`. + (f, t) if moments::buffer(f).is_some() && moments::reads_as_double(t) => { + DataType::Float64 + } (f, t) => return Err(unsupported(format!("{f:?} of {t}"))), }) } @@ -241,6 +278,18 @@ impl AggregateExpr { Field::new(n("value"), out, true), Field::new(n("valueSet"), DataType::Boolean, false), ], + ( + F::StddevSamp + | F::StddevPop + | F::VarSamp + | F::VarPop + | F::Skewness + | F::Kurtosis + | F::Corr + | F::CovarSamp + | F::CovarPop, + _, + ) => moments::state_fields(self.function, &self.name), }) } @@ -297,6 +346,10 @@ pub struct AggregateSpec { /// Reject input weights other than 1 (`DISTINCT`, `UNION`: set /// semantics over Z-sets need IVM). pub require_unit_weights: bool, + /// `spark.sql.legacy.statisticalAggregate`: where Spark divides by zero, + /// the sample moment aggregates (`stddev_samp`, `var_samp`, `covar_samp`, + /// `corr`, `skewness`, `kurtosis`) give NaN instead of NULL. + pub legacy_statistical: bool, } impl AggregateSpec { diff --git a/crates/krishiv-exec/src/aggregate/moments.rs b/crates/krishiv-exec/src/aggregate/moments.rs new file mode 100644 index 000000000..b75bb1a7a --- /dev/null +++ b/crates/krishiv-exec/src/aggregate/moments.rs @@ -0,0 +1,427 @@ +// SPDX-License-Identifier: Apache-2.0 +//! The moment aggregates: `stddev`, `variance`, `skewness`, `kurtosis`, +//! `covar_*` and `corr`, with Spark 4.2's buffers and arithmetic +//! (`CentralMomentAgg`, `Covariance` and `PearsonCorrelation` of +//! `org.apache.spark.sql.catalyst.expressions.aggregate`). +//! +//! A group's state is Spark's aggregation buffer, `double`s all: a count +//! `n`, the running mean(s) and sums of powers of the deviations from the +//! mean(s): +//! +//! | function | buffer | +//! |---|---| +//! | `stddev_*`, `var_*` | `n`, `avg`, `m2` | +//! | `skewness` | `n`, `avg`, `m2`, `m3` | +//! | `kurtosis` | `n`, `avg`, `m2`, `m3`, `m4` | +//! | `covar_*` | `n`, `xAvg`, `yAvg`, `ck` | +//! | `corr` | `n`, `xAvg`, `yAvg`, `ck`, `xMk`, `yMk` | +//! +//! A row folds in with the one-pass update of Meng's "Simpler Online +//! Updates for Arbitrary-Order Central Moments" ([`update`]); two buffers +//! combine with the pairwise formulas of the same paper ([`merge`]). The +//! operations are applied in Spark's order, so one pass over rows in one +//! order gives Spark's bits. The naive sum-of-squares formulas are not +//! used: they lose all precision for values far from zero. +//! +//! [`update`] and [`merge`] are the only arithmetic: the default and the +//! row-layout tables fold rows and merge spilled or partial states through +//! the same [`Moments`], so they cannot disagree. +//! +//! # Floating-point tolerance +//! One pass in a fixed order, and any fixed partitioning merged in a fixed +//! order, are deterministic to the bit. Different partitionings (split the +//! input in two, in four, merge in another order) give the same value only +//! up to rounding: the merge reassociates additions and products, as it does +//! in Spark. For well-conditioned data (the mean is not many orders of +//! magnitude above the spread) the relative difference is a few ulps; the +//! tests allow `1e-9`. +//! +//! # NULL, NaN, infinity +//! A row with a NULL in any argument is skipped. NaN and infinities flow +//! through the arithmetic as IEEE doubles do (`stddev` of `1, NaN` is NaN; +//! an infinity makes the deviations NaN); a NaN count of rows never occurs. +//! +//! # Division by zero +//! Where Spark divides by zero for want of rows or spread (one row of a +//! sample statistic, no variance for `skewness` and `kurtosis`) the value is +//! NULL, or NaN under `spark.sql.legacy.statisticalAggregate`. `corr` of +//! several rows of which one column is constant is not governed by that +//! flag: it is `DIVIDE_BY_ZERO` under ANSI and NULL without. + +use arrow_array::{ArrayRef, Float64Array}; +use arrow_schema::{DataType, Field}; + +use krishiv_common::error::{DivideByZero, KrishivError}; + +use super::{ANSI_CONF, AggregateFunction}; +use crate::error::{Result, internal}; + +/// The widest buffer (`corr`'s). +const WIDTH: usize = 6; + +/// One group's buffer: the first [`Family::width`] entries are used. +type Row = [f64; WIDTH]; + +/// The buffer attributes of `f`, in Spark's order; `None` for an aggregate +/// that is not a moment aggregate. +pub(super) fn buffer(f: AggregateFunction) -> Option<&'static [&'static str]> { + use AggregateFunction as F; + Some(match f { + F::StddevSamp | F::StddevPop | F::VarSamp | F::VarPop => &["n", "avg", "m2"], + F::Skewness => &["n", "avg", "m2", "m3"], + F::Kurtosis => &["n", "avg", "m2", "m3", "m4"], + F::CovarSamp | F::CovarPop => &["n", "xAvg", "yAvg", "ck"], + F::Corr => &["n", "xAvg", "yAvg", "ck", "xMk", "yMk"], + _ => return None, + }) +} + +/// Whether a moment aggregate reads a column of type `t` (as doubles). +pub(super) fn reads_as_double(t: &DataType) -> bool { + matches!( + t, + DataType::Float32 + | DataType::Float64 + | DataType::Int8 + | DataType::Int16 + | DataType::Int32 + | DataType::Int64 + ) +} + +/// The partial-state columns of moment aggregate `f` (names derive from +/// the aggregate's `name`): Spark's buffer, non-nullable `double`s. +pub(super) fn state_fields(f: AggregateFunction, name: &str) -> Vec { + buffer(f) + .unwrap_or_default() + .iter() + .map(|attr| Field::new(format!("{name}#{attr}"), DataType::Float64, false)) + .collect() +} + +/// The arguments of an aggregate other than `count`: two for the pair +/// statistics (`corr`, `covar_*`), else one. +pub(crate) fn arity(f: AggregateFunction) -> usize { + match f { + AggregateFunction::CovarSamp | AggregateFunction::CovarPop | AggregateFunction::Corr => 2, + _ => 1, + } +} + +/// The buffer's shape: which update and merge apply. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum Family { + /// `CentralMomentAgg` of this order (2, 3 or 4). + Central(usize), + /// `Covariance` (`marginals`: false) or `PearsonCorrelation` (true). + Pair { marginals: bool }, +} + +impl Family { + fn of(f: AggregateFunction) -> Option { + let width = buffer(f)?.len(); + Some(if arity(f) == 2 { + Self::Pair { + marginals: width == WIDTH, + } + } else { + Self::Central(width - 1) + }) + } + + fn width(self) -> usize { + match self { + Self::Central(order) => order + 1, + Self::Pair { marginals: false } => 4, + Self::Pair { marginals: true } => WIDTH, + } + } +} + +/// Folds one row into a buffer: `x` (and `y` for the pair statistics). +/// Spark's `updateExpressions`. +fn update(family: Family, s: &mut Row, x: f64, y: f64) { + match family { + Family::Central(order) => { + let [n, avg, m2, m3, m4, _] = *s; + let new_n = n + 1.0; + let delta = x - avg; + let delta_n = delta / new_n; + let new_avg = avg + delta_n; + let new_m2 = m2 + delta * (delta - delta_n); + let delta2 = delta * delta; + let delta_n2 = delta_n * delta_n; + let new_m3 = if order >= 3 { + m3 - 3.0 * delta_n * new_m2 + delta * (delta2 - delta_n2) + } else { + 0.0 + }; + let new_m4 = if order >= 4 { + m4 - 4.0 * delta_n * new_m3 - 6.0 * delta_n2 * new_m2 + + delta * (delta * delta2 - delta_n * delta_n2) + } else { + 0.0 + }; + *s = [new_n, new_avg, new_m2, new_m3, new_m4, 0.0]; + } + Family::Pair { marginals } => { + let [n, x_avg, y_avg, ck, x_mk, y_mk] = *s; + let new_n = n + 1.0; + let dx = x - x_avg; + let dx_n = dx / new_n; + let dy = y - y_avg; + let dy_n = dy / new_n; + let new_x_avg = x_avg + dx_n; + let new_y_avg = y_avg + dy_n; + let new_ck = ck + dx * (y - new_y_avg); + let (new_x_mk, new_y_mk) = if marginals { + (x_mk + dx * (x - new_x_avg), y_mk + dy * (y - new_y_avg)) + } else { + (0.0, 0.0) + }; + *s = [new_n, new_x_avg, new_y_avg, new_ck, new_x_mk, new_y_mk]; + } + } +} + +/// Combines buffer `b` into buffer `a`. Spark's `mergeExpressions` (`a` is +/// the left buffer, `b` the right one). Applied even when a side is empty, +/// as Spark does. +fn merge(family: Family, a: &mut Row, b: &Row) { + let new_n = a[0] + b[0]; + match family { + Family::Central(order) => { + let [n1, avg1, m2_1, m3_1, m4_1, _] = *a; + let [n2, avg2, m2_2, m3_2, m4_2, _] = *b; + let delta = avg2 - avg1; + let delta_n = if new_n == 0.0 { 0.0 } else { delta / new_n }; + let new_avg = avg1 + delta_n * n2; + let new_m2 = m2_1 + m2_2 + delta * delta_n * n1 * n2; + let new_m3 = if order >= 3 { + m3_1 + m3_2 + + delta_n * delta_n * delta * n1 * n2 * (n1 - n2) + + 3.0 * delta_n * (n1 * m2_2 - n2 * m2_1) + } else { + 0.0 + }; + let new_m4 = if order >= 4 { + m4_1 + m4_2 + + delta_n * delta_n * delta_n * delta * n1 * n2 * (n1 * n1 - n1 * n2 + n2 * n2) + + 6.0 * delta_n * delta_n * (n1 * n1 * m2_2 + n2 * n2 * m2_1) + + 4.0 * delta_n * (n1 * m3_2 - n2 * m3_1) + } else { + 0.0 + }; + *a = [new_n, new_avg, new_m2, new_m3, new_m4, 0.0]; + } + Family::Pair { marginals } => { + let [n1, x_avg1, y_avg1, ck1, x_mk1, y_mk1] = *a; + let [n2, x_avg2, y_avg2, ck2, x_mk2, y_mk2] = *b; + let dx = x_avg2 - x_avg1; + let dx_n = if new_n == 0.0 { 0.0 } else { dx / new_n }; + let dy = y_avg2 - y_avg1; + let dy_n = if new_n == 0.0 { 0.0 } else { dy / new_n }; + let new_x_avg = x_avg1 + dx_n * n2; + let new_y_avg = y_avg1 + dy_n * n2; + let new_ck = ck1 + ck2 + dx * dy_n * n1 * n2; + let (new_x_mk, new_y_mk) = if marginals { + ( + x_mk1 + x_mk2 + dx * dx_n * n1 * n2, + y_mk1 + y_mk2 + dy * dy_n * n1 * n2, + ) + } else { + (0.0, 0.0) + }; + *a = [new_n, new_x_avg, new_y_avg, new_ck, new_x_mk, new_y_mk]; + } + } +} + +/// The value of `f` over buffer `s`: NULL for no rows; where Spark divides +/// by zero, NULL, or NaN with `legacy` (`spark.sql.legacy.statisticalAggregate`). +/// `corr` of several rows over a constant column is the exception: a +/// division by zero whatever `legacy` says, `DIVIDE_BY_ZERO` with `ansi` and +/// NULL without. Spark's `evaluateExpression`. +fn finish(f: AggregateFunction, s: &Row, legacy: bool, ansi: bool) -> Result> { + use AggregateFunction as F; + let [n, _, m2, _, m4, _] = *s; + if n == 0.0 { + return Ok(None); + } + let divide_by_zero = legacy.then_some(f64::NAN); + Ok(match f { + F::StddevPop => Some((m2 / n).sqrt()), + F::StddevSamp | F::VarSamp | F::CovarSamp | F::Corr if n == 1.0 => divide_by_zero, + F::StddevSamp => Some((m2 / (n - 1.0)).sqrt()), + F::VarPop => Some(m2 / n), + F::VarSamp => Some(m2 / (n - 1.0)), + F::Skewness | F::Kurtosis if m2 == 0.0 => divide_by_zero, + F::Skewness => Some(n.sqrt() * s[3] / (m2 * m2 * m2).sqrt()), + F::Kurtosis => Some(n * m4 / (m2 * m2) - 3.0), + // The pair statistics keep their `ck` in the fourth slot and + // `xMk`, `yMk` in the last two. + F::CovarPop => Some(s[3] / n), + F::CovarSamp => Some(s[3] / (n - 1.0)), + F::Corr if s[4] * s[5] == 0.0 => { + if ansi { + return Err(KrishivError::from(DivideByZero { + config: ANSI_CONF.to_owned(), + })); + } + None + } + F::Corr => Some(s[3] / (s[4] * s[5]).sqrt()), + _ => None, + }) +} + +/// The states of one moment aggregate for every group, one vector per +/// buffer attribute (column-oriented, as the other accumulators). +#[derive(Debug)] +pub(super) struct Moments { + function: AggregateFunction, + family: Family, + /// NaN instead of NULL where Spark divides by zero. + legacy: bool, + /// `corr` over a constant column is `DIVIDE_BY_ZERO`, not NULL. + ansi: bool, + cols: Vec>, +} + +impl Moments { + /// No groups yet, for moment aggregate `function`. + pub(super) fn new(function: AggregateFunction, ansi: bool) -> Result { + let family = Family::of(function) + .ok_or_else(|| internal(format!("{function:?} is not a moment aggregate")))?; + Ok(Self { + function, + family, + legacy: false, + ansi, + cols: vec![Vec::new(); family.width()], + }) + } + + pub(super) fn legacy(&mut self, on: bool) { + self.legacy = on; + } + + /// The number of state columns. + pub(super) fn width(&self) -> usize { + self.family.width() + } + + /// Each vector as (length, capacity, element size). + pub(super) fn vecs(&self) -> Vec<(usize, usize, usize)> { + self.cols + .iter() + .map(|c| (c.len(), c.capacity(), size_of::())) + .collect() + } + + /// Adds groups up to `n` with empty states. + pub(super) fn resize(&mut self, n: usize) { + self.cols.iter_mut().for_each(|c| c.resize(n, 0.0)); + } + + /// Reserves room for `new` more groups. + pub(super) fn reserve(&mut self, new: usize) { + self.cols.iter_mut().for_each(|c| c.reserve(new)); + } + + fn load(&self, g: u32) -> Result { + let g = g as usize; + let mut row = [0.0; WIDTH]; + for (slot, col) in row.iter_mut().zip(&self.cols) { + *slot = *col + .get(g) + .ok_or_else(|| internal(format!("aggregate group {g} out of range")))?; + } + Ok(row) + } + + fn store(&mut self, g: u32, row: &Row) { + for (col, v) in self.cols.iter_mut().zip(row) { + if let Some(slot) = col.get_mut(g as usize) { + *slot = *v; + } + } + } + + /// Folds rows into `groups` (one group id per row). `args` are the + /// argument columns as doubles; `valid(i)` says whether row `i` has no + /// NULL argument. + pub(super) fn update( + &mut self, + groups: &[u32], + args: &[&[f64]], + valid: impl Fn(usize) -> bool, + ) -> Result<()> { + let (Some(x), y) = (args.first(), args.get(1)) else { + return Err(internal("a moment aggregate without arguments".to_owned())); + }; + if x.len() != groups.len() + || args.len() != arity(self.function) + || y.is_some_and(|y| y.len() != groups.len()) + { + return Err(internal( + "a moment aggregate over columns of other lengths than its rows".to_owned(), + )); + } + for (i, &g) in groups.iter().enumerate() { + if valid(i) { + let mut s = self.load(g)?; + update(self.family, &mut s, x[i], y.map_or(0.0, |y| y[i])); + self.store(g, &s); + } + } + Ok(()) + } + + /// Merges partial states: `states` are the state columns, in buffer + /// order, one row per entry of `groups`. + pub(super) fn merge(&mut self, groups: &[u32], states: &[&[f64]]) -> Result<()> { + if states.len() != self.width() || states.iter().any(|c| c.len() != groups.len()) { + return Err(internal( + "moment states of another shape than the rows merged".to_owned(), + )); + } + for (i, &g) in groups.iter().enumerate() { + let mut incoming = [0.0; WIDTH]; + for (slot, col) in incoming.iter_mut().zip(states) { + *slot = col[i]; + } + let mut s = self.load(g)?; + merge(self.family, &mut s, &incoming); + self.store(g, &s); + } + Ok(()) + } + + /// The partial states of groups `ids`, one `double` column per buffer + /// attribute. + pub(super) fn state(&self, ids: &[u32]) -> Result> { + let rows = ids + .iter() + .map(|&g| self.load(g)) + .collect::>>()?; + Ok((0..self.width()) + .map(|k| { + let v: Float64Array = rows.iter().map(|r| r[k]).collect(); + std::sync::Arc::new(v) as ArrayRef + }) + .collect()) + } + + /// The final values of groups `ids`. + pub(super) fn evaluate(&self, ids: &[u32]) -> Result { + let v = ids + .iter() + .map(|&g| finish(self.function, &self.load(g)?, self.legacy, self.ansi)) + .collect::>()?; + Ok(std::sync::Arc::new(v)) + } +} + +#[cfg(test)] +mod tests; diff --git a/crates/krishiv-exec/src/aggregate/moments/tests.rs b/crates/krishiv-exec/src/aggregate/moments/tests.rs new file mode 100644 index 000000000..67134ae96 --- /dev/null +++ b/crates/krishiv-exec/src/aggregate/moments/tests.rs @@ -0,0 +1,766 @@ +// SPDX-License-Identifier: Apache-2.0 +//! The moment aggregates through their accumulators: edge inputs under both +//! values of `spark.sql.legacy.statisticalAggregate`, Spark's recorded +//! results reproduced bit for bit under Spark's partitioning, merges over +//! any split, and the checks of malformed input. +#![cfg(not(loom))] + +use std::ops::Range; +use std::sync::Arc; + +use arrow_array::cast::AsArray; +use arrow_array::types::Float64Type; +use arrow_array::{Array, ArrayRef, Float32Array, Float64Array, Int32Array}; +use arrow_buffer::NullBuffer; +use arrow_schema::DataType; +use krishiv_common::error::ErrorClass; +use krishiv_types::batch::Weights; + +use super::arity; +use crate::aggregate::{Accumulator, AggInput, AggregateExpr, AggregateFunction}; + +use AggregateFunction::{ + Corr, CovarPop, CovarSamp, Kurtosis, Skewness, StddevPop, StddevSamp, VarPop, VarSamp, +}; + +const ALL: [AggregateFunction; 9] = [ + StddevSamp, StddevPop, VarSamp, VarPop, Skewness, Kurtosis, Corr, CovarSamp, CovarPop, +]; + +/// The functions NULL (or NaN, with the legacy flag) for one row. +const SINGLE_ROW_DIVIDES: [AggregateFunction; 4] = [StddevSamp, VarSamp, CovarSamp, Corr]; + +fn expr(f: AggregateFunction) -> AggregateExpr { + AggregateExpr { + function: f, + args: (0..arity(f)).collect(), + input: AggInput::Raw, + arg_type: DataType::Float64, + name: "m".to_owned(), + } +} + +fn doubles(v: &[Option]) -> ArrayRef { + Arc::new(Float64Array::from(v.to_vec())) +} + +fn some(v: &[f64]) -> Vec> { + v.iter().copied().map(Some).collect() +} + +/// An accumulator over `groups` groups, without ANSI. +fn acc(f: AggregateFunction, legacy: bool, groups: usize) -> Accumulator { + acc_ansi(f, legacy, false, groups) +} + +fn acc_ansi(f: AggregateFunction, legacy: bool, ansi: bool, groups: usize) -> Accumulator { + let mut a = Accumulator::new(&expr(f), ansi) + .unwrap() + .legacy_statistical(legacy); + a.resize(groups); + a +} + +/// The columns an aggregate reads: `xs`, and `ys` for the pair statistics. +fn columns(f: AggregateFunction, xs: &[Option], ys: &[Option]) -> Vec { + let mut cols = vec![doubles(xs)]; + if arity(f) == 2 { + cols.push(doubles(ys)); + } + cols +} + +/// One global group over rows `xs` (and `ys`), in one pass. +fn global( + f: AggregateFunction, + legacy: bool, + xs: &[Option], + ys: &[Option], +) -> Option { + let mut a = acc(f, legacy, 1); + a.update(&vec![0; xs.len()], &columns(f, xs, ys), &Weights::Const(1)) + .unwrap(); + value(&a, 0) +} + +fn value(a: &Accumulator, g: u32) -> Option { + let out = a.evaluate(&[g]).unwrap(); + assert_eq!(out.data_type(), &DataType::Float64); + let out = out.as_primitive::(); + (!out.is_null(0)).then(|| out.value(0)) +} + +/// `Some(NaN)` and `Some(v)` compared by bits, `None` as is. +fn same(a: Option, b: Option) -> bool { + a.map(f64::to_bits) == b.map(f64::to_bits) +} + +fn close(a: Option, b: Option, tolerance: f64) -> bool { + match (a, b) { + (None, None) => true, + (Some(a), Some(b)) if a.is_nan() || b.is_nan() => a.is_nan() && b.is_nan(), + (Some(a), Some(b)) => (a - b).abs() <= tolerance * a.abs().max(b.abs()).max(1e-300), + _ => false, + } +} + +fn is_nan(v: Option) -> bool { + v.is_some_and(f64::is_nan) +} + +/// A deterministic pseudo-random sequence in `[-50, 50)` (an LCG). +fn sequence(n: usize, seed: u64) -> Vec { + let mut s = seed; + (0..n) + .map(|_| { + s = s + .wrapping_mul(6_364_136_223_846_793_005) + .wrapping_add(1_442_695_040_888_963_407); + #[expect(clippy::cast_precision_loss, reason = "a test value")] + let v = (s >> 33) as f64 / f64::from(1_u32 << 31) * 100.0 - 50.0; + v + }) + .collect() +} + +#[test] +fn no_rows_is_null() { + for f in ALL { + for legacy in [false, true] { + assert_eq!(global(f, legacy, &[], &[]), None, "{f:?} legacy={legacy}"); + // A group that only saw NULLs has no rows either. + let nulls = [None, None]; + assert_eq!(global(f, legacy, &nulls, &nulls), None, "{f:?}"); + } + } +} + +#[test] +fn one_row_divides_by_zero_where_spark_does() { + let (x, y) = (some(&[5.0]), some(&[7.0])); + for f in ALL { + let strict = global(f, false, &x, &y); + let legacy = global(f, true, &x, &y); + if [Skewness, Kurtosis].contains(&f) || SINGLE_ROW_DIVIDES.contains(&f) { + assert_eq!(strict, None, "{f:?}"); + assert!(is_nan(legacy), "{f:?}: {legacy:?}"); + } else { + // The population statistics of one row are 0. + assert_eq!(strict, Some(0.0), "{f:?}"); + assert_eq!(legacy, Some(0.0), "{f:?}"); + } + } +} + +#[test] +fn constant_values_have_no_variance() { + let (x, y) = (some(&[3.0; 4]), some(&[2.0; 4])); + for f in ALL { + let strict = global(f, false, &x, &y); + let legacy = global(f, true, &x, &y); + match f { + Skewness | Kurtosis => { + assert_eq!(strict, None, "{f:?}"); + assert!(is_nan(legacy), "{f:?}"); + } + // Several rows over a constant column: NULL without ANSI under + // both settings of the legacy flag, never NaN (see + // `corr_of_a_constant_column_is_a_division_by_zero_under_ansi`). + Corr => { + assert_eq!((strict, legacy), (None, None), "{f:?}"); + } + _ => { + assert_eq!(strict, Some(0.0), "{f:?}"); + assert_eq!(legacy, Some(0.0), "{f:?}"); + } + } + } +} + +#[test] +fn corr_of_a_constant_column_is_a_division_by_zero_under_ansi() { + let varying = some(&[1.0, 2.0, 4.0]); + let constant = some(&[3.0; 3]); + for (xs, ys) in [ + (&constant, &varying), + (&varying, &constant), + (&constant, &constant), + ] { + for legacy in [false, true] { + let mut a = acc_ansi(Corr, legacy, true, 1); + a.update(&[0; 3], &columns(Corr, xs, ys), &Weights::Const(1)) + .unwrap(); + let e = a.evaluate(&[0]).unwrap_err(); + assert_eq!(e.error_class(), ErrorClass::DivideByZero, "legacy={legacy}"); + // Without ANSI the same buffer is NULL. + assert_eq!(global(Corr, legacy, xs, ys), None); + } + } + // Rows with variance in both columns, and a single row, never raise + // (the single row is NULL, or NaN with the legacy flag). + for legacy in [false, true] { + let mut a = acc_ansi(Corr, legacy, true, 2); + a.update( + &[0, 0, 0, 1], + &columns( + Corr, + &some(&[1.0, 2.0, 4.0, 5.0]), + &some(&[2.0, 1.0, 7.0, 5.0]), + ), + &Weights::Const(1), + ) + .unwrap(); + assert!(a.evaluate(&[0]).is_ok()); + assert_eq!(value(&a, 1).is_some(), legacy); + } + // A group that raises does not stop the evaluation of the others when + // they are evaluated apart: only the selected groups are read. + let mut a = acc_ansi(Corr, false, true, 2); + a.update( + &[0, 0, 1, 1], + &columns( + Corr, + &some(&[1.0, 2.0, 3.0, 3.0]), + &some(&[1.0, 2.0, 1.0, 2.0]), + ), + &Weights::Const(1), + ) + .unwrap(); + assert!(value(&a, 0).is_some()); + assert!(a.evaluate(&[1]).is_err()); +} + +#[test] +fn rows_with_a_null_argument_are_skipped() { + let xs = [Some(1.0), None, Some(3.0), Some(5.0), Some(9.5), None]; + let ys = [Some(2.0), Some(9.0), None, Some(10.0), Some(1.5), None]; + for f in ALL { + let complete: Vec = (0..xs.len()) + .filter(|&i| xs[i].is_some() && (arity(f) == 1 || ys[i].is_some())) + .collect(); + let pick = |v: &[Option]| complete.iter().map(|&i| v[i]).collect::>(); + for legacy in [false, true] { + let got = global(f, legacy, &xs, &ys); + let want = global(f, legacy, &pick(&xs), &pick(&ys)); + assert!(same(got, want), "{f:?}: {got:?} vs {want:?}"); + assert!(got.is_some(), "{f:?}"); + } + } +} + +#[test] +fn values_under_null_slots_never_count() { + // The slots under NULL hold NaN and infinities: never read. + let values = vec![1.0, f64::NAN, 3.0, f64::INFINITY, 5.0, f64::NEG_INFINITY]; + let valid = NullBuffer::from(vec![true, false, true, false, true, false]); + let x: ArrayRef = Arc::new(Float64Array::new( + values.clone().into(), + Some(valid.clone()), + )); + let y: ArrayRef = Arc::new(Float64Array::new(values.into(), Some(valid))); + for f in ALL { + let mut a = acc(f, false, 1); + let cols = if arity(f) == 2 { + vec![Arc::clone(&x), Arc::clone(&y)] + } else { + vec![Arc::clone(&x)] + }; + a.update(&[0; 6], &cols, &Weights::Const(1)).unwrap(); + let (clean_x, clean_y) = (some(&[1.0, 3.0, 5.0]), some(&[1.0, 3.0, 5.0])); + let want = global(f, false, &clean_x, &clean_y); + assert!(same(value(&a, 0), want), "{f:?}"); + } +} + +#[test] +fn nan_and_infinity_propagate() { + let ys = some(&[1.0, 2.0, 3.0]); + for bad in [f64::NAN, f64::INFINITY, f64::NEG_INFINITY] { + let xs = some(&[1.0, bad, 3.0]); + for f in ALL { + for legacy in [false, true] { + let v = global(f, legacy, &xs, &ys); + assert!(is_nan(v), "{f:?}({bad}): {v:?}"); + } + } + } + // The pair statistics see a NaN in either argument. + for f in [Corr, CovarPop, CovarSamp] { + let v = global(f, false, &ys, &some(&[1.0, f64::NAN, 3.0])); + assert!(is_nan(v), "{f:?}"); + } + // A NaN is a value, not a NULL: it counts as a row. One row still + // divides by zero first (Spark checks the count before the value). + let nan = some(&[f64::NAN]); + assert!(is_nan(global(StddevPop, false, &nan, &[]))); + assert_eq!(global(StddevSamp, false, &nan, &[]), None); + assert!(is_nan(global(StddevSamp, true, &nan, &[]))); +} + +/// A two-pass reference: the central moments and co-moments of `xs`, `ys`. +struct Reference { + n: f64, + m2: f64, + m3: f64, + m4: f64, + ck: f64, + y_m2: f64, +} + +fn reference(xs: &[f64], ys: &[f64]) -> Reference { + #[expect(clippy::cast_precision_loss, reason = "a test count")] + let n = xs.len() as f64; + let (mx, my) = (xs.iter().sum::() / n, ys.iter().sum::() / n); + let pow = |k: i32| xs.iter().map(|x| (x - mx).powi(k)).sum::(); + Reference { + n, + m2: pow(2), + m3: pow(3), + m4: pow(4), + ck: xs.iter().zip(ys).map(|(x, y)| (x - mx) * (y - my)).sum(), + y_m2: ys.iter().map(|y| (y - my).powi(2)).sum(), + } +} + +#[test] +fn matches_the_textbook_definitions() { + let (xs, ys) = (sequence(400, 7), sequence(400, 99)); + let (x, y) = (some(&xs), some(&ys)); + let r = reference(&xs, &ys); + let n = r.n; + let want = |f| match f { + StddevSamp => (r.m2 / (n - 1.0)).sqrt(), + StddevPop => (r.m2 / n).sqrt(), + VarSamp => r.m2 / (n - 1.0), + VarPop => r.m2 / n, + Skewness => n.sqrt() * r.m3 / r.m2.powf(1.5), + Kurtosis => n * r.m4 / (r.m2 * r.m2) - 3.0, + Corr => r.ck / (r.m2 * r.y_m2).sqrt(), + CovarSamp => r.ck / (n - 1.0), + CovarPop => r.ck / n, + _ => f64::NAN, + }; + for f in ALL { + let got = global(f, false, &x, &y); + assert!( + close(got, Some(want(f)), 1e-9), + "{f:?}: {got:?} vs {}", + want(f) + ); + } +} + +#[test] +fn values_far_from_zero_keep_their_precision() { + // The spread is 0..=4 around 1e9: the sum-of-squares formula loses it + // all (1e18 magnitudes, 1e-16 precision). + let xs: Vec = (0..1000_u32).map(|i| 1e9 + f64::from(i % 5)).collect(); + let ys: Vec = (0..1000_u32) + .map(|i| 3e9 - f64::from(i % 5) * 2.0) + .collect(); + let (x, y) = (some(&xs), some(&ys)); + // The sum-of-squares formula is hopeless here (it answers 0). + #[expect(clippy::cast_precision_loss, reason = "a test count")] + let n = xs.len() as f64; + let naive = xs.iter().map(|x| x * x).sum::() / n - (xs.iter().sum::() / n).powi(2); + assert!((naive - 2.0).abs() > 1.0, "{naive}"); + // The values 0..=4 equally often have variance 2. + let within = |v, want| close(v, Some(want), 1e-6); + assert!(within(global(VarPop, false, &x, &[]), 2.0)); + assert!(within(global(StddevPop, false, &x, &[]), 2.0_f64.sqrt())); + assert!(within(global(CovarPop, false, &x, &y), -4.0)); + assert!(within(global(Corr, false, &x, &y), -1.0)); +} + +/// A recorded result: the function, its arguments and Spark's value. +type Recorded<'a> = (AggregateFunction, &'a [Option], &'a [Option], f64); + +/// The order partial states are merged in, by position. +type Order = Box usize>; + +/// Spark's slicing of `rows` rows into `p` partitions +/// (`ParallelCollectionRDD.slice`). +fn spark_slices(rows: usize, p: usize) -> Vec> { + (0..p).map(|i| i * rows / p..(i + 1) * rows / p).collect() +} + +/// Spark's evaluation: each partition folds its rows into a partial buffer; +/// the final stage merges the partials, in order, into an empty buffer. +fn spark_style( + f: AggregateFunction, + legacy: bool, + xs: &[Option], + ys: &[Option], + partitions: usize, +) -> Option { + let mut last = acc(f, legacy, 1); + for r in spark_slices(xs.len(), partitions) { + let mut partial = acc(f, legacy, 1); + let (px, py) = (&xs[r.clone()], ys.get(r.clone()).unwrap_or_default()); + partial + .update(&vec![0; r.len()], &columns(f, px, py), &Weights::Const(1)) + .unwrap(); + last.merge(&[0], &partial.state(&[0]).unwrap()).unwrap(); + } + value(&last, 0) +} + +/// Results Spark 4.2.0 recorded, reproduced to the bit by Spark's own +/// partitioning of the rows (which the digits depend on: one pass over the +/// first query's rows differs in the last digits). The partition counts are +/// the ones under which the recorded digits come out: `SQLQueryTestSuite` +/// runs two partitions, the function matrix one row per partition. +#[test] +fn recorded_spark_results_are_reproduced_to_the_bit() { + let t = |v: &[f64]| some(v); + let none: Vec> = Vec::new(); + // group-by.sql: `SELECT SKEWNESS(a), KURTOSIS(a), ..., VARIANCE(a), + // STDDEV(a) ... FROM testData`, two partitions. + let a = [ + Some(1.0), + Some(1.0), + Some(2.0), + Some(2.0), + Some(3.0), + Some(3.0), + None, + Some(3.0), + None, + ]; + for (f, want) in [ + (Skewness, -0.272_380_105_814_572_9), + (Kurtosis, -1.506_920_415_224_913_4), + (VarSamp, 0.809_523_809_523_809_4), + (StddevSamp, 0.899_735_410_842_437_2), + ] { + let got = spark_style(f, false, &a, &none, 2); + assert!(same(got, Some(want)), "{f:?}: {got:?} vs {want:?}"); + } + // linear-regression.sql: `SELECT k, corr(y, x) ... GROUP BY k`, group + // k = 2 of the six-row table cut in two partitions: its rows + // (y, x) = (10, 11), (20, 22) | (25, NULL), (30, 35), (NULL, 40). + let mut last = acc(Corr, false, 1); + for (ys, xs) in [ + (vec![Some(10.0), Some(20.0)], vec![Some(11.0), Some(22.0)]), + ( + vec![Some(25.0), Some(30.0), None], + vec![None, Some(35.0), Some(40.0)], + ), + ] { + let mut partial = acc(Corr, false, 1); + partial + .update( + &[0; 3][..ys.len()], + &columns(Corr, &ys, &xs), + &Weights::Const(1), + ) + .unwrap(); + last.merge(&[0], &partial.state(&[0]).unwrap()).unwrap(); + } + assert!(same(value(&last, 0), Some(0.998_844_598_112_153_3))); + // The function matrix (`function-matrix/cases/agg.jsonl`), one row per + // partition. + let (m10, m100) = ( + t(&[-10.0, -20.0, 100.0, 1000.0]), + t(&[1.0, 10.0, 100.0, 10.0, 1.0]), + ); + let m1000 = t(&[-1000.0, -100.0, 10.0, 20.0]); + let one_two_three = t(&[1.0, 2.0, 3.0]); + let pairs_x = t(&[3.0, 3.0, 6.0]); + let pairs_y = t(&[2.0, 3.0, 4.0]); + let cases: [Recorded<'_>; 11] = [ + (Skewness, &m10, &none, 1.113_565_746_902_201_3), + (Skewness, &m1000, &none, -1.113_565_746_902_201_1), + (Kurtosis, &m10, &none, -0.701_436_804_752_961_8), + (Kurtosis, &m100, &none, 0.194_323_231_916_989_86), + (StddevSamp, &one_two_three, &none, 1.0), + (StddevPop, &one_two_three, &none, 0.816_496_580_927_726), + (VarSamp, &one_two_three, &none, 1.0), + (VarPop, &one_two_three, &none, 0.666_666_666_666_666_6), + (CovarSamp, &one_two_three, &one_two_three, 1.0), + ( + CovarPop, + &one_two_three, + &one_two_three, + 0.666_666_666_666_666_6, + ), + (Corr, &pairs_x, &pairs_y, 0.866_025_403_784_438_7), + ]; + for (f, x, y, want) in cases { + let got = spark_style(f, false, x, y, x.len()); + assert!( + same(got, Some(want)), + "{f:?} over {x:?}: {got:?} vs {want:?}" + ); + } + // A pass over one partition is another rounding of the same value. + let one_pass = global(Skewness, false, &m10, &none); + assert!(close(one_pass, Some(1.113_565_746_902_201_3), 1e-14)); + assert!(!same(one_pass, Some(1.113_565_746_902_201_3))); +} + +/// `values` in `parts` contiguous pieces, each folded on its own and the +/// partial states merged into a fresh accumulator in `order`. +fn split_merge( + f: AggregateFunction, + xs: &[Option], + ys: &[Option], + parts: usize, + order: &dyn Fn(usize) -> usize, +) -> Option { + let slices = spark_slices(xs.len(), parts); + let states: Vec> = slices + .iter() + .map(|r| { + let mut a = acc(f, false, 1); + let (px, py) = (&xs[r.clone()], ys.get(r.clone()).unwrap_or_default()); + a.update(&vec![0; r.len()], &columns(f, px, py), &Weights::Const(1)) + .unwrap(); + a.state(&[0]).unwrap() + }) + .collect(); + let mut total = acc(f, false, 1); + for i in 0..parts { + total.merge(&[0], &states[order(i)]).unwrap(); + } + value(&total, 0) +} + +#[test] +fn any_split_merged_in_any_order_agrees_with_one_pass() { + // NULLs in the middle, so some pieces are empty or all-NULL. + let mut xs = some(&sequence(120, 3)); + let mut ys = some(&sequence(120, 4)); + for i in (0..120).step_by(7) { + xs[i] = None; + } + for i in (3..120).step_by(11) { + ys[i] = None; + } + for f in ALL { + let once = global(f, false, &xs, &ys); + for parts in 1..=9 { + let orders: [(&str, Order); 3] = [ + ("forward", Box::new(|i| i)), + ("reverse", Box::new(move |i| parts - 1 - i)), + ("rotated", Box::new(move |i| (i + 2) % parts)), + ]; + for (name, order) in orders { + let got = split_merge(f, &xs, &ys, parts, &order); + assert!( + close(got, once, 1e-9), + "{f:?} {parts} {name}: {got:?} vs {once:?}" + ); + // The same split in the same order is deterministic. + assert!(same(got, split_merge(f, &xs, &ys, parts, &order))); + } + } + } +} + +#[test] +fn empty_partial_states_merge_to_the_same_value() { + let (xs, ys) = (some(&sequence(31, 5)), some(&sequence(31, 6))); + for f in ALL { + let want = global(f, false, &xs, &ys); + let mut full = acc(f, false, 1); + full.update(&[0; 31], &columns(f, &xs, &ys), &Weights::Const(1)) + .unwrap(); + let empty = acc(f, false, 1).state(&[0]).unwrap(); + let state = full.state(&[0]).unwrap(); + // Spark's final stage merges every partial into an empty buffer; + // an empty partial merged last changes nothing. + let mut first = acc(f, false, 1); + first.merge(&[0], &state).unwrap(); + let mut last = acc(f, false, 1); + last.merge(&[0], &state).unwrap(); + last.merge(&[0], &empty).unwrap(); + for a in [&first, &last] { + assert!( + same(value(a, 0), want), + "{f:?}: {:?} vs {want:?}", + value(a, 0) + ); + } + // Only empty partials: no rows. + let mut none = acc(f, false, 1); + none.merge(&[0], &empty).unwrap(); + none.merge(&[0], &empty).unwrap(); + assert_eq!(value(&none, 0), None, "{f:?}"); + } +} + +#[test] +fn groups_are_independent() { + let xs = some(&[1.0, 10.0, 2.0, 20.0, 3.0, 30.0, 5.0]); + let ys = some(&[2.0, 5.0, 4.0, 4.0, 7.0, 3.0, 1.0]); + let groups = [0_u32, 1, 0, 1, 0, 1, 2]; + for f in ALL { + let mut a = acc(f, false, 3); + a.update(&groups, &columns(f, &xs, &ys), &Weights::Const(1)) + .unwrap(); + let pick = |v: &[Option], g: u32| { + groups + .iter() + .zip(v) + .filter(|(k, _)| **k == g) + .map(|(_, v)| *v) + .collect::>() + }; + for g in 0..3 { + let want = global(f, false, &pick(&xs, g), &pick(&ys, g)); + assert!(same(value(&a, g), want), "{f:?} group {g}"); + } + // Several ids at once, in any order. + let all = a.evaluate(&[2, 0, 1]).unwrap(); + assert_eq!(all.len(), 3); + assert!(same( + all.as_primitive::().iter().next().unwrap(), + value(&a, 2) + )); + } +} + +#[test] +fn integer_and_float_inputs_are_read_as_doubles() { + let ints: ArrayRef = Arc::new(Int32Array::from(vec![Some(1), None, Some(2), Some(4)])); + let floats: ArrayRef = Arc::new(Float32Array::from(vec![ + Some(1.0), + None, + Some(2.0), + Some(4.0), + ])); + let want = global(VarSamp, false, &some(&[1.0, 2.0, 4.0]), &[]); + for (col, t) in [(ints, DataType::Int32), (floats, DataType::Float32)] { + let e = AggregateExpr { + arg_type: t, + ..expr(VarSamp) + }; + assert_eq!(e.result_type().unwrap(), DataType::Float64); + let mut a = Accumulator::new(&e, true).unwrap(); + a.resize(1); + a.update(&[0; 4], &[col], &Weights::Const(1)).unwrap(); + assert!(same(value(&a, 0), want)); + } +} + +#[test] +fn state_columns_are_sparks_buffers() { + let names = |f| -> Vec { + expr(f) + .state_fields() + .unwrap() + .iter() + .map(|c| { + assert_eq!(c.data_type(), &DataType::Float64); + assert!(!c.is_nullable(), "a buffer is never NULL"); + c.name().clone() + }) + .collect() + }; + assert_eq!(names(StddevSamp), ["m#n", "m#avg", "m#m2"]); + assert_eq!(names(VarPop), ["m#n", "m#avg", "m#m2"]); + assert_eq!(names(Skewness), ["m#n", "m#avg", "m#m2", "m#m3"]); + assert_eq!(names(Kurtosis), ["m#n", "m#avg", "m#m2", "m#m3", "m#m4"]); + assert_eq!(names(CovarPop), ["m#n", "m#xAvg", "m#yAvg", "m#ck"]); + assert_eq!( + names(Corr), + ["m#n", "m#xAvg", "m#yAvg", "m#ck", "m#xMk", "m#yMk"] + ); + for f in ALL { + let a = acc(f, false, 2); + let state = a.state(&[0, 1]).unwrap(); + assert_eq!(state.len(), names(f).len(), "{f:?}"); + assert!(state.iter().all(|c| c.len() == 2 && c.null_count() == 0)); + assert_eq!(a.state_len(), state.len()); + // The merging form reads those columns, wherever they sit. + let merging = expr(f).merging(4); + assert_eq!(merging.input, AggInput::State); + assert_eq!(merging.args, (4..4 + state.len()).collect::>()); + } +} + +#[test] +fn only_numbers_are_read_and_the_result_is_double() { + for f in ALL { + for t in [ + DataType::Int8, + DataType::Int64, + DataType::Float32, + DataType::Float64, + ] { + let e = AggregateExpr { + arg_type: t, + ..expr(f) + }; + assert_eq!(e.result_type().unwrap(), DataType::Float64, "{f:?}"); + } + // The analyzer casts everything else to double before it gets here. + for t in [ + DataType::Utf8View, + DataType::Decimal128(10, 2), + DataType::Boolean, + ] { + let e = AggregateExpr { + arg_type: t, + ..expr(f) + }; + let err = e.result_type().unwrap_err(); + assert_eq!(err.error_class(), ErrorClass::KrishivUnsupportedFeature); + assert!(e.state_fields().is_err()); + } + } +} + +#[test] +fn row_weights_other_than_one_are_refused() { + for f in ALL { + let mut a = acc(f, false, 1); + let cols = columns(f, &some(&[1.0, 2.0]), &some(&[3.0, 4.0])); + let err = a.update(&[0, 0], &cols, &Weights::Const(2)).unwrap_err(); + assert_eq!(err.error_class(), ErrorClass::KrishivUnsupportedFeature); + let column = Weights::Column(arrow_array::Int64Array::from(vec![1, -1])); + assert!(a.update(&[0, 0], &cols, &column).is_err()); + assert!(a.update(&[0, 0], &cols, &Weights::Const(1)).is_ok()); + // Nothing of the refused batches was folded in. + assert_eq!( + a.state(&[0]).unwrap()[0] + .as_primitive::() + .value(0), + 2.0 + ); + } +} + +#[test] +fn malformed_inputs_are_errors_not_panics() { + for f in ALL { + let mut a = acc(f, false, 1); + let cols = columns(f, &some(&[1.0, 2.0]), &some(&[1.0, 2.0])); + // A group that does not exist. + let err = a.update(&[0, 5], &cols, &Weights::Const(1)).unwrap_err(); + assert_eq!(err.error_class(), ErrorClass::InternalError, "{f:?}"); + assert!(a.evaluate(&[3]).is_err()); + assert!(a.state(&[3]).is_err()); + // Columns of another length than the rows. + assert!(a.update(&[0], &cols, &Weights::Const(1)).is_err()); + // A running fold has no moment form. + let run = crate::aggregate::Running { + parts: &[0, 0], + snapshots: &[crate::aggregate::Running::NONE; 2], + }; + assert!(a.update_running(&[0, 0], &cols, &run).is_err()); + // States: too few columns, NULLs, another type, another length. + let good = acc(f, false, 1).state(&[0]).unwrap(); + assert!(a.merge(&[0], &good[..good.len() - 1]).is_err()); + let mut nulled = good.clone(); + nulled[0] = doubles(&[None]); + assert!(a.merge(&[0], &nulled).is_err()); + let mut mistyped = good.clone(); + mistyped[0] = Arc::new(Int32Array::from(vec![1])); + assert!(a.merge(&[0], &mistyped).is_err()); + assert!(a.merge(&[0, 0], &good).is_err()); + } + // A moment kernel only exists for moment aggregates. + assert!(super::Moments::new(AggregateFunction::Count, true).is_err()); +} diff --git a/crates/krishiv-exec/src/aggregate/node.rs b/crates/krishiv-exec/src/aggregate/node.rs index 8822ad13d..2c1dc1265 100644 --- a/crates/krishiv-exec/src/aggregate/node.rs +++ b/crates/krishiv-exec/src/aggregate/node.rs @@ -157,6 +157,7 @@ pub fn union_distinct( ansi: true, spill, require_unit_weights: true, + legacy_statistical: false, }; Ok(AggregateNode::new(input, HashAggregate::new(spec)?)) } diff --git a/crates/krishiv-exec/src/aggregate/rowtable.rs b/crates/krishiv-exec/src/aggregate/rowtable.rs index 26999852e..c4078beba 100644 --- a/crates/krishiv-exec/src/aggregate/rowtable.rs +++ b/crates/krishiv-exec/src/aggregate/rowtable.rs @@ -1477,7 +1477,7 @@ fn new_accs(core: &AggCore) -> Result> { core.spec .aggregates .iter() - .map(|a| Accumulator::new(a, core.spec.ansi)) + .map(|a| core.accumulator(a)) .collect() } @@ -1709,6 +1709,7 @@ mod tests { ansi: true, spill: None, require_unit_weights: false, + legacy_statistical: false, }; Arc::new(AggCore::new(spec).unwrap()) } @@ -1742,6 +1743,7 @@ mod tests { ansi: true, spill: None, require_unit_weights: false, + legacy_statistical: false, }; Arc::new(AggCore::new(spec).unwrap()) } @@ -2202,6 +2204,7 @@ mod tests { ansi: true, spill: None, require_unit_weights: false, + legacy_statistical: false, }; let core = Arc::new(AggCore::new(spec).unwrap()); let seed = HashSeed::new(Some(6)); diff --git a/crates/krishiv-exec/src/aggregate/table.rs b/crates/krishiv-exec/src/aggregate/table.rs index 72491087d..91e7bf6a5 100644 --- a/crates/krishiv-exec/src/aggregate/table.rs +++ b/crates/krishiv-exec/src/aggregate/table.rs @@ -12,7 +12,7 @@ use krishiv_types::batch::{Batch, Weights}; use super::accumulator::Accumulator; use super::rowtable::{RowPrepared, RowTable}; -use super::{AggregateOutput, AggregateSpec}; +use super::{AggregateExpr, AggregateOutput, AggregateSpec}; use crate::error::{Result, internal}; use crate::hash::{HashSeed, KeyEncoder, KeyTable, ProbeCount, ProbeTotals, spill_partition}; use crate::rows::RowLayout; @@ -75,6 +75,11 @@ impl AggCore { self.spec.group_by.len() } + /// The accumulator of aggregate `a`, with the spec's semantics flags. + pub(crate) fn accumulator(&self, a: &AggregateExpr) -> Result { + Ok(Accumulator::new(a, self.spec.ansi)?.legacy_statistical(self.spec.legacy_statistical)) + } + /// Row indices of a state batch per spill partition at `level`. pub(crate) fn partitions_of( &self, @@ -152,7 +157,7 @@ impl AggTable { .spec .aggregates .iter() - .map(|a| Accumulator::new(a, core.spec.ansi)) + .map(|a| core.accumulator(a)) .collect::>>()?; let mut t = Self { core: Arc::clone(core), @@ -341,7 +346,7 @@ impl AggTable { let ids: Vec = (0..u32::try_from(n).map_err(|_| internal("batch too large".to_owned()))?).collect(); for a in &core.spec.aggregates { - let mut acc = Accumulator::new(a, core.spec.ansi)?; + let mut acc = core.accumulator(a)?; acc.resize(n); acc.update(&ids, cols, batch.weights())?; out.extend(acc.state(&ids)?); diff --git a/crates/krishiv-exec/src/aggregate/tests.rs b/crates/krishiv-exec/src/aggregate/tests.rs index 43193e659..687fe0513 100644 --- a/crates/krishiv-exec/src/aggregate/tests.rs +++ b/crates/krishiv-exec/src/aggregate/tests.rs @@ -237,6 +237,7 @@ fn spec(s: &SchemaRef, group_by: Vec, aggregates: Vec) -> ansi: true, spill: None, require_unit_weights: false, + legacy_statistical: false, } } @@ -670,9 +671,20 @@ fn two_phase( policy: BypassPolicy, spill: Option, ) -> (Arc, Arc) { - let mut first = spec(s, vec![0], grouped_aggs(s)); + two_phase_of(input, spec(s, vec![0], grouped_aggs(s)), policy, spill) +} + +/// [`two_phase`] of the aggregation `first` (raw input, final output). +fn two_phase_of( + input: Arc, + mut first: AggregateSpec, + policy: BypassPolicy, + spill: Option, +) -> (Arc, Arc) { first.output = AggregateOutput::Partial; - let second = second_phase(&first, spill); + let mut second = second_phase(&first, spill); + second.legacy_statistical = first.legacy_statistical; + second.ansi = first.ansi; let partial = PartialAggregate::new(first, policy, Some(7)).unwrap(); let plan = node(PartialNode::new(input, Arc::clone(&partial)), second); (plan, partial) @@ -2577,3 +2589,451 @@ async fn a_refused_widening_merges_in_one_task() { assert_eq!(sink_metric(&agg, "row_final_tasks"), 1); assert_eq!(sink_metric(&agg, "row_final_spills"), 0); } + +// ---- the moment aggregates ---- + +const MOMENTS: [AggregateFunction; 9] = [ + AggregateFunction::StddevSamp, + AggregateFunction::StddevPop, + AggregateFunction::VarSamp, + AggregateFunction::VarPop, + AggregateFunction::Skewness, + AggregateFunction::Kurtosis, + AggregateFunction::Corr, + AggregateFunction::CovarSamp, + AggregateFunction::CovarPop, +]; + +/// Every moment aggregate over `x` (column 1) and, for the pair +/// statistics, `y` (column 2). +fn moment_aggs() -> Vec { + MOMENTS + .iter() + .map(|&f| AggregateExpr { + function: f, + args: if super::arity(f) == 2 { + vec![1, 2] + } else { + vec![1] + }, + input: AggInput::Raw, + arg_type: DataType::Float64, + name: format!("{f:?}"), + }) + .collect() +} + +fn f64s(v: Vec>) -> ArrayRef { + Arc::new(Float64Array::from(v)) +} + +/// `n` rows `k = i % keys`, a `double` x (NULL every 7th) and a `double` y +/// (NULL every 11th), in batches of 1000. +fn moment_input(n: i64, keys: i64) -> (SchemaRef, Vec>) { + let s = schema(vec![ + ("k", DataType::Int64), + ("x", DataType::Float64), + ("y", DataType::Float64), + ]); + #[expect(clippy::cast_precision_loss, reason = "small test values")] + let batches = (0..n) + .step_by(1000) + .map(|start| { + let r = start..(start + 1000).min(n); + let k: Vec<_> = r.clone().map(|i| Some(i % keys)).collect(); + let x = r + .clone() + .map(|i| (i % 7 != 0).then(|| (i * 37 % 101) as f64 / 3.0)); + let y = r.map(|i| (i % 11 != 0).then(|| (i * 53 % 97) as f64 / 7.0 - (i % 5) as f64)); + vec![i64s(&k), f64s(x.collect()), f64s(y.collect())] + }) + .collect(); + (s, batches) +} + +type Moments = BTreeMap>>; + +/// The output of a moment aggregation by key (`""` without keys). +fn moment_rows(batches: &[Batch], keyed: bool) -> Moments { + rows(batches) + .iter() + .map(|r| { + let mut cells = r.split('|'); + let key = if keyed { cells.next().unwrap() } else { "" }; + let values = cells + .map(|c| (c != "NULL").then(|| c.parse::().unwrap())) + .collect(); + (key.to_owned(), values) + }) + .collect() +} + +/// The aggregates of every group folded in one pass by one accumulator +/// each: what a single driver computes. +fn moment_oracle(b: &[Vec], keys: i64, keyed: bool, legacy: bool) -> Moments { + let groups = if keyed { + usize::try_from(keys).unwrap() + } else { + 1 + }; + let mut accs: Vec = moment_aggs() + .iter() + .map(|a| { + let mut acc = Accumulator::new(a, true) + .unwrap() + .legacy_statistical(legacy); + acc.resize(groups); + acc + }) + .collect(); + for cols in b { + let ids: Vec = cols[0] + .as_primitive::() + .values() + .iter() + .map(|&k| if keyed { u32::try_from(k).unwrap() } else { 0 }) + .collect(); + for a in &mut accs { + a.update(&ids, cols, &Weights::Const(1)).unwrap(); + } + } + let ids: Vec = (0..u32::try_from(groups).unwrap()).collect(); + let columns: Vec = accs.iter().map(|a| a.evaluate(&ids).unwrap()).collect(); + ids.iter() + .map(|&g| { + let values = columns + .iter() + .map(|c| { + let c = c.as_primitive::(); + (!c.is_null(g as usize)).then(|| c.value(g as usize)) + }) + .collect(); + (if keyed { g.to_string() } else { String::new() }, values) + }) + .collect() +} + +/// The values agree to `1e-9` (merging partial states reassociates the +/// arithmetic), NULLs and NaNs exactly. +fn assert_moments_close(got: &Moments, want: &Moments, what: &str) { + assert_eq!( + got.keys().collect::>(), + want.keys().collect::>(), + "{what}" + ); + for (key, want) in want { + for (i, (g, w)) in got[key].iter().zip(want).enumerate() { + let ok = match (g, w) { + (None, None) => true, + (Some(g), Some(w)) if g.is_nan() || w.is_nan() => g.is_nan() && w.is_nan(), + (Some(g), Some(w)) => (g - w).abs() <= 1e-9 * g.abs().max(w.abs()).max(1.0), + _ => false, + }; + assert!(ok, "{what}: key {key:?} column {i}: {g:?} vs {w:?}"); + } + } +} + +/// Runs `plan` on each layout and checks both against `want`. +async fn check_both_layouts( + make: impl Fn() -> Arc, + o: Opts, + keyed: bool, + want: &Moments, +) { + for rows_layout in [false, true] { + let (out, _) = exec_with(make(), o, rows_layout).await.unwrap(); + assert_moments_close( + &moment_rows(&out, keyed), + want, + &format!("row layout {rows_layout}"), + ); + } +} + +/// Grouped and global moment aggregates, on both table layouts, against a +/// single-pass reference. +#[tokio::test] +async fn moments_grouped_and_global() { + let (s, b) = moment_input(8_000, 13); + let o = Opts { + both_layouts: false, + ..Opts::default() + }; + let grouped = || node(scan(&s, b.clone()), spec(&s, vec![0], moment_aggs())); + check_both_layouts(grouped, o, true, &moment_oracle(&b, 13, true, false)).await; + let global = || node(scan(&s, b.clone()), spec(&s, vec![], moment_aggs())); + check_both_layouts(global, o, false, &moment_oracle(&b, 13, false, false)).await; + // A global aggregate of no rows is one row of NULLs; a grouped one has none. + let empty = |keys: Vec| node(scan(&s, vec![]), spec(&s, keys, moment_aggs())); + let (out, _) = exec(empty(vec![]), o).await.unwrap(); + assert_eq!(rows(&out), ["NULL|NULL|NULL|NULL|NULL|NULL|NULL|NULL|NULL"]); + assert!(rows(&exec(empty(vec![0]), o).await.unwrap().0).is_empty()); +} + +/// `spark.sql.legacy.statisticalAggregate` reaches the accumulators of +/// both layouts and of the final phase of a two-phase aggregation: groups +/// of one row give NULL, or NaN with the flag. +#[tokio::test] +async fn moments_legacy_flag_reaches_every_accumulator() { + // Every group has one row, so the sample statistics divide by zero. + let (s, b) = moment_input(40, 40); + let o = Opts { + both_layouts: false, + ..Opts::default() + }; + for legacy in [false, true] { + let mut sp = spec(&s, vec![0], moment_aggs()); + sp.legacy_statistical = legacy; + let want = moment_oracle(&b, 40, true, legacy); + // The oracle itself shows the flag: some groups' cells are NaN or NULL. + let nan_cells = want + .values() + .flatten() + .filter(|v| v.is_some_and(f64::is_nan)) + .count(); + assert_eq!(nan_cells > 0, legacy); + let single = || node(scan(&s, b.clone()), sp.clone()); + check_both_layouts(single, o, true, &want).await; + let phased = || { + two_phase_of( + scan(&s, b.clone()), + sp.clone(), + BypassPolicy::default(), + None, + ) + .0 + }; + check_both_layouts(phased, o, true, &want).await; + } +} + +type Plan<'a> = Box]) -> Arc + 'a>; + +/// `corr` of several rows over a constant column: `DIVIDE_BY_ZERO` with ANSI +/// and NULL without, on both layouts and across a partial and a final +/// phase, whatever the legacy flag says; the other groups are unaffected +/// when none is constant. +#[tokio::test] +async fn corr_of_a_constant_column_follows_ansi() { + let s = schema(vec![ + ("k", DataType::Int64), + ("x", DataType::Float64), + ("y", DataType::Float64), + ]); + let batch = |k: &[i64], x: &[f64], y: &[f64]| { + let some = |v: &[f64]| f64s(v.iter().map(|&v| Some(v)).collect()); + let k: Vec<_> = k.iter().map(|&k| Some(k)).collect(); + vec![i64s(&k), some(x), some(y)] + }; + let corr = || vec![moment_aggs().remove(6)]; + let o = Opts { + both_layouts: false, + ..Opts::default() + }; + let varying = batch( + &[1, 1, 2, 2, 2], + &[1.0, 2.0, 1.0, 2.0, 4.0], + &[3.0, 1.0, 2.0, 5.0, 6.0], + ); + let constant = batch( + &[1, 1, 2, 2, 2], + &[1.0, 2.0, 7.0, 7.0, 7.0], + &[3.0, 1.0, 2.0, 5.0, 6.0], + ); + for legacy in [false, true] { + for ansi in [false, true] { + let mut sp = spec(&s, vec![0], corr()); + sp.ansi = ansi; + sp.legacy_statistical = legacy; + let plans: [Plan<'_>; 2] = [ + Box::new(|b| node(scan(&s, b.to_vec()), sp.clone())), + Box::new(|b| { + two_phase_of( + scan(&s, b.to_vec()), + sp.clone(), + BypassPolicy::default(), + None, + ) + .0 + }), + ]; + for plan in &plans { + for rows_layout in [false, true] { + let ok = exec_with(plan(std::slice::from_ref(&varying)), o, rows_layout) + .await + .unwrap() + .0; + assert!(moment_rows(&ok, true).values().all(|v| v[0].is_some())); + let got = + exec_with(plan(std::slice::from_ref(&constant)), o, rows_layout).await; + if ansi { + let e = got.unwrap_err(); + assert_eq!(e.error_class(), ErrorClass::DivideByZero, "{e}"); + } else { + let m = moment_rows(&got.unwrap().0, true); + assert_eq!(m["2"], [None], "legacy={legacy} layout={rows_layout}"); + assert!(m["1"][0].is_some()); + } + } + } + } + } +} + +/// Partial → final over partial states equals one phase, also when the +/// partial phase flushes and bypasses. +#[tokio::test] +async fn moments_two_phase_equals_single_phase() { + let (s, b) = moment_input(60_000, 700); + let want = moment_oracle(&b, 700, true, false); + let o = Opts { + both_layouts: false, + ..Opts::default() + }; + let tiny = BypassPolicy { + max_table_bytes: 64 << 10, + min_rows: u64::MAX, + ..BypassPolicy::default() + }; + let policies = [BypassPolicy::default(), tiny]; + for policy in policies { + let make = || { + two_phase_of( + scan(&s, b.clone()), + spec(&s, vec![0], moment_aggs()), + policy, + None, + ) + .0 + }; + check_both_layouts(make, o, true, &want).await; + } + // Unique keys: the partial phase passes rows through, one state per row. + let (s, b) = moment_input(40_000, 1_000_000); + let want = moment_oracle(&b, 40_000, true, false); + let bypass = BypassPolicy { + min_rows: 10_000, + ..BypassPolicy::default() + }; + let (plan, partial) = two_phase_of( + scan(&s, b.clone()), + spec(&s, vec![0], moment_aggs()), + bypass, + None, + ); + let (out, _) = exec_with(plan, o, false).await.unwrap(); + assert!(partial.bypassed_drivers() > 0); + assert_moments_close(&moment_rows(&out, true), &want, "bypassed"); +} + +/// Under a small budget the moment states spill with the others, merge from +/// the runs and give the in-memory result. +#[tokio::test] +async fn moments_spill_and_merge() { + let (s, b) = moment_input(30_000, 25_000); + // Three of the aggregates (one of each state width): the budget below + // holds about a third of their groups. + let aggs: Vec = [0, 4, 6] + .into_iter() + .map(|i| moment_aggs().remove(i)) + .collect(); + let keep = |m: Moments| -> Moments { + m.into_iter() + .map(|(k, v)| (k, [0, 4, 6].iter().map(|&i| v[i]).collect())) + .collect() + }; + let want = keep(moment_oracle(&b, 25_000, true, false)); + let (_m, scope) = spill_dir(); + let mut sp = spec(&s, vec![0], aggs); + sp.spill = Some(scope.clone()); + let o = Opts { + limit: 2 << 20, + both_layouts: false, + ..Opts::default() + }; + for rows_layout in [false, true] { + let (out, m) = exec_with(node(scan(&s, b.clone()), sp.clone()), o, rows_layout) + .await + .unwrap(); + assert!( + extra(&m, "HashAggregate", "spill_runs") > 0, + "rows {rows_layout}" + ); + assert_eq!(extra(&m, "HashAggregate", "groups"), 25_000); + assert_moments_close(&moment_rows(&out, true), &want, "spilled"); + } + let (plan, _) = two_phase_of( + scan(&s, b.clone()), + sp, + BypassPolicy::default(), + Some(scope), + ); + let (out, _) = exec_with(plan, o, false).await.unwrap(); + assert_moments_close(&moment_rows(&out, true), &want, "spilled two-phase"); +} + +/// `DISTINCT` pairs: `corr(DISTINCT x, y)` over the distinct non-NULL +/// pairs, next to a plain `count(*)`. +#[tokio::test] +async fn moments_over_distinct_arguments() { + let s = schema(vec![ + ("k", DataType::Int64), + ("x", DataType::Float64), + ("y", DataType::Float64), + ]); + let k = i64s(&[Some(1); 6]); + let x = f64s(vec![ + Some(1.0), + Some(1.0), + Some(2.0), + Some(3.0), + None, + Some(3.0), + ]); + let y = f64s(vec![ + Some(2.0), + Some(2.0), + Some(4.5), + Some(5.0), + Some(1.0), + Some(5.0), + ]); + let input = || { + scan( + &s, + vec![vec![Arc::clone(&k), Arc::clone(&x), Arc::clone(&y)]], + ) + }; + let mut corr = moment_aggs().remove(6); + assert_eq!(corr.function, AggregateFunction::Corr); + corr.name = "c".to_owned(); + let flags = [true, false]; + let plan = distinct_plan( + input(), + spec( + &s, + vec![], + vec![corr.clone(), AggregateExpr::count_star("n")], + ), + &flags, + ) + .unwrap(); + // The distinct pairs (1,2), (2,4.5), (3,5), with a correlation of its own. + let pairs = vec![Some(1.0), Some(2.0), Some(3.0)]; + let want = moment_oracle( + &[vec![ + i64s(&[Some(1); 3]), + f64s(pairs), + f64s(vec![Some(2.0), Some(4.5), Some(5.0)]), + ]], + 1, + false, + false, + ); + let got = moment_rows(&exec(plan, Opts::default()).await.unwrap().0, false); + let corr_at = |m: &Moments, i: usize| m[""][i]; + let want_corr = corr_at(&want, 6).unwrap(); + assert!((got[""][0].unwrap() - want_corr).abs() < 1e-12, "{got:?}"); + assert_eq!(got[""][1], Some(6.0)); +} diff --git a/crates/krishiv-exec/src/physical/instantiate.rs b/crates/krishiv-exec/src/physical/instantiate.rs index fed03b8a0..e45d40db4 100644 --- a/crates/krishiv-exec/src/physical/instantiate.rs +++ b/crates/krishiv-exec/src/physical/instantiate.rs @@ -365,8 +365,7 @@ impl Instantiation<'_> { } PhysNode::Aggregate { input, aggregate } => { let input = self.node(input)?; - let spec = - aggregate_spec(&input.schema, aggregate, self.env.ansi, self.spill.clone()); + let spec = aggregate_spec(&input.schema, aggregate, &self.env, self.spill.clone()); let schema = spec.output_schema()?; let node = aggregate_node(input.node, spec, &aggregate.distinct)?; Ok(Built { node, schema }) @@ -706,7 +705,7 @@ pub(crate) fn scan_filter( pub(crate) fn aggregate_spec( input: &SchemaRef, a: &AggregatePlan, - ansi: bool, + env: &EvalEnv, spill: Option, ) -> AggregateSpec { AggregateSpec { @@ -714,9 +713,10 @@ pub(crate) fn aggregate_spec( group_by: a.group_by.clone(), aggregates: a.aggregates.clone(), output: a.output, - ansi, + ansi: env.ansi, spill, require_unit_weights: false, + legacy_statistical: env.legacy_statistical_aggregate, } } diff --git a/crates/krishiv-exec/src/physical/mod.rs b/crates/krishiv-exec/src/physical/mod.rs index 17c825681..ed98a85a6 100644 --- a/crates/krishiv-exec/src/physical/mod.rs +++ b/crates/krishiv-exec/src/physical/mod.rs @@ -150,6 +150,9 @@ pub enum JoinDistribution { pub struct ExprEnv { /// `spark.sql.ansi.enabled`. pub ansi: bool, + /// `spark.sql.legacy.statisticalAggregate`: the sample moment aggregates + /// give NaN instead of NULL where Spark divides by zero. + pub legacy_statistical_aggregate: bool, /// The session time zone as a fixed offset in seconds, when it is one. pub zone: Option, /// The query's clock (microseconds since the epoch, UTC), read once at @@ -469,7 +472,7 @@ plan_data!( Partitioned, } ); -plan_data!(struct ExprEnv { ansi, zone, query_start_micros, user }); +plan_data!(struct ExprEnv { ansi, legacy_statistical_aggregate, zone, query_start_micros, user }); plan_data!(enum PhysNode { Scan(scan), OneRow, diff --git a/crates/krishiv-exec/src/planner/aggregate.rs b/crates/krishiv-exec/src/planner/aggregate.rs index e5a7d4ae4..41a3cfd28 100644 --- a/crates/krishiv-exec/src/planner/aggregate.rs +++ b/crates/krishiv-exec/src/planner/aggregate.rs @@ -17,7 +17,7 @@ use krishiv_plan::{ use super::build::{Planned, Planner, column_of, internal, list, unsupported}; use super::properties::Distribution; -use crate::aggregate::{AggregateExpr, AggregateFunction, AggregateOutput}; +use crate::aggregate::{AggregateExpr, AggregateFunction, AggregateOutput, arity}; use crate::error::Result; use crate::physical::{AggregatePlan, PhysNode, aggregate_node, aggregate_spec}; @@ -73,6 +73,15 @@ impl Planner<'_> { "avg" => AggregateFunction::Avg, "min" => AggregateFunction::Min, "max" => AggregateFunction::Max, + "stddev" | "std" | "stddev_samp" => AggregateFunction::StddevSamp, + "stddev_pop" => AggregateFunction::StddevPop, + "variance" | "var_samp" => AggregateFunction::VarSamp, + "var_pop" => AggregateFunction::VarPop, + "skewness" => AggregateFunction::Skewness, + "kurtosis" => AggregateFunction::Kurtosis, + "corr" => AggregateFunction::Corr, + "covar_samp" => AggregateFunction::CovarSamp, + "covar_pop" => AggregateFunction::CovarPop, // `first(x, ignoreNulls)`: the analyzer checked the flag is a // boolean literal. name if picks(name).is_some() => { @@ -94,8 +103,8 @@ impl Planner<'_> { return Err(unsupported("aggregate FILTER (WHERE ...) clauses")); } let args: Vec = columns[at.clone()].to_vec(); - if function != AggregateFunction::Count && args.len() != 1 { - return Err(internal("an aggregate of several arguments")); + if function != AggregateFunction::Count && args.len() != arity(function) { + return Err(internal("an aggregate with the wrong number of arguments")); } // Spark's STRING and BINARY are views here whatever carrier a // registered table uses; the accumulators read any carrier. @@ -146,7 +155,7 @@ impl Planner<'_> { }; // Validates the aggregate (over an input that is never run) and // gives its schema. - let spec = aggregate_spec(&input.schema, &plan, self.env.ansi, None); + let spec = aggregate_spec(&input.schema, &plan, &self.env, None); let schema = spec.output_schema()?; aggregate_node(Arc::new(crate::physical::Unbuilt), spec, &plan.distinct)?; let mut layout: Vec = columns[..k].iter().map(|c| input.layout[*c]).collect(); diff --git a/crates/krishiv-exec/src/planner/eval/mod.rs b/crates/krishiv-exec/src/planner/eval/mod.rs index 71fa27969..71c598918 100644 --- a/crates/krishiv-exec/src/planner/eval/mod.rs +++ b/crates/krishiv-exec/src/planner/eval/mod.rs @@ -35,7 +35,9 @@ use std::sync::Arc; use arrow_array::cast::AsArray; use arrow_array::{Array, ArrayRef, BooleanArray, RecordBatch, UInt32Array, new_null_array}; use arrow_schema::{DataType as ArrowType, SchemaRef}; -use krishiv_common::config::{ANSI_ENABLED, Config, SESSION_TIME_ZONE}; +use krishiv_common::config::{ + ANSI_ENABLED, Config, LEGACY_STATISTICAL_AGGREGATE, SESSION_TIME_ZONE, +}; use krishiv_kernels::KernelCtx; use krishiv_plan::{ AttrId, Builtin, Expr, ExprNode, FunctionId, FunctionKind, Literal, ResolvedExpr, @@ -153,6 +155,8 @@ pub(crate) trait PhysExpr: Send + Sync + fmt::Debug { pub(crate) struct EvalEnv { /// `spark.sql.ansi.enabled`. pub(crate) ansi: bool, + /// `spark.sql.legacy.statisticalAggregate`. + pub(crate) legacy_statistical_aggregate: bool, /// The session time zone as a fixed offset in seconds, when it is one /// (`UTC`, `+hh:mm`); region zones need the tz database (P2-T021b). pub(crate) zone: Option, @@ -171,6 +175,10 @@ impl EvalEnv { .get(&ANSI_ENABLED) .and_then(krishiv_common::config::ConfigValue::as_bool) .unwrap_or(true); + let legacy_statistical_aggregate = config + .get(&LEGACY_STATISTICAL_AGGREGATE) + .and_then(krishiv_common::config::ConfigValue::as_bool) + .unwrap_or(false); let zone = config .get(&SESSION_TIME_ZONE) .and_then(|v| v.as_str()) @@ -182,6 +190,7 @@ impl EvalEnv { .unwrap_or_default(); Self { ansi, + legacy_statistical_aggregate, zone, kernels: KernelCtx::new(), query_start_micros, @@ -193,6 +202,7 @@ impl EvalEnv { pub(crate) fn from_plan(env: &ExprEnv) -> Self { Self { ansi: env.ansi, + legacy_statistical_aggregate: env.legacy_statistical_aggregate, zone: env.zone, kernels: KernelCtx::new(), query_start_micros: env.query_start_micros, @@ -204,6 +214,7 @@ impl EvalEnv { pub(crate) fn to_plan(&self) -> ExprEnv { ExprEnv { ansi: self.ansi, + legacy_statistical_aggregate: self.legacy_statistical_aggregate, zone: self.zone, query_start_micros: self.query_start_micros, user: Arc::clone(&self.user), diff --git a/crates/krishiv-exec/src/profile/tests.rs b/crates/krishiv-exec/src/profile/tests.rs index e8d42dd65..e704ea1ab 100644 --- a/crates/krishiv-exec/src/profile/tests.rs +++ b/crates/krishiv-exec/src/profile/tests.rs @@ -17,6 +17,7 @@ use crate::physical::{ExprEnv, PhysNode}; fn env() -> ExprEnv { ExprEnv { ansi: true, + legacy_statistical_aggregate: false, zone: Some(0), query_start_micros: 0, user: Arc::from("u"), diff --git a/crates/krishiv-exec/tests/aggregate_tpch.rs b/crates/krishiv-exec/tests/aggregate_tpch.rs index a8a594d2d..c869f90bb 100644 --- a/crates/krishiv-exec/tests/aggregate_tpch.rs +++ b/crates/krishiv-exec/tests/aggregate_tpch.rs @@ -206,6 +206,7 @@ fn spec(s: &SchemaRef, group_by: Vec, aggregates: Vec) -> ansi: true, spill: None, require_unit_weights: false, + legacy_statistical: false, } } diff --git a/crates/krishiv-exec/tests/planner/physical.rs b/crates/krishiv-exec/tests/planner/physical.rs index 0b9bba51c..0e49eec92 100644 --- a/crates/krishiv-exec/tests/planner/physical.rs +++ b/crates/krishiv-exec/tests/planner/physical.rs @@ -163,6 +163,7 @@ fn instantiation_checks_the_tables_still_match_the_plan() { fn env_data() -> ExprEnv { ExprEnv { ansi: true, + legacy_statistical_aggregate: false, zone: Some(0), query_start_micros: 0, user: Arc::from("u"), diff --git a/crates/krishiv-exec/tests/planner/write.rs b/crates/krishiv-exec/tests/planner/write.rs index 0fe986ad4..7c84b5091 100644 --- a/crates/krishiv-exec/tests/planner/write.rs +++ b/crates/krishiv-exec/tests/planner/write.rs @@ -123,6 +123,7 @@ fn query_results_are_written_committed_and_read_back() { ExplainNode::leaf("WriteFiles", ""), ExprEnv { ansi: true, + legacy_statistical_aggregate: false, zone: None, query_start_micros: 0, user: Arc::from("u"), diff --git a/crates/krishiv-functions/src/registry/defs.rs b/crates/krishiv-functions/src/registry/defs.rs index 2a7bb8078..b913c58ab 100644 --- a/crates/krishiv-functions/src/registry/defs.rs +++ b/crates/krishiv-functions/src/registry/defs.rs @@ -457,8 +457,15 @@ fn last(c: &CallTypes<'_>) -> Result { pick(c, "last") } +/// The moment aggregates (`CentralMomentAgg`, `Covariance`, `Corr`): every +/// argument is cast to `DOUBLE`; the result is a nullable `DOUBLE`. +fn moment(_: &CallTypes<'_>) -> Result { + Ok(ArgType::new(DataType::Double, true)) +} + fn aggregates() -> Vec { let numeric = || Inputs::Expects(vec![InputType::Numeric { intervals: true }]); + let doubles = || Inputs::Expects(vec![exact(DataType::Double)]); vec![ aggregate("count", Inputs::Any, count, Aggregate::Count), aggregate("sum", numeric(), sum, Aggregate::Sum), @@ -470,6 +477,18 @@ fn aggregates() -> Vec { aggregate("first_value", Inputs::Any, first, Aggregate::First), aggregate("last", Inputs::Any, last, Aggregate::Last), aggregate("last_value", Inputs::Any, last, Aggregate::Last), + aggregate("stddev", doubles(), moment, Aggregate::StddevSamp), + aggregate("std", doubles(), moment, Aggregate::StddevSamp), + aggregate("stddev_samp", doubles(), moment, Aggregate::StddevSamp), + aggregate("stddev_pop", doubles(), moment, Aggregate::StddevPop), + aggregate("variance", doubles(), moment, Aggregate::VarSamp), + aggregate("var_samp", doubles(), moment, Aggregate::VarSamp), + aggregate("var_pop", doubles(), moment, Aggregate::VarPop), + aggregate("skewness", doubles(), moment, Aggregate::Skewness), + aggregate("kurtosis", doubles(), moment, Aggregate::Kurtosis), + aggregate("corr", doubles(), moment, Aggregate::Corr), + aggregate("covar_samp", doubles(), moment, Aggregate::CovarSamp), + aggregate("covar_pop", doubles(), moment, Aggregate::CovarPop), ] } diff --git a/crates/krishiv-functions/src/registry/mod.rs b/crates/krishiv-functions/src/registry/mod.rs index c6af46e3e..6a3fdc11a 100644 --- a/crates/krishiv-functions/src/registry/mod.rs +++ b/crates/krishiv-functions/src/registry/mod.rs @@ -66,6 +66,24 @@ pub enum Aggregate { First, /// `last` / `last_value`; `ignoreNulls` is the second argument. Last, + /// `stddev` / `std` / `stddev_samp`. + StddevSamp, + /// `stddev_pop`. + StddevPop, + /// `variance` / `var_samp`. + VarSamp, + /// `var_pop`. + VarPop, + /// `skewness`. + Skewness, + /// `kurtosis`. + Kurtosis, + /// `corr`. + Corr, + /// `covar_samp`. + CovarSamp, + /// `covar_pop`. + CovarPop, } /// A built-in function. diff --git a/crates/krishiv-functions/src/registry/tests.rs b/crates/krishiv-functions/src/registry/tests.rs index 6388993b5..019e16a0e 100644 --- a/crates/krishiv-functions/src/registry/tests.rs +++ b/crates/krishiv-functions/src/registry/tests.rs @@ -1203,3 +1203,66 @@ fn regex_signatures() { let r = result("regexp_count", &[s, null_ty()], &[]).unwrap(); assert_eq!(r, ArgType::new(DataType::Integer, true)); } + +#[test] +fn moment_aggregates_are_registered_as_doubles() { + let r = builtins(); + for (name, agg, arguments) in [ + ("stddev", Aggregate::StddevSamp, 1), + ("std", Aggregate::StddevSamp, 1), + ("stddev_samp", Aggregate::StddevSamp, 1), + ("stddev_pop", Aggregate::StddevPop, 1), + ("variance", Aggregate::VarSamp, 1), + ("var_samp", Aggregate::VarSamp, 1), + ("var_pop", Aggregate::VarPop, 1), + ("skewness", Aggregate::Skewness, 1), + ("kurtosis", Aggregate::Kurtosis, 1), + ("corr", Aggregate::Corr, 2), + ("covar_samp", Aggregate::CovarSamp, 2), + ("covar_pop", Aggregate::CovarPop, 2), + ] { + let d = def(name); + assert_eq!(d.kind, FunctionKind::Aggregate, "{name}"); + assert_eq!(d.execution, Execution::Aggregate(agg), "{name}"); + assert!(d.deterministic, "{name}"); + assert!( + d.arity.accepts(arguments) && !d.arity.accepts(3 - arguments), + "{name}" + ); + // Found in every profile that has the name, and only there. + let since = r.spark_function(name).unwrap().since; + for p in SparkProfile::ALL { + assert_eq!( + matches!(r.lookup(name, p), Lookup::Found(_)), + p >= since, + "{name} {p:?}" + ); + } + // Every argument is cast to DOUBLE: an integer, a decimal and a + // string become one; the result is a nullable DOUBLE whatever the + // arguments (no rows, one row). + let args: Vec = (0..arguments).map(|_| ty(DataType::Integer)).collect(); + assert_eq!( + d.inputs.expected(&CallTypes::new(&args, &[])), + Some(vec![InputType::Exact(DataType::Double)]), + "{name}" + ); + assert_eq!( + result(name, &args, &[]).unwrap(), + ArgType::new(DataType::Double, true), + "{name}" + ); + } +} + +#[test] +fn moment_aggregates_reject_the_wrong_number_of_arguments() { + for (name, wrong) in [ + ("stddev", 2), + ("corr", 1), + ("covar_pop", 3), + ("kurtosis", 0), + ] { + assert!(!def(name).arity.accepts(wrong), "{name}({wrong})"); + } +} diff --git a/crates/krishiv-plan-codec/src/decode.rs b/crates/krishiv-plan-codec/src/decode.rs index c7469866c..e0113c7b5 100644 --- a/crates/krishiv-plan-codec/src/decode.rs +++ b/crates/krishiv-plan-codec/src/decode.rs @@ -288,6 +288,9 @@ fn plan(m: &pb::PhysicalPlan, limits: &Limits) -> Result { .ok_or_else(|| corrupt("no expression environment"))?; let env = ExprEnv { ansi: env.ansi, + // Not in the format: only the moment aggregates read it, and the + // encoder refuses those. + legacy_statistical_aggregate: false, zone: env.zone, query_start_micros: env.query_start_micros, user: Arc::from(env.user.as_str()), diff --git a/crates/krishiv-plan-codec/src/encode.rs b/crates/krishiv-plan-codec/src/encode.rs index f8d0cceb6..32d8dc601 100644 --- a/crates/krishiv-plan-codec/src/encode.rs +++ b/crates/krishiv-plan-codec/src/encode.rs @@ -313,7 +313,11 @@ impl Encoder { Ok(pb::AggregateNode { input: self.node_index(input)?, group_by: indexes(&plan.group_by), - aggregates: plan.aggregates.iter().map(aggregate_expr).collect(), + aggregates: plan + .aggregates + .iter() + .map(aggregate_expr) + .collect::>()?, distinct: plan.distinct.clone(), output: match plan.output { AggregateOutput::Final => pb::AggregateOutput::Final, @@ -611,7 +615,7 @@ fn window_function(f: &WindowFunctionPlan) -> Result { WindowFunction::PercentRank => Function::PercentRank(pb::Empty {}), WindowFunction::CumeDist => Function::CumeDist(pb::Empty {}), WindowFunction::Ntile(buckets) => Function::Ntile(*buckets), - WindowFunction::Aggregate(a) => Function::Aggregate(aggregate_expr(a)), + WindowFunction::Aggregate(a) => Function::Aggregate(aggregate_expr(a)?), }; Ok(pb::WindowFunctionPlan { function: Some(function), @@ -641,7 +645,7 @@ fn frame_bound(b: FrameBound) -> Result { } } -fn aggregate_expr(a: &AggregateExpr) -> pb::AggregateExpr { +fn aggregate_expr(a: &AggregateExpr) -> Result { let (function, ignore_nulls) = match a.function { AggregateFunction::Count => (pb::AggregateFunction::Count, false), AggregateFunction::Sum => (pb::AggregateFunction::Sum, false), @@ -650,8 +654,22 @@ fn aggregate_expr(a: &AggregateExpr) -> pb::AggregateExpr { AggregateFunction::Max => (pb::AggregateFunction::Max, false), AggregateFunction::First { ignore_nulls } => (pb::AggregateFunction::First, ignore_nulls), AggregateFunction::Last { ignore_nulls } => (pb::AggregateFunction::Last, ignore_nulls), + // The format has no code for the moment aggregates yet. + f @ (AggregateFunction::StddevSamp + | AggregateFunction::StddevPop + | AggregateFunction::VarSamp + | AggregateFunction::VarPop + | AggregateFunction::Skewness + | AggregateFunction::Kurtosis + | AggregateFunction::Corr + | AggregateFunction::CovarSamp + | AggregateFunction::CovarPop) => { + return Err(CodecError::Unsupported(format!( + "the aggregate function {f:?}" + ))); + } }; - pb::AggregateExpr { + Ok(pb::AggregateExpr { function: function as i32, ignore_nulls, args: indexes(&a.args), @@ -661,7 +679,7 @@ fn aggregate_expr(a: &AggregateExpr) -> pb::AggregateExpr { } as i32, arg_type: type_bytes(&a.arg_type), name: a.name.clone(), - } + }) } fn read_source(source: &ReadSource) -> Result { diff --git a/crates/krishiv-plan-codec/src/features.rs b/crates/krishiv-plan-codec/src/features.rs index af80ab004..135363234 100644 --- a/crates/krishiv-plan-codec/src/features.rs +++ b/crates/krishiv-plan-codec/src/features.rs @@ -217,7 +217,18 @@ pub fn feature_of(construct: Construct<'_>) -> Option { | AggregateFunction::Min | AggregateFunction::Max | AggregateFunction::First { .. } - | AggregateFunction::Last { .. } => None, + | AggregateFunction::Last { .. } + // The moment aggregates are not encodable yet: the encoder + // refuses them, so none reaches a worker. + | AggregateFunction::StddevSamp + | AggregateFunction::StddevPop + | AggregateFunction::VarSamp + | AggregateFunction::VarPop + | AggregateFunction::Skewness + | AggregateFunction::Kurtosis + | AggregateFunction::Corr + | AggregateFunction::CovarSamp + | AggregateFunction::CovarPop => None, }, } } diff --git a/crates/krishiv-plan-codec/tests/common/inventory.rs b/crates/krishiv-plan-codec/tests/common/inventory.rs index c42eb2ec9..fde31c577 100644 --- a/crates/krishiv-plan-codec/tests/common/inventory.rs +++ b/crates/krishiv-plan-codec/tests/common/inventory.rs @@ -33,6 +33,7 @@ use krishiv_types::{DataType, DecimalType}; pub(crate) fn env() -> ExprEnv { ExprEnv { ansi: true, + legacy_statistical_aggregate: false, zone: Some(-28_800), query_start_micros: 1_760_000_000_123_456, user: Arc::from("analyst"), diff --git a/crates/krishiv-plan-codec/tests/ratchet.rs b/crates/krishiv-plan-codec/tests/ratchet.rs index a8f2c437c..72df01f39 100644 --- a/crates/krishiv-plan-codec/tests/ratchet.rs +++ b/crates/krishiv-plan-codec/tests/ratchet.rs @@ -12,7 +12,7 @@ mod suites; use krishiv_exec::physical::{PhysicalPlan, WorkerContext, instantiate}; -use krishiv_plan_codec::{decode, encode}; +use krishiv_plan_codec::{CodecError, decode, encode}; use suites::{Suite, queries}; @@ -21,10 +21,17 @@ use suites::{Suite, queries}; /// fails when a query on it plans, so its entry is removed in the change /// that makes it plan. const TPCDS_NOT_PLANNED: &[&str] = &[ - "q10", "q16", "q17", "q30", "q32", "q35", "q39", "q41", "q45", "q50", "q51", "q62", "q70", - "q72", "q86", "q92", "q94", "q95", "q97", "q99", + "q10", "q16", "q30", "q32", "q35", "q41", "q45", "q50", "q51", "q62", "q70", "q72", "q86", + "q92", "q94", "q95", "q97", "q99", ]; +/// TPC-DS queries that plan but use a moment aggregate (`stddev_samp`), +/// which `krishiv.plan.v1` has no code for yet: the encoder refuses them +/// with `CodecError::Unsupported`, so none reaches a worker. The list may +/// only shrink: the test fails when one of them encodes, so its entry is +/// removed in the change that adds the aggregate to the format. +const TPCDS_NOT_ENCODABLE: &[&str] = &["q17", "q39"]; + fn assert_round_trips(name: &str, plan: &PhysicalPlan) -> Vec { let bytes = encode(plan).unwrap_or_else(|e| panic!("{name}: encode: {e}")); let back = decode(&bytes).unwrap_or_else(|e| panic!("{name}: decode: {e}")); @@ -61,14 +68,22 @@ fn tpcds_plans_round_trip() { let all = queries("tpcds"); assert_eq!(all.len(), 99); let mut not_planned = Vec::new(); + let mut not_encodable = Vec::new(); for (name, sql) in &all { match suite.plan(sql) { + Ok(plan) if TPCDS_NOT_ENCODABLE.contains(&name.as_str()) => match encode(&plan) { + Err(CodecError::Unsupported(what)) if what.contains("aggregate function") => { + not_encodable.push(name.as_str()); + } + other => panic!("tpcds {name}: expected a refused aggregate: {other:?}"), + }, Ok(plan) => { assert_round_trips(&format!("tpcds {name}"), &plan); } Err(_) => not_planned.push(name.as_str()), } } + assert_eq!(not_encodable, TPCDS_NOT_ENCODABLE); assert_eq!( not_planned, TPCDS_NOT_PLANNED, "the TPC-DS queries that do not plan changed: update TPCDS_NOT_PLANNED (it may only shrink)" diff --git a/crates/krishiv-session/tests/aggregate_moments.rs b/crates/krishiv-session/tests/aggregate_moments.rs new file mode 100644 index 000000000..11017972d --- /dev/null +++ b/crates/krishiv-session/tests/aggregate_moments.rs @@ -0,0 +1,407 @@ +// SPDX-License-Identifier: Apache-2.0 +//! The moment aggregates (`stddev`, `variance`, `skewness`, `kurtosis`, +//! `corr`, `covar_*`) end to end against Spark 4.2 `SQLQueryTestSuite` +//! goldens. Each claimed statement is found by its exact text in a vendored +//! result file and must give the golden's schema and rows; the two knobs +//! that change a result (`spark.sql.ansi.enabled` and +//! `spark.sql.legacy.statisticalAggregate`) are checked on top. +//! +//! The session has no `CREATE VIEW` or inline `VALUES` yet, so the driver +//! registers the golden's tables by hand. +//! +//! Spark computes these goldens over two partitions: each folds its rows +//! into a partial buffer and the buffers are merged. A session aggregates +//! one partition in one pass, which rounds the same values differently in +//! the last digits, so the double columns are compared to 1e-14 (relative) +//! and everything else exactly. The digits themselves are reproduced bit +//! for bit, under Spark's partitioning, by the accumulator tests of +//! `krishiv-exec` (`recorded_spark_results_are_reproduced_to_the_bit`). +//! +//! Statements of the three goldens that use these functions and are not +//! claimed (see `UNCLAIMED`) are listed with the reason, and the list is +//! checked against the files, so a new statement cannot go unnoticed. +#![forbid(unsafe_code)] + +use std::path::Path; +use std::sync::Arc; + +use arrow_array::cast::AsArray; +use arrow_array::types::{Float64Type, Int32Type, Int64Type}; +use arrow_array::{Array, ArrayRef, Int32Array, RecordBatch}; +use arrow_schema::{DataType, Field, Schema}; +use futures::executor::block_on; +use krishiv_common::error::{ErrorClass, KrishivError}; +use krishiv_session::Session; + +type R = Result>; + +/// One golden statement: the query and its expected schema and output. +struct Case { + query: String, + schema: String, + output: String, +} + +fn load(file: &str) -> R> { + let path = Path::new(env!("CARGO_MANIFEST_DIR")) + .join("../../tests/spark-golden/v4.2.0/results") + .join(file); + let text = std::fs::read_to_string(path)?; + let total = text.split("-- !query\n").skip(1).count(); + let cases: Vec = text + .split("-- !query\n") + .skip(1) + .filter_map(|block| { + let (query, rest) = block.split_once("\n-- !query schema\n")?; + let (schema, output) = rest.split_once("\n-- !query output\n")?; + Some(Case { + query: query.trim().to_owned(), + schema: schema.trim().to_owned(), + output: output.trim_end().to_owned(), + }) + }) + .collect(); + if cases.len() != total { + return Err(format!("dropped {} statements of {file}", total - cases.len()).into()); + } + Ok(cases) +} + +fn find<'a>(cases: &'a [Case], query: &str) -> R<&'a Case> { + cases + .iter() + .find(|c| c.query == query) + .ok_or_else(|| format!("the golden has no statement {query:?}").into()) +} + +fn spark_type(t: &DataType) -> &'static str { + match t { + DataType::Int32 => "int", + DataType::Int64 => "bigint", + DataType::Float64 => "double", + _ => "other", + } +} + +fn cell(c: &dyn Array, row: usize) -> String { + if c.is_null(row) { + return "NULL".into(); + } + match c.data_type() { + DataType::Int32 => c.as_primitive::().value(row).to_string(), + DataType::Int64 => c.as_primitive::().value(row).to_string(), + DataType::Float64 => format!("{:?}", c.as_primitive::().value(row)), + other => format!("<{other}>"), + } +} + +/// Spark's rendering of a result: `struct` and the tab +/// separated rows. +fn render(schema: &Schema, batches: &[RecordBatch]) -> (String, Vec) { + let fields = schema + .fields() + .iter() + .map(|f| format!("{}:{}", f.name(), spark_type(f.data_type()))) + .collect::>() + .join(","); + let rows = batches + .iter() + .flat_map(|b| { + (0..b.num_rows()).map(|r| { + b.columns() + .iter() + .map(|c| cell(c.as_ref(), r)) + .collect::>() + .join("\t") + }) + }) + .collect(); + (format!("struct<{fields}>"), rows) +} + +/// A table of nullable `int` columns. +fn table(s: &Session, name: &str, columns: &[&str], rows: &[&[Option]]) -> R { + let fields: Vec = columns + .iter() + .map(|c| Field::new(*c, DataType::Int32, true)) + .collect(); + let arrays: Vec = (0..columns.len()) + .map(|c| { + Arc::new(Int32Array::from( + rows.iter().map(|r| r[c]).collect::>(), + )) as ArrayRef + }) + .collect(); + let schema = Arc::new(Schema::new(fields)); + let batch = RecordBatch::try_new(Arc::clone(&schema), arrays)?; + s.register_batches(name, schema, vec![batch])?; + Ok(()) +} + +fn run(s: &Session, query: &str) -> Result<(Arc, Vec), KrishivError> { + block_on(async { + let df = s.sql(query).await?; + let schema = df.schema().await?; + Ok((schema, df.collect().await?)) + }) +} + +/// The golden's `testData` view of `group-by.sql`. +const TEST_DATA: [[Option; 2]; 9] = [ + [Some(1), Some(1)], + [Some(1), Some(2)], + [Some(2), Some(1)], + [Some(2), Some(2)], + [Some(3), Some(1)], + [Some(3), Some(2)], + [None, Some(1)], + [Some(3), None], + [None, None], +]; + +/// The golden's `testRegression` view of `linear-regression.sql`. +const REGRESSION: [[Option; 3]; 6] = [ + [Some(1), Some(10), None], + [Some(2), Some(10), Some(11)], + [Some(2), Some(20), Some(22)], + [Some(2), Some(25), None], + [Some(2), Some(30), Some(35)], + [Some(2), None, Some(40)], +]; + +fn session() -> R { + let s = Session::new()?; + let rows: Vec<&[Option]> = TEST_DATA.iter().map(|r| &r[..]).collect(); + table(&s, "testData", &["a", "b"], &rows)?; + let rows: Vec<&[Option]> = REGRESSION.iter().map(|r| &r[..]).collect(); + table(&s, "testRegression", &["k", "y", "x"], &rows)?; + table( + &s, + "t", + &["x", "y"], + &[ + &[Some(1), Some(1)], + &[Some(2), Some(2)], + &[Some(2), Some(2)], + ], + )?; + Ok(s) +} + +/// Whether two rows agree: cells equal as text, or doubles (written with a +/// fraction or an exponent) within 1e-14 of each other, relative. +fn agree(got: &str, want: &str) -> bool { + let (got, want): (Vec<_>, Vec<_>) = (got.split('\t').collect(), want.split('\t').collect()); + let real = |c: &str| c.contains(['.', 'E']); + got.len() == want.len() + && got.iter().zip(&want).all(|(g, w)| { + g == w + || (real(g) + && real(w) + && g.parse::().is_ok_and(|g| { + w.parse::() + .is_ok_and(|w| (g - w).abs() <= 1e-14 * g.abs().max(w.abs())) + })) + }) +} + +fn assert_rows(mut got: Vec, mut want: Vec, what: &str) { + got.sort(); + want.sort(); + assert_eq!(got.len(), want.len(), "{what}: {got:?} vs {want:?}"); + for (g, w) in got.iter().zip(&want) { + assert!(agree(g, w), "{what}: {g:?} vs {w:?}"); + } +} + +/// Runs the golden statement `query` of `file` and compares schema and rows. +fn replay(file: &str, query: &str) -> R { + let cases = load(file)?; + let case = find(&cases, query)?; + let (schema, batches) = run(&session()?, &case.query)?; + let (schema, rows) = render(&schema, &batches); + assert_eq!(schema, case.schema, "{query}"); + assert_rows( + rows, + case.output.lines().map(str::to_owned).collect(), + query, + ); + Ok(()) +} + +#[test] +fn group_by_golden_skewness_through_stddev_over_nulls() -> R { + replay( + "group-by.sql.out", + "SELECT SKEWNESS(a), KURTOSIS(a), MIN(a), MAX(a), AVG(a), VARIANCE(a), STDDEV(a), SUM(a), COUNT(a)\nFROM testData", + ) +} + +/// `corr(DISTINCT x, y)` of the golden's pairs; the golden statement also +/// asks for `corr(DISTINCT y, x)`, a second argument list that `DISTINCT` +/// aggregates do not support yet. +#[test] +fn distinct_pair_statistics_of_the_group_by_golden_rows() -> R { + let cases = load("group-by.sql.out")?; + let golden = find( + &cases, + "SELECT corr(DISTINCT x, y), corr(DISTINCT y, x), count(*)\n FROM (VALUES (1, 1), (2, 2), (2, 2)) t(x, y)", + )?; + assert_eq!(golden.output, "1.0\t1.0\t3"); + let (schema, batches) = run(&session()?, "SELECT corr(DISTINCT x, y), count(*) FROM t")?; + let (schema, rows) = render(&schema, &batches); + assert_eq!(schema, "struct"); + assert_eq!(rows, ["1.0\t3"]); + Ok(()) +} + +/// The `corr` columns of the `linear-regression` goldens, without their +/// `regr_r2` column (the `regr_*` family is not implemented): the golden's +/// own rows for `corr`. +#[test] +fn linear_regression_golden_corr_columns() -> R { + let cases = load("linear-regression.sql.out")?; + let golden = find( + &cases, + "SELECT k, corr(y, x), regr_r2(y, x) FROM testRegression GROUP BY k", + )?; + // `k corr regr_r2`: keep the first two cells. + let want: Vec = golden + .output + .lines() + .map(|l| l.split('\t').take(2).collect::>().join("\t")) + .collect(); + assert_eq!(want, ["1\tNULL", "2\t0.9988445981121533"]); + let (schema, batches) = run( + &session()?, + "SELECT k, corr(y, x) FROM testRegression GROUP BY k", + )?; + let (schema, rows) = render(&schema, &batches); + assert_eq!(schema, "struct"); + assert_rows(rows, want, "corr by k"); + Ok(()) +} + +#[test] +fn moments_over_the_golden_tables_follow_the_flags() -> R { + let s = session()?; + let one = |query: &str| -> R { + let (schema, batches) = run(&s, query)?; + Ok(render(&schema, &batches).1.join("|")) + }; + // One row per group in `testData` grouped by `a`, `b`: the sample + // statistics divide by zero. + let q = "SELECT stddev(a), var_samp(a), covar_samp(a, b), corr(a, b), skewness(a), kurtosis(a), \ + stddev_pop(a), var_pop(a), covar_pop(a, b) FROM testData WHERE a = 1 AND b = 1"; + assert_eq!(one(q)?, "NULL\tNULL\tNULL\tNULL\tNULL\tNULL\t0.0\t0.0\t0.0"); + s.conf() + .set("spark.sql.legacy.statisticalAggregate", "true")?; + assert_eq!(one(q)?, "NaN\tNaN\tNaN\tNaN\tNaN\tNaN\t0.0\t0.0\t0.0"); + // The setting does not change a group with spread. + let spread = "SELECT stddev(a), covar_pop(a, b) FROM testData"; + let with = one(spread)?; + s.conf() + .set("spark.sql.legacy.statisticalAggregate", "false")?; + assert_eq!(one(spread)?, with); + Ok(()) +} + +#[test] +fn corr_of_a_constant_column_follows_ansi() -> R { + let s = session()?; + // `a` is 3 in every row of the filtered input, `b` varies. + let q = "SELECT corr(a, b) FROM testData WHERE a = 3 AND b IS NOT NULL"; + for legacy in ["false", "true"] { + s.conf() + .set("spark.sql.legacy.statisticalAggregate", legacy)?; + s.conf().set("spark.sql.ansi.enabled", "false")?; + let (schema, batches) = run(&s, q)?; + assert_eq!(render(&schema, &batches).1, ["NULL"], "legacy={legacy}"); + s.conf().set("spark.sql.ansi.enabled", "true")?; + let e = run(&s, q).expect_err("DIVIDE_BY_ZERO under ANSI"); + assert_eq!(e.error_class(), ErrorClass::DivideByZero, "legacy={legacy}"); + } + Ok(()) +} + +/// The window form is not implemented: it is refused, never computed +/// as another function. +#[test] +fn moment_window_aggregates_are_refused() -> R { + let s = session()?; + let e = run(&s, "SELECT stddev(a) OVER (PARTITION BY b) FROM testData") + .expect_err("not implemented"); + assert_eq!(e.error_class(), ErrorClass::KrishivUnsupportedFeature); + Ok(()) +} + +/// Statements of the three goldens that use these functions and are not +/// claimed, each with the reason. +const UNCLAIMED: [(&str, &str, &str); 4] = [ + ( + "group-by.sql.out", + "SELECT corr(DISTINCT x, y), corr(DISTINCT y, x), count(*)", + "DISTINCT aggregates over different argument lists; checked above with one", + ), + ( + "window.sql.out", + "stddev(val) OVER w AS stddev,", + "window aggregates of the moment functions, with approx_count_distinct, collect_list and collect_set in the same statement", + ), + ( + "linear-regression.sql.out", + "SELECT k, corr(y, x), regr_r2(y, x) FROM testRegression GROUP BY k", + "regr_r2 (the regr_* family is out of scope); its corr column is checked above", + ), + ( + "linear-regression.sql.out", + "SELECT k, corr(y, x) FILTER (WHERE x IS NOT NULL), regr_r2(y, x) FROM testRegression GROUP BY k", + "regr_r2 (out of scope) and aggregate FILTER clauses", + ), +]; + +#[test] +fn every_statement_using_a_moment_function_is_claimed_or_listed() -> R { + let names = [ + "stddev(", + "stddev_samp(", + "stddev_pop(", + "std(", + "variance(", + "var_samp(", + "var_pop(", + "corr(", + "covar_samp(", + "covar_pop(", + "skewness(", + "kurtosis(", + ]; + let claimed = [("group-by.sql.out", "SELECT SKEWNESS(a)")]; + for file in [ + "group-by.sql.out", + "window.sql.out", + "linear-regression.sql.out", + ] { + for case in load(file)? { + let lower = case.query.to_ascii_lowercase(); + if !names.iter().any(|n| lower.contains(n)) { + continue; + } + let known = claimed + .iter() + .any(|(f, q)| *f == file && case.query.starts_with(q)) + || UNCLAIMED + .iter() + .any(|(f, q, _)| *f == file && case.query.contains(q)); + assert!(known, "{file}: unlisted statement {:?}", case.query); + } + } + // Every listed statement is still in its file. + for (file, q, _) in UNCLAIMED { + assert!( + load(file)?.iter().any(|c| c.query.contains(q)), + "{file}: {q}" + ); + } + Ok(()) +} diff --git a/crates/krishiv-testkit/tests/type_coercion.rs b/crates/krishiv-testkit/tests/type_coercion.rs index f9cecc84c..a349d27ad 100644 --- a/crates/krishiv-testkit/tests/type_coercion.rs +++ b/crates/krishiv-testkit/tests/type_coercion.rs @@ -161,8 +161,8 @@ fn failures(root: &Path) -> (BTreeMap<(String, String), String>, usize) { fn type_coercion_goldens_analyze_as_spark_does() { let root = krishiv_testkit::tests_dir(); let (failures, queries) = failures(&root); - // Of 2,858 query statements, 2,635 resolve as Spark resolves them. - assert_eq!((queries, queries - failures.len()), (2858, 2635)); + // Of 2,858 query statements, 2,641 resolve as Spark resolves them. + assert_eq!((queries, queries - failures.len()), (2858, 2641)); let path = root.join(MANIFEST); if std::env::var_os("KRISHIV_UPDATE_MANIFEST").is_some() { let entries: Vec = failures diff --git a/tests/fixtures/analysis/registry-matrix-manifest.json b/tests/fixtures/analysis/registry-matrix-manifest.json index dc142a826..943bed73d 100644 --- a/tests/fixtures/analysis/registry-matrix-manifest.json +++ b/tests/fixtures/analysis/registry-matrix-manifest.json @@ -1,5 +1,5 @@ { - "cases": 12512, + "cases": 12594, "entries": [ { "id": "c1a16bec609/ansi_off", @@ -1929,6 +1929,14 @@ "id": "conv#ec214d3068bf/ansi_on", "reason": "SELECT conv('100', 2, array(1, 2)) | KRISHIV_UNSUPPORTED_FEATURE: function array where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" }, + { + "id": "corr#d345280d5c4a/ansi_off", + "reason": "SELECT corr(c1, c2) FROM VALUES (3, 2), (3, 3), (6, 4) as tab(c1, c2) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "corr#d345280d5c4a/ansi_on", + "reason": "SELECT corr(c1, c2) FROM VALUES (3, 2), (3, 3), (6, 4) as tab(c1, c2) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, { "id": "cos#0aa1d83e940a/ansi_off", "reason": "SELECT cos(INTERVAL '1-2' YEAR TO MONTH) | parse: KRISHIV_UNSUPPORTED_FEATURE: INTERVAL literals where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" @@ -2097,6 +2105,22 @@ "id": "count#f7eefd2913d9/ansi_on", "reason": "SELECT count(*) FROM VALUES (NULL), (5), (5), (20) AS tab(col) WHERE false | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" }, + { + "id": "covar_pop#7ce2b5725450/ansi_off", + "reason": "SELECT covar_pop(c1, c2) FROM VALUES (1,1), (2,2), (3,3) AS tab(c1, c2) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "covar_pop#7ce2b5725450/ansi_on", + "reason": "SELECT covar_pop(c1, c2) FROM VALUES (1,1), (2,2), (3,3) AS tab(c1, c2) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "covar_samp#0d8013c6ae21/ansi_off", + "reason": "SELECT covar_samp(c1, c2) FROM VALUES (1,1), (2,2), (3,3) AS tab(c1, c2) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "covar_samp#0d8013c6ae21/ansi_on", + "reason": "SELECT covar_samp(c1, c2) FROM VALUES (1,1), (2,2), (3,3) AS tab(c1, c2) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, { "id": "crc32#2651371c29d2/ansi_off", "reason": "SELECT crc32(array(1, 2)) | KRISHIV_UNSUPPORTED_FEATURE: function array where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" @@ -5185,6 +5209,46 @@ "id": "isnull#ca1c1fc4fce7/ansi_on", "reason": "SELECT isnull(named_struct('a', 1, 'b', 'x')) | KRISHIV_UNSUPPORTED_FEATURE: function named_struct" }, + { + "id": "kurtosis#302fdce6faba/ansi_off", + "reason": "SELECT kurtosis(col) FROM VALUES (CAST(NULL AS int)), (CAST(NULL AS int)) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "kurtosis#302fdce6faba/ansi_on", + "reason": "SELECT kurtosis(col) FROM VALUES (CAST(NULL AS int)), (CAST(NULL AS int)) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "kurtosis#45aed161bbcd/ansi_off", + "reason": "SELECT kurtosis(col) FROM VALUES (-10), (-20), (100), (1000) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "kurtosis#45aed161bbcd/ansi_on", + "reason": "SELECT kurtosis(col) FROM VALUES (-10), (-20), (100), (1000) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "kurtosis#5142258ab80d/ansi_off", + "reason": "SELECT kurtosis(col) FROM VALUES (CAST(2147483647 AS INT)), (CAST(2147483647 AS INT)) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "kurtosis#5142258ab80d/ansi_on", + "reason": "SELECT kurtosis(col) FROM VALUES (CAST(2147483647 AS INT)), (CAST(2147483647 AS INT)) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "kurtosis#8d3a95e5b5b6/ansi_off", + "reason": "SELECT kurtosis(col) FROM VALUES (1), (10), (100), (10), (1) as tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "kurtosis#8d3a95e5b5b6/ansi_on", + "reason": "SELECT kurtosis(col) FROM VALUES (1), (10), (100), (10), (1) as tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "kurtosis#db80cb99845a/ansi_off", + "reason": "SELECT kurtosis(col) FROM VALUES (-10), (-20), (100), (1000) AS tab(col) WHERE false | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "kurtosis#db80cb99845a/ansi_on", + "reason": "SELECT kurtosis(col) FROM VALUES (-10), (-20), (100), (1000) AS tab(col) WHERE false | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, { "id": "last#03f68fa9443d/ansi_off", "reason": "SELECT last(col, CAST(NULL AS boolean)) FROM VALUES (10), (5), (NULL) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES where Spark has _LEGACY_ERROR_TEMP_1210" @@ -9641,6 +9705,46 @@ "id": "sinh#feb84311c281/ansi_on", "reason": "SELECT sinh(INTERVAL '1 02:03:04' DAY TO SECOND) | parse: KRISHIV_UNSUPPORTED_FEATURE: INTERVAL literals where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" }, + { + "id": "skewness#3703f2caff00/ansi_off", + "reason": "SELECT skewness(col) FROM VALUES (-10), (-20), (100), (1000) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "skewness#3703f2caff00/ansi_on", + "reason": "SELECT skewness(col) FROM VALUES (-10), (-20), (100), (1000) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "skewness#66c085d23cc2/ansi_off", + "reason": "SELECT skewness(col) FROM VALUES (-1000), (-100), (10), (20) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "skewness#66c085d23cc2/ansi_on", + "reason": "SELECT skewness(col) FROM VALUES (-1000), (-100), (10), (20) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "skewness#68781d23ea57/ansi_off", + "reason": "SELECT skewness(col) FROM VALUES (CAST(2147483647 AS INT)), (CAST(2147483647 AS INT)) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "skewness#68781d23ea57/ansi_on", + "reason": "SELECT skewness(col) FROM VALUES (CAST(2147483647 AS INT)), (CAST(2147483647 AS INT)) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "skewness#9984fabe879b/ansi_off", + "reason": "SELECT skewness(col) FROM VALUES (CAST(NULL AS int)), (CAST(NULL AS int)) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "skewness#9984fabe879b/ansi_on", + "reason": "SELECT skewness(col) FROM VALUES (CAST(NULL AS int)), (CAST(NULL AS int)) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "skewness#c6c2c9d249b1/ansi_off", + "reason": "SELECT skewness(col) FROM VALUES (-10), (-20), (100), (1000) AS tab(col) WHERE false | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "skewness#c6c2c9d249b1/ansi_on", + "reason": "SELECT skewness(col) FROM VALUES (-10), (-20), (100), (1000) AS tab(col) WHERE false | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, { "id": "sqrt#16036e88e7c1/ansi_off", "reason": "SELECT sqrt(INTERVAL '1-2' YEAR TO MONTH) | parse: KRISHIV_UNSUPPORTED_FEATURE: INTERVAL literals where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" @@ -9761,6 +9865,134 @@ "id": "startswith#bf1ab845092b/ansi_on", "reason": "SELECT startswith('Spark SQL', array(1, 2)) | KRISHIV_UNSUPPORTED_FEATURE: function array where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" }, + { + "id": "std#4432de39277c/ansi_off", + "reason": "SELECT std(col) FROM VALUES (CAST(2147483647 AS INT)), (CAST(2147483647 AS INT)) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "std#4432de39277c/ansi_on", + "reason": "SELECT std(col) FROM VALUES (CAST(2147483647 AS INT)), (CAST(2147483647 AS INT)) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "std#61ca896dc0f6/ansi_off", + "reason": "SELECT std(col) FROM VALUES (1), (2), (3) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "std#61ca896dc0f6/ansi_on", + "reason": "SELECT std(col) FROM VALUES (1), (2), (3) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "std#6b002a4cb79e/ansi_off", + "reason": "SELECT std(col) FROM VALUES (CAST(NULL AS int)), (CAST(NULL AS int)) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "std#6b002a4cb79e/ansi_on", + "reason": "SELECT std(col) FROM VALUES (CAST(NULL AS int)), (CAST(NULL AS int)) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "std#bc26cc485619/ansi_off", + "reason": "SELECT std(col) FROM VALUES (1), (2), (3) AS tab(col) WHERE false | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "std#bc26cc485619/ansi_on", + "reason": "SELECT std(col) FROM VALUES (1), (2), (3) AS tab(col) WHERE false | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "stddev#31f480f9a15c/ansi_off", + "reason": "SELECT stddev(col) FROM VALUES (CAST(NULL AS int)), (CAST(NULL AS int)) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "stddev#31f480f9a15c/ansi_on", + "reason": "SELECT stddev(col) FROM VALUES (CAST(NULL AS int)), (CAST(NULL AS int)) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "stddev#3e5793561511/ansi_off", + "reason": "SELECT stddev(col) FROM VALUES (CAST(2147483647 AS INT)), (CAST(2147483647 AS INT)) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "stddev#3e5793561511/ansi_on", + "reason": "SELECT stddev(col) FROM VALUES (CAST(2147483647 AS INT)), (CAST(2147483647 AS INT)) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "stddev#4ec3f9fff7c6/ansi_off", + "reason": "SELECT stddev(col) FROM VALUES (1), (2), (3) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "stddev#4ec3f9fff7c6/ansi_on", + "reason": "SELECT stddev(col) FROM VALUES (1), (2), (3) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "stddev#e98195c18179/ansi_off", + "reason": "SELECT stddev(col) FROM VALUES (1), (2), (3) AS tab(col) WHERE false | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "stddev#e98195c18179/ansi_on", + "reason": "SELECT stddev(col) FROM VALUES (1), (2), (3) AS tab(col) WHERE false | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "stddev_pop#53ab5668ef2d/ansi_off", + "reason": "SELECT stddev_pop(col) FROM VALUES (1), (2), (3) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "stddev_pop#53ab5668ef2d/ansi_on", + "reason": "SELECT stddev_pop(col) FROM VALUES (1), (2), (3) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "stddev_pop#877e1511c802/ansi_off", + "reason": "SELECT stddev_pop(col) FROM VALUES (CAST(NULL AS int)), (CAST(NULL AS int)) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "stddev_pop#877e1511c802/ansi_on", + "reason": "SELECT stddev_pop(col) FROM VALUES (CAST(NULL AS int)), (CAST(NULL AS int)) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "stddev_pop#f14e79c2136c/ansi_off", + "reason": "SELECT stddev_pop(col) FROM VALUES (CAST(2147483647 AS INT)), (CAST(2147483647 AS INT)) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "stddev_pop#f14e79c2136c/ansi_on", + "reason": "SELECT stddev_pop(col) FROM VALUES (CAST(2147483647 AS INT)), (CAST(2147483647 AS INT)) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "stddev_pop#fa5d44f3124b/ansi_off", + "reason": "SELECT stddev_pop(col) FROM VALUES (1), (2), (3) AS tab(col) WHERE false | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "stddev_pop#fa5d44f3124b/ansi_on", + "reason": "SELECT stddev_pop(col) FROM VALUES (1), (2), (3) AS tab(col) WHERE false | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "stddev_samp#0fe5c33e1041/ansi_off", + "reason": "SELECT stddev_samp(col) FROM VALUES (CAST(2147483647 AS INT)), (CAST(2147483647 AS INT)) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "stddev_samp#0fe5c33e1041/ansi_on", + "reason": "SELECT stddev_samp(col) FROM VALUES (CAST(2147483647 AS INT)), (CAST(2147483647 AS INT)) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "stddev_samp#2f720e7da5ac/ansi_off", + "reason": "SELECT stddev_samp(col) FROM VALUES (1), (2), (3) AS tab(col) WHERE false | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "stddev_samp#2f720e7da5ac/ansi_on", + "reason": "SELECT stddev_samp(col) FROM VALUES (1), (2), (3) AS tab(col) WHERE false | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "stddev_samp#5561eced0f8f/ansi_off", + "reason": "SELECT stddev_samp(col) FROM VALUES (CAST(NULL AS int)), (CAST(NULL AS int)) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "stddev_samp#5561eced0f8f/ansi_on", + "reason": "SELECT stddev_samp(col) FROM VALUES (CAST(NULL AS int)), (CAST(NULL AS int)) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "stddev_samp#bcc33eac2f58/ansi_off", + "reason": "SELECT stddev_samp(col) FROM VALUES (1), (2), (3) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "stddev_samp#bcc33eac2f58/ansi_on", + "reason": "SELECT stddev_samp(col) FROM VALUES (1), (2), (3) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, { "id": "substr#0423abe35937/ansi_off", "reason": "SELECT substr(named_struct('a', 1, 'b', 'x'), 5, 1) | KRISHIV_UNSUPPORTED_FEATURE: function named_struct where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" @@ -11897,6 +12129,102 @@ "id": "upper#dc8ae8109118/ansi_on", "reason": "SELECT upper(array(1, 2)) | KRISHIV_UNSUPPORTED_FEATURE: function array where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" }, + { + "id": "var_pop#0004df2b0e9c/ansi_off", + "reason": "SELECT var_pop(col) FROM VALUES (1), (2), (3) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "var_pop#0004df2b0e9c/ansi_on", + "reason": "SELECT var_pop(col) FROM VALUES (1), (2), (3) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "var_pop#38811a56a315/ansi_off", + "reason": "SELECT var_pop(col) FROM VALUES (CAST(2147483647 AS INT)), (CAST(2147483647 AS INT)) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "var_pop#38811a56a315/ansi_on", + "reason": "SELECT var_pop(col) FROM VALUES (CAST(2147483647 AS INT)), (CAST(2147483647 AS INT)) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "var_pop#d570e747d528/ansi_off", + "reason": "SELECT var_pop(col) FROM VALUES (CAST(NULL AS int)), (CAST(NULL AS int)) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "var_pop#d570e747d528/ansi_on", + "reason": "SELECT var_pop(col) FROM VALUES (CAST(NULL AS int)), (CAST(NULL AS int)) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "var_pop#e1825a716015/ansi_off", + "reason": "SELECT var_pop(col) FROM VALUES (1), (2), (3) AS tab(col) WHERE false | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "var_pop#e1825a716015/ansi_on", + "reason": "SELECT var_pop(col) FROM VALUES (1), (2), (3) AS tab(col) WHERE false | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "var_samp#2bce1ead63a4/ansi_off", + "reason": "SELECT var_samp(col) FROM VALUES (1), (2), (3) AS tab(col) WHERE false | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "var_samp#2bce1ead63a4/ansi_on", + "reason": "SELECT var_samp(col) FROM VALUES (1), (2), (3) AS tab(col) WHERE false | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "var_samp#57d96c2d29ae/ansi_off", + "reason": "SELECT var_samp(col) FROM VALUES (CAST(NULL AS int)), (CAST(NULL AS int)) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "var_samp#57d96c2d29ae/ansi_on", + "reason": "SELECT var_samp(col) FROM VALUES (CAST(NULL AS int)), (CAST(NULL AS int)) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "var_samp#de022352efb6/ansi_off", + "reason": "SELECT var_samp(col) FROM VALUES (CAST(2147483647 AS INT)), (CAST(2147483647 AS INT)) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "var_samp#de022352efb6/ansi_on", + "reason": "SELECT var_samp(col) FROM VALUES (CAST(2147483647 AS INT)), (CAST(2147483647 AS INT)) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "var_samp#f7e8c8549809/ansi_off", + "reason": "SELECT var_samp(col) FROM VALUES (1), (2), (3) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "var_samp#f7e8c8549809/ansi_on", + "reason": "SELECT var_samp(col) FROM VALUES (1), (2), (3) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "variance#0058ce8b433d/ansi_off", + "reason": "SELECT variance(col) FROM VALUES (1), (2), (3) AS tab(col) WHERE false | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "variance#0058ce8b433d/ansi_on", + "reason": "SELECT variance(col) FROM VALUES (1), (2), (3) AS tab(col) WHERE false | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "variance#03e02e1a0f84/ansi_off", + "reason": "SELECT variance(col) FROM VALUES (1), (2), (3) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "variance#03e02e1a0f84/ansi_on", + "reason": "SELECT variance(col) FROM VALUES (1), (2), (3) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "variance#39a600813155/ansi_off", + "reason": "SELECT variance(col) FROM VALUES (CAST(2147483647 AS INT)), (CAST(2147483647 AS INT)) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "variance#39a600813155/ansi_on", + "reason": "SELECT variance(col) FROM VALUES (CAST(2147483647 AS INT)), (CAST(2147483647 AS INT)) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "variance#e6cdd43a60e6/ansi_off", + "reason": "SELECT variance(col) FROM VALUES (CAST(NULL AS int)), (CAST(NULL AS int)) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, + { + "id": "variance#e6cdd43a60e6/ansi_on", + "reason": "SELECT variance(col) FROM VALUES (CAST(NULL AS int)), (CAST(NULL AS int)) AS tab(col) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" + }, { "id": "xxhash64#04d801e25e35/ansi_off", "reason": "SELECT xxhash64(CAST(NULL AS string), array(123), 2) | KRISHIV_UNSUPPORTED_FEATURE: function array" @@ -12712,6 +13040,11 @@ "function": "conv", "pass": 128 }, + { + "cases": 2, + "function": "corr", + "pass": 0 + }, { "cases": 50, "function": "cos", @@ -12732,6 +13065,16 @@ "function": "count", "pass": 0 }, + { + "cases": 2, + "function": "covar_pop", + "pass": 0 + }, + { + "cases": 2, + "function": "covar_samp", + "pass": 0 + }, { "cases": 54, "function": "crc32", @@ -12917,6 +13260,11 @@ "function": "isnull", "pass": 38 }, + { + "cases": 10, + "function": "kurtosis", + "pass": 0 + }, { "cases": 54, "function": "last", @@ -13227,6 +13575,11 @@ "function": "sinh", "pass": 40 }, + { + "cases": 10, + "function": "skewness", + "pass": 0 + }, { "cases": 50, "function": "sqrt", @@ -13237,6 +13590,26 @@ "function": "startswith", "pass": 106 }, + { + "cases": 8, + "function": "std", + "pass": 0 + }, + { + "cases": 8, + "function": "stddev", + "pass": 0 + }, + { + "cases": 8, + "function": "stddev_pop", + "pass": 0 + }, + { + "cases": 8, + "function": "stddev_samp", + "pass": 0 + }, { "cases": 266, "function": "substr", @@ -13332,6 +13705,21 @@ "function": "user", "pass": 2 }, + { + "cases": 8, + "function": "var_pop", + "pass": 0 + }, + { + "cases": 8, + "function": "var_samp", + "pass": 0 + }, + { + "cases": 8, + "function": "variance", + "pass": 0 + }, { "cases": 150, "function": "xxhash64", diff --git a/tests/fixtures/analysis/type-coercion-manifest.json b/tests/fixtures/analysis/type-coercion-manifest.json index da02e7887..1fc6c85ac 100644 --- a/tests/fixtures/analysis/type-coercion-manifest.json +++ b/tests/fixtures/analysis/type-coercion-manifest.json @@ -935,36 +935,6 @@ "query": "SELECT cast('2017-12-11 09:30:00.0' as timestamp) - '1' FROM t", "reason": "KRISHIV_UNSUPPORTED_FEATURE: date and timestamp arithmetic" }, - { - "file": "typeCoercion/native/promoteStrings", - "query": "SELECT kurtosis('1') FROM t", - "reason": "KRISHIV_UNSUPPORTED_FEATURE: function kurtosis" - }, - { - "file": "typeCoercion/native/promoteStrings", - "query": "SELECT skewness('1') FROM t", - "reason": "KRISHIV_UNSUPPORTED_FEATURE: function skewness" - }, - { - "file": "typeCoercion/native/promoteStrings", - "query": "SELECT stddev_pop('1') FROM t", - "reason": "KRISHIV_UNSUPPORTED_FEATURE: function stddev_pop" - }, - { - "file": "typeCoercion/native/promoteStrings", - "query": "SELECT stddev_samp('1') FROM t", - "reason": "KRISHIV_UNSUPPORTED_FEATURE: function stddev_samp" - }, - { - "file": "typeCoercion/native/promoteStrings", - "query": "SELECT var_pop('1') FROM t", - "reason": "KRISHIV_UNSUPPORTED_FEATURE: function var_pop" - }, - { - "file": "typeCoercion/native/promoteStrings", - "query": "SELECT var_samp('1') FROM t", - "reason": "KRISHIV_UNSUPPORTED_FEATURE: function var_samp" - }, { "file": "typeCoercion/native/stringCastAndExpressions", "query": "select cast(a as array) from t", From c879f28eb08dae06f30ca46de4ed1dcb0a79de10 Mon Sep 17 00:00:00 2001 From: techysneha14 <343708768+shehatathe@users.noreply.github.com> Date: Sun, 25 Oct 2026 02:14:05 +0100 Subject: [PATCH 2/5] Carry legacy_statistical_aggregate on the wire; share the aggregate name table The codec dropped the legacy_statistical_aggregate environment flag on decode. It is now field 5 of ExprEnv under a new Semantic feature, read and written by the codec, derived from the plan in plan_features, with an inventory case and golden bytes. Older bytes still decode the flag as false; readers without the feature refuse plans that set it. Moment aggregates remain unencodable, as pinned by the TPC-DS ratchet. The grouping and window planners shared no aggregate name table: mean planned in windows but was refused in grouping, and the window arity check was hardcoded. Both planners now use aggregate_function with the shared arity, and the window planner refuses moment aggregates through is_moment. Moments::store returned nothing for an out-of-range group while load returned INTERNAL_ERROR; store now reports the same error. The session golden sweep covers all ten result files with moment statements: the PostgreSQL aggregates_part1 plain, decimal-cast, pair and single-row moments and the promoteStrings literals replay exactly, and the udf, udaf and window statements stay listed with their reasons. Signed-off-by: techysneha14 <343708768+shehatathe@users.noreply.github.com> --- crates/krishiv-exec/src/aggregate/mod.rs | 7 + crates/krishiv-exec/src/aggregate/moments.rs | 12 +- crates/krishiv-exec/src/planner/aggregate.rs | 78 +++--- crates/krishiv-exec/src/planner/window.rs | 30 +-- crates/krishiv-plan-codec/src/decode.rs | 8 +- crates/krishiv-plan-codec/src/encode.rs | 1 + crates/krishiv-plan-codec/src/features.rs | 8 +- .../src/proto/krishiv.plan.v1.pb.rs | 5 + .../tests/common/inventory.rs | 19 +- .../one_row_legacy_statistical_aggregate.bin | Bin 0 -> 420 bytes crates/krishiv-plan-codec/tests/roundtrip.rs | 2 +- .../tests/aggregate_moments.rs | 231 +++++++++++++++++- proto/krishiv/plan/v1/plan.proto | 4 + 13 files changed, 335 insertions(+), 70 deletions(-) create mode 100644 crates/krishiv-plan-codec/tests/golden/legacy_statistical_aggregate/one_row_legacy_statistical_aggregate.bin diff --git a/crates/krishiv-exec/src/aggregate/mod.rs b/crates/krishiv-exec/src/aggregate/mod.rs index d0a080821..644c29137 100644 --- a/crates/krishiv-exec/src/aggregate/mod.rs +++ b/crates/krishiv-exec/src/aggregate/mod.rs @@ -89,6 +89,13 @@ pub use sink::{AggregateMetrics, HashAggregate}; use crate::error::{Result, internal}; +/// Whether `f` is a moment aggregate (`stddev`, `variance`, `skewness`, +/// `kurtosis`, `covar_*`, `corr`): it reads its arguments as doubles, and +/// the window planner refuses it. +pub(crate) fn is_moment(f: AggregateFunction) -> bool { + moments::buffer(f).is_some() +} + /// An aggregate function. #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] pub enum AggregateFunction { diff --git a/crates/krishiv-exec/src/aggregate/moments.rs b/crates/krishiv-exec/src/aggregate/moments.rs index b75bb1a7a..ecce68e72 100644 --- a/crates/krishiv-exec/src/aggregate/moments.rs +++ b/crates/krishiv-exec/src/aggregate/moments.rs @@ -340,12 +340,12 @@ impl Moments { Ok(row) } - fn store(&mut self, g: u32, row: &Row) { + fn store(&mut self, g: u32, row: &Row) -> Result<()> { for (col, v) in self.cols.iter_mut().zip(row) { - if let Some(slot) = col.get_mut(g as usize) { - *slot = *v; - } + *col.get_mut(g as usize) + .ok_or_else(|| internal(format!("aggregate group {g} out of range")))? = *v; } + Ok(()) } /// Folds rows into `groups` (one group id per row). `args` are the @@ -372,7 +372,7 @@ impl Moments { if valid(i) { let mut s = self.load(g)?; update(self.family, &mut s, x[i], y.map_or(0.0, |y| y[i])); - self.store(g, &s); + self.store(g, &s)?; } } Ok(()) @@ -393,7 +393,7 @@ impl Moments { } let mut s = self.load(g)?; merge(self.family, &mut s, &incoming); - self.store(g, &s); + self.store(g, &s)?; } Ok(()) } diff --git a/crates/krishiv-exec/src/planner/aggregate.rs b/crates/krishiv-exec/src/planner/aggregate.rs index 41a3cfd28..b4aa6e877 100644 --- a/crates/krishiv-exec/src/planner/aggregate.rs +++ b/crates/krishiv-exec/src/planner/aggregate.rs @@ -67,37 +67,11 @@ impl Planner<'_> { let mut aggregates = Vec::with_capacity(calls.len()); let mut distinct = Vec::with_capacity(calls.len()); for (c, at) in calls.iter().zip(&arg_at) { - let function = match c.f.id.name() { - "count" => AggregateFunction::Count, - "sum" => AggregateFunction::Sum, - "avg" => AggregateFunction::Avg, - "min" => AggregateFunction::Min, - "max" => AggregateFunction::Max, - "stddev" | "std" | "stddev_samp" => AggregateFunction::StddevSamp, - "stddev_pop" => AggregateFunction::StddevPop, - "variance" | "var_samp" => AggregateFunction::VarSamp, - "var_pop" => AggregateFunction::VarPop, - "skewness" => AggregateFunction::Skewness, - "kurtosis" => AggregateFunction::Kurtosis, - "corr" => AggregateFunction::Corr, - "covar_samp" => AggregateFunction::CovarSamp, - "covar_pop" => AggregateFunction::CovarPop, - // `first(x, ignoreNulls)`: the analyzer checked the flag is a - // boolean literal. - name if picks(name).is_some() => { - let ignore_nulls = matches!( - c.f.args.get(1).map(|a| &a.node), - Some(ExprNode::Literal(krishiv_plan::Literal::Boolean(true))) - ); - if picks(name) == Some(true) { - AggregateFunction::Last { ignore_nulls } - } else { - AggregateFunction::First { ignore_nulls } - } - } - other => { - return Err(unsupported(&format!("the aggregate function {other}"))); - } + let Some(function) = aggregate_function(c.f.id.name(), ignore_nulls_arg(&c.f)) else { + return Err(unsupported(&format!( + "the aggregate function {}", + c.f.id.name() + ))); }; if c.f.filter.is_some() { return Err(unsupported("aggregate FILTER (WHERE ...) clauses")); @@ -273,6 +247,48 @@ fn replace( }) } +/// The executor's function for aggregate name `name` (`first`/`last` +/// read `ignore_nulls`), or `None` for a name neither planner handles. +/// The grouping and window planners share this table so a name never +/// plans in one and is refused in the other. +pub(super) fn aggregate_function(name: &str, ignore_nulls: bool) -> Option { + Some(match name { + "count" => AggregateFunction::Count, + "sum" => AggregateFunction::Sum, + "avg" | "mean" => AggregateFunction::Avg, + "min" => AggregateFunction::Min, + "max" => AggregateFunction::Max, + "stddev" | "std" | "stddev_samp" => AggregateFunction::StddevSamp, + "stddev_pop" => AggregateFunction::StddevPop, + "variance" | "var_samp" => AggregateFunction::VarSamp, + "var_pop" => AggregateFunction::VarPop, + "skewness" => AggregateFunction::Skewness, + "kurtosis" => AggregateFunction::Kurtosis, + "corr" => AggregateFunction::Corr, + "covar_samp" => AggregateFunction::CovarSamp, + "covar_pop" => AggregateFunction::CovarPop, + // `first(x, ignoreNulls)`: the analyzer checked the flag is a + // boolean literal. + name if picks(name).is_some() => { + if picks(name) == Some(true) { + AggregateFunction::Last { ignore_nulls } + } else { + AggregateFunction::First { ignore_nulls } + } + } + _ => return None, + }) +} + +/// Whether `first`/`last` call `f` ignores NULLs: the analyzer checked its +/// second argument is a boolean literal when present. +pub(super) fn ignore_nulls_arg(f: &ResolvedFunction) -> bool { + matches!( + f.args.get(1).map(|a| &a.node), + Some(ExprNode::Literal(krishiv_plan::Literal::Boolean(true))) + ) +} + /// `Some(true)` for `last`/`last_value`, `Some(false)` for `first`/ /// `first_value`: the aggregates that pick a row's value. pub(super) fn picks(name: &str) -> Option { diff --git a/crates/krishiv-exec/src/planner/window.rs b/crates/krishiv-exec/src/planner/window.rs index 092c45dff..a994eeb1c 100644 --- a/crates/krishiv-exec/src/planner/window.rs +++ b/crates/krishiv-exec/src/planner/window.rs @@ -11,10 +11,10 @@ use krishiv_plan::{ SortDirection, SortOrder, WindowSpecRef, }; -use super::aggregate::{non_null_literal, picks}; +use super::aggregate::{aggregate_function, ignore_nulls_arg, non_null_literal, picks}; use super::build::{Planned, Planner, internal, list, orders, unsupported}; use super::properties::Distribution; -use crate::aggregate::{AggInput, AggregateExpr, AggregateFunction}; +use crate::aggregate::{AggInput, AggregateExpr, AggregateFunction, arity, is_moment}; use crate::error::Result; use crate::operator::Source; use crate::operators::SortKey; @@ -259,31 +259,21 @@ impl Planner<'_> { other => return Err(unsupported(&format!("the window function {other}"))), }); } - let function = match name { - "count" => AggregateFunction::Count, - "sum" => AggregateFunction::Sum, - "avg" | "mean" => AggregateFunction::Avg, - "min" => AggregateFunction::Min, - "max" => AggregateFunction::Max, - name if picks(name).is_some() => { - let ignore_nulls = matches!( - c.f.args.get(1).map(|a| &a.node), - Some(ExprNode::Literal(krishiv_plan::Literal::Boolean(true))) - ); - if picks(name) == Some(true) { - AggregateFunction::Last { ignore_nulls } - } else { - AggregateFunction::First { ignore_nulls } - } + let function = match aggregate_function(name, ignore_nulls_arg(c.f)) { + // The moment aggregates have no window form yet: refused, never + // computed as another function. + Some(function) if is_moment(function) => { + return Err(unsupported(&format!("the window aggregate {name}"))); } - other => return Err(unsupported(&format!("the window aggregate {other}"))), + Some(function) => function, + None => return Err(unsupported(&format!("the window aggregate {name}"))), }; // The analyzer rejects `count(DISTINCT x) OVER (...)` as Spark does; // a plan that still holds one must not be computed as `count(x)`. if c.f.distinct { return Err(unsupported("DISTINCT window aggregates")); } - if function != AggregateFunction::Count && args.len() != 1 { + if function != AggregateFunction::Count && args.len() != arity(function) { return Err(internal("a window aggregate of several arguments")); } // Spark's STRING and BINARY are views here whatever carrier the diff --git a/crates/krishiv-plan-codec/src/decode.rs b/crates/krishiv-plan-codec/src/decode.rs index e0113c7b5..2245a2854 100644 --- a/crates/krishiv-plan-codec/src/decode.rs +++ b/crates/krishiv-plan-codec/src/decode.rs @@ -288,9 +288,11 @@ fn plan(m: &pb::PhysicalPlan, limits: &Limits) -> Result { .ok_or_else(|| corrupt("no expression environment"))?; let env = ExprEnv { ansi: env.ansi, - // Not in the format: only the moment aggregates read it, and the - // encoder refuses those. - legacy_statistical_aggregate: false, + // The wire carries the flag since the `legacy_statistical_aggregate` + // feature: its absence in older bytes decodes as false, which is + // what every plan such an encoder could write carried (a plan that + // sets it needs a reader with the feature). + legacy_statistical_aggregate: env.legacy_statistical_aggregate, zone: env.zone, query_start_micros: env.query_start_micros, user: Arc::from(env.user.as_str()), diff --git a/crates/krishiv-plan-codec/src/encode.rs b/crates/krishiv-plan-codec/src/encode.rs index 32d8dc601..30495a4a7 100644 --- a/crates/krishiv-plan-codec/src/encode.rs +++ b/crates/krishiv-plan-codec/src/encode.rs @@ -115,6 +115,7 @@ fn env(env: &ExprEnv) -> pb::ExprEnv { zone: env.zone, query_start_micros: env.query_start_micros, user: env.user.to_string(), + legacy_statistical_aggregate: env.legacy_statistical_aggregate, } } diff --git a/crates/krishiv-plan-codec/src/features.rs b/crates/krishiv-plan-codec/src/features.rs index 135363234..5c39516e5 100644 --- a/crates/krishiv-plan-codec/src/features.rs +++ b/crates/krishiv-plan-codec/src/features.rs @@ -116,6 +116,7 @@ macro_rules! registry { registry! { (expand, Expand, 5, Semantic, "The `Expand` node and its `NullOf` project item (ROLLUP, CUBE, GROUPING SETS)."), (join_build_preserved, JoinBuildPreserved, POST_MIGRATION_READER_VERSION, Optimization, "Joins that build their preserved side (`RightSemi`, `RightAnti`, `RightOuter`); without it the planner builds the other side (`LeftSemi`, `LeftAnti`, `LeftOuter`)."), + (legacy_statistical_aggregate, LegacyStatisticalAggregate, POST_MIGRATION_READER_VERSION, Semantic, "The `legacy_statistical_aggregate` expression environment flag (`spark.sql.legacy.statisticalAggregate`): the sample moment aggregates give NaN instead of NULL where Spark divides by zero."), (limit_offset, LimitOffset, 3, Semantic, "A `Limit` that skips leading rows (a nonzero offset)."), (window, Window, 4, Semantic, "The `Window` node and its functions."), (window_group_limit, WindowGroupLimit, POST_MIGRATION_READER_VERSION, Optimization, "The `WindowGroupLimit` node (the rows of each partition within a rank)."), @@ -271,10 +272,15 @@ fn constructs_of<'a>(node: &'a PhysNode, f: &mut impl FnMut(Construct<'a>)) { } /// The features `plan` uses: a pure function of its content. Subquery plans -/// count; nothing recurses. +/// count; nothing recurses. The expression environment counts too: a plan +/// whose session sets `spark.sql.legacy.statisticalAggregate` needs a reader +/// that carries the flag, or it would compute NULL where Spark gives NaN. #[must_use] pub fn plan_features(plan: &PhysicalPlan) -> BTreeSet { let mut out = BTreeSet::new(); + if plan.env().legacy_statistical_aggregate { + out.insert(Feature::LegacyStatisticalAggregate); + } let mut stack = vec![plan.root()]; while let Some(node) = stack.pop() { constructs_of(node, &mut |c| out.extend(feature_of(c))); diff --git a/crates/krishiv-plan-codec/src/proto/krishiv.plan.v1.pb.rs b/crates/krishiv-plan-codec/src/proto/krishiv.plan.v1.pb.rs index 5962bfb41..0bda609c7 100644 --- a/crates/krishiv-plan-codec/src/proto/krishiv.plan.v1.pb.rs +++ b/crates/krishiv-plan-codec/src/proto/krishiv.plan.v1.pb.rs @@ -67,6 +67,11 @@ pub struct ExprEnv { pub query_start_micros: i64, #[prost(string, tag = "4")] pub user: ::prost::alloc::string::String, + /// feature: legacy_statistical_aggregate + /// `spark.sql.legacy.statisticalAggregate`: the sample moment aggregates + /// give NaN instead of NULL where Spark divides by zero. + #[prost(bool, tag = "5")] + pub legacy_statistical_aggregate: bool, } #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] pub struct ExplainNode { diff --git a/crates/krishiv-plan-codec/tests/common/inventory.rs b/crates/krishiv-plan-codec/tests/common/inventory.rs index fde31c577..f9c17f1e7 100644 --- a/crates/krishiv-plan-codec/tests/common/inventory.rs +++ b/crates/krishiv-plan-codec/tests/common/inventory.rs @@ -629,6 +629,13 @@ pub(crate) fn nodes() -> Vec<(&'static str, PhysNode)> { /// A plan of `root`, with a two-level explain tree and a warning. pub(crate) fn plan_of(root: PhysNode) -> PhysicalPlan { + plan_of_env(root, env()) +} + +/// [`plan_of`] with `env`: the golden for the +/// `legacy_statistical_aggregate` feature sets the flag the wire now +/// carries. +pub(crate) fn plan_of_env(root: PhysNode, env: ExprEnv) -> PhysicalPlan { let explain = ExplainNode::new( "Root", "details", @@ -637,8 +644,7 @@ pub(crate) fn plan_of(root: PhysNode) -> PhysicalPlan { ExplainNode::new("Right", "", vec![ExplainNode::leaf("Leaf", "x")]), ], ); - PhysicalPlan::new(root, schema(), explain, env()) - .with_warnings(vec!["a hint was ignored".into()]) + PhysicalPlan::new(root, schema(), explain, env).with_warnings(vec!["a hint was ignored".into()]) } /// Every kind of expression and literal, in one projection. @@ -761,5 +767,14 @@ pub(crate) fn cases() -> Vec<(String, PhysicalPlan)> { .map(|(name, node)| (name.to_owned(), plan_of(node))) .collect(); out.push(("expressions".to_owned(), plan_of(expressions()))); + // The expression environment flag the `legacy_statistical_aggregate` + // feature carries: the same root as `one_row`, under a session with + // `spark.sql.legacy.statisticalAggregate` set. + let mut legacy = env(); + legacy.legacy_statistical_aggregate = true; + out.push(( + "one_row_legacy_statistical_aggregate".to_owned(), + plan_of_env(PhysNode::OneRow, legacy), + )); out } diff --git a/crates/krishiv-plan-codec/tests/golden/legacy_statistical_aggregate/one_row_legacy_statistical_aggregate.bin b/crates/krishiv-plan-codec/tests/golden/legacy_statistical_aggregate/one_row_legacy_statistical_aggregate.bin new file mode 100644 index 0000000000000000000000000000000000000000..0c0ac12745bd9b8942be01d32cd85cdaca68cd9b GIT binary patch literal 420 zcmZWiu}Z{15S`u2C2JH7t3?P`ELLKf%2F)F>PivL;gZWOgamFE^g2uJEc^g}z)!Fc z{2c!w&TOs`XLxUBcIVAg(=0BehVxf>004>-Alfl^0N;dX!X<&{yI~k+RCymjq$X6R zm~=5EhYZXl2%hlvSC{Y&4n>#L&XZMqmo*h3gA)S4QE*nF&uJdLA?H7iDCDSfG>^K& ze>%gXp}`9a~&!@mQwmbZ(oqFN7$M(`erJUME%U z(YCITATK6Y8qL~v-aMc)w(0s!U0oMSWW(1ae|i~yK0n^$(*!$gx6Yr5+gV#TxZEzB TM{k|COKca|G#h&L^}YB7e-k)v literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/roundtrip.rs b/crates/krishiv-plan-codec/tests/roundtrip.rs index e48fcbd68..2a26176ad 100644 --- a/crates/krishiv-plan-codec/tests/roundtrip.rs +++ b/crates/krishiv-plan-codec/tests/roundtrip.rs @@ -293,7 +293,7 @@ fn readers_that_predate_features_still_refuse_what_they_cannot_read() { "window_ranking" | "window_aggregates" => 4, n if n.starts_with("expand") => 5, // A feature added after the freeze: above every frozen version. - "window_group_limit" => 6, + "window_group_limit" | "one_row_legacy_statistical_aggregate" => 6, n if n.starts_with("hash_join_right") => POST_MIGRATION_READER_VERSION, _ => MIN_READER_VERSION, }; diff --git a/crates/krishiv-session/tests/aggregate_moments.rs b/crates/krishiv-session/tests/aggregate_moments.rs index 11017972d..87ca0ae6a 100644 --- a/crates/krishiv-session/tests/aggregate_moments.rs +++ b/crates/krishiv-session/tests/aggregate_moments.rs @@ -26,8 +26,8 @@ use std::path::Path; use std::sync::Arc; use arrow_array::cast::AsArray; -use arrow_array::types::{Float64Type, Int32Type, Int64Type}; -use arrow_array::{Array, ArrayRef, Int32Array, RecordBatch}; +use arrow_array::types::{Float32Type, Float64Type, Int32Type, Int64Type}; +use arrow_array::{Array, ArrayRef, Float32Array, Int32Array, RecordBatch}; use arrow_schema::{DataType, Field, Schema}; use futures::executor::block_on; use krishiv_common::error::{ErrorClass, KrishivError}; @@ -78,6 +78,7 @@ fn spark_type(t: &DataType) -> &'static str { match t { DataType::Int32 => "int", DataType::Int64 => "bigint", + DataType::Float32 => "float", DataType::Float64 => "double", _ => "other", } @@ -90,6 +91,7 @@ fn cell(c: &dyn Array, row: usize) -> String { match c.data_type() { DataType::Int32 => c.as_primitive::().value(row).to_string(), DataType::Int64 => c.as_primitive::().value(row).to_string(), + DataType::Float32 => format!("{:?}", c.as_primitive::().value(row)), DataType::Float64 => format!("{:?}", c.as_primitive::().value(row)), other => format!("<{other}>"), } @@ -138,6 +140,33 @@ fn table(s: &Session, name: &str, columns: &[&str], rows: &[&[Option]]) -> Ok(()) } +/// The PostgreSQL `aggtest` table of the goldens (`a int`, `b float`). +fn aggtest(s: &Session) -> R { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Float32, true), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![ + Some(56), + Some(100), + Some(0), + Some(42), + ])) as ArrayRef, + Arc::new(Float32Array::from(vec![ + Some(7.8_f32), + Some(99.097_f32), + Some(0.09561_f32), + Some(324.78_f32), + ])) as ArrayRef, + ], + )?; + s.register_batches("aggtest", schema, vec![batch])?; + Ok(()) +} + fn run(s: &Session, query: &str) -> Result<(Arc, Vec), KrishivError> { block_on(async { let df = s.sql(query).await?; @@ -185,6 +214,7 @@ fn session() -> R { &[Some(2), Some(2)], ], )?; + aggtest(&s)?; Ok(s) } @@ -282,6 +312,82 @@ fn linear_regression_golden_corr_columns() -> R { Ok(()) } +/// The `aggtest` moments of `postgreSQL/aggregates_part1.sql.out`: the +/// plain, decimal-cast and pair statements, replayed exactly. The +/// decimal casts exercise the analyzer-cast-to-double path end to end. +#[test] +fn postgresql_aggregates_part1_moments() -> R { + for query in [ + "SELECT stddev_pop(b) FROM aggtest", + "SELECT stddev_samp(b) FROM aggtest", + "SELECT var_pop(b) FROM aggtest", + "SELECT var_samp(b) FROM aggtest", + "SELECT stddev_pop(CAST(b AS Decimal(38,0))) FROM aggtest", + "SELECT stddev_samp(CAST(b AS Decimal(38,0))) FROM aggtest", + "SELECT var_pop(CAST(b AS Decimal(38,0))) FROM aggtest", + "SELECT var_samp(CAST(b AS Decimal(38,0))) FROM aggtest", + "SELECT covar_pop(b, a), covar_samp(b, a) FROM aggtest", + "SELECT corr(b, a) FROM aggtest", + ] { + replay("postgreSQL/aggregates_part1.sql.out", query)?; + } + Ok(()) +} + +/// The single-row statements of `postgreSQL/aggregates_part1.sql.out`: a +/// population statistic over one row is its value, a sample statistic is +/// NULL. +#[test] +fn postgresql_aggregates_part1_single_row_moments() -> R { + for query in [ + "SELECT var_pop(1.0), var_samp(2.0)", + "SELECT stddev_pop(CAST(3.0 AS Decimal(38,0))), stddev_samp(CAST(4.0 AS Decimal(38,0)))", + ] { + replay("postgreSQL/aggregates_part1.sql.out", query)?; + } + Ok(()) +} + +/// The string-literal moments of `promoteStrings.sql.out`: `'1'` promotes +/// to double, so one row gives the population value and NULL for the +/// sample statistics, skewness and kurtosis. +#[test] +fn promote_strings_moment_literals() -> R { + let s = Session::new()?; + table(&s, "t", &["v"], &[&[Some(1)]])?; + let cases = load("typeCoercion/native/promoteStrings.sql.out")?; + for query in [ + "SELECT stddev_pop('1') FROM t", + "SELECT stddev_samp('1') FROM t", + "SELECT var_pop('1') FROM t", + "SELECT var_samp('1') FROM t", + "SELECT skewness('1') FROM t", + "SELECT kurtosis('1') FROM t", + ] { + let case = find(&cases, query)?; + let (schema, batches) = run(&s, &case.query)?; + let (schema, rows) = render(&schema, &batches); + assert_eq!(schema, case.schema, "{query}"); + assert_rows( + rows, + case.output.lines().map(str::to_owned).collect(), + query, + ); + } + Ok(()) +} + +/// `mean` is Spark's alias of `avg`: the grouping and window planners +/// share one name table, so it plans in both. +#[test] +fn mean_plans_as_avg_in_group_by() -> R { + let (schema, batches) = run(&session()?, "SELECT mean(a), avg(a) FROM testData")?; + let (schema, rows) = render(&schema, &batches); + assert_eq!(schema, "struct"); + assert_eq!(rows, ["2.142857142857143\t2.142857142857143"]); + Ok(()) +} + #[test] fn moments_over_the_golden_tables_follow_the_flags() -> R { let s = session()?; @@ -335,9 +441,11 @@ fn moment_window_aggregates_are_refused() -> R { Ok(()) } -/// Statements of the three goldens that use these functions and are not -/// claimed, each with the reason. -const UNCLAIMED: [(&str, &str, &str); 4] = [ +/// Statements of the goldens that use these functions and are not +/// claimed, each with the reason. The `udf`/`udaf` suites need SQL UDFs and +/// user-defined aggregates, which are not implemented; the window suites +/// need the window form of the moment aggregates, which is refused. +const UNCLAIMED: [(&str, &str, &str); 10] = [ ( "group-by.sql.out", "SELECT corr(DISTINCT x, y), corr(DISTINCT y, x), count(*)", @@ -358,6 +466,36 @@ const UNCLAIMED: [(&str, &str, &str); 4] = [ "SELECT k, corr(y, x) FILTER (WHERE x IS NOT NULL), regr_r2(y, x) FROM testRegression GROUP BY k", "regr_r2 (out of scope) and aggregate FILTER clauses", ), + ( + "postgreSQL/aggregates_part1.sql.out", + "SELECT avg(CAST(x AS DOUBLE)), var_pop(CAST(x AS DOUBLE))", + "inline VALUES, which the session runner does not support; infinite inputs are covered by the accumulator unit tests", + ), + ( + "postgreSQL/window_part4.sql.out", + "OVER (ORDER BY i ROWS BETWEEN CURRENT ROW AND UNBOUNDED FOLLOWING)", + "window form of the moment aggregates, which is refused", + ), + ( + "udaf/udaf-group-by.sql.out", + "udaf(", + "user-defined aggregate functions are not implemented; its moment columns repeat the group-by golden checked above", + ), + ( + "udf/postgreSQL/udf-aggregates_part1.sql.out", + "udf(", + "SQL UDFs are not implemented; the same statements without udf are checked above", + ), + ( + "udf/udf-group-by.sql.out", + "udf(", + "SQL UDFs are not implemented; the DISTINCT pair without udf is checked above", + ), + ( + "udf/udf-window.sql.out", + "udf(", + "window form of the moment aggregates over SQL UDFs, neither of which is implemented", + ), ]; #[test] @@ -376,11 +514,92 @@ fn every_statement_using_a_moment_function_is_claimed_or_listed() -> R { "skewness(", "kurtosis(", ]; - let claimed = [("group-by.sql.out", "SELECT SKEWNESS(a)")]; + let claimed = [ + ("group-by.sql.out", "SELECT SKEWNESS(a)"), + ( + "postgreSQL/aggregates_part1.sql.out", + "SELECT stddev_pop(b) FROM aggtest", + ), + ( + "postgreSQL/aggregates_part1.sql.out", + "SELECT stddev_samp(b) FROM aggtest", + ), + ( + "postgreSQL/aggregates_part1.sql.out", + "SELECT var_pop(b) FROM aggtest", + ), + ( + "postgreSQL/aggregates_part1.sql.out", + "SELECT var_samp(b) FROM aggtest", + ), + ( + "postgreSQL/aggregates_part1.sql.out", + "SELECT stddev_pop(CAST(b AS Decimal(38,0))) FROM aggtest", + ), + ( + "postgreSQL/aggregates_part1.sql.out", + "SELECT stddev_samp(CAST(b AS Decimal(38,0))) FROM aggtest", + ), + ( + "postgreSQL/aggregates_part1.sql.out", + "SELECT var_pop(CAST(b AS Decimal(38,0))) FROM aggtest", + ), + ( + "postgreSQL/aggregates_part1.sql.out", + "SELECT var_samp(CAST(b AS Decimal(38,0))) FROM aggtest", + ), + ( + "postgreSQL/aggregates_part1.sql.out", + "SELECT var_pop(1.0), var_samp(2.0)", + ), + ( + "postgreSQL/aggregates_part1.sql.out", + "SELECT stddev_pop(CAST(3.0 AS Decimal(38,0)))", + ), + ( + "postgreSQL/aggregates_part1.sql.out", + "SELECT covar_pop(b, a), covar_samp(b, a) FROM aggtest", + ), + ( + "postgreSQL/aggregates_part1.sql.out", + "SELECT corr(b, a) FROM aggtest", + ), + ( + "typeCoercion/native/promoteStrings.sql.out", + "SELECT stddev_pop('1') FROM t", + ), + ( + "typeCoercion/native/promoteStrings.sql.out", + "SELECT stddev_samp('1') FROM t", + ), + ( + "typeCoercion/native/promoteStrings.sql.out", + "SELECT var_pop('1') FROM t", + ), + ( + "typeCoercion/native/promoteStrings.sql.out", + "SELECT var_samp('1') FROM t", + ), + ( + "typeCoercion/native/promoteStrings.sql.out", + "SELECT skewness('1') FROM t", + ), + ( + "typeCoercion/native/promoteStrings.sql.out", + "SELECT kurtosis('1') FROM t", + ), + ]; for file in [ "group-by.sql.out", "window.sql.out", "linear-regression.sql.out", + "postgreSQL/aggregates_part1.sql.out", + "postgreSQL/window_part4.sql.out", + "udaf/udaf-group-by.sql.out", + "typeCoercion/native/promoteStrings.sql.out", + "udf/postgreSQL/udf-aggregates_part1.sql.out", + "udf/udf-group-by.sql.out", + "udf/udf-window.sql.out", ] { for case in load(file)? { let lower = case.query.to_ascii_lowercase(); diff --git a/proto/krishiv/plan/v1/plan.proto b/proto/krishiv/plan/v1/plan.proto index 9fa38e5bb..1bfef6cab 100644 --- a/proto/krishiv/plan/v1/plan.proto +++ b/proto/krishiv/plan/v1/plan.proto @@ -74,6 +74,10 @@ message ExprEnv { optional sint32 zone = 2; sint64 query_start_micros = 3; string user = 4; + // feature: legacy_statistical_aggregate + // `spark.sql.legacy.statisticalAggregate`: the sample moment aggregates + // give NaN instead of NULL where Spark divides by zero. + bool legacy_statistical_aggregate = 5; } message ExplainNode { From e2c10e44247a332e2a86ede097baf417bd499c81 Mon Sep 17 00:00:00 2001 From: techysneha14 <343708768+shehatathe@users.noreply.github.com> Date: Mon, 26 Oct 2026 03:20:53 +0100 Subject: [PATCH 3/5] test(session): moment aggregates state their Spark oracle and cover the native legacy key The flag tests say which Spark 4.2.0 queries produced the expected values, and set the native key spelling of the legacy statistical-aggregate flag so the config registry coverage check finds a test for it. Signed-off-by: techysneha14 <343708768+shehatathe@users.noreply.github.com> --- crates/krishiv-session/tests/aggregate_moments.rs | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/crates/krishiv-session/tests/aggregate_moments.rs b/crates/krishiv-session/tests/aggregate_moments.rs index 87ca0ae6a..ccbbe54cf 100644 --- a/crates/krishiv-session/tests/aggregate_moments.rs +++ b/crates/krishiv-session/tests/aggregate_moments.rs @@ -390,6 +390,10 @@ fn mean_plans_as_avg_in_group_by() -> R { #[test] fn moments_over_the_golden_tables_follow_the_flags() -> R { + // Oracle: these exact queries (over the same `testData` rows) were run on + // a Spark 4.2.0 server with `spark.sql.legacy.statisticalAggregate` unset + // and set to true; the expected rows below are what it returned (NULLs, + // then NaNs, with the population statistics 0.0 in both). let s = session()?; let one = |query: &str| -> R { let (schema, batches) = run(&s, query)?; @@ -409,11 +413,21 @@ fn moments_over_the_golden_tables_follow_the_flags() -> R { s.conf() .set("spark.sql.legacy.statisticalAggregate", "false")?; assert_eq!(one(spread)?, with); + // The native key is the same setting as the Spark alias. + s.conf() + .set("krishiv.sql.legacy.statisticalAggregate", "true")?; + assert_eq!(one(q)?, "NaN\tNaN\tNaN\tNaN\tNaN\tNaN\t0.0\t0.0\t0.0"); + s.conf() + .set("krishiv.sql.legacy.statisticalAggregate", "false")?; + assert_eq!(one(q)?, "NULL\tNULL\tNULL\tNULL\tNULL\tNULL\t0.0\t0.0\t0.0"); Ok(()) } #[test] fn corr_of_a_constant_column_follows_ansi() -> R { + // Oracle: on a Spark 4.2.0 server this query returns NULL with + // `spark.sql.ansi.enabled=false` and fails with DIVIDE_BY_ZERO with it + // true, for `spark.sql.legacy.statisticalAggregate` false and true alike. let s = session()?; // `a` is 3 in every row of the filtered input, `b` varies. let q = "SELECT corr(a, b) FROM testData WHERE a = 3 AND b IS NOT NULL"; From dbc322e410741dd00c020ccaf6d70ef13ef38502 Mon Sep 17 00:00:00 2001 From: techysneha14 <343708768+shehatathe@users.noreply.github.com> Date: Mon, 26 Oct 2026 05:21:08 +0100 Subject: [PATCH 4/5] python: export std, PySpark's alias of stddev_samp pyspark.sql.functions.std(col) exists in Spark 4.2; the signature table listed it as not exported. Signed-off-by: techysneha14 <343708768+shehatathe@users.noreply.github.com> --- crates/krishiv-functions/src/python_wrappers/table.rs | 2 +- python/krishiv/sql/_generated_functions.py | 6 ++++++ python/krishiv/sql/_generated_functions.pyi | 2 ++ python/tests/_pyspark_signatures.py | 1 + 4 files changed, 10 insertions(+), 1 deletion(-) diff --git a/crates/krishiv-functions/src/python_wrappers/table.rs b/crates/krishiv-functions/src/python_wrappers/table.rs index 9bcf1802c..2d4317faf 100644 --- a/crates/krishiv-functions/src/python_wrappers/table.rs +++ b/crates/krishiv-functions/src/python_wrappers/table.rs @@ -30,7 +30,6 @@ pub(super) const NOT_EXPORTED: &[(&str, &str)] = &[ ("min_by", "no PySpark function"), ("some", "no PySpark function"), ("space", "no PySpark function"), - ("std", "no PySpark function; use `stddev`"), ( "mean", "an alias of `avg`; the exported `mean` is a row that sends `avg`", @@ -144,6 +143,7 @@ pub(super) fn rows() -> Vec { "skewness", "soundex", "sqrt", + "std", "stddev", "stddev_pop", "stddev_samp", diff --git a/python/krishiv/sql/_generated_functions.py b/python/krishiv/sql/_generated_functions.py index 63f56d1ef..9b29680ee 100644 --- a/python/krishiv/sql/_generated_functions.py +++ b/python/krishiv/sql/_generated_functions.py @@ -202,6 +202,7 @@ "split_part", "sqrt", "startswith", + "std", "stddev", "stddev_pop", "stddev_samp", @@ -1157,6 +1158,11 @@ def startswith(str: ColumnOrName, prefix: ColumnOrName) -> Column: return _call("startswith", _named(str, "str"), _named(prefix, "prefix")) +def std(col: ColumnOrName) -> Column: + """Calls the registry function ``std``.""" + return _call("std", _named(col, "col")) + + def stddev(col: ColumnOrName) -> Column: """Calls the registry function ``stddev``.""" return _call("stddev", _named(col, "col")) diff --git a/python/krishiv/sql/_generated_functions.pyi b/python/krishiv/sql/_generated_functions.pyi index 58cc8e45a..a3d5f33ef 100644 --- a/python/krishiv/sql/_generated_functions.pyi +++ b/python/krishiv/sql/_generated_functions.pyi @@ -183,6 +183,7 @@ __all__ = [ "split_part", "sqrt", "startswith", + "std", "stddev", "stddev_pop", "stddev_samp", @@ -415,6 +416,7 @@ def soundex(col: ColumnOrName) -> Column: ... def split_part(src: ColumnOrName, delimiter: ColumnOrName, partNum: ColumnOrName) -> Column: ... def sqrt(col: ColumnOrName) -> Column: ... def startswith(str: ColumnOrName, prefix: ColumnOrName) -> Column: ... +def std(col: ColumnOrName) -> Column: ... def stddev(col: ColumnOrName) -> Column: ... def stddev_pop(col: ColumnOrName) -> Column: ... def stddev_samp(col: ColumnOrName) -> Column: ... diff --git a/python/tests/_pyspark_signatures.py b/python/tests/_pyspark_signatures.py index 680021b47..2f85fccce 100644 --- a/python/tests/_pyspark_signatures.py +++ b/python/tests/_pyspark_signatures.py @@ -198,6 +198,7 @@ "split_part": "src, delimiter, partNum", "sqrt": "col", "startswith": "str, prefix", + "std": "col", "stddev": "col", "stddev_pop": "col", "stddev_samp": "col", From 9182ff5e876a4fab895a059ea21bad282c890f29 Mon Sep 17 00:00:00 2001 From: techysneha14 <343708768+shehatathe@users.noreply.github.com> Date: Mon, 26 Oct 2026 09:21:55 +0100 Subject: [PATCH 5/5] plan-codec: gate legacy_statistical_aggregate on a moment aggregate plan_features listed the feature for every plan from a session with spark.sql.legacy.statisticalAggregate set, even with no aggregate in the plan. The encoder refuses every moment aggregate, so the flag could change no encodable plan; sessions with the flag set produced plans that readers without the feature refused for every query. List the feature only when the plan both sets the flag and holds a moment aggregate (aggregate node, window function or subquery plan). The decoder compares the header against the same plan_features, so both sides agree. The flag-only inventory case now pins that such plans carry no feature, round-trip, and decode on readers without it; its golden moved to base/. The two registry-wide invariants exempt the dormant feature until a moment aggregate is encodable. Signed-off-by: techysneha14 <343708768+shehatathe@users.noreply.github.com> --- crates/krishiv-exec/src/aggregate/mod.rs | 2 +- crates/krishiv-plan-codec/src/decode.rs | 7 +- crates/krishiv-plan-codec/src/features.rs | 25 ++-- .../tests/common/inventory.rs | 6 +- crates/krishiv-plan-codec/tests/features.rs | 128 +++++++++++++++++- crates/krishiv-plan-codec/tests/golden.rs | 10 +- .../one_row_legacy_statistical_aggregate.bin | Bin 420 -> 390 bytes crates/krishiv-plan-codec/tests/roundtrip.rs | 2 +- 8 files changed, 162 insertions(+), 18 deletions(-) rename crates/krishiv-plan-codec/tests/golden/{legacy_statistical_aggregate => base}/one_row_legacy_statistical_aggregate.bin (76%) diff --git a/crates/krishiv-exec/src/aggregate/mod.rs b/crates/krishiv-exec/src/aggregate/mod.rs index 4b47cbdc7..db4f9a937 100644 --- a/crates/krishiv-exec/src/aggregate/mod.rs +++ b/crates/krishiv-exec/src/aggregate/mod.rs @@ -95,7 +95,7 @@ use crate::error::{Result, internal}; /// Whether `f` is a moment aggregate (`stddev`, `variance`, `skewness`, /// `kurtosis`, `covar_*`, `corr`): it reads its arguments as doubles, and /// the window planner refuses it. -pub(crate) fn is_moment(f: AggregateFunction) -> bool { +pub fn is_moment(f: AggregateFunction) -> bool { moments::buffer(f).is_some() } diff --git a/crates/krishiv-plan-codec/src/decode.rs b/crates/krishiv-plan-codec/src/decode.rs index e5083c8ef..b6fc323d0 100644 --- a/crates/krishiv-plan-codec/src/decode.rs +++ b/crates/krishiv-plan-codec/src/decode.rs @@ -288,10 +288,11 @@ fn plan(m: &pb::PhysicalPlan, limits: &Limits) -> Result { .ok_or_else(|| corrupt("no expression environment"))?; let env = ExprEnv { ansi: env.ansi, - // The wire carries the flag since the `legacy_statistical_aggregate` + // The wire carries the flag for the `legacy_statistical_aggregate` // feature: its absence in older bytes decodes as false, which is - // what every plan such an encoder could write carried (a plan that - // sets it needs a reader with the feature). + // what every plan such an encoder could write carried (only a plan + // that sets it around a moment aggregate needs a reader with the + // feature). legacy_statistical_aggregate: env.legacy_statistical_aggregate, zone: env.zone, query_start_micros: env.query_start_micros, diff --git a/crates/krishiv-plan-codec/src/features.rs b/crates/krishiv-plan-codec/src/features.rs index 640d68bc9..68db74a84 100644 --- a/crates/krishiv-plan-codec/src/features.rs +++ b/crates/krishiv-plan-codec/src/features.rs @@ -39,7 +39,7 @@ use std::fmt; use std::sync::Arc; use krishiv_exec::JoinType; -use krishiv_exec::aggregate::AggregateFunction; +use krishiv_exec::aggregate::{AggregateFunction, is_moment}; use krishiv_exec::physical::{Partitioning, PhysNode, PhysicalPlan, ProjectItemPlan}; use krishiv_exec::window::{FrameBound, WindowFrame, WindowFunction}; use krishiv_plan::{Expr, ExprNode, Literal, ResolvedExpr}; @@ -441,23 +441,32 @@ fn types_of<'a>(node: &'a PhysNode) -> Vec<&'a DataType> { } /// The features `plan` uses: a pure function of its content. Subquery plans -/// count; nothing recurses. The expression environment counts too: a plan -/// whose session sets `spark.sql.legacy.statisticalAggregate` needs a reader -/// that carries the flag, or it would compute NULL where Spark gives NaN. +/// count; nothing recurses. The expression environment counts too, but only +/// around a moment aggregate: a plan whose session sets +/// `spark.sql.legacy.statisticalAggregate` needs a reader that carries the +/// flag, or it would compute NULL where Spark gives NaN. A plan without one +/// runs the same everywhere, so the flag alone lists no feature and readers +/// without it keep reading the plan. #[must_use] pub fn plan_features(plan: &PhysicalPlan) -> BTreeSet { let mut out = BTreeSet::new(); - if plan.env().legacy_statistical_aggregate { - out.insert(Feature::LegacyStatisticalAggregate); - } + let mut has_moment = false; let mut stack = vec![plan.root()]; while let Some(node) = stack.pop() { - constructs_of(node, &mut |c| out.extend(feature_of(c))); + constructs_of(node, &mut |c| { + if matches!(c, Construct::Aggregate(function) if is_moment(function)) { + has_moment = true; + } + out.extend(feature_of(c)); + }); stack.extend(node.children()); } for field in plan.schema().fields() { out.extend(feature_of(Construct::ArrowType(field.data_type()))); } + if has_moment && plan.env().legacy_statistical_aggregate { + out.insert(Feature::LegacyStatisticalAggregate); + } out } diff --git a/crates/krishiv-plan-codec/tests/common/inventory.rs b/crates/krishiv-plan-codec/tests/common/inventory.rs index 8c8d28a84..f63a2bc28 100644 --- a/crates/krishiv-plan-codec/tests/common/inventory.rs +++ b/crates/krishiv-plan-codec/tests/common/inventory.rs @@ -949,8 +949,10 @@ pub(crate) fn cases() -> Vec<(String, PhysicalPlan)> { .collect(); out.push(("expressions".to_owned(), plan_of(expressions()))); // The expression environment flag the `legacy_statistical_aggregate` - // feature carries: the same root as `one_row`, under a session with - // `spark.sql.legacy.statisticalAggregate` set. + // feature is about: the same root as `one_row`, under a session with + // `spark.sql.legacy.statisticalAggregate` set. With no moment aggregate + // in the plan it lists no feature, so readers without it keep reading + // the plan; the flag itself still travels on the wire. let mut legacy = env(); legacy.legacy_statistical_aggregate = true; out.push(( diff --git a/crates/krishiv-plan-codec/tests/features.rs b/crates/krishiv-plan-codec/tests/features.rs index 18b832dbc..b469d734b 100644 --- a/crates/krishiv-plan-codec/tests/features.rs +++ b/crates/krishiv-plan-codec/tests/features.rs @@ -15,7 +15,7 @@ use krishiv_plan_codec::{ }; use prost::Message as _; -use inventory::{cases, plan_of}; +use inventory::{cases, plan_of, plan_of_env}; fn message(name: &str) -> pb::PhysicalPlan { let (_, plan) = cases().into_iter().find(|(n, _)| n == name).unwrap(); @@ -49,8 +49,15 @@ fn the_header_lists_exactly_the_features_in_use() { limit: 1, }); assert!(plan_features(&zero).is_empty()); - // Every registered feature occurs in some inventory case. + // Every registered feature occurs in some inventory case, except + // `legacy_statistical_aggregate`: it fires only around a moment + // aggregate, and the encoder refuses every moment aggregate (pinned by + // the TPC-DS ratchet), so no encodable case can use it. The test below + // pins its `plan_features` rule on plans that are never encoded. for f in Feature::ALL { + if *f == Feature::LegacyStatisticalAggregate { + continue; + } assert!( cases().iter().any(|(_, p)| plan_features(p).contains(f)), "{f}: no inventory case uses it" @@ -80,6 +87,123 @@ fn a_feature_in_a_subquery_plan_counts() { assert_eq!(decode(&encode(&plan).unwrap()).unwrap(), plan); } +/// `legacy_statistical_aggregate` fires only for a plan that sets the flag +/// around a moment aggregate. The flag alone changes nothing the encoder +/// can write (it refuses every moment aggregate), so such plans list no +/// feature: a reader without it keeps reading them, and the flag itself +/// still round-trips. +#[test] +fn legacy_statistical_aggregate_needs_a_moment_aggregate() { + use krishiv_exec::aggregate::{AggregateExpr, AggregateFunction, AggregateOutput}; + use krishiv_exec::physical::{AggregatePlan, SubqueryPlan}; + use krishiv_exec::window::{ + FrameUnits, WindowFrame, WindowFunction, WindowFunctionPlan, WindowPlan, + }; + + let moment = || { + AggregateExpr::new( + AggregateFunction::StddevSamp, + 0, + arrow_schema::DataType::Float64, + "s", + ) + }; + let aggregate_node = || PhysNode::Aggregate { + input: Box::new(inventory::scan()), + aggregate: AggregatePlan { + group_by: Vec::new(), + aggregates: vec![moment()], + distinct: vec![false], + output: AggregateOutput::Final, + }, + }; + let legacy = || { + let mut env = inventory::env(); + env.legacy_statistical_aggregate = true; + env + }; + + // The flag alone lists no feature. + let (_, flag_only) = cases() + .into_iter() + .find(|(n, _)| n == "one_row_legacy_statistical_aggregate") + .unwrap(); + assert!(plan_features(&flag_only).is_empty()); + assert!(to_message(&flag_only).unwrap().features.is_empty()); + // A reader without the feature still reads such a plan, flag included. + let bytes = encode(&flag_only).unwrap(); + let without = ReaderFeatures::only([]); + assert_eq!( + decode_with(&bytes, &Limits::default(), &without).unwrap(), + flag_only + ); + // A moment aggregate without the flag lists none either. + assert!(plan_features(&plan_of(aggregate_node())).is_empty()); + // Both together list the feature: in an aggregate node, ... + let with_both = plan_of_env(aggregate_node(), legacy()); + assert_eq!( + plan_features(&with_both).into_iter().collect::>(), + [Feature::LegacyStatisticalAggregate] + ); + // ... in a window function, ... + let window = plan_of_env( + PhysNode::Window { + input: Box::new(inventory::scan()), + window: WindowPlan { + partition_by: Vec::new(), + order_by: Vec::new(), + functions: vec![WindowFunctionPlan { + function: WindowFunction::Aggregate(moment()), + frame: WindowFrame::whole(FrameUnits::Rows), + filter: None, + name: "w".to_owned(), + }], + }, + }, + legacy(), + ); + assert_eq!( + plan_features(&window).into_iter().collect::>(), + [Feature::LegacyStatisticalAggregate, Feature::Window] + ); + // ... and in a subquery plan ... + let subquery = plan_of_env( + PhysNode::Filter { + input: Box::new(PhysNode::OneRow), + layout: Vec::new(), + predicate: inventory::predicate(), + subqueries: vec![SubqueryPlan { + plan: aggregate_node(), + ..inventory::subquery() + }], + }, + legacy(), + ); + assert_eq!( + plan_features(&subquery).into_iter().collect::>(), + [Feature::LegacyStatisticalAggregate] + ); + // ... but never for another aggregate under the flag. + let plain = plan_of_env( + PhysNode::Aggregate { + input: Box::new(inventory::scan()), + aggregate: AggregatePlan { + group_by: Vec::new(), + aggregates: vec![AggregateExpr::new( + AggregateFunction::Sum, + 0, + arrow_schema::DataType::Float64, + "s", + )], + distinct: vec![false], + output: AggregateOutput::Final, + }, + }, + legacy(), + ); + assert!(plan_features(&plain).is_empty()); +} + #[test] fn constructs_map_to_one_feature_or_the_baseline() { let (_, window) = cases() diff --git a/crates/krishiv-plan-codec/tests/golden.rs b/crates/krishiv-plan-codec/tests/golden.rs index 8af7a0ec0..818759e32 100644 --- a/crates/krishiv-plan-codec/tests/golden.rs +++ b/crates/krishiv-plan-codec/tests/golden.rs @@ -188,7 +188,15 @@ fn the_registry_is_sorted_and_every_feature_has_goldens() { assert!(!f.meaning().is_empty(), "{id}: no meaning"); let dir = root().join(id); let files = std::fs::read_dir(&dir).map_or(0, Iterator::count); - assert!(files > 0, "{id}: no golden files in {}", dir.display()); + // `legacy_statistical_aggregate` fires only around a moment + // aggregate, and the encoder refuses every moment aggregate (pinned + // by the TPC-DS ratchet), so no encodable plan can carry it yet: it + // has no golden files until one can. + assert!( + files > 0 || *f == Feature::LegacyStatisticalAggregate, + "{id}: no golden files in {}", + dir.display() + ); } // `goldens` panics on a directory that is none of the three kinds. assert!(!goldens().is_empty()); diff --git a/crates/krishiv-plan-codec/tests/golden/legacy_statistical_aggregate/one_row_legacy_statistical_aggregate.bin b/crates/krishiv-plan-codec/tests/golden/base/one_row_legacy_statistical_aggregate.bin similarity index 76% rename from crates/krishiv-plan-codec/tests/golden/legacy_statistical_aggregate/one_row_legacy_statistical_aggregate.bin rename to crates/krishiv-plan-codec/tests/golden/base/one_row_legacy_statistical_aggregate.bin index 0c0ac12745bd9b8942be01d32cd85cdaca68cd9b..f0dd1f6c231ed3a81cfbbc8ea34f99ea4033387c 100644 GIT binary patch delta 16 XcmZ3&+{Vnp!79MCk)?r=DS;6H8;k<& delta 46 zcmZo;Uc$`6!79MEk)?rAAxb7EH9ax8GQPMZu_Uv&Br`cNCq6Mfy$C2;lA6E>068-c A1ONa4 diff --git a/crates/krishiv-plan-codec/tests/roundtrip.rs b/crates/krishiv-plan-codec/tests/roundtrip.rs index 106c7ca54..ecff4b2a3 100644 --- a/crates/krishiv-plan-codec/tests/roundtrip.rs +++ b/crates/krishiv-plan-codec/tests/roundtrip.rs @@ -293,7 +293,7 @@ fn readers_that_predate_features_still_refuse_what_they_cannot_read() { "window_ranking" | "window_aggregates" => 4, n if n.starts_with("expand") => 5, // A feature added after the freeze: above every frozen version. - "window_group_limit" | "nested_types" | "one_row_legacy_statistical_aggregate" => 6, + "window_group_limit" | "nested_types" => 6, n if n.starts_with("hash_join_right") => POST_MIGRATION_READER_VERSION, "hash_join_existence" | "hash_join_single_match" | "nested_loop_existence" => { POST_MIGRATION_READER_VERSION