diff --git a/sqlglot/dialects/duckdb.py b/sqlglot/dialects/duckdb.py index 5837d2ff0..b7f372d63 100644 --- a/sqlglot/dialects/duckdb.py +++ b/sqlglot/dialects/duckdb.py @@ -5,6 +5,7 @@ from sqlglot import exp, generator, parser, tokens, transforms from sqlglot.dialects.dialect import ( Dialect, + JSON_EXTRACT_TYPE, NormalizationStrategy, approx_count_distinct_sql, arg_max_or_min_no_count, @@ -156,6 +157,13 @@ def _unix_to_time_sql(self: DuckDB.Generator, expression: exp.UnixToTime) -> str return self.func("TO_TIMESTAMP", exp.Div(this=timestamp, expression=exp.func("POW", 10, scale))) +def _arrow_json_extract_sql(self: DuckDB.Generator, expression: JSON_EXTRACT_TYPE) -> str: + arrow_sql = arrow_json_extract_sql(self, expression) + if not expression.same_parent and isinstance(expression.parent, exp.Binary): + arrow_sql = self.wrap(arrow_sql) + return arrow_sql + + class DuckDB(Dialect): NULL_ORDERING = "nulls_are_last" SUPPORTS_USER_DEFINED_TYPES = False @@ -405,8 +413,8 @@ class Generator(generator.Generator): exp.IntDiv: lambda self, e: self.binary(e, "//"), exp.IsInf: rename_func("ISINF"), exp.IsNan: rename_func("ISNAN"), - exp.JSONExtract: arrow_json_extract_sql, - exp.JSONExtractScalar: arrow_json_extract_sql, + exp.JSONExtract: _arrow_json_extract_sql, + exp.JSONExtractScalar: _arrow_json_extract_sql, exp.JSONFormat: _json_format_sql, exp.LogicalOr: rename_func("BOOL_OR"), exp.LogicalAnd: rename_func("BOOL_AND"), diff --git a/tests/dialects/test_duckdb.py b/tests/dialects/test_duckdb.py index c44177838..6eb23939e 100644 --- a/tests/dialects/test_duckdb.py +++ b/tests/dialects/test_duckdb.py @@ -271,6 +271,14 @@ def test_duckdb(self): self.validate_identity( "SELECT * FROM x LEFT JOIN UNNEST(y)", "SELECT * FROM x LEFT JOIN UNNEST(y) ON TRUE" ) + self.validate_identity( + "SELECT JSON_EXTRACT_STRING(c, '$.k1') = 'v1'", + "SELECT (c ->> '$.k1') = 'v1'", + ) + self.validate_identity( + "SELECT JSON_EXTRACT(c, '$.k1') = 'v1'", + "SELECT (c -> '$.k1') = 'v1'", + ) self.validate_identity( """SELECT '{"foo": [1, 2, 3]}' -> 'foo' -> 0""", """SELECT '{"foo": [1, 2, 3]}' -> '$.foo' -> '$[0]'""",