From 37c3ce84873a03a00d33d291e43b15a62f8d8541 Mon Sep 17 00:00:00 2001 From: John Ed Quinn Date: Tue, 12 Dec 2023 13:51:59 -0800 Subject: [PATCH] Adds evaluation of TUPLEUNION --- .../org/partiql/eval/internal/Compiler.kt | 6 + .../internal/operator/rex/ExprTupleUnion.kt | 39 +++++ .../eval/internal/PartiQLEngineDefaultTest.kt | 139 ++++++++++++++++++ .../kotlin/org/partiql/value/PartiQLValue.kt | 1 + 4 files changed, 185 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 6ad48a1472..46dc70845a 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 @@ -11,6 +11,7 @@ import org.partiql.eval.internal.operator.rex.ExprLiteral import org.partiql.eval.internal.operator.rex.ExprPathKey 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 @@ -83,6 +84,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 visitRexOpPath(node: Rex.Op.Path, ctx: Unit): Operator { val root = visitRex(node.root, ctx) var path = root 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..4e1082548a --- /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.fields!! }.asSequence()) + } +} 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 dd1da13ff4..488c97cb1a 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 @@ -7,11 +7,14 @@ import org.partiql.parser.PartiQLParserBuilder import org.partiql.planner.PartiQLPlanner import org.partiql.planner.PartiQLPlannerBuilder import org.partiql.value.BagValue +import org.partiql.value.PartiQLValue import org.partiql.value.PartiQLValueExperimental import org.partiql.value.bagValue import org.partiql.value.boolValue import org.partiql.value.int32Value +import org.partiql.value.missingValue import org.partiql.value.nullValue +import org.partiql.value.stringValue import org.partiql.value.structValue import kotlin.test.assertEquals @@ -98,4 +101,140 @@ class PartiQLEngineDefaultTest { val expected = bagValue(sequenceOf(structValue(sequenceOf("a" to int32Value(1), "b" to nullValue())))) assertEquals(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( + sequenceOf( + "a" to int32Value(1), + "b" to boolValue(true), + "c" to stringValue("hello") + ) + ) + assertEquals(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) + @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( + sequenceOf( + "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( + sequenceOf( + structValue( + sequenceOf( + "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 eb6eda092f..09667a63a2 100644 --- a/partiql-types/src/main/kotlin/org/partiql/value/PartiQLValue.kt +++ b/partiql-types/src/main/kotlin/org/partiql/value/PartiQLValue.kt @@ -557,6 +557,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() }