[AURON #2090] Convert CoalesceExec to native implementation. #2090

Author: guihuawen

spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/auron/ShimsImpl.scala

@@ -283,6 +283,9 @@ class ShimsImpl extends Shims with Logging {
   override def createNativeFilterExec(condition: Expression, child: SparkPlan): NativeFilterBase =
     NativeFilterExec(condition, child)
 
+  def createNativeCoalesceExec(numPartitions: Int, child: SparkPlan): NativeCoalesceBase =
+    NativeCoalesceExec(numPartitions, child)
+
   override def createNativeGenerateExec(
       generator: Generator,
       requiredChildOutput: Seq[Attribute],

spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeCoalesceExec.scala

@@ -0,0 +1,36 @@
+/*
+ * 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.auron.plan
+
+import org.apache.spark.sql.catalyst.expressions.Attribute
+import org.apache.spark.sql.execution.SparkPlan
+
+import org.apache.auron.sparkver
+
+case class NativeCoalesceExec(numPartitions: Int, child: SparkPlan)
+    extends NativeCoalesceBase(numPartitions: Int, child: SparkPlan) {
+  @sparkver("3.2 / 3.3 / 3.4 / 3.5 / 4.0 / 4.1")
+  override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan =
+    copy(child = newChild)
+
+  @sparkver("3.0 / 3.1")
+  override def withNewChildren(newChildren: Seq[SparkPlan]): SparkPlan =
+    copy(child = newChildren.head)
+
+  override def output: Seq[Attribute] =
+    child.output
+}

spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronNativeCoalesceExecSuite.scala

@@ -0,0 +1,41 @@
+/*
+ * 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.auron
+
+import org.apache.spark.sql.{AuronQueryTest, Row}
+import org.apache.spark.sql.execution.auron.plan.NativeCoalesceExec
+
+class AuronNativeCoalesceExecSuite extends AuronQueryTest with BaseAuronSQLSuite {
+  import testImplicits._
+
+  test("test CoalesceExec to native") {
+    withSQLConf("spark.auron.enable.coalesce" -> "true") {
+      Seq((1, 2, "test test"))
+        .toDF("c1", "c2", "part")
+        .createOrReplaceTempView("coalesce_table1")
+      val df =
+        spark.sql("select /*+ coalesce(2)*/ a.c1, a.c2 from coalesce_table1 a ")
+
+      checkAnswer(df, Seq(Row(1, 2)))
+      assert(collectFirst(df.queryExecution.executedPlan) {
+        case coalesceExec: NativeCoalesceExec =>
+          coalesceExec
+      }.isDefined)
+
+    }
+  }
+}

spark-extension/src/main/java/org/apache/auron/spark/configuration/SparkAuronConfiguration.java

@@ -411,6 +411,11 @@ public class SparkAuronConfiguration extends AuronConfiguration {
             .withDescription("Enable AggregateExec operation conversion to native Auron implementations.")
             .withDefaultValue(true);
 
+    public static final ConfigOption<Boolean> ENABLE_COALESEC = new SQLConfOption<>(Boolean.class)
+            .withKey("auron.enable.coalesce")
+            .withCategory("Operator Supports")
+            .withDescription("Enable CoalesceExec operation conversion to native Auron implementations.")
+            .withDefaultValue(true);
     public static final ConfigOption<Boolean> ENABLE_EXPAND = new SQLConfOption<>(Boolean.class)
             .withKey("auron.enable.expand")
             .withCategory("Operator Supports")

spark-extension/src/main/scala/org/apache/spark/sql/auron/AuronConvertStrategy.scala

@@ -165,6 +165,8 @@ object AuronConvertStrategy extends Logging {
         e.setTagValue(convertStrategyTag, AlwaysConvert)
       case e: GenerateExec if isNative(e.child) =>
         e.setTagValue(convertStrategyTag, AlwaysConvert)
+      case e: CoalesceExec if isNative(e.child) =>
+        e.setTagValue(convertStrategyTag, AlwaysConvert)
       case e: ObjectHashAggregateExec if isNative(e.child) =>
         e.setTagValue(convertStrategyTag, AlwaysConvert)
       case e: LocalTableScanExec =>

spark-extension/src/main/scala/org/apache/spark/sql/auron/AuronConverters.scala

@@ -95,6 +95,7 @@ object AuronConverters extends Logging {
     SparkAuronConfiguration.ENABLE_TAKE_ORDERED_AND_PROJECT.get()
   def enableCollectLimit: Boolean = SparkAuronConfiguration.ENABLE_COLLECT_LIMIT.get()
   def enableAggr: Boolean = SparkAuronConfiguration.ENABLE_AGGR.get()
+  def enableCoalesec: Boolean = SparkAuronConfiguration.ENABLE_COALESEC.get()
   def enableExpand: Boolean = SparkAuronConfiguration.ENABLE_EXPAND.get()
   def enableWindow: Boolean = SparkAuronConfiguration.ENABLE_WINDOW.get()
   def enableWindowGroupLimit: Boolean = SparkAuronConfiguration.ENABLE_WINDOW_GROUP_LIMIT.get()
@@ -201,6 +202,10 @@ object AuronConverters extends Logging {
         }
         convertedAgg
 
+      case e: CoalesceExec if enableCoalesec => // coalesec
+        val convertedCoalesce = tryConvert(e, convertCoalesceExec)
+        convertedCoalesce
+
       case e: ObjectHashAggregateExec if enableAggr => // object hash aggregate
         val convertedAgg = tryConvert(e, convertObjectHashAggregateExec)
         if (!e.getTagValue(convertibleTag).contains(true)) {
@@ -308,6 +313,8 @@ object AuronConverters extends Logging {
           "Conversion disabled: spark.auron.enable.local.table.scan=false."
         case _: DataWritingCommandExec if !enableDataWriting =>
           "Conversion disabled: spark.auron.enable.data.writing=false."
+        case _: CoalesceExec if !enableCoalesec =>
+          "Conversion disabled: spark.auron.enable.coalesce=false."
         case _ =>
           s"${exec.getClass.getSimpleName} is not supported yet."
       }
@@ -745,6 +752,10 @@ object AuronConverters extends Logging {
     Shims.get.createNativeCollectLimitExec(limit, offset, exec.child)
   }
 
+  def convertCoalesceExec(exec: CoalesceExec): SparkPlan = {
+    Shims.get.createNativeCoalesceExec(exec.numPartitions, exec.child)
+  }
+
   def convertHashAggregateExec(exec: HashAggregateExec): SparkPlan = {
     // split non-trivial children exprs in partial-agg to a ProjectExec
     // for enabling filter-project optimization in native side

spark-extension/src/main/scala/org/apache/spark/sql/auron/NativeRDD.scala

@@ -26,7 +26,7 @@ import org.apache.spark.Partitioner
 import org.apache.spark.SparkContext
 import org.apache.spark.TaskContext
 import org.apache.spark.internal.Logging
-import org.apache.spark.rdd.RDD
+import org.apache.spark.rdd.{CoalescedRDDPartition, RDD}
 import org.apache.spark.sql.catalyst.InternalRow
 
 import org.apache.auron.metric.SparkMetricNode
@@ -78,6 +78,32 @@ class NativeRDD(
   }
 }
 
+class CoalesceNativeRDD(
+    @transient private val rddSparkContext: SparkContext,
+    rddDependencies: Seq[Dependency[_]],
+    partitions: Array[Partition],
+    friendlyName: String)
+    extends NativeRDD(
+      rddSparkContext,
+      metrics = SparkMetricNode(Map.empty, Seq(), None),
+      rddPartitions = partitions,
+      rddPartitioner = None,
+      rddDependencies,
+      rddShuffleReadFull = false,
+      nativePlan = (_, _) => null,
+      friendlyName)
+    with Logging
+    with Serializable {
+
+  override protected def getPartitions: Array[Partition] = partitions
+
+  override def compute(split: Partition, context: TaskContext): Iterator[InternalRow] = {
+    split.asInstanceOf[CoalescedRDDPartition].parents.iterator.flatMap { parentPartition =>
+      firstParent[InternalRow].iterator(parentPartition, context)
+    }
+  }
+}
+
 class EmptyNativeRDD(@transient private val rddSparkContext: SparkContext)
     extends NativeRDD(
       rddSparkContext = rddSparkContext,

spark-extension/src/main/scala/org/apache/spark/sql/auron/Shims.scala

@@ -139,6 +139,8 @@ abstract class Shims {
       offset: Int,
       child: SparkPlan): NativeCollectLimitBase
 
+  def createNativeCoalesceExec(numPartitions: Int, child: SparkPlan): NativeCoalesceBase
+
   def createNativeParquetInsertIntoHiveTableExec(
       cmd: InsertIntoHiveTable,
       child: SparkPlan): NativeParquetInsertIntoHiveTableBase