Expose datafusion table provider as python binding (#1324)

## Which issue does this PR close?

<!--
We generally require a GitHub issue to be filed for all bug fixes and
enhancements and this helps us generate change logs for our releases.
You can link an issue to this PR using the GitHub syntax. For example
`Closes #123` indicates that this PR will close issue #123.
-->

- Closes #865

## What changes are included in this PR?
This PR creates a new `IcebergDataFusionTable` python class and exposes
it through the new `pyiceberg_core.datafusion` module.
```
from pyiceberg_core.datafusion import IcebergDataFusionTable
```

The goal of exposing `IcebergDataFusionTable` is to be able to register
the Iceberg table provider to datafusion-python, using the
`register_table_provider` API.
See the usage example in
`bindings/python/tests/test_datafusion_table_provider.py`

The integration relies on the `FFI_TableProvider` API as described in
https://datafusion.apache.org/python/user-guide/io/table_provider.html

Note that this integration only works for `datafusion >= 45` due to this
issue apache/datafusion#13851


<!--
Provide a summary of the modifications in this PR. List the main changes
such as new features, bug fixes, refactoring, or any other updates.
-->

## Are these changes tested?

<!--
Specify what test covers (unit test, integration test, etc.).

If tests are not included in your PR, please explain why (for example,
are they covered by existing tests)?
-->
Yes, unit tests. 

To build and test locally:
```
cd bindings/python
hatch run dev:develop
hatch run dev:test
```

Author: kevinjqliu

bindings/python/Cargo.lock

@@ -34,3 +34,6 @@ crate-type = ["cdylib"]
 arrow = { version = "55", features = ["pyarrow", "chrono-tz"] }
 iceberg = { path = "../../crates/iceberg" }
 pyo3 = { version = "0.24", features = ["extension-module", "abi3-py39"] }
+iceberg-datafusion = { path = "../../crates/integrations/datafusion" }
+datafusion-ffi = { version = "47" }
+tokio = { version = "1.44", default-features = false }

bindings/python/Cargo.toml

@@ -44,7 +44,7 @@ python-source = "python"
 ignore = ["F403", "F405"]
 
 [tool.hatch.envs.dev]
-dependencies = ["maturin>=1.0,<2.0", "pytest>=8.3.2", "pyarrow>=17.0.0"]
+dependencies = ["maturin>=1.0,<2.0", "pytest>=8.3.2", "pyarrow>=17.0.0", "datafusion>=45", "pyiceberg[sql-sqlite]>=0.9.1"]
 
 [tool.hatch.envs.dev.scripts]
 build = "maturin build --out dist --sdist"

bindings/python/pyproject.toml

@@ -0,0 +1,107 @@
+// 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::collections::HashMap;
+use std::ffi::CString;
+use std::sync::Arc;
+
+use datafusion_ffi::table_provider::FFI_TableProvider;
+use iceberg::io::FileIO;
+use iceberg::table::StaticTable;
+use iceberg::TableIdent;
+use iceberg_datafusion::table::IcebergTableProvider;
+use pyo3::exceptions::PyRuntimeError;
+use pyo3::prelude::*;
+use pyo3::types::PyCapsule;
+
+use crate::runtime::runtime;
+
+#[pyclass(name = "IcebergDataFusionTable")]
+pub struct PyIcebergDataFusionTable {
+    inner: Arc<IcebergTableProvider>,
+}
+
+#[pymethods]
+impl PyIcebergDataFusionTable {
+    #[new]
+    fn new(
+        identifier: Vec<String>,
+        metadata_location: String,
+        file_io_properties: Option<HashMap<String, String>>,
+    ) -> PyResult<Self> {
+        let runtime = runtime();
+
+        let provider = runtime.block_on(async {
+            let table_ident = TableIdent::from_strs(identifier)
+                .map_err(|e| PyRuntimeError::new_err(format!("Invalid table identifier: {e}")))?;
+
+            let mut builder = FileIO::from_path(&metadata_location)
+                .map_err(|e| PyRuntimeError::new_err(format!("Failed to init FileIO: {e}")))?;
+
+            if let Some(props) = file_io_properties {
+                builder = builder.with_props(props);
+            }
+
+            let file_io = builder
+                .build()
+                .map_err(|e| PyRuntimeError::new_err(format!("Failed to build FileIO: {e}")))?;
+
+            let static_table =
+                StaticTable::from_metadata_file(&metadata_location, table_ident, file_io)
+                    .await
+                    .map_err(|e| {
+                        PyRuntimeError::new_err(format!("Failed to load static table: {e}"))
+                    })?;
+
+            let table = static_table.into_table();
+
+            IcebergTableProvider::try_new_from_table(table)
+                .await
+                .map_err(|e| {
+                    PyRuntimeError::new_err(format!("Failed to create table provider: {e}"))
+                })
+        })?;
+
+        Ok(Self {
+            inner: Arc::new(provider),
+        })
+    }
+
+    fn __datafusion_table_provider__<'py>(
+        &self,
+        py: Python<'py>,
+    ) -> PyResult<Bound<'py, PyCapsule>> {
+        let capsule_name = CString::new("datafusion_table_provider").unwrap();
+
+        let ffi_provider = FFI_TableProvider::new(self.inner.clone(), false, Some(runtime()));
+
+        PyCapsule::new(py, ffi_provider, Some(capsule_name))
+    }
+}
+
+pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
+    let this = PyModule::new(py, "datafusion")?;
+
+    this.add_class::<PyIcebergDataFusionTable>()?;
+
+    m.add_submodule(&this)?;
+    py.import("sys")?
+        .getattr("modules")?
+        .set_item("pyiceberg_core.datafusion", this)?;
+
+    Ok(())
+}

bindings/python/src/datafusion_table_provider.rs

@@ -17,11 +17,14 @@
 
 use pyo3::prelude::*;
 
+mod datafusion_table_provider;
 mod error;
+mod runtime;
 mod transform;
 
 #[pymodule]
 fn pyiceberg_core_rust(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
+    datafusion_table_provider::register_module(py, m)?;
     transform::register_module(py, m)?;
     Ok(())
 }

bindings/python/src/lib.rs

@@ -0,0 +1,32 @@
+// 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::OnceLock;
+
+use tokio::runtime::{Handle, Runtime};
+
+static RUNTIME: OnceLock<Runtime> = OnceLock::new();
+
+pub fn runtime() -> Handle {
+    match Handle::try_current() {
+        Ok(h) => h.clone(),
+        _ => {
+            let rt = RUNTIME.get_or_init(|| Runtime::new().unwrap());
+            rt.handle().clone()
+        }
+    }
+}

bindings/python/src/runtime.rs

@@ -0,0 +1,160 @@
+# 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 datetime import date, datetime
+import uuid
+import pytest
+from pyiceberg_core.datafusion import IcebergDataFusionTable
+from datafusion import SessionContext
+from pyiceberg.catalog import Catalog, load_catalog
+import pyarrow as pa
+from pathlib import Path
+import datafusion
+
+assert (
+    datafusion.__version__ >= "45"
+)  # iceberg table provider only works for datafusion >= 45
+
+
+@pytest.fixture(scope="session")
+def warehouse(tmp_path_factory: pytest.TempPathFactory) -> Path:
+    return tmp_path_factory.mktemp("warehouse")
+
+
+@pytest.fixture(scope="session")
+def catalog(warehouse: Path) -> Catalog:
+    catalog = load_catalog(
+        "default",
+        **{
+            "uri": f"sqlite:///{warehouse}/pyiceberg_catalog.db",
+            "warehouse": f"file://{warehouse}",
+        },
+    )
+    return catalog
+
+
+@pytest.fixture(scope="session")
+def arrow_table_with_null() -> "pa.Table":
+    """Pyarrow table with all kinds of columns."""
+    import pyarrow as pa
+
+    return pa.Table.from_pydict(
+        {
+            "bool": [False, None, True],
+            "string": ["a", None, "z"],
+            # Go over the 16 bytes to kick in truncation
+            "string_long": ["a" * 22, None, "z" * 22],
+            "int": [1, None, 9],
+            "long": [1, None, 9],
+            "float": [0.0, None, 0.9],
+            "double": [0.0, None, 0.9],
+            # 'time': [1_000_000, None, 3_000_000],  # Example times: 1s, none, and 3s past midnight #Spark does not support time fields
+            "timestamp": [
+                datetime(2023, 1, 1, 19, 25, 00),
+                None,
+                datetime(2023, 3, 1, 19, 25, 00),
+            ],
+            # "timestamptz": [
+            #     datetime(2023, 1, 1, 19, 25, 00, tzinfo=timezone.utc),
+            #     None,
+            #     datetime(2023, 3, 1, 19, 25, 00, tzinfo=timezone.utc),
+            # ],
+            "date": [date(2023, 1, 1), None, date(2023, 3, 1)],
+            # Not supported by Spark
+            # 'time': [time(1, 22, 0), None, time(19, 25, 0)],
+            # Not natively supported by Arrow
+            # 'uuid': [uuid.UUID('00000000-0000-0000-0000-000000000000').bytes, None, uuid.UUID('11111111-1111-1111-1111-111111111111').bytes],
+            "binary": [b"\01", None, b"\22"],
+            "fixed": [
+                uuid.UUID("00000000-0000-0000-0000-000000000000").bytes,
+                None,
+                uuid.UUID("11111111-1111-1111-1111-111111111111").bytes,
+            ],
+        },
+    )
+
+
+def test_register_iceberg_table_provider(
+    catalog: Catalog, arrow_table_with_null: pa.Table
+) -> None:
+    catalog.create_namespace_if_not_exists("default")
+    iceberg_table = catalog.create_table_if_not_exists(
+        "default.dataset",
+        schema=arrow_table_with_null.schema,
+    )
+    iceberg_table.append(arrow_table_with_null)
+
+    iceberg_table_provider = IcebergDataFusionTable(
+        identifier=iceberg_table.name(),
+        metadata_location=iceberg_table.metadata_location,
+        file_io_properties=iceberg_table.io.properties,
+    )
+
+    ctx = SessionContext()
+    ctx.register_table_provider("test", iceberg_table_provider)
+
+    datafusion_table = ctx.table("test")
+    assert datafusion_table is not None
+
+    # check that the schema is the same
+    from pyiceberg.io.pyarrow import _pyarrow_schema_ensure_small_types
+
+    assert _pyarrow_schema_ensure_small_types(
+        datafusion_table.schema()
+    ) == _pyarrow_schema_ensure_small_types(iceberg_table.schema().as_arrow())
+    # large/small type mismatches, fixed in pyiceberg 0.10.0
+    # assert datafusion_table.schema() == iceberg_table.schema().as_arrow()
+
+    # check that the data is the same
+    assert (
+        datafusion_table.to_arrow_table().to_pylist()
+        == iceberg_table.scan().to_arrow().to_pylist()
+    )
+    # large/small type mismatches, fixed in pyiceberg 0.10.0
+    # assert datafusion_table.to_arrow_table() == iceberg_table.scan().to_arrow()
+
+
+def test_register_pyiceberg_table(
+    catalog: Catalog, arrow_table_with_null: pa.Table
+) -> None:
+    from types import MethodType
+
+    catalog.create_namespace_if_not_exists("default")
+    iceberg_table = catalog.create_table_if_not_exists(
+        "default.dataset",
+        schema=arrow_table_with_null.schema,
+    )
+    iceberg_table.append(arrow_table_with_null)
+
+    # monkey patch the __datafusion_table_provider__ method to the iceberg table
+    def __datafusion_table_provider__(self):
+        return IcebergDataFusionTable(
+            identifier=self.name(),
+            metadata_location=self.metadata_location,
+            file_io_properties=self.io.properties,
+        ).__datafusion_table_provider__()
+
+    iceberg_table.__datafusion_table_provider__ = MethodType(
+        __datafusion_table_provider__, iceberg_table
+    )
+
+    ctx = SessionContext()
+    ctx.register_table_provider("test", iceberg_table)
+
+    datafusion_table = ctx.table("test")
+    assert datafusion_table is not None

bindings/python/tests/test_datafusion_table_provider.py

@@ -23,6 +23,6 @@ pub use error::*;
 
 mod physical_plan;
 mod schema;
-mod table;
+pub mod table;
 pub use table::table_provider_factory::IcebergTableProviderFactory;
 pub use table::*;