diff --git a/AGENTS.md b/AGENTS.md index 1853a84cd..f6fdfbd90 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -25,3 +25,20 @@ Skills follow the [Agent Skills](https://agentskills.io) open standard. Each ski - `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 5e61b71be..8bb927718 100644 --- a/crates/core/src/functions.rs +++ b/crates/core/src/functions.rs @@ -93,6 +93,50 @@ 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 = (string, delimiter, null_string=None))] +fn string_to_array(string: PyExpr, delimiter: PyExpr, null_string: Option) -> 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 { + 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, values: Vec) -> PyExpr { let keys = keys.into_iter().map(|x| x.into()).collect(); @@ -681,6 +725,10 @@ 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!(cardinality, array); array_fn!(flatten, array); array_fn!(range, start stop step); @@ -1152,6 +1200,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 3febb44e3..1b267731e 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 @@ 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. @@ -2891,6 +2926,69 @@ 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(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 @@ 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. @@ -3429,6 +3545,227 @@ 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 | 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 f25c6e78c..2100da9ae 100644 --- a/python/tests/test_functions.py +++ b/python/tests/test_functions.py @@ -330,6 +330,10 @@ def py_flatten(arr): f.empty, 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], @@ -354,18 +358,54 @@ def py_flatten(arr): lambda col: f.array_has(col, literal(1.0)), 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], @@ -418,10 +458,18 @@ def py_flatten(arr): f.array_pop_back, 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 df_with_nulls(): 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]