fix(duckdb): support postgres JSON/JSONB_OBJECT_AGG to duckdb JSON_GR…

…OUP_OBJECT (#4677) * fix(duckdb): support postgres JSON/JSONB_OBJECT_AGG to duckdb JSON_GROUP_OBJECT * PR feedback 1
tobymao · Jan 29, 2025 · df75edd · df75edd
1 parent c45f174
commit df75edd
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 0 deletions.
diff --git a/sqlglot/dialects/duckdb.py b/sqlglot/dialects/duckdb.py
@@ -662,6 +662,8 @@ class Generator(generator.Generator):
             exp.Levenshtein: unsupported_args("ins_cost", "del_cost", "sub_cost", "max_dist")(
                 rename_func("LEVENSHTEIN")
             ),
+            exp.JSONObjectAgg: rename_func("JSON_GROUP_OBJECT"),
+            exp.JSONBObjectAgg: rename_func("JSON_GROUP_OBJECT"),
         }
 
         SUPPORTED_JSON_PATH_PARTS = {

diff --git a/sqlglot/dialects/postgres.py b/sqlglot/dialects/postgres.py
@@ -393,6 +393,8 @@ class Parser(parser.Parser):
             "SHA384": lambda args: exp.SHA2(this=seq_get(args, 0), length=exp.Literal.number(384)),
             "SHA512": lambda args: exp.SHA2(this=seq_get(args, 0), length=exp.Literal.number(512)),
             "LEVENSHTEIN_LESS_EQUAL": _build_levenshtein_less_equal,
+            "JSON_OBJECT_AGG": lambda args: exp.JSONObjectAgg(expressions=args),
+            "JSONB_OBJECT_AGG": exp.JSONBObjectAgg.from_arg_list,
         }
 
         NO_PAREN_FUNCTIONS = {
@@ -617,6 +619,8 @@ class Generator(generator.Generator):
             exp.Unicode: rename_func("ASCII"),
             exp.UnixToTime: _unix_to_time_sql,
             exp.Levenshtein: _levenshtein_sql,
+            exp.JSONObjectAgg: rename_func("JSON_OBJECT_AGG"),
+            exp.JSONBObjectAgg: rename_func("JSONB_OBJECT_AGG"),
         }
 
         TRANSFORMS.pop(exp.CommentColumnConstraint)

diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
@@ -6068,6 +6068,11 @@ class JSONObjectAgg(AggFunc):
     }
 
 
+# https://www.postgresql.org/docs/9.5/functions-aggregate.html
+class JSONBObjectAgg(AggFunc):
+    arg_types = {"this": True, "expression": True}
+
+
 # https://docs.oracle.com/en/database/oracle/oracle-database/19/sqlrf/JSON_ARRAY.html
 class JSONArray(Func):
     arg_types = {

diff --git a/tests/dialects/test_postgres.py b/tests/dialects/test_postgres.py
@@ -841,6 +841,22 @@ def test_postgres(self):
             "SELECT PERCENTILE_CONT(0.5) WITHIN GROUP (ORDER BY a) FILTER(WHERE CAST(b AS BOOLEAN)) AS mean_value FROM (VALUES (0, 't')) AS fake_data(a, b)"
         )
 
+        self.validate_all(
+            "SELECT JSON_OBJECT_AGG(k, v) FROM t",
+            write={
+                "postgres": "SELECT JSON_OBJECT_AGG(k, v) FROM t",
+                "duckdb": "SELECT JSON_GROUP_OBJECT(k, v) FROM t",
+            },
+        )
+
+        self.validate_all(
+            "SELECT JSONB_OBJECT_AGG(k, v) FROM t",
+            write={
+                "postgres": "SELECT JSONB_OBJECT_AGG(k, v) FROM t",
+                "duckdb": "SELECT JSON_GROUP_OBJECT(k, v) FROM t",
+            },
+        )
+
     def test_ddl(self):
         # Checks that user-defined types are parsed into DataType instead of Identifier
         self.parse_one("CREATE TABLE t (a udt)").this.expressions[0].args["kind"].assert_is(