Add support for FILTER clause in LISTAGG

ksobolew · Nov 23, 2023 · 93c0040 · 93c0040
1 parent 257c8ad
commit 93c0040
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 2 deletions.
diff --git a/core/trino-grammar/src/main/antlr4/io/trino/grammar/sql/SqlBase.g4 b/core/trino-grammar/src/main/antlr4/io/trino/grammar/sql/SqlBase.g4
@@ -566,7 +566,8 @@ primaryExpression
     | ROW '(' expression (',' expression)* ')'                                            #rowConstructor
     | name=LISTAGG '(' setQuantifier? expression (',' string)?
         (ON OVERFLOW listAggOverflowBehavior)? ')'
-        (WITHIN GROUP '(' ORDER BY sortItem (',' sortItem)* ')')                          #listagg
+        (WITHIN GROUP '(' ORDER BY sortItem (',' sortItem)* ')')
+        filter?                                                                           #listagg
     | processingMode? qualifiedName '(' (label=identifier '.')? ASTERISK ')'
         filter? over?                                                                     #functionCall
     | processingMode? qualifiedName '(' (setQuantifier? expression (',' expression)*)?

diff --git a/core/trino-main/src/test/java/io/trino/sql/query/TestListagg.java b/core/trino-main/src/test/java/io/trino/sql/query/TestListagg.java
@@ -391,4 +391,17 @@ public void testListaggQueryGroupingOverflowTruncateWithCountAndWithOverflowFill
                         "   (1, VARCHAR '" + largeValue + ",everything,.....(2)')," +
                         "   (2, VARCHAR 'listagg,string joiner')");
     }
+
+    @Test
+    void testFilter()
+    {
+        assertThat(assertions.query(
+                """
+                SELECT listagg(value, ',') WITHIN GROUP (ORDER BY id) FILTER (WHERE id % 2 = 0)
+                FROM (
+                     VALUES (1, 'a'), (2, 'b'), (3, 'c'), (4, 'd')
+                ) t(id, value)
+                 """))
+                .matches("VALUES VARCHAR 'b,d'");
+    }
 }
diff --git a/core/trino-parser/src/main/java/io/trino/sql/ExpressionFormatter.java b/core/trino-parser/src/main/java/io/trino/sql/ExpressionFormatter.java
@@ -1023,6 +1023,10 @@ private String visitListagg(FunctionCall node)
                         .append(')');
             }
 
+            if (node.getFilter().isPresent()) {
+                builder.append(" FILTER ").append(visitFilter(node.getFilter().get(), null));
+            }
+
             return builder.toString();
         }
     }

diff --git a/core/trino-parser/src/main/java/io/trino/sql/parser/AstBuilder.java b/core/trino-parser/src/main/java/io/trino/sql/parser/AstBuilder.java
@@ -2457,7 +2457,7 @@ else if (listaggCountIndicationContext.WITHOUT() != null) {
                 Optional.of(getLocation(context)),
                 QualifiedName.of("LISTAGG"),
                 window,
-                Optional.empty(),
+                visitIfPresent(context.filter(), Expression.class),
                 Optional.of(orderBy),
                 distinct,
                 Optional.empty(),