fix

apache · Aug 2, 2024 · 416c801 · 416c801
1 parent 0753524
commit 416c801
Show file tree

Hide file tree

Showing 6 changed files with 50 additions and 8 deletions.
diff --git a/backends-velox/src/main/scala/org/apache/gluten/extension/CollectRewriteRule.scala b/backends-velox/src/main/scala/org/apache/gluten/extension/CollectRewriteRule.scala
@@ -19,6 +19,7 @@ package org.apache.gluten.extension
 import org.apache.gluten.GlutenConfig
 import org.apache.gluten.expression.ExpressionMappings
 import org.apache.gluten.expression.aggregate.{VeloxCollectList, VeloxCollectSet}
+import org.apache.gluten.sql.shims.SparkShimLoader
 import org.apache.gluten.utils.LogicalPlanSelector
 
 import org.apache.spark.sql.SparkSession
@@ -55,10 +56,10 @@ case class CollectRewriteRule(spark: SparkSession) extends Rule[LogicalPlan] {
     case PhysicalAggregation(_, aggregateExpr, _, _)
         if !GlutenConfig.getConf.veloxObjectHashAggCollectRewriteEnabled =>
       val aggregateExpressions = aggregateExpr.map(expr => expr.asInstanceOf[AggregateExpression])
-      val useHash = Aggregate.supportsHashAggregate(
+      val useHash = SparkShimLoader.getSparkShims.supportsHashAggregate(
         aggregateExpressions.flatMap(_.aggregateFunction.aggBufferAttributes))
       val useObjectHash = plan.conf.useObjectHashAggregation &&
-        Aggregate.supportsObjectHashAggregate(aggregateExpressions)
+        SparkShimLoader.getSparkShims.supportsObjectHashAggregate(aggregateExpressions)
       useHash || !useObjectHash
     case _ => true
   }

diff --git a/shims/common/src/main/scala/org/apache/gluten/sql/shims/SparkShims.scala b/shims/common/src/main/scala/org/apache/gluten/sql/shims/SparkShims.scala
@@ -28,7 +28,7 @@ import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.catalog.BucketSpec
 import org.apache.spark.sql.catalyst.csv.CSVOptions
 import org.apache.spark.sql.catalyst.expressions.{Attribute, BinaryExpression, Expression}
-import org.apache.spark.sql.catalyst.expressions.aggregate.TypedImperativeAggregate
+import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, TypedImperativeAggregate}
 import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
 import org.apache.spark.sql.catalyst.plans.physical.{Distribution, Partitioning}
 import org.apache.spark.sql.catalyst.rules.Rule
@@ -267,4 +267,8 @@ trait SparkShims {
       DecimalType(math.min(integralLeastNumDigits + newScale, 38), newScale)
     }
   }
+
+  def supportsHashAggregate(aggregateBufferAttributes: Seq[Attribute]): Boolean
+
+  def supportsObjectHashAggregate(aggregateExpressions: Seq[AggregateExpression]): Boolean
 }
diff --git a/shims/spark32/src/main/scala/org/apache/gluten/sql/shims/spark32/Spark32Shims.scala b/shims/spark32/src/main/scala/org/apache/gluten/sql/shims/spark32/Spark32Shims.scala
@@ -28,7 +28,7 @@ import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.catalog.BucketSpec
 import org.apache.spark.sql.catalyst.csv.CSVOptions
 import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, BinaryExpression, Expression, InputFileBlockLength, InputFileBlockStart, InputFileName}
-import org.apache.spark.sql.catalyst.expressions.aggregate.TypedImperativeAggregate
+import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, TypedImperativeAggregate}
 import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
 import org.apache.spark.sql.catalyst.plans.physical.{Distribution, HashClusteredDistribution}
 import org.apache.spark.sql.catalyst.rules.Rule
@@ -37,6 +37,7 @@ import org.apache.spark.sql.catalyst.util.RebaseDateTime.RebaseSpec
 import org.apache.spark.sql.connector.catalog.Table
 import org.apache.spark.sql.connector.expressions.Transform
 import org.apache.spark.sql.execution.{FileSourceScanExec, PartitionedFileUtil, SparkPlan}
+import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, ObjectHashAggregateExec}
 import org.apache.spark.sql.execution.datasources._
 import org.apache.spark.sql.execution.datasources.FileFormatWriter.Empty2Null
 import org.apache.spark.sql.execution.datasources.parquet.ParquetFilters
@@ -284,4 +285,13 @@ class Spark32Shims extends SparkShims {
     val s = decimalType.scale
     DecimalType(p, if (toScale > s) s else toScale)
   }
+
+  override def supportsHashAggregate(aggregateBufferAttributes: Seq[Attribute]): Boolean = {
+    HashAggregateExec.supportsAggregate(aggregateBufferAttributes)
+  }
+
+  override def supportsObjectHashAggregate(
+      aggregateExpressions: Seq[AggregateExpression]): Boolean = {
+    ObjectHashAggregateExec.supportsAggregate(aggregateExpressions)
+  }
 }
diff --git a/shims/spark33/src/main/scala/org/apache/gluten/sql/shims/spark33/Spark33Shims.scala b/shims/spark33/src/main/scala/org/apache/gluten/sql/shims/spark33/Spark33Shims.scala
@@ -29,8 +29,8 @@ import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.catalog.BucketSpec
 import org.apache.spark.sql.catalyst.csv.CSVOptions
 import org.apache.spark.sql.catalyst.expressions._
-import org.apache.spark.sql.catalyst.expressions.aggregate.{BloomFilterAggregate, RegrR2, TypedImperativeAggregate}
-import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
+import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, BloomFilterAggregate, RegrR2, TypedImperativeAggregate}
+import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LogicalPlan}
 import org.apache.spark.sql.catalyst.plans.physical.{ClusteredDistribution, Distribution}
 import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.catalyst.util.CaseInsensitiveMap
@@ -365,4 +365,13 @@ class Spark33Shims extends SparkShims {
       RebaseSpec(LegacyBehaviorPolicy.CORRECTED)
     )
   }
+
+  override def supportsHashAggregate(aggregateBufferAttributes: Seq[Attribute]): Boolean = {
+    Aggregate.supportsHashAggregate(aggregateBufferAttributes)
+  }
+
+  override def supportsObjectHashAggregate(
+      aggregateExpressions: Seq[AggregateExpression]): Boolean = {
+    Aggregate.supportsObjectHashAggregate(aggregateExpressions)
+  }
 }
diff --git a/shims/spark34/src/main/scala/org/apache/gluten/sql/shims/spark34/Spark34Shims.scala b/shims/spark34/src/main/scala/org/apache/gluten/sql/shims/spark34/Spark34Shims.scala
@@ -32,7 +32,7 @@ import org.apache.spark.sql.catalyst.catalog.BucketSpec
 import org.apache.spark.sql.catalyst.csv.CSVOptions
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.expressions.aggregate._
-import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
+import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LogicalPlan}
 import org.apache.spark.sql.catalyst.plans.physical.{ClusteredDistribution, Distribution, KeyGroupedPartitioning, Partitioning}
 import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.catalyst.util.{CaseInsensitiveMap, InternalRowComparableWrapper, TimestampFormatter}
@@ -493,4 +493,13 @@ class Spark34Shims extends SparkShims {
       RebaseSpec(LegacyBehaviorPolicy.CORRECTED)
     )
   }
+
+  override def supportsHashAggregate(aggregateBufferAttributes: Seq[Attribute]): Boolean = {
+    Aggregate.supportsHashAggregate(aggregateBufferAttributes)
+  }
+
+  override def supportsObjectHashAggregate(
+      aggregateExpressions: Seq[AggregateExpression]): Boolean = {
+    Aggregate.supportsObjectHashAggregate(aggregateExpressions)
+  }
 }
diff --git a/shims/spark35/src/main/scala/org/apache/gluten/sql/shims/spark35/Spark35Shims.scala b/shims/spark35/src/main/scala/org/apache/gluten/sql/shims/spark35/Spark35Shims.scala
@@ -31,7 +31,7 @@ import org.apache.spark.sql.catalyst.catalog.BucketSpec
 import org.apache.spark.sql.catalyst.csv.CSVOptions
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.expressions.aggregate._
-import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
+import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LogicalPlan}
 import org.apache.spark.sql.catalyst.plans.physical.{ClusteredDistribution, Distribution, KeyGroupedPartitioning, Partitioning}
 import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.catalyst.types.DataTypeUtils
@@ -518,4 +518,13 @@ class Spark35Shims extends SparkShims {
       RebaseSpec(LegacyBehaviorPolicy.CORRECTED)
     )
   }
+
+  override def supportsHashAggregate(aggregateBufferAttributes: Seq[Attribute]): Boolean = {
+    Aggregate.supportsHashAggregate(aggregateBufferAttributes)
+  }
+
+  override def supportsObjectHashAggregate(
+      aggregateExpressions: Seq[AggregateExpression]): Boolean = {
+    Aggregate.supportsObjectHashAggregate(aggregateExpressions)
+  }
 }