Copilot commented on code in PR #12511:
URL: https://github.com/apache/gluten/pull/12511#discussion_r3581197207
##########
gluten-substrait/src/main/scala/org/apache/spark/sql/execution/ColumnarShuffleExchangeExec.scala:
##########
@@ -16,29 +16,217 @@
*/
package org.apache.spark.sql.execution
+import org.apache.gluten.backendsapi.BackendsApiManager
+import org.apache.gluten.config.{GpuHashShuffleWriterType, ShuffleWriterType}
+import org.apache.gluten.execution.{CPUStageMode, GPUStageMode,
StageExecutionMode, ValidatablePlan, ValidationResult}
+import org.apache.gluten.extension.columnar.transition.Convention
import org.apache.gluten.sql.shims.SparkShimLoader
+import org.apache.spark._
import org.apache.spark.internal.Logging
+import org.apache.spark.rdd.RDD
+import org.apache.spark.serializer.Serializer
+import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions.Attribute
-import org.apache.spark.sql.catalyst.plans.physical._
+import org.apache.spark.sql.catalyst.plans.logical.Statistics
+import org.apache.spark.sql.catalyst.plans.physical.{SinglePartition, _}
+import org.apache.spark.sql.catalyst.util.truncatedString
import org.apache.spark.sql.execution.exchange._
+import org.apache.spark.sql.execution.metric.SQLShuffleWriteMetricsReporter
+import org.apache.spark.sql.metric.SQLColumnarShuffleReadMetricsReporter
+import org.apache.spark.sql.vectorized.ColumnarBatch
+
+import scala.concurrent.Future
case class ColumnarShuffleExchangeExec(
override val outputPartitioning: Partitioning,
child: SparkPlan,
shuffleOrigin: ShuffleOrigin = ENSURE_REQUIREMENTS,
projectOutputAttributes: Seq[Attribute],
- advisoryPartitionSize: Option[Long] = None)
- extends ColumnarShuffleExchangeExecBase(outputPartitioning, child,
projectOutputAttributes) {
+ advisoryPartitionSize: Option[Long] = None,
+ mapperStageMode: Option[StageExecutionMode] = None,
+ reducerStageMode: Option[StageExecutionMode] = None)
+ extends ShuffleExchangeLike
+ with ValidatablePlan {
+
+ override def nodeName: String = "ColumnarShuffleExchange" + {
+ if (mapperStageMode.isDefined) {
+ if (conf.adaptiveExecutionEnabled) {
+ // In AQE, the reducer stage mode is set in the downstream query stage.
+ // It is shown in the ColumnarAQEShuffleReaderExec node.
+ s"(${mapperStageMode.get.name})"
+ } else {
+ // Mapper and reducer stage modes should be set together when AQE is
disabled.
+ if (reducerStageMode.isEmpty) {
+ throw new IllegalStateException(
+ "Reducer stage mode is not defined in ColumnarShuffleExchangeExec
when AQE is disabled")
+ }
+ s"(${mapperStageMode.get.name}, ${reducerStageMode.get.name})"
+ }
+ } else {
+ ""
+ }
+ }
+
+ private[sql] lazy val writeMetrics =
+ SQLShuffleWriteMetricsReporter.createShuffleWriteMetrics(sparkContext)
+
+ private[sql] lazy val readMetrics =
+
SQLColumnarShuffleReadMetricsReporter.createShuffleReadMetrics(sparkContext)
+
+ lazy val shuffleWriterType: ShuffleWriterType = getShuffleWriterType
+
+ // super.stringArgs ++ Iterator(output.map(o =>
s"${o}#${o.dataType.simpleString}"))
+ lazy val serializer: Serializer =
BackendsApiManager.getSparkPlanExecApiInstance
+ .createColumnarBatchSerializer(schema, metrics, shuffleWriterType)
+
+ // Note: "metrics" is made transient to avoid sending driver-side metrics to
tasks.
+ @transient override lazy val metrics =
+ BackendsApiManager.getMetricsApiInstance
+ .genColumnarShuffleExchangeMetrics(
+ sparkContext,
+ shuffleWriterType) ++ readMetrics ++ writeMetrics
+
+ @transient lazy val inputColumnarRDD: RDD[ColumnarBatch] =
child.executeColumnar()
+
+ // 'mapOutputStatisticsFuture' is only needed when enable AQE.
+ @transient override lazy val mapOutputStatisticsFuture:
Future[MapOutputStatistics] = {
+ if (inputColumnarRDD.getNumPartitions == 0) {
+ Future.successful(null)
+ } else {
+ sparkContext.submitMapStage(columnarShuffleDependency)
+ }
+ }
- override def nodeName: String = "ColumnarExchange"
+ /**
+ * A [[ShuffleDependency]] that will partition rows of its child based on
the partitioning scheme
+ * defined in `newPartitioning`. Those partitions of the returned
ShuffleDependency will be the
+ * input of shuffle.
+ */
+ @transient
+ lazy val columnarShuffleDependency: ShuffleDependency[Int, ColumnarBatch,
ColumnarBatch] = {
+ BackendsApiManager.getSparkPlanExecApiInstance.genShuffleDependency(
+ inputColumnarRDD,
+ child.output,
+ projectOutputAttributes,
+ outputPartitioning,
+ serializer,
+ writeMetrics,
+ metrics,
+ shuffleWriterType)
+ }
+
+ var cachedShuffleRDD: ShuffledColumnarBatchRDD = _
+
+ override protected def doValidateInternal(): ValidationResult = {
+ val validation = BackendsApiManager.getValidatorApiInstance
+ .doColumnarShuffleExchangeExecValidate(output, outputPartitioning, child)
+ if (validation.nonEmpty) {
+ return ValidationResult.failed(
+ s"Found schema check failure for schema ${child.schema} due to:
${validation.get}")
+ }
+ outputPartitioning match {
+ case _: HashPartitioning => ValidationResult.succeeded
+ case _: RangePartitioning => ValidationResult.succeeded
+ case SinglePartition => ValidationResult.succeeded
+ case _: RoundRobinPartitioning => ValidationResult.succeeded
+ case _ =>
+ ValidationResult.failed(
+ s"Unsupported partitioning
${outputPartitioning.getClass.getSimpleName}")
+ }
+ }
+
+ override def numMappers: Int = inputColumnarRDD.getNumPartitions
+
+ override def numPartitions: Int =
columnarShuffleDependency.partitioner.numPartitions
+
+ override def runtimeStatistics: Statistics = {
+ val dataSize = metrics("dataSize").value
+ val rowCount =
metrics(SQLShuffleWriteMetricsReporter.SHUFFLE_RECORDS_WRITTEN).value
+ Statistics(dataSize, Some(rowCount))
+ }
+
+ def getShuffleWriterType: ShuffleWriterType = {
+ mapperStageMode match {
+ case Some(GPUStageMode) =>
+ GpuHashShuffleWriterType
+ case _ =>
+ BackendsApiManager.getSparkPlanExecApiInstance.getShuffleWriterType(
+ outputPartitioning,
+ output)
+ }
+ }
Review Comment:
`getShuffleWriterType` forces `GpuHashShuffleWriterType` whenever
`mapperStageMode` is `GPUStageMode`, ignoring the backend/config-selected
shuffle writer type (e.g. Sort/RSS/Celeborn). This can unintentionally switch
away from a required/selected writer implementation and break shuffle-manager
integrations.
Consider only switching to `GpuHashShuffleWriterType` when the backend
already selected `HashShuffleWriterType`, and otherwise preserve the backend
decision.
##########
gluten-substrait/src/main/scala/org/apache/gluten/execution/ColumnarCollectTailBaseExec.scala:
##########
@@ -101,7 +101,9 @@ abstract class ColumnarCollectTailBaseExec(
metrics,
shuffleWriterType
),
- readMetrics
+ readMetrics,
+ // FIXME: pass proper StageExecutionMode
+ CPUStageMode
)
Review Comment:
This shuffle path is now stage-mode aware, but `ColumnarCollectTailBaseExec`
still hard-codes `CPUStageMode` (and has a `FIXME`). If this operator is part
of a GPU stage, it will force the CPU shuffle reader/writer/resize behavior and
can lead to inconsistent stage execution modes across shuffle boundaries.
Consider threading the correct `StageExecutionMode` into this operator (or
deriving it from stage tags) and passing it into `ShuffledColumnarBatchRDD`.
##########
gluten-substrait/src/main/scala/org/apache/gluten/execution/ColumnarCollectLimitBaseExec.scala:
##########
@@ -102,7 +102,9 @@ abstract class ColumnarCollectLimitBaseExec(
metrics,
shuffleWriterType
),
- readMetrics
+ readMetrics,
+ // FIXME: pass proper StageExecutionMode
+ CPUStageMode
)
Review Comment:
This shuffle path is now stage-mode aware, but
`ColumnarCollectLimitBaseExec` still hard-codes `CPUStageMode` (and has a
`FIXME`). If this operator runs inside a GPU stage, it will force the CPU
shuffle reader/writer/resize behavior and can reintroduce inconsistent shuffle
execution modes.
It would be better to thread the relevant `StageExecutionMode` into this
operator (or infer it consistently from the surrounding stage tags) and pass it
into `ShuffledColumnarBatchRDD` instead of hard-coding CPU.
##########
gluten-substrait/src/main/scala/org/apache/spark/sql/execution/adaptive/ColumnarAQEShuffleReadExec.scala:
##########
@@ -0,0 +1,288 @@
+/*
+ * 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.execution.adaptive
+
+import org.apache.gluten.execution.StageExecutionMode
+
+import org.apache.spark.SparkException
+import org.apache.spark.rdd.RDD
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression}
+import org.apache.spark.sql.catalyst.plans.physical.{CoalescedBoundary,
CoalescedHashPartitioning, HashPartitioning, Partitioning, RangePartitioning,
RoundRobinPartitioning, SinglePartition, UnknownPartitioning}
+import org.apache.spark.sql.catalyst.trees.CurrentOrigin
+import org.apache.spark.sql.execution._
+import org.apache.spark.sql.execution.exchange.{ReusedExchangeExec,
ShuffleExchangeLike}
+import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics}
+import org.apache.spark.sql.vectorized.ColumnarBatch
+
+import scala.collection.mutable.ArrayBuffer
+
+/**
+ * A wrapper of shuffle query stage, which follows the given partition
arrangement.
+ *
+ * @param child
+ * It is usually `ShuffleQueryStageExec`, but can be the shuffle exchange
node during
+ * canonicalization.
+ * @param partitionSpecs
+ * The partition specs that defines the arrangement, requires at least one
partition.
+ */
+case class ColumnarAQEShuffleReadExec private (
+ child: SparkPlan,
+ partitionSpecs: Seq[ShufflePartitionSpec],
+ executionMode: StageExecutionMode,
+ isWrapper: Boolean)
+ extends UnaryExecNode {
+ assert(partitionSpecs.nonEmpty, s"${getClass.getSimpleName} requires at
least one partition")
+
+ // If this is to read shuffle files locally, then all partition specs should
be
+ // `PartialMapperPartitionSpec`.
+ if (partitionSpecs.exists(_.isInstanceOf[PartialMapperPartitionSpec])) {
+ assert(partitionSpecs.forall(_.isInstanceOf[PartialMapperPartitionSpec]))
+ }
+
+ override def nodeName: String = super.nodeName + s"(${executionMode.name})"
+
+ override def supportsColumnar: Boolean = child.supportsColumnar
+
+ override def output: Seq[Attribute] = child.output
+
+ override lazy val outputPartitioning: Partitioning = {
+ // If it is a local shuffle read with one mapper per task, then the output
partitioning is
+ // the same as the plan before shuffle.
+ // TODO this check is based on assumptions of callers' behavior but is
sufficient for now.
+ if (
+ partitionSpecs.forall(_.isInstanceOf[PartialMapperPartitionSpec]) &&
+
partitionSpecs.map(_.asInstanceOf[PartialMapperPartitionSpec].mapIndex).toSet.size
==
+ partitionSpecs.length
+ ) {
+ child match {
+ case ShuffleQueryStageExec(_, s: ShuffleExchangeLike, _) =>
+ s.child.outputPartitioning
+ case ShuffleQueryStageExec(_, r @ ReusedExchangeExec(_, s:
ShuffleExchangeLike), _) =>
+ s.child.outputPartitioning match {
+ case e: Expression => r.updateAttr(e).asInstanceOf[Partitioning]
+ case other => other
+ }
+ case _ =>
+ throw new IllegalStateException("operating on canonicalization plan")
+ }
+ } else if (isCoalescedRead) {
+ // For coalesced shuffle read, the data distribution is not changed,
only the number of
+ // partitions is changed.
+ child.outputPartitioning match {
+ case h: HashPartitioning =>
+ val partitions = partitionSpecs.map {
+ case CoalescedPartitionSpec(start, end, _) =>
CoalescedBoundary(start, end)
+ // Can not happend due to isCoalescedRead
Review Comment:
Typo in comment: "Can not happend" should be "Cannot happen".
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
