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"]
193
194
195class StrtokToArray(Expression, Func):
196    arg_types = {"this": True, "expression": False}
197
198
199# Higher-order / lambda
200
201
202class Apply(Expression, Func):
203    arg_types = {"this": True, "expression": True}
204
205
206class Reduce(Expression, Func):
207    arg_types = {"this": True, "initial": True, "merge": True, "finish": False}
208
209
210class Transform(Expression, Func):
211    arg_types = {"this": True, "expression": True}
212
213
214# Table-valued / UDTF
215
216
217class GenerateSeries(Expression, Func):
218    arg_types = {"start": True, "end": True, "step": False, "is_end_exclusive": False}
219
220
221class ExplodingGenerateSeries(GenerateSeries):
222    pass
223
224
225class Generator(Expression, Func, UDTF):
226    arg_types = {"rowcount": False, "timelimit": False}
227
228
229class Explode(Expression, Func, UDTF):
230    arg_types = {"this": True, "expressions": False}
231    is_var_len_args = True
232
233
234class Inline(Expression, Func):
235    pass
236
237
238@trait
239class ExplodeOuter(Expr):
240    pass
241
242
243class _ExplodeOuter(Explode, ExplodeOuter):
244    _sql_names = ["EXPLODE_OUTER"]
245
246
247class Posexplode(Explode):
248    pass
249
250
251class PosexplodeOuter(Posexplode, ExplodeOuter):
252    pass
253
254
255class PositionalColumn(Expression):
256    pass
257
258
259class Unnest(Expression, Func, UDTF):
260    arg_types = {
261        "expressions": True,
262        "alias": False,
263        "offset": False,
264        "explode_array": False,
265    }
266
267    @property
268    def selects(self) -> list[Expr]:
269        columns = super().selects
270        offset = self.args.get("offset")
271        if offset:
272            columns = columns + [to_identifier("offset") if offset is True else offset]
273        return columns
274
275
276# Map
277
278
279class Map(Expression, Func):
280    arg_types = {"keys": False, "values": False}
281
282    @property
283    def keys(self) -> list[Expr]:
284        keys = self.args.get("keys")
285        return keys.expressions if keys else []
286
287    @property
288    def values(self) -> list[Expr]:
289        values = self.args.get("values")
290        return values.expressions if values else []
291
292
293class MapCat(Expression, Func):
294    arg_types = {"this": True, "expression": True}
295
296
297class MapContainsKey(Expression, Func):
298    arg_types = {"this": True, "key": True}
299
300
301class MapDelete(Expression, Func):
302    arg_types = {"this": True, "expressions": True}
303    is_var_len_args = True
304
305
306class MapFromEntries(Expression, Func):
307    pass
308
309
310class MapInsert(Expression, Func):
311    arg_types = {"this": True, "key": False, "value": True, "update_flag": False}
312
313
314class MapKeys(Expression, Func):
315    pass
316
317
318class MapPick(Expression, Func):
319    arg_types = {"this": True, "expressions": True}
320    is_var_len_args = True
321
322
323class MapSize(Expression, Func):
324    pass
325
326
327class StarMap(Expression, Func):
328    pass
329
330
331class ToMap(Expression, Func):
332    pass
333
334
335class VarMap(Expression, Func):
336    arg_types = {"keys": True, "values": True}
337    is_var_len_args = True
338
339    @property
340    def keys(self) -> list[Expr]:
341        return self.args["keys"].expressions
342
343    @property
344    def values(self) -> list[Expr]:
345        return self.args["values"].expressions
346
347
348# Struct
349
350
351class Struct(Expression, Func):
352    arg_types = {"expressions": False}
353    is_var_len_args = True
354
355
356class StructExtract(Expression, Func):
357    arg_types = {"this": True, "expression": True}
358
359
360# Geospatial
361
362
363class StDistance(Expression, Func):
364    arg_types = {"this": True, "expression": True, "use_spheroid": False}
365
366
367class StPoint(Expression, Func):
368    arg_types = {"this": True, "expression": True, "z": False, "m": False}
369    _sql_names = ["ST_POINT", "ST_MAKEPOINT"]