Skip to content

Commit

Permalink
Updates modeling of set operations in plan
Browse files Browse the repository at this point in the history
  • Loading branch information
@johnedquinn
johnedquinn committed Apr 18, 2024
1 parent 103cd82 commit 85d1525
Show file tree
Hide file tree
Showing 8 changed files with 266 additions and 139 deletions.
30 changes: 22 additions & 8 deletions partiql-eval/src/main/kotlin/org/partiql/eval/internal/Compiler.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -261,16 +261,30 @@ internal class Compiler(
}
}

override fun visitRelOpSet(node: Rel.Op.Set, ctx: StaticType?): Operator {
override fun visitRelOpSetExcept(node: Rel.Op.Set.Except, ctx: StaticType?): Operator {
val lhs = visitRel(node.lhs, ctx)
val rhs = visitRel(node.rhs, ctx)
return when (node.type) {
Rel.Op.Set.Type.UNION_ALL -> RelUnionAll(lhs, rhs)
Rel.Op.Set.Type.UNION_DISTINCT -> RelUnionDistinct(lhs, rhs)
Rel.Op.Set.Type.INTERSECT_ALL -> RelIntersectAll(lhs, rhs)
Rel.Op.Set.Type.INTERSECT_DISTINCT -> RelIntersectDistinct(lhs, rhs)
Rel.Op.Set.Type.EXCEPT_ALL -> RelExceptAll(lhs, rhs)
Rel.Op.Set.Type.EXCEPT_DISTINCT -> RelExceptDistinct(lhs, rhs)
return when (node.quantifier) {
Rel.Op.Set.Quantifier.ALL -> RelExceptAll(lhs, rhs)
Rel.Op.Set.Quantifier.DISTINCT -> RelExceptDistinct(lhs, rhs)
}
}

override fun visitRelOpSetIntersect(node: Rel.Op.Set.Intersect, ctx: StaticType?): Operator {
val lhs = visitRel(node.lhs, ctx)
val rhs = visitRel(node.rhs, ctx)
return when (node.quantifier) {
Rel.Op.Set.Quantifier.ALL -> RelIntersectAll(lhs, rhs)
Rel.Op.Set.Quantifier.DISTINCT -> RelIntersectDistinct(lhs, rhs)
}
}

override fun visitRelOpSetUnion(node: Rel.Op.Set.Union, ctx: StaticType?): Operator {
val lhs = visitRel(node.lhs, ctx)
val rhs = visitRel(node.rhs, ctx)
return when (node.quantifier) {
Rel.Op.Set.Quantifier.ALL -> RelUnionAll(lhs, rhs)
Rel.Op.Set.Quantifier.DISTINCT -> RelUnionDistinct(lhs, rhs)
}
}

Expand Down
35 changes: 22 additions & 13 deletions partiql-plan/src/main/resources/partiql_plan.ion
View file
Original file line number Diff line number Diff line change
Expand Up @@ -261,20 +261,29 @@ rel::{
],
},

// From Substrait: "The set operation encompasses several set-level operations that support combining datasets, ...
// possibly excluding records based on various types of record level matching."
set::{
lhs: rel,
rhs: rel,
type: [
UNION_ALL,
UNION_DISTINCT,
INTERSECT_ALL,
INTERSECT_DISTINCT,
EXCEPT_ALL,
EXCEPT_DISTINCT
set::[
union::{
quantifier: quantifier,
lhs: rel,
rhs: rel,
},

intersect::{
quantifier: quantifier,
lhs: rel,
rhs: rel,
},

except::{
quantifier: quantifier,
lhs: rel,
rhs: rel,
},

_::[
quantifier::[ ALL, DISTINCT ],
]
},
],

limit::{
input: rel,
Expand Down
97 changes: 73 additions & 24 deletions partiql-planner/src/main/kotlin/org/partiql/planner/internal/ir/Nodes.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,9 @@ import org.partiql.planner.internal.ir.builder.RelOpOffsetBuilder
import org.partiql.planner.internal.ir.builder.RelOpProjectBuilder
import org.partiql.planner.internal.ir.builder.RelOpScanBuilder
import org.partiql.planner.internal.ir.builder.RelOpScanIndexedBuilder
import org.partiql.planner.internal.ir.builder.RelOpSetBuilder
import org.partiql.planner.internal.ir.builder.RelOpSetExceptBuilder
import org.partiql.planner.internal.ir.builder.RelOpSetIntersectBuilder
import org.partiql.planner.internal.ir.builder.RelOpSetUnionBuilder
import org.partiql.planner.internal.ir.builder.RelOpSortBuilder
import org.partiql.planner.internal.ir.builder.RelOpSortSpecBuilder
import org.partiql.planner.internal.ir.builder.RelOpUnpivotBuilder
Expand Down Expand Up @@ -1016,35 +1018,82 @@ internal data class Rel(
internal fun builder(): RelOpSortBuilder = RelOpSortBuilder()
}
}
internal sealed class Set : Op() {

internal data class Set(
@JvmField internal val lhs: Rel,
@JvmField internal val rhs: Rel,
@JvmField internal val type: Type,
@JvmField internal val isOuter: Boolean
) : Op() {
public override val children: List<PlanNode> by lazy {
val kids = mutableListOf<PlanNode?>()
kids.add(lhs)
kids.add(rhs)
kids.filterNotNull()
public override fun <R, C> accept(visitor: PlanVisitor<R, C>, ctx: C): R = when (this) {
is Union -> visitor.visitRelOpSetUnion(this, ctx)
is Intersect -> visitor.visitRelOpSetIntersect(this, ctx)
is Except -> visitor.visitRelOpSetExcept(this, ctx)
}

public override fun <R, C> accept(visitor: PlanVisitor<R, C>, ctx: C): R =
visitor.visitRelOpSet(this, ctx)
internal data class Union(
@JvmField internal val quantifier: Quantifier,
@JvmField internal val lhs: Rel,
@JvmField internal val rhs: Rel,
@JvmField internal val isOuter: Boolean,
) : Set() {
public override val children: List<PlanNode> by lazy {
val kids = mutableListOf<PlanNode?>()
kids.add(lhs)
kids.add(rhs)
kids.filterNotNull()
}

internal companion object {
@JvmStatic
internal fun builder(): RelOpSetBuilder = RelOpSetBuilder()
public override fun <R, C> accept(visitor: PlanVisitor<R, C>, ctx: C): R =
visitor.visitRelOpSetUnion(this, ctx)

internal companion object {
@JvmStatic
internal fun builder(): RelOpSetUnionBuilder = RelOpSetUnionBuilder()
}
}

internal enum class Type {
UNION_ALL,
UNION_DISTINCT,
INTERSECT_ALL,
INTERSECT_DISTINCT,
EXCEPT_ALL,
EXCEPT_DISTINCT
internal data class Intersect(
@JvmField internal val quantifier: Quantifier,
@JvmField internal val lhs: Rel,
@JvmField internal val rhs: Rel,
@JvmField internal val isOuter: Boolean,
) : Set() {
public override val children: List<PlanNode> by lazy {
val kids = mutableListOf<PlanNode?>()
kids.add(lhs)
kids.add(rhs)
kids.filterNotNull()
}

public override fun <R, C> accept(visitor: PlanVisitor<R, C>, ctx: C): R =
visitor.visitRelOpSetIntersect(this, ctx)

internal companion object {
@JvmStatic
internal fun builder(): RelOpSetIntersectBuilder = RelOpSetIntersectBuilder()
}
}

internal data class Except(
@JvmField internal val quantifier: Quantifier,
@JvmField internal val lhs: Rel,
@JvmField internal val rhs: Rel,
@JvmField internal val isOuter: Boolean,
) : Set() {
public override val children: List<PlanNode> by lazy {
val kids = mutableListOf<PlanNode?>()
kids.add(lhs)
kids.add(rhs)
kids.filterNotNull()
}

public override fun <R, C> accept(visitor: PlanVisitor<R, C>, ctx: C): R =
visitor.visitRelOpSetExcept(this, ctx)

internal companion object {
@JvmStatic
internal fun builder(): RelOpSetExceptBuilder = RelOpSetExceptBuilder()
}
}

internal enum class Quantifier {
ALL, DISTINCT
}
}

Expand Down
28 changes: 19 additions & 9 deletions partiql-planner/src/main/kotlin/org/partiql/planner/internal/transforms/PlanTransform.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -328,19 +328,29 @@ internal class PlanTransform(
}
)

override fun visitRelOpSet(node: Rel.Op.Set, ctx: Unit) = org.partiql.plan.Rel.Op.Set(
override fun visitRelOpSetExcept(node: Rel.Op.Set.Except, ctx: Unit) = org.partiql.plan.Rel.Op.Set.Except(
lhs = visitRel(node.lhs, ctx),
rhs = visitRel(node.rhs, ctx),
type = when (node.type) {
Rel.Op.Set.Type.UNION_ALL -> org.partiql.plan.Rel.Op.Set.Type.UNION_ALL
Rel.Op.Set.Type.UNION_DISTINCT -> org.partiql.plan.Rel.Op.Set.Type.UNION_DISTINCT
Rel.Op.Set.Type.EXCEPT_ALL -> org.partiql.plan.Rel.Op.Set.Type.EXCEPT_ALL
Rel.Op.Set.Type.EXCEPT_DISTINCT -> org.partiql.plan.Rel.Op.Set.Type.EXCEPT_DISTINCT
Rel.Op.Set.Type.INTERSECT_ALL -> org.partiql.plan.Rel.Op.Set.Type.INTERSECT_ALL
Rel.Op.Set.Type.INTERSECT_DISTINCT -> org.partiql.plan.Rel.Op.Set.Type.INTERSECT_DISTINCT
}
quantifier = visitRelOpSetQuantifier(node.quantifier)
)

override fun visitRelOpSetIntersect(node: Rel.Op.Set.Intersect, ctx: Unit) = org.partiql.plan.Rel.Op.Set.Intersect(
lhs = visitRel(node.lhs, ctx),
rhs = visitRel(node.rhs, ctx),
quantifier = visitRelOpSetQuantifier(node.quantifier)
)

override fun visitRelOpSetUnion(node: Rel.Op.Set.Union, ctx: Unit) = org.partiql.plan.Rel.Op.Set.Union(
lhs = visitRel(node.lhs, ctx),
rhs = visitRel(node.rhs, ctx),
quantifier = visitRelOpSetQuantifier(node.quantifier)
)

private fun visitRelOpSetQuantifier(node: Rel.Op.Set.Quantifier) = when (node) {
Rel.Op.Set.Quantifier.ALL -> org.partiql.plan.Rel.Op.Set.Quantifier.ALL
Rel.Op.Set.Quantifier.DISTINCT -> org.partiql.plan.Rel.Op.Set.Quantifier.DISTINCT
}

override fun visitRelOpLimit(node: Rel.Op.Limit, ctx: Unit) = org.partiql.plan.Rel.Op.Limit(
input = visitRel(node.input, ctx),
limit = visitRex(node.limit, ctx),
Expand Down
26 changes: 8 additions & 18 deletions partiql-planner/src/main/kotlin/org/partiql/planner/internal/transforms/RelConverter.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -432,9 +432,6 @@ internal object RelConverter {

/**
* Append SQL set operator if present
*
* TODO combine/compare schemas
* TODO set quantifier
*/
private fun convertSetOp(input: Rel, setOp: Expr.SFW.SetOp?): Rel {
if (setOp == null) {
Expand All @@ -443,22 +440,15 @@ internal object RelConverter {
val type = input.type.copy(props = emptySet())
val lhs = input
val rhs = visitExprSFW(setOp.operand, nil)

val setType = when (setOp.type.type) {
SetOp.Type.UNION -> when (setOp.type.setq) {
SetQuantifier.ALL -> Rel.Op.Set.Type.UNION_ALL
null, SetQuantifier.DISTINCT -> Rel.Op.Set.Type.UNION_DISTINCT
}
SetOp.Type.EXCEPT -> when (setOp.type.setq) {
SetQuantifier.ALL -> Rel.Op.Set.Type.EXCEPT_ALL
null, SetQuantifier.DISTINCT -> Rel.Op.Set.Type.EXCEPT_DISTINCT
}
SetOp.Type.INTERSECT -> when (setOp.type.setq) {
SetQuantifier.ALL -> Rel.Op.Set.Type.INTERSECT_ALL
null, SetQuantifier.DISTINCT -> Rel.Op.Set.Type.INTERSECT_DISTINCT
}
val quantifier = when (setOp.type.setq) {
SetQuantifier.ALL -> Rel.Op.Set.Quantifier.ALL
null, SetQuantifier.DISTINCT -> Rel.Op.Set.Quantifier.DISTINCT
}
val op = when (setOp.type.type) {
SetOp.Type.UNION -> Rel.Op.Set.Union(quantifier, lhs, rhs, false)
SetOp.Type.EXCEPT -> Rel.Op.Set.Except(quantifier, lhs, rhs, false)
SetOp.Type.INTERSECT -> Rel.Op.Set.Intersect(quantifier, lhs, rhs, false)
}
val op = Rel.Op.Set(lhs, rhs, setType, isOuter = false)
return rel(type, op)
}

Expand Down
23 changes: 9 additions & 14 deletions partiql-planner/src/main/kotlin/org/partiql/planner/internal/transforms/RexConverter.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -832,21 +832,16 @@ internal object RexConverter {
type = Rel.Type(listOf(Rel.Binding("_1", StaticType.ANY)), props = emptySet()),
op = Rel.Op.Scan(visitExpr(node.rhs, ctx))
)
val type = when (node.type.type) {
SetOp.Type.UNION -> when (node.type.setq) {
SetQuantifier.ALL -> Rel.Op.Set.Type.UNION_ALL
null, SetQuantifier.DISTINCT -> Rel.Op.Set.Type.UNION_DISTINCT
}
SetOp.Type.EXCEPT -> when (node.type.setq) {
SetQuantifier.ALL -> Rel.Op.Set.Type.EXCEPT_ALL
null, SetQuantifier.DISTINCT -> Rel.Op.Set.Type.EXCEPT_DISTINCT
}
SetOp.Type.INTERSECT -> when (node.type.setq) {
SetQuantifier.ALL -> Rel.Op.Set.Type.INTERSECT_ALL
null, SetQuantifier.DISTINCT -> Rel.Op.Set.Type.INTERSECT_DISTINCT
}
val quantifier = when (node.type.setq) {
SetQuantifier.ALL -> Rel.Op.Set.Quantifier.ALL
null, SetQuantifier.DISTINCT -> Rel.Op.Set.Quantifier.DISTINCT
}
val isOuter = node.outer == true
val op = when (node.type.type) {
SetOp.Type.UNION -> Rel.Op.Set.Union(quantifier, lhs, rhs, isOuter)
SetOp.Type.EXCEPT -> Rel.Op.Set.Except(quantifier, lhs, rhs, isOuter)
SetOp.Type.INTERSECT -> Rel.Op.Set.Intersect(quantifier, lhs, rhs, isOuter)
}
val op = Rel.Op.Set(lhs, rhs, type, isOuter = node.outer == true)
val rel = Rel(
type = Rel.Type(listOf(Rel.Binding("_0", StaticType.ANY)), props = emptySet()),
op = op
Expand Down
Loading

0 comments on commit 85d1525

Please sign in to comment.