xiangfu0 commented on code in PR #19988:
URL: https://github.com/apache/pinot/pull/19988#discussion_r4384519436
##########
pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotApproximateAggregateRewriteRule.java:
##########
@@ -105,15 +108,30 @@ public void onMatch(RelOptRuleCall call) {
// and AggregatePushdownRule inline the literal only through
findImmediateProjects(input), so anything that puts a
// node between them would leave the leaf with a column reference where
the function expects a parameter string.
// Nothing in the later phases does: AggregateProjectMergeRule needs an
all-input-ref project and the literals
- // block it, and ProjectMergeRule collapses into a single project that
stays adjacent.
+ // block it, and ProjectMergeRule collapses into a single project that
stays adjacent. The parameters are appended
+ // to an existing project beneath the aggregate, so the literals already
in it stay inlinable even when
+ // ProjectMerge is skipped.
RelNode newInput = input;
int distinctCountParamsIndex = -1;
int percentileParamsIndex = -1;
if (!distinctCountParams.isEmpty() || !percentileParams.isEmpty()) {
RexBuilder rexBuilder = aggRel.getCluster().getRexBuilder();
- List<RexNode> projects = new ArrayList<>(numInputFields + 2);
- for (int i = 0; i < numInputFields; i++) {
- projects.add(rexBuilder.makeInputRef(input, i));
+ RelNode unboxedInput = PinotRuleUtils.unboxRel(input);
+ List<RexNode> projects;
+ RelNode projectInput;
+ List<RelHint> hints;
+ if (unboxedInput instanceof Project) {
Review Comment:
Fixed in f51d298a91: both rewrites now use
PinotRuleUtils.appendAggregateLiterals, which preserves existing literals and
recurses through nested UNION inputs when necessary. The approximate-parameter
test covers a three-branch UNION with ProjectMerge disabled under both planners.
_🤖 Addressed by [Claude Code](https://claude.com/claude-code)_
##########
pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotAggregateTypeArgumentRewriteRule.java:
##########
@@ -72,9 +77,56 @@ public boolean matches(RelOptRuleCall call) {
public void onMatch(RelOptRuleCall call) {
Aggregate aggRel = call.rel(0);
RelNode input = aggRel.getInput();
- RelNode unboxedInput = PinotRuleUtils.unboxRel(input);
RexBuilder rexBuilder = aggRel.getCluster().getRexBuilder();
+ // Appending the literals after the input fields keeps the existing field
indexes, so the group set and the other
+ // calls are unaffected.
+ int numInputFields = input.getRowType().getFieldCount();
+ List<RexNode> literals = new ArrayList<>();
+ Map<String, Integer> literalIndexByDataType = new HashMap<>();
+ List<AggregateCall> orgAggCalls = aggRel.getAggCallList();
+ List<AggregateCall> aggCalls = new ArrayList<>(orgAggCalls.size());
+ for (AggregateCall aggCall : orgAggCalls) {
+ String dataType = getDataTypeArgument(aggCall, input);
+ if (dataType == null) {
+ aggCalls.add(aggCall);
+ continue;
+ }
+ int literalIndex = literalIndexByDataType.computeIfAbsent(dataType, type
-> {
+ literals.add(rexBuilder.makeLiteral(type));
+ return numInputFields + literals.size() - 1;
+ });
+ List<Integer> argList = new ArrayList<>(aggCall.getArgList());
+ if (isArrayAgg(aggCall)) {
+ argList.add(1, literalIndex);
+ } else {
+ argList.add(literalIndex);
+ }
+ // Keeps the original result type, so the rewrite does not change the
result schema.
+ aggCalls.add(aggCall.withArgList(argList));
+ }
+
+ // Under a union the literals go into every input, because a project above
the union would hide the literal
+ // arguments already projected in its inputs.
+ RelNode unboxedInput = PinotRuleUtils.unboxRel(input);
+ RelNode newInput;
+ if (unboxedInput instanceof Union) {
Review Comment:
Fixed in f51d298a91: the shared helper recurses through nested unions when
existing projected literals require branch projection; otherwise it keeps one
project above the union. The new plan case verifies one flattened three-input
union with a single type-literal project, and the runtime case checks
FIRST_WITH_TIME, ARRAY_AGG and PERCENTILE together.
_🤖 Addressed by [Claude Code](https://claude.com/claude-code)_
--
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]