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(table.clone(), write_input));
222
223 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#[derive(Debug, Clone)]
244pub struct IcebergStaticTableProvider {
245 table: Table,
247 snapshot_id: Option<i64>,
249 schema: ArrowSchemaRef,
251}
252
253impl IcebergStaticTableProvider {
254 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 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 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 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 #[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 let result = ctx.sql("INSERT INTO mytable VALUES (1, 2, 3)").await;
488
489 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 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 #[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 let provider =
523 IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), table_name.clone())
524 .await
525 .unwrap();
526
527 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 let df = ctx.sql("SELECT * FROM test_table").await.unwrap();
549
550 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 let result = ctx.sql("INSERT INTO test_table VALUES (1, 'test')").await;
575
576 assert!(result.is_ok());
578
579 let df = result.unwrap();
581 let execution_result = df.collect().await;
582
583 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 let df = ctx.sql("SELECT id, name FROM test_table").await.unwrap();
602
603 let logical_schema = df.schema().clone();
605
606 let physical_plan = df.create_physical_plan().await.unwrap();
608 let physical_schema = physical_plan.schema();
609
610 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 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 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 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 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 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 let scan_plan = table_provider
819 .scan(&state, None, &[], Some(10))
820 .await
821 .unwrap();
822
823 let iceberg_scan = scan_plan
825 .downcast_ref::<IcebergTableScan>()
826 .expect("Expected IcebergTableScan");
827
828 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 let scan_plan = provider.scan(&state, None, &[], Some(5)).await.unwrap();
852
853 let iceberg_scan = scan_plan
855 .downcast_ref::<IcebergTableScan>()
856 .expect("Expected IcebergTableScan");
857
858 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 let scan_plan = table_provider.scan(&state, None, &[], None).await.unwrap();
880
881 let iceberg_scan = scan_plan
883 .downcast_ref::<IcebergTableScan>()
884 .expect("Expected IcebergTableScan");
885
886 assert_eq!(
888 iceberg_scan.limit(),
889 None,
890 "Limit should be None when not specified"
891 );
892 }
893}