Skip to content

Commit 92cc57a

Browse files
alancai98alancai98
authored
Add COALESCE and NULLIF to the logical plan (#1404)
1 parent 6fbfa78 commit 92cc57a

File tree

12 files changed

+557
-39
lines changed

12 files changed

+557
-39
lines changed

‎CHANGELOG.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@ Thank you to all who have contributed!
3232

3333
### Changed
3434
- Change `StaticType.AnyOfType`'s `.toString` to not perform `.flatten()`
35+
- Change modeling of `COALESCE` and `NULLIF` to dedicated nodes in logical plan
3536

3637
### Deprecated
3738
- 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

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -135,6 +135,15 @@ rex::{
135135
],
136136
},
137137

138+
nullif::{
139+
value: rex,
140+
nullifier: rex
141+
},
142+
143+
coalesce::{
144+
args: list::[rex]
145+
},
146+
138147
collection::{
139148
values: list::[rex],
140149
},

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

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,10 +51,12 @@ import org.partiql.planner.internal.ir.builder.RexOpCallDynamicCandidateBuilder
5151
import org.partiql.planner.internal.ir.builder.RexOpCallStaticBuilder
5252
import org.partiql.planner.internal.ir.builder.RexOpCaseBranchBuilder
5353
import org.partiql.planner.internal.ir.builder.RexOpCaseBuilder
54+
import org.partiql.planner.internal.ir.builder.RexOpCoalesceBuilder
5455
import org.partiql.planner.internal.ir.builder.RexOpCollectionBuilder
5556
import org.partiql.planner.internal.ir.builder.RexOpErrBuilder
5657
import org.partiql.planner.internal.ir.builder.RexOpGlobalBuilder
5758
import org.partiql.planner.internal.ir.builder.RexOpLitBuilder
59+
import org.partiql.planner.internal.ir.builder.RexOpNullifBuilder
5860
import org.partiql.planner.internal.ir.builder.RexOpPathIndexBuilder
5961
import org.partiql.planner.internal.ir.builder.RexOpPathKeyBuilder
6062
import org.partiql.planner.internal.ir.builder.RexOpPathSymbolBuilder
@@ -312,6 +314,8 @@ internal data class Rex(
312314
is Path -> visitor.visitRexOpPath(this, ctx)
313315
is Call -> visitor.visitRexOpCall(this, ctx)
314316
is Case -> visitor.visitRexOpCase(this, ctx)
317+
is Nullif -> visitor.visitRexOpNullif(this, ctx)
318+
is Coalesce -> visitor.visitRexOpCoalesce(this, ctx)
315319
is Collection -> visitor.visitRexOpCollection(this, ctx)
316320
is Struct -> visitor.visitRexOpStruct(this, ctx)
317321
is Pivot -> visitor.visitRexOpPivot(this, ctx)
@@ -567,6 +571,47 @@ internal data class Rex(
567571
}
568572
}
569573

574+
internal data class Nullif(
575+
@JvmField
576+
internal val value: Rex,
577+
@JvmField
578+
internal val nullifier: Rex,
579+
) : Op() {
580+
internal override val children: List<PlanNode> by lazy {
581+
val kids = mutableListOf<PlanNode?>()
582+
kids.add(value)
583+
kids.add(nullifier)
584+
kids.filterNotNull()
585+
}
586+
587+
internal override fun <R, C> accept(visitor: PlanVisitor<R, C>, ctx: C): R =
588+
visitor.visitRexOpNullif(this, ctx)
589+
590+
internal companion object {
591+
@JvmStatic
592+
internal fun builder(): RexOpNullifBuilder = RexOpNullifBuilder()
593+
}
594+
}
595+
596+
internal data class Coalesce(
597+
@JvmField
598+
internal val args: List<Rex>,
599+
) : Op() {
600+
override val children: List<PlanNode> by lazy {
601+
val kids = mutableListOf<PlanNode?>()
602+
kids.addAll(args)
603+
kids.filterNotNull()
604+
}
605+
606+
override fun <R, C> accept(visitor: PlanVisitor<R, C>, ctx: C): R =
607+
visitor.visitRexOpCoalesce(this, ctx)
608+
609+
internal companion object {
610+
@JvmStatic
611+
internal fun builder(): RexOpCoalesceBuilder = RexOpCoalesceBuilder()
612+
}
613+
}
614+
570615
internal data class Collection(
571616
@JvmField internal val values: List<Rex>,
572617
) : Op() {

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

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -180,6 +180,17 @@ internal object PlanTransform : PlanBaseVisitor<PlanNode, ProblemCallback>() {
180180
branches = node.branches.map { visitRexOpCaseBranch(it, ctx) }, default = visitRex(node.default, ctx)
181181
)
182182

183+
override fun visitRexOpNullif(node: Rex.Op.Nullif, ctx: ProblemCallback) =
184+
org.partiql.plan.Rex.Op.Nullif(
185+
value = visitRex(node.value, ctx),
186+
nullifier = visitRex(node.nullifier, ctx),
187+
)
188+
189+
override fun visitRexOpCoalesce(node: Rex.Op.Coalesce, ctx: ProblemCallback) =
190+
org.partiql.plan.Rex.Op.Coalesce(
191+
args = node.args.map { visitRex(it, ctx) }
192+
)
193+
183194
override fun visitRexOpCaseBranch(node: Rex.Op.Case.Branch, ctx: ProblemCallback) =
184195
org.partiql.plan.Rex.Op.Case.Branch(
185196
condition = visitRex(node.condition, ctx), rex = visitRex(node.rex, ctx)

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

Lines changed: 13 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -30,8 +30,10 @@ import org.partiql.planner.internal.ir.identifierQualified
3030
import org.partiql.planner.internal.ir.identifierSymbol
3131
import org.partiql.planner.internal.ir.rex
3232
import org.partiql.planner.internal.ir.rexOpCallStatic
33+
import org.partiql.planner.internal.ir.rexOpCoalesce
3334
import org.partiql.planner.internal.ir.rexOpCollection
3435
import org.partiql.planner.internal.ir.rexOpLit
36+
import org.partiql.planner.internal.ir.rexOpNullif
3537
import org.partiql.planner.internal.ir.rexOpPathIndex
3638
import org.partiql.planner.internal.ir.rexOpPathKey
3739
import org.partiql.planner.internal.ir.rexOpPathSymbol
@@ -107,7 +109,7 @@ internal object RexConverter {
107109
private fun visitExprCoerce(node: Expr, ctx: Env, coercion: Rex.Op.Subquery.Coercion = Rex.Op.Subquery.Coercion.SCALAR): Rex {
108110
val rex = super.visitExpr(node, ctx)
109111
return when (rex.op is Rex.Op.Select) {
110-
true -> rex(StaticType.ANY, rexOpSubquery(rex.op as Rex.Op.Select, coercion))
112+
true -> rex(StaticType.ANY, rexOpSubquery(rex.op, coercion))
111113
else -> rex
112114
}
113115
}
@@ -439,44 +441,21 @@ internal object RexConverter {
439441
return rex(type, call)
440442
}
441443

442-
// coalesce(expr1, expr2, ... exprN) ->
443-
// CASE
444-
// WHEN expr1 IS NOT NULL THEN EXPR1
445-
// ...
446-
// WHEN exprn is NOT NULL THEN exprn
447-
// ELSE NULL END
448-
override fun visitExprCoalesce(node: Expr.Coalesce, ctx: Env): Rex = plan {
444+
override fun visitExprCoalesce(node: Expr.Coalesce, ctx: Env): Rex {
449445
val type = StaticType.ANY
450-
val createBranch: (Rex) -> Rex.Op.Case.Branch = { expr: Rex ->
451-
val updatedCondition = rex(type, negate(call("is_null", expr)))
452-
rexOpCaseBranch(updatedCondition, expr)
446+
val args = node.args.map { arg ->
447+
visitExprCoerce(arg, ctx)
453448
}
454-
455-
val branches = node.args.map {
456-
createBranch(visitExpr(it, ctx))
457-
}.toMutableList()
458-
459-
val defaultRex = rex(type = StaticType.NULL, op = rexOpLit(value = nullValue()))
460-
val op = rexOpCase(branches, defaultRex)
461-
rex(type, op)
449+
val op = rexOpCoalesce(args)
450+
return rex(type, op)
462451
}
463452

464-
// nullIf(expr1, expr2) ->
465-
// CASE
466-
// WHEN expr1 = expr2 THEN NULL
467-
// ELSE expr1 END
468-
override fun visitExprNullIf(node: Expr.NullIf, ctx: Env): Rex = plan {
453+
override fun visitExprNullIf(node: Expr.NullIf, ctx: Env): Rex {
469454
val type = StaticType.ANY
470-
val expr1 = visitExpr(node.value, ctx)
471-
val expr2 = visitExpr(node.nullifier, ctx)
472-
val id = identifierSymbol(Expr.Binary.Op.EQ.name.lowercase(), Identifier.CaseSensitivity.SENSITIVE)
473-
val fn = fnUnresolved(id, true)
474-
val call = rexOpCallStatic(fn, listOf(expr1, expr2))
475-
val branches = listOf(
476-
rexOpCaseBranch(rex(type, call), rex(type = StaticType.NULL, op = rexOpLit(value = nullValue()))),
477-
)
478-
val op = rexOpCase(branches.toMutableList(), expr1)
479-
rex(type, op)
455+
val value = visitExprCoerce(node.value, ctx)
456+
val nullifier = visitExprCoerce(node.nullifier, ctx)
457+
val op = rexOpNullif(value, nullifier)
458+
return rex(type, op)
480459
}
481460

482461
/**

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

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -79,7 +79,15 @@ internal class DynamicTyper {
7979
*/
8080
fun accumulate(type: StaticType) {
8181
val nonAbsentTypes = mutableSetOf<StaticType>()
82-
for (t in type.flatten().allTypes) {
82+
val flatType = type.flatten()
83+
if (flatType == StaticType.ANY) {
84+
// Use ANY runtime; do not expand ANY
85+
types.add(flatType)
86+
args.add(ANY)
87+
calculate(ANY)
88+
return
89+
}
90+
for (t in flatType.allTypes) {
8391
when (t) {
8492
is NullType -> nullable = true
8593
is MissingType -> missable = true
@@ -121,7 +129,7 @@ internal class DynamicTyper {
121129
if (missable) modifiers.add(StaticType.MISSING)
122130
// If at top supertype, then return union of all accumulated types
123131
if (supertype == ANY) {
124-
return StaticType.unionOf(types + modifiers) to null
132+
return StaticType.unionOf(types + modifiers).flatten() to null
125133
}
126134
// If a collection, then return union of all accumulated types as these coercion rules are not defined by SQL.
127135
if (supertype == STRUCT || supertype == BAG || supertype == LIST || supertype == SEXP) {

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

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,9 +55,11 @@ import org.partiql.planner.internal.ir.rexOpCallDynamic
5555
import org.partiql.planner.internal.ir.rexOpCallDynamicCandidate
5656
import org.partiql.planner.internal.ir.rexOpCallStatic
5757
import org.partiql.planner.internal.ir.rexOpCaseBranch
58+
import org.partiql.planner.internal.ir.rexOpCoalesce
5859
import org.partiql.planner.internal.ir.rexOpCollection
5960
import org.partiql.planner.internal.ir.rexOpErr
6061
import org.partiql.planner.internal.ir.rexOpLit
62+
import org.partiql.planner.internal.ir.rexOpNullif
6163
import org.partiql.planner.internal.ir.rexOpPathIndex
6264
import org.partiql.planner.internal.ir.rexOpPathKey
6365
import org.partiql.planner.internal.ir.rexOpPathSymbol
@@ -760,6 +762,54 @@ internal class PlanTyper(
760762
return rex(type, op)
761763
}
762764

765+
// COALESCE(v1, v2,..., vN)
766+
// ==
767+
// CASE
768+
// WHEN v1 IS NOT NULL THEN v1 -- WHEN branch always a boolean
769+
// WHEN v2 IS NOT NULL THEN v2 -- WHEN branch always a boolean
770+
// ... -- similarly for v3..vN-1
771+
// ELSE vN
772+
// END
773+
// --> minimal common supertype of(<type v1>, <type v2>, ..., <type v3>)
774+
override fun visitRexOpCoalesce(node: Rex.Op.Coalesce, ctx: StaticType?): Rex {
775+
val args = node.args.map { visitRex(it, it.type) }.toMutableList()
776+
val typer = DynamicTyper()
777+
args.forEach { v ->
778+
typer.accumulate(v.type)
779+
}
780+
val (type, mapping) = typer.mapping()
781+
if (mapping != null) {
782+
assert(mapping.size == args.size) { "Coercion mappings `len ${mapping.size}` did not match the number of COALESCE arguments `len ${args.size}`" }
783+
for (i in args.indices) {
784+
val (operand, target) = mapping[i]
785+
if (operand == target) continue // skip; no coercion needed
786+
val cast = env.fnResolver.cast(operand, target)
787+
val rex = rex(type, rexOpCallStatic(fnResolved(cast), listOf(args[i])))
788+
args[i] = rex
789+
}
790+
}
791+
val op = rexOpCoalesce(args)
792+
return rex(type, op)
793+
}
794+
795+
// NULLIF(v1, v2)
796+
// ==
797+
// CASE
798+
// WHEN v1 = v2 THEN NULL -- WHEN branch always a boolean
799+
// ELSE v1
800+
// END
801+
// --> minimal common supertype of (NULL, <type v1>)
802+
override fun visitRexOpNullif(node: Rex.Op.Nullif, ctx: StaticType?): Rex {
803+
val value = visitRex(node.value, node.value.type)
804+
val nullifier = visitRex(node.nullifier, node.nullifier.type)
805+
val typer = DynamicTyper()
806+
typer.accumulate(NULL)
807+
typer.accumulate(value.type)
808+
val (type, _) = typer.mapping()
809+
val op = rexOpNullif(value, nullifier)
810+
return rex(type, op)
811+
}
812+
763813
/**
764814
* In this context, Boolean means PartiQLValueType Bool, which can be nullable.
765815
* Hence, we permit Static Type BOOL, Static Type NULL, Static Type Missing here.

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

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -158,6 +158,15 @@ rex::{
158158
],
159159
},
160160

161+
nullif::{
162+
value: rex,
163+
nullifier: rex
164+
},
165+
166+
coalesce::{
167+
args: list::[rex]
168+
},
169+
161170
collection::{
162171
values: list::[rex],
163172
},

0 commit comments

Comments
 (0)