xiangfu0 commented on code in PR #19988:
URL: https://github.com/apache/pinot/pull/19988#discussion_r4382082116
##########
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:
Done in b9d77debd1: any ARRAY_AGG with the DISTINCT flag now gets `true` as
isDistinct. That also fixes `ARRAY_AGG(DISTINCT x, false)` and the explicit
`ARRAY_AGG(DISTINCT x, 'INT')`, and plan and runtime tests cover both.
_🤖 Addressed by [Claude Code](https://claude.com/claude-code)_
##########
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:
Done in b9d77debd1: the early instance now runs just before
`AggregateCaseToFilter`, as suggested. New INTERSECT ALL (with FIRSTWITHTIME)
and EXCEPT ALL runtime cases fail with the previous placement, and no plans
changed.
_🤖 Addressed by [Claude Code](https://claude.com/claude-code)_
##########
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:
Done in b9d77debd1: the DECIMAL element type now comes from
`RelToPlanNodeConverter.convertToColumnDataType(aggCall.getType())`, so it
follows `resolveDecimal` without copying its thresholds.
_🤖 Addressed by [Claude Code](https://claude.com/claude-code)_
##########
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:
Done in b9d77debd1: the FIRST/LAST_WITH_TIME and ARRAY_AGG skip-leaf cases
now use `is_skip_leaf_stage_group_by`, and the ARRAY_AGG case no longer
aggregates multi-value columns. The skip-leaf multi-value bug from master is
not fixed here.
_🤖 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]
