Add missing array functions (#1468)
* 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) <noreply@anthropic.com>
* 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) <noreply@anthropic.com>
* Improve array function APIs: optional params, better naming, restore comment
- Make null_string optional in string_to_array/string_to_list
- Make step optional in gen_series/generate_series
- Rename second_array to element in array_contains/list_has/list_contains
- Restore # Window Functions section comment in __all__
- Add tests for optional parameter variants
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Consolidate array/list function tests using pytest parametrize
Reduce 26 individual tests to 14 test functions with parametrized
cases, eliminating boilerplate while maintaining full coverage.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Move list alias tests into existing test_array_functions parametrize block
Merge standalone tests for list_empty, list_pop_back, list_pop_front,
list_has, array_contains, list_contains, list_has_all, and list_has_any
into the existing parametrized test_array_functions block alongside
their array_* counterparts.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Merge test_array_any_value into parametrized test_any_value_aliases
Use the richer multi-row dataset (including all-nulls case) for both
array_any_value and list_any_value via the parametrized test.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Add arrays_overlap and list_overlap as aliases for array_has_any
These aliases match the upstream DataFusion SQL-level aliases, completing
the set of missing array functions from issue #1452.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Add docstring examples for optional params in string_to_array and gen_series
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Update AGENTS file to demonstrate preferred method of documenting python functions
---------
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
diff --git a/AGENTS.md b/AGENTS.md
index 1853a84..f6fdfbd 100644
--- a/AGENTS.md
+++ b/AGENTS.md
@@ -25,3 +25,20 @@
- `SKILL.md` — The skill definition with YAML frontmatter (name, description, argument-hint) and detailed instructions.
- Additional supporting files as needed.
+
+## Python Function Docstrings
+
+Every Python function must include a docstring with usage examples.
+
+- **Examples are required**: Each function needs at least one doctest-style example
+ demonstrating basic usage.
+- **Optional parameters**: If a function has optional parameters, include separate
+ examples that show usage both without and with the optional arguments. Pass
+ optional arguments using their keyword name (e.g., `step=dfn.lit(3)`) so readers
+ can immediately see which parameter is being demonstrated.
+- **Reuse input data**: Use the same input data across examples wherever possible.
+ The examples should demonstrate how different optional arguments change the output
+ for the same input, making the effect of each option easy to understand.
+- **Alias functions**: Functions that are simple aliases (e.g., `list_sort` aliasing
+ `array_sort`) only need a one-line description and a `See Also` reference to the
+ primary function. They do not need their own examples.
diff --git a/crates/core/src/functions.rs b/crates/core/src/functions.rs
index 5e61b71..8bb9277 100644
--- a/crates/core/src/functions.rs
+++ b/crates/core/src/functions.rs
@@ -94,6 +94,50 @@
}
#[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>) -> PyExpr {
+ let exprs = exprs.into_iter().map(|x| x.into()).collect();
+ datafusion::functions_nested::expr_fn::arrays_zip(exprs).into()
+}
+
+#[pyfunction]
+#[pyo3(signature = (string, delimiter, null_string=None))]
+fn string_to_array(string: PyExpr, delimiter: PyExpr, null_string: Option<PyExpr>) -> PyExpr {
+ let mut args = vec![string.into(), delimiter.into()];
+ if let Some(null_string) = null_string {
+ args.push(null_string.into());
+ }
+ Expr::ScalarFunction(datafusion::logical_expr::expr::ScalarFunction::new_udf(
+ datafusion::functions_nested::string::string_to_array_udf(),
+ args,
+ ))
+ .into()
+}
+
+#[pyfunction]
+#[pyo3(signature = (start, stop, step=None))]
+fn gen_series(start: PyExpr, stop: PyExpr, step: Option<PyExpr>) -> PyExpr {
+ let mut args = vec![start.into(), stop.into()];
+ if let Some(step) = step {
+ args.push(step.into());
+ }
+ Expr::ScalarFunction(datafusion::logical_expr::expr::ScalarFunction::new_udf(
+ datafusion::functions_nested::range::gen_series_udf(),
+ args,
+ ))
+ .into()
+}
+
+#[pyfunction]
fn make_map(keys: Vec<PyExpr>, values: Vec<PyExpr>) -> PyExpr {
let keys = keys.into_iter().map(|x| x.into()).collect();
let values = values.into_iter().map(|x| x.into()).collect();
@@ -681,6 +725,10 @@
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!(cardinality, array);
array_fn!(flatten, array);
array_fn!(range, start stop step);
@@ -1152,6 +1200,14 @@
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 3febb44..1b26773 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,13 @@
"array_replace_all",
"array_replace_n",
"array_resize",
+ "array_reverse",
"array_slice",
"array_sort",
"array_to_string",
"array_union",
+ "arrays_overlap",
+ "arrays_zip",
"arrow_cast",
"arrow_typeof",
"ascii",
@@ -153,6 +161,8 @@
"floor",
"from_unixtime",
"gcd",
+ "gen_series",
+ "generate_series",
"greatest",
"ifnull",
"in_list",
@@ -167,19 +177,31 @@
"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_overlap",
+ "list_pop_back",
+ "list_pop_front",
"list_position",
"list_positions",
"list_prepend",
@@ -193,10 +215,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",
@@ -273,6 +297,8 @@
"stddev_pop",
"stddev_samp",
"string_agg",
+ "string_to_array",
+ "string_to_list",
"strpos",
"struct",
"substr",
@@ -2794,6 +2820,15 @@
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.
@@ -2891,6 +2926,69 @@
return Expr(f.array_has_any(first_array.expr, second_array.expr))
+def array_contains(array: Expr, element: Expr) -> Expr:
+ """Returns true if the element appears in the array, otherwise false.
+
+ See Also:
+ This is an alias for :py:func:`array_has`.
+ """
+ return array_has(array, element)
+
+
+def list_has(array: Expr, element: Expr) -> Expr:
+ """Returns true if the element appears in the array, otherwise false.
+
+ See Also:
+ This is an alias for :py:func:`array_has`.
+ """
+ return array_has(array, element)
+
+
+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 arrays_overlap(first_array: Expr, second_array: Expr) -> Expr:
+ """Returns true if any element appears in both arrays.
+
+ See Also:
+ This is an alias for :py:func:`array_has_any`.
+ """
+ return array_has_any(first_array, second_array)
+
+
+def list_overlap(first_array: Expr, second_array: Expr) -> Expr:
+ """Returns true if any element appears in both arrays.
+
+ See Also:
+ This is an alias for :py:func:`array_has_any`.
+ """
+ return array_has_any(first_array, second_array)
+
+
+def list_contains(array: Expr, element: Expr) -> Expr:
+ """Returns true if the element appears in the array, otherwise false.
+
+ See Also:
+ This is an alias for :py:func:`array_has`.
+ """
+ return array_has(array, element)
+
+
def array_position(array: Expr, element: Expr, index: int | None = 1) -> Expr:
"""Return the position of the first occurrence of ``element`` in ``array``.
@@ -3058,6 +3156,24 @@
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.
@@ -3429,6 +3545,227 @@
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 | None = None
+) -> Expr:
+ """Splits a string based on a delimiter and returns an array of parts.
+
+ Any parts matching the optional ``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(","),
+ ... ).alias("result"))
+ >>> result.collect_column("result")[0].as_py()
+ ['hello', 'world']
+
+ Replace parts matching a ``null_string`` with ``NULL``:
+
+ >>> result = df.select(
+ ... dfn.functions.string_to_array(
+ ... dfn.col("a"), dfn.lit(","), null_string=dfn.lit("world"),
+ ... ).alias("result"))
+ >>> result.collect_column("result")[0].as_py()
+ ['hello', None]
+ """
+ null_expr = null_string.expr if null_string is not None else None
+ return Expr(f.string_to_array(string.expr, delimiter.expr, null_expr))
+
+
+def string_to_list(
+ string: Expr, delimiter: Expr, null_string: Expr | None = None
+) -> 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 | None = None) -> 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),
+ ... ).alias("result"))
+ >>> result.collect_column("result")[0].as_py()
+ [1, 2, 3, 4, 5]
+
+ Specify a custom ``step``:
+
+ >>> result = df.select(
+ ... dfn.functions.gen_series(
+ ... dfn.lit(1), dfn.lit(10), step=dfn.lit(3),
+ ... ).alias("result"))
+ >>> result.collect_column("result")[0].as_py()
+ [1, 4, 7, 10]
+ """
+ step_expr = step.expr if step is not None else None
+ return Expr(f.gen_series(start.expr, stop.expr, step_expr))
+
+
+def generate_series(start: Expr, stop: Expr, step: Expr | None = None) -> 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.
diff --git a/python/tests/test_functions.py b/python/tests/test_functions.py
index f25c6e7..2100da9 100644
--- a/python/tests/test_functions.py
+++ b/python/tests/test_functions.py
@@ -331,6 +331,10 @@
lambda data: [len(r) == 0 for r in data],
),
(
+ f.list_empty,
+ lambda data: [len(r) == 0 for r in data],
+ ),
+ (
lambda col: f.array_extract(col, literal(1)),
lambda data: [r[0] for r in data],
),
@@ -355,18 +359,54 @@
lambda data: [1.0 in r for r in data],
),
(
+ lambda col: f.list_has(col, literal(1.0)),
+ lambda data: [1.0 in r for r in data],
+ ),
+ (
+ lambda col: f.array_contains(col, literal(1.0)),
+ lambda data: [1.0 in r for r in data],
+ ),
+ (
+ lambda col: f.list_contains(col, literal(1.0)),
+ lambda data: [1.0 in r for r in data],
+ ),
+ (
lambda col: f.array_has_all(
col, f.make_array(*[literal(v) for v in [1.0, 3.0, 5.0]])
),
lambda data: [np.all([v in r for v in [1.0, 3.0, 5.0]]) for r in data],
),
(
+ lambda col: f.list_has_all(
+ col, f.make_array(*[literal(v) for v in [1.0, 3.0, 5.0]])
+ ),
+ lambda data: [np.all([v in r for v in [1.0, 3.0, 5.0]]) for r in data],
+ ),
+ (
lambda col: f.array_has_any(
col, f.make_array(*[literal(v) for v in [1.0, 3.0, 5.0]])
),
lambda data: [np.any([v in r for v in [1.0, 3.0, 5.0]]) for r in data],
),
(
+ lambda col: f.list_has_any(
+ col, f.make_array(*[literal(v) for v in [1.0, 3.0, 5.0]])
+ ),
+ lambda data: [np.any([v in r for v in [1.0, 3.0, 5.0]]) for r in data],
+ ),
+ (
+ lambda col: f.arrays_overlap(
+ col, f.make_array(*[literal(v) for v in [1.0, 3.0, 5.0]])
+ ),
+ lambda data: [np.any([v in r for v in [1.0, 3.0, 5.0]]) for r in data],
+ ),
+ (
+ lambda col: f.list_overlap(
+ col, f.make_array(*[literal(v) for v in [1.0, 3.0, 5.0]])
+ ),
+ lambda data: [np.any([v in r for v in [1.0, 3.0, 5.0]]) for r in data],
+ ),
+ (
lambda col: f.array_position(col, literal(1.0)),
lambda data: [py_indexof(r, 1.0) for r in data],
),
@@ -419,10 +459,18 @@
lambda data: [arr[:-1] for arr in data],
),
(
+ f.list_pop_back,
+ lambda data: [arr[:-1] for arr in data],
+ ),
+ (
f.array_pop_front,
lambda data: [arr[1:] for arr in data],
),
(
+ f.list_pop_front,
+ lambda data: [arr[1:] for arr in data],
+ ),
+ (
lambda col: f.array_remove(col, literal(3.0)),
lambda data: [py_arr_remove(arr, 3.0, 1) for arr in data],
),
@@ -1760,3 +1808,92 @@
def test_conditional_functions(df_with_nulls, expr, expected):
result = df_with_nulls.select(expr.alias("result")).collect()[0]
assert result.column(0) == expected
+
+
+@pytest.mark.parametrize("func", [f.array_any_value, f.list_any_value])
+def test_any_value_aliases(func):
+ ctx = SessionContext()
+ df = ctx.from_pydict({"a": [[None, 2, 3], [None, None, None], [1, 2, 3]]})
+ result = df.select(func(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
+
+
+@pytest.mark.parametrize("func", [f.array_distance, f.list_distance])
+def test_array_distance_aliases(func):
+ ctx = SessionContext()
+ df = ctx.from_pydict({"a": [[1.0, 2.0]], "b": [[1.0, 4.0]]})
+ result = df.select(func(column("a"), column("b")).alias("v")).collect()
+ assert result[0].column(0)[0].as_py() == pytest.approx(2.0)
+
+
+@pytest.mark.parametrize(
+ ("func", "expected"),
+ [
+ (f.array_max, [5, 10]),
+ (f.list_max, [5, 10]),
+ (f.array_min, [1, 2]),
+ (f.list_min, [1, 2]),
+ ],
+)
+def test_array_min_max(func, expected):
+ ctx = SessionContext()
+ df = ctx.from_pydict({"a": [[1, 5, 3], [10, 2]]})
+ result = df.select(func(column("a")).alias("v")).collect()
+ values = [row.as_py() for row in result[0].column(0)]
+ assert values == expected
+
+
+@pytest.mark.parametrize("func", [f.array_reverse, f.list_reverse])
+def test_array_reverse_aliases(func):
+ ctx = SessionContext()
+ df = ctx.from_pydict({"a": [[1, 2, 3], [4, 5]]})
+ result = df.select(func(column("a")).alias("v")).collect()
+ values = [row.as_py() for row in result[0].column(0)]
+ assert values == [[3, 2, 1], [5, 4]]
+
+
+@pytest.mark.parametrize("func", [f.arrays_zip, f.list_zip])
+def test_arrays_zip_aliases(func):
+ ctx = SessionContext()
+ df = ctx.from_pydict({"a": [[1, 2]], "b": [[3, 4]]})
+ result = df.select(func(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}]
+
+
+@pytest.mark.parametrize("func", [f.string_to_array, f.string_to_list])
+def test_string_to_array_aliases(func):
+ ctx = SessionContext()
+ df = ctx.from_pydict({"a": ["hello,world,foo"]})
+ result = df.select(func(column("a"), literal(",")).alias("v")).collect()
+ assert result[0].column(0)[0].as_py() == ["hello", "world", "foo"]
+
+
+def test_string_to_array_with_null_string():
+ ctx = SessionContext()
+ df = ctx.from_pydict({"a": ["hello,NA,world"]})
+ result = df.select(
+ f.string_to_array(column("a"), literal(","), literal("NA")).alias("v")
+ ).collect()
+ values = result[0].column(0)[0].as_py()
+ assert values == ["hello", None, "world"]
+
+
+@pytest.mark.parametrize("func", [f.gen_series, f.generate_series])
+def test_gen_series_aliases(func):
+ ctx = SessionContext()
+ df = ctx.from_pydict({"a": [0]})
+ result = df.select(func(literal(1), literal(5)).alias("v")).collect()
+ assert result[0].column(0)[0].as_py() == [1, 2, 3, 4, 5]
+
+
+def test_gen_series_with_step():
+ ctx = SessionContext()
+ df = ctx.from_pydict({"a": [0]})
+ result = df.select(
+ f.gen_series(literal(1), literal(10), literal(3)).alias("v")
+ ).collect()
+ assert result[0].column(0)[0].as_py() == [1, 4, 7, 10]
