marin-ma commented on code in PR #12511:
URL: https://github.com/apache/gluten/pull/12511#discussion_r3602190627
##########
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:
`GpuHashShuffleWriterType` no longer represents a shuffle writer type and
should be removed. Will do the cleanup in the following pr.
--
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]
