huaxingao commented on a change in pull request #32049:
URL: https://github.com/apache/spark/pull/32049#discussion_r648604302
##########
File path:
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala
##########
@@ -17,61 +17,203 @@
package org.apache.spark.sql.execution.datasources.v2
-import org.apache.spark.sql.catalyst.expressions.{And, Expression,
NamedExpression, ProjectionOverSchema, SubqueryExpression}
-import org.apache.spark.sql.catalyst.planning.ScanOperation
-import org.apache.spark.sql.catalyst.plans.logical.{Filter, LogicalPlan,
Project}
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression
+import org.apache.spark.sql.catalyst.planning.{OperationHelper, ScanOperation}
+import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Filter,
LeafNode, LogicalPlan, Project}
import org.apache.spark.sql.catalyst.rules.Rule
-import org.apache.spark.sql.connector.read.{Scan, V1Scan}
+import org.apache.spark.sql.catalyst.util.toPrettySQL
+import org.apache.spark.sql.connector.expressions.Aggregation
+import org.apache.spark.sql.connector.read.{Scan, ScanBuilder,
SupportsPushDownAggregates, SupportsPushDownFilters, V1Scan}
import org.apache.spark.sql.execution.datasources.DataSourceStrategy
import org.apache.spark.sql.sources
import org.apache.spark.sql.types.StructType
-object V2ScanRelationPushDown extends Rule[LogicalPlan] {
+object V2ScanRelationPushDown extends Rule[LogicalPlan] with AliasHelper
+ with OperationHelper with PredicateHelper {
import DataSourceV2Implicits._
- override def apply(plan: LogicalPlan): LogicalPlan = plan transformDown {
- case ScanOperation(project, filters, relation: DataSourceV2Relation) =>
- val scanBuilder =
relation.table.asReadable.newScanBuilder(relation.options)
+ def apply(plan: LogicalPlan): LogicalPlan = {
+
applyColumnPruning(pushdownAggregate(pushDownFilters(createScanBuilder(plan))))
+ }
+
+ private def createScanBuilder(plan: LogicalPlan) = plan.transform {
+ case r: DataSourceV2Relation =>
+ ScanBuilderHolder(r.output, r,
r.table.asReadable.newScanBuilder(r.options))
+ }
- val normalizedFilters = DataSourceStrategy.normalizeExprs(filters,
relation.output)
+ private def pushDownFilters(plan: LogicalPlan) = plan.transform {
+ // update the scan builder with filter push down and return a new plan
with filter pushed
+ case filter @ Filter(_, sHolder: ScanBuilderHolder) =>
+ val (filters, _, _) = collectFilters(filter).get
+
+ val normalizedFilters =
+ DataSourceStrategy.normalizeExprs(filters, sHolder.relation.output)
val (normalizedFiltersWithSubquery, normalizedFiltersWithoutSubquery) =
normalizedFilters.partition(SubqueryExpression.hasSubquery)
// `pushedFilters` will be pushed down and evaluated in the underlying
data sources.
// `postScanFilters` need to be evaluated after the scan.
// `postScanFilters` and `pushedFilters` can overlap, e.g. the parquet
row group filter.
val (pushedFilters, postScanFiltersWithoutSubquery) =
PushDownUtils.pushFilters(
- scanBuilder, normalizedFiltersWithoutSubquery)
+ sHolder.builder, normalizedFiltersWithoutSubquery)
val postScanFilters = postScanFiltersWithoutSubquery ++
normalizedFiltersWithSubquery
+ val output = sHolder.output
+
+ logInfo(
+ s"""
+ |Pushing operators to
${sHolder.relation.asInstanceOf[DataSourceV2Relation].name}
+ |Pushed Filters: ${pushedFilters.mkString(", ")}
+ |Post-Scan Filters: ${postScanFilters.mkString(",")}
+ """.stripMargin)
+ val projectionOverSchema = ProjectionOverSchema(output.toStructType)
+ val projectionFunc = (expr: Expression) => expr transformDown {
+ case projectionOverSchema(newExpr) => newExpr
+ }
+
+ val filterCondition = postScanFilters.reduceLeftOption(And)
+ val newFilterCondition = filterCondition.map(projectionFunc)
+ newFilterCondition.map(Filter(_, sHolder)).getOrElse(sHolder)
+ }
+
+ def pushdownAggregate(plan: LogicalPlan): LogicalPlan = plan.transform {
+ // update the scan builder with agg pushdown and return a new plan with
agg pushed
+ case aggNode @ Aggregate(groupingExpressions, resultExpressions, child) =>
+ child match {
+ case ScanOperation(project, filters, sHolder: ScanBuilderHolder) =>
+ sHolder.builder match {
+ case r: SupportsPushDownAggregates =>
+ if (filters.isEmpty || r.supportsPushDownAggregateWithFilter()) {
+ if (r.supportsGlobalAggregatePushDownOnly() &&
groupingExpressions.nonEmpty) {
+ aggNode // return original plan node
+ } else {
+ val aggregates = getAggregateExpression(resultExpressions,
project, sHolder)
+ val pushedAggregates = PushDownUtils
+ .pushAggregates(sHolder.builder, aggregates,
groupingExpressions)
+ logInfo(
+ s"""
+ |Pushed Aggregates:
${pushedAggregates.aggregateExpressions.mkString(", ")}
+ """.stripMargin)
+ aggNode
Review comment:
Changed.
--
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.
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For additional commands, e-mail: reviews-h...@spark.apache.org
