Pass codec capsule to table providers

Author: timsaucer

examples/datafusion-ffi-example/python/tests/_test_table_function.py

@@ -29,7 +29,7 @@
 
 def test_ffi_table_function_register():
     ctx = SessionContext()
-    table_func = MyTableFunction(ctx)
+    table_func = MyTableFunction()
     table_udtf = udtf(table_func, "my_table_func")
     ctx.register_udtf(table_udtf)
     result = ctx.sql("select * from my_table_func()").collect()
@@ -49,7 +49,7 @@ def test_ffi_table_function_register():
 
 def test_ffi_table_function_call_directly():
     ctx = SessionContext()
-    table_func = MyTableFunction(ctx)
+    table_func = MyTableFunction()
     table_udtf = udtf(table_func, "my_table_func")
 
     my_table = table_udtf()
@@ -77,9 +77,6 @@ class PythonTableFunction:
     provider, and this function takes no arguments
     """
 
-    def __init__(self, ctx: SessionContext) -> None:
-        self._ctx = ctx
-
     def __call__(
         self, num_cols: Expr, num_rows: Expr, num_batches: Expr
     ) -> TableProviderExportable:
@@ -88,7 +85,7 @@ def __call__(
             num_rows.to_variant().value_i64(),
             num_batches.to_variant().value_i64(),
         ]
-        return MyTableProvider(self._ctx, *args)
+        return MyTableProvider(*args)
 
 
 def common_table_function_test(test_ctx: SessionContext) -> None:
@@ -111,7 +108,7 @@ def common_table_function_test(test_ctx: SessionContext) -> None:
 
 def test_python_table_function():
     ctx = SessionContext()
-    table_func = PythonTableFunction(ctx)
+    table_func = PythonTableFunction()
     table_udtf = udtf(table_func, "my_table_func")
     ctx.register_udtf(table_udtf)
 
@@ -130,7 +127,7 @@ def my_udtf(
             num_rows.to_variant().value_i64(),
             num_batches.to_variant().value_i64(),
         ]
-        return MyTableProvider(ctx, *args)
+        return MyTableProvider(*args)
 
     ctx.register_udtf(my_udtf)
 

examples/datafusion-ffi-example/python/tests/_test_table_provider.py

@@ -24,7 +24,7 @@
 
 def test_table_loading():
     ctx = SessionContext()
-    table = MyTableProvider(ctx, 3, 2, 4)
+    table = MyTableProvider(3, 2, 4)
     ctx.register_table("t", table)
     result = ctx.table("t").collect()
 

examples/datafusion-ffi-example/src/table_function.rs

@@ -20,7 +20,6 @@ use std::sync::Arc;
 use datafusion_catalog::{TableFunctionImpl, TableProvider};
 use datafusion_common::error::Result as DataFusionResult;
 use datafusion_expr::Expr;
-use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec;
 use datafusion_ffi::udtf::FFI_TableFunction;
 use pyo3::types::PyCapsule;
 use pyo3::{pyclass, pymethods, Bound, PyAny, PyResult, Python};
@@ -30,37 +29,33 @@ use crate::utils::ffi_logical_codec_from_pycapsule;
 
 #[pyclass(name = "MyTableFunction", module = "datafusion_ffi_example", subclass)]
 #[derive(Debug, Clone)]
-pub(crate) struct MyTableFunction {
-    logical_codec: FFI_LogicalExtensionCodec,
-}
+pub(crate) struct MyTableFunction {}
 
 #[pymethods]
 impl MyTableFunction {
     #[new]
-    fn new(session: &Bound<PyAny>) -> PyResult<Self> {
-        let logical_codec = ffi_logical_codec_from_pycapsule(session)?;
-
-        Ok(Self { logical_codec })
+    fn new() -> Self {
+        Self {}
     }
 
     fn __datafusion_table_function__<'py>(
         &self,
         py: Python<'py>,
+        session: &Bound<PyAny>,
     ) -> PyResult<Bound<'py, PyCapsule>> {
         let name = cr"datafusion_table_function".into();
 
         let func = self.clone();
-        let provider =
-            FFI_TableFunction::new_with_ffi_codec(Arc::new(func), None, self.logical_codec.clone());
+        let codec = ffi_logical_codec_from_pycapsule(session)?;
+        let provider = FFI_TableFunction::new_with_ffi_codec(Arc::new(func), None, codec);
 
         PyCapsule::new(py, provider, Some(name))
     }
 }
 
 impl TableFunctionImpl for MyTableFunction {
     fn call(&self, _args: &[Expr]) -> DataFusionResult<Arc<dyn TableProvider>> {
-        let provider = MyTableProvider::new_from_ffi_session(self.logical_codec.clone(), 4, 3, 2)
-            .create_table()?;
+        let provider = MyTableProvider::new(4, 3, 2).create_table()?;
         Ok(Arc::new(provider))
     }
 }

examples/datafusion-ffi-example/src/table_provider.rs

@@ -21,7 +21,6 @@ use arrow_array::{ArrayRef, RecordBatch};
 use arrow_schema::{DataType, Field, Schema};
 use datafusion_catalog::MemTable;
 use datafusion_common::error::{DataFusionError, Result as DataFusionResult};
-use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec;
 use datafusion_ffi::table_provider::FFI_TableProvider;
 use pyo3::exceptions::PyRuntimeError;
 use pyo3::types::PyCapsule;
@@ -34,7 +33,6 @@ use crate::utils::ffi_logical_codec_from_pycapsule;
 #[pyclass(name = "MyTableProvider", module = "datafusion_ffi_example", subclass)]
 #[derive(Clone)]
 pub(crate) struct MyTableProvider {
-    logical_codec: FFI_LogicalExtensionCodec,
     num_cols: usize,
     num_rows: usize,
     num_batches: usize,
@@ -81,55 +79,31 @@ impl MyTableProvider {
     }
 }
 
-impl MyTableProvider {
-    pub fn new_from_ffi_session(
-        logical_codec: FFI_LogicalExtensionCodec,
-        num_cols: usize,
-        num_rows: usize,
-        num_batches: usize,
-    ) -> Self {
-        Self {
-            logical_codec,
-            num_cols,
-            num_rows,
-            num_batches,
-        }
-    }
-}
-
 #[pymethods]
 impl MyTableProvider {
     #[new]
-    pub fn new(
-        session: &Bound<PyAny>,
-        num_cols: usize,
-        num_rows: usize,
-        num_batches: usize,
-    ) -> PyResult<Self> {
-        let logical_codec = ffi_logical_codec_from_pycapsule(session)?;
-        Ok(Self {
-            logical_codec,
+    pub fn new(num_cols: usize, num_rows: usize, num_batches: usize) -> Self {
+        Self {
             num_cols,
             num_rows,
             num_batches,
-        })
+        }
     }
 
     pub fn __datafusion_table_provider__<'py>(
         &self,
         py: Python<'py>,
+        session: &Bound<PyAny>,
     ) -> PyResult<Bound<'py, PyCapsule>> {
         let name = cr"datafusion_table_provider".into();
 
         let provider = self
             .create_table()
             .map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
-        let provider = FFI_TableProvider::new_with_ffi_codec(
-            Arc::new(provider),
-            false,
-            None,
-            self.logical_codec.clone(),
-        );
+
+        let codec = ffi_logical_codec_from_pycapsule(session)?;
+        let provider =
+            FFI_TableProvider::new_with_ffi_codec(Arc::new(provider), false, None, codec);
 
         PyCapsule::new(py, provider, Some(name))
     }

python/datafusion/catalog.py

@@ -27,7 +27,7 @@
 if TYPE_CHECKING:
     import pyarrow as pa
 
-    from datafusion import DataFrame
+    from datafusion import DataFrame, SessionContext
     from datafusion.context import TableProviderExportable
 
 try:
@@ -65,9 +65,9 @@ def schema_names(self) -> set[str]:
         return self.catalog.schema_names()
 
     @staticmethod
-    def memory_catalog() -> Catalog:
+    def memory_catalog(ctx: SessionContext | None = None) -> Catalog:
         """Create an in-memory catalog provider."""
-        catalog = df_internal.catalog.RawCatalog.memory_catalog()
+        catalog = df_internal.catalog.RawCatalog.memory_catalog(ctx)
         return Catalog(catalog)
 
     def schema(self, name: str = "public") -> Schema:
@@ -112,9 +112,9 @@ def __repr__(self) -> str:
         return self._raw_schema.__repr__()
 
     @staticmethod
-    def memory_schema() -> Schema:
+    def memory_schema(ctx: SessionContext | None = None) -> Schema:
         """Create an in-memory schema provider."""
-        schema = df_internal.catalog.RawSchema.memory_schema()
+        schema = df_internal.catalog.RawSchema.memory_schema(ctx)
         return Schema(schema)
 
     def names(self) -> set[str]:
@@ -163,10 +163,12 @@ class Table:
     __slots__ = ("_inner",)
 
     def __init__(
-        self, table: Table | TableProviderExportable | DataFrame | pa.dataset.Dataset
+        self,
+        table: Table | TableProviderExportable | DataFrame | pa.dataset.Dataset,
+        ctx: SessionContext | None = None,
     ) -> None:
         """Constructor."""
-        self._inner = df_internal.catalog.RawTable(table)
+        self._inner = df_internal.catalog.RawTable(table, ctx)
 
     def __repr__(self) -> str:
         """Print a string representation of the table."""

python/datafusion/user_defined.py

@@ -27,6 +27,7 @@
 import pyarrow as pa
 
 import datafusion._internal as df_internal
+from datafusion import SessionContext
 from datafusion.expr import Expr
 
 if TYPE_CHECKING:
@@ -923,16 +924,14 @@ class TableFunction:
     """
 
     def __init__(
-        self,
-        name: str,
-        func: Callable[[], any],
+        self, name: str, func: Callable[[], any], ctx: SessionContext | None = None
     ) -> None:
         """Instantiate a user-defined table function (UDTF).
 
         See :py:func:`udtf` for a convenience function and argument
         descriptions.
         """
-        self._udtf = df_internal.TableFunction(name, func)
+        self._udtf = df_internal.TableFunction(name, func, ctx)
 
     def __call__(self, *args: Expr) -> Any:
         """Execute the UDTF and return a table provider."""