feat: enhance Dataset to support any Python object implementing __arrow_c_stream__ interface

Author: kosiew

src/dataset.rs

@@ -17,47 +17,46 @@
 
 use datafusion::catalog::Session;
 use pyo3::exceptions::PyValueError;
-/// Implements a Datafusion TableProvider that delegates to a PyArrow Dataset
-/// This allows us to use PyArrow Datasets as Datafusion tables while pushing down projections and filters
+/// Implements a Datafusion TableProvider that delegates to any Python object
+/// implementing the ``__arrow_c_stream__`` interface. This allows DataFusion to
+/// scan objects from libraries such as PyArrow, nanoarrow, Polars, or DuckDB
+/// without requiring a concrete ``pyarrow.dataset.Dataset`` instance.
 use pyo3::prelude::*;
-use pyo3::types::PyType;
 
 use std::any::Any;
 use std::sync::Arc;
 
 use async_trait::async_trait;
 
+use arrow::array::RecordBatchReader;
+use arrow::ffi_stream::ArrowArrayStreamReader;
+use arrow::pyarrow::FromPyArrow;
 use datafusion::arrow::datatypes::SchemaRef;
-use datafusion::arrow::pyarrow::PyArrowType;
 use datafusion::datasource::{TableProvider, TableType};
 use datafusion::error::{DataFusionError, Result as DFResult};
 use datafusion::logical_expr::Expr;
 use datafusion::logical_expr::TableProviderFilterPushDown;
 use datafusion::physical_plan::ExecutionPlan;
 
 use crate::dataset_exec::DatasetExec;
-use crate::pyarrow_filter_expression::PyArrowFilterExpression;
 
-// Wraps a pyarrow.dataset.Dataset class and implements a Datafusion TableProvider around it
+// Wraps a Python object implementing ``__arrow_c_stream__`` and exposes it as a
+// DataFusion TableProvider.
 #[derive(Debug)]
 pub(crate) struct Dataset {
     dataset: PyObject,
 }
 
 impl Dataset {
-    // Creates a Python PyArrow.Dataset
-    pub fn new(dataset: &Bound<'_, PyAny>, py: Python) -> PyResult<Self> {
-        // Ensure that we were passed an instance of pyarrow.dataset.Dataset
-        let ds = PyModule::import(py, "pyarrow.dataset")?;
-        let ds_attr = ds.getattr("Dataset")?;
-        let ds_type = ds_attr.downcast::<PyType>()?;
-        if dataset.is_instance(ds_type)? {
+    // Creates a Python Dataset wrapper
+    pub fn new(dataset: &Bound<'_, PyAny>, _py: Python) -> PyResult<Self> {
+        if dataset.hasattr("__arrow_c_stream__")? {
             Ok(Dataset {
                 dataset: dataset.clone().unbind(),
             })
         } else {
             Err(PyValueError::new_err(
-                "dataset argument must be a pyarrow.dataset.Dataset object",
+                "dataset argument must implement __arrow_c_stream__",
             ))
         }
     }
@@ -75,15 +74,9 @@ impl TableProvider for Dataset {
     fn schema(&self) -> SchemaRef {
         Python::with_gil(|py| {
             let dataset = self.dataset.bind(py);
-            // This can panic but since we checked that self.dataset is a pyarrow.dataset.Dataset it should never
-            Arc::new(
-                dataset
-                    .getattr("schema")
-                    .unwrap()
-                    .extract::<PyArrowType<_>>()
-                    .unwrap()
-                    .0,
-            )
+            let reader = ArrowArrayStreamReader::from_pyarrow_bound(dataset)
+                .expect("dataset must implement __arrow_c_stream__");
+            Arc::new(reader.schema().as_ref().clone())
         })
     }
 
@@ -100,7 +93,7 @@ impl TableProvider for Dataset {
         &self,
         _ctx: &dyn Session,
         projection: Option<&Vec<usize>>,
-        filters: &[Expr],
+        _filters: &[Expr],
         // limit can be used to reduce the amount scanned
         // from the datasource as a performance optimization.
         // If set, it contains the amount of rows needed by the `LogicalPlan`,
@@ -109,7 +102,7 @@ impl TableProvider for Dataset {
     ) -> DFResult<Arc<dyn ExecutionPlan>> {
         Python::with_gil(|py| {
             let plan: Arc<dyn ExecutionPlan> = Arc::new(
-                DatasetExec::new(py, self.dataset.bind(py), projection.cloned(), filters)
+                DatasetExec::new(self.dataset.bind(py), projection.cloned())
                     .map_err(|err| DataFusionError::External(Box::new(err)))?,
             );
             Ok(plan)
@@ -124,10 +117,7 @@ impl TableProvider for Dataset {
     ) -> DFResult<Vec<TableProviderFilterPushDown>> {
         filter
             .iter()
-            .map(|&f| match PyArrowFilterExpression::try_from(f) {
-                Ok(_) => Ok(TableProviderFilterPushDown::Exact),
-                _ => Ok(TableProviderFilterPushDown::Unsupported),
-            })
+            .map(|_| Ok(TableProviderFilterPushDown::Unsupported))
             .collect()
     }
 }