Merge branch 'main' of github.com:apache/iceberg-rust into fd-avro-pyiceberg

Author: Fokko

bindings/python/Cargo.lock

@@ -35,3 +35,6 @@ arrow = { version = "55", features = ["pyarrow", "chrono-tz"] }
 iceberg = { path = "../../crates/iceberg" }
 pyo3 = { version = "0.24", features = ["extension-module", "abi3-py39"] }
 serde_json = "1.0.138"
+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,12 +17,15 @@
 
 use pyo3::prelude::*;
 
+mod datafusion_table_provider;
 mod error;
+mod runtime;
 mod transform;
 mod manifest;
 
 #[pymodule]
 fn pyiceberg_core_rust(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
+    datafusion_table_provider::register_module(py, m)?;
     transform::register_module(py, m)?;
     manifest::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

@@ -172,7 +172,7 @@ impl Catalog for GlueCatalog {
     ///
     /// - Errors from `validate_namespace` if the namespace identifier does not
     /// meet validation criteria.
-    /// - Errors from `convert_to_database` if the properties cannot be  
+    /// - Errors from `convert_to_database` if the properties cannot be
     /// successfully converted into a database configuration.
     /// - Errors from the underlying database creation process, converted using
     /// `from_sdk_error`.
@@ -259,7 +259,7 @@ impl Catalog for GlueCatalog {
     /// Asynchronously updates properties of an existing namespace.
     ///
     /// Converts the given namespace identifier and properties into a database
-    /// representation and then attempts to update the corresponding namespace  
+    /// representation and then attempts to update the corresponding namespace
     /// in the Glue Catalog.
     ///
     /// # Returns
@@ -295,7 +295,7 @@ impl Catalog for GlueCatalog {
     /// # Returns
     /// A `Result<()>` indicating the outcome:
     /// - `Ok(())` signifies successful namespace deletion.
-    /// - `Err(...)` signifies failure to drop the namespace due to validation  
+    /// - `Err(...)` signifies failure to drop the namespace due to validation
     /// errors, connectivity issues, or Glue Catalog constraints.
     async fn drop_namespace(&self, namespace: &NamespaceIdent) -> Result<()> {
         let db_name = validate_namespace(namespace)?;
@@ -322,7 +322,7 @@ impl Catalog for GlueCatalog {
     /// A `Result<Vec<TableIdent>>`, which is:
     /// - `Ok(vec![...])` containing a vector of `TableIdent` instances, each
     /// representing a table within the specified namespace.
-    /// - `Err(...)` if an error occurs during namespace validation or while  
+    /// - `Err(...)` if an error occurs during namespace validation or while
     /// querying the database.
     async fn list_tables(&self, namespace: &NamespaceIdent) -> Result<Vec<TableIdent>> {
         let db_name = validate_namespace(namespace)?;
@@ -375,7 +375,7 @@ impl Catalog for GlueCatalog {
     async fn create_table(
         &self,
         namespace: &NamespaceIdent,
-        creation: TableCreation,
+        mut creation: TableCreation,
     ) -> Result<Table> {
         let db_name = validate_namespace(namespace)?;
         let table_name = creation.name.clone();
@@ -384,10 +384,12 @@ impl Catalog for GlueCatalog {
             Some(location) => location.clone(),
             None => {
                 let ns = self.get_namespace(namespace).await?;
-                get_default_table_location(&ns, &db_name, &table_name, &self.config.warehouse)
+                let location =
+                    get_default_table_location(&ns, &db_name, &table_name, &self.config.warehouse);
+                creation.location = Some(location.clone());
+                location
             }
         };
-
         let metadata = TableMetadataBuilder::from_table_creation(creation)?
             .build()?
             .metadata;