Add NULLIF and COALESCE to the logical plan

CHANGELOG.md

@@ -32,6 +32,7 @@ Thank you to all who have contributed!
 
 ### Changed
 - Change `StaticType.AnyOfType`'s `.toString` to not perform `.flatten()`
+- Change modeling of `COALESCE` and `NULLIF` to dedicated nodes in logical plan
 
 ### Deprecated
 
@@ -43,7 +44,7 @@ Thank you to all who have contributed!
 
 ### Contributors
 Thank you to all who have contributed!
-- @<your-username>
+- @alancai98
 
 ## [0.14.4]
 

partiql-plan/src/main/resources/partiql_plan.ion

@@ -135,6 +135,15 @@ rex::{
       ],
     },
 
+    nullif::{
+      v1: rex,
+      v2: rex
+    },
+
+    coalesce::{
+      values: list::[rex]
+    },
+
     collection::{
       values: list::[rex],
     },

partiql-planner/src/main/kotlin/org/partiql/planner/internal/ir/Nodes.kt

@@ -51,10 +51,12 @@ import org.partiql.planner.internal.ir.builder.RexOpCallDynamicCandidateBuilder
 import org.partiql.planner.internal.ir.builder.RexOpCallStaticBuilder
 import org.partiql.planner.internal.ir.builder.RexOpCaseBranchBuilder
 import org.partiql.planner.internal.ir.builder.RexOpCaseBuilder
+import org.partiql.planner.internal.ir.builder.RexOpCoalesceBuilder
 import org.partiql.planner.internal.ir.builder.RexOpCollectionBuilder
 import org.partiql.planner.internal.ir.builder.RexOpErrBuilder
 import org.partiql.planner.internal.ir.builder.RexOpGlobalBuilder
 import org.partiql.planner.internal.ir.builder.RexOpLitBuilder
+import org.partiql.planner.internal.ir.builder.RexOpNullifBuilder
 import org.partiql.planner.internal.ir.builder.RexOpPathIndexBuilder
 import org.partiql.planner.internal.ir.builder.RexOpPathKeyBuilder
 import org.partiql.planner.internal.ir.builder.RexOpPathSymbolBuilder
@@ -312,6 +314,8 @@ internal data class Rex(
             is Path -> visitor.visitRexOpPath(this, ctx)
             is Call -> visitor.visitRexOpCall(this, ctx)
             is Case -> visitor.visitRexOpCase(this, ctx)
+            is Nullif -> visitor.visitRexOpNullif(this, ctx)
+            is Coalesce -> visitor.visitRexOpCoalesce(this, ctx)
             is Collection -> visitor.visitRexOpCollection(this, ctx)
             is Struct -> visitor.visitRexOpStruct(this, ctx)
             is Pivot -> visitor.visitRexOpPivot(this, ctx)
@@ -567,6 +571,47 @@ internal data class Rex(
             }
         }
 
+        internal data class Nullif(
+            @JvmField
+            internal val v1: Rex,
+            @JvmField
+            internal val v2: Rex,
+        ) : Op() {
+            internal override val children: List<PlanNode> by lazy {
+                val kids = mutableListOf<PlanNode?>()
+                kids.add(v1)
+                kids.add(v2)
+                kids.filterNotNull()
+            }
+
+            internal override fun <R, C> accept(visitor: PlanVisitor<R, C>, ctx: C): R =
+                visitor.visitRexOpNullif(this, ctx)
+
+            internal companion object {
+                @JvmStatic
+                internal fun builder(): RexOpNullifBuilder = RexOpNullifBuilder()
+            }
+        }
+
+        internal data class Coalesce(
+            @JvmField
+            internal val values: List<Rex>,
+        ) : Op() {
+            override val children: List<PlanNode> by lazy {
+                val kids = mutableListOf<PlanNode?>()
+                kids.addAll(values)
+                kids.filterNotNull()
+            }
+
+            override fun <R, C> accept(visitor: PlanVisitor<R, C>, ctx: C): R =
+                visitor.visitRexOpCoalesce(this, ctx)
+
+            internal companion object {
+                @JvmStatic
+                internal fun builder(): RexOpCoalesceBuilder = RexOpCoalesceBuilder()
+            }
+        }
+
         internal data class Collection(
             @JvmField internal val values: List<Rex>,
         ) : Op() {

partiql-planner/src/main/kotlin/org/partiql/planner/internal/transforms/PlanTransform.kt

@@ -180,6 +180,17 @@ internal object PlanTransform : PlanBaseVisitor<PlanNode, ProblemCallback>() {
         branches = node.branches.map { visitRexOpCaseBranch(it, ctx) }, default = visitRex(node.default, ctx)
         )
 
+        override fun visitRexOpNullif(node: Rex.Op.Nullif, ctx: ProblemCallback) =
+            org.partiql.plan.Rex.Op.Nullif(
+                v1 = visitRex(node.v1, ctx),
+                v2 = visitRex(node.v2, ctx),
+            )
+
+        override fun visitRexOpCoalesce(node: Rex.Op.Coalesce, ctx: ProblemCallback) =
+            org.partiql.plan.Rex.Op.Coalesce(
+                values = node.values.map { visitRex(it, ctx) }
+            )
+
         override fun visitRexOpCaseBranch(node: Rex.Op.Case.Branch, ctx: ProblemCallback) =
             org.partiql.plan.Rex.Op.Case.Branch(
                 condition = visitRex(node.condition, ctx), rex = visitRex(node.rex, ctx)

partiql-planner/src/main/kotlin/org/partiql/planner/internal/transforms/RexConverter.kt

@@ -30,8 +30,10 @@ import org.partiql.planner.internal.ir.identifierQualified
 import org.partiql.planner.internal.ir.identifierSymbol
 import org.partiql.planner.internal.ir.rex
 import org.partiql.planner.internal.ir.rexOpCallStatic
+import org.partiql.planner.internal.ir.rexOpCoalesce
 import org.partiql.planner.internal.ir.rexOpCollection
 import org.partiql.planner.internal.ir.rexOpLit
+import org.partiql.planner.internal.ir.rexOpNullif
 import org.partiql.planner.internal.ir.rexOpPathIndex
 import org.partiql.planner.internal.ir.rexOpPathKey
 import org.partiql.planner.internal.ir.rexOpPathSymbol
@@ -107,7 +109,7 @@ internal object RexConverter {
         private fun visitExprCoerce(node: Expr, ctx: Env, coercion: Rex.Op.Subquery.Coercion = Rex.Op.Subquery.Coercion.SCALAR): Rex {
             val rex = super.visitExpr(node, ctx)
             return when (rex.op is Rex.Op.Select) {
-                true -> rex(StaticType.ANY, rexOpSubquery(rex.op as Rex.Op.Select, coercion))
+                true -> rex(StaticType.ANY, rexOpSubquery(rex.op, coercion))
                 else -> rex
             }
         }
@@ -439,44 +441,21 @@ internal object RexConverter {
             return rex(type, call)
         }
 
-        // coalesce(expr1, expr2, ... exprN) ->
-        //   CASE
-        //     WHEN expr1 IS NOT NULL THEN EXPR1
-        //     ...
-        //     WHEN exprn is NOT NULL THEN exprn
-        //     ELSE NULL END
-        override fun visitExprCoalesce(node: Expr.Coalesce, ctx: Env): Rex = plan {
+        override fun visitExprCoalesce(node: Expr.Coalesce, ctx: Env): Rex {
             val type = StaticType.ANY
-            val createBranch: (Rex) -> Rex.Op.Case.Branch = { expr: Rex ->
-                val updatedCondition = rex(type, negate(call("is_null", expr)))
-                rexOpCaseBranch(updatedCondition, expr)
+            val values = node.args.map { arg ->
+                visitExprCoerce(arg, ctx)
             }
-
-            val branches = node.args.map {
-                createBranch(visitExpr(it, ctx))
-            }.toMutableList()
-
-            val defaultRex = rex(type = StaticType.NULL, op = rexOpLit(value = nullValue()))
-            val op = rexOpCase(branches, defaultRex)
-            rex(type, op)
+            val op = rexOpCoalesce(values)
+            return rex(type, op)
         }
 
-        // nullIf(expr1, expr2) ->
-        //   CASE
-        //     WHEN expr1 = expr2 THEN NULL
-        //     ELSE expr1 END
-        override fun visitExprNullIf(node: Expr.NullIf, ctx: Env): Rex = plan {
+        override fun visitExprNullIf(node: Expr.NullIf, ctx: Env): Rex {
             val type = StaticType.ANY
-            val expr1 = visitExpr(node.value, ctx)
-            val expr2 = visitExpr(node.nullifier, ctx)
-            val id = identifierSymbol(Expr.Binary.Op.EQ.name.lowercase(), Identifier.CaseSensitivity.SENSITIVE)
-            val fn = fnUnresolved(id, true)
-            val call = rexOpCallStatic(fn, listOf(expr1, expr2))
-            val branches = listOf(
-                rexOpCaseBranch(rex(type, call), rex(type = StaticType.NULL, op = rexOpLit(value = nullValue()))),
-            )
-            val op = rexOpCase(branches.toMutableList(), expr1)
-            rex(type, op)
+            val expr1 = visitExprCoerce(node.value, ctx)
+            val expr2 = visitExprCoerce(node.nullifier, ctx)
+            val op = rexOpNullif(expr1, expr2)
+            return rex(type, op)
         }
 
         /**

partiql-planner/src/main/kotlin/org/partiql/planner/internal/typer/PlanTyper.kt

@@ -55,9 +55,11 @@ import org.partiql.planner.internal.ir.rexOpCallDynamic
 import org.partiql.planner.internal.ir.rexOpCallDynamicCandidate
 import org.partiql.planner.internal.ir.rexOpCallStatic
 import org.partiql.planner.internal.ir.rexOpCaseBranch
+import org.partiql.planner.internal.ir.rexOpCoalesce
 import org.partiql.planner.internal.ir.rexOpCollection
 import org.partiql.planner.internal.ir.rexOpErr
 import org.partiql.planner.internal.ir.rexOpLit
+import org.partiql.planner.internal.ir.rexOpNullif
 import org.partiql.planner.internal.ir.rexOpPathIndex
 import org.partiql.planner.internal.ir.rexOpPathKey
 import org.partiql.planner.internal.ir.rexOpPathSymbol
@@ -760,6 +762,46 @@ internal class PlanTyper(
             return rex(type, op)
         }
 
+        override fun visitRexOpCoalesce(node: Rex.Op.Coalesce, ctx: StaticType?): Rex {
+            val values = node.values.map { visitRex(it, it.type) }.toMutableList()
+            val typer = DynamicTyper()
+            values.forEach { v ->
+                typer.accumulate(v.type)
+            }
+            val (type, mapping) = typer.mapping()
+            if (mapping != null) {
+                assert(mapping.size == values.size)
+                for (i in values.indices) {
+                    val (operand, target) = mapping[i]
+                    if (operand == target) continue // skip
+                    val cast = env.fnResolver.cast(operand, target)
+                    val rex = rex(type, rexOpCallStatic(fnResolved(cast), listOf(values[i])))
+                    values[i] = rex
+                }
+            }
+            val op = rexOpCoalesce(values)
+            return rex(type, op)
+        }
+
+        // NULLIF(v1, v2)
+        // ==
+        // CASE
+        //     WHEN V1 = V2 THEN NULL -- WHEN branch always a boolean
+        //     ELSE V1
+        // END
+        // --> union(null, <type of V1>)
+        override fun visitRexOpNullif(node: Rex.Op.Nullif, ctx: StaticType?): Rex {
+            val v1 = visitRex(node.v1, node.v1.type)
+            val v2 = visitRex(node.v2, node.v2.type)
+            // reuse dynamic typing logic
+            val typer = DynamicTyper()
+            typer.accumulate(NULL)
+            typer.accumulate(v1.type)
+            val (type, _) = typer.mapping()
+            val op = rexOpNullif(v1, v2)
+            return rex(type, op)
+        }
+
         /**
          * In this context, Boolean means PartiQLValueType Bool, which can be nullable.
          * Hence, we permit Static Type BOOL, Static Type NULL, Static Type Missing here.

partiql-planner/src/main/resources/partiql_plan_internal.ion

@@ -158,6 +158,15 @@ rex::{
       ],
     },
 
+    nullif::{
+      v1: rex,
+      v2: rex
+    },
+
+    coalesce::{
+      values: list::[rex]
+    },
+
     collection::{
       values: list::[rex],
     },