feat(isthmus): support exists and unique set predicates (#354)

Author: nielspardon

isthmus/src/main/java/io/substrait/isthmus/expression/ExpressionRexConverter.java

@@ -6,6 +6,7 @@
 import io.substrait.expression.Expression;
 import io.substrait.expression.Expression.FailureBehavior;
 import io.substrait.expression.Expression.ScalarSubquery;
+import io.substrait.expression.Expression.SetPredicate;
 import io.substrait.expression.Expression.SingleOrList;
 import io.substrait.expression.Expression.Switch;
 import io.substrait.expression.FieldReference;
@@ -545,4 +546,20 @@ public RexNode visit(ScalarSubquery expr) throws RuntimeException {
     RelNode inputRelnode = expr.input().accept(relNodeConverter);
     return RexSubQuery.scalar(inputRelnode);
   }
+
+  @Override
+  public RexNode visit(SetPredicate expr) throws RuntimeException {
+    RelNode inputRelnode = expr.tuples().accept(relNodeConverter);
+    switch (expr.predicateOp()) {
+      case PREDICATE_OP_EXISTS:
+        return RexSubQuery.exists(inputRelnode);
+      case PREDICATE_OP_UNIQUE:
+        return RexSubQuery.unique(inputRelnode);
+      case PREDICATE_OP_UNSPECIFIED:
+      default:
+        throw new UnsupportedOperationException(
+            String.format(
+                "Cannot handle SetPredicate when PredicateOp is %s.", expr.predicateOp().name()));
+    }
+  }
 }

isthmus/src/test/java/io/substrait/isthmus/SubstraitExpressionConverterTest.java

@@ -2,6 +2,7 @@
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertInstanceOf;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 
 import io.substrait.dsl.SubstraitBuilder;
 import io.substrait.expression.Expression;
@@ -53,16 +54,7 @@ public void switchExpression() {
 
   @Test
   public void scalarSubQuery() {
-    Rel subQueryRel =
-        b.project(
-            input -> List.of(b.fieldReference(input, 0)),
-            Remap.of(List.of(3)),
-            b.filter(
-                input ->
-                    b.equal(
-                        b.fieldReference(input, 2),
-                        Expression.StrLiteral.builder().nullable(false).value("EUROPE").build()),
-                commonTable));
+    Rel subQueryRel = createSubQueryRel();
 
     Expression.ScalarSubquery expr =
         Expression.ScalarSubquery.builder()
@@ -81,6 +73,92 @@ public void scalarSubQuery() {
     assertEquals(SqlKind.SCALAR_QUERY, calciteProjectExpr.get(0).getKind());
   }
 
+  @Test
+  public void existsSetPredicate() {
+    Rel subQueryRel = createSubQueryRel();
+
+    Expression.SetPredicate expr =
+        Expression.SetPredicate.builder()
+            .predicateOp(Expression.PredicateOp.PREDICATE_OP_EXISTS)
+            .tuples(subQueryRel)
+            .build();
+
+    Project query = b.project(input -> List.of(expr), b.emptyScan());
+
+    SubstraitToCalcite substraitToCalcite = new SubstraitToCalcite(extensions, typeFactory);
+    RelNode calciteRel = substraitToCalcite.convert(query);
+
+    assertInstanceOf(LogicalProject.class, calciteRel);
+    List<RexNode> calciteProjectExpr = ((LogicalProject) calciteRel).getProjects();
+    assertEquals(1, calciteProjectExpr.size());
+    assertEquals(SqlKind.EXISTS, calciteProjectExpr.get(0).getKind());
+  }
+
+  @Test
+  public void uniqueSetPredicate() {
+    Rel subQueryRel = createSubQueryRel();
+
+    Expression.SetPredicate expr =
+        Expression.SetPredicate.builder()
+            .predicateOp(Expression.PredicateOp.PREDICATE_OP_UNIQUE)
+            .tuples(subQueryRel)
+            .build();
+
+    Project query = b.project(input -> List.of(expr), b.emptyScan());
+
+    SubstraitToCalcite substraitToCalcite = new SubstraitToCalcite(extensions, typeFactory);
+    RelNode calciteRel = substraitToCalcite.convert(query);
+
+    assertInstanceOf(LogicalProject.class, calciteRel);
+    List<RexNode> calciteProjectExpr = ((LogicalProject) calciteRel).getProjects();
+    assertEquals(1, calciteProjectExpr.size());
+    assertEquals(SqlKind.UNIQUE, calciteProjectExpr.get(0).getKind());
+  }
+
+  @Test
+  public void unspecifiedSetPredicate() {
+    Rel subQueryRel = createSubQueryRel();
+
+    Expression.SetPredicate expr =
+        Expression.SetPredicate.builder()
+            .predicateOp(Expression.PredicateOp.PREDICATE_OP_UNSPECIFIED)
+            .tuples(subQueryRel)
+            .build();
+
+    Project query = b.project(input -> List.of(expr), b.emptyScan());
+
+    SubstraitToCalcite substraitToCalcite = new SubstraitToCalcite(extensions, typeFactory);
+    Exception exception =
+        assertThrows(
+            UnsupportedOperationException.class,
+            () -> {
+              substraitToCalcite.convert(query);
+            });
+
+    assertEquals(
+        "Cannot handle SetPredicate when PredicateOp is PREDICATE_OP_UNSPECIFIED.",
+        exception.getMessage());
+  }
+
+  /**
+   * Creates a Substrait {@link Rel} equivalent to the following SQL query:
+   *
+   * <p>select a from example where c = 'EUROPE'
+   *
+   * @return the Substrait {@link Rel} equivalent of the above SQL query
+   */
+  Rel createSubQueryRel() {
+    return b.project(
+        input -> List.of(b.fieldReference(input, 0)),
+        Remap.of(List.of(3)),
+        b.filter(
+            input ->
+                b.equal(
+                    b.fieldReference(input, 2),
+                    Expression.StrLiteral.builder().nullable(false).value("EUROPE").build()),
+            commonTable));
+  }
+
   void assertTypeMatch(RelDataType actual, Type expected) {
     Type type = TypeConverter.DEFAULT.toSubstrait(actual);
     assertEquals(expected, type);