This is an automated email from the ASF dual-hosted git repository.
agrove pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-python.git
The following commit(s) were added to refs/heads/main by this push:
new f00b8ee Expose array sort (#764)
f00b8ee is described below
commit f00b8ee1959f354950660b0b0c784d4f48f420c6
Author: Tim Saucer <[email protected]>
AuthorDate: Sat Jul 20 11:21:46 2024 -0400
Expose array sort (#764)
---
python/datafusion/functions.py | 22 ++++++++++++++++++++++
python/datafusion/tests/test_functions.py | 8 ++++++++
src/functions.rs | 4 ++++
3 files changed, 34 insertions(+)
diff --git a/python/datafusion/functions.py b/python/datafusion/functions.py
index ad77712..46d2a2f 100644
--- a/python/datafusion/functions.py
+++ b/python/datafusion/functions.py
@@ -1126,6 +1126,28 @@ def list_replace_all(array: Expr, from_val: Expr,
to_val: Expr) -> Expr:
return array_replace_all(array, from_val, to_val)
+def array_sort(array: Expr, descending: bool = False, null_first: bool =
False) -> Expr:
+ """Sort an array.
+
+ Args:
+ array: The input array to sort.
+ descending: If True, sorts in descending order.
+ null_first: If True, nulls will be returned at the beginning of the
array.
+ """
+ desc = "DESC" if descending else "ASC"
+ nulls_first = "NULLS FIRST" if null_first else "NULLS LAST"
+ return Expr(
+ f.array_sort(
+ array.expr, Expr.literal(desc).expr, Expr.literal(nulls_first).expr
+ )
+ )
+
+
+def list_sort(array: Expr, descending: bool = False, null_first: bool = False)
-> Expr:
+ """This is an alias for ``array_sort``."""
+ return array_sort(array, descending=descending, null_first=null_first)
+
+
def array_slice(
array: Expr, begin: Expr, end: Expr, stride: Expr | None = None
) -> Expr:
diff --git a/python/datafusion/tests/test_functions.py
b/python/datafusion/tests/test_functions.py
index 2384b6a..25d7de1 100644
--- a/python/datafusion/tests/test_functions.py
+++ b/python/datafusion/tests/test_functions.py
@@ -453,6 +453,14 @@ def py_flatten(arr):
lambda col: f.list_replace_all(col, literal(3.0), literal(4.0)),
lambda data: [py_arr_replace(arr, 3.0, 4.0) for arr in data],
],
+ [
+ lambda col: f.array_sort(col, descending=True, null_first=True),
+ lambda data: [np.sort(arr)[::-1] for arr in data],
+ ],
+ [
+ lambda col: f.list_sort(col, descending=False, null_first=False),
+ lambda data: [np.sort(arr) for arr in data],
+ ],
[
lambda col: f.array_slice(col, literal(2), literal(4)),
lambda data: [arr[1:4] for arr in data],
diff --git a/src/functions.rs b/src/functions.rs
index 1e75dd5..74eb48a 100644
--- a/src/functions.rs
+++ b/src/functions.rs
@@ -673,6 +673,8 @@ array_fn!(array_replace_n, array from to max);
array_fn!(list_replace_n, array_replace_n, array from to max);
array_fn!(array_replace_all, array from to);
array_fn!(list_replace_all, array_replace_all, array from to);
+array_fn!(array_sort, array desc null_first);
+array_fn!(list_sort, array_sort, array desc null_first);
array_fn!(array_intersect, first_array second_array);
array_fn!(list_intersect, array_intersect, first_array second_array);
array_fn!(array_union, array1 array2);
@@ -936,6 +938,8 @@ pub(crate) fn init_module(m: &Bound<'_, 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_sort))?;
+ m.add_wrapped(wrap_pyfunction!(list_sort))?;
m.add_wrapped(wrap_pyfunction!(array_slice))?;
m.add_wrapped(wrap_pyfunction!(list_slice))?;
m.add_wrapped(wrap_pyfunction!(flatten))?;
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
