Repository: tobymao/sqlglot
Pull request: fix(duckdb): Slice + Array bug
Fixes #3136 

Caught 3 consecutive bugs:

- When matching `R_BRACKET` or `R_BRACE` inside `parse_bracket`, the short-circuit should be flipped, otherwise a single `parse_bracket` call can match 2 closing tokens, e.g. in `{'x': arr[i]}`, the parsing of `arr[i]` will also match the following `}` 

- When transforming `exp.Slice` to `exp.PropertyEq`, the `exp.Slice` name was used instead of `exp.Slice.this` name

- DuckDB requires wrapped arrays if the indexing is in place, i.e `SELECT [1, 2][1]` does not work, but `SELECT ([1, 2])[1]` does. The function `LIST_VALUE()` is parsed and generated as an `exp.Array` so it's also affected by this behavior; thus, the DuckDB `exp.Bracket` generation was enhanced to wrap arrays.

Docs
----------
- https://duckdb.org/docs/sql/data_types/list.html#:~:text=Lists%20can%20be%20created%20using,use%20the%20list%20aggregate%20function

Problem statement:
valid duckdb query fails to parse
**Fully reproducible code snippet**

This query cannot be parsed by sqlglot, but it is a valid duckdb query:

```
In [18]: import sqlglot as sg, sqlglot.expressions as sge

In [19]: import duckdb

In [20]: sql = """
    ...: SELECT
    ...:   LIST_APPLY(
    ...:     RANGE(
    ...:       1,
    ...:       GREATEST(
    ...:         LENGTH(LIST_VALUE(CAST(1 AS TINYINT), CAST(2 AS TINYINT), CAST(3 AS TINYINT))),
    ...:         LENGTH(LIST_VALUE(CAST(1 AS TINYINT), CAST(2 AS TINYINT), CAST(3 AS TINYINT)))
    ...:       ) + 1
    ...:     ),
    ...:     i -> {'f1': LIST_VALUE(CAST(1 AS TINYINT), CAST(2 AS TINYINT), CAST(3 AS TINYINT))[i], 'f2': LIST_VALUE(CAST(1 AS TINYINT), CAST(2 AS TINYINT), CAST(3 AS TINYINT))[i]}
    ...:   ) AS "ArrayZip()"
    ...: """

In [21]: duckdb.execute(sql).df()
Out[21]:
                                          ArrayZip()
0  [{'f1': 1, 'f2': 1}, {'f1': 2, 'f2': 2}, {'f1'...

In [22]: sg.parse_one(sql, read="duckdb")
...
ParseError: Expected }. Line 12, Col: 2.
  AS TINYINT))[i], 'f2': LIST_VALUE(CAST(1 AS TINYINT), CAST(2 AS TINYINT), CAST(3 AS TINYINT))[i]}
  ) AS "ArrayZip()"
```

Hints from issue discussion:
A smaller example to reproduce this issue `{x: arr[i]}`.
Ah, thanks!

Commit message: Refactoring duckdb's bracket_sql

Patch to review:
diff --git a/sqlglot/dialects/duckdb.py b/sqlglot/dialects/duckdb.py
index dab0775068..9837025ace 100644
--- a/sqlglot/dialects/duckdb.py
+++ b/sqlglot/dialects/duckdb.py
@@ -571,3 +571,9 @@ def generateseries_sql(self, expression: exp.GenerateSeries) -> str:
                 return rename_func("RANGE")(self, expression)
 
             return super().generateseries_sql(expression)
+
+        def bracket_sql(self, expression: exp.Bracket) -> str:
+            if isinstance(expression.this, exp.Array):
+                expression.this.replace(exp.paren(expression.this))
+
+            return super().bracket_sql(expression)
diff --git a/sqlglot/parser.py b/sqlglot/parser.py
index 85422b3c7d..0eaa845030 100644
--- a/sqlglot/parser.py
+++ b/sqlglot/parser.py
@@ -4316,7 +4316,7 @@ def _kv_to_prop_eq(self, expressions: t.List[exp.Expression]) -> t.List[exp.Expr
 
                 if not isinstance(e, exp.PropertyEQ):
                     e = self.expression(
-                        exp.PropertyEQ, this=exp.to_identifier(e.name), expression=e.expression
+                        exp.PropertyEQ, this=exp.to_identifier(e.this.name), expression=e.expression
                     )
 
                 if isinstance(e.this, exp.Column):
@@ -4735,9 +4735,9 @@ def _parse_bracket(self, this: t.Optional[exp.Expression]) -> t.Optional[exp.Exp
             lambda: self._parse_bracket_key_value(is_map=bracket_kind == TokenType.L_BRACE)
         )
 
-        if not self._match(TokenType.R_BRACKET) and bracket_kind == TokenType.L_BRACKET:
+        if bracket_kind == TokenType.L_BRACKET and not self._match(TokenType.R_BRACKET):
             self.raise_error("Expected ]")
-        elif not self._match(TokenType.R_BRACE) and bracket_kind == TokenType.L_BRACE:
+        elif bracket_kind == TokenType.L_BRACE and not self._match(TokenType.R_BRACE):
             self.raise_error("Expected }")
 
         # https://duckdb.org/docs/sql/data_types/struct.html#creating-structs
diff --git a/tests/dialects/test_duckdb.py b/tests/dialects/test_duckdb.py
index 8362ed0902..50ee03a2a1 100644
--- a/tests/dialects/test_duckdb.py
+++ b/tests/dialects/test_duckdb.py
@@ -701,6 +701,18 @@ def test_array_index(self):
             self.validate_identity(
                 "[x.STRING_SPLIT(' ')[1] FOR x IN ['1', '2', 3] IF x.CONTAINS('1')]"
             )
+            self.validate_identity(
+                """SELECT LIST_VALUE(1)[1]""",
+                """SELECT ([1])[1]""",
+            )
+            self.validate_identity(
+                """{'x': LIST_VALUE(1)[1]}""",
+                """{'x': ([1])[1]}""",
+            )
+            self.validate_identity(
+                """SELECT LIST_APPLY(RANGE(1, 4), i -> {'f1': LIST_VALUE(1, 2, 3)[i], 'f2': LIST_VALUE(1, 2, 3)[i]})""",
+                """SELECT LIST_APPLY(RANGE(1, 4), i -> {'f1': ([1, 2, 3])[i], 'f2': ([1, 2, 3])[i]})""",
+            )
 
             self.assertEqual(
                 cm.output,
@@ -711,6 +723,10 @@ def test_array_index(self):
                     "WARNING:sqlglot:Applying array index offset (1)",
                     "WARNING:sqlglot:Applying array index offset (-1)",
                     "WARNING:sqlglot:Applying array index offset (1)",
+                    "WARNING:sqlglot:Applying array index offset (-1)",
+                    "WARNING:sqlglot:Applying array index offset (1)",
+                    "WARNING:sqlglot:Applying array index offset (-1)",
+                    "WARNING:sqlglot:Applying array index offset (1)",
                 ],
             )
 


Code review comments:
diff_hunk: @@ -711,6 +723,10 @@ def test_array_index(self):
                     "WARNING:sqlglot:Applying array index offset (1)",
                     "WARNING:sqlglot:Applying array index offset (-1)",
                     "WARNING:sqlglot:Applying array index offset (1)",
+                    "WARNING:sqlglot:Applying array index offset (-1)",
+                    "WARNING:sqlglot:Applying array index offset (1)",
+                    "WARNING:sqlglot:Applying array index offset (-1)",
+                    "WARNING:sqlglot:Applying array index offset (1)",
path: tests/dialects/test_duckdb.py
text: @user1:
I think this assertion is confusing, the sequence by itself doesn't explain much besides the fact that there are a bunch of warnings logged. I'd rather either 1) get rid of it altogether (so won't need the `as cm` part above), or 2) reduce the scope of the assertion to only a single test so that we observe the dec/increment of the index once.

@author:
By changing the array indexes to be an identifier instead of an integer literal we can get rid of these warnings and maintain the tests, would that work?

@user1:
I'd just do (2) instead of augmenting this list every time we add a new test tbh