vendor/iceberg-datafusion-arrow59/src/table/mod.rs
// 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.
//! Iceberg table providers for DataFusion.
//!
//! This module provides two table provider implementations:
//!
//! - [`IcebergTableProvider`]: Catalog-backed provider with automatic metadata refresh.
//! Use for write operations and when you need to see the latest table state.
//!
//! - [`IcebergStaticTableProvider`]: Static provider for read-only access to a specific
//! table snapshot. Use for consistent analytical queries or time-travel scenarios.
pub mod metadata_table;
pub mod table_provider_factory;
use std::num::NonZeroUsize;
use std::sync::Arc;
use async_trait::async_trait;
use datafusion::arrow::datatypes::SchemaRef as ArrowSchemaRef;
use datafusion::catalog::Session;
use datafusion::common::{
DFSchemaRef, DataFusionError, config_datafusion_err, exec_datafusion_err, not_impl_err,
};
use datafusion::datasource::{TableProvider, TableType};
use datafusion::error::Result;
use datafusion::logical_expr::dml::{InsertOp, MergeIntoClause};
use datafusion::logical_expr::{Expr, TableProviderFilterPushDown};
use datafusion::physical_plan::ExecutionPlan;
use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec;
use iceberg::arrow::schema_to_arrow_schema;
use iceberg::inspect::MetadataTableType;
use iceberg::spec::TableProperties;
use iceberg::table::Table;
use iceberg::{Catalog, Error, ErrorKind, NamespaceIdent, TableIdent};
use metadata_table::IcebergMetadataTableProvider;
use crate::error::to_datafusion_error;
use crate::physical_plan::commit::IcebergCommitExec;
use crate::physical_plan::merge::{IcebergMergeExec, compile_clauses, equi_keys_from_on};
use crate::physical_plan::project::project_with_partition;
use crate::physical_plan::repartition::repartition;
use crate::physical_plan::scan::IcebergTableScan;
use crate::physical_plan::sort::sort_by_partition;
use crate::physical_plan::write::IcebergWriteExec;
/// Catalog-backed table provider with automatic metadata refresh.
///
/// This provider loads fresh table metadata from the catalog on every scan and write
/// operation, ensuring you always see the latest table state. Use this when you need
/// write operations or want to see the most up-to-date data.
///
/// For read-only access to a specific snapshot without catalog overhead, use
/// [`IcebergStaticTableProvider`] instead.
#[derive(Debug, Clone)]
pub struct IcebergTableProvider {
/// The catalog that manages this table
catalog: Arc<dyn Catalog>,
/// The table identifier (namespace + name)
table_ident: TableIdent,
/// A reference-counted arrow `Schema` (cached at construction)
schema: ArrowSchemaRef,
}
impl IcebergTableProvider {
/// Creates a new catalog-backed table provider.
///
/// Loads the table once to get the initial schema, then stores the catalog
/// reference for future metadata refreshes on each operation.
pub(crate) async fn try_new(
catalog: Arc<dyn Catalog>,
namespace: NamespaceIdent,
name: impl Into<String>,
) -> Result<Self> {
let table_ident = TableIdent::new(namespace, name.into());
// Load table once to get initial schema
let table = catalog
.load_table(&table_ident)
.await
.map_err(to_datafusion_error)?;
let schema = Arc::new(
schema_to_arrow_schema(table.metadata().current_schema())
.map_err(to_datafusion_error)?,
);
Ok(IcebergTableProvider {
catalog,
table_ident,
schema,
})
}
pub(crate) async fn metadata_table(
&self,
r#type: MetadataTableType,
) -> Result<IcebergMetadataTableProvider> {
// Load fresh table metadata for metadata table access
let table = self
.catalog
.load_table(&self.table_ident)
.await
.map_err(to_datafusion_error)?;
Ok(IcebergMetadataTableProvider { table, r#type })
}
}
#[async_trait]
impl TableProvider for IcebergTableProvider {
fn schema(&self) -> ArrowSchemaRef {
self.schema.clone()
}
fn table_type(&self) -> TableType {
TableType::Base
}
async fn scan(
&self,
_state: &dyn Session,
projection: Option<&Vec<usize>>,
filters: &[Expr],
limit: Option<usize>,
) -> Result<Arc<dyn ExecutionPlan>> {
// Load fresh table metadata from catalog
let table = self
.catalog
.load_table(&self.table_ident)
.await
.map_err(to_datafusion_error)?;
// Create scan with fresh metadata (always use current snapshot)
// Plans the file scan tasks here (async), one DataFusion partition per
// file — see physical_plan/scan.rs.
Ok(Arc::new(
IcebergTableScan::new(
table,
None, // Always use current snapshot for catalog-backed provider
self.schema.clone(),
projection.map(|v| v.as_slice()),
filters,
limit,
)
.await?,
))
}
fn supports_filters_pushdown(
&self,
filters: &[&Expr],
) -> Result<Vec<TableProviderFilterPushDown>> {
// Push down all filters, as a single source of truth, the scanner will drop the filters which couldn't be push down
Ok(vec![TableProviderFilterPushDown::Inexact; filters.len()])
}
async fn insert_into(
&self,
state: &dyn Session,
input: Arc<dyn ExecutionPlan>,
_insert_op: InsertOp,
) -> Result<Arc<dyn ExecutionPlan>> {
if _insert_op != InsertOp::Append {
return not_impl_err!(
"IcebergTableProvider supports only append inserts, got {_insert_op}"
);
}
// Load fresh table metadata from catalog
let table = self
.catalog
.load_table(&self.table_ident)
.await
.map_err(to_datafusion_error)?;
let partition_spec = table.metadata().default_partition_spec();
// Step 1: Project partition values for partitioned tables
let plan_with_partition = if !partition_spec.is_unpartitioned() {
project_with_partition(input, &table)?
} else {
input
};
// Step 2: Repartition for parallel processing
let target_partitions = NonZeroUsize::new(state.config().target_partitions())
.ok_or_else(|| {
config_datafusion_err!("target_partitions must be greater than 0")
})?;
let repartitioned_plan =
repartition(plan_with_partition, table.metadata_ref(), target_partitions)?;
// Apply sort node when it's not fanout mode
let fanout_enabled = table
.metadata()
.properties()
.get(TableProperties::PROPERTY_DATAFUSION_WRITE_FANOUT_ENABLED)
.map(|value| {
value.parse::<bool>().map_err(|e| {
config_datafusion_err!(
"Invalid value for {}, expected 'true' or 'false': {e}",
TableProperties::PROPERTY_DATAFUSION_WRITE_FANOUT_ENABLED
)
})
})
.transpose()?
.unwrap_or(TableProperties::PROPERTY_DATAFUSION_WRITE_FANOUT_ENABLED_DEFAULT);
let write_input = if fanout_enabled {
repartitioned_plan
} else {
sort_by_partition(repartitioned_plan)?
};
let write_plan = Arc::new(IcebergWriteExec::new(table.clone(), write_input));
// Merge the outputs of write_plan into one so we can commit all files together
let coalesce_partitions = Arc::new(CoalescePartitionsExec::new(write_plan));
Ok(Arc::new(IcebergCommitExec::new(
table,
self.catalog.clone(),
coalesce_partitions,
self.schema.clone(),
)))
}
/// `MERGE INTO` — the row-level write path, by **copy-on-write**.
///
/// DataFusion plans the statement and hands this method the combined
/// target-then-source schema, the `ON` condition and the `WHEN` clauses;
/// see [`crate::physical_plan::merge`] for what happens to them. Every
/// expression is compiled to a [`PhysicalExpr`] HERE, against `state`,
/// because that is the only point where a user-defined function inside a
/// clause can still be resolved — the returned plan runs with only a
/// `TaskContext`.
///
/// [`PhysicalExpr`]: datafusion::physical_expr::PhysicalExpr
async fn merge_into(
&self,
state: &dyn Session,
source: Arc<dyn ExecutionPlan>,
merge_schema: DFSchemaRef,
on: Expr,
clauses: Vec<MergeIntoClause>,
) -> Result<Arc<dyn ExecutionPlan>> {
let table = self
.catalog
.load_table(&self.table_ident)
.await
.map_err(to_datafusion_error)?;
// A rewritten row needs its partition values recomputed, and the
// projection that does that (`project_with_partition`) sits on the
// INSERT path, ahead of a writer this path does not use. Refuse by
// name rather than write files into the wrong partition.
if !table.metadata().default_partition_spec().is_unpartitioned() {
return Err(to_datafusion_error(Error::new(
ErrorKind::FeatureUnsupported,
format!(
"MERGE INTO is not supported on the partitioned table `{}`: a rewritten row \
needs its partition values recomputed, which this path does not do.",
self.table_ident
),
)));
}
let n_target = self.schema.fields().len();
if merge_schema.fields().len() != n_target + source.schema().fields().len() {
return Err(DataFusionError::Internal(format!(
"MERGE combined schema has {} fields, expected {} target + {} source",
merge_schema.fields().len(),
n_target,
source.schema().fields().len()
)));
}
let equi = equi_keys_from_on(&on, &merge_schema, n_target)?;
let compiled = compile_clauses(clauses, &self.schema, |e| {
state.create_physical_expr(e, merge_schema.as_ref())
})?;
Ok(Arc::new(IcebergMergeExec::new(
self.catalog.clone(),
self.table_ident.clone(),
self.schema.clone(),
merge_schema.inner().clone(),
source,
equi,
compiled,
)))
}
}
/// Static table provider for read-only snapshot access.
///
/// This provider holds a cached table instance and does not refresh metadata or support
/// write operations. Use this for consistent analytical queries, time-travel scenarios,
/// or when you want to avoid catalog overhead.
///
/// For catalog-backed tables with write support and automatic refresh, use
/// [`IcebergTableProvider`] instead.
#[derive(Debug, Clone)]
pub struct IcebergStaticTableProvider {
/// The static table instance (never refreshed)
table: Table,
/// Optional snapshot ID for this static view
snapshot_id: Option<i64>,
/// A reference-counted arrow `Schema`
schema: ArrowSchemaRef,
}
impl IcebergStaticTableProvider {
/// Creates a static provider from a table instance.
///
/// Uses the table's current snapshot for all queries. Does not support write operations.
pub async fn try_new_from_table(table: Table) -> Result<Self> {
let schema = Arc::new(
schema_to_arrow_schema(table.metadata().current_schema())
.map_err(to_datafusion_error)?,
);
Ok(IcebergStaticTableProvider {
table,
snapshot_id: None,
schema,
})
}
/// Creates a static provider for a specific table snapshot.
///
/// Queries the specified snapshot for all operations. Useful for time-travel queries.
/// Does not support write operations.
pub async fn try_new_from_table_snapshot(
table: Table,
snapshot_id: i64,
) -> Result<Self> {
let snapshot = table
.metadata()
.snapshot_by_id(snapshot_id)
.ok_or_else(|| {
exec_datafusion_err!(
"snapshot id {snapshot_id} not found in table {}",
table.identifier().name()
)
})?;
let table_schema = snapshot
.schema(table.metadata())
.map_err(to_datafusion_error)?;
let schema =
Arc::new(schema_to_arrow_schema(&table_schema).map_err(to_datafusion_error)?);
Ok(IcebergStaticTableProvider {
table,
snapshot_id: Some(snapshot_id),
schema,
})
}
}
#[async_trait]
impl TableProvider for IcebergStaticTableProvider {
fn schema(&self) -> ArrowSchemaRef {
self.schema.clone()
}
fn table_type(&self) -> TableType {
TableType::Base
}
async fn scan(
&self,
_state: &dyn Session,
projection: Option<&Vec<usize>>,
filters: &[Expr],
limit: Option<usize>,
) -> Result<Arc<dyn ExecutionPlan>> {
// Use cached table (no refresh)
// Plans the file scan tasks here (async), one DataFusion partition per
// file — see physical_plan/scan.rs.
Ok(Arc::new(
IcebergTableScan::new(
self.table.clone(),
self.snapshot_id,
self.schema.clone(),
projection.map(|v| v.as_slice()),
filters,
limit,
)
.await?,
))
}
fn supports_filters_pushdown(
&self,
filters: &[&Expr],
) -> Result<Vec<TableProviderFilterPushDown>> {
// Push down all filters, as a single source of truth, the scanner will drop the filters which couldn't be push down
Ok(vec![TableProviderFilterPushDown::Inexact; filters.len()])
}
async fn insert_into(
&self,
_state: &dyn Session,
_input: Arc<dyn ExecutionPlan>,
_insert_op: InsertOp,
) -> Result<Arc<dyn ExecutionPlan>> {
not_impl_err!(
"Write operations are not supported on IcebergStaticTableProvider. \
Use IcebergTableProvider with a catalog for write support."
)
}
}
#[cfg(test)]
mod tests {
use std::collections::HashMap;
use std::sync::Arc;
use datafusion::common::Column;
use datafusion::error::DataFusionError;
use datafusion::physical_plan::ExecutionPlan;
use datafusion::prelude::SessionContext;
use iceberg::io::FileIO;
use iceberg::memory::{MEMORY_CATALOG_WAREHOUSE, MemoryCatalogBuilder};
use iceberg::spec::{NestedField, PrimitiveType, Schema, Type};
use iceberg::table::{StaticTable, Table};
use iceberg::{Catalog, CatalogBuilder, NamespaceIdent, TableCreation, TableIdent};
use tempfile::TempDir;
use super::*;
async fn get_test_table_from_metadata_file() -> Table {
let metadata_file_name = "TableMetadataV2Valid.json";
let metadata_file_path = format!(
"{}/tests/test_data/{}",
env!("CARGO_MANIFEST_DIR"),
metadata_file_name
);
let file_io = FileIO::new_with_fs();
let static_identifier =
TableIdent::from_strs(["static_ns", "static_table"]).unwrap();
let static_table = StaticTable::from_metadata_file(
&metadata_file_path,
static_identifier,
file_io,
)
.await
.unwrap();
static_table.into_table()
}
/// A real, empty (no-snapshot) table loaded from a memory catalog. Unlike
/// `get_test_table_from_metadata_file` (whose `TableMetadataV2Valid.json`
/// fixture references manifest lists like `s3://a/b/2.avro` that do not
/// exist on disk), this table is actually scannable: `IcebergTableScan::new`
/// plans file tasks eagerly (skade's per-file partitioned scan), so the
/// static-provider scan tests need a table whose manifests can be read. The
/// returned `TempDir` must be kept alive for the warehouse to stay valid.
async fn get_test_scannable_static_table() -> (Table, TempDir) {
let (catalog, namespace, table_name, temp_dir) = get_test_catalog_and_table().await;
let ident = TableIdent::new(namespace, table_name);
let table = catalog.load_table(&ident).await.unwrap();
(table, temp_dir)
}
async fn get_test_catalog_and_table()
-> (Arc<dyn Catalog>, NamespaceIdent, String, TempDir) {
let temp_dir = TempDir::new().unwrap();
let warehouse_path = temp_dir.path().to_str().unwrap().to_string();
let catalog = MemoryCatalogBuilder::default()
.load(
"memory",
HashMap::from([(
MEMORY_CATALOG_WAREHOUSE.to_string(),
warehouse_path.clone(),
)]),
)
.await
.unwrap();
let namespace = NamespaceIdent::new("test_ns".to_string());
catalog
.create_namespace(&namespace, HashMap::new())
.await
.unwrap();
let schema = Schema::builder()
.with_schema_id(0)
.with_fields(vec![
NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int))
.into(),
NestedField::required(2, "name", Type::Primitive(PrimitiveType::String))
.into(),
])
.build()
.unwrap();
let table_creation = TableCreation::builder()
.name("test_table".to_string())
.location(format!("{warehouse_path}/test_table"))
.schema(schema)
.properties(HashMap::new())
.build();
catalog
.create_table(&namespace, table_creation)
.await
.unwrap();
(
Arc::new(catalog),
namespace,
"test_table".to_string(),
temp_dir,
)
}
// Tests for IcebergStaticTableProvider
#[tokio::test]
async fn test_static_provider_from_table() {
let table = get_test_table_from_metadata_file().await;
let table_provider =
IcebergStaticTableProvider::try_new_from_table(table.clone())
.await
.unwrap();
let ctx = SessionContext::new();
ctx.register_table("mytable", Arc::new(table_provider))
.unwrap();
let df = ctx.sql("SELECT * FROM mytable").await.unwrap();
let df_schema = df.schema();
let df_columns = df_schema.fields();
assert_eq!(df_columns.len(), 3);
let x_column = df_columns.first().unwrap();
let column_data = format!(
"{:?}:{:?}",
x_column.name(),
x_column.data_type().to_string()
);
assert_eq!(column_data, "\"x\":\"Int64\"");
let has_column = df_schema.has_column(&Column::from_name("z"));
assert!(has_column);
}
#[tokio::test]
async fn test_static_provider_from_snapshot() {
let table = get_test_table_from_metadata_file().await;
let snapshot_id = table.metadata().snapshots().next().unwrap().snapshot_id();
let table_provider = IcebergStaticTableProvider::try_new_from_table_snapshot(
table.clone(),
snapshot_id,
)
.await
.unwrap();
let ctx = SessionContext::new();
ctx.register_table("mytable", Arc::new(table_provider))
.unwrap();
let df = ctx.sql("SELECT * FROM mytable").await.unwrap();
let df_schema = df.schema();
let df_columns = df_schema.fields();
assert_eq!(df_columns.len(), 3);
let x_column = df_columns.first().unwrap();
let column_data = format!(
"{:?}:{:?}",
x_column.name(),
x_column.data_type().to_string()
);
assert_eq!(column_data, "\"x\":\"Int64\"");
let has_column = df_schema.has_column(&Column::from_name("z"));
assert!(has_column);
}
#[tokio::test]
async fn test_static_provider_rejects_writes() {
let table = get_test_table_from_metadata_file().await;
let table_provider =
IcebergStaticTableProvider::try_new_from_table(table.clone())
.await
.unwrap();
let ctx = SessionContext::new();
ctx.register_table("mytable", Arc::new(table_provider))
.unwrap();
// Attempt to insert into the static provider should fail
let result = ctx.sql("INSERT INTO mytable VALUES (1, 2, 3)").await;
// The error should occur during planning or execution
// We expect an error indicating write operations are not supported
assert!(
result.is_err() || {
let df = result.unwrap();
df.collect().await.is_err()
}
);
}
/// **The scan tells DataFusion how many rows it holds.** Every physical
/// optimiser decision that weighs one side against another — hash-join
/// build/probe selection above all — reads `Statistics`, and this scan
/// returned `new_unknown` for everything. MEASURED (skade-flight, TPC-H
/// SF=1, 2026-09-10): Q17 built its hash table on the 6 001 215-row
/// lineitem and probed with the ~2 000-row filtered part, because with both
/// sides unknown JoinSelection cannot swap.
///
/// RED direction: delete `statistics_from_inputs` from `IcebergTableScan`
/// and `num_rows` is `Absent`.
#[tokio::test]
async fn the_scan_reports_row_and_byte_counts_from_the_planned_files() {
use datafusion::common::stats::Precision;
use datafusion::physical_plan::statistics::{StatisticsArgs, StatisticsContext};
let (catalog, namespace, table_name, _temp_dir) = get_test_catalog_and_table().await;
let provider =
IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), table_name.clone())
.await
.unwrap();
let ctx = SessionContext::new();
ctx.register_table("test_table", Arc::new(provider)).unwrap();
// Two commits, so the scan plans two files and has something to sum.
ctx.sql("INSERT INTO test_table VALUES (1, 'a'), (2, 'b'), (3, 'c')")
.await
.unwrap()
.collect()
.await
.unwrap();
ctx.sql("INSERT INTO test_table VALUES (4, 'd'), (5, 'e')")
.await
.unwrap()
.collect()
.await
.unwrap();
let plan = ctx
.sql("SELECT * FROM test_table")
.await
.unwrap()
.create_physical_plan()
.await
.unwrap();
// Whole plan, unfiltered: the file record counts are the truth, exactly.
let stats = StatisticsContext::new().compute(plan.as_ref(), &StatisticsArgs::new()).unwrap();
assert_eq!(stats.num_rows, Precision::Exact(5), "{:?}", stats.num_rows);
// Deliberately unknown: the Parquet length on disk is not the Arrow
// output size, and reporting it as such was measured to mislead the
// hash-join mode selection (see `statistics_from_inputs`).
assert_eq!(stats.total_byte_size, Precision::Absent, "{:?}", stats.total_byte_size);
// Column statistics, folded from the manifest across the two files:
// id 1..=5 → min 1, max 5, no nulls, and a distinct ceiling of 5 for
// an integer column (the number the join estimator reads).
let id = &stats.column_statistics[0];
assert_eq!(id.min_value, Precision::Exact(datafusion::common::ScalarValue::Int32(Some(1))), "{:?}", id.min_value);
assert_eq!(id.max_value, Precision::Exact(datafusion::common::ScalarValue::Int32(Some(5))), "{:?}", id.max_value);
assert_eq!(id.null_count, Precision::Exact(0));
assert_eq!(id.distinct_count, Precision::Inexact(5), "{:?}", id.distinct_count);
let name = &stats.column_statistics[1];
assert_eq!(name.min_value, Precision::Exact(datafusion::common::ScalarValue::Utf8(Some("a".into()))), "{:?}", name.min_value);
assert_eq!(name.distinct_count, Precision::Absent, "no range ceiling for a string");
// A filtered scan: the files still say 5, but a residual filter means
// the count is a ceiling, not the answer.
let plan = ctx
.sql("SELECT * FROM test_table WHERE id > 3")
.await
.unwrap()
.create_physical_plan()
.await
.unwrap();
// The scan is under a FilterExec; ask the scan itself.
fn find_scan(p: &Arc<dyn ExecutionPlan>) -> Option<Arc<dyn ExecutionPlan>> {
if p.name() == "IcebergTableScan" {
return Some(Arc::clone(p));
}
p.children().into_iter().find_map(find_scan)
}
let scan = find_scan(&plan).expect("an IcebergTableScan under the filter");
let stats = StatisticsContext::new().compute(scan.as_ref(), &StatisticsArgs::new()).unwrap();
assert!(matches!(stats.num_rows, Precision::Inexact(n) if n <= 5 && n > 0), "{:?}", stats.num_rows);
}
#[tokio::test]
async fn test_static_provider_scan() {
let (table, _temp_dir) = get_test_scannable_static_table().await;
let table_provider =
IcebergStaticTableProvider::try_new_from_table(table.clone())
.await
.unwrap();
let ctx = SessionContext::new();
ctx.register_table("mytable", Arc::new(table_provider))
.unwrap();
// Test that scan operations work correctly
let df = ctx.sql("SELECT count(*) FROM mytable").await.unwrap();
let physical_plan = df.create_physical_plan().await;
assert!(physical_plan.is_ok());
}
// Tests for IcebergTableProvider
#[tokio::test]
async fn test_catalog_backed_provider_creation() {
let (catalog, namespace, table_name, _temp_dir) =
get_test_catalog_and_table().await;
// Test creating a catalog-backed provider
let provider = IcebergTableProvider::try_new(
catalog.clone(),
namespace.clone(),
table_name.clone(),
)
.await
.unwrap();
// Verify the schema is loaded correctly
let schema = provider.schema();
assert_eq!(schema.fields().len(), 2);
assert_eq!(schema.field(0).name(), "id");
assert_eq!(schema.field(1).name(), "name");
}
#[tokio::test]
async fn test_catalog_backed_provider_scan() {
let (catalog, namespace, table_name, _temp_dir) =
get_test_catalog_and_table().await;
let provider = IcebergTableProvider::try_new(
catalog.clone(),
namespace.clone(),
table_name.clone(),
)
.await
.unwrap();
let ctx = SessionContext::new();
ctx.register_table("test_table", Arc::new(provider))
.unwrap();
// Test that scan operations work correctly
let df = ctx.sql("SELECT * FROM test_table").await.unwrap();
// Verify the schema in the query result
let df_schema = df.schema();
assert_eq!(df_schema.fields().len(), 2);
assert_eq!(df_schema.field(0).name(), "id");
assert_eq!(df_schema.field(1).name(), "name");
let physical_plan = df.create_physical_plan().await;
assert!(physical_plan.is_ok());
}
#[tokio::test]
async fn test_catalog_backed_provider_insert() {
let (catalog, namespace, table_name, _temp_dir) =
get_test_catalog_and_table().await;
let provider = IcebergTableProvider::try_new(
catalog.clone(),
namespace.clone(),
table_name.clone(),
)
.await
.unwrap();
let ctx = SessionContext::new();
ctx.register_table("test_table", Arc::new(provider))
.unwrap();
// Test that insert operations work correctly
let result = ctx.sql("INSERT INTO test_table VALUES (1, 'test')").await;
// Insert should succeed (or at least not fail during planning)
assert!(result.is_ok());
// Try to execute the insert plan
let df = result.unwrap();
let execution_result = df.collect().await;
// The execution should succeed
assert!(execution_result.is_ok());
}
#[tokio::test]
async fn test_physical_input_schema_consistent_with_logical_input_schema() {
let (catalog, namespace, table_name, _temp_dir) =
get_test_catalog_and_table().await;
let provider = IcebergTableProvider::try_new(
catalog.clone(),
namespace.clone(),
table_name.clone(),
)
.await
.unwrap();
let ctx = SessionContext::new();
ctx.register_table("test_table", Arc::new(provider))
.unwrap();
// Create a query plan
let df = ctx.sql("SELECT id, name FROM test_table").await.unwrap();
// Get logical schema before consuming df
let logical_schema = df.schema().clone();
// Get physical plan (this consumes df)
let physical_plan = df.create_physical_plan().await.unwrap();
let physical_schema = physical_plan.schema();
// Verify that logical and physical schemas are consistent
assert_eq!(
logical_schema.fields().len(),
physical_schema.fields().len()
);
for (logical_field, physical_field) in logical_schema
.fields()
.iter()
.zip(physical_schema.fields().iter())
{
assert_eq!(logical_field.name(), physical_field.name());
assert_eq!(logical_field.data_type(), physical_field.data_type());
}
}
async fn get_partitioned_test_catalog_and_table(
fanout_enabled: Option<bool>,
) -> (Arc<dyn Catalog>, NamespaceIdent, String, TempDir) {
use iceberg::spec::Transform;
let temp_dir = TempDir::new().unwrap();
let warehouse_path = temp_dir.path().to_str().unwrap().to_string();
let catalog = MemoryCatalogBuilder::default()
.load(
"memory",
HashMap::from([(
MEMORY_CATALOG_WAREHOUSE.to_string(),
warehouse_path.clone(),
)]),
)
.await
.unwrap();
let namespace = NamespaceIdent::new("test_ns".to_string());
catalog
.create_namespace(&namespace, HashMap::new())
.await
.unwrap();
let schema = Schema::builder()
.with_schema_id(0)
.with_fields(vec![
NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int))
.into(),
NestedField::required(
2,
"category",
Type::Primitive(PrimitiveType::String),
)
.into(),
])
.build()
.unwrap();
// iceberg-rust main deab069 made `UnboundPartitionField`'s builder
// crate-private; bind the spec by column name and unbind it instead.
let partition_spec = iceberg::spec::PartitionSpec::builder(schema.clone())
.with_spec_id(0)
.add_partition_field("category", "category", Transform::Identity)
.unwrap()
.build()
.unwrap()
.into_unbound();
let mut properties = HashMap::new();
if let Some(enabled) = fanout_enabled {
properties.insert(
TableProperties::PROPERTY_DATAFUSION_WRITE_FANOUT_ENABLED.to_string(),
enabled.to_string(),
);
}
let table_creation = TableCreation::builder()
.name("partitioned_table".to_string())
.location(format!("{warehouse_path}/partitioned_table"))
.schema(schema)
.partition_spec(partition_spec)
.properties(properties)
.build();
catalog
.create_table(&namespace, table_creation)
.await
.unwrap();
(
Arc::new(catalog),
namespace,
"partitioned_table".to_string(),
temp_dir,
)
}
/// Helper to check if a plan contains a SortExec node
fn plan_contains_sort(plan: &Arc<dyn ExecutionPlan>) -> bool {
if plan.name() == "SortExec" {
return true;
}
for child in plan.children() {
if plan_contains_sort(child) {
return true;
}
}
false
}
#[tokio::test]
async fn test_catalog_backed_provider_rejects_non_append_op() {
use datafusion::physical_plan::empty::EmptyExec;
let (catalog, namespace, table_name, _temp_dir) =
get_test_catalog_and_table().await;
let provider = IcebergTableProvider::try_new(catalog, namespace, table_name)
.await
.unwrap();
let ctx = SessionContext::new();
for (insert_op, expected_message) in [
(
InsertOp::Overwrite,
"IcebergTableProvider supports only append inserts, got Insert Overwrite",
),
(
InsertOp::Replace,
"IcebergTableProvider supports only append inserts, got Replace Into",
),
] {
let input =
Arc::new(EmptyExec::new(provider.schema())) as Arc<dyn ExecutionPlan>;
let error = provider
.insert_into(&ctx.state(), input, insert_op)
.await
.expect_err("non-append inserts should be rejected");
assert!(
matches!(
error,
DataFusionError::NotImplemented(ref message) if message == expected_message
),
"unexpected error: {error}"
);
}
}
#[tokio::test]
async fn test_insert_plan_fanout_enabled_no_sort() {
use datafusion::datasource::TableProvider;
use datafusion::logical_expr::dml::InsertOp;
use datafusion::physical_plan::empty::EmptyExec;
// When fanout is enabled (default), no sort node should be added
let (catalog, namespace, table_name, _temp_dir) =
get_partitioned_test_catalog_and_table(Some(true)).await;
let provider = IcebergTableProvider::try_new(
catalog.clone(),
namespace.clone(),
table_name.clone(),
)
.await
.unwrap();
let ctx = SessionContext::new();
let input_schema = provider.schema();
let input = Arc::new(EmptyExec::new(input_schema)) as Arc<dyn ExecutionPlan>;
let state = ctx.state();
let insert_plan = provider
.insert_into(&state, input, InsertOp::Append)
.await
.unwrap();
// With fanout enabled, there should be no SortExec in the plan
assert!(
!plan_contains_sort(&insert_plan),
"Plan should NOT contain SortExec when fanout is enabled"
);
}
#[tokio::test]
async fn test_insert_plan_fanout_disabled_has_sort() {
use datafusion::datasource::TableProvider;
use datafusion::logical_expr::dml::InsertOp;
use datafusion::physical_plan::empty::EmptyExec;
// When fanout is disabled, a sort node should be added
let (catalog, namespace, table_name, _temp_dir) =
get_partitioned_test_catalog_and_table(Some(false)).await;
let provider = IcebergTableProvider::try_new(
catalog.clone(),
namespace.clone(),
table_name.clone(),
)
.await
.unwrap();
let ctx = SessionContext::new();
let input_schema = provider.schema();
let input = Arc::new(EmptyExec::new(input_schema)) as Arc<dyn ExecutionPlan>;
let state = ctx.state();
let insert_plan = provider
.insert_into(&state, input, InsertOp::Append)
.await
.unwrap();
// With fanout disabled, there should be a SortExec in the plan
assert!(
plan_contains_sort(&insert_plan),
"Plan should contain SortExec when fanout is disabled"
);
}
#[tokio::test]
async fn test_limit_pushdown_static_provider() {
use datafusion::datasource::TableProvider;
let (table, _temp_dir) = get_test_scannable_static_table().await;
let table_provider =
IcebergStaticTableProvider::try_new_from_table(table.clone())
.await
.unwrap();
let ctx = SessionContext::new();
let state = ctx.state();
// Test scan with limit
let scan_plan = table_provider
.scan(&state, None, &[], Some(10))
.await
.unwrap();
// Verify that the scan plan is an IcebergTableScan
let iceberg_scan = scan_plan
.downcast_ref::<IcebergTableScan>()
.expect("Expected IcebergTableScan");
// Verify the limit is set
assert_eq!(
iceberg_scan.limit(),
Some(10),
"Limit should be set to 10 in the scan plan"
);
}
#[tokio::test]
async fn test_limit_pushdown_catalog_backed_provider() {
use datafusion::datasource::TableProvider;
let (catalog, namespace, table_name, _temp_dir) =
get_test_catalog_and_table().await;
let provider = IcebergTableProvider::try_new(
catalog.clone(),
namespace.clone(),
table_name.clone(),
)
.await
.unwrap();
let ctx = SessionContext::new();
let state = ctx.state();
// Test scan with limit
let scan_plan = provider.scan(&state, None, &[], Some(5)).await.unwrap();
// Verify that the scan plan is an IcebergTableScan
let iceberg_scan = scan_plan
.downcast_ref::<IcebergTableScan>()
.expect("Expected IcebergTableScan");
// Verify the limit is set
assert_eq!(
iceberg_scan.limit(),
Some(5),
"Limit should be set to 5 in the scan plan"
);
}
#[tokio::test]
async fn test_no_limit_pushdown() {
use datafusion::datasource::TableProvider;
let (table, _temp_dir) = get_test_scannable_static_table().await;
let table_provider =
IcebergStaticTableProvider::try_new_from_table(table.clone())
.await
.unwrap();
let ctx = SessionContext::new();
let state = ctx.state();
// Test scan without limit
let scan_plan = table_provider.scan(&state, None, &[], None).await.unwrap();
// Verify that the scan plan is an IcebergTableScan
let iceberg_scan = scan_plan
.downcast_ref::<IcebergTableScan>()
.expect("Expected IcebergTableScan");
// Verify the limit is None
assert_eq!(
iceberg_scan.limit(),
None,
"Limit should be None when not specified"
);
}
}