From 3fd8b16b8badc1b0ea9ddd4fe18b003aed06ce12 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Mon, 30 Mar 2026 16:06:41 -0400 Subject: [PATCH 1/2] Add missing array/list functions and aliases (#1452) Add new array functions from upstream DataFusion v53: array_any_value, array_distance, array_max, array_min, array_reverse, arrays_zip, string_to_array, and gen_series. Add corresponding list_* aliases and missing list_* aliases for existing functions (list_empty, list_pop_back, list_pop_front, list_has, list_has_all, list_has_any). Also add array_contains/list_contains as aliases for array_has, generate_series as alias for gen_series, and string_to_list as alias for string_to_array. Co-Authored-By: Claude Opus 4.6 (1M context) --- crates/core/src/functions.rs | 30 ++++ python/datafusion/functions.py | 294 ++++++++++++++++++++++++++++++++- 2 files changed, 323 insertions(+), 1 deletion(-) diff --git a/crates/core/src/functions.rs b/crates/core/src/functions.rs index c32134054..8584b09f2 100644 --- a/crates/core/src/functions.rs +++ b/crates/core/src/functions.rs @@ -93,6 +93,22 @@ fn array_cat(exprs: Vec) -> PyExpr { array_concat(exprs) } +#[pyfunction] +fn array_distance(array1: PyExpr, array2: PyExpr) -> PyExpr { + let args = vec![array1.into(), array2.into()]; + Expr::ScalarFunction(datafusion::logical_expr::expr::ScalarFunction::new_udf( + datafusion::functions_nested::distance::array_distance_udf(), + args, + )) + .into() +} + +#[pyfunction] +fn arrays_zip(exprs: Vec) -> PyExpr { + let exprs = exprs.into_iter().map(|x| x.into()).collect(); + datafusion::functions_nested::expr_fn::arrays_zip(exprs).into() +} + #[pyfunction] #[pyo3(signature = (array, element, index=None))] fn array_position(array: PyExpr, element: PyExpr, index: Option) -> PyExpr { @@ -661,6 +677,12 @@ array_fn!(array_intersect, first_array second_array); array_fn!(array_union, array1 array2); array_fn!(array_except, first_array second_array); array_fn!(array_resize, array size value); +array_fn!(array_any_value, array); +array_fn!(array_max, array); +array_fn!(array_min, array); +array_fn!(array_reverse, array); +array_fn!(string_to_array, string delimiter null_string); +array_fn!(gen_series, start stop step); array_fn!(cardinality, array); array_fn!(flatten, array); array_fn!(range, start stop step); @@ -1121,6 +1143,14 @@ pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_wrapped(wrap_pyfunction!(array_replace_all))?; m.add_wrapped(wrap_pyfunction!(array_sort))?; m.add_wrapped(wrap_pyfunction!(array_slice))?; + m.add_wrapped(wrap_pyfunction!(array_any_value))?; + m.add_wrapped(wrap_pyfunction!(array_distance))?; + m.add_wrapped(wrap_pyfunction!(array_max))?; + m.add_wrapped(wrap_pyfunction!(array_min))?; + m.add_wrapped(wrap_pyfunction!(array_reverse))?; + m.add_wrapped(wrap_pyfunction!(arrays_zip))?; + m.add_wrapped(wrap_pyfunction!(string_to_array))?; + m.add_wrapped(wrap_pyfunction!(gen_series))?; m.add_wrapped(wrap_pyfunction!(flatten))?; m.add_wrapped(wrap_pyfunction!(cardinality))?; diff --git a/python/datafusion/functions.py b/python/datafusion/functions.py index f062cbfce..f466989d9 100644 --- a/python/datafusion/functions.py +++ b/python/datafusion/functions.py @@ -53,10 +53,13 @@ "approx_percentile_cont_with_weight", "array", "array_agg", + "array_any_value", "array_append", "array_cat", "array_concat", + "array_contains", "array_dims", + "array_distance", "array_distinct", "array_element", "array_empty", @@ -69,6 +72,8 @@ "array_intersect", "array_join", "array_length", + "array_max", + "array_min", "array_ndims", "array_pop_back", "array_pop_front", @@ -85,10 +90,12 @@ "array_replace_all", "array_replace_n", "array_resize", + "array_reverse", "array_slice", "array_sort", "array_to_string", "array_union", + "arrays_zip", "arrow_cast", "arrow_typeof", "ascii", @@ -149,6 +156,8 @@ "floor", "from_unixtime", "gcd", + "gen_series", + "generate_series", "in_list", "initcap", "isnan", @@ -160,19 +169,30 @@ "left", "length", "levenshtein", + "list_any_value", "list_append", "list_cat", "list_concat", + "list_contains", "list_dims", + "list_distance", "list_distinct", "list_element", + "list_empty", "list_except", "list_extract", + "list_has", + "list_has_all", + "list_has_any", "list_indexof", "list_intersect", "list_join", "list_length", + "list_max", + "list_min", "list_ndims", + "list_pop_back", + "list_pop_front", "list_position", "list_positions", "list_prepend", @@ -186,10 +206,12 @@ "list_replace_all", "list_replace_n", "list_resize", + "list_reverse", "list_slice", "list_sort", "list_to_string", "list_union", + "list_zip", "ln", "log", "log2", @@ -259,6 +281,8 @@ "stddev_pop", "stddev_samp", "string_agg", + "string_to_array", + "string_to_list", "strpos", "struct", "substr", @@ -289,7 +313,6 @@ "var_samp", "var_sample", "when", - # Window Functions "window", ] @@ -2668,6 +2691,15 @@ def array_empty(array: Expr) -> Expr: return Expr(f.array_empty(array.expr)) +def list_empty(array: Expr) -> Expr: + """Returns a boolean indicating whether the array is empty. + + See Also: + This is an alias for :py:func:`array_empty`. + """ + return array_empty(array) + + def array_extract(array: Expr, n: Expr) -> Expr: """Extracts the element with the index n from the array. @@ -2765,6 +2797,51 @@ def array_has_any(first_array: Expr, second_array: Expr) -> Expr: return Expr(f.array_has_any(first_array.expr, second_array.expr)) +def array_contains(first_array: Expr, second_array: Expr) -> Expr: + """Returns true if the element appears in the first array, otherwise false. + + See Also: + This is an alias for :py:func:`array_has`. + """ + return array_has(first_array, second_array) + + +def list_has(first_array: Expr, second_array: Expr) -> Expr: + """Returns true if the element appears in the first array, otherwise false. + + See Also: + This is an alias for :py:func:`array_has`. + """ + return array_has(first_array, second_array) + + +def list_has_all(first_array: Expr, second_array: Expr) -> Expr: + """Determines if there is complete overlap ``second_array`` in ``first_array``. + + See Also: + This is an alias for :py:func:`array_has_all`. + """ + return array_has_all(first_array, second_array) + + +def list_has_any(first_array: Expr, second_array: Expr) -> Expr: + """Determine if there is an overlap between ``first_array`` and ``second_array``. + + See Also: + This is an alias for :py:func:`array_has_any`. + """ + return array_has_any(first_array, second_array) + + +def list_contains(first_array: Expr, second_array: Expr) -> Expr: + """Returns true if the element appears in the first array, otherwise false. + + See Also: + This is an alias for :py:func:`array_has`. + """ + return array_has(first_array, second_array) + + def array_position(array: Expr, element: Expr, index: int | None = 1) -> Expr: """Return the position of the first occurrence of ``element`` in ``array``. @@ -2932,6 +3009,24 @@ def array_pop_front(array: Expr) -> Expr: return Expr(f.array_pop_front(array.expr)) +def list_pop_back(array: Expr) -> Expr: + """Returns the array without the last element. + + See Also: + This is an alias for :py:func:`array_pop_back`. + """ + return array_pop_back(array) + + +def list_pop_front(array: Expr) -> Expr: + """Returns the array without the first element. + + See Also: + This is an alias for :py:func:`array_pop_front`. + """ + return array_pop_front(array) + + def array_remove(array: Expr, element: Expr) -> Expr: """Removes the first element from the array equal to the given value. @@ -3303,6 +3398,203 @@ def list_resize(array: Expr, size: Expr, value: Expr) -> Expr: return array_resize(array, size, value) +def array_any_value(array: Expr) -> Expr: + """Returns the first non-null element in the array. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [[None, 2, 3]]}) + >>> result = df.select( + ... dfn.functions.array_any_value(dfn.col("a")).alias("result")) + >>> result.collect_column("result")[0].as_py() + 2 + """ + return Expr(f.array_any_value(array.expr)) + + +def list_any_value(array: Expr) -> Expr: + """Returns the first non-null element in the array. + + See Also: + This is an alias for :py:func:`array_any_value`. + """ + return array_any_value(array) + + +def array_distance(array1: Expr, array2: Expr) -> Expr: + """Returns the Euclidean distance between two numeric arrays. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [[1.0, 2.0]], "b": [[1.0, 4.0]]}) + >>> result = df.select( + ... dfn.functions.array_distance( + ... dfn.col("a"), dfn.col("b"), + ... ).alias("result")) + >>> result.collect_column("result")[0].as_py() + 2.0 + """ + return Expr(f.array_distance(array1.expr, array2.expr)) + + +def list_distance(array1: Expr, array2: Expr) -> Expr: + """Returns the Euclidean distance between two numeric arrays. + + See Also: + This is an alias for :py:func:`array_distance`. + """ + return array_distance(array1, array2) + + +def array_max(array: Expr) -> Expr: + """Returns the maximum value in the array. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [[1, 2, 3]]}) + >>> result = df.select( + ... dfn.functions.array_max(dfn.col("a")).alias("result")) + >>> result.collect_column("result")[0].as_py() + 3 + """ + return Expr(f.array_max(array.expr)) + + +def list_max(array: Expr) -> Expr: + """Returns the maximum value in the array. + + See Also: + This is an alias for :py:func:`array_max`. + """ + return array_max(array) + + +def array_min(array: Expr) -> Expr: + """Returns the minimum value in the array. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [[1, 2, 3]]}) + >>> result = df.select( + ... dfn.functions.array_min(dfn.col("a")).alias("result")) + >>> result.collect_column("result")[0].as_py() + 1 + """ + return Expr(f.array_min(array.expr)) + + +def list_min(array: Expr) -> Expr: + """Returns the minimum value in the array. + + See Also: + This is an alias for :py:func:`array_min`. + """ + return array_min(array) + + +def array_reverse(array: Expr) -> Expr: + """Reverses the order of elements in the array. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [[1, 2, 3]]}) + >>> result = df.select( + ... dfn.functions.array_reverse(dfn.col("a")).alias("result")) + >>> result.collect_column("result")[0].as_py() + [3, 2, 1] + """ + return Expr(f.array_reverse(array.expr)) + + +def list_reverse(array: Expr) -> Expr: + """Reverses the order of elements in the array. + + See Also: + This is an alias for :py:func:`array_reverse`. + """ + return array_reverse(array) + + +def arrays_zip(*arrays: Expr) -> Expr: + """Combines multiple arrays into a single array of structs. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [[1, 2]], "b": [[3, 4]]}) + >>> result = df.select( + ... dfn.functions.arrays_zip(dfn.col("a"), dfn.col("b")).alias("result")) + >>> result.collect_column("result")[0].as_py() + [{'c0': 1, 'c1': 3}, {'c0': 2, 'c1': 4}] + """ + args = [a.expr for a in arrays] + return Expr(f.arrays_zip(args)) + + +def list_zip(*arrays: Expr) -> Expr: + """Combines multiple arrays into a single array of structs. + + See Also: + This is an alias for :py:func:`arrays_zip`. + """ + return arrays_zip(*arrays) + + +def string_to_array(string: Expr, delimiter: Expr, null_string: Expr) -> Expr: + """Splits a string based on a delimiter and returns an array of parts. + + Any parts matching the ``null_string`` will be replaced with ``NULL``. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["hello,world"]}) + >>> result = df.select( + ... dfn.functions.string_to_array( + ... dfn.col("a"), dfn.lit(","), dfn.lit(""), + ... ).alias("result")) + >>> result.collect_column("result")[0].as_py() + ['hello', 'world'] + """ + return Expr(f.string_to_array(string.expr, delimiter.expr, null_string.expr)) + + +def string_to_list(string: Expr, delimiter: Expr, null_string: Expr) -> Expr: + """Splits a string based on a delimiter and returns an array of parts. + + See Also: + This is an alias for :py:func:`string_to_array`. + """ + return string_to_array(string, delimiter, null_string) + + +def gen_series(start: Expr, stop: Expr, step: Expr) -> Expr: + """Creates a list of values in the range between start and stop. + + Unlike :py:func:`range`, this includes the upper bound. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [0]}) + >>> result = df.select( + ... dfn.functions.gen_series( + ... dfn.lit(1), dfn.lit(5), dfn.lit(1), + ... ).alias("result")) + >>> result.collect_column("result")[0].as_py() + [1, 2, 3, 4, 5] + """ + return Expr(f.gen_series(start.expr, stop.expr, step.expr)) + + +def generate_series(start: Expr, stop: Expr, step: Expr) -> Expr: + """Creates a list of values in the range between start and stop. + + Unlike :py:func:`range`, this includes the upper bound. + + See Also: + This is an alias for :py:func:`gen_series`. + """ + return gen_series(start, stop, step) + + def flatten(array: Expr) -> Expr: """Flattens an array of arrays into a single array. From 5b592dc7a28c4bbb6100f783993e1fb9192674aa Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Mon, 30 Mar 2026 16:14:13 -0400 Subject: [PATCH 2/2] Add unit tests for new array/list functions and aliases Tests cover all functions and aliases added in the previous commit: array_any_value, array_distance, array_max, array_min, array_reverse, arrays_zip, string_to_array, gen_series, generate_series, array_contains, list_contains, list_empty, list_pop_back, list_pop_front, list_has, list_has_all, list_has_any, and list_* aliases for the new functions. Co-Authored-By: Claude Opus 4.6 (1M context) --- python/tests/test_functions.py | 189 +++++++++++++++++++++++++++++++++ 1 file changed, 189 insertions(+) diff --git a/python/tests/test_functions.py b/python/tests/test_functions.py index 37d349c58..cd6dca685 100644 --- a/python/tests/test_functions.py +++ b/python/tests/test_functions.py @@ -1435,3 +1435,192 @@ def test_coalesce(df): assert result.column(0) == pa.array( ["Hello", "fallback", "!"], type=pa.string_view() ) + + +def test_array_any_value(): + ctx = SessionContext() + df = ctx.from_pydict({"a": [[None, 2, 3], [None, None, None], [1, 2, 3]]}) + result = df.select(f.array_any_value(column("a")).alias("v")).collect() + values = [row.as_py() for row in result[0].column(0)] + assert values[0] == 2 + assert values[1] is None + assert values[2] == 1 + + +def test_list_any_value(): + ctx = SessionContext() + df = ctx.from_pydict({"a": [[None, 5]]}) + result = df.select(f.list_any_value(column("a")).alias("v")).collect() + assert result[0].column(0)[0].as_py() == 5 + + +def test_array_distance(): + ctx = SessionContext() + df = ctx.from_pydict({"a": [[1.0, 2.0]], "b": [[1.0, 4.0]]}) + result = df.select(f.array_distance(column("a"), column("b")).alias("v")).collect() + assert result[0].column(0)[0].as_py() == pytest.approx(2.0) + + +def test_list_distance(): + ctx = SessionContext() + df = ctx.from_pydict({"a": [[3.0, 0.0]], "b": [[0.0, 4.0]]}) + result = df.select(f.list_distance(column("a"), column("b")).alias("v")).collect() + assert result[0].column(0)[0].as_py() == pytest.approx(5.0) + + +def test_array_max(): + ctx = SessionContext() + df = ctx.from_pydict({"a": [[1, 5, 3], [10, 2]]}) + result = df.select(f.array_max(column("a")).alias("v")).collect() + values = [row.as_py() for row in result[0].column(0)] + assert values == [5, 10] + + +def test_list_max(): + ctx = SessionContext() + df = ctx.from_pydict({"a": [[7, 2, 9]]}) + result = df.select(f.list_max(column("a")).alias("v")).collect() + assert result[0].column(0)[0].as_py() == 9 + + +def test_array_min(): + ctx = SessionContext() + df = ctx.from_pydict({"a": [[1, 5, 3], [10, 2]]}) + result = df.select(f.array_min(column("a")).alias("v")).collect() + values = [row.as_py() for row in result[0].column(0)] + assert values == [1, 2] + + +def test_list_min(): + ctx = SessionContext() + df = ctx.from_pydict({"a": [[7, 2, 9]]}) + result = df.select(f.list_min(column("a")).alias("v")).collect() + assert result[0].column(0)[0].as_py() == 2 + + +def test_array_reverse(): + ctx = SessionContext() + df = ctx.from_pydict({"a": [[1, 2, 3], [4, 5]]}) + result = df.select(f.array_reverse(column("a")).alias("v")).collect() + values = [row.as_py() for row in result[0].column(0)] + assert values == [[3, 2, 1], [5, 4]] + + +def test_list_reverse(): + ctx = SessionContext() + df = ctx.from_pydict({"a": [[10, 20, 30]]}) + result = df.select(f.list_reverse(column("a")).alias("v")).collect() + assert result[0].column(0)[0].as_py() == [30, 20, 10] + + +def test_arrays_zip(): + ctx = SessionContext() + df = ctx.from_pydict({"a": [[1, 2]], "b": [[3, 4]]}) + result = df.select(f.arrays_zip(column("a"), column("b")).alias("v")).collect() + values = result[0].column(0)[0].as_py() + assert values == [{"c0": 1, "c1": 3}, {"c0": 2, "c1": 4}] + + +def test_list_zip(): + ctx = SessionContext() + df = ctx.from_pydict({"a": [[1, 2]], "b": [[3, 4]]}) + result = df.select(f.list_zip(column("a"), column("b")).alias("v")).collect() + values = result[0].column(0)[0].as_py() + assert values == [{"c0": 1, "c1": 3}, {"c0": 2, "c1": 4}] + + +def test_string_to_array(): + ctx = SessionContext() + df = ctx.from_pydict({"a": ["hello,world,foo"]}) + result = df.select( + f.string_to_array(column("a"), literal(","), literal("")).alias("v") + ).collect() + assert result[0].column(0)[0].as_py() == ["hello", "world", "foo"] + + +def test_string_to_list(): + ctx = SessionContext() + df = ctx.from_pydict({"a": ["a-b-c"]}) + result = df.select( + f.string_to_list(column("a"), literal("-"), literal("")).alias("v") + ).collect() + assert result[0].column(0)[0].as_py() == ["a", "b", "c"] + + +def test_gen_series(): + ctx = SessionContext() + df = ctx.from_pydict({"a": [0]}) + result = df.select( + f.gen_series(literal(1), literal(5), literal(1)).alias("v") + ).collect() + assert result[0].column(0)[0].as_py() == [1, 2, 3, 4, 5] + + +def test_generate_series(): + ctx = SessionContext() + df = ctx.from_pydict({"a": [0]}) + result = df.select( + f.generate_series(literal(1), literal(3), literal(1)).alias("v") + ).collect() + assert result[0].column(0)[0].as_py() == [1, 2, 3] + + +def test_array_contains(): + ctx = SessionContext() + df = ctx.from_pydict({"a": [[1, 2, 3]]}) + result = df.select(f.array_contains(column("a"), literal(2)).alias("v")).collect() + assert result[0].column(0)[0].as_py() is True + + +def test_list_contains(): + ctx = SessionContext() + df = ctx.from_pydict({"a": [[1, 2, 3]]}) + result = df.select(f.list_contains(column("a"), literal(99)).alias("v")).collect() + assert result[0].column(0)[0].as_py() is False + + +def test_list_empty(): + ctx = SessionContext() + df = ctx.from_pydict({"a": [[], [1, 2]]}) + result = df.select(f.list_empty(column("a")).alias("v")).collect() + values = [row.as_py() for row in result[0].column(0)] + assert values == [True, False] + + +def test_list_pop_back(): + ctx = SessionContext() + df = ctx.from_pydict({"a": [[1, 2, 3]]}) + result = df.select(f.list_pop_back(column("a")).alias("v")).collect() + assert result[0].column(0)[0].as_py() == [1, 2] + + +def test_list_pop_front(): + ctx = SessionContext() + df = ctx.from_pydict({"a": [[1, 2, 3]]}) + result = df.select(f.list_pop_front(column("a")).alias("v")).collect() + assert result[0].column(0)[0].as_py() == [2, 3] + + +def test_list_has(): + ctx = SessionContext() + df = ctx.from_pydict({"a": [[1, 2, 3]]}) + result = df.select(f.list_has(column("a"), literal(2)).alias("v")).collect() + assert result[0].column(0)[0].as_py() is True + + +def test_list_has_all(): + ctx = SessionContext() + df = ctx.from_pydict({"a": [[1, 2, 3]]}) + result = df.select( + f.list_has_all(column("a"), f.make_array(literal(1), literal(2))).alias("v") + ).collect() + assert result[0].column(0)[0].as_py() is True + + +def test_list_has_any(): + ctx = SessionContext() + df = ctx.from_pydict({"a": [[1, 2, 3]]}) + result = df.select( + f.list_has_any(column("a"), f.make_array(literal(5), literal(2))).alias("v") + ).collect() + assert result[0].column(0)[0].as_py() is True