Exposing FFI to python

Exposing FFI to python

Workin progress on python catalog

Flushing out schema and catalog providers

Adding implementation of python based catalog and schema providers

Small updates after rebase

Author: renato2099

Cargo.lock

@@ -37,6 +37,7 @@ substrait = ["dep:datafusion-substrait"]
 tokio = { version = "1.45", features = ["macros", "rt", "rt-multi-thread", "sync"] }
 pyo3 = { version = "0.24", features = ["extension-module", "abi3", "abi3-py39"] }
 pyo3-async-runtimes = { version = "0.24", features = ["tokio-runtime"]}
+pyo3-log = "0.12.4"
 arrow = { version = "55.1.0", features = ["pyarrow"] }
 datafusion = { version = "48.0.0", features = ["avro", "unicode_expressions"] }
 datafusion-substrait = { version = "48.0.0", optional = true }
@@ -49,6 +50,7 @@ async-trait = "0.1.88"
 futures = "0.3"
 object_store = { version = "0.12.1", features = ["aws", "gcp", "azure", "http"] }
 url = "2"
+log = "0.4.27"
 
 [build-dependencies]
 prost-types = "0.13.1" # keep in line with `datafusion-substrait`

Cargo.toml

@@ -27,6 +27,7 @@ pyo3 = { version = "0.23", features = ["extension-module", "abi3", "abi3-py39"]
 arrow = { version = "55.0.0" }
 arrow-array = { version = "55.0.0" }
 arrow-schema = { version = "55.0.0" }
+async-trait = "0.1.88"
 
 [build-dependencies]
 pyo3-build-config = "0.23"

examples/datafusion-ffi-example/Cargo.lock

@@ -0,0 +1,60 @@
+# 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.
+
+from __future__ import annotations
+
+import pyarrow as pa
+from datafusion import SessionContext
+from datafusion_ffi_example import MyCatalogProvider
+
+
+def test_catalog_provider():
+    ctx = SessionContext()
+
+    my_catalog_name = "my_catalog"
+    expected_schema_name = "my_schema"
+    expected_table_name = "my_table"
+    expected_table_columns = ["units", "price"]
+
+    catalog_provider = MyCatalogProvider()
+    ctx.register_catalog_provider(my_catalog_name, catalog_provider)
+    my_catalog = ctx.catalog(my_catalog_name)
+
+    my_catalog_schemas = my_catalog.names()
+    assert expected_schema_name in my_catalog_schemas
+    my_database = my_catalog.database(expected_schema_name)
+    assert expected_table_name in my_database.names()
+    my_table = my_database.table(expected_table_name)
+    assert expected_table_columns == my_table.schema.names
+
+    result = ctx.table(
+        f"{my_catalog_name}.{expected_schema_name}.{expected_table_name}"
+    ).collect()
+    assert len(result) == 2
+
+    col0_result = [r.column(0) for r in result]
+    col1_result = [r.column(1) for r in result]
+    expected_col0 = [
+        pa.array([10, 20, 30], type=pa.int32()),
+        pa.array([5, 7], type=pa.int32()),
+    ]
+    expected_col1 = [
+        pa.array([1, 2, 5], type=pa.float64()),
+        pa.array([1.5, 2.5], type=pa.float64()),
+    ]
+    assert col0_result == expected_col0
+    assert col1_result == expected_col1

examples/datafusion-ffi-example/Cargo.toml

@@ -0,0 +1,179 @@
+// 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 pyo3::{pyclass, pymethods, Bound, PyResult, Python};
+use std::{any::Any, fmt::Debug, sync::Arc};
+
+use arrow::datatypes::Schema;
+use async_trait::async_trait;
+use datafusion::{
+    catalog::{
+        CatalogProvider, MemoryCatalogProvider, MemorySchemaProvider, SchemaProvider, TableProvider,
+    },
+    common::exec_err,
+    datasource::MemTable,
+    error::{DataFusionError, Result},
+};
+use datafusion_ffi::catalog_provider::FFI_CatalogProvider;
+use pyo3::types::PyCapsule;
+
+pub fn my_table() -> Arc<dyn TableProvider + 'static> {
+    use arrow::datatypes::{DataType, Field};
+    use datafusion::common::record_batch;
+
+    let schema = Arc::new(Schema::new(vec![
+        Field::new("units", DataType::Int32, true),
+        Field::new("price", DataType::Float64, true),
+    ]));
+
+    let partitions = vec![
+        record_batch!(
+            ("units", Int32, vec![10, 20, 30]),
+            ("price", Float64, vec![1.0, 2.0, 5.0])
+        )
+        .unwrap(),
+        record_batch!(
+            ("units", Int32, vec![5, 7]),
+            ("price", Float64, vec![1.5, 2.5])
+        )
+        .unwrap(),
+    ];
+
+    Arc::new(MemTable::try_new(schema, vec![partitions]).unwrap())
+}
+
+#[derive(Debug)]
+pub struct FixedSchemaProvider {
+    inner: MemorySchemaProvider,
+}
+
+impl Default for FixedSchemaProvider {
+    fn default() -> Self {
+        let inner = MemorySchemaProvider::new();
+
+        let table = my_table();
+
+        let _ = inner.register_table("my_table".to_string(), table).unwrap();
+
+        Self { inner }
+    }
+}
+
+#[async_trait]
+impl SchemaProvider for FixedSchemaProvider {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn table_names(&self) -> Vec<String> {
+        self.inner.table_names()
+    }
+
+    async fn table(&self, name: &str) -> Result<Option<Arc<dyn TableProvider>>, DataFusionError> {
+        self.inner.table(name).await
+    }
+
+    fn register_table(
+        &self,
+        name: String,
+        table: Arc<dyn TableProvider>,
+    ) -> Result<Option<Arc<dyn TableProvider>>> {
+        self.inner.register_table(name, table)
+    }
+
+    fn deregister_table(&self, name: &str) -> Result<Option<Arc<dyn TableProvider>>> {
+        self.inner.deregister_table(name)
+    }
+
+    fn table_exist(&self, name: &str) -> bool {
+        self.inner.table_exist(name)
+    }
+}
+
+/// This catalog provider is intended only for unit tests. It prepopulates with one
+/// schema and only allows for schemas named after four types of fruit.
+#[pyclass(
+    name = "MyCatalogProvider",
+    module = "datafusion_ffi_example",
+    subclass
+)]
+#[derive(Debug)]
+pub(crate) struct MyCatalogProvider {
+    inner: MemoryCatalogProvider,
+}
+
+impl Default for MyCatalogProvider {
+    fn default() -> Self {
+        let inner = MemoryCatalogProvider::new();
+
+        let schema_name: &str = "my_schema";
+        let _ = inner.register_schema(schema_name, Arc::new(FixedSchemaProvider::default()));
+
+        Self { inner }
+    }
+}
+
+impl CatalogProvider for MyCatalogProvider {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn schema_names(&self) -> Vec<String> {
+        self.inner.schema_names()
+    }
+
+    fn schema(&self, name: &str) -> Option<Arc<dyn SchemaProvider>> {
+        self.inner.schema(name)
+    }
+
+    fn register_schema(
+        &self,
+        name: &str,
+        schema: Arc<dyn SchemaProvider>,
+    ) -> Result<Option<Arc<dyn SchemaProvider>>> {
+        self.inner.register_schema(name, schema)
+    }
+
+    fn deregister_schema(
+        &self,
+        name: &str,
+        cascade: bool,
+    ) -> Result<Option<Arc<dyn SchemaProvider>>> {
+        self.inner.deregister_schema(name, cascade)
+    }
+}
+
+#[pymethods]
+impl MyCatalogProvider {
+    #[new]
+    pub fn new() -> Self {
+        Self {
+            inner: Default::default(),
+        }
+    }
+
+    pub fn __datafusion_catalog_provider__<'py>(
+        &self,
+        py: Python<'py>,
+    ) -> PyResult<Bound<'py, PyCapsule>> {
+        let name = cr"datafusion_catalog_provider".into();
+        let catalog_provider =
+            FFI_CatalogProvider::new(Arc::new(MyCatalogProvider::default()), None);
+
+        PyCapsule::new(py, catalog_provider, Some(name))
+    }
+}

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

@@ -15,16 +15,19 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use crate::catalog_provider::MyCatalogProvider;
 use crate::table_function::MyTableFunction;
 use crate::table_provider::MyTableProvider;
 use pyo3::prelude::*;
 
+pub(crate) mod catalog_provider;
 pub(crate) mod table_function;
 pub(crate) mod table_provider;
 
 #[pymodule]
 fn datafusion_ffi_example(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<MyTableProvider>()?;
     m.add_class::<MyTableFunction>()?;
+    m.add_class::<MyCatalogProvider>()?;
     Ok(())
 }

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

@@ -32,7 +32,7 @@
 
 from datafusion.col import col, column
 
-from . import functions, object_store, substrait, unparser
+from . import catalog, functions, object_store, substrait, unparser
 
 # The following imports are okay to remain as opaque to the user.
 from ._internal import Config
@@ -91,6 +91,7 @@
     "TableFunction",
     "WindowFrame",
     "WindowUDF",
+    "catalog",
     "col",
     "column",
     "common",