Merge pull request #1307 from partiql/eval-case

Adds support for evaluation of CASE WHEN THEN

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
@@ -102,6 +103,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

@@ -166,6 +166,30 @@ class PartiQLEngineDefaultTest {
         assertEquals(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() {
@@ -212,6 +236,31 @@ class PartiQLEngineDefaultTest {
         val output = result.value
 
         val expected = structValue<PartiQLValue>(null)
+
+        assertEquals(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)
     }
 
@@ -245,11 +294,35 @@ class PartiQLEngineDefaultTest {
         val prepared = engine.prepare(plan.plan)
         val result = engine.execute(prepared) as PartiQLResult.Value
         val output = result.value
-
         val expected = missingValue()
         assertEquals(expected, output)
     }
 
+    @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)
+    }
+
     @OptIn(PartiQLValueExperimental::class)
     @Test
     fun testTupleUnionDuplicates() {
@@ -259,6 +332,7 @@ class PartiQLEngineDefaultTest {
                 { 'b': TRUE },
                 { 'c': 'hello' }
             );
+            
         """.trimIndent()
         val statement = parser.parse(source).root
         val session = PartiQLPlanner.Session("q", "u")
@@ -277,6 +351,29 @@ class PartiQLEngineDefaultTest {
         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)
+    }
+
     @OptIn(PartiQLValueExperimental::class)
     @Test
     fun testSelectStarTupleUnion() {