diff --git a/AGENTS.md b/AGENTS.md
index 103ff4c..69b8a18 100644
--- a/AGENTS.md
+++ b/AGENTS.md
@@ -73,6 +73,8 @@ Tests are self-contained (in-memory catalogs, temp dirs, checked-in metadata JSO
`IcebergTableScan` converts the projection to column names and the filters to a single Iceberg `Predicate` (`physical_plan/expr_to_predicate.rs`). It applies `limit` in-stream, has one output partition, and delegates the actual read to iceberg's `TableScan::to_arrow()`.
+With the session option `iceberg.execution.preserve_data_ordering` (`config.rs`; off by default), `TableProvider::scan` lists the data files up front. If there are at most `iceberg.execution.max_merge_files_per_partition` of them and all record the same resolvable sort order, `IcebergDataSource` reports that order's leading fields as its output ordering and k-way merges one reader per file in `open()`. Only identity transforms of projected top-level columns are reported, stopping at the first float, double, or UUID field, because Arrow orders -0.0 and NaN differently from Spark-written files. Otherwise the source is built exactly as without the option. A source rebuilt with `IcebergDataSource::builder` must be given what `sorted_tasks()` returns through `with_sorted_tasks`, or it will read the files unordered under a plan that relies on their order. The providers build a sorted source the same way: `plan_sorted_tasks` lists the files, and they pass its result to `with_sorted_tasks`. The order is also not reported unless the provider's declared schema matches the snapshot schema the files are read with, since DataFusion ignores the null order of a column declared non-nullable.
+
Both providers report every filter as `Inexact`, so DataFusion re-applies the original filters after the scan. The pushed-down predicate therefore only has to be *implied by* the original filter. It may match extra rows but must never drop a row that matches. Unconvertible parts are dropped: an `AND` keeps whichever side converted, while an `OR` needs both sides. Preserve this soundness rule when extending `expr_to_predicate.rs`; for example, it is why date casts and some NaN arithmetic are deliberately not pushed down.
### Write path (`IcebergTableProvider::insert_into`)
@@ -98,11 +100,12 @@ Library code returns `datafusion::error::Result`. Convert iceberg errors with `t
## sqllogictest harness (`crates/sqllogictest`)
- It's a custom `harness = false` runner using libtest-mimic. Each TOML file in `testdata/schedules/` is one test: it declares engines and ordered `[[steps]]` that point at `.slt` files under `testdata/slts/`. An `.slt` file that no schedule references never runs.
-- Each schedule gets a fresh `SessionContext` (`target_partitions = 4`, information_schema enabled) and a fresh in-memory Iceberg catalog. The catalog is registered as `default` with namespace `default`, so SQL addresses tables as `default.default.
`. `datafusion` is the only engine type. The `catalog` setting in a schedule is parsed but ignored.
+- Each schedule gets a fresh `SessionContext` (`target_partitions = 4`, information_schema enabled, `IcebergDataFusionConfig` registered) and a fresh in-memory Iceberg catalog. The catalog is registered as `default` with namespace `default`, so SQL addresses tables as `default.default.`. `datafusion` is the only engine type. The `catalog` setting in a schedule is parsed but ignored.
- `engine/datafusion.rs` pre-creates the tables SQL can't express yet:
- `test_partitioned_table`: identity-partitioned on `category`
- `test_binary_table`
- `test_encrypted_round_trip`: format V3, encrypted with in-memory KMS key `test-master-key`
+ - `test_sorted_table`: sorted by `id`, with two data files whose ids overlap, each recording that sort order
All other tables are created by `CREATE TABLE` inside the `.slt` files.
- Several `.slt` files assert `EXPLAIN` output, such as `IcebergTableScan projection:[...] predicate:[...] limit:[...]`. Changes to the scan's `DisplayAs`, to pushdown, or to the DataFusion version require updating those expectations by hand; the harness has no auto-complete mode.
diff --git a/crates/datafusion/src/config.rs b/crates/datafusion/src/config.rs
new file mode 100644
index 0000000..e96b95c
--- /dev/null
+++ b/crates/datafusion/src/config.rs
@@ -0,0 +1,92 @@
+// 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.
+
+//! Session options for Iceberg tables, set with `SET iceberg..`.
+
+use datafusion::common::config::ConfigExtension;
+use datafusion::common::extensions_options;
+
+extensions_options! {
+ /// Session options for reading Iceberg tables through DataFusion.
+ ///
+ /// `SET iceberg.. = ` only works once this extension
+ /// is registered on the session, with
+ /// [`SessionConfig::with_option_extension`](datafusion::prelude::SessionConfig::with_option_extension).
+ /// Without it, every option keeps its default. DataFusion lists these
+ /// options in `information_schema.df_settings` by their last name alone,
+ /// such as `max_merge_files_per_partition`, and `SHOW` and `RESET` do not
+ /// support them.
+ ///
+ /// ```
+ /// # #[tokio::main(flavor = "current_thread")]
+ /// # async fn main() -> datafusion::error::Result<()> {
+ /// use datafusion::prelude::{SessionConfig, SessionContext};
+ /// use iceberg_datafusion::IcebergDataFusionConfig;
+ ///
+ /// let config =
+ /// SessionConfig::new().with_option_extension(IcebergDataFusionConfig::default());
+ /// let ctx = SessionContext::new_with_config(config);
+ /// ctx.sql("SET iceberg.execution.preserve_data_ordering = true")
+ /// .await?
+ /// .collect()
+ /// .await?;
+ ///
+ /// let state = ctx.state();
+ /// let options = state.config().options();
+ /// let iceberg = options.extensions.get::().unwrap();
+ /// assert!(iceberg.execution.preserve_data_ordering);
+ /// # Ok(())
+ /// # }
+ /// ```
+ pub struct IcebergDataFusionConfig {
+ /// Options for executing table scans.
+ pub execution: IcebergExecutionConfig, default = IcebergExecutionConfig::default()
+ }
+}
+
+extensions_options! {
+ /// Options for executing Iceberg table scans, under `iceberg.execution`.
+ pub struct IcebergExecutionConfig {
+ /// When true, a scan lists its data files while it is planned and, if
+ /// every file records the same sort order, reports that order to
+ /// DataFusion so that it can skip sorting the scan's output. The scan
+ /// then merges the files in that order, with every file open at once.
+ ///
+ /// Only the leading sort fields that are identity transforms of
+ /// projected top-level columns are reported, stopping at the first
+ /// floating point or UUID field, whose orders in Iceberg writers and in
+ /// DataFusion can differ.
+ ///
+ /// The scan reports each sort field's null order as the table declares
+ /// it. Writers commonly sort ascending fields with nulls first, the
+ /// default of Iceberg Java's `SortOrderBuilder.asc`, while DataFusion's
+ /// `ORDER BY c` puts them last. On a nullable column sorted that way,
+ /// only `ORDER BY c NULLS FIRST`, or a session with
+ /// `datafusion.sql_parser.default_null_ordering = 'nulls_min'`, can
+ /// skip its sort. Other queries sort the merged output again.
+ pub preserve_data_ordering: bool, default = false
+ /// The most data files a scan merges in one output partition to
+ /// preserve their sort order. A scan with more files in a partition
+ /// reports no order, and DataFusion sorts its output instead, which
+ /// can spill to disk where the merge cannot.
+ pub max_merge_files_per_partition: usize, default = 64
+ }
+}
+
+impl ConfigExtension for IcebergDataFusionConfig {
+ const PREFIX: &'static str = "iceberg";
+}
diff --git a/crates/datafusion/src/lib.rs b/crates/datafusion/src/lib.rs
index 41cbc8b..733c6dd 100644
--- a/crates/datafusion/src/lib.rs
+++ b/crates/datafusion/src/lib.rs
@@ -18,6 +18,9 @@
mod catalog;
pub use catalog::*;
+mod config;
+pub use config::{IcebergDataFusionConfig, IcebergExecutionConfig};
+
mod error;
pub use error::*;
diff --git a/crates/datafusion/src/read/scan.rs b/crates/datafusion/src/read/scan.rs
index afcf3ee..aedc11e 100644
--- a/crates/datafusion/src/read/scan.rs
+++ b/crates/datafusion/src/read/scan.rs
@@ -19,23 +19,37 @@ use std::pin::Pin;
use std::sync::Arc;
use datafusion::arrow::array::RecordBatch;
-use datafusion::arrow::datatypes::SchemaRef as ArrowSchemaRef;
-use datafusion::common::Statistics;
+use datafusion::arrow::compute::SortOptions;
+use datafusion::arrow::datatypes::{Schema as ArrowSchema, SchemaRef as ArrowSchemaRef};
+use datafusion::catalog::Session;
use datafusion::common::tree_node::TreeNodeRecursion;
+use datafusion::common::{Statistics, plan_err};
use datafusion::datasource::source::DataSource;
use datafusion::error::Result;
+use datafusion::execution::memory_pool::MemoryConsumer;
use datafusion::execution::{SendableRecordBatchStream, TaskContext};
+use datafusion::physical_expr::expressions::Column;
use datafusion::physical_expr::projection::ProjectionExprs;
-use datafusion::physical_expr::{EquivalenceProperties, PhysicalExpr};
+use datafusion::physical_expr::{
+ EquivalenceProperties, LexOrdering, PhysicalExpr, PhysicalSortExpr,
+};
+use datafusion::physical_plan::common::spawn_buffered;
+use datafusion::physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet};
+use datafusion::physical_plan::sorts::streaming_merge::StreamingMergeBuilder;
use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
use datafusion::physical_plan::{DisplayFormatType, Partitioning};
use datafusion::prelude::Expr;
-use futures::{Stream, TryStreamExt};
+use futures::{Stream, StreamExt, TryStreamExt};
+use iceberg::arrow::schema_to_arrow_schema;
use iceberg::expr::Predicate;
+use iceberg::scan::{FileScanTask, TableScan};
+use iceberg::spec::{
+ NullOrder, PrimitiveType, Schema, SortDirection, SortOrder, Transform, Type,
+};
use iceberg::table::Table;
use super::predicate::convert_filters_to_predicate;
-use crate::to_datafusion_error;
+use crate::{IcebergDataFusionConfig, to_datafusion_error};
/// Manages the scanning process of an Iceberg [`Table`], encapsulating the
/// necessary details and computed properties required for execution planning.
@@ -53,6 +67,19 @@ pub struct IcebergDataSource {
predicate: Option,
/// Optional limit on the number of rows to return
fetch: Option,
+ /// The data files to merge in the order the source reports, listed while
+ /// the scan was planned. `None` lists the files when the source runs and
+ /// reports no order.
+ sorted: Option,
+ metrics: ExecutionPlanMetricsSet,
+}
+
+/// Data files that all record the same sort order, and the part of that order
+/// the source reports, over the columns of its output.
+#[derive(Debug, Clone)]
+struct SortedTasks {
+ tasks: Vec,
+ ordering: LexOrdering,
}
impl IcebergDataSource {
@@ -67,6 +94,7 @@ impl IcebergDataSource {
snapshot_id: None,
predicate: None,
fetch: None,
+ sorted_tasks: None,
}
}
@@ -93,14 +121,73 @@ impl IcebergDataSource {
pub fn fetch(&self) -> Option {
self.fetch
}
+
+ /// The data files the source merges in the order it reports, listed while
+ /// the scan was planned, or `None` if it reports no order and lists its
+ /// files when it runs. A source rebuilt from this one must be given them
+ /// with [`IcebergDataSourceBuilder::with_sorted_tasks`], because a plan
+ /// built over a sorted source may rely on its order.
+ pub fn sorted_tasks(&self) -> Option<&[FileScanTask]> {
+ self.sorted.as_ref().map(|sorted| sorted.tasks.as_slice())
+ }
+
+ /// Reads each sorted data file as its own stream and merges the streams in
+ /// the order the source reports.
+ fn merge_sorted(
+ &self,
+ sorted: &SortedTasks,
+ partition: usize,
+ context: &TaskContext,
+ ) -> Result {
+ // The builder defaults to the reader settings `TableScan::to_arrow` uses.
+ let reader = self.table.reader_builder().build();
+ let mut streams = sorted
+ .tasks
+ .iter()
+ .map(|task| {
+ let tasks = futures::stream::iter([Ok(task.clone())]).boxed();
+ let batches = reader
+ .clone()
+ .read(tasks)
+ .map_err(to_datafusion_error)?
+ .stream()
+ .map_err(to_datafusion_error);
+ Ok(
+ Box::pin(RecordBatchStreamAdapter::new(self.schema(), batches))
+ as SendableRecordBatchStream,
+ )
+ })
+ .collect::>>()?;
+ if streams.len() == 1 {
+ return Ok(streams.remove(0));
+ }
+
+ let reservation = MemoryConsumer::new(format!("IcebergDataSource[{partition}]"))
+ .register(&context.runtime_env().memory_pool);
+ StreamingMergeBuilder::new()
+ .with_streams(
+ streams
+ .into_iter()
+ .map(|stream| spawn_buffered(stream, 1))
+ .collect(),
+ )
+ .with_schema(self.schema())
+ .with_expressions(&sorted.ordering)
+ .with_metrics(BaselineMetrics::new(&self.metrics, partition))
+ .with_batch_size(context.session_config().batch_size())
+ .with_fetch(self.fetch)
+ .with_reservation(reservation)
+ .build()
+ }
}
/// Builds an [`IcebergDataSource`]. Created by [`IcebergDataSource::builder`].
///
/// A table provider passes the table's schema and selects columns with
/// [`Self::with_projection`]. To rebuild a source, such as in another process,
-/// pass its [`IcebergDataSource::schema`], which is already projected, and no
-/// projection.
+/// pass its [`IcebergDataSource::schema`], which is already projected, no
+/// projection, and its [`IcebergDataSource::sorted_tasks`] to
+/// [`Self::with_sorted_tasks`].
#[derive(Debug, Clone)]
pub struct IcebergDataSourceBuilder {
table: Table,
@@ -109,6 +196,7 @@ pub struct IcebergDataSourceBuilder {
snapshot_id: Option,
predicate: Option,
fetch: Option,
+ sorted_tasks: Option>,
}
impl IcebergDataSourceBuilder {
@@ -145,11 +233,24 @@ impl IcebergDataSourceBuilder {
self
}
+ /// Merges these data files in the sort order they all record, which the
+ /// source then reports to DataFusion, or lists the files when the source
+ /// runs and reports no order if `None` (the default). A plan built over a
+ /// source with [`IcebergDataSource::sorted_tasks`] may rely on their order,
+ /// so a source rebuilt from it must be given them.
+ pub fn with_sorted_tasks(mut self, sorted_tasks: Option>) -> Self {
+ self.sorted_tasks = sorted_tasks;
+ self
+ }
+
/// Builds the source, which reads the columns of its projected schema by name.
///
/// # Errors
///
- /// Returns an error if the projection is out of bounds for the schema.
+ /// Returns an error if the projection is out of bounds for the schema, or
+ /// if the sorted tasks do not all record a sort order the source can
+ /// report over its schema, or are read with a snapshot schema whose
+ /// columns do not match it.
pub fn build(self) -> Result {
let schema = match &self.projection {
None => self.schema,
@@ -160,6 +261,19 @@ impl IcebergDataSourceBuilder {
.iter()
.map(|field| field.name().clone())
.collect();
+ let sorted = match self.sorted_tasks {
+ None => None,
+ Some(tasks) => {
+ let Some(ordering) = shared_ordering(&tasks, &schema) else {
+ return plan_err!(
+ "IcebergDataSource cannot report an order for the given data \
+ files: they do not all record the same sort order, or are \
+ read with a schema that does not match the source's"
+ );
+ };
+ Some(SortedTasks { tasks, ordering })
+ }
+ };
Ok(IcebergDataSource {
table: self.table,
snapshot_id: self.snapshot_id,
@@ -167,6 +281,8 @@ impl IcebergDataSourceBuilder {
projection,
predicate: self.predicate,
fetch: self.fetch,
+ sorted,
+ metrics: ExecutionPlanMetricsSet::new(),
})
}
}
@@ -181,16 +297,22 @@ impl DataSource for IcebergDataSource {
fn open(
&self,
- _partition: usize,
- _context: Arc,
+ partition: usize,
+ context: Arc,
) -> Result {
- let fut = get_batch_stream(
- self.table.clone(),
- self.snapshot_id,
- Some(self.projection.clone()),
- self.predicate.clone(),
- );
- let stream = futures::stream::once(fut).try_flatten();
+ let stream: Pin> + Send>> =
+ match &self.sorted {
+ Some(sorted) => self.merge_sorted(sorted, partition, &context)?,
+ None => {
+ let fut = get_batch_stream(
+ self.table.clone(),
+ self.snapshot_id,
+ Some(self.projection.clone()),
+ self.predicate.clone(),
+ );
+ Box::pin(futures::stream::once(fut).try_flatten())
+ }
+ };
// Apply the fetch if specified.
let limited_stream: Pin> + Send>> =
@@ -228,6 +350,14 @@ impl DataSource for IcebergDataSource {
if let Some(predicate) = &self.predicate {
write!(f, ", predicate={predicate}")?;
}
+ if let Some(sorted) = &self.sorted {
+ write!(
+ f,
+ ", output_ordering=[{}], files={}",
+ sorted.ordering,
+ sorted.tasks.len()
+ )?;
+ }
if let Some(fetch) = self.fetch {
write!(f, ", fetch={fetch}")?;
}
@@ -239,7 +369,13 @@ impl DataSource for IcebergDataSource {
}
fn eq_properties(&self) -> EquivalenceProperties {
- EquivalenceProperties::new(self.schema.clone())
+ match &self.sorted {
+ Some(sorted) => EquivalenceProperties::new_with_orderings(
+ self.schema.clone(),
+ [sorted.ordering.clone()],
+ ),
+ None => EquivalenceProperties::new(self.schema.clone()),
+ }
}
fn partition_statistics(&self, _partition: Option) -> Result> {
@@ -256,6 +392,10 @@ impl DataSource for IcebergDataSource {
self.fetch
}
+ fn metrics(&self) -> ExecutionPlanMetricsSet {
+ self.metrics.clone()
+ }
+
fn try_swapping_with_projection(
&self,
_projection: &ProjectionExprs,
@@ -264,6 +404,64 @@ impl DataSource for IcebergDataSource {
}
}
+/// The data files that a source built by `builder` can merge in the sort order
+/// they all record, for [`IcebergDataSourceBuilder::with_sorted_tasks`], or
+/// `None` to have the source list them when it runs and report no order. They
+/// are listed only if the session enables
+/// `iceberg.execution.preserve_data_ordering`, and returned only if they are
+/// few enough to merge and the source can report their order.
+pub(crate) async fn plan_sorted_tasks(
+ builder: &IcebergDataSourceBuilder,
+ state: &dyn Session,
+) -> Result>> {
+ let config = state
+ .config()
+ .options()
+ .extensions
+ .get::()
+ .map(|config| config.execution.clone())
+ .unwrap_or_default();
+ if !config.preserve_data_ordering {
+ return Ok(None);
+ }
+
+ let source = builder.clone().build()?;
+ let metadata = source.table.metadata();
+ let snapshot = match source.snapshot_id {
+ Some(snapshot_id) => metadata.snapshot_by_id(snapshot_id),
+ None => metadata.current_snapshot(),
+ };
+ let Some(snapshot) = snapshot else {
+ return Ok(None);
+ };
+ let schema = snapshot.schema(metadata).map_err(to_datafusion_error)?;
+ let output = source.schema();
+ if !metadata
+ .sort_orders_iter()
+ .any(|order| lex_ordering(order, &schema, &output).is_some())
+ {
+ return Ok(None);
+ }
+
+ let tasks: Vec = build_table_scan(
+ &source.table,
+ source.snapshot_id,
+ Some(source.projection.clone()),
+ source.predicate.clone(),
+ )?
+ .plan_files()
+ .await
+ .map_err(to_datafusion_error)?
+ .take(config.max_merge_files_per_partition.saturating_add(1))
+ .try_collect()
+ .await
+ .map_err(to_datafusion_error)?;
+ if tasks.len() > config.max_merge_files_per_partition {
+ return Ok(None);
+ }
+ Ok(shared_ordering(&tasks, &output).map(|_| tasks))
+}
+
/// Asynchronously retrieves a stream of [`RecordBatch`] instances
/// from a given table.
///
@@ -275,6 +473,24 @@ async fn get_batch_stream(
column_names: Option>,
predicate: Option,
) -> Result> + Send>>> {
+ let table_scan = build_table_scan(&table, snapshot_id, column_names, predicate)?;
+
+ let stream = table_scan
+ .to_arrow()
+ .await
+ .map_err(to_datafusion_error)?
+ .map_err(to_datafusion_error);
+ Ok(Box::pin(stream))
+}
+
+/// Builds the [`TableScan`] of the columns `column_names` of a snapshot of
+/// `table`, or of its current snapshot, filtered by `predicate`.
+fn build_table_scan(
+ table: &Table,
+ snapshot_id: Option,
+ column_names: Option>,
+ predicate: Option,
+) -> Result {
let scan_builder = match snapshot_id {
Some(snapshot_id) => table.scan().snapshot_id(snapshot_id),
None => table.scan(),
@@ -287,12 +503,298 @@ async fn get_batch_stream(
if let Some(predicate) = predicate {
scan_builder = scan_builder.with_filter(predicate);
}
- let table_scan = scan_builder.build().map_err(to_datafusion_error)?;
+ scan_builder.build().map_err(to_datafusion_error)
+}
- let stream = table_scan
- .to_arrow()
- .await
- .map_err(to_datafusion_error)?
- .map_err(to_datafusion_error);
- Ok(Box::pin(stream))
+/// The order of the columns of `output` that every one of `tasks` is sorted
+/// in, or `None` if they record different sort orders, one records none, or
+/// the schema they are read with does not match `output`.
+fn shared_ordering(tasks: &[FileScanTask], output: &ArrowSchema) -> Option {
+ let first = tasks.first()?;
+ let order = first.sort_order()?;
+ if !tasks.iter().all(|task| task.sort_order() == Some(order)) {
+ return None;
+ }
+ if !reads_as_declared(first.schema(), output) {
+ return None;
+ }
+ lex_ordering(order, first.schema(), output)
+}
+
+/// Whether reading `schema`, the snapshot schema the tasks are read with,
+/// yields every column of `output` with the type it declares, and no nulls
+/// where it declares none. A provider's schema can be older than the
+/// snapshot's, and DataFusion ignores the null order of a non-nullable
+/// column, while the merge requires its input to match `output`.
+fn reads_as_declared(schema: &Schema, output: &ArrowSchema) -> bool {
+ let Ok(read) = schema_to_arrow_schema(schema) else {
+ return false;
+ };
+ output.fields().iter().all(|field| {
+ read.field_with_name(field.name()).is_ok_and(|read| {
+ read.data_type() == field.data_type()
+ && (field.is_nullable() || !read.is_nullable())
+ })
+ })
+}
+
+/// The longest leading part of `order` that DataFusion can merge and report as
+/// an order of the columns of `output`, or `None` if that is empty. `schema`
+/// is the Iceberg schema the scan reads.
+fn lex_ordering(
+ order: &SortOrder,
+ schema: &Schema,
+ output: &ArrowSchema,
+) -> Option {
+ let sort_exprs = order.fields.iter().map_while(|field| {
+ if field.transform != Transform::Identity {
+ return None;
+ }
+ let source = schema.as_struct().field_by_id(field.source_id)?;
+ if !orders_like_writers(&source.field_type) {
+ return None;
+ }
+ let index = output.index_of(&source.name).ok()?;
+ Some(PhysicalSortExpr::new(
+ Arc::new(Column::new(&source.name, index)),
+ SortOptions {
+ descending: field.direction == SortDirection::Descending,
+ nulls_first: field.null_order == NullOrder::First,
+ },
+ ))
+ });
+ LexOrdering::new(sort_exprs)
+}
+
+/// Whether DataFusion orders values of `field_type` as the engines that write
+/// sorted Iceberg files do. For floating point, DataFusion uses the order the
+/// spec defines (), with -0.0
+/// before 0.0 and negative NaN first. Spark treats the zeros as equal and
+/// every NaN as the largest value, so files it writes may not be sorted in the
+/// spec's order. The spec leaves the order of UUIDs undefined
+/// (apache/iceberg#14216).
+fn orders_like_writers(field_type: &Type) -> bool {
+ matches!(
+ field_type,
+ Type::Primitive(
+ PrimitiveType::Boolean
+ | PrimitiveType::Int
+ | PrimitiveType::Long
+ | PrimitiveType::Decimal { .. }
+ | PrimitiveType::Date
+ | PrimitiveType::Time
+ | PrimitiveType::Timestamp
+ | PrimitiveType::Timestamptz
+ | PrimitiveType::TimestampNs
+ | PrimitiveType::TimestamptzNs
+ | PrimitiveType::String
+ | PrimitiveType::Fixed(_)
+ | PrimitiveType::Binary
+ )
+ )
+}
+
+#[cfg(test)]
+mod tests {
+ use iceberg::spec::{NestedField, SortField, StructType};
+
+ use super::*;
+
+ fn schema() -> Schema {
+ Schema::builder()
+ .with_fields(vec![
+ NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int))
+ .into(),
+ NestedField::optional(2, "name", Type::Primitive(PrimitiveType::String))
+ .into(),
+ NestedField::optional(3, "ts", Type::Primitive(PrimitiveType::Timestamp))
+ .into(),
+ NestedField::optional(
+ 4,
+ "amount",
+ Type::Primitive(PrimitiveType::Decimal {
+ precision: 10,
+ scale: 2,
+ }),
+ )
+ .into(),
+ NestedField::optional(5, "price", Type::Primitive(PrimitiveType::Double))
+ .into(),
+ NestedField::optional(6, "ratio", Type::Primitive(PrimitiveType::Float))
+ .into(),
+ NestedField::optional(7, "uid", Type::Primitive(PrimitiveType::Uuid))
+ .into(),
+ NestedField::optional(
+ 8,
+ "s",
+ Type::Struct(StructType::new(vec![
+ NestedField::optional(
+ 9,
+ "inner",
+ Type::Primitive(PrimitiveType::Int),
+ )
+ .into(),
+ ])),
+ )
+ .into(),
+ ])
+ .build()
+ .unwrap()
+ }
+
+ fn field(
+ source_id: i32,
+ transform: Transform,
+ direction: SortDirection,
+ null_order: NullOrder,
+ ) -> SortField {
+ SortField::builder()
+ .source_id(source_id)
+ .transform(transform)
+ .direction(direction)
+ .null_order(null_order)
+ .build()
+ }
+
+ fn asc(source_id: i32) -> SortField {
+ field(
+ source_id,
+ Transform::Identity,
+ SortDirection::Ascending,
+ NullOrder::First,
+ )
+ }
+
+ /// The order reported for `fields` over the columns `projection` of
+ /// [`schema`], rendered as DataFusion displays it.
+ fn ordering(fields: Vec, projection: &[&str]) -> Option {
+ let schema = schema();
+ let mut builder = SortOrder::builder();
+ builder.with_order_id(1);
+ for field in fields {
+ builder.with_sort_field(field);
+ }
+ let order = builder.build(&schema).unwrap();
+ let arrow = schema_to_arrow_schema(&schema).unwrap();
+ let indices: Vec = projection
+ .iter()
+ .map(|name| arrow.index_of(name).unwrap())
+ .collect();
+ let output = arrow.project(&indices).unwrap();
+ lex_ordering(&order, &schema, &output).map(|ordering| ordering.to_string())
+ }
+
+ const ALL: &[&str] = &["id", "name", "ts", "amount", "price", "ratio", "uid", "s"];
+
+ #[test]
+ fn test_identity_fields_keep_direction_and_null_order() {
+ let fields = vec![
+ field(
+ 1,
+ Transform::Identity,
+ SortDirection::Ascending,
+ NullOrder::First,
+ ),
+ field(
+ 2,
+ Transform::Identity,
+ SortDirection::Ascending,
+ NullOrder::Last,
+ ),
+ field(
+ 3,
+ Transform::Identity,
+ SortDirection::Descending,
+ NullOrder::First,
+ ),
+ field(
+ 4,
+ Transform::Identity,
+ SortDirection::Descending,
+ NullOrder::Last,
+ ),
+ ];
+ assert_eq!(
+ ordering(fields, ALL).as_deref(),
+ Some("id@0 ASC, name@1 ASC NULLS LAST, ts@2 DESC, amount@3 DESC NULLS LAST")
+ );
+ }
+
+ #[test]
+ fn test_ordering_stops_at_non_identity_transform() {
+ for transform in [Transform::Bucket(4), Transform::Truncate(2), Transform::Day] {
+ let source_id = if transform == Transform::Day { 3 } else { 2 };
+ let fields = vec![
+ asc(1),
+ field(
+ source_id,
+ transform,
+ SortDirection::Ascending,
+ NullOrder::First,
+ ),
+ asc(4),
+ ];
+ assert_eq!(
+ ordering(fields, ALL).as_deref(),
+ Some("id@0 ASC"),
+ "{transform}"
+ );
+ }
+ let fields = vec![field(
+ 1,
+ Transform::Bucket(4),
+ SortDirection::Ascending,
+ NullOrder::First,
+ )];
+ assert_eq!(ordering(fields, ALL), None);
+ }
+
+ #[test]
+ fn test_ordering_stops_at_floating_point_and_uuid() {
+ for source_id in [5, 6, 7] {
+ assert_eq!(
+ ordering(vec![asc(1), asc(source_id), asc(2)], ALL).as_deref(),
+ Some("id@0 ASC")
+ );
+ assert_eq!(
+ ordering(vec![asc(1), asc(source_id)], ALL).as_deref(),
+ Some("id@0 ASC")
+ );
+ assert_eq!(ordering(vec![asc(source_id), asc(1)], ALL), None);
+ }
+ }
+
+ #[test]
+ fn test_ordering_stops_at_nested_field() {
+ assert_eq!(
+ ordering(vec![asc(1), asc(9), asc(2)], ALL).as_deref(),
+ Some("id@0 ASC")
+ );
+ assert_eq!(ordering(vec![asc(9)], ALL), None);
+ }
+
+ #[test]
+ fn test_ordering_stops_at_unprojected_field() {
+ assert_eq!(
+ ordering(vec![asc(1), asc(2), asc(3)], &["ts", "id"]).as_deref(),
+ Some("id@1 ASC")
+ );
+ assert_eq!(ordering(vec![asc(2), asc(1)], &["id"]), None);
+ }
+
+ #[test]
+ fn test_ordering_stops_at_field_missing_from_schema() {
+ let mut builder = SortOrder::builder();
+ builder.with_order_id(1);
+ builder.with_sort_field(asc(1));
+ builder.with_sort_field(asc(42));
+ builder.with_sort_field(asc(2));
+ let order = builder.build_unbound().unwrap();
+ let schema = schema();
+ let output = schema_to_arrow_schema(&schema).unwrap();
+ assert_eq!(
+ lex_ordering(&order, &schema, &output).map(|o| o.to_string()),
+ Some("id@0 ASC".to_string())
+ );
+ }
}
diff --git a/crates/datafusion/src/table_provider/catalog_table_provider.rs b/crates/datafusion/src/table_provider/catalog_table_provider.rs
index b804f56..f7f6963 100644
--- a/crates/datafusion/src/table_provider/catalog_table_provider.rs
+++ b/crates/datafusion/src/table_provider/catalog_table_provider.rs
@@ -37,6 +37,7 @@ use iceberg::{Catalog, NamespaceIdent, TableIdent};
use crate::error::to_datafusion_error;
use crate::metadata::table::IcebergMetadataTableProvider;
use crate::read::IcebergDataSource;
+use crate::read::scan::plan_sorted_tasks;
use crate::write::commit::IcebergCommitExec;
use crate::write::exec::IcebergWriteExec;
use crate::write::project::project_with_partition;
@@ -126,7 +127,7 @@ impl TableProvider for IcebergCatalogTableProvider {
async fn scan(
&self,
- _state: &dyn Session,
+ state: &dyn Session,
projection: Option<&Vec>,
filters: &[Expr],
limit: Option,
@@ -139,11 +140,12 @@ impl TableProvider for IcebergCatalogTableProvider {
.map_err(to_datafusion_error)?;
// Create scan with fresh metadata (always use current snapshot)
- let source = IcebergDataSource::builder(table, self.schema.clone())
+ let builder = IcebergDataSource::builder(table, self.schema.clone())
.with_projection(projection)
.with_filters(filters)
- .with_fetch(limit)
- .build()?;
+ .with_fetch(limit);
+ let sorted_tasks = plan_sorted_tasks(&builder, state).await?;
+ let source = builder.with_sorted_tasks(sorted_tasks).build()?;
Ok(DataSourceExec::from_data_source(source))
}
diff --git a/crates/datafusion/src/table_provider/static_table_provider.rs b/crates/datafusion/src/table_provider/static_table_provider.rs
index 5ebabc8..c095fdd 100644
--- a/crates/datafusion/src/table_provider/static_table_provider.rs
+++ b/crates/datafusion/src/table_provider/static_table_provider.rs
@@ -32,6 +32,7 @@ use iceberg::table::Table;
use crate::error::to_datafusion_error;
use crate::read::IcebergDataSource;
+use crate::read::scan::plan_sorted_tasks;
/// Static table provider for read-only snapshot access.
///
@@ -120,18 +121,19 @@ impl TableProvider for IcebergStaticTableProvider {
async fn scan(
&self,
- _state: &dyn Session,
+ state: &dyn Session,
projection: Option<&Vec>,
filters: &[Expr],
limit: Option,
) -> Result> {
// Use cached table (no refresh)
- let source = IcebergDataSource::builder(self.table.clone(), self.schema.clone())
+ let builder = IcebergDataSource::builder(self.table.clone(), self.schema.clone())
.with_snapshot_id(self.snapshot_id)
.with_projection(projection)
.with_filters(filters)
- .with_fetch(limit)
- .build()?;
+ .with_fetch(limit);
+ let sorted_tasks = plan_sorted_tasks(&builder, state).await?;
+ let source = builder.with_sorted_tasks(sorted_tasks).build()?;
Ok(DataSourceExec::from_data_source(source))
}
diff --git a/crates/datafusion/tests/integration_datafusion_test.rs b/crates/datafusion/tests/integration_datafusion_test.rs
index 42d350a..5454c87 100644
--- a/crates/datafusion/tests/integration_datafusion_test.rs
+++ b/crates/datafusion/tests/integration_datafusion_test.rs
@@ -920,6 +920,7 @@ fn rebuild_scan(scan: &IcebergDataSource) -> IcebergDataSource {
.with_snapshot_id(scan.snapshot_id())
.with_predicate(scan.predicate().cloned())
.with_fetch(scan.fetch())
+ .with_sorted_tasks(scan.sorted_tasks().map(<[_]>::to_vec))
.build()
.unwrap()
}
diff --git a/crates/datafusion/tests/sorted_iceberg_scan.rs b/crates/datafusion/tests/sorted_iceberg_scan.rs
new file mode 100644
index 0000000..1b751c5
--- /dev/null
+++ b/crates/datafusion/tests/sorted_iceberg_scan.rs
@@ -0,0 +1,1002 @@
+// 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.
+
+//! Integration tests for sorted scans: with
+//! `iceberg.execution.preserve_data_ordering`, a scan of files that all record
+//! the same sort order reports that order to DataFusion and merges the files.
+
+use std::collections::HashMap;
+use std::error::Error;
+use std::str::FromStr;
+use std::sync::Arc;
+
+use datafusion::arrow::array::{
+ AsArray, Int32Array, Int64Array, ListArray, RecordBatch, StringArray, StructArray,
+};
+use datafusion::arrow::buffer::OffsetBuffer;
+use datafusion::arrow::compute::concat_batches;
+use datafusion::arrow::datatypes::{DataType, Field, Int32Type, Schema as ArrowSchema};
+use datafusion::datasource::TableProvider;
+use datafusion::datasource::source::{DataSource, DataSourceExec};
+use datafusion::execution::context::{SessionConfig, SessionContext};
+use datafusion::parquet::arrow::{ArrowWriter, PARQUET_FIELD_ID_META_KEY};
+use datafusion::physical_plan::common::collect;
+use datafusion::physical_plan::joins::SortMergeJoinExec;
+use datafusion::physical_plan::sorts::sort::SortExec;
+use datafusion::physical_plan::{ExecutionPlan, displayable};
+use futures::{StreamExt, TryStreamExt};
+use iceberg::arrow::schema_to_arrow_schema;
+use iceberg::io::LocalFsStorageFactory;
+use iceberg::memory::{MEMORY_CATALOG_WAREHOUSE, MemoryCatalogBuilder};
+use iceberg::scan::{FileScanTask, FileScanTaskDeleteFile};
+use iceberg::spec::{
+ DataContentType, DataFileBuilder, DataFileFormat, ListType, NestedField, NullOrder,
+ PrimitiveType, Schema, SortDirection, SortField, SortOrder, SortOrderRef, Struct,
+ StructType, Transform, Type,
+};
+use iceberg::table::Table;
+use iceberg::transaction::{ApplyTransactionAction, Transaction};
+use iceberg::{
+ Catalog, CatalogBuilder, MemoryCatalog, MetadataLocation, NamespaceIdent,
+ Result as IcebergResult, TableCreation,
+};
+use iceberg_datafusion::read::IcebergDataSource;
+use iceberg_datafusion::{
+ IcebergCatalogTableProvider, IcebergDataFusionConfig, IcebergStaticTableProvider,
+};
+use tempfile::TempDir;
+
+fn temp_path() -> String {
+ let temp_dir = TempDir::new().unwrap();
+ temp_dir.path().to_str().unwrap().to_string()
+}
+
+async fn get_iceberg_catalog() -> MemoryCatalog {
+ MemoryCatalogBuilder::default()
+ .with_storage_factory(Arc::new(LocalFsStorageFactory))
+ .load(
+ "memory",
+ HashMap::from([(MEMORY_CATALOG_WAREHOUSE.to_string(), temp_path())]),
+ )
+ .await
+ .unwrap()
+}
+
+async fn set_test_namespace(
+ catalog: &MemoryCatalog,
+ namespace: &NamespaceIdent,
+) -> IcebergResult<()> {
+ let properties = HashMap::new();
+
+ catalog.create_namespace(namespace, properties).await?;
+
+ Ok(())
+}
+
+/// Executes `plan`, which must have a single partition, and returns its rows.
+async fn run_batches(
+ plan: &dyn ExecutionPlan,
+ ctx: &SessionContext,
+) -> Result, Box> {
+ assert_eq!(plan.properties().partitioning.partition_count(), 1);
+ let stream = plan.execute(0, ctx.task_ctx())?;
+ Ok(collect(stream).await?)
+}
+
+/// Returns the first node of type `T` in `plan`, depth first.
+fn find_node(plan: &Arc) -> Option<&T> {
+ plan.downcast_ref::()
+ .or_else(|| plan.children().into_iter().find_map(find_node::))
+}
+
+/// Returns the Iceberg source of the first `DataSourceExec` in `plan`.
+fn iceberg_source(plan: &Arc) -> Option<&IcebergDataSource> {
+ find_node::(plan)?
+ .data_source()
+ .downcast_ref::()
+}
+
+/// Rebuilds `scan` from its accessors alone, as a codec would.
+fn rebuild_scan(scan: &IcebergDataSource) -> IcebergDataSource {
+ rebuild_scan_with_tasks(scan, scan.sorted_tasks().map(<[_]>::to_vec)).unwrap()
+}
+
+/// Rebuilds `scan` from its accessors, with `sorted_tasks` instead of its own.
+fn rebuild_scan_with_tasks(
+ scan: &IcebergDataSource,
+ sorted_tasks: Option>,
+) -> datafusion::error::Result {
+ IcebergDataSource::builder(scan.table().clone(), scan.schema())
+ .with_snapshot_id(scan.snapshot_id())
+ .with_predicate(scan.predicate().cloned())
+ .with_fetch(scan.fetch())
+ .with_sorted_tasks(sorted_tasks)
+ .build()
+}
+
+/// The id of the sort order of the tables `create_sorted_table` creates.
+const SORTED: Option = Some(1);
+
+/// Creates table `name`, sorted by `id` ascending, with columns `id`, `data`,
+/// a struct `info` and a list `tags`, whose values are derived from `id`. Each
+/// entry of `files` becomes one data file of the given ids, in the given order,
+/// which records the given sort order id.
+async fn create_sorted_table(
+ catalog: &Arc,
+ namespace: &NamespaceIdent,
+ name: &str,
+ files: &[(Option, &[i32])],
+) -> Result> {
+ let sort_order = SortOrder::builder()
+ .with_sort_field(
+ SortField::builder()
+ .source_id(1)
+ .transform(Transform::Identity)
+ .direction(SortDirection::Ascending)
+ .null_order(NullOrder::First)
+ .build(),
+ )
+ .build_unbound()?;
+ let creation = TableCreation::builder()
+ .location(temp_path())
+ .name(name.to_string())
+ .schema(
+ Schema::builder()
+ .with_fields(vec![
+ NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int))
+ .into(),
+ NestedField::required(
+ 2,
+ "data",
+ Type::Primitive(PrimitiveType::String),
+ )
+ .into(),
+ NestedField::optional(
+ 3,
+ "info",
+ Type::Struct(StructType::new(vec![
+ NestedField::optional(
+ 4,
+ "tag",
+ Type::Primitive(PrimitiveType::String),
+ )
+ .into(),
+ ])),
+ )
+ .into(),
+ NestedField::optional(
+ 5,
+ "tags",
+ Type::List(ListType::new(
+ NestedField::list_element(
+ 6,
+ Type::Primitive(PrimitiveType::Int),
+ false,
+ )
+ .into(),
+ )),
+ )
+ .into(),
+ ])
+ .build()?,
+ )
+ .sort_order(sort_order)
+ .build();
+ let table = catalog.create_table(namespace, creation).await?;
+ assert_eq!(
+ Some(table.metadata().default_sort_order_id() as i32),
+ SORTED
+ );
+ let files: Vec<_> = files
+ .iter()
+ .map(|(sort_order_id, ids)| {
+ (*sort_order_id, ids.iter().copied().map(Some).collect())
+ })
+ .collect();
+ append_sorted_files(catalog, &table, &files).await
+}
+
+/// Commits one data file of `table` per entry of `files`, with the given ids,
+/// in the given order, recording the given sort order id. The other columns'
+/// values are derived from `id`.
+async fn append_sorted_files(
+ catalog: &Arc,
+ table: &Table,
+ files: &[(Option, Vec>)],
+) -> Result> {
+ if files.is_empty() {
+ return Ok(table.clone());
+ }
+ let arrow_schema =
+ Arc::new(schema_to_arrow_schema(table.metadata().current_schema())?);
+ let (DataType::Struct(info_fields), DataType::List(tags_field)) = (
+ arrow_schema.field(2).data_type().clone(),
+ arrow_schema.field(3).data_type().clone(),
+ ) else {
+ unreachable!("info is a struct and tags a list");
+ };
+ let data_dir = format!("{}/data", table.metadata().location());
+ std::fs::create_dir_all(&data_dir)?;
+ let snapshots = table.metadata().snapshots().count();
+ let mut data_files = Vec::new();
+ for (i, (sort_order_id, ids)) in files.iter().enumerate() {
+ let path = format!("{data_dir}/{snapshots}-{i}.parquet");
+ let label = |prefix: &str, id: &Option| match id {
+ Some(id) => format!("{prefix} {id}"),
+ None => format!("{prefix} null"),
+ };
+ let batch = RecordBatch::try_new(
+ arrow_schema.clone(),
+ vec![
+ Arc::new(Int32Array::from(ids.clone())),
+ Arc::new(StringArray::from_iter_values(
+ ids.iter().map(|id| label("row", id)),
+ )),
+ Arc::new(StructArray::new(
+ info_fields.clone(),
+ vec![Arc::new(StringArray::from_iter_values(
+ ids.iter().map(|id| label("tag", id)),
+ ))],
+ None,
+ )),
+ Arc::new(ListArray::new(
+ tags_field.clone(),
+ OffsetBuffer::from_lengths(ids.iter().map(|_| 2)),
+ Arc::new(Int32Array::from_iter_values(ids.iter().flat_map(|id| {
+ let id = id.unwrap_or_default();
+ [id, -id]
+ }))),
+ None,
+ )),
+ ],
+ )?;
+ let mut writer = ArrowWriter::try_new(
+ std::fs::File::create(&path)?,
+ arrow_schema.clone(),
+ None,
+ )?;
+ writer.write(&batch)?;
+ writer.close()?;
+
+ let mut builder = DataFileBuilder::default();
+ builder
+ .content(DataContentType::Data)
+ .file_path(path.clone())
+ .file_format(DataFileFormat::Parquet)
+ .file_size_in_bytes(std::fs::metadata(&path)?.len())
+ .record_count(ids.len() as u64)
+ .partition_spec_id(0)
+ .partition(Struct::empty());
+ if let Some(sort_order_id) = sort_order_id {
+ builder.sort_order_id(*sort_order_id);
+ }
+ data_files.push(builder.build()?);
+ }
+ let tx = Transaction::new(table);
+ let tx = tx.fast_append().add_data_files(data_files).apply(tx)?;
+ Ok(tx.commit(catalog.as_ref()).await?)
+}
+
+/// Stands in for another engine committing, without writing data, the current
+/// schema of `table` with the field of the same id as `field` replaced by it.
+/// iceberg-rust has no action that makes a column optional or promotes its
+/// type, so this rewrites the table's current metadata file in place.
+async fn replace_field_externally(
+ table: &Table,
+ field: NestedField,
+) -> Result<(), Box> {
+ let current = table.metadata().current_schema();
+ let schema = Schema::builder()
+ .with_fields(current.as_struct().fields().iter().map(|current| {
+ if current.id == field.id {
+ Arc::new(field.clone())
+ } else {
+ current.clone()
+ }
+ }))
+ .build()?;
+ let metadata = table
+ .metadata()
+ .clone()
+ .into_builder(None)
+ .add_current_schema(schema)?
+ .build()?
+ .metadata;
+ let location = MetadataLocation::from_str(table.metadata_location_result()?)?;
+ metadata.write_to(table.file_io(), &location).await?;
+ Ok(())
+}
+
+/// The values, or nulls, of the `id` column, which must be the first, of
+/// `batches`.
+fn nullable_ids(batches: &[RecordBatch]) -> Vec> {
+ batches
+ .iter()
+ .flat_map(|batch| batch.column(0).as_primitive::().iter())
+ .collect()
+}
+
+/// A session that registers the Iceberg options, with
+/// `iceberg.execution.preserve_data_ordering` set to `preserve`.
+async fn ordering_session(
+ preserve: bool,
+ target_partitions: usize,
+) -> Result> {
+ let config = SessionConfig::new()
+ .with_target_partitions(target_partitions)
+ .with_option_extension(IcebergDataFusionConfig::default());
+ let ctx = SessionContext::new_with_config(config);
+ ctx.sql(&format!(
+ "SET iceberg.execution.preserve_data_ordering = {preserve}"
+ ))
+ .await?
+ .collect()
+ .await?;
+ Ok(ctx)
+}
+
+/// Plans `table` as a full scan in `ctx`.
+async fn plan_sorted_scan(
+ table: &Table,
+ ctx: &SessionContext,
+) -> Result, Box> {
+ let provider = IcebergStaticTableProvider::try_new_from_table(table.clone()).await?;
+ Ok(provider.scan(&ctx.state(), None, &[], None).await?)
+}
+
+/// The values of the `id` column, which must be the first, of `batches`.
+fn ids(batches: &[RecordBatch]) -> Vec {
+ batches
+ .iter()
+ .flat_map(|batch| {
+ batch
+ .column(0)
+ .as_primitive::()
+ .values()
+ .to_vec()
+ })
+ .collect()
+}
+
+async fn sorted_namespace() -> Result<(Arc, NamespaceIdent), Box>
+{
+ let catalog = get_iceberg_catalog().await;
+ let namespace = NamespaceIdent::new("sorted".to_string());
+ set_test_namespace(&catalog, &namespace).await?;
+ Ok((Arc::new(catalog), namespace))
+}
+
+#[tokio::test]
+async fn test_sorted_scan_merges_files_with_overlapping_ranges()
+-> Result<(), Box> {
+ let (catalog, namespace) = sorted_namespace().await?;
+ let table = create_sorted_table(
+ &catalog,
+ &namespace,
+ "t",
+ &[
+ (SORTED, &[1, 4, 7, 7]),
+ (SORTED, &[2, 4, 8]),
+ (SORTED, &[3, 5, 9]),
+ ],
+ )
+ .await?;
+ let ctx = ordering_session(true, 4).await?;
+
+ let plan = plan_sorted_scan(&table, &ctx).await?;
+ let scan = iceberg_source(&plan).unwrap();
+ assert_eq!(scan.sorted_tasks().map(<[_]>::len), Some(3));
+ assert_eq!(
+ plan.properties()
+ .output_ordering()
+ .map(|ordering| ordering.to_string()),
+ Some("id@0 ASC".to_string())
+ );
+ insta::assert_snapshot!(
+ displayable(plan.as_ref()).one_line().to_string().trim_end(),
+ @"DataSourceExec: format=iceberg, projection=[id, data, info, tags], output_ordering=[id@0 ASC], files=3"
+ );
+
+ let batches = run_batches(plan.as_ref(), &ctx).await?;
+ assert_eq!(ids(&batches), vec![1, 2, 3, 4, 4, 5, 7, 7, 8, 9]);
+ let batch = concat_batches(&batches[0].schema(), &batches)?;
+ for row in 0..batch.num_rows() {
+ let id = batch.column(0).as_primitive::().value(row);
+ assert_eq!(
+ batch.column(1).as_string::().value(row),
+ format!("row {id}")
+ );
+ assert_eq!(
+ batch
+ .column(2)
+ .as_struct()
+ .column(0)
+ .as_string::()
+ .value(row),
+ format!("tag {id}")
+ );
+ assert_eq!(
+ batch
+ .column(3)
+ .as_list::()
+ .value(row)
+ .as_primitive::()
+ .values(),
+ &[id, -id]
+ );
+ }
+ Ok(())
+}
+
+/// Each file holds more rows than a reader batch, so the merge advances every
+/// per-file stream across batch boundaries, where the order within a stream
+/// comes from the reader rather than the merge.
+#[tokio::test]
+async fn test_sorted_scan_merges_files_longer_than_a_batch() -> Result<(), Box>
+{
+ let (catalog, namespace) = sorted_namespace().await?;
+ let evens: Vec = (0..6000).step_by(2).collect();
+ let odds: Vec = (1..6000).step_by(2).collect();
+ let table = create_sorted_table(
+ &catalog,
+ &namespace,
+ "t",
+ &[(SORTED, &evens), (SORTED, &odds)],
+ )
+ .await?;
+ let ctx = ordering_session(true, 4).await?;
+ let plan = plan_sorted_scan(&table, &ctx).await?;
+ let scan = iceberg_source(&plan).unwrap();
+
+ for task in scan.sorted_tasks().unwrap() {
+ let batches: Vec = table
+ .reader_builder()
+ .build()
+ .read(futures::stream::iter([Ok(task.clone())]).boxed())?
+ .stream()
+ .try_collect()
+ .await?;
+ assert!(batches.len() > 1, "{}", task.data_file_path());
+ }
+ assert_eq!(
+ ids(&run_batches(plan.as_ref(), &ctx).await?),
+ (0..6000).collect::>()
+ );
+ Ok(())
+}
+
+#[tokio::test]
+async fn test_sorted_scan_merge_honors_limit() -> Result<(), Box> {
+ let (catalog, namespace) = sorted_namespace().await?;
+ let table = create_sorted_table(
+ &catalog,
+ &namespace,
+ "t",
+ &[
+ (SORTED, &[1, 4, 7]),
+ (SORTED, &[2, 5, 8]),
+ (SORTED, &[3, 6, 9]),
+ ],
+ )
+ .await?;
+ let ctx = ordering_session(true, 4).await?;
+ let provider = IcebergStaticTableProvider::try_new_from_table(table).await?;
+
+ let plan = provider.scan(&ctx.state(), None, &[], Some(4)).await?;
+ assert!(plan.properties().output_ordering().is_some());
+ assert_eq!(
+ ids(&run_batches(plan.as_ref(), &ctx).await?),
+ vec![1, 2, 3, 4]
+ );
+ Ok(())
+}
+
+/// A scan reports no order, and is planned as it is without the option, when
+/// its files do not all record the same resolvable sort order.
+#[tokio::test]
+async fn test_sorted_scan_reports_no_order_unless_files_agree()
+-> Result<(), Box> {
+ let (catalog, namespace) = sorted_namespace().await?;
+ let cases: [(&str, [Option; 2]); 5] = [
+ ("unresolvable", [SORTED, Some(99)]),
+ ("missing", [SORTED, None]),
+ ("unsorted", [SORTED, Some(0)]),
+ ("all_unsorted", [Some(0), Some(0)]),
+ ("all_missing", [None, None]),
+ ];
+ let on = ordering_session(true, 4).await?;
+ let off = ordering_session(false, 4).await?;
+ for (name, [first, second]) in cases {
+ let table = create_sorted_table(
+ &catalog,
+ &namespace,
+ name,
+ &[(first, &[1, 3]), (second, &[2, 4])],
+ )
+ .await?;
+
+ let plan = plan_sorted_scan(&table, &on).await?;
+ let scan = iceberg_source(&plan).unwrap();
+ assert!(scan.sorted_tasks().is_none(), "{name}");
+ assert!(plan.properties().output_ordering().is_none(), "{name}");
+ assert_eq!(plan.metrics().map(|m| m.iter().count()), Some(0), "{name}");
+ let unordered = plan_sorted_scan(&table, &off).await?;
+ assert_eq!(
+ displayable(plan.as_ref()).indent(true).to_string(),
+ displayable(unordered.as_ref()).indent(true).to_string(),
+ "{name}"
+ );
+
+ let mut rows = ids(&run_batches(plan.as_ref(), &on).await?);
+ rows.sort();
+ assert_eq!(rows, vec![1, 2, 3, 4], "{name}");
+ }
+ Ok(())
+}
+
+#[tokio::test]
+async fn test_sorted_scan_is_opt_in_and_capped() -> Result<(), Box> {
+ let (catalog, namespace) = sorted_namespace().await?;
+ let table = create_sorted_table(
+ &catalog,
+ &namespace,
+ "t",
+ &[(SORTED, &[1]), (SORTED, &[2]), (SORTED, &[3])],
+ )
+ .await?;
+
+ let unregistered = SessionContext::new();
+ let plan = plan_sorted_scan(&table, &unregistered).await?;
+ assert!(plan.properties().output_ordering().is_none());
+
+ let off = ordering_session(false, 4).await?;
+ let plan = plan_sorted_scan(&table, &off).await?;
+ assert!(plan.properties().output_ordering().is_none());
+
+ let capped = ordering_session(true, 4).await?;
+ capped
+ .sql("SET iceberg.execution.max_merge_files_per_partition = 2")
+ .await?
+ .collect()
+ .await?;
+ let plan = plan_sorted_scan(&table, &capped).await?;
+ assert!(plan.properties().output_ordering().is_none());
+
+ capped
+ .sql("SET iceberg.execution.max_merge_files_per_partition = 3")
+ .await?
+ .collect()
+ .await?;
+ let plan = plan_sorted_scan(&table, &capped).await?;
+ assert!(plan.properties().output_ordering().is_some());
+ Ok(())
+}
+
+/// Plans and runs `sql` in `ctx`, returning the plan that ran and the ids it
+/// returned.
+async fn plan_and_run(
+ ctx: &SessionContext,
+ sql: &str,
+) -> Result<(Arc, Vec), Box> {
+ let plan = ctx.sql(sql).await?.create_physical_plan().await?;
+ let batches =
+ datafusion::physical_plan::collect(plan.clone(), ctx.task_ctx()).await?;
+ Ok((plan, ids(&batches)))
+}
+
+#[tokio::test]
+async fn test_sorted_scan_removes_sort_from_order_by() -> Result<(), Box> {
+ let (catalog, namespace) = sorted_namespace().await?;
+ let table = create_sorted_table(
+ &catalog,
+ &namespace,
+ "t",
+ &[
+ (SORTED, &[1, 4, 7]),
+ (SORTED, &[2, 5, 8]),
+ (SORTED, &[3, 6, 9]),
+ ],
+ )
+ .await?;
+ let provider = Arc::new(IcebergStaticTableProvider::try_new_from_table(table).await?);
+ let on = ordering_session(true, 4).await?;
+ on.register_table("t", provider.clone())?;
+ let off = ordering_session(false, 4).await?;
+ off.register_table("t", provider)?;
+
+ let (plan, rows) = plan_and_run(&on, "SELECT id FROM t ORDER BY id").await?;
+ assert!(find_node::(&plan).is_none());
+ assert_eq!(rows, (1..=9).collect::>());
+
+ let (plan, rows) = plan_and_run(&on, "SELECT id FROM t ORDER BY id LIMIT 4").await?;
+ assert!(find_node::(&plan).is_none());
+ assert_eq!(rows, vec![1, 2, 3, 4]);
+
+ let (plan, rows) = plan_and_run(&on, "SELECT id FROM t ORDER BY id DESC").await?;
+ assert!(find_node::(&plan).is_some());
+ assert_eq!(rows, (1..=9).rev().collect::>());
+
+ let (plan, rows) = plan_and_run(&off, "SELECT id FROM t ORDER BY id").await?;
+ assert!(find_node::(&plan).is_some());
+ assert_eq!(rows, (1..=9).collect::>());
+ Ok(())
+}
+
+#[tokio::test]
+async fn test_sorted_scan_removes_sort_from_sort_merge_join() -> Result<(), Box>
+{
+ let (catalog, namespace) = sorted_namespace().await?;
+ let left = create_sorted_table(
+ &catalog,
+ &namespace,
+ "l",
+ &[(SORTED, &[1, 3, 5]), (SORTED, &[2, 4, 6])],
+ )
+ .await?;
+ let right = create_sorted_table(
+ &catalog,
+ &namespace,
+ "r",
+ &[(SORTED, &[2, 5]), (SORTED, &[3, 6, 7])],
+ )
+ .await?;
+ let sql = "SELECT l.id FROM l JOIN r ON l.id = r.id ORDER BY l.id";
+ for preserve in [true, false] {
+ let ctx = ordering_session(preserve, 4).await?;
+ ctx.sql("SET datafusion.optimizer.prefer_hash_join = false")
+ .await?
+ .collect()
+ .await?;
+ for (name, table) in [("l", &left), ("r", &right)] {
+ ctx.register_table(
+ name,
+ Arc::new(
+ IcebergStaticTableProvider::try_new_from_table(table.clone()).await?,
+ ),
+ )?;
+ }
+
+ let (plan, rows) = plan_and_run(&ctx, sql).await?;
+ assert!(find_node::(&plan).is_some());
+ assert_eq!(find_node::(&plan).is_none(), preserve);
+ assert_eq!(rows, vec![2, 3, 5, 6]);
+ }
+ Ok(())
+}
+
+/// A scan rebuilt from the accessors of a sorted scan reports and keeps the
+/// same order.
+#[tokio::test]
+async fn test_rebuilt_sorted_scan_keeps_order() -> Result<(), Box> {
+ let (catalog, namespace) = sorted_namespace().await?;
+ let table = create_sorted_table(
+ &catalog,
+ &namespace,
+ "t",
+ &[(SORTED, &[1, 3]), (SORTED, &[2, 4])],
+ )
+ .await?;
+ let ctx = ordering_session(true, 4).await?;
+ let plan = plan_sorted_scan(&table, &ctx).await?;
+ let scan = iceberg_source(&plan).unwrap();
+ assert_eq!(scan.sorted_tasks().map(<[_]>::len), Some(2));
+
+ let rebuilt = rebuild_scan(scan);
+ assert_eq!(
+ rebuilt.eq_properties().output_ordering().as_ref(),
+ plan.properties().output_ordering()
+ );
+ let rebuilt = DataSourceExec::from_data_source(rebuilt);
+ assert_eq!(
+ ids(&run_batches(rebuilt.as_ref(), &ctx).await?),
+ vec![1, 2, 3, 4]
+ );
+
+ // Without the tasks a rebuilt source lists its files when it runs, and
+ // reports no order.
+ let unsorted = rebuild_scan_with_tasks(scan, None)?;
+ assert!(unsorted.sorted_tasks().is_none());
+ assert!(unsorted.eq_properties().output_ordering().is_none());
+
+ assert!(rebuild_scan_with_tasks(scan, Some(vec![])).is_err());
+ Ok(())
+}
+
+/// Returns a copy of `task` with `deletes` and `sort_order` in place of its own.
+fn copy_task(
+ task: &FileScanTask,
+ deletes: Vec,
+ sort_order: Option,
+) -> Result> {
+ Ok(FileScanTask::builder()
+ .with_file_size_in_bytes(task.file_size_in_bytes())
+ .with_start(task.start())
+ .with_length(task.length())
+ .with_record_count(task.record_count())
+ .with_data_file_path(task.data_file_path().to_string())
+ .with_data_file_format(task.data_file_format())
+ .with_schema(task.schema_ref())
+ .with_project_field_ids(task.project_field_ids().to_vec())
+ .with_predicate(task.predicate().cloned())
+ .with_deletes(deletes)
+ .with_partition(task.partition().cloned())
+ .with_partition_spec(task.partition_spec().cloned())
+ .with_name_mapping(task.name_mapping().cloned())
+ .with_sort_order_id(task.sort_order_id())
+ .with_sort_order(sort_order)
+ .with_case_sensitive(task.case_sensitive())
+ .build()?)
+}
+
+#[tokio::test]
+async fn test_sorted_scan_applies_position_deletes() -> Result<(), Box> {
+ let (catalog, namespace) = sorted_namespace().await?;
+ let table = create_sorted_table(
+ &catalog,
+ &namespace,
+ "t",
+ &[(SORTED, &[1, 4, 7]), (SORTED, &[2, 5, 8])],
+ )
+ .await?;
+ let ctx = ordering_session(true, 4).await?;
+ let plan = plan_sorted_scan(&table, &ctx).await?;
+ let scan = iceberg_source(&plan).unwrap();
+ let tasks = scan.sorted_tasks().unwrap();
+
+ // Deletes row 1 of each data file, ids 4 and 5.
+ let delete_schema = Arc::new(ArrowSchema::new(vec![
+ Field::new("file_path", DataType::Utf8, false).with_metadata(HashMap::from([(
+ PARQUET_FIELD_ID_META_KEY.to_string(),
+ "2147483546".to_string(),
+ )])),
+ Field::new("pos", DataType::Int64, false).with_metadata(HashMap::from([(
+ PARQUET_FIELD_ID_META_KEY.to_string(),
+ "2147483545".to_string(),
+ )])),
+ ]));
+ let mut paths: Vec<&str> = tasks.iter().map(|task| task.data_file_path()).collect();
+ paths.sort();
+ let delete_path = format!("{}/data/deletes.parquet", table.metadata().location());
+ let mut writer = ArrowWriter::try_new(
+ std::fs::File::create(&delete_path)?,
+ delete_schema.clone(),
+ None,
+ )?;
+ writer.write(&RecordBatch::try_new(
+ delete_schema,
+ vec![
+ Arc::new(StringArray::from(paths)),
+ Arc::new(Int64Array::from(vec![1, 1])),
+ ],
+ )?)?;
+ writer.close()?;
+ let delete = FileScanTaskDeleteFile::builder()
+ .with_file_path(delete_path.clone())
+ .with_file_size_in_bytes(std::fs::metadata(&delete_path)?.len())
+ .with_file_type(DataContentType::PositionDeletes)
+ .with_file_format(DataFileFormat::Parquet)
+ .with_partition_spec_id(0)
+ .build()?;
+ let tasks = tasks
+ .iter()
+ .map(|task| copy_task(task, vec![delete.clone()], task.sort_order().cloned()))
+ .collect::, _>>()?;
+
+ let rebuilt =
+ DataSourceExec::from_data_source(rebuild_scan_with_tasks(scan, Some(tasks))?);
+ assert_eq!(
+ ids(&run_batches(rebuilt.as_ref(), &ctx).await?),
+ vec![1, 2, 7, 8]
+ );
+ Ok(())
+}
+
+/// A provider created while `id` is required declares it non-nullable, and
+/// DataFusion ignores the null order of a non-nullable column. Once another
+/// engine makes `id` optional and writes a NULL, the scan must report no order,
+/// or `ORDER BY id`, which puts NULLs last, would keep the file's NULL first.
+#[tokio::test]
+async fn test_sorted_scan_reports_no_order_for_stale_nullability()
+-> Result<(), Box> {
+ let (catalog, namespace) = sorted_namespace().await?;
+ let table = create_sorted_table(&catalog, &namespace, "t", &[]).await?;
+ let provider = Arc::new(
+ IcebergCatalogTableProvider::try_new(catalog.clone(), namespace.clone(), "t")
+ .await?,
+ );
+ let id = NestedField::optional(1, "id", Type::Primitive(PrimitiveType::Int));
+ replace_field_externally(&table, id).await?;
+ let table = catalog.load_table(table.identifier()).await?;
+ append_sorted_files(&catalog, &table, &[(SORTED, vec![None, Some(0), Some(7)])])
+ .await?;
+
+ for preserve in [false, true] {
+ let ctx = ordering_session(preserve, 4).await?;
+ ctx.register_table("t", provider.clone())?;
+ let plan = ctx
+ .sql("SELECT id FROM t ORDER BY id")
+ .await?
+ .create_physical_plan()
+ .await?;
+ let scan = iceberg_source(&plan).unwrap();
+ assert!(scan.sorted_tasks().is_none(), "preserve = {preserve}");
+ let batches =
+ datafusion::physical_plan::collect(plan.clone(), ctx.task_ctx()).await?;
+ assert_eq!(
+ nullable_ids(&batches),
+ vec![Some(0), Some(7), None],
+ "preserve = {preserve}"
+ );
+ }
+ Ok(())
+}
+
+/// Another engine promotes `id` from int to long without writing data. A
+/// provider created afterwards declares a long `id`, while the scan reads the
+/// snapshot's int, which the merge could not compare as declared.
+#[tokio::test]
+async fn test_sorted_scan_reports_no_order_for_promoted_type()
+-> Result<(), Box> {
+ let (catalog, namespace) = sorted_namespace().await?;
+ let table = create_sorted_table(
+ &catalog,
+ &namespace,
+ "t",
+ &[(SORTED, &[1, 3]), (SORTED, &[2, 4])],
+ )
+ .await?;
+ let id = NestedField::required(1, "id", Type::Primitive(PrimitiveType::Long));
+ replace_field_externally(&table, id).await?;
+ let provider = Arc::new(
+ IcebergCatalogTableProvider::try_new(catalog.clone(), namespace, "t").await?,
+ );
+
+ for preserve in [false, true] {
+ let ctx = ordering_session(preserve, 4).await?;
+ ctx.register_table("t", provider.clone())?;
+ let plan = ctx
+ .sql("SELECT * FROM t")
+ .await?
+ .create_physical_plan()
+ .await?;
+ let scan = iceberg_source(&plan).unwrap();
+ assert!(scan.sorted_tasks().is_none(), "preserve = {preserve}");
+ let batches =
+ datafusion::physical_plan::collect(plan.clone(), ctx.task_ctx()).await?;
+ let rows: usize = batches.iter().map(RecordBatch::num_rows).sum();
+ assert_eq!(rows, 4, "preserve = {preserve}");
+ }
+ Ok(())
+}
+
+/// Iceberg sorts ascending fields with NULLs first by default, while
+/// DataFusion's `ORDER BY id` puts them last, so on a nullable sort column only
+/// `ORDER BY id NULLS FIRST` can use the order the scan reports.
+#[tokio::test]
+async fn test_sorted_scan_on_nullable_key() -> Result<(), Box> {
+ let (catalog, namespace) = sorted_namespace().await?;
+ let table = create_sorted_table(&catalog, &namespace, "t", &[]).await?;
+ let id = NestedField::optional(1, "id", Type::Primitive(PrimitiveType::Int));
+ replace_field_externally(&table, id).await?;
+ let table = catalog.load_table(table.identifier()).await?;
+ let table = append_sorted_files(
+ &catalog,
+ &table,
+ &[
+ (SORTED, vec![None, Some(1), Some(4)]),
+ (SORTED, vec![None, Some(2), Some(3)]),
+ ],
+ )
+ .await?;
+ let ctx = ordering_session(true, 4).await?;
+ ctx.register_table(
+ "t",
+ Arc::new(IcebergStaticTableProvider::try_new_from_table(table).await?),
+ )?;
+
+ for (sql, sorts, expected) in [
+ (
+ "SELECT id FROM t ORDER BY id NULLS FIRST",
+ false,
+ vec![None, None, Some(1), Some(2), Some(3), Some(4)],
+ ),
+ (
+ "SELECT id FROM t ORDER BY id",
+ true,
+ vec![Some(1), Some(2), Some(3), Some(4), None, None],
+ ),
+ ] {
+ let plan = ctx.sql(sql).await?.create_physical_plan().await?;
+ let scan = iceberg_source(&plan).unwrap();
+ assert_eq!(scan.sorted_tasks().map(<[_]>::len), Some(2), "{sql}");
+ assert_eq!(find_node::(&plan).is_some(), sorts, "{sql}");
+ let batches =
+ datafusion::physical_plan::collect(plan.clone(), ctx.task_ctx()).await?;
+ assert_eq!(nullable_ids(&batches), expected, "{sql}");
+ }
+ Ok(())
+}
+
+/// A sorted scan of one file reports its order and returns its rows without a
+/// merge, and a scan of a sorted table with no files reports no order.
+#[tokio::test]
+async fn test_sorted_scan_of_one_file_and_of_no_files() -> Result<(), Box> {
+ let (catalog, namespace) = sorted_namespace().await?;
+ let ctx = ordering_session(true, 4).await?;
+
+ let one =
+ create_sorted_table(&catalog, &namespace, "one", &[(SORTED, &[1, 2, 3])]).await?;
+ let plan = plan_sorted_scan(&one, &ctx).await?;
+ assert_eq!(
+ iceberg_source(&plan)
+ .unwrap()
+ .sorted_tasks()
+ .map(<[_]>::len),
+ Some(1)
+ );
+ assert_eq!(
+ plan.properties()
+ .output_ordering()
+ .map(|ordering| ordering.to_string()),
+ Some("id@0 ASC".to_string())
+ );
+ assert_eq!(ids(&run_batches(plan.as_ref(), &ctx).await?), vec![1, 2, 3]);
+
+ let empty = create_sorted_table(&catalog, &namespace, "empty", &[]).await?;
+ let plan = plan_sorted_scan(&empty, &ctx).await?;
+ assert!(iceberg_source(&plan).unwrap().sorted_tasks().is_none());
+ assert!(plan.properties().output_ordering().is_none());
+ assert!(ids(&run_batches(plan.as_ref(), &ctx).await?).is_empty());
+ Ok(())
+}
+
+/// Tasks given to the builder must all record the same sort order, not only
+/// the same sort order id, because the merge uses one order for all of them.
+#[tokio::test]
+async fn test_rebuilt_sorted_scan_rejects_tasks_with_different_orders()
+-> Result<(), Box> {
+ let (catalog, namespace) = sorted_namespace().await?;
+ let table = create_sorted_table(
+ &catalog,
+ &namespace,
+ "t",
+ &[(SORTED, &[1, 3]), (SORTED, &[2, 4])],
+ )
+ .await?;
+ let ctx = ordering_session(true, 4).await?;
+ let plan = plan_sorted_scan(&table, &ctx).await?;
+ let scan = iceberg_source(&plan).unwrap();
+ let mut tasks = scan.sorted_tasks().unwrap().to_vec();
+
+ let descending = SortOrder::builder()
+ .with_order_id(1)
+ .with_sort_field(
+ SortField::builder()
+ .source_id(1)
+ .transform(Transform::Identity)
+ .direction(SortDirection::Descending)
+ .null_order(NullOrder::Last)
+ .build(),
+ )
+ .build_unbound()?;
+ tasks[1] = copy_task(&tasks[1], vec![], Some(Arc::new(descending)))?;
+ assert_eq!(tasks[1].sort_order_id(), tasks[0].sort_order_id());
+ assert_ne!(tasks[1].sort_order(), tasks[0].sort_order());
+
+ assert!(rebuild_scan_with_tasks(scan, Some(tasks)).is_err());
+ Ok(())
+}
diff --git a/crates/playground/src/main.rs b/crates/playground/src/main.rs
index 853494c..308d297 100644
--- a/crates/playground/src/main.rs
+++ b/crates/playground/src/main.rs
@@ -27,6 +27,7 @@ use datafusion_cli::exec;
use datafusion_cli::object_storage::instrumented::InstrumentedObjectStoreRegistry;
use datafusion_cli::print_format::PrintFormat;
use datafusion_cli::print_options::{MaxRows, PrintOptions};
+use iceberg_datafusion::IcebergDataFusionConfig;
use iceberg_playground::{ICEBERG_PLAYGROUND_VERSION, IcebergCatalogList};
#[derive(Debug, Parser, PartialEq)]
@@ -80,7 +81,9 @@ async fn main_inner() -> anyhow::Result<()> {
println!("ICEBERG PLAYGROUND v{ICEBERG_PLAYGROUND_VERSION}");
}
- let session_config = SessionConfig::from_env()?.with_information_schema(true);
+ let session_config = SessionConfig::from_env()?
+ .with_information_schema(true)
+ .with_option_extension(IcebergDataFusionConfig::default());
let rt_builder = RuntimeEnvBuilder::new();
diff --git a/crates/sqllogictest/src/engine/datafusion.rs b/crates/sqllogictest/src/engine/datafusion.rs
index 6617023..85d0de2 100644
--- a/crates/sqllogictest/src/engine/datafusion.rs
+++ b/crates/sqllogictest/src/engine/datafusion.rs
@@ -19,17 +19,22 @@ use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::Arc;
+use datafusion::arrow::array::{Int32Array, RecordBatch, StringArray};
use datafusion::catalog::CatalogProvider;
+use datafusion::parquet::arrow::ArrowWriter;
use datafusion::prelude::{SessionConfig, SessionContext};
use datafusion_sqllogictest::DataFusion;
+use iceberg::arrow::schema_to_arrow_schema;
use iceberg::encryption::kms::MemoryKmsClientFactory;
use iceberg::memory::{MEMORY_CATALOG_WAREHOUSE, MemoryCatalogBuilder};
use iceberg::spec::{
- NestedField, PrimitiveType, Schema, TableProperties, Transform, Type,
- UnboundPartitionField, UnboundPartitionSpec,
+ DataContentType, DataFileBuilder, DataFileFormat, NestedField, NullOrder,
+ PrimitiveType, Schema, SortDirection, SortField, SortOrder, Struct, TableProperties,
+ Transform, Type, UnboundPartitionField, UnboundPartitionSpec,
};
+use iceberg::transaction::{ApplyTransactionAction, Transaction};
use iceberg::{Catalog, CatalogBuilder, NamespaceIdent, TableCreation};
-use iceberg_datafusion::IcebergCatalogProvider;
+use iceberg_datafusion::{IcebergCatalogProvider, IcebergDataFusionConfig};
use indicatif::ProgressBar;
use crate::engine::{DatafusionCatalogConfig, EngineRunner, run_slt_with_runner};
@@ -67,7 +72,8 @@ impl DataFusionEngine {
pub async fn new(catalog_config: Option) -> Result {
let session_config = SessionConfig::new()
.with_target_partitions(4)
- .with_information_schema(true);
+ .with_information_schema(true)
+ .with_option_extension(IcebergDataFusionConfig::default());
let ctx = SessionContext::new_with_config(session_config);
ctx.register_catalog(
"default",
@@ -107,6 +113,7 @@ impl DataFusionEngine {
Self::create_partitioned_table(&catalog, &namespace).await?;
Self::create_binary_table(&catalog, &namespace).await?;
Self::create_encrypted_table(&catalog, &namespace).await?;
+ Self::create_sorted_table(&catalog, &namespace).await?;
Ok(Arc::new(
IcebergCatalogProvider::try_new(Arc::new(catalog)).await?,
@@ -219,4 +226,86 @@ impl DataFusionEngine {
Ok(())
}
+
+ /// Create a table sorted by `id`, with two data files whose ids overlap,
+ /// each recording that sort order
+ /// TODO: this can be removed when INSERT writes sorted data files
+ async fn create_sorted_table(
+ catalog: &impl Catalog,
+ namespace: &NamespaceIdent,
+ ) -> anyhow::Result<()> {
+ let schema = Schema::builder()
+ .with_fields(vec![
+ NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int))
+ .into(),
+ NestedField::optional(2, "name", Type::Primitive(PrimitiveType::String))
+ .into(),
+ ])
+ .build()?;
+ let sort_order = SortOrder::builder()
+ .with_sort_field(
+ SortField::builder()
+ .source_id(1)
+ .transform(Transform::Identity)
+ .direction(SortDirection::Ascending)
+ .null_order(NullOrder::First)
+ .build(),
+ )
+ .build_unbound()?;
+ let table = catalog
+ .create_table(
+ namespace,
+ TableCreation::builder()
+ .name("test_sorted_table".to_string())
+ .schema(schema)
+ .sort_order(sort_order)
+ .build(),
+ )
+ .await?;
+
+ let arrow_schema =
+ Arc::new(schema_to_arrow_schema(table.metadata().current_schema())?);
+ let mut data_files = Vec::new();
+ for (i, ids) in [[1, 3, 5], [2, 3, 4]].iter().enumerate() {
+ let batch = RecordBatch::try_new(
+ arrow_schema.clone(),
+ vec![
+ Arc::new(Int32Array::from(ids.to_vec())),
+ Arc::new(StringArray::from_iter_values(
+ ids.iter().map(|id| format!("name {id}")),
+ )),
+ ],
+ )?;
+ let mut bytes = Vec::new();
+ let mut writer =
+ ArrowWriter::try_new(&mut bytes, arrow_schema.clone(), None)?;
+ writer.write(&batch)?;
+ writer.close()?;
+
+ let path = format!("{}/data/sorted-{i}.parquet", table.metadata().location());
+ let file_size_in_bytes = bytes.len() as u64;
+ table
+ .file_io()
+ .new_output(&path)?
+ .write(bytes.into())
+ .await?;
+ data_files.push(
+ DataFileBuilder::default()
+ .content(DataContentType::Data)
+ .file_path(path)
+ .file_format(DataFileFormat::Parquet)
+ .file_size_in_bytes(file_size_in_bytes)
+ .record_count(ids.len() as u64)
+ .partition_spec_id(0)
+ .partition(Struct::empty())
+ .sort_order_id(table.metadata().default_sort_order_id() as i32)
+ .build()?,
+ );
+ }
+ let tx = Transaction::new(&table);
+ let tx = tx.fast_append().add_data_files(data_files).apply(tx)?;
+ tx.commit(catalog).await?;
+
+ Ok(())
+ }
}
diff --git a/crates/sqllogictest/testdata/schedules/df_sort_order_reporting.toml b/crates/sqllogictest/testdata/schedules/df_sort_order_reporting.toml
new file mode 100644
index 0000000..bd34be3
--- /dev/null
+++ b/crates/sqllogictest/testdata/schedules/df_sort_order_reporting.toml
@@ -0,0 +1,23 @@
+# 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.
+
+[engines]
+df = { type = "datafusion" }
+
+[[steps]]
+engine = "df"
+slt = "df_test/sort_order_reporting.slt"
diff --git a/crates/sqllogictest/testdata/slts/df_test/show_tables.slt b/crates/sqllogictest/testdata/slts/df_test/show_tables.slt
index 29cc3cd..93c2a84 100644
--- a/crates/sqllogictest/testdata/slts/df_test/show_tables.slt
+++ b/crates/sqllogictest/testdata/slts/df_test/show_tables.slt
@@ -37,6 +37,10 @@ default default test_partitioned_table BASE TABLE
default default test_partitioned_table$history BASE TABLE
default default test_partitioned_table$manifests BASE TABLE
default default test_partitioned_table$snapshots BASE TABLE
+default default test_sorted_table BASE TABLE
+default default test_sorted_table$history BASE TABLE
+default default test_sorted_table$manifests BASE TABLE
+default default test_sorted_table$snapshots BASE TABLE
default information_schema columns VIEW
default information_schema df_settings VIEW
default information_schema parameters VIEW
diff --git a/crates/sqllogictest/testdata/slts/df_test/sort_order_reporting.slt b/crates/sqllogictest/testdata/slts/df_test/sort_order_reporting.slt
new file mode 100644
index 0000000..22b6a0f
--- /dev/null
+++ b/crates/sqllogictest/testdata/slts/df_test/sort_order_reporting.slt
@@ -0,0 +1,109 @@
+# 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.
+
+# test_sorted_table is sorted by id, and its two data files hold ids 1, 3, 5
+# and 2, 3, 4, each recording that sort order.
+
+# Without iceberg.execution.preserve_data_ordering the scan reports no order.
+query TT
+EXPLAIN SELECT id, name FROM default.default.test_sorted_table ORDER BY id
+----
+logical_plan
+01)Sort: default.default.test_sorted_table.id ASC NULLS LAST
+02)--TableScan: default.default.test_sorted_table projection=[id, name]
+physical_plan
+01)SortExec: expr=[id@0 ASC NULLS LAST], preserve_partitioning=[false]
+02)--CooperativeExec
+03)----DataSourceExec: format=iceberg, projection=[id, name]
+
+statement ok
+SET iceberg.execution.preserve_data_ordering = true
+
+# With it, the scan merges the files in their sort order, and reports it.
+query TT
+EXPLAIN SELECT id, name FROM default.default.test_sorted_table ORDER BY id
+----
+logical_plan
+01)Sort: default.default.test_sorted_table.id ASC NULLS LAST
+02)--TableScan: default.default.test_sorted_table projection=[id, name]
+physical_plan
+01)CooperativeExec
+02)--DataSourceExec: format=iceberg, projection=[id, name], output_ordering=[id@0 ASC], files=2
+
+query IT
+SELECT id, name FROM default.default.test_sorted_table ORDER BY id
+----
+1 name 1
+2 name 2
+3 name 3
+3 name 3
+4 name 4
+5 name 5
+
+query IT
+SELECT id, name FROM default.default.test_sorted_table ORDER BY id LIMIT 3
+----
+1 name 1
+2 name 2
+3 name 3
+
+# A descending sort cannot use the scan's order, so DataFusion sorts the
+# merged files again.
+query TT
+EXPLAIN SELECT id FROM default.default.test_sorted_table ORDER BY id DESC
+----
+logical_plan
+01)Sort: default.default.test_sorted_table.id DESC NULLS FIRST
+02)--TableScan: default.default.test_sorted_table projection=[id]
+physical_plan
+01)SortExec: expr=[id@0 DESC], preserve_partitioning=[false]
+02)--CooperativeExec
+03)----DataSourceExec: format=iceberg, projection=[id], output_ordering=[id@0 ASC], files=2
+
+# Tables created through SQL have no sort order, so scans of them report none.
+statement ok
+CREATE TABLE default.default.test_unsorted_table (id INT NOT NULL, name STRING)
+
+statement ok
+INSERT INTO default.default.test_unsorted_table VALUES (2, 'b'), (1, 'a')
+
+query TT
+EXPLAIN SELECT id FROM default.default.test_unsorted_table ORDER BY id
+----
+logical_plan
+01)Sort: default.default.test_unsorted_table.id ASC NULLS LAST
+02)--TableScan: default.default.test_unsorted_table projection=[id]
+physical_plan
+01)SortExec: expr=[id@0 ASC NULLS LAST], preserve_partitioning=[false]
+02)--CooperativeExec
+03)----DataSourceExec: format=iceberg, projection=[id]
+
+# A scan of more files than iceberg.execution.max_merge_files_per_partition reports
+# no order.
+statement ok
+SET iceberg.execution.max_merge_files_per_partition = 1
+
+query TT
+EXPLAIN SELECT id, name FROM default.default.test_sorted_table ORDER BY id
+----
+logical_plan
+01)Sort: default.default.test_sorted_table.id ASC NULLS LAST
+02)--TableScan: default.default.test_sorted_table projection=[id, name]
+physical_plan
+01)SortExec: expr=[id@0 ASC NULLS LAST], preserve_partitioning=[false]
+02)--CooperativeExec
+03)----DataSourceExec: format=iceberg, projection=[id, name]