Skip to content

Commit 9459631

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 fde4f03 commit 9459631

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
@@ -832,6 +832,16 @@ default void visit(LambdaExpression lambdaExpression) {
832832
this.visit(lambdaExpression, null);
833833
}
834834

835+
default <S> T visit(ListComprehension listComprehension, S context) {
836+
visitExpression(listComprehension.getExpression(), context);
837+
visitExpression(listComprehension.getList(), context);
838+
return visitExpression(listComprehension.getCondition(), context);
839+
}
840+
841+
default void visit(ListComprehension listComprehension) {
842+
this.visit(listComprehension, null);
843+
}
844+
835845
default <S> T visit(ColumnsExpression columnsExpression, S context) {
836846
for (Expression expression : columnsExpression.getAllExpressions()) {
837847
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
@@ -915,6 +915,12 @@ public <S> T visit(LambdaExpression lambdaExpression, S context) {
915915
return lambdaExpression.getExpression().accept(this, context);
916916
}
917917

918+
@Override
919+
public <S> T visit(ListComprehension listComprehension, S context) {
920+
return visitExpressions(listComprehension, context, listComprehension.getExpression(),
921+
listComprehension.getList(), listComprehension.getCondition());
922+
}
923+
918924
@Override
919925
public <S> T visit(ColumnsExpression columnsExpression, S context) {
920926
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
@@ -2345,6 +2345,16 @@ public <S> Void visit(LambdaExpression lambdaExpression, S context) {
23452345
return null;
23462346
}
23472347

2348+
@Override
2349+
public <S> Void visit(ListComprehension listComprehension, S context) {
2350+
listComprehension.getExpression().accept(this, context);
2351+
listComprehension.getList().accept(this, context);
2352+
if (listComprehension.getCondition() != null) {
2353+
listComprehension.getCondition().accept(this, context);
2354+
}
2355+
return null;
2356+
}
2357+
23482358
@Override
23492359
public <S> Void visit(ColumnsExpression columnsExpression, S context) {
23502360
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;
@@ -1872,6 +1873,12 @@ public <S> StringBuilder visit(LambdaExpression lambdaExpression, S context) {
18721873
return lambdaExpression.appendTo(builder, expression -> expression.accept(this, context));
18731874
}
18741875

1876+
@Override
1877+
public <S> StringBuilder visit(ListComprehension listComprehension, S context) {
1878+
return listComprehension.appendTo(builder,
1879+
expression -> expression.accept(this, context));
1880+
}
1881+
18751882
@Override
18761883
public <S> StringBuilder visit(ColumnsExpression columnsExpression, S context) {
18771884
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;
@@ -1250,6 +1251,16 @@ public <S> Void visit(LambdaExpression lambdaExpression, S context) {
12501251
return null;
12511252
}
12521253

1254+
@Override
1255+
public <S> Void visit(ListComprehension listComprehension, S context) {
1256+
listComprehension.getExpression().accept(this, context);
1257+
listComprehension.getList().accept(this, context);
1258+
if (listComprehension.getCondition() != null) {
1259+
listComprehension.getCondition().accept(this, context);
1260+
}
1261+
return null;
1262+
}
1263+
12531264
@Override
12541265
public <S> Void visit(ColumnsExpression columnsExpression, S context) {
12551266
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
@@ -1659,6 +1659,30 @@ public class CCJSqlParser extends AbstractJSqlParser<CCJSqlParser> {
16591659
|| next.kind >= MIN_NON_RESERVED_WORD && next.kind <= MAX_NON_RESERVED_WORD;
16601660
}
16611661

1662+
/** Whether a {@code [} opens DuckDB's {@code [expr FOR x IN list]}: FOR at its top level. */
1663+
private boolean isListComprehensionAhead() {
1664+
if (!"[".equals(getToken(1).image)) {
1665+
return false;
1666+
}
1667+
int depth = 0;
1668+
for (int i = 1;; i++) {
1669+
Token token = getToken(i);
1670+
if (token == null || token.kind == EOF || ";".equals(token.image)) {
1671+
return false;
1672+
}
1673+
if ("(".equals(token.image) || "[".equals(token.image) || "{".equals(token.image)) {
1674+
depth++;
1675+
} else if (")".equals(token.image) || "]".equals(token.image)
1676+
|| "}".equals(token.image)) {
1677+
if (--depth == 0) {
1678+
return false;
1679+
}
1680+
} else if (depth == 1 && token.kind == K_FOR) {
1681+
return true;
1682+
}
1683+
}
1684+
}
1685+
16621686
private boolean isAliasAhead() {
16631687
if (isAsOfJoinAhead() || isSemiAntiJoinAhead() || isApproxNearestAhead()) {
16641688
return false;
@@ -11687,6 +11711,8 @@ Expression PrimaryExpression() #PrimaryExpression:
1168711711

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

11714+
| LOOKAHEAD({ !interrupted && isListComprehensionAhead() }) retval=ListComprehension()
11715+
1169011716
| LOOKAHEAD(1, {!interrupted}) retval=ArrayConstructor(false)
1169111717

1169211718
| LOOKAHEAD(2, {!interrupted}) retval = NextValExpression()
@@ -11995,6 +12021,24 @@ RangeExpression RangeExpression(Expression startExpression):
1199512021
}
1199612022
}
1199712023

12024+
/** DuckDB's list comprehension {@code [expr FOR x[, i] IN list [IF condition]]}. */
12025+
ListComprehension ListComprehension() : {
12026+
Expression expression;
12027+
List<String> variables = new ArrayList<String>();
12028+
String variable;
12029+
Expression list;
12030+
Expression condition = null;
12031+
} {
12032+
"["
12033+
expression=Expression()
12034+
<K_FOR> variable=RelObjectName() { variables.add(variable); }
12035+
[ "," variable=RelObjectName() { variables.add(variable); } ]
12036+
<K_IN> list=Expression()
12037+
[ <K_IF> condition=Expression() ]
12038+
"]"
12039+
{ return new ListComprehension(expression, variables, list, condition); }
12040+
}
12041+
1199812042
ArrayConstructor ArrayConstructor(boolean arrayKeyword) : {
1199912043
ExpressionList<Expression> expList = new ExpressionList<Expression>();
1200012044
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
@@ -986,4 +986,28 @@ private static String tree(net.sf.jsqlparser.expression.Expression expression) {
986986
}
987987
return expression.toString();
988988
}
989+
990+
@Test
991+
void testListComprehension() throws JSQLParserException {
992+
for (String sqlStr : new String[] {"SELECT [x * 2 FOR x IN l] FROM t",
993+
"SELECT [x FOR x IN l IF x > 1] FROM t", "SELECT [x + i FOR x, i IN l] FROM t",
994+
"SELECT [[y FOR y IN x] FOR x IN [[1], [2]]]",
995+
"SELECT [x FOR x IN l IF x IN (1, 2)] FROM t",
996+
"SELECT [x FOR x IN (SELECT list(a) FROM t)]",
997+
"SELECT [x FOR x IN l][1] FROM t"}) {
998+
TestUtils.assertSqlCanBeParsedAndDeparsed(sqlStr, true);
999+
}
1000+
1001+
net.sf.jsqlparser.expression.ListComprehension comprehension =
1002+
(net.sf.jsqlparser.expression.ListComprehension) CCJSqlParserUtil
1003+
.parseExpression("[x + i FOR x, i IN l IF x > 1]");
1004+
Assertions.assertEquals("x + i", comprehension.getExpression().toString());
1005+
Assertions.assertEquals(Arrays.asList("x", "i"), comprehension.getVariables());
1006+
Assertions.assertEquals("l", comprehension.getList().toString());
1007+
Assertions.assertEquals("x > 1", comprehension.getCondition().toString());
1008+
1009+
Assertions.assertEquals(new HashSet<>(Arrays.asList("t", "u")),
1010+
new HashSet<>(new TablesNamesFinder<>().getTables(CCJSqlParserUtil.parse(
1011+
"SELECT [x FOR x IN (SELECT list(a) FROM u)] FROM t"))));
1012+
}
9891013
}

0 commit comments

Comments
 (0)