diff --git a/datafusion_iceberg/src/table/mod.rs b/datafusion_iceberg/src/table/mod.rs index 3850f466..87bcfbb9 100644 --- a/datafusion_iceberg/src/table/mod.rs +++ b/datafusion_iceberg/src/table/mod.rs @@ -45,7 +45,10 @@ use crate::{ pruning_statistics::{transform_predicate, PruneDataFiles, PruneManifests}, statistics::manifest_statistics, }; +use datafusion::arrow::compute::SortOptions; use datafusion::common::{NullEquality, Statistics}; +use datafusion::datasource::physical_plan::FileScanConfig; +use datafusion::physical_expr::{LexOrdering, PhysicalSortExpr}; use datafusion::physical_plan::empty::EmptyExec; use datafusion::physical_plan::ColumnStatistics; use datafusion::{ @@ -84,7 +87,13 @@ use datafusion::{ scalar::ScalarValue, sql::parser::DFParserBuilder, }; -use iceberg_rust::spec::{manifest::DataFile, schema::Schema, view_metadata::ViewRepresentation}; +use iceberg_rust::spec::{ + manifest::DataFile, + partition::Transform, + schema::Schema, + sort::{NullOrder, SortDirection, SortOrder}, + view_metadata::ViewRepresentation, +}; use iceberg_rust::{ catalog::tabular::Tabular, error::Error, materialized_view::MaterializedView, table::Table, view::View, @@ -494,6 +503,15 @@ async fn table_scan( let file_schema: SchemaRef = Arc::new((schema.fields()).try_into().unwrap()); + // The ordering the table declares, if any, expressed on the file schema. + // Files attest it individually (manifest `sort_order_id`), so the claim is + // only ever made for the files that carry the attestation. + let declared_ordering = table + .metadata() + .default_sort_order() + .ok() + .and_then(|order| declared_output_ordering(order, &schema, &file_schema)); + // If no projection was specified default to projecting all the fields let projection = projection .cloned() @@ -693,13 +711,13 @@ async fn table_scan( }); } - let file_source = { - let table_schema = TableSchema::new( - file_schema.clone(), - table_partition_cols.iter().cloned().map(Arc::new).collect(), - ); - Arc::new(ParquetSource::new(table_schema)) - }; + let table_schema = TableSchema::new( + file_schema.clone(), + table_partition_cols.iter().cloned().map(Arc::new).collect(), + ); + // File schema plus partition columns: what scan orderings are expressed on. + let scan_schema: SchemaRef = table_schema.table_schema().clone(); + let file_source = Arc::new(ParquetSource::new(table_schema)); // Create plan for every partition with delete files let mut plans = stream::iter(delete_file_groups.into_iter()) @@ -927,47 +945,89 @@ async fn table_scan( .try_collect::>() .await?; - // Create plan for partitions without delete files - let file_groups: Vec<_> = data_file_groups - .into_values() - .map(|x| { - x.into_iter() - .map(|x| { - let last_updated_ms = table.metadata().last_updated_ms; - let manifest_path = if enable_manifest_file_path_column { - Some(x.0) - } else { - None - }; - generate_partitioned_file( - &schema, - &x.1, - last_updated_ms, - enable_data_file_path_column, - manifest_path, - ) - .unwrap() - }) - .collect() - }) - .collect(); + // Create plan for partitions without delete files. + // + // Files that attest the table's declared sort order (manifest + // `sort_order_id`) are scanned separately from those that do not, so the + // ordering claim covers exactly the files that honor it: the attested scan + // carries `output_ordering` (and can be regrouped by statistics into + // non-overlapping groups), the unattested scan claims nothing. A mixed + // table therefore keeps its explicit sort while a fully attested one can + // drop it — never the other way round. + let mut attested_groups: Vec = Vec::new(); + let mut unattested_groups: Vec = Vec::new(); + for entries in data_file_groups.into_values() { + let mut attested = Vec::new(); + let mut unattested = Vec::new(); + for (manifest_path, entry) in entries { + let last_updated_ms = table.metadata().last_updated_ms; + let manifest_path = if enable_manifest_file_path_column { + Some(manifest_path) + } else { + None + }; + let is_attested = declared_ordering + .as_ref() + .is_some_and(|(order_id, _)| entry.data_file().sort_order_id() == &Some(*order_id)); + let file = generate_partitioned_file( + &schema, + &entry, + last_updated_ms, + enable_data_file_path_column, + manifest_path, + )?; + if is_attested { + attested.push(file); + } else { + unattested.push(file); + } + } + if !attested.is_empty() { + attested_groups.push(FileGroup::new(attested)); + } + if !unattested.is_empty() { + unattested_groups.push(FileGroup::new(unattested)); + } + } + + if !unattested_groups.is_empty() { + let file_scan_config = + FileScanConfigBuilder::new(object_store_url.clone(), file_source.clone()) + .with_file_groups(unattested_groups) + .with_statistics(statistics.clone()) + .with_projection_indices(Some(projection.clone()))? + .with_limit(limit) + .build(); + + let other_plan = ParquetFormat::default() + .create_physical_plan(session, file_scan_config) + .instrument(tracing::debug_span!( + "datafusion_iceberg::create_physical_plan_scan_data_files" + )) + .await?; - if !file_groups.is_empty() { + plans.push(other_plan); + } + + if let (false, Some((_, ordering))) = (attested_groups.is_empty(), &declared_ordering) { + let file_groups = + regroup_attested_files_by_statistics(session, &scan_schema, attested_groups, ordering); let file_scan_config = FileScanConfigBuilder::new(object_store_url, file_source) .with_file_groups(file_groups) .with_statistics(statistics) + .with_output_ordering(vec![ordering.clone()]) .with_projection_indices(Some(projection.clone()))? .with_limit(limit) .build(); - let other_plan = ParquetFormat::default() + let sorted_plan = ParquetFormat::default() .create_physical_plan(session, file_scan_config) .instrument(tracing::debug_span!( - "datafusion_iceberg::create_physical_plan_scan_data_files" + "datafusion_iceberg::create_physical_plan_scan_sorted_data_files" )) .await?; - plans.push(other_plan); + plans.push(sorted_plan); } match plans.len() { @@ -980,6 +1040,79 @@ async fn table_scan( } } +/// Maps a table sort order onto a DataFusion ordering over `file_schema`. +/// +/// Only the leading run of identity-transformed fields is claimed: rows +/// sorted by `(a, b, bucket(c))` are sorted by `(a, b)`, but Parquet +/// statistics cannot reason about a transformed value, and a claim on `c` +/// itself would be false. Returns `None` for an order without a usable +/// leading field (including the unsorted order), paired with the order id +/// files must attest to be covered by the claim. +fn declared_output_ordering( + sort_order: &SortOrder, + schema: &Schema, + file_schema: &SchemaRef, +) -> Option<(i32, LexOrdering)> { + let sort_exprs: Vec = sort_order + .fields + .iter() + .take_while(|field| field.transform == Transform::Identity) + .map(|field| { + let name = &schema.get(field.source_id as usize)?.name; + let index = file_schema.index_of(name).ok()?; + Some(PhysicalSortExpr::new( + Arc::new(Column::new(name, index)), + SortOptions { + descending: field.direction == SortDirection::Descending, + nulls_first: field.null_order == NullOrder::First, + }, + )) + }) + .map_while(|expr| expr) + .collect(); + LexOrdering::new(sort_exprs).map(|ordering| (sort_order.order_id, ordering)) +} + +/// Regroups attested files into non-overlapping, statistics-ordered groups +/// when the session asks for it (`split_file_groups_by_statistics`), the same +/// way DataFusion's listing table does. Falls back to the partition-shaped +/// groups if the split is off, fails (e.g. a file without bounds on a sort +/// column), or would exceed `target_partitions`; the ordering claim stays +/// valid either way because DataFusion re-validates it against the file +/// statistics of each group at plan time. +fn regroup_attested_files_by_statistics( + session: &SessionState, + table_schema: &SchemaRef, + file_groups: Vec, + ordering: &LexOrdering, +) -> Vec { + let options = session.config_options(); + if !options.execution.split_file_groups_by_statistics { + return file_groups; + } + let target_partitions = options.execution.target_partitions; + match FileScanConfig::split_groups_by_statistics_with_target_partitions( + table_schema, + &file_groups, + ordering, + target_partitions, + ) { + Ok(new_groups) if new_groups.len() <= target_partitions => new_groups, + Ok(new_groups) => { + tracing::debug!( + groups = new_groups.len(), + target_partitions, + "statistics split produced more file groups than target partitions; keeping partition groups" + ); + file_groups + } + Err(error) => { + tracing::debug!(%error, "failed to split attested file groups by statistics"); + file_groups + } + } +} + fn datafusion_partition_columns( partition_fields: &[BoundPartitionField<'_>], ) -> Result, DataFusionError> { diff --git a/datafusion_iceberg/tests/sort_order_scan.rs b/datafusion_iceberg/tests/sort_order_scan.rs new file mode 100644 index 00000000..d8b45e9c --- /dev/null +++ b/datafusion_iceberg/tests/sort_order_scan.rs @@ -0,0 +1,239 @@ +//! The DataFusion provider claims the table's declared sort order only for +//! files that attest it, so a fully attested table can drop a redundant sort +//! while a table holding any unattested file keeps it — and stays correct. + +use std::sync::Arc; + +use datafusion::arrow::array::{Int64Array, StringArray}; +use datafusion::arrow::datatypes::{DataType, Field, Schema as ArrowSchema}; +use datafusion::arrow::record_batch::RecordBatch; +use datafusion::physical_plan::displayable; +use datafusion::prelude::{SessionConfig, SessionContext}; +use futures::stream; + +use datafusion_iceberg::DataFusionTable; +use iceberg_rust::arrow::write::{write_parquet_partitioned, write_sorted_parquet_partitioned}; +use iceberg_rust::catalog::Catalog; +use iceberg_rust::object_store::ObjectStoreBuilder; +use iceberg_rust::spec::partition::Transform; +use iceberg_rust::spec::schema::Schema; +use iceberg_rust::spec::sort::{NullOrder, SortDirection, SortField, SortOrderBuilder}; +use iceberg_rust::spec::types::{PrimitiveType, StructField, Type}; +use iceberg_rust::table::Table; +use iceberg_sql_catalog::SqlCatalog; + +fn schema() -> Schema { + let mut builder = Schema::builder(); + builder + .with_struct_field(StructField { + id: 1, + name: "id".to_string(), + required: true, + field_type: Type::Primitive(PrimitiveType::Long), + doc: None, + initial_default: None, + write_default: None, + }) + .with_struct_field(StructField { + id: 2, + name: "name".to_string(), + required: true, + field_type: Type::Primitive(PrimitiveType::String), + doc: None, + initial_default: None, + write_default: None, + }); + builder.build().unwrap() +} + +fn arrow_schema() -> Arc { + Arc::new(ArrowSchema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("name", DataType::Utf8, false), + ])) +} + +fn batch(ids: &[i64]) -> RecordBatch { + RecordBatch::try_new( + arrow_schema(), + vec![ + Arc::new(Int64Array::from(ids.to_vec())), + Arc::new(StringArray::from( + ids.iter().map(|id| format!("row-{id}")).collect::>(), + )), + ], + ) + .unwrap() +} + +/// An unpartitioned table sorted by `id ASC`. +async fn sorted_table(name: &str) -> Table { + let catalog: Arc = Arc::new( + SqlCatalog::new("sqlite://", "warehouse", ObjectStoreBuilder::memory()) + .await + .unwrap(), + ); + let sort_order = SortOrderBuilder::default() + .with_order_id(1) + .with_sort_field(SortField { + source_id: 1, + transform: Transform::Identity, + direction: SortDirection::Ascending, + null_order: NullOrder::First, + }) + .build() + .unwrap(); + Table::builder() + .with_name(name) + .with_location(format!("/test/{name}")) + .with_schema(schema()) + .with_sort_order(sort_order) + .build(&["test".to_owned()], catalog) + .await + .expect("Failed to create table") +} + +async fn append_sorted(table: &mut Table, ids: &[i64]) { + let files = write_sorted_parquet_partitioned(table, stream::iter(vec![Ok(batch(ids))]), None) + .await + .unwrap(); + assert!(files.iter().all(|f| *f.sort_order_id() == Some(1))); + table + .new_transaction(None) + .append_data(files) + .commit() + .await + .unwrap(); +} + +async fn append_unsorted(table: &mut Table, ids: &[i64]) { + let files = write_parquet_partitioned(table, stream::iter(vec![Ok(batch(ids))]), None) + .await + .unwrap(); + assert!(files.iter().all(|f| f.sort_order_id().is_none())); + table + .new_transaction(None) + .append_data(files) + .commit() + .await + .unwrap(); +} + +fn context(split_file_groups_by_statistics: bool) -> SessionContext { + let mut config = SessionConfig::new().with_target_partitions(4); + config + .options_mut() + .execution + .split_file_groups_by_statistics = split_file_groups_by_statistics; + SessionContext::new_with_config(config) +} + +async fn physical_plan(ctx: &SessionContext, table: &Table, sql: &str) -> String { + ctx.register_table("t", Arc::new(DataFusionTable::from(table.clone()))) + .unwrap(); + let plan = ctx + .sql(sql) + .await + .unwrap() + .create_physical_plan() + .await + .unwrap(); + let rendered = displayable(plan.as_ref()).indent(true).to_string(); + ctx.deregister_table("t").unwrap(); + rendered +} + +async fn ids(ctx: &SessionContext, table: &Table, sql: &str) -> Vec { + ctx.register_table("t", Arc::new(DataFusionTable::from(table.clone()))) + .unwrap(); + let batches = ctx.sql(sql).await.unwrap().collect().await.unwrap(); + ctx.deregister_table("t").unwrap(); + batches + .iter() + .flat_map(|b| { + b.column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect() +} + +#[tokio::test] +async fn fully_attested_scan_declares_the_order_and_elides_the_sort() { + let mut table = sorted_table("attested").await; + append_sorted(&mut table, &[1, 2, 3, 4]).await; + append_sorted(&mut table, &[5, 6, 7, 8]).await; + + let ctx = context(true); + let plan = physical_plan(&ctx, &table, "SELECT id FROM t ORDER BY id ASC").await; + assert!( + plan.contains("output_ordering=[id@0 ASC]"), + "scan must declare the table's sort order:\n{plan}" + ); + assert!( + !plan.contains("SortExec"), + "non-overlapping attested files satisfy ORDER BY without a sort:\n{plan}" + ); + + assert_eq!( + ids(&ctx, &table, "SELECT id FROM t ORDER BY id ASC").await, + vec![1, 2, 3, 4, 5, 6, 7, 8] + ); + assert_eq!( + ids(&ctx, &table, "SELECT id FROM t ORDER BY id DESC LIMIT 3").await, + vec![8, 7, 6] + ); +} + +#[tokio::test] +async fn one_unattested_file_keeps_the_sort_and_the_results_exact() { + let mut table = sorted_table("mixed").await; + append_sorted(&mut table, &[1, 2, 3, 4]).await; + append_sorted(&mut table, &[9, 10, 11, 12]).await; + // Rows that interleave with both attested files, written without a sort + // and without attestation. + append_unsorted(&mut table, &[6, 5, 8, 7]).await; + + let ctx = context(true); + let plan = physical_plan(&ctx, &table, "SELECT id FROM t ORDER BY id ASC").await; + assert!( + plan.contains("SortExec") || plan.contains("SortPreservingMergeExec"), + "a mixed table must keep an explicit sort:\n{plan}" + ); + + assert_eq!( + ids(&ctx, &table, "SELECT id FROM t ORDER BY id ASC").await, + (1..=12).collect::>() + ); + assert_eq!( + ids(&ctx, &table, "SELECT id FROM t ORDER BY id DESC LIMIT 5").await, + vec![12, 11, 10, 9, 8] + ); +} + +#[tokio::test] +async fn overlapping_attested_files_keep_the_sort_and_the_results_exact() { + let mut table = sorted_table("overlapping").await; + // Each file is sorted, but their ranges overlap: reading them back to back + // is not sorted, and DataFusion must not pretend otherwise. + append_sorted(&mut table, &[1, 3, 5, 7]).await; + append_sorted(&mut table, &[2, 4, 6, 8]).await; + + for split in [true, false] { + let ctx = context(split); + assert_eq!( + ids(&ctx, &table, "SELECT id FROM t ORDER BY id ASC").await, + (1..=8).collect::>(), + "split_file_groups_by_statistics={split}" + ); + assert_eq!( + ids(&ctx, &table, "SELECT id FROM t ORDER BY id DESC LIMIT 3").await, + vec![8, 7, 6], + "split_file_groups_by_statistics={split}" + ); + } +} diff --git a/iceberg-rust-spec/src/spec/table_metadata.rs b/iceberg-rust-spec/src/spec/table_metadata.rs index e15732bd..b8ea45c8 100644 --- a/iceberg-rust-spec/src/spec/table_metadata.rs +++ b/iceberg-rust-spec/src/spec/table_metadata.rs @@ -255,6 +255,21 @@ impl TableMetadata { .ok_or_else(|| Error::InvalidFormat("partition spec".to_string())) } + /// Gets the default sort order for the table + /// + /// Order id `0` is the spec's "unsorted" order; callers that want to know + /// whether the table declares an ordering at all should check + /// [`SortOrder::fields`] for emptiness rather than the id. + /// + /// # Returns + /// * `Result<&SortOrder, Error>` - The default sort order, or an error if it cannot be found + #[inline] + pub fn default_sort_order(&self) -> Result<&SortOrder, Error> { + self.sort_orders + .get(&self.default_sort_order_id) + .ok_or_else(|| Error::InvalidFormat("sort order".to_string())) + } + /// Gets the current partition fields, binding them to their source schema fields /// /// # Returns diff --git a/iceberg-rust/src/arrow/write.rs b/iceberg-rust/src/arrow/write.rs index 8be01942..f8a8adab 100644 --- a/iceberg-rust/src/arrow/write.rs +++ b/iceberg-rust/src/arrow/write.rs @@ -49,7 +49,13 @@ use arrow::{datatypes::Schema as ArrowSchema, error::ArrowError, record_batch::R use futures::Stream; use iceberg_rust_spec::{ partition::BoundPartitionField, - spec::{manifest::DataFile, schema::Schema, values::Value}, + spec::{ + manifest::DataFile, + partition::Transform, + schema::Schema, + sort::{NullOrder, SortDirection, SortOrder}, + values::Value, + }, table_metadata::{ self, WRITE_DATA_PATH, WRITE_METADATA_METRICS_DISTINCT_COUNTS_ENABLED, WRITE_OBJECT_STORAGE_ENABLED, WRITE_PARQUET_BLOOM_FILTER_ENABLED_COLUMN_PREFIX, @@ -63,10 +69,10 @@ use iceberg_rust_spec::{ util::strip_prefix, }; use parquet::{ - arrow::AsyncArrowWriter, + arrow::{ArrowSchemaConverter, AsyncArrowWriter}, basic::{BrotliLevel, Compression, GzipLevel, ZstdLevel}, file::{ - metadata::{KeyValue, ParquetMetaData}, + metadata::{KeyValue, ParquetMetaData, SortingColumn}, properties::{EnabledStatistics, WriterProperties, WriterVersion}, }, schema::types::ColumnPath, @@ -75,7 +81,10 @@ use uuid::Uuid; use crate::{ error::Error, - file_format::parquet::{parquet_to_datafile, ICEBERG_ESTIMATE_INT64_DISTINCT_COUNT_META_KEY}, + file_format::parquet::{ + parquet_to_datafile, ICEBERG_ESTIMATE_INT64_DISTINCT_COUNT_META_KEY, + ICEBERG_SORT_ORDER_ID_META_KEY, + }, object_store::Bucket, table::Table, }; @@ -128,7 +137,59 @@ pub async fn write_parquet_partitioned( batches: impl Stream> + Send + 'static, branch: Option<&str>, ) -> Result, ArrowError> { - store_parquet_partitioned(table, batches, branch, None).await + store_parquet_partitioned(table, batches, branch, None, InputOrdering::Unspecified).await +} + +#[instrument(skip(table, batches), fields(table_name = %table.identifier().name()))] +/// Writes Arrow record batches that are already sorted by the table's default +/// sort order as partitioned Parquet files, attesting that order on every file. +/// +/// The caller promises that the rows of `batches`, taken in stream order, are +/// sorted by the table's default [`SortOrder`]. Rows are routed to partitions +/// and rolled into files in stream order, so every written file is sorted too; +/// each file records the order in its Parquet footer (per-row-group +/// `sorting_columns` for identity-transformed fields plus a +/// `iceberg.sort-order-id` key/value entry) and its manifest entry carries the +/// matching `sort_order_id`. Readers use that attestation to skip sorts, so a +/// caller that cannot guarantee the order MUST use +/// [`write_parquet_partitioned`] instead: a false attestation yields wrong +/// query results, not a slower query. +/// +/// A table whose default sort order has no fields is written unattested, +/// exactly like [`write_parquet_partitioned`]. +/// +/// # Arguments +/// * `table` - The Iceberg table to write data for +/// * `batches` - Stream of record batches, sorted by the table's default sort order +/// * `branch` - Optional branch name to write to +/// +/// # Returns +/// * `Result, ArrowError>` - List of metadata for the written data files +/// +/// # Errors +/// Returns an error under the same conditions as [`write_parquet_partitioned`]. +pub async fn write_sorted_parquet_partitioned( + table: &Table, + batches: impl Stream> + Send + 'static, + branch: Option<&str>, +) -> Result, ArrowError> { + store_parquet_partitioned( + table, + batches, + branch, + None, + InputOrdering::SortedByDefaultOrder, + ) + .await +} + +/// What the caller promises about the order of the rows handed to a write. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum InputOrdering { + /// Nothing is known about the order; files are written unattested. + Unspecified, + /// Rows are sorted by the table's default sort order; files attest it. + SortedByDefaultOrder, } #[instrument(skip(table, batches), fields(table_name = %table.identifier().name(), equality_ids = ?equality_ids))] @@ -159,7 +220,14 @@ pub async fn write_equality_deletes_parquet_partitioned( branch: Option<&str>, equality_ids: &[i32], ) -> Result, ArrowError> { - store_parquet_partitioned(table, batches, branch, Some(equality_ids)).await + store_parquet_partitioned( + table, + batches, + branch, + Some(equality_ids), + InputOrdering::Unspecified, + ) + .await } #[instrument(skip(table, batches), fields(table_name = %table.identifier().name(), equality_ids = ?equality_ids))] @@ -189,10 +257,20 @@ async fn store_parquet_partitioned( batches: impl Stream> + Send + 'static, branch: Option<&str>, equality_ids: Option<&[i32]>, + input_ordering: InputOrdering, ) -> Result, ArrowError> { let metadata = table.metadata(); let object_store = table.object_store(); let schema = Arc::new(metadata.current_schema().map_err(Error::from)?.clone()); + // Only data files can attest an order: delete files are projected onto the + // equality columns, and an order without fields says nothing. + let sort_order: Option> = match input_ordering { + InputOrdering::SortedByDefaultOrder if equality_ids.is_none() => { + let order = metadata.default_sort_order().map_err(Error::from)?; + (!order.fields.is_empty()).then(|| Arc::new(order.clone())) + } + _ => None, + }; // project the schema on to the equality_ids for equality deletes let schema = if let Some(equality_ids) = equality_ids { Arc::new(schema.project(equality_ids)) @@ -238,6 +316,7 @@ async fn store_parquet_partitioned( object_store.clone(), equality_ids, &metadata.properties, + sort_order.as_deref(), ) .await?; Ok(files) @@ -276,6 +355,7 @@ async fn store_parquet_partitioned( let partition_spec = partition_spec.clone(); let equality_ids = equality_ids.map(Vec::from); let table_properties = table_properties.clone(); + let sort_order = sort_order.clone(); let partition_path = if metadata .properties .get(WRITE_OBJECT_STORAGE_ENABLED) @@ -302,6 +382,7 @@ async fn store_parquet_partitioned( object_store.clone(), equality_ids.as_deref(), &table_properties, + sort_order.as_deref(), ) .await?; Ok::<_, Error>(files) @@ -343,6 +424,7 @@ type ArrowReciever = Receiver<(String, ParquetMetaData)>; /// * `batches` - Stream of record batches to write /// * `object_store` - Object store to write files to /// * `equality_ids` - Optional list of field IDs for equality deletes +/// * `sort_order` - The sort order the batches honor and the files attest, if any /// /// # Returns /// * `Result, ArrowError>` - List of metadata for the written files @@ -364,10 +446,15 @@ async fn write_parquet_files( object_store: Arc, equality_ids: Option<&[i32]>, table_properties: &HashMap, + sort_order: Option<&SortOrder>, ) -> Result, ArrowError> { let bucket = Bucket::from_path(data_location)?; let (mut writer_sender, writer_reciever): (ArrowSender, ArrowReciever) = channel(0); let table_properties_owned = Arc::new(table_properties.clone()); + let sort_attestation = sort_order + .map(|order| SortAttestation::new(order, schema, arrow_schema)) + .transpose()? + .map(Arc::new); // Create initial writer let initial_writer = create_arrow_writer( @@ -376,6 +463,7 @@ async fn write_parquet_files( arrow_schema, object_store.clone(), table_properties, + sort_attestation.as_deref(), ) .await?; @@ -400,6 +488,7 @@ async fn write_parquet_files( let arrow_schema = arrow_schema.clone(); let mut writer_sender = writer_sender.clone(); let table_properties = table_properties_owned.clone(); + let sort_attestation = sort_attestation.clone(); async move { // Roll on the file's real on-disk size: what the writer has @@ -424,6 +513,7 @@ async fn write_parquet_files( &arrow_schema, object_store, &table_properties, + sort_attestation.as_deref(), ) .await?; @@ -517,6 +607,8 @@ pub fn generate_partition_path( /// * `partition_path` - Optional partition path component /// * `schema` - Arrow schema for the record batches /// * `object_store` - Object store to write files to +/// * `table_properties` - Table properties that tune the Parquet writer +/// * `sort_attestation` - The sort order to record in the file footer, if any /// /// # Returns /// * `Result<(String, AsyncArrowWriter), ArrowError>` - The file path and configured writer @@ -532,6 +624,7 @@ async fn create_arrow_writer( schema: &arrow::datatypes::Schema, object_store: Arc, table_properties: &HashMap, + sort_attestation: Option<&SortAttestation>, ) -> Result<(String, AsyncArrowWriter), ArrowError> { let parquet_path = generate_file_path(data_location, partition_path); @@ -547,11 +640,22 @@ async fn create_arrow_writer( props_builder = apply_writer_properties(props_builder, table_properties); props_builder = apply_bloom_filter_properties(props_builder, table_properties); props_builder = apply_column_write_properties(props_builder, table_properties); + let mut key_value_metadata = Vec::new(); if estimate_distinct_count { - props_builder = props_builder.set_key_value_metadata(Some(vec![KeyValue::new( + key_value_metadata.push(KeyValue::new( ICEBERG_ESTIMATE_INT64_DISTINCT_COUNT_META_KEY.to_owned(), "true".to_owned(), - )])); + )); + } + if let Some(attestation) = sort_attestation { + key_value_metadata.push(KeyValue::new( + ICEBERG_SORT_ORDER_ID_META_KEY.to_owned(), + attestation.order_id.to_string(), + )); + props_builder = props_builder.set_sorting_columns(attestation.sorting_columns.clone()); + } + if !key_value_metadata.is_empty() { + props_builder = props_builder.set_key_value_metadata(Some(key_value_metadata)); } Ok(( @@ -564,6 +668,77 @@ async fn create_arrow_writer( )) } +/// How a declared sort order is recorded in the Parquet files of one write. +/// +/// The `iceberg.sort-order-id` footer entry always names the order; the +/// per-row-group `sorting_columns` are only emitted when every sort field is +/// an identity transform on a top-level primitive column, since Parquet has +/// no way to describe an order over transformed values or nested leaves. A +/// partial `sorting_columns` list would claim a *different* (prefix) order +/// than the one the file honors, so it is all or nothing. +#[derive(Debug, Clone, PartialEq, Eq)] +struct SortAttestation { + order_id: i32, + sorting_columns: Option>, +} + +impl SortAttestation { + fn new( + sort_order: &SortOrder, + schema: &Schema, + arrow_schema: &ArrowSchema, + ) -> Result { + Ok(Self { + order_id: sort_order.order_id, + sorting_columns: sorting_columns_for(sort_order, schema, arrow_schema)?, + }) + } +} + +/// Maps a sort order onto Parquet `sorting_columns`, or `None` when the order +/// cannot be expressed in Parquet terms (see [`SortAttestation`]). +/// +/// # Errors +/// Returns an error if a sort field references a column missing from the +/// schema, since that means the caller's promise cannot even be stated. +fn sorting_columns_for( + sort_order: &SortOrder, + schema: &Schema, + arrow_schema: &ArrowSchema, +) -> Result>, ArrowError> { + let descriptor = ArrowSchemaConverter::new().convert(arrow_schema)?; + let mut columns = Vec::with_capacity(sort_order.fields.len()); + for field in &sort_order.fields { + if field.transform != Transform::Identity { + return Ok(None); + } + let name = schema + .get(field.source_id as usize) + .map(|f| f.name.as_str()) + .ok_or_else(|| { + ArrowError::SchemaError(format!( + "sort order {} references field id {} that is not in the table schema", + sort_order.order_id, field.source_id + )) + })?; + // A top-level primitive column is exactly one Parquet leaf whose path + // is its own name; anything nested has a longer path. + let Some(column_idx) = descriptor + .columns() + .iter() + .position(|column| column.path().parts() == [name]) + else { + return Ok(None); + }; + columns.push(SortingColumn { + column_idx: column_idx as i32, + descending: field.direction == SortDirection::Descending, + nulls_first: field.null_order == NullOrder::First, + }); + } + Ok(Some(columns)) +} + /// Applies per-column bloom-filter table properties to the writer builder. /// /// Honors the standard Iceberg properties diff --git a/iceberg-rust/src/catalog/create.rs b/iceberg-rust/src/catalog/create.rs index 674212d9..51f1f930 100644 --- a/iceberg-rust/src/catalog/create.rs +++ b/iceberg-rust/src/catalog/create.rs @@ -134,6 +134,14 @@ impl CreateTableBuilder { // Validate sort order references valid schema fields if let Some(Some(order)) = &self.write_order { + // Order id 0 is reserved for the unsorted order, so a sort order + // with fields must carry its own id. + if order.order_id == DEFAULT_SORT_ORDER_ID && !order.fields.is_empty() { + return Err(Error::InvalidFormat(format!( + "Sort order for table '{}' has fields but uses order id {}, which is reserved for the unsorted order", + name, DEFAULT_SORT_ORDER_ID + ))); + } for field in &order.fields { let source_id = field.source_id; if !schema.fields().iter().any(|f| f.id == source_id) { @@ -189,6 +197,8 @@ impl TryInto for CreateTable { .and_then(|x| x.fields().iter().map(|x| *x.field_id()).max()) .unwrap_or(0); + let write_order = self.write_order.unwrap_or_default(); + Ok(TableMetadata { format_version: Default::default(), table_uuid: Uuid::new_v4(), @@ -214,11 +224,10 @@ impl TryInto for CreateTable { snapshots: HashMap::new(), snapshot_log: Vec::new(), metadata_log: Vec::new(), - sort_orders: HashMap::from_iter(vec![( - DEFAULT_SORT_ORDER_ID, - self.write_order.unwrap_or_default(), - )]), - default_sort_order_id: DEFAULT_SORT_ORDER_ID, + // The declared order keeps its own id and becomes the default, + // so a table created with a sort order reports it as such. + default_sort_order_id: write_order.order_id, + sort_orders: HashMap::from_iter(vec![(write_order.order_id, write_order)]), refs: HashMap::new(), next_row_id: 0, }) diff --git a/iceberg-rust/src/file_format/parquet.rs b/iceberg-rust/src/file_format/parquet.rs index 08ccc1c6..d58d2103 100644 --- a/iceberg-rust/src/file_format/parquet.rs +++ b/iceberg-rust/src/file_format/parquet.rs @@ -36,6 +36,27 @@ use crate::file_format::metrics::{ pub const ICEBERG_ESTIMATE_INT64_DISTINCT_COUNT_META_KEY: &str = "iceberg.estimate-int64-distinct-count"; +/// Parquet file-level KV metadata key recording the id of the table sort order +/// the file's rows honor. Written by the sorted write path together with the +/// row groups' `sorting_columns`, and lifted into the manifest entry's +/// `sort_order_id` by [`parquet_to_datafile`], so the attestation travels with +/// the file itself rather than being asserted at commit time. +pub const ICEBERG_SORT_ORDER_ID_META_KEY: &str = "iceberg.sort-order-id"; + +/// The sort order id a Parquet file attests in its footer, if any. +/// +/// Returns `None` when the file carries no [`ICEBERG_SORT_ORDER_ID_META_KEY`] +/// entry or when the entry does not parse as an integer. +pub fn attested_sort_order_id(file_metadata: &ParquetMetaData) -> Option { + file_metadata + .file_metadata() + .key_value_metadata()? + .iter() + .find(|kv| kv.key == ICEBERG_SORT_ORDER_ID_META_KEY) + .and_then(|kv| kv.value.as_deref()) + .and_then(|value| value.parse::().ok()) +} + /// Read datafile statistics from parquetfile #[instrument(name = "iceberg_rust::file_format::parquet::parquet_to_datafile", level = "debug", skip(file_metadata, schema, partition_fields, table_properties), fields( location = location, @@ -421,6 +442,13 @@ pub fn parquet_to_datafile( builder.with_equality_ids(Some(equality_ids.to_vec())); } + // A file written under a declared sort order attests it in its footer; + // carry that into the manifest so readers can trust the ordering without + // opening the file. + if let Some(sort_order_id) = attested_sort_order_id(file_metadata) { + builder.with_sort_order_id(Some(sort_order_id)); + } + let content = builder.build()?; Ok(content) } diff --git a/iceberg-rust/src/table/transaction/mod.rs b/iceberg-rust/src/table/transaction/mod.rs index 857e0a85..72438d17 100644 --- a/iceberg-rust/src/table/transaction/mod.rs +++ b/iceberg-rust/src/table/transaction/mod.rs @@ -18,7 +18,9 @@ use std::collections::HashMap; use tracing::{debug, instrument}; -use iceberg_rust_spec::spec::{manifest::DataFile, schema::Schema, snapshot::SnapshotReference}; +use iceberg_rust_spec::spec::{ + manifest::DataFile, schema::Schema, snapshot::SnapshotReference, sort::SortOrder, +}; use crate::table::transaction::append::append_summary; use crate::table::transaction::operation::SequenceGroup; @@ -39,8 +41,9 @@ pub(crate) static OVERWRITE_INDEX: usize = 5; pub(crate) static UPDATE_PROPERTIES_INDEX: usize = 6; pub(crate) static SET_SNAPSHOT_REF_INDEX: usize = 7; pub(crate) static EXPIRE_SNAPSHOTS_INDEX: usize = 8; +pub(crate) static REPLACE_SORT_ORDER_INDEX: usize = 9; -pub(crate) static NUM_OPERATIONS: usize = 9; +pub(crate) static NUM_OPERATIONS: usize = 10; /// A transaction that can perform multiple operations on a table atomically /// @@ -102,6 +105,24 @@ impl<'table> TableTransaction<'table> { self.operations[SET_DEFAULT_SPEC_INDEX] = Some(Operation::SetDefaultSpec(spec_id)); self } + /// Declares a sort order on the table and makes it the default + /// + /// The order is added to the table metadata under its own `order_id` + /// (replacing any existing order with that id) and becomes the default + /// order that subsequent sorted writes attest their files with. Existing + /// data files keep whatever `sort_order_id` they were written with — this + /// operation changes the declared intent, not the files. + /// + /// # Arguments + /// * `sort_order` - The sort order to declare; must not use order id `0` + /// with fields, since that id is reserved for the unsorted order + /// + /// # Returns + /// * `Self` - The transaction builder for method chaining + pub fn replace_sort_order(mut self, sort_order: SortOrder) -> Self { + self.operations[REPLACE_SORT_ORDER_INDEX] = Some(Operation::ReplaceSortOrder(sort_order)); + self + } /// Appends new data files to the table /// /// This operation adds new data files to the table's current snapshot. Multiple diff --git a/iceberg-rust/src/table/transaction/operation.rs b/iceberg-rust/src/table/transaction/operation.rs index ffdae33c..0c4838f5 100644 --- a/iceberg-rust/src/table/transaction/operation.rs +++ b/iceberg-rust/src/table/transaction/operation.rs @@ -22,6 +22,7 @@ use iceberg_rust_spec::spec::{ snapshot::{ generate_snapshot_id, SnapshotBuilder, SnapshotReference, SnapshotRetention, Summary, }, + sort::{SortOrder, DEFAULT_SORT_ORDER_ID}, }; use iceberg_rust_spec::table_metadata::FormatVersion; use iceberg_rust_spec::util::strip_prefix; @@ -65,8 +66,8 @@ pub enum Operation { UpdateProperties(Vec<(String, String)>), /// Set Ref SetSnapshotRef((String, SnapshotReference)), - /// Replace the sort order - // ReplaceSortOrder, + /// Declare a sort order and make it the default + ReplaceSortOrder(SortOrder), // /// Update the table location // UpdateLocation, /// Append new files to the table @@ -845,6 +846,38 @@ impl Operation { debug!("Executing SetDefaultSpec operation: spec_id={}", spec_id); Ok((None, vec![TableUpdate::SetDefaultSpec { spec_id }])) } + Operation::ReplaceSortOrder(sort_order) => { + debug!( + "Executing ReplaceSortOrder operation: order_id={}", + sort_order.order_id + ); + if sort_order.order_id == DEFAULT_SORT_ORDER_ID && !sort_order.fields.is_empty() { + return Err(Error::InvalidFormat(format!( + "sort order with fields uses order id {DEFAULT_SORT_ORDER_ID}, which is reserved for the unsorted order" + ))); + } + let schema = table_metadata.current_schema()?; + for field in &sort_order.fields { + if !schema.fields().iter().any(|f| f.id == field.source_id) { + return Err(Error::NotFound(format!( + "Sort order field references non-existent schema field ID {}", + field.source_id + ))); + } + } + let sort_order_id = sort_order.order_id; + Ok(( + // Guard against a concurrent writer having moved the + // default in the meantime. + Some(TableRequirement::AssertDefaultSortOrderId { + default_sort_order_id: table_metadata.default_sort_order_id, + }), + vec![ + TableUpdate::AddSortOrder { sort_order }, + TableUpdate::SetDefaultSortOrder { sort_order_id }, + ], + )) + } Operation::ExpireSnapshots { older_than, retain_last, diff --git a/iceberg-rust/tests/sort_order_test.rs b/iceberg-rust/tests/sort_order_test.rs new file mode 100644 index 00000000..d4bc2b03 --- /dev/null +++ b/iceberg-rust/tests/sort_order_test.rs @@ -0,0 +1,311 @@ +//! Declared sort orders end to end: a table declares its default sort order, +//! sorted writes attest it per file (Parquet footer + manifest entry), and +//! unsorted writes stay honest by not attesting anything. + +use std::sync::Arc; + +use arrow::array::{Int64Array, StringArray}; +use arrow::datatypes::{DataType, Field, Schema as ArrowSchema}; +use arrow::record_batch::RecordBatch; +use bytes::Bytes; +use futures::stream; +use object_store::{path::Path, ObjectStore, ObjectStoreExt}; +use parquet::file::metadata::SortingColumn; +use parquet::file::reader::{FileReader, SerializedFileReader}; + +use iceberg_rust::arrow::write::{write_parquet_partitioned, write_sorted_parquet_partitioned}; +use iceberg_rust::catalog::Catalog; +use iceberg_rust::file_format::parquet::{attested_sort_order_id, ICEBERG_SORT_ORDER_ID_META_KEY}; +use iceberg_rust::object_store::ObjectStoreBuilder; +use iceberg_rust::spec::manifest::DataFile; +use iceberg_rust::table::Table; +use iceberg_rust_spec::spec::partition::{PartitionField, PartitionSpec, Transform}; +use iceberg_rust_spec::spec::schema::Schema; +use iceberg_rust_spec::spec::sort::{ + NullOrder, SortDirection, SortField, SortOrder, SortOrderBuilder, +}; +use iceberg_rust_spec::spec::types::{PrimitiveType, StructField, Type}; +use iceberg_rust_spec::util::strip_prefix; +use iceberg_sql_catalog::SqlCatalog; + +const SORT_ORDER_ID: i32 = 1; + +fn schema() -> Schema { + let mut builder = Schema::builder(); + builder + .with_struct_field(StructField { + id: 1, + name: "id".to_string(), + required: true, + field_type: Type::Primitive(PrimitiveType::Long), + doc: None, + initial_default: None, + write_default: None, + }) + .with_struct_field(StructField { + id: 2, + name: "region".to_string(), + required: true, + field_type: Type::Primitive(PrimitiveType::String), + doc: None, + initial_default: None, + write_default: None, + }) + .with_struct_field(StructField { + id: 3, + name: "value".to_string(), + required: false, + field_type: Type::Primitive(PrimitiveType::Long), + doc: None, + initial_default: None, + write_default: None, + }); + builder.build().unwrap() +} + +fn region_partition_spec() -> PartitionSpec { + PartitionSpec::builder() + .with_partition_field(PartitionField::new(2, 1000, "region", Transform::Identity)) + .build() + .unwrap() +} + +/// `(id ASC NULLS FIRST, value DESC NULLS LAST)`. +fn sort_order() -> SortOrder { + SortOrderBuilder::default() + .with_order_id(SORT_ORDER_ID) + .with_sort_field(SortField { + source_id: 1, + transform: Transform::Identity, + direction: SortDirection::Ascending, + null_order: NullOrder::First, + }) + .with_sort_field(SortField { + source_id: 3, + transform: Transform::Identity, + direction: SortDirection::Descending, + null_order: NullOrder::Last, + }) + .build() + .unwrap() +} + +/// The Parquet rendering of [`sort_order`]: leaf indices in the file schema. +fn expected_sorting_columns() -> Vec { + vec![ + SortingColumn { + column_idx: 0, + descending: false, + nulls_first: true, + }, + SortingColumn { + column_idx: 2, + descending: true, + nulls_first: false, + }, + ] +} + +fn arrow_schema() -> ArrowSchema { + ArrowSchema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("region", DataType::Utf8, false), + Field::new("value", DataType::Int64, true), + ]) +} + +/// Rows sorted by [`sort_order`], spanning two partitions. +fn sorted_batch() -> RecordBatch { + RecordBatch::try_new( + Arc::new(arrow_schema()), + vec![ + Arc::new(Int64Array::from(vec![1, 2, 3, 4])), + Arc::new(StringArray::from(vec![ + "us-east", "us-west", "us-east", "us-west", + ])), + Arc::new(Int64Array::from(vec![ + Some(100), + Some(200), + Some(300), + None, + ])), + ], + ) + .unwrap() +} + +async fn catalog() -> Arc { + Arc::new( + SqlCatalog::new("sqlite://", "warehouse", ObjectStoreBuilder::memory()) + .await + .unwrap(), + ) +} + +async fn create_table(catalog: Arc, name: &str, order: Option) -> Table { + let mut builder = Table::builder(); + builder + .with_name(name) + .with_location(format!("/test/{name}")) + .with_schema(schema()) + .with_partition_spec(region_partition_spec()); + if let Some(order) = order { + builder.with_sort_order(order); + } + builder + .build(&["test".to_owned()], catalog) + .await + .expect("Failed to create table") +} + +async fn footer( + store: &Arc, + file: &DataFile, +) -> parquet::file::metadata::ParquetMetaData { + let path = Path::from(strip_prefix(file.file_path())); + let bytes: Bytes = store.get(&path).await.unwrap().bytes().await.unwrap(); + SerializedFileReader::new(bytes).unwrap().metadata().clone() +} + +#[tokio::test] +async fn table_created_with_sort_order_declares_it_as_default() { + let table = create_table(catalog().await, "declared", Some(sort_order())).await; + + let metadata = table.metadata(); + assert_eq!(metadata.default_sort_order_id, SORT_ORDER_ID); + assert_eq!(metadata.default_sort_order().unwrap(), &sort_order()); +} + +#[tokio::test] +async fn sorted_write_attests_the_order_on_every_file() { + let table = create_table(catalog().await, "sorted", Some(sort_order())).await; + + let files = + write_sorted_parquet_partitioned(&table, stream::iter(vec![Ok(sorted_batch())]), None) + .await + .expect("sorted write"); + assert_eq!(files.len(), 2, "one file per region partition"); + + let store = table.object_store(); + for file in &files { + assert_eq!( + *file.sort_order_id(), + Some(SORT_ORDER_ID), + "manifest entry must carry the attested sort order id" + ); + + let metadata = footer(&store, file).await; + assert_eq!(attested_sort_order_id(&metadata), Some(SORT_ORDER_ID)); + let key_value = metadata + .file_metadata() + .key_value_metadata() + .and_then(|kvs| { + kvs.iter() + .find(|kv| kv.key == ICEBERG_SORT_ORDER_ID_META_KEY) + }) + .expect("footer records the sort order id"); + assert_eq!(key_value.value.as_deref(), Some("1")); + assert!(!metadata.row_groups().is_empty()); + for row_group in metadata.row_groups() { + assert_eq!( + row_group.sorting_columns(), + Some(&expected_sorting_columns()), + "every row group declares the sort order's columns" + ); + } + } +} + +#[tokio::test] +async fn unsorted_write_attests_nothing() { + let table = create_table(catalog().await, "unsorted", Some(sort_order())).await; + + let files = write_parquet_partitioned(&table, stream::iter(vec![Ok(sorted_batch())]), None) + .await + .expect("plain write"); + assert!(!files.is_empty()); + + let store = table.object_store(); + for file in &files { + assert_eq!(*file.sort_order_id(), None); + let metadata = footer(&store, file).await; + assert_eq!(attested_sort_order_id(&metadata), None); + for row_group in metadata.row_groups() { + assert_eq!(row_group.sorting_columns(), None); + } + } +} + +#[tokio::test] +async fn sorted_write_on_a_table_without_declared_order_attests_nothing() { + let table = create_table(catalog().await, "undeclared", None).await; + assert!(table + .metadata() + .default_sort_order() + .unwrap() + .fields + .is_empty()); + + let files = + write_sorted_parquet_partitioned(&table, stream::iter(vec![Ok(sorted_batch())]), None) + .await + .expect("sorted write"); + for file in &files { + assert_eq!(*file.sort_order_id(), None); + } +} + +#[tokio::test] +async fn replace_sort_order_declares_the_order_on_an_existing_table() { + let catalog = catalog().await; + let mut table = create_table(catalog.clone(), "upgraded", None).await; + assert_eq!(table.metadata().default_sort_order_id, 0); + + table + .new_transaction(None) + .replace_sort_order(sort_order()) + .commit() + .await + .expect("replace sort order"); + + let metadata = table.metadata(); + assert_eq!(metadata.default_sort_order_id, SORT_ORDER_ID); + assert_eq!(metadata.default_sort_order().unwrap(), &sort_order()); + + // Idempotent: declaring the same order again is a no-op on the metadata. + table + .new_transaction(None) + .replace_sort_order(sort_order()) + .commit() + .await + .expect("replace sort order again"); + assert_eq!(table.metadata().default_sort_order_id, SORT_ORDER_ID); + assert_eq!( + table.metadata().sort_orders.len(), + 2, + "unsorted order 0 plus the declared one" + ); + + // Files written after the declaration attest it. + let files = + write_sorted_parquet_partitioned(&table, stream::iter(vec![Ok(sorted_batch())]), None) + .await + .expect("sorted write"); + for file in &files { + assert_eq!(*file.sort_order_id(), Some(SORT_ORDER_ID)); + } +} + +#[tokio::test] +async fn replace_sort_order_rejects_the_reserved_unsorted_id() { + let mut table = create_table(catalog().await, "reserved", None).await; + let mut order = sort_order(); + order.order_id = 0; + + let result = table + .new_transaction(None) + .replace_sort_order(order) + .commit() + .await; + assert!(result.is_err(), "order id 0 with fields must be rejected"); +}