You are viewing a plain text version of this content. The canonical link for it is here.
Posted to reviews@spark.apache.org by GitBox <gi...@apache.org> on 2022/07/18 05:17:37 UTC
[GitHub] [spark] cloud-fan commented on a diff in pull request #37195: [SPARK-39148][SQL] DS V2 aggregate push down can work with OFFSET or LIMIT
cloud-fan commented on code in PR #37195:
URL: https://github.com/apache/spark/pull/37195#discussion_r922977591
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala:
##########
@@ -92,189 +93,199 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan] with PredicateHelper wit
def pushDownAggregates(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)
- if filters.isEmpty && CollapseProject.canCollapseExpressions(
- resultExpressions, project, alwaysInline = true) =>
- sHolder.builder match {
- case r: SupportsPushDownAggregates =>
- val aliasMap = getAliasMap(project)
- val actualResultExprs = resultExpressions.map(replaceAliasButKeepName(_, aliasMap))
- val actualGroupExprs = groupingExpressions.map(replaceAlias(_, aliasMap))
-
- val aggExprToOutputOrdinal = mutable.HashMap.empty[Expression, Int]
- val aggregates = collectAggregates(actualResultExprs, aggExprToOutputOrdinal)
- val normalizedAggregates = DataSourceStrategy.normalizeExprs(
- aggregates, sHolder.relation.output).asInstanceOf[Seq[AggregateExpression]]
- val normalizedGroupingExpressions = DataSourceStrategy.normalizeExprs(
- actualGroupExprs, sHolder.relation.output)
- val translatedAggregates = DataSourceStrategy.translateAggregation(
- normalizedAggregates, normalizedGroupingExpressions)
- val (finalResultExpressions, finalAggregates, finalTranslatedAggregates) = {
- if (translatedAggregates.isEmpty ||
- r.supportCompletePushDown(translatedAggregates.get) ||
- translatedAggregates.get.aggregateExpressions().forall(!_.isInstanceOf[Avg])) {
- (actualResultExprs, aggregates, translatedAggregates)
- } else {
- // scalastyle:off
- // The data source doesn't support the complete push-down of this aggregation.
- // Here we translate `AVG` to `SUM / COUNT`, so that it's more likely to be
- // pushed, completely or partially.
- // e.g. TABLE t (c1 INT, c2 INT, c3 INT)
- // SELECT avg(c1) FROM t GROUP BY c2;
- // The original logical plan is
- // Aggregate [c2#10],[avg(c1#9) AS avg(c1)#19]
- // +- ScanOperation[...]
- //
- // After convert avg(c1#9) to sum(c1#9)/count(c1#9)
- // we have the following
- // Aggregate [c2#10],[sum(c1#9)/count(c1#9) AS avg(c1)#19]
- // +- ScanOperation[...]
- // scalastyle:on
- val newResultExpressions = actualResultExprs.map { expr =>
- expr.transform {
- case AggregateExpression(avg: aggregate.Average, _, isDistinct, _, _) =>
- val sum = aggregate.Sum(avg.child).toAggregateExpression(isDistinct)
- val count = aggregate.Count(avg.child).toAggregateExpression(isDistinct)
- avg.evaluateExpression transform {
- case a: Attribute if a.semanticEquals(avg.sum) =>
- addCastIfNeeded(sum, avg.sum.dataType)
- case a: Attribute if a.semanticEquals(avg.count) =>
- addCastIfNeeded(count, avg.count.dataType)
- }
- }
- }.asInstanceOf[Seq[NamedExpression]]
- // Because aggregate expressions changed, translate them again.
- aggExprToOutputOrdinal.clear()
- val newAggregates =
- collectAggregates(newResultExpressions, aggExprToOutputOrdinal)
- val newNormalizedAggregates = DataSourceStrategy.normalizeExprs(
- newAggregates, sHolder.relation.output).asInstanceOf[Seq[AggregateExpression]]
- (newResultExpressions, newAggregates, DataSourceStrategy.translateAggregation(
- newNormalizedAggregates, normalizedGroupingExpressions))
+ case agg: Aggregate => rewriteAggregate(agg)
+ }
+
+ private def rewriteAggregate(agg: Aggregate): LogicalPlan = agg.child match {
+ case ScanOperation(project, Nil, holder@ScanBuilderHolder(_, _, r: SupportsPushDownAggregates))
+ if CollapseProject.canCollapseExpressions(
+ agg.aggregateExpressions, project, alwaysInline = true) =>
+ val aliasMap = getAliasMap(project)
+ val actualResultExprs = agg.aggregateExpressions.map(replaceAliasButKeepName(_, aliasMap))
+ val actualGroupExprs = agg.groupingExpressions.map(replaceAlias(_, aliasMap))
+
+ val aggExprToOutputOrdinal = mutable.HashMap.empty[Expression, Int]
+ val aggregates = collectAggregates(actualResultExprs, aggExprToOutputOrdinal)
+ val normalizedAggExprs = DataSourceStrategy.normalizeExprs(
+ aggregates, holder.relation.output).asInstanceOf[Seq[AggregateExpression]]
+ val normalizedGroupingExpr = DataSourceStrategy.normalizeExprs(
+ actualGroupExprs, holder.relation.output)
+ val translatedAggOpt = DataSourceStrategy.translateAggregation(
+ normalizedAggExprs, normalizedGroupingExpr)
+ if (translatedAggOpt.isEmpty) {
+ // Cannot translate the catalyst aggregate, return the query plan unchanged.
+ return agg
+ }
+
+ val (finalResultExprs, finalAggExprs, translatedAgg, canCompletePushDown) = {
+ if (r.supportCompletePushDown(translatedAggOpt.get)) {
+ (actualResultExprs, normalizedAggExprs, translatedAggOpt.get, true)
+ } else if (!translatedAggOpt.get.aggregateExpressions().exists(_.isInstanceOf[Avg])) {
+ (actualResultExprs, normalizedAggExprs, translatedAggOpt.get, false)
+ } else {
+ // scalastyle:off
+ // The data source doesn't support the complete push-down of this aggregation.
+ // Here we translate `AVG` to `SUM / COUNT`, so that it's more likely to be
+ // pushed, completely or partially.
+ // e.g. TABLE t (c1 INT, c2 INT, c3 INT)
+ // SELECT avg(c1) FROM t GROUP BY c2;
+ // The original logical plan is
+ // Aggregate [c2#10],[avg(c1#9) AS avg(c1)#19]
+ // +- ScanOperation[...]
+ //
+ // After convert avg(c1#9) to sum(c1#9)/count(c1#9)
+ // we have the following
+ // Aggregate [c2#10],[sum(c1#9)/count(c1#9) AS avg(c1)#19]
+ // +- ScanOperation[...]
+ // scalastyle:on
+ val newResultExpressions = actualResultExprs.map { expr =>
+ expr.transform {
+ case AggregateExpression(avg: aggregate.Average, _, isDistinct, _, _) =>
+ val sum = aggregate.Sum(avg.child).toAggregateExpression(isDistinct)
+ val count = aggregate.Count(avg.child).toAggregateExpression(isDistinct)
+ avg.evaluateExpression transform {
+ case a: Attribute if a.semanticEquals(avg.sum) =>
+ addCastIfNeeded(sum, avg.sum.dataType)
+ case a: Attribute if a.semanticEquals(avg.count) =>
+ addCastIfNeeded(count, avg.count.dataType)
}
- }
+ }
+ }.asInstanceOf[Seq[NamedExpression]]
+ // Because aggregate expressions changed, translate them again.
+ aggExprToOutputOrdinal.clear()
+ val newAggregates =
+ collectAggregates(newResultExpressions, aggExprToOutputOrdinal)
+ val newNormalizedAggExprs = DataSourceStrategy.normalizeExprs(
+ newAggregates, holder.relation.output).asInstanceOf[Seq[AggregateExpression]]
+ val newTranslatedAggOpt = DataSourceStrategy.translateAggregation(
+ newNormalizedAggExprs, normalizedGroupingExpr)
+ if (newTranslatedAggOpt.isEmpty) {
+ // Ideally we should never reach here. But if we end up with not able to translate
+ // new aggregate with AVG replaced by SUM/COUNT, revert to the original one.
+ (actualResultExprs, normalizedAggExprs, translatedAggOpt.get, false)
+ } else {
+ (newResultExpressions, newNormalizedAggExprs, newTranslatedAggOpt.get,
+ r.supportCompletePushDown(newTranslatedAggOpt.get))
+ }
+ }
+ }
- if (finalTranslatedAggregates.isEmpty) {
- aggNode // return original plan node
- } else if (!r.supportCompletePushDown(finalTranslatedAggregates.get) &&
- !supportPartialAggPushDown(finalTranslatedAggregates.get)) {
- aggNode // return original plan node
- } else {
- val pushedAggregates = finalTranslatedAggregates.filter(r.pushAggregation)
- if (pushedAggregates.isEmpty) {
- aggNode // return original plan node
- } else {
- // No need to do column pruning because only the aggregate columns are used as
- // DataSourceV2ScanRelation output columns. All the other columns are not
- // included in the output.
- val scan = sHolder.builder.build()
-
- // scalastyle:off
- // use the group by columns and aggregate columns as the output columns
- // e.g. TABLE t (c1 INT, c2 INT, c3 INT)
- // SELECT min(c1), max(c1) FROM t GROUP BY c2;
- // Use c2, min(c1), max(c1) as output for DataSourceV2ScanRelation
- // We want to have the following logical plan:
- // == Optimized Logical Plan ==
- // Aggregate [c2#10], [min(min(c1)#21) AS min(c1)#17, max(max(c1)#22) AS max(c1)#18]
- // +- RelationV2[c2#10, min(c1)#21, max(c1)#22]
- // scalastyle:on
- val newOutput = scan.readSchema().toAttributes
- assert(newOutput.length == groupingExpressions.length + finalAggregates.length)
- val groupByExprToOutputOrdinal = mutable.HashMap.empty[Expression, Int]
- val groupAttrs = normalizedGroupingExpressions.zip(newOutput).zipWithIndex.map {
- case ((a: Attribute, b: Attribute), _) => b.withExprId(a.exprId)
- case ((expr, attr), ordinal) =>
- if (!groupByExprToOutputOrdinal.contains(expr.canonicalized)) {
- groupByExprToOutputOrdinal(expr.canonicalized) = ordinal
- }
- attr
- }
- val aggOutput = newOutput.drop(groupAttrs.length)
- val output = groupAttrs ++ aggOutput
-
- logInfo(
- s"""
- |Pushing operators to ${sHolder.relation.name}
- |Pushed Aggregate Functions:
- | ${pushedAggregates.get.aggregateExpressions.mkString(", ")}
- |Pushed Group by:
- | ${pushedAggregates.get.groupByExpressions.mkString(", ")}
- |Output: ${output.mkString(", ")}
- """.stripMargin)
-
- val wrappedScan = getWrappedScan(scan, sHolder, pushedAggregates)
- val scanRelation =
- DataSourceV2ScanRelation(sHolder.relation, wrappedScan, output)
- if (r.supportCompletePushDown(pushedAggregates.get)) {
- val projectExpressions = finalResultExpressions.map { expr =>
- expr.transformDown {
- case agg: AggregateExpression =>
- val ordinal = aggExprToOutputOrdinal(agg.canonicalized)
- val child =
- addCastIfNeeded(aggOutput(ordinal), agg.resultAttribute.dataType)
- Alias(child, agg.resultAttribute.name)(agg.resultAttribute.exprId)
- case expr if groupByExprToOutputOrdinal.contains(expr.canonicalized) =>
- val ordinal = groupByExprToOutputOrdinal(expr.canonicalized)
- addCastIfNeeded(groupAttrs(ordinal), expr.dataType)
- }
- }.asInstanceOf[Seq[NamedExpression]]
- Project(projectExpressions, scanRelation)
+ if (!canCompletePushDown && !supportPartialAggPushDown(translatedAgg)) {
+ return agg
+ }
+ if (!r.pushAggregation(translatedAgg)) {
+ return agg
+ }
+
+ // scalastyle:off
+ // use the group by columns and aggregate columns as the output columns
+ // e.g. TABLE t (c1 INT, c2 INT, c3 INT)
+ // SELECT min(c1), max(c1) FROM t GROUP BY c2;
+ // Use c2, min(c1), max(c1) as output for DataSourceV2ScanRelation
+ // We want to have the following logical plan:
+ // == Optimized Logical Plan ==
+ // Aggregate [c2#10], [min(min(c1)#21) AS min(c1)#17, max(max(c1)#22) AS max(c1)#18]
+ // +- RelationV2[c2#10, min(c1)#21, max(c1)#22]
+ // scalastyle:on
+ val newOutput = normalizedGroupingExpr.zipWithIndex.map { case (e, i) =>
+ AttributeReference(s"group_col_$i", e.dataType)()
+ } ++ finalAggExprs.zipWithIndex.map { case (e, i) =>
+ AttributeReference(s"agg_func_$i", e.dataType)()
+ }
+ val groupOutput = newOutput.take(normalizedGroupingExpr.length)
+ val aggOutput = newOutput.drop(normalizedGroupingExpr.length)
+ val groupByExprToOutputOrdinal = mutable.HashMap.empty[Expression, Int]
+ normalizedGroupingExpr.zipWithIndex.foreach { case (expr, ordinal) =>
+ if (!groupByExprToOutputOrdinal.contains(expr.canonicalized)) {
+ groupByExprToOutputOrdinal(expr.canonicalized) = ordinal
+ }
+ }
+
+ holder.pushedAggregate = Some(translatedAgg)
+ holder.output = newOutput
+ logInfo(
+ s"""
+ |Pushing operators to ${holder.relation.name}
+ |Pushed Aggregate Functions:
+ | ${translatedAgg.aggregateExpressions().mkString(", ")}
+ |Pushed Group by:
+ | ${translatedAgg.groupByExpressions.mkString(", ")}
Review Comment:
Now this is completed inferred by Spark, according to the pushed agg exprs and group by exprs.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org