Skip to content

feat: support kurtosis aggregate - #4818

Open
andygrove wants to merge 3 commits into
apache:mainfrom
andygrove:feat/kurtosis-native-support
Open

feat: support kurtosis aggregate#4818
andygrove wants to merge 3 commits into
apache:mainfrom
andygrove:feat/kurtosis-native-support

Conversation

@andygrove

@andygrove andygrove commented Jul 3, 2026

Copy link
Copy Markdown
Member

Which issue does this PR close?

Closes #.

Rationale for this change

kurtosis is a standard SQL statistical aggregate and one of the last remaining CentralMomentAgg siblings that Comet didn't run natively (variance and stddev are already supported). Adding it lets queries using kurtosis stay in Comet's native path instead of falling back to Spark for the whole aggregate.

What changes are included in this PR?

  • Adds a Comet-owned Kurtosis UDAF (native/spark-expr/src/agg_funcs/kurtosis.rs) with a per-row KurtosisAccumulator. Intermediate state is [n, avg, m2, m3, m4] Float64, mirroring Spark's CentralMomentAgg (momentOrder = 4) wire format so a Spark-produced Partial and a Comet-produced Final can share bytes without conversion. Update and merge kernels are direct ports of Spark's updateExpressionsDef and mergeExpressions (Meng 2015 recurrence).
  • Threads nullOnDivideByZero through the proto so spark.sql.legacy.statisticalAggregate behaves the same on both engines: the default returns NULL when m2 == 0, legacy mode returns NaN.
  • Adds a Kurtosis protobuf message and wires it through the native planner.
  • Adds CometKurtosis in spark/src/main/scala/org/apache/comet/serde/aggregates.scala and registers it in QueryPlanSerde.aggrSerdeMap. supportsMixedPartialFinal is left at false to match the conservative policy already used by Variance and Stddev in the same file.
  • Documents the audit in docs/source/contributor-guide/expression-audits/agg_funcs.md; flips the support status in docs/source/user-guide/latest/expressions.md from planned to supported.
  • Window use (kurtosis(x) OVER (...)) still falls back because the Comet window path doesn't wire kurtosis today. Captured as an expect_fallback case rather than left as an implicit gap.
  • Threads ansiEnabled through the proto as well. Spark's evaluate expression divides by m2 * m2, and that Divide captures its eval mode from the session, so a divisor that underflows to zero while m2 itself is non-zero raises DIVIDE_BY_ZERO under ANSI and returns NULL otherwise. This is a separate switch from nullOnDivideByZero, which comes from legacyStatisticalAggregate and governs only the m2 == 0 branch.
  • The order-4 moment recurrences live in native/spark-expr/src/agg_funcs/welford.rs as moments4_update / moments4_merge, alongside the existing variance_* and covariance_* helpers, so all the central-moment algebra stays in one module and skewness can reuse it.
  • Grouped kurtosis is adapter-backed. There is no GroupsAccumulator, so grouped aggregation runs through DataFusion's generic GroupsAccumulatorAdapter at the cost of one boxed Accumulator and a ScalarValue round trip per group per batch. This is a deliberate gap relative to VarianceGroupsAccumulator / StddevGroupsAccumulator, recorded here and in a comment at the accumulator site rather than left implicit. The vectorized version should land with skewness: both are an evaluate over the same [n, avg, m2, m3, m4] state, so one flat-state accumulator can serve all three instead of being written twice.

Scaffolding produced by the implement-comet-expression skill; the audit-comet-expression skill drove the audit and produced the extra fallback and coverage-gap tests.

How are these changes tested?

  • Rust unit tests for KurtosisAccumulator covering empty group, single-value divide-by-zero in both nullOnDivideByZero modes, divisor underflow (1e-100 / 2e-100, where m2 is non-zero but m2 * m2 is not) under both ANSI settings, both of Spark's own ExpressionDescription examples (-0.7014368047529627 and 0.19432323191699075), and a two-partition state merge that reproduces the single-batch result.
  • New Comet SQL test at spark/src/test/resources/sql-tests/expressions/aggregate/kurtosis.sql running under ConfigMatrix: parquet.enable.dictionary=false,true. Covers Spark's documented examples, GROUP BY, global aggregate, empty table, integer/decimal/float/bigint inputs, literal argument, FILTER (WHERE ...), mixed with other aggregates, and an expect_fallback case for window use of kurtosis. Numerical-stress cases (NaN, Infinity, -Infinity, 1e15 magnitudes) use spark_answer_only mode.
  • Separate SQL file kurtosis_legacy.sql gated with Config: spark.sql.legacy.statisticalAggregate=true exercises the NaN path for single-value and all-equal groups.
  • ./mvnw test -Dsuites="org.apache.comet.CometSqlFileTestSuite kurtosis" -Dtest=none passes under both the default (Spark 3.5) and -Pspark-4.0 profiles.
  • cd native && cargo clippy --all-targets --workspace -- -D warnings passes.

Adds a Comet-owned `Kurtosis` UDAF and `CometKurtosis` serde so that
Spark's excess-kurtosis aggregate runs natively. The native accumulator
stores `[n, avg, m2, m3, m4]` `Float64` state to mirror Spark's
`CentralMomentAgg` (momentOrder = 4) buffer, and the update/merge
kernels are a direct port of Spark's `updateExpressionsDef` /
`mergeExpressions` expressions. `nullOnDivideByZero` is threaded
through so `spark.sql.legacy.statisticalAggregate` both branches
behave the same as Spark (NULL vs NaN when `m2 == 0`).

Window use falls back to Spark today: the window path doesn't wire the
Comet aggregate for kurtosis. Captured as an `expect_fallback` case in
the SQL test.

Scaffolding produced by the `implement-comet-expression` skill; audit
performed by the `audit-comet-expression` skill.
@andygrove andygrove added this to the 1.0.0 milestone Jul 3, 2026
@andygrove andygrove mentioned this pull request Jul 6, 2026
27 tasks
@andygrove andygrove modified the milestones: 1.0.0, 1.1.0 Jul 20, 2026

@mbutrovich mbutrovich left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

First pass, thanks @andygrove!

])
}

fn default_value(&self, _data_type: &DataType) -> Result<ScalarValue> {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

default_value returns ScalarValue::Float64(None), which is what the AggregateUDFImpl trait default already produces for a Float64 return type via ScalarValue::try_from(data_type) (datafusion/expr/src/udaf.rs:828-830). This override can be deleted.

This is the same redundancy already flagged on mode.rs:116 in #4782. Worth fixing in both so the pattern does not propagate to the next aggregate.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Removed in dcabe65. 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.

/// Online update for the first four central moments. Direct port of Spark's
/// `CentralMomentAgg.updateExpressionsDef` for `momentOrder = 4`.
#[inline]
fn kurtosis_update(

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There is no GroupsAccumulator here, so grouped kurtosis runs through DataFusion's generic GroupsAccumulatorAdapter, one boxed Accumulator and a ScalarValue round trip per group per batch.

That is a defensible starting point, but it is worth stating explicitly why, because the neighbouring central-moment aggregates in this same crate do have one: VarianceGroupsAccumulator (variance.rs:268) keeps flat Vec<f64> state and StddevGroupsAccumulator (stddev.rs:206) reuses it. 4817 in this same batch also ships a vectorized grouped accumulator and benchmarks it. Please either add one or record in the PR description that grouped kurtosis is adapter-backed, so the gap is a decision rather than an oversight.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Taking the second option you offered: recorded as a decision, not an oversight. In dcabe65 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.

}
}

/// Online update for the first four central moments. Direct port of Spark's

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

kurtosis_update and kurtosis_merge are free functions in this file, but welford.rs is exactly the module this crate uses to share moment recurrences across aggregates: it already holds variance_update, variance_merge, covariance_update, covariance_merge and finalize_moments (welford.rs:27-141), consumed by both variance.rs and covariance.rs.

The order-4 recurrence subsumes the order-2 one, and skewness is the obvious next sibling and needs [n, avg, m2, m3] from the same recurrence. Putting these two functions in welford.rs alongside the existing ones keeps all the moment math in one place and means skewness is a momentOrder = 3 evaluate on top of shared state rather than a third copy of the same algebra.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Moved in dcabe65. 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.

-- becomes a plain query.)
-- ============================================================

query expect_fallback(unsupported Spark aggregate function: skewness)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

expect_fallback(unsupported Spark aggregate function: skewness) pins Comet's current lack of skewness support as expected behavior in the kurtosis fixture.

That is the wrong place for it. It couples an unrelated expression's status to this file, and when skewness lands, this assertion fails in a fixture whose name gives no hint why. The comment above at :161 already anticipates this. Either drop the query or move it to a fixture named for skewness.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Dropped in dcabe65, 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.

@andygrove andygrove added enhancement New feature or request area:aggregation Hash aggregates, aggregate expressions labels Sep 6, 2026
Some(f64::NAN)
}
} else {
Some(self.n * self.m4 / (self.m2 * self.m2) - 3.0)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For inputs 1e-100 and 2e-100, m2 is nonzero but m2 * m2 underflows to zero. The native aggregate returns NaN, while Spark returns NULL with ANSI off and raises DIVIDE_BY_ZERO with ANSI on. Could you preserve Spark's division semantics here and add a regression test?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in dcabe65. 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 —

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.

@rich7420

rich7420 commented Sep 7, 2026

Copy link
Copy Markdown
Contributor

@andygrove thanks for the patch!!!

Also moves the order-4 moment recurrences into welford.rs alongside the
existing variance/covariance ones, drops the redundant default_value
override, and removes the skewness fallback assertion from the kurtosis
fixture.
@github-actions github-actions Bot added the area:expressions Expression evaluation label Sep 9, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:aggregation Hash aggregates, aggregate expressions area:expressions Expression evaluation enhancement New feature or request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants