From 0bcc3e29e1ad9bb87af6fb1195e738986487809b Mon Sep 17 00:00:00 2001
From: John Ed Quinn <johqunn@amazon.com>
Date: Wed, 17 Apr 2024 16:19:17 -0700
Subject: [PATCH] Adds support for NULLIF and COALESCE

---
 .../org/partiql/eval/internal/Compiler.kt     | 13 ++++++++++
 .../internal/operator/rex/ExprCoalesce.kt     | 24 +++++++++++++++++
 .../eval/internal/operator/rex/ExprNullIf.kt  | 26 +++++++++++++++++++
 3 files changed, 63 insertions(+)
 create mode 100644 partiql-eval/src/main/kotlin/org/partiql/eval/internal/operator/rex/ExprCoalesce.kt
 create mode 100644 partiql-eval/src/main/kotlin/org/partiql/eval/internal/operator/rex/ExprNullIf.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 07d091a636..d096a765c7 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
@@ -23,8 +23,10 @@ import org.partiql.eval.internal.operator.rex.ExprCallDynamic
 import org.partiql.eval.internal.operator.rex.ExprCallStatic
 import org.partiql.eval.internal.operator.rex.ExprCase
 import org.partiql.eval.internal.operator.rex.ExprCast
+import org.partiql.eval.internal.operator.rex.ExprCoalesce
 import org.partiql.eval.internal.operator.rex.ExprCollection
 import org.partiql.eval.internal.operator.rex.ExprLiteral
+import org.partiql.eval.internal.operator.rex.ExprNullIf
 import org.partiql.eval.internal.operator.rex.ExprPathIndex
 import org.partiql.eval.internal.operator.rex.ExprPathKey
 import org.partiql.eval.internal.operator.rex.ExprPathSymbol
@@ -132,6 +134,17 @@ internal class Compiler(
         }
     }
 
+    override fun visitRexOpCoalesce(node: Rex.Op.Coalesce, ctx: StaticType?): Operator {
+        val args = Array(node.args.size) { visitRex(node.args[it], node.args[it].type) }
+        return ExprCoalesce(args)
+    }
+
+    override fun visitRexOpNullif(node: Rex.Op.Nullif, ctx: StaticType?): Operator {
+        val value = visitRex(node.value, node.value.type)
+        val nullifier = visitRex(node.nullifier, node.value.type)
+        return ExprNullIf(value, nullifier)
+    }
+
     /**
      * All variables from the local scope have a depth of 0.
      *
diff --git a/partiql-eval/src/main/kotlin/org/partiql/eval/internal/operator/rex/ExprCoalesce.kt b/partiql-eval/src/main/kotlin/org/partiql/eval/internal/operator/rex/ExprCoalesce.kt
new file mode 100644
index 0000000000..09b2438fe8
--- /dev/null
+++ b/partiql-eval/src/main/kotlin/org/partiql/eval/internal/operator/rex/ExprCoalesce.kt
@@ -0,0 +1,24 @@
+package org.partiql.eval.internal.operator.rex
+
+import org.partiql.eval.internal.Environment
+import org.partiql.eval.internal.operator.Operator
+import org.partiql.value.PartiQLValue
+import org.partiql.value.PartiQLValueExperimental
+import org.partiql.value.PartiQLValueType
+import org.partiql.value.nullValue
+
+internal class ExprCoalesce(
+    private val args: Array<Operator.Expr>
+) : Operator.Expr {
+
+    @PartiQLValueExperimental
+    override fun eval(env: Environment): PartiQLValue {
+        for (arg in args) {
+            val result = arg.eval(env)
+            if (!result.isNull && result.type != PartiQLValueType.MISSING) {
+                return result
+            }
+        }
+        return nullValue()
+    }
+}
diff --git a/partiql-eval/src/main/kotlin/org/partiql/eval/internal/operator/rex/ExprNullIf.kt b/partiql-eval/src/main/kotlin/org/partiql/eval/internal/operator/rex/ExprNullIf.kt
new file mode 100644
index 0000000000..63b994af94
--- /dev/null
+++ b/partiql-eval/src/main/kotlin/org/partiql/eval/internal/operator/rex/ExprNullIf.kt
@@ -0,0 +1,26 @@
+package org.partiql.eval.internal.operator.rex
+
+import org.partiql.eval.internal.Environment
+import org.partiql.eval.internal.operator.Operator
+import org.partiql.value.PartiQLValue
+import org.partiql.value.PartiQLValueExperimental
+import org.partiql.value.nullValue
+
+internal class ExprNullIf(
+    private val valueExpr: Operator.Expr,
+    private val nullifierExpr: Operator.Expr
+) : Operator.Expr {
+
+    @OptIn(PartiQLValueExperimental::class)
+    private val comparator = PartiQLValue.comparator()
+
+    @PartiQLValueExperimental
+    override fun eval(env: Environment): PartiQLValue {
+        val value = valueExpr.eval(env)
+        val nullifier = nullifierExpr.eval(env)
+        return when (comparator.compare(value, nullifier)) {
+            0 -> nullValue()
+            else -> value
+        }
+    }
+}