1pub 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#[derive(Debug, Clone)]
68pub struct IcebergTableProvider {
69 catalog: Arc<dyn Catalog>,
71 table_ident: TableIdent,
73 schema: ArrowSchemaRef,
75}
76
77impl IcebergTableProvider {
78 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 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 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 let table = self
129 .catalog
130 .load_table(&self.table_ident)
131 .await
132 .map_err(to_datafusion_error)?;
133
134 Ok(Arc::new(IcebergTableScan::new(
136 table,
137 None, 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 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 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 let plan_with_partition = if !partition_spec.is_unpartitioned() {
176 project_with_partition(input, &table)?
177 } else {
178 input
179 };
180
181 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 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(
222 table.clone(),
223 write_input,
224 self.schema.clone(),
225 ));
226
227 let coalesce_partitions = Arc::new(CoalescePartitionsExec::new(write_plan));
229
230 Ok(Arc::new(IcebergCommitExec::new(
231 table,
232 self.catalog.clone(),
233 coalesce_partitions,
234 self.schema.clone(),
235 )))
236 }
237}
238
239#[derive(Debug, Clone)]
248pub struct IcebergStaticTableProvider {
249 table: Table,
251 snapshot_id: Option<i64>,
253 schema: ArrowSchemaRef,
255}
256
257impl IcebergStaticTableProvider {
258 pub async fn try_new_from_table(table: Table) -> Result<Self> {
262 let schema = Arc::new(schema_to_arrow_schema(table.metadata().current_schema())?);
263 Ok(IcebergStaticTableProvider {
264 table,
265 snapshot_id: None,
266 schema,
267 })
268 }
269
270 pub async fn try_new_from_table_snapshot(table: Table, snapshot_id: i64) -> Result<Self> {
275 let snapshot = table
276 .metadata()
277 .snapshot_by_id(snapshot_id)
278 .ok_or_else(|| {
279 Error::new(
280 ErrorKind::Unexpected,
281 format!(
282 "snapshot id {snapshot_id} not found in table {}",
283 table.identifier().name()
284 ),
285 )
286 })?;
287 let table_schema = snapshot.schema(table.metadata())?;
288 let schema = Arc::new(schema_to_arrow_schema(&table_schema)?);
289 Ok(IcebergStaticTableProvider {
290 table,
291 snapshot_id: Some(snapshot_id),
292 schema,
293 })
294 }
295}
296
297#[async_trait]
298impl TableProvider for IcebergStaticTableProvider {
299 fn schema(&self) -> ArrowSchemaRef {
300 self.schema.clone()
301 }
302
303 fn table_type(&self) -> TableType {
304 TableType::Base
305 }
306
307 async fn scan(
308 &self,
309 _state: &dyn Session,
310 projection: Option<&Vec<usize>>,
311 filters: &[Expr],
312 limit: Option<usize>,
313 ) -> DFResult<Arc<dyn ExecutionPlan>> {
314 Ok(Arc::new(IcebergTableScan::new(
316 self.table.clone(),
317 self.snapshot_id,
318 self.schema.clone(),
319 projection,
320 filters,
321 limit,
322 )))
323 }
324
325 fn supports_filters_pushdown(
326 &self,
327 filters: &[&Expr],
328 ) -> DFResult<Vec<TableProviderFilterPushDown>> {
329 Ok(vec![TableProviderFilterPushDown::Inexact; filters.len()])
331 }
332
333 async fn insert_into(
334 &self,
335 _state: &dyn Session,
336 _input: Arc<dyn ExecutionPlan>,
337 _insert_op: InsertOp,
338 ) -> DFResult<Arc<dyn ExecutionPlan>> {
339 Err(to_datafusion_error(Error::new(
340 ErrorKind::FeatureUnsupported,
341 "Write operations are not supported on IcebergStaticTableProvider. \
342 Use IcebergTableProvider with a catalog for write support."
343 .to_string(),
344 )))
345 }
346}
347
348#[cfg(test)]
349mod tests {
350 use std::collections::HashMap;
351 use std::sync::Arc;
352
353 use datafusion::common::Column;
354 use datafusion::physical_plan::ExecutionPlan;
355 use datafusion::prelude::SessionContext;
356 use iceberg::io::FileIO;
357 use iceberg::memory::{MEMORY_CATALOG_WAREHOUSE, MemoryCatalogBuilder};
358 use iceberg::spec::{NestedField, PrimitiveType, Schema, Type};
359 use iceberg::table::{StaticTable, Table};
360 use iceberg::{Catalog, CatalogBuilder, NamespaceIdent, TableCreation, TableIdent};
361 use tempfile::TempDir;
362
363 use super::*;
364
365 async fn get_test_table_from_metadata_file() -> Table {
366 let metadata_file_name = "TableMetadataV2Valid.json";
367 let metadata_file_path = format!(
368 "{}/tests/test_data/{}",
369 env!("CARGO_MANIFEST_DIR"),
370 metadata_file_name
371 );
372 let file_io = FileIO::new_with_fs();
373 let static_identifier = TableIdent::from_strs(["static_ns", "static_table"]).unwrap();
374 let static_table =
375 StaticTable::from_metadata_file(&metadata_file_path, static_identifier, file_io)
376 .await
377 .unwrap();
378 static_table.into_table()
379 }
380
381 async fn get_test_catalog_and_table() -> (Arc<dyn Catalog>, NamespaceIdent, String, TempDir) {
382 let temp_dir = TempDir::new().unwrap();
383 let warehouse_path = temp_dir.path().to_str().unwrap().to_string();
384
385 let catalog = MemoryCatalogBuilder::default()
386 .load(
387 "memory",
388 HashMap::from([(MEMORY_CATALOG_WAREHOUSE.to_string(), warehouse_path.clone())]),
389 )
390 .await
391 .unwrap();
392
393 let namespace = NamespaceIdent::new("test_ns".to_string());
394 catalog
395 .create_namespace(&namespace, HashMap::new())
396 .await
397 .unwrap();
398
399 let schema = Schema::builder()
400 .with_schema_id(0)
401 .with_fields(vec![
402 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
403 NestedField::required(2, "name", Type::Primitive(PrimitiveType::String)).into(),
404 ])
405 .build()
406 .unwrap();
407
408 let table_creation = TableCreation::builder()
409 .name("test_table".to_string())
410 .location(format!("{warehouse_path}/test_table"))
411 .schema(schema)
412 .properties(HashMap::new())
413 .build();
414
415 catalog
416 .create_table(&namespace, table_creation)
417 .await
418 .unwrap();
419
420 (
421 Arc::new(catalog),
422 namespace,
423 "test_table".to_string(),
424 temp_dir,
425 )
426 }
427
428 #[tokio::test]
431 async fn test_static_provider_from_table() {
432 let table = get_test_table_from_metadata_file().await;
433 let table_provider = IcebergStaticTableProvider::try_new_from_table(table.clone())
434 .await
435 .unwrap();
436 let ctx = SessionContext::new();
437 ctx.register_table("mytable", Arc::new(table_provider))
438 .unwrap();
439 let df = ctx.sql("SELECT * FROM mytable").await.unwrap();
440 let df_schema = df.schema();
441 let df_columns = df_schema.fields();
442 assert_eq!(df_columns.len(), 3);
443 let x_column = df_columns.first().unwrap();
444 let column_data = format!(
445 "{:?}:{:?}",
446 x_column.name(),
447 x_column.data_type().to_string()
448 );
449 assert_eq!(column_data, "\"x\":\"Int64\"");
450 let has_column = df_schema.has_column(&Column::from_name("z"));
451 assert!(has_column);
452 }
453
454 #[tokio::test]
455 async fn test_static_provider_from_snapshot() {
456 let table = get_test_table_from_metadata_file().await;
457 let snapshot_id = table.metadata().snapshots().next().unwrap().snapshot_id();
458 let table_provider =
459 IcebergStaticTableProvider::try_new_from_table_snapshot(table.clone(), snapshot_id)
460 .await
461 .unwrap();
462 let ctx = SessionContext::new();
463 ctx.register_table("mytable", Arc::new(table_provider))
464 .unwrap();
465 let df = ctx.sql("SELECT * FROM mytable").await.unwrap();
466 let df_schema = df.schema();
467 let df_columns = df_schema.fields();
468 assert_eq!(df_columns.len(), 3);
469 let x_column = df_columns.first().unwrap();
470 let column_data = format!(
471 "{:?}:{:?}",
472 x_column.name(),
473 x_column.data_type().to_string()
474 );
475 assert_eq!(column_data, "\"x\":\"Int64\"");
476 let has_column = df_schema.has_column(&Column::from_name("z"));
477 assert!(has_column);
478 }
479
480 #[tokio::test]
481 async fn test_static_provider_rejects_writes() {
482 let table = get_test_table_from_metadata_file().await;
483 let table_provider = IcebergStaticTableProvider::try_new_from_table(table.clone())
484 .await
485 .unwrap();
486 let ctx = SessionContext::new();
487 ctx.register_table("mytable", Arc::new(table_provider))
488 .unwrap();
489
490 let result = ctx.sql("INSERT INTO mytable VALUES (1, 2, 3)").await;
492
493 assert!(
496 result.is_err() || {
497 let df = result.unwrap();
498 df.collect().await.is_err()
499 }
500 );
501 }
502
503 #[tokio::test]
504 async fn test_static_provider_scan() {
505 let table = get_test_table_from_metadata_file().await;
506 let table_provider = IcebergStaticTableProvider::try_new_from_table(table.clone())
507 .await
508 .unwrap();
509 let ctx = SessionContext::new();
510 ctx.register_table("mytable", Arc::new(table_provider))
511 .unwrap();
512
513 let df = ctx.sql("SELECT count(*) FROM mytable").await.unwrap();
515 let physical_plan = df.create_physical_plan().await;
516 assert!(physical_plan.is_ok());
517 }
518
519 #[tokio::test]
522 async fn test_catalog_backed_provider_creation() {
523 let (catalog, namespace, table_name, _temp_dir) = get_test_catalog_and_table().await;
524
525 let provider =
527 IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), table_name.clone())
528 .await
529 .unwrap();
530
531 let schema = provider.schema();
533 assert_eq!(schema.fields().len(), 2);
534 assert_eq!(schema.field(0).name(), "id");
535 assert_eq!(schema.field(1).name(), "name");
536 }
537
538 #[tokio::test]
539 async fn test_catalog_backed_provider_scan() {
540 let (catalog, namespace, table_name, _temp_dir) = get_test_catalog_and_table().await;
541
542 let provider =
543 IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), table_name.clone())
544 .await
545 .unwrap();
546
547 let ctx = SessionContext::new();
548 ctx.register_table("test_table", Arc::new(provider))
549 .unwrap();
550
551 let df = ctx.sql("SELECT * FROM test_table").await.unwrap();
553
554 let df_schema = df.schema();
556 assert_eq!(df_schema.fields().len(), 2);
557 assert_eq!(df_schema.field(0).name(), "id");
558 assert_eq!(df_schema.field(1).name(), "name");
559
560 let physical_plan = df.create_physical_plan().await;
561 assert!(physical_plan.is_ok());
562 }
563
564 #[tokio::test]
565 async fn test_catalog_backed_provider_insert() {
566 let (catalog, namespace, table_name, _temp_dir) = get_test_catalog_and_table().await;
567
568 let provider =
569 IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), table_name.clone())
570 .await
571 .unwrap();
572
573 let ctx = SessionContext::new();
574 ctx.register_table("test_table", Arc::new(provider))
575 .unwrap();
576
577 let result = ctx.sql("INSERT INTO test_table VALUES (1, 'test')").await;
579
580 assert!(result.is_ok());
582
583 let df = result.unwrap();
585 let execution_result = df.collect().await;
586
587 assert!(execution_result.is_ok());
589 }
590
591 #[tokio::test]
592 async fn test_physical_input_schema_consistent_with_logical_input_schema() {
593 let (catalog, namespace, table_name, _temp_dir) = get_test_catalog_and_table().await;
594
595 let provider =
596 IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), table_name.clone())
597 .await
598 .unwrap();
599
600 let ctx = SessionContext::new();
601 ctx.register_table("test_table", Arc::new(provider))
602 .unwrap();
603
604 let df = ctx.sql("SELECT id, name FROM test_table").await.unwrap();
606
607 let logical_schema = df.schema().clone();
609
610 let physical_plan = df.create_physical_plan().await.unwrap();
612 let physical_schema = physical_plan.schema();
613
614 assert_eq!(
616 logical_schema.fields().len(),
617 physical_schema.fields().len()
618 );
619
620 for (logical_field, physical_field) in logical_schema
621 .fields()
622 .iter()
623 .zip(physical_schema.fields().iter())
624 {
625 assert_eq!(logical_field.name(), physical_field.name());
626 assert_eq!(logical_field.data_type(), physical_field.data_type());
627 }
628 }
629
630 async fn get_partitioned_test_catalog_and_table(
631 fanout_enabled: Option<bool>,
632 ) -> (Arc<dyn Catalog>, NamespaceIdent, String, TempDir) {
633 use iceberg::spec::{Transform, UnboundPartitionSpec};
634
635 let temp_dir = TempDir::new().unwrap();
636 let warehouse_path = temp_dir.path().to_str().unwrap().to_string();
637
638 let catalog = MemoryCatalogBuilder::default()
639 .load(
640 "memory",
641 HashMap::from([(MEMORY_CATALOG_WAREHOUSE.to_string(), warehouse_path.clone())]),
642 )
643 .await
644 .unwrap();
645
646 let namespace = NamespaceIdent::new("test_ns".to_string());
647 catalog
648 .create_namespace(&namespace, HashMap::new())
649 .await
650 .unwrap();
651
652 let schema = Schema::builder()
653 .with_schema_id(0)
654 .with_fields(vec![
655 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
656 NestedField::required(2, "category", Type::Primitive(PrimitiveType::String)).into(),
657 ])
658 .build()
659 .unwrap();
660
661 let partition_spec = UnboundPartitionSpec::builder()
662 .with_spec_id(0)
663 .add_partition_field(2, "category", Transform::Identity)
664 .unwrap()
665 .build();
666
667 let mut properties = HashMap::new();
668 if let Some(enabled) = fanout_enabled {
669 properties.insert(
670 TableProperties::PROPERTY_DATAFUSION_WRITE_FANOUT_ENABLED.to_string(),
671 enabled.to_string(),
672 );
673 }
674
675 let table_creation = TableCreation::builder()
676 .name("partitioned_table".to_string())
677 .location(format!("{warehouse_path}/partitioned_table"))
678 .schema(schema)
679 .partition_spec(partition_spec)
680 .properties(properties)
681 .build();
682
683 catalog
684 .create_table(&namespace, table_creation)
685 .await
686 .unwrap();
687
688 (
689 Arc::new(catalog),
690 namespace,
691 "partitioned_table".to_string(),
692 temp_dir,
693 )
694 }
695
696 fn plan_contains_sort(plan: &Arc<dyn ExecutionPlan>) -> bool {
698 if plan.name() == "SortExec" {
699 return true;
700 }
701 for child in plan.children() {
702 if plan_contains_sort(child) {
703 return true;
704 }
705 }
706 false
707 }
708
709 #[tokio::test]
710 async fn test_catalog_backed_provider_rejects_non_append_op() {
711 use datafusion::physical_plan::empty::EmptyExec;
712
713 let (catalog, namespace, table_name, _temp_dir) = get_test_catalog_and_table().await;
714 let provider = IcebergTableProvider::try_new(catalog, namespace, table_name)
715 .await
716 .unwrap();
717 let ctx = SessionContext::new();
718
719 for (insert_op, expected_message) in [
720 (
721 InsertOp::Overwrite,
722 "IcebergTableProvider supports only append inserts, got Insert Overwrite",
723 ),
724 (
725 InsertOp::Replace,
726 "IcebergTableProvider supports only append inserts, got Replace Into",
727 ),
728 ] {
729 let input = Arc::new(EmptyExec::new(provider.schema())) as Arc<dyn ExecutionPlan>;
730 let error = provider
731 .insert_into(&ctx.state(), input, insert_op)
732 .await
733 .expect_err("non-append inserts should be rejected");
734
735 assert!(
736 matches!(
737 error,
738 DataFusionError::NotImplemented(ref message) if message == expected_message
739 ),
740 "unexpected error: {error}"
741 );
742 }
743 }
744
745 #[tokio::test]
746 async fn test_insert_plan_fanout_enabled_no_sort() {
747 use datafusion::datasource::TableProvider;
748 use datafusion::logical_expr::dml::InsertOp;
749 use datafusion::physical_plan::empty::EmptyExec;
750
751 let (catalog, namespace, table_name, _temp_dir) =
753 get_partitioned_test_catalog_and_table(Some(true)).await;
754
755 let provider =
756 IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), table_name.clone())
757 .await
758 .unwrap();
759
760 let ctx = SessionContext::new();
761 let input_schema = provider.schema();
762 let input = Arc::new(EmptyExec::new(input_schema)) as Arc<dyn ExecutionPlan>;
763
764 let state = ctx.state();
765 let insert_plan = provider
766 .insert_into(&state, input, InsertOp::Append)
767 .await
768 .unwrap();
769
770 assert!(
772 !plan_contains_sort(&insert_plan),
773 "Plan should NOT contain SortExec when fanout is enabled"
774 );
775 }
776
777 #[tokio::test]
778 async fn test_insert_plan_fanout_disabled_has_sort() {
779 use datafusion::datasource::TableProvider;
780 use datafusion::logical_expr::dml::InsertOp;
781 use datafusion::physical_plan::empty::EmptyExec;
782
783 let (catalog, namespace, table_name, _temp_dir) =
785 get_partitioned_test_catalog_and_table(Some(false)).await;
786
787 let provider =
788 IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), table_name.clone())
789 .await
790 .unwrap();
791
792 let ctx = SessionContext::new();
793 let input_schema = provider.schema();
794 let input = Arc::new(EmptyExec::new(input_schema)) as Arc<dyn ExecutionPlan>;
795
796 let state = ctx.state();
797 let insert_plan = provider
798 .insert_into(&state, input, InsertOp::Append)
799 .await
800 .unwrap();
801
802 assert!(
804 plan_contains_sort(&insert_plan),
805 "Plan should contain SortExec when fanout is disabled"
806 );
807 }
808
809 #[tokio::test]
810 async fn test_limit_pushdown_static_provider() {
811 use datafusion::datasource::TableProvider;
812
813 let table = get_test_table_from_metadata_file().await;
814 let table_provider = IcebergStaticTableProvider::try_new_from_table(table.clone())
815 .await
816 .unwrap();
817
818 let ctx = SessionContext::new();
819 let state = ctx.state();
820
821 let scan_plan = table_provider
823 .scan(&state, None, &[], Some(10))
824 .await
825 .unwrap();
826
827 let iceberg_scan = scan_plan
829 .downcast_ref::<IcebergTableScan>()
830 .expect("Expected IcebergTableScan");
831
832 assert_eq!(
834 iceberg_scan.limit(),
835 Some(10),
836 "Limit should be set to 10 in the scan plan"
837 );
838 }
839
840 #[tokio::test]
841 async fn test_limit_pushdown_catalog_backed_provider() {
842 use datafusion::datasource::TableProvider;
843
844 let (catalog, namespace, table_name, _temp_dir) = get_test_catalog_and_table().await;
845
846 let provider =
847 IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), table_name.clone())
848 .await
849 .unwrap();
850
851 let ctx = SessionContext::new();
852 let state = ctx.state();
853
854 let scan_plan = provider.scan(&state, None, &[], Some(5)).await.unwrap();
856
857 let iceberg_scan = scan_plan
859 .downcast_ref::<IcebergTableScan>()
860 .expect("Expected IcebergTableScan");
861
862 assert_eq!(
864 iceberg_scan.limit(),
865 Some(5),
866 "Limit should be set to 5 in the scan plan"
867 );
868 }
869
870 #[tokio::test]
871 async fn test_no_limit_pushdown() {
872 use datafusion::datasource::TableProvider;
873
874 let table = get_test_table_from_metadata_file().await;
875 let table_provider = IcebergStaticTableProvider::try_new_from_table(table.clone())
876 .await
877 .unwrap();
878
879 let ctx = SessionContext::new();
880 let state = ctx.state();
881
882 let scan_plan = table_provider.scan(&state, None, &[], None).await.unwrap();
884
885 let iceberg_scan = scan_plan
887 .downcast_ref::<IcebergTableScan>()
888 .expect("Expected IcebergTableScan");
889
890 assert_eq!(
892 iceberg_scan.limit(),
893 None,
894 "Limit should be None when not specified"
895 );
896 }
897}