Merge pull request #1301 from partiql/eval-tuple-union

Adds evaluation of TUPLEUNION

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

@@ -12,6 +12,7 @@ import org.partiql.eval.internal.operator.rex.ExprCollection
 import org.partiql.eval.internal.operator.rex.ExprLiteral
 import org.partiql.eval.internal.operator.rex.ExprSelect
 import org.partiql.eval.internal.operator.rex.ExprStruct
+import org.partiql.eval.internal.operator.rex.ExprTupleUnion
 import org.partiql.eval.internal.operator.rex.ExprVar
 import org.partiql.plan.PartiQLPlan
 import org.partiql.plan.PlanNode
@@ -84,6 +85,11 @@ internal object Compiler {
             return super.visitRexOp(node.op, ctx) as Operator.Expr
         }
 
+        override fun visitRexOpTupleUnion(node: Rex.Op.TupleUnion, ctx: Unit): Operator {
+            val args = node.args.map { visitRex(it, ctx) }.toTypedArray()
+            return ExprTupleUnion(args)
+        }
+
         override fun visitRelOpJoin(node: Rel.Op.Join, ctx: Unit): Operator {
             val lhs = visitRel(node.lhs, ctx)
             val rhs = visitRel(node.rhs, ctx)

partiql-eval/src/main/kotlin/org/partiql/eval/internal/operator/rex/ExprTupleUnion.kt

@@ -0,0 +1,39 @@
+package org.partiql.eval.internal.operator.rex
+
+import org.partiql.eval.internal.Record
+import org.partiql.eval.internal.operator.Operator
+import org.partiql.value.NullValue
+import org.partiql.value.PartiQLValue
+import org.partiql.value.PartiQLValueExperimental
+import org.partiql.value.StructValue
+import org.partiql.value.missingValue
+import org.partiql.value.structValue
+
+internal class ExprTupleUnion(
+    val args: Array<Operator.Expr>
+) : Operator.Expr {
+
+    @OptIn(PartiQLValueExperimental::class)
+    override fun eval(record: Record): PartiQLValue {
+        // Return MISSING on Mistyping Case
+        val tuples = args.map {
+            when (val arg = it.eval(record)) {
+                is StructValue<*> -> arg
+                is NullValue -> structValue(null)
+                else -> when (arg.isNull) {
+                    true -> structValue<PartiQLValue>(null)
+                    false -> return missingValue()
+                }
+            }
+        }
+
+        // Return NULL if any arguments are NULL
+        tuples.forEach {
+            if (it.isNull) {
+                return structValue<PartiQLValue>(null)
+            }
+        }
+
+        return structValue(tuples.flatMap { it.entries })
+    }
+}

partiql-eval/src/test/kotlin/org/partiql/eval/internal/PartiQLEngineDefaultTest.kt

@@ -14,7 +14,9 @@ import org.partiql.value.bagValue
 import org.partiql.value.boolValue
 import org.partiql.value.int32Value
 import org.partiql.value.io.PartiQLValueIonWriterBuilder
+import org.partiql.value.missingValue
 import org.partiql.value.nullValue
+import org.partiql.value.stringValue
 import org.partiql.value.structValue
 import java.io.ByteArrayOutputStream
 import kotlin.test.assertEquals
@@ -112,10 +114,12 @@ class PartiQLEngineDefaultTest {
     fun testJoinOuterFull() {
         val statement =
             parser.parse("SELECT a, b FROM << { 'a': 1 } >> t FULL OUTER JOIN << { 'b': 2 } >> s ON false;").root
+
         val session = PartiQLPlanner.Session("q", "u")
         val plan = planner.plan(statement, session)
 
         val prepared = engine.prepare(plan.plan)
+
         val result = engine.execute(prepared)
         if (result is PartiQLResult.Error) {
             throw result.cause
@@ -136,15 +140,43 @@ class PartiQLEngineDefaultTest {
         assertEquals(expected, output, comparisonString(expected, output))
     }
 
+    @OptIn(PartiQLValueExperimental::class)
+    @Test
+    fun testTupleUnion() {
+        val source = """
+            TUPLEUNION(
+                { 'a': 1 },
+                { 'b': TRUE },
+                { 'c': 'hello' }
+            );
+        """.trimIndent()
+        val statement = parser.parse(source).root
+        val session = PartiQLPlanner.Session("q", "u")
+        val plan = planner.plan(statement, session)
+
+        val prepared = engine.prepare(plan.plan)
+        val result = engine.execute(prepared) as PartiQLResult.Value
+        val output = result.value
+
+        val expected = structValue(
+            "a" to int32Value(1),
+            "b" to boolValue(true),
+            "c" to stringValue("hello")
+        )
+        assertEquals(expected, output)
+    }
+
     @OptIn(PartiQLValueExperimental::class)
     @Test
     fun testJoinOuterFullOnTrue() {
         val statement =
             parser.parse("SELECT a, b FROM << { 'a': 1 } >> t FULL OUTER JOIN << { 'b': 2 } >> s ON TRUE;").root
+
         val session = PartiQLPlanner.Session("q", "u")
         val plan = planner.plan(statement, session)
 
         val prepared = engine.prepare(plan.plan)
+
         val result = engine.execute(prepared)
         if (result is PartiQLResult.Error) {
             throw result.cause
@@ -161,6 +193,28 @@ class PartiQLEngineDefaultTest {
         assertEquals(expected, output, comparisonString(expected, output))
     }
 
+    @OptIn(PartiQLValueExperimental::class)
+    @Test
+    fun testTupleUnionNullInput() {
+        val source = """
+            TUPLEUNION(
+                { 'a': 1 },
+                NULL,
+                { 'c': 'hello' }
+            );
+        """.trimIndent()
+        val statement = parser.parse(source).root
+        val session = PartiQLPlanner.Session("q", "u")
+        val plan = planner.plan(statement, session)
+
+        val prepared = engine.prepare(plan.plan)
+        val result = engine.execute(prepared) as PartiQLResult.Value
+        val output = result.value
+
+        val expected = structValue<PartiQLValue>(null)
+        assertEquals(expected, output)
+    }
+
     @OptIn(PartiQLValueExperimental::class)
     private fun comparisonString(expected: PartiQLValue, actual: PartiQLValue): String {
         val expectedBuffer = ByteArrayOutputStream()
@@ -173,4 +227,84 @@ class PartiQLEngineDefaultTest {
             appendLine("Actual   : $expectedBuffer")
         }
     }
+
+    @OptIn(PartiQLValueExperimental::class)
+    @Test
+    fun testTupleUnionBadInput() {
+        val source = """
+            TUPLEUNION(
+                { 'a': 1 },
+                5,
+                { 'c': 'hello' }
+            );
+        """.trimIndent()
+        val statement = parser.parse(source).root
+        val session = PartiQLPlanner.Session("q", "u")
+        val plan = planner.plan(statement, session)
+
+        val prepared = engine.prepare(plan.plan)
+        val result = engine.execute(prepared) as PartiQLResult.Value
+        val output = result.value
+
+        val expected = missingValue()
+        assertEquals(expected, output)
+    }
+
+    @OptIn(PartiQLValueExperimental::class)
+    @Test
+    fun testTupleUnionDuplicates() {
+        val source = """
+            TUPLEUNION(
+                { 'a': 1, 'b': FALSE },
+                { 'b': TRUE },
+                { 'c': 'hello' }
+            );
+        """.trimIndent()
+        val statement = parser.parse(source).root
+        val session = PartiQLPlanner.Session("q", "u")
+        val plan = planner.plan(statement, session)
+
+        val prepared = engine.prepare(plan.plan)
+        val result = engine.execute(prepared) as PartiQLResult.Value
+        val output = result.value
+
+        val expected = structValue(
+            "a" to int32Value(1),
+            "b" to boolValue(false),
+            "b" to boolValue(true),
+            "c" to stringValue("hello")
+        )
+        assertEquals(expected, output)
+    }
+
+    @OptIn(PartiQLValueExperimental::class)
+    @Test
+    fun testSelectStarTupleUnion() {
+        // As SELECT * gets converted to TUPLEUNION, this is a sanity check
+        val source = """
+            SELECT * FROM
+            <<
+                { 'a': 1, 'b': FALSE }
+            >> AS t,
+            <<
+                { 'b': TRUE }
+            >> AS s
+        """.trimIndent()
+        val statement = parser.parse(source).root
+        val session = PartiQLPlanner.Session("q", "u")
+        val plan = planner.plan(statement, session)
+
+        val prepared = engine.prepare(plan.plan)
+        val result = engine.execute(prepared) as PartiQLResult.Value
+        val output = result.value
+
+        val expected = bagValue(
+            structValue(
+                "a" to int32Value(1),
+                "b" to boolValue(false),
+                "b" to boolValue(true)
+            )
+        )
+        assertEquals(expected, output)
+    }
 }

partiql-types/src/main/kotlin/org/partiql/value/PartiQLValue.kt

@@ -558,6 +558,7 @@ public abstract class MissingValue : PartiQLValue {
 public fun PartiQLValue.toIon(): IonElement = accept(ToIon, Unit)
 
 @PartiQLValueExperimental
+@Throws(TypeCheckException::class)
 public inline fun <reified T : PartiQLValue> PartiQLValue.check(): T {
     if (this is T) return this else throw TypeCheckException()
 }