Skip to content

Commit

Permalink
docs: Fix create_udf example (#14405)
Browse files Browse the repository at this point in the history
  • Loading branch information
nuno-faria authored Feb 2, 2025
1 parent 1e0531f commit d5428b2
Showing 1 changed file with 14 additions and 13 deletions.
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

0 comments on commit d5428b2

Please sign in to comment.