You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by su...@apache.org on 2026/08/06 22:27:40 UTC
(spark) branch branch-4.x updated: [SPARK-58378][SQL] Fuse compatible approximate percentile sketches
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 <ch...@openai.com>
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 <ch...@openai.com>
Signed-off-by: Chao Sun <ch...@openai.com>
(cherry picked from commit 316c6efe053173d044eaeef7b8945e4342bd4de0)
Signed-off-by: Chao Sun <ch...@openai.com>
---
.../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: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org