Skip to content

Commit

Permalink
Move around top-level fns; make EXCLUDE input private
Browse files Browse the repository at this point in the history
  • Loading branch information
alancai98 committed Dec 21, 2023
1 parent 5bc681e commit 9e9a040
Show file tree
Hide file tree
Showing 3 changed files with 156 additions and 154 deletions.
Original file line number Diff line number Diff line change
@@ -1,5 +1,8 @@
package org.partiql.eval.internal

import org.partiql.eval.internal.exclude.CompiledExcludeItem
import org.partiql.eval.internal.exclude.ExcludeFieldCase
import org.partiql.eval.internal.exclude.ExcludeStep
import org.partiql.eval.internal.operator.Operator
import org.partiql.eval.internal.operator.rel.RelDistinct
import org.partiql.eval.internal.operator.rel.RelExclude
Expand All @@ -11,7 +14,6 @@ import org.partiql.eval.internal.operator.rel.RelJoinRight
import org.partiql.eval.internal.operator.rel.RelProject
import org.partiql.eval.internal.operator.rel.RelScan
import org.partiql.eval.internal.operator.rel.RelScanIndexed
import org.partiql.eval.internal.operator.rel.compileExcludeItems
import org.partiql.eval.internal.operator.rex.ExprCase
import org.partiql.eval.internal.operator.rex.ExprCollection
import org.partiql.eval.internal.operator.rex.ExprGlobal
Expand All @@ -24,6 +26,7 @@ import org.partiql.eval.internal.operator.rex.ExprSelect
import org.partiql.eval.internal.operator.rex.ExprStruct
import org.partiql.eval.internal.operator.rex.ExprTupleUnion
import org.partiql.eval.internal.operator.rex.ExprVar
import org.partiql.plan.Identifier
import org.partiql.plan.PartiQLPlan
import org.partiql.plan.PlanNode
import org.partiql.plan.Rel
Expand Down Expand Up @@ -188,6 +191,38 @@ internal class Compiler(
return RelFilter(input, condition)
}

/**
* Creates a list of [CompiledExcludeItem] with each index of the resulting list corresponding to a different
* exclude path root.
*/
internal fun compileExcludeItems(excludeExprs: List<Rel.Op.Exclude.Item>): List<CompiledExcludeItem> {
val compiledExcludeItems = excludeExprs
.groupBy { it.root }
.map { (root, exclusions) ->
exclusions.fold(CompiledExcludeItem.empty(root.ref)) { acc, exclusion ->
acc.addNode(exclusion.steps.map { it.toCompiledExcludeStep() })
acc
}
}
return compiledExcludeItems
}

private fun Rel.Op.Exclude.Step.toCompiledExcludeStep(): ExcludeStep {
return when (this) {
is Rel.Op.Exclude.Step.StructField -> ExcludeStep.StructField(this.symbol.symbol, this.symbol.caseSensitivity.toCompiledExcludeStepCase())
is Rel.Op.Exclude.Step.StructWildcard -> ExcludeStep.StructWildcard
is Rel.Op.Exclude.Step.CollIndex -> ExcludeStep.CollIndex(this.index)
is Rel.Op.Exclude.Step.CollWildcard -> ExcludeStep.CollWildcard
}
}

private fun Identifier.CaseSensitivity.toCompiledExcludeStepCase(): ExcludeFieldCase {
return when (this) {
Identifier.CaseSensitivity.SENSITIVE -> ExcludeFieldCase.SENSITIVE
Identifier.CaseSensitivity.INSENSITIVE -> ExcludeFieldCase.INSENSITIVE
}
}

override fun visitRelOpExclude(node: Rel.Op.Exclude, ctx: Unit): Operator {
val input = visitRel(node.input, ctx)
val compiledExcludeExprs = compileExcludeItems(node.items)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -6,8 +6,6 @@ import org.partiql.eval.internal.exclude.ExcludeFieldCase
import org.partiql.eval.internal.exclude.ExcludeNode
import org.partiql.eval.internal.exclude.ExcludeStep
import org.partiql.eval.internal.operator.Operator
import org.partiql.plan.Identifier
import org.partiql.plan.Rel
import org.partiql.value.BagValue
import org.partiql.value.CollectionValue
import org.partiql.value.ListValue
Expand All @@ -22,8 +20,8 @@ import org.partiql.value.sexpValue
import org.partiql.value.structValue

internal class RelExclude(
val input: Operator.Relation,
private val compiledExcludeItems: List<CompiledExcludeItem>
private val input: Operator.Relation,
private val exclusions: List<CompiledExcludeItem>
) : Operator.Relation {

override fun open() {
Expand All @@ -33,7 +31,7 @@ internal class RelExclude(
override fun next(): Record? {
while (true) {
val row = input.next() ?: return null
val newRecord = compiledExcludeItems.fold(row) { curRecord, expr ->
val newRecord = exclusions.fold(row) { curRecord, expr ->
excludeOnRecord(curRecord, expr)
}
return newRecord
Expand All @@ -43,170 +41,138 @@ internal class RelExclude(
override fun close() {
input.close()
}
}

@OptIn(PartiQLValueExperimental::class)
private fun excludeOnRecord(
record: Record,
exclusions: CompiledExcludeItem
): Record {
val values = record.values
val value = values.getOrNull(exclusions.root)
val newValues = if (value != null) {
values[exclusions.root] = excludeOnPartiQLValue(value, exclusions)
values
} else {
values
}
return Record(newValues)
}

@OptIn(PartiQLValueExperimental::class)
private fun excludeOnStructValue(
structValue: StructValue<*>,
exclusions: ExcludeNode
): PartiQLValue {
val leavesSteps = exclusions.leaves.map { leaf -> leaf.step }
val branches = exclusions.branches
if (leavesSteps.any { it is ExcludeStep.StructWildcard }) {
// tuple wildcard at current level. return empty struct
return structValue<PartiQLValue>()
}
val attrsToRemove = leavesSteps.filterIsInstance<ExcludeStep.StructField>()
.map { it.attr }
.toSet()
val entriesWithRemoved = structValue.entries.filter { structField ->
!attrsToRemove.contains(structField.first)
}
val finalStruct = entriesWithRemoved.map { structField ->
val name = structField.first
var expr = structField.second
// apply case-sensitive tuple attr exclusions
val structFieldCaseSensitiveKey = ExcludeStep.StructField(name, ExcludeFieldCase.SENSITIVE)
branches.find {
it.step == structFieldCaseSensitiveKey
}?.let {
expr = excludeOnPartiQLValue(expr, it)
}
// apply case-insensitive tuple attr exclusions
val structFieldCaseInsensitiveKey = ExcludeStep.StructField(name, ExcludeFieldCase.INSENSITIVE)
branches.find {
it.step == structFieldCaseInsensitiveKey
}?.let {
expr = excludeOnPartiQLValue(expr, it)
}
// apply tuple wildcard exclusions
val tupleWildcardKey = ExcludeStep.StructWildcard
branches.find {
it.step == tupleWildcardKey
}?.let {
expr = excludeOnPartiQLValue(expr, it)
@OptIn(PartiQLValueExperimental::class)
private fun excludeOnRecord(
record: Record,
exclusions: CompiledExcludeItem
): Record {
val values = record.values
val value = values.getOrNull(exclusions.root)
val newValues = if (value != null) {
values[exclusions.root] = excludeOnPartiQLValue(value, exclusions)
values
} else {
values
}
Pair(name, expr)
return Record(newValues)
}
return structValue(finalStruct)
}

/**
* Returns a [PartiQLValue] created from an iterable of [coll]. Requires [type] to be a collection type
* (i.e. [PartiQLValueType.LIST], [PartiQLValueType.BAG], or [PartiQLValueType.SEXP]).
*/
@OptIn(PartiQLValueExperimental::class)
private fun newCollValue(type: PartiQLValueType, coll: Iterable<PartiQLValue>): PartiQLValue {
return when (type) {
PartiQLValueType.LIST -> listValue(coll)
PartiQLValueType.BAG -> bagValue(coll)
PartiQLValueType.SEXP -> sexpValue(coll)
else -> error("Collection type required")
}
}

@OptIn(PartiQLValueExperimental::class)
private fun excludeOnCollValue(
coll: CollectionValue<*>,
type: PartiQLValueType,
exclusions: ExcludeNode
): PartiQLValue {
val leavesSteps = exclusions.leaves.map { leaf -> leaf.step }
val branches = exclusions.branches
if (leavesSteps.any { it is ExcludeStep.CollWildcard }) {
// collection wildcard at current level. return empty collection
return newCollValue(type, emptyList())
} else {
val indexesToRemove = leavesSteps.filterIsInstance<ExcludeStep.CollIndex>()
.map { it.index }
@OptIn(PartiQLValueExperimental::class)
private fun excludeOnStructValue(
structValue: StructValue<*>,
exclusions: ExcludeNode
): PartiQLValue {
val leavesSteps = exclusions.leaves.map { leaf -> leaf.step }
val branches = exclusions.branches
if (leavesSteps.any { it is ExcludeStep.StructWildcard }) {
// tuple wildcard at current level. return empty struct
return structValue<PartiQLValue>()
}
val attrsToRemove = leavesSteps.filterIsInstance<ExcludeStep.StructField>()
.map { it.attr }
.toSet()
val collWithRemoved = when (coll) {
is BagValue -> coll
is ListValue, is SexpValue -> coll.filterIndexed { index, _ ->
!indexesToRemove.contains(index)
}
val entriesWithRemoved = structValue.entries.filter { structField ->
!attrsToRemove.contains(structField.first)
}
val finalColl = collWithRemoved.mapIndexed { index, element ->
var expr = element
if (coll is ListValue || coll is SexpValue) {
// apply collection index exclusions for lists and sexps
val elementKey = ExcludeStep.CollIndex(index)
branches.find {
it.step == elementKey
}?.let {
expr = excludeOnPartiQLValue(element, it)
}
val finalStruct = entriesWithRemoved.map { structField ->
val name = structField.first
var expr = structField.second
// apply case-sensitive tuple attr exclusions
val structFieldCaseSensitiveKey = ExcludeStep.StructField(name, ExcludeFieldCase.SENSITIVE)
branches.find {
it.step == structFieldCaseSensitiveKey
}?.let {
expr = excludeOnPartiQLValue(expr, it)
}
// apply case-insensitive tuple attr exclusions
val structFieldCaseInsensitiveKey = ExcludeStep.StructField(name, ExcludeFieldCase.INSENSITIVE)
branches.find {
it.step == structFieldCaseInsensitiveKey
}?.let {
expr = excludeOnPartiQLValue(expr, it)
}
// apply collection wildcard exclusions for lists, bags, and sexps
val collectionWildcardKey = ExcludeStep.CollWildcard
// apply tuple wildcard exclusions
val tupleWildcardKey = ExcludeStep.StructWildcard
branches.find {
it.step == collectionWildcardKey
it.step == tupleWildcardKey
}?.let {
expr = excludeOnPartiQLValue(expr, it)
}
expr
Pair(name, expr)
}
return newCollValue(type, finalColl)
return structValue(finalStruct)
}
}

@OptIn(PartiQLValueExperimental::class)
private fun excludeOnPartiQLValue(initialPartiQLValue: PartiQLValue, exclusions: ExcludeNode): PartiQLValue {
return when (initialPartiQLValue) {
is StructValue<*> -> excludeOnStructValue(initialPartiQLValue, exclusions)
is BagValue<*> -> excludeOnCollValue(initialPartiQLValue, PartiQLValueType.BAG, exclusions)
is ListValue<*> -> excludeOnCollValue(initialPartiQLValue, PartiQLValueType.LIST, exclusions)
is SexpValue<*> -> excludeOnCollValue(initialPartiQLValue, PartiQLValueType.SEXP, exclusions)
else -> {
initialPartiQLValue
/**
* Returns a [PartiQLValue] created from an iterable of [coll]. Requires [type] to be a collection type
* (i.e. [PartiQLValueType.LIST], [PartiQLValueType.BAG], or [PartiQLValueType.SEXP]).
*/
@OptIn(PartiQLValueExperimental::class)
private fun newCollValue(type: PartiQLValueType, coll: Iterable<PartiQLValue>): PartiQLValue {
return when (type) {
PartiQLValueType.LIST -> listValue(coll)
PartiQLValueType.BAG -> bagValue(coll)
PartiQLValueType.SEXP -> sexpValue(coll)
else -> error("Collection type required")
}
}
}

/**
* Creates a list of [CompiledExcludeItem] with each index of the resulting list corresponding to a different
* exclude path root.
*/
internal fun compileExcludeItems(excludeExprs: List<Rel.Op.Exclude.Item>): List<CompiledExcludeItem> {
val compiledExcludeItems = excludeExprs
.groupBy { it.root }
.map { (root, exclusions) ->
exclusions.fold(CompiledExcludeItem.empty(root.ref)) { acc, exclusion ->
acc.addNode(exclusion.steps.map { it.toCompiledExcludeStep() })
acc
@OptIn(PartiQLValueExperimental::class)
private fun excludeOnCollValue(
coll: CollectionValue<*>,
type: PartiQLValueType,
exclusions: ExcludeNode
): PartiQLValue {
val leavesSteps = exclusions.leaves.map { leaf -> leaf.step }
val branches = exclusions.branches
if (leavesSteps.any { it is ExcludeStep.CollWildcard }) {
// collection wildcard at current level. return empty collection
return newCollValue(type, emptyList())
} else {
val indexesToRemove = leavesSteps.filterIsInstance<ExcludeStep.CollIndex>()
.map { it.index }
.toSet()
val collWithRemoved = when (coll) {
is BagValue -> coll
is ListValue, is SexpValue -> coll.filterIndexed { index, _ ->
!indexesToRemove.contains(index)
}
}
val finalColl = collWithRemoved.mapIndexed { index, element ->
var expr = element
if (coll is ListValue || coll is SexpValue) {
// apply collection index exclusions for lists and sexps
val elementKey = ExcludeStep.CollIndex(index)
branches.find {
it.step == elementKey
}?.let {
expr = excludeOnPartiQLValue(element, it)
}
}
// apply collection wildcard exclusions for lists, bags, and sexps
val collectionWildcardKey = ExcludeStep.CollWildcard
branches.find {
it.step == collectionWildcardKey
}?.let {
expr = excludeOnPartiQLValue(expr, it)
}
expr
}
return newCollValue(type, finalColl)
}
return compiledExcludeItems
}

private fun Rel.Op.Exclude.Step.toCompiledExcludeStep(): ExcludeStep {
return when (this) {
is Rel.Op.Exclude.Step.StructField -> ExcludeStep.StructField(this.symbol.symbol, this.symbol.caseSensitivity.toCompiledExcludeStepCase())
is Rel.Op.Exclude.Step.StructWildcard -> ExcludeStep.StructWildcard
is Rel.Op.Exclude.Step.CollIndex -> ExcludeStep.CollIndex(this.index)
is Rel.Op.Exclude.Step.CollWildcard -> ExcludeStep.CollWildcard
}
}

private fun Identifier.CaseSensitivity.toCompiledExcludeStepCase(): ExcludeFieldCase {
return when (this) {
Identifier.CaseSensitivity.SENSITIVE -> ExcludeFieldCase.SENSITIVE
Identifier.CaseSensitivity.INSENSITIVE -> ExcludeFieldCase.INSENSITIVE
@OptIn(PartiQLValueExperimental::class)
private fun excludeOnPartiQLValue(initialPartiQLValue: PartiQLValue, exclusions: ExcludeNode): PartiQLValue {
return when (initialPartiQLValue) {
is StructValue<*> -> excludeOnStructValue(initialPartiQLValue, exclusions)
is BagValue<*> -> excludeOnCollValue(initialPartiQLValue, PartiQLValueType.BAG, exclusions)
is ListValue<*> -> excludeOnCollValue(initialPartiQLValue, PartiQLValueType.LIST, exclusions)
is SexpValue<*> -> excludeOnCollValue(initialPartiQLValue, PartiQLValueType.SEXP, exclusions)
else -> {
initialPartiQLValue
}
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ package org.partiql.eval.internal.exclude
import org.junit.Assert.assertEquals
import org.junit.jupiter.params.ParameterizedTest
import org.junit.jupiter.params.provider.ArgumentsSource
import org.partiql.eval.internal.operator.rel.compileExcludeItems
import org.partiql.eval.internal.Compiler
import org.partiql.lang.util.ArgumentsProviderBase
import org.partiql.parser.PartiQLParser
import org.partiql.plan.Rel
Expand All @@ -24,9 +24,10 @@ class CompiledExcludeItemTest {
private fun testExcludeExprSubsumption(tc: SubsumptionTC) {
val statement = parser.parse("SELECT * EXCLUDE ${tc.excludeExprStr} FROM <<>> AS s, <<>> AS t;").root
val session = PartiQLPlanner.Session("q", "u")
val plan = planner.plan(statement, session)
val excludeClause = getExcludeClause(plan.plan.statement)
val actualExcludeExprs = compileExcludeItems(excludeClause.items)
val plan = planner.plan(statement, session).plan
val compiler = Compiler(plan, emptyMap())
val excludeClause = getExcludeClause(plan.statement)
val actualExcludeExprs = compiler.compileExcludeItems(excludeClause.items)
assertEquals(tc.expectedExcludeExprs, actualExcludeExprs)
}

Expand Down

0 comments on commit 9e9a040

Please sign in to comment.