Adds compilation/evaluation support for UNION/INTERSECT/EXCEPT ALL/DISTINCT

partiql-eval/src/main/kotlin/org/partiql/eval/internal/Compiler.kt

@@ -4,8 +4,12 @@ import org.partiql.eval.PartiQLEngine
 import org.partiql.eval.internal.operator.Operator
 import org.partiql.eval.internal.operator.rel.RelAggregate
 import org.partiql.eval.internal.operator.rel.RelDistinct
+import org.partiql.eval.internal.operator.rel.RelExceptAll
+import org.partiql.eval.internal.operator.rel.RelExceptDistinct
 import org.partiql.eval.internal.operator.rel.RelExclude
 import org.partiql.eval.internal.operator.rel.RelFilter
+import org.partiql.eval.internal.operator.rel.RelIntersectAll
+import org.partiql.eval.internal.operator.rel.RelIntersectDistinct
 import org.partiql.eval.internal.operator.rel.RelJoinInner
 import org.partiql.eval.internal.operator.rel.RelJoinLeft
 import org.partiql.eval.internal.operator.rel.RelJoinOuterFull
@@ -18,6 +22,8 @@ import org.partiql.eval.internal.operator.rel.RelScanIndexed
 import org.partiql.eval.internal.operator.rel.RelScanIndexedPermissive
 import org.partiql.eval.internal.operator.rel.RelScanPermissive
 import org.partiql.eval.internal.operator.rel.RelSort
+import org.partiql.eval.internal.operator.rel.RelUnionAll
+import org.partiql.eval.internal.operator.rel.RelUnionDistinct
 import org.partiql.eval.internal.operator.rel.RelUnpivot
 import org.partiql.eval.internal.operator.rex.ExprCallDynamic
 import org.partiql.eval.internal.operator.rex.ExprCallStatic
@@ -308,6 +314,33 @@ internal class Compiler(
         }
     }
 
+    override fun visitRelOpSetExcept(node: Rel.Op.Set.Except, ctx: StaticType?): Operator {
+        val lhs = visitRel(node.lhs, ctx)
+        val rhs = visitRel(node.rhs, ctx)
+        return when (node.quantifier) {
+            Rel.Op.Set.Quantifier.ALL -> RelExceptAll(lhs, rhs)
+            Rel.Op.Set.Quantifier.DISTINCT -> RelExceptDistinct(lhs, rhs)
+        }
+    }
+
+    override fun visitRelOpSetIntersect(node: Rel.Op.Set.Intersect, ctx: StaticType?): Operator {
+        val lhs = visitRel(node.lhs, ctx)
+        val rhs = visitRel(node.rhs, ctx)
+        return when (node.quantifier) {
+            Rel.Op.Set.Quantifier.ALL -> RelIntersectAll(lhs, rhs)
+            Rel.Op.Set.Quantifier.DISTINCT -> RelIntersectDistinct(lhs, rhs)
+        }
+    }
+
+    override fun visitRelOpSetUnion(node: Rel.Op.Set.Union, ctx: StaticType?): Operator {
+        val lhs = visitRel(node.lhs, ctx)
+        val rhs = visitRel(node.rhs, ctx)
+        return when (node.quantifier) {
+            Rel.Op.Set.Quantifier.ALL -> RelUnionAll(lhs, rhs)
+            Rel.Op.Set.Quantifier.DISTINCT -> RelUnionDistinct(lhs, rhs)
+        }
+    }
+
     override fun visitRelOpLimit(node: Rel.Op.Limit, ctx: StaticType?): Operator {
         val input = visitRel(node.input, ctx)
         val limit = visitRex(node.limit, ctx)

partiql-eval/src/main/kotlin/org/partiql/eval/internal/helpers/IteratorChain.kt

@@ -0,0 +1,27 @@
+package org.partiql.eval.internal.helpers
+
+internal class IteratorChain<T>(
+    iterators: Array<Iterator<T>>
+) : IteratorPeeking<T>() {
+
+    private var iterator: Iterator<Iterator<T>> = when (iterators.isEmpty()) {
+        true -> listOf(emptyList<T>().iterator()).iterator()
+        false -> iterators.iterator()
+    }
+    private var current: Iterator<T> = iterator.next()
+
+    override fun peek(): T? {
+        return when (current.hasNext()) {
+            true -> current.next()
+            false -> {
+                while (iterator.hasNext()) {
+                    current = iterator.next()
+                    if (current.hasNext()) {
+                        return current.next()
+                    }
+                }
+                return null
+            }
+        }
+    }
+}

partiql-eval/src/main/kotlin/org/partiql/eval/internal/helpers/IteratorPeeking.kt

@@ -0,0 +1,36 @@
+package org.partiql.eval.internal.helpers
+
+/**
+ * For [Iterator]s that MUST materialize data in order to execute [hasNext], this abstract class caches the
+ * result of [peek] to implement both [hasNext] and [next].
+ *
+ * With this implementation, invoking hasNext() multiple times will not iterate unnecessarily. Invoking next() without
+ * invoking hasNext() is allowed -- however, it is highly recommended to avoid doing so.
+ */
+internal abstract class IteratorPeeking<T> : Iterator<T> {
+
+    internal var next: T? = null
+
+    /**
+     * @return NULL when there is not another [T] to be produced. Returns a [T] when able to.
+     *
+     * @see IteratorPeeking
+     */
+    abstract fun peek(): T?
+
+    override fun hasNext(): Boolean {
+        if (next != null) {
+            return true
+        }
+        this.next = peek()
+        return this.next != null
+    }
+
+    override fun next(): T {
+        val next = next
+            ?: peek()
+            ?: error("There was no more elements, however, next() was called. Please use hasNext() beforehand.")
+        this.next = null
+        return next
+    }
+}

partiql-eval/src/main/kotlin/org/partiql/eval/internal/operator/rel/RelExceptAll.kt

@@ -0,0 +1,55 @@
+package org.partiql.eval.internal.operator.rel
+
+import org.partiql.eval.internal.Environment
+import org.partiql.eval.internal.Record
+import org.partiql.eval.internal.operator.Operator
+
+internal class RelExceptAll(
+    private val lhs: Operator.Relation,
+    private val rhs: Operator.Relation,
+) : RelPeeking() {
+
+    private val seen: MutableMap<Record, Int> = mutableMapOf()
+    private var init: Boolean = false
+
+    override fun open(env: Environment) {
+        lhs.open(env)
+        rhs.open(env)
+        init = false
+        seen.clear()
+        super.open(env)
+    }
+
+    override fun peek(): Record? {
+        if (!init) {
+            seed()
+        }
+        for (row in lhs) {
+            val remaining = seen[row] ?: 0
+            if (remaining > 0) {
+                seen[row] = remaining - 1
+                continue
+            }
+            return row
+        }
+        return null
+    }
+
+    override fun close() {
+        lhs.close()
+        rhs.close()
+        seen.clear()
+        super.close()
+    }
+
+    /**
+     * Read the entire right-hand-side into our search structure.
+     */
+    private fun seed() {
+        init = true
+        for (row in rhs) {
+            val n = seen[row] ?: 0
+            seen[row] = n + 1
+        }
+    }
+}

partiql-eval/src/main/kotlin/org/partiql/eval/internal/operator/rel/RelExcept.kt → partiql-eval/src/main/kotlin/org/partiql/eval/internal/operator/rel/RelExceptDistinct.kt

@@ -10,7 +10,7 @@ import org.partiql.eval.internal.operator.Operator
  * @property lhs
  * @property rhs
  */
-internal class RelExcept(
+internal class RelExceptDistinct(
     private val lhs: Operator.Relation,
     private val rhs: Operator.Relation,
 ) : RelPeeking() {

partiql-eval/src/main/kotlin/org/partiql/eval/internal/operator/rel/RelIntersectAll.kt

@@ -0,0 +1,56 @@
+package org.partiql.eval.internal.operator.rel
+
+import org.partiql.eval.internal.Environment
+import org.partiql.eval.internal.Record
+import org.partiql.eval.internal.operator.Operator
+
+internal class RelIntersectAll(
+    private val lhs: Operator.Relation,
+    private val rhs: Operator.Relation,
+) : RelPeeking() {
+
+    private val seen: MutableMap<Record, Int> = mutableMapOf()
+    private var init: Boolean = false
+
+    override fun open(env: Environment) {
+        lhs.open(env)
+        rhs.open(env)
+        init = false
+        seen.clear()
+        super.open(env)
+    }
+
+    override fun peek(): Record? {
+        if (!init) {
+            seed()
+        }
+        for (row in rhs) {
+            seen.computeIfPresent(row) { _, y ->
+                when (y) {
+                    0 -> null
+                    else -> y - 1
+                }
+            }?.let { return row }
+        }
+        return null
+    }
+
+    override fun close() {
+        lhs.close()
+        rhs.close()
+        seen.clear()
+        super.close()
+    }
+
+    /**
+     * Read the entire left-hand-side into our search structure.
+     */
+    private fun seed() {
+        init = true
+        for (row in lhs) {
+            seen.computeIfPresent(row) { _, y ->
+                y + 1
+            } ?: seen.put(row, 1)
+        }
+    }
+}

partiql-eval/src/main/kotlin/org/partiql/eval/internal/operator/rel/RelIntersect.kt → partiql-eval/src/main/kotlin/org/partiql/eval/internal/operator/rel/RelIntersectDistinct.kt

@@ -4,19 +4,19 @@ import org.partiql.eval.internal.Environment
 import org.partiql.eval.internal.Record
 import org.partiql.eval.internal.operator.Operator
 
-internal class RelIntersect(
+internal class RelIntersectDistinct(
     private val lhs: Operator.Relation,
     private val rhs: Operator.Relation,
 ) : RelPeeking() {
 
-    private var seen: MutableSet<Record> = mutableSetOf()
+    private val seen: MutableSet<Record> = mutableSetOf()
     private var init: Boolean = false
 
     override fun open(env: Environment) {
         lhs.open(env)
         rhs.open(env)
         init = false
-        seen = mutableSetOf()
+        seen.clear()
         super.open(env)
     }
 
@@ -25,7 +25,7 @@ internal class RelIntersect(
             seed()
         }
         for (row in rhs) {
-            if (seen.contains(row)) {
+            if (seen.remove(row)) {
                 return row
             }
         }
@@ -44,8 +44,7 @@ internal class RelIntersect(
      */
     private fun seed() {
         init = true
-        while (true) {
-            val row = lhs.next() ?: break
+        for (row in lhs) {
             seen.add(row)
         }
     }

partiql-eval/src/main/kotlin/org/partiql/eval/internal/operator/rel/RelPeeking.kt

@@ -2,44 +2,20 @@ package org.partiql.eval.internal.operator.rel
 
 import org.partiql.eval.internal.Environment
 import org.partiql.eval.internal.Record
+import org.partiql.eval.internal.helpers.IteratorPeeking
 import org.partiql.eval.internal.operator.Operator
 
 /**
  * For [Operator.Relation]'s that MUST materialize data in order to execute [hasNext], this abstract class caches the
  * result of [peek] to implement both [hasNext] and [next].
  */
-internal abstract class RelPeeking : Operator.Relation {
-
-    private var _next: Record? = null
-
-    /**
-     * @return Null when there is not another record to be produced. Returns a [Record] when able to.
-     *
-     * @see RelPeeking
-     */
-    abstract fun peek(): Record?
+internal abstract class RelPeeking : Operator.Relation, IteratorPeeking<Record>() {
 
     override fun open(env: Environment) {
-        _next = null
-    }
-
-    override fun hasNext(): Boolean {
-        if (_next != null) {
-            return true
-        }
-        this._next = peek()
-        return this._next != null
-    }
-
-    override fun next(): Record {
-        val next = _next
-            ?: peek()
-            ?: error("There was not a record to be produced, however, next() was called. Please use hasNext() beforehand.")
-        this._next = null
-        return next
+        next = null
     }
 
     override fun close() {
-        _next = null
+        next = null
     }
 }

partiql-eval/src/main/kotlin/org/partiql/eval/internal/operator/rel/RelUnion.kt → partiql-eval/src/main/kotlin/org/partiql/eval/internal/operator/rel/RelUnionAll.kt

@@ -4,7 +4,7 @@ import org.partiql.eval.internal.Environment
 import org.partiql.eval.internal.Record
 import org.partiql.eval.internal.operator.Operator
 
-internal class RelUnion(
+internal class RelUnionAll(
     private val lhs: Operator.Relation,
     private val rhs: Operator.Relation,
 ) : Operator.Relation {

partiql-eval/src/main/kotlin/org/partiql/eval/internal/operator/rel/RelUnionDistinct.kt

@@ -0,0 +1,40 @@
+package org.partiql.eval.internal.operator.rel
+
+import org.partiql.eval.internal.Environment
+import org.partiql.eval.internal.Record
+import org.partiql.eval.internal.helpers.IteratorChain
+import org.partiql.eval.internal.operator.Operator
+
+internal class RelUnionDistinct(
+    private val lhs: Operator.Relation,
+    private val rhs: Operator.Relation,
+) : RelPeeking() {
+
+    private val seen: MutableSet<Record> = mutableSetOf()
+    private lateinit var input: Iterator<Record>
+
+    override fun open(env: Environment) {
+        lhs.open(env)
+        rhs.open(env)
+        seen.clear()
+        input = IteratorChain(arrayOf(lhs, rhs))
+        super.open(env)
+    }
+
+    override fun peek(): Record? {
+        for (record in input) {
+            if (!seen.contains(record)) {
+                seen.add(record)
+                return record
+            }
+        }
+        return null
+    }
+
+    override fun close() {
+        lhs.close()
+        rhs.close()
+        seen.clear()
+        super.close()
+    }
+}