sqlglot.parsers.duckdb

  1from __future__ import annotations
  2
  3import typing as t
  4
  5from sqlglot import exp, parser
  6from sqlglot.trie import new_trie
  7from sqlglot.dialects.dialect import (
  8    binary_from_function,
  9    build_default_decimal_type,
 10    build_formatted_time,
 11    build_regexp_extract,
 12    date_trunc_to_time,
 13    pivot_column_names,
 14)
 15from sqlglot.helper import seq_get
 16from sqlglot.parser import binary_range_parser
 17from sqlglot.tokens import TokenType
 18from collections.abc import Collection
 19
 20
 21def _build_sort_array_desc(args: list) -> exp.Expr:
 22    return exp.SortArray(this=seq_get(args, 0), asc=exp.false())
 23
 24
 25def _build_array_prepend(args: list) -> exp.Expr:
 26    return exp.ArrayPrepend(this=seq_get(args, 1), expression=seq_get(args, 0))
 27
 28
 29def _build_date_diff(args: list) -> exp.Expr:
 30    return exp.DateDiff(this=seq_get(args, 2), expression=seq_get(args, 1), unit=seq_get(args, 0))
 31
 32
 33def _build_generate_series(end_exclusive: bool = False) -> t.Callable[[list], exp.GenerateSeries]:
 34    def _builder(args: list) -> exp.GenerateSeries:
 35        # Check https://duckdb.org/docs/sql/functions/nested.html#range-functions
 36        if len(args) == 1:
 37            # DuckDB uses 0 as a default for the series' start when it's omitted
 38            args.insert(0, exp.Literal.number("0"))
 39
 40        gen_series = exp.GenerateSeries.from_arg_list(args)
 41        gen_series.set("is_end_exclusive", end_exclusive)
 42
 43        return gen_series
 44
 45    return _builder
 46
 47
 48def _build_make_timestamp(args: list) -> exp.Expr:
 49    if len(args) == 1:
 50        return exp.UnixToTime(this=seq_get(args, 0), scale=exp.UnixToTime.MICROS)
 51
 52    return exp.TimestampFromParts(
 53        year=seq_get(args, 0),
 54        month=seq_get(args, 1),
 55        day=seq_get(args, 2),
 56        hour=seq_get(args, 3),
 57        min=seq_get(args, 4),
 58        sec=seq_get(args, 5),
 59    )
 60
 61
 62def _show_parser(*args: t.Any, **kwargs: t.Any) -> t.Callable[[DuckDBParser], exp.Show]:
 63    def _parse(self: DuckDBParser) -> exp.Show:
 64        return self._parse_show_duckdb(*args, **kwargs)
 65
 66    return _parse
 67
 68
 69def _convert_text_type(dtype: exp.DataType) -> exp.DataType:
 70    dtype.set("expressions", None)
 71    return dtype
 72
 73
 74class DuckDBParser(parser.Parser):
 75    MAP_KEYS_ARE_ARBITRARY_EXPRESSIONS = True
 76    PIVOT_COLUMN_NAMING = "agg_name_if_aliased_or_multiple"
 77
 78    NO_PAREN_FUNCTIONS = {
 79        **parser.Parser.NO_PAREN_FUNCTIONS,
 80        TokenType.LOCALTIME: exp.Localtime,
 81        TokenType.LOCALTIMESTAMP: exp.Localtimestamp,
 82        TokenType.CURRENT_CATALOG: exp.CurrentCatalog,
 83        TokenType.SESSION_USER: exp.SessionUser,
 84    }
 85
 86    BITWISE = {k: v for k, v in parser.Parser.BITWISE.items() if k != TokenType.CARET}
 87
 88    COLUMN_OPERATORS = {
 89        k: v
 90        for k, v in parser.Parser.COLUMN_OPERATORS.items()
 91        if k not in (TokenType.ARROW, TokenType.DARROW)
 92    }
 93
 94    JSON_OPERATORS = {
 95        TokenType.ARROW: parser.build_json_extract,
 96        TokenType.DARROW: parser.build_json_extract_scalar,
 97    }
 98
 99    RANGE_PARSERS = {
100        **parser.Parser.RANGE_PARSERS,
101        TokenType.DAMP: binary_range_parser(exp.ArrayOverlaps),
102        TokenType.CARET_AT: binary_range_parser(exp.StartsWith),
103        TokenType.TILDE: binary_range_parser(exp.RegexpFullMatch),
104    }
105
106    EXPONENT = {
107        **parser.Parser.EXPONENT,
108        TokenType.CARET: exp.Pow,
109        TokenType.DSTAR: exp.Pow,
110    }
111
112    FUNCTIONS_WITH_ALIASED_ARGS = {*parser.Parser.FUNCTIONS_WITH_ALIASED_ARGS, "STRUCT_PACK"}
113
114    SHOW_PARSERS = {
115        "TABLES": _show_parser("TABLES"),
116        "ALL TABLES": _show_parser("ALL TABLES"),
117    }
118
119    FUNCTIONS = {
120        **{k: v for k, v in parser.Parser.FUNCTIONS.items() if k not in ("DATE_SUB", "GLOB")},
121        "ANY_VALUE": lambda args: exp.IgnoreNulls(this=exp.AnyValue.from_arg_list(args)),
122        "ARRAY_PREPEND": _build_array_prepend,
123        "ARRAY_REVERSE_SORT": _build_sort_array_desc,
124        "ARRAY_INTERSECT": lambda args: exp.ArrayIntersect(expressions=args),
125        "ARRAY_SORT": exp.SortArray.from_arg_list,
126        "BIT_AND": exp.BitwiseAndAgg.from_arg_list,
127        "BIT_OR": exp.BitwiseOrAgg.from_arg_list,
128        "BIT_XOR": exp.BitwiseXorAgg.from_arg_list,
129        "CURRENT_LOCALTIMESTAMP": exp.Localtimestamp.from_arg_list,
130        "DATEDIFF": _build_date_diff,
131        "DATE_DIFF": _build_date_diff,
132        "DATE_TRUNC": date_trunc_to_time,
133        "DATETRUNC": date_trunc_to_time,
134        "DECODE": lambda args: exp.Decode(
135            this=seq_get(args, 0), charset=exp.Literal.string("utf-8")
136        ),
137        "EDITDIST3": exp.Levenshtein.from_arg_list,
138        "ENCODE": lambda args: exp.Encode(
139            this=seq_get(args, 0), charset=exp.Literal.string("utf-8")
140        ),
141        "EPOCH": exp.TimeToUnix.from_arg_list,
142        "EPOCH_MS": lambda args: exp.UnixToTime(this=seq_get(args, 0), scale=exp.UnixToTime.MILLIS),
143        "FROM_HEX": exp.Unhex.from_arg_list,
144        "GENERATE_SERIES": _build_generate_series(),
145        "GET_CURRENT_TIME": exp.CurrentTime.from_arg_list,
146        "GET_BIT": lambda args: exp.Getbit(
147            this=seq_get(args, 0), expression=seq_get(args, 1), zero_is_msb=True
148        ),
149        "JARO_WINKLER_SIMILARITY": exp.JarowinklerSimilarity.from_arg_list,
150        "JSON": exp.ParseJSON.from_arg_list,
151        "JSON_ARRAY": lambda args: exp.JSONArray(expressions=args),
152        "JSON_EXTRACT_PATH": parser.build_extract_json_with_path(exp.JSONExtract),
153        "JSON_EXTRACT_STRING": parser.build_extract_json_with_path(exp.JSONExtractScalar),
154        "LIST": exp.ArrayAgg.from_arg_list,
155        "LIST_DISTINCT": exp.ArrayDistinct.from_arg_list,
156        "LIST_APPEND": exp.ArrayAppend.from_arg_list,
157        "LIST_CONCAT": parser.build_array_concat,
158        "LIST_CONTAINS": exp.ArrayContains.from_arg_list,
159        "LIST_COSINE_DISTANCE": exp.CosineDistance.from_arg_list,
160        "LIST_DISTANCE": exp.EuclideanDistance.from_arg_list,
161        "LIST_FILTER": exp.ArrayFilter.from_arg_list,
162        "LIST_HAS": exp.ArrayContains.from_arg_list,
163        "LIST_HAS_ANY": exp.ArrayOverlaps.from_arg_list,
164        "LIST_MAX": exp.ArrayMax.from_arg_list,
165        "LIST_MIN": exp.ArrayMin.from_arg_list,
166        "LIST_PREPEND": _build_array_prepend,
167        "LIST_REVERSE_SORT": _build_sort_array_desc,
168        "LIST_SORT": exp.SortArray.from_arg_list,
169        "LIST_TRANSFORM": exp.Transform.from_arg_list,
170        "LIST_VALUE": lambda args: exp.Array(expressions=args),
171        "MAKE_DATE": exp.DateFromParts.from_arg_list,
172        "MAKE_TIME": exp.TimeFromParts.from_arg_list,
173        "MAKE_TIMESTAMP": _build_make_timestamp,
174        "QUANTILE_CONT": exp.PercentileCont.from_arg_list,
175        "QUANTILE_DISC": exp.PercentileDisc.from_arg_list,
176        "RANGE": _build_generate_series(end_exclusive=True),
177        "REGEXP_EXTRACT": build_regexp_extract(exp.RegexpExtract),
178        "REGEXP_EXTRACT_ALL": build_regexp_extract(exp.RegexpExtractAll),
179        "REGEXP_MATCHES": exp.RegexpLike.from_arg_list,
180        "REGEXP_REPLACE": lambda args: exp.RegexpReplace(
181            this=seq_get(args, 0),
182            expression=seq_get(args, 1),
183            replacement=seq_get(args, 2),
184            modifiers=seq_get(args, 3),
185            single_replace=True,
186        ),
187        "SHA256": lambda args: exp.SHA2(this=seq_get(args, 0), length=exp.Literal.number(256)),
188        "STRFTIME": build_formatted_time(exp.TimeToStr),
189        "STRING_SPLIT": exp.Split.from_arg_list,
190        "STRING_SPLIT_REGEX": exp.RegexpSplit.from_arg_list,
191        "STRING_TO_ARRAY": exp.Split.from_arg_list,
192        "STRPTIME": build_formatted_time(exp.StrToTime),
193        "STRUCT_PACK": exp.Struct.from_arg_list,
194        "STR_SPLIT": exp.Split.from_arg_list,
195        "STR_SPLIT_REGEX": exp.RegexpSplit.from_arg_list,
196        "TODAY": exp.CurrentDate.from_arg_list,
197        "TIME_BUCKET": exp.DateBin.from_arg_list,
198        "TO_TIMESTAMP": exp.UnixToTime.from_arg_list,
199        "UNNEST": exp.Explode.from_arg_list,
200        "VERSION": exp.CurrentVersion.from_arg_list,
201        "XOR": binary_from_function(exp.BitwiseXor),
202    }
203
204    FUNCTION_PARSERS = {
205        **{k: v for k, v in parser.Parser.FUNCTION_PARSERS.items() if k != "DECODE"},
206        **dict.fromkeys(
207            ("GROUP_CONCAT", "LISTAGG", "STRINGAGG"), lambda self: self._parse_string_agg()
208        ),
209        "APPROX_QUANTILE": lambda self: self._parse_distinct_arg_function(exp.ApproxQuantile),
210        "QUANTILE": lambda self: self._parse_distinct_arg_function(exp.Quantile),
211        "QUANTILE_CONT": lambda self: self._parse_distinct_arg_function(exp.PercentileCont),
212        "QUANTILE_DISC": lambda self: self._parse_distinct_arg_function(exp.PercentileDisc),
213    }
214
215    NO_PAREN_FUNCTION_PARSERS = {
216        **parser.Parser.NO_PAREN_FUNCTION_PARSERS,
217        "MAP": lambda self: self._parse_map(),
218        "@": lambda self: exp.Abs(this=self._parse_bitwise()),
219    }
220
221    PLACEHOLDER_PARSERS = {
222        **parser.Parser.PLACEHOLDER_PARSERS,
223        TokenType.PARAMETER: lambda self: (
224            self.expression(exp.Placeholder(this=self._prev.text))
225            if self._match(TokenType.NUMBER) or self._match_set(self.ID_VAR_TOKENS)
226            else None
227        ),
228    }
229
230    TYPE_CONVERTERS = {
231        # https://duckdb.org/docs/sql/data_types/numeric
232        exp.DType.DECIMAL: build_default_decimal_type(precision=18, scale=3),
233        # https://duckdb.org/docs/sql/data_types/text
234        exp.DType.TEXT: _convert_text_type,
235    }
236
237    STATEMENT_PARSERS = {
238        **parser.Parser.STATEMENT_PARSERS,
239        TokenType.ATTACH: lambda self: self._parse_attach_detach(),
240        TokenType.DETACH: lambda self: self._parse_attach_detach(is_attach=False),
241        TokenType.FORCE: lambda self: self._parse_force(),
242        TokenType.INSTALL: lambda self: self._parse_install(),
243        TokenType.SHOW: lambda self: self._parse_show(),
244    }
245
246    SET_PARSERS = {
247        **parser.Parser.SET_PARSERS,
248        "VARIABLE": lambda self: self._parse_set_item_assignment("VARIABLE"),
249    }
250
251    SHOW_TRIE = new_trie(key.split(" ") for key in SHOW_PARSERS)
252    SET_TRIE = new_trie(key.split(" ") for key in SET_PARSERS)
253
254    def _parse_function_properties(self) -> exp.Properties | None:
255        if self._match(TokenType.TABLE):
256            return exp.Properties(
257                expressions=[
258                    exp.ReturnsProperty(
259                        this=exp.Schema(this=exp.var("TABLE")),
260                        is_table=True,
261                    ),
262                ]
263            )
264        return super()._parse_function_properties()
265
266    def _parse_lambda(self, alias: bool = False) -> exp.Expr | None:
267        index = self._index
268        if not self._match_text_seq("LAMBDA"):
269            return super()._parse_lambda(alias=alias)
270
271        expressions = self._parse_csv(self._parse_lambda_arg)
272        if not self._match(TokenType.COLON):
273            self._retreat(index)
274            return None
275
276        this = self._replace_lambda(self._parse_assignment(), expressions)
277        return self.expression(exp.Lambda(this=this, expressions=expressions, colon=True))
278
279    def _parse_expression(self) -> exp.Expr | None:
280        # DuckDB supports prefix aliases, e.g. foo: 1
281        if self._next.token_type == TokenType.COLON:
282            alias = self._parse_id_var(tokens=self.ALIAS_TOKENS)
283            self._match(TokenType.COLON)
284            comments = self._prev_comments
285
286            this = self._parse_assignment()
287            if isinstance(this, exp.Expr):
288                # Moves the comment next to the alias in `alias: expr /* comment */`
289                comments += this.pop_comments() or []
290
291            return self.expression(exp.Alias(this=this, alias=alias), comments=comments)
292
293        return super()._parse_expression()
294
295    def _parse_table(
296        self,
297        schema: bool = False,
298        joins: bool = False,
299        alias_tokens: Collection[TokenType] | None = None,
300        parse_bracket: bool = False,
301        is_db_reference: bool = False,
302        parse_partition: bool = False,
303        consume_pipe: bool = False,
304    ) -> exp.Expr | None:
305        # DuckDB supports prefix aliases, e.g. FROM foo: bar
306        if self._next.token_type == TokenType.COLON:
307            alias = self._parse_table_alias(alias_tokens=alias_tokens or self.TABLE_ALIAS_TOKENS)
308            self._match(TokenType.COLON)
309            comments = self._prev_comments
310        else:
311            alias = None
312            comments = []
313
314        table = super()._parse_table(
315            schema=schema,
316            joins=joins,
317            alias_tokens=alias_tokens,
318            parse_bracket=parse_bracket,
319            is_db_reference=is_db_reference,
320            parse_partition=parse_partition,
321        )
322        if isinstance(table, exp.Expr) and isinstance(alias, exp.TableAlias):
323            # Moves the comment next to the alias in `alias: table /* comment */`
324            comments += table.pop_comments() or []
325            alias.comments = alias.pop_comments() + comments
326            table.set("alias", alias)
327
328        return table
329
330    def _parse_table_sample(self, as_modifier: bool = False) -> exp.TableSample | None:
331        # https://duckdb.org/docs/sql/samples.html
332        sample = super()._parse_table_sample(as_modifier=as_modifier)
333        if sample and not sample.args.get("method"):
334            if sample.args.get("size"):
335                sample.set("method", exp.var("RESERVOIR"))
336            else:
337                sample.set("method", exp.var("SYSTEM"))
338
339        return sample
340
341    def _parse_bracket(self, this: exp.Expr | None = None) -> exp.Expr | None:
342        bracket = super()._parse_bracket(this)
343
344        if self.dialect.version < (1, 2) and isinstance(bracket, exp.Bracket):
345            # https://duckdb.org/2025/02/05/announcing-duckdb-120.html#breaking-changes
346            bracket.set("returns_list_for_maps", True)
347
348        return bracket
349
350    def _parse_map(self) -> exp.ToMap | exp.Map:
351        if self._match(TokenType.L_BRACE, advance=False):
352            return self.expression(exp.ToMap(this=self._parse_bracket()))
353
354        args = self._parse_wrapped_csv(self._parse_assignment)
355        return self.expression(exp.Map(keys=seq_get(args, 0), values=seq_get(args, 1)))
356
357    def _parse_struct_types(self, type_required: bool = False) -> exp.Expr | None:
358        return self._parse_field_def()
359
360    def _pivot_column_names(self, aggregations: list[exp.Expr]) -> list[str]:
361        if len(aggregations) == 1:
362            return super()._pivot_column_names(aggregations)
363        return pivot_column_names(aggregations, dialect="duckdb")
364
365    def _parse_attach_detach(self, is_attach: bool = True) -> exp.Attach | exp.Detach:
366        def _parse_attach_option() -> exp.AttachOption:
367            return self.expression(
368                exp.AttachOption(
369                    this=self._parse_var(any_token=True),
370                    expression=self._parse_field(any_token=True),
371                )
372            )
373
374        self._match(TokenType.DATABASE)
375        exists = self._parse_exists(not_=is_attach)
376        this = self._parse_alias(self._parse_primary_or_var(), explicit=True)
377
378        if self._match(TokenType.L_PAREN, advance=False):
379            expressions = self._parse_wrapped_csv(_parse_attach_option)
380        else:
381            expressions = None
382
383        return (
384            self.expression(exp.Attach(this=this, exists=exists, expressions=expressions))
385            if is_attach
386            else self.expression(exp.Detach(this=this, exists=exists))
387        )
388
389    def _parse_show_duckdb(self, this: str) -> exp.Show:
390        from_ = self._parse_table(schema=True) if self._match(TokenType.FROM) else None
391        return self.expression(exp.Show(this=this, from_=from_))
392
393    def _parse_force(self) -> exp.Install | exp.Command:
394        # FORCE can only be followed by INSTALL or CHECKPOINT
395        # In the case of CHECKPOINT, we fallback
396        if not self._match(TokenType.INSTALL):
397            return self._parse_as_command(self._prev)
398
399        return self._parse_install(force=True)
400
401    def _parse_install(self, force: bool = False) -> exp.Install:
402        return self.expression(
403            exp.Install(
404                this=self._parse_id_var(),
405                from_=self._parse_var_or_string() if self._match(TokenType.FROM) else None,
406                force=force,
407            )
408        )
409
410    def _parse_primary(self) -> exp.Expr | None:
411        if self._match_pair(TokenType.HASH, TokenType.NUMBER):
412            return exp.PositionalColumn(this=exp.Literal.number(self._prev.text))
413
414        return super()._parse_primary()