Skip to content

Commit

Permalink
feat: Add flatten array function (#562)
Browse files Browse the repository at this point in the history
* Updated functions.rs, test_functions.py - Flatten

* Updated functions.rs - Formatting

* Updated test_functions.py - Converted test array to a literal for testing

* Updated test_functions.py - Linting
  • Loading branch information
mobley-trent authored Feb 25, 2024
1 parent 697ca2c commit 27a9264
Show file tree
Hide file tree
Showing 2 changed files with 12 additions and 0 deletions.
10 changes: 10 additions & 0 deletions datafusion/tests/test_functions.py
Original file line number Diff line number Diff line change
Expand Up @@ -238,6 +238,15 @@ def py_arr_replace(arr, from_, to, n=None):

return new_arr

def py_flatten(arr):
result = []
for elem in arr:
if isinstance(elem, list):
result.extend(py_flatten(elem))
else:
result.append(elem)
return result

col = column("arr")
test_items = [
[
Expand Down Expand Up @@ -432,6 +441,7 @@ def py_arr_replace(arr, from_, to, n=None):
f.list_slice(col, literal(-1), literal(2)),
lambda: [arr[-1:2] for arr in data],
],
[f.flatten(literal(data)), lambda: [py_flatten(data)]],
]

for stmt, py_expr in test_items:
Expand Down
2 changes: 2 additions & 0 deletions src/functions.rs
Original file line number Diff line number Diff line change
Expand Up @@ -409,6 +409,7 @@ scalar_function!(array_replace_all, ArrayReplaceAll);
scalar_function!(list_replace_all, ArrayReplaceAll);
scalar_function!(array_slice, ArraySlice);
scalar_function!(list_slice, ArraySlice);
scalar_function!(flatten, Flatten);

aggregate_function!(approx_distinct, ApproxDistinct);
aggregate_function!(approx_median, ApproxMedian);
Expand Down Expand Up @@ -651,6 +652,7 @@ pub(crate) fn init_module(m: &PyModule) -> PyResult<()> {
m.add_wrapped(wrap_pyfunction!(list_replace_all))?;
m.add_wrapped(wrap_pyfunction!(array_slice))?;
m.add_wrapped(wrap_pyfunction!(list_slice))?;
m.add_wrapped(wrap_pyfunction!(flatten))?;

Ok(())
}

0 comments on commit 27a9264

Please sign in to comment.