fn pyarrow_function_to_rust()

in src/udf.rs [36:63]

• 23 lines of code
• 1 McCabe index (conditional complexity)

fn pyarrow_function_to_rust(
    func: PyObject,
) -> impl Fn(&[ArrayRef]) -> Result<ArrayRef, DataFusionError> {
    move |args: &[ArrayRef]| -> Result<ArrayRef, DataFusionError> {
        Python::with_gil(|py| {
            // 1. cast args to Pyarrow arrays
            let py_args = args
                .iter()
                .map(|arg| {
                    arg.into_data()
                        .to_pyarrow(py)
                        .map_err(|e| DataFusionError::Execution(format!("{e:?}")))
                })
                .collect::<Result<Vec<_>, _>>()?;
            let py_args = PyTuple::new(py, py_args).map_err(to_datafusion_err)?;

            // 2. call function
            let value = func
                .call(py, py_args, None)
                .map_err(|e| DataFusionError::Execution(format!("{e:?}")))?;

            // 3. cast to arrow::array::Array
            let array_data = ArrayData::from_pyarrow_bound(value.bind(py))
                .map_err(|e| DataFusionError::Execution(format!("{e:?}")))?;
            Ok(make_array(array_data))
        })
    }
}