diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala index bea86501e6f3a..83fd1863e1508 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala @@ -23,6 +23,7 @@ import java.util.concurrent.TimeUnit._ import scala.collection.mutable import scala.concurrent.ExecutionContext import scala.concurrent.duration.Duration +import scala.reflect.ClassTag import org.apache.spark.{InterruptibleIterator, SparkException, TaskContext} import org.apache.spark.internal.LogKeys @@ -1011,7 +1012,7 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup // Intersect across all children, anchored on the first child's set. Every surviving member // is shared by all children, so their `numPartitions` agree; a `PartitioningCollection` // built from a subset of one child's members therefore keeps its uniform-numPartitions - // invariant, and the co-located `doExecute` arm's invariant holds. + // invariant, and the co-located arm in `unionRDDs` keeps its invariant. val head = candidateSets.head val intersection = head.filter { c => candidateSets.tail.forall(_.exists(comparePartitioning(c, _))) @@ -1024,8 +1025,8 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup } // True when the codegen path applies: `outputPartitioning` is `UnknownPartitioning`, - // and `unionedInputRDD` matches the semantics of `sparkContext.union(...)` in `doExecute`. - // A `KeyedPartitioning` union also uses `sparkContext.union(...)` in `doExecute`, but + // and `unionedInputRDD` matches the semantics of `sparkContext.union(...)` in `unionRDDs`. + // A `KeyedPartitioning` union also uses `sparkContext.union(...)` in `unionRDDs`, but // codegen is disabled for it (`supportCodegenFailureReason` reports "partitioning-aware"): // the per-partition key descriptor is consumed by a downstream `GroupPartitionsExec`, and // keeping these unions out of whole-stage codegen matches the `HashPartitioning` union case. @@ -1234,8 +1235,14 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup // decides which output columns to materialize. override def usedInputs: AttributeSet = AttributeSet.empty - protected override def doExecute(): RDD[InternalRow] = { - outputPartitioning match { + // Shared by `doExecute` and `doExecuteColumnar` so the two cannot report one partitioning and + // build another. `outputPartitioning` is read once, before the children execute: a child's + // partitioning can sharpen once it has run, as `InMemoryTableScanExec` does over a + // not-yet-materialized AQE cached plan. + private def unionRDDs[T: ClassTag](executeChild: SparkPlan => RDD[T]): RDD[T] = { + val partitioning = outputPartitioning + val rdds = children.map(executeChild) + partitioning match { case _: UnknownPartitioning | _: KeyedPartitioning => // An `UnknownPartitioning` union simply concatenates its children. A // `KeyedPartitioning` union does the same: its merged partition keys describe the @@ -1243,24 +1250,23 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup // `GroupPartitionsExec` regroups partitions that share a key. This differs from an // index-co-locatable partitioning (e.g. `HashPartitioning`), where a partitioning-aware // union RDD interleaves same-index partitions across children. - sparkContext.union(children.map(_.execute())) + sparkContext.union(rdds) case _ => // This union has a known, index-co-locatable partitioning, i.e., its children have the // same partitioning in semantics so this union can choose not to change the partitioning // by using a custom partitioning aware union RDD. - val nonEmptyRdds = children.map(_.execute()).filter(!_.partitions.isEmpty) new SQLPartitioningAwareUnionRDD( - sparkContext, nonEmptyRdds, outputPartitioning.numPartitions) + sparkContext, rdds.filter(!_.partitions.isEmpty), partitioning.numPartitions) } } + protected override def doExecute(): RDD[InternalRow] = unionRDDs(_.execute()) + override def supportsColumnar: Boolean = children.forall(_.supportsColumnar) override def supportsRowBased: Boolean = children.forall(_.supportsRowBased) - protected override def doExecuteColumnar(): RDD[ColumnarBatch] = { - sparkContext.union(children.map(_.executeColumnar())) - } + protected override def doExecuteColumnar(): RDD[ColumnarBatch] = unionRDDs(_.executeColumnar()) override protected def withNewChildrenInternal(newChildren: IndexedSeq[SparkPlan]): UnionExec = copy(children = newChildren) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala index 32b6da3a30b15..4cce6129dc5ac 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala @@ -1772,6 +1772,40 @@ class DataFrameSetOperationsSuite extends SharedSparkSession with AdaptiveSparkP } } + test("SPARK-59141: columnar union interleaves the partitions it reports as co-located") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BUCKETED_SCAN_ENABLED.key -> "false") { + withTable("t1", "t2") { + spark.range(0, 20, 1, 1).selectExpr("id % 5 AS k") + .write.bucketBy(4, "k").saveAsTable("t1") + spark.range(20, 40, 1, 1).selectExpr("id % 5 AS k") + .write.bucketBy(4, "k").saveAsTable("t2") + + val df = sql( + "SELECT k, count(*) c FROM (SELECT k FROM t1 UNION ALL SELECT k FROM t2) GROUP BY k") + val plan = df.queryExecution.executedPlan + // `FileSourceScanExec` overrides `supportsColumnar` but not `supportsRowBased`, whose + // default is its negation, so a bucketed batch-readable scan is columnar-only and so is + // this union. It therefore runs `doExecuteColumnar` under a `ColumnarToRowExec` while + // reporting the bucketed `HashPartitioning` that lets the aggregate drop its exchange. + val unionExec = plan.collect { case u: UnionExec => u } + assert(unionExec.size == 1) + assert(unionExec.head.supportsColumnar && !unionExec.head.supportsRowBased, + "this shape must take the columnar path, or the test exercises nothing") + assert(unionExec.head.outputPartitioning.isInstanceOf[HashPartitioning], + "this shape must report a co-locatable partitioning, or the test exercises nothing") + assert(plan.collect { case s: ShuffleExchangeExec => s }.isEmpty, + "the aggregate's exchange must have been dropped, or the test exercises nothing") + + // Interleaving puts a key's bucket from each table in one partition, so its eight rows + // land together; concatenating would split them and the exchange-free aggregate would + // report the key twice with four. + checkAnswer(df, (0L until 5L).map(k => Row(k, 8L))) + } + } + } + test("SPARK-57881: union partitioning - keyed partitioning") { withSQLConf("spark.sql.catalog.testcat" -> classOf[InMemoryCatalog].getName) { sql("CREATE TABLE testcat.ns.t1 (id bigint, data string) PARTITIONED BY (id)")