[multistage][test] adding tests for IN/NOT-IN operation

pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/TransformOperandFactory.java

@@ -46,12 +46,6 @@ private static TransformOperand getTransformOperand(RexExpression.FunctionCall f
     int numOperands = operands.size();
     String canonicalName = OperatorUtils.canonicalizeFunctionName(functionCall.getFunctionName());
     switch (canonicalName) {
-      case "IN":
-        Preconditions.checkState(numOperands >= 2, "IN takes >=2 arguments, got: %s", numOperands);
-        return new FilterOperand.In(operands, dataSchema, false);
-      case "NOT_IN":
-        Preconditions.checkState(numOperands >= 2, "NOT_IN takes >=2 arguments, got: %s", numOperands);
-        return new FilterOperand.In(operands, dataSchema, true);
       case "AND":
         Preconditions.checkState(numOperands >= 2, "AND takes >=2 arguments, got: %s", numOperands);
         return new FilterOperand.And(operands, dataSchema);
@@ -61,6 +55,12 @@ private static TransformOperand getTransformOperand(RexExpression.FunctionCall f
       case "NOT":
         Preconditions.checkState(numOperands == 1, "NOT takes one argument, got: %s", numOperands);
         return new FilterOperand.Not(operands.get(0), dataSchema);
+      case "IN":
+        Preconditions.checkState(numOperands >= 2, "IN takes >=2 arguments, got: %s", numOperands);
+        return new FilterOperand.In(operands, dataSchema, false);
+      case "NOT_IN":
+        Preconditions.checkState(numOperands >= 2, "NOT_IN takes >=2 arguments, got: %s", numOperands);
+        return new FilterOperand.In(operands, dataSchema, true);
       case "ISTRUE":
         Preconditions.checkState(numOperands == 1, "IS_TRUE takes one argument, got: %s", numOperands);
         return new FilterOperand.IsTrue(operands.get(0), dataSchema);

pinot-query-runtime/src/test/resources/queries/Case.json

@@ -98,6 +98,10 @@
       {
         "sql": "SELECT {tbl1}.primary_key, SUM(CASE WHEN {tbl2}.attribute = 'chocolate' THEN 1 ELSE 0 END) as chocolate_count FROM {tbl1} JOIN {tbl2} ON {tbl1}.primary_key = {tbl2}.primary_key GROUP BY {tbl1}.primary_key",
         "description": "Joins the two tables and aggregates the number of times 'chocolate' appears as an attribute in tbl2"
+      },
+      {
+        "sql": "SELECT primary_key, CASE WHEN description IN ('Item one', 'Item two') THEN attribute ELSE description END AS description, CASE WHEN description NOT IN ('Item three', 'Item four') THEN attribute ELSE description END AS attribute FROM ( select {tbl1}.primary_key, {tbl1}.description, {tbl2}.attribute FROM {tbl1} JOIN {tbl2} ON {tbl1}.primary_key = {tbl2}.primary_key) tmp WHERE attribute IN ('A','B','C','D') limit 10",
+        "description": "Joins the two tables and selects either the attribute using IN/NOT-IN clause"
       }
     ]
   }

pinot-query-runtime/src/test/resources/queries/FilterAggregates.json

@@ -75,6 +75,9 @@
       },
       {
         "sql": "SELECT bool_col, COALESCE(min(double_col) FILTER (WHERE string_col = 'a' OR string_col = 'b'), 0), COALESCE(max(double_col) FILTER (WHERE string_col = 'a' OR int_col > 10), 0), avg(double_col), sum(double_col), count(double_col), count(distinct(double_col)) FILTER (WHERE string_col = 'b' OR int_col > 10), count(string_col) FROM {tbl} WHERE string_col='b' GROUP BY bool_col"
+      },
+      {
+        "sql": "SELECT string_col, count(bool_col) FILTER ( WHERE double_col NOT IN (1, 3, 5, 7)) FROM {tbl} WHERE double_col < 10 AND int_col BETWEEN 1 AND 1 AND int_col <> 1 GROUP BY string_col"
       }
     ]
   },