Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

docs: Fix create_udf examples #14405

Merged
merged 1 commit into from
Feb 2, 2025
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 14 additions & 13 deletions docs/source/library-user-guide/adding-udfs.md
Original file line number Diff line number Diff line change
Expand Up @@ -123,21 +123,21 @@ There is a an older, more concise, but also more limited API [`create_udf`] avai
use std::sync::Arc;

use datafusion::arrow::array::{ArrayRef, Int64Array};
use datafusion::common::Result;
use datafusion::common::cast::as_int64_array;
use datafusion::physical_plan::functions::columnar_values_to_array;
use datafusion::common::Result;
use datafusion::logical_expr::ColumnarValue;

pub fn add_one(args: &[ColumnarValue]) -> Result<ArrayRef> {
pub fn add_one(args: &[ColumnarValue]) -> Result<ColumnarValue> {
// Error handling omitted for brevity
let args = columnar_values_to_array(args)?;
let args = ColumnarValue::values_to_arrays(args)?;
let i64s = as_int64_array(&args[0])?;

let new_array = i64s
.iter()
.map(|array_elem| array_elem.map(|value| value + 1))
.collect::<Int64Array>();

Ok(Arc::new(new_array))
Ok(ColumnarValue::from(Arc::new(new_array) as ArrayRef))
}
```

Expand All @@ -146,10 +146,11 @@ This "works" in isolation, i.e. if you have a slice of `ArrayRef`s, you can call

```rust
let input = vec![Some(1), None, Some(3)];
let input = Arc::new(Int64Array::from(input)) as ArrayRef;
let input = ColumnarValue::from(Arc::new(Int64Array::from(input)) as ArrayRef);

let result = add_one( & [input]).unwrap();
let result = result.as_any().downcast_ref::<Int64Array>().unwrap();
let result = add_one(&[input]).unwrap();
let binding = result.into_array(1).unwrap();
let result = binding.as_any().downcast_ref::<Int64Array>().unwrap();

assert_eq!(result, &Int64Array::from(vec![Some(2), None, Some(4)]));
```
Expand All @@ -169,11 +170,11 @@ use datafusion::arrow::datatypes::DataType;
use std::sync::Arc;

let udf = create_udf(
"add_one",
vec![DataType::Int64],
Arc::new(DataType::Int64),
Volatility::Immutable,
Arc::new(add_one),
"add_one",
vec![DataType::Int64],
DataType::Int64,
Volatility::Immutable,
Arc::new(add_one),
);
```

Expand Down