Review with Devin contd.

Author: malhotrashivam

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

@@ -4,32 +4,26 @@
 package io.deephaven.iceberg.location;
 
 import io.deephaven.api.SortColumn;
-import io.deephaven.base.verify.Require;
 import io.deephaven.engine.table.impl.locations.TableKey;
 import io.deephaven.engine.table.impl.locations.TableLocation;
 import io.deephaven.parquet.table.ParquetInstructions;
 import io.deephaven.parquet.table.location.ParquetTableLocation;
 import org.jetbrains.annotations.NotNull;
-import org.jetbrains.annotations.Nullable;
 
 import java.util.List;
 
-public class IcebergTableParquetLocation extends ParquetTableLocation implements TableLocation {
+class IcebergTableParquetLocation extends ParquetTableLocation implements TableLocation {
 
-    @Nullable
-    private final List<SortColumn> sortedColumns;
-
-    public IcebergTableParquetLocation(
+    IcebergTableParquetLocation(
             @NotNull final TableKey tableKey,
             @NotNull final IcebergTableParquetLocationKey tableLocationKey,
             @NotNull final ParquetInstructions readInstructions) {
         super(tableKey, tableLocationKey, readInstructions);
-        sortedColumns = Require.neqNull(tableLocationKey.sortedColumns(), "tableLocationKey.sortedColumns()");
     }
 
     @Override
     @NotNull
     public List<SortColumn> getSortedColumns() {
-        return sortedColumns;
+        return ((IcebergTableParquetLocationKey) getKey()).sortedColumns();
     }
 }

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

@@ -8,6 +8,7 @@
 import io.deephaven.engine.table.impl.locations.TableLocationKey;
 import io.deephaven.parquet.table.ParquetInstructions;
 import io.deephaven.parquet.table.location.ParquetTableLocationKey;
+import io.deephaven.util.annotations.InternalUseOnly;
 import org.apache.iceberg.DataFile;
 import org.apache.iceberg.ManifestFile;
 import io.deephaven.util.channel.SeekableChannelsProvider;
@@ -25,6 +26,7 @@
 /**
  * {@link TableLocationKey} implementation for use with data stored in Iceberg tables in the parquet format.
  */
+@InternalUseOnly
 public class IcebergTableParquetLocationKey extends ParquetTableLocationKey implements IcebergTableLocationKey {
 
     private static final String IMPLEMENTATION_NAME = IcebergTableParquetLocationKey.class.getSimpleName();
@@ -119,7 +121,7 @@ public IcebergTableParquetLocationKey(
         manifestSequenceNumber = manifestFile.sequenceNumber();
 
         this.readInstructions = readInstructions;
-        this.sortedColumns = sortedColumns;
+        this.sortedColumns = Require.neqNull(sortedColumns, "sortedColumns");
     }
 
     @Override

extensions/iceberg/src/test/java/io/deephaven/iceberg/junit5/SqliteCatalogBase.java

@@ -1051,28 +1051,23 @@ void testAutoRefreshingPartitionedAppend() throws InterruptedException {
     /**
      * Verify that the sort order for the data files in the table match the expected sort order.
      */
-    private void verifySortOrder(
+    private static void verifySortOrder(
             final IcebergTableAdapter tableAdapter,
             final List<List<SortColumn>> expectedSortOrders) {
         verifySortOrder(tableAdapter, expectedSortOrders,
                 ParquetInstructions.EMPTY.withTableDefinition(tableAdapter.definition()));
     }
 
-    private void verifySortOrder(
+    private static void verifySortOrder(
             @NotNull final IcebergTableAdapter tableAdapter,
             @NotNull final List<List<SortColumn>> expectedSortOrders,
             @NotNull final ParquetInstructions readInstructions) {
         final org.apache.iceberg.Table icebergTable = tableAdapter.icebergTable();
-        final Map<ManifestFile, List<DataFile>> manifestToDataFiles =
-                IcebergUtils.manifestToDataFiles(icebergTable, icebergTable.currentSnapshot());
         final List<List<SortColumn>> actualSortOrders = new ArrayList<>();
-        for (final Map.Entry<ManifestFile, List<DataFile>> entry : manifestToDataFiles.entrySet()) {
-            final List<DataFile> dataFiles = entry.getValue();
-            for (final DataFile dataFile : dataFiles) {
-                actualSortOrders.add(computeSortedColumns(icebergTable, dataFile, readInstructions));
-            }
-        }
-        assertThat(actualSortOrders).containsExactlyInAnyOrderElementsOf(expectedSortOrders);
+        IcebergUtils.allDataFiles(icebergTable, icebergTable.currentSnapshot())
+                .forEach(dataFile -> actualSortOrders
+                        .add(computeSortedColumns(icebergTable, dataFile, readInstructions)));
+        assertThat(actualSortOrders).isEqualTo(expectedSortOrders);
     }
 
     @Test
@@ -1110,8 +1105,8 @@ void testApplyDefaultSortOrder() {
 
         // Verify that the new data file is sorted
         verifySortOrder(tableAdapter, List.of(
-                List.of(),
-                List.of(SortColumn.asc(ColumnName.of("intCol")))));
+                List.of(SortColumn.asc(ColumnName.of("intCol"))),
+                List.of()));
 
         // Append more unsorted data to the table without enforcing sort order
         tableWriterWithoutSorting.append(IcebergWriteInstructions.builder()