viirya commented on code in PR #37195:
URL: https://github.com/apache/spark/pull/37195#discussion_r924067529
##########
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)
Review Comment:
Got it. `canCompletePushDown` is false but it's still possible to be partial
pushdown. It will be checked later.
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
