Removing Unnest feature

somu-imply · Aug 24, 2023 · b4a8c62 · b4a8c62
1 parent d3c1e76
commit b4a8c62
Show file tree

Hide file tree

Showing 11 changed files with 125 additions and 46 deletions.
diff --git a/...ns-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/DataSourcePlan.java b/...ns-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/DataSourcePlan.java
@@ -140,16 +140,29 @@ public static DataSourcePlan forDataSource(
       checkQuerySegmentSpecIsEternity(dataSource, querySegmentSpec);
       return forLookup((LookupDataSource) dataSource, broadcast);
     } else if (dataSource instanceof FilteredDataSource) {
-      return forFilteredDataSource(queryKit, queryId, queryContext,
-                                   (FilteredDataSource) dataSource, querySegmentSpec, maxWorkerCount, minStageNumber, broadcast);
+      return forFilteredDataSource(
+          queryKit,
+          queryId,
+          queryContext,
+          (FilteredDataSource) dataSource,
+          querySegmentSpec,
+          maxWorkerCount,
+          minStageNumber,
+          broadcast
+      );
     } else if (dataSource instanceof UnnestDataSource) {
       checkQuerySegmentSpecIsEternity(dataSource, querySegmentSpec);
       return forUnnest(
-          queryKit, queryId, queryContext,
-          (UnnestDataSource) dataSource, querySegmentSpec, maxWorkerCount, minStageNumber, broadcast
+          queryKit,
+          queryId,
+          queryContext,
+          (UnnestDataSource) dataSource,
+          querySegmentSpec,
+          maxWorkerCount,
+          minStageNumber,
+          broadcast
       );
-    }
-    else if (dataSource instanceof QueryDataSource) {
+    } else if (dataSource instanceof QueryDataSource) {
       checkQuerySegmentSpecIsEternity(dataSource, querySegmentSpec);
       return forQuery(
           queryKit,
@@ -390,12 +403,15 @@ private static DataSourcePlan forFilteredDataSource(
 
     int shift = basePlan.getInputSpecs().size();
     newDataSource = FilteredDataSource.create(shiftInputNumbers(newDataSource, shift), dataSource.getFilter());
-    return new DataSourcePlan(newDataSource,
-                              inputSpecs,
-                              broadcast ? IntOpenHashSet.of(0) : IntSets.emptySet(),
-                              subQueryDefBuilder);
+    return new DataSourcePlan(
+        newDataSource,
+        inputSpecs,
+        broadcast ? IntOpenHashSet.of(0) : IntSets.emptySet(),
+        subQueryDefBuilder
+    );
 
   }
+
   /**
    * Build a plan for Unnest data source
    */
@@ -428,11 +444,17 @@ private static DataSourcePlan forUnnest(
     final List<InputSpec> inputSpecs = new ArrayList<>(basePlan.getInputSpecs());
 
     int shift = basePlan.getInputSpecs().size();
-    newDataSource = UnnestDataSource.create(shiftInputNumbers(newDataSource, shift ), dataSource.getVirtualColumn(), dataSource.getUnnestFilter());
-    return new DataSourcePlan(newDataSource,
-                              inputSpecs,
-                              broadcast ? IntOpenHashSet.of(0) : IntSets.emptySet(),
-                              subQueryDefBuilder);
+    newDataSource = UnnestDataSource.create(
+        shiftInputNumbers(newDataSource, shift),
+        dataSource.getVirtualColumn(),
+        dataSource.getUnnestFilter()
+    );
+    return new DataSourcePlan(
+        newDataSource,
+        inputSpecs,
+        broadcast ? IntOpenHashSet.of(0) : IntSets.emptySet(),
+        subQueryDefBuilder
+    );
   }
 
   /**

diff --git a/...sions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/MSQTaskSqlEngine.java b/...sions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/MSQTaskSqlEngine.java
@@ -114,7 +114,7 @@ public boolean featureAvailable(EngineFeature feature, PlannerContext plannerCon
       case TIME_BOUNDARY_QUERY:
       case GROUPING_SETS:
       case WINDOW_FUNCTIONS:
-      case UNNEST:
+        return false;
       case CAN_SELECT:
       case CAN_INSERT:
       case CAN_REPLACE:

diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQSelectTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQSelectTest.java
@@ -140,7 +140,7 @@ public static Collection<Object[]> data()
     };*/
 
     Object[][] data = new Object[][]{
-        {DEFAULT, UNNEST_CONTEXT}
+        {DEFAULT, DEFAULT_MSQ_CONTEXT}
     };
     return Arrays.asList(data);
   }
@@ -2312,6 +2312,64 @@ public void testSelectOnFoo5()
         .verifyResults();
   }
 
+  @Test
+  public void testSelectOnFoo6()
+  {
+    RowSignature resultSignature = RowSignature.builder()
+                                               .add("j0.unnest", ColumnType.STRING)
+                                               .build();
+
+    RowSignature outputSignature = RowSignature.builder()
+                                               .add("d3", ColumnType.STRING)
+                                               .build();
+
+    final ColumnMappings expectedColumnMappings = new ColumnMappings(
+        ImmutableList.of(
+            new ColumnMapping("j0.unnest", "d3")
+        )
+    );
+
+    testSelectQuery()
+        .setSql("SELECT COUNT(*) FROM foo")
+        .setExpectedMSQSpec(
+            MSQSpec.builder()
+                   .query(newScanQueryBuilder()
+                              .dataSource(UnnestDataSource.create(
+                                  new TableDataSource(CalciteTests.DATASOURCE1),
+                                  expressionVirtualColumn("j0.unnest", "\"dim3\"", ColumnType.STRING),
+                                  null
+                              ))
+                              .intervals(querySegmentSpec(Filtration.eternity()))
+                              .resultFormat(ScanQuery.ResultFormat.RESULT_FORMAT_COMPACTED_LIST)
+                              .legacy(false)
+                              .context(defaultScanQueryContext(
+                                  context,
+                                  resultSignature
+                              ))
+                              .columns(ImmutableList.of("j0.unnest"))
+                              .build())
+                   .columnMappings(expectedColumnMappings)
+                   .tuningConfig(MSQTuningConfig.defaultConfig())
+                   .destination(isDurableStorageDestination()
+                                ? DurableStorageMSQDestination.INSTANCE
+                                : TaskReportMSQDestination.INSTANCE)
+                   .build()
+        )
+        .setExpectedRowSignature(outputSignature)
+        .setQueryContext(context)
+        .setExpectedResultRows(ImmutableList.of(
+            new Object[]{"a"},
+            new Object[]{"b"},
+            new Object[]{"b"},
+            new Object[]{"c"},
+            new Object[]{"d"},
+            new Object[]{""},
+            new Object[]{null},
+            new Object[]{null}
+        ))
+        .verifyResults();
+  }
+
   @Nonnull
   private List<Object[]> expectedMultiValueFooRowsGroup()
   {

diff --git a/...-stage-query/src/test/java/org/apache/druid/msq/test/CalciteArraysSelectQueryMSQTest.java b/...-stage-query/src/test/java/org/apache/druid/msq/test/CalciteArraysSelectQueryMSQTest.java
@@ -28,7 +28,6 @@
 import org.apache.druid.query.groupby.TestGroupByBuffers;
 import org.apache.druid.server.QueryLifecycleFactory;
 import org.apache.druid.sql.calcite.CalciteArraysQueryTest;
-import org.apache.druid.sql.calcite.CalciteQueryTest;
 import org.apache.druid.sql.calcite.QueryTestBuilder;
 import org.apache.druid.sql.calcite.run.SqlEngine;
 import org.junit.After;

diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestBase.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestBase.java
@@ -154,7 +154,6 @@
 import org.apache.druid.sql.calcite.planner.CalciteRulesManager;
 import org.apache.druid.sql.calcite.planner.CatalogResolver;
 import org.apache.druid.sql.calcite.planner.PlannerConfig;
-import org.apache.druid.sql.calcite.planner.PlannerContext;
 import org.apache.druid.sql.calcite.planner.PlannerFactory;
 import org.apache.druid.sql.calcite.rel.DruidQuery;
 import org.apache.druid.sql.calcite.run.SqlEngine;
@@ -269,12 +268,6 @@ public class MSQTestBase extends BaseCalciteQueryTest
                                           .put(GroupByQueryConfig.CTX_KEY_ENABLE_MULTI_VALUE_UNNESTING, false)
                                           .build();
 
-  public static final Map<String, Object> UNNEST_CONTEXT =
-      ImmutableMap.<String, Object>builder()
-                  .putAll(DEFAULT_MSQ_CONTEXT)
-                  .put(PlannerContext.CTX_ENABLE_UNNEST, true)
-                  .build();
-
   public static final String FAULT_TOLERANCE = "fault_tolerance";
   public static final String DURABLE_STORAGE = "durable_storage";
   public static final String DEFAULT = "default";

diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerContext.java b/sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerContext.java
@@ -82,11 +82,6 @@ public class PlannerContext
    */
   public static final String CTX_ENABLE_WINDOW_FNS = "windowsAreForClosers";
 
-  /**
-   * Undocumented context key, used to enable {@link org.apache.calcite.sql.fun.SqlStdOperatorTable#UNNEST}.
-   */
-  public static final String CTX_ENABLE_UNNEST = "enableUnnest";
-
   public static final String CTX_SQL_USE_BOUNDS_AND_SELECTORS = "sqlUseBoundAndSelectors";
   public static final boolean DEFAULT_SQL_USE_BOUNDS_AND_SELECTORS = NullHandling.replaceWithDefault();
 
@@ -528,11 +523,11 @@ public boolean featureAvailable(final EngineFeature feature)
       return false;
     }
 
-    if (feature == EngineFeature.UNNEST &&
+    /*if (feature == EngineFeature.UNNEST &&
         !QueryContexts.getAsBoolean(CTX_ENABLE_UNNEST, queryContext.get(CTX_ENABLE_UNNEST), false)) {
       // Short-circuit: feature requires context flag.
       return false;
-    }
+    }*/
 
     return engine.featureAvailable(feature, this);
   }

diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/rule/DruidRules.java b/sql/src/main/java/org/apache/druid/sql/calcite/rule/DruidRules.java
@@ -115,13 +115,13 @@ public static List<RelOptRule> rules(PlannerContext plannerContext)
       retVal.add(DruidOuterQueryRule.WINDOW);
     }
 
-    if (plannerContext.featureAvailable(EngineFeature.UNNEST)) {
-      retVal.add(new DruidUnnestRule(plannerContext));
-      retVal.add(new DruidCorrelateUnnestRule(plannerContext));
-      retVal.add(CoreRules.PROJECT_CORRELATE_TRANSPOSE);
-      retVal.add(DruidFilterUnnestRule.instance());
-      retVal.add(DruidFilterUnnestRule.DruidProjectOnUnnestRule.instance());
-    }
+    // Adding unnest specific rules
+    retVal.add(new DruidUnnestRule(plannerContext));
+    retVal.add(new DruidCorrelateUnnestRule(plannerContext));
+    retVal.add(CoreRules.PROJECT_CORRELATE_TRANSPOSE);
+    retVal.add(DruidFilterUnnestRule.instance());
+    retVal.add(DruidFilterUnnestRule.DruidProjectOnUnnestRule.instance());
+
 
     return retVal;
   }

diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/run/NativeSqlEngine.java b/sql/src/main/java/org/apache/druid/sql/calcite/run/NativeSqlEngine.java
@@ -103,7 +103,6 @@ public boolean featureAvailable(EngineFeature feature, PlannerContext plannerCon
       case TOPN_QUERY:
       case GROUPING_SETS:
       case WINDOW_FUNCTIONS:
-      case UNNEST:
       case ALLOW_BROADCAST_RIGHTY_JOIN:
         return true;
       case TIME_BOUNDARY_QUERY:

diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/view/ViewSqlEngine.java b/sql/src/main/java/org/apache/druid/sql/calcite/view/ViewSqlEngine.java
@@ -62,9 +62,6 @@ public boolean featureAvailable(EngineFeature feature, PlannerContext plannerCon
       case SCAN_ORDER_BY_NON_TIME:
       case GROUPING_SETS:
       case WINDOW_FUNCTIONS:
-      case UNNEST:
-        return true;
-
       // Views can't sit on top of INSERT or REPLACE.
       case CAN_INSERT:
       case CAN_REPLACE:

diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/BaseCalciteQueryTest.java b/sql/src/test/java/org/apache/druid/sql/calcite/BaseCalciteQueryTest.java
@@ -203,7 +203,6 @@ public static void setupNullValues()
       ImmutableMap.<String, Object>builder()
                   .putAll(QUERY_CONTEXT_DEFAULT)
                   .put(QueryContexts.CTX_SQL_STRINGIFY_ARRAYS, false)
-                  .put(PlannerContext.CTX_ENABLE_UNNEST, true)
                   .build();
 
   public static final Map<String, Object> QUERY_CONTEXT_NO_STRINGIFY_ARRAY_USE_EQUALITY =