Skip to content

Commit

Permalink
Add array_slice and list_slice
Browse files Browse the repository at this point in the history
  • Loading branch information
ongchi committed Jan 9, 2024
1 parent f093386 commit e5649c9
Show file tree
Hide file tree
Showing 2 changed files with 12 additions and 0 deletions.
8 changes: 8 additions & 0 deletions datafusion/tests/test_functions.py
Original file line number Diff line number Diff line change
Expand Up @@ -424,6 +424,14 @@ def py_arr_replace(arr, from_, to, n=None):
f.list_replace_all(col, literal(3.0), literal(4.0)),
lambda: [py_arr_replace(arr, 3.0, 4.0) for arr in data],
],
[
f.array_slice(col, literal(2), literal(4)),
lambda: [arr[1:4] for arr in data],
],
[
f.list_slice(col, literal(-1), literal(2)),
lambda: [arr[-1:2] for arr in data],
],
]

for stmt, py_expr in test_items:
Expand Down
4 changes: 4 additions & 0 deletions src/functions.rs
Original file line number Diff line number Diff line change
Expand Up @@ -406,6 +406,8 @@ scalar_function!(array_replace_n, ArrayReplaceN);
scalar_function!(list_replace_n, ArrayReplaceN);
scalar_function!(array_replace_all, ArrayReplaceAll);
scalar_function!(list_replace_all, ArrayReplaceAll);
scalar_function!(array_slice, ArraySlice);
scalar_function!(list_slice, ArraySlice);

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

Ok(())
}

0 comments on commit e5649c9

Please sign in to comment.