yashmayya commented on code in PR #19988:
URL: https://github.com/apache/pinot/pull/19988#discussion_r4383935406
##########
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:
Regression: with approximate rewrite on and params set
(`useApproximateFunction=true`,
`pinot.broker.approximate.function.distinct.count.params=threshold=1000`), this
works on master and fails here in v1, v2 and lite mode:
```sql
SELECT FIRSTWITHTIME(x, t), COUNT(DISTINCT s) FROM (
SELECT int_col x, string_col s, ts_col t FROM tbl
UNION ALL SELECT int_col, string_col, ts_col FROM tbl2)
```
The error is `FIRST_WITH_TIME expects the 3rd argument to be literal, got:
IDENTIFIER`. The type rule now puts `'INT'` into the union branches. Then this
rule wraps the `Union` in a new project for its params, and that project hides
`'INT'`. On master both literals ended up in one project.
Suggestion: use one shared, union-aware helper for both rules (see the next
comment).
##########
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:
Two plan regressions for no-type FIRST/LAST_WITH_TIME over UNION ALL,
compared with master:
- **3 or more branches:** only the outer `Union` is handled. The nested
`Union` branch gets wrapped in a project, so `UnionMergeRule` can no longer
flatten it. In v1 this adds one union stage and one more exchange (33 probe
plans). All rows of the first N-1 branches then go through an extra exchange.
- **Minor, 2 branches:** the type literals moved from one project above the
union into every branch. So each row now carries extra constant columns through
the branch exchanges.
Suggestion:
- Make `appendLiterals` recurse into nested `Union` inputs, and let
`PinotApproximateAggregateRewriteRule` call the same helper. A local patch that
did both fixed the approximate query above, and it removed the extra union and
exchange. All planner tests, `ResourceBasedQueriesTest` and `QueryRunnerTest`
passed, and no query that works on master failed.
- Better still, push into the branches only when another call's literal
argument is already there. That is the case that failed on master. Otherwise
keep master's project above the union.
--
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]
