Adds support for evaluation of CASE WHEN THEN

Author: johnedquinn

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

@@ -8,6 +8,7 @@ import org.partiql.eval.internal.operator.rel.RelJoinOuterFull
 import org.partiql.eval.internal.operator.rel.RelJoinRight
 import org.partiql.eval.internal.operator.rel.RelProject
 import org.partiql.eval.internal.operator.rel.RelScan
+import org.partiql.eval.internal.operator.rex.ExprCase
 import org.partiql.eval.internal.operator.rex.ExprCollection
 import org.partiql.eval.internal.operator.rex.ExprLiteral
 import org.partiql.eval.internal.operator.rex.ExprSelect
@@ -96,6 +97,14 @@ internal object Compiler {
             }
         }
 
+        override fun visitRexOpCase(node: Rex.Op.Case, ctx: Unit): Operator {
+            val branches = node.branches.map { branch ->
+                visitRex(branch.condition, ctx) to visitRex(branch.rex, ctx)
+            }
+            val default = visitRex(node.default, ctx)
+            return ExprCase(branches, default)
+        }
+
         // TODO: Re-look at
         override fun visitPartiQLPlan(node: PartiQLPlan, ctx: Unit): Operator.Expr {
             return visitStatement(node.statement, ctx) as Operator.Expr

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

@@ -0,0 +1,29 @@
+package org.partiql.eval.internal.operator.rex
+
+import org.partiql.eval.internal.Record
+import org.partiql.eval.internal.operator.Operator
+import org.partiql.value.BoolValue
+import org.partiql.value.PartiQLValue
+import org.partiql.value.PartiQLValueExperimental
+
+internal class ExprCase(
+    private val branches: List<Pair<Operator.Expr, Operator.Expr>>,
+    private val default: Operator.Expr
+) : Operator.Expr {
+
+    @OptIn(PartiQLValueExperimental::class)
+    override fun eval(record: Record): PartiQLValue {
+        branches.forEach { branch ->
+            val condition = branch.first.eval(record)
+            if (condition.isTrue()) {
+                return branch.second.eval(record)
+            }
+        }
+        return default.eval(record)
+    }
+
+    @OptIn(PartiQLValueExperimental::class)
+    private fun PartiQLValue.isTrue(): Boolean {
+        return this is BoolValue && this.value == true
+    }
+}

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

@@ -15,6 +15,7 @@ import org.partiql.value.boolValue
 import org.partiql.value.int32Value
 import org.partiql.value.io.PartiQLValueIonWriterBuilder
 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 +113,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,6 +139,30 @@ class PartiQLEngineDefaultTest {
         assertEquals(expected, output, comparisonString(expected, output))
     }
 
+    @OptIn(PartiQLValueExperimental::class)
+    @Test
+    fun testCaseLiteral00() {
+        val source = """
+            CASE
+                WHEN NULL THEN 'isNull'
+                WHEN MISSING THEN 'isMissing'
+                WHEN FALSE THEN 'isFalse'
+                WHEN TRUE THEN 'isTrue'
+            END
+            ;
+        """.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 = stringValue("isTrue")
+        assertEquals(expected, output)
+    }
+
     @OptIn(PartiQLValueExperimental::class)
     @Test
     fun testJoinOuterFullOnTrue() {
@@ -161,6 +188,30 @@ class PartiQLEngineDefaultTest {
         assertEquals(expected, output, comparisonString(expected, output))
     }
 
+    @OptIn(PartiQLValueExperimental::class)
+    @Test
+    fun testCaseLiteral01() {
+        val source = """
+            CASE
+                WHEN NULL THEN 'isNull'
+                WHEN MISSING THEN 'isMissing'
+                WHEN FALSE THEN 'isFalse'
+            END
+            ;
+        """.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 = nullValue()
+        assertEquals(expected, output)
+    }
+
     @OptIn(PartiQLValueExperimental::class)
     private fun comparisonString(expected: PartiQLValue, actual: PartiQLValue): String {
         val expectedBuffer = ByteArrayOutputStream()
@@ -173,4 +224,53 @@ class PartiQLEngineDefaultTest {
             appendLine("Actual   : $expectedBuffer")
         }
     }
+
+    @Disabled("This is disabled because FN EQUALS is not yet implemented.")
+    @OptIn(PartiQLValueExperimental::class)
+    @Test
+    fun testCaseLiteral02() {
+        val source = """
+            CASE (1)
+                WHEN NULL THEN 'isNull'
+                WHEN MISSING THEN 'isMissing'
+                WHEN 2 THEN 'isTwo'
+                WHEN 1 THEN 'isOne'
+            END
+            ;
+        """.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 = stringValue("isOne")
+        assertEquals(expected, output)
+    }
+
+    @Disabled("This is disabled because FN EQUALS is not yet implemented.")
+    @OptIn(PartiQLValueExperimental::class)
+    @Test
+    fun testCaseLiteral03() {
+        val source = """
+            CASE (1)
+                WHEN NULL THEN 'isNull'
+                WHEN MISSING THEN 'isMissing'
+                WHEN 2 THEN 'isTwo'
+            END
+            ;
+        """.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 = nullValue()
+        assertEquals(expected, output)
+    }
 }