Skip to content

Add COALESCE and NULLIF to the logical plan #1404

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Mar 29, 2024
Merged

Add COALESCE and NULLIF to the logical plan #1404

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
808bbf7
Add NULLIF and COALESCE to the logical plan
alancai98 Mar 28, 2024
b30605d
Merge branch 'main' into add-coalesce-nullif-to-plan
alancai98 Mar 29, 2024
0ce27b7
Align plan modeling of NULLIF to ast; fix expansion of Any in dynamic…
alancai98 Mar 29, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion CHANGELOG.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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]

Expand Down
9 changes: 9 additions & 0 deletions partiql-plan/src/main/resources/partiql_plan.ion
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -135,6 +135,15 @@ rex::{
],
},

nullif::{
v1: rex,
v2: rex
},

coalesce::{
args: list::[rex]
},

collection::{
values: list::[rex],
},
Expand Down
45 changes: 45 additions & 0 deletions partiql-planner/src/main/kotlin/org/partiql/planner/internal/ir/Nodes.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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 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() {
Expand Down
11 changes: 11 additions & 0 deletions partiql-planner/src/main/kotlin/org/partiql/planner/internal/transforms/PlanTransform.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
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)
Expand Down
47 changes: 13 additions & 34 deletions partiql-planner/src/main/kotlin/org/partiql/planner/internal/transforms/RexConverter.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
}
}
Expand Down Expand Up @@ -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 expr1 = visitExprCoerce(node.value, ctx)
val expr2 = visitExprCoerce(node.nullifier, ctx)
val op = rexOpNullif(expr1, expr2)
return rex(type, op)
}

/**
Expand Down
50 changes: 50 additions & 0 deletions partiql-planner/src/main/kotlin/org/partiql/planner/internal/typer/PlanTyper.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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 v1 = visitRex(node.v1, node.v1.type)
val v2 = visitRex(node.v2, node.v2.type)
val typer = DynamicTyper()
typer.accumulate(NULL)
typer.accumulate(v1.type)
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Only add v1's type here. I found that posgresql actually does something different here. Per their docs:

The result has the same type as the first argument — but there is a subtlety. What is actually returned is the first argument of the implied = operator, and in some cases that will have been promoted to match the second argument's type. For example, NULLIF(1, 2.2) yields numeric, because there is no integer = numeric operator, only numeric = numeric.

Per my read of SQL-99, I can't find any mention of this promotion behavior, so leaving the result type as just the minimal common supertype of (<type v1>, null). Postgresql would do something like the minimal common supertype of (<type v1>, <type v2>, null).

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I believe we can skip the dynamic typer here. It's just type of union(v1, null).flatten() or something as simple as that

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah I'll simplify it. I initially included the dynamic typer in case we wanted to follow what postgresql did, which is easy to do w/ the dynamic typer.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Followup from offline discussion, will keep the dynamic typer here to ensure consistency with the typing of CASE WHEN and COALESCE.

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.
Expand Down
9 changes: 9 additions & 0 deletions partiql-planner/src/main/resources/partiql_plan_internal.ion
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -158,6 +158,15 @@ rex::{
],
},

nullif::{
v1: rex,
v2: rex
},

coalesce::{
args: list::[rex]
},

collection::{
values: list::[rex],
},
Expand Down
Loading