Make 'valueIn' MV transform function work with the multi-stage query …

…engine
apache · Jun 20, 2024 · 4243b4f · 4243b4f
1 parent 99f6934
commit 4243b4f
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 1 deletion.
diff --git a/pinot-common/src/main/java/org/apache/pinot/common/function/TransformFunctionType.java b/pinot-common/src/main/java/org/apache/pinot/common/function/TransformFunctionType.java
@@ -30,6 +30,7 @@
 import org.apache.calcite.sql.SqlOperatorBinding;
 import org.apache.calcite.sql.type.OperandTypes;
 import org.apache.calcite.sql.type.ReturnTypes;
+import org.apache.calcite.sql.type.SqlOperandCountRanges;
 import org.apache.calcite.sql.type.SqlOperandTypeChecker;
 import org.apache.calcite.sql.type.SqlReturnTypeInference;
 import org.apache.calcite.sql.type.SqlTypeFamily;
@@ -196,7 +197,8 @@ public enum TransformFunctionType {
   ARRAY_SUM_INT("arraySumInt", ReturnTypes.INTEGER, OperandTypes.family(SqlTypeFamily.ARRAY), "array_sum_int"),
   ARRAY_SUM_LONG("arraySumLong", ReturnTypes.BIGINT, OperandTypes.family(SqlTypeFamily.ARRAY), "array_sum_long"),
 
-  VALUE_IN("valueIn", "value_in"),
+  VALUE_IN("valueIn", ReturnTypes.ARG0_FORCE_NULLABLE, OperandTypes.variadic(SqlOperandCountRanges.from(2)),
+      "value_in"),
   MAP_VALUE("mapValue", ReturnTypes.cascade(opBinding ->
       opBinding.getOperandType(2).getComponentType(), SqlTypeTransforms.FORCE_NULLABLE),
       OperandTypes.family(ImmutableList.of(SqlTypeFamily.ANY, SqlTypeFamily.ANY, SqlTypeFamily.ANY)),

diff --git a/...main/java/org/apache/pinot/core/operator/transform/function/ValueInTransformFunction.java b/...main/java/org/apache/pinot/core/operator/transform/function/ValueInTransformFunction.java
@@ -47,6 +47,11 @@
 import org.apache.pinot.spi.data.FieldSpec.DataType;
 
 
+/**
+ * This class implements the valueIn function for multi-valued columns. It takes at least 2 arguments, where the first
+ * argument is a multi-valued column, and the following arguments are constant values. The transform function will
+ * filter the values from the multi-valued column with the given constant values.
+ */
 public class ValueInTransformFunction extends BaseTransformFunction {
   public static final String FUNCTION_NAME = "valueIn";
 

diff --git a/...sts/src/test/java/org/apache/pinot/integration/tests/MultiStageEngineIntegrationTest.java b/...sts/src/test/java/org/apache/pinot/integration/tests/MultiStageEngineIntegrationTest.java
@@ -693,6 +693,17 @@ public void testMultiValueColumnGroupBy()
     Assert.assertEquals(jsonNode.get("resultTable").get("rows").size(), 154);
   }
 
+  @Test
+  public void testVariadicFunction() throws Exception {
+    String sqlQuery = "SELECT ARRAY_TO_MV(VALUE_IN(RandomAirports, 'MFR', 'SUN', 'GTR')) as airport, count(*) "
+        + "FROM mytable WHERE ARRAY_TO_MV(RandomAirports) IN ('MFR', 'SUN', 'GTR') GROUP BY airport";
+    JsonNode jsonNode = postQuery(sqlQuery);
+    assertNoError(jsonNode);
+    assertEquals(jsonNode.get("resultTable").get("dataSchema").get("columnDataTypes").get(0).asText(), "STRING");
+    assertEquals(jsonNode.get("resultTable").get("dataSchema").get("columnDataTypes").get(1).asText(), "LONG");
+    assertEquals(jsonNode.get("numRowsResultSet").asInt(), 3);
+  }
+
   @Test
   public void testMultiValueColumnGroupByOrderBy()
       throws Exception {