feat: DH-18143: Improve handling of sort order for Iceberg tables

extensions/iceberg/src/main/java/io/deephaven/iceberg/base/IcebergUtils.java

@@ -75,8 +75,19 @@ public final class IcebergUtils {
      */
     public static Stream<DataFile> allDataFiles(@NotNull final Table table, @NotNull final Snapshot snapshot) {
         return allManifestFiles(table, snapshot)
-                .map(manifestFile -> ManifestFiles.read(manifestFile, table.io()))
-                .flatMap(IcebergUtils::toStream);
+                .flatMap(manifestFile -> allDataFiles(table, manifestFile));
+    }
+
+    /**
+     * Get a stream of all {@link DataFile} objects from the given {@link Table} and {@link ManifestFile}.
+     *
+     * @param table The {@link Table} to retrieve data files for.
+     * @param manifestFile The {@link ManifestFile} to retrieve data files from.
+     *
+     * @return A stream of {@link DataFile} objects.
+     */
+    public static Stream<DataFile> allDataFiles(@NotNull final Table table, @NotNull ManifestFile manifestFile) {
+        return toStream(ManifestFiles.read(manifestFile, table.io()));
     }
 
     /**
@@ -117,6 +128,14 @@ static List<ManifestFile> allManifests(@NotNull final Table table, @NotNull fina
         }
     }
 
+    public static Map<ManifestFile, List<DataFile>> manifestToDataFiles(
+            @NotNull final Table table, @NotNull final Snapshot snapshot) {
+        return allManifestFiles(table, snapshot)
+                .collect(Collectors.toMap(
+                        manifestFile -> manifestFile,
+                        manifestFile -> allDataFiles(table, manifestFile).collect(Collectors.toList())));
+    }
+
     /**
      * Convert a {@link org.apache.iceberg.io.CloseableIterable} to a {@link Stream} that will close the iterable when
      * the stream is closed.

extensions/iceberg/src/main/java/io/deephaven/iceberg/location/IcebergTableLocationFactory.java

@@ -7,25 +7,30 @@
 import io.deephaven.engine.table.impl.locations.TableLocation;
 import io.deephaven.engine.table.impl.locations.impl.TableLocationFactory;
 import io.deephaven.engine.table.impl.locations.util.TableDataRefreshService;
+import io.deephaven.iceberg.util.IcebergTableAdapter;
 import io.deephaven.parquet.table.ParquetInstructions;
-import io.deephaven.parquet.table.location.ParquetTableLocation;
-import io.deephaven.parquet.table.location.ParquetTableLocationKey;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
 /**
  * {@link TableLocationFactory} for Iceberg {@link TableLocation}s.
  */
 public final class IcebergTableLocationFactory implements TableLocationFactory<TableKey, IcebergTableLocationKey> {
-    public IcebergTableLocationFactory() {}
+
+    final IcebergTableAdapter tableAdapter;
+
+    public IcebergTableLocationFactory(@NotNull final IcebergTableAdapter tableAdapter) {
+        this.tableAdapter = tableAdapter;
+    }
 
     @Override
     @NotNull
-    public TableLocation makeLocation(@NotNull final TableKey tableKey,
+    public TableLocation makeLocation(
+            @NotNull final TableKey tableKey,
             @NotNull final IcebergTableLocationKey locationKey,
             @Nullable final TableDataRefreshService refreshService) {
         if (locationKey instanceof IcebergTableParquetLocationKey) {
-            return new ParquetTableLocation(tableKey, (ParquetTableLocationKey) locationKey,
+            return new IcebergTableParquetLocation(tableAdapter, tableKey, (IcebergTableParquetLocationKey) locationKey,
                     (ParquetInstructions) locationKey.readInstructions());
         }
         throw new UnsupportedOperationException("Unsupported location key type: " + locationKey.getClass());

extensions/iceberg/src/main/java/io/deephaven/iceberg/location/IcebergTableParquetLocation.java

@@ -0,0 +1,83 @@
+//
+// Copyright (c) 2016-2025 Deephaven Data Labs and Patent Pending
+//
+package io.deephaven.iceberg.location;
+
+import io.deephaven.api.ColumnName;
+import io.deephaven.api.SortColumn;
+import io.deephaven.engine.table.impl.locations.TableKey;
+import io.deephaven.engine.table.impl.locations.TableLocation;
+import io.deephaven.iceberg.util.IcebergTableAdapter;
+import io.deephaven.parquet.table.ParquetInstructions;
+import io.deephaven.parquet.table.location.ParquetTableLocation;
+import org.apache.iceberg.DataFile;
+import org.apache.iceberg.NullOrder;
+import org.apache.iceberg.Schema;
+import org.apache.iceberg.SortDirection;
+import org.apache.iceberg.SortField;
+import org.apache.iceberg.SortOrder;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+
+public class IcebergTableParquetLocation extends ParquetTableLocation implements TableLocation {
+
+    @Nullable
+    private final List<SortColumn> sortedColumns;
+
+    public IcebergTableParquetLocation(
+            @NotNull final IcebergTableAdapter tableAdapter,
+            @NotNull final TableKey tableKey,
+            @NotNull final IcebergTableParquetLocationKey tableLocationKey,
+            @NotNull final ParquetInstructions readInstructions) {
+        super(tableKey, tableLocationKey, readInstructions);
+        sortedColumns = computeSortedColumns(tableAdapter, tableLocationKey.dataFile());
+    }
+
+    @Override
+    @NotNull
+    public List<SortColumn> getSortedColumns() {
+        return sortedColumns == null ? super.getSortedColumns() : sortedColumns;
+    }
+
+    @Nullable
+    private static List<SortColumn> computeSortedColumns(
+            @NotNull final IcebergTableAdapter tableAdapter,
+            @NotNull final DataFile dataFile) {
+        final Integer sortOrderId = dataFile.sortOrderId();
+        // If sort order is missing or unknown, we cannot determine the sorted columns from the metadata and will
+        // check the underlying parquet file for the sorted columns, when the user asks for them.
+        if (sortOrderId == null) {
+            return null;
+        }
+        final SortOrder sortOrder = tableAdapter.icebergTable().sortOrders().get(sortOrderId);
+        if (sortOrder == null) {
+            return null;
+        }
+        if (sortOrder.isUnsorted()) {
+            return Collections.emptyList();
+        }
+        final Schema schema = sortOrder.schema();
+        final List<SortColumn> sortColumns = new ArrayList<>(sortOrder.fields().size());
+        for (final SortField field : sortOrder.fields()) {
+            if (!field.transform().isIdentity()) {
+                // TODO (DH-18160): Improve support for handling non-identity transforms
+                break;
+            }
+            final ColumnName columnName = ColumnName.of(schema.findColumnName(field.sourceId()));
+            final SortColumn sortColumn;
+            if (field.nullOrder() == NullOrder.NULLS_FIRST && field.direction() == SortDirection.ASC) {
+                sortColumn = SortColumn.asc(columnName);
+            } else if (field.nullOrder() == NullOrder.NULLS_LAST && field.direction() == SortDirection.DESC) {
+                sortColumn = SortColumn.desc(columnName);
+            } else {
+                break;
+            }
+            sortColumns.add(sortColumn);
+        }
+        return Collections.unmodifiableList(sortColumns);
+    }
+}

extensions/iceberg/src/main/java/io/deephaven/iceberg/location/IcebergTableParquetLocationKey.java

@@ -41,6 +41,9 @@ public class IcebergTableParquetLocationKey extends ParquetTableLocationKey impl
     @Nullable
     private final TableIdentifier tableIdentifier;
 
+    @NotNull
+    private final DataFile dataFile;
+
     /**
      * The {@link DataFile#dataSequenceNumber()} of the data file backing this keyed location.
      */
@@ -102,6 +105,8 @@ public IcebergTableParquetLocationKey(
         // tableUUID was null
         this.tableIdentifier = tableUuid != null ? null : tableIdentifier;
 
+        this.dataFile = dataFile;
+
         // Files with unknown sequence numbers should be ordered first
         dataSequenceNumber = dataFile.dataSequenceNumber() != null ? dataFile.dataSequenceNumber() : Long.MIN_VALUE;
         fileSequenceNumber = dataFile.fileSequenceNumber() != null ? dataFile.fileSequenceNumber() : Long.MIN_VALUE;
@@ -124,6 +129,11 @@ public ParquetInstructions readInstructions() {
         return readInstructions;
     }
 
+    @NotNull
+    DataFile dataFile() {
+        return dataFile;
+    }
+
     /**
      * When comparing with another {@link IcebergTableParquetLocationKey}, precedence-wise this implementation compares:
      * <ul>

extensions/iceberg/src/main/java/io/deephaven/iceberg/util/IcebergTableAdapter.java

@@ -401,7 +401,7 @@ public IcebergTable table(@NotNull final IcebergReadInstructions readInstruction
                     new IcebergStaticTableLocationProvider<>(
                             StandaloneTableKey.getInstance(),
                             keyFinder,
-                            new IcebergTableLocationFactory(),
+                            new IcebergTableLocationFactory(this),
                             tableIdentifier);
 
             return new IcebergTableImpl(
@@ -419,14 +419,14 @@ public IcebergTable table(@NotNull final IcebergReadInstructions readInstruction
             locationProvider = new IcebergManualRefreshTableLocationProvider<>(
                     StandaloneTableKey.getInstance(),
                     keyFinder,
-                    new IcebergTableLocationFactory(),
+                    new IcebergTableLocationFactory(this),
                     this,
                     tableIdentifier);
         } else {
             locationProvider = new IcebergAutoRefreshTableLocationProvider<>(
                     StandaloneTableKey.getInstance(),
                     keyFinder,
-                    new IcebergTableLocationFactory(),
+                    new IcebergTableLocationFactory(this),
                     TableDataRefreshService.getSharedRefreshService(),
                     updatedInstructions.updateMode().autoRefreshMs(),
                     this,