Skip to content

Commit bb0342d

Browse files
alancai98alancai98
authored
Add EXCLUDE to evaluator; modify modeling of EXCLUDE in the plan (#1320)
1 parent da03142 commit bb0342d

File tree

18 files changed

+1518
-483
lines changed

18 files changed

+1518
-483
lines changed

‎partiql-eval/src/main/kotlin/org/partiql/eval/internal/Compiler.kt

+6
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package org.partiql.eval.internal
22

33
import org.partiql.eval.internal.operator.Operator
44
import org.partiql.eval.internal.operator.rel.RelDistinct
5+
import org.partiql.eval.internal.operator.rel.RelExclude
56
import org.partiql.eval.internal.operator.rel.RelFilter
67
import org.partiql.eval.internal.operator.rel.RelJoinInner
78
import org.partiql.eval.internal.operator.rel.RelJoinLeft
@@ -185,4 +186,9 @@ internal class Compiler(
185186
val condition = visitRex(node.predicate, ctx)
186187
return RelFilter(input, condition)
187188
}
189+
190+
override fun visitRelOpExclude(node: Rel.Op.Exclude, ctx: Unit): Operator {
191+
val input = visitRel(node.input, ctx)
192+
return RelExclude(input, node.paths)
193+
}
188194
}

‎partiql-eval/src/main/kotlin/org/partiql/eval/internal/operator/rel/RelExclude.kt

+189
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,189 @@
1+
package org.partiql.eval.internal.operator.rel
2+
3+
import org.partiql.eval.internal.Record
4+
import org.partiql.eval.internal.operator.Operator
5+
import org.partiql.plan.Rel
6+
import org.partiql.plan.relOpExcludeTypeCollIndex
7+
import org.partiql.plan.relOpExcludeTypeCollWildcard
8+
import org.partiql.plan.relOpExcludeTypeStructKey
9+
import org.partiql.plan.relOpExcludeTypeStructSymbol
10+
import org.partiql.plan.relOpExcludeTypeStructWildcard
11+
import org.partiql.value.BagValue
12+
import org.partiql.value.CollectionValue
13+
import org.partiql.value.ListValue
14+
import org.partiql.value.PartiQLValue
15+
import org.partiql.value.PartiQLValueExperimental
16+
import org.partiql.value.PartiQLValueType
17+
import org.partiql.value.SexpValue
18+
import org.partiql.value.StructValue
19+
import org.partiql.value.bagValue
20+
import org.partiql.value.listValue
21+
import org.partiql.value.sexpValue
22+
import org.partiql.value.structValue
23+
24+
internal class RelExclude(
25+
private val input: Operator.Relation,
26+
private val exclusions: List<Rel.Op.Exclude.Path>
27+
) : Operator.Relation {
28+
29+
override fun open() {
30+
input.open()
31+
}
32+
33+
@OptIn(PartiQLValueExperimental::class)
34+
override fun next(): Record? {
35+
val record = input.next() ?: return null
36+
exclusions.forEach { path ->
37+
val root = path.root.ref
38+
val value = record.values[root]
39+
record.values[root] = exclude(value, path.steps)
40+
}
41+
return record
42+
}
43+
44+
override fun close() {
45+
input.close()
46+
}
47+
48+
@OptIn(PartiQLValueExperimental::class)
49+
private fun exclude(
50+
structValue: StructValue<*>,
51+
exclusions: List<Rel.Op.Exclude.Step>
52+
): PartiQLValue {
53+
val structSymbolsToRemove = mutableSetOf<String>()
54+
val structKeysToRemove = mutableSetOf<String>() // keys stored as lowercase strings
55+
val branches = mutableMapOf<Rel.Op.Exclude.Type, List<Rel.Op.Exclude.Step>>()
56+
exclusions.forEach { exclusion ->
57+
when (exclusion.substeps.isEmpty()) {
58+
true -> {
59+
when (val leafType = exclusion.type) {
60+
is Rel.Op.Exclude.Type.StructWildcard -> {
61+
// struct wildcard at current level. return empty struct
62+
return structValue<PartiQLValue>()
63+
}
64+
is Rel.Op.Exclude.Type.StructSymbol -> structSymbolsToRemove.add(leafType.symbol)
65+
is Rel.Op.Exclude.Type.StructKey -> structKeysToRemove.add(leafType.key.lowercase())
66+
else -> { /* coll step; do nothing */ }
67+
}
68+
}
69+
false -> {
70+
when (exclusion.type) {
71+
is Rel.Op.Exclude.Type.StructWildcard, is Rel.Op.Exclude.Type.StructSymbol, is Rel.Op.Exclude.Type.StructKey -> branches[exclusion.type] =
72+
exclusion.substeps
73+
else -> { /* coll step; do nothing */ }
74+
}
75+
}
76+
}
77+
}
78+
val finalStruct = structValue.entries.mapNotNull { structField ->
79+
if (structSymbolsToRemove.contains(structField.first) || structKeysToRemove.contains(structField.first.lowercase())) {
80+
// struct attr is to be removed at current level
81+
null
82+
} else {
83+
// deeper level exclusions
84+
val name = structField.first
85+
var value = structField.second
86+
// apply struct key exclusions at deeper levels
87+
val structKey = relOpExcludeTypeStructKey(name)
88+
branches[structKey]?.let {
89+
value = exclude(value, it)
90+
}
91+
// apply struct symbol exclusions at deeper levels
92+
val structSymbol = relOpExcludeTypeStructSymbol(name)
93+
branches[structSymbol]?.let {
94+
value = exclude(value, it)
95+
}
96+
// apply struct wildcard exclusions at deeper levels
97+
val structWildcard = relOpExcludeTypeStructWildcard()
98+
branches[structWildcard]?.let {
99+
value = exclude(value, it)
100+
}
101+
Pair(name, value)
102+
}
103+
}
104+
return structValue(finalStruct)
105+
}
106+
107+
/**
108+
* Returns a [PartiQLValue] created from an iterable of [coll]. Requires [type] to be a collection type
109+
* (i.e. [PartiQLValueType.LIST], [PartiQLValueType.BAG], or [PartiQLValueType.SEXP]).
110+
*/
111+
@OptIn(PartiQLValueExperimental::class)
112+
private fun newCollValue(type: PartiQLValueType, coll: Iterable<PartiQLValue>): PartiQLValue {
113+
return when (type) {
114+
PartiQLValueType.LIST -> listValue(coll)
115+
PartiQLValueType.BAG -> bagValue(coll)
116+
PartiQLValueType.SEXP -> sexpValue(coll)
117+
else -> error("Collection type required")
118+
}
119+
}
120+
121+
@OptIn(PartiQLValueExperimental::class)
122+
private fun exclude(
123+
coll: CollectionValue<*>,
124+
type: PartiQLValueType,
125+
exclusions: List<Rel.Op.Exclude.Step>
126+
): PartiQLValue {
127+
val indexesToRemove = mutableSetOf<Int>()
128+
val branches = mutableMapOf<Rel.Op.Exclude.Type, List<Rel.Op.Exclude.Step>>()
129+
exclusions.forEach { exclusion ->
130+
when (exclusion.substeps.isEmpty()) {
131+
true -> {
132+
when (val leafType = exclusion.type) {
133+
is Rel.Op.Exclude.Type.CollWildcard -> {
134+
// collection wildcard at current level. return empty collection
135+
return newCollValue(type, emptyList())
136+
}
137+
is Rel.Op.Exclude.Type.CollIndex -> {
138+
indexesToRemove.add(leafType.index)
139+
}
140+
else -> { /* struct step; do nothing */ }
141+
}
142+
}
143+
false -> {
144+
when (exclusion.type) {
145+
is Rel.Op.Exclude.Type.CollWildcard, is Rel.Op.Exclude.Type.CollIndex -> branches[exclusion.type] =
146+
exclusion.substeps
147+
else -> { /* struct step; do nothing */ }
148+
}
149+
}
150+
}
151+
}
152+
val finalColl = coll.mapIndexedNotNull { index, element ->
153+
if (indexesToRemove.contains(index)) {
154+
// coll index is to be removed at current level
155+
null
156+
} else {
157+
// deeper level exclusions
158+
var value = element
159+
if (coll is ListValue || coll is SexpValue) {
160+
// apply collection index exclusions at deeper levels for lists and sexps
161+
val collIndex = relOpExcludeTypeCollIndex(index)
162+
branches[collIndex]?.let {
163+
value = exclude(element, it)
164+
}
165+
}
166+
// apply collection wildcard exclusions at deeper levels for lists, bags, and sexps
167+
val collWildcard = relOpExcludeTypeCollWildcard()
168+
branches[collWildcard]?.let {
169+
value = exclude(value, it)
170+
}
171+
value
172+
}
173+
}
174+
return newCollValue(type, finalColl)
175+
}
176+
177+
@OptIn(PartiQLValueExperimental::class)
178+
private fun exclude(initialPartiQLValue: PartiQLValue, exclusions: List<Rel.Op.Exclude.Step>): PartiQLValue {
179+
return when (initialPartiQLValue) {
180+
is StructValue<*> -> exclude(initialPartiQLValue, exclusions)
181+
is BagValue<*> -> exclude(initialPartiQLValue, PartiQLValueType.BAG, exclusions)
182+
is ListValue<*> -> exclude(initialPartiQLValue, PartiQLValueType.LIST, exclusions)
183+
is SexpValue<*> -> exclude(initialPartiQLValue, PartiQLValueType.SEXP, exclusions)
184+
else -> {
185+
initialPartiQLValue
186+
}
187+
}
188+
}
189+
}

‎partiql-eval/src/test/kotlin/org/partiql/eval/internal/PartiQLEngineDefaultTest.kt

+66
Original file line numberDiff line numberDiff line change
@@ -226,6 +226,72 @@ class PartiQLEngineDefaultTest {
226226
"c" to stringValue("z"),
227227
)
228228
),
229+
SuccessTestCase(
230+
input = """
231+
SELECT t
232+
EXCLUDE t.a.b
233+
FROM <<
234+
{'a': {'b': 2}, 'foo': 'bar', 'foo2': 'bar2'}
235+
>> AS t
236+
""".trimIndent(),
237+
expected = bagValue(
238+
structValue(
239+
"t" to structValue(
240+
"a" to structValue<PartiQLValue>(
241+
// field `b` excluded
242+
),
243+
"foo" to stringValue("bar"),
244+
"foo2" to stringValue("bar2")
245+
)
246+
),
247+
)
248+
),
249+
SuccessTestCase(
250+
input = """
251+
SELECT *
252+
EXCLUDE
253+
t.a.b.c[*].field_x
254+
FROM [{
255+
'a': {
256+
'b': {
257+
'c': [
258+
{ -- c[0]; field_x to be removed
259+
'field_x': 0,
260+
'field_y': 0
261+
},
262+
{ -- c[1]; field_x to be removed
263+
'field_x': 1,
264+
'field_y': 1
265+
},
266+
{ -- c[2]; field_x to be removed
267+
'field_x': 2,
268+
'field_y': 2
269+
}
270+
]
271+
}
272+
}
273+
}] AS t
274+
""".trimIndent(),
275+
expected = bagValue(
276+
structValue(
277+
"a" to structValue(
278+
"b" to structValue(
279+
"c" to bagValue( // TODO: should be ListValue; currently, Rex.ExprCollection doesn't return lists
280+
structValue(
281+
"field_y" to int32Value(0)
282+
),
283+
structValue(
284+
"field_y" to int32Value(1)
285+
),
286+
structValue(
287+
"field_y" to int32Value(2)
288+
)
289+
)
290+
)
291+
)
292+
)
293+
)
294+
)
229295
)
230296
}
231297
public class SuccessTestCase @OptIn(PartiQLValueExperimental::class) constructor(

0 commit comments

Comments
 (0)