[spark] branch master updated: [SPARK-39964][SQL] DS V2 pushdown should unify the translate path

[we...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

wenchen pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git


The following commit(s) were added to refs/heads/master by this push:
     new de1d366b275 [SPARK-39964][SQL] DS V2 pushdown should unify the translate path
de1d366b275 is described below

commit de1d366b2756f5a5c9f3c6c5bc4a58ae313c2db4
Author: Jiaan Geng <be...@163.com>
AuthorDate: Fri Aug 5 18:38:50 2022 +0800

    [SPARK-39964][SQL] DS V2 pushdown should unify the translate path
    
    ### What changes were proposed in this pull request?
    Currently, DS V2 pushdown have two translate path `DataSourceStrategy.translateAggregate` used to translate aggregate functions and `V2ExpressionBuilder` used to translate other functions and expressions, we can unify them.
    
    After this PR, the translate have only one code path, developers will easy to coding and reading.
    
    ### Why are the changes needed?
    Unify the translate path for DS V2 pushdown.
    
    ### Does this PR introduce _any_ user-facing change?
    'No'.
    Just update the inner implementation.
    
    ### How was this patch tested?
    N/A
    
    Closes #37391 from beliefer/SPARK-39964.
    
    Authored-by: Jiaan Geng <be...@163.com>
    Signed-off-by: Wenchen Fan <we...@databricks.com>
---
 .../sql/catalyst/util/V2ExpressionBuilder.scala    | 53 +++++++++++++++++++
 .../execution/datasources/DataSourceStrategy.scala | 59 ++--------------------
 2 files changed, 57 insertions(+), 55 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/catalyst/util/V2ExpressionBuilder.scala b/sql/core/src/main/scala/org/apache/spark/sql/catalyst/util/V2ExpressionBuilder.scala
index b274fd38e99..89115bf7ab5 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/catalyst/util/V2ExpressionBuilder.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/catalyst/util/V2ExpressionBuilder.scala
@@ -18,8 +18,11 @@
 package org.apache.spark.sql.catalyst.util
 
 import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, AggregateFunction, Complete}
 import org.apache.spark.sql.connector.expressions.{Cast => V2Cast, Expression => V2Expression, Extract => V2Extract, FieldReference, GeneralScalarExpression, LiteralValue, UserDefinedScalarFunc}
+import org.apache.spark.sql.connector.expressions.aggregate.{AggregateFunc, Avg, Count, CountStar, GeneralAggregateFunc, Max, Min, Sum, UserDefinedAggregateFunc}
 import org.apache.spark.sql.connector.expressions.filter.{AlwaysFalse, AlwaysTrue, And => V2And, Not => V2Not, Or => V2Or, Predicate => V2Predicate}
+import org.apache.spark.sql.execution.datasources.PushableExpression
 import org.apache.spark.sql.types.{BooleanType, IntegerType}
 
 /**
@@ -91,6 +94,8 @@ class V2ExpressionBuilder(e: Expression, isPredicate: Boolean = false) {
     case Cast(child, dataType, _, ansiEnabled)
         if ansiEnabled || Cast.canUpCast(child.dataType, dataType) =>
       generateExpression(child).map(v => new V2Cast(v, dataType))
+    case AggregateExpression(aggregateFunction, Complete, isDistinct, None, _) =>
+      generateAggregateFunc(aggregateFunction, isDistinct)
     case Abs(child, true) => generateExpressionWithName("ABS", Seq(child))
     case Coalesce(children) => generateExpressionWithName("COALESCE", children)
     case Greatest(children) => generateExpressionWithName("GREATEST", children)
@@ -276,6 +281,54 @@ class V2ExpressionBuilder(e: Expression, isPredicate: Boolean = false) {
     case _ => None
   }
 
+  private def generateAggregateFunc(
+      aggregateFunction: AggregateFunction,
+      isDistinct: Boolean): Option[AggregateFunc] = aggregateFunction match {
+    case aggregate.Min(PushableExpression(expr)) => Some(new Min(expr))
+    case aggregate.Max(PushableExpression(expr)) => Some(new Max(expr))
+    case count: aggregate.Count if count.children.length == 1 =>
+      count.children.head match {
+        // COUNT(any literal) is the same as COUNT(*)
+        case Literal(_, _) => Some(new CountStar())
+        case PushableExpression(expr) => Some(new Count(expr, isDistinct))
+        case _ => None
+      }
+    case aggregate.Sum(PushableExpression(expr), _) => Some(new Sum(expr, isDistinct))
+    case aggregate.Average(PushableExpression(expr), _) => Some(new Avg(expr, isDistinct))
+    case aggregate.VariancePop(PushableExpression(expr), _) =>
+      Some(new GeneralAggregateFunc("VAR_POP", isDistinct, Array(expr)))
+    case aggregate.VarianceSamp(PushableExpression(expr), _) =>
+      Some(new GeneralAggregateFunc("VAR_SAMP", isDistinct, Array(expr)))
+    case aggregate.StddevPop(PushableExpression(expr), _) =>
+      Some(new GeneralAggregateFunc("STDDEV_POP", isDistinct, Array(expr)))
+    case aggregate.StddevSamp(PushableExpression(expr), _) =>
+      Some(new GeneralAggregateFunc("STDDEV_SAMP", isDistinct, Array(expr)))
+    case aggregate.CovPopulation(PushableExpression(left), PushableExpression(right), _) =>
+      Some(new GeneralAggregateFunc("COVAR_POP", isDistinct, Array(left, right)))
+    case aggregate.CovSample(PushableExpression(left), PushableExpression(right), _) =>
+      Some(new GeneralAggregateFunc("COVAR_SAMP", isDistinct, Array(left, right)))
+    case aggregate.Corr(PushableExpression(left), PushableExpression(right), _) =>
+      Some(new GeneralAggregateFunc("CORR", isDistinct, Array(left, right)))
+    case aggregate.RegrIntercept(PushableExpression(left), PushableExpression(right)) =>
+      Some(new GeneralAggregateFunc("REGR_INTERCEPT", isDistinct, Array(left, right)))
+    case aggregate.RegrR2(PushableExpression(left), PushableExpression(right)) =>
+      Some(new GeneralAggregateFunc("REGR_R2", isDistinct, Array(left, right)))
+    case aggregate.RegrSlope(PushableExpression(left), PushableExpression(right)) =>
+      Some(new GeneralAggregateFunc("REGR_SLOPE", isDistinct, Array(left, right)))
+    case aggregate.RegrSXY(PushableExpression(left), PushableExpression(right)) =>
+      Some(new GeneralAggregateFunc("REGR_SXY", isDistinct, Array(left, right)))
+    // TODO supports other aggregate functions
+    case aggregate.V2Aggregator(aggrFunc, children, _, _) =>
+      val translatedExprs = children.flatMap(PushableExpression.unapply(_))
+      if (translatedExprs.length == children.length) {
+        Some(new UserDefinedAggregateFunc(aggrFunc.name(),
+          aggrFunc.canonicalName(), isDistinct, translatedExprs.toArray[V2Expression]))
+      } else {
+        None
+      }
+    case _ => None
+  }
+
   private def flipComparisonOperatorName(operatorName: String): String = {
     operatorName match {
       case ">" => "<"
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSourceStrategy.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSourceStrategy.scala
index 27e4e146177..a9d5c6da384 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSourceStrategy.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSourceStrategy.scala
@@ -42,7 +42,7 @@ import org.apache.spark.sql.catalyst.util.{ResolveDefaultColumns, V2ExpressionBu
 import org.apache.spark.sql.connector.catalog.SupportsRead
 import org.apache.spark.sql.connector.catalog.TableCapability._
 import org.apache.spark.sql.connector.expressions.{Expression => V2Expression, NullOrdering, SortDirection, SortOrder => V2SortOrder, SortValue}
-import org.apache.spark.sql.connector.expressions.aggregate.{AggregateFunc, Aggregation, Avg, Count, CountStar, GeneralAggregateFunc, Max, Min, Sum, UserDefinedAggregateFunc}
+import org.apache.spark.sql.connector.expressions.aggregate.{AggregateFunc, Aggregation}
 import org.apache.spark.sql.errors.QueryCompilationErrors
 import org.apache.spark.sql.execution.{RowDataSourceScanExec, SparkPlan}
 import org.apache.spark.sql.execution.command._
@@ -694,57 +694,6 @@ object DataSourceStrategy
     (nonconvertiblePredicates ++ unhandledPredicates, pushedFilters, handledFilters)
   }
 
-  protected[sql] def translateAggregate(agg: AggregateExpression): Option[AggregateFunc] = {
-    if (agg.filter.isEmpty) {
-      agg.aggregateFunction match {
-        case aggregate.Min(PushableExpression(expr)) => Some(new Min(expr))
-        case aggregate.Max(PushableExpression(expr)) => Some(new Max(expr))
-        case count: aggregate.Count if count.children.length == 1 =>
-          count.children.head match {
-            // COUNT(any literal) is the same as COUNT(*)
-            case Literal(_, _) => Some(new CountStar())
-            case PushableExpression(expr) => Some(new Count(expr, agg.isDistinct))
-            case _ => None
-          }
-        case aggregate.Sum(PushableExpression(expr), _) => Some(new Sum(expr, agg.isDistinct))
-        case aggregate.Average(PushableExpression(expr), _) => Some(new Avg(expr, agg.isDistinct))
-        case aggregate.VariancePop(PushableExpression(expr), _) =>
-          Some(new GeneralAggregateFunc("VAR_POP", agg.isDistinct, Array(expr)))
-        case aggregate.VarianceSamp(PushableExpression(expr), _) =>
-          Some(new GeneralAggregateFunc("VAR_SAMP", agg.isDistinct, Array(expr)))
-        case aggregate.StddevPop(PushableExpression(expr), _) =>
-          Some(new GeneralAggregateFunc("STDDEV_POP", agg.isDistinct, Array(expr)))
-        case aggregate.StddevSamp(PushableExpression(expr), _) =>
-          Some(new GeneralAggregateFunc("STDDEV_SAMP", agg.isDistinct, Array(expr)))
-        case aggregate.CovPopulation(PushableExpression(left), PushableExpression(right), _) =>
-          Some(new GeneralAggregateFunc("COVAR_POP", agg.isDistinct, Array(left, right)))
-        case aggregate.CovSample(PushableExpression(left), PushableExpression(right), _) =>
-          Some(new GeneralAggregateFunc("COVAR_SAMP", agg.isDistinct, Array(left, right)))
-        case aggregate.Corr(PushableExpression(left), PushableExpression(right), _) =>
-          Some(new GeneralAggregateFunc("CORR", agg.isDistinct, Array(left, right)))
-        case aggregate.RegrIntercept(PushableExpression(left), PushableExpression(right)) =>
-          Some(new GeneralAggregateFunc("REGR_INTERCEPT", agg.isDistinct, Array(left, right)))
-        case aggregate.RegrR2(PushableExpression(left), PushableExpression(right)) =>
-          Some(new GeneralAggregateFunc("REGR_R2", agg.isDistinct, Array(left, right)))
-        case aggregate.RegrSlope(PushableExpression(left), PushableExpression(right)) =>
-          Some(new GeneralAggregateFunc("REGR_SLOPE", agg.isDistinct, Array(left, right)))
-        case aggregate.RegrSXY(PushableExpression(left), PushableExpression(right)) =>
-          Some(new GeneralAggregateFunc("REGR_SXY", agg.isDistinct, Array(left, right)))
-        case aggregate.V2Aggregator(aggrFunc, children, _, _) =>
-          val translatedExprs = children.flatMap(PushableExpression.unapply(_))
-          if (translatedExprs.length == children.length) {
-            Some(new UserDefinedAggregateFunc(aggrFunc.name(),
-              aggrFunc.canonicalName(), agg.isDistinct, translatedExprs.toArray[V2Expression]))
-          } else {
-            None
-          }
-        case _ => None
-      }
-    } else {
-      None
-    }
-  }
-
   /**
    * Translate aggregate expressions and group by expressions.
    *
@@ -753,13 +702,13 @@ object DataSourceStrategy
   protected[sql] def translateAggregation(
       aggregates: Seq[AggregateExpression], groupBy: Seq[Expression]): Option[Aggregation] = {
 
-    def translateGroupBy(e: Expression): Option[V2Expression] = e match {
+    def translate(e: Expression): Option[V2Expression] = e match {
       case PushableExpression(expr) => Some(expr)
       case _ => None
     }
 
-    val translatedAggregates = aggregates.flatMap(translateAggregate)
-    val translatedGroupBys = groupBy.flatMap(translateGroupBy)
+    val translatedAggregates = aggregates.flatMap(translate).asInstanceOf[Seq[AggregateFunc]]
+    val translatedGroupBys = groupBy.flatMap(translate)
 
     if (translatedAggregates.length != aggregates.length ||
       translatedGroupBys.length != groupBy.length) {


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org