-
Notifications
You must be signed in to change notification settings - Fork 44
/
arrays.py
362 lines (272 loc) · 10.5 KB
/
arrays.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
from pysparkling.sql.expressions.expressions import Expression, UnaryExpression
from pysparkling.sql.utils import AnalysisException
class ArraysOverlap(Expression):
def __init__(self, array1, array2):
super(ArraysOverlap, self).__init__(array1, array2)
self.array1 = array1
self.array2 = array2
def eval(self, row, schema):
set1 = set(self.array1.eval(row, schema))
set2 = set(self.array2.eval(row, schema))
overlap = set1 & set2
if len(overlap) > 1 or (len(overlap) == 1 and None not in overlap):
return True
if set1 and set2 and (None in set1 or None in set2):
return None
return False
def __str__(self):
return "array_overlap({0}, {1})".format(self.array1, self.array2)
class ArrayContains(Expression):
def __init__(self, array, value):
self.array = array
self.value = value # not a column
super(ArrayContains, self).__init__(array)
def eval(self, row, schema):
array_eval = self.array.eval(row, schema)
if array_eval is None:
return None
return self.value in array_eval
def __str__(self):
return "array_contains({0}, {1})".format(self.array, self.value)
class ArrayColumn(Expression):
def __init__(self, columns):
super(ArrayColumn, self).__init__(columns)
self.columns = columns
def eval(self, row, schema):
return [col.eval(row, schema) for col in self.columns]
def __str__(self):
return "array({0})".format(", ".join(str(col) for col in self.columns))
class MapColumn(Expression):
def __init__(self, columns):
super(MapColumn, self).__init__(columns)
self.columns = columns
self.keys = columns[::2]
self.values = columns[1::2]
def eval(self, row, schema):
return dict(
(key.eval(row, schema), value.eval(row, schema))
for key, value in zip(self.keys, self.values)
)
def __str__(self):
return "map({0})".format(", ".join(str(col) for col in self.columns))
class MapFromArraysColumn(Expression):
def __init__(self, keys, values):
super(MapFromArraysColumn, self).__init__(keys, values)
self.keys = keys
self.values = values
def eval(self, row, schema):
return dict(
zip(self.keys.eval(row, schema), self.values.eval(row, schema))
)
def __str__(self):
return "map_from_arrays({0}, {1})".format(
self.keys,
self.values
)
class Size(UnaryExpression):
def eval(self, row, schema):
column_value = self.column.eval(row, schema)
if isinstance(column_value, (list, dict)):
return len(column_value)
raise AnalysisException(
"{0} value should be an array or a map, got {1}".format(
self.column,
type(column_value)
)
)
def __str__(self):
return "size({0})".format(self.column)
class ArraySort(UnaryExpression):
def eval(self, row, schema):
return sorted(self.column.eval(row, schema))
def __str__(self):
return "array_sort({0})".format(self.column)
class ArrayMin(UnaryExpression):
def eval(self, row, schema):
return min(self.column.eval(row, schema))
def __str__(self):
return "array_min({0})".format(self.column)
class ArrayMax(UnaryExpression):
def eval(self, row, schema):
return max(self.column.eval(row, schema))
def __str__(self):
return "array_max({0})".format(self.column)
class Slice(Expression):
def __init__(self, x, start, length):
self.x = x
self.start = start
self.length = length
super(Slice, self).__init__(x)
def eval(self, row, schema):
return self.x.eval(row, schema)[self.start, self.start + self.length]
def __str__(self):
return "slice({0}, {1}, {2})".format(self.x, self.start, self.length)
class ArrayRepeat(Expression):
def __init__(self, col, count):
super(ArrayRepeat, self).__init__(col)
self.col = col
self.count = count
def eval(self, row, schema):
value = self.col.eval(row, schema)
return [value for _ in range(self.count)]
def __str__(self):
return "array_repeat({0}, {1})".format(self.col, self.count)
class Sequence(Expression):
def __init__(self, start, stop, step):
super(Sequence, self).__init__(start, stop, step)
self.start = start
self.stop = stop
self.step = step
def eval(self, row, schema):
start_value = self.start.eval(row, schema)
stop_value = self.stop.eval(row, schema)
if self.step is not None:
step_value = self.step.eval(row, schema)
if ((step_value < stop_value and step_value <= 0) or
(step_value > stop_value and step_value >= 0)):
raise Exception(
"requirement failed: Illegal sequence boundaries: "
"{0} to {1} by {2}".format(
start_value,
stop_value,
step_value
)
)
else:
step_value = 1 if start_value < stop_value else -1
return list(range(start_value, stop_value, step_value))
def __str__(self):
return "array_join({0}, {1}{2})".format(
self.start,
self.stop,
# Spark use the same logic of not displaying step
# if it is None, even if it was explicitly set
", {0}".format(self.step) if self.step is not None else ""
)
class ArrayJoin(Expression):
def __init__(self, column, delimiter, nullReplacement):
super(ArrayJoin, self).__init__(column)
self.column = column
self.delimiter = delimiter
self.nullReplacement = nullReplacement
def eval(self, row, schema):
column_eval = self.column.eval(row, schema)
return self.delimiter.join(
value if value is not None else self.nullReplacement for value in column_eval
)
def __str__(self):
return "array_join({0}, {1}{2})".format(
self.column,
self.delimiter,
# Spark use the same logic of not displaying nullReplacement
# if it is None, even if it was explicitly set
", {0}".format(self.nullReplacement) if self.nullReplacement is not None else ""
)
class SortArray(Expression):
def __init__(self, col, asc):
super(SortArray, self).__init__(col)
self.col = col
self.asc = asc
def eval(self, row, schema):
return sorted(self.col.eval(row, schema), reverse=not self.asc)
def __str__(self):
return "sort_array({0}, {1})".format(
self.col,
self.asc
)
class ArraysZip(Expression):
def __init__(self, cols):
super(ArraysZip, self).__init__(*cols)
self.cols = cols
def eval(self, row, schema):
return [
list(combination)
for combination in zip(
*(c.eval(row, schema) for c in self.cols)
)
]
def __str__(self):
return "arrays_zip({0})".format(", ".join(self.cols))
class Flatten(UnaryExpression):
def eval(self, row, schema):
return [
value
for array in self.column.eval(row, schema)
for value in array
]
def __str__(self):
return "flatten({0})".format(self.column)
class ArrayPosition(Expression):
def __init__(self, col, value):
super(ArrayPosition, self).__init__(col)
self.col = col
self.value = value
def eval(self, row, schema):
if self.value is None:
return None
col_eval = self.col.eval(row, schema)
if col_eval is None:
return None
return col_eval.find(self.value) + 1
def __str__(self):
return "array_position({0}, {1})".format(self.col, self.value)
class ElementAt(Expression):
def __init__(self, col, extraction):
super(ElementAt, self).__init__(col)
self.col = col
self.extraction = extraction
def eval(self, row, schema):
col_eval = self.col.eval(row, schema)
if isinstance(col_eval, list):
return col_eval[self.extraction - 1]
return col_eval.get(self.extraction)
def __str__(self):
return "element_at({0}, {1})".format(self.col, self.extraction)
class ArrayRemove(Expression):
def __init__(self, col, element):
super(ArrayRemove, self).__init__(col, element)
self.col = col
self.element = element
def eval(self, row, schema):
array = self.col.eval(row, schema)
return [value for value in array if value != self.element]
def __str__(self):
return "array_remove({0}, {1})".format(self.col, self.element)
class ArrayDistinct(UnaryExpression):
def eval(self, row, schema):
return list(set(self.column.eval(row, schema)))
def __str__(self):
return "array_distinct({0})".format(self.column)
class ArrayIntersect(Expression):
def __init__(self, col1, col2):
super(ArrayIntersect, self).__init__(col1, col2)
self.col1 = col1
self.col2 = col2
def eval(self, row, schema):
return list(set(self.col1.eval(row, schema)) & set(self.col2.eval(row, schema)))
def __str__(self):
return "array_intersect({0}, {1})".format(self.col1, self.col2)
class ArrayUnion(Expression):
def __init__(self, col1, col2):
super(ArrayUnion, self).__init__(col1, col2)
self.col1 = col1
self.col2 = col2
def eval(self, row, schema):
return list(set(self.col1.eval(row, schema)) | set(self.col2.eval(row, schema)))
def __str__(self):
return "array_union({0}, {1})".format(self.col1, self.col2)
class ArrayExcept(Expression):
def __init__(self, col1, col2):
super(ArrayExcept, self).__init__(col1, col2)
self.col1 = col1
self.col2 = col2
def eval(self, row, schema):
return list(set(self.col1.eval(row, schema)) - set(self.col2.eval(row, schema)))
def __str__(self):
return "array_except({0}, {1})".format(self.col1, self.col2)
__all__ = [
"ArraysZip", "ArrayRepeat", "Flatten", "ArrayMax", "ArrayMin", "SortArray", "Size",
"ArrayExcept", "ArrayUnion", "ArrayIntersect", "ArrayDistinct", "ArrayRemove", "ArraySort",
"ElementAt", "ArrayPosition", "ArrayJoin", "ArraysOverlap", "ArrayContains",
"MapFromArraysColumn", "MapColumn", "ArrayColumn", "Slice", "Sequence"
]