gianm commented on code in PR #16388:
URL: https://github.com/apache/druid/pull/16388#discussion_r1597077523
##########
sql/src/test/java/org/apache/druid/sql/calcite/CalciteQueryTest.java:
##########
@@ -5634,6 +5719,245 @@ public void testInIsNotTrueAndLessThanFilter()
);
}
+ @Test
+ public void testInExpression()
+ {
+ // Cannot vectorize scalar_in_array expression.
+ cannotVectorize();
+
+ testQuery(
+ "SELECT dim1 IN ('abc', 'def', 'ghi'), COUNT(*)\n"
+ + "FROM druid.foo\n"
+ + "GROUP BY 1",
+ ImmutableList.of(
+ GroupByQuery.builder()
+ .setDataSource(CalciteTests.DATASOURCE1)
+ .setInterval(querySegmentSpec(Filtration.eternity()))
+ .setGranularity(Granularities.ALL)
+ .setVirtualColumns(
+ expressionVirtualColumn(
+ "v0",
+
"scalar_in_array(\"dim1\",array('abc','def','ghi'))",
+ ColumnType.LONG
+ )
+ )
+ .setDimensions(dimensions(new
DefaultDimensionSpec("v0", "d0", ColumnType.LONG)))
+ .setAggregatorSpecs(new CountAggregatorFactory("a0"))
+ .setContext(QUERY_CONTEXT_DEFAULT)
+ .build()
+ ),
+ ImmutableList.of(
+ new Object[]{false, 4L},
+ new Object[]{true, 2L}
+ )
+ );
+ }
+
+ @Test
+ public void testInExpressionBelowThreshold()
+ {
+ // Cannot vectorize scalar_in_array expression.
+ cannotVectorize();
+
+ testQuery(
+ "SELECT dim1 IN ('abc', 'def', 'ghi'), COUNT(*)\n"
+ + "FROM druid.foo\n"
+ + "GROUP BY 1",
+ QueryContexts.override(QUERY_CONTEXT_DEFAULT,
QueryContexts.IN_FUNCTION_EXPR_THRESHOLD, 100),
+ ImmutableList.of(
+ GroupByQuery.builder()
+ .setDataSource(CalciteTests.DATASOURCE1)
+ .setInterval(querySegmentSpec(Filtration.eternity()))
+ .setGranularity(Granularities.ALL)
+ .setVirtualColumns(
+ expressionVirtualColumn(
+ "v0",
+ "((\"dim1\" == 'abc') || (\"dim1\" == 'def')
|| (\"dim1\" == 'ghi'))",
+ ColumnType.LONG
+ )
+ )
+ .setDimensions(dimensions(new
DefaultDimensionSpec("v0", "d0", ColumnType.LONG)))
+ .setAggregatorSpecs(new CountAggregatorFactory("a0"))
+ .setContext(QUERY_CONTEXT_DEFAULT)
+ .build()
+ ),
+ ImmutableList.of(
+ new Object[]{false, 4L},
+ new Object[]{true, 2L}
+ )
+ );
+ }
+
+ @Test
+ public void testInOrIsNullExpression()
+ {
+ // Cannot vectorize scalar_in_array expression.
+ cannotVectorize();
+
+ testQuery(
+ "SELECT dim1 IN ('abc', 'def', 'ghi') OR dim1 IS NULL, COUNT(*)\n"
+ + "FROM druid.foo\n"
+ + "GROUP BY 1",
+ ImmutableList.of(
+ GroupByQuery.builder()
+ .setDataSource(CalciteTests.DATASOURCE1)
+ .setInterval(querySegmentSpec(Filtration.eternity()))
+ .setGranularity(Granularities.ALL)
+ .setVirtualColumns(
+ expressionVirtualColumn(
+ "v0",
+ "(isnull(\"dim1\") ||
scalar_in_array(\"dim1\",array('abc','def','ghi')))",
+ ColumnType.LONG
+ )
+ )
+ .setDimensions(dimensions(new
DefaultDimensionSpec("v0", "d0", ColumnType.LONG)))
+ .setAggregatorSpecs(new CountAggregatorFactory("a0"))
+ .setContext(QUERY_CONTEXT_DEFAULT)
+ .build()
+ ),
+ ImmutableList.of(
+ new Object[]{false, NullHandling.sqlCompatible() ? 4L : 3L},
+ new Object[]{true, NullHandling.sqlCompatible() ? 2L : 3L}
+ )
+ );
+ }
+
+ @Test
+ public void testNotInOrIsNullExpression()
+ {
+ // Cannot vectorize scalar_in_array expression.
+ cannotVectorize();
+
+ testQuery(
+ "SELECT NOT (dim1 IN ('abc', 'def', 'ghi') OR dim1 IS NULL),
COUNT(*)\n"
+ + "FROM druid.foo\n"
+ + "GROUP BY 1",
+ ImmutableList.of(
+ GroupByQuery.builder()
+ .setDataSource(CalciteTests.DATASOURCE1)
+ .setInterval(querySegmentSpec(Filtration.eternity()))
+ .setGranularity(Granularities.ALL)
+ .setVirtualColumns(
+ expressionVirtualColumn(
+ "v0",
+ "(notnull(\"dim1\") && (!
scalar_in_array(\"dim1\",array('abc','def','ghi'))))",
+ ColumnType.LONG
+ )
+ )
+ .setDimensions(dimensions(new
DefaultDimensionSpec("v0", "d0", ColumnType.LONG)))
+ .setAggregatorSpecs(new CountAggregatorFactory("a0"))
+ .setContext(QUERY_CONTEXT_DEFAULT)
+ .build()
+ ),
+ ImmutableList.of(
+ new Object[]{false, NullHandling.sqlCompatible() ? 2L : 3L},
+ new Object[]{true, NullHandling.sqlCompatible() ? 4L : 3L}
+ )
+ );
+ }
+
+ @Test
+ public void testNotInAndIsNotNullExpression()
+ {
+ // Cannot vectorize scalar_in_array expression.
+ cannotVectorize();
+
+ testQuery(
+ "SELECT dim1 NOT IN ('abc', 'def', 'ghi') AND dim1 IS NOT NULL,
COUNT(*)\n"
+ + "FROM druid.foo\n"
+ + "GROUP BY 1",
+ ImmutableList.of(
+ GroupByQuery.builder()
+ .setDataSource(CalciteTests.DATASOURCE1)
+ .setInterval(querySegmentSpec(Filtration.eternity()))
+ .setGranularity(Granularities.ALL)
+ .setVirtualColumns(
+ expressionVirtualColumn(
+ "v0",
+ "(notnull(\"dim1\") && (!
scalar_in_array(\"dim1\",array('abc','def','ghi'))))",
+ ColumnType.LONG
+ )
+ )
+ .setDimensions(dimensions(new
DefaultDimensionSpec("v0", "d0", ColumnType.LONG)))
+ .setAggregatorSpecs(new CountAggregatorFactory("a0"))
+ .setContext(QUERY_CONTEXT_DEFAULT)
+ .build()
+ ),
+ ImmutableList.of(
+ new Object[]{false, NullHandling.sqlCompatible() ? 2L : 3L},
+ new Object[]{true, NullHandling.sqlCompatible() ? 4L : 3L}
+ )
+ );
+ }
+
+ @Test
+ public void testInOrGreaterThanExpression()
+ {
+ // Cannot vectorize scalar_in_array expression.
+ cannotVectorize();
+
+ testQuery(
+ "SELECT dim1 IN ('abc', 'def', 'ghi') OR dim1 > 'zzz', COUNT(*)\n"
+ + "FROM druid.foo\n"
+ + "GROUP BY 1",
+ ImmutableList.of(
+ GroupByQuery.builder()
+ .setDataSource(CalciteTests.DATASOURCE1)
+ .setInterval(querySegmentSpec(Filtration.eternity()))
+ .setGranularity(Granularities.ALL)
+ .setVirtualColumns(
+ expressionVirtualColumn(
+ "v0",
+
"(scalar_in_array(\"dim1\",array('abc','def','ghi')) || (\"dim1\" > 'zzz'))",
+ ColumnType.LONG
+ )
+ )
+ .setDimensions(dimensions(new
DefaultDimensionSpec("v0", "d0", ColumnType.LONG)))
+ .setAggregatorSpecs(new CountAggregatorFactory("a0"))
+ .setContext(QUERY_CONTEXT_DEFAULT)
+ .build()
+ ),
+ ImmutableList.of(
+ new Object[]{false, 4L},
+ new Object[]{true, 2L}
+ )
+ );
+ }
+
+ @Test
+ public void testNotInAndLessThanExpression()
+ {
+ // Cannot vectorize scalar_in_array expression.
+ cannotVectorize();
+
+ testQuery(
+ "SELECT dim1 NOT IN ('abc', 'def', 'ghi') AND dim1 < 'zzz', COUNT(*)\n"
Review Comment:
Cool idea. I added a test for this as well:
`testNotInOrEqualToOneOfThemExpression`.
--
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: commits-unsubscr...@druid.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscr...@druid.apache.org
For additional commands, e-mail: commits-h...@druid.apache.org
