[SPARK-58378][SQL] Fuse compatible approximate percentile sketches

[sunchao]

Why are the changes needed?

This addresses SPARK-58378.

An approximate percentile sketch summarizes a distribution. Once the sketch exists,
Spark can read p50, p90, p95, or any other requested percentile from the same
state. Nevertheless, a query that writes those percentiles as separate scalar
aggregates currently constructs and maintains a separate sketch for each one.

For example:

SELECT
  service,
  percentile_approx(latency_ms, 0.50, 10000) AS p50,
  percentile_approx(latency_ms, 0.90, 10000) AS p90,
  percentile_approx(latency_ms, 0.95, 10000) AS p95
FROM request_metrics
GROUP BY service;

All three aggregates receive the same values and use the same accuracy. The
percentile only affects how the completed sketch is queried. Despite that, Spark
currently allocates three aggregate buffers per group, inserts every qualifying
input into all three, serializes three partial sketches, and merges three sketches
during final aggregation.

Callers can avoid that work by rewriting the query to use
percentile_approx(latency_ms, array(0.50, 0.90, 0.95), 10000) and extracting
individual array elements. However, requiring that rewrite is awkward for existing
SQL, generated reports, named scalar output columns, and expressions that consume
individual percentile values.

What changes were proposed in this pull request?

Teach Catalyst to recognize scalar approximate percentile aggregates that describe
the same distribution, calculate that distribution once, and extract each original
scalar result from the shared array-valued aggregate.

For the example above, the resulting aggregation is equivalent to:

percentile_approx(latency_ms, array(0.50, 0.90, 0.95), 10000)

Catalyst returns elements 0, 1, and 2 under the original p50, p90, and
p95 expressions. Callers do not change their query. Both partial and final
aggregation maintain one percentile sketch instead of three.

Fusion is intentionally limited to deterministic scalar aggregates with the same
input expression, evaluated accuracy, aggregation mode, distinctness, and filter.
Inputs and filters are compared structurally so that floating-point evaluation and
ANSI overflow behavior do not change. Accuracy is compared after evaluation, which
remains correct when ConstantFolding is excluded. Existing array-valued
aggregates, incompatible distributions, and Structured Streaming aggregates are
left unchanged; excluding streaming preserves existing checkpoint state schemas.

Does this PR introduce any user-facing change?

No SQL syntax, public API, output schema, or percentile result changes. Eligible
batch queries use fewer percentile aggregation buffers and can perform less sketch
update, serialization, shuffle, and merge work.

How was this patch tested?

Added Catalyst optimizer tests, end-to-end physical-plan and result tests, a
Spark 5 TIME regression, and streaming checkpoint recovery tests. The regression
coverage includes distinct and filtered aggregates, differently associated
floating-point expressions, ANSI overflow, null and empty inputs, grouping sets,
decimal and temporal types, aliases, repeated percentile values, pre-existing
array-valued aggregates, and differently evaluated accuracy expressions.

All 54 focused tests and all three Scala style checks pass:

build/sbt \
  'catalyst/testOnly org.apache.spark.sql.catalyst.optimizer.CombineApproximatePercentilesSuite' \
  'sql/testOnly org.apache.spark.sql.ApproximatePercentileQuerySuite' \
  'sql/testOnly org.apache.spark.sql.streaming.StreamingAggregationSuite -- -z "approximate percentiles preserve existing streaming checkpoints"' \
  'catalyst/scalastyle' \
  'sql/scalastyle' \
  'sql/Test/scalastyle'

Was this patch authored or co-authored using generative AI tooling?

Generated-by: OpenAI Codex (GPT-5)

[sunchao]

cc @viirya @cloud-fan @dongjoon-hyun @peter-toth

[viirya]

Nice optimization, and the correctness reasoning holds up under scrutiny. I traced the sketch semantics and the rewrite end to end.

The core equivalence is sound because the digest is fully decoupled from the requested percentile: ApproximatePercentile is deterministic (no Random/seed, and QuantileSummaries — Greenwald-Khanna — has no randomness), so three scalar percentile_approx(x, {0.5,0.9,0.95}, acc) build three bit-identical digests, and querying one shared digest three times equals querying three identical digests once each. The getPercentiles(array(...)) path preserves input order, so array element index maps exactly to the percentile at expressions(index), and GetArrayItem(combined, index) extracts it. I checked the GK algorithm is order-sensitive in its intermediate state, but that's a non-issue here: fusion doesn't change the input sequence, partitioning, or the partial/final merge topology any buffer sees — it just collapses N identical computations into one — so the fused digest is bit-identical to the originals, not merely within the error bound.

A few specific things I verified:

• Empty/null input: eval returns null for both the scalar and array cases (the whole array is null, not an array of nulls), and GetArrayItem(null, i) short-circuits to null via BinaryExpression.eval, so the fused result matches the original scalar null.
• Result type/nullability: the original scalar type is child.dataType; the fused ArrayType(child.dataType) unwrapped by GetArrayItem gives back child.dataType, unchanged.
• Idempotency: after fusion the inner aggregate is array-valued, so the percentageExpression.dataType == DoubleType guard excludes it on a second pass.
• The child.deterministic / filter.deterministic guard is exactly the right scope. The sketch itself is deterministic, so the only way fusion could change results is a non-deterministic input expression (e.g. percentile_approx(x + rand(), ...)), where collapsing N aggregates into one would evaluate rand() once instead of N times. Excluding those is the precise condition, and structural (not semantic) comparison of inputs/filters keeps FP-associativity and ANSI-overflow behavior intact — it can only miss fusions, never make an unsafe one.

Excluding streaming to preserve checkpoint value schemas is the right call, and the optimizer test suite is thorough — it targets each of these risk areas directly (duplicate/non-monotonic order, nested expressions, structural input/filter equivalence, differently-associated FP, distinct, empty/null, existing array percentiles, streaming).

One non-blocking thought: the rule is a standalone Once batch, so it won't re-fire if a later fixed-point batch ever exposes newly-fusible percentiles. That's almost certainly fine (percentiles rarely materialize late in optimization) and cheaper — just confirming Once is the intended choice rather than folding it into an existing fixed-point batch.

LGTM.

[uros-b]

LGTM, thank you @sunchao and @viirya!

[dongjoon-hyun]

Nice optimization, @sunchao . +1, LGTM for code.

@viirya's review already covers the logical-plan-level correctness (digest determinism, null/empty inputs, type/nullability, idempotency, the deterministic guard, and the streaming exclusion), so this comment focuses only on points not raised yet.

1. The physical-dedup interaction deserves an inline comment

The most subtle part of this rule is the PhysicalCompatibilityKey guard, and its rationale isn't written down anywhere in the code. PhysicalAggregation dedups aggregate expressions semantically via EquivalentExpressions (sql/catalyst/.../planning/patterns.scala), so two structurally-different-but-canonically-equal aggregates already share one buffer today. Fusion must not change those dedup pairings — e.g. fusing one group into an array while its canonical twin stays scalar (or letting two fused arrays newly dedup against each other) would change which expression structure gets evaluated, and thus FP-associativity / ANSI-overflow behavior. The canonical-key guard blocks exactly those cases, and I couldn't find a gap in it — but today the "why" only lives implicitly in the e2e tests. A short comment in combine() referencing PhysicalAggregation's semantic dedup would protect this from being "simplified" away later.

Two smaller notes in the same area:

• accuracyExpression.eval() is safe here because ApproximatePercentile.checkInputDataTypes already evaluates and validates accuracy (non-null, range) at analysis time — worth a one-line comment as well.
• The physical key ignores the percentage expression, so fusion is also skipped in cases where no physical dedup could actually occur (canonically-equal groups whose percentages all differ). That's harmlessly conservative; just noting the intent in the comment would suffice.

2. Document the batch-ordering constraints

The placement is correct — after the operator-optimization fixed point (so accuracies are constant-folded before comparison on the normal path) and before RewriteDistinctAggregates (so a fused distinct array percentile is rewritten properly). Neither constraint is recorded at the batch declaration, so a future batch reshuffle could silently break one of them. A one-line comment would help.

[peter-toth]

Thanks for the PR, @sunchao!

I reviewed this independently before reading the existing approvals, and I reached the same conclusion on the core equivalence: the digest is fully decoupled from the requested percentage, so collapsing N scalar percentile_approx into one array-valued aggregate and projecting with GetArrayItem is sound. I checked the one thing that reasoning rests on — that QuantileSummaries.query(Seq(p1..pn)) returns exactly what N separate query(Seq(pi)) calls return — since the batched path carries (index, minRank) forward across ascending percentiles while the scalar path restarts at index 0 each time. Exhaustive brute force over GK-valid summaries plus an end-to-end differential fuzz (distinct, filter, grouping sets, rollup/cube, ANSI, nulls, accuracies 3..10000) found no divergence. I also confirmed the streaming test is effective: dropping !aggregate.isStreaming makes it fail. My only real finding is a regression case in the fusion gate.

Blocking

• 1. Fusion of duplicate percentages adds work: the gate only requires 2+ compatible aggregates, not 2+ distinct percentages, so percentile_approx(col, 0.5) twice fuses into array(0.5, 0.5) — on base PhysicalAggregation already deduped those to a single digest queried once, so this is a pure regression. Gate on distinct evaluated percentages. [inline: sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/CombineApproximatePercentiles.scala:103]

Non-blocking

• 2. Reused AggregateExpression instance gets the last ordinal: replacements is keyed by resultId, so when the same instance appears twice in a group the second put overwrites the first. It is currently harmless — a repeated instance always carries the same percentage — but it is an invariant nobody states, and it silently depends on percentages being positionally aligned with expressions. [inline: sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/CombineApproximatePercentiles.scala:121]

I also agree with @dongjoon-hyun's two documentation asks (the PhysicalCompatibilityKey rationale and the batch-ordering constraints) and with @viirya's Once-vs-fixed-point question — no need to restate them.

[peter-toth]

Finding 1. expressions.sizeCompare(1) > 0 fires on 2+ compatible aggregates regardless of whether their percentages differ, so a query repeating one percentage gets more work than on base, not less.

PhysicalAggregation already dedups semantically-equal aggregates via EquivalentExpressions (sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/planning/patterns.scala:300), so two identical scalar percentiles share one digest and query it once today. After fusion they become one digest queried N times. Measured on the PR head, counting ApproximatePercentile buffers and percentages per ObjectHashAggregateExec:

SELECT percentile_approx(col, 0.5D) a, percentile_approx(col, 0.5D) b FROM t
  base : digests=1  percentagesQueried=1
  fused: digests=1  percentagesQueried=2     <- pure loss

SELECT percentile_approx(col, 0.5D) a, percentile_approx(col, 0.5D) b,
       percentile_approx(col, 0.5D) c FROM t
  base : digests=1  percentagesQueried=1
  fused: digests=1  percentagesQueried=3     <- pure loss

The extra getPercentiles work is small, but the plan also gets wider (array(0.5, 0.5, 0.5) plus three GetArrayItem) for zero gain, and this is a default-on rule. The mixed case is still a win and should keep firing:

SELECT percentile_approx(col, 0.5D) a, percentile_approx(col, 0.5D) b,
       percentile_approx(col, 0.9D) c FROM t
  base : digests=2  percentagesQueried=2
  fused: digests=1  percentagesQueried=3     <- still worth it

So the condition wants distinct evaluated percentages, not distinct expressions — array(0.5, 0.5) and the 0.25 + 0.25 / 0.5 pair are both no-gain. Since checkInputDataTypes has already validated that every percentage is foldable and non-null, evaluating is safe here:

expressions.sizeCompare(1) > 0 && expressions.map { expression =>
  expression.aggregateFunction
    .asInstanceOf[ApproximatePercentile]
    .percentageExpression.eval()
}.distinct.sizeCompare(1) > 0 && expressions.forall { expression =>

Note this must stay a gate on whether to fuse, not a dedup of the array contents: preserve duplicate and non-monotonic percentile order depends on the fused array keeping one slot per original aggregate.

Worth a test for the all-duplicates case — the suite covers repeated percentages only in preserve duplicate and non-monotonic percentile order, where a third distinct percentage makes fusion profitable anyway.

[peter-toth]

Finding 2. replacements is keyed by resultId, so if the same AggregateExpression instance is collected twice into one group, the second put overwrites the first and both references resolve to the last ordinal.

This happens whenever a Column is reused, e.g. df.agg(p50.as("a"), p50.as("b"), p90.as("c")) — aggregate.aggregateExpressions.foreach(_.foreach ...) walks each occurrence, so the buffer is [p50, p50, p90] while replacements maps that one resultId to ordinal 1 only:

Aggregate [percentile_approx(col, [0.5,0.5,0.9], ...)[1] AS a,
           percentile_approx(col, [0.5,0.5,0.9], ...)[1] AS b,
           percentile_approx(col, [0.5,0.5,0.9], ...)[2] AS c]

Results are correct today, and I could not construct a wrong answer: a repeated resultId implies the identical instance, hence the identical percentage, so every ordinal it could have pointed at holds the same value. But that is an unstated invariant that the positional expressions/percentages alignment depends on, and it would break quietly if a future change ever derived the array from anything other than one-slot-per-occurrence. Deduping the group by resultId before building the array makes the intent explicit and drops the dead slot:

}.foreach { case (_, allExpressions) =>
  val expressions = allExpressions.distinctBy(_.resultId)

(If finding 1 lands, this dedup should happen before the distinct-percentage check so the reused instance doesn't count twice toward "2+ compatible aggregates".)