diff --git a/crates/krishiv-common/src/config/registry.rs b/crates/krishiv-common/src/config/registry.rs index c2e04ce01..1b67d1dce 100644 --- a/crates/krishiv-common/src/config/registry.rs +++ b/crates/krishiv-common/src/config/registry.rs @@ -61,6 +61,7 @@ pub static REGISTRY: &[&ConfigEntry] = &[ &RUNTIME_FILTER_ADAPTIVE_MAX_PASS_PERCENT, &RUNTIME_FILTER_ADAPTIVE_MIN_ROWS, &WINDOW_GROUP_LIMIT_THRESHOLD, + &LEGACY_STATISTICAL_AGGREGATE, ]; /// Every key and Spark alias with its entry, sorted by (length, name) for @@ -165,6 +166,10 @@ pub static INDEX: &[(&str, &ConfigEntry)] = &[ "spark.sql.leafNodeDefaultParallelism", &LEAF_NODE_DEFAULT_PARALLELISM, ), + ( + "spark.sql.legacy.statisticalAggregate", + &LEGACY_STATISTICAL_AGGREGATE, + ), ( "krishiv.runtimeFilter.adaptive.minRows", &RUNTIME_FILTER_ADAPTIVE_MIN_ROWS, @@ -181,6 +186,10 @@ pub static INDEX: &[(&str, &ConfigEntry)] = &[ "spark.sql.sources.bucketing.maxBuckets", &BUCKETING_MAX_BUCKETS, ), + ( + "krishiv.sql.legacy.statisticalAggregate", + &LEGACY_STATISTICAL_AGGREGATE, + ), ( "krishiv.runtimeFilter.exactSet.maxValues", &RUNTIME_FILTER_EXACT_SET_MAX_VALUES, @@ -859,3 +868,16 @@ pub static WINDOW_GROUP_LIMIT_THRESHOLD: ConfigEntry = ConfigEntry { version_added: "0.1.0", spark_alias: Some("spark.sql.optimizer.windowGroupLimitThreshold"), }; + +/// `spark.sql.legacy.statisticalAggregate`. +pub static LEGACY_STATISTICAL_AGGREGATE: ConfigEntry = ConfigEntry { + key: "krishiv.sql.legacy.statisticalAggregate", + value_type: ValueType::Bool, + default: Some(ConfigValue::Bool(false)), + doc: "When true, `stddev_samp`, `var_samp`, `covar_samp` and `corr` of a single row, and \ + `skewness` and `kurtosis` of values without variance, return NaN; when false they \ + return NULL, as Spark 4 does.", + scope: Scope::Session, + version_added: "0.1.0", + spark_alias: Some("spark.sql.legacy.statisticalAggregate"), +}; diff --git a/crates/krishiv-exec/benches/aggregate.rs b/crates/krishiv-exec/benches/aggregate.rs index 5f4ecfbf3..3790f5b61 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 cae285b12..f994004c1 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 332792e87..2ed3f86cb 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, @@ -252,6 +253,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. @@ -547,6 +550,18 @@ fn counts_of(a: &ArrayRef) -> Result<&[i64]> { .ok_or_else(|| internal(format!("count state of type {}", a.data_type()))) } +/// A moment aggregate's state column: `double`s, none NULL (a buffer is +/// never NULL; a NULL would silently skip a partial's rows). +fn state_doubles(a: &ArrayRef) -> Result<&[f64]> { + match a.as_primitive_opt::() { + Some(c) if c.null_count() == 0 => Ok(c.values().as_ref()), + _ => Err(internal(format!( + "moment state of type {} with NULLs or not double", + a.data_type() + ))), + } +} + /// A `boolean` state column. fn booleans_of(a: &ArrayRef) -> Result<&BooleanArray> { a.as_boolean_opt() @@ -917,6 +932,18 @@ impl Accumulator { null: Vec::new(), heap: 0, }, + ( + f @ (F::StddevSamp + | F::StddevPop + | F::VarSamp + | F::VarPop + | F::Skewness + | F::Kurtosis + | F::Corr + | F::CovarSamp + | F::CovarPop), + _, + ) => Kind::Moments(Moments::new(f, ansi)?), }; let sum_bound = match (expr.function, expr.avg_sum_type()) { (F::Sum | F::Avg, DataType::Decimal64(p, _) | DataType::Decimal128(p, _)) => { @@ -935,6 +962,17 @@ impl Accumulator { }) } + /// Whether the sample moment aggregates give NaN instead of NULL where + /// Spark divides by zero (`spark.sql.legacy.statisticalAggregate`); the + /// other aggregates ignore it. + #[must_use] + pub(crate) fn legacy_statistical(mut self, on: bool) -> Self { + if let Kind::Moments(m) = &mut self.kind { + m.legacy(on); + } + self + } + /// Bytes one group's state takes (heap bytes of string values aside). pub(crate) fn group_bytes(&self) -> usize { self.vecs().iter().map(|&(_, _, e)| e).sum() @@ -967,6 +1005,7 @@ impl Accumulator { null, .. } => vec![values.vec(), keys.vec(), v(set), v(null)], + Kind::Moments(m) => m.vecs(), } } @@ -1068,6 +1107,7 @@ impl Accumulator { null.resize(n, false); values.resize(n); } + Kind::Moments(m) => m.resize(n), Kind::By { values, keys, @@ -1123,6 +1163,7 @@ impl Accumulator { null.reserve(new); values.reserve(new); } + Kind::Moments(m) => m.reserve(new), Kind::By { values, keys, @@ -1215,7 +1256,7 @@ impl Accumulator { if (state || matches!( self.kind, - Kind::MinMax { .. } | Kind::Pick { .. } | Kind::By { .. } + Kind::MinMax { .. } | Kind::Pick { .. } | Kind::By { .. } | Kind::Moments(_) )) && !is_unit(w) { @@ -1251,6 +1292,7 @@ impl Accumulator { self.kind, Kind::SumDec { .. } | Kind::AvgDec { .. } + | Kind::Moments(_) | Kind::MinMax { values: Values::Bytes(_), .. @@ -1526,6 +1568,23 @@ impl Accumulator { } } } + (Kind::Moments(m), false) => { + let cols = args.iter().map(|a| floats(a)).collect::>>()?; + 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)?; + } ( Kind::By { max, @@ -1592,6 +1651,7 @@ impl Accumulator { | Kind::AvgDec { .. } | Kind::Pick { .. } | Kind::By { .. } => 2, + Kind::Moments(ref m) => m.width(), } } @@ -1651,6 +1711,7 @@ impl Accumulator { ids.iter().map(|i| Some(set[pick(i)])), )), ], + Kind::Moments(m) => m.state(ids)?, Kind::By { values, keys, @@ -1796,6 +1857,7 @@ impl Accumulator { } self.decimal(out.into_iter(), &self.result)? } + Kind::Moments(m) => m.evaluate(ids)?, // Only the value: the ordering key is state, not a result. Kind::By { values, set, null, .. diff --git a/crates/krishiv-exec/src/aggregate/accumulator/differential.rs b/crates/krishiv-exec/src/aggregate/accumulator/differential.rs index 0a4265153..6cbbb40d1 100644 --- a/crates/krishiv-exec/src/aggregate/accumulator/differential.rs +++ b/crates/krishiv-exec/src/aggregate/accumulator/differential.rs @@ -331,6 +331,48 @@ fn min_max_on_decimal64() { } } +/// The moment aggregates: `(group, x, y, x is valid, y is valid)`. +type MomentRow = (u32, f64, f64, bool, bool); + +fn moment_acc(f: AggregateFunction, groups: u32) -> Accumulator { + let args = if crate::aggregate::arity(f) == 2 { + vec![0, 1] + } else { + vec![0] + }; + let e = AggregateExpr { + function: f, + args, + input: AggInput::Raw, + arg_type: DataType::Float64, + name: "m".to_owned(), + }; + // Without ANSI: `corr` of a constant column is NULL, not an error. + let mut a = Accumulator::new(&e, false).unwrap(); + a.resize(groups as usize); + a +} + +/// The columns of `rows` as the arguments of a moment aggregate (garbage +/// under the NULL slots). +fn moment_columns(rows: &[MomentRow]) -> Vec { + 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() +} + /// Keys of the `max_by`/`min_by` merge test: NaN (the largest), the /// infinities, and `-0.0 = 0.0` (a tie). const KEYS: [f64; 7] = [ @@ -387,6 +429,74 @@ fn by_expr(f: AggregateFunction, strings: (bool, bool)) -> AggregateExpr { } } +proptest! { + #![proptest_config(ProptestConfig { + cases: 96, + failure_persistence: None, + ..ProptestConfig::default() + })] + + /// Partial states of any split of the input merge, in any order and + /// through a second level, to the single pass's answer (up to the + /// reassociation of the floating-point operations). + #[test] + fn moments_merge_differential( + groups in 1..6_u32, + rows in proptest::collection::vec( + (0..6_u32, -1000.0..1000.0_f64, -1000.0..1000.0_f64, + proptest::bool::weighted(0.85), proptest::bool::weighted(0.85)), + 0..160, + ), + parts in 1..7_usize, + rotate in 0..7_usize, + reverse in proptest::bool::ANY, + ) { + let rows: Vec = 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), + } + } + } + } +} + proptest! { #![proptest_config(ProptestConfig { cases: 256, @@ -464,3 +574,15 @@ proptest! { } } } + +const MOMENT_FUNCTIONS: [AggregateFunction; 9] = [ + AggregateFunction::StddevSamp, + AggregateFunction::StddevPop, + AggregateFunction::VarSamp, + AggregateFunction::VarPop, + AggregateFunction::Skewness, + AggregateFunction::Kurtosis, + AggregateFunction::Corr, + AggregateFunction::CovarSamp, + AggregateFunction::CovarPop, +]; diff --git a/crates/krishiv-exec/src/aggregate/mod.rs b/crates/krishiv-exec/src/aggregate/mod.rs index 869539766..db4f9a937 100644 --- a/crates/krishiv-exec/src/aggregate/mod.rs +++ b/crates/krishiv-exec/src/aggregate/mod.rs @@ -1,9 +1,10 @@ // SPDX-License-Identifier: Apache-2.0 //! Hash aggregation (ADR 0011): `count`, `sum`, `avg`, `min`, `max`, -//! `first`, `last`, `any_value`, `max_by`, `min_by` with Spark 4.2 result -//! types, in two phases, with spill and `DISTINCT`. `count_if` and the -//! boolean aggregates are `count` and `min`/`max` over a rewritten argument -//! (the physical planner). +//! `first`, `last`, `any_value`, `max_by`, `min_by` and the moment +//! aggregates (`stddev`, `variance`, `skewness`, `kurtosis`, `covar_*`, +//! `corr`) with Spark 4.2 result types, in two phases, with spill and +//! `DISTINCT`. `count_if` and the boolean aggregates are `count` and +//! `min`/`max` over a rewritten argument (the physical planner). //! //! # Operators //! - [`HashAggregate`] is a pipeline breaker. Each driver aggregates into @@ -53,12 +54,14 @@ //! 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`, `first`, `last`, -//! `any_value`, `max_by` and `min_by` need retractions to be maintained -//! incrementally (P4 IVM), so they and `DISTINCT` reject any weight other -//! than 1 with `KRISHIV_UNSUPPORTED_FEATURE`. Output rows have weight 1; -//! a group exists once it received a row. +//! `any_value`, `max_by`, `min_by` and the moment aggregates need +//! retractions to be maintained incrementally (P4 IVM), so they and +//! `DISTINCT` reject any weight other than 1 with +//! `KRISHIV_UNSUPPORTED_FEATURE`. Output rows have weight 1; a group exists +//! once it received a row. mod accumulator; +mod moments; mod node; mod partial; mod rowtable; @@ -80,6 +83,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)] @@ -88,6 +92,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 fn is_moment(f: AggregateFunction) -> bool { + moments::buffer(f).is_some() +} + /// An aggregate function. #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] pub enum AggregateFunction { @@ -116,6 +127,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, /// `any_value(x)`: some row's value of the group; here the first row's, /// as `first` (a legal choice); NULL for no value. The input type. /// In-order merges keep the first row in input order; a `DISTINCT` @@ -256,6 +294,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::MaxBy | F::MinBy, _) => match self.by_types() { Some((v, k)) if accumulator::min_max_supported(v) && accumulator::min_max_supported(k) => @@ -298,6 +341,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), // The key is NULL until a row with a key was taken: it says // whether the value (which may be NULL) counts. (F::MaxBy | F::MinBy, _) => { @@ -371,6 +426,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..ecce68e72 --- /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) -> Result<()> { + for (col, v) in self.cols.iter_mut().zip(row) { + *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 + /// 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 3ac57e97f..9ed89da49 100644 --- a/crates/krishiv-exec/src/aggregate/rowtable.rs +++ b/crates/krishiv-exec/src/aggregate/rowtable.rs @@ -1481,7 +1481,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() } @@ -1713,6 +1713,7 @@ mod tests { ansi: true, spill: None, require_unit_weights: false, + legacy_statistical: false, }; Arc::new(AggCore::new(spec).unwrap()) } @@ -1746,6 +1747,7 @@ mod tests { ansi: true, spill: None, require_unit_weights: false, + legacy_statistical: false, }; Arc::new(AggCore::new(spec).unwrap()) } @@ -2206,6 +2208,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 5f3d44c59..6912e3af0 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), @@ -345,7 +350,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 00fb1ed42..182463c30 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) @@ -2580,3 +2592,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 c71bd8259..f30c5aecc 100644 --- a/crates/krishiv-exec/src/physical/instantiate.rs +++ b/crates/krishiv-exec/src/physical/instantiate.rs @@ -387,8 +387,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 }) @@ -760,7 +759,7 @@ pub(crate) fn scan_filter( pub(crate) fn aggregate_spec( input: &SchemaRef, a: &AggregatePlan, - ansi: bool, + env: &EvalEnv, spill: Option, ) -> AggregateSpec { AggregateSpec { @@ -768,9 +767,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 e2ad0bc01..7e81f0f64 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 @@ -493,7 +496,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 81ba3c3a5..ff7d94f0b 100644 --- a/crates/krishiv-exec/src/planner/aggregate.rs +++ b/crates/krishiv-exec/src/planner/aggregate.rs @@ -19,7 +19,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}; @@ -97,7 +97,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(); @@ -259,6 +259,19 @@ pub(super) fn lower(f: &ResolvedFunction) -> Option { "any_value" => (AggregateFunction::AnyValue { ignore_nulls }, value), "max_by" if f.args.len() == 2 => (AggregateFunction::MaxBy, f.args.clone()), "min_by" if f.args.len() == 2 => (AggregateFunction::MinBy, f.args.clone()), + "stddev" | "std" | "stddev_samp" if f.args.len() == 1 => { + (AggregateFunction::StddevSamp, f.args.clone()) + } + "stddev_pop" if f.args.len() == 1 => (AggregateFunction::StddevPop, f.args.clone()), + "variance" | "var_samp" if f.args.len() == 1 => { + (AggregateFunction::VarSamp, f.args.clone()) + } + "var_pop" if f.args.len() == 1 => (AggregateFunction::VarPop, f.args.clone()), + "skewness" if f.args.len() == 1 => (AggregateFunction::Skewness, f.args.clone()), + "kurtosis" if f.args.len() == 1 => (AggregateFunction::Kurtosis, f.args.clone()), + "corr" if f.args.len() == 2 => (AggregateFunction::Corr, f.args.clone()), + "covar_samp" if f.args.len() == 2 => (AggregateFunction::CovarSamp, f.args.clone()), + "covar_pop" if f.args.len() == 2 => (AggregateFunction::CovarPop, f.args.clone()), _ => return None, }; Some(Lowered { function, args }) @@ -323,8 +336,8 @@ pub(super) fn aggregate_expr( name, )); } - 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")); } let arg_type = args.first().map_or(arrow_schema::DataType::Null, typed); Ok(AggregateExpr { diff --git a/crates/krishiv-exec/src/planner/eval/mod.rs b/crates/krishiv-exec/src/planner/eval/mod.rs index 9d6510f25..a56a5cb4d 100644 --- a/crates/krishiv-exec/src/planner/eval/mod.rs +++ b/crates/krishiv-exec/src/planner/eval/mod.rs @@ -36,7 +36,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, @@ -154,6 +156,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, @@ -172,6 +176,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()) @@ -183,6 +191,7 @@ impl EvalEnv { .unwrap_or_default(); Self { ansi, + legacy_statistical_aggregate, zone, kernels: KernelCtx::new(), query_start_micros, @@ -194,6 +203,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, @@ -205,6 +215,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/planner/window.rs b/crates/krishiv-exec/src/planner/window.rs index 92504cc3e..7c230087b 100644 --- a/crates/krishiv-exec/src/planner/window.rs +++ b/crates/krishiv-exec/src/planner/window.rs @@ -14,6 +14,7 @@ use krishiv_plan::{ use super::aggregate::{aggregate_expr, lower}; use super::build::{Planned, Planner, internal, list, orders, unsupported}; use super::properties::Distribution; +use crate::aggregate::is_moment; use crate::error::Result; use crate::operator::Source; use crate::operators::SortKey; @@ -289,9 +290,14 @@ impl Planner<'_> { other => return Err(unsupported(&format!("the window function {other}"))), }); } - let function = lower(c.f) - .ok_or_else(|| unsupported(&format!("the window aggregate {name}")))? - .function; + let lowered = + lower(c.f).ok_or_else(|| unsupported(&format!("the window aggregate {name}")))?; + // The moment aggregates have no window form yet: refused, never + // computed as another function. + if is_moment(lowered.function) { + return Err(unsupported(&format!("the window aggregate {name}"))); + } + let function = lowered.function; // 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 { 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/src/window/tests.rs b/crates/krishiv-exec/src/window/tests.rs index b6693d683..60dee3cd7 100644 --- a/crates/krishiv-exec/src/window/tests.rs +++ b/crates/krishiv-exec/src/window/tests.rs @@ -479,7 +479,16 @@ fn reference( } AggregateFunction::AnyValue { .. } | AggregateFunction::MaxBy - | AggregateFunction::MinBy => { + | AggregateFunction::MinBy + | AggregateFunction::StddevSamp + | AggregateFunction::StddevPop + | AggregateFunction::VarSamp + | AggregateFunction::VarPop + | AggregateFunction::Skewness + | AggregateFunction::Kurtosis + | AggregateFunction::Corr + | AggregateFunction::CovarSamp + | AggregateFunction::CovarPop => { unreachable!("no window test aggregates with this function") } } 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 2413b776f..2f1ee3862 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/python_wrappers/table.rs b/crates/krishiv-functions/src/python_wrappers/table.rs index 20d57b5f1..2d4317faf 100644 --- a/crates/krishiv-functions/src/python_wrappers/table.rs +++ b/crates/krishiv-functions/src/python_wrappers/table.rs @@ -110,6 +110,7 @@ pub(super) fn rows() -> Vec { "isnan", "isnotnull", "isnull", + "kurtosis", "length", "ln", "log10", @@ -139,8 +140,13 @@ pub(super) fn rows() -> Vec { "signum", "sin", "sinh", + "skewness", "soundex", "sqrt", + "std", + "stddev", + "stddev_pop", + "stddev_samp", "sum", "tan", "tanh", @@ -148,6 +154,9 @@ pub(super) fn rows() -> Vec { "unhex", "unix_date", "upper", + "var_pop", + "var_samp", + "variance", "weekday", "weekofyear", "year", @@ -172,7 +181,16 @@ pub(super) fn rows() -> Vec { for name in ["bit_get", "getbit"] { rows.push(Row::new(name, "col: ColumnOrName, pos: ColumnOrName")); } - for name in ["equal_null", "ifnull", "nanvl", "nullif", "nvl"] { + for name in [ + "corr", + "covar_pop", + "covar_samp", + "equal_null", + "ifnull", + "nanvl", + "nullif", + "nvl", + ] { rows.push(Row::new(name, "col1: ColumnOrName, col2: ColumnOrName")); } for name in ["to_char", "to_number", "to_varchar", "try_to_number"] { diff --git a/crates/krishiv-functions/src/registry/defs.rs b/crates/krishiv-functions/src/registry/defs.rs index 1f73f7ac7..7c5a80329 100644 --- a/crates/krishiv-functions/src/registry/defs.rs +++ b/crates/krishiv-functions/src/registry/defs.rs @@ -613,6 +613,12 @@ fn any_value(c: &CallTypes<'_>) -> Result { pick(c, "any_value") } +/// 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)) +} + /// `bool_and`, `bool_or` and their aliases: `boolean`, NULL for no non-NULL /// input. fn nullable_boolean(_: &CallTypes<'_>) -> Result { @@ -635,6 +641,7 @@ fn by(c: &CallTypes<'_>) -> Result { fn aggregates() -> Vec { let numeric = || Inputs::Expects(vec![InputType::Numeric { intervals: true }]); let boolean_arg = || Inputs::Expects(vec![exact(DataType::Boolean)]); + let doubles = || Inputs::Expects(vec![exact(DataType::Double)]); vec![ aggregate("count", Inputs::Any, count, Aggregate::Count), aggregate("sum", numeric(), sum, Aggregate::Sum), @@ -647,6 +654,18 @@ fn aggregates() -> Vec { aggregate("last", Inputs::Any, last, Aggregate::Last), aggregate("last_value", Inputs::Any, last, Aggregate::Last), aggregate("any_value", Inputs::Any, any_value, Aggregate::AnyValue), + 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), // Spark runs these as `count`, `min` and `max` of a rewritten argument. aggregate("count_if", boolean_arg(), count, Aggregate::Count), aggregate("bool_and", boolean_arg(), nullable_boolean, Aggregate::Min), diff --git a/crates/krishiv-functions/src/registry/mod.rs b/crates/krishiv-functions/src/registry/mod.rs index 2498f9863..0705189da 100644 --- a/crates/krishiv-functions/src/registry/mod.rs +++ b/crates/krishiv-functions/src/registry/mod.rs @@ -67,6 +67,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, /// `any_value`; `ignoreNulls` is the second argument. AnyValue, /// `max_by(x, y)`. diff --git a/crates/krishiv-functions/src/registry/tests.rs b/crates/krishiv-functions/src/registry/tests.rs index a3406d0ad..13716bcbc 100644 --- a/crates/krishiv-functions/src/registry/tests.rs +++ b/crates/krishiv-functions/src/registry/tests.rs @@ -1519,4 +1519,67 @@ fn regex_signatures() { 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})"); + } +} + mod text; diff --git a/crates/krishiv-plan-codec/src/decode.rs b/crates/krishiv-plan-codec/src/decode.rs index 5e1347fc9..b6fc323d0 100644 --- a/crates/krishiv-plan-codec/src/decode.rs +++ b/crates/krishiv-plan-codec/src/decode.rs @@ -288,6 +288,12 @@ 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 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 (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, 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 f77911ca1..ff4e6a658 100644 --- a/crates/krishiv-plan-codec/src/encode.rs +++ b/crates/krishiv-plan-codec/src/encode.rs @@ -130,6 +130,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, } } @@ -717,13 +718,22 @@ fn aggregate_expr(a: &AggregateExpr, nested: &Cell) -> Result (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 ids for these yet. - AggregateFunction::AnyValue { .. } + // The format has no code for these yet. + f @ (AggregateFunction::StddevSamp + | AggregateFunction::StddevPop + | AggregateFunction::VarSamp + | AggregateFunction::VarPop + | AggregateFunction::Skewness + | AggregateFunction::Kurtosis + | AggregateFunction::Corr + | AggregateFunction::CovarSamp + | AggregateFunction::CovarPop + | AggregateFunction::AnyValue { .. } | AggregateFunction::MaxBy - | AggregateFunction::MinBy => { - return Err(CodecError::Unsupported( - "an aggregate function the format does not know".to_owned(), - )); + | AggregateFunction::MinBy) => { + return Err(CodecError::Unsupported(format!( + "the aggregate function {f:?}" + ))); } }; Ok(pb::AggregateExpr { diff --git a/crates/krishiv-plan-codec/src/features.rs b/crates/krishiv-plan-codec/src/features.rs index b44d5efe2..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}; @@ -121,6 +121,7 @@ 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`)."), (join_single_match, JoinSingleMatch, POST_MIGRATION_READER_VERSION, Semantic, "The single-match flag of a `LeftOuter` hash join (a scalar subquery: more than one match per probe row is an error)."), + (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)."), (nested_types, NestedTypes, POST_MIGRATION_READER_VERSION, Semantic, "Array, map and struct types (`DataType.KIND_NESTED`, written as Spark's JSON form of the type)."), (window, Window, 4, Semantic, "The `Window` node and its functions."), @@ -282,6 +283,17 @@ pub fn feature_of(construct: Construct<'_>) -> Option { | AggregateFunction::Max | AggregateFunction::First { .. } | AggregateFunction::Last { .. } + // These 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 | AggregateFunction::AnyValue { .. } | AggregateFunction::MaxBy | AggregateFunction::MinBy => None, @@ -429,18 +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. +/// 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(); + 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/src/proto/krishiv.plan.v1.pb.rs b/crates/krishiv-plan-codec/src/proto/krishiv.plan.v1.pb.rs index 22342b7f6..72c16f7b8 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 bd61a0c0d..f63a2bc28 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, parse_data_type}; 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"), @@ -764,6 +765,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", @@ -772,8 +780,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. @@ -941,6 +948,17 @@ 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 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(( + "one_row_legacy_statistical_aggregate".to_owned(), + plan_of_env(PhysNode::OneRow, legacy), + )); out.push(("nested_types".to_owned(), plan_of(nested_types()))); out } 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/base/one_row_legacy_statistical_aggregate.bin b/crates/krishiv-plan-codec/tests/golden/base/one_row_legacy_statistical_aggregate.bin new file mode 100644 index 000000000..f0dd1f6c2 Binary files /dev/null and b/crates/krishiv-plan-codec/tests/golden/base/one_row_legacy_statistical_aggregate.bin differ 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..ccbbe54cf --- /dev/null +++ b/crates/krishiv-session/tests/aggregate_moments.rs @@ -0,0 +1,640 @@ +// 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::{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}; +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::Float32 => "float", + 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::Float32 => format!("{:?}", c.as_primitive::().value(row)), + 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(()) +} + +/// 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?; + 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)], + ], + )?; + aggtest(&s)?; + 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(()) +} + +/// 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 { + // 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)?; + 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); + // 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"; + 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 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(*)", + "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", + ), + ( + "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] +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)"), + ( + "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(); + 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 6477f108a..eb0fd759e 100644 --- a/crates/krishiv-testkit/tests/type_coercion.rs +++ b/crates/krishiv-testkit/tests/type_coercion.rs @@ -208,8 +208,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,645 resolve as Spark resolves them. - assert_eq!((queries, queries - failures.len()), (2858, 2645)); + // Of 2,858 query statements, 2,651 resolve as Spark resolves them. + assert_eq!((queries, queries - failures.len()), (2858, 2651)); let dir = root.join(MANIFEST_DIR); if std::env::var_os("KRISHIV_UPDATE_MANIFEST").is_some() { let mut by_file: BTreeMap<&str, Vec<(&String, &String)>> = BTreeMap::new(); diff --git a/proto/krishiv/plan/v1/plan.proto b/proto/krishiv/plan/v1/plan.proto index 6d641f545..fe4207e12 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 { diff --git a/python/krishiv/sql/_generated_functions.py b/python/krishiv/sql/_generated_functions.py index 085f75a1f..9b29680ee 100644 --- a/python/krishiv/sql/_generated_functions.py +++ b/python/krishiv/sql/_generated_functions.py @@ -56,9 +56,12 @@ "concat_ws", "contains", "conv", + "corr", "cos", "cosh", "cot", + "covar_pop", + "covar_samp", "crc32", "csc", "cume_dist", @@ -108,6 +111,7 @@ "isnan", "isnotnull", "isnull", + "kurtosis", "lag", "last", "last_day", @@ -193,10 +197,15 @@ "signum", "sin", "sinh", + "skewness", "soundex", "split_part", "sqrt", "startswith", + "std", + "stddev", + "stddev_pop", + "stddev_samp", "substr", "substring", "substring_index", @@ -231,6 +240,9 @@ "url_encode", "user", "validate_utf8", + "var_pop", + "var_samp", + "variance", "weekday", "weekofyear", "xxhash64", @@ -395,6 +407,11 @@ def conv(col: ColumnOrName, fromBase: int, toBase: int) -> Column: return _call("conv", _named(col, "col"), _lit(fromBase, "fromBase"), _lit(toBase, "toBase")) +def corr(col1: ColumnOrName, col2: ColumnOrName) -> Column: + """Calls the registry function ``corr``.""" + return _call("corr", _named(col1, "col1"), _named(col2, "col2")) + + def cos(col: ColumnOrName) -> Column: """Calls the registry function ``cos``.""" return _call("cos", _named(col, "col")) @@ -410,6 +427,16 @@ def cot(col: ColumnOrName) -> Column: return _call("cot", _named(col, "col")) +def covar_pop(col1: ColumnOrName, col2: ColumnOrName) -> Column: + """Calls the registry function ``covar_pop``.""" + return _call("covar_pop", _named(col1, "col1"), _named(col2, "col2")) + + +def covar_samp(col1: ColumnOrName, col2: ColumnOrName) -> Column: + """Calls the registry function ``covar_samp``.""" + return _call("covar_samp", _named(col1, "col1"), _named(col2, "col2")) + + def crc32(col: ColumnOrName) -> Column: """Calls the registry function ``crc32``.""" return _call("crc32", _named(col, "col")) @@ -656,6 +683,11 @@ def isnull(col: ColumnOrName) -> Column: return _call("isnull", _named(col, "col")) +def kurtosis(col: ColumnOrName) -> Column: + """Calls the registry function ``kurtosis``.""" + return _call("kurtosis", _named(col, "col")) + + def lag(col: ColumnOrName, offset: int = 1, default: Any | None = None) -> Column: """Calls the registry function ``lag``.""" return _call("lag", _named(col, "col"), _lit(offset, "offset"), *_opt(_lit, default, "default")) @@ -1101,6 +1133,11 @@ def sinh(col: ColumnOrName) -> Column: return _call("sinh", _named(col, "col")) +def skewness(col: ColumnOrName) -> Column: + """Calls the registry function ``skewness``.""" + return _call("skewness", _named(col, "col")) + + def soundex(col: ColumnOrName) -> Column: """Calls the registry function ``soundex``.""" return _call("soundex", _named(col, "col")) @@ -1121,6 +1158,26 @@ 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")) + + +def stddev_pop(col: ColumnOrName) -> Column: + """Calls the registry function ``stddev_pop``.""" + return _call("stddev_pop", _named(col, "col")) + + +def stddev_samp(col: ColumnOrName) -> Column: + """Calls the registry function ``stddev_samp``.""" + return _call("stddev_samp", _named(col, "col")) + + def substr(str: ColumnOrName, pos: ColumnOrName, len: ColumnOrName | None = None) -> Column: """Calls the registry function ``substr``.""" return _call("substr", _named(str, "str"), _named(pos, "pos"), *_opt(_named, len, "len")) @@ -1293,6 +1350,21 @@ def validate_utf8(str: ColumnOrName) -> Column: return _call("validate_utf8", _named(str, "str")) +def var_pop(col: ColumnOrName) -> Column: + """Calls the registry function ``var_pop``.""" + return _call("var_pop", _named(col, "col")) + + +def var_samp(col: ColumnOrName) -> Column: + """Calls the registry function ``var_samp``.""" + return _call("var_samp", _named(col, "col")) + + +def variance(col: ColumnOrName) -> Column: + """Calls the registry function ``variance``.""" + return _call("variance", _named(col, "col")) + + def weekday(col: ColumnOrName) -> Column: """Calls the registry function ``weekday``.""" return _call("weekday", _named(col, "col")) diff --git a/python/krishiv/sql/_generated_functions.pyi b/python/krishiv/sql/_generated_functions.pyi index 82a0a317c..a3d5f33ef 100644 --- a/python/krishiv/sql/_generated_functions.pyi +++ b/python/krishiv/sql/_generated_functions.pyi @@ -37,9 +37,12 @@ __all__ = [ "concat_ws", "contains", "conv", + "corr", "cos", "cosh", "cot", + "covar_pop", + "covar_samp", "crc32", "csc", "cume_dist", @@ -89,6 +92,7 @@ __all__ = [ "isnan", "isnotnull", "isnull", + "kurtosis", "lag", "last", "last_day", @@ -174,10 +178,15 @@ __all__ = [ "signum", "sin", "sinh", + "skewness", "soundex", "split_part", "sqrt", "startswith", + "std", + "stddev", + "stddev_pop", + "stddev_samp", "substr", "substring", "substring_index", @@ -212,6 +221,9 @@ __all__ = [ "url_encode", "user", "validate_utf8", + "var_pop", + "var_samp", + "variance", "weekday", "weekofyear", "xxhash64", @@ -250,9 +262,12 @@ def coalesce(*cols: ColumnOrName) -> Column: ... def concat_ws(sep: str, *cols: ColumnOrName) -> Column: ... def contains(left: ColumnOrName, right: ColumnOrName) -> Column: ... def conv(col: ColumnOrName, fromBase: int, toBase: int) -> Column: ... +def corr(col1: ColumnOrName, col2: ColumnOrName) -> Column: ... def cos(col: ColumnOrName) -> Column: ... def cosh(col: ColumnOrName) -> Column: ... def cot(col: ColumnOrName) -> Column: ... +def covar_pop(col1: ColumnOrName, col2: ColumnOrName) -> Column: ... +def covar_samp(col1: ColumnOrName, col2: ColumnOrName) -> Column: ... def crc32(col: ColumnOrName) -> Column: ... def csc(col: ColumnOrName) -> Column: ... def cume_dist() -> Column: ... @@ -302,6 +317,7 @@ def is_valid_utf8(str: ColumnOrName) -> Column: ... def isnan(col: ColumnOrName) -> Column: ... def isnotnull(col: ColumnOrName) -> Column: ... def isnull(col: ColumnOrName) -> Column: ... +def kurtosis(col: ColumnOrName) -> Column: ... def lag(col: ColumnOrName, offset: int = 1, default: Any | None = None) -> Column: ... def last(col: ColumnOrName, ignorenulls: bool = False) -> Column: ... def last_day(date: ColumnOrName) -> Column: ... @@ -395,10 +411,15 @@ def sign(col: ColumnOrName) -> Column: ... def signum(col: ColumnOrName) -> Column: ... def sin(col: ColumnOrName) -> Column: ... def sinh(col: ColumnOrName) -> Column: ... +def skewness(col: ColumnOrName) -> Column: ... 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: ... def substr(str: ColumnOrName, pos: ColumnOrName, len: ColumnOrName | None = None) -> Column: ... def substring(str: ColumnOrName, pos: ColumnOrName | int, len: ColumnOrName | int) -> Column: ... def substring_index(str: ColumnOrName, delim: str, count: int) -> Column: ... @@ -433,6 +454,9 @@ def url_decode(str: ColumnOrName) -> Column: ... def url_encode(str: ColumnOrName) -> Column: ... def user() -> Column: ... def validate_utf8(str: ColumnOrName) -> Column: ... +def var_pop(col: ColumnOrName) -> Column: ... +def var_samp(col: ColumnOrName) -> Column: ... +def variance(col: ColumnOrName) -> Column: ... def weekday(col: ColumnOrName) -> Column: ... def weekofyear(col: ColumnOrName) -> Column: ... def xxhash64(*cols: ColumnOrName) -> Column: ... diff --git a/python/tests/_pyspark_signatures.py b/python/tests/_pyspark_signatures.py index adf626e58..2f85fccce 100644 --- a/python/tests/_pyspark_signatures.py +++ b/python/tests/_pyspark_signatures.py @@ -45,12 +45,15 @@ "contains": "left, right", "conv": "col, fromBase, toBase", "convert_timezone": "sourceTz, targetTz, sourceTs", + "corr": "col1, col2", "cos": "col", "cosh": "col", "cot": "col", "count": "col", "countDistinct": "col, *cols", "count_distinct": "col, *cols", + "covar_pop": "col1, col2", + "covar_samp": "col1, col2", "crc32": "col", "csc": "col", "cume_dist": "", @@ -102,6 +105,7 @@ "isnan": "col", "isnotnull": "col", "isnull": "col", + "kurtosis": "col", "lag": "col, offset=1, default=None", "last": "col, ignorenulls=False", "last_day": "date", @@ -189,10 +193,15 @@ "signum": "col", "sin": "col", "sinh": "col", + "skewness": "col", "soundex": "col", "split_part": "src, delimiter, partNum", "sqrt": "col", "startswith": "str, prefix", + "std": "col", + "stddev": "col", + "stddev_pop": "col", + "stddev_samp": "col", "substr": "str, pos, len=None", "substring": "str, pos, len", "substring_index": "str, delim, count", @@ -226,6 +235,9 @@ "url_decode": "str", "url_encode": "str", "user": "", + "var_pop": "col", + "var_samp": "col", + "variance": "col", "weekday": "col", "weekofyear": "col", "validate_utf8": "str", diff --git a/tests/fixtures/analysis/type-coercion/typeCoercion/native/promoteStrings.json b/tests/fixtures/analysis/type-coercion/typeCoercion/native/promoteStrings.json index b3b39d1db..a9a838955 100644 --- a/tests/fixtures/analysis/type-coercion/typeCoercion/native/promoteStrings.json +++ b/tests/fixtures/analysis/type-coercion/typeCoercion/native/promoteStrings.json @@ -31,30 +31,6 @@ { "query": "SELECT cast('2017-12-11 09:30:00.0' as timestamp) - '1' FROM t", "reason": "KRISHIV_UNSUPPORTED_FEATURE: date and timestamp arithmetic" - }, - { - "query": "SELECT kurtosis('1') FROM t", - "reason": "KRISHIV_UNSUPPORTED_FEATURE: function kurtosis" - }, - { - "query": "SELECT skewness('1') FROM t", - "reason": "KRISHIV_UNSUPPORTED_FEATURE: function skewness" - }, - { - "query": "SELECT stddev_pop('1') FROM t", - "reason": "KRISHIV_UNSUPPORTED_FEATURE: function stddev_pop" - }, - { - "query": "SELECT stddev_samp('1') FROM t", - "reason": "KRISHIV_UNSUPPORTED_FEATURE: function stddev_samp" - }, - { - "query": "SELECT var_pop('1') FROM t", - "reason": "KRISHIV_UNSUPPORTED_FEATURE: function var_pop" - }, - { - "query": "SELECT var_samp('1') FROM t", - "reason": "KRISHIV_UNSUPPORTED_FEATURE: function var_samp" } ], "file": "typeCoercion/native/promoteStrings", diff --git a/tests/spark-golden/known-failures-cases/postgreSQL/aggregates_part1.json b/tests/spark-golden/known-failures-cases/postgreSQL/aggregates_part1.json index d34f4a44a..b53f6b317 100644 --- a/tests/spark-golden/known-failures-cases/postgreSQL/aggregates_part1.json +++ b/tests/spark-golden/known-failures-cases/postgreSQL/aggregates_part1.json @@ -7,25 +7,13 @@ 48 ], "fn-aggregate": [ - 9, - 10, - 11, - 12, - 13, - 14, - 15, - 16, - 17, - 18, 35, 36, 37, 38, 39, 40, - 41, - 42, - 43 + 41 ], "query-features": [ 19, @@ -38,12 +26,7 @@ 26, 27, 28, - 29, - 30, - 31, - 32, - 33, - 34 + 33 ], "subqueries": [ 54, diff --git a/tests/spark-golden/known-failures-cases/typeCoercion/native/promoteStrings.json b/tests/spark-golden/known-failures-cases/typeCoercion/native/promoteStrings.json index c322eeace..dbd1c378a 100644 --- a/tests/spark-golden/known-failures-cases/typeCoercion/native/promoteStrings.json +++ b/tests/spark-golden/known-failures-cases/typeCoercion/native/promoteStrings.json @@ -36,13 +36,7 @@ 290, 291, 302, - 303, - 308, - 309, - 312, - 313, - 314, - 315 + 303 ] }, "file": "typeCoercion/native/promoteStrings",