yashmayya commented on code in PR #20091:
URL: https://github.com/apache/pinot/pull/20091#discussion_r4416889552


##########
pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java:
##########
@@ -838,6 +838,18 @@ public static class Broker {
     /// Calcite itself stops treating an IN list as a scalar predicate.
     public static final int DEFAULT_SEALED_IN_LIST_THRESHOLD = 20;
 
+    /// Whether the multi-stage engine resolves a simple name in `GROUP BY` to 
a `SELECT` alias before a column of the
+    /// same name. Calcite 1.42 (CALCITE-7052) resolves the column first, as 
MySQL and PostgreSQL do. Before that, Pinot
+    /// resolved the alias first. Enable this to keep the old behavior for 
queries such as
+    /// `SELECT CASE WHEN c = 'a' THEN 'x' ELSE c END AS c, COUNT(*) FROM t 
GROUP BY c`.
+    ///
+    /// This value can always be overridden by the 
[Request.QueryOptionKey#PREFER_SELECT_ALIAS_IN_GROUP_BY] query
+    /// option. Planning outside the broker's multi-stage request handler (for 
example for the controller `/sql`
+    /// endpoint) does not read this broker config and uses the default. The 
query option works everywhere.

Review Comment:
   Non-blocking. I checked this on a quickstart with only the broker config 
set. A query that fails with `Expression '...' is not being grouped` by default 
still fails through the controller `/sql` endpoint. The controller compiles the 
query with the default to find its tables, and the validation error stops the 
request there. The same query works through the broker. Queries that only 
return different rows pass, because their compile succeeds.
   
   Table names do not depend on how `GROUP BY` resolves. So 
`PinotQueryResource.getTableNames` can retry the compile with 
`preferSelectAliasInGroupBy=true` after a validation error, and let the broker 
decide.



##########
pinot-query-planner/src/main/java/org/apache/pinot/query/validate/Validator.java:
##########
@@ -62,11 +71,108 @@ public boolean isNonStrictGroupBy() {
         }
       };
 
-  public Validator(SqlOperatorTable opTab, SqlValidatorCatalogReader 
catalogReader, RelDataTypeFactory typeFactory) {
+  private final boolean _preferSelectAliasInGroupBy;
+  /// `GROUP BY` names that match a `SELECT` alias, per `SELECT`. See 
[#rewriteGroupByAliases].
+  private final Map<SqlSelect, List<String>> _groupByAliasNames = new 
IdentityHashMap<>();
+  /// Whether [BrokerMeter#QUERIES_WITH_GROUP_BY_ALIAS_SHADOWED_BY_COLUMN] was 
already marked for this query.
+  private boolean _groupByAliasShadowedByColumnFound;
+
+  /// @param preferSelectAliasInGroupBy whether a simple name in `GROUP BY` 
resolves to a `SELECT` alias before a
+  ///                                   column of the same name. See 
[#rewriteGroupByAliases].
+  public Validator(SqlOperatorTable opTab, SqlValidatorCatalogReader 
catalogReader, RelDataTypeFactory typeFactory,
+      boolean preferSelectAliasInGroupBy) {
     super(opTab, catalogReader, typeFactory,
         Config.DEFAULT.withConformance(PINOT_CONFORMANCE)
             .withIdentifierExpansion(true)
             .withTypeCoercionFactory(PinotTypeCoercion::new));
+    _preferSelectAliasInGroupBy = preferSelectAliasInGroupBy;
+  }
+
+  @Override
+  protected @PolyNull SqlNode performUnconditionalRewrites(@PolyNull SqlNode 
node, boolean underFrom) {
+    SqlNode rewritten = super.performUnconditionalRewrites(node, underFrom);
+    if (rewritten instanceof SqlSelect) {
+      rewriteGroupByAliases((SqlSelect) rewritten);
+    }
+    return rewritten;
+  }
+
+  /// Finds each simple name in the `GROUP BY` list that matches exactly one 
`SELECT` alias, unless the `SELECT` item is
+  /// the column of the same name. When [#_preferSelectAliasInGroupBy] is set, 
replaces the name with the ordinal of
+  /// that `SELECT` item.
+  ///
+  /// Calcite 1.42 (CALCITE-7052) resolves a `GROUP BY` name to a column of 
the `FROM` clause first, and to a `SELECT`
+  /// alias only when no column (or more than one) has that name. Calcite 1.40 
resolved the alias first. Given
+  /// `SELECT CASE WHEN c = 'a' THEN 'x' ELSE c END AS c, COUNT(*) FROM t 
GROUP BY c`, Calcite 1.40 groups by the
+  /// `CASE` expression and Calcite 1.42 groups by the column `c`. An ordinal 
always refers to the `SELECT` item, so
+  /// the rewrite restores the Calcite 1.40 result. It also makes `HAVING` 
resolve the name to the alias, as before.
+  ///
+  /// The names found are kept in [#_groupByAliasNames]. 
[#validateGroupClause] later checks which of them are also
+  /// columns. That check needs the scopes, which do not exist yet at this 
point.
+  ///
+  /// Only top-level `GROUP BY` items are considered. Names inside `ROLLUP`, 
`CUBE` and `GROUPING SETS` keep the
+  /// Calcite 1.42 resolution.
+  private void rewriteGroupByAliases(SqlSelect select) {
+    SqlNodeList groupList = select.getGroup();
+    if (groupList == null) {
+      return;
+    }
+    SqlNameMatcher nameMatcher = getCatalogReader().nameMatcher();
+    List<SqlNode> selectList = select.getSelectList().getList();
+    List<String> aliasNames = new ArrayList<>();
+    for (int i = 0; i < groupList.size(); i++) {
+      SqlNode groupItem = groupList.get(i);
+      if (!(groupItem instanceof SqlIdentifier) || !((SqlIdentifier) 
groupItem).isSimple()) {
+        continue;
+      }
+      String name = ((SqlIdentifier) groupItem).getSimple();
+      int matchIndex = -1;
+      int matchCount = 0;
+      for (int j = 0; j < selectList.size(); j++) {
+        String alias = SqlValidatorUtil.alias(selectList.get(j));
+        if (alias != null && nameMatcher.matches(alias, name)) {
+          matchIndex = j;
+          matchCount++;
+        }
+      }
+      // With more than one match, Calcite 1.40 failed with an ambiguity 
error. Keep the Calcite 1.42 resolution.
+      if (matchCount != 1) {
+        continue;
+      }
+      SqlNode expr = SqlUtil.stripAs(selectList.get(matchIndex));
+      if (expr instanceof SqlIdentifier && 
nameMatcher.matches(Util.last(((SqlIdentifier) expr).names), name)) {
+        continue;
+      }
+      aliasNames.add(name);
+      if (_preferSelectAliasInGroupBy) {
+        groupList.set(i, 
SqlLiteral.createExactNumeric(String.valueOf(matchIndex + 1), 
groupItem.getParserPosition()));
+      }
+    }
+    if (!aliasNames.isEmpty()) {
+      _groupByAliasNames.put(select, aliasNames);
+    }
+  }
+
+  /// Marks [BrokerMeter#QUERIES_WITH_GROUP_BY_ALIAS_SHADOWED_BY_COLUMN] when 
a name found by [#rewriteGroupByAliases]
+  /// is also exactly one column of the `FROM` clause. Calcite 1.40 and 
Calcite 1.42 resolve such a name differently.
+  ///
+  /// This runs before the validation of the `SELECT` list, so the meter is 
also marked when the query then fails with
+  /// `Expression '...' is not being grouped`.
+  @Override
+  protected void validateGroupClause(SqlSelect select) {
+    List<String> aliasNames = _groupByAliasNames.get(select);
+    if (aliasNames != null && !_groupByAliasShadowedByColumnFound) {
+      SelectScope selectScope = getRawSelectScope(select);
+      SqlNameMatcher nameMatcher = getCatalogReader().nameMatcher();
+      for (String name : aliasNames) {
+        if (selectScope != null && selectScope.findQualifyingTableNames(name, 
select, nameMatcher).size() == 1) {
+          _groupByAliasShadowedByColumnFound = true;
+          
BrokerMetrics.get().addMeteredGlobalValue(BrokerMeter.QUERIES_WITH_GROUP_BY_ALIAS_SHADOWED_BY_COLUMN,
 1);

Review Comment:
   Non-blocking. This also counts single-stage queries. These paths compile 
single-stage queries with this `Validator`:
   - the migration check (`pinot.broker.enable.multistage.migration.metric`), 
which compiles every single-stage query in the background;
   - `EmptyResponseUtils.fillEmptyResponseSchema` with 
`useMSEToFillEmptyResponseSchema`;
   - `ParserUtils.canCompileWithMultiStageEngine` after a single-stage compile 
error.
   
   The single-stage engine still groups by the alias, so nothing changed for 
those queries. But the meter, which the Javadoc describes as a count of 
multi-stage queries, goes up, and operators cannot tell the two cases apart. 
One option: mark the meter only when planning for 
`MultiStageBrokerRequestHandler`, for example with a `QueryEnvironment.Config` 
flag that only its `getQueryEnvConf` sets. A separate meter for the migration 
check is also useful: those single-stage queries return different rows after 
they move to the multi-stage engine.



-- 
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