datafusion_table_provider.rs

// 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::proto::logical_extension_codec::FFI_LogicalExtensionCodec;
use datafusion_ffi::table_provider::FFI_TableProvider;
use iceberg::TableIdent;
use iceberg::io::{FileIOBuilder, StorageFactory};
use iceberg::table::StaticTable;
use iceberg_datafusion::table::IcebergStaticTableProvider;
use iceberg_storage_opendal::OpenDalStorageFactory;
use pyo3::exceptions::{PyRuntimeError, PyValueError};
use pyo3::prelude::{PyAnyMethods, PyCapsuleMethods, *};
use pyo3::types::{PyAny, PyCapsule};

use crate::runtime::runtime;

/// Parse the scheme from a URL and return the appropriate StorageFactory.
fn storage_factory_from_path(path: &str) -> PyResult<Arc<dyn StorageFactory>> {
    let scheme = path
        .split("://")
        .next()
        .ok_or_else(|| PyRuntimeError::new_err(format!("Invalid path, missing scheme: {path}")))?;

    let factory: Arc<dyn StorageFactory> = match scheme {
        "file" | "" => Arc::new(OpenDalStorageFactory::Fs),
        "s3" | "s3a" => Arc::new(OpenDalStorageFactory::S3 {
            configured_scheme: scheme.to_string(),
            customized_credential_load: None,
        }),
        "memory" => Arc::new(OpenDalStorageFactory::Memory),
        _ => {
            return Err(PyRuntimeError::new_err(format!(
                "Unsupported storage scheme: {scheme}"
            )));
        }
    };

    Ok(factory)
}

pub(crate) fn validate_pycapsule(capsule: &Bound<PyCapsule>, name: &str) -> PyResult<()> {
    let capsule_name = capsule.name()?;
    if capsule_name.is_none() {
        return Err(PyValueError::new_err(format!(
            "Expected {name} PyCapsule to have name set."
        )));
    }

    let capsule_name = capsule_name.unwrap().to_str()?;
    if capsule_name != name {
        return Err(PyValueError::new_err(format!(
            "Expected name '{name}' in PyCapsule, instead got '{capsule_name}'"
        )));
    }

    Ok(())
}

pub(crate) fn ffi_logical_codec_from_pycapsule(
    obj: Bound<PyAny>,
) -> PyResult<FFI_LogicalExtensionCodec> {
    let attr_name = "__datafusion_logical_extension_codec__";
    let capsule = if obj.hasattr(attr_name)? {
        obj.getattr(attr_name)?.call0()?
    } else {
        obj
    };

    let capsule = capsule.downcast::<PyCapsule>()?;
    validate_pycapsule(capsule, "datafusion_logical_extension_codec")?;

    let codec = unsafe { capsule.reference::<FFI_LogicalExtensionCodec>() };
    Ok(codec.clone())
}

#[pyclass(name = "IcebergDataFusionTable")]
pub struct PyIcebergDataFusionTable {
    inner: Arc<IcebergStaticTableProvider>,
}

#[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 factory = storage_factory_from_path(&metadata_location)?;

            let mut builder = FileIOBuilder::new(factory);

            if let Some(props) = file_io_properties {
                builder = builder.with_props(props);
            }

            let file_io = builder
                .build(None)
                .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();

            IcebergStaticTableProvider::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>,
        session: Bound<'py, PyAny>,
    ) -> PyResult<Bound<'py, PyCapsule>> {
        let capsule_name = CString::new("datafusion_table_provider").unwrap();

        let logical_codec = ffi_logical_codec_from_pycapsule(session)?;

        let ffi_provider = FFI_TableProvider::new_with_ffi_codec(
            self.inner.clone(),
            false,
            Some(runtime()),
            logical_codec,
        );

        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(())
}