andygrove commented on code in PR #4818: URL: https://github.com/apache/datafusion-comet/pull/4818#discussion_r3969622263
########## native/spark-expr/src/agg_funcs/kurtosis.rs: ########## @@ -0,0 +1,363 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Spark-compatible excess-kurtosis aggregate. +//! +//! Spark's `Kurtosis` is a `CentralMomentAgg` (`DeclarativeAggregate`) whose +//! intermediate buffer is `[n, avg, m2, m3, m4]` of Float64. This accumulator +//! mirrors that buffer exactly, using the same higher-order online update / +//! merge recurrences (Meng 2015) that `CentralMomentAgg` compiles into +//! catalyst expressions. Matching the wire format lets Spark's Partial and +//! Comet's Final (or vice versa) share intermediate state without a cast. +//! +//! Result formula (excess kurtosis, Fisher definition): +//! +//! * `n == 0` -> NULL +//! * `m2 == 0` -> NULL when `null_on_divide_by_zero`, else NaN +//! * otherwise -> `n * m4 / (m2 * m2) - 3.0` + +use std::mem::size_of; +use std::sync::Arc; + +use arrow::array::{ArrayRef, Float64Array}; +use arrow::datatypes::{DataType, Field, FieldRef}; +use datafusion::common::{downcast_value, Result, ScalarValue}; +use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs}; +use datafusion::logical_expr::Volatility::Immutable; +use datafusion::logical_expr::{Accumulator, AggregateUDFImpl, Signature}; +use datafusion::physical_expr::expressions::format_state_name; + +#[derive(Debug, PartialEq, Eq)] +pub struct Kurtosis { + name: String, + signature: Signature, + null_on_divide_by_zero: bool, +} + +impl std::hash::Hash for Kurtosis { + fn hash<H: std::hash::Hasher>(&self, state: &mut H) { + self.name.hash(state); + self.signature.hash(state); + self.null_on_divide_by_zero.hash(state); + } +} + +impl Kurtosis { + pub fn new(name: impl Into<String>, null_on_divide_by_zero: bool) -> Self { + Self { + name: name.into(), + signature: Signature::numeric(1, Immutable), + null_on_divide_by_zero, + } + } +} + +impl AggregateUDFImpl for Kurtosis { + fn name(&self) -> &str { + &self.name + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> { + Ok(DataType::Float64) + } + + fn accumulator(&self, _acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> { + Ok(Box::new(KurtosisAccumulator::new( + self.null_on_divide_by_zero, + ))) + } + + // Fields ordered to match Spark's `[n, avg, m2, m3, m4]` buffer so that a + // Spark-produced Partial state can be merged into a Comet-produced Final + // (and vice versa) without a schema conversion. + fn state_fields(&self, _args: StateFieldsArgs) -> Result<Vec<FieldRef>> { + Ok(vec![ + Arc::new(Field::new( + format_state_name(&self.name, "n"), + DataType::Float64, + true, + )), + Arc::new(Field::new( + format_state_name(&self.name, "avg"), + DataType::Float64, + true, + )), + Arc::new(Field::new( + format_state_name(&self.name, "m2"), + DataType::Float64, + true, + )), + Arc::new(Field::new( + format_state_name(&self.name, "m3"), + DataType::Float64, + true, + )), + Arc::new(Field::new( + format_state_name(&self.name, "m4"), + DataType::Float64, + true, + )), + ]) + } + + fn default_value(&self, _data_type: &DataType) -> Result<ScalarValue> { Review Comment: Removed in dcabe65d0. `return_type` already returns `Float64`, which is exactly what the trait default feeds to `ScalarValue::try_from`, so the override was equivalent. `empty_group_returns_null` still passes. Fixed on the `mode.rs` side too, in #4782. ########## native/spark-expr/src/agg_funcs/kurtosis.rs: ########## @@ -0,0 +1,363 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Spark-compatible excess-kurtosis aggregate. +//! +//! Spark's `Kurtosis` is a `CentralMomentAgg` (`DeclarativeAggregate`) whose +//! intermediate buffer is `[n, avg, m2, m3, m4]` of Float64. This accumulator +//! mirrors that buffer exactly, using the same higher-order online update / +//! merge recurrences (Meng 2015) that `CentralMomentAgg` compiles into +//! catalyst expressions. Matching the wire format lets Spark's Partial and +//! Comet's Final (or vice versa) share intermediate state without a cast. +//! +//! Result formula (excess kurtosis, Fisher definition): +//! +//! * `n == 0` -> NULL +//! * `m2 == 0` -> NULL when `null_on_divide_by_zero`, else NaN +//! * otherwise -> `n * m4 / (m2 * m2) - 3.0` + +use std::mem::size_of; +use std::sync::Arc; + +use arrow::array::{ArrayRef, Float64Array}; +use arrow::datatypes::{DataType, Field, FieldRef}; +use datafusion::common::{downcast_value, Result, ScalarValue}; +use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs}; +use datafusion::logical_expr::Volatility::Immutable; +use datafusion::logical_expr::{Accumulator, AggregateUDFImpl, Signature}; +use datafusion::physical_expr::expressions::format_state_name; + +#[derive(Debug, PartialEq, Eq)] +pub struct Kurtosis { + name: String, + signature: Signature, + null_on_divide_by_zero: bool, +} + +impl std::hash::Hash for Kurtosis { + fn hash<H: std::hash::Hasher>(&self, state: &mut H) { + self.name.hash(state); + self.signature.hash(state); + self.null_on_divide_by_zero.hash(state); + } +} + +impl Kurtosis { + pub fn new(name: impl Into<String>, null_on_divide_by_zero: bool) -> Self { + Self { + name: name.into(), + signature: Signature::numeric(1, Immutable), + null_on_divide_by_zero, + } + } +} + +impl AggregateUDFImpl for Kurtosis { + fn name(&self) -> &str { + &self.name + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> { + Ok(DataType::Float64) + } + + fn accumulator(&self, _acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> { + Ok(Box::new(KurtosisAccumulator::new( + self.null_on_divide_by_zero, + ))) + } + + // Fields ordered to match Spark's `[n, avg, m2, m3, m4]` buffer so that a + // Spark-produced Partial state can be merged into a Comet-produced Final + // (and vice versa) without a schema conversion. + fn state_fields(&self, _args: StateFieldsArgs) -> Result<Vec<FieldRef>> { + Ok(vec![ + Arc::new(Field::new( + format_state_name(&self.name, "n"), + DataType::Float64, + true, + )), + Arc::new(Field::new( + format_state_name(&self.name, "avg"), + DataType::Float64, + true, + )), + Arc::new(Field::new( + format_state_name(&self.name, "m2"), + DataType::Float64, + true, + )), + Arc::new(Field::new( + format_state_name(&self.name, "m3"), + DataType::Float64, + true, + )), + Arc::new(Field::new( + format_state_name(&self.name, "m4"), + DataType::Float64, + true, + )), + ]) + } + + fn default_value(&self, _data_type: &DataType) -> Result<ScalarValue> { + Ok(ScalarValue::Float64(None)) + } +} + +/// Online update for the first four central moments. Direct port of Spark's Review Comment: Moved in dcabe65d0. Both functions now live in `welford.rs` next to `variance_update` / `variance_merge` / `covariance_*`, renamed to `moments4_update` and `moments4_merge` since they are no longer kurtosis-specific once they are shared. I took your framing about skewness into the doc comment, so the reason they are named for the moment order rather than the aggregate is on the page: > The order-4 recurrence subsumes the order-2 one above, so `skewness` (`momentOrder = 3`) is an `evaluate` on top of this same state rather than another copy of the algebra. ########## native/spark-expr/src/agg_funcs/kurtosis.rs: ########## @@ -0,0 +1,363 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Spark-compatible excess-kurtosis aggregate. +//! +//! Spark's `Kurtosis` is a `CentralMomentAgg` (`DeclarativeAggregate`) whose +//! intermediate buffer is `[n, avg, m2, m3, m4]` of Float64. This accumulator +//! mirrors that buffer exactly, using the same higher-order online update / +//! merge recurrences (Meng 2015) that `CentralMomentAgg` compiles into +//! catalyst expressions. Matching the wire format lets Spark's Partial and +//! Comet's Final (or vice versa) share intermediate state without a cast. +//! +//! Result formula (excess kurtosis, Fisher definition): +//! +//! * `n == 0` -> NULL +//! * `m2 == 0` -> NULL when `null_on_divide_by_zero`, else NaN +//! * otherwise -> `n * m4 / (m2 * m2) - 3.0` + +use std::mem::size_of; +use std::sync::Arc; + +use arrow::array::{ArrayRef, Float64Array}; +use arrow::datatypes::{DataType, Field, FieldRef}; +use datafusion::common::{downcast_value, Result, ScalarValue}; +use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs}; +use datafusion::logical_expr::Volatility::Immutable; +use datafusion::logical_expr::{Accumulator, AggregateUDFImpl, Signature}; +use datafusion::physical_expr::expressions::format_state_name; + +#[derive(Debug, PartialEq, Eq)] +pub struct Kurtosis { + name: String, + signature: Signature, + null_on_divide_by_zero: bool, +} + +impl std::hash::Hash for Kurtosis { + fn hash<H: std::hash::Hasher>(&self, state: &mut H) { + self.name.hash(state); + self.signature.hash(state); + self.null_on_divide_by_zero.hash(state); + } +} + +impl Kurtosis { + pub fn new(name: impl Into<String>, null_on_divide_by_zero: bool) -> Self { + Self { + name: name.into(), + signature: Signature::numeric(1, Immutable), + null_on_divide_by_zero, + } + } +} + +impl AggregateUDFImpl for Kurtosis { + fn name(&self) -> &str { + &self.name + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> { + Ok(DataType::Float64) + } + + fn accumulator(&self, _acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> { + Ok(Box::new(KurtosisAccumulator::new( + self.null_on_divide_by_zero, + ))) + } + + // Fields ordered to match Spark's `[n, avg, m2, m3, m4]` buffer so that a + // Spark-produced Partial state can be merged into a Comet-produced Final + // (and vice versa) without a schema conversion. + fn state_fields(&self, _args: StateFieldsArgs) -> Result<Vec<FieldRef>> { + Ok(vec![ + Arc::new(Field::new( + format_state_name(&self.name, "n"), + DataType::Float64, + true, + )), + Arc::new(Field::new( + format_state_name(&self.name, "avg"), + DataType::Float64, + true, + )), + Arc::new(Field::new( + format_state_name(&self.name, "m2"), + DataType::Float64, + true, + )), + Arc::new(Field::new( + format_state_name(&self.name, "m3"), + DataType::Float64, + true, + )), + Arc::new(Field::new( + format_state_name(&self.name, "m4"), + DataType::Float64, + true, + )), + ]) + } + + fn default_value(&self, _data_type: &DataType) -> Result<ScalarValue> { + Ok(ScalarValue::Float64(None)) + } +} + +/// Online update for the first four central moments. Direct port of Spark's +/// `CentralMomentAgg.updateExpressionsDef` for `momentOrder = 4`. +#[inline] +fn kurtosis_update( Review Comment: Taking the second option you offered: recorded as a decision, not an oversight. In dcabe65d0 there is a comment at the `accumulator` site naming what it costs and why the vectorized version is not here: > No `GroupsAccumulator`: grouped `kurtosis` deliberately runs through DataFusion's generic `GroupsAccumulatorAdapter`, which costs one boxed `Accumulator` and a `ScalarValue` round trip per group per batch. This is a gap relative to the neighbouring central-moment aggregates [...] The vectorized version wants to land with `skewness`, since both are an `evaluate` over the same `[n, avg, m2, m3, m4]` state that `moments4_update` already maintains, and one flat-state accumulator should then serve all three. That is also why it pairs with the `welford.rs` move rather than being independent of it — building the flat-state accumulator once against shared moment math is a much better trade than building a kurtosis-only one now and a skewness-only one after. I'll add the same note to the PR description. ########## spark/src/test/resources/sql-tests/expressions/aggregate/kurtosis.sql: ########## @@ -0,0 +1,273 @@ +-- Licensed to the Apache Software Foundation (ASF) under one +-- or more contributor license agreements. See the NOTICE file +-- distributed with this work for additional information +-- regarding copyright ownership. The ASF licenses this file +-- to you under the Apache License, Version 2.0 (the +-- "License"); you may not use this file except in compliance +-- with the License. You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, +-- software distributed under the License is distributed on an +-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +-- KIND, either express or implied. See the License for the +-- specific language governing permissions and limitations +-- under the License. + +-- ConfigMatrix: parquet.enable.dictionary=false,true + +-- ============================================================ +-- Setup +-- ============================================================ + +statement +CREATE TABLE k_dbl(v double, grp string) USING parquet + +statement +INSERT INTO k_dbl VALUES + (-10.0, 'g1'), (-20.0, 'g1'), (100.0, 'g1'), (1000.0, 'g1'), + (1.0, 'g2'), (10.0, 'g2'), (100.0, 'g2'), (10.0, 'g2'), (1.0, 'g2'), + (42.0, 'g3'), + (NULL, 'g4'), (NULL, 'g4'), + (7.0, 'g5'), (7.0, 'g5'), (7.0, 'g5') + +statement +CREATE TABLE k_int(v int, grp string) USING parquet + +statement +INSERT INTO k_int VALUES + (1, 'g1'), (10, 'g1'), (100, 'g1'), (10, 'g1'), (1, 'g1'), + (NULL, 'g2'), (5, 'g2') + +statement +CREATE TABLE k_dec(v decimal(10,2), grp string) USING parquet + +statement +INSERT INTO k_dec VALUES + (1.50, 'g1'), (2.50, 'g1'), (3.50, 'g1'), (4.50, 'g1') + +statement +CREATE TABLE k_empty(v double) USING parquet + +statement +CREATE TABLE k_spark_ex1(v double) USING parquet + +statement +INSERT INTO k_spark_ex1 VALUES (-10.0), (-20.0), (100.0), (1000.0) + +statement +CREATE TABLE k_spark_ex2(v double) USING parquet + +statement +INSERT INTO k_spark_ex2 VALUES (1.0), (10.0), (100.0), (10.0), (1.0) + +statement +CREATE TABLE k_single(v double) USING parquet + +statement +INSERT INTO k_single VALUES (42.0) + +statement +CREATE TABLE k_const(v double) USING parquet + +statement +INSERT INTO k_const VALUES (7.0), (7.0), (7.0) + +statement +CREATE TABLE k_lit(x int) USING parquet + +statement +INSERT INTO k_lit VALUES (1) + +-- ============================================================ +-- Spark's own example: matches -0.7014368047529627. +-- ============================================================ + +query +SELECT kurtosis(v) FROM k_spark_ex1 + +-- Spark's second example: matches 0.19432323191699075. +query +SELECT kurtosis(v) FROM k_spark_ex2 + +-- ============================================================ +-- GROUP BY over doubles: covers a "normal" group (g1), a heavier +-- group (g2), a single-value group (g3, m2=0 => NULL by default), +-- an all-NULL group (g4 => NULL), and constants (g5, m2=0). +-- ============================================================ + +query +SELECT grp, kurtosis(v) FROM k_dbl GROUP BY grp ORDER BY grp + +-- ============================================================ +-- Global aggregate (no GROUP BY). +-- ============================================================ + +query +SELECT kurtosis(v) FROM k_dbl + +-- Empty table returns NULL. +query +SELECT kurtosis(v) FROM k_empty + +-- ============================================================ +-- Integer input: promoted to Double by Spark's ImplicitCastInputTypes. +-- ============================================================ + +query +SELECT grp, kurtosis(v) FROM k_int GROUP BY grp ORDER BY grp + +-- ============================================================ +-- Decimal input. +-- ============================================================ + +query +SELECT grp, kurtosis(v) FROM k_dec GROUP BY grp ORDER BY grp + +-- ============================================================ +-- Literal argument (constant folded; still exercises planning). +-- ============================================================ + +query +SELECT kurtosis(1.0) FROM k_lit + +query +SELECT kurtosis(NULL) FROM k_lit + +-- ============================================================ +-- Divide-by-zero cases under default (nullOnDivideByZero=true): +-- single-value and all-equal groups both yield NULL. See +-- kurtosis_legacy.sql for the `legacyStatisticalAggregate=true` +-- variant that returns NaN instead. +-- ============================================================ + +query +SELECT kurtosis(v) FROM k_single + +query +SELECT kurtosis(v) FROM k_const + +-- ============================================================ +-- FILTER (WHERE ...) — Partial only carries the filter. +-- ============================================================ + +query +SELECT grp, kurtosis(v) FILTER (WHERE v > 0) FROM k_dbl GROUP BY grp ORDER BY grp + +-- ============================================================ +-- Skewness is Spark's sibling in CentralMomentAgg; we don't +-- implement it here, so it should fall back. (This documents the +-- boundary; if we add skewness later, the expect_fallback +-- becomes a plain query.) +-- ============================================================ + +query expect_fallback(unsupported Spark aggregate function: skewness) Review Comment: Dropped in dcabe65d0, along with the comment block above it. You're right that it was the wrong place: the fixture's name gives no hint why a skewness assertion would be the thing that broke, and pinning another expression's absence as expected behaviour makes implementing it look like a regression. I removed the query rather than relocating it. A fixture named for skewness that exists only to assert skewness is unsupported would have the same problem in a thinner disguise — the natural home for that assertion is the skewness PR, where it turns into real coverage. ########## native/spark-expr/src/agg_funcs/kurtosis.rs: ########## @@ -0,0 +1,363 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Spark-compatible excess-kurtosis aggregate. +//! +//! Spark's `Kurtosis` is a `CentralMomentAgg` (`DeclarativeAggregate`) whose +//! intermediate buffer is `[n, avg, m2, m3, m4]` of Float64. This accumulator +//! mirrors that buffer exactly, using the same higher-order online update / +//! merge recurrences (Meng 2015) that `CentralMomentAgg` compiles into +//! catalyst expressions. Matching the wire format lets Spark's Partial and +//! Comet's Final (or vice versa) share intermediate state without a cast. +//! +//! Result formula (excess kurtosis, Fisher definition): +//! +//! * `n == 0` -> NULL +//! * `m2 == 0` -> NULL when `null_on_divide_by_zero`, else NaN +//! * otherwise -> `n * m4 / (m2 * m2) - 3.0` + +use std::mem::size_of; +use std::sync::Arc; + +use arrow::array::{ArrayRef, Float64Array}; +use arrow::datatypes::{DataType, Field, FieldRef}; +use datafusion::common::{downcast_value, Result, ScalarValue}; +use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs}; +use datafusion::logical_expr::Volatility::Immutable; +use datafusion::logical_expr::{Accumulator, AggregateUDFImpl, Signature}; +use datafusion::physical_expr::expressions::format_state_name; + +#[derive(Debug, PartialEq, Eq)] +pub struct Kurtosis { + name: String, + signature: Signature, + null_on_divide_by_zero: bool, +} + +impl std::hash::Hash for Kurtosis { + fn hash<H: std::hash::Hasher>(&self, state: &mut H) { + self.name.hash(state); + self.signature.hash(state); + self.null_on_divide_by_zero.hash(state); + } +} + +impl Kurtosis { + pub fn new(name: impl Into<String>, null_on_divide_by_zero: bool) -> Self { + Self { + name: name.into(), + signature: Signature::numeric(1, Immutable), + null_on_divide_by_zero, + } + } +} + +impl AggregateUDFImpl for Kurtosis { + fn name(&self) -> &str { + &self.name + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> { + Ok(DataType::Float64) + } + + fn accumulator(&self, _acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> { + Ok(Box::new(KurtosisAccumulator::new( + self.null_on_divide_by_zero, + ))) + } + + // Fields ordered to match Spark's `[n, avg, m2, m3, m4]` buffer so that a + // Spark-produced Partial state can be merged into a Comet-produced Final + // (and vice versa) without a schema conversion. + fn state_fields(&self, _args: StateFieldsArgs) -> Result<Vec<FieldRef>> { + Ok(vec![ + Arc::new(Field::new( + format_state_name(&self.name, "n"), + DataType::Float64, + true, + )), + Arc::new(Field::new( + format_state_name(&self.name, "avg"), + DataType::Float64, + true, + )), + Arc::new(Field::new( + format_state_name(&self.name, "m2"), + DataType::Float64, + true, + )), + Arc::new(Field::new( + format_state_name(&self.name, "m3"), + DataType::Float64, + true, + )), + Arc::new(Field::new( + format_state_name(&self.name, "m4"), + DataType::Float64, + true, + )), + ]) + } + + fn default_value(&self, _data_type: &DataType) -> Result<ScalarValue> { + Ok(ScalarValue::Float64(None)) + } +} + +/// Online update for the first four central moments. Direct port of Spark's +/// `CentralMomentAgg.updateExpressionsDef` for `momentOrder = 4`. +#[inline] +fn kurtosis_update( + n: f64, + avg: f64, + m2: f64, + m3: f64, + m4: f64, + value: f64, +) -> (f64, f64, f64, f64, f64) { + let new_n = n + 1.0; + let delta = value - 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 = m3 - 3.0 * delta_n * new_m2 + delta * (delta2 - delta_n2); + let new_m4 = m4 - 4.0 * delta_n * new_m3 - 6.0 * delta_n2 * new_m2 + + delta * (delta * delta2 - delta_n * delta_n2); + (new_n, new_avg, new_m2, new_m3, new_m4) +} + +/// Merge two partial states. Direct port of Spark's +/// `CentralMomentAgg.mergeExpressions` for `momentOrder = 4`. +#[inline] +#[allow(clippy::too_many_arguments)] +fn kurtosis_merge( + n1: f64, + avg1: f64, + m2_1: f64, + m3_1: f64, + m4_1: f64, + n2: f64, + avg2: f64, + m2_2: f64, + m3_2: f64, + m4_2: f64, +) -> (f64, f64, f64, f64, f64) { + let new_n = n1 + n2; + 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 = m3_1 + + m3_2 + + delta_n * delta_n * delta * n1 * n2 * (n1 - n2) + + 3.0 * delta_n * (n1 * m2_2 - n2 * m2_1); + let new_m4 = 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); + (new_n, new_avg, new_m2, new_m3, new_m4) +} + +#[derive(Debug)] +pub struct KurtosisAccumulator { + n: f64, + avg: f64, + m2: f64, + m3: f64, + m4: f64, + null_on_divide_by_zero: bool, +} + +impl KurtosisAccumulator { + pub fn new(null_on_divide_by_zero: bool) -> Self { + Self { + n: 0.0, + avg: 0.0, + m2: 0.0, + m3: 0.0, + m4: 0.0, + null_on_divide_by_zero, + } + } +} + +impl Accumulator for KurtosisAccumulator { + fn state(&mut self) -> Result<Vec<ScalarValue>> { + Ok(vec![ + ScalarValue::from(self.n), + ScalarValue::from(self.avg), + ScalarValue::from(self.m2), + ScalarValue::from(self.m3), + ScalarValue::from(self.m4), + ]) + } + + fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { + let arr = downcast_value!(&values[0], Float64Array).iter().flatten(); + for value in arr { + let (n, avg, m2, m3, m4) = + kurtosis_update(self.n, self.avg, self.m2, self.m3, self.m4, value); + self.n = n; + self.avg = avg; + self.m2 = m2; + self.m3 = m3; + self.m4 = m4; + } + Ok(()) + } + + fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> { + let ns = downcast_value!(states[0], Float64Array); + let avgs = downcast_value!(states[1], Float64Array); + let m2s = downcast_value!(states[2], Float64Array); + let m3s = downcast_value!(states[3], Float64Array); + let m4s = downcast_value!(states[4], Float64Array); + + for i in 0..ns.len() { + let n2 = ns.value(i); + if n2 == 0.0 { + // Empty partial state contributes nothing and would produce + // divide-by-zero garbage in `delta_n`; skip it. + continue; + } + let (n, avg, m2, m3, m4) = kurtosis_merge( + self.n, + self.avg, + self.m2, + self.m3, + self.m4, + n2, + avgs.value(i), + m2s.value(i), + m3s.value(i), + m4s.value(i), + ); + self.n = n; + self.avg = avg; + self.m2 = m2; + self.m3 = m3; + self.m4 = m4; + } + Ok(()) + } + + fn evaluate(&mut self) -> Result<ScalarValue> { + Ok(ScalarValue::Float64(if self.n == 0.0 { + None + } else if self.m2 == 0.0 { + if self.null_on_divide_by_zero { + None + } else { + Some(f64::NAN) + } + } else { + Some(self.n * self.m4 / (self.m2 * self.m2) - 3.0) Review Comment: Fixed in dcabe65d0. Good catch — this one slips through precisely because the guard above it looks like it already handles the case. Spark's guard is on `m2`, but the division is by `m2 * m2`. Your inputs give an `m2` of 5e-201, which is finite and non-zero so the `m2 === 0` branch does not fire, and whose square underflows to exactly 0. Spark's `Divide` then sees a zero divisor and applies **its** rule — which is the session's ANSI setting, a different switch from `nullOnDivideByZero` (that one comes from `legacyStatisticalAggregate`). Native code was doing a plain IEEE divide and returning NaN. `evaluate` now computes the divisor once and branches on it, so ANSI off gives NULL and ANSI on raises DIVIDE_BY_ZERO. That needed the ANSI flag, which was not previously plumbed for kurtosis, so there is a new `ansi_enabled` field on the `Kurtosis` proto set from `conf.ansiEnabled` in the serde. I noted on the proto field why it is separate from `null_on_divide_by_zero`, since having two divide-by-zero switches on one expression is otherwise very easy to misread. The regression test is `divisor_underflow_follows_spark_division_semantics`. It asserts the premise before the behaviour — ```rust assert_ne!(probe.m2, 0.0, "m2 must be non-zero for this case to bite"); assert_eq!(probe.m2 * probe.m2, 0.0, "m2 * m2 must underflow to zero"); ``` — so it cannot quietly stop testing anything if the moment math changes, and then covers NULL for both values of `null_on_divide_by_zero` with ANSI off, plus the DIVIDE_BY_ZERO error with ANSI on. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
