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, "asc": False, "nulls_first": False}
102
103
104class Shuffle(Expression, Func):
105    pass
106
107
108# Array predicates / search
109
110
111class ArrayAll(Expression, Func):
112    arg_types = {"this": True, "expression": True}
113
114
115class ArrayAny(Expression, Func):
116    arg_types = {"this": True, "expression": True}
117
118
119class ArrayContains(Expression, Binary, Predicate, Func):
120    arg_types = {"this": True, "expression": True, "ensure_variant": False, "check_null": False}
121    _sql_names = ["ARRAY_CONTAINS", "ARRAY_HAS"]
122
123
124class ArrayContainsAll(Expression, Binary, Predicate, Func):
125    _sql_names = ["ARRAY_CONTAINS_ALL", "ARRAY_HAS_ALL"]
126
127
128class ArrayContainedBy(Expression, Binary, Predicate, Func):
129    pass
130
131
132class ArrayExcept(Expression, Func):
133    arg_types = {"this": True, "expression": True, "is_multiset": False}
134
135
136class ArrayIntersect(Expression, Func):
137    arg_types = {"expressions": True, "is_multiset": False}
138    is_var_len_args = True
139    _sql_names = ["ARRAY_INTERSECT", "ARRAY_INTERSECTION"]
140
141
142class ArrayOverlaps(Expression, Binary, Predicate, Func):
143    arg_types = {"this": True, "expression": True, "null_safe": False}
144
145
146class ArrayPosition(Expression, Binary, Func):
147    arg_types = {"this": True, "expression": True, "zero_based": False}
148
149
150# Array properties
151
152
153class ArrayDistinct(Expression, Func):
154    arg_types = {"this": True, "check_null": False}
155
156
157class ArrayFirst(Expression, Func):
158    arg_types = {"this": True, "expression": False}
159
160
161class ArrayLast(Expression, Func):
162    pass
163
164
165class ArrayMax(Expression, Func):
166    pass
167
168
169class ArrayMin(Expression, Func):
170    pass
171
172
173class ArraySize(Expression, Func):
174    arg_types = {"this": True, "expression": False}
175    _sql_names = ["ARRAY_SIZE", "ARRAY_LENGTH"]
176
177
178class ArraySum(Expression, Func):
179    arg_types = {"this": True, "expression": False}
180
181
182# Array conversion / utility
183
184
185class ArraysZip(Expression, Func):
186    arg_types = {"expressions": False}
187    is_var_len_args = True
188
189
190class ArrayToString(Expression, Func):
191    arg_types = {
192        "this": True,
193        "expression": True,
194        "null": False,
195        "null_is_empty": False,
196        "null_delim_is_null": False,
197    }
198    _sql_names = ["ARRAY_TO_STRING", "ARRAY_JOIN"]
199
200
201class Flatten(Expression, Func):
202    pass
203
204
205class StringToArray(Expression, Func):
206    arg_types = {"this": True, "expression": False, "null": False}
207    _sql_names = ["STRING_TO_ARRAY", "SPLIT_BY_STRING"]
208
209
210class StrtokToArray(Expression, Func):
211    arg_types = {"this": True, "expression": False}
212
213
214# Higher-order / lambda
215
216
217class Apply(Expression, Func):
218    arg_types = {"this": True, "expression": True}
219
220
221class Reduce(Expression, Func):
222    arg_types = {"this": True, "initial": True, "merge": True, "finish": False}
223
224
225class Transform(Expression, Func):
226    arg_types = {"this": True, "expression": True}
227
228
229# Table-valued / UDTF
230
231
232class GenerateSeries(Expression, Func):
233    arg_types = {"start": True, "end": True, "step": False, "is_end_exclusive": False}
234
235
236class ExplodingGenerateSeries(GenerateSeries):
237    pass
238
239
240class Generator(Expression, Func, UDTF):
241    arg_types = {"rowcount": False, "timelimit": False}
242
243
244class Explode(Expression, Func, UDTF):
245    arg_types = {"this": True, "expressions": False}
246    is_var_len_args = True
247
248
249class Inline(Expression, Func, UDTF):
250    pass
251
252
253class SplitToTable(Expression, Func, UDTF):
254    arg_types = {"this": True, "expression": True}
255
256
257class Stack(Expression, Func, UDTF):
258    arg_types = {"this": True, "expressions": True}
259    is_var_len_args = True
260
261
262@trait
263class ExplodeOuter(Expr):
264    pass
265
266
267class _ExplodeOuter(Explode, ExplodeOuter):
268    _sql_names = ["EXPLODE_OUTER"]
269
270
271class Posexplode(Explode):
272    pass
273
274
275class PosexplodeOuter(Posexplode, ExplodeOuter):
276    pass
277
278
279class PositionalColumn(Expression):
280    pass
281
282
283class Unnest(Expression, Func, UDTF):
284    arg_types = {
285        "expressions": True,
286        "alias": False,
287        "offset": False,
288        "explode_array": False,
289    }
290
291    @property
292    def selects(self) -> list[Expr]:
293        columns = super().selects
294        offset = self.args.get("offset")
295        if offset:
296            columns = columns + [to_identifier("offset") if offset is True else offset]
297        return columns
298
299
300# Map
301
302
303class Map(Expression, Func):
304    arg_types = {"keys": False, "values": False}
305
306    @property
307    def keys(self) -> list[Expr]:
308        keys = self.args.get("keys")
309        return keys.expressions if keys else []
310
311    @property
312    def values(self) -> list[Expr]:
313        values = self.args.get("values")
314        return values.expressions if values else []
315
316
317class MapCat(Expression, Func):
318    arg_types = {"this": True, "expression": True}
319
320
321class MapContainsKey(Expression, Func):
322    arg_types = {"this": True, "key": True}
323
324
325class MapDelete(Expression, Func):
326    arg_types = {"this": True, "expressions": True}
327    is_var_len_args = True
328
329
330class MapFromEntries(Expression, Func):
331    pass
332
333
334class MapInsert(Expression, Func):
335    arg_types = {"this": True, "key": False, "value": True, "update_flag": False}
336
337
338class MapKeys(Expression, Func):
339    pass
340
341
342class MapPick(Expression, Func):
343    arg_types = {"this": True, "expressions": True}
344    is_var_len_args = True
345
346
347class MapSize(Expression, Func):
348    pass
349
350
351class StarMap(Expression, Func):
352    pass
353
354
355class ToMap(Expression, Func):
356    pass
357
358
359class VarMap(Expression, Func):
360    arg_types = {"keys": True, "values": True}
361    is_var_len_args = True
362    var_len_arg_key = "values"
363
364    @property
365    def keys(self) -> list[Expr]:
366        return self.args["keys"].expressions
367
368    @property
369    def values(self) -> list[Expr]:
370        return self.args["values"].expressions
371
372
373# Struct
374
375
376class Struct(Expression, Func):
377    arg_types = {"expressions": False}
378    is_var_len_args = True
379
380
381class StructExtract(Expression, Func):
382    arg_types = {"this": True, "expression": True}
383
384
385# Geospatial
386
387
388class StDistance(Expression, Func):
389    arg_types = {"this": True, "expression": True, "use_spheroid": False}
390
391
392class StPoint(Expression, Func):
393    arg_types = {"this": True, "expression": True, "z": False, "m": False}
394    _sql_names = ["ST_POINT", "ST_MAKEPOINT"]