Initial commit of udtf work

Author: timsaucer

python/datafusion/__init__.py

@@ -51,7 +51,16 @@
 from .io import read_avro, read_csv, read_json, read_parquet
 from .plan import ExecutionPlan, LogicalPlan
 from .record_batch import RecordBatch, RecordBatchStream
-from .udf import Accumulator, AggregateUDF, ScalarUDF, WindowUDF, udaf, udf, udwf
+from .udf import (
+    Accumulator,
+    AggregateUDF,
+    ScalarUDF,
+    TableFunction,
+    WindowUDF,
+    udaf,
+    udf,
+    udwf,
+)
 
 __version__ = importlib_metadata.version(__name__)
 
@@ -74,6 +83,7 @@
     "SessionConfig",
     "SessionContext",
     "Table",
+    "TableFunction",
     "WindowFrame",
     "WindowUDF",
     "col",

python/datafusion/context.py

@@ -30,7 +30,7 @@
 from datafusion.dataframe import DataFrame
 from datafusion.expr import Expr, SortExpr, sort_list_to_raw_sort_list
 from datafusion.record_batch import RecordBatchStream
-from datafusion.udf import AggregateUDF, ScalarUDF, WindowUDF
+from datafusion.udf import AggregateUDF, ScalarUDF, TableFunction, WindowUDF
 
 from ._internal import RuntimeEnvBuilder as RuntimeEnvBuilderInternal
 from ._internal import SessionConfig as SessionConfigInternal
@@ -752,6 +752,10 @@ def register_table_provider(
         """
         self.ctx.register_table_provider(name, provider)
 
+    def register_udtf(self, name: str, func: TableFunction) -> None:
+        """Register a user defined table function."""
+        self.ctx.register_udtf(name, func._udtf)
+
     def register_record_batches(
         self, name: str, partitions: list[list[pa.RecordBatch]]
     ) -> None:

python/datafusion/udf.py

@@ -760,8 +760,72 @@ def wrapper(*args: Any, **kwargs: Any) -> Expr:
         return decorator
 
 
+class TableFunction:
+    """Class for performing user-defined table functions (UDTF).
+
+    Table functions generate new table providers based on the
+    input expressions.
+    """
+
+    def __init__(
+        self,
+        name: str,
+        func: Callable[[], any],
+    ) -> None:
+        """Instantiate a user-defined table function (UDTF).
+
+        See :py:func:`udtf` for a convenience function and argument
+        descriptions.
+        """
+        self._udtf = df_internal.user_defined.TableFunction(name, func)
+
+    def __call__(self, *args: Expr) -> Any:
+        """Execute the UDTF and return a table provider."""
+        args_raw = [arg.expr for arg in args]
+        return Expr(self._udtf.__call__(*args_raw))
+
+    @overload
+    @staticmethod
+    def udtf(
+        name: str,
+    ) -> Callable[..., Any]: ...
+
+    @overload
+    @staticmethod
+    def udtf(
+        func: Callable[[], Any],
+        name: str,
+    ) -> TableFunction: ...
+
+    @staticmethod
+    def udtf(*args: Any, **kwargs: Any):
+        """Create a new User-Defined Table Function (UDTF)."""
+        if args and callable(args[0]):
+            # Case 1: Used as a function, require the first parameter to be callable
+            return TableFunction._create_table_udf(*args, **kwargs)
+        # Case 2: Used as a decorator with parameters
+        return TableFunction._create_table_udf_decorator(*args, **kwargs)
+
+    @staticmethod
+    def _create_table_udf(
+        func: Callable[..., Any],
+        name: str,
+    ) -> TableFunction:
+        """Create a TableFunction instance from function arguments."""
+        if not callable(func):
+            msg = "`func` must be callable."
+            raise TypeError(msg)
+
+        return TableFunction(name, func)
+
+    def __repr__(self) -> str:
+        """User printable representation."""
+        return self._udtf.__repr__()
+
+
 # Convenience exports so we can import instead of treating as
 # variables at the package root
 udf = ScalarUDF.udf
 udaf = AggregateUDF.udaf
 udwf = WindowUDF.udwf
+udtf = TableFunction.udtf

src/context.rs

@@ -43,6 +43,7 @@ use crate::sql::logical::PyLogicalPlan;
 use crate::store::StorageContexts;
 use crate::udaf::PyAggregateUDF;
 use crate::udf::PyScalarUDF;
+use crate::udtf::PyTableFunction;
 use crate::udwf::PyWindowUDF;
 use crate::utils::{get_global_ctx, get_tokio_runtime, validate_pycapsule, wait_for_future};
 use datafusion::arrow::datatypes::{DataType, Schema, SchemaRef};
@@ -390,6 +391,12 @@ impl PySessionContext {
         Ok(())
     }
 
+    pub fn register_udtf(&mut self, func: PyTableFunction) {
+        let name = func.name.clone();
+        let func = Arc::new(func);
+        self.ctx.register_udtf(&name, func);
+    }
+
     /// Returns a PyDataFrame whose plan corresponds to the SQL statement.
     pub fn sql(&mut self, query: &str, py: Python) -> PyDataFusionResult<PyDataFrame> {
         let result = self.ctx.sql(query);

src/lib.rs

@@ -60,6 +60,7 @@ pub mod substrait;
 mod udaf;
 #[allow(clippy::borrow_deref_ref)]
 mod udf;
+pub mod udtf;
 mod udwf;
 pub mod utils;
 
@@ -88,6 +89,7 @@ fn _internal(py: Python, m: Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<udf::PyScalarUDF>()?;
     m.add_class::<udaf::PyAggregateUDF>()?;
     m.add_class::<udwf::PyWindowUDF>()?;
+    m.add_class::<udtf::PyTableFunction>()?;
     m.add_class::<config::PyConfig>()?;
     m.add_class::<sql::logical::PyLogicalPlan>()?;
     m.add_class::<physical_plan::PyExecutionPlan>()?;

src/udtf.rs

@@ -0,0 +1,89 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::sync::Arc;
+
+use pyo3::prelude::*;
+
+use crate::dataframe::PyTableProvider;
+use crate::errors::py_datafusion_err;
+use crate::expr::PyExpr;
+use crate::utils::validate_pycapsule;
+use datafusion::catalog::{TableFunctionImpl, TableProvider};
+use datafusion::common::exec_err;
+use datafusion::logical_expr::Expr;
+use datafusion_ffi::udtf::{FFI_TableFunction, ForeignTableFunction};
+use pyo3::types::PyCapsule;
+
+/// Represents a user defined table function
+#[pyclass(name = "TableFunction", module = "datafusion")]
+#[derive(Debug, Clone)]
+pub struct PyTableFunction {
+    pub(crate) name: String,
+    pub(crate) inner: PyTableFunctionInner,
+}
+
+// TODO: Implement pure python based user defined table functions
+#[derive(Debug, Clone)]
+pub(crate) enum PyTableFunctionInner {
+    // PythonFunction(Arc<PyObject>),
+    FFIFunction(Arc<dyn TableFunctionImpl>),
+}
+
+#[pymethods]
+impl PyTableFunction {
+    #[new]
+    #[pyo3(signature=(name, func))]
+    pub fn new(name: &str, func: Bound<'_, PyAny>) -> PyResult<Self> {
+        if func.hasattr("__datafusion_table_function__")? {
+            let capsule = func.getattr("__datafusion_table_function__")?.call0()?;
+            let capsule = capsule.downcast::<PyCapsule>().map_err(py_datafusion_err)?;
+            validate_pycapsule(capsule, "datafusion_table_function")?;
+
+            let ffi_func = unsafe { capsule.reference::<FFI_TableFunction>() };
+            let foreign_func: ForeignTableFunction = ffi_func.to_owned().into();
+
+            Ok(Self {
+                name: name.to_string(),
+                inner: PyTableFunctionInner::FFIFunction(Arc::new(foreign_func)),
+            })
+        } else {
+            exec_err!("Python based Table Functions are not yet implemented")
+                .map_err(py_datafusion_err)
+        }
+    }
+
+    #[pyo3(signature = (*args))]
+    pub fn __call__(&self, args: Vec<PyExpr>) -> PyResult<PyTableProvider> {
+        let args: Vec<Expr> = args.iter().map(|e| e.expr.clone()).collect();
+        let table_provider = self.call(&args).map_err(py_datafusion_err)?;
+
+        Ok(PyTableProvider::new(table_provider))
+    }
+
+    fn __repr__(&self) -> PyResult<String> {
+        Ok(format!("TableUDF({})", self.name))
+    }
+}
+
+impl TableFunctionImpl for PyTableFunction {
+    fn call(&self, args: &[Expr]) -> datafusion::common::Result<Arc<dyn TableProvider>> {
+        match &self.inner {
+            PyTableFunctionInner::FFIFunction(func) => func.call(args),
+        }
+    }
+}