Skip to content

Commit

Permalink
Adds support for evaluation of CASE WHEN THEN
Browse files Browse the repository at this point in the history
  • Loading branch information
johnedquinn committed Dec 18, 2023
1 parent ce8d9be commit fa08f65
Show file tree
Hide file tree
Showing 3 changed files with 138 additions and 0 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
@@ -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
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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() {
Expand All @@ -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()
Expand All @@ -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)
}
}

0 comments on commit fa08f65

Please sign in to comment.