[MOREL-60] Push elem, notElem and not ... elem down to Calcite …

…(as SQL `IN` and `NOT IN`)

Fixes #60

Test that `union`, `intersect`, `except` are pushed down to Calcite.

src/main/java/net/hydromatic/morel/Main.java

@@ -123,7 +123,8 @@ public void run() {
           out.write("\n");
         }
         final CompiledStatement compiled =
-            Compiles.prepareStatement(typeSystem, session, env, statement);
+            Compiles.prepareStatement(typeSystem, session, env, statement,
+                null);
         compiled.eval(session, env, lines, bindings);
         for (String line : lines) {
           out.write(line);

src/main/java/net/hydromatic/morel/Shell.java

@@ -244,7 +244,8 @@ public void run() {
           try {
             statement = smlParser.statementSemicolon();
             final CompiledStatement compiled =
-                Compiles.prepareStatement(typeSystem, session, env, statement);
+                Compiles.prepareStatement(typeSystem, session, env, statement,
+                    null);
             compiled.eval(session, env, lines, bindings);
             printAll(lines);
             terminal.writer().flush();

src/main/java/net/hydromatic/morel/compile/CalciteCompiler.java

@@ -57,6 +57,7 @@
 import org.apache.calcite.rel.type.RelDataType;
 import org.apache.calcite.rel.type.RelDataTypeFactory;
 import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.rex.RexSubQuery;
 import org.apache.calcite.sql.SqlAggFunction;
 import org.apache.calcite.sql.SqlOperator;
 import org.apache.calcite.sql.fun.SqlStdOperatorTable;
@@ -82,10 +83,18 @@
 
 import static net.hydromatic.morel.ast.CoreBuilder.core;
 
+import static java.util.Objects.requireNonNull;
+
 /** Compiles an expression to code that can be evaluated. */
 public class CalciteCompiler extends Compiler {
-  /** Morel operators and their exact equivalents in Calcite. */
-  static final Map<BuiltIn, SqlOperator> INFIX_OPERATORS =
+  /** Morel prefix and suffix operators and their exact equivalents in Calcite. */
+  static final Map<BuiltIn, SqlOperator> UNARY_OPERATORS =
+      ImmutableMap.<BuiltIn, SqlOperator>builder()
+          .put(BuiltIn.NOT, SqlStdOperatorTable.NOT)
+          .build();
+
+  /** Morel infix operators and their exact equivalents in Calcite. */
+  static final Map<BuiltIn, SqlOperator> BINARY_OPERATORS =
       ImmutableMap.<BuiltIn, SqlOperator>builder()
           .put(BuiltIn.OP_EQ, SqlStdOperatorTable.EQUALS)
           .put(BuiltIn.OP_NE, SqlStdOperatorTable.NOT_EQUALS)
@@ -94,6 +103,8 @@ public class CalciteCompiler extends Compiler {
           .put(BuiltIn.OP_GT, SqlStdOperatorTable.GREATER_THAN)
           .put(BuiltIn.OP_GE, SqlStdOperatorTable.GREATER_THAN_OR_EQUAL)
           .put(BuiltIn.OP_NEGATE, SqlStdOperatorTable.UNARY_MINUS)
+          .put(BuiltIn.OP_ELEM, SqlStdOperatorTable.IN)
+          .put(BuiltIn.OP_NOT_ELEM, SqlStdOperatorTable.NOT_IN)
           .put(BuiltIn.Z_NEGATE_INT, SqlStdOperatorTable.UNARY_MINUS)
           .put(BuiltIn.Z_NEGATE_REAL, SqlStdOperatorTable.UNARY_MINUS)
           .put(BuiltIn.OP_PLUS, SqlStdOperatorTable.PLUS)
@@ -118,7 +129,7 @@ public class CalciteCompiler extends Compiler {
 
   public CalciteCompiler(TypeSystem typeSystem, Calcite calcite) {
     super(typeSystem);
-    this.calcite = calcite;
+    this.calcite = requireNonNull(calcite, "calcite");
   }
 
   public @Nullable RelNode toRel(Environment env, Core.Exp expression) {
@@ -542,11 +553,34 @@ record = toRecord(cx, id);
       switch (apply.fn.op) {
       case FN_LITERAL:
         BuiltIn op = (BuiltIn) ((Core.Literal) apply.fn).value;
-        final SqlOperator operator = INFIX_OPERATORS.get(op);
-        if (operator != null) {
+
+        // Is it a unary operator with a Calcite equivalent? E.g. not => NOT
+        final SqlOperator unaryOp = UNARY_OPERATORS.get(op);
+        if (unaryOp != null) {
+          return cx.relBuilder.call(unaryOp, translate(cx, apply.arg));
+        }
+
+        // Is it a binary operator with a Calcite equivalent? E.g. + => PLUS
+        final SqlOperator binaryOp = BINARY_OPERATORS.get(op);
+        if (binaryOp != null) {
           assert apply.arg.op == Op.TUPLE;
-          return cx.relBuilder.call(operator,
-              translateList(cx, ((Core.Tuple) apply.arg).args));
+          final List<Core.Exp> args = ((Core.Tuple) apply.arg).args;
+          switch (op) {
+          case OP_ELEM:
+          case OP_NOT_ELEM:
+            final RelNode r = toRel2(cx, args.get(1));
+            if (r != null) {
+              final RexNode e = translate(cx, args.get(0));
+              final RexSubQuery in = RexSubQuery.in(r, ImmutableList.of(e));
+              switch (op) {
+              case OP_NOT_ELEM:
+                return cx.relBuilder.not(in);
+              default:
+                return in;
+              }
+            }
+          }
+          return cx.relBuilder.call(binaryOp, translateList(cx, args));
         }
       }
       if (apply.fn instanceof Core.RecordSelector

src/main/java/net/hydromatic/morel/compile/Compiler.java

@@ -60,14 +60,16 @@
 import static net.hydromatic.morel.ast.CoreBuilder.core;
 import static net.hydromatic.morel.util.Static.toImmutableList;
 
+import static java.util.Objects.requireNonNull;
+
 /** Compiles an expression to code that can be evaluated. */
 public class Compiler {
   protected static final EvalEnv EMPTY_ENV = Codes.emptyEnv();
 
   protected final TypeSystem typeSystem;
 
   public Compiler(TypeSystem typeSystem) {
-    this.typeSystem = typeSystem;
+    this.typeSystem = requireNonNull(typeSystem, "typeSystem");
   }
 
   CompiledStatement compileStatement(Environment env, Core.Decl decl) {

src/main/java/net/hydromatic/morel/compile/Compiles.java

@@ -36,6 +36,7 @@
 
 import java.util.List;
 import java.util.Map;
+import javax.annotation.Nullable;
 
 import static net.hydromatic.morel.ast.AstBuilder.ast;
 
@@ -57,22 +58,24 @@ public static TypeResolver.Resolved validateExpression(AstNode statement,
    * compiles it to code that can be evaluated by the interpreter.
    */
   public static CompiledStatement prepareStatement(TypeSystem typeSystem,
-      Session session, Environment env, AstNode statement) {
+      Session session, Environment env, AstNode statement,
+      @Nullable Calcite calcite) {
     Ast.Decl decl;
     if (statement instanceof Ast.Exp) {
       decl = toValDecl((Ast.Exp) statement);
     } else {
       decl = (Ast.Decl) statement;
     }
-    return prepareDecl(typeSystem, session, env, decl);
+    return prepareDecl(typeSystem, session, env, calcite, decl);
   }
 
   /**
    * Validates and compiles a declaration, and compiles it to
    * code that can be evaluated by the interpreter.
    */
   private static CompiledStatement prepareDecl(TypeSystem typeSystem,
-      Session session, Environment env, Ast.Decl decl) {
+      Session session, Environment env, @Nullable Calcite calcite,
+      Ast.Decl decl) {
     final TypeResolver.Resolved resolved =
         TypeResolver.deduceType(env, decl, typeSystem);
     final boolean hybrid = Prop.HYBRID.booleanValue(session.map);
@@ -102,7 +105,9 @@ private static CompiledStatement prepareDecl(TypeSystem typeSystem,
     }
     final Compiler compiler;
     if (hybrid) {
-      final Calcite calcite = Calcite.withDataSets(ImmutableMap.of());
+      if (calcite == null) {
+        calcite = Calcite.withDataSets(ImmutableMap.of());
+      }
       compiler = new CalciteCompiler(typeSystem, calcite);
     } else {
       compiler = new Compiler(typeSystem);

src/main/java/net/hydromatic/morel/foreign/Calcite.java

@@ -26,17 +26,22 @@
 import net.hydromatic.morel.type.Type;
 import net.hydromatic.morel.util.ThreadLocals;
 
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import org.apache.calcite.DataContext;
 import org.apache.calcite.adapter.java.JavaTypeFactory;
 import org.apache.calcite.interpreter.Interpreter;
 import org.apache.calcite.jdbc.CalciteSchema;
 import org.apache.calcite.linq4j.Enumerable;
 import org.apache.calcite.linq4j.QueryProvider;
+import org.apache.calcite.plan.RelOptPlanner;
 import org.apache.calcite.plan.RelOptUtil;
+import org.apache.calcite.plan.RelTraitSet;
 import org.apache.calcite.rel.RelNode;
 import org.apache.calcite.schema.SchemaPlus;
 import org.apache.calcite.tools.Frameworks;
+import org.apache.calcite.tools.Program;
+import org.apache.calcite.tools.Programs;
 import org.apache.calcite.tools.RelBuilder;
 
 import java.util.List;
@@ -77,7 +82,15 @@ public RelBuilder relBuilder() {
   /** Creates a {@code Code} that evaluates a Calcite relational expression,
    * converting it to Morel list type {@code type}. */
   public Code code(Environment env, RelNode rel, Type type) {
-    return new CalciteCode(dataContext, rel, type, env);
+    // Transform the relational expression, converting sub-queries. For example,
+    // RexSubQuery.IN becomes a Join.
+    final Program program = Programs.SUB_QUERY_PROGRAM;
+    final RelOptPlanner planner = rel.getCluster().getPlanner();
+    final RelTraitSet traitSet = rel.getCluster().traitSet();
+    final RelNode rel2 = program.run(planner, rel, traitSet,
+        ImmutableList.of(), ImmutableList.of());
+
+    return new CalciteCode(dataContext, rel2, type, env);
   }
 
   /** Extension to Calcite context that remembers the foreign value

src/test/java/net/hydromatic/morel/AlgebraTest.java

@@ -22,9 +22,13 @@
 
 import org.junit.jupiter.api.Test;
 
+import java.util.function.UnaryOperator;
 import java.util.stream.Stream;
 
 import static net.hydromatic.morel.Matchers.equalsOrdered;
+import static net.hydromatic.morel.Matchers.equalsUnordered;
+import static net.hydromatic.morel.Matchers.isCode;
+import static net.hydromatic.morel.Matchers.isFullyCalcite;
 import static net.hydromatic.morel.Matchers.list;
 import static net.hydromatic.morel.Ml.ml;
 
@@ -265,7 +269,7 @@ public class AlgebraTest {
         .with(Prop.HYBRID, true)
         .assertType("{d5:int, deptno:int, empno:int} list")
         .assertEvalIter(equalsOrdered(list(25, 20, 7369)))
-        .assertPlan(is(plan));
+        .assertPlan(isCode(plan));
   }
 
   /** Tests a query that can be fully executed in Calcite. */
@@ -315,7 +319,7 @@ private void checkFullCalcite(String ml) {
         .with(Prop.HYBRID, true)
         .assertType("{d5:int, deptno:int, empno:int} list")
         .assertEvalIter(equalsOrdered(list(25, 20, 7369), list(35, 30, 7499)))
-        .assertPlan(is(plan));
+        .assertPlan(isCode(plan));
   }
 
   /** Tests a query that is "from" over no variables. The result has one row
@@ -326,7 +330,7 @@ private void checkFullCalcite(String ml) {
     ml(ml)
         .with(Prop.HYBRID, true)
         .assertType("unit list")
-        .assertPlan(is(plan))
+        .assertPlan(isCode(plan))
         .assertEvalIter(equalsOrdered(list()));
   }
 
@@ -389,7 +393,7 @@ private void checkCalciteWithVariable(int inlinePassCount, String plan) {
         .with(Prop.HYBRID, true)
         .with(Prop.INLINE_PASS_COUNT, inlinePassCount)
         .assertType("{d5:int, deptno:int, empno:int} list")
-        .assertPlan(is(plan))
+        .assertPlan(isCode(plan))
         .assertEvalIter(equalsOrdered(list(25, 20, 7369), list(35, 30, 7499)));
   }
 
@@ -419,7 +423,7 @@ private void checkCalciteWithVariable(int inlinePassCount, String plan) {
         .with(Prop.HYBRID, true)
         .with(Prop.INLINE_PASS_COUNT, 0)
         .assertType("int list")
-        .assertPlan(is(plan))
+        .assertPlan(isCode(plan))
         .assertEvalIter(equalsOrdered(20, 40, 60, 80));
   }
 
@@ -457,9 +461,92 @@ private void checkCalciteWithVariable(int inlinePassCount, String plan) {
         .with(Prop.HYBRID, true)
         .with(Prop.INLINE_PASS_COUNT, 0)
         .assertType("int list")
-        .assertPlan(is(plan))
+        .assertPlan(isCode(plan))
         .assertEvalIter(equalsOrdered(15, 25, 35, 45));
   }
+
+  /** Tests that we can send {@code union} to Calcite. */
+  @Test void testUnion() {
+    final String ml = "from x in (\n"
+        + "(from e in scott.emp where e.job = \"CLERK\" yield e.deptno)\n"
+        + "union\n"
+        + "(from d in scott.dept yield d.deptno))\n";
+    ml(ml)
+        .withBinding("scott", BuiltInDataSet.SCOTT)
+        .with(Prop.HYBRID, true)
+        .assertType("int list")
+        .assertPlan(isFullyCalcite())
+        .assertEvalIter(equalsUnordered(20, 20, 20, 40, 10, 10, 30, 30));
+  }
+
+  /** Tests that we can send {@code except} to Calcite. */
+  @Test void testExcept() {
+    final String ml = "from x in (\n"
+        + "(from d in scott.dept yield d.deptno)"
+        + "except\n"
+        + "(from e in scott.emp where e.job = \"CLERK\" yield e.deptno))\n";
+    ml(ml)
+        .withBinding("scott", BuiltInDataSet.SCOTT)
+        .with(Prop.HYBRID, true)
+        .assertType("int list")
+        .assertPlan(isFullyCalcite())
+        .assertEvalIter(equalsUnordered(40));
+  }
+
+  /** Tests that we can send {@code intersect} to Calcite. */
+  @Test void testIntersect() {
+    final String ml = "from x in (\n"
+        + "(from e in scott.emp where e.job = \"CLERK\" yield e.deptno)\n"
+        + "intersect\n"
+        + "(from d in scott.dept yield d.deptno))\n";
+    ml(ml)
+        .withBinding("scott", BuiltInDataSet.SCOTT)
+        .with(Prop.HYBRID, true)
+        .assertType("int list")
+        .assertPlan(isFullyCalcite())
+        .assertEvalIter(equalsUnordered(10, 20, 30));
+  }
+
+  /** Tests that we can send (what in SQL would be) an uncorrelated {@code IN}
+   * sub-query to Calcite. */
+  @Test void testElem() {
+    final String ml = "from d in scott.dept\n"
+        + "where d.deptno elem (from e in scott.emp\n"
+        + "    where e.job elem [\"ANALYST\", \"PRESIDENT\"]\n"
+        + "    yield e.deptno)\n"
+        + "yield d.dname";
+    ml(ml)
+        .withBinding("scott", BuiltInDataSet.SCOTT)
+        .with(Prop.HYBRID, true)
+        .assertType("string list")
+        .assertPlan(isFullyCalcite())
+        .assertEvalIter(equalsUnordered("ACCOUNTING", "RESEARCH"));
+  }
+
+  /** Tests that we can send (what in SQL would be) an uncorrelated {@code IN}
+   * sub-query to Calcite. */
+  @Test void testNotElem() {
+    final UnaryOperator<Ml> fn = ml ->
+        ml.withBinding("scott", BuiltInDataSet.SCOTT)
+            .with(Prop.HYBRID, true)
+            .assertType("string list")
+            .assertPlan(isFullyCalcite())
+            .assertEvalIter(equalsUnordered("SALES", "OPERATIONS"));
+
+    final String ml0 = "from d in scott.dept\n"
+        + "where not (d.deptno elem\n"
+        + "    (from e in scott.emp\n"
+        + "        where e.job elem [\"ANALYST\", \"PRESIDENT\"]\n"
+        + "        yield e.deptno))\n"
+        + "yield d.dname";
+    final String ml1 = "from d in scott.dept\n"
+        + "where d.deptno notElem (from e in scott.emp\n"
+        + "    where e.job elem [\"ANALYST\", \"PRESIDENT\"]\n"
+        + "    yield e.deptno)\n"
+        + "yield d.dname";
+    fn.apply(ml(ml0));
+    fn.apply(ml(ml1));
+  }
 }
 
 // End AlgebraTest.java

src/test/java/net/hydromatic/morel/InlineTest.java

@@ -22,6 +22,7 @@
 
 import org.junit.jupiter.api.Test;
 
+import static net.hydromatic.morel.Matchers.isCode;
 import static net.hydromatic.morel.Matchers.isUnordered;
 import static net.hydromatic.morel.Matchers.list;
 import static net.hydromatic.morel.Matchers.whenAppliedTo;
@@ -63,7 +64,7 @@ public class InlineTest {
     final String plan = "match(x, apply(fnValue +, argCode "
         + "tuple(apply(fnValue +, argCode tuple(get(name x), constant(1))), "
         + "constant(2))))";
-    ml(ml).assertPlan(is(plan));
+    ml(ml).assertPlan(isCode(plan));
   }
 
   @Test void testInlineFn() {
@@ -75,7 +76,7 @@ public class InlineTest {
         + "  end";
     final String plan =
         "match(x, apply(fnValue +, argCode tuple(get(name x), constant(1))))";
-    ml(ml).assertPlan(is(plan))
+    ml(ml).assertPlan(isCode(plan))
         .assertEval(whenAppliedTo(2, is(3)));
   }