Skip to main content

iceberg_datafusion/table/
mod.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
18//! Iceberg table providers for DataFusion.
19//!
20//! This module provides two table provider implementations:
21//!
22//! - [`IcebergTableProvider`]: Catalog-backed provider with automatic metadata refresh.
23//!   Use for write operations and when you need to see the latest table state.
24//!
25//! - [`IcebergStaticTableProvider`]: Static provider for read-only access to a specific
26//!   table snapshot. Use for consistent analytical queries or time-travel scenarios.
27
28pub mod metadata_table;
29pub mod table_provider_factory;
30
31use std::num::NonZeroUsize;
32use std::sync::Arc;
33
34use async_trait::async_trait;
35use datafusion::arrow::datatypes::SchemaRef as ArrowSchemaRef;
36use datafusion::catalog::Session;
37use datafusion::common::DataFusionError;
38use datafusion::datasource::{TableProvider, TableType};
39use datafusion::error::Result as DFResult;
40use datafusion::logical_expr::dml::InsertOp;
41use datafusion::logical_expr::{Expr, TableProviderFilterPushDown};
42use datafusion::physical_plan::ExecutionPlan;
43use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec;
44use iceberg::arrow::schema_to_arrow_schema;
45use iceberg::inspect::MetadataTableType;
46use iceberg::spec::TableProperties;
47use iceberg::table::Table;
48use iceberg::{Catalog, Error, ErrorKind, NamespaceIdent, Result, TableIdent};
49use metadata_table::IcebergMetadataTableProvider;
50
51use crate::error::to_datafusion_error;
52use crate::physical_plan::commit::IcebergCommitExec;
53use crate::physical_plan::project::project_with_partition;
54use crate::physical_plan::repartition::repartition;
55use crate::physical_plan::scan::IcebergTableScan;
56use crate::physical_plan::sort::sort_by_partition;
57use crate::physical_plan::write::IcebergWriteExec;
58
59/// Catalog-backed table provider with automatic metadata refresh.
60///
61/// This provider loads fresh table metadata from the catalog on every scan and write
62/// operation, ensuring you always see the latest table state. Use this when you need
63/// write operations or want to see the most up-to-date data.
64///
65/// For read-only access to a specific snapshot without catalog overhead, use
66/// [`IcebergStaticTableProvider`] instead.
67#[derive(Debug, Clone)]
68pub struct IcebergTableProvider {
69    /// The catalog that manages this table
70    catalog: Arc<dyn Catalog>,
71    /// The table identifier (namespace + name)
72    table_ident: TableIdent,
73    /// A reference-counted arrow `Schema` (cached at construction)
74    schema: ArrowSchemaRef,
75}
76
77impl IcebergTableProvider {
78    /// Creates a new catalog-backed table provider.
79    ///
80    /// Loads the table once to get the initial schema, then stores the catalog
81    /// reference for future metadata refreshes on each operation.
82    pub(crate) async fn try_new(
83        catalog: Arc<dyn Catalog>,
84        namespace: NamespaceIdent,
85        name: impl Into<String>,
86    ) -> Result<Self> {
87        let table_ident = TableIdent::new(namespace, name.into());
88
89        // Load table once to get initial schema
90        let table = catalog.load_table(&table_ident).await?;
91        let schema = Arc::new(schema_to_arrow_schema(table.metadata().current_schema())?);
92
93        Ok(IcebergTableProvider {
94            catalog,
95            table_ident,
96            schema,
97        })
98    }
99
100    pub(crate) async fn metadata_table(
101        &self,
102        r#type: MetadataTableType,
103    ) -> Result<IcebergMetadataTableProvider> {
104        // Load fresh table metadata for metadata table access
105        let table = self.catalog.load_table(&self.table_ident).await?;
106        Ok(IcebergMetadataTableProvider { table, r#type })
107    }
108}
109
110#[async_trait]
111impl TableProvider for IcebergTableProvider {
112    fn schema(&self) -> ArrowSchemaRef {
113        self.schema.clone()
114    }
115
116    fn table_type(&self) -> TableType {
117        TableType::Base
118    }
119
120    async fn scan(
121        &self,
122        _state: &dyn Session,
123        projection: Option<&Vec<usize>>,
124        filters: &[Expr],
125        limit: Option<usize>,
126    ) -> DFResult<Arc<dyn ExecutionPlan>> {
127        // Load fresh table metadata from catalog
128        let table = self
129            .catalog
130            .load_table(&self.table_ident)
131            .await
132            .map_err(to_datafusion_error)?;
133
134        // Create scan with fresh metadata (always use current snapshot)
135        Ok(Arc::new(IcebergTableScan::new(
136            table,
137            None, // Always use current snapshot for catalog-backed provider
138            self.schema.clone(),
139            projection,
140            filters,
141            limit,
142        )))
143    }
144
145    fn supports_filters_pushdown(
146        &self,
147        filters: &[&Expr],
148    ) -> DFResult<Vec<TableProviderFilterPushDown>> {
149        // Push down all filters, as a single source of truth, the scanner will drop the filters which couldn't be push down
150        Ok(vec![TableProviderFilterPushDown::Inexact; filters.len()])
151    }
152
153    async fn insert_into(
154        &self,
155        state: &dyn Session,
156        input: Arc<dyn ExecutionPlan>,
157        _insert_op: InsertOp,
158    ) -> DFResult<Arc<dyn ExecutionPlan>> {
159        if _insert_op != InsertOp::Append {
160            return Err(DataFusionError::NotImplemented(format!(
161                "IcebergTableProvider supports only append inserts, got {_insert_op}"
162            )));
163        }
164
165        // Load fresh table metadata from catalog
166        let table = self
167            .catalog
168            .load_table(&self.table_ident)
169            .await
170            .map_err(to_datafusion_error)?;
171
172        let partition_spec = table.metadata().default_partition_spec();
173
174        // Step 1: Project partition values for partitioned tables
175        let plan_with_partition = if !partition_spec.is_unpartitioned() {
176            project_with_partition(input, &table)?
177        } else {
178            input
179        };
180
181        // Step 2: Repartition for parallel processing
182        let target_partitions =
183            NonZeroUsize::new(state.config().target_partitions()).ok_or_else(|| {
184                DataFusionError::Configuration(
185                    "target_partitions must be greater than 0".to_string(),
186                )
187            })?;
188
189        let repartitioned_plan =
190            repartition(plan_with_partition, table.metadata_ref(), target_partitions)?;
191
192        // Apply sort node when it's not fanout mode
193        let fanout_enabled = table
194            .metadata()
195            .properties()
196            .get(TableProperties::PROPERTY_DATAFUSION_WRITE_FANOUT_ENABLED)
197            .map(|value| {
198                value
199                    .parse::<bool>()
200                    .map_err(|e| {
201                        Error::new(
202                            ErrorKind::DataInvalid,
203                            format!(
204                                "Invalid value for {}, expected 'true' or 'false'",
205                                TableProperties::PROPERTY_DATAFUSION_WRITE_FANOUT_ENABLED
206                            ),
207                        )
208                        .with_source(e)
209                    })
210                    .map_err(to_datafusion_error)
211            })
212            .transpose()?
213            .unwrap_or(TableProperties::PROPERTY_DATAFUSION_WRITE_FANOUT_ENABLED_DEFAULT);
214
215        let write_input = if fanout_enabled {
216            repartitioned_plan
217        } else {
218            sort_by_partition(repartitioned_plan)?
219        };
220
221        let write_plan = Arc::new(IcebergWriteExec::new(table.clone(), write_input));
222
223        // Merge the outputs of write_plan into one so we can commit all files together
224        let coalesce_partitions = Arc::new(CoalescePartitionsExec::new(write_plan));
225
226        Ok(Arc::new(IcebergCommitExec::new(
227            table,
228            self.catalog.clone(),
229            coalesce_partitions,
230            self.schema.clone(),
231        )))
232    }
233}
234
235/// Static table provider for read-only snapshot access.
236///
237/// This provider holds a cached table instance and does not refresh metadata or support
238/// write operations. Use this for consistent analytical queries, time-travel scenarios,
239/// or when you want to avoid catalog overhead.
240///
241/// For catalog-backed tables with write support and automatic refresh, use
242/// [`IcebergTableProvider`] instead.
243#[derive(Debug, Clone)]
244pub struct IcebergStaticTableProvider {
245    /// The static table instance (never refreshed)
246    table: Table,
247    /// Optional snapshot ID for this static view
248    snapshot_id: Option<i64>,
249    /// A reference-counted arrow `Schema`
250    schema: ArrowSchemaRef,
251}
252
253impl IcebergStaticTableProvider {
254    /// Creates a static provider from a table instance.
255    ///
256    /// Uses the table's current snapshot for all queries. Does not support write operations.
257    pub async fn try_new_from_table(table: Table) -> Result<Self> {
258        let schema = Arc::new(schema_to_arrow_schema(table.metadata().current_schema())?);
259        Ok(IcebergStaticTableProvider {
260            table,
261            snapshot_id: None,
262            schema,
263        })
264    }
265
266    /// Creates a static provider for a specific table snapshot.
267    ///
268    /// Queries the specified snapshot for all operations. Useful for time-travel queries.
269    /// Does not support write operations.
270    pub async fn try_new_from_table_snapshot(table: Table, snapshot_id: i64) -> Result<Self> {
271        let snapshot = table
272            .metadata()
273            .snapshot_by_id(snapshot_id)
274            .ok_or_else(|| {
275                Error::new(
276                    ErrorKind::Unexpected,
277                    format!(
278                        "snapshot id {snapshot_id} not found in table {}",
279                        table.identifier().name()
280                    ),
281                )
282            })?;
283        let table_schema = snapshot.schema(table.metadata())?;
284        let schema = Arc::new(schema_to_arrow_schema(&table_schema)?);
285        Ok(IcebergStaticTableProvider {
286            table,
287            snapshot_id: Some(snapshot_id),
288            schema,
289        })
290    }
291}
292
293#[async_trait]
294impl TableProvider for IcebergStaticTableProvider {
295    fn schema(&self) -> ArrowSchemaRef {
296        self.schema.clone()
297    }
298
299    fn table_type(&self) -> TableType {
300        TableType::Base
301    }
302
303    async fn scan(
304        &self,
305        _state: &dyn Session,
306        projection: Option<&Vec<usize>>,
307        filters: &[Expr],
308        limit: Option<usize>,
309    ) -> DFResult<Arc<dyn ExecutionPlan>> {
310        // Use cached table (no refresh)
311        Ok(Arc::new(IcebergTableScan::new(
312            self.table.clone(),
313            self.snapshot_id,
314            self.schema.clone(),
315            projection,
316            filters,
317            limit,
318        )))
319    }
320
321    fn supports_filters_pushdown(
322        &self,
323        filters: &[&Expr],
324    ) -> DFResult<Vec<TableProviderFilterPushDown>> {
325        // Push down all filters, as a single source of truth, the scanner will drop the filters which couldn't be push down
326        Ok(vec![TableProviderFilterPushDown::Inexact; filters.len()])
327    }
328
329    async fn insert_into(
330        &self,
331        _state: &dyn Session,
332        _input: Arc<dyn ExecutionPlan>,
333        _insert_op: InsertOp,
334    ) -> DFResult<Arc<dyn ExecutionPlan>> {
335        Err(to_datafusion_error(Error::new(
336            ErrorKind::FeatureUnsupported,
337            "Write operations are not supported on IcebergStaticTableProvider. \
338             Use IcebergTableProvider with a catalog for write support."
339                .to_string(),
340        )))
341    }
342}
343
344#[cfg(test)]
345mod tests {
346    use std::collections::HashMap;
347    use std::sync::Arc;
348
349    use datafusion::common::Column;
350    use datafusion::physical_plan::ExecutionPlan;
351    use datafusion::prelude::SessionContext;
352    use iceberg::io::FileIO;
353    use iceberg::memory::{MEMORY_CATALOG_WAREHOUSE, MemoryCatalogBuilder};
354    use iceberg::spec::{NestedField, PrimitiveType, Schema, Type};
355    use iceberg::table::{StaticTable, Table};
356    use iceberg::{Catalog, CatalogBuilder, NamespaceIdent, TableCreation, TableIdent};
357    use tempfile::TempDir;
358
359    use super::*;
360
361    async fn get_test_table_from_metadata_file() -> Table {
362        let metadata_file_name = "TableMetadataV2Valid.json";
363        let metadata_file_path = format!(
364            "{}/tests/test_data/{}",
365            env!("CARGO_MANIFEST_DIR"),
366            metadata_file_name
367        );
368        let file_io = FileIO::new_with_fs();
369        let static_identifier = TableIdent::from_strs(["static_ns", "static_table"]).unwrap();
370        let static_table =
371            StaticTable::from_metadata_file(&metadata_file_path, static_identifier, file_io)
372                .await
373                .unwrap();
374        static_table.into_table()
375    }
376
377    async fn get_test_catalog_and_table() -> (Arc<dyn Catalog>, NamespaceIdent, String, TempDir) {
378        let temp_dir = TempDir::new().unwrap();
379        let warehouse_path = temp_dir.path().to_str().unwrap().to_string();
380
381        let catalog = MemoryCatalogBuilder::default()
382            .load(
383                "memory",
384                HashMap::from([(MEMORY_CATALOG_WAREHOUSE.to_string(), warehouse_path.clone())]),
385            )
386            .await
387            .unwrap();
388
389        let namespace = NamespaceIdent::new("test_ns".to_string());
390        catalog
391            .create_namespace(&namespace, HashMap::new())
392            .await
393            .unwrap();
394
395        let schema = Schema::builder()
396            .with_schema_id(0)
397            .with_fields(vec![
398                NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
399                NestedField::required(2, "name", Type::Primitive(PrimitiveType::String)).into(),
400            ])
401            .build()
402            .unwrap();
403
404        let table_creation = TableCreation::builder()
405            .name("test_table".to_string())
406            .location(format!("{warehouse_path}/test_table"))
407            .schema(schema)
408            .properties(HashMap::new())
409            .build();
410
411        catalog
412            .create_table(&namespace, table_creation)
413            .await
414            .unwrap();
415
416        (
417            Arc::new(catalog),
418            namespace,
419            "test_table".to_string(),
420            temp_dir,
421        )
422    }
423
424    // Tests for IcebergStaticTableProvider
425
426    #[tokio::test]
427    async fn test_static_provider_from_table() {
428        let table = get_test_table_from_metadata_file().await;
429        let table_provider = IcebergStaticTableProvider::try_new_from_table(table.clone())
430            .await
431            .unwrap();
432        let ctx = SessionContext::new();
433        ctx.register_table("mytable", Arc::new(table_provider))
434            .unwrap();
435        let df = ctx.sql("SELECT * FROM mytable").await.unwrap();
436        let df_schema = df.schema();
437        let df_columns = df_schema.fields();
438        assert_eq!(df_columns.len(), 3);
439        let x_column = df_columns.first().unwrap();
440        let column_data = format!(
441            "{:?}:{:?}",
442            x_column.name(),
443            x_column.data_type().to_string()
444        );
445        assert_eq!(column_data, "\"x\":\"Int64\"");
446        let has_column = df_schema.has_column(&Column::from_name("z"));
447        assert!(has_column);
448    }
449
450    #[tokio::test]
451    async fn test_static_provider_from_snapshot() {
452        let table = get_test_table_from_metadata_file().await;
453        let snapshot_id = table.metadata().snapshots().next().unwrap().snapshot_id();
454        let table_provider =
455            IcebergStaticTableProvider::try_new_from_table_snapshot(table.clone(), snapshot_id)
456                .await
457                .unwrap();
458        let ctx = SessionContext::new();
459        ctx.register_table("mytable", Arc::new(table_provider))
460            .unwrap();
461        let df = ctx.sql("SELECT * FROM mytable").await.unwrap();
462        let df_schema = df.schema();
463        let df_columns = df_schema.fields();
464        assert_eq!(df_columns.len(), 3);
465        let x_column = df_columns.first().unwrap();
466        let column_data = format!(
467            "{:?}:{:?}",
468            x_column.name(),
469            x_column.data_type().to_string()
470        );
471        assert_eq!(column_data, "\"x\":\"Int64\"");
472        let has_column = df_schema.has_column(&Column::from_name("z"));
473        assert!(has_column);
474    }
475
476    #[tokio::test]
477    async fn test_static_provider_rejects_writes() {
478        let table = get_test_table_from_metadata_file().await;
479        let table_provider = IcebergStaticTableProvider::try_new_from_table(table.clone())
480            .await
481            .unwrap();
482        let ctx = SessionContext::new();
483        ctx.register_table("mytable", Arc::new(table_provider))
484            .unwrap();
485
486        // Attempt to insert into the static provider should fail
487        let result = ctx.sql("INSERT INTO mytable VALUES (1, 2, 3)").await;
488
489        // The error should occur during planning or execution
490        // We expect an error indicating write operations are not supported
491        assert!(
492            result.is_err() || {
493                let df = result.unwrap();
494                df.collect().await.is_err()
495            }
496        );
497    }
498
499    #[tokio::test]
500    async fn test_static_provider_scan() {
501        let table = get_test_table_from_metadata_file().await;
502        let table_provider = IcebergStaticTableProvider::try_new_from_table(table.clone())
503            .await
504            .unwrap();
505        let ctx = SessionContext::new();
506        ctx.register_table("mytable", Arc::new(table_provider))
507            .unwrap();
508
509        // Test that scan operations work correctly
510        let df = ctx.sql("SELECT count(*) FROM mytable").await.unwrap();
511        let physical_plan = df.create_physical_plan().await;
512        assert!(physical_plan.is_ok());
513    }
514
515    // Tests for IcebergTableProvider
516
517    #[tokio::test]
518    async fn test_catalog_backed_provider_creation() {
519        let (catalog, namespace, table_name, _temp_dir) = get_test_catalog_and_table().await;
520
521        // Test creating a catalog-backed provider
522        let provider =
523            IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), table_name.clone())
524                .await
525                .unwrap();
526
527        // Verify the schema is loaded correctly
528        let schema = provider.schema();
529        assert_eq!(schema.fields().len(), 2);
530        assert_eq!(schema.field(0).name(), "id");
531        assert_eq!(schema.field(1).name(), "name");
532    }
533
534    #[tokio::test]
535    async fn test_catalog_backed_provider_scan() {
536        let (catalog, namespace, table_name, _temp_dir) = get_test_catalog_and_table().await;
537
538        let provider =
539            IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), table_name.clone())
540                .await
541                .unwrap();
542
543        let ctx = SessionContext::new();
544        ctx.register_table("test_table", Arc::new(provider))
545            .unwrap();
546
547        // Test that scan operations work correctly
548        let df = ctx.sql("SELECT * FROM test_table").await.unwrap();
549
550        // Verify the schema in the query result
551        let df_schema = df.schema();
552        assert_eq!(df_schema.fields().len(), 2);
553        assert_eq!(df_schema.field(0).name(), "id");
554        assert_eq!(df_schema.field(1).name(), "name");
555
556        let physical_plan = df.create_physical_plan().await;
557        assert!(physical_plan.is_ok());
558    }
559
560    #[tokio::test]
561    async fn test_catalog_backed_provider_insert() {
562        let (catalog, namespace, table_name, _temp_dir) = get_test_catalog_and_table().await;
563
564        let provider =
565            IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), table_name.clone())
566                .await
567                .unwrap();
568
569        let ctx = SessionContext::new();
570        ctx.register_table("test_table", Arc::new(provider))
571            .unwrap();
572
573        // Test that insert operations work correctly
574        let result = ctx.sql("INSERT INTO test_table VALUES (1, 'test')").await;
575
576        // Insert should succeed (or at least not fail during planning)
577        assert!(result.is_ok());
578
579        // Try to execute the insert plan
580        let df = result.unwrap();
581        let execution_result = df.collect().await;
582
583        // The execution should succeed
584        assert!(execution_result.is_ok());
585    }
586
587    #[tokio::test]
588    async fn test_physical_input_schema_consistent_with_logical_input_schema() {
589        let (catalog, namespace, table_name, _temp_dir) = get_test_catalog_and_table().await;
590
591        let provider =
592            IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), table_name.clone())
593                .await
594                .unwrap();
595
596        let ctx = SessionContext::new();
597        ctx.register_table("test_table", Arc::new(provider))
598            .unwrap();
599
600        // Create a query plan
601        let df = ctx.sql("SELECT id, name FROM test_table").await.unwrap();
602
603        // Get logical schema before consuming df
604        let logical_schema = df.schema().clone();
605
606        // Get physical plan (this consumes df)
607        let physical_plan = df.create_physical_plan().await.unwrap();
608        let physical_schema = physical_plan.schema();
609
610        // Verify that logical and physical schemas are consistent
611        assert_eq!(
612            logical_schema.fields().len(),
613            physical_schema.fields().len()
614        );
615
616        for (logical_field, physical_field) in logical_schema
617            .fields()
618            .iter()
619            .zip(physical_schema.fields().iter())
620        {
621            assert_eq!(logical_field.name(), physical_field.name());
622            assert_eq!(logical_field.data_type(), physical_field.data_type());
623        }
624    }
625
626    async fn get_partitioned_test_catalog_and_table(
627        fanout_enabled: Option<bool>,
628    ) -> (Arc<dyn Catalog>, NamespaceIdent, String, TempDir) {
629        use iceberg::spec::{Transform, UnboundPartitionSpec};
630
631        let temp_dir = TempDir::new().unwrap();
632        let warehouse_path = temp_dir.path().to_str().unwrap().to_string();
633
634        let catalog = MemoryCatalogBuilder::default()
635            .load(
636                "memory",
637                HashMap::from([(MEMORY_CATALOG_WAREHOUSE.to_string(), warehouse_path.clone())]),
638            )
639            .await
640            .unwrap();
641
642        let namespace = NamespaceIdent::new("test_ns".to_string());
643        catalog
644            .create_namespace(&namespace, HashMap::new())
645            .await
646            .unwrap();
647
648        let schema = Schema::builder()
649            .with_schema_id(0)
650            .with_fields(vec![
651                NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
652                NestedField::required(2, "category", Type::Primitive(PrimitiveType::String)).into(),
653            ])
654            .build()
655            .unwrap();
656
657        let partition_spec = UnboundPartitionSpec::builder()
658            .with_spec_id(0)
659            .add_partition_field(2, "category", Transform::Identity)
660            .unwrap()
661            .build();
662
663        let mut properties = HashMap::new();
664        if let Some(enabled) = fanout_enabled {
665            properties.insert(
666                TableProperties::PROPERTY_DATAFUSION_WRITE_FANOUT_ENABLED.to_string(),
667                enabled.to_string(),
668            );
669        }
670
671        let table_creation = TableCreation::builder()
672            .name("partitioned_table".to_string())
673            .location(format!("{warehouse_path}/partitioned_table"))
674            .schema(schema)
675            .partition_spec(partition_spec)
676            .properties(properties)
677            .build();
678
679        catalog
680            .create_table(&namespace, table_creation)
681            .await
682            .unwrap();
683
684        (
685            Arc::new(catalog),
686            namespace,
687            "partitioned_table".to_string(),
688            temp_dir,
689        )
690    }
691
692    /// Helper to check if a plan contains a SortExec node
693    fn plan_contains_sort(plan: &Arc<dyn ExecutionPlan>) -> bool {
694        if plan.name() == "SortExec" {
695            return true;
696        }
697        for child in plan.children() {
698            if plan_contains_sort(child) {
699                return true;
700            }
701        }
702        false
703    }
704
705    #[tokio::test]
706    async fn test_catalog_backed_provider_rejects_non_append_op() {
707        use datafusion::physical_plan::empty::EmptyExec;
708
709        let (catalog, namespace, table_name, _temp_dir) = get_test_catalog_and_table().await;
710        let provider = IcebergTableProvider::try_new(catalog, namespace, table_name)
711            .await
712            .unwrap();
713        let ctx = SessionContext::new();
714
715        for (insert_op, expected_message) in [
716            (
717                InsertOp::Overwrite,
718                "IcebergTableProvider supports only append inserts, got Insert Overwrite",
719            ),
720            (
721                InsertOp::Replace,
722                "IcebergTableProvider supports only append inserts, got Replace Into",
723            ),
724        ] {
725            let input = Arc::new(EmptyExec::new(provider.schema())) as Arc<dyn ExecutionPlan>;
726            let error = provider
727                .insert_into(&ctx.state(), input, insert_op)
728                .await
729                .expect_err("non-append inserts should be rejected");
730
731            assert!(
732                matches!(
733                    error,
734                    DataFusionError::NotImplemented(ref message) if message == expected_message
735                ),
736                "unexpected error: {error}"
737            );
738        }
739    }
740
741    #[tokio::test]
742    async fn test_insert_plan_fanout_enabled_no_sort() {
743        use datafusion::datasource::TableProvider;
744        use datafusion::logical_expr::dml::InsertOp;
745        use datafusion::physical_plan::empty::EmptyExec;
746
747        // When fanout is enabled (default), no sort node should be added
748        let (catalog, namespace, table_name, _temp_dir) =
749            get_partitioned_test_catalog_and_table(Some(true)).await;
750
751        let provider =
752            IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), table_name.clone())
753                .await
754                .unwrap();
755
756        let ctx = SessionContext::new();
757        let input_schema = provider.schema();
758        let input = Arc::new(EmptyExec::new(input_schema)) as Arc<dyn ExecutionPlan>;
759
760        let state = ctx.state();
761        let insert_plan = provider
762            .insert_into(&state, input, InsertOp::Append)
763            .await
764            .unwrap();
765
766        // With fanout enabled, there should be no SortExec in the plan
767        assert!(
768            !plan_contains_sort(&insert_plan),
769            "Plan should NOT contain SortExec when fanout is enabled"
770        );
771    }
772
773    #[tokio::test]
774    async fn test_insert_plan_fanout_disabled_has_sort() {
775        use datafusion::datasource::TableProvider;
776        use datafusion::logical_expr::dml::InsertOp;
777        use datafusion::physical_plan::empty::EmptyExec;
778
779        // When fanout is disabled, a sort node should be added
780        let (catalog, namespace, table_name, _temp_dir) =
781            get_partitioned_test_catalog_and_table(Some(false)).await;
782
783        let provider =
784            IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), table_name.clone())
785                .await
786                .unwrap();
787
788        let ctx = SessionContext::new();
789        let input_schema = provider.schema();
790        let input = Arc::new(EmptyExec::new(input_schema)) as Arc<dyn ExecutionPlan>;
791
792        let state = ctx.state();
793        let insert_plan = provider
794            .insert_into(&state, input, InsertOp::Append)
795            .await
796            .unwrap();
797
798        // With fanout disabled, there should be a SortExec in the plan
799        assert!(
800            plan_contains_sort(&insert_plan),
801            "Plan should contain SortExec when fanout is disabled"
802        );
803    }
804
805    #[tokio::test]
806    async fn test_limit_pushdown_static_provider() {
807        use datafusion::datasource::TableProvider;
808
809        let table = get_test_table_from_metadata_file().await;
810        let table_provider = IcebergStaticTableProvider::try_new_from_table(table.clone())
811            .await
812            .unwrap();
813
814        let ctx = SessionContext::new();
815        let state = ctx.state();
816
817        // Test scan with limit
818        let scan_plan = table_provider
819            .scan(&state, None, &[], Some(10))
820            .await
821            .unwrap();
822
823        // Verify that the scan plan is an IcebergTableScan
824        let iceberg_scan = scan_plan
825            .downcast_ref::<IcebergTableScan>()
826            .expect("Expected IcebergTableScan");
827
828        // Verify the limit is set
829        assert_eq!(
830            iceberg_scan.limit(),
831            Some(10),
832            "Limit should be set to 10 in the scan plan"
833        );
834    }
835
836    #[tokio::test]
837    async fn test_limit_pushdown_catalog_backed_provider() {
838        use datafusion::datasource::TableProvider;
839
840        let (catalog, namespace, table_name, _temp_dir) = get_test_catalog_and_table().await;
841
842        let provider =
843            IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), table_name.clone())
844                .await
845                .unwrap();
846
847        let ctx = SessionContext::new();
848        let state = ctx.state();
849
850        // Test scan with limit
851        let scan_plan = provider.scan(&state, None, &[], Some(5)).await.unwrap();
852
853        // Verify that the scan plan is an IcebergTableScan
854        let iceberg_scan = scan_plan
855            .downcast_ref::<IcebergTableScan>()
856            .expect("Expected IcebergTableScan");
857
858        // Verify the limit is set
859        assert_eq!(
860            iceberg_scan.limit(),
861            Some(5),
862            "Limit should be set to 5 in the scan plan"
863        );
864    }
865
866    #[tokio::test]
867    async fn test_no_limit_pushdown() {
868        use datafusion::datasource::TableProvider;
869
870        let table = get_test_table_from_metadata_file().await;
871        let table_provider = IcebergStaticTableProvider::try_new_from_table(table.clone())
872            .await
873            .unwrap();
874
875        let ctx = SessionContext::new();
876        let state = ctx.state();
877
878        // Test scan without limit
879        let scan_plan = table_provider.scan(&state, None, &[], None).await.unwrap();
880
881        // Verify that the scan plan is an IcebergTableScan
882        let iceberg_scan = scan_plan
883            .downcast_ref::<IcebergTableScan>()
884            .expect("Expected IcebergTableScan");
885
886        // Verify the limit is None
887        assert_eq!(
888            iceberg_scan.limit(),
889            None,
890            "Limit should be None when not specified"
891        );
892    }
893}