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