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