Merge branch 'main' into feat/expose-agg-fns

Author: timsaucer

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.

README.md

@@ -275,7 +275,7 @@ needing to activate the virtual environment:
 
 ```bash
 uv run --no-project maturin develop --uv
-uv run --no-project pytest .
+uv run --no-project pytest
 ```
 
 To run the FFI tests within the examples folder, after you have built

crates/core/src/functions.rs

@@ -93,6 +93,57 @@ fn array_cat(exprs: Vec<PyExpr>) -> 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>) -> 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();
+    datafusion::functions_nested::map::map(keys, values).into()
+}
+
 #[pyfunction]
 #[pyo3(signature = (array, element, index=None))]
 fn array_position(array: PyExpr, element: PyExpr, index: Option<i64>) -> PyExpr {
@@ -644,8 +695,29 @@ expr_fn_vec!(named_struct);
 expr_fn!(from_unixtime, unixtime);
 expr_fn!(arrow_typeof, arg_1);
 expr_fn!(arrow_cast, arg_1 datatype);
+expr_fn_vec!(arrow_metadata);
+expr_fn!(union_tag, arg1);
 expr_fn!(random);
 
+#[pyfunction]
+fn get_field(expr: PyExpr, name: PyExpr) -> PyExpr {
+    functions::core::get_field()
+        .call(vec![expr.into(), name.into()])
+        .into()
+}
+
+#[pyfunction]
+fn union_extract(union_expr: PyExpr, field_name: PyExpr) -> PyExpr {
+    functions::core::union_extract()
+        .call(vec![union_expr.into(), field_name.into()])
+        .into()
+}
+
+#[pyfunction]
+fn version() -> PyExpr {
+    functions::core::version().call(vec![]).into()
+}
+
 // Array Functions
 array_fn!(array_append, array element);
 array_fn!(array_to_string, array delimiter);
@@ -674,10 +746,20 @@ 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);
 
+// Map Functions
+array_fn!(map_keys, map);
+array_fn!(map_values, map);
+array_fn!(map_extract, map key);
+array_fn!(map_entries, map);
+
 aggregate_function!(array_agg);
 aggregate_function!(max);
 aggregate_function!(min);
@@ -968,6 +1050,7 @@ pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_wrapped(wrap_pyfunction!(array_agg))?;
     m.add_wrapped(wrap_pyfunction!(arrow_typeof))?;
     m.add_wrapped(wrap_pyfunction!(arrow_cast))?;
+    m.add_wrapped(wrap_pyfunction!(arrow_metadata))?;
     m.add_wrapped(wrap_pyfunction!(ascii))?;
     m.add_wrapped(wrap_pyfunction!(asin))?;
     m.add_wrapped(wrap_pyfunction!(asinh))?;
@@ -1096,6 +1179,10 @@ pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_wrapped(wrap_pyfunction!(trim))?;
     m.add_wrapped(wrap_pyfunction!(trunc))?;
     m.add_wrapped(wrap_pyfunction!(upper))?;
+    m.add_wrapped(wrap_pyfunction!(get_field))?;
+    m.add_wrapped(wrap_pyfunction!(union_extract))?;
+    m.add_wrapped(wrap_pyfunction!(union_tag))?;
+    m.add_wrapped(wrap_pyfunction!(version))?;
     m.add_wrapped(wrap_pyfunction!(self::uuid))?; // Use self to avoid name collision
     m.add_wrapped(wrap_pyfunction!(var_pop))?;
     m.add_wrapped(wrap_pyfunction!(var_sample))?;
@@ -1154,9 +1241,24 @@ 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))?;
 
+    // Map Functions
+    m.add_wrapped(wrap_pyfunction!(make_map))?;
+    m.add_wrapped(wrap_pyfunction!(map_keys))?;
+    m.add_wrapped(wrap_pyfunction!(map_values))?;
+    m.add_wrapped(wrap_pyfunction!(map_extract))?;
+    m.add_wrapped(wrap_pyfunction!(map_entries))?;
+
     // Window Functions
     m.add_wrapped(wrap_pyfunction!(lead))?;
     m.add_wrapped(wrap_pyfunction!(lag))?;