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