Skip to content

Commit 71c4a0b

Browse files
lucaxchenlucaxchen
authored
Merge 156d2a5 into ccefcfd 
2 parents ccefcfd + 156d2a5 commit 71c4a0b

File tree

2 files changed

+292
-0
lines changed
  • partiql-eval/src/test/kotlin/org/partiql/eval/internal
  • partiql-planner/src/main/kotlin/org/partiql/planner/internal/transforms

2 files changed

+292
-0
lines changed

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

Lines changed: 259 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,259 @@
1+
package org.partiql.eval.internal
2+
3+
import org.junit.jupiter.api.Disabled
4+
import org.junit.jupiter.api.Test
5+
import org.junit.jupiter.api.parallel.Execution
6+
import org.junit.jupiter.api.parallel.ExecutionMode
7+
import org.junit.jupiter.params.ParameterizedTest
8+
import org.junit.jupiter.params.provider.MethodSource
9+
import org.partiql.spi.value.Datum
10+
import org.partiql.spi.value.Field
11+
12+
/**
13+
* This test file exercises the `LET` clause in PartiQL.
14+
*/
15+
class LetTests {
16+
17+
@ParameterizedTest
18+
@MethodSource("successTestCases")
19+
@Execution(ExecutionMode.CONCURRENT)
20+
fun successTests(tc: SuccessTestCase) = tc.run()
21+
22+
@ParameterizedTest
23+
@MethodSource("failureTestCases")
24+
@Execution(ExecutionMode.CONCURRENT)
25+
fun failureTests(tc: FailureTestCase) = tc.run()
26+
27+
companion object {
28+
29+
@JvmStatic
30+
fun successTestCases() = listOf(
31+
SuccessTestCase(
32+
name = "Basic LET usage 1",
33+
input = """
34+
SELECT t.a, c
35+
FROM <<{ 'a': 1 , 'b': 2}>> AS t
36+
LET t.a*5 AS c
37+
""".trimIndent(),
38+
expected = Datum.bagVararg(
39+
Datum.struct(
40+
Field.of("a", Datum.integer(1)),
41+
Field.of("c", Datum.integer(5))
42+
)
43+
)
44+
),
45+
46+
SuccessTestCase(
47+
name = "Basic LET usage 2",
48+
input = """
49+
SELECT t.x, t.y, t.z * 2 AS double_z
50+
FROM (
51+
SELECT A AS x, B AS y, new_val AS z
52+
FROM <<{ 'A': 1, 'B': 2, 'C': 3}>>
53+
LET B + C AS new_val
54+
) AS t;
55+
""".trimIndent(),
56+
expected = Datum.bagVararg(
57+
Datum.struct(
58+
Field.of("x", Datum.integer(1)),
59+
Field.of("y", Datum.integer(2)),
60+
Field.of("double_z", Datum.integer(10))
61+
)
62+
)
63+
),
64+
65+
SuccessTestCase(
66+
name = "LET with JOIN operation",
67+
input = """
68+
SELECT t.customer_name, t.order_total
69+
FROM (
70+
SELECT
71+
c.name AS customer_name,
72+
total AS order_total
73+
FROM <<
74+
{ 'id': 1, 'name': 'Alice' },
75+
{ 'id': 2, 'name': 'Bob' }
76+
>> AS c
77+
JOIN <<
78+
{ 'customer_id': 1, 'amount': 100 },
79+
{ 'customer_id': 1, 'amount': 200 },
80+
{ 'customer_id': 2, 'amount': 150 }
81+
>> AS o
82+
ON c.id = o.customer_id
83+
LET o.amount * c.id AS total
84+
) AS t;
85+
""".trimIndent(),
86+
expected = Datum.bagVararg(
87+
Datum.struct(
88+
Field.of("customer_name", Datum.string("Alice")),
89+
Field.of("order_total", Datum.integer(100))
90+
),
91+
Datum.struct(
92+
Field.of("customer_name", Datum.string("Alice")),
93+
Field.of("order_total", Datum.integer(200))
94+
),
95+
Datum.struct(
96+
Field.of("customer_name", Datum.string("Bob")),
97+
Field.of("order_total", Datum.integer(300))
98+
)
99+
)
100+
),
101+
102+
SuccessTestCase(
103+
name = "LET with multiple items in data",
104+
input = """
105+
SELECT t.x, t.y, t.z AS total
106+
FROM (
107+
SELECT A AS x, B AS y, sum_val AS z
108+
FROM <<
109+
{ 'A': 1, 'B': 2, 'C': 3 },
110+
{ 'A': 10, 'B': 20, 'C': 30 }
111+
>>
112+
LET B + C AS sum_val
113+
) AS t;
114+
""".trimIndent(),
115+
expected = Datum.bagVararg(
116+
Datum.struct(
117+
Field.of("x", Datum.integer(1)),
118+
Field.of("y", Datum.integer(2)),
119+
Field.of("total", Datum.integer(5))
120+
),
121+
Datum.struct(
122+
Field.of("x", Datum.integer(10)),
123+
Field.of("y", Datum.integer(20)),
124+
Field.of("total", Datum.integer(50))
125+
)
126+
)
127+
),
128+
SuccessTestCase(
129+
name = "LET referencing prior expressions",
130+
input = """
131+
SELECT t.x, t.sum_val, t.double_sum
132+
FROM (
133+
SELECT
134+
A AS x,
135+
sum_val,
136+
sum_val * 2 AS double_sum
137+
FROM <<
138+
{ 'A': 3, 'B': 5 },
139+
{ 'A': 10, 'B': 2 }
140+
>>
141+
LET A + B AS sum_val
142+
) AS t;
143+
""".trimIndent(),
144+
expected = Datum.bagVararg(
145+
Datum.struct(
146+
Field.of("x", Datum.integer(3)),
147+
Field.of("sum_val", Datum.integer(8)),
148+
Field.of("double_sum", Datum.integer(16))
149+
),
150+
Datum.struct(
151+
Field.of("x", Datum.integer(10)),
152+
Field.of("sum_val", Datum.integer(12)),
153+
Field.of("double_sum", Datum.integer(24))
154+
)
155+
)
156+
),
157+
SuccessTestCase(
158+
name = "LET with multiple LET clauses 1",
159+
input = """
160+
SELECT t.a, b, c
161+
FROM << { 'a': 1 }>> AS t
162+
LET t.a + 2 AS b, t.a * 3 AS c
163+
""".trimIndent(),
164+
expected = Datum.bagVararg(
165+
Datum.struct(
166+
Field.of("a", Datum.integer(1)),
167+
Field.of("b", Datum.integer(3)),
168+
Field.of("c", Datum.integer(3))
169+
)
170+
)
171+
),
172+
SuccessTestCase(
173+
name = "LET with multiple LET clauses 2",
174+
input = """
175+
SELECT a, b, c, d, e
176+
FROM << { 'a': 1 , 'b':2}>> AS t
177+
LET t.a + 5 AS c, t.b+ 10 AS d, t.a + 15 AS e
178+
""".trimIndent(),
179+
expected = Datum.bagVararg(
180+
Datum.struct(
181+
Field.of("a", Datum.integer(1)),
182+
Field.of("b", Datum.integer(2)),
183+
Field.of("c", Datum.integer(6)),
184+
Field.of("d", Datum.integer(12)),
185+
Field.of("e", Datum.integer(16))
186+
)
187+
)
188+
)
189+
)
190+
191+
@JvmStatic
192+
fun failureTestCases() = listOf(
193+
FailureTestCase(
194+
name = "LET referencing undefined variable",
195+
input = """
196+
SELECT t.x
197+
FROM (
198+
SELECT A AS x
199+
FROM << { 'A': 1, 'B': 2 } >>
200+
LET nonexistent + B AS new_val
201+
) AS t;
202+
""".trimIndent()
203+
),
204+
FailureTestCase(
205+
name = "LET clause with ambiguous reference",
206+
input = """
207+
SELECT t.z
208+
FROM (
209+
SELECT new_val AS z
210+
FROM << { 'A': 1, 'B': 2 } >>
211+
-- 'new_val' references itself in LET, which is not allowed
212+
LET new_val + B AS new_val
213+
) AS t;
214+
""".trimIndent()
215+
),
216+
FailureTestCase(
217+
name = "Outside clauses referencing subquery's LET bindings",
218+
input = """
219+
SELECT t.x, t.y, new_val
220+
FROM (
221+
SELECT A AS x, B AS y
222+
FROM <<{ 'A': 1, 'B': 2, 'C': 3}>>
223+
LET B + C AS new_val
224+
) AS t;
225+
""".trimIndent()
226+
),
227+
FailureTestCase(
228+
name = "LET with invalid JOIN reference",
229+
input = """
230+
SELECT t.customer_name, t.calculated_total
231+
FROM (
232+
SELECT
233+
c.name AS customer_name,
234+
total AS calculated_total
235+
FROM <<
236+
{ 'id': 1, 'name': 'Alice' },
237+
{ 'id': 2, 'name': 'Bob' }
238+
>> AS c
239+
LEFT JOIN <<
240+
{ 'customer_id': 1, 'amount': 100 },
241+
{ 'customer_id': 2, 'amount': 150 }
242+
>> AS o
243+
ON c.id = o.customer_id
244+
-- This should fail because we're trying to reference 'missing_field'
245+
-- which doesn't exist in either joined table
246+
LET missing_field + o.amount AS total
247+
) AS t;
248+
""".trimIndent(),
249+
)
250+
)
251+
}
252+
253+
// Example of a test that might need special handling or a skip
254+
@Test
255+
@Disabled("Demonstration of a scenario needing further investigation.")
256+
fun disabledTestExample() {
257+
// Implementation left blank or used for demonstration
258+
}
259+
}

‎partiql-planner/src/main/kotlin/org/partiql/planner/internal/transforms/RelConverter.kt

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,8 @@ import org.partiql.ast.GroupBy
3131
import org.partiql.ast.GroupByStrategy
3232
import org.partiql.ast.Identifier
3333
import org.partiql.ast.JoinType
34+
import org.partiql.ast.Let
35+
import org.partiql.ast.Let.Binding
3436
import org.partiql.ast.Literal.intNum
3537
import org.partiql.ast.Nulls
3638
import org.partiql.ast.Order
@@ -192,6 +194,7 @@ internal object RelConverter {
192194
rel = convertOffset(rel, offset)
193195
rel = convertLimit(rel, limit)
194196
rel = convertExclude(rel, sel.exclude)
197+
rel = convertLet(rel, sel.let)
195198
// append SQL projection if present
196199
rel = when (val projection = sel.select) {
197200
is SelectValue -> {
@@ -383,6 +386,13 @@ internal object RelConverter {
383386
return binding to rex
384387
}
385388

389+
private fun convertLetBinding(binding: Binding): Pair<Rel.Binding, Rex> {
390+
val name = binding.asAlias.text
391+
val rex = RexConverter.apply(binding.expr, env)
392+
val newBinding = relBinding(name, rex.type)
393+
return newBinding to rex
394+
}
395+
386396
/**
387397
* Append [Rel.Op.Filter] only if a WHERE condition exists
388398
*/
@@ -582,6 +592,29 @@ internal object RelConverter {
582592
return rel(type, op)
583593
}
584594

595+
/**
596+
* Concatenate bindings in LET clause with existing env bindings from input
597+
*/
598+
private fun convertLet(input: Rel, let: Let?): Rel {
599+
if (let == null) {
600+
return input
601+
}
602+
val schema = input.type.schema.toMutableList()
603+
val props = input.type.props
604+
val projections = mutableListOf<Rex>()
605+
repeat(input.type.schema.size) { index ->
606+
projections.add(rex(ANY, rexOpVarLocal(0, index)))
607+
}
608+
let.bindings.forEach {
609+
val (newBinding, projection) = convertLetBinding(it)
610+
schema.add(newBinding)
611+
projections.add(projection)
612+
}
613+
val type = relType(schema, props)
614+
val op = relOpProject(input, projections)
615+
return rel(type, op)
616+
}
617+
585618
/**
586619
* Append [Rel.Op.Offset] if there is an OFFSET
587620
*/

0 commit comments

Comments
 (0)