Skip to content

Commit 50de032

Browse files
alancai98alancai98
committed
Reduce repeated looping; additional test for EXCLUDE
1 parent 9e9a040 commit 50de032

File tree

2 files changed

+104
-52
lines changed

2 files changed

+104
-52
lines changed

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

Lines changed: 58 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -63,43 +63,48 @@ internal class RelExclude(
6363
structValue: StructValue<*>,
6464
exclusions: ExcludeNode
6565
): PartiQLValue {
66-
val leavesSteps = exclusions.leaves.map { leaf -> leaf.step }
67-
val branches = exclusions.branches
68-
if (leavesSteps.any { it is ExcludeStep.StructWildcard }) {
69-
// tuple wildcard at current level. return empty struct
70-
return structValue<PartiQLValue>()
71-
}
72-
val attrsToRemove = leavesSteps.filterIsInstance<ExcludeStep.StructField>()
73-
.map { it.attr }
74-
.toSet()
75-
val entriesWithRemoved = structValue.entries.filter { structField ->
76-
!attrsToRemove.contains(structField.first)
77-
}
78-
val finalStruct = entriesWithRemoved.map { structField ->
79-
val name = structField.first
80-
var expr = structField.second
81-
// apply case-sensitive tuple attr exclusions
82-
val structFieldCaseSensitiveKey = ExcludeStep.StructField(name, ExcludeFieldCase.SENSITIVE)
83-
branches.find {
84-
it.step == structFieldCaseSensitiveKey
85-
}?.let {
86-
expr = excludeOnPartiQLValue(expr, it)
87-
}
88-
// apply case-insensitive tuple attr exclusions
89-
val structFieldCaseInsensitiveKey = ExcludeStep.StructField(name, ExcludeFieldCase.INSENSITIVE)
90-
branches.find {
91-
it.step == structFieldCaseInsensitiveKey
92-
}?.let {
93-
expr = excludeOnPartiQLValue(expr, it)
66+
val attrsToRemove = exclusions.leaves.mapNotNull { leaf ->
67+
when (val leafStep = leaf.step) {
68+
is ExcludeStep.StructWildcard -> {
69+
// tuple wildcard at current level. return empty struct
70+
return structValue<PartiQLValue>()
71+
}
72+
is ExcludeStep.StructField -> leafStep.attr
73+
is ExcludeStep.CollIndex, is ExcludeStep.CollWildcard -> null
9474
}
95-
// apply tuple wildcard exclusions
96-
val tupleWildcardKey = ExcludeStep.StructWildcard
97-
branches.find {
98-
it.step == tupleWildcardKey
99-
}?.let {
100-
expr = excludeOnPartiQLValue(expr, it)
75+
}.toSet()
76+
val branches = exclusions.branches
77+
val finalStruct = structValue.entries.mapNotNull { structField ->
78+
if (attrsToRemove.contains(structField.first)) {
79+
// struct attr is to be removed at current level
80+
null
81+
} else {
82+
// deeper level exclusions
83+
val name = structField.first
84+
var expr = structField.second
85+
// apply case-sensitive tuple attr exclusions at deeper levels
86+
val structFieldCaseSensitiveKey = ExcludeStep.StructField(name, ExcludeFieldCase.SENSITIVE)
87+
branches.find {
88+
it.step == structFieldCaseSensitiveKey
89+
}?.let {
90+
expr = excludeOnPartiQLValue(expr, it)
91+
}
92+
// apply case-insensitive tuple attr exclusions at deeper levels
93+
val structFieldCaseInsensitiveKey = ExcludeStep.StructField(name, ExcludeFieldCase.INSENSITIVE)
94+
branches.find {
95+
it.step == structFieldCaseInsensitiveKey
96+
}?.let {
97+
expr = excludeOnPartiQLValue(expr, it)
98+
}
99+
// apply tuple wildcard exclusions at deeper levels
100+
val tupleWildcardKey = ExcludeStep.StructWildcard
101+
branches.find {
102+
it.step == tupleWildcardKey
103+
}?.let {
104+
expr = excludeOnPartiQLValue(expr, it)
105+
}
106+
Pair(name, expr)
101107
}
102-
Pair(name, expr)
103108
}
104109
return structValue(finalStruct)
105110
}
@@ -124,33 +129,34 @@ internal class RelExclude(
124129
type: PartiQLValueType,
125130
exclusions: ExcludeNode
126131
): PartiQLValue {
127-
val leavesSteps = exclusions.leaves.map { leaf -> leaf.step }
128-
val branches = exclusions.branches
129-
if (leavesSteps.any { it is ExcludeStep.CollWildcard }) {
130-
// collection wildcard at current level. return empty collection
131-
return newCollValue(type, emptyList())
132-
} else {
133-
val indexesToRemove = leavesSteps.filterIsInstance<ExcludeStep.CollIndex>()
134-
.map { it.index }
135-
.toSet()
136-
val collWithRemoved = when (coll) {
137-
is BagValue -> coll
138-
is ListValue, is SexpValue -> coll.filterIndexed { index, _ ->
139-
!indexesToRemove.contains(index)
132+
val indexesToRemove = exclusions.leaves.mapNotNull { leaf ->
133+
when (val leafStep = leaf.step) {
134+
is ExcludeStep.CollWildcard -> {
135+
// collection wildcard at current level. return empty collection
136+
return newCollValue(type, emptyList())
140137
}
138+
is ExcludeStep.CollIndex -> leafStep.index
139+
is ExcludeStep.StructField, is ExcludeStep.StructWildcard -> null
141140
}
142-
val finalColl = collWithRemoved.mapIndexed { index, element ->
141+
}.toSet()
142+
val branches = exclusions.branches
143+
val finalColl = coll.mapIndexedNotNull { index, element ->
144+
if (indexesToRemove.contains(index)) {
145+
// coll index is to be removed at current level
146+
null
147+
} else {
148+
// deeper level exclusions
143149
var expr = element
144150
if (coll is ListValue || coll is SexpValue) {
145-
// apply collection index exclusions for lists and sexps
151+
// apply collection index exclusions at deeper levels for lists and sexps
146152
val elementKey = ExcludeStep.CollIndex(index)
147153
branches.find {
148154
it.step == elementKey
149155
}?.let {
150156
expr = excludeOnPartiQLValue(element, it)
151157
}
152158
}
153-
// apply collection wildcard exclusions for lists, bags, and sexps
159+
// apply collection wildcard exclusions at deeper levels for lists, bags, and sexps
154160
val collectionWildcardKey = ExcludeStep.CollWildcard
155161
branches.find {
156162
it.step == collectionWildcardKey
@@ -159,8 +165,8 @@ internal class RelExclude(
159165
}
160166
expr
161167
}
162-
return newCollValue(type, finalColl)
163168
}
169+
return newCollValue(type, finalColl)
164170
}
165171

166172
@OptIn(PartiQLValueExperimental::class)

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

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -245,6 +245,52 @@ class PartiQLEngineDefaultTest {
245245
)
246246
),
247247
)
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+
)
248294
)
249295
)
250296
}

0 commit comments

Comments
 (0)