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, UDTF):
244    pass
245
246
247class Stack(Expression, Func, UDTF):
248    arg_types = {"this": True, "expressions": True}
249    is_var_len_args = True
250
251
252@trait
253class ExplodeOuter(Expr):
254    pass
255
256
257class _ExplodeOuter(Explode, ExplodeOuter):
258    _sql_names = ["EXPLODE_OUTER"]
259
260
261class Posexplode(Explode):
262    pass
263
264
265class PosexplodeOuter(Posexplode, ExplodeOuter):
266    pass
267
268
269class PositionalColumn(Expression):
270    pass
271
272
273class Unnest(Expression, Func, UDTF):
274    arg_types = {
275        "expressions": True,
276        "alias": False,
277        "offset": False,
278        "explode_array": False,
279    }
280
281    @property
282    def selects(self) -> list[Expr]:
283        columns = super().selects
284        offset = self.args.get("offset")
285        if offset:
286            columns = columns + [to_identifier("offset") if offset is True else offset]
287        return columns
288
289
290# Map
291
292
293class Map(Expression, Func):
294    arg_types = {"keys": False, "values": False}
295
296    @property
297    def keys(self) -> list[Expr]:
298        keys = self.args.get("keys")
299        return keys.expressions if keys else []
300
301    @property
302    def values(self) -> list[Expr]:
303        values = self.args.get("values")
304        return values.expressions if values else []
305
306
307class MapCat(Expression, Func):
308    arg_types = {"this": True, "expression": True}
309
310
311class MapContainsKey(Expression, Func):
312    arg_types = {"this": True, "key": True}
313
314
315class MapDelete(Expression, Func):
316    arg_types = {"this": True, "expressions": True}
317    is_var_len_args = True
318
319
320class MapFromEntries(Expression, Func):
321    pass
322
323
324class MapInsert(Expression, Func):
325    arg_types = {"this": True, "key": False, "value": True, "update_flag": False}
326
327
328class MapKeys(Expression, Func):
329    pass
330
331
332class MapPick(Expression, Func):
333    arg_types = {"this": True, "expressions": True}
334    is_var_len_args = True
335
336
337class MapSize(Expression, Func):
338    pass
339
340
341class StarMap(Expression, Func):
342    pass
343
344
345class ToMap(Expression, Func):
346    pass
347
348
349class VarMap(Expression, Func):
350    arg_types = {"keys": True, "values": True}
351    is_var_len_args = True
352    var_len_arg_key = "values"
353
354    @property
355    def keys(self) -> list[Expr]:
356        return self.args["keys"].expressions
357
358    @property
359    def values(self) -> list[Expr]:
360        return self.args["values"].expressions
361
362
363# Struct
364
365
366class Struct(Expression, Func):
367    arg_types = {"expressions": False}
368    is_var_len_args = True
369
370
371class StructExtract(Expression, Func):
372    arg_types = {"this": True, "expression": True}
373
374
375# Geospatial
376
377
378class StDistance(Expression, Func):
379    arg_types = {"this": True, "expression": True, "use_spheroid": False}
380
381
382class StPoint(Expression, Func):
383    arg_types = {"this": True, "expression": True, "z": False, "m": False}
384    _sql_names = ["ST_POINT", "ST_MAKEPOINT"]