fallback SQL IN filter to expression filter when VirtualColumnRegistr…

…y is null (#16836)
apache · Aug 5, 2024 · c7eacd0 · c7eacd0
1 parent 31b4375
commit c7eacd0
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 0 deletions.
diff --git a/...java/org/apache/druid/sql/calcite/expression/builtin/ScalarInArrayOperatorConversion.java b/...java/org/apache/druid/sql/calcite/expression/builtin/ScalarInArrayOperatorConversion.java
@@ -94,6 +94,9 @@ public DimFilter toDruidFilter(
     } else if (scalarExpression.isSimpleExtraction() && plannerContext.isUseLegacyInFilter()) {
       scalarColumn = scalarExpression.getSimpleExtraction().getColumn();
       scalarExtractionFn = scalarExpression.getSimpleExtraction().getExtractionFn();
+    } else if (virtualColumnRegistry == null) {
+      // virtual column registry unavailable, fallback to expression filter
+      return null;
     } else {
       scalarColumn = virtualColumnRegistry.getOrCreateVirtualColumnForExpression(
           scalarExpression,

diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/CalciteSubqueryTest.java b/sql/src/test/java/org/apache/druid/sql/calcite/CalciteSubqueryTest.java
@@ -58,8 +58,12 @@
 import org.apache.druid.query.dimension.ExtractionDimensionSpec;
 import org.apache.druid.query.expression.TestExprMacroTable;
 import org.apache.druid.query.extraction.SubstringDimExtractionFn;
+import org.apache.druid.query.filter.DimFilter;
+import org.apache.druid.query.filter.InDimFilter;
+import org.apache.druid.query.filter.TypedInFilter;
 import org.apache.druid.query.groupby.GroupByQuery;
 import org.apache.druid.query.groupby.orderby.DefaultLimitSpec;
+import org.apache.druid.query.groupby.orderby.NoopLimitSpec;
 import org.apache.druid.query.groupby.orderby.OrderByColumnSpec;
 import org.apache.druid.query.ordering.StringComparators;
 import org.apache.druid.query.scan.ScanQuery;
@@ -92,6 +96,7 @@
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.stream.Collectors;
@@ -1601,6 +1606,36 @@ public void testTimeseriesSubqueryWithEarliestAggregator(String testName, Map<St
     );
   }
 
+  @MethodSource("constructorFeeder")
+  @ParameterizedTest(name = "{0}")
+  public void testScalarInArrayToUseHavingFilter(String testName, Map<String, Object> queryContext)
+  {
+    DimFilter filter = NullHandling.replaceWithDefault()
+                       ? new InDimFilter("v0", new HashSet<>(Arrays.asList("1", "17")))
+                       : new TypedInFilter("v0", ColumnType.LONG, null, ImmutableList.of(1, 17), null);
+    testQuery(
+        "select countryName from "
+        + "(select countryName, length(countryName) as cname from wikipedia group by countryName) "
+        + "where SCALAR_IN_ARRAY(cname, ARRAY[17, 1])",
+        queryContext,
+        ImmutableList.of(
+            GroupByQuery.builder()
+                .setDataSource(new TableDataSource(CalciteTests.WIKIPEDIA))
+                .setInterval(querySegmentSpec(Intervals.ETERNITY))
+                .setVirtualColumns(expressionVirtualColumn("v0", "strlen(\"countryName\")", ColumnType.LONG))
+                .setDimFilter(filter)
+                .setGranularity(Granularities.ALL)
+                .setDimensions(new DefaultDimensionSpec("countryName", "d0", ColumnType.STRING))
+                .setLimitSpec(NoopLimitSpec.instance())
+                .setContext(QUERY_CONTEXT_DEFAULT)
+                .build()
+        ),
+        ImmutableList.of(
+            new Object[]{"Republic of Korea"}
+        )
+    );
+  }
+
   public static class SubqueryComponentSupplier extends SqlTestFramework.StandardComponentSupplier
   {