liurenjie1024 commented on code in PR #600: URL: https://github.com/apache/iceberg-rust/pull/600#discussion_r1797634121
########## crates/integrations/datafusion/src/table/table_provider_factory.rs: ########## @@ -0,0 +1,301 @@ +// 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::borrow::Cow; +use std::collections::HashMap; +use std::sync::Arc; + +use async_trait::async_trait; +use datafusion::catalog::{Session, TableProvider, TableProviderFactory}; +use datafusion::error::Result as DFResult; +use datafusion::logical_expr::CreateExternalTable; +use datafusion::sql::TableReference; +use iceberg::arrow::schema_to_arrow_schema; +use iceberg::io::FileIO; +use iceberg::table::StaticTable; +use iceberg::{Error, ErrorKind, Result, TableIdent}; + +use super::IcebergTableProvider; +use crate::to_datafusion_error; + +/// A factory that implements DataFusion's `TableProviderFactory` to create `IcebergTableProvider` instances. +/// +/// # Example +/// +/// The following example demonstrates how to create an Iceberg external table using SQL in +/// a DataFusion session with `IcebergTableProviderFactory`: +/// +/// ``` +/// use std::sync::Arc; +/// +/// use datafusion::execution::session_state::SessionStateBuilder; +/// use datafusion::prelude::*; +/// use datafusion::sql::TableReference; +/// use iceberg_datafusion::IcebergTableProviderFactory; +/// +/// #[tokio::main] +/// async fn main() { +/// // Create a new session context +/// let mut state = SessionStateBuilder::new().with_default_features().build(); +/// +/// // Register the IcebergTableProviderFactory in the session +/// state.table_factories_mut().insert( +/// "ICEBERG".to_string(), +/// Arc::new(IcebergTableProviderFactory::new()), +/// ); +/// +/// let ctx = SessionContext::new_with_state(state); +/// +/// // Define the table reference and the location of the Iceberg metadata file +/// let table_ref = TableReference::bare("my_iceberg_table"); +/// // /path/to/iceberg/metadata +/// let metadata_file_path = format!( +/// "{}/testdata/table_metadata/{}", +/// env!("CARGO_MANIFEST_DIR"), +/// "TableMetadataV2.json" +/// ); +/// +/// // SQL command to create the Iceberg external table +/// let sql = format!( +/// "CREATE EXTERNAL TABLE {} STORED AS ICEBERG LOCATION '{}'", +/// table_ref, metadata_file_path +/// ); +/// +/// // Execute the SQL to create the external table +/// ctx.sql(&sql).await.expect("Failed to create table"); +/// +/// // Verify the table was created by retrieving the table provider +/// let table_provider = ctx +/// .table_provider(table_ref) +/// .await +/// .expect("Table not found"); +/// +/// println!("Iceberg external table created successfully."); +/// } +/// ``` +/// +/// # Note +/// This factory is designed to work with the DataFusion query engine, +/// specifically for handling Iceberg tables in external table commands. +/// Currently, this implementation supports only reading Iceberg tables, with +/// the creation of new tables not yet available. +/// +/// # Errors +/// An error will be returned if any unsupported feature, such as partition columns, +/// order expressions, constraints, or column defaults, is detected in the table creation command. +#[derive(Default)] +#[non_exhaustive] Review Comment: Sorry, I don't quite understand why we need to add this annotation here? ########## crates/integrations/datafusion/src/table/table_provider_factory.rs: ########## @@ -0,0 +1,312 @@ +// 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::borrow::Cow; +use std::collections::HashMap; +use std::sync::Arc; + +use async_trait::async_trait; +use datafusion::catalog::{Session, TableProvider, TableProviderFactory}; +use datafusion::error::Result as DFResult; +use datafusion::logical_expr::CreateExternalTable; +use datafusion::sql::TableReference; +use iceberg::arrow::schema_to_arrow_schema; +use iceberg::io::FileIO; +use iceberg::table::StaticTable; +use iceberg::{Error, ErrorKind, Result, TableIdent}; + +use super::IcebergTableProvider; +use crate::to_datafusion_error; + +/// A factory that implements DataFusion's `TableProviderFactory` to create `IcebergTableProvider` instances. +/// +/// # Example +/// +/// The following example demonstrates how to create an Iceberg external table using SQL in +/// a DataFusion session with `IcebergTableProviderFactory`: +/// +/// ``` +/// use std::sync::Arc; +/// +/// use datafusion::prelude::*; +/// use iceberg_datafusion::IcebergTableProviderFactory; +/// +/// #[tokio::main] +/// async fn main() { +/// // Create a new session context +/// let mut state = SessionStateBuilder::new().with_default_features().build(); +/// +/// // Register the IcebergTableProviderFactory in the session +/// state.table_factories_mut().insert( +/// "ICEBERG".to_string(), +/// Arc::new(IcebergTableProviderFactory::new()), +/// ); +/// +/// let ctx = SessionContext::new_with_state(state); +/// +/// // Define the table reference and the location of the Iceberg metadata file +/// let table_ref = TableReference::bare("my_iceberg_table"); +/// let metadata_file_path = "/path/to/iceberg/metadata"; +/// +/// // SQL command to create the Iceberg external table +/// let sql = format!( +/// "CREATE EXTERNAL TABLE {} STORED AS ICEBERG LOCATION '{}'", +/// table_ref, metadata_file_path +/// ); +/// +/// // Execute the SQL to create the external table +/// ctx.sql(&sql).await.expect("Failed to create table"); +/// +/// // Verify the table was created by retrieving the table provider +/// let table_provider = ctx +/// .table_provider(table_ref) +/// .await +/// .expect("Table not found"); +/// +/// println!("Iceberg external table created successfully."); +/// } +/// ``` +/// +/// # Note +/// This factory is designed to work with the DataFusion query engine, +/// specifically for handling Iceberg tables in external table commands. +/// Currently, this implementation supports only reading Iceberg tables, with +/// the creation of new tables not yet available. +/// +/// # Errors +/// An error will be returned if any unsupported feature, such as partition columns, +/// order expressions, constraints, or column defaults, is detected in the table creation command. +#[derive(Default)] +#[non_exhaustive] +pub struct IcebergTableProviderFactory {} + +impl IcebergTableProviderFactory { + pub fn new() -> Self { + Self {} + } +} + +#[async_trait] +impl TableProviderFactory for IcebergTableProviderFactory { + async fn create( + &self, + _state: &dyn Session, + cmd: &CreateExternalTable, + ) -> DFResult<Arc<dyn TableProvider>> { + check_cmd(cmd).map_err(to_datafusion_error)?; + + let table_name = &cmd.name; + let metadata_file_path = &cmd.location; + let options = &cmd.options; + + let table_name_with_ns = complement_namespace_if_necessary(table_name); + + let table = create_static_table(table_name_with_ns, metadata_file_path, options) + .await + .map_err(to_datafusion_error)? + .into_table(); + + let schema = schema_to_arrow_schema(table.metadata().current_schema()) + .map_err(to_datafusion_error)?; + + Ok(Arc::new(IcebergTableProvider::new(table, Arc::new(schema)))) + } +} + +fn check_cmd(cmd: &CreateExternalTable) -> Result<()> { + println!("Checking command: {:?}", cmd); + + let CreateExternalTable { + schema, + table_partition_cols, + order_exprs, + constraints, + column_defaults, + .. + } = cmd; + + // Check if any of the fields violate the constraints in a single condition + let is_invalid = !schema.fields().is_empty() + || !table_partition_cols.is_empty() + || !order_exprs.is_empty() + || !constraints.is_empty() + || !column_defaults.is_empty(); + + if is_invalid { + return Err(Error::new(ErrorKind::FeatureUnsupported, "Currently we only support reading existing icebergs tables in external table command. To create new table, please use catalog provider.")); + } + + Ok(()) +} + +/// Complements the namespace of a table name if necessary. +/// +/// # Example +/// +/// ``` +/// use datafusion::sql::TableReference; +/// use iceberg_datafusion::complement_namespace_if_necessary; +/// +/// let table_name = TableReference::bare("table_name"); +/// let table_name_with_ns = complement_namespace_if_necessary(&table_name); +/// +/// assert_eq!( +/// table_name_with_ns, +/// TableReference::partial("default", "table_name") +/// ); +/// ``` +/// +/// # Note +/// If the table name is a bare name, it will be complemented with the 'default' namespace. +/// Otherwise, it will be returned as is. Because Iceberg tables are always namespaced, but DataFusion +/// external table commands not include the namespace, this function ensures that the namespace is always present. +/// +/// # See also +/// - [`iceberg::NamespaceIdent`] +/// - [`datafusion::sql::planner::SqlToRel::external_table_to_plan`] +fn complement_namespace_if_necessary(table_name: &TableReference) -> Cow<'_, TableReference> { + match table_name { + TableReference::Bare { table } => { + Cow::Owned(TableReference::partial("default", table.as_ref())) + } + other => Cow::Borrowed(other), + } +} Review Comment: > Does anyone have a good approach for handling compatibility? I'm not quite familiar with datafusion, but typicall it should be registered with current namespace? For example, a user has executed `use ns1` before registering this table, then ideally this table should be register under `ns1`? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: issues-unsubscr...@iceberg.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org --------------------------------------------------------------------- To unsubscribe, e-mail: issues-unsubscr...@iceberg.apache.org For additional commands, e-mail: issues-h...@iceberg.apache.org
