Skip to content

Commit b178180

Browse files
alancai98alancai98
committed
Add NULLIF and COALESCE to the logical plan
1 parent 6da7496 commit b178180

File tree

11 files changed

+475
-38
lines changed

11 files changed

+475
-38
lines changed

‎CHANGELOG.md

+2-1
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

@@ -43,7 +44,7 @@ Thank you to all who have contributed!
4344

4445
### Contributors
4546
Thank you to all who have contributed!
46-
- @<your-username>
47+
- @alancai98
4748

4849
## [0.14.4]
4950

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

+9
Original file line numberDiff line numberDiff line change
@@ -135,6 +135,15 @@ rex::{
135135
],
136136
},
137137

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

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

+45
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 v1: Rex,
577+
@JvmField
578+
internal val v2: Rex,
579+
) : Op() {
580+
internal override val children: List<PlanNode> by lazy {
581+
val kids = mutableListOf<PlanNode?>()
582+
kids.add(v1)
583+
kids.add(v2)
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 values: List<Rex>,
599+
) : Op() {
600+
override val children: List<PlanNode> by lazy {
601+
val kids = mutableListOf<PlanNode?>()
602+
kids.addAll(values)
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

+11
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+
v1 = visitRex(node.v1, ctx),
186+
v2 = visitRex(node.v2, ctx),
187+
)
188+
189+
override fun visitRexOpCoalesce(node: Rex.Op.Coalesce, ctx: ProblemCallback) =
190+
org.partiql.plan.Rex.Op.Coalesce(
191+
values = node.values.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

+13-34
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 values = 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(values)
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 expr1 = visitExprCoerce(node.value, ctx)
456+
val expr2 = visitExprCoerce(node.nullifier, ctx)
457+
val op = rexOpNullif(expr1, expr2)
458+
return rex(type, op)
480459
}
481460

482461
/**

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

+42
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,46 @@ internal class PlanTyper(
760762
return rex(type, op)
761763
}
762764

765+
override fun visitRexOpCoalesce(node: Rex.Op.Coalesce, ctx: StaticType?): Rex {
766+
val values = node.values.map { visitRex(it, it.type) }.toMutableList()
767+
val typer = DynamicTyper()
768+
values.forEach { v ->
769+
typer.accumulate(v.type)
770+
}
771+
val (type, mapping) = typer.mapping()
772+
if (mapping != null) {
773+
assert(mapping.size == values.size)
774+
for (i in values.indices) {
775+
val (operand, target) = mapping[i]
776+
if (operand == target) continue // skip
777+
val cast = env.fnResolver.cast(operand, target)
778+
val rex = rex(type, rexOpCallStatic(fnResolved(cast), listOf(values[i])))
779+
values[i] = rex
780+
}
781+
}
782+
val op = rexOpCoalesce(values)
783+
return rex(type, op)
784+
}
785+
786+
// NULLIF(v1, v2)
787+
// ==
788+
// CASE
789+
// WHEN V1 = V2 THEN NULL -- WHEN branch always a boolean
790+
// ELSE V1
791+
// END
792+
// --> union(null, <type of V1>)
793+
override fun visitRexOpNullif(node: Rex.Op.Nullif, ctx: StaticType?): Rex {
794+
val v1 = visitRex(node.v1, node.v1.type)
795+
val v2 = visitRex(node.v2, node.v2.type)
796+
// reuse dynamic typing logic
797+
val typer = DynamicTyper()
798+
typer.accumulate(NULL)
799+
typer.accumulate(v1.type)
800+
val (type, _) = typer.mapping()
801+
val op = rexOpNullif(v1, v2)
802+
return rex(type, op)
803+
}
804+
763805
/**
764806
* In this context, Boolean means PartiQLValueType Bool, which can be nullable.
765807
* Hence, we permit Static Type BOOL, Static Type NULL, Static Type Missing here.

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

+9
Original file line numberDiff line numberDiff line change
@@ -158,6 +158,15 @@ rex::{
158158
],
159159
},
160160

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

0 commit comments

Comments
 (0)