Spark 4.1: Add MergingSortedRowDataReader for k-way merge of sorted files by anuragmantri · Pull Request #14948 · apache/iceberg

anuragmantri · 2025-12-31T00:20:06Z

This is first PR to report ordering to Spark.

This PR adds MergingSortedRowDataReader, a PartitionReader that merges rows from multiple sorted data files into a single sorted stream using a k-way heap merge (SortedMerge).

This reader is not wired up yet. A follow-up PR will integrate it with the SupportsReportOrdering DSv2 API to enable Spark's sort elimination optimization.

How it works:

Each file in a task group is opened as a separate RowDataReader
Sort key columns absent from Spark's projection are temporarily added to the read schema so the comparator can access them
SortOrderComparators.forSchema() with InternalRowWrapper handles all transform types (identity, bucket, truncate), ASC/DESC directions, and null ordering
Extra columns are stripped via ProjectingInternalRow before returning rows to Spark

Constraints:

Nested struct sort keys are not supported (to be handled in a follow-up)
Row-based only (not compatible with vectorized reads)

AI Usage: I used Claude Opus 4.6 for code generation and writing tests. I manually reviewed the generated code.

anuragmantri · 2026-01-21T22:10:47Z

Moved the changes to Spark 4.1 since it is now the latest version. Marked this PR as WIP as there is a prerequisite PR #14683 that is also in review.

peter-toth · 2026-02-19T16:14:59Z

My concern from Spark PoV is that unnecessary partition grouping can cause performance degradations. SPARK-55092 is a ticket about the problem and apache/spark#53859 / apache/spark#54330 PRs try to fix the problem.

If this PR disables bin packing then the above PRs won't be able to fix the issue.

Bin-packing of file scan tasks is disabled when ordering is required since Spark will discard ordering if multiple input partitions exist with the same grouping key.

So I would suggest keeping bin packing and reporting sort order for those packed partitions (i.e. the partitions might not be unique by key, but they are locally sorted), and when partition grouping is actually needed then Spark should merge the sorted partitions with the same key using k-way merge.

peter-toth · 2026-02-19T18:04:49Z

As we discussed offline, a long term (after apache/spark#54330) solution could be to improve the new GroupPartitionsExec operator to not only coalesce partitions with the same key, but k-way merge them to keep their sorted order.

github-actions · 2026-03-22T00:25:58Z

This pull request has been marked as stale due to 30 days of inactivity. It will be closed in 1 week if no further activity occurs. If you think that’s incorrect or this pull request requires a review, please simply write any comment. If closed, you can revive the PR at any time and @mention a reviewer or discuss it on the dev@iceberg.apache.org list. Thank you for your contributions.

anuragmantri · 2026-03-23T19:25:03Z

This PR is not stale. We are waiting waiting for the pre-requisite PR to be merged. I will update this PR after that one is merged.

anuragmantri · 2026-03-31T00:48:41Z

I rebased the PR after #15150. This is ready for review now.

@RussellSpitzer @aokolnychyi @szehon-ho - Could you take a look please?

anuragmantri · 2026-03-31T00:52:42Z

As we discussed offline, a long term (after apache/spark#54330) solution could be to improve the new GroupPartitionsExec operator to not only coalesce partitions with the same key, but k-way merge them to keep their sorted order.

Thanks @peter-toth , this makes sense. I think this PR is still needed and would still be valuable for tables with decently sized partitions. Since it's gated by a flag, I think it's safe to implement.

peter-toth · 2026-03-31T18:19:31Z

As we discussed offline, a long term (after apache/spark#54330) solution could be to improve the new GroupPartitionsExec operator to not only coalesce partitions with the same key, but k-way merge them to keep their sorted order.

Thanks @peter-toth , this makes sense. I think this PR is still needed and would still be valuable for tables with decently sized partitions. Since it's gated by a flag, I think it's safe to implement.

Absolutely.
FYI apache/spark#54330 has been merged. apache/spark#55116 will do the Spark side k-way merge to keep full ordering, but it requires this PR to report ordering and provide decently sized partitions.

Hugo-WB · 2026-04-27T22:04:46Z

Excited for this! Out of curiosity would this fix/help rewrite_data_files of sorted files/partitions?

anuragmantri · 2026-05-02T01:43:43Z

Out of curiosity would this fix/help rewrite_data_files of sorted files/partitions?

@Hugo-WB - This is only a read side optimization for partitioned tables using the SparkScan. Rewrite data files uses a different staged file scan which bypasses this for more control of the shuffling and sorting.

I do have another change that uses the same K-way merge logic to do compactions. I will submit a PR for that soon.

anuragmantri · 2026-05-12T23:03:51Z

@Hugo-WB - #16305 is a PR that leverages sort metadata for rewriting data.

Hugo-WB · 2026-05-13T08:36:44Z

amazing! Tysm. Looking forward to these changes! we are compacting a lot of sorted files and keen to eliminate shuffles where possible.

manuzhang · 2026-05-15T04:16:01Z

+        continue;
+      }
+
+      FileScanTask fileTask = (FileScanTask) task;


fileTask can be replaced with pattern variable

manuzhang · 2026-05-15T04:19:11Z

  }

+  /** Returns whether sort ordering was reported for this batch's scan. */
+  private boolean isOrderingEnabled() {


why do we need this private method?

Inlined it with orderingEnabled boolean.

manuzhang · 2026-05-15T04:25:20Z

+   * enabled at the table level (validated by {@link #isOrderingEnabled()}, multiple files in the
+   * group, and all tasks being {@link FileScanTask}s.
+   */
+  private boolean shouldUseMergingSortedReader(ScanTaskGroup<?> taskGroup) {


We can cache the result instead of caching this every time.

You are right. We are anyway disabling bin packing so a task group is 1:1 with an input partition. I refactored this check for all readers anyGroupNeedsMergingReader().

dramaticlly · 2026-07-22T23:42:05Z

+
+  MergingSortedRowDataReader(
+      Table table,
+      org.apache.iceberg.io.FileIO io,


nit: can we import FileIO? I don't see a existing class conflict

dramaticlly · 2026-07-22T23:43:13Z

+        }
+      }
+      Preconditions.checkState(
+          found, "Projection field id=%s not found in merge read schema — this is a bug", fieldId);


can probably drop — this is a bug, I don't think we usually include such in precondition check?

dramaticlly · 2026-07-22T23:44:39Z

+      boolean caseSensitive,
+      boolean cacheDeleteFilesOnExecutors) {
+    SortOrder sortOrder = table.sortOrder();
+


nit, remove the extra newline

dramaticlly · 2026-07-22T23:50:42Z

+ * A {@link PartitionReader} that reads multiple sorted files and merges them into a single sorted
+ * stream using a k-way heap merge ({@link SortedMerge}).
+ *
+ * <p>This reader is used when {@code preserve-data-ordering} is enabled and the task group contains


I think this refer to 2nd part of the PR in https://github.com/apache/iceberg/pull/16750/changes? Probably can follow up with the actual link to the table property. WDYT?

Yes, removed the reference.

dramaticlly · 2026-07-22T23:58:44Z

+    List<Types.NestedField> mergeColumns = mergeSchema.columns();
+    List<Object> positions = Lists.newArrayListWithCapacity(projection.columns().size());
+
+    for (int i = 0; i < projection.columns().size(); i++) {
+      int fieldId = projection.columns().get(i).fieldId();
+      boolean found = false;
+      for (int j = 0; j < mergeColumns.size(); j++) {
+        if (mergeColumns.get(j).fieldId() == fieldId) {
+          positions.add(j);
+          found = true;
+          break;
+        }
+      }
+      Preconditions.checkState(
+          found, "Projection field id=%s not found in merge read schema — this is a bug", fieldId);
+    }


instead of manually find projection with mergedSchema with O(n·m) nested scan, I think we can probably leverage existing utility?

List<Object> positions = Lists.newArrayListWithCapacity(projection.columns().size()); for (Types.NestedField column : projection.columns()) { Accessor<StructLike> accessor = mergeSchema.accessorForField(column.fieldId()); Preconditions.checkState( accessor != null, "Projection field id=%s not found in merge read schema", column.fieldId()); positions.add(Accessors.toPosition(accessor)); }

Good idea. Done.

dramaticlly · 2026-07-23T00:02:31Z

+              @Override
+              public void close() {
+                // No-op. The RowDataReaders are owned by the enclosing CloseableGroup
+                // (resources) and closed exactly once from close(). SortedMerge cannot be the
+                // sole owner because it filters out empty iterators (for example a file whose
+                // rows are all removed by deletes), so those readers would never be closed
+                // through the merge. Closing here as well would double-close the readers that
+                // SortedMerge does drain.
+              }


do we still need this override?

Just kept it for the comment. I don't have a preference though.

dramaticlly · 2026-07-23T00:03:23Z

+ * <p>Sort key columns absent from the requested projection are temporarily added to the read schema
+ * so that {@link SortOrderComparators} can access them during the merge. The extra columns are
+ * stripped from each row before it is returned to Spark.
+ */


I feel this is implementation details, do we really need those?

I clean it up a bit.

dramaticlly · 2026-07-23T00:05:43Z

+   * directions, and null ordering. The two {@link InternalRowWrapper} instances are allocated once
+   * and reused — {@code wrap()} just updates an internal reference.


I think last sentence is also a bit explaining the details, maybe just mention reuse of InternalRowWrapper for the purpose of comparator.

I changed to have less details.

anuragmantri

Thanks @dramaticlly. I addressed your comments and made some improvements in the latest revision.

anuragmantri · 2026-07-31T06:19:31Z

+              @Override
+              public void close() {
+                // No-op. The RowDataReaders are owned by the enclosing CloseableGroup
+                // (resources) and closed exactly once from close(). SortedMerge cannot be the
+                // sole owner because it filters out empty iterators (for example a file whose
+                // rows are all removed by deletes), so those readers would never be closed
+                // through the merge. Closing here as well would double-close the readers that
+                // SortedMerge does drain.
+              }


Just kept it for the comment. I don't have a preference though.

anuragmantri · 2026-07-31T06:20:11Z

+   * directions, and null ordering. The two {@link InternalRowWrapper} instances are allocated once
+   * and reused — {@code wrap()} just updates an internal reference.


I changed to have less details.

anuragmantri · 2026-07-31T06:20:21Z

+ * <p>Sort key columns absent from the requested projection are temporarily added to the read schema
+ * so that {@link SortOrderComparators} can access them during the merge. The extra columns are
+ * stripped from each row before it is returned to Spark.
+ */


I clean it up a bit.

anuragmantri · 2026-07-31T06:20:31Z

+    List<Types.NestedField> mergeColumns = mergeSchema.columns();
+    List<Object> positions = Lists.newArrayListWithCapacity(projection.columns().size());
+
+    for (int i = 0; i < projection.columns().size(); i++) {
+      int fieldId = projection.columns().get(i).fieldId();
+      boolean found = false;
+      for (int j = 0; j < mergeColumns.size(); j++) {
+        if (mergeColumns.get(j).fieldId() == fieldId) {
+          positions.add(j);
+          found = true;
+          break;
+        }
+      }
+      Preconditions.checkState(
+          found, "Projection field id=%s not found in merge read schema — this is a bug", fieldId);
+    }


Good idea. Done.

anuragmantri · 2026-07-31T06:20:41Z

+ * A {@link PartitionReader} that reads multiple sorted files and merges them into a single sorted
+ * stream using a k-way heap merge ({@link SortedMerge}).
+ *
+ * <p>This reader is used when {@code preserve-data-ordering} is enabled and the task group contains


Yes, removed the reference.

anuragmantri · 2026-07-31T06:20:46Z

+      boolean caseSensitive,
+      boolean cacheDeleteFilesOnExecutors) {
+    SortOrder sortOrder = table.sortOrder();
+


anuragmantri · 2026-07-31T06:20:50Z

+        }
+      }
+      Preconditions.checkState(
+          found, "Projection field id=%s not found in merge read schema — this is a bug", fieldId);


anuragmantri · 2026-07-31T06:20:58Z

+
+  MergingSortedRowDataReader(
+      Table table,
+      org.apache.iceberg.io.FileIO io,


dramaticlly · 2026-08-03T19:16:26Z

+          fieldId,
+          table.name());
+      Preconditions.checkArgument(
+          TypeUtil.ancestorFields(tableSchema, fieldId).isEmpty(),


ancestorFields is a full schema walk per sort field, I think if we can just check with tableSchema.asStruct().field(fieldId) != null is enough for verifying top level (non-nested) field here. As it depends on the sortKey

if this is already projected we can read fine even nested.

else I think we need fails as we don't use join with nested sortKey

dramaticlly · 2026-08-03T19:34:17Z

+    if (projection.columns().size() == mergeSchema.columns().size()) {
+      return null;
+    }


looks like column count check is a proxy to evaluate whether the projection is the same

Suggested change

if (projection.columns().size() == mergeSchema.columns().size()) {

return null;

}

if (projection.sameSchema(mergeSchema)) {

return null;

}

Yes, the same. Done.

dramaticlly · 2026-08-03T22:18:02Z

+    StructType sparkSchema = SparkSchemaUtil.convert(mergeReadSchema);
+    Comparator<StructLike> keyComparator =
+        SortOrderComparators.forSchema(mergeReadSchema, sortOrder);


From Claude, it seems for UUIDType we might have different ordering, the merge comparator use iceberg ordering

iceberg/api/src/main/java/org/apache/iceberg/types/Comparators.java

Line 49 in 58d5c37

.put(Types.UUIDType.get(), Comparator.naturalOrder())

while files written by spark might be sorted by spark's ordering.

iceberg/spark/v4.1/spark/src/main/java/org/apache/iceberg/spark/TypeToSparkType.java

Lines 149 to 151 in 58d5c37

case UUID:

// use String

return StringType$.MODULE$;

See example repro in https://gist.github.com/dramaticlly/c57b7382f4e57cc9bf90f4b9e8366e9a

Found a related issue in #14216, I think at minimal we shall exclude the sortKey with UUID type.

Good catch, I looked at this more. In the merging reader, we are comparing the output of the transforms. So only impacted case seems to be identity(UUID). I added a pre-condition to ensure the result of the transform cannot be UUID. But the actual check should be upstream when creating the MergingSortedRowReader. I will add that check in the wiring PR #16750

dramaticlly · 2026-08-03T22:23:41Z

+    Schema mergeReadSchema = mergeReadSchema(projection, sortOrder, table);
+    this.projectingRow = buildProjectingRow(projection, mergeReadSchema);
+
+    this.resources = new CloseableGroup();


do we want to enable this.closeableGroup.setSuppressCloseFailure(true)?

dramaticlly · 2026-08-03T22:24:57Z

+  }
+
+  @Test
+  void mergeWithFileFullyRemovedByDeletes() throws IOException {


can we add a 3 files here so that file1 is fully removed and file2 and file3 still enter the merge code path after apply the row to be positional deleted?

Added a test

…entions for SupportsReportOrdering

…g ordering decision to SparkRowReaderFactory

Revert integration code (SortOrderAnalyzer, SparkPartitioningAwareScan, SparkBatch wiring, config flags) to a follow-up PR. Keep only the k-way merge reader and its direct unit tests.

Fix NPE on files with a null sort order ID, avoid re-setting InputFileBlockHolder on every row, use TypeUtil.ancestorFields to distinguish nested sort keys from missing ones, and resolve projection positions via Accessors. Add tests for the preconditions, sort order evolution, and a file fully removed by deletes.

anuragmantri · 2026-08-06T22:03:05Z

This is ready for another round of review.

@szehon-ho @huaxingao @RussellSpitzer - It will be good to include this in the 1.12 release as this has been under review for a while now. Could you all also take a look, please?

anuragmantri · 2026-08-06T22:20:59Z

@manuzhang - Just FYI for your Spark 4.2 work. This PR is relevant for Spark 4.2 as well. However the plumbing PR #16750 will be slightly less restrictive as we can keep the bin-packing on Iceberg and spark can do another k-way merge.

dramaticlly · 2026-08-13T22:59:03Z

+
+  @Test
+  void mergeDescendingOrder() throws IOException {
+    catalog.dropTable(TableIdentifier.of("default", "test_merging_reader"));


do we need this? I thought we already dropped in afterEach ? same for the rest of 5 repetition below

dramaticlly · 2026-08-13T23:02:28Z

+    catalog.dropTable(TableIdentifier.of("default", "test_merging_reader"));
+
+    Schema nestedSchema =
+        new Schema(
+            required(1, "id", Types.IntegerType.get()),
+            required(
+                2, "location", Types.StructType.of(required(3, "city", Types.StringType.get()))));
+
+    table = catalog.createTable(TableIdentifier.of("default", "test_merging_reader"), nestedSchema);


nit: maybe update schema is easier than drop and recreate?

dramaticlly · 2026-08-13T23:47:12Z

+                  hasNext();
+                }
+                advanced = false;
+                return new TaggedRow(reader.get().copy(), block);


Actually I think there might be more problem as this is not a deep copy for Parquet collections of complex elements. We might need toUnsafe.apply(reader.get()).copy() via UnsafeProjection over the merge schema.

I think primitive type is generally ok, but for array or map of struct it might lead to row corruption. Might worth a try for this UT https://gist.github.com/dramaticlly/f91eb9d4186ae4b3101834bd4c5d20dd

github-actions Bot added spark core labels Dec 31, 2025

anuragmantri changed the title ~~[WIP] Spark 4.0: Implement SupportsReportOrdering DSv2 API~~ Spark 4.0: Implement SupportsReportOrdering DSv2 API Dec 31, 2025

anuragmantri mentioned this pull request Dec 31, 2025

Spark (4.0, 3.5): Set data file sort_order_id in manifest for writes from Spark #14683

Closed

anuragmantri force-pushed the supports-report-ordering branch from cc08ff2 to b4fde94 Compare January 21, 2026 22:08

anuragmantri changed the title ~~Spark 4.0: Implement SupportsReportOrdering DSv2 API~~ Spark 4.1: Implement SupportsReportOrdering DSv2 API Jan 21, 2026

anuragmantri changed the title ~~Spark 4.1: Implement SupportsReportOrdering DSv2 API~~ [WIP] Spark 4.1: Implement SupportsReportOrdering DSv2 API Jan 21, 2026

github-actions Bot added the stale label Mar 22, 2026

github-actions Bot removed the stale label Mar 24, 2026

peter-toth mentioned this pull request Mar 26, 2026

[SPARK-56241][SQL] Derive outputOrdering from KeyedPartitioning key expressions apache/spark#55036

Closed

anuragmantri force-pushed the supports-report-ordering branch from b4fde94 to dbf3ea9 Compare March 31, 2026 00:47

anuragmantri changed the title ~~[WIP] Spark 4.1: Implement SupportsReportOrdering DSv2 API~~ Spark 4.1: Implement SupportsReportOrdering DSv2 API Mar 31, 2026

anuragmantri force-pushed the supports-report-ordering branch from dbf3ea9 to a6f0205 Compare April 3, 2026 03:29

peter-toth mentioned this pull request Apr 3, 2026

[SPARK-55715][SQL] Keep outputOrdering when GroupPartitionsExec coalesces partitions apache/spark#55116

Closed

anuragmantri force-pushed the supports-report-ordering branch from a6f0205 to 51b4150 Compare May 2, 2026 01:33

manuzhang reviewed May 15, 2026

View reviewed changes

dramaticlly reviewed Jul 23, 2026

View reviewed changes

nssalian added this to the Iceberg 1.12.0 milestone Jul 29, 2026

nssalian requested a review from huaxingao July 29, 2026 21:50

anuragmantri force-pushed the supports-report-ordering branch from 7110906 to cb987a4 Compare July 31, 2026 06:18

anuragmantri commented Jul 31, 2026

View reviewed changes

dramaticlly reviewed Aug 3, 2026

View reviewed changes

anuragmantri and others added 18 commits August 6, 2026 14:57

Spark 4.1: Implement SupportsReportOrdering DSv2 API

6cd9a35

Avoid redundant checking of sort order

59283d0

Fix reader lifecycle, use UncheckedIOException, and improve test conv…

d9b05f4

…entions for SupportsReportOrdering

Address review comments

c212283

Address review comments

b74a6fa

Address review comments

ddc51ef

Fix checkstyle

a26d604

Use ProjectingInternalRow to avoid copying values

fc04804

Spotless apply

862bcea

Decouple MergingSortedRowDataReader from SparkInputPartition by movin…

dfa5f47

…g ordering decision to SparkRowReaderFactory

Extract MergingSortedRowDataReader into standalone PR with focused tests

7844f24

Revert integration code (SortOrderAnalyzer, SparkPartitioningAwareScan, SparkBatch wiring, config flags) to a follow-up PR. Keep only the k-way merge reader and its direct unit tests.

Add fileReaders to closable group

344aad4

Spotless apply

570e077

Address review comments

65de2a6

Do not project duplicate ids

5ee4029

Use TaggedRow and address review comments

900c130

Address review comments

d21c514

anuragmantri force-pushed the supports-report-ordering branch from cb987a4 to d21c514 Compare August 6, 2026 21:58

This was referenced Aug 10, 2026

Sort-aware Iceberg reads in Comet via a per-partition streaming merge apache/datafusion-comet#5323

Open

feat: Sort-aware Iceberg reads in Comet via a per-partition streaming merge apache/datafusion-comet#5331

Open

dramaticlly reviewed Aug 13, 2026

View reviewed changes

		* directions, and null ordering. The two {@link InternalRowWrapper} instances are allocated once
		* and reused — {@code wrap()} just updates an internal reference.

Uh oh!

Conversation

anuragmantri commented Dec 31, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

anuragmantri commented Jan 21, 2026

Uh oh!

peter-toth commented Feb 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

peter-toth commented Feb 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions Bot commented Mar 22, 2026

Uh oh!

anuragmantri commented Mar 23, 2026

Uh oh!

anuragmantri commented Mar 31, 2026

Uh oh!

anuragmantri commented Mar 31, 2026

Uh oh!

peter-toth commented Mar 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Hugo-WB commented Apr 27, 2026

Uh oh!

anuragmantri commented May 2, 2026

Uh oh!

anuragmantri commented May 12, 2026

Uh oh!

Hugo-WB commented May 13, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

manuzhang May 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

anuragmantri left a comment

Choose a reason for hiding this comment

anuragmantri commented Dec 31, 2025 •

edited

Loading

peter-toth commented Feb 19, 2026 •

edited

Loading

peter-toth commented Feb 19, 2026 •

edited

Loading

peter-toth commented Mar 31, 2026 •

edited

Loading

manuzhang May 15, 2026 •

edited

Loading