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 were 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/RelDistinct.kt

@@ -10,9 +10,8 @@ internal class RelDistinct(
 
     private val seen = mutableSetOf<Record>()
 
-    override fun open(env: Environment) {
+    override fun openPeeking(env: Environment) {
         input.open(env)
-        super.open(env)
     }
 
     override fun peek(): Record? {
@@ -25,9 +24,8 @@ internal class RelDistinct(
         return null
     }
 
-    override fun close() {
+    override fun closePeeking() {
         seen.clear()
         input.close()
-        super.close()
     }
 }

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

@@ -0,0 +1,53 @@
+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 openPeeking(env: Environment) {
+        lhs.open(env)
+        rhs.open(env)
+        init = false
+        seen.clear()
+    }
+
+    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 closePeeking() {
+        lhs.close()
+        rhs.close()
+        seen.clear()
+    }
+
+    /**
+     * 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,20 +10,19 @@ 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() {
 
     private var seen: MutableSet<Record> = mutableSetOf()
     private var init: Boolean = false
 
-    override fun open(env: Environment) {
+    override fun openPeeking(env: Environment) {
         lhs.open(env)
         rhs.open(env)
         init = false
         seen = mutableSetOf()
-        super.open(env)
     }
 
     override fun peek(): Record? {
@@ -38,11 +37,10 @@ internal class RelExcept(
         return null
     }
 
-    override fun close() {
+    override fun closePeeking() {
         lhs.close()
         rhs.close()
         seen.clear()
-        super.close()
     }
 
     /**

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

@@ -13,10 +13,9 @@ internal class RelFilter(
 
     private lateinit var env: Environment
 
-    override fun open(env: Environment) {
+    override fun openPeeking(env: Environment) {
         this.env = env
         input.open(env)
-        super.open(env)
     }
 
     override fun peek(): Record? {
@@ -28,9 +27,8 @@ internal class RelFilter(
         return null
     }
 
-    override fun close() {
+    override fun closePeeking() {
         input.close()
-        super.close()
     }
 
     @OptIn(PartiQLValueExperimental::class)

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

@@ -0,0 +1,54 @@
+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 openPeeking(env: Environment) {
+        lhs.open(env)
+        rhs.open(env)
+        init = false
+        seen.clear()
+    }
+
+    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 closePeeking() {
+        lhs.close()
+        rhs.close()
+        seen.clear()
+    }
+
+    /**
+     * 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,48 +4,45 @@ 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) {
+    override fun openPeeking(env: Environment) {
         lhs.open(env)
         rhs.open(env)
         init = false
-        seen = mutableSetOf()
-        super.open(env)
+        seen.clear()
     }
 
     override fun peek(): Record? {
         if (!init) {
             seed()
         }
         for (row in rhs) {
-            if (seen.contains(row)) {
+            if (seen.remove(row)) {
                 return row
             }
         }
         return null
     }
 
-    override fun close() {
+    override fun closePeeking() {
         lhs.close()
         rhs.close()
         seen.clear()
-        super.close()
     }
 
     /**
      * Read the entire left-hand-side into our search structure.
      */
     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/RelJoinNestedLoop.kt

@@ -19,15 +19,14 @@ internal abstract class RelJoinNestedLoop : RelPeeking() {
     private var lhsRecord: Record? = null
     private lateinit var env: Environment
 
-    override fun open(env: Environment) {
+    override fun openPeeking(env: Environment) {
         this.env = env
         lhs.open(env)
         if (lhs.hasNext().not()) {
             return
         }
         lhsRecord = lhs.next()
         rhs.open(env.push(lhsRecord!!))
-        super.open(env)
     }
 
     abstract fun join(condition: Boolean, lhs: Record, rhs: Record): Record?
@@ -69,10 +68,9 @@ internal abstract class RelJoinNestedLoop : RelPeeking() {
         return toReturn
     }
 
-    override fun close() {
+    override fun closePeeking() {
         lhs.close()
         rhs.close()
-        super.close()
     }
 
     @OptIn(PartiQLValueExperimental::class)