feat: add 'array_agg', 'max_by' function

prqlc/bindings/prqlc-python/src/lib.rs

@@ -232,15 +232,15 @@ mod test {
     fn debug_prql_lineage() {
         assert_snapshot!(
             debug::prql_lineage(r#"from a | select { beta, gamma }"#).unwrap(),
-            @r#"{"frames":[["1:9-31",{"columns":[{"Single":{"name":["a","beta"],"target_id":121,"target_name":null}},{"Single":{"name":["a","gamma"],"target_id":122,"target_name":null}}],"inputs":[{"id":119,"name":"a","table":["default_db","a"]}]}]],"nodes":[{"id":119,"kind":"Ident","span":"1:0-6","ident":{"Ident":["default_db","a"]},"parent":124},{"id":121,"kind":"Ident","span":"1:18-22","ident":{"Ident":["this","a","beta"]},"targets":[119],"parent":123},{"id":122,"kind":"Ident","span":"1:24-29","ident":{"Ident":["this","a","gamma"]},"targets":[119],"parent":123},{"id":123,"kind":"Tuple","span":"1:16-31","children":[121,122],"parent":124},{"id":124,"kind":"TransformCall: Select","span":"1:9-31","children":[119,123]}],"ast":{"name":"Project","stmts":[{"VarDef":{"kind":"Main","name":"main","value":{"Pipeline":{"exprs":[{"FuncCall":{"name":{"Ident":["from"],"span":"1:0-4"},"args":[{"Ident":["a"],"span":"1:5-6"}]},"span":"1:0-6"},{"FuncCall":{"name":{"Ident":["select"],"span":"1:9-15"},"args":[{"Tuple":[{"Ident":["beta"],"span":"1:18-22"},{"Ident":["gamma"],"span":"1:24-29"}],"span":"1:16-31"}]},"span":"1:9-31"}]},"span":"1:0-31"}},"span":"1:0-31"}]}}"#
+            @r#"{"frames":[["1:9-31",{"columns":[{"Single":{"name":["a","beta"],"target_id":123,"target_name":null}},{"Single":{"name":["a","gamma"],"target_id":124,"target_name":null}}],"inputs":[{"id":121,"name":"a","table":["default_db","a"]}]}]],"nodes":[{"id":121,"kind":"Ident","span":"1:0-6","ident":{"Ident":["default_db","a"]},"parent":126},{"id":123,"kind":"Ident","span":"1:18-22","ident":{"Ident":["this","a","beta"]},"targets":[121],"parent":125},{"id":124,"kind":"Ident","span":"1:24-29","ident":{"Ident":["this","a","gamma"]},"targets":[121],"parent":125},{"id":125,"kind":"Tuple","span":"1:16-31","children":[123,124],"parent":126},{"id":126,"kind":"TransformCall: Select","span":"1:9-31","children":[121,125]}],"ast":{"name":"Project","stmts":[{"VarDef":{"kind":"Main","name":"main","value":{"Pipeline":{"exprs":[{"FuncCall":{"name":{"Ident":["from"],"span":"1:0-4"},"args":[{"Ident":["a"],"span":"1:5-6"}]},"span":"1:0-6"},{"FuncCall":{"name":{"Ident":["select"],"span":"1:9-15"},"args":[{"Tuple":[{"Ident":["beta"],"span":"1:18-22"},{"Ident":["gamma"],"span":"1:24-29"}],"span":"1:16-31"}]},"span":"1:9-31"}]},"span":"1:0-31"}},"span":"1:0-31"}]}}"#
         );
     }
 
     #[test]
     fn debug_pl_to_lineage() {
         assert_snapshot!(
             prql_to_pl(r#"from a | select { beta, gamma }"#).and_then(|x| debug::pl_to_lineage(&x)).unwrap(),
-            @r#"{"frames":[["1:9-31",{"columns":[{"Single":{"name":["a","beta"],"target_id":121,"target_name":null}},{"Single":{"name":["a","gamma"],"target_id":122,"target_name":null}}],"inputs":[{"id":119,"name":"a","table":["default_db","a"]}]}]],"nodes":[{"id":119,"kind":"Ident","span":"1:0-6","ident":{"Ident":["default_db","a"]},"parent":124},{"id":121,"kind":"Ident","span":"1:18-22","ident":{"Ident":["this","a","beta"]},"targets":[119],"parent":123},{"id":122,"kind":"Ident","span":"1:24-29","ident":{"Ident":["this","a","gamma"]},"targets":[119],"parent":123},{"id":123,"kind":"Tuple","span":"1:16-31","children":[121,122],"parent":124},{"id":124,"kind":"TransformCall: Select","span":"1:9-31","children":[119,123]}],"ast":{"name":"Project","stmts":[{"VarDef":{"kind":"Main","name":"main","value":{"Pipeline":{"exprs":[{"FuncCall":{"name":{"Ident":["from"],"span":"1:0-4"},"args":[{"Ident":["a"],"span":"1:5-6"}]},"span":"1:0-6"},{"FuncCall":{"name":{"Ident":["select"],"span":"1:9-15"},"args":[{"Tuple":[{"Ident":["beta"],"span":"1:18-22"},{"Ident":["gamma"],"span":"1:24-29"}],"span":"1:16-31"}]},"span":"1:9-31"}]},"span":"1:0-31"}},"span":"1:0-31"}]}}"#
+            @r#"{"frames":[["1:9-31",{"columns":[{"Single":{"name":["a","beta"],"target_id":123,"target_name":null}},{"Single":{"name":["a","gamma"],"target_id":124,"target_name":null}}],"inputs":[{"id":121,"name":"a","table":["default_db","a"]}]}]],"nodes":[{"id":121,"kind":"Ident","span":"1:0-6","ident":{"Ident":["default_db","a"]},"parent":126},{"id":123,"kind":"Ident","span":"1:18-22","ident":{"Ident":["this","a","beta"]},"targets":[121],"parent":125},{"id":124,"kind":"Ident","span":"1:24-29","ident":{"Ident":["this","a","gamma"]},"targets":[121],"parent":125},{"id":125,"kind":"Tuple","span":"1:16-31","children":[123,124],"parent":126},{"id":126,"kind":"TransformCall: Select","span":"1:9-31","children":[121,125]}],"ast":{"name":"Project","stmts":[{"VarDef":{"kind":"Main","name":"main","value":{"Pipeline":{"exprs":[{"FuncCall":{"name":{"Ident":["from"],"span":"1:0-4"},"args":[{"Ident":["a"],"span":"1:5-6"}]},"span":"1:0-6"},{"FuncCall":{"name":{"Ident":["select"],"span":"1:9-15"},"args":[{"Tuple":[{"Ident":["beta"],"span":"1:18-22"},{"Ident":["gamma"],"span":"1:24-29"}],"span":"1:16-31"}]},"span":"1:9-31"}]},"span":"1:0-31"}},"span":"1:0-31"}]}}"#
         );
     }
 }

prqlc/prqlc/src/cli/test.rs

@@ -434,61 +434,61 @@ fn debug() {
           name:
           - tracks
           - artist
-          target_id: 121
+          target_id: 123
           target_name: null
         - !Single
           name:
           - tracks
           - album
-          target_id: 122
+          target_id: 124
           target_name: null
         inputs:
-        - id: 119
+        - id: 121
           name: tracks
           table:
           - default_db
           - tracks
     nodes:
-    - id: 119
+    - id: 121
       kind: Ident
       span: 1:0-11
       ident: !Ident
       - default_db
       - tracks
-      parent: 124
-    - id: 121
+      parent: 126
+    - id: 123
       kind: Ident
       span: 1:22-28
       ident: !Ident
       - this
       - tracks
       - artist
       targets:
-      - 119
-      parent: 123
-    - id: 122
+      - 121
+      parent: 125
+    - id: 124
       kind: Ident
       span: 1:30-35
       ident: !Ident
       - this
       - tracks
       - album
       targets:
-      - 119
-      parent: 123
-    - id: 123
+      - 121
+      parent: 125
+    - id: 125
       kind: Tuple
       span: 1:21-36
       children:
-      - 121
-      - 122
-      parent: 124
-    - id: 124
+      - 123
+      - 124
+      parent: 126
+    - id: 126
       kind: 'TransformCall: Select'
       span: 1:14-36
       children:
-      - 119
-      - 123
+      - 121
+      - 125
     ast:
       name: Project
       stmts:

prqlc/prqlc/src/semantic/resolver/snapshots/prqlc__semantic__resolver__test__append_union_different_tables.snap

@@ -7,21 +7,21 @@ columns:
       name:
         - employees
         - name
-      target_id: 135
+      target_id: 137
       target_name: ~
   - Single:
       name:
         - employees
         - salary
-      target_id: 136
+      target_id: 138
       target_name: ~
 inputs:
-  - id: 133
+  - id: 135
     name: employees
     table:
       - default_db
       - employees
-  - id: 122
+  - id: 124
     name: managers
     table:
       - default_db

prqlc/prqlc/src/semantic/resolver/snapshots/prqlc__semantic__resolver__test__frames_and_names-2.snap

@@ -1,22 +1,21 @@
 ---
 source: prqlc/prqlc/src/semantic/resolver/mod.rs
-assertion_line: 190
 expression: "resolve_lineage(r#\"\n            from table_1\n            join customers (==customer_no)\n            \"#).unwrap()"
 ---
 columns:
   - All:
-      input_id: 120
+      input_id: 122
       except: []
   - All:
-      input_id: 117
+      input_id: 119
       except: []
 inputs:
-  - id: 120
+  - id: 122
     name: table_1
     table:
       - default_db
       - table_1
-  - id: 117
+  - id: 119
     name: customers
     table:
       - default_db

prqlc/prqlc/src/semantic/resolver/snapshots/prqlc__semantic__resolver__test__frames_and_names-3.snap

@@ -1,33 +1,32 @@
 ---
 source: prqlc/prqlc/src/semantic/resolver/mod.rs
-assertion_line: 198
 expression: "resolve_lineage(r#\"\n            from e = employees\n            join salaries (==emp_no)\n            group {e.emp_no, e.gender} (\n                aggregate {\n                    emp_salary = average salaries.salary\n                }\n            )\n            \"#).unwrap()"
 ---
 columns:
   - Single:
       name:
         - e
         - emp_no
-      target_id: 130
+      target_id: 132
       target_name: ~
   - Single:
       name:
         - e
         - gender
-      target_id: 131
+      target_id: 133
       target_name: ~
   - Single:
       name:
         - emp_salary
-      target_id: 149
+      target_id: 151
       target_name: ~
 inputs:
-  - id: 123
+  - id: 125
     name: e
     table:
       - default_db
       - employees
-  - id: 120
+  - id: 122
     name: salaries
     table:
       - default_db

prqlc/prqlc/src/semantic/resolver/snapshots/prqlc__semantic__resolver__test__frames_and_names.snap

@@ -7,26 +7,26 @@ columns:
       name:
         - orders
         - customer_no
-      target_id: 124
+      target_id: 126
       target_name: ~
   - Single:
       name:
         - orders
         - gross
-      target_id: 125
+      target_id: 127
       target_name: ~
   - Single:
       name:
         - orders
         - tax
-      target_id: 126
+      target_id: 128
       target_name: ~
   - Single:
       name: ~
-      target_id: 127
+      target_id: 129
       target_name: ~
 inputs:
-  - id: 122
+  - id: 124
     name: orders
     table:
       - default_db

prqlc/prqlc/src/semantic/resolver/snapshots/prqlc__semantic__resolver__transforms__tests__aggregate_positional_arg-2.snap

@@ -21,10 +21,10 @@ TransformCall:
     lineage:
       columns:
         - All:
-            input_id: 119
+            input_id: 121
             except: []
       inputs:
-        - id: 119
+        - id: 121
           name: c_invoice
           table:
             - default_db
@@ -94,14 +94,14 @@ lineage:
         name:
           - c_invoice
           - issued_at
-        target_id: 121
+        target_id: 123
         target_name: ~
     - Single:
         name: ~
-        target_id: 137
+        target_id: 139
         target_name: ~
   inputs:
-    - id: 119
+    - id: 121
       name: c_invoice
       table:
         - default_db

prqlc/prqlc/src/semantic/std.prql

@@ -154,6 +154,10 @@ let any = column <array> -> <bool> internal std.any
 
 let concat_array = column <array> -> <text> internal std.concat_array
 
+let array_agg = column <array> -> <array> internal std.array_agg
+
+let max_by = column <array> by <array> -> internal std.max_by
+
 # Counts number of items in the column.
 # Note that the count will include null values.
 let count = column <array> -> <int> internal count

prqlc/prqlc/src/sql/std.sql.prql

@@ -39,6 +39,12 @@ let any = column -> s"BOOL_OR({column:0})"
 @{window_frame=true, coalesce="''"}
 let concat_array = column -> s"STRING_AGG({column:0}, '')"
 
+@{window_frame=true}
+let array_agg = column -> s"ARRAY_AGG({column:0})"
+
+@{window_frame=true}
+let max_by = column by -> null
+
 @{window_frame=true}
 let count = column -> s"COUNT(*)"
 
@@ -213,6 +219,9 @@ module ansi {
 }
 
 module bigquery {
+  @{window_frame=true}
+  let max_by = column by -> s"MAX_BY({column:0}, {by:0})"
+
   @{binding_strength=11}
   let div_f = l r -> s"({l} * 1.0 / {r:12})"
 
@@ -235,6 +244,12 @@ module bigquery {
 }
 
 module clickhouse {
+  @{window_frame=true}
+  let array_agg = column -> s"groupArray({column:0})"
+
+  @{window_frame=true}
+  let max_by = column by -> s"argMax({column:0}, {by:0})"
+
   # https://clickhouse.com/docs/en/sql-reference/functions/arithmetic-functions#divide
   @{binding_strength=11}
   let div_f = l r -> s"({l} / {r:12})"
@@ -260,6 +275,12 @@ module clickhouse {
 }
 
 module duckdb {
+  @{window_frame=true, coalesce="[]"}
+  let array_agg = column -> s"ARRAY_AGG({column:0})"
+
+  @{window_frame=true}
+  let max_by = column by -> s"arg_max({column:0}, {by:0})"
+
   @{binding_strength=11}
   let div_f = l r -> s"({l} / {r:12})"
 
@@ -315,6 +336,9 @@ module mssql {
 }
 
 module mysql {
+  @{window_frame=true}
+  let array_agg = column -> s"JSON_ARRAYAGG({column:0})"
+
   @{binding_strength=11}
   let div_f = l r -> s"({l} / {r:12})"
 
@@ -341,6 +365,9 @@ module mysql {
 }
 
 module postgres {
+  @{window_frame=true}
+  let max_by = column by -> s"(ARRAY_AGG({column:0} ORDER BY {by:0} DESC NULLS LAST))[1]"
+
   @{binding_strength=11}
   let div_f = l r -> s"({l} * 1.0 / {r:12})"
 
@@ -415,6 +442,9 @@ module sqlite {
   @{window_frame=true, coalesce="''"}
   let concat_array = column -> s"GROUP_CONCAT({column:0}, '')"
 
+  @{window_frame=true, coalesce="'[]'"}
+  let array_agg = column -> s"json_group_array({column:0})"
+
   @{binding_strength=11}
   let div_f = l r -> s"({l} * 1.0 / {r:12})"
 
@@ -442,6 +472,9 @@ module sqlite {
 }
 
 module snowflake {
+  @{window_frame=true}
+  let max_by = column by -> s"MAX_BY({column:0}, {by:0})"
+
   # https://docs.snowflake.com/en/sql-reference/operators-arithmetic#division
   @{binding_strength=11}
   let div_f = l r -> s"({l} / {r:12})"

prqlc/prqlc/tests/integration/queries/aggregation.prql

@@ -4,4 +4,4 @@
 from tracks
 filter genre_id == 100
 derive empty_name = name == ''
-aggregate {sum track_id, concat_array name, all empty_name, any empty_name}
+aggregate {sum track_id, concat_array name, array_agg name, all empty_name, any empty_name}

prqlc/prqlc/tests/integration/snapshots/integration__queries__compile__aggregation.snap

@@ -1,11 +1,12 @@
 ---
 source: prqlc/prqlc/tests/integration/queries.rs
-expression: "# mysql:skip\n# clickhouse:skip\n# glaredb:skip (the string_agg function is not supported)\nfrom tracks\nfilter genre_id == 100\nderive empty_name = name == ''\naggregate {sum track_id, concat_array name, all empty_name, any empty_name}\n"
+expression: "# mysql:skip\n# clickhouse:skip\n# glaredb:skip (the string_agg function is not supported)\nfrom tracks\nfilter genre_id == 100\nderive empty_name = name == ''\naggregate {sum track_id, concat_array name, array_agg name, all empty_name, any empty_name}\n"
 input_file: prqlc/prqlc/tests/integration/queries/aggregation.prql
 ---
 SELECT
   COALESCE(SUM(track_id), 0),
   COALESCE(STRING_AGG(name, ''), ''),
+  ARRAY_AGG(name),
   COALESCE(BOOL_AND(name = ''), TRUE),
   COALESCE(BOOL_OR(name = ''), FALSE)
 FROM