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]

Reply via email to