Skip to content

Commit 0f50a31

Browse files
committed
Merge branch 'main' into feat/add-missing-scalar-fns
2 parents df1ead1 + 99bc960 commit 0f50a31

File tree

5 files changed

+826
-1
lines changed

5 files changed

+826
-1
lines changed

AGENTS.md

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,3 +25,20 @@ Skills follow the [Agent Skills](https://agentskills.io) open standard. Each ski
2525

2626
- `SKILL.md` — The skill definition with YAML frontmatter (name, description, argument-hint) and detailed instructions.
2727
- Additional supporting files as needed.
28+
29+
## Python Function Docstrings
30+
31+
Every Python function must include a docstring with usage examples.
32+
33+
- **Examples are required**: Each function needs at least one doctest-style example
34+
demonstrating basic usage.
35+
- **Optional parameters**: If a function has optional parameters, include separate
36+
examples that show usage both without and with the optional arguments. Pass
37+
optional arguments using their keyword name (e.g., `step=dfn.lit(3)`) so readers
38+
can immediately see which parameter is being demonstrated.
39+
- **Reuse input data**: Use the same input data across examples wherever possible.
40+
The examples should demonstrate how different optional arguments change the output
41+
for the same input, making the effect of each option easy to understand.
42+
- **Alias functions**: Functions that are simple aliases (e.g., `list_sort` aliasing
43+
`array_sort`) only need a one-line description and a `See Also` reference to the
44+
primary function. They do not need their own examples.

README.md

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -275,7 +275,7 @@ needing to activate the virtual environment:
275275

276276
```bash
277277
uv run --no-project maturin develop --uv
278-
uv run --no-project pytest .
278+
uv run --no-project pytest
279279
```
280280

281281
To run the FFI tests within the examples folder, after you have built

crates/core/src/functions.rs

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -93,6 +93,57 @@ fn array_cat(exprs: Vec<PyExpr>) -> PyExpr {
9393
array_concat(exprs)
9494
}
9595

96+
#[pyfunction]
97+
fn array_distance(array1: PyExpr, array2: PyExpr) -> PyExpr {
98+
let args = vec![array1.into(), array2.into()];
99+
Expr::ScalarFunction(datafusion::logical_expr::expr::ScalarFunction::new_udf(
100+
datafusion::functions_nested::distance::array_distance_udf(),
101+
args,
102+
))
103+
.into()
104+
}
105+
106+
#[pyfunction]
107+
fn arrays_zip(exprs: Vec<PyExpr>) -> PyExpr {
108+
let exprs = exprs.into_iter().map(|x| x.into()).collect();
109+
datafusion::functions_nested::expr_fn::arrays_zip(exprs).into()
110+
}
111+
112+
#[pyfunction]
113+
#[pyo3(signature = (string, delimiter, null_string=None))]
114+
fn string_to_array(string: PyExpr, delimiter: PyExpr, null_string: Option<PyExpr>) -> PyExpr {
115+
let mut args = vec![string.into(), delimiter.into()];
116+
if let Some(null_string) = null_string {
117+
args.push(null_string.into());
118+
}
119+
Expr::ScalarFunction(datafusion::logical_expr::expr::ScalarFunction::new_udf(
120+
datafusion::functions_nested::string::string_to_array_udf(),
121+
args,
122+
))
123+
.into()
124+
}
125+
126+
#[pyfunction]
127+
#[pyo3(signature = (start, stop, step=None))]
128+
fn gen_series(start: PyExpr, stop: PyExpr, step: Option<PyExpr>) -> PyExpr {
129+
let mut args = vec![start.into(), stop.into()];
130+
if let Some(step) = step {
131+
args.push(step.into());
132+
}
133+
Expr::ScalarFunction(datafusion::logical_expr::expr::ScalarFunction::new_udf(
134+
datafusion::functions_nested::range::gen_series_udf(),
135+
args,
136+
))
137+
.into()
138+
}
139+
140+
#[pyfunction]
141+
fn make_map(keys: Vec<PyExpr>, values: Vec<PyExpr>) -> PyExpr {
142+
let keys = keys.into_iter().map(|x| x.into()).collect();
143+
let values = values.into_iter().map(|x| x.into()).collect();
144+
datafusion::functions_nested::map::map(keys, values).into()
145+
}
146+
96147
#[pyfunction]
97148
#[pyo3(signature = (array, element, index=None))]
98149
fn array_position(array: PyExpr, element: PyExpr, index: Option<i64>) -> PyExpr {
@@ -695,10 +746,20 @@ array_fn!(array_intersect, first_array second_array);
695746
array_fn!(array_union, array1 array2);
696747
array_fn!(array_except, first_array second_array);
697748
array_fn!(array_resize, array size value);
749+
array_fn!(array_any_value, array);
750+
array_fn!(array_max, array);
751+
array_fn!(array_min, array);
752+
array_fn!(array_reverse, array);
698753
array_fn!(cardinality, array);
699754
array_fn!(flatten, array);
700755
array_fn!(range, start stop step);
701756

757+
// Map Functions
758+
array_fn!(map_keys, map);
759+
array_fn!(map_values, map);
760+
array_fn!(map_extract, map key);
761+
array_fn!(map_entries, map);
762+
702763
aggregate_function!(array_agg);
703764
aggregate_function!(max);
704765
aggregate_function!(min);
@@ -1165,9 +1226,24 @@ pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
11651226
m.add_wrapped(wrap_pyfunction!(array_replace_all))?;
11661227
m.add_wrapped(wrap_pyfunction!(array_sort))?;
11671228
m.add_wrapped(wrap_pyfunction!(array_slice))?;
1229+
m.add_wrapped(wrap_pyfunction!(array_any_value))?;
1230+
m.add_wrapped(wrap_pyfunction!(array_distance))?;
1231+
m.add_wrapped(wrap_pyfunction!(array_max))?;
1232+
m.add_wrapped(wrap_pyfunction!(array_min))?;
1233+
m.add_wrapped(wrap_pyfunction!(array_reverse))?;
1234+
m.add_wrapped(wrap_pyfunction!(arrays_zip))?;
1235+
m.add_wrapped(wrap_pyfunction!(string_to_array))?;
1236+
m.add_wrapped(wrap_pyfunction!(gen_series))?;
11681237
m.add_wrapped(wrap_pyfunction!(flatten))?;
11691238
m.add_wrapped(wrap_pyfunction!(cardinality))?;
11701239

1240+
// Map Functions
1241+
m.add_wrapped(wrap_pyfunction!(make_map))?;
1242+
m.add_wrapped(wrap_pyfunction!(map_keys))?;
1243+
m.add_wrapped(wrap_pyfunction!(map_values))?;
1244+
m.add_wrapped(wrap_pyfunction!(map_extract))?;
1245+
m.add_wrapped(wrap_pyfunction!(map_entries))?;
1246+
11711247
// Window Functions
11721248
m.add_wrapped(wrap_pyfunction!(lead))?;
11731249
m.add_wrapped(wrap_pyfunction!(lag))?;

0 commit comments

Comments
 (0)