iceberg_datafusion/table/
table_provider_factory.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use std::borrow::Cow;
19use std::collections::HashMap;
20use std::sync::Arc;
21
22use async_trait::async_trait;
23use datafusion::catalog::{Session, TableProvider, TableProviderFactory};
24use datafusion::common::TableReference;
25use datafusion::error::Result as DFResult;
26use datafusion::logical_expr::CreateExternalTable;
27use iceberg::io::{FileIOBuilder, LocalFsStorageFactory, StorageFactory};
28use iceberg::table::StaticTable;
29use iceberg::{Error, ErrorKind, Result, TableIdent};
30
31use super::IcebergStaticTableProvider;
32use crate::to_datafusion_error;
33
34/// A factory that implements DataFusion's `TableProviderFactory` to create `IcebergTableProvider` instances.
35///
36/// # Example
37///
38/// The following example demonstrates how to create an Iceberg external table using SQL in
39/// a DataFusion session with `IcebergTableProviderFactory`:
40///
41/// ```
42/// use std::sync::Arc;
43///
44/// use datafusion::common::TableReference;
45/// use datafusion::execution::session_state::SessionStateBuilder;
46/// use datafusion::prelude::*;
47/// use iceberg_datafusion::IcebergTableProviderFactory;
48///
49/// #[tokio::main]
50/// async fn main() {
51///     // Create a new session context
52///     let mut state = SessionStateBuilder::new().with_default_features().build();
53///
54///     // Register the IcebergTableProviderFactory in the session
55///     state.table_factories_mut().insert(
56///         "ICEBERG".to_string(),
57///         Arc::new(IcebergTableProviderFactory::new()),
58///     );
59///
60///     let ctx = SessionContext::new_with_state(state);
61///
62///     // Define the table reference and the location of the Iceberg metadata file
63///     let table_ref = TableReference::bare("my_iceberg_table");
64///     // /path/to/iceberg/metadata
65///     let metadata_file_path = format!(
66///         "{}/testdata/table_metadata/{}",
67///         env!("CARGO_MANIFEST_DIR"),
68///         "TableMetadataV2.json"
69///     );
70///
71///     // SQL command to create the Iceberg external table
72///     let sql = format!(
73///         "CREATE EXTERNAL TABLE {} STORED AS ICEBERG LOCATION '{}'",
74///         table_ref, metadata_file_path
75///     );
76///
77///     // Execute the SQL to create the external table
78///     ctx.sql(&sql).await.expect("Failed to create table");
79///
80///     // Verify the table was created by retrieving the table provider
81///     let table_provider = ctx
82///         .table_provider(table_ref)
83///         .await
84///         .expect("Table not found");
85///
86///     println!("Iceberg external table created successfully.");
87/// }
88/// ```
89///
90/// # Note
91/// This factory is designed to work with the DataFusion query engine,
92/// specifically for handling Iceberg tables in external table commands.
93/// Currently, this implementation supports only reading Iceberg tables, with
94/// the creation of new tables not yet available.
95///
96/// # Errors
97/// An error will be returned if any unsupported feature, such as partition columns,
98/// order expressions, constraints, or column defaults, is detected in the table creation command.
99#[derive(Debug, Default)]
100pub struct IcebergTableProviderFactory {
101    storage_factory: Option<Arc<dyn StorageFactory>>,
102}
103
104impl IcebergTableProviderFactory {
105    pub fn new() -> Self {
106        Self {
107            storage_factory: None,
108        }
109    }
110
111    /// Create a new factory with a custom storage factory for creating FileIO instances.
112    pub fn new_with_storage_factory(storage_factory: Arc<dyn StorageFactory>) -> Self {
113        Self {
114            storage_factory: Some(storage_factory),
115        }
116    }
117}
118
119#[async_trait]
120impl TableProviderFactory for IcebergTableProviderFactory {
121    async fn create(
122        &self,
123        _state: &dyn Session,
124        cmd: &CreateExternalTable,
125    ) -> DFResult<Arc<dyn TableProvider>> {
126        let metadata_file_path = check_cmd(cmd).map_err(to_datafusion_error)?;
127
128        let table_name = &cmd.name;
129        let options = &cmd.options;
130
131        let table_name_with_ns = complement_namespace_if_necessary(table_name);
132
133        let storage_factory = self
134            .storage_factory
135            .clone()
136            .unwrap_or_else(|| Arc::new(LocalFsStorageFactory));
137
138        let table = create_static_table(
139            table_name_with_ns,
140            metadata_file_path,
141            options,
142            storage_factory,
143        )
144        .await
145        .map_err(to_datafusion_error)?
146        .into_table();
147
148        let provider = IcebergStaticTableProvider::try_new_from_table(table)
149            .await
150            .map_err(to_datafusion_error)?;
151
152        Ok(Arc::new(provider))
153    }
154}
155
156fn check_cmd(cmd: &CreateExternalTable) -> Result<&str> {
157    let CreateExternalTable {
158        schema,
159        table_partition_cols,
160        order_exprs,
161        constraints,
162        column_defaults,
163        ..
164    } = cmd;
165
166    // Check if any of the fields violate the constraints in a single condition
167    let is_invalid = !schema.fields().is_empty()
168        || !table_partition_cols.is_empty()
169        || !order_exprs.is_empty()
170        || !constraints.is_empty()
171        || !column_defaults.is_empty();
172
173    if is_invalid {
174        return Err(Error::new(
175            ErrorKind::FeatureUnsupported,
176            "Currently we only support reading existing icebergs tables in external table command. To create new table, please use catalog provider.",
177        ));
178    }
179
180    match cmd.locations.as_slice() {
181        [location] => Ok(location),
182        _ => Err(Error::new(
183            ErrorKind::FeatureUnsupported,
184            "Iceberg external tables require exactly one metadata location.",
185        )),
186    }
187}
188
189/// Complements the namespace of a table name if necessary.
190///
191/// # Note
192/// If the table name is a bare name, it will be complemented with the 'default' namespace.
193/// Otherwise, it will be returned as is. Because Iceberg tables are always namespaced, but DataFusion
194/// external table commands maybe not include the namespace, this function ensures that the namespace is always present.
195///
196/// # See also
197/// - [`iceberg::NamespaceIdent`]
198/// - [`datafusion::sql::planner::SqlToRel::external_table_to_plan`]
199fn complement_namespace_if_necessary(table_name: &TableReference) -> Cow<'_, TableReference> {
200    match table_name {
201        TableReference::Bare { table } => {
202            Cow::Owned(TableReference::partial("default", table.as_ref()))
203        }
204        other => Cow::Borrowed(other),
205    }
206}
207
208async fn create_static_table(
209    table_name: Cow<'_, TableReference>,
210    metadata_file_path: &str,
211    props: &HashMap<String, String>,
212    storage_factory: Arc<dyn StorageFactory>,
213) -> Result<StaticTable> {
214    let table_ident = TableIdent::from_strs(table_name.to_vec())?;
215    let file_io = FileIOBuilder::new(storage_factory)
216        .with_props(props)
217        .build();
218    StaticTable::from_metadata_file(metadata_file_path, table_ident, file_io).await
219}
220
221#[cfg(test)]
222mod tests {
223
224    use datafusion::arrow::datatypes::{DataType, Field, Schema};
225    use datafusion::catalog::TableProviderFactory;
226    use datafusion::common::{Constraints, DFSchema, TableReference};
227    use datafusion::execution::session_state::SessionStateBuilder;
228    use datafusion::logical_expr::CreateExternalTable;
229    use datafusion::parquet::arrow::PARQUET_FIELD_ID_META_KEY;
230    use datafusion::prelude::SessionContext;
231
232    use super::*;
233
234    fn table_metadata_v2_schema() -> Schema {
235        Schema::new(vec![
236            Field::new("x", DataType::Int64, false).with_metadata(HashMap::from([(
237                PARQUET_FIELD_ID_META_KEY.to_string(),
238                "1".to_string(),
239            )])),
240            Field::new("y", DataType::Int64, false).with_metadata(HashMap::from([(
241                PARQUET_FIELD_ID_META_KEY.to_string(),
242                "2".to_string(),
243            )])),
244            Field::new("z", DataType::Int64, false).with_metadata(HashMap::from([(
245                PARQUET_FIELD_ID_META_KEY.to_string(),
246                "3".to_string(),
247            )])),
248        ])
249    }
250
251    fn table_metadata_location() -> String {
252        format!(
253            "{}/testdata/table_metadata/{}",
254            env!("CARGO_MANIFEST_DIR"),
255            "TableMetadataV2.json"
256        )
257    }
258
259    fn create_external_table_cmd() -> CreateExternalTable {
260        let metadata_file_path = table_metadata_location();
261
262        CreateExternalTable {
263            name: TableReference::partial("static_ns", "static_table"),
264            locations: vec![metadata_file_path],
265            schema: Arc::new(DFSchema::empty()),
266            file_type: "iceberg".to_string(),
267            options: Default::default(),
268            table_partition_cols: Default::default(),
269            order_exprs: Default::default(),
270            constraints: Constraints::default(),
271            column_defaults: Default::default(),
272            if_not_exists: Default::default(),
273            or_replace: false,
274            temporary: false,
275            definition: Default::default(),
276            unbounded: Default::default(),
277        }
278    }
279
280    #[tokio::test]
281    async fn test_schema_of_created_table() {
282        let factory = IcebergTableProviderFactory::new();
283
284        let state = SessionStateBuilder::new().build();
285        let cmd = create_external_table_cmd();
286
287        let table_provider = factory
288            .create(&state, &cmd)
289            .await
290            .expect("create table failed");
291
292        let expected_schema = table_metadata_v2_schema();
293        let actual_schema = table_provider.schema();
294
295        assert_eq!(actual_schema.as_ref(), &expected_schema);
296    }
297
298    #[tokio::test]
299    async fn test_schema_of_created_external_table_sql() {
300        let mut state = SessionStateBuilder::new().with_default_features().build();
301        state.table_factories_mut().insert(
302            "ICEBERG".to_string(),
303            Arc::new(IcebergTableProviderFactory::new()),
304        );
305        let ctx = SessionContext::new_with_state(state);
306
307        // All external tables in DataFusion use bare names.
308        // See https://github.com/apache/datafusion/blob/main/datafusion/sql/src/statement.rs#L1038-#L1039
309        let table_ref = TableReference::bare("static_table");
310
311        // Create the external table
312        let sql = format!(
313            "CREATE EXTERNAL TABLE {} STORED AS ICEBERG LOCATION '{}'",
314            table_ref,
315            table_metadata_location()
316        );
317        let _df = ctx.sql(&sql).await.expect("create table failed");
318
319        // Get the created external table
320        let table_provider = ctx
321            .table_provider(table_ref)
322            .await
323            .expect("table not found");
324
325        // Check the schema of the created table
326        let expected_schema = table_metadata_v2_schema();
327        let actual_schema = table_provider.schema();
328
329        assert_eq!(actual_schema.as_ref(), &expected_schema);
330    }
331}
iceberg_datafusion/table/table_provider_factory.rs

iceberg_datafusion/table/
table_provider_factory.rs
