You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by do...@apache.org on 2026/03/20 14:26:54 UTC
(spark) branch master updated: [SPARK-56102][SQL] `UnionEstimation` code cleanup
This is an automated email from the ASF dual-hosted git repository.
dongjoon pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 5bf8d95beb83 [SPARK-56102][SQL] `UnionEstimation` code cleanup
5bf8d95beb83 is described below
commit 5bf8d95beb83f1e0979c6e1705466459b304bc11
Author: Peter Toth <pe...@gmail.com>
AuthorDate: Fri Mar 20 07:26:43 2026 -0700
[SPARK-56102][SQL] `UnionEstimation` code cleanup
### What changes were proposed in this pull request?
Refactored UnionEstimation to consolidate column stats computation into a single, cleaner pass:
- Three compute* methods merged into one computeColumnStats: computeMinMaxStats, computeNullCountStats, and computeDistinctCountStats are replaced by a single method that computes all three
stat types in one traversal over the union's output columns.
- Stats lookup moved upfront: Instead of repeatedly calling child.stats.attributeStats.get(attr) inside each fold, stats are resolved once via c.output.map(c.stats.attributeStats.get)
before transposing, yielding a Seq[Option[ColumnStat]] per output column (maybeColStats). Folds then operate directly on Option[ColumnStat] values, eliminating redundant index-based lookups
and attrs.zip(union.children) re-pairings.
- estimate simplified: The multi-pass layering of newMinMaxStats ++ withNullCount ++ withDistinctCount with intermediate AttributeMap merges is replaced by a direct
AttributeMap(computeColumnStats(...)) call.
- Short-circuiting folds: Each stat type uses a foldLeft(Option(...)) that propagates None as soon as any child is missing the required stat, avoiding separate filter+map passes and double
attributeStats lookups.
- collect over filter+map: Used throughout to combine guard and value extraction in one step (e.g. .collect { case s if s.hasMinMaxStats => (s.min.get, s.max.get) }).
- zip(union.output) / zip(union.children): Replaced zipWithIndex + index-based access (unionOutput(outputIndex), union.children(childIndex)) with direct zips, making data relationships
explicit.
### Why are the changes needed?
The previous implementation had several readability and efficiency issues:
- attributeStats was looked up twice per attribute per child (once in the filter guard, once in the computation body).
- Three separate traversals over the same column structure were performed independently.
- Index-based access made the relationship between attributes and their source children implicit.
- The stats merging in estimate duplicated newMinMaxStats entries multiple times, relying on AttributeMap position-based deduplication.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Existing unit tests in `UnionEstimationSuite` cover the changed logic. No new tests were added as this is a pure refactoring with no behavior change.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Sonnet 4.6
Closes #54918 from peter-toth/SPARK-56102-unionestimation-code-cleanup.
Authored-by: Peter Toth <pe...@gmail.com>
Signed-off-by: Dongjoon Hyun <do...@apache.org>
---
.../logical/statsEstimation/UnionEstimation.scala | 163 +++++++--------------
1 file changed, 50 insertions(+), 113 deletions(-)
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/statsEstimation/UnionEstimation.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/statsEstimation/UnionEstimation.scala
index 4602744ce8a4..7ad05ee3ad6b 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/statsEstimation/UnionEstimation.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/statsEstimation/UnionEstimation.scala
@@ -46,99 +46,18 @@ object UnionEstimation {
val sizeInBytes = union.children.map(_.stats.sizeInBytes).sum
val outputRows = if (rowCountsExist(union.children: _*)) {
Some(union.children.map(_.stats.rowCount.get).sum)
- } else {
- None
- }
-
- val newMinMaxStats = computeMinMaxStats(union)
- val newNullCountStats = computeNullCountStats(union)
- val newDistinctCountStats = computeDistinctCountStats(union, outputRows)
- val newAttrStats = {
- val baseStats = AttributeMap(newMinMaxStats)
- // Layer nullCount on top of min/max
- val withNullCount = newNullCountStats.map { case attrStat @ (attr, stat) =>
- baseStats.get(attr).map { baseStat =>
- attr -> baseStat.copy(nullCount = stat.nullCount)
- }.getOrElse(attrStat)
- }
- val merged = AttributeMap(newMinMaxStats ++ withNullCount)
- // Layer distinctCount on top of min/max + nullCount
- val withDistinctCount = newDistinctCountStats.map { case attrStat @ (attr, stat) =>
- merged.get(attr).map { baseStat =>
- attr -> baseStat.copy(distinctCount = stat.distinctCount)
- }.getOrElse(attrStat)
- }
- AttributeMap(newMinMaxStats ++ withNullCount ++ withDistinctCount)
- }
-
- Some(
- Statistics(
- sizeInBytes = sizeInBytes,
- rowCount = outputRows,
- attributeStats = newAttrStats))
- }
-
- private def computeMinMaxStats(union: Union): Seq[(Attribute, ColumnStat)] = {
- val unionOutput = union.output
- val attrToComputeMinMaxStats = union.children.map(_.output).transpose.zipWithIndex.filter {
- case (attrs, outputIndex) => isTypeSupported(unionOutput(outputIndex).dataType) &&
- // checks if all the children has min/max stats for an attribute
- attrs.zipWithIndex.forall {
- case (attr, childIndex) =>
- val attrStats = union.children(childIndex).stats.attributeStats
- attrStats.get(attr).isDefined && attrStats(attr).hasMinMaxStats
- }
- }
- attrToComputeMinMaxStats.map {
- case (attrs, outputIndex) =>
- val dataType = unionOutput(outputIndex).dataType
- val statComparator = createStatComparator(dataType)
- val minMaxValue = attrs.zipWithIndex.foldLeft[(Option[Any], Option[Any])]((None, None)) {
- case ((minVal, maxVal), (attr, childIndex)) =>
- val colStat = union.children(childIndex).stats.attributeStats(attr)
- val min = if (minVal.isEmpty || statComparator(colStat.min.get, minVal.get)) {
- colStat.min
- } else {
- minVal
- }
- val max = if (maxVal.isEmpty || statComparator(maxVal.get, colStat.max.get)) {
- colStat.max
- } else {
- maxVal
- }
- (min, max)
- }
- val newStat = ColumnStat(min = minMaxValue._1, max = minMaxValue._2)
- unionOutput(outputIndex) -> newStat
- }
- }
-
- private def computeNullCountStats(union: Union): Seq[(Attribute, ColumnStat)] = {
- val unionOutput = union.output
- val attrToComputeNullCount = union.children.map(_.output).transpose.zipWithIndex.filter {
- case (attrs, _) => attrs.zipWithIndex.forall {
- case (attr, childIndex) =>
- val attrStats = union.children(childIndex).stats.attributeStats
- attrStats.get(attr).isDefined && attrStats(attr).nullCount.isDefined
- }
- }
- attrToComputeNullCount.map {
- case (attrs, outputIndex) =>
- val firstStat = union.children.head.stats.attributeStats(attrs.head)
- val firstNullCount = firstStat.nullCount.get
- val colWithNullStatValues = attrs.zipWithIndex.tail.foldLeft[BigInt](firstNullCount) {
- case (totalNullCount, (attr, childIndex)) =>
- val colStat = union.children(childIndex).stats.attributeStats(attr)
- totalNullCount + colStat.nullCount.get
- }
- val newStat = ColumnStat(nullCount = Some(colWithNullStatValues))
- unionOutput(outputIndex) -> newStat
- }
+ } else None
+ Some(Statistics(
+ sizeInBytes = sizeInBytes,
+ rowCount = outputRows,
+ attributeStats = AttributeMap(computeColumnStats(union, outputRows))))
}
/**
- * For each column, if all children have distinctCount, propagate the max distinctCount
- * across children. The result is capped by outputRows when available.
+ * For each output column, compute min/max, nullCount, and distinctCount in a single pass.
+ * Min and max are computed as the overall min/max across children (only for supported types),
+ * nullCount is summed, and distinctCount is estimated as the max across children (capped by
+ * outputRows).
*
* For UNION ALL, true distinctCount satisfies:
* max(dc_i) <= true_dc <= min(sum(dc_i), rowCount)
@@ -147,31 +66,49 @@ object UnionEstimation {
* (e.g. web_sales and catalog_sales reference the same date keys),
* so max is a reasonable approximation in the common case.
*/
- private def computeDistinctCountStats(
+ private def computeColumnStats(
union: Union,
outputRows: Option[BigInt]): Seq[(Attribute, ColumnStat)] = {
- val unionOutput = union.output
- val attrToComputeDistinctCount = union.children.map(_.output).transpose.zipWithIndex.filter {
- case (attrs, _) => attrs.zipWithIndex.forall {
- case (attr, childIndex) =>
- val attrStats = union.children(childIndex).stats.attributeStats
- attrStats.get(attr).isDefined && attrStats(attr).distinctCount.isDefined
+ // For each child, look up the ColumnStat for each of its output attributes.
+ // After transposing, maybeColStats(i) holds the ColumnStat (if available) contributed
+ // by the i-th child for the current output column.
+ union.output.zip(union.children.map(c => c.output.map(c.stats.attributeStats.get)).transpose)
+ .flatMap {
+ case (outputAttr, maybeColStats) =>
+ val maybeMinMax: Option[(Any, Any)] = if (isTypeSupported(outputAttr.dataType)) {
+ val statComparator = createStatComparator(outputAttr.dataType)
+ val initial = maybeColStats.head
+ .collect { case s if s.hasMinMaxStats => (s.min.get, s.max.get) }
+ maybeColStats.tail.foldLeft(initial) {
+ case (Some((minVal, maxVal)), Some(s)) if s.hasMinMaxStats =>
+ Some((
+ if (statComparator(s.min.get, minVal)) s.min.get else minVal,
+ if (statComparator(maxVal, s.max.get)) s.max.get else maxVal
+ ))
+ case _ => None
+ }
+ } else None
+
+ val maybeNullCount = maybeColStats.foldLeft(Option(BigInt(0))) {
+ case (Some(total), Some(s)) => s.nullCount.map(total + _)
+ case _ => None
+ }
+
+ val maybeDistinctCount = maybeColStats.foldLeft(Option(BigInt(0))) {
+ case (Some(currentMax), Some(s)) => s.distinctCount.map(currentMax.max)
+ case _ => None
+ }
+ // Cap distinctCount by outputRows if available
+ val cappedDistinctCount = maybeDistinctCount.map(max => outputRows.fold(max)(max.min))
+
+ if (maybeMinMax.isDefined || maybeNullCount.isDefined || cappedDistinctCount.isDefined) {
+ Some(outputAttr -> ColumnStat(
+ min = maybeMinMax.map(_._1),
+ max = maybeMinMax.map(_._2),
+ nullCount = maybeNullCount,
+ distinctCount = cappedDistinctCount
+ ))
+ } else None
}
- }
- attrToComputeDistinctCount.map {
- case (attrs, outputIndex) =>
- val maxDistinctCount = attrs.zipWithIndex.foldLeft[BigInt](0) {
- case (currentMax, (attr, childIndex)) =>
- val colStat = union.children(childIndex).stats.attributeStats(attr)
- currentMax.max(colStat.distinctCount.get)
- }
- // Cap distinctCount by outputRows if available
- val cappedDistinctCount = outputRows match {
- case Some(rows) => maxDistinctCount.min(rows)
- case None => maxDistinctCount
- }
- val newStat = ColumnStat(distinctCount = Some(cappedDistinctCount))
- unionOutput(outputIndex) -> newStat
- }
}
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org