feat: Sort-aware Iceberg reads in Comet via a per-partition streaming merge

.github/workflows/pr_build_linux.yml

@@ -315,6 +315,7 @@ jobs:
               org.apache.comet.CometIcebergEncryptionSuite
               org.apache.comet.CometIcebergRewriteActionSuite
               org.apache.comet.CometIcebergWriteActionSuite
+              org.apache.comet.CometIcebergSortMergeReadSuite
               org.apache.comet.iceberg.IcebergReflectionSuite
               org.apache.comet.csv.CometCsvNativeReadSuite
               org.apache.comet.CometFuzzTestSuite

.github/workflows/pr_build_macos.yml

@@ -131,6 +131,7 @@ jobs:
               org.apache.comet.CometIcebergEncryptionSuite
               org.apache.comet.CometIcebergRewriteActionSuite
               org.apache.comet.CometIcebergWriteActionSuite
+              org.apache.comet.CometIcebergSortMergeReadSuite
               org.apache.comet.iceberg.IcebergReflectionSuite
               org.apache.comet.csv.CometCsvNativeReadSuite
               org.apache.comet.CometFuzzTestSuite

native/core/src/execution/operators/iceberg_scan.rs

@@ -28,7 +28,7 @@ use arrow::datatypes::SchemaRef;
 use datafusion::common::{DataFusionError, Result as DFResult};
 use datafusion::execution::{RecordBatchStream, SendableRecordBatchStream, TaskContext};
 use datafusion::physical_expr::expressions::Column;
-use datafusion::physical_expr::{EquivalenceProperties, PhysicalExpr};
+use datafusion::physical_expr::{EquivalenceProperties, LexOrdering, PhysicalExpr};
 use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
 use datafusion::physical_plan::metrics::{
     BaselineMetrics, Count, ExecutionPlanMetricsSet, MetricBuilder, MetricsSet,
@@ -86,6 +86,24 @@ pub struct IcebergScanExec {
     tasks: Vec<FileScanTask>,
     /// Number of data files to read concurrently
     data_file_concurrency_limit: usize,
+    /// FileIO (and, for S3, the JVM credential bridge behind it) built once at plan time and shared
+    /// across partitions. FileIO is cheap to clone (Arc-backed), so each `execute` clones this
+    /// rather than rebuilding the storage factory + credential bridge. This matters in the ordered
+    /// path, where the scan is one partition per file and `execute` is called once per file.
+    file_io: FileIO,
+    /// Table sort order Iceberg reported, translated against `output_schema`. `Some` makes this a
+    /// multi-partition scan: one sorted stream per task, which a SortPreservingMergeExec above
+    /// merges back into one sorted partition. It is also advertised in `plan_properties`. `None`
+    /// keeps the old single-partition unordered read (all tasks streamed together).
+    ///
+    /// Concurrency note: in the ordered path each partition reads exactly one task, so
+    /// `data_file_concurrency_limit` no longer bounds cross-file concurrency; instead the wrapping
+    /// SortPreservingMergeExec drives one reader per file to merge them. That fan-out (files per
+    /// Spark partition) is intrinsic to a k-way merge of per-file sorted streams -- the files must
+    /// be read as separate streams to stay individually sorted -- and is the natural granularity
+    /// for a sorted Iceberg table. `data_file_concurrency_limit` still bounds delete-file stats and
+    /// the unordered path.
+    ordering: Option<LexOrdering>,
     /// Metrics
     metrics: ExecutionPlanMetricsSet,
 }
@@ -98,12 +116,31 @@ impl IcebergScanExec {
         catalog_name: String,
         tasks: Vec<FileScanTask>,
         data_file_concurrency_limit: usize,
+        ordering: Option<LexOrdering>,
     ) -> Result<Self, ExecutionError> {
         let output_schema = schema;
-        let plan_properties = Self::compute_properties(Arc::clone(&output_schema), 1);
+        // With an ordering, read each task as its own sorted stream on its own partition, so the
+        // SortPreservingMergeExec above can merge them. Without one, keep the single partition that
+        // reads every task. Comet only drives execute(0), so a multi-partition leaf with no merge
+        // above it would read only the first task.
+        let num_partitions = if ordering.is_some() {
+            tasks.len().max(1)
+        } else {
+            1
+        };
+        let plan_properties = Self::compute_properties(
+            Arc::clone(&output_schema),
+            num_partitions,
+            ordering.as_ref(),
+        );
 
         let metrics = ExecutionPlanMetricsSet::new();
 
+        // Build FileIO (and the S3 credential bridge) once here rather than per `execute`. In the
+        // ordered path `execute` is called once per file, so rebuilding it there would repeat the
+        // JNI/reflection credential-bridge construction for every file in the partition.
+        let file_io = Self::load_file_io(&catalog_properties, &metadata_location, &catalog_name)?;
+
         Ok(Self {
             metadata_location,
             output_schema,
@@ -112,13 +149,26 @@ impl IcebergScanExec {
             catalog_name,
             tasks,
             data_file_concurrency_limit,
+            file_io,
+            ordering,
             metrics,
         })
     }
 
-    fn compute_properties(schema: SchemaRef, num_partitions: usize) -> Arc<PlanProperties> {
+    fn compute_properties(
+        schema: SchemaRef,
+        num_partitions: usize,
+        ordering: Option<&LexOrdering>,
+    ) -> Arc<PlanProperties> {
+        let eq_properties = match ordering {
+            Some(lex) => EquivalenceProperties::new_with_orderings(
+                Arc::clone(&schema),
+                std::iter::once(lex.iter().cloned()),
+            ),
+            None => EquivalenceProperties::new(Arc::clone(&schema)),
+        };
         Arc::new(PlanProperties::new(
-            EquivalenceProperties::new(schema),
+            eq_properties,
             Partitioning::UnknownPartitioning(num_partitions),
             EmissionType::Incremental,
             Boundedness::Bounded,
@@ -152,10 +202,18 @@ impl ExecutionPlan for IcebergScanExec {
 
     fn execute(
         &self,
-        _partition: usize,
+        partition: usize,
         context: Arc<TaskContext>,
     ) -> DFResult<SendableRecordBatchStream> {
-        self.execute_with_tasks(self.tasks.clone(), context)
+        // With a reported ordering this is a multi-partition operator: partition `i` reads only
+        // task `i` as its own sorted stream, and the SortPreservingMergeExec above merges them.
+        // Without one, the single partition reads every task together (legacy unordered path).
+        let tasks = if self.ordering.is_some() {
+            self.tasks.get(partition).cloned().into_iter().collect()
+        } else {
+            self.tasks.clone()
+        };
+        self.execute_with_tasks(tasks, context)
     }
 
     fn metrics(&self) -> Option<MetricsSet> {
@@ -172,11 +230,7 @@ impl IcebergScanExec {
         context: Arc<TaskContext>,
     ) -> DFResult<SendableRecordBatchStream> {
         let output_schema = Arc::clone(&self.output_schema);
-        let file_io = Self::load_file_io(
-            &self.catalog_properties,
-            &self.metadata_location,
-            &self.catalog_name,
-        )?;
+        let file_io = self.file_io.clone();
         let batch_size = context.session_config().batch_size();
 
         let metrics = IcebergScanMetrics::new(&self.metrics);
@@ -640,6 +694,7 @@ mod tests {
     use iceberg_storage_opendal::OpenDalStorageFactory;
 
     use super::IcebergScanExec;
+    use datafusion::physical_plan::ExecutionPlan;
 
     fn fs_file_io() -> FileIO {
         FileIOBuilder::new(Arc::new(OpenDalStorageFactory::Fs)).build()
@@ -741,6 +796,59 @@ mod tests {
             .unwrap();
     }
 
+    fn int_schema() -> arrow::datatypes::SchemaRef {
+        use arrow::datatypes::{DataType, Field, Schema};
+        Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]))
+    }
+
+    fn single_col_ordering() -> Option<datafusion::physical_expr::LexOrdering> {
+        use arrow::compute::SortOptions;
+        use datafusion::physical_expr::expressions::Column;
+        use datafusion::physical_expr::{LexOrdering, PhysicalSortExpr};
+        LexOrdering::new(vec![PhysicalSortExpr {
+            expr: Arc::new(Column::new("a", 0)),
+            options: SortOptions::default(),
+        }])
+    }
+
+    // Builds a scan over three empty-delete tasks with the given reported ordering.
+    fn exec_with_ordering(
+        ordering: Option<datafusion::physical_expr::LexOrdering>,
+    ) -> IcebergScanExec {
+        use std::collections::HashMap;
+        let tasks = vec![
+            task_with_deletes(vec![]),
+            task_with_deletes(vec![]),
+            task_with_deletes(vec![]),
+        ];
+        IcebergScanExec::new(
+            "metadata.json".to_string(),
+            int_schema(),
+            HashMap::new(),
+            "cat".to_string(),
+            tasks,
+            1,
+            ordering,
+        )
+        .unwrap()
+    }
+
+    // A reported ordering turns the scan into a multi-partition operator (one partition per task)
+    // so a SortPreservingMergeExec above can k-way merge the per-file sorted streams.
+    #[test]
+    fn reported_ordering_makes_scan_multi_partition() {
+        let exec = exec_with_ordering(single_col_ordering());
+        assert_eq!(exec.properties().partitioning.partition_count(), 3);
+    }
+
+    // Without a reported ordering the scan stays single-partition (Comet drives only execute(0),
+    // which must read every task), preserving the legacy unordered behaviour.
+    #[test]
+    fn no_ordering_keeps_single_partition() {
+        let exec = exec_with_ordering(None);
+        assert_eq!(exec.properties().partitioning.partition_count(), 1);
+    }
+
     fn from_hex(s: &str) -> Vec<u8> {
         assert!(s.len().is_multiple_of(2), "odd-length hex string");
         (0..s.len())

native/core/src/execution/planner.rs

@@ -70,6 +70,7 @@ use datafusion::{
         limit::LocalLimitExec,
         projection::ProjectionExec,
         sorts::sort::SortExec,
+        sorts::sort_preserving_merge::SortPreservingMergeExec,
         ExecutionPlan,
     },
     prelude::SessionContext,
@@ -1757,24 +1758,51 @@ impl PhysicalPlanner {
                 let tasks = parse_file_scan_tasks_from_common(common, &scan.file_scan_tasks)?;
                 let data_file_concurrency_limit = common.data_file_concurrency_limit as usize;
 
-                let iceberg_scan = IcebergScanExec::new(
+                // Table sort order Iceberg reported. Empty unless sortMerge is on and the order
+                // passed the identity gate in CometIcebergNativeScan. The SortOrder children are
+                // bound references into required_schema, so build the LexOrdering against it.
+                let ordering: Option<LexOrdering> = if common.table_sort_orders.is_empty() {
+                    None
+                } else {
+                    let exprs = common
+                        .table_sort_orders
+                        .iter()
+                        .map(|expr| self.create_sort_expr(expr, Arc::clone(&required_schema)))
+                        .collect::<Result<Vec<PhysicalSortExpr>, ExecutionError>>()?;
+                    LexOrdering::new(exprs)
+                };
+
+                let iceberg_scan: Arc<dyn ExecutionPlan> = Arc::new(IcebergScanExec::new(
                     metadata_location,
                     required_schema,
                     catalog_properties,
                     catalog_name,
                     tasks,
                     data_file_concurrency_limit,
-                )?;
+                    ordering.clone(),
+                )?);
 
-                Ok((
-                    vec![],
-                    vec![],
-                    Arc::new(SparkPlan::new(
-                        spark_plan.plan_id,
-                        Arc::new(iceberg_scan),
-                        vec![],
-                    )),
-                ))
+                // With an ordering, the scan is multi-partition (one sorted stream per file). Wrap
+                // it in a SortPreservingMergeExec so each Spark partition comes out sorted. Keep the
+                // scan as an additional native plan so its metrics (num_splits, bytes_scanned) still
+                // roll up. Without an ordering, the scan stays a single-partition leaf, unchanged.
+                let result_plan = match ordering {
+                    Some(lex) => {
+                        let spm: Arc<dyn ExecutionPlan> = Arc::new(
+                            SortPreservingMergeExec::new(lex, Arc::clone(&iceberg_scan))
+                                .with_round_robin_repartition(false),
+                        );
+                        SparkPlan::new_with_additional(
+                            spark_plan.plan_id,
+                            spm,
+                            vec![],
+                            vec![iceberg_scan],
+                        )
+                    }
+                    None => SparkPlan::new(spark_plan.plan_id, iceberg_scan, vec![]),
+                };
+
+                Ok((vec![], vec![], Arc::new(result_plan)))
             }
             OpStruct::ShuffleWriter(writer) => {
                 assert_eq!(children.len(), 1);

native/proto/src/proto/operator.proto

@@ -304,6 +304,12 @@ message IcebergScanCommon {
   // projections or filters, so PlanDataInjector keys planning data by
   // (metadata_location, scan_hash_code) rather than metadata_location alone.
   int32 scan_hash_code = 15;
+
+  // Table sort order Iceberg reported (SupportsReportOrdering), against required_schema. Each Expr
+  // wraps a SortOrder (child, direction, null_ordering). When set, the native scan reads each file
+  // as its own sorted stream and wraps the scan in a SortPreservingMergeExec, so each Spark
+  // partition comes out sorted. Empty means no reported order: read unordered, as before.
+  repeated spark.spark_expression.Expr table_sort_orders = 16;
 }
 
 message IcebergScan {

spark/src/main/scala/org/apache/comet/CometConf.scala

@@ -121,6 +121,30 @@ object CometConf extends ShimCometConf {
       .booleanConf
       .createWithDefault(true)
 
+  val COMET_ICEBERG_SORT_MERGE_ENABLED: ConfigEntry[Boolean] =
+    conf("spark.comet.scan.icebergNative.sortMerge.enabled")
+      .category(CATEGORY_SCAN)
+      .doc("Whether the native Iceberg scan reports the table sort order and performs a " +
+        "per-partition streaming merge of already-sorted files. When enabled and Iceberg reports " +
+        "an ordering (requires Iceberg's spark.sql.iceberg.planning.preserve-data-ordering), " +
+        "each Spark partition reads its files as separate sorted streams merged into one sorted " +
+        "output, and the ordering is surfaced to Spark so redundant sorts are eliminated. When " +
+        "disabled, files are read unordered as before.")
+      .booleanConf
+      .createWithDefault(true)
+
+  val COMET_ICEBERG_REPORT_PARTITIONING_ENABLED: ConfigEntry[Boolean] =
+    conf("spark.comet.scan.icebergNative.reportPartitioning.enabled")
+      .category(CATEGORY_SCAN)
+      .doc("Whether the native Iceberg scan reports Iceberg's key-grouped partitioning (requires " +
+        "spark.sql.sources.v2.bucketing.enabled) so storage-partitioned joins avoid a shuffle, " +
+        "which is what lets a reported sort order eliminate a join's sort. Independent of " +
+        "sortMerge.enabled. Defaults to disabled: the AQE + push-down-partition-values path is " +
+        "not yet verified, so reporting is off until it is exercised. When disabled, " +
+        "partitioning is reported as unknown, as before.")
+      .booleanConf
+      .createWithDefault(false)
+
   val COMET_ICEBERG_WRITE_SPLIT_OPERATOR_ENABLED: ConfigEntry[Boolean] =
     conf("spark.comet.write.iceberg.splitOperator.enabled")
       .category(CATEGORY_TESTING)