[multistage] Support IN and NOT-IN Clauses

pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RexExpression.java

@@ -18,7 +18,6 @@
  */
 package org.apache.pinot.query.planner.logical;
 
-import com.google.common.base.Preconditions;
 import java.math.BigDecimal;
 import java.util.List;
 import java.util.stream.Collectors;
@@ -29,7 +28,6 @@
 import org.apache.calcite.rex.RexLiteral;
 import org.apache.calcite.rex.RexNode;
 import org.apache.calcite.sql.SqlKind;
-import org.apache.calcite.sql.type.SqlTypeName;
 import org.apache.calcite.util.NlsString;
 import org.apache.pinot.common.utils.PinotDataType;
 import org.apache.pinot.query.planner.serde.ProtoProperties;
@@ -52,30 +50,24 @@ static RexExpression toRexExpression(RexNode rexNode) {
     } else if (rexNode instanceof RexLiteral) {
       RexLiteral rexLiteral = ((RexLiteral) rexNode);
       FieldSpec.DataType dataType = toDataType(rexLiteral.getType());
-      return new RexExpression.Literal(dataType, rexLiteral.getTypeName(),
-          toRexValue(dataType, rexLiteral.getValue()));
+      return new RexExpression.Literal(dataType, toRexValue(dataType, rexLiteral.getValue()));
     } else if (rexNode instanceof RexCall) {
       RexCall rexCall = (RexCall) rexNode;
-      List<RexExpression> operands = rexCall.getOperands().stream().map(RexExpression::toRexExpression)
-          .collect(Collectors.toList());
-      return toRexExpression(rexCall, operands);
+      return toRexExpression(rexCall);
     } else {
       throw new IllegalArgumentException("Unsupported RexNode type with SqlKind: " + rexNode.getKind());
     }
   }
 
-  static RexExpression toRexExpression(RexCall rexCall, List<RexExpression> operands) {
+  static RexExpression toRexExpression(RexCall rexCall) {
     switch (rexCall.getKind()) {
       case CAST:
-        // CAST is being rewritten into "rexCall.CAST<targetType>(inputValue)",
-        //   - e.g. result type has already been converted into the CAST RexCall, so we assert single operand.
-        Preconditions.checkState(operands.size() == 1, "CAST takes exactly 2 arguments");
-        RelDataType castType = rexCall.getType();
-        // add the 2nd argument as the source type info.
-        operands.add(new Literal(FieldSpec.DataType.STRING, rexCall.getOperands().get(0).getType().getSqlTypeName(),
-            toPinotDataType(rexCall.getOperands().get(0).getType()).name()));
-        return new RexExpression.FunctionCall(rexCall.getKind(), toDataType(rexCall.getType()), "CAST", operands);
+        return RexExpressionUtils.handleCast(rexCall);
+      case SEARCH:
+        return RexExpressionUtils.handleSearch(rexCall);
       default:
+        List<RexExpression> operands =
+            rexCall.getOperands().stream().map(RexExpression::toRexExpression).collect(Collectors.toList());
         return new RexExpression.FunctionCall(rexCall.getKind(), toDataType(rexCall.getType()),
             rexCall.getOperator().getName(), operands);
     }
@@ -186,7 +178,7 @@ class Literal implements RexExpression {
     public Literal() {
     }
 
-    public Literal(FieldSpec.DataType dataType, SqlTypeName sqlTypeName, @Nullable Object value) {
+    public Literal(FieldSpec.DataType dataType, @Nullable Object value) {
       _sqlKind = SqlKind.LITERAL;
       _dataType = dataType;
       _value = value;

pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RexExpressionUtils.java

@@ -0,0 +1,84 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.query.planner.logical;
+
+import com.google.common.base.Preconditions;
+import com.google.common.collect.Range;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+import java.util.stream.Collectors;
+import org.apache.calcite.rel.type.RelDataType;
+import org.apache.calcite.rex.RexCall;
+import org.apache.calcite.rex.RexInputRef;
+import org.apache.calcite.rex.RexLiteral;
+import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.sql.SqlKind;
+import org.apache.calcite.util.Sarg;
+import org.apache.commons.lang3.NotImplementedException;
+import org.apache.pinot.spi.data.FieldSpec;
+
+
+public class RexExpressionUtils {
+
+  private RexExpressionUtils() {
+  }
+
+  static RexExpression handleCast(RexCall rexCall) {
+    // CAST is being rewritten into "rexCall.CAST<targetType>(inputValue)",
+    //   - e.g. result type has already been converted into the CAST RexCall, so we assert single operand.
+    List<RexExpression> operands =
+        rexCall.getOperands().stream().map(RexExpression::toRexExpression).collect(Collectors.toList());
+    Preconditions.checkState(operands.size() == 1, "CAST takes exactly 2 arguments");
+    RelDataType castType = rexCall.getType();
+    // add the 2nd argument as the source type info.
+    operands.add(new RexExpression.Literal(FieldSpec.DataType.STRING,
+        RexExpression.toPinotDataType(rexCall.getOperands().get(0).getType()).name()));
+    return new RexExpression.FunctionCall(rexCall.getKind(), RexExpression.toDataType(rexCall.getType()), "CAST",
+        operands);
+  }
+
+  // TODO: Add support for range filter expressions (e.g. a > 0 and a < 30)
+  static RexExpression handleSearch(RexCall rexCall) {
+    List<RexNode> operands = rexCall.getOperands();
+    RexInputRef rexInputRef = (RexInputRef) operands.get(0);
+    RexLiteral rexLiteral = (RexLiteral) operands.get(1);
+    FieldSpec.DataType dataType = RexExpression.toDataType(rexLiteral.getType());
+    Sarg sarg = rexLiteral.getValueAs(Sarg.class);
+    if (sarg.isPoints()) {
+      return new RexExpression.FunctionCall(SqlKind.IN, dataType, SqlKind.IN.name(), toFunctionOperands(rexInputRef,
+          sarg.rangeSet.asRanges(), dataType));
+    } else if (sarg.isComplementedPoints()) {
+      return new RexExpression.FunctionCall(SqlKind.NOT_IN, dataType, SqlKind.NOT_IN.name(),
+          toFunctionOperands(rexInputRef, sarg.rangeSet.complement().asRanges(), dataType));
+    } else {
+      throw new NotImplementedException("Range is not implemented yet");
+    }
+  }
+
+  private static List<RexExpression> toFunctionOperands(RexInputRef rexInputRef, Set<Range> ranges,
+      FieldSpec.DataType dataType) {
+    List<RexExpression> result = new ArrayList<>(ranges.size() + 1);
+    result.add(RexExpression.toRexExpression(rexInputRef));
+    for (Range range : ranges) {
+      result.add(new RexExpression.Literal(dataType, RexExpression.toRexValue(dataType, range.lowerEndpoint())));
+    }
+    return result;
+  }
+}

pinot-query-planner/src/test/java/org/apache/pinot/query/QueryCompilationTest.java

@@ -255,6 +255,10 @@ private Object[][] provideQueriesWithException() {
         new Object[]{"SELECT b.col1 - a.col3 FROM a JOIN c ON a.col1 = c.col3", "Table 'b' not found"},
         // non-agg column not being grouped
         new Object[]{"SELECT a.col1, SUM(a.col3) FROM a", "'a.col1' is not being grouped"},
+        // empty IN clause fails compilation
+        new Object[]{"SELECT a.col1 FROM a WHERE a.col1 IN ()", "Encountered \"\" at line"},
+        // range filter queries are not supported right now
+        new Object[]{"SELECT a.col1 FROM a WHERE a.col1 > 'x' AND a.col1 < 'y'", "Range is not implemented yet"}
     };
   }
 

pinot-query-planner/src/test/java/org/apache/pinot/query/QueryEnvironmentTestBase.java

@@ -67,6 +67,8 @@ protected Object[][] provideQueries() {
         new Object[]{"SELECT dateTrunc('DAY', a.ts + b.ts) FROM a JOIN b on a.col1 = b.col1 AND a.col2 = b.col2"},
         new Object[]{"SELECT a.col2, a.col3 FROM a JOIN b ON a.col1 = b.col1 "
             + " WHERE a.col3 >= 0 GROUP BY a.col2, a.col3"},
+        new Object[]{"SELECT a.col1, b.col2 FROM a JOIN b ON a.col1 = b.col1 WHERE a.col2 IN ('foo', 'bar') AND"
+            + " b.col2 NOT IN ('alice', 'charlie')"},
     };
   }
 }

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

@@ -25,7 +25,6 @@
 import java.util.List;
 import java.util.Map;
 import javax.annotation.Nullable;
-import org.apache.calcite.sql.type.SqlTypeName;
 import org.apache.pinot.common.request.context.ExpressionContext;
 import org.apache.pinot.common.utils.DataSchema;
 import org.apache.pinot.core.common.Operator;
@@ -98,8 +97,7 @@ public AggregateOperator(BaseOperator<TransferableBlock> inputOperator, DataSche
   private RexExpression toAggregationFunctionOperand(RexExpression rexExpression) {
     List<RexExpression> functionOperands = ((RexExpression.FunctionCall) rexExpression).getFunctionOperands();
     Preconditions.checkState(functionOperands.size() < 2);
-    return functionOperands.size() > 0 ? functionOperands.get(0)
-        : new RexExpression.Literal(FieldSpec.DataType.INT, SqlTypeName.INTEGER, 1);
+    return functionOperands.size() > 0 ? functionOperands.get(0) : new RexExpression.Literal(FieldSpec.DataType.INT, 1);
   }
 
   @Override

pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/QueryRunnerTest.java

@@ -109,6 +109,23 @@ private Object[][] provideTestSqlAndRowCount() {
         new Object[]{"SELECT a.col1, a.ts, b.col2, b.col3 FROM a JOIN b ON a.col1 = b.col2 "
             + " WHERE a.col3 >= 0 AND a.col2 = 'alice' AND b.col3 >= 0", 3},
 
+        // Join query with IN and Not-IN clause. Table A's side of join will return 9 rows and Table B's side will
+        // return 2 rows. Join will be only on col1=bar and since A will return 3 rows with that value and B will return
+        // 1 row, the final output will have 3 rows.
+        new Object[]{"SELECT a.col1, b.col2 FROM a JOIN b ON a.col1 = b.col1 "
+            + " WHERE a.col1 IN ('foo', 'bar', 'alice') AND b.col2 NOT IN ('foo', 'alice')", 3},
+
+        // Same query as above but written using OR/AND instead of IN.
+        new Object[]{"SELECT a.col1, b.col2 FROM a JOIN b ON a.col1 = b.col1 "
+            + " WHERE (a.col1 = 'foo' OR a.col1 = 'bar' OR a.col1 = 'alice') AND b.col2 != 'foo'"
+            + " AND b.col2 != 'alice'", 3},
+
+        // Same as above but with single argument IN clauses. Left side of the join returns 3 rows, and the right side
+        // returns 5 rows. Only key where join succeeds is col1=foo, and since table B has only 1 row with that value,
+        // the number of rows should be 3.
+        new Object[]{"SELECT a.col1, b.col2 FROM a JOIN b ON a.col1 = b.col1 "
+            + " WHERE a.col1 IN ('foo') AND b.col2 NOT IN ('')", 3},
+
         // Projection pushdown
         new Object[]{"SELECT a.col1, a.col3 + a.col3 FROM a WHERE a.col3 >= 0 AND a.col2 = 'alice'", 3},