You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@pinot.apache.org by xi...@apache.org on 2021/10/26 23:43:00 UTC

[pinot] branch master updated: Reject query with identifiers not in schema (#7590)

This is an automated email from the ASF dual-hosted git repository.

xiangfu pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git


The following commit(s) were added to refs/heads/master by this push:
     new 23f7002  Reject query with identifiers not in schema (#7590)
23f7002 is described below

commit 23f7002831c00e206b14de9c084465f278495470
Author: Xiang Fu <xi...@gmail.com>
AuthorDate: Tue Oct 26 16:42:43 2021 -0700

    Reject query with identifiers not in schema (#7590)
    
    * Reject query with identifiers not in schema
    
    * Adding error code for unknown column exception
    
    * On the fly build alias map and check on that as well
---
 .../requesthandler/BaseBrokerRequestHandler.java   | 98 ++++++++++++++++------
 .../broker/requesthandler/QueryValidationTest.java | 82 ++++++++++++++++++
 .../pinot/common/exception/QueryException.java     |  4 +
 .../apache/pinot/common/metrics/BrokerMeter.java   |  2 +
 .../pinot/common/utils/helix/TableCache.java       | 11 +--
 .../pinot/sql/parsers/rewriter/AliasApplier.java   | 25 ------
 .../pinot/sql/parsers/CalciteSqlCompilerTest.java  |  4 +-
 7 files changed, 169 insertions(+), 57 deletions(-)

diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandler.java b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandler.java
index 2ac2f2a..34c746e 100644
--- a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandler.java
+++ b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandler.java
@@ -87,6 +87,7 @@ import org.apache.pinot.spi.config.table.TableConfig;
 import org.apache.pinot.spi.config.table.TableType;
 import org.apache.pinot.spi.data.Schema;
 import org.apache.pinot.spi.env.PinotConfiguration;
+import org.apache.pinot.spi.exception.BadQueryRequestException;
 import org.apache.pinot.spi.utils.BytesUtils;
 import org.apache.pinot.spi.utils.CommonConstants;
 import org.apache.pinot.spi.utils.CommonConstants.Broker;
@@ -246,8 +247,20 @@ public abstract class BaseBrokerRequestHandler implements BrokerRequestHandler {
     requestStatistics.setTableName(rawTableName);
 
     try {
-      updateColumnNames(rawTableName, pinotQuery);
+      boolean isCaseInsensitive = _tableCache.isCaseInsensitive();
+      Map<String, String> columnNameMap = _tableCache.getColumnNameMap(rawTableName);
+      if (columnNameMap != null) {
+        updateColumnNames(rawTableName, pinotQuery, isCaseInsensitive, columnNameMap);
+      }
     } catch (Exception e) {
+      // Throw exceptions with column in-existence error.
+      if (e instanceof BadQueryRequestException) {
+        LOGGER.info("Caught exception while checking column names in request, {}: {}, {}", requestId, query,
+            e.getMessage());
+        requestStatistics.setErrorCode(QueryException.UNKNOWN_COLUMN_ERROR_CODE);
+        _brokerMetrics.addMeteredTableValue(rawTableName, BrokerMeter.UNKNOWN_COLUMN_EXCEPTIONS, 1);
+        return new BrokerResponseNative(QueryException.getException(QueryException.UNKNOWN_COLUMN_ERROR, e));
+      }
       LOGGER.warn("Caught exception while updating column names in request {}: {}, {}", requestId, query,
           e.getMessage());
     }
@@ -1282,33 +1295,35 @@ public abstract class BaseBrokerRequestHandler implements BrokerRequestHandler {
   /**
    * Fixes the column names to the actual column names in the given SQL query.
    */
-  private void updateColumnNames(String rawTableName, PinotQuery pinotQuery) {
-    Map<String, String> columnNameMap =
-        _tableCache.isCaseInsensitive() ? _tableCache.getColumnNameMap(rawTableName) : null;
+  @VisibleForTesting
+  static void updateColumnNames(String rawTableName, PinotQuery pinotQuery, boolean isCaseInsensitive,
+      Map<String, String> columnNameMap) {
+    Map<String, String> aliasMap = new HashMap<>();
     if (pinotQuery != null) {
       for (Expression expression : pinotQuery.getSelectList()) {
-        fixColumnName(rawTableName, expression, columnNameMap);
+        fixColumnName(rawTableName, expression, columnNameMap, aliasMap, isCaseInsensitive);
       }
       Expression filterExpression = pinotQuery.getFilterExpression();
       if (filterExpression != null) {
-        fixColumnName(rawTableName, filterExpression, columnNameMap);
+        fixColumnName(rawTableName, filterExpression, columnNameMap, aliasMap, isCaseInsensitive);
       }
       List<Expression> groupByList = pinotQuery.getGroupByList();
       if (groupByList != null) {
         for (Expression expression : groupByList) {
-          fixColumnName(rawTableName, expression, columnNameMap);
+          fixColumnName(rawTableName, expression, columnNameMap, aliasMap, isCaseInsensitive);
         }
       }
       List<Expression> orderByList = pinotQuery.getOrderByList();
       if (orderByList != null) {
         for (Expression expression : orderByList) {
           // NOTE: Order-by is always a Function with the ordering of the Expression
-          fixColumnName(rawTableName, expression.getFunctionCall().getOperands().get(0), columnNameMap);
+          fixColumnName(rawTableName, expression.getFunctionCall().getOperands().get(0), columnNameMap, aliasMap,
+              isCaseInsensitive);
         }
       }
       Expression havingExpression = pinotQuery.getHavingExpression();
       if (havingExpression != null) {
-        fixColumnName(rawTableName, havingExpression, columnNameMap);
+        fixColumnName(rawTableName, havingExpression, columnNameMap, aliasMap, isCaseInsensitive);
       }
     }
   }
@@ -1385,7 +1400,8 @@ public abstract class BaseBrokerRequestHandler implements BrokerRequestHandler {
       @Nullable Map<String, String> columnNameMap) {
     TransformExpressionTree.ExpressionType expressionType = expression.getExpressionType();
     if (expressionType == TransformExpressionTree.ExpressionType.IDENTIFIER) {
-      expression.setValue(getActualColumnName(rawTableName, expression.getValue(), columnNameMap));
+      expression.setValue(getActualColumnName(rawTableName, expression.getValue(), columnNameMap, null,
+          _tableCache.isCaseInsensitive()));
     } else if (expressionType == TransformExpressionTree.ExpressionType.FUNCTION) {
       for (TransformExpressionTree child : expression.getChildren()) {
         fixColumnName(rawTableName, child, columnNameMap);
@@ -1396,14 +1412,33 @@ public abstract class BaseBrokerRequestHandler implements BrokerRequestHandler {
   /**
    * Fixes the column names to the actual column names in the given SQL expression.
    */
-  private void fixColumnName(String rawTableName, Expression expression, @Nullable Map<String, String> columnNameMap) {
+  private static void fixColumnName(String rawTableName, Expression expression, Map<String, String> columnNameMap,
+      Map<String, String> aliasMap, boolean isCaseInsensitive) {
     ExpressionType expressionType = expression.getType();
     if (expressionType == ExpressionType.IDENTIFIER) {
       Identifier identifier = expression.getIdentifier();
-      identifier.setName(getActualColumnName(rawTableName, identifier.getName(), columnNameMap));
+      identifier.setName(
+          getActualColumnName(rawTableName, identifier.getName(), columnNameMap, aliasMap, isCaseInsensitive));
     } else if (expressionType == ExpressionType.FUNCTION) {
-      for (Expression operand : expression.getFunctionCall().getOperands()) {
-        fixColumnName(rawTableName, operand, columnNameMap);
+      final Function functionCall = expression.getFunctionCall();
+      switch (functionCall.getOperator()) {
+        case "AS":
+          fixColumnName(rawTableName, functionCall.getOperands().get(0), columnNameMap, aliasMap, isCaseInsensitive);
+          final Expression rightAsExpr = functionCall.getOperands().get(1);
+          if (rightAsExpr.isSetIdentifier()) {
+            String rightColumn = rightAsExpr.getIdentifier().getName();
+            if (isCaseInsensitive) {
+              aliasMap.put(rightColumn.toLowerCase(), rightColumn);
+            } else {
+              aliasMap.put(rightColumn, rightColumn);
+            }
+          }
+          break;
+        default:
+          for (Expression operand : functionCall.getOperands()) {
+            fixColumnName(rawTableName, operand, columnNameMap, aliasMap, isCaseInsensitive);
+          }
+          break;
       }
     }
   }
@@ -1413,25 +1448,40 @@ public abstract class BaseBrokerRequestHandler implements BrokerRequestHandler {
    * - Case-insensitive cluster
    * - Column name in the format of [table_name].[column_name]
    */
-  private String getActualColumnName(String rawTableName, String columnName,
-      @Nullable Map<String, String> columnNameMap) {
+  private static String getActualColumnName(String rawTableName, String columnName,
+      @Nullable Map<String, String> columnNameMap, @Nullable Map<String, String> aliasMap, boolean isCaseInsensitive) {
+    if ("*".equals(columnName)) {
+      return columnName;
+    }
     // Check if column is in the format of [table_name].[column_name]
     String[] splits = StringUtils.split(columnName, ".", 2);
-    if (_tableCache.isCaseInsensitive()) {
+    String columnNameToCheck;
+    if (isCaseInsensitive) {
       if (splits.length == 2 && rawTableName.equalsIgnoreCase(splits[0])) {
-        columnName = splits[1];
-      }
-      if (columnNameMap != null) {
-        return columnNameMap.getOrDefault(columnName.toLowerCase(), columnName);
+        columnNameToCheck = splits[1].toLowerCase();
       } else {
-        return columnName;
+        columnNameToCheck = columnName.toLowerCase();
       }
     } else {
       if (splits.length == 2 && rawTableName.equals(splits[0])) {
-        columnName = splits[1];
+        columnNameToCheck = splits[1];
+      } else {
+        columnNameToCheck = columnName;
+      }
+    }
+    if (columnNameMap != null) {
+      String actualColumnName = columnNameMap.get(columnNameToCheck);
+      if (actualColumnName != null) {
+        return actualColumnName;
+      }
+    }
+    if (aliasMap != null) {
+      String actualAlias = aliasMap.get(columnNameToCheck);
+      if (actualAlias != null) {
+        return actualAlias;
       }
-      return columnName;
     }
+    throw new BadQueryRequestException("Unknown columnName '" + columnName + "' found in the query");
   }
 
   private static Map<String, String> getOptionsFromJson(JsonNode request, String optionsKey) {
diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/QueryValidationTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/QueryValidationTest.java
index 96e7768..a713eb4 100644
--- a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/QueryValidationTest.java
+++ b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/QueryValidationTest.java
@@ -19,6 +19,8 @@
 
 package org.apache.pinot.broker.requesthandler;
 
+import com.google.common.collect.ImmutableMap;
+import java.util.Map;
 import org.apache.pinot.common.request.BrokerRequest;
 import org.apache.pinot.common.request.PinotQuery;
 import org.apache.pinot.pql.parsers.Pql2Compiler;
@@ -120,6 +122,65 @@ public class QueryValidationTest {
     testUnsupportedPQLQuery(pql, "Aggregation functions cannot be used with DISTINCT");
   }
 
+  @Test
+  public void testUnsupportedNonExistColumnsQueries() {
+    String sql = "SELECT DISTINCT(col1, col2) FROM foo OPTION(groupByMode=sql,responseFormat=sql)";
+    testNonExistedColumnInSQLQuery("foo", false, ImmutableMap.of("col1", "col1"), sql,
+        "Unknown columnName 'col2' found in the query");
+    testNonExistedColumnInSQLQuery("foo", false, ImmutableMap.of("col2", "col2"), sql,
+        "Unknown columnName 'col1' found in the query");
+    testExistedColumnInSQLQuery("foo", false, ImmutableMap.of("col2", "col2", "col1", "col1"), sql);
+    sql = "SELECT sum(Col1) FROM foo OPTION(groupByMode=sql,responseFormat=sql)";
+    testNonExistedColumnInSQLQuery("foo", false, ImmutableMap.of("col1", "col1"), sql,
+        "Unknown columnName 'Col1' found in the query");
+    testExistedColumnInSQLQuery("foo", false, ImmutableMap.of("Col1", "Col1"), sql);
+    testExistedColumnInSQLQuery("foo", true, ImmutableMap.of("col1", "col1"), sql);
+    testExistedColumnInSQLQuery("foo", true, ImmutableMap.of("col1", "Col1"), sql);
+    sql = "SELECT sum(Col1) AS sum_col1 FROM foo OPTION(groupByMode=sql,responseFormat=sql)";
+    testNonExistedColumnInSQLQuery("foo", false, ImmutableMap.of("col1", "col1"), sql,
+        "Unknown columnName 'Col1' found in the query");
+    testExistedColumnInSQLQuery("foo", false, ImmutableMap.of("Col1", "Col1"), sql);
+    testExistedColumnInSQLQuery("foo", true, ImmutableMap.of("col1", "col1"), sql);
+    testExistedColumnInSQLQuery("foo", true, ImmutableMap.of("col1", "Col1"), sql);
+    sql = "SELECT sum(Col1) AS sum_col1 FROM foo HAVING sum_col1 > 10 OPTION(groupByMode=sql,responseFormat=sql)";
+    testNonExistedColumnInSQLQuery("foo", false, ImmutableMap.of("col1", "col1"), sql,
+        "Unknown columnName 'Col1' found in the query");
+    testNonExistedColumnInSQLQuery("foo", false, ImmutableMap.of("col1", "cOL1"), sql,
+        "Unknown columnName 'Col1' found in the query");
+    testExistedColumnInSQLQuery("foo", false, ImmutableMap.of("Col1", "Col1"), sql);
+    testExistedColumnInSQLQuery("foo", true, ImmutableMap.of("col1", "col1"), sql);
+    testExistedColumnInSQLQuery("foo", true, ImmutableMap.of("col1", "Col1"), sql);
+    testExistedColumnInSQLQuery("foo", true, ImmutableMap.of("col1", "cOL1"), sql);
+    sql = "SELECT sum(Col1) AS sum_col1, b AS B, c as D FROM foo GROUP BY B, D OPTION(groupByMode=sql,"
+        + "responseFormat=sql)";
+    testNonExistedColumnInSQLQuery("foo", false, ImmutableMap.of("col1", "col1", "b", "b", "c", "c"), sql,
+        "Unknown columnName 'Col1' found in the query");
+    testNonExistedColumnInSQLQuery("foo", false, ImmutableMap.of("Col1", "Col1", "B", "B", "c", "c"), sql,
+        "Unknown columnName 'b' found in the query");
+    testNonExistedColumnInSQLQuery("foo", false, ImmutableMap.of("Col1", "Col1", "c", "c"), sql,
+        "Unknown columnName 'b' found in the query");
+    testNonExistedColumnInSQLQuery("foo", false, ImmutableMap.of("Col1", "Col1", "b", "b", "C", "C"), sql,
+        "Unknown columnName 'c' found in the query");
+    testExistedColumnInSQLQuery("foo", false, ImmutableMap.of("Col1", "Col1", "b", "b", "c", "c"), sql);
+    testExistedColumnInSQLQuery("foo", true, ImmutableMap.of("col1", "col1", "b", "b", "c", "c"), sql);
+    testExistedColumnInSQLQuery("foo", true, ImmutableMap.of("col1", "COL1", "b", "B", "c", "C"), sql);
+    sql = "SELECT sum(Col1) AS sum_col1, b AS B, c as D FROM foo GROUP BY 2, 3 OPTION(groupByMode=sql,"
+        + "responseFormat=sql)";
+    testNonExistedColumnInSQLQuery("foo", false, ImmutableMap.of("col1", "col1", "B", "B", "c", "c", "D", "D"), sql,
+        "Unknown columnName 'Col1' found in the query");
+    testNonExistedColumnInSQLQuery("foo", false, ImmutableMap.of("col1", "col1", "b", "b", "c", "c"), sql,
+        "Unknown columnName 'Col1' found in the query");
+    testNonExistedColumnInSQLQuery("foo", false, ImmutableMap.of("Col1", "Col1", "B", "B", "c", "c"), sql,
+        "Unknown columnName 'b' found in the query");
+    testNonExistedColumnInSQLQuery("foo", false, ImmutableMap.of("Col1", "Col1", "c", "c"), sql,
+        "Unknown columnName 'b' found in the query");
+    testNonExistedColumnInSQLQuery("foo", false, ImmutableMap.of("Col1", "Col1", "b", "b", "C", "C"), sql,
+        "Unknown columnName 'c' found in the query");
+    testExistedColumnInSQLQuery("foo", false, ImmutableMap.of("Col1", "Col1", "b", "b", "c", "c", "D", "D"), sql);
+    testExistedColumnInSQLQuery("foo", true, ImmutableMap.of("col1", "col1", "b", "b", "c", "c", "d", "d"), sql);
+    testExistedColumnInSQLQuery("foo", true, ImmutableMap.of("col1", "COL1", "b", "B", "c", "C"), sql);
+  }
+
   private void testUnsupportedPQLQuery(String query, String errorMessage) {
     try {
       BrokerRequest brokerRequest = PQL_COMPILER.compileToBrokerRequest(query);
@@ -139,4 +200,25 @@ public class QueryValidationTest {
       Assert.assertEquals(e.getMessage(), errorMessage);
     }
   }
+
+  private void testNonExistedColumnInSQLQuery(String rawTableName, boolean isCaseInsensitive,
+      Map<String, String> columnNameMap, String query, String errorMessage) {
+    try {
+      PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(query);
+      BaseBrokerRequestHandler.updateColumnNames(rawTableName, pinotQuery, isCaseInsensitive, columnNameMap);
+      Assert.fail("Query should have failed");
+    } catch (Exception e) {
+      Assert.assertEquals(errorMessage, e.getMessage());
+    }
+  }
+
+  private void testExistedColumnInSQLQuery(String rawTableName, boolean isCaseInsensitive,
+      Map<String, String> columnNameMap, String query) {
+    try {
+      PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(query);
+      BaseBrokerRequestHandler.updateColumnNames(rawTableName, pinotQuery, isCaseInsensitive, columnNameMap);
+    } catch (Exception e) {
+      Assert.fail("Query should have succeed");
+    }
+  }
 }
diff --git a/pinot-common/src/main/java/org/apache/pinot/common/exception/QueryException.java b/pinot-common/src/main/java/org/apache/pinot/common/exception/QueryException.java
index 0cacd51..2fd75a5 100644
--- a/pinot-common/src/main/java/org/apache/pinot/common/exception/QueryException.java
+++ b/pinot-common/src/main/java/org/apache/pinot/common/exception/QueryException.java
@@ -75,6 +75,7 @@ public class QueryException {
   public static final int FEDERATED_BROKER_UNAVAILABLE_ERROR_CODE = 550;
   public static final int COMBINE_GROUP_BY_EXCEPTION_ERROR_CODE = 600;
   public static final int QUERY_VALIDATION_ERROR_CODE = 700;
+  public static final int UNKNOWN_COLUMN_ERROR_CODE = 710;
   public static final int UNKNOWN_ERROR_CODE = 1000;
   // NOTE: update isClientError() method appropriately when new codes are added
 
@@ -117,6 +118,7 @@ public class QueryException {
   public static final ProcessingException COMBINE_GROUP_BY_EXCEPTION_ERROR =
       new ProcessingException(COMBINE_GROUP_BY_EXCEPTION_ERROR_CODE);
   public static final ProcessingException QUERY_VALIDATION_ERROR = new ProcessingException(QUERY_VALIDATION_ERROR_CODE);
+  public static final ProcessingException UNKNOWN_COLUMN_ERROR = new ProcessingException(UNKNOWN_COLUMN_ERROR_CODE);
   public static final ProcessingException UNKNOWN_ERROR = new ProcessingException(UNKNOWN_ERROR_CODE);
   public static final ProcessingException QUOTA_EXCEEDED_ERROR = new ProcessingException(TOO_MANY_REQUESTS_ERROR_CODE);
 
@@ -145,6 +147,7 @@ public class QueryException {
     FEDERATED_BROKER_UNAVAILABLE_ERROR.setMessage("FederatedBrokerUnavailableError");
     COMBINE_GROUP_BY_EXCEPTION_ERROR.setMessage("CombineGroupByExceptionError");
     QUERY_VALIDATION_ERROR.setMessage("QueryValidationError");
+    UNKNOWN_COLUMN_ERROR.setMessage("UnknownColumnError");
     UNKNOWN_ERROR.setMessage("UnknownError");
     QUOTA_EXCEEDED_ERROR.setMessage("QuotaExceededError");
   }
@@ -200,6 +203,7 @@ public class QueryException {
       case QueryException.JSON_COMPILATION_ERROR_CODE:
       case QueryException.JSON_PARSING_ERROR_CODE:
       case QueryException.QUERY_VALIDATION_ERROR_CODE:
+      case QueryException.UNKNOWN_COLUMN_ERROR_CODE:
       case QueryException.PQL_PARSING_ERROR_CODE:
       case QueryException.TOO_MANY_REQUESTS_ERROR_CODE:
       case QueryException.TABLE_DOES_NOT_EXIST_ERROR_CODE:
diff --git a/pinot-common/src/main/java/org/apache/pinot/common/metrics/BrokerMeter.java b/pinot-common/src/main/java/org/apache/pinot/common/metrics/BrokerMeter.java
index 183106d..46ce867 100644
--- a/pinot-common/src/main/java/org/apache/pinot/common/metrics/BrokerMeter.java
+++ b/pinot-common/src/main/java/org/apache/pinot/common/metrics/BrokerMeter.java
@@ -38,6 +38,8 @@ public enum BrokerMeter implements AbstractMetrics.Meter {
   RESOURCE_MISSING_EXCEPTIONS("exceptions", true),
   // Query validation phase.
   QUERY_VALIDATION_EXCEPTIONS("exceptions", false),
+  // Query validation phase.
+  UNKNOWN_COLUMN_EXCEPTIONS("exceptions", false),
   // Scatter phase.
   NO_SERVER_FOUND_EXCEPTIONS("exceptions", false),
   REQUEST_TIMEOUT_BEFORE_SCATTERED_EXCEPTIONS("exceptions", false),
diff --git a/pinot-common/src/main/java/org/apache/pinot/common/utils/helix/TableCache.java b/pinot-common/src/main/java/org/apache/pinot/common/utils/helix/TableCache.java
index 58ca9ee..b94cd4c 100644
--- a/pinot-common/src/main/java/org/apache/pinot/common/utils/helix/TableCache.java
+++ b/pinot-common/src/main/java/org/apache/pinot/common/utils/helix/TableCache.java
@@ -130,9 +130,8 @@ public class TableCache {
    */
   @Nullable
   public Map<String, String> getColumnNameMap(String rawTableName) {
-    Preconditions.checkState(_caseInsensitive, "TableCache is not case-insensitive");
     String schemaName = _schemaNameMap.getOrDefault(rawTableName, rawTableName);
-    SchemaInfo schemaInfo = _schemaInfoMap.get(schemaName);
+    SchemaInfo schemaInfo = _schemaInfoMap.getOrDefault(schemaName, _schemaInfoMap.get(rawTableName));
     return schemaInfo != null ? schemaInfo._columnNameMap : null;
   }
 
@@ -248,15 +247,17 @@ public class TableCache {
       throws IOException {
     Schema schema = SchemaUtils.fromZNRecord(znRecord);
     String schemaName = schema.getSchemaName();
+    Map<String, String> columnNameMap = new HashMap<>();
     if (_caseInsensitive) {
-      Map<String, String> columnNameMap = new HashMap<>();
       for (String columnName : schema.getColumnNames()) {
         columnNameMap.put(columnName.toLowerCase(), columnName);
       }
-      _schemaInfoMap.put(schemaName, new SchemaInfo(schema, columnNameMap));
     } else {
-      _schemaInfoMap.put(schemaName, new SchemaInfo(schema, null));
+      for (String columnName : schema.getColumnNames()) {
+        columnNameMap.put(columnName, columnName);
+      }
     }
+    _schemaInfoMap.put(schemaName, new SchemaInfo(schema, columnNameMap));
   }
 
   private void removeSchema(String path) {
diff --git a/pinot-common/src/main/java/org/apache/pinot/sql/parsers/rewriter/AliasApplier.java b/pinot-common/src/main/java/org/apache/pinot/sql/parsers/rewriter/AliasApplier.java
index 1acf671..680b2f8 100644
--- a/pinot-common/src/main/java/org/apache/pinot/sql/parsers/rewriter/AliasApplier.java
+++ b/pinot-common/src/main/java/org/apache/pinot/sql/parsers/rewriter/AliasApplier.java
@@ -112,30 +112,5 @@ public class AliasApplier implements QueryRewriter {
         throw new SqlCompilationException("Duplicated alias name found.");
       }
     }
-    for (Expression selectExpr : pinotQuery.getSelectList()) {
-      matchIdentifierInAliasMap(selectExpr, aliasKeys);
-    }
-  }
-
-  private static void matchIdentifierInAliasMap(Expression selectExpr, Set<String> aliasKeys)
-      throws SqlCompilationException {
-    Function functionCall = selectExpr.getFunctionCall();
-    if (functionCall != null) {
-      if (functionCall.getOperator().equalsIgnoreCase(SqlKind.AS.toString())) {
-        matchIdentifierInAliasMap(functionCall.getOperands().get(0), aliasKeys);
-      } else {
-        if (functionCall.getOperandsSize() > 0) {
-          for (Expression operand : functionCall.getOperands()) {
-            matchIdentifierInAliasMap(operand, aliasKeys);
-          }
-        }
-      }
-    }
-    if (selectExpr.getIdentifier() != null) {
-      if (aliasKeys.contains(selectExpr.getIdentifier().getName().toLowerCase())) {
-        throw new SqlCompilationException(
-            "Alias " + selectExpr.getIdentifier().getName() + " cannot be referred in SELECT Clause");
-      }
-    }
   }
 }
diff --git a/pinot-common/src/test/java/org/apache/pinot/sql/parsers/CalciteSqlCompilerTest.java b/pinot-common/src/test/java/org/apache/pinot/sql/parsers/CalciteSqlCompilerTest.java
index bce5c5e..5e1c278 100644
--- a/pinot-common/src/test/java/org/apache/pinot/sql/parsers/CalciteSqlCompilerTest.java
+++ b/pinot-common/src/test/java/org/apache/pinot/sql/parsers/CalciteSqlCompilerTest.java
@@ -1295,10 +1295,8 @@ public class CalciteSqlCompilerTest {
     try {
       sql = "SELECT C1 AS ALIAS_C1, C2 AS ALIAS_C2, ADD(alias_c1, alias_c2) FROM Foo";
       CalciteSqlParser.compileToPinotQuery(sql);
-      Assert.fail("Query should have failed compilation");
     } catch (Exception e) {
-      Assert.assertTrue(e instanceof SqlCompilationException);
-      Assert.assertTrue(e.getMessage().contains("cannot be referred in SELECT Clause"));
+      Assert.fail("Query compilation shouldn't fail");
     }
   }
 

---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org