cluster iceberg rewrite_manifests output by order-preserving partition ranges

Signed-off-by: dontknow9179 <clin56322@gmail.com>

Author: dontknow9179

fe/fe-core/src/main/java/com/starrocks/connector/iceberg/procedure/RewriteManifestsProcedure.java

@@ -14,29 +14,45 @@
 
 package com.starrocks.connector.iceberg.procedure;
 
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.Iterables;
 import com.starrocks.connector.exception.StarRocksConnectorException;
 import com.starrocks.connector.iceberg.IcebergTableOperation;
 import com.starrocks.qe.ShowResultSet;
 import com.starrocks.sql.optimizer.operator.scalar.ConstantOperator;
+import org.apache.iceberg.DataFile;
 import org.apache.iceberg.ManifestFile;
+import org.apache.iceberg.ManifestFiles;
+import org.apache.iceberg.PartitionSpec;
 import org.apache.iceberg.RewriteManifests;
 import org.apache.iceberg.Snapshot;
+import org.apache.iceberg.StructLike;
 import org.apache.iceberg.Table;
+import org.apache.iceberg.io.CloseableIterable;
+import org.apache.iceberg.io.FileIO;
+import org.apache.iceberg.types.Comparators;
+import org.apache.iceberg.types.Type;
 import org.apache.iceberg.types.Types;
+import org.apache.iceberg.util.PartitionUtil;
 import org.apache.iceberg.util.PropertyUtil;
-import org.apache.iceberg.util.StructLikeWrapper;
 
+import java.io.IOException;
 import java.util.Collections;
+import java.util.Comparator;
 import java.util.List;
 import java.util.Map;
-import java.util.Objects;
+import java.util.NavigableMap;
+import java.util.NavigableSet;
+import java.util.TreeMap;
+import java.util.TreeSet;
 
 import static org.apache.iceberg.TableProperties.MANIFEST_TARGET_SIZE_BYTES;
 import static org.apache.iceberg.TableProperties.MANIFEST_TARGET_SIZE_BYTES_DEFAULT;
 
 public class RewriteManifestsProcedure extends IcebergTableProcedure {
     private static final String PROCEDURE_NAME = "rewrite_manifests";
     private static final int MAX_MANIFEST_CLUSTERS = 100;
+    private static final int NULL_PARTITION_CLUSTER = -1;
 
     private static final RewriteManifestsProcedure INSTANCE = new RewriteManifestsProcedure();
 
@@ -65,8 +81,8 @@ public ShowResultSet execute(IcebergTableProcedureContext context, Map<String, C
             return null;
         }
 
-        List<ManifestFile> manifests = currentSnapshot.allManifests(icebergTable.io());
-        if (manifests.isEmpty()) {
+        List<ManifestFile> dataManifests = currentSnapshot.dataManifests(icebergTable.io());
+        if (dataManifests.isEmpty()) {
             return null;
         }
 
@@ -76,28 +92,122 @@ public ShowResultSet execute(IcebergTableProcedureContext context, Map<String, C
             manifestTargetSizeBytes = MANIFEST_TARGET_SIZE_BYTES_DEFAULT;
         }
 
-        if (manifests.size() == 1 && manifests.get(0).length() < manifestTargetSizeBytes) {
+        if (dataManifests.size() == 1 && dataManifests.get(0).length() < manifestTargetSizeBytes) {
             return null;
         }
 
-        long totalManifestsSize = manifests.stream().mapToLong(ManifestFile::length).sum();
+        long totalManifestsSize = dataManifests.stream().mapToLong(ManifestFile::length).sum();
         // Having too many open manifest writers can potentially cause OOM on the coordinator
-        // Floor to at least 1 to avoid modulo-by-zero when totalManifestsSize is 0 (e.g., manifests report length 0)
-        long targetManifestClusters = Math.max(1, Math.min(
+        long targetManifestClusters = Math.min(
                 ((totalManifestsSize + manifestTargetSizeBytes - 1) / manifestTargetSizeBytes),
-                MAX_MANIFEST_CLUSTERS));
+                MAX_MANIFEST_CLUSTERS);
+
+        Types.StructType partitionType = icebergTable.spec().partitionType();
+        Map<Integer, PartitionSpec> specsById = icebergTable.specs();
+
+        // Assign each distinct top-level partition value an order-preserving bucket so that
+        // sorted-adjacent partitions land in the same manifest. This keeps each manifest's
+        // partition min/max range narrow (effective manifest pruning at read time) while bounding
+        // the number of concurrently open writers.
+        NavigableMap<Object, Integer> partitionValueToCluster;
+        if (icebergTable.spec().isPartitioned() && targetManifestClusters > 1) {
+            partitionValueToCluster = buildClusteredPartitionValues(
+                    icebergTable, currentSnapshot, (int) targetManifestClusters);
+        } else {
+            partitionValueToCluster = Collections.emptyNavigableMap();
+        }
 
         RewriteManifests rewriteManifests = context.transaction().rewriteManifests();
-        Types.StructType structType = icebergTable.spec().partitionType();
 
         rewriteManifests
                 .clusterBy(file -> {
-                    // Cluster by partitions for better locality when reading data files
-                    StructLikeWrapper partitionWrapper = StructLikeWrapper.forType(structType).set(file.partition());
-                    // Limit the number of clustering buckets to avoid creating too many small manifest files
-                    return Integer.toUnsignedLong(Objects.hash(partitionWrapper)) % targetManifestClusters;
+                    if (partitionValueToCluster.isEmpty()) {
+                        return 0;
+                    }
+                    PartitionSpec fileSpec = specsById.get(file.specId());
+                    if (fileSpec == null) {
+                        // Unknown spec: the partition spec evolved under a concurrent commit. Evolution
+                        // is rare, so abort and let the next maintenance run re-cluster against the new
+                        // spec rather than guessing an order for it.
+                        throw new StarRocksConnectorException(
+                                "rewrite_manifests aborted: the table's partition spec changed under a " +
+                                        "concurrent commit; it will be retried on the next maintenance run");
+                    }
+                    // Coerce into the current spec's partition type so files written under an older
+                    // partition spec still cluster correctly.
+                    StructLike partition =
+                            PartitionUtil.coercePartition(partitionType, fileSpec, file.partition());
+                    Object value = partition.get(0, Object.class);
+                    if (value == null) {
+                        return NULL_PARTITION_CLUSTER;
+                    }
+                    Integer cluster = partitionValueToCluster.get(value);
+                    if (cluster != null) {
+                        return cluster;
+                    }
+                    // Value not seen up front (typically a new partition added by a concurrent write to
+                    // the same spec). Slot it next to its sort-order neighbor.
+                    Map.Entry<Object, Integer> neighbor = partitionValueToCluster.floorEntry(value);
+                    if (neighbor == null) {
+                        neighbor = partitionValueToCluster.firstEntry();
+                    }
+                    return neighbor.getValue();
                 })
                 .commit();
         return null;
     }
+
+    // Collect the distinct top-level partition values across all data manifests of the snapshot,
+    // sort them, and assign each a bucket in [0, targetClusters) grouping sorted-adjacent values
+    // into the same bucket. Clustering is limited to the top-level partition field because it is
+    // usually the most selective for read-time filters and keeps the distinct-value set small.
+    private static NavigableMap<Object, Integer> buildClusteredPartitionValues(
+            Table table, Snapshot snapshot, int targetClusters) {
+        PartitionSpec spec = table.spec();
+        Type.PrimitiveType firstFieldType =
+                spec.partitionType().fields().get(0).type().asPrimitiveType();
+        Comparator<Object> comparator = partitionValueComparator(firstFieldType);
+        Types.StructType partitionType = spec.partitionType();
+        Map<Integer, PartitionSpec> specsById = table.specs();
+        FileIO io = table.io();
+        List<ManifestFile> dataManifests = snapshot.dataManifests(io);
+        Iterable<CloseableIterable<DataFile>> perManifest = Iterables.transform(
+                dataManifests,
+                manifest -> (CloseableIterable<DataFile>)
+                        ManifestFiles.read(manifest, io, specsById).select(ImmutableList.of("partition")));
+
+        // Comparator-based dedup mirrors Iceberg's own equality for partition value types.
+        NavigableSet<Object> uniqueValues = new TreeSet<>(comparator);
+        try (CloseableIterable<DataFile> dataFiles = CloseableIterable.concat(perManifest)) {
+            for (DataFile file : dataFiles) {
+                StructLike partition = PartitionUtil.coercePartition(
+                        partitionType, specsById.get(file.specId()), file.partition());
+                Object value = partition.get(0, Object.class);
+                if (value != null) {
+                    uniqueValues.add(value);
+                }
+            }
+        } catch (IOException e) {
+            throw new StarRocksConnectorException(
+                    "Unable to read manifests while clustering partitions for snapshot " + snapshot.snapshotId(), e);
+        }
+
+        if (uniqueValues.isEmpty()) {
+            return Collections.emptyNavigableMap();
+        }
+
+        NavigableMap<Object, Integer> valueToCluster = new TreeMap<>(comparator);
+        int index = 0;
+        int size = uniqueValues.size();
+        for (Object value : uniqueValues) {
+            valueToCluster.put(value, (int) ((long) index * targetClusters / size));
+            index++;
+        }
+        return valueToCluster;
+    }
+
+    @SuppressWarnings("unchecked")
+    private static Comparator<Object> partitionValueComparator(Type.PrimitiveType type) {
+        return (Comparator<Object>) (Comparator<?>) Comparators.forType(type);
+    }
 }

fe/fe-core/src/test/java/com/starrocks/connector/iceberg/procedure/RewriteManifestsProcedureTest.java

@@ -17,25 +17,39 @@
 import com.starrocks.connector.HdfsEnvironment;
 import com.starrocks.connector.exception.StarRocksConnectorException;
 import com.starrocks.connector.iceberg.IcebergTableOperation;
+import com.starrocks.connector.iceberg.TestTables;
 import com.starrocks.connector.iceberg.hive.IcebergHiveCatalog;
 import com.starrocks.qe.ConnectContext;
 import com.starrocks.sql.ast.AlterTableOperationClause;
 import com.starrocks.sql.ast.AlterTableStmt;
 import com.starrocks.sql.optimizer.operator.scalar.ConstantOperator;
+import org.apache.iceberg.AppendFiles;
+import org.apache.iceberg.DataFile;
+import org.apache.iceberg.DataFiles;
 import org.apache.iceberg.ManifestFile;
 import org.apache.iceberg.PartitionSpec;
 import org.apache.iceberg.RewriteManifests;
+import org.apache.iceberg.Schema;
 import org.apache.iceberg.Snapshot;
 import org.apache.iceberg.Table;
+import org.apache.iceberg.TableProperties;
 import org.apache.iceberg.Transaction;
+import org.apache.iceberg.types.Types;
+import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.TempDir;
+import org.mockito.ArgumentCaptor;
 import org.mockito.Mockito;
 
+import java.io.File;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.function.Function;
 
+import static org.apache.iceberg.types.Types.NestedField.required;
 import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
@@ -51,6 +65,19 @@ public class RewriteManifestsProcedureTest {
 
     public static final HdfsEnvironment HDFS_ENVIRONMENT = new HdfsEnvironment();
 
+    private static final Schema SCHEMA =
+            new Schema(required(1, "k1", Types.IntegerType.get()), required(2, "k2", Types.IntegerType.get()));
+    // Identity partition on an int column, so partition values have a natural order to cluster by.
+    private static final PartitionSpec SPEC = PartitionSpec.builderFor(SCHEMA).identity("k2").build();
+
+    @TempDir
+    public File warehouse;
+
+    @AfterEach
+    public void clearNativeTables() {
+        TestTables.clearTables();
+    }
+
     @Test
     void testProcedureCreation() {
         RewriteManifestsProcedure procedure = RewriteManifestsProcedure.getInstance();
@@ -93,7 +120,7 @@ void testExecuteWithNoSnapshot() {
     void testExecuteWithEmptyManifests() {
         RewriteManifestsProcedure procedure = RewriteManifestsProcedure.getInstance();
         Snapshot snapshot = Mockito.mock(Snapshot.class);
-        when(snapshot.allManifests(any())).thenReturn(Collections.emptyList());
+        when(snapshot.dataManifests(any())).thenReturn(Collections.emptyList());
 
         Table table = Mockito.mock(Table.class);
         when(table.currentSnapshot()).thenReturn(snapshot);
@@ -112,7 +139,7 @@ void testExecuteWithSingleSmallManifest() {
         when(smallManifest.length()).thenReturn(100L);
 
         Snapshot snapshot = Mockito.mock(Snapshot.class);
-        when(snapshot.allManifests(any())).thenReturn(List.of(smallManifest));
+        when(snapshot.dataManifests(any())).thenReturn(List.of(smallManifest));
 
         Table table = Mockito.mock(Table.class);
         when(table.currentSnapshot()).thenReturn(snapshot);
@@ -136,7 +163,7 @@ void testExecuteRewriteSuccess() {
         when(m2.length()).thenReturn(10 * 1024 * 1024L);
 
         Snapshot snapshot = Mockito.mock(Snapshot.class);
-        when(snapshot.allManifests(any())).thenReturn(List.of(m1, m2));
+        when(snapshot.dataManifests(any())).thenReturn(List.of(m1, m2));
 
         Table table = Mockito.mock(Table.class);
         when(table.currentSnapshot()).thenReturn(snapshot);
@@ -167,7 +194,7 @@ void testExecuteWithSingleLargeManifestTriggersRewrite() {
         when(largeManifest.length()).thenReturn(10 * 1024 * 1024L);
 
         Snapshot snapshot = Mockito.mock(Snapshot.class);
-        when(snapshot.allManifests(any())).thenReturn(List.of(largeManifest));
+        when(snapshot.dataManifests(any())).thenReturn(List.of(largeManifest));
 
         Table table = Mockito.mock(Table.class);
         when(table.currentSnapshot()).thenReturn(snapshot);
@@ -190,6 +217,96 @@ void testExecuteWithSingleLargeManifestTriggersRewrite() {
         verify(rewriteManifests).commit();
     }
 
+    @Test
+    void testPartitionedManifestsClusteredByOrderPreservingRange() throws Exception {
+        // 200 distinct partitions, cluster count capped at MAX_MANIFEST_CLUSTERS(=100),
+        // so the order-preserving assignment is bucket = rank * 100 / 200 = rank / 2.
+        int numPartitions = 200;
+        int expectedClusters = 100;
+
+        Function<DataFile, Object> clusterFn = runAndCaptureClusterFunction(numPartitions);
+
+        int[] buckets = new int[numPartitions];
+        for (int i = 0; i < numPartitions; i++) {
+            Object bucket = clusterFn.apply(dataFileForPartition(i));
+            assertNotNull(bucket);
+            buckets[i] = ((Number) bucket).intValue();
+        }
+
+        for (int i = 0; i < numPartitions; i++) {
+            // partition value equals its rank here, so the expected bucket is a pure function of i
+            assertEquals(i / 2, buckets[i], "partition k2=" + i + " landed in an unexpected bucket");
+            assertTrue(buckets[i] >= 0 && buckets[i] < expectedClusters, "bucket out of range: " + buckets[i]);
+            if (i > 0) {
+                // monotonic non-decreasing in partition-value order: the property a hash would break
+                assertTrue(buckets[i] >= buckets[i - 1], "bucketing is not order-preserving");
+            }
+        }
+        // grouping actually happened and stayed within the cap
+        assertEquals(expectedClusters, (int) Arrays.stream(buckets).distinct().count());
+    }
+
+    @Test
+    void testUnmappedPartitionValueSlotsIntoNeighborBucket() throws Exception {
+        int numPartitions = 10;
+        Function<DataFile, Object> clusterFn = runAndCaptureClusterFunction(numPartitions);
+        // A value not seen at planning time (e.g. a new partition from a concurrent write to the same
+        // spec) must not abort and must not be hash-scattered: an above-max value slots into the last
+        // known partition's bucket, preserving the order-preserving clustering.
+        Object maxBucket = clusterFn.apply(dataFileForPartition(numPartitions - 1));
+        Object unmappedBucket = clusterFn.apply(dataFileForPartition(9999));
+        assertNotNull(unmappedBucket);
+        assertEquals(maxBucket, unmappedBucket);
+    }
+
+    @Test
+    void testUnknownSpecAbortsRewrite() throws Exception {
+        Function<DataFile, Object> clusterFn = runAndCaptureClusterFunction(10);
+        // A file under a spec not seen at planning time (concurrent partition-spec evolution) aborts,
+        // deferring to the next maintenance run rather than guessing an order.
+        DataFile unknownSpecFile = Mockito.mock(DataFile.class);
+        when(unknownSpecFile.specId()).thenReturn(999);
+        assertThrows(StarRocksConnectorException.class, () -> clusterFn.apply(unknownSpecFile));
+    }
+
+    // Builds a real partitioned table with `numPartitions` single-file partitions, runs the
+    // procedure against it, and returns the clustering Function it hands to RewriteManifests.
+    // The manifest target size is forced to 1 byte so the cluster count is pinned to the cap.
+    private Function<DataFile, Object> runAndCaptureClusterFunction(int numPartitions) throws Exception {
+        TestTables.TestTable table = TestTables.create(warehouse, "t_range_" + numPartitions, SCHEMA, SPEC, 2);
+        AppendFiles append = table.newAppend();
+        for (int i = 0; i < numPartitions; i++) {
+            append.appendFile(dataFileForPartition(i));
+        }
+        append.commit();
+        table.updateProperties().set(TableProperties.MANIFEST_TARGET_SIZE_BYTES, "1").commit();
+
+        RewriteManifests rewriteManifests = Mockito.mock(RewriteManifests.class);
+        when(rewriteManifests.clusterBy(any())).thenReturn(rewriteManifests);
+        doNothing().when(rewriteManifests).commit();
+        Transaction txn = Mockito.mock(Transaction.class);
+        when(txn.rewriteManifests()).thenReturn(rewriteManifests);
+
+        IcebergTableProcedureContext context = createContext(table, txn);
+        RewriteManifestsProcedure procedure = RewriteManifestsProcedure.getInstance();
+        assertDoesNotThrow(() -> procedure.execute(context, Collections.emptyMap()));
+
+        @SuppressWarnings("unchecked")
+        ArgumentCaptor<Function<DataFile, Object>> captor = ArgumentCaptor.forClass(Function.class);
+        verify(rewriteManifests).clusterBy(captor.capture());
+        verify(rewriteManifests).commit();
+        return captor.getValue();
+    }
+
+    private static DataFile dataFileForPartition(int k2) {
+        return DataFiles.builder(SPEC)
+                .withPath("/path/to/data-k2-" + k2 + ".parquet")
+                .withFileSizeInBytes(10)
+                .withPartitionPath("k2=" + k2)
+                .withRecordCount(1)
+                .build();
+    }
+
     private IcebergTableProcedureContext createContext(Table table, Transaction transaction) {
         IcebergHiveCatalog catalog = Mockito.mock(IcebergHiveCatalog.class);
         ConnectContext ctx = Mockito.mock(ConnectContext.class);