Skip to content

Commit dcdb70e

Browse files
guizmaiiguizmaii
committed
DuckDB: support list comprehensions
`[x * 2 FOR x IN l IF x > 1]` and `[x + i FOR x, i IN l]` parse as the new `ListComprehension`. A `[` opens one when a `FOR` follows at its top level; any other `[` is still an array. `visit(ListComprehension)` is a default method that visits the element, list and condition, so existing `ExpressionVisitor` implementations keep compiling.
1 parent 49b5ea9 commit dcdb70e

8 files changed

Lines changed: 210 additions & 0 deletions

File tree

‎src/main/java/net/sf/jsqlparser/expression/ExpressionVisitor.java‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -822,6 +822,16 @@ default void visit(LambdaExpression lambdaExpression) {
822822
this.visit(lambdaExpression, null);
823823
}
824824

825+
default <S> T visit(ListComprehension listComprehension, S context) {
826+
visitExpression(listComprehension.getExpression(), context);
827+
visitExpression(listComprehension.getList(), context);
828+
return visitExpression(listComprehension.getCondition(), context);
829+
}
830+
831+
default void visit(ListComprehension listComprehension) {
832+
this.visit(listComprehension, null);
833+
}
834+
825835
default <S> T visit(ColumnsExpression columnsExpression, S context) {
826836
for (Expression expression : columnsExpression.getAllExpressions()) {
827837
if (expression != null) {

‎src/main/java/net/sf/jsqlparser/expression/ExpressionVisitorAdapter.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -909,6 +909,12 @@ public <S> T visit(LambdaExpression lambdaExpression, S context) {
909909
return lambdaExpression.getExpression().accept(this, context);
910910
}
911911

912+
@Override
913+
public <S> T visit(ListComprehension listComprehension, S context) {
914+
return visitExpressions(listComprehension, context, listComprehension.getExpression(),
915+
listComprehension.getList(), listComprehension.getCondition());
916+
}
917+
912918
@Override
913919
public <S> T visit(ColumnsExpression columnsExpression, S context) {
914920
return visitExpressions(columnsExpression, context, columnsExpression.getAllExpressions());

‎src/main/java/net/sf/jsqlparser/expression/ListComprehension.java‎

Lines changed: 98 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,98 @@
1+
/*-
2+
* #%L
3+
* JSQLParser library
4+
* %%
5+
* Copyright (C) 2004 - 2026 JSQLParser
6+
* %%
7+
* Dual licensed under GNU LGPL 2.1 or Apache License 2.0
8+
* #L%
9+
*/
10+
package net.sf.jsqlparser.expression;
11+
12+
import net.sf.jsqlparser.parser.ASTNodeAccessImpl;
13+
14+
import java.util.List;
15+
import java.util.function.Consumer;
16+
17+
/**
18+
* DuckDB's list comprehension {@code [expression FOR x[, i] IN list [IF condition]]}, where
19+
* {@code i} is the 1-based index of {@code x}.
20+
*
21+
* @see <a href="https://duckdb.org/docs/stable/sql/functions/list#list-comprehension">List
22+
* comprehension</a>
23+
*/
24+
public class ListComprehension extends ASTNodeAccessImpl implements Expression {
25+
private Expression expression;
26+
private List<String> variables;
27+
private Expression list;
28+
private Expression condition;
29+
30+
public ListComprehension(Expression expression, List<String> variables, Expression list,
31+
Expression condition) {
32+
this.expression = expression;
33+
this.variables = variables;
34+
this.list = list;
35+
this.condition = condition;
36+
}
37+
38+
public Expression getExpression() {
39+
return expression;
40+
}
41+
42+
public ListComprehension setExpression(Expression expression) {
43+
this.expression = expression;
44+
return this;
45+
}
46+
47+
/** The element variable, followed by the index variable if any. */
48+
public List<String> getVariables() {
49+
return variables;
50+
}
51+
52+
public ListComprehension setVariables(List<String> variables) {
53+
this.variables = variables;
54+
return this;
55+
}
56+
57+
public Expression getList() {
58+
return list;
59+
}
60+
61+
public ListComprehension setList(Expression list) {
62+
this.list = list;
63+
return this;
64+
}
65+
66+
/** The {@code IF} filter, or {@code null}. */
67+
public Expression getCondition() {
68+
return condition;
69+
}
70+
71+
public ListComprehension setCondition(Expression condition) {
72+
this.condition = condition;
73+
return this;
74+
}
75+
76+
public StringBuilder appendTo(StringBuilder builder, Consumer<Expression> expressionRenderer) {
77+
builder.append("[");
78+
expressionRenderer.accept(expression);
79+
builder.append(" FOR ").append(String.join(", ", variables)).append(" IN ");
80+
expressionRenderer.accept(list);
81+
if (condition != null) {
82+
builder.append(" IF ");
83+
expressionRenderer.accept(condition);
84+
}
85+
return builder.append("]");
86+
}
87+
88+
@Override
89+
public String toString() {
90+
StringBuilder builder = new StringBuilder();
91+
return appendTo(builder, builder::append).toString();
92+
}
93+
94+
@Override
95+
public <T, S> T accept(ExpressionVisitor<T> expressionVisitor, S context) {
96+
return expressionVisitor.visit(this, context);
97+
}
98+
}

‎src/main/java/net/sf/jsqlparser/util/TablesNamesFinder.java‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2338,6 +2338,16 @@ public <S> Void visit(LambdaExpression lambdaExpression, S context) {
23382338
return null;
23392339
}
23402340

2341+
@Override
2342+
public <S> Void visit(ListComprehension listComprehension, S context) {
2343+
listComprehension.getExpression().accept(this, context);
2344+
listComprehension.getList().accept(this, context);
2345+
if (listComprehension.getCondition() != null) {
2346+
listComprehension.getCondition().accept(this, context);
2347+
}
2348+
return null;
2349+
}
2350+
23412351
@Override
23422352
public <S> Void visit(ColumnsExpression columnsExpression, S context) {
23432353
for (Expression expression : columnsExpression.getAllExpressions()) {

‎src/main/java/net/sf/jsqlparser/util/deparser/ExpressionDeParser.java‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,7 @@
5151
import net.sf.jsqlparser.expression.KeepExpression;
5252
import net.sf.jsqlparser.expression.KeyExpression;
5353
import net.sf.jsqlparser.expression.LambdaExpression;
54+
import net.sf.jsqlparser.expression.ListComprehension;
5455
import net.sf.jsqlparser.expression.LongValue;
5556
import net.sf.jsqlparser.expression.LowExpression;
5657
import net.sf.jsqlparser.expression.MapExpression;
@@ -1856,6 +1857,12 @@ public <S> StringBuilder visit(LambdaExpression lambdaExpression, S context) {
18561857
return lambdaExpression.appendTo(builder, expression -> expression.accept(this, context));
18571858
}
18581859

1860+
@Override
1861+
public <S> StringBuilder visit(ListComprehension listComprehension, S context) {
1862+
return listComprehension.appendTo(builder,
1863+
expression -> expression.accept(this, context));
1864+
}
1865+
18591866
@Override
18601867
public <S> StringBuilder visit(ColumnsExpression columnsExpression, S context) {
18611868
columnsExpression.getColumns().accept(this, context);

‎src/main/java/net/sf/jsqlparser/util/validation/validator/ExpressionValidator.java‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@
4848
import net.sf.jsqlparser.expression.KeepExpression;
4949
import net.sf.jsqlparser.expression.KeyExpression;
5050
import net.sf.jsqlparser.expression.LambdaExpression;
51+
import net.sf.jsqlparser.expression.ListComprehension;
5152
import net.sf.jsqlparser.expression.LongValue;
5253
import net.sf.jsqlparser.expression.LowExpression;
5354
import net.sf.jsqlparser.expression.MapExpression;
@@ -1243,6 +1244,16 @@ public <S> Void visit(LambdaExpression lambdaExpression, S context) {
12431244
return null;
12441245
}
12451246

1247+
@Override
1248+
public <S> Void visit(ListComprehension listComprehension, S context) {
1249+
listComprehension.getExpression().accept(this, context);
1250+
listComprehension.getList().accept(this, context);
1251+
if (listComprehension.getCondition() != null) {
1252+
listComprehension.getCondition().accept(this, context);
1253+
}
1254+
return null;
1255+
}
1256+
12461257
@Override
12471258
public <S> Void visit(ColumnsExpression columnsExpression, S context) {
12481259
for (Expression expression : columnsExpression.getAllExpressions()) {

‎src/main/jjtree/net/sf/jsqlparser/parser/JSqlParserCC.jjt‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1650,6 +1650,30 @@ public class CCJSqlParser extends AbstractJSqlParser<CCJSqlParser> {
16501650
|| next.kind >= MIN_NON_RESERVED_WORD && next.kind <= MAX_NON_RESERVED_WORD;
16511651
}
16521652

1653+
/** Whether a {@code [} opens DuckDB's {@code [expr FOR x IN list]}: FOR at its top level. */
1654+
private boolean isListComprehensionAhead() {
1655+
if (!"[".equals(getToken(1).image)) {
1656+
return false;
1657+
}
1658+
int depth = 0;
1659+
for (int i = 1;; i++) {
1660+
Token token = getToken(i);
1661+
if (token == null || token.kind == EOF || ";".equals(token.image)) {
1662+
return false;
1663+
}
1664+
if ("(".equals(token.image) || "[".equals(token.image) || "{".equals(token.image)) {
1665+
depth++;
1666+
} else if (")".equals(token.image) || "]".equals(token.image)
1667+
|| "}".equals(token.image)) {
1668+
if (--depth == 0) {
1669+
return false;
1670+
}
1671+
} else if (depth == 1 && token.kind == K_FOR) {
1672+
return true;
1673+
}
1674+
}
1675+
}
1676+
16531677
private boolean isAliasAhead() {
16541678
if (isAsOfJoinAhead() || isSemiAntiJoinAhead() || isApproxNearestAhead()) {
16551679
return false;
@@ -11663,6 +11687,8 @@ Expression PrimaryExpression() #PrimaryExpression:
1166311687

1166411688
| LOOKAHEAD(3, {!interrupted}) <K_ARRAY_LITERAL> [ "<" type=ColDataType() ">" ] retval=ArrayConstructor(true) { if (type!=null) ((ArrayConstructor) retval).setDataType(type); }
1166511689

11690+
| LOOKAHEAD({ !interrupted && isListComprehensionAhead() }) retval=ListComprehension()
11691+
1166611692
| LOOKAHEAD(1, {!interrupted}) retval=ArrayConstructor(false)
1166711693

1166811694
| LOOKAHEAD(2, {!interrupted}) retval = NextValExpression()
@@ -11968,6 +11994,24 @@ RangeExpression RangeExpression(Expression startExpression):
1196811994
}
1196911995
}
1197011996

11997+
/** DuckDB's list comprehension {@code [expr FOR x[, i] IN list [IF condition]]}. */
11998+
ListComprehension ListComprehension() : {
11999+
Expression expression;
12000+
List<String> variables = new ArrayList<String>();
12001+
String variable;
12002+
Expression list;
12003+
Expression condition = null;
12004+
} {
12005+
"["
12006+
expression=Expression()
12007+
<K_FOR> variable=RelObjectName() { variables.add(variable); }
12008+
[ "," variable=RelObjectName() { variables.add(variable); } ]
12009+
<K_IN> list=Expression()
12010+
[ <K_IF> condition=Expression() ]
12011+
"]"
12012+
{ return new ListComprehension(expression, variables, list, condition); }
12013+
}
12014+
1197112015
ArrayConstructor ArrayConstructor(boolean arrayKeyword) : {
1197212016
ExpressionList<Expression> expList = new ExpressionList<Expression>();
1197312017
Expression exp;

‎src/test/java/net/sf/jsqlparser/statement/select/DuckDBTest.java‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -903,4 +903,28 @@ void testTableSampleWithoutOrWithReservoirMethod() throws JSQLParserException {
903903
.parse("SELECT a FROM t TABLESAMPLE 10 OFFSET 3");
904904
Assertions.assertEquals("3", select.getOffset().getOffset().toString());
905905
}
906+
907+
@Test
908+
void testListComprehension() throws JSQLParserException {
909+
for (String sqlStr : new String[] {"SELECT [x * 2 FOR x IN l] FROM t",
910+
"SELECT [x FOR x IN l IF x > 1] FROM t", "SELECT [x + i FOR x, i IN l] FROM t",
911+
"SELECT [[y FOR y IN x] FOR x IN [[1], [2]]]",
912+
"SELECT [x FOR x IN l IF x IN (1, 2)] FROM t",
913+
"SELECT [x FOR x IN (SELECT list(a) FROM t)]",
914+
"SELECT [x FOR x IN l][1] FROM t"}) {
915+
TestUtils.assertSqlCanBeParsedAndDeparsed(sqlStr, true);
916+
}
917+
918+
net.sf.jsqlparser.expression.ListComprehension comprehension =
919+
(net.sf.jsqlparser.expression.ListComprehension) CCJSqlParserUtil
920+
.parseExpression("[x + i FOR x, i IN l IF x > 1]");
921+
Assertions.assertEquals("x + i", comprehension.getExpression().toString());
922+
Assertions.assertEquals(Arrays.asList("x", "i"), comprehension.getVariables());
923+
Assertions.assertEquals("l", comprehension.getList().toString());
924+
Assertions.assertEquals("x > 1", comprehension.getCondition().toString());
925+
926+
Assertions.assertEquals(new HashSet<>(Arrays.asList("t", "u")),
927+
new HashSet<>(new TablesNamesFinder<>().getTables(CCJSqlParserUtil.parse(
928+
"SELECT [x FOR x IN (SELECT list(a) FROM u)] FROM t"))));
929+
}
906930
}

0 commit comments

Comments
 (0)