Add COALESCE and NULLIF to the logical plan (#1404)

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
 - The current SqlBlock, SqlDialect, and SqlLayout are marked as deprecated and will be slightly changed in the next release. 

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

@@ -135,6 +135,15 @@ rex::{
       ],
     },
 
+    nullif::{
+      value: rex,
+      nullifier: rex
+    },
+
+    coalesce::{
+      args: 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 value: Rex,
+            @JvmField
+            internal val nullifier: Rex,
+        ) : Op() {
+            internal override val children: List<PlanNode> by lazy {
+                val kids = mutableListOf<PlanNode?>()
+                kids.add(value)
+                kids.add(nullifier)
+                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 args: List<Rex>,
+        ) : Op() {
+            override val children: List<PlanNode> by lazy {
+                val kids = mutableListOf<PlanNode?>()
+                kids.addAll(args)
+                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(
+                value = visitRex(node.value, ctx),
+                nullifier = visitRex(node.nullifier, ctx),
+            )
+
+        override fun visitRexOpCoalesce(node: Rex.Op.Coalesce, ctx: ProblemCallback) =
+            org.partiql.plan.Rex.Op.Coalesce(
+                args = node.args.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 args = 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(args)
+            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 value = visitExprCoerce(node.value, ctx)
+            val nullifier = visitExprCoerce(node.nullifier, ctx)
+            val op = rexOpNullif(value, nullifier)
+            return rex(type, op)
         }
 
         /**

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

@@ -79,7 +79,15 @@ internal class DynamicTyper {
      */
     fun accumulate(type: StaticType) {
         val nonAbsentTypes = mutableSetOf<StaticType>()
-        for (t in type.flatten().allTypes) {
+        val flatType = type.flatten()
+        if (flatType == StaticType.ANY) {
+            // Use ANY runtime; do not expand ANY
+            types.add(flatType)
+            args.add(ANY)
+            calculate(ANY)
+            return
+        }
+        for (t in flatType.allTypes) {
             when (t) {
                 is NullType -> nullable = true
                 is MissingType -> missable = true
@@ -121,7 +129,7 @@ internal class DynamicTyper {
         if (missable) modifiers.add(StaticType.MISSING)
         // If at top supertype, then return union of all accumulated types
         if (supertype == ANY) {
-            return StaticType.unionOf(types + modifiers) to null
+            return StaticType.unionOf(types + modifiers).flatten() to null
         }
         // If a collection, then return union of all accumulated types as these coercion rules are not defined by SQL.
         if (supertype == STRUCT || supertype == BAG || supertype == LIST || supertype == SEXP) {

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,54 @@ internal class PlanTyper(
             return rex(type, op)
         }
 
+        // COALESCE(v1, v2,..., vN)
+        // ==
+        // CASE
+        //     WHEN v1 IS NOT NULL THEN v1  -- WHEN branch always a boolean
+        //     WHEN v2 IS NOT NULL THEN v2  -- WHEN branch always a boolean
+        //     ... -- similarly for v3..vN-1
+        //     ELSE vN
+        // END
+        // --> minimal common supertype of(<type v1>, <type v2>, ..., <type v3>)
+        override fun visitRexOpCoalesce(node: Rex.Op.Coalesce, ctx: StaticType?): Rex {
+            val args = node.args.map { visitRex(it, it.type) }.toMutableList()
+            val typer = DynamicTyper()
+            args.forEach { v ->
+                typer.accumulate(v.type)
+            }
+            val (type, mapping) = typer.mapping()
+            if (mapping != null) {
+                assert(mapping.size == args.size) { "Coercion mappings `len ${mapping.size}` did not match the number of COALESCE arguments `len ${args.size}`" }
+                for (i in args.indices) {
+                    val (operand, target) = mapping[i]
+                    if (operand == target) continue // skip; no coercion needed
+                    val cast = env.fnResolver.cast(operand, target)
+                    val rex = rex(type, rexOpCallStatic(fnResolved(cast), listOf(args[i])))
+                    args[i] = rex
+                }
+            }
+            val op = rexOpCoalesce(args)
+            return rex(type, op)
+        }
+
+        // NULLIF(v1, v2)
+        // ==
+        // CASE
+        //     WHEN v1 = v2 THEN NULL -- WHEN branch always a boolean
+        //     ELSE v1
+        // END
+        // --> minimal common supertype of (NULL, <type v1>)
+        override fun visitRexOpNullif(node: Rex.Op.Nullif, ctx: StaticType?): Rex {
+            val value = visitRex(node.value, node.value.type)
+            val nullifier = visitRex(node.nullifier, node.nullifier.type)
+            val typer = DynamicTyper()
+            typer.accumulate(NULL)
+            typer.accumulate(value.type)
+            val (type, _) = typer.mapping()
+            val op = rexOpNullif(value, nullifier)
+            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::{
+      value: rex,
+      nullifier: rex
+    },
+
+    coalesce::{
+      args: list::[rex]
+    },
+
     collection::{
       values: list::[rex],
     },