datafusion_python/
udtf.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::catalog::{Session, TableFunctionArgs, TableFunctionImpl, TableProvider};
22use datafusion::error::{DataFusionError, Result as DataFusionResult};
23use datafusion::execution::context::SessionContext;
24use datafusion::execution::session_state::SessionState;
25use datafusion::logical_expr::Expr;
26use datafusion_ffi::udtf::FFI_TableFunction;
27use pyo3::IntoPyObjectExt;
28use pyo3::exceptions::{PyImportError, PyTypeError};
29use pyo3::prelude::*;
30use pyo3::types::{PyCapsule, PyDict, PyTuple, PyType};
31
32use crate::context::PySessionContext;
33use crate::errors::{py_datafusion_err, to_datafusion_err};
34use crate::expr::PyExpr;
35use crate::table::PyTable;
36
37/// A pure-Python UDTF callable plus the metadata we discovered about it
38/// at registration time.
39#[derive(Debug, Clone)]
40pub(crate) struct PythonTableFunctionCallable {
41    pub(crate) callable: Arc<Py<PyAny>>,
42    /// When true, the calling :class:`SessionContext` is passed to the
43    /// callable as a ``session`` keyword argument on every invocation.
44    /// Opt-in at registration time via ``with_session=True`` on the
45    /// Python wrapper.
46    pub(crate) inject_session_on_call: bool,
47}
48
49/// Represents a user defined table function
50#[pyclass(from_py_object, frozen, name = "TableFunction", module = "datafusion")]
51#[derive(Debug, Clone)]
52pub struct PyTableFunction {
53    pub(crate) name: String,
54    pub(crate) inner: PyTableFunctionInner,
55}
56
57#[derive(Debug, Clone)]
58pub(crate) enum PyTableFunctionInner {
59    PythonFunction(PythonTableFunctionCallable),
60    FFIFunction(Arc<dyn TableFunctionImpl>),
61}
62
63#[pymethods]
64impl PyTableFunction {
65    #[new]
66    #[pyo3(signature=(name, func, session, inject_session_on_call=false))]
67    pub fn new(
68        name: &str,
69        func: Bound<'_, PyAny>,
70        session: Option<Bound<PyAny>>,
71        inject_session_on_call: bool,
72    ) -> PyResult<Self> {
73        let inner = if func.hasattr("__datafusion_table_function__")? {
74            let py = func.py();
75            let session = match session {
76                Some(session) => session,
77                None => PySessionContext::global_ctx()?.into_bound_py_any(py)?,
78            };
79            let capsule = func
80                .getattr("__datafusion_table_function__")?
81                .call1((session,)).map_err(|err| {
82                if err.get_type(py).is(PyType::new::<PyTypeError>(py)) {
83                    PyImportError::new_err("Incompatible libraries. DataFusion 52.0.0 introduced an incompatible signature change for table functions. Either downgrade DataFusion or upgrade your function library.")
84                } else {
85                    err
86                }
87            })?;
88            let capsule = capsule.cast::<PyCapsule>()?;
89            let data: NonNull<FFI_TableFunction> = capsule
90                .pointer_checked(Some(c"datafusion_table_function"))?
91                .cast();
92            let ffi_func = unsafe { data.as_ref() };
93            let foreign_func: Arc<dyn TableFunctionImpl> = ffi_func.to_owned().into();
94
95            PyTableFunctionInner::FFIFunction(foreign_func)
96        } else {
97            PyTableFunctionInner::PythonFunction(PythonTableFunctionCallable {
98                callable: Arc::new(func.unbind()),
99                inject_session_on_call,
100            })
101        };
102
103        Ok(Self {
104            name: name.to_string(),
105            inner,
106        })
107    }
108
109    #[pyo3(signature = (*args))]
110    pub fn __call__(&self, args: Vec<PyExpr>) -> PyResult<PyTable> {
111        let args: Vec<Expr> = args.iter().map(|e| e.expr.clone()).collect();
112        let global = PySessionContext::global_ctx()?;
113        let state = global.ctx.state();
114        let table_provider = self
115            .call_with_args(TableFunctionArgs::new(&args, &state))
116            .map_err(py_datafusion_err)?;
117
118        Ok(PyTable::from(table_provider))
119    }
120
121    fn __repr__(&self) -> PyResult<String> {
122        Ok(format!("TableUDF({})", self.name))
123    }
124}
125
126/// Materialize a fresh :class:`PySessionContext` from the borrowed
127/// ``&dyn Session`` handed in at call time.
128///
129/// Upstream invokes ``call_with_args`` with a trait-object reference
130/// rather than an owned context; we downcast it to the canonical
131/// :class:`SessionState` impl and rebuild a :class:`SessionContext`
132/// (sharing the same registries via the Arc-heavy interior of
133/// :class:`SessionState`).
134///
135/// The downcast is defensive. Every path that reaches a pure-Python
136/// UDTF today hands us a `SessionState`: the SQL planner builds the
137/// args from its own `SessionState`, and `PyTableFunction::__call__`
138/// uses the global context's state. A non-`SessionState` session
139/// (e.g. a `ForeignSession`) would only arrive if this UDTF were
140/// exported across the FFI boundary to a foreign-library consumer,
141/// which datafusion-python does not do. Should that change, this
142/// returns an error rather than silently misbehaving.
143fn py_session_from_session(session: &dyn Session) -> DataFusionResult<PySessionContext> {
144    let state = session
145        .as_any()
146        .downcast_ref::<SessionState>()
147        .ok_or_else(|| {
148            DataFusionError::Execution(
149                "Cannot expose this UDTF's calling session to Python: the \
150                 session is not a SessionState. Drop the `session` keyword \
151                 from the callback signature to fall back to the \
152                 expression-only call form."
153                    .to_string(),
154            )
155        })?;
156    Ok(PySessionContext::from(SessionContext::new_with_state(
157        state.clone(),
158    )))
159}
160
161#[allow(clippy::result_large_err)]
162fn call_python_table_function(
163    func: &PythonTableFunctionCallable,
164    args: TableFunctionArgs,
165) -> DataFusionResult<Arc<dyn TableProvider>> {
166    let py_session = if func.inject_session_on_call {
167        Some(py_session_from_session(args.session())?)
168    } else {
169        None
170    };
171    let py_exprs = args
172        .exprs()
173        .iter()
174        .map(|arg| PyExpr::from(arg.clone()))
175        .collect::<Vec<_>>();
176
177    Python::attach(|py| {
178        let py_args = PyTuple::new(py, py_exprs)?;
179        let provider_obj = if let Some(session) = py_session {
180            let kwargs = PyDict::new(py);
181            kwargs.set_item("session", session.into_pyobject(py)?)?;
182            func.callable.call(py, py_args, Some(&kwargs))?
183        } else {
184            func.callable.call1(py, py_args)?
185        };
186        let provider = provider_obj.bind(py).clone();
187
188        Ok::<Arc<dyn TableProvider>, PyErr>(PyTable::new(provider, None)?.table)
189    })
190    .map_err(to_datafusion_err)
191}
192
193impl TableFunctionImpl for PyTableFunction {
194    fn call_with_args(&self, args: TableFunctionArgs) -> DataFusionResult<Arc<dyn TableProvider>> {
195        match &self.inner {
196            PyTableFunctionInner::FFIFunction(func) => func.call_with_args(args),
197            PyTableFunctionInner::PythonFunction(callable) => {
198                call_python_table_function(callable, args)
199            }
200        }
201    }
202}
datafusion_python/udtf.rs

datafusion_python/
udtf.rs
