This is an automated email from the ASF dual-hosted git repository.
sunchao pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new 2e8911d7b362 [SPARK-58378][SQL] Fuse compatible approximate percentile
sketches
2e8911d7b362 is described below
commit 2e8911d7b3621bd30213c87698214372566c4682
Author: Chao Sun <[email protected]>
AuthorDate: Thu Aug 6 15:27:12 2026 -0700
[SPARK-58378][SQL] Fuse compatible approximate percentile sketches
### Why are the changes needed?
This addresses
[SPARK-58378](https://issues.apache.org/jira/browse/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:
```sql
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 PR?
Add an opt-in Catalyst optimization that recognizes scalar approximate
percentile
aggregates describing the same distribution, calculates that distribution
once,
and extracts each original scalar result from the shared array-valued
aggregate.
The optimization is controlled by the public SQL configuration
`spark.sql.optimizer.combineApproximatePercentiles.enabled`, which defaults
to
`false`. Existing queries are unchanged until a user explicitly enables it.
For the example above, the resulting aggregation is equivalent to:
```sql
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 only
runs when
it removes a sketch that Spark would not already share: repeating the same
scalar
percentile therefore does not add unnecessary array construction or
lookups. It
also runs after subplan merging, which can bring independently written
scalar
subqueries together into a newly eligible aggregate.
The fused percentage array also preserves the identity of the original
scalar
aggregates. Otherwise, retaining only the first aggregate's accuracy
expression
could make previously distinct plans appear identical; for example, `100`
and
`100L` evaluate to the same accuracy but remain different Catalyst literals.
`PercentileFusionArray` wraps a `PercentileFusionIdentity` containing those
original aggregate structures, and `ConstantFolding` explicitly leaves that
internal expression intact. This prevents fusion from introducing incorrect
exchange or subquery reuse.
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 retain their original identity, incompatible distributions are
not
combined, and Structured Streaming aggregates remain unchanged to preserve
existing checkpoint state schemas.
### Does this PR introduce _any_ user-facing change?
Yes. It introduces the public SQL configuration
`spark.sql.optimizer.combineApproximatePercentiles.enabled`, which defaults
to
`false`. The default preserves existing query behavior and lets users
evaluate
the optimization explicitly before enabling it more broadly.
When enabled, eligible batch queries use fewer percentile aggregation
buffers and
can perform less sketch update, serialization, shuffle, and merge work.
Preserving
the original aggregate structures can also prevent an existing, unsafe
exchange
or subquery reuse between canonically equivalent but differently evaluated
expressions; in those cases, enabling fusion can correct a result that was
already
wrong without the optimization. SQL syntax, public APIs, and output schemas
are
unchanged.
### How was this PR tested?
Added focused Catalyst optimizer tests, end-to-end physical-plan and result
tests,
and streaming checkpoint-recovery coverage. The regressions cover the
opt-in flag
and default-off behavior, duplicate percentages and aggregate identities,
structural input/filter matching, ANSI overflow, evaluated accuracy and
percentage
collisions, existing array aggregates, mixed DISTINCT/non-DISTINCT groups,
DISTINCT rewriting and removal, scalar subqueries merged into a shared
aggregate,
exchange/subquery reuse with differently typed accuracy literals,
compressed and
merged low-accuracy sketches across multiple partitions, empty inputs, and
streaming checkpoint compatibility across both existing state-format
versions.
```text
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)
Closes #57576 from
sunchao/agent/fuse-compatible-approximate-percentile-sketches.
Authored-by: Chao Sun <[email protected]>
Signed-off-by: Chao Sun <[email protected]>
(cherry picked from commit 316c6efe053173d044eaeef7b8945e4342bd4de0)
Signed-off-by: Chao Sun <[email protected]>
---
.../optimizer/CombineApproximatePercentiles.scala | 225 ++++++++++
.../spark/sql/catalyst/optimizer/Optimizer.scala | 4 +
.../spark/sql/catalyst/optimizer/expressions.scala | 3 +
.../sql/catalyst/rules/RuleIdCollection.scala | 1 +
.../org/apache/spark/sql/internal/SQLConf.scala | 9 +
.../CombineApproximatePercentilesSuite.scala | 304 ++++++++++++++
.../spark/sql/execution/SparkOptimizer.scala | 1 +
.../sql/ApproximatePercentileQuerySuite.scala | 466 ++++++++++++++++++++-
.../sql/streaming/StreamingAggregationSuite.scala | 32 ++
9 files changed, 1043 insertions(+), 2 deletions(-)
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/CombineApproximatePercentiles.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/CombineApproximatePercentiles.scala
new file mode 100644
index 000000000000..fb6f50627840
--- /dev/null
+++
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/CombineApproximatePercentiles.scala
@@ -0,0 +1,225 @@
+/*
+ * 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.
+ */
+
+package org.apache.spark.sql.catalyst.optimizer
+
+import scala.collection.mutable
+
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{AttributeReference,
Expression, ExprId, GetArrayItem, LeafExpression, Literal, NamedExpression}
+import
org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression,
AggregateMode, ApproximatePercentile}
+import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback
+import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LogicalPlan}
+import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.catalyst.trees.TreePattern.AGGREGATE
+import org.apache.spark.sql.catalyst.util.GenericArrayData
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{ArrayType, DoubleType}
+
+private[optimizer] case class PercentileFusionIdentity(
+ aggregateFunctions: Seq[Expression],
+ mode: AggregateMode,
+ isDistinct: Boolean,
+ filter: Option[Expression],
+ percentageBits: Seq[Long])
+
+/**
+ * Foldable percentage array that retains the original scalar aggregate
structures in equality.
+ *
+ * Fusion removes those structures from the physical aggregate. Keeping them
here prevents
+ * subquery or exchange reuse from equating plans that were distinct before
fusion.
+ */
+private[optimizer] case class PercentileFusionArray(identity:
PercentileFusionIdentity)
+ extends LeafExpression with CodegenFallback {
+ override def foldable: Boolean = true
+ override def nullable: Boolean = false
+ override def dataType: ArrayType = ArrayType(DoubleType, containsNull =
false)
+
+ private lazy val value = new GenericArrayData(
+ identity.percentageBits.map(java.lang.Double.longBitsToDouble))
+ private lazy val literal = Literal(value, dataType)
+
+ override def eval(input: InternalRow): Any = value
+ override def toString: String = literal.toString
+ override def sql: String = literal.sql
+}
+
+/**
+ * Combines scalar approximate percentiles that can share the same percentile
digest.
+ *
+ * An approximate percentile digest depends on its input, accuracy, filter,
distinctness, and
+ * aggregate mode, but not on the percentile requested from the completed
digest. Consequently,
+ * compatible scalar percentiles can be calculated by one array-valued
aggregate and projected
+ * back to their original scalar outputs.
+ *
+ * Inputs and filters must retain their original expression structure so that
floating-point
+ * evaluation and ANSI overflow behavior are preserved. Streaming aggregates
are left unchanged
+ * to preserve the value schemas of existing checkpoints.
+ */
+object CombineApproximatePercentiles extends Rule[LogicalPlan] {
+
+ private case class CompatibilityKey(
+ child: Expression,
+ accuracy: Long,
+ mode: AggregateMode,
+ isDistinct: Boolean,
+ filter: Option[Expression])
+
+ private case class PhysicalCompatibilityKey(
+ child: Expression,
+ percentage: Expression,
+ accuracy: Expression,
+ mode: AggregateMode,
+ isDistinct: Boolean,
+ filter: Option[Expression])
+
+ private def structurallyNormalize(
+ expression: Expression,
+ inputOrdinals: scala.collection.Map[ExprId, Int]): Expression =
expression.transformUp {
+ case attribute: AttributeReference =>
+ inputOrdinals.get(attribute.exprId) match {
+ case Some(ordinal) => AttributeReference("none",
attribute.dataType)(ExprId(ordinal))
+ case None => attribute
+ }
+ }
+
+ private def physicalCompatibilityKey(
+ key: CompatibilityKey,
+ percentile: ApproximatePercentile): PhysicalCompatibilityKey =
PhysicalCompatibilityKey(
+ key.child.canonicalized,
+ percentile.percentageExpression.canonicalized,
+ percentile.accuracyExpression.canonicalized,
+ key.mode,
+ key.isDistinct,
+ key.filter.map(_.canonicalized))
+
+ private def hasSafePhysicalFusion(
+ expressions: scala.collection.Iterable[AggregateExpression]): Boolean = {
+ val physicalGroups = expressions.groupBy(_.canonicalized)
+ // PhysicalAggregation already shares a digest within each canonical
group. Fusion must both
+ // remove a digest and preserve cases where canonical percentages evaluate
differently.
+ physicalGroups.sizeCompare(1) > 0 && physicalGroups.values.forall { group
=>
+ val percentages = group.iterator.map { expression =>
+ expression.aggregateFunction
+ .asInstanceOf[ApproximatePercentile]
+ .percentageExpression
+ .eval()
+ }
+ val first = percentages.next()
+ percentages.forall(_ == first)
+ }
+ }
+
+ override def apply(plan: LogicalPlan): LogicalPlan = {
+ if (!conf.getConf(SQLConf.COMBINE_APPROXIMATE_PERCENTILES_ENABLED)) return
plan
+
+ plan.transformUpWithPruning(_.containsPattern(AGGREGATE), ruleId) {
+ case aggregate: Aggregate if aggregate.resolved &&
!aggregate.isStreaming =>
+ combine(aggregate)
+ }
+ }
+
+ private def combine(aggregate: Aggregate): Aggregate = {
+ val compatible = mutable.LinkedHashMap.empty[
+ CompatibilityKey, mutable.ArrayBuffer[AggregateExpression]]
+ // PhysicalAggregation deduplicates semantically equivalent aggregates.
Track every logical
+ // key that shares a physical key so fusion does not change that existing
deduplication.
+ val physicalCompatibilityKeys = mutable.HashMap.empty[
+ PhysicalCompatibilityKey, mutable.HashSet[CompatibilityKey]]
+
+ aggregate.aggregateExpressions.foreach(_.foreach {
+ case expression @ AggregateExpression(
+ percentile: ApproximatePercentile, mode, isDistinct, filter, _)
+ if percentile.child.deterministic &&
+ filter.forall(_.deterministic) =>
+ val key = CompatibilityKey(
+ percentile.child,
+ // Analysis already validates that accuracy is foldable, non-null,
and in range.
+ percentile.accuracyExpression.eval().asInstanceOf[Number].longValue,
+ mode,
+ isDistinct,
+ filter)
+ physicalCompatibilityKeys.getOrElseUpdate(
+ physicalCompatibilityKey(key, percentile),
+ mutable.HashSet.empty) += key
+ if (percentile.percentageExpression.dataType == DoubleType) {
+ compatible.getOrElseUpdate(key, mutable.ArrayBuffer.empty) +=
expression
+ }
+ case _ =>
+ })
+
+ val replacements = mutable.HashMap.empty[ExprId, (AggregateExpression,
Int)]
+ lazy val inputOrdinals = {
+ val ordinals = mutable.HashMap.empty[ExprId, Int]
+ aggregate.child.output.zipWithIndex.foreach { case (attribute, ordinal)
=>
+ ordinals.getOrElseUpdate(attribute.exprId, ordinal)
+ }
+ ordinals
+ }
+ compatible.iterator.map { case (key, expressions) =>
+ key -> expressions.distinctBy(_.resultId)
+ }.filter { case (key, expressions) =>
+ hasSafePhysicalFusion(expressions) && expressions.forall { expression =>
+ val percentile =
expression.aggregateFunction.asInstanceOf[ApproximatePercentile]
+ val physicalKey = physicalCompatibilityKey(key, percentile)
+ // OptimizeOneRowPlan can erase DISTINCT after fusion. Across
distinctness boundaries,
+ // canonical matches are safe only when their original inputs and
filters also match.
+ physicalCompatibilityKeys(physicalKey).sizeCompare(1) == 0 &&
+ physicalCompatibilityKeys
+ .get(physicalKey.copy(isDistinct = !physicalKey.isDistinct))
+ .forall(_.forall(other => other.child == key.child && other.filter
== key.filter))
+ }
+ }.foreach { case (key, expressions) =>
+ val first = expressions.head
+ val percentile =
first.aggregateFunction.asInstanceOf[ApproximatePercentile]
+ val percentages = expressions.map { expression =>
+ expression.aggregateFunction
+ .asInstanceOf[ApproximatePercentile]
+ .percentageExpression
+ }
+ val percentageValues =
percentages.map(_.eval().asInstanceOf[Double]).toSeq
+ val identity = PercentileFusionIdentity(
+ expressions.map { expression =>
+ structurallyNormalize(expression.aggregateFunction, inputOrdinals)
+ }.toSeq,
+ key.mode,
+ key.isDistinct,
+ key.filter.map(structurallyNormalize(_, inputOrdinals)),
+ percentageValues.map(java.lang.Double.doubleToRawLongBits))
+ val combinedFunction = percentile.copy(percentageExpression =
PercentileFusionArray(identity))
+ combinedFunction.copyTagsFrom(percentile)
+ val combined = first.copy(
+ aggregateFunction = combinedFunction, resultId =
NamedExpression.newExprId)
+ expressions.zipWithIndex.foreach { case (expression, index) =>
+ replacements.put(expression.resultId, (combined, index))
+ }
+ }
+
+ if (replacements.isEmpty) {
+ aggregate
+ } else {
+ val rewrittenExpressions = aggregate.aggregateExpressions.map {
expression =>
+ expression.transformUp {
+ case original: AggregateExpression if
replacements.contains(original.resultId) =>
+ val (combined, index) = replacements(original.resultId)
+ GetArrayItem(combined, Literal(index), failOnError = false)
+ }.asInstanceOf[NamedExpression]
+ }
+ aggregate.copy(aggregateExpressions = rewrittenExpressions)
+ }
+ }
+}
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala
index c673c08fba15..25edd59412e9 100644
---
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala
+++
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala
@@ -263,6 +263,10 @@ abstract class Optimizer(catalogManager: CatalogManager)
Batch("Eliminate Sorts", Once,
EliminateSorts,
RemoveRedundantSorts),
+ // Run after operator optimization normally folds accuracy expressions and
before
+ // RewriteDistinctAggregates so fused distinct percentiles are rewritten
correctly.
+ Batch("Combine Approximate Percentiles", Once,
+ CombineApproximatePercentiles),
Batch("Decimal Optimizations", fixedPoint,
DecimalAggregates),
// This batch must run after "Decimal Optimizations", as that one may
change the
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/expressions.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/expressions.scala
index 13fa3b09c23c..580c92f6eed5 100644
---
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/expressions.scala
+++
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/expressions.scala
@@ -84,6 +84,9 @@ object ConstantFolding extends Rule[LogicalPlan] {
// object and running eval unnecessarily.
case l: Literal => l
+ // This foldable expression carries planning identity that must survive
later optimizer batches.
+ case p: PercentileFusionArray => p
+
case Size(c: CreateArray, _) if c.children.forall(hasNoSideEffect) =>
Literal(c.children.length)
case Size(c: CreateMap, _) if c.children.forall(hasNoSideEffect) =>
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/rules/RuleIdCollection.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/rules/RuleIdCollection.scala
index d2c1ab9a8bc2..9221637d4db5 100644
---
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/rules/RuleIdCollection.scala
+++
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/rules/RuleIdCollection.scala
@@ -127,6 +127,7 @@ object RuleIdCollection {
"org.apache.spark.sql.catalyst.optimizer.CollapseRepartition" ::
"org.apache.spark.sql.catalyst.optimizer.CollapseWindow" ::
"org.apache.spark.sql.catalyst.optimizer.ColumnPruning" ::
+ "org.apache.spark.sql.catalyst.optimizer.CombineApproximatePercentiles"
::
"org.apache.spark.sql.catalyst.optimizer.CombineConcats" ::
"org.apache.spark.sql.catalyst.optimizer.CombineFilters" ::
"org.apache.spark.sql.catalyst.optimizer.CombineTypedFilters" ::
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
index cce753a68fff..16175a1a8c18 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
@@ -983,6 +983,15 @@ object SQLConf {
.booleanConf
.createWithDefault(false)
+ val COMBINE_APPROXIMATE_PERCENTILES_ENABLED =
+ buildConf("spark.sql.optimizer.combineApproximatePercentiles.enabled")
+ .doc("When true, combines compatible scalar approximate percentile
aggregates into a " +
+ "single array-valued aggregate so they share one percentile digest.")
+ .version("5.0.0")
+ .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE)
+ .booleanConf
+ .createWithDefault(false)
+
val EXPRESSION_PROJECTION_CANDIDATE_LIMIT =
buildConf("spark.sql.optimizer.expressionProjectionCandidateLimit")
.doc("The maximum number of the candidate of output expressions whose
alias are replaced." +
diff --git
a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/CombineApproximatePercentilesSuite.scala
b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/CombineApproximatePercentilesSuite.scala
new file mode 100644
index 000000000000..96bf9c7db256
--- /dev/null
+++
b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/CombineApproximatePercentilesSuite.scala
@@ -0,0 +1,304 @@
+/*
+ * 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.
+ */
+
+package org.apache.spark.sql.catalyst.optimizer
+
+import org.apache.spark.sql.catalyst.analysis.FunctionRegistry
+import org.apache.spark.sql.catalyst.dsl.expressions._
+import org.apache.spark.sql.catalyst.dsl.plans._
+import org.apache.spark.sql.catalyst.expressions.{Alias, CreateArray,
Expression, GetArrayItem, Literal}
+import
org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression,
ApproximatePercentile}
+import org.apache.spark.sql.catalyst.plans.PlanTest
+import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LocalRelation,
LogicalPlan}
+import org.apache.spark.sql.catalyst.rules.RuleExecutor
+import org.apache.spark.sql.catalyst.util.ArrayData
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.IntegerType
+
+class CombineApproximatePercentilesSuite extends PlanTest {
+
+ private object Optimize extends RuleExecutor[LogicalPlan] {
+ override val batches: Seq[Batch] =
+ Batch("Combine Approximate Percentiles", Once,
CombineApproximatePercentiles) :: Nil
+ }
+
+ private val relation = LocalRelation($"value".int, $"other".int,
$"group".int)
+ private val value = relation.output(0)
+ private val other = relation.output(1)
+ private val group = relation.output(2)
+
+ private def percentile(
+ child: Expression,
+ percentage: Double,
+ accuracy: Int = 10000,
+ isDistinct: Boolean = false,
+ filter: Option[Expression] = None): AggregateExpression = {
+ new ApproximatePercentile(child, Literal(percentage), Literal(accuracy))
+ .toAggregateExpression(isDistinct = isDistinct, filter = filter)
+ }
+
+ private def optimizedAggregate(expressions: Alias*): Aggregate = {
+ val plan = Aggregate(Seq.empty, expressions, relation)
+ withSQLConf(SQLConf.COMBINE_APPROXIMATE_PERCENTILES_ENABLED.key -> "true")
{
+ Optimize.execute(plan.analyze).asInstanceOf[Aggregate]
+ }
+ }
+
+ private def percentileAggregates(aggregate: Aggregate):
Seq[AggregateExpression] = {
+ aggregate.aggregateExpressions.flatMap(_.collect {
+ case expression @ AggregateExpression(_: ApproximatePercentile, _, _, _,
_) => expression
+ })
+ }
+
+ private def percentageValues(percentile: ApproximatePercentile): Seq[Double]
= {
+
percentile.percentageExpression.eval().asInstanceOf[ArrayData].toDoubleArray().toSeq
+ }
+
+ private def ordinals(aggregate: Aggregate): Seq[Option[Int]] = {
+ aggregate.aggregateExpressions.map(_.collectFirst {
+ case GetArrayItem(_, Literal(index: Int, _), false) => index
+ })
+ }
+
+ private def assertNotCombined(aggregate: Aggregate): Unit = {
+ val expressions = percentileAggregates(aggregate)
+ assert(expressions.map(_.resultId).distinct.size == expressions.size)
+
assert(!aggregate.aggregateExpressions.exists(_.exists(_.isInstanceOf[GetArrayItem])))
+ }
+
+ test("do not combine approximate percentiles when disabled") {
+ assert(!new
SQLConf().getConf(SQLConf.COMBINE_APPROXIMATE_PERCENTILES_ENABLED))
+ val original = Aggregate(
+ Seq.empty,
+ Seq(
+ Alias(percentile(value, 0.5), "p50")(),
+ Alias(percentile(value, 0.9), "p90")()),
+ relation).analyze.asInstanceOf[Aggregate]
+
+ withSQLConf(SQLConf.COMBINE_APPROXIMATE_PERCENTILES_ENABLED.key ->
"false") {
+ assertNotCombined(Optimize.execute(original).asInstanceOf[Aggregate])
+ }
+ }
+
+ test("combine compatible percentiles and preserve output shape") {
+ val firstPercentile = percentile(value, 0.9)
+ firstPercentile.aggregateFunction.setTagValue(FunctionRegistry.FUNC_ALIAS,
"approx_percentile")
+ val aliases = Seq(
+ Alias(firstPercentile, "first")(),
+ Alias(percentile(value, 0.5), "second")(),
+ Alias(percentile(value, 0.9), "third")())
+ val original = Aggregate(Seq(group), group +: aliases, relation).analyze
+ .asInstanceOf[Aggregate]
+ val result =
withSQLConf(SQLConf.COMBINE_APPROXIMATE_PERCENTILES_ENABLED.key -> "true") {
+ Optimize.execute(original).asInstanceOf[Aggregate]
+ }
+ val aggregates = percentileAggregates(result)
+ val combined =
aggregates.head.aggregateFunction.asInstanceOf[ApproximatePercentile]
+
+ assert(aggregates.map(_.resultId).distinct.size == 1)
+ assert(aggregates.head.resultId != firstPercentile.resultId)
+ assert(percentageValues(combined) == Seq(0.9, 0.5, 0.9))
+ assert(combined.percentageExpression.toString == "[0.9,0.5,0.9]")
+ assert(combined.percentageExpression.sql == "ARRAY(0.9D, 0.5D, 0.9D)")
+ assert(combined.prettyName == "approx_percentile")
+ assert(ordinals(result) == Seq(None, Some(0), Some(1), Some(2)))
+ assert(result.groupingExpressions == original.groupingExpressions)
+ assert(result.output.map(_.exprId) == original.output.map(_.exprId))
+ }
+
+ test("preserve fusion identity through later constant folding") {
+ val result = optimizedAggregate(
+ Alias(percentile(value, 0.5), "p50")(),
+ Alias(percentile(value, 0.9), "p90")())
+ val folded = ConstantFolding(result).asInstanceOf[Aggregate]
+ val combined = percentileAggregates(folded).head.aggregateFunction
+ .asInstanceOf[ApproximatePercentile]
+
+ assert(combined.percentageExpression.isInstanceOf[PercentileFusionArray])
+ assert(!combined.percentageExpression.contextIndependentFoldable)
+ }
+
+ test("avoid redundant entries for duplicate physical aggregates") {
+ val duplicateOnly = optimizedAggregate(
+ Alias(percentile(value, 0.5), "first")(),
+ Alias(percentile(value, 0.5), "second")())
+ assertNotCombined(duplicateOnly)
+
+ val p50 = percentile(value, 0.5)
+ val mixed = optimizedAggregate(
+ Alias(p50, "first")(),
+ Alias(p50, "second")(),
+ Alias(percentile(value, 0.9), "third")())
+ val combined = percentileAggregates(mixed).head.aggregateFunction
+ .asInstanceOf[ApproximatePercentile]
+ assert(percentageValues(combined) == Seq(0.5, 0.9))
+ assert(ordinals(mixed) == Seq(Some(0), Some(0), Some(1)))
+ }
+
+ test("respect basic compatibility boundaries") {
+ val incompatible = Seq(
+ "input" -> (
+ percentile(value, 0.5),
+ percentile(other, 0.9)),
+ "accuracy" -> (
+ percentile(value, 0.5, accuracy = 1000),
+ percentile(value, 0.9, accuracy = 10000)),
+ "filter" -> (
+ percentile(value, 0.5),
+ percentile(value, 0.9, filter = Some(value > Literal(0)))),
+ "distinct" -> (
+ percentile(value, 0.5),
+ percentile(value, 0.9, isDistinct = true)))
+
+ incompatible.foreach { case (name, (first, second)) =>
+ withClue(name) {
+ assertNotCombined(optimizedAggregate(
+ Alias(first, "first")(),
+ Alias(second, "second")()))
+ }
+ }
+
+ val filter = value > Literal(0)
+ val compatible = optimizedAggregate(
+ Alias(percentile(value, 0.5, filter = Some(filter)), "filtered_p50")(),
+ Alias(percentile(value, 0.9, filter = Some(filter)), "filtered_p90")(),
+ Alias(percentile(value, 0.5, isDistinct = true, filter = Some(filter)),
"distinct_p50")(),
+ Alias(percentile(value, 0.9, isDistinct = true, filter = Some(filter)),
"distinct_p90")())
+ assert(percentileAggregates(compatible).map(_.resultId).distinct.size == 2)
+ }
+
+ test("require structural equality for inputs and filters") {
+ val firstInput = value + (other + group)
+ val secondInput = (value + other) + group
+ val inputResult = optimizedAggregate(
+ Alias(percentile(firstInput, 0.5), "first_p50")(),
+ Alias(percentile(secondInput, 0.5), "second_p50")(),
+ Alias(percentile(secondInput, 0.9), "second_p90")(),
+ Alias(percentile(firstInput, 0.9), "first_p90")())
+
+ assert(firstInput != secondInput)
+ assert(firstInput.canonicalized == secondInput.canonicalized)
+ assertNotCombined(inputResult)
+
+ val disjointInputPercentages = optimizedAggregate(
+ Alias(percentile(firstInput, 0.5), "first_p50")(),
+ Alias(percentile(firstInput, 0.9), "first_p90")(),
+ Alias(percentile(secondInput, 0.25), "second_p25")(),
+ Alias(percentile(secondInput, 0.75), "second_p75")())
+
assert(percentileAggregates(disjointInputPercentages).map(_.resultId).distinct.size
== 2)
+ assert(ordinals(disjointInputPercentages) ==
+ Seq(Some(0), Some(1), Some(0), Some(1)))
+
+ val crossDistinctInput = optimizedAggregate(
+ Alias(percentile(firstInput, 0.5), "first_p50")(),
+ Alias(percentile(firstInput, 0.9), "first_p90")(),
+ Alias(percentile(secondInput, 0.5, isDistinct = true), "distinct_p50")(),
+ Alias(percentile(secondInput, 0.9, isDistinct = true), "distinct_p90")())
+ assertNotCombined(crossDistinctInput)
+
+ val firstFilter = firstInput > Literal(0)
+ val secondFilter = secondInput > Literal(0)
+ val filterResult = optimizedAggregate(
+ Alias(percentile(value, 0.5, filter = Some(firstFilter)), "first_p50")(),
+ Alias(percentile(value, 0.5, filter = Some(secondFilter)),
"second_p50")(),
+ Alias(percentile(value, 0.9, filter = Some(secondFilter)),
"second_p90")(),
+ Alias(percentile(value, 0.9, filter = Some(firstFilter)), "first_p90")())
+ assertNotCombined(filterResult)
+
+ val crossDistinctFilter = optimizedAggregate(
+ Alias(percentile(value, 0.5, filter = Some(firstFilter)), "first_p50")(),
+ Alias(percentile(value, 0.9, filter = Some(firstFilter)), "first_p90")(),
+ Alias(percentile(
+ value, 0.5, isDistinct = true, filter = Some(secondFilter)),
"distinct_p50")(),
+ Alias(percentile(
+ value, 0.9, isDistinct = true, filter = Some(secondFilter)),
"distinct_p90")())
+ assertNotCombined(crossDistinctFilter)
+ }
+
+ test("do not fuse canonically equivalent but differently evaluated
accuracies") {
+ val firstAccuracy =
+ ((Literal(1.0e16) + Literal(-1.0e16)) + Literal(3.0)).cast(IntegerType)
+ val secondAccuracy =
+ (Literal(1.0e16) + (Literal(-1.0e16) + Literal(3.0))).cast(IntegerType)
+
+ def percentileWithAccuracy(
+ percentage: Double,
+ accuracy: Expression): AggregateExpression = {
+ new ApproximatePercentile(value, Literal(percentage),
accuracy).toAggregateExpression()
+ }
+
+ val result = optimizedAggregate(
+ Alias(percentileWithAccuracy(0.5, firstAccuracy), "first_p50")(),
+ Alias(percentileWithAccuracy(0.5, secondAccuracy), "second_p50")(),
+ Alias(percentileWithAccuracy(0.9, secondAccuracy), "second_p90")(),
+ Alias(percentileWithAccuracy(0.9, firstAccuracy), "first_p90")())
+
+ assert(firstAccuracy.eval() == 3)
+ assert(secondAccuracy.eval() == 4)
+ assert(firstAccuracy.canonicalized == secondAccuracy.canonicalized)
+ assertNotCombined(result)
+ }
+
+ test("do not fuse canonically equal percentages that evaluate differently") {
+ val firstPercentage =
+ (Literal(1.0e16) + Literal(-1.0e16)) + Literal(0.5)
+ val secondPercentage =
+ Literal(1.0e16) + (Literal(-1.0e16) + Literal(0.5))
+ def percentileWithPercentage(percentage: Expression): AggregateExpression
= {
+ new ApproximatePercentile(value, percentage,
Literal(10000)).toAggregateExpression()
+ }
+
+ val scalarResult = optimizedAggregate(
+ Alias(percentileWithPercentage(firstPercentage), "p50")(),
+ Alias(percentileWithPercentage(secondPercentage), "p0")(),
+ Alias(percentile(value, 0.9), "p90")())
+ assert(firstPercentage.eval() == 0.5d)
+ assert(secondPercentage.eval() == 0.0d)
+ assert(firstPercentage.canonicalized == secondPercentage.canonicalized)
+ assertNotCombined(scalarResult)
+
+ val arrayPercentile = new ApproximatePercentile(
+ value,
+ CreateArray(Seq(firstPercentage, Literal(0.9))),
+ Literal(10000)).toAggregateExpression()
+ val arrayResult = optimizedAggregate(
+ Alias(arrayPercentile, "percentiles")(),
+ Alias(percentileWithPercentage(secondPercentage), "p0")(),
+ Alias(percentile(value, 0.9), "p90")())
+ val arrayAggregates = percentileAggregates(arrayResult)
+ assert(arrayAggregates.map(_.resultId).distinct.size == 2)
+ assert(percentageValues(arrayAggregates.head.aggregateFunction
+ .asInstanceOf[ApproximatePercentile]) == Seq(0.5, 0.9))
+ assert(percentageValues(arrayAggregates(1).aggregateFunction
+ .asInstanceOf[ApproximatePercentile]) == Seq(0.0, 0.9))
+ assert(ordinals(arrayResult) == Seq(None, Some(0), Some(1)))
+ }
+
+ test("preserve existing arrays while combining compatible scalars") {
+ val arrayPercentile = new ApproximatePercentile(
+ value,
+ CreateArray(Seq(Literal(0.25), Literal(0.75))),
+ Literal(10000)).toAggregateExpression()
+ val result = optimizedAggregate(
+ Alias(arrayPercentile, "percentiles")(),
+ Alias(percentile(value, 0.5), "p50")(),
+ Alias(percentile(value, 0.9), "p90")())
+
+ assert(percentileAggregates(result).map(_.resultId).distinct.size == 2)
+
assert(result.aggregateExpressions.drop(1).forall(_.exists(_.isInstanceOf[GetArrayItem])))
+ }
+}
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkOptimizer.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkOptimizer.scala
index 85723550821b..9dd7bef549fa 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkOptimizer.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkOptimizer.scala
@@ -72,6 +72,7 @@ class SparkOptimizer(
InjectRuntimeFilter),
Batch("MergeSubplans", Once,
MergeSubplans,
+ CombineApproximatePercentiles,
RewriteDistinctAggregates),
Batch("Pushdown Filters from PartitionPruning", fixedPoint,
PushDownPredicates),
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/ApproximatePercentileQuerySuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/ApproximatePercentileQuerySuite.scala
index ae24a21538f1..a7705de0b4b7 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/ApproximatePercentileQuerySuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/ApproximatePercentileQuerySuite.scala
@@ -20,12 +20,19 @@ package org.apache.spark.sql
import java.sql.{Date, Timestamp}
import java.time.{Duration, LocalDateTime, LocalTime, Period}
-import
org.apache.spark.sql.catalyst.expressions.aggregate.ApproximatePercentile
+import org.apache.spark.SparkArithmeticException
+import
org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression,
+ ApproximatePercentile, Final, Partial}
import
org.apache.spark.sql.catalyst.expressions.aggregate.ApproximatePercentile.DEFAULT_PERCENTILE_ACCURACY
import
org.apache.spark.sql.catalyst.expressions.aggregate.ApproximatePercentile.PercentileDigest
+import org.apache.spark.sql.catalyst.plans.logical.Expand
import org.apache.spark.sql.catalyst.util.DateTimeUtils
+import org.apache.spark.sql.execution.ReusedSubqueryExec
+import org.apache.spark.sql.execution.aggregate.ObjectHashAggregateExec
+import org.apache.spark.sql.execution.exchange.ReusedExchangeExec
+import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.test.SharedSparkSession
-import org.apache.spark.sql.types.TimeType
+import org.apache.spark.sql.types.{ArrayType, DoubleType, TimeType}
import org.apache.spark.tags.SlowSQLTest
/**
@@ -36,6 +43,461 @@ class ApproximatePercentileQuerySuite extends
SharedSparkSession {
import testImplicits._
private val table = "percentile_approx"
+ private val constantFoldingRule =
+ "org.apache.spark.sql.catalyst.optimizer.ConstantFolding"
+
+ private def fusionTest(name: String)(body: => Unit): Unit = test(name) {
+ withSQLConf(SQLConf.COMBINE_APPROXIMATE_PERCENTILES_ENABLED.key -> "true")
{
+ body
+ }
+ }
+
+ private def excludedRules: Seq[String] = {
+ spark.sessionState.conf.optimizerExcludedRules.toSeq
+ .flatMap(_.split(","))
+ .map(_.trim)
+ .filter(_.nonEmpty)
+ }
+
+ private def assertPercentileDigestCount(query: DataFrame, expected: Int):
Unit = {
+ val counts = query.queryExecution.sparkPlan.collect {
+ case aggregate: ObjectHashAggregateExec =>
+ aggregate.aggregateExpressions.count(
+ _.aggregateFunction.isInstanceOf[ApproximatePercentile])
+ }
+ assert(counts.nonEmpty)
+ assert(counts.forall(_ == expected), counts)
+ }
+
+ private def checkMatchesUnfusedBaseline(sql: String, expectedDigests: Int):
Unit = {
+ val withoutConstantFolding = (excludedRules :+
constantFoldingRule).distinct
+ val baseline = withSQLConf(
+ SQLConf.OPTIMIZER_EXCLUDED_RULES.key ->
withoutConstantFolding.mkString(","),
+ SQLConf.COMBINE_APPROXIMATE_PERCENTILES_ENABLED.key -> "false") {
+ spark.sql(sql).collect().toSeq
+ }
+
+ withSQLConf(
+ SQLConf.OPTIMIZER_EXCLUDED_RULES.key ->
withoutConstantFolding.mkString(",")) {
+ val query = spark.sql(sql)
+ checkAnswer(query, baseline)
+ assertPercentileDigestCount(query, expectedDigests)
+ }
+ }
+
+ test("approximate percentile fusion can be disabled") {
+ withSQLConf(SQLConf.COMBINE_APPROXIMATE_PERCENTILES_ENABLED.key ->
"false") {
+ val query = spark.sql(
+ "SELECT percentile_approx(id, 0.5D), percentile_approx(id, 0.9D) FROM
range(10)")
+ checkAnswer(query, Row(4L, 8L))
+ assertPercentileDigestCount(query, 2)
+ }
+ }
+
+ fusionTest("compatible scalar percentiles share one physical percentile
digest") {
+ withTempView(table) {
+ (1 to 1000).toDF("col").createOrReplaceTempView(table)
+ val query = spark.sql(
+ s"""SELECT
+ | approx_percentile(col, 0.5, 10000),
+ | approx_percentile(col, 0.9, 10000),
+ | approx_percentile(col, 0.95, 10000)
+ |FROM $table
+ |""".stripMargin)
+
+ checkAnswer(query, Row(500, 900, 950))
+ assertPercentileDigestCount(query, 1)
+ val optimizedPercentiles =
query.queryExecution.optimizedPlan.expressions.flatMap {
+ _.collect { case percentile: ApproximatePercentile => percentile }
+ }
+ assert(optimizedPercentiles.nonEmpty)
+ assert(optimizedPercentiles.forall(_.prettyName == "approx_percentile"))
+ val modes = query.queryExecution.sparkPlan.collect {
+ case aggregate: ObjectHashAggregateExec =>
+ aggregate.aggregateExpressions.collect {
+ case expression @ AggregateExpression(_: ApproximatePercentile, _,
_, _, _) =>
+ expression.mode
+ }
+ }.flatten.toSet
+ assert(modes == Set(Partial, Final))
+ }
+ }
+
+ fusionTest("do not fuse duplicate percentages already shared by physical
planning") {
+ withTempView(table) {
+ (1 to 1000).toDF("col").createOrReplaceTempView(table)
+ val query = spark.sql(
+ s"""SELECT
+ | percentile_approx(col, 0.5D),
+ | percentile_approx(col, 0.25D + 0.25D)
+ |FROM $table
+ |""".stripMargin)
+
+ checkAnswer(query, Row(500, 500))
+ val percentiles = query.queryExecution.sparkPlan.collect {
+ case aggregate: ObjectHashAggregateExec =>
+ aggregate.aggregateExpressions.collect {
+ case AggregateExpression(
+ percentile: ApproximatePercentile, _, _, _, _) => percentile
+ }
+ }.flatten
+ assert(percentiles.nonEmpty)
+ assert(percentiles.forall(_.percentageExpression.dataType == DoubleType))
+ }
+ }
+
+ fusionTest("preserve structural input and filter evaluation") {
+ checkAnswer(
+ spark.sql(
+ """SELECT
+ | percentile_approx((a + b) + c, 0.5D),
+ | percentile_approx(a + (b + c), 0.9D)
+ |FROM VALUES (
+ | CAST(10000000000000000 AS DOUBLE),
+ | CAST(-10000000000000000 AS DOUBLE),
+ | CAST(1 AS DOUBLE)
+ |) AS t(a, b, c)
+ |""".stripMargin),
+ Row(1.0d, 0.0d))
+
+ checkAnswer(
+ spark.sql(
+ """SELECT
+ | percentile_approx(v, 0.5D)
+ | FILTER (WHERE (a + b) + c = 1D),
+ | percentile_approx(v, 0.9D)
+ | FILTER (WHERE a + (b + c) = 1D)
+ |FROM VALUES (
+ | 7,
+ | CAST(10000000000000000 AS DOUBLE),
+ | CAST(-10000000000000000 AS DOUBLE),
+ | CAST(1 AS DOUBLE)
+ |) AS t(v, a, b, c)
+ |""".stripMargin),
+ Row(7, null))
+
+ withSQLConf(SQLConf.ANSI_ENABLED.key -> "true") {
+ val exception = intercept[SparkArithmeticException] {
+ spark.sql(
+ """SELECT
+ | percentile_approx(a + (b + c), 0.5D),
+ | percentile_approx((a + b) + c, 0.9D)
+ |FROM VALUES (
+ | CAST(2147483647 AS INT),
+ | CAST(1 AS INT),
+ | CAST(-1 AS INT)
+ |) AS t(a, b, c)
+ |""".stripMargin).collect()
+ }
+ assert(exception.getCondition == "ARITHMETIC_OVERFLOW")
+ }
+ }
+
+ fusionTest("do not fuse canonical input or filter collisions") {
+ checkMatchesUnfusedBaseline(
+ """SELECT
+ | percentile_approx(a + (b + c), 0.5D),
+ | percentile_approx((a + b) + c, 0.5D),
+ | percentile_approx((a + b) + c, 0.9D),
+ | percentile_approx(a + (b + c), 0.9D)
+ |FROM VALUES (
+ | CAST(10000000000000000 AS DOUBLE),
+ | CAST(-10000000000000000 AS DOUBLE),
+ | CAST(1 AS DOUBLE)
+ |) AS t(a, b, c)
+ |""".stripMargin,
+ expectedDigests = 2)
+
+ val crossDistinctCollision =
+ """SELECT
+ | percentile_approx(DISTINCT a + (b + c), 0.5D),
+ | percentile_approx(DISTINCT a + (b + c), 0.9D),
+ | percentile_approx((a + b) + c, 0.5D),
+ | percentile_approx((a + b) + c, 0.9D)
+ |FROM VALUES (
+ | CAST(10000000000000000 AS DOUBLE),
+ | CAST(-10000000000000000 AS DOUBLE),
+ | CAST(1 AS DOUBLE)
+ |) AS t(a, b, c)
+ |""".stripMargin
+ val unfusedBaseline = withSQLConf(
+ SQLConf.COMBINE_APPROXIMATE_PERCENTILES_ENABLED.key -> "false") {
+ spark.sql(crossDistinctCollision).collect().toSeq
+ }
+ checkAnswer(spark.sql(crossDistinctCollision), unfusedBaseline)
+
+ val filteredArrayCollision = spark.sql(
+ """SELECT
+ | percentile_approx(v, array(0.5D, 0.9D))
+ | FILTER (WHERE a + (b + c) = 0D),
+ | percentile_approx(v, 0.5D)
+ | FILTER (WHERE (a + b) + c = 0D),
+ | percentile_approx(v, 0.9D)
+ | FILTER (WHERE (a + b) + c = 0D)
+ |FROM VALUES (
+ | 7,
+ | CAST(10000000000000000 AS DOUBLE),
+ | CAST(-10000000000000000 AS DOUBLE),
+ | CAST(1 AS DOUBLE)
+ |) AS t(v, a, b, c)
+ |""".stripMargin)
+ checkAnswer(filteredArrayCollision, Row(Seq(7, 7), null, null))
+ assertPercentileDigestCount(filteredArrayCollision, 2)
+ }
+
+ fusionTest("preserve canonically colliding parameters") {
+ val cases = Seq(
+ (
+ """SELECT
+ | percentile_approx(
+ | v, 0.5D, CAST((1e16D + -1e16D) + 3D AS INT)),
+ | percentile_approx(
+ | v, 0.5D, CAST(1e16D + (-1e16D + 3D) AS INT)),
+ | percentile_approx(
+ | v, 0.9D, CAST(1e16D + (-1e16D + 3D) AS INT)),
+ | percentile_approx(
+ | v, 0.9D, CAST((1e16D + -1e16D) + 3D AS INT))
+ |FROM range(100) AS t(v)
+ |""".stripMargin,
+ 2),
+ (
+ """SELECT
+ | percentile_approx(
+ | id, array((1e16D + -1e16D) + 0.5D, 0.9D)),
+ | percentile_approx(
+ | id, 1e16D + (-1e16D + 0.5D)),
+ | percentile_approx(id, 0.9D)
+ |FROM range(100)
+ |""".stripMargin,
+ 2))
+
+ cases.foreach { case (sql, expectedDigests) =>
+ checkMatchesUnfusedBaseline(sql, expectedDigests)
+ }
+ }
+
+ fusionTest("preserve existing arrays that collide after distinct removal") {
+ val query = spark.sql(
+ """SELECT
+ | percentile_approx(
+ | DISTINCT a + (b + c), array(0.5D, 0.9D)),
+ | percentile_approx(DISTINCT a + (b + c), 0.1D),
+ | percentile_approx((a + b) + c, 0.5D),
+ | percentile_approx((a + b) + c, 0.9D)
+ |FROM VALUES (
+ | CAST(10000000000000000 AS DOUBLE),
+ | CAST(-10000000000000000 AS DOUBLE),
+ | CAST(1 AS DOUBLE)
+ |) AS t(a, b, c)
+ |""".stripMargin)
+ checkAnswer(query, Row(Seq(0.0d, 0.0d), 0.0d, 1.0d, 1.0d))
+ assertPercentileDigestCount(query, 3)
+ }
+
+ fusionTest("fused percentiles use fresh result IDs across CTE references") {
+ val query = spark.sql(
+ """WITH c AS (
+ | SELECT
+ | percentile_approx(v, 0.5D) AS a,
+ | percentile_approx(v, 0.9D) AS b
+ | FROM range(1, 6) AS t(v)
+ |)
+ |SELECT c1.a, c1.b, c2.a
+ |FROM c AS c1 JOIN c AS c2
+ |""".stripMargin)
+ checkAnswer(query, Row(3L, 5L, 3L))
+ val optimizedPlan = query.queryExecution.optimizedPlan
+
assert(optimizedPlan.subqueriesAll.exists(_.exists(_.expressions.exists(_.exists
{
+ case percentile: ApproximatePercentile =>
+ percentile.percentageExpression.dataType.isInstanceOf[ArrayType]
+ case _ => false
+ }))), optimizedPlan.numberedTreeString)
+ }
+
+ fusionTest("preserve percentile inputs across scalar subquery reuse") {
+ val query =
+ """SELECT
+ | (SELECT named_struct(
+ | 'p50', percentile_approx((a + b) + c, 0.5D),
+ | 'p90', percentile_approx((a + b) + c, 0.9D))
+ | FROM VALUES (
+ | CAST(10000000000000000 AS DOUBLE),
+ | CAST(-10000000000000000 AS DOUBLE),
+ | CAST(1 AS DOUBLE)
+ | ) AS t(a, b, c)),
+ | (SELECT named_struct(
+ | 'p50', get(percentile_approx(
+ | a + (b + c), array(0.5D, 0.9D)), 0),
+ | 'p90', get(percentile_approx(
+ | a + (b + c), array(0.5D, 0.9D)), 1))
+ | FROM VALUES (
+ | CAST(10000000000000000 AS DOUBLE),
+ | CAST(-10000000000000000 AS DOUBLE),
+ | CAST(1 AS DOUBLE)
+ | ) AS t(a, b, c))
+ |""".stripMargin
+
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.SUBQUERY_REUSE_ENABLED.key -> "true",
+ SQLConf.EXCHANGE_REUSE_ENABLED.key -> "false") {
+ checkAnswer(
+ spark.sql(query),
+ Row(Row(1.0d, 1.0d), Row(0.0d, 0.0d)))
+ }
+ }
+
+ fusionTest("preserve pre-fusion identity across exchange reuse") {
+ val parameterPairs = Seq(
+ ("(1e16D + -1e16D) + 0.5D", "100"),
+ ("0.5D", "CAST((1e16D + -1e16D) + 100D AS INT)"),
+ ("0.5D", "100L"))
+ val activeRules = (excludedRules :+
constantFoldingRule).distinct.mkString(",")
+
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.SUBQUERY_REUSE_ENABLED.key -> "false",
+ SQLConf.EXCHANGE_REUSE_ENABLED.key -> "true",
+ SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> activeRules) {
+ parameterPairs.foreach { case (firstPercentage, secondAccuracy) =>
+ val query = spark.sql(
+ s"""SELECT
+ | (a + b) + c,
+ | array(
+ | percentile_approx(1D, $firstPercentage, 100),
+ | percentile_approx(1D, 0.9D, $secondAccuracy))
+ |FROM VALUES (
+ | CAST(10000000000000000 AS DOUBLE),
+ | CAST(-10000000000000000 AS DOUBLE),
+ | CAST(1 AS DOUBLE)
+ |) AS t1(a, b, c)
+ |GROUP BY (a + b) + c
+ |UNION ALL
+ |SELECT
+ | a + (b + c),
+ | array(
+ | percentile_approx(1D, 0.5D, 100),
+ | percentile_approx(1D, 0.9D, 100))
+ |FROM VALUES (
+ | CAST(10000000000000000 AS DOUBLE),
+ | CAST(-10000000000000000 AS DOUBLE),
+ | CAST(1 AS DOUBLE)
+ |) AS t2(a, b, c)
+ |GROUP BY a + (b + c)
+ |""".stripMargin)
+
+ checkAnswer(
+ query,
+ Seq(
+ Row(1.0d, Seq(1.0d, 1.0d)),
+ Row(0.0d, Seq(1.0d, 1.0d))))
+ assert(query.queryExecution.executedPlan.collect {
+ case _: ReusedExchangeExec => true
+ }.isEmpty)
+ }
+ }
+ }
+
+ fusionTest("identical fused percentiles remain reusable") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.SUBQUERY_REUSE_ENABLED.key -> "true") {
+ val subquery =
+ """(SELECT array(
+ | percentile_approx(v, 0D),
+ | percentile_approx(v, 1D))
+ | FROM VALUES (1), (2), (3) AS t(v))
+ |""".stripMargin
+ val query = spark.sql(s"SELECT $subquery, $subquery")
+
+ checkAnswer(query, Row(Seq(1, 3), Seq(1, 3)))
+ assert(query.queryExecution.executedPlan.collectWithSubqueries {
+ case _: ReusedSubqueryExec => true
+ }.nonEmpty)
+ }
+ }
+
+ fusionTest("fused distinct percentiles keep a single distinct group") {
+ val query = spark.sql(
+ """SELECT
+ | count(DISTINCT id),
+ | percentile_approx(DISTINCT id, 0.5D),
+ | percentile_approx(DISTINCT id, 0.9D)
+ |FROM range(10)
+ |""".stripMargin)
+
+ checkAnswer(query, Row(10L, 4L, 8L))
+ assertPercentileDigestCount(query, 1)
+ assert(query.queryExecution.optimizedPlan.collect {
+ case _: Expand => true
+ }.isEmpty)
+ }
+
+ fusionTest("combined scalar percentiles preserve empty-input nulls") {
+ withTempView(table) {
+ Seq.empty[Int].toDF("col").createOrReplaceTempView(table)
+ val query = spark.sql(
+ s"""SELECT
+ | percentile_approx(col, 0.5D),
+ | percentile_approx(col, 0.9D)
+ |FROM $table
+ |""".stripMargin)
+ checkAnswer(query, Row(null, null))
+ assertPercentileDigestCount(query, 1)
+ }
+ }
+
+ fusionTest("preserve compressed and merged low-accuracy percentile digests")
{
+ val sql =
+ """SELECT
+ | percentile_approx(id, 0.1D, 100),
+ | percentile_approx(id, 0.5D, 100),
+ | percentile_approx(id, 0.9D, 100)
+ |FROM range(0, 50000, 1, 4)
+ |""".stripMargin
+ val baseline = withSQLConf(
+ SQLConf.COMBINE_APPROXIMATE_PERCENTILES_ENABLED.key -> "false") {
+ spark.sql(sql).collect().toSeq
+ }
+
+ val query = spark.sql(sql)
+ checkAnswer(query, baseline)
+ assertPercentileDigestCount(query, 1)
+ }
+
+ fusionTest("fuse compatible percentiles introduced by merged scalar
subqueries") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.SUBQUERY_REUSE_ENABLED.key -> "true") {
+ val query = spark.sql(
+ """SELECT
+ | (SELECT percentile_approx(id, 0.5D) FROM range(10)),
+ | (SELECT percentile_approx(id, 0.9D) FROM range(10))
+ |""".stripMargin)
+
+ checkAnswer(query, Row(4L, 8L))
+ val digestCounts =
query.queryExecution.executedPlan.collectWithSubqueries {
+ case aggregate: ObjectHashAggregateExec =>
+ aggregate.aggregateExpressions.count(
+ _.aggregateFunction.isInstanceOf[ApproximatePercentile])
+ }
+ assert(digestCounts.nonEmpty && digestCounts.forall(_ == 1),
digestCounts)
+ }
+ }
+
+ fusionTest("fuse canonical input groups with disjoint percentages
independently") {
+ val query = spark.sql(
+ """SELECT
+ | percentile_approx(a + b, 0.5D),
+ | percentile_approx(a + b, 0.9D),
+ | percentile_approx(b + a, 0.25D),
+ | percentile_approx(b + a, 0.75D)
+ |FROM VALUES (1D, 2D), (3D, 4D), (5D, 6D) AS t(a, b)
+ |""".stripMargin)
+
+ checkAnswer(query, Row(7.0d, 11.0d, 3.0d, 11.0d))
+ assertPercentileDigestCount(query, 2)
+ }
test("percentile_approx, single percentile value") {
withTempView(table) {
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationSuite.scala
index f065f1de5cdc..09dffb4f65cc 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationSuite.scala
@@ -110,6 +110,38 @@ class StreamingAggregationSuite extends
StateStoreMetricsTest with Assertions {
)
}
+ testWithAllStateVersions("approximate percentiles preserve existing
streaming checkpoints") {
+ withTempDir { checkpointDir =>
+ val inputData = MemoryStream[(Int, Int)]
+ val aggregated = inputData.toDF()
+ .groupBy($"_1")
+ .agg(
+ expr("percentile_approx(_2, 0.5)").as("p50"),
+ expr("percentile_approx(_2, 0.9)").as("p90"))
+ .as[(Int, Int, Int)]
+
+ // Create the two-digest checkpoint used before percentile fusion was
introduced.
+ testStream(aggregated, Update)(
+ StartStream(
+ checkpointLocation = checkpointDir.getAbsolutePath,
+ additionalConfs = Map(
+ SQLConf.COMBINE_APPROXIMATE_PERCENTILES_ENABLED.key -> "false")),
+ AddData(inputData, (0, 1), (0, 2), (0, 3)),
+ CheckLastBatch((0, 2, 3)),
+ StopStream)
+
+ // Restart with percentile fusion enabled to verify the existing
checkpoint stays valid.
+ testStream(aggregated, Update)(
+ StartStream(
+ checkpointLocation = checkpointDir.getAbsolutePath,
+ additionalConfs = Map(
+ SQLConf.COMBINE_APPROXIMATE_PERCENTILES_ENABLED.key -> "true")),
+ AddData(inputData, (0, 4)),
+ CheckLastBatch((0, 2, 4)),
+ StopStream)
+ }
+ }
+
testWithAllStateVersions("count distinct") {
val inputData = MemoryStream[(Int, Seq[Int])]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
