add schema fallback tests

apache · Jan 21, 2025 · 6e6f371 · 6e6f371
1 parent fe1c228
commit 6e6f371
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 10 deletions.
diff --git a/...-tests/src/test/java/org/apache/pinot/integration/tests/EmptyResponseIntegrationTest.java b/...-tests/src/test/java/org/apache/pinot/integration/tests/EmptyResponseIntegrationTest.java
@@ -221,18 +221,43 @@ public void testCompiledByV2AggregatedFields() throws Exception {
   }
 
   @Test
-  public void testSchemaFallbackStarField() {
-    // TODO: find an example that won't compile for V2
+  public void testSchemaFallbackStarField() throws Exception {
+    String sqlQuery = "SELECT * FROM myTable WHERE AirlineID = 0 AND add(AirTime, AirTime, ArrTime) > 0 LIMIT 1";
+    JsonNode response = postQuery(sqlQuery);
+    assertNoRowsReturned(response);
+    assertDataTypes(response, "INT", "INT", "LONG", "INT", "FLOAT", "DOUBLE", "INT", "STRING", "INT", "INT", "INT",
+        "INT", "STRING", "INT", "STRING", "INT", "INT", "INT", "INT", "INT", "DOUBLE", "FLOAT", "INT", "STRING", "INT",
+        "STRING", "INT", "INT", "INT", "STRING", "STRING", "INT", "STRING", "INT", "INT", "INT", "INT", "INT", "INT",
+        "INT", "STRING", "INT", "INT", "FLOAT", "INT", "STRING", "LONG", "INT", "INT", "INT", "INT", "STRING", "INT",
+        "INT", "INT", "INT", "INT", "INT", "STRING", "INT", "INT", "INT", "STRING", "STRING", "INT", "STRING", "INT",
+        "INT", "STRING", "INT", "STRING", "INT", "INT", "INT", "STRING", "INT", "INT", "INT", "INT");
+  }
+
+  @Test
+  public void testSchemaFallbackSelectionFields() throws Exception {
+    String sqlQuery = "SELECT AirlineID, ArrTime, ArrTimeBlk FROM myTable"
+      + " WHERE AirlineID = 0 AND add(ArrTime, ArrTime, ArrTime) > 0 LIMIT 1";
+    JsonNode response = postQuery(sqlQuery);
+    assertNoRowsReturned(response);
+    assertDataTypes(response, "LONG", "INT", "STRING");
   }
 
   @Test
-  public void testSchemaFallbackSelectionFields() {
-    // TODO: find an example that won't compile for V2
+  public void testSchemaFallbackTransformedFields() throws Exception {
+    String sqlQuery = "SELECT AirlineID, ArrTime, ArrTime+1 FROM myTable"
+        + " WHERE AirlineID = 0 AND add(ArrTime, ArrTime, ArrTime) > 0 LIMIT 1";
+    JsonNode response = postQuery(sqlQuery);
+    assertNoRowsReturned(response);
+    assertDataTypes(response, "LONG", "INT", "STRING");
   }
 
   @Test
-  public void testSchemaFallbackNonSelectionFields() {
-    // TODO: find an example that won't compile for V2
+  public void testSchemaFallbackAggregatedFields() throws Exception {
+    String sqlQuery = "SELECT AirlineID, avg(ArrTime) FROM myTable"
+        + " WHERE AirlineID = 0 AND add(ArrTime, ArrTime, ArrTime) > 0 GROUP BY AirlineID LIMIT 1";
+    JsonNode response = postQuery(sqlQuery);
+    assertNoRowsReturned(response);
+    assertDataTypes(response, "LONG", "DOUBLE");
   }
 
   private void assertNoRowsReturned(JsonNode response) {

diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java
@@ -287,8 +287,9 @@ public boolean canCompileQuery(String query) {
    */
   @Nullable
   public RelRoot getRelRootIfCanCompile(String query) {
-    SqlNodeAndOptions sqlNodeAndOptions = CalciteSqlParser.compileToSqlNodeAndOptions(query);
-    try (PlannerContext plannerContext = getPlannerContext(sqlNodeAndOptions)) {
+    try {
+      SqlNodeAndOptions sqlNodeAndOptions = CalciteSqlParser.compileToSqlNodeAndOptions(query);
+      PlannerContext plannerContext = getPlannerContext(sqlNodeAndOptions);
       SqlNode sqlNode = sqlNodeAndOptions.getSqlNode();
       if (sqlNode.getKind().equals(SqlKind.EXPLAIN)) {
         sqlNode = ((SqlExplain) sqlNode).getExplicandum();

diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/parser/utils/ParserUtils.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/parser/utils/ParserUtils.java
@@ -111,10 +111,12 @@ public static void fillEmptyResponseSchema(
       for (int i = 0; i < responseSchema.size(); i++) {
         FieldSpec spec = schema.getFieldSpecFor(responseSchema.getColumnName(i));
         try {
-          resolved = DataSchema.ColumnDataType.fromDataType(spec.getDataType(), false);
+          // try single value first
+          resolved = DataSchema.ColumnDataType.fromDataType(spec.getDataType(), true);
         } catch (Exception e) {
           try {
-            resolved = DataSchema.ColumnDataType.fromDataType(spec.getDataType(), true);
+            // fallback to multi value
+            resolved = DataSchema.ColumnDataType.fromDataType(spec.getDataType(), false);
           } catch (Exception e2) {
             resolved = DataSchema.ColumnDataType.UNKNOWN;
           }