feat: add TableProvider support and related functionality

Author: kosiew

python/datafusion/__init__.py

@@ -33,7 +33,7 @@
 from . import functions, object_store, substrait, unparser
 
 # The following imports are okay to remain as opaque to the user.
-from ._internal import Config
+from ._internal import Config, TableProvider
 from .catalog import Catalog, Database, Table
 from .col import col, column
 from .common import (
@@ -90,6 +90,7 @@
     "SessionContext",
     "Table",
     "TableFunction",
+    "TableProvider",
     "WindowFrame",
     "WindowUDF",
     "catalog",

python/datafusion/dataframe.py

@@ -40,6 +40,7 @@
 from datafusion._internal import DataFrame as DataFrameInternal
 from datafusion._internal import ParquetColumnOptions as ParquetColumnOptionsInternal
 from datafusion._internal import ParquetWriterOptions as ParquetWriterOptionsInternal
+from datafusion._internal import TableProvider as TableProviderInternal
 from datafusion.expr import Expr, SortExpr, sort_or_default
 from datafusion.plan import ExecutionPlan, LogicalPlan
 from datafusion.record_batch import RecordBatchStream
@@ -307,8 +308,8 @@ def __init__(self, df: DataFrameInternal) -> None:
         """
         self.df = df
 
-    def into_view(self) -> pa.Table:
-        """Convert DataFrame as a ViewTable which can be used in register_table."""
+    def into_view(self) -> TableProviderInternal:
+        """Convert ``DataFrame`` into a ``TableProvider`` view for registration."""
         return self.df.into_view()
 
     def __getitem__(self, key: str | list[str]) -> DataFrame:

src/dataframe.rs

@@ -36,20 +36,19 @@ use datafusion::error::DataFusionError;
 use datafusion::execution::SendableRecordBatchStream;
 use datafusion::parquet::basic::{BrotliLevel, Compression, GzipLevel, ZstdLevel};
 use datafusion::prelude::*;
-use datafusion_ffi::table_provider::FFI_TableProvider;
 use futures::{StreamExt, TryStreamExt};
 use pyo3::exceptions::PyValueError;
 use pyo3::prelude::*;
 use pyo3::pybacked::PyBackedStr;
 use pyo3::types::{PyCapsule, PyList, PyTuple, PyTupleMethods};
 use tokio::task::JoinHandle;
 
-use crate::catalog::PyTable;
 use crate::errors::{py_datafusion_err, to_datafusion_err, PyDataFusionError};
 use crate::expr::sort_expr::to_sort_expressions;
 use crate::physical_plan::PyExecutionPlan;
 use crate::record_batch::PyRecordBatchStream;
 use crate::sql::logical::PyLogicalPlan;
+use crate::table::PyTableProvider;
 use crate::utils::{
     get_tokio_runtime, is_ipython_env, py_obj_to_scalar_value, validate_pycapsule, wait_for_future,
 };
@@ -58,40 +57,6 @@ use crate::{
     expr::{sort_expr::PySortExpr, PyExpr},
 };
 
-// https://github.com/apache/datafusion-python/pull/1016#discussion_r1983239116
-// - we have not decided on the table_provider approach yet
-// this is an interim implementation
-#[pyclass(name = "TableProvider", module = "datafusion")]
-pub struct PyTableProvider {
-    provider: Arc<dyn TableProvider + Send>,
-}
-
-impl PyTableProvider {
-    pub fn new(provider: Arc<dyn TableProvider>) -> Self {
-        Self { provider }
-    }
-
-    pub fn as_table(&self) -> PyTable {
-        let table_provider: Arc<dyn TableProvider> = self.provider.clone();
-        PyTable::new(table_provider)
-    }
-}
-
-#[pymethods]
-impl PyTableProvider {
-    fn __datafusion_table_provider__<'py>(
-        &self,
-        py: Python<'py>,
-    ) -> PyResult<Bound<'py, PyCapsule>> {
-        let name = CString::new("datafusion_table_provider").unwrap();
-
-        let runtime = get_tokio_runtime().0.handle().clone();
-        let provider = FFI_TableProvider::new(Arc::clone(&self.provider), false, Some(runtime));
-
-        PyCapsule::new(py, provider, Some(name.clone()))
-    }
-}
-
 /// Configuration for DataFrame display formatting
 #[derive(Debug, Clone)]
 pub struct FormatterConfig {
@@ -303,6 +268,10 @@ impl PyDataFrame {
         }
     }
 
+    pub(crate) fn into_view_provider(&self) -> Arc<dyn TableProvider> {
+        self.df.as_ref().clone().into_view()
+    }
+
     fn prepare_repr_string(&mut self, py: Python, as_html: bool) -> PyDataFusionResult<String> {
         // Get the Python formatter and config
         let PythonFormatter { formatter, config } = get_python_formatter_with_config(py)?;
@@ -436,14 +405,12 @@ impl PyDataFrame {
     /// https://github.com/apache/datafusion-python/pull/1016#discussion_r1983239116
     /// - we have not decided on the table_provider approach yet
     #[allow(clippy::wrong_self_convention)]
-    fn into_view(&self) -> PyDataFusionResult<PyTable> {
+    fn into_view(&self) -> PyDataFusionResult<PyTableProvider> {
         // Call the underlying Rust DataFrame::into_view method.
         // Note that the Rust method consumes self; here we clone the inner Arc<DataFrame>
         // so that we don’t invalidate this PyDataFrame.
-        let table_provider = self.df.as_ref().clone().into_view();
-        let table_provider = PyTableProvider::new(table_provider);
-
-        Ok(table_provider.as_table())
+        let table_provider = self.into_view_provider();
+        Ok(PyTableProvider::new(table_provider))
     }
 
     #[pyo3(signature = (*args))]

src/lib.rs

@@ -52,6 +52,7 @@ pub mod pyarrow_util;
 mod record_batch;
 pub mod sql;
 pub mod store;
+pub mod table;
 pub mod unparser;
 
 #[cfg(feature = "substrait")]
@@ -97,6 +98,7 @@ fn _internal(py: Python, m: Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<physical_plan::PyExecutionPlan>()?;
     m.add_class::<record_batch::PyRecordBatch>()?;
     m.add_class::<record_batch::PyRecordBatchStream>()?;
+    m.add_class::<table::PyTableProvider>()?;
 
     let catalog = PyModule::new(py, "catalog")?;
     catalog::init_module(&catalog)?;

src/table.rs

@@ -0,0 +1,74 @@
+// 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::ffi::CString;
+use std::sync::Arc;
+
+use datafusion::datasource::TableProvider;
+use datafusion_ffi::table_provider::{FFI_TableProvider, ForeignTableProvider};
+use pyo3::prelude::*;
+use pyo3::types::PyCapsule;
+
+use crate::dataframe::PyDataFrame;
+use crate::errors::{py_datafusion_err, PyDataFusionResult};
+use crate::utils::{get_tokio_runtime, validate_pycapsule};
+
+/// Represents a table provider that can be registered with DataFusion
+#[pyclass(name = "TableProvider", module = "datafusion")]
+pub struct PyTableProvider {
+    pub(crate) provider: Arc<dyn TableProvider + Send>,
+}
+
+impl PyTableProvider {
+    pub(crate) fn new(provider: Arc<dyn TableProvider>) -> Self {
+        Self { provider }
+    }
+}
+
+#[pymethods]
+impl PyTableProvider {
+    /// Create a `TableProvider` from a PyCapsule containing an FFI pointer
+    #[staticmethod]
+    pub fn from_capsule(capsule: Bound<'_, PyAny>) -> PyResult<Self> {
+        let capsule = capsule.downcast::<PyCapsule>().map_err(py_datafusion_err)?;
+        validate_pycapsule(&capsule, "datafusion_table_provider")?;
+
+        let provider = unsafe { capsule.reference::<FFI_TableProvider>() };
+        let provider: ForeignTableProvider = provider.into();
+
+        Ok(Self::new(Arc::new(provider)))
+    }
+
+    /// Create a `TableProvider` from a DataFrame by converting it into a view
+    #[staticmethod]
+    pub fn from_view(df: &PyDataFrame) -> PyDataFusionResult<Self> {
+        let table_provider = df.into_view_provider();
+        Ok(Self::new(table_provider))
+    }
+
+    fn __datafusion_table_provider__<'py>(
+        &self,
+        py: Python<'py>,
+    ) -> PyResult<Bound<'py, PyCapsule>> {
+        let name = CString::new("datafusion_table_provider").unwrap();
+
+        let runtime = get_tokio_runtime().0.handle().clone();
+        let provider = FFI_TableProvider::new(Arc::clone(&self.provider), false, Some(runtime));
+
+        PyCapsule::new(py, provider, Some(name.clone()))
+    }
+}

src/udtf.rs

@@ -18,9 +18,9 @@
 use pyo3::prelude::*;
 use std::sync::Arc;
 
-use crate::dataframe::PyTableProvider;
 use crate::errors::{py_datafusion_err, to_datafusion_err};
 use crate::expr::PyExpr;
+use crate::table::PyTableProvider;
 use crate::utils::validate_pycapsule;
 use datafusion::catalog::{TableFunctionImpl, TableProvider};
 use datafusion::error::Result as DataFusionResult;