yashmayya commented on code in PR #19988:
URL: https://github.com/apache/pinot/pull/19988#discussion_r4382018923
##########
pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java:
##########
@@ -600,12 +600,17 @@ private static HepProgram getOptProgram(PinotRuleSet
ruleSet, Set<String> skipRu
List<RelOptRule> pruneRules =
filterRuleList(ruleSet.rulesFor(Phase.PRUNE), skipRuleSet, useRuleSet,
defaultDisabledRuleSet);
+ // Supplies the data type argument before AggregateCaseToFilter, which
rewrites only one-argument calls: its
+ // ARRAY_AGG(expression) FILTER (WHERE condition) form of ARRAY_AGG(CASE
WHEN ...) fails at runtime.
+
hepProgramBuilder.addRuleInstance(PinotAggregateTypeArgumentRewriteRule.INSTANCE);
Review Comment:
This instance now runs before `ProjectSetOpTransposeRule`. That rule pushes
the new literal project into the set-op inputs. `findImmediateProjects` in
`PinotAggregateExchangeNodeInsertRule` (and `AggregatePushdownRule` in v2) only
looks through `Union`, so for INTERSECT ALL / EXCEPT ALL the type argument
becomes an input ref.
This breaks FIRST_WITH_TIME from #19986, which works on master:
```sql
SELECT FIRSTWITHTIME(x, CAST(t AS TIMESTAMP)) FROM (
SELECT int_col x, long_col t FROM tbl INTERSECT ALL SELECT int_col,
long_col FROM tbl)
```
This fails with `FIRST_WITH_TIME expects the 3rd argument to be literal,
got: IDENTIFIER` in v1, v2 and lite mode. `ARRAY_AGG(x)` over INTERSECT ALL and
EXCEPT ALL fails the same way. For UNION ALL the query works, but the literal
is now a column in every branch and goes through the exchange.
Suggestion: add the early instance just before `AGGREGATE_CASE_TO_FILTER`,
not before all basic rules:
```java
for (RelOptRule relOptRule : basicRules) {
if
(relOptRule.toString().equals(PlannerRuleNames.AGGREGATE_CASE_TO_FILTER)) {
hepProgramBuilder.addRuleInstance(PinotAggregateTypeArgumentRewriteRule.INSTANCE);
}
hepProgramBuilder.addRuleInstance(relOptRule);
}
```
I tested this locally. The set-op queries pass, your three fixes still work,
and all 938 planner tests pass with no plan changes.
##########
pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotAggregateTypeArgumentRewriteRule.java:
##########
@@ -97,40 +100,76 @@ public void onMatch(RelOptRuleCall call) {
List<AggregateCall> orgAggCalls = aggRel.getAggCallList();
List<AggregateCall> aggCalls = new ArrayList<>(orgAggCalls.size());
Map<String, Integer> literalIndexByDataType = new HashMap<>();
+ int isDistinctLiteralIndex = -1;
for (AggregateCall aggCall : orgAggCalls) {
String dataType = getDataTypeArgument(aggCall, input);
if (dataType == null) {
aggCalls.add(aggCall);
continue;
}
- List<Integer> argList = new ArrayList<>(aggCall.getArgList());
- argList.add(literalIndexByDataType.computeIfAbsent(dataType, type -> {
+ int literalIndex = literalIndexByDataType.computeIfAbsent(dataType, type
-> {
projects.add(rexBuilder.makeLiteral(type));
return projects.size() - 1;
- }));
+ });
+ List<Integer> argList = new ArrayList<>(aggCall.getArgList());
+ AggregateCall newAggCall = aggCall;
+ if (isArrayAgg(aggCall)) {
+ argList.add(1, literalIndex);
+ // The aggregate exchange drops the DISTINCT flag, so
ARRAY_AGG(DISTINCT expression) passes it as isDistinct
+ if (aggCall.isDistinct() && argList.size() == 2) {
Review Comment:
`ARRAY_AGG(DISTINCT x, false)` still returns duplicates with no error
(CARDINALITY 4 instead of 3 over {1, 2, 2, 3}). The DISTINCT flag stays, the
isDistinct argument is `false`, and the exchange drops DISTINCT. On master this
form failed on the server, so the silent wrong result is new.
Suggestion: when the DISTINCT keyword is set, always pass `true` as
isDistinct. Or handle ARRAY_AGG DISTINCT in `buildAggCall`, as LISTAGG does.
That also fixes the explicit `ARRAY_AGG(DISTINCT x, 'INT')`, which drops
DISTINCT on master too.
##########
pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotAggregateTypeArgumentRewriteRule.java:
##########
@@ -143,6 +182,14 @@ private static String getDataTypeArgument(AggregateCall
aggCall, RelNode input)
case DOUBLE:
return "DOUBLE";
case DECIMAL:
+ if (name.equals(AggregationFunctionType.ARRAYAGG.name())) {
+ // The planner types an array of DECIMAL as INT_ARRAY or LONG_ARRAY
for scale 0 and precision up to 38, and as
+ // DOUBLE_ARRAY otherwise, see RelToPlanNodeConverter.resolveDecimal
+ if (dataType.getScale() == 0 && dataType.getPrecision() <= 38) {
Review Comment:
Minor: this copies the thresholds of `resolveDecimal`. #19790 changed them,
and it was reverted in #19949. If it lands again, this copy drifts with no test
failure. For example, DECIMAL(10,0) gets 'INT' while the plan gives LONG_ARRAY.
Suggestion: take the element type from
`RelToPlanNodeConverter.convertToColumnDataType(aggCall.getType())`, like the
ANY_VALUE fix in #19990.
##########
pinot-query-runtime/src/test/resources/queries/UDFAggregates.json:
##########
@@ -285,5 +285,66 @@
"outputs": [["a", 1, 1, 2, 2], ["b", 2, 2, 2, 2], ["c", 3, 3, 3, 3]]
}
]
+ },
+ "pinot_array_agg_tests": {
+ "tables": {
+ "tbl": {
+ "schema": [
+ {"name": "int_col", "type": "INT"},
+ {"name": "long_col", "type": "LONG"},
+ {"name": "double_col", "type": "DOUBLE"},
+ {"name": "decimal_col", "type": "BIG_DECIMAL"},
+ {"name": "string_col", "type": "STRING"},
+ {"name": "json_col", "type": "JSON"},
+ {"name": "int_mv_col", "type": "INT", "isSingleValue": false},
+ {"name": "string_mv_col", "type": "STRING", "isSingleValue": false}
+ ],
+ "inputs": [
+ [1, 10, 1.5, "1.5", "a", "{\"k\":1}", [1, 2], ["x", "y"]],
+ [2, 20, 2.5, "2.5", "a", "{\"k\":2}", [2, 3], ["y"]],
+ [2, 20, 2.5, "2.5", "b", "{\"k\":2}", [3], ["z", "x"]],
+ [3, 30, 3.5, "3.5", "b", "{\"k\":3}", [4, 5], ["z"]]
+ ]
+ }
+ },
+ "queries": [
+ {
+ "comment": "The data type argument of ARRAY_AGG is inferred when
omitted",
+ "sql": "SELECT arraySortInt(ARRAY_AGG(int_col)),
arraySortInt(ARRAY_AGG(int_col, true)), arraySortString(ARRAY_AGG(string_col,
true)), arraySortString(ARRAY_AGG(json_col)), CARDINALITY(ARRAY_AGG(long_col,
true)), CARDINALITY(ARRAY_AGG(double_col)), CARDINALITY(ARRAY_AGG(decimal_col,
true)) FROM {tbl}",
+ "outputs": [[[1, 2, 2, 3], [1, 2, 3], ["a", "b"], ["{\"k\":1}",
"{\"k\":2}", "{\"k\":2}", "{\"k\":3}"], 3, 4, 3]]
+ },
+ {
+ "sql": "SELECT string_col, arraySortInt(ARRAY_AGG(int_col + 1, true)),
arraySortInt(ARRAY_AGG(int_mv_col, true)),
arraySortString(ARRAY_AGG(string_mv_col)) FROM {tbl} GROUP BY string_col",
+ "outputs": [["a", [2, 3], [1, 2, 3], ["x", "y", "y"]], ["b", [3, 4],
[3, 4, 5], ["x", "z", "z"]]]
+ },
+ {
+ "sql": "SELECT /*+ aggOptions(is_skip_leaf_stage_aggregate='true') */
string_col, arraySortInt(ARRAY_AGG(int_col + 1, true)),
arraySortInt(ARRAY_AGG(int_mv_col, true)),
arraySortString(ARRAY_AGG(string_mv_col)) FROM {tbl} GROUP BY string_col",
Review Comment:
Minor: `is_skip_leaf_stage_aggregate` is not a hint (the real one is
`is_skip_leaf_stage_group_by`), so this case does not test skip-leaf. With the
real hint, ARRAY_AGG over an MV column gives wrong results in all modes. v1 and
lite return the group key (for example `[a, a]`), and v2 returns nested arrays
or `[I cannot be cast to Number`. This happens on master too, for the explicit
form.
Suggestion: use the real hint, and drop the MV columns from this case or
mark them as a known bug. The FIRST/LAST_WITH_TIME cases at lines 38, 50 and 63
have the same wrong hint name.
--
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]
