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