1"""sqlglot expressions - array, map, struct, and table-valued functions."""
  2
  3from __future__ import annotations
  4
  5
  6from sqlglot.expressions.core import (
  7    Expression,
  8    Expr,
  9    Func,
 10    Binary,
 11    Predicate,
 12    to_identifier,
 13)
 14from sqlglot.helper import trait
 15from sqlglot.expressions.query import UDTF
 16
 17
 18# Array creation / construction
 19
 20
 21class Array(Expression, Func):
 22    arg_types = {
 23        "expressions": False,
 24        "bracket_notation": False,
 25        "struct_name_inheritance": False,
 26        "value_constructor": False,
 27    }
 28    is_var_len_args = True
 29
 30
 31class ArrayConstructCompact(Expression, Func):
 32    arg_types = {"expressions": False}
 33    is_var_len_args = True
 34
 35
 36class List(Expression, Func):
 37    arg_types = {"expressions": False}
 38    is_var_len_args = True
 39
 40
 41class ToArray(Expression, Func):
 42    pass
 43
 44
 45# Array manipulation
 46
 47
 48class ArrayAppend(Expression, Func):
 49    arg_types = {"this": True, "expression": True, "null_propagation": False}
 50
 51
 52class ArrayCompact(Expression, Func):
 53    pass
 54
 55
 56class ArrayConcat(Expression, Func):
 57    _sql_names = ["ARRAY_CONCAT", "ARRAY_CAT"]
 58    arg_types = {"this": True, "expressions": False, "null_propagation": False}
 59    is_var_len_args = True
 60
 61
 62class ArrayFilter(Expression, Func):
 63    arg_types = {"this": True, "expression": True}
 64    _sql_names = ["FILTER", "ARRAY_FILTER"]
 65
 66
 67class ArrayInsert(Expression, Func):
 68    arg_types = {"this": True, "position": True, "expression": True, "offset": False}
 69
 70
 71class ArrayPrepend(Expression, Func):
 72    arg_types = {"this": True, "expression": True, "null_propagation": False}
 73
 74
 75class ArrayRemove(Expression, Func):
 76    arg_types = {"this": True, "expression": True, "null_propagation": False}
 77
 78
 79class ArrayRemoveAt(Expression, Func):
 80    arg_types = {"this": True, "position": True}
 81
 82
 83class ArrayReverse(Expression, Func):
 84    pass
 85
 86
 87class ArraySlice(Expression, Func):
 88    """Slice by end index, or by length when is_length is set."""
 89
 90    arg_types = {
 91        "this": True,
 92        "start": True,
 93        "end": False,
 94        "step": False,
 95        "zero_based": False,
 96        "is_length": False,
 97    }
 98
 99
100class ArraySort(Expression, Func):
101    arg_types = {"this": True, "expression": False}
102
103
104class Shuffle(Expression, Func):
105    pass
106
107
108class SortArray(Expression, Func):
109    arg_types = {"this": True, "asc": False, "nulls_first": False}
110
111
112# Array predicates / search
113
114
115class ArrayAll(Expression, Func):
116    arg_types = {"this": True, "expression": True}
117
118
119class ArrayAny(Expression, Func):
120    arg_types = {"this": True, "expression": True}
121
122
123class ArrayContains(Expression, Binary, Predicate, Func):
124    arg_types = {"this": True, "expression": True, "ensure_variant": False, "check_null": False}
125    _sql_names = ["ARRAY_CONTAINS", "ARRAY_HAS"]
126
127
128class ArrayContainsAll(Expression, Binary, Predicate, Func):
129    _sql_names = ["ARRAY_CONTAINS_ALL", "ARRAY_HAS_ALL"]
130
131
132class ArrayContainedBy(Expression, Binary, Predicate, Func):
133    pass
134
135
136class ArrayExcept(Expression, Func):
137    arg_types = {"this": True, "expression": True, "is_multiset": False}
138
139
140class ArrayIntersect(Expression, Func):
141    arg_types = {"expressions": True, "is_multiset": False}
142    is_var_len_args = True
143    _sql_names = ["ARRAY_INTERSECT", "ARRAY_INTERSECTION"]
144
145
146class ArrayOverlaps(Expression, Binary, Predicate, Func):
147    arg_types = {"this": True, "expression": True, "null_safe": False}
148
149
150class ArrayPosition(Expression, Binary, Func):
151    arg_types = {"this": True, "expression": True, "zero_based": False}
152
153
154# Array properties
155
156
157class ArrayDistinct(Expression, Func):
158    arg_types = {"this": True, "check_null": False}
159
160
161class ArrayFirst(Expression, Func):
162    arg_types = {"this": True, "expression": False}
163
164
165class ArrayLast(Expression, Func):
166    pass
167
168
169class ArrayMax(Expression, Func):
170    pass
171
172
173class ArrayMin(Expression, Func):
174    pass
175
176
177class ArraySize(Expression, Func):
178    arg_types = {"this": True, "expression": False}
179    _sql_names = ["ARRAY_SIZE", "ARRAY_LENGTH"]
180
181
182class ArraySum(Expression, Func):
183    arg_types = {"this": True, "expression": False}
184
185
186# Array conversion / utility
187
188
189class ArraysZip(Expression, Func):
190    arg_types = {"expressions": False}
191    is_var_len_args = True
192
193
194class ArrayToString(Expression, Func):
195    arg_types = {
196        "this": True,
197        "expression": True,
198        "null": False,
199        "null_is_empty": False,
200        "null_delim_is_null": False,
201    }
202    _sql_names = ["ARRAY_TO_STRING", "ARRAY_JOIN"]
203
204
205class Flatten(Expression, Func):
206    pass
207
208
209class StringToArray(Expression, Func):
210    arg_types = {"this": True, "expression": False, "null": False}
211    _sql_names = ["STRING_TO_ARRAY", "SPLIT_BY_STRING"]
212
213
214class StrtokToArray(Expression, Func):
215    arg_types = {"this": True, "expression": False}
216
217
218# Higher-order / lambda
219
220
221class Apply(Expression, Func):
222    arg_types = {"this": True, "expression": True}
223
224
225class Reduce(Expression, Func):
226    arg_types = {"this": True, "initial": True, "merge": True, "finish": False}
227
228
229class Transform(Expression, Func):
230    arg_types = {"this": True, "expression": True}
231
232
233# Table-valued / UDTF
234
235
236class GenerateSeries(Expression, Func):
237    arg_types = {"start": True, "end": True, "step": False, "is_end_exclusive": False}
238
239
240class ExplodingGenerateSeries(GenerateSeries):
241    pass
242
243
244class Generator(Expression, Func, UDTF):
245    arg_types = {"rowcount": False, "timelimit": False}
246
247
248class Explode(Expression, Func, UDTF):
249    arg_types = {"this": True, "expressions": False}
250    is_var_len_args = True
251
252
253class Inline(Expression, Func, UDTF):
254    pass
255
256
257class SplitToTable(Expression, Func, UDTF):
258    arg_types = {"this": True, "expression": True}
259
260
261class Stack(Expression, Func, UDTF):
262    arg_types = {"this": True, "expressions": True}
263    is_var_len_args = True
264
265
266@trait
267class ExplodeOuter(Expr):
268    pass
269
270
271class _ExplodeOuter(Explode, ExplodeOuter):
272    _sql_names = ["EXPLODE_OUTER"]
273
274
275class Posexplode(Explode):
276    pass
277
278
279class PosexplodeOuter(Posexplode, ExplodeOuter):
280    pass
281
282
283class PositionalColumn(Expression):
284    pass
285
286
287class Unnest(Expression, Func, UDTF):
288    arg_types = {
289        "expressions": True,
290        "alias": False,
291        "offset": False,
292        "explode_array": False,
293    }
294
295    @property
296    def selects(self) -> list[Expr]:
297        columns = super().selects
298        offset = self.args.get("offset")
299        if offset:
300            columns = columns + [to_identifier("offset") if offset is True else offset]
301        return columns
302
303
304# Map
305
306
307class Map(Expression, Func):
308    arg_types = {"keys": False, "values": False}
309
310    @property
311    def keys(self) -> list[Expr]:
312        keys = self.args.get("keys")
313        return keys.expressions if keys else []
314
315    @property
316    def values(self) -> list[Expr]:
317        values = self.args.get("values")
318        return values.expressions if values else []
319
320
321class MapCat(Expression, Func):
322    arg_types = {"this": True, "expression": True}
323
324
325class MapContainsKey(Expression, Func):
326    arg_types = {"this": True, "key": True}
327
328
329class MapDelete(Expression, Func):
330    arg_types = {"this": True, "expressions": True}
331    is_var_len_args = True
332
333
334class MapFromEntries(Expression, Func):
335    pass
336
337
338class MapInsert(Expression, Func):
339    arg_types = {"this": True, "key": False, "value": True, "update_flag": False}
340
341
342class MapKeys(Expression, Func):
343    pass
344
345
346class MapPick(Expression, Func):
347    arg_types = {"this": True, "expressions": True}
348    is_var_len_args = True
349
350
351class MapSize(Expression, Func):
352    pass
353
354
355class StarMap(Expression, Func):
356    pass
357
358
359class ToMap(Expression, Func):
360    pass
361
362
363class VarMap(Expression, Func):
364    arg_types = {"keys": True, "values": True}
365    is_var_len_args = True
366    var_len_arg_key = "values"
367
368    @property
369    def keys(self) -> list[Expr]:
370        return self.args["keys"].expressions
371
372    @property
373    def values(self) -> list[Expr]:
374        return self.args["values"].expressions
375
376
377# Struct
378
379
380class Struct(Expression, Func):
381    arg_types = {"expressions": False}
382    is_var_len_args = True
383
384
385class StructExtract(Expression, Func):
386    arg_types = {"this": True, "expression": True}
387
388
389# Geospatial
390
391
392class StDistance(Expression, Func):
393    arg_types = {"this": True, "expression": True, "use_spheroid": False}
394
395
396class StPoint(Expression, Func):
397    arg_types = {"this": True, "expression": True, "z": False, "m": False}
398    _sql_names = ["ST_POINT", "ST_MAKEPOINT"]