From dd63f1da0c5cc16aa19195b4596c9efcd24d2159 Mon Sep 17 00:00:00 2001 From: John Ed Quinn Date: Mon, 18 Dec 2023 14:25:50 -0800 Subject: [PATCH] Adds evaluation of TUPLEUNION --- .../org/partiql/eval/internal/Compiler.kt | 6 + .../internal/operator/rex/ExprTupleUnion.kt | 39 +++++ .../eval/internal/PartiQLEngineDefaultTest.kt | 134 ++++++++++++++++++ .../kotlin/org/partiql/value/PartiQLValue.kt | 1 + 4 files changed, 180 insertions(+) create mode 100644 partiql-eval/src/main/kotlin/org/partiql/eval/internal/operator/rex/ExprTupleUnion.kt diff --git a/partiql-eval/src/main/kotlin/org/partiql/eval/internal/Compiler.kt b/partiql-eval/src/main/kotlin/org/partiql/eval/internal/Compiler.kt index a8ba36b7a8..f9e7909f75 100644 --- a/partiql-eval/src/main/kotlin/org/partiql/eval/internal/Compiler.kt +++ b/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) diff --git a/partiql-eval/src/main/kotlin/org/partiql/eval/internal/operator/rex/ExprTupleUnion.kt b/partiql-eval/src/main/kotlin/org/partiql/eval/internal/operator/rex/ExprTupleUnion.kt new file mode 100644 index 0000000000..2bd9926037 --- /dev/null +++ b/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 { + + @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(null) + false -> return missingValue() + } + } + } + + // Return NULL if any arguments are NULL + tuples.forEach { + if (it.isNull) { + return structValue(null) + } + } + + return structValue(tuples.flatMap { it.entries }) + } +} diff --git a/partiql-eval/src/test/kotlin/org/partiql/eval/internal/PartiQLEngineDefaultTest.kt b/partiql-eval/src/test/kotlin/org/partiql/eval/internal/PartiQLEngineDefaultTest.kt index f6699eb733..91d8eae0ca 100644 --- a/partiql-eval/src/test/kotlin/org/partiql/eval/internal/PartiQLEngineDefaultTest.kt +++ b/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(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) + } } diff --git a/partiql-types/src/main/kotlin/org/partiql/value/PartiQLValue.kt b/partiql-types/src/main/kotlin/org/partiql/value/PartiQLValue.kt index bca12407cb..be2ba90a55 100644 --- a/partiql-types/src/main/kotlin/org/partiql/value/PartiQLValue.kt +++ b/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 PartiQLValue.check(): T { if (this is T) return this else throw TypeCheckException() }