Skip to content

Commit bd0dc8f

Browse files
author
Yang Chen
committed
Add implementation of LET clause and test cases
1 parent ccefcfd commit bd0dc8f

File tree

2 files changed

+275
-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

+275
-0
lines changed

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

Lines changed: 242 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,242 @@
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",
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+
)
173+
174+
@JvmStatic
175+
fun failureTestCases() = listOf(
176+
FailureTestCase(
177+
name = "LET referencing undefined variable",
178+
input = """
179+
SELECT t.x
180+
FROM (
181+
SELECT A AS x
182+
FROM << { 'A': 1, 'B': 2 } >>
183+
LET nonexistent + B AS new_val
184+
) AS t;
185+
""".trimIndent()
186+
),
187+
FailureTestCase(
188+
name = "LET clause with ambiguous reference",
189+
input = """
190+
SELECT t.z
191+
FROM (
192+
SELECT new_val AS z
193+
FROM << { 'A': 1, 'B': 2 } >>
194+
-- 'new_val' references itself in LET, which is not allowed
195+
LET new_val + B AS new_val
196+
) AS t;
197+
""".trimIndent()
198+
),
199+
FailureTestCase(
200+
name = "Outside clauses referencing subquery's LET bindings",
201+
input = """
202+
SELECT t.x, t.y, new_val
203+
FROM (
204+
SELECT A AS x, B AS y
205+
FROM <<{ 'A': 1, 'B': 2, 'C': 3}>>
206+
LET B + C AS new_val
207+
) AS t;
208+
""".trimIndent()
209+
),
210+
FailureTestCase(
211+
name = "LET with invalid JOIN reference",
212+
input = """
213+
SELECT t.customer_name, t.calculated_total
214+
FROM (
215+
SELECT
216+
c.name AS customer_name,
217+
total AS calculated_total
218+
FROM <<
219+
{ 'id': 1, 'name': 'Alice' },
220+
{ 'id': 2, 'name': 'Bob' }
221+
>> AS c
222+
LEFT JOIN <<
223+
{ 'customer_id': 1, 'amount': 100 },
224+
{ 'customer_id': 2, 'amount': 150 }
225+
>> AS o
226+
ON c.id = o.customer_id
227+
-- This should fail because we're trying to reference 'missing_field'
228+
-- which doesn't exist in either joined table
229+
LET missing_field + o.amount AS total
230+
) AS t;
231+
""".trimIndent(),
232+
)
233+
)
234+
}
235+
236+
// Example of a test that might need special handling or a skip
237+
@Test
238+
@Disabled("Demonstration of a scenario needing further investigation.")
239+
fun disabledTestExample() {
240+
// Implementation left blank or used for demonstration
241+
}
242+
}

‎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)