Implement IN and NOT_IN filters in FilterOperand class (apache#12285 …

…fix) (apache#12305)
walterddr · Feb 1, 2024 · 04cf96e · 04cf96e
1 parent a8b1fa3
commit 04cf96e
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 0 deletions.
diff --git a/...runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/FilterOperand.java b/...runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/FilterOperand.java
@@ -110,6 +110,38 @@ public Integer apply(Object[] row) {
     }
   }
 
+  public static class In extends FilterOperand {
+    List<TransformOperand> _childOperands;
+    boolean _isNotIn;
+
+    public In(List<RexExpression> children, DataSchema dataSchema, boolean isNotIn) {
+      _childOperands = new ArrayList<>(children.size());
+      for (RexExpression child : children) {
+        _childOperands.add(TransformOperandFactory.getTransformOperand(child, dataSchema));
+      }
+      _isNotIn = isNotIn;
+    }
+
+    @Nullable
+    @Override
+    public Integer apply(Object[] row) {
+      Object firstResult = _childOperands.get(0).apply(row);
+      if (firstResult == null) {
+        return null;
+      }
+      for (int i = 1; i < _childOperands.size(); i++) {
+        Object result = _childOperands.get(i).apply(row);
+        if (result == null) {
+          return null;
+        }
+        if (firstResult.equals(result)) {
+          return _isNotIn ? 0 : 1;
+        }
+      }
+      return _isNotIn ? 1 : 0;
+    }
+  }
+
   public static class IsTrue extends FilterOperand {
     TransformOperand _childOperand;
 

diff --git a/...c/main/java/org/apache/pinot/query/runtime/operator/operands/TransformOperandFactory.java b/...c/main/java/org/apache/pinot/query/runtime/operator/operands/TransformOperandFactory.java
@@ -46,6 +46,12 @@ 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);