[SPARK-52921][SQL] Specify outputPartitioning for UnionExec for same output partitoning as children operators #51623

viirya · 2025-07-22T21:34:41Z

What changes were proposed in this pull request?

This patch updates outputPartitioning for UnionExec operator for the cases that the output partitionings of its children are the same. So the output partitioning can be known.

Why are the changes needed?

Currently the output partitioning of UnionExec is simply unknown. But if the output partitionings of its children are known to be the same, we can make the union output as the same output partitioning with the children.

But different to the RDD-level PartitionerAwareUnionRDD, which only considers the RDD partitioner, SQL operators' outputPartitioning doesn't really rely on RDD's partitioner.

Thus, this patch introduces SQLPartitioningAwareUnionRDD which is a specified union RDD only for SQL UnionExec if the output partitioning is to be the same as its children. Similar to PartitionerAwareUnionRDD, it groups the partitions of parent RDDs at corresponding index together but it doesn't require that parent RDDs to have same partitioner.

Does this PR introduce any user-facing change?

No

How was this patch tested?

Unit test.

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

No

dongjoon-hyun · 2025-07-22T21:39:01Z

core/src/main/scala/org/apache/spark/SparkContext.scala

    new ReliableCheckpointRDD[T](this, path)
  }

+  protected[spark] def isPartitionerAwareUnion[T: ClassTag](rdds: Seq[RDD[T]]): Boolean = {


Could you add a comment about the assumption, rdds.filter(!_.partitions.isEmpty)? Otherwise, it may cause correctness issues later if we use this blindly.

Otherwise, we had better include the assumption inside this method.

Added comment and a check.

dongjoon-hyun · 2025-07-22T21:40:48Z

sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala

+  private lazy val childrenRDDs = children.map(_.execute())
+
+  override def outputPartitioning: Partitioning = {
+    val nonEmptyRdds = childrenRDDs.filter(!_.partitions.isEmpty)


ditto. We can remove this too if isPartitionerAwareUnion has the logic.

Because SparkContext.union uses nonEmptyRdds, so I didn't move nonEmptyRdds logic into isPartitionerAwareUnion. I leave to the callers to pass in non empty rdds.

Got it~ Thank you for the explanation.

dongjoon-hyun · 2025-07-22T21:41:00Z

cc @peter-toth

dongjoon-hyun · 2025-07-22T21:56:22Z

core/src/main/scala/org/apache/spark/SparkContext.scala


+  // Note that input rdds must be all non-empty, i.e., rdds.filter(_.partitions.isEmpty).isEmpty
+  protected[spark] def isPartitionerAwareUnion[T: ClassTag](rdds: Seq[RDD[T]]): Boolean = {
+    assert(!rdds.exists(_.partitions.isEmpty), "Must not have empty RDDs")


dongjoon-hyun

+1, LGTM. Thank you, @viirya .

viirya · 2025-07-23T06:58:07Z

sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala

+      // child operator will be replaced by Spark in query planning later, in other
+      // words, `execute` won't be actually called on them during the execution of
+      // this plan. So we can safely return the default partitioning.
+      case e if NonFatal(e) => super.outputPartitioning


This handles nodes that don't implement execute method. The reason is described like the comment said.

peter-toth · 2025-07-23T10:08:44Z

core/src/main/scala/org/apache/spark/SparkContext.scala

+  protected[spark] def isPartitionerAwareUnion[T: ClassTag](rdds: Seq[RDD[T]]): Boolean = {
+    assert(!rdds.exists(_.partitions.isEmpty), "Must not have empty RDDs")
+    val partitioners = rdds.flatMap(_.partitioner).toSet
+    rdds.forall(_.partitioner.isDefined) && partitioners.size == 1


It seems we don't need the partitioners set before the forall isDefined check.

peter-toth

LGTM, just a minor nit.

viirya · 2025-07-23T16:20:58Z

Hmm, there are a few test failures, I will take a look.

viirya · 2025-07-23T20:11:39Z

sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala

+        "default partitioning.")
+      .version("4.1.0")
+      .booleanConf
+      .createWithDefault(true)


For safety, added an internal config for it.

dongjoon-hyun · 2025-07-25T18:12:31Z

Could you re-trigger ThriftServer tests?

[info] *** 3 TESTS FAILED ***
[error] Failed tests:
[error] 	org.apache.spark.sql.hive.thriftserver.ThriftServerQueryTestSuite

viirya · 2025-07-25T18:46:35Z

Could you re-trigger ThriftServer tests?

[info] *** 3 TESTS FAILED ***
[error] Failed tests:
[error] 	org.apache.spark.sql.hive.thriftserver.ThriftServerQueryTestSuite

They are related test failures. I'm investigating them. Thanks.

viirya · 2025-07-26T00:47:17Z

sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/BroadcastExchangeExec.scala

+  override def resetMetrics(): Unit = {
+    // no-op
+    // BroadcastExchangeExec after materialized won't be materialized again, so we should not
+    // reset the metrics. Otherwise, we will lose the metrics collected in the broadcast job.
+  }


I spent a lot time debugging the remaining test failures. When there is broadcast exchange operator, AQE empty relation propagation rule will produce incorrect query plan around it. It is caused by this reset metrics method.

I think it is valuable to be a separate PR: #51673

#51673 also has related test case.

viirya · 2025-07-26T01:46:32Z

sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala

+
+      try {
+        val nonEmptyRdds = childrenRDDs.filter(!_.partitions.isEmpty)
+        if (sparkContext.isPartitionerAwareUnion(nonEmptyRdds)) {


Actually this only covers limited cases like reused shuffles they have the same partitioner. I would like to extend this to SQL cases, i.e., the outputPartitioning is same or compatible for Union's children. But I will leave to follow up works.

violet-nspct · 2025-07-26T03:22:09Z

@viirya Should unit tests be added in DataFrameSetOperationsSuite.scala to cover the following scenarios?

// Core Functionality
test("union partitioning - different partitioners") {
  // Covers: Different partitioner scenarios
}

// Mixed Partitioners
test("union partitioning - mixed partitioners") {
  // Covers: Mixed partitioner scenarios
}

// Command Handling
test("union partitioning with commands") {
  // Covers: Command plan interactions
}

// Error Handling
test("union partitioning error handling") {
  // Covers: Error scenarios and fallback behavior
}

viirya · 2025-07-26T06:10:45Z

@viirya Should unit tests be added in DataFrameSetOperationsSuite.scala to cover the following scenarios?

Those cases are covered by existing tests. For example, there were test failures on union with commands before. I added the logic to handle command case to pass these tests.

violet-nspct · 2025-07-26T06:43:09Z

Thanks, @viirya
Can you please point me which tests cover the remaining three scenarios?

// Core Functionality
test("union partitioning - different partitioners") {
  // Covers: Different partitioner scenarios
}

// Mixed Partitioners
test("union partitioning - mixed partitioners") {
  // Covers: Mixed partitioner scenarios
}

// Error Handling
test("union partitioning error handling") {
  // Covers: Error scenarios and fallback behavior
}

viirya · 2025-07-26T06:51:17Z

Thanks, @viirya Can you please point me which tests cover the remaining three scenarios?

// Core Functionality
test("union partitioning - different partitioners") {
  // Covers: Different partitioner scenarios
}

// Mixed Partitioners
test("union partitioning - mixed partitioners") {
  // Covers: Mixed partitioner scenarios
}

// Error Handling
test("union partitioning error handling") {
  // Covers: Error scenarios and fallback behavior
}

They can be found in previous CI failures. For example:

HiveCompatibilitySuite's add_part_multiple covers union with commands:
https://github.com/viirya/spark-1/actions/runs/16457342432/job/46519111118

DynamicPartitionPruningHiveScanSuiteAEOn's SPARK-39338: Remove dynamic pruning subquery if pruningKey's references is empty covers error handling case:
https://github.com/viirya/spark-1/actions/runs/16457342432/job/46519111126

I think most union tests cover different/mixed partitioner case, because most union queries don't have same (rdd) partitioner.

peter-toth · 2025-07-26T09:20:04Z

sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala

    }
  }

+  private lazy val childrenRDDs = children.map(_.execute())


Actaully, does this mean that children executions are triggered to get outputPartitioning of an Union?
E.g. a simple explain to show the physical plan can now trigger execuion of union children?

Yes, so this approach has this drawback. So as I mentioned #51623 (comment), this doesn't cover SQL cases generally. I plan to extend this to deal with outputPartitioning of children, i.e., no need to invoke execute on children.

It will be done in follow up works.

I actually have the next PR ready locally. After this gets merged, I will open a new PR to improve it and get rid of this execute calls.

IMO that's a serious drawback. But if we can fix it in a follow-up PR right after this PR then I'm ok with merging. Or just update this PR with you local changes.

Okay, I updated to this PR.

Thanks. I can check it tomorrow.

github-actions bot added SQL CORE labels Jul 22, 2025

viirya marked this pull request as draft July 22, 2025 21:35

dongjoon-hyun reviewed Jul 22, 2025

View reviewed changes

viirya marked this pull request as ready for review July 22, 2025 23:00

viirya changed the title ~~[SPARK-XXXXX][SQL] Specify outputPartitioning for UnionExec for partitioner aware case~~ [SPARK-52921][SQL] Specify outputPartitioning for UnionExec for partitioner aware case Jul 22, 2025

dongjoon-hyun approved these changes Jul 22, 2025

View reviewed changes

viirya commented Jul 23, 2025

View reviewed changes

peter-toth reviewed Jul 23, 2025

View reviewed changes

peter-toth approved these changes Jul 23, 2025

View reviewed changes

viirya commented Jul 23, 2025

View reviewed changes

dongjoon-hyun approved these changes Jul 25, 2025

View reviewed changes

viirya commented Jul 26, 2025

View reviewed changes

viirya mentioned this pull request Jul 26, 2025

[SPARK-52962][SQL] BroadcastExchangeExec should not reset metrics #51673

Closed

viirya commented Jul 26, 2025

View reviewed changes

peter-toth reviewed Jul 26, 2025

View reviewed changes

viirya added 3 commits July 26, 2025 13:28

Specify outputPartitioning for UnionExec for partitioner aware case

9ea07fe

add comment

dee167f

add test

f19e1bf

abellina mentioned this pull request Jan 6, 2026

[FEA][AUDIT][SPARK-52921][SQL] Specify outputPartitioning for UnionExec for same output partitoning as children operators NVIDIA/spark-rapids#14083

Open

baibaichen mentioned this pull request Jan 7, 2026

[GLUTEN-11343][CORE][VL] Support Spark 4.1 UT apache/incubator-gluten#11353

Merged

baibaichen mentioned this pull request Jan 13, 2026

[VL] Track on Spark-4.1.x failed unit tests apache/incubator-gluten#11400

Open

This was referenced Jan 22, 2026

[RAS] GroupLeafExec does not preserve outputPartitioning apache/incubator-gluten#11468

Open

[GLUTEN-11400][CORE] Implement partitioning-aware union for ColumnarUnionExec apache/incubator-gluten#11455

Merged

[SPARK-52921][SQL] Specify outputPartitioning for UnionExec for same output partitoning as children operators #51623

[SPARK-52921][SQL] Specify outputPartitioning for UnionExec for same output partitoning as children operators #51623

Uh oh!

Conversation

viirya commented Jul 22, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What changes were proposed in this pull request?

Why are the changes needed?

Does this PR introduce any user-facing change?

How was this patch tested?

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

Uh oh!

dongjoon-hyun Jul 22, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dongjoon-hyun commented Jul 22, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dongjoon-hyun left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

peter-toth left a comment

Choose a reason for hiding this comment

Uh oh!

viirya commented Jul 23, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dongjoon-hyun commented Jul 25, 2025

Uh oh!

viirya commented Jul 25, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

viirya Jul 26, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

violet-nspct commented Jul 26, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

viirya commented Jul 26, 2025

Uh oh!

violet-nspct commented Jul 26, 2025

Uh oh!

viirya commented Jul 26, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

peter-toth Jul 26, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

viirya Jul 26, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

viirya commented Jul 22, 2025 •

edited

Loading

dongjoon-hyun Jul 22, 2025 •

edited

Loading

viirya Jul 26, 2025 •

edited

Loading

violet-nspct commented Jul 26, 2025 •

edited

Loading

viirya commented Jul 26, 2025 •

edited

Loading

peter-toth Jul 26, 2025 •

edited

Loading

viirya Jul 26, 2025 •

edited

Loading