datafusion_python/
udaf.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use std::ptr::NonNull;
19use std::sync::Arc;
20
21use datafusion::arrow::array::ArrayRef;
22use datafusion::arrow::datatypes::{DataType, Field, FieldRef};
23use datafusion::arrow::pyarrow::{PyArrowType, ToPyArrow};
24use datafusion::common::ScalarValue;
25use datafusion::error::{DataFusionError, Result};
26use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs};
27use datafusion::logical_expr::{
28    Accumulator, AccumulatorFactoryFunction, AggregateUDF, AggregateUDFImpl, Signature, Volatility,
29};
30use datafusion_ffi::udaf::FFI_AggregateUDF;
31use datafusion_python_util::parse_volatility;
32use pyo3::prelude::*;
33use pyo3::types::{PyCapsule, PyTuple};
34
35use crate::common::data_type::PyScalarValue;
36use crate::errors::{PyDataFusionResult, py_datafusion_err, to_datafusion_err};
37use crate::expr::PyExpr;
38
39#[derive(Debug)]
40struct RustAccumulator {
41    accum: Py<PyAny>,
42}
43
44impl RustAccumulator {
45    fn new(accum: Py<PyAny>) -> Self {
46        Self { accum }
47    }
48}
49
50impl Accumulator for RustAccumulator {
51    fn state(&mut self) -> Result<Vec<ScalarValue>> {
52        Python::attach(|py| -> PyResult<Vec<ScalarValue>> {
53            let values = self.accum.bind(py).call_method0("state")?;
54            let mut scalars = Vec::new();
55            for item in values.try_iter()? {
56                let item: Bound<'_, PyAny> = item?;
57                let scalar = item.extract::<PyScalarValue>()?.0;
58                scalars.push(scalar);
59            }
60            Ok(scalars)
61        })
62        .map_err(|e| DataFusionError::Execution(format!("{e}")))
63    }
64
65    fn evaluate(&mut self) -> Result<ScalarValue> {
66        Python::attach(|py| -> PyResult<ScalarValue> {
67            let value = self.accum.bind(py).call_method0("evaluate")?;
68            value.extract::<PyScalarValue>().map(|v| v.0)
69        })
70        .map_err(|e| DataFusionError::Execution(format!("{e}")))
71    }
72
73    fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
74        Python::attach(|py| {
75            // 1. cast args to Pyarrow array
76            let py_args = values
77                .iter()
78                .map(|arg| arg.to_data().to_pyarrow(py).unwrap())
79                .collect::<Vec<_>>();
80            let py_args = PyTuple::new(py, py_args).map_err(to_datafusion_err)?;
81
82            // 2. call function
83            self.accum
84                .bind(py)
85                .call_method1("update", py_args)
86                .map_err(|e| DataFusionError::Execution(format!("{e}")))?;
87
88            Ok(())
89        })
90    }
91
92    fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> {
93        Python::attach(|py| {
94            // // 1. cast states to Pyarrow arrays
95            let py_states: Result<Vec<Bound<'_, PyAny>>> = states
96                .iter()
97                .map(|state| {
98                    state
99                        .to_data()
100                        .to_pyarrow(py)
101                        .map_err(|e| DataFusionError::Execution(format!("{e}")))
102                })
103                .collect();
104
105            // 2. call merge
106            self.accum
107                .bind(py)
108                .call_method1("merge", (py_states?,))
109                .map_err(|e| DataFusionError::Execution(format!("{e}")))?;
110
111            Ok(())
112        })
113    }
114
115    fn size(&self) -> usize {
116        std::mem::size_of_val(self)
117    }
118
119    fn retract_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
120        Python::attach(|py| {
121            // 1. cast args to Pyarrow array
122            let py_args = values
123                .iter()
124                .map(|arg| arg.to_data().to_pyarrow(py).unwrap())
125                .collect::<Vec<_>>();
126            let py_args = PyTuple::new(py, py_args).map_err(to_datafusion_err)?;
127
128            // 2. call function
129            self.accum
130                .bind(py)
131                .call_method1("retract_batch", py_args)
132                .map_err(|e| DataFusionError::Execution(format!("{e}")))?;
133
134            Ok(())
135        })
136    }
137
138    fn supports_retract_batch(&self) -> bool {
139        Python::attach(
140            |py| match self.accum.bind(py).call_method0("supports_retract_batch") {
141                Ok(x) => x.extract().unwrap_or(false),
142                Err(_) => false,
143            },
144        )
145    }
146}
147
148fn instantiate_accumulator(accum: &Py<PyAny>) -> Result<Box<dyn Accumulator>> {
149    let instance = Python::attach(|py| {
150        accum
151            .call0(py)
152            .map_err(|e| DataFusionError::Execution(format!("{e}")))
153    })?;
154    Ok(Box::new(RustAccumulator::new(instance)))
155}
156
157/// Wrap a Python accumulator factory in an `AccumulatorFactoryFunction`.
158///
159/// Retained for downstream callers that previously consumed this
160/// helper to build a [`AccumulatorFactoryFunction`] for `create_udaf`
161/// or similar factory-based APIs. New in-crate code should construct
162/// a [`PythonFunctionAggregateUDF`] directly so the codec can downcast
163/// and ship it inline.
164pub fn to_rust_accumulator(accum: Py<PyAny>) -> AccumulatorFactoryFunction {
165    Arc::new(move |_args| instantiate_accumulator(&accum))
166}
167
168/// Named-struct `AggregateUDFImpl` for Python-defined aggregate UDFs.
169/// Holds the Python accumulator factory directly so the codec can
170/// downcast and cloudpickle it across process boundaries.
171#[derive(Debug)]
172pub(crate) struct PythonFunctionAggregateUDF {
173    name: String,
174    accumulator: Py<PyAny>,
175    signature: Signature,
176    return_type: DataType,
177    state_fields: Vec<FieldRef>,
178}
179
180impl PythonFunctionAggregateUDF {
181    fn new(
182        name: String,
183        accumulator: Py<PyAny>,
184        input_types: Vec<DataType>,
185        return_type: DataType,
186        state_types: Vec<DataType>,
187        volatility: Volatility,
188    ) -> Self {
189        let signature = Signature::exact(input_types, volatility);
190        let state_fields = state_types
191            .into_iter()
192            .enumerate()
193            .map(|(i, t)| Arc::new(Field::new(format!("state_{i}"), t, true)))
194            .collect();
195        Self {
196            name,
197            accumulator,
198            signature,
199            return_type,
200            state_fields,
201        }
202    }
203
204    /// Stored Python callable that returns a fresh accumulator instance
205    /// per partition. Consumed by the codec to cloudpickle the factory
206    /// across process boundaries.
207    pub(crate) fn accumulator(&self) -> &Py<PyAny> {
208        &self.accumulator
209    }
210
211    pub(crate) fn return_type(&self) -> &DataType {
212        &self.return_type
213    }
214
215    pub(crate) fn state_fields_ref(&self) -> &[FieldRef] {
216        &self.state_fields
217    }
218
219    /// Reconstruct a `PythonFunctionAggregateUDF` from the parts emitted
220    /// by the codec. `state_fields` carries the full state schema
221    /// (names, data types, nullability, metadata) — the codec extracts
222    /// it from the IPC payload, so the post-decode state schema is
223    /// identical to the pre-encode one. Use [`Self::new`] when only
224    /// `Vec<DataType>` is available (e.g. the Python constructor path,
225    /// where field names are synthesized).
226    pub(crate) fn from_parts(
227        name: String,
228        accumulator: Py<PyAny>,
229        input_types: Vec<DataType>,
230        return_type: DataType,
231        state_fields: Vec<FieldRef>,
232        volatility: Volatility,
233    ) -> Self {
234        Self {
235            name,
236            accumulator,
237            signature: Signature::exact(input_types, volatility),
238            return_type,
239            state_fields,
240        }
241    }
242}
243
244impl Eq for PythonFunctionAggregateUDF {}
245impl PartialEq for PythonFunctionAggregateUDF {
246    fn eq(&self, other: &Self) -> bool {
247        self.name == other.name
248            && self.signature == other.signature
249            && self.return_type == other.return_type
250            && self.state_fields == other.state_fields
251            // Pointer-identity fast path: `Arc`-shared clones of the
252            // same UDF skip the GIL roundtrip. Falls through to Python
253            // `__eq__` only for two distinct callables.
254            && (self.accumulator.as_ptr() == other.accumulator.as_ptr()
255                || Python::attach(|py| {
256                    // See `PythonFunctionScalarUDF::eq` for the
257                    // rationale on swallowing the exception as `false`
258                    // and logging at `debug`. FIXME: revisit if
259                    // upstream `AggregateUDFImpl` exposes a fallible
260                    // `PartialEq`.
261                    self.accumulator
262                        .bind(py)
263                        .eq(other.accumulator.bind(py))
264                        .unwrap_or_else(|e| {
265                            log::debug!(
266                                target: "datafusion_python::udaf",
267                                "PythonFunctionAggregateUDF {:?} __eq__ raised; treating as unequal: {e}",
268                                self.name,
269                            );
270                            false
271                        })
272                }))
273    }
274}
275
276impl std::hash::Hash for PythonFunctionAggregateUDF {
277    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
278        // See `PythonFunctionScalarUDF`'s `Hash` impl for the
279        // rationale: hash the identifying header only and let
280        // `PartialEq` disambiguate callables.
281        self.name.hash(state);
282        self.signature.hash(state);
283        self.return_type.hash(state);
284        for f in &self.state_fields {
285            f.hash(state);
286        }
287    }
288}
289
290impl AggregateUDFImpl for PythonFunctionAggregateUDF {
291    fn name(&self) -> &str {
292        &self.name
293    }
294
295    fn signature(&self) -> &Signature {
296        &self.signature
297    }
298
299    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
300        Ok(self.return_type.clone())
301    }
302
303    fn accumulator(&self, _acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> {
304        instantiate_accumulator(&self.accumulator)
305    }
306
307    fn state_fields(&self, _args: StateFieldsArgs) -> Result<Vec<FieldRef>> {
308        Ok(self.state_fields.clone())
309    }
310}
311
312fn aggregate_udf_from_capsule(capsule: &Bound<'_, PyCapsule>) -> PyDataFusionResult<AggregateUDF> {
313    let data: NonNull<FFI_AggregateUDF> = capsule
314        .pointer_checked(Some(c"datafusion_aggregate_udf"))?
315        .cast();
316    let udaf = unsafe { data.as_ref() };
317    let udaf: Arc<dyn AggregateUDFImpl> = udaf.into();
318
319    Ok(AggregateUDF::new_from_shared_impl(udaf))
320}
321
322/// Represents an AggregateUDF
323#[pyclass(
324    from_py_object,
325    frozen,
326    name = "AggregateUDF",
327    module = "datafusion",
328    subclass
329)]
330#[derive(Debug, Clone)]
331pub struct PyAggregateUDF {
332    pub(crate) function: AggregateUDF,
333}
334
335#[pymethods]
336impl PyAggregateUDF {
337    #[new]
338    #[pyo3(signature=(name, accumulator, input_type, return_type, state_type, volatility))]
339    fn new(
340        name: &str,
341        accumulator: Py<PyAny>,
342        input_type: PyArrowType<Vec<DataType>>,
343        return_type: PyArrowType<DataType>,
344        state_type: PyArrowType<Vec<DataType>>,
345        volatility: &str,
346    ) -> PyResult<Self> {
347        let py_udf = PythonFunctionAggregateUDF::new(
348            name.to_string(),
349            accumulator,
350            input_type.0,
351            return_type.0,
352            state_type.0,
353            parse_volatility(volatility)?,
354        );
355        let function = AggregateUDF::new_from_impl(py_udf);
356        Ok(Self { function })
357    }
358
359    #[staticmethod]
360    pub fn from_pycapsule(func: Bound<'_, PyAny>) -> PyDataFusionResult<Self> {
361        if func.is_instance_of::<PyCapsule>() {
362            let capsule = func.cast::<PyCapsule>().map_err(py_datafusion_err)?;
363            let function = aggregate_udf_from_capsule(capsule)?;
364            return Ok(Self { function });
365        }
366
367        if func.hasattr("__datafusion_aggregate_udf__")? {
368            let capsule = func.getattr("__datafusion_aggregate_udf__")?.call0()?;
369            let capsule = capsule.cast::<PyCapsule>().map_err(py_datafusion_err)?;
370            let function = aggregate_udf_from_capsule(capsule)?;
371            return Ok(Self { function });
372        }
373
374        Err(crate::errors::PyDataFusionError::Common(
375            "__datafusion_aggregate_udf__ does not exist on AggregateUDF object.".to_string(),
376        ))
377    }
378
379    /// creates a new PyExpr with the call of the udf
380    #[pyo3(signature = (*args))]
381    fn __call__(&self, args: Vec<PyExpr>) -> PyResult<PyExpr> {
382        let args = args.iter().map(|e| e.expr.clone()).collect();
383        Ok(self.function.call(args).into())
384    }
385
386    fn __repr__(&self) -> PyResult<String> {
387        Ok(format!("AggregateUDF({})", self.function.name()))
388    }
389
390    #[getter]
391    fn name(&self) -> &str {
392        self.function.name()
393    }
394}
datafusion_python/udaf.rs

datafusion_python/
udaf.rs
