[SPARK-39148][SQL] DS V2 aggregate push down can work with OFFSET or LIMIT

[cloud-fan]

What changes were proposed in this pull request?

This PR refactors the v2 agg pushdown code. The main change is, now we don't build the Scan immediately when pushing agg. We did it so before because we want to know the data schema with agg pushed, then we can add cast when rewriting the query plan after pushdown. But the problem is, we build Scan too early and can't push down any more operators, while it's common to see LIMIT/OFFSET after agg.

The idea of the refactor is, we don't need to know the data schema with agg pushed. We just give an expectation (the data type should be the same of Spark agg functions), use it to define the output of ScanBuilderHolder, and then rewrite the query plan. Later on, when we build the Scan and replace ScanBuilderHolder with DataSourceV2ScanRelation, we check the actual data schema and add a Project to do type cast if necessary.

Why are the changes needed?

support pushing down LIMIT/OFFSET after agg.

Does this PR introduce any user-facing change?

no

How was this patch tested?

updated tests

[beliefer]

Suggested change

	case ScanOperation(project, Nil, holder@ScanBuilderHolder(_, _, r: SupportsPushDownAggregates))
	case ScanOperation(project, Nil, holder @ ScanBuilderHolder(_, _, r: SupportsPushDownAggregates))

[beliefer]

Suggested change

	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 groupOutput = normalizedGroupingExpr.zipWithIndex.map { case (e, i) =>
	AttributeReference(s"group_col_$i", e.dataType)()
	}
	val aggOutput = finalAggExprs.zipWithIndex.map { case (e, i) =>
	AttributeReference(s"agg_func_$i", e.dataType)()
	}
	val newOutput = groupOutput ++ aggOutput

[beliefer]

Good decoupling !

[beliefer]

It seems we not display Output information here now.

[cloud-fan]

Now this is completed inferred by Spark, according to the pushed agg exprs and group by exprs.

[beliefer]

Got it.

[beliefer]

Good decoupling !

[cloud-fan]

cc @huaxingao @viirya

[viirya]

Why for this case canCompletePushDown is false? Previously it is treated like supportCompletePushDown, no?

[cloud-fan]

It was false before this PR as well. The previous code will invoke r.supportCompletePushDown one more time which returns false.

The logic is also quite clear:

1. the Aggregate can't be completely pushed
2. the Aggregate can't be rewritten (avg -> sum/count)

The above two means this Aggregate can't be completely pushed.

[viirya]

Got it. canCompletePushDown is false but it's still possible to be partial pushdown. It will be checked later.

[viirya]

Hmm, is this regression? I.e. previously we can pushdown, but now we cannot?

[cloud-fan]

Previously we can't push down either. After agg pushdown, Scan is already built, so limit/offset/topN pushdown won't apply anyway.

[viirya]

I see. I was confused by the added comment and thought it was supported before.

[huaxingao]

I think this example is not accurate any more. Do we need to update this?

[cloud-fan]

thanks for review, merging to master!