[server][client] Support per-partition bucket count for partitioned tables

fluss-client/src/main/java/org/apache/fluss/client/lookup/AbstractLookupQuery.java

@@ -39,22 +39,33 @@ public abstract class AbstractLookupQuery<T> {
      */
     private final @Nullable String originalPartitionName;
 
+    private final int bucketCount;
     private int retries;
     private long nextRetryTimeMs;
 
     public AbstractLookupQuery(TablePath tablePath, TableBucket tableBucket, byte[] key) {
-        this(tablePath, tableBucket, key, null);
+        this(tablePath, tableBucket, key, null, 0);
     }
 
     public AbstractLookupQuery(
             TablePath tablePath,
             TableBucket tableBucket,
             byte[] key,
             @Nullable String originalPartitionName) {
+        this(tablePath, tableBucket, key, originalPartitionName, 0);
+    }
+
+    public AbstractLookupQuery(
+            TablePath tablePath,
+            TableBucket tableBucket,
+            byte[] key,
+            @Nullable String originalPartitionName,
+            int bucketCount) {
         this.tablePath = tablePath;
         this.tableBucket = tableBucket;
         this.key = key;
         this.originalPartitionName = originalPartitionName;
+        this.bucketCount = bucketCount;
         this.retries = 0;
         this.nextRetryTimeMs = 0;
     }
@@ -75,6 +86,11 @@ public TableBucket tableBucket() {
         return originalPartitionName;
     }
 
+    /** The bucket count used to calculate this lookup's bucketId, or 0 if unknown (legacy). */
+    public int bucketCount() {
+        return bucketCount;
+    }
+
     public int retries() {
         return retries;
     }

fluss-client/src/main/java/org/apache/fluss/client/lookup/AbstractLookuper.java

@@ -23,6 +23,7 @@
 import org.apache.fluss.metadata.SchemaGetter;
 import org.apache.fluss.metadata.SchemaInfo;
 import org.apache.fluss.metadata.TableInfo;
+import org.apache.fluss.metadata.TablePartition;
 import org.apache.fluss.row.InternalRow;
 import org.apache.fluss.row.decode.FixedSchemaDecoder;
 import org.apache.fluss.utils.CopyOnWriteMap;
@@ -75,6 +76,19 @@ abstract class AbstractLookuper implements Lookuper {
                         tableInfo.getTableConfig().getKvFormat(), tableInfo.getSchema()));
     }
 
+    /**
+     * Resolves the effective bucket count for the target partition: the per-partition bucket count
+     * ({@code bucket.num.actual}) when the cluster metadata has it, falling back to the table-level
+     * bucket count otherwise (non-partitioned tables or partitions created by older versions).
+     */
+    protected int resolvePartitionBucketCount(
+            TablePartition tablePartition, int tableLevelNumBuckets) {
+        return metadataUpdater
+                .getCluster()
+                .getBucketCount(tablePartition)
+                .orElse(tableLevelNumBuckets);
+    }
+
     protected void handleLookupResponse(
             List<byte[]> result, CompletableFuture<LookupResult> lookupFuture) {
         List<MemorySegment> valueList = new ArrayList<>(result.size());

fluss-client/src/main/java/org/apache/fluss/client/lookup/LookupBatch.java

@@ -34,9 +34,12 @@ public class LookupBatch {
 
     private final List<LookupQuery> lookups;
 
-    LookupBatch(LookupBatchKey lookupBatchKey) {
+    private final int bucketCount;
+
+    LookupBatch(LookupBatchKey lookupBatchKey, int bucketCount) {
         this.lookupBatchKey = lookupBatchKey;
         this.lookups = new ArrayList<>();
+        this.bucketCount = bucketCount;
     }
 
     public void addLookup(LookupQuery lookup) {
@@ -55,6 +58,11 @@ public TableBucket tableBucket() {
         return lookupBatchKey.originalPartitionName();
     }
 
+    /** The bucket count the bucketId was calculated with, or 0 if unknown (legacy). */
+    public int getBucketCount() {
+        return bucketCount;
+    }
+
     LookupBatchKey lookupBatchKey() {
         return lookupBatchKey;
     }

fluss-client/src/main/java/org/apache/fluss/client/lookup/LookupClient.java

@@ -113,17 +113,24 @@ public CompletableFuture<byte[]> lookup(
             TableBucket tableBucket,
             byte[] keyBytes,
             boolean insertIfNotExists,
-            @Nullable String originalPartitionName) {
+            @Nullable String originalPartitionName,
+            int bucketCount) {
         LookupQuery lookup =
                 new LookupQuery(
-                        tablePath, tableBucket, keyBytes, insertIfNotExists, originalPartitionName);
+                        tablePath,
+                        tableBucket,
+                        keyBytes,
+                        insertIfNotExists,
+                        originalPartitionName,
+                        bucketCount);
         lookupQueue.appendLookup(lookup);
         return lookup.future();
     }
 
     public CompletableFuture<List<byte[]>> prefixLookup(
-            TablePath tablePath, TableBucket tableBucket, byte[] keyBytes) {
-        PrefixLookupQuery prefixLookup = new PrefixLookupQuery(tablePath, tableBucket, keyBytes);
+            TablePath tablePath, TableBucket tableBucket, byte[] keyBytes, int bucketCount) {
+        PrefixLookupQuery prefixLookup =
+                new PrefixLookupQuery(tablePath, tableBucket, keyBytes, bucketCount);
         lookupQueue.appendLookup(prefixLookup);
         return prefixLookup.future();
     }

fluss-client/src/main/java/org/apache/fluss/client/lookup/LookupQuery.java

@@ -41,15 +41,25 @@ public class LookupQuery extends AbstractLookupQuery<byte[]> {
             TableBucket tableBucket,
             byte[] key,
             boolean insertIfNotExists,
-            @Nullable String originalPartitionName) {
-        super(tablePath, tableBucket, key, originalPartitionName);
+            @Nullable String originalPartitionName,
+            int bucketCount) {
+        super(tablePath, tableBucket, key, originalPartitionName, bucketCount);
         this.future = new CompletableFuture<>();
         this.insertIfNotExists = insertIfNotExists;
     }
 
+    LookupQuery(
+            TablePath tablePath,
+            TableBucket tableBucket,
+            byte[] key,
+            boolean insertIfNotExists,
+            @Nullable String originalPartitionName) {
+        this(tablePath, tableBucket, key, insertIfNotExists, originalPartitionName, 0);
+    }
+
     @VisibleForTesting
     LookupQuery(TablePath tablePath, TableBucket tableBucket, byte[] key) {
-        this(tablePath, tableBucket, key, false, null);
+        this(tablePath, tableBucket, key, false, null, 0);
     }
 
     @Override

fluss-client/src/main/java/org/apache/fluss/client/lookup/LookupSender.java

@@ -39,6 +39,7 @@
 import org.apache.fluss.rpc.messages.PrefixLookupRequest;
 import org.apache.fluss.rpc.messages.PrefixLookupResponse;
 import org.apache.fluss.rpc.protocol.ApiError;
+import org.apache.fluss.rpc.protocol.Errors;
 import org.apache.fluss.utils.ExponentialBackoff;
 import org.apache.fluss.utils.types.Tuple2;
 
@@ -221,7 +222,7 @@ private void sendLookupRequest(
             LookupBatchKey batchKey = new LookupBatchKey(tb, lookup.originalPartitionName());
             lookupByTableId
                     .computeIfAbsent(tableId, k -> new LinkedHashMap<>())
-                    .computeIfAbsent(batchKey, k -> new LookupBatch(batchKey))
+                    .computeIfAbsent(batchKey, k -> new LookupBatch(batchKey, lookup.bucketCount()))
                     .addLookup(lookup);
         }
 
@@ -299,7 +300,7 @@ private void sendPrefixLookupRequest(
             long tableId = tb.getTableId();
             lookupByTableId
                     .computeIfAbsent(tableId, k -> new HashMap<>())
-                    .computeIfAbsent(tb, k -> new PrefixLookupBatch(tb))
+                    .computeIfAbsent(tb, k -> new PrefixLookupBatch(tb, prefixLookup.bucketCount()))
                     .addLookup(prefixLookup);
         }
 
@@ -565,6 +566,13 @@ private void handleLookupError(
             invalidTableOrPartitions(tableOrPartitions);
         }
 
+        if (error.error() == Errors.STALE_METADATA) {
+            for (AbstractLookupQuery<?> lookup : lookups) {
+                lookup.future().completeExceptionally(exception);
+            }
+            return;
+        }
+
         for (AbstractLookupQuery<?> lookup : lookups) {
             if (canRetry(lookup, exception)) {
                 long retryDelayMs = prepareRetry(lookup, exception);

fluss-client/src/main/java/org/apache/fluss/client/lookup/PrefixKeyLookuper.java

@@ -25,6 +25,7 @@
 import org.apache.fluss.metadata.SchemaGetter;
 import org.apache.fluss.metadata.TableBucket;
 import org.apache.fluss.metadata.TableInfo;
+import org.apache.fluss.metadata.TablePartition;
 import org.apache.fluss.row.InternalRow;
 import org.apache.fluss.row.encode.KeyEncoder;
 import org.apache.fluss.types.RowType;
@@ -165,9 +166,9 @@ public CompletableFuture<LookupResult> lookup(InternalRow prefixKey) {
                 prefixKeyEncoder == bucketKeyEncoder
                         ? prefixKeyBytes
                         : bucketKeyEncoder.encodeKey(prefixKey);
-        int bucketId = bucketingFunction.bucketing(bucketKeyBytes, numBuckets);
 
         Long partitionId = null;
+        int effectiveNumBuckets = numBuckets;
         if (partitionGetter != null) {
             try {
                 partitionId =
@@ -176,15 +177,25 @@ public CompletableFuture<LookupResult> lookup(InternalRow prefixKey) {
                                 partitionGetter,
                                 tableInfo.getTablePath(),
                                 metadataUpdater);
+                if (partitionId != null) {
+                    effectiveNumBuckets =
+                            resolvePartitionBucketCount(
+                                    new TablePartition(tableInfo.getTableId(), partitionId),
+                                    numBuckets);
+                }
             } catch (PartitionNotExistException e) {
                 return CompletableFuture.completedFuture(new LookupResult(Collections.emptyList()));
             }
         }
 
+        // Compute bucket ID after partition resolution — needs per-partition bucket count
+        int bucketId = bucketingFunction.bucketing(bucketKeyBytes, effectiveNumBuckets);
+
         CompletableFuture<LookupResult> lookupFuture = new CompletableFuture<>();
         TableBucket tableBucket = new TableBucket(tableInfo.getTableId(), partitionId, bucketId);
         lookupClient
-                .prefixLookup(tableInfo.getTablePath(), tableBucket, prefixKeyBytes)
+                .prefixLookup(
+                        tableInfo.getTablePath(), tableBucket, prefixKeyBytes, effectiveNumBuckets)
                 .whenComplete(
                         (result, error) -> {
                             if (error != null) {

fluss-client/src/main/java/org/apache/fluss/client/lookup/PrefixLookupBatch.java

@@ -34,10 +34,13 @@ public class PrefixLookupBatch {
     /** The table bucket that the lookup operations should fall into. */
     private final TableBucket tableBucket;
 
+    private final int bucketCount;
+
     private final List<PrefixLookupQuery> prefixLookups;
 
-    public PrefixLookupBatch(TableBucket tableBucket) {
+    public PrefixLookupBatch(TableBucket tableBucket, int bucketCount) {
         this.tableBucket = tableBucket;
+        this.bucketCount = bucketCount;
         this.prefixLookups = new ArrayList<>();
     }
 
@@ -53,6 +56,11 @@ public TableBucket tableBucket() {
         return tableBucket;
     }
 
+    /** The bucket count the bucketId was calculated with, or 0 if unknown (legacy). */
+    public int getBucketCount() {
+        return bucketCount;
+    }
+
     public void complete(List<List<byte[]>> values) {
         if (values.size() != prefixLookups.size()) {
             completeExceptionally(

fluss-client/src/main/java/org/apache/fluss/client/lookup/PrefixLookupQuery.java

@@ -32,11 +32,16 @@
 public class PrefixLookupQuery extends AbstractLookupQuery<List<byte[]>> {
     private final CompletableFuture<List<byte[]>> future;
 
-    PrefixLookupQuery(TablePath tablePath, TableBucket tableBucket, byte[] prefixKey) {
-        super(tablePath, tableBucket, prefixKey);
+    PrefixLookupQuery(
+            TablePath tablePath, TableBucket tableBucket, byte[] prefixKey, int bucketCount) {
+        super(tablePath, tableBucket, prefixKey, null, bucketCount);
         this.future = new CompletableFuture<>();
     }
 
+    PrefixLookupQuery(TablePath tablePath, TableBucket tableBucket, byte[] prefixKey) {
+        this(tablePath, tableBucket, prefixKey, 0);
+    }
+
     @Override
     public CompletableFuture<List<byte[]>> future() {
         return future;

fluss-client/src/main/java/org/apache/fluss/client/lookup/PrimaryKeyLookuper.java

@@ -26,6 +26,7 @@
 import org.apache.fluss.metadata.SchemaGetter;
 import org.apache.fluss.metadata.TableBucket;
 import org.apache.fluss.metadata.TableInfo;
+import org.apache.fluss.metadata.TablePartition;
 import org.apache.fluss.row.InternalRow;
 import org.apache.fluss.row.encode.KeyEncoder;
 import org.apache.fluss.types.RowType;
@@ -127,6 +128,7 @@ public CompletableFuture<LookupResult> lookup(InternalRow lookupKey) {
         int bucketId = bucketingFunction.bucketing(bkBytes, numBuckets);
         Long partitionId = null;
         String originalPartitionName = null;
+        int effectiveNumBuckets = numBuckets;
         if (partitionGetter != null) {
             originalPartitionName = partitionGetter.getPartition(lookupKey);
             if (confirmedHistoricalPartitions.contains(originalPartitionName)) {
@@ -139,13 +141,31 @@ public CompletableFuture<LookupResult> lookup(InternalRow lookupKey) {
                                 partitionGetter,
                                 tableInfo.getTablePath(),
                                 metadataUpdater);
+                if (partitionId != null) {
+                    effectiveNumBuckets =
+                            resolvePartitionBucketCount(
+                                    new TablePartition(tableInfo.getTableId(), partitionId),
+                                    numBuckets);
+                }
             } catch (PartitionNotExistException e) {
                 return mayFallbackToHistoricalLookup(bucketId, pkBytes, originalPartitionName);
             }
         }
 
+        // A partition created before ALTER bucket.num keeps its own layout, so re-route by the
+        // partition's actual count. The historical lookups above keep the table-level bucketId
+        // because the historical partition is resolved on its own path.
+        if (effectiveNumBuckets != numBuckets) {
+            bucketId = bucketingFunction.bucketing(bkBytes, effectiveNumBuckets);
+        }
         TableBucket tableBucket = new TableBucket(tableInfo.getTableId(), partitionId, bucketId);
-        return lookupBucket(tableBucket, pkBytes, insertIfNotExists, false, originalPartitionName);
+        return lookupBucket(
+                tableBucket,
+                pkBytes,
+                insertIfNotExists,
+                false,
+                originalPartitionName,
+                effectiveNumBuckets);
     }
 
     /**
@@ -181,6 +201,22 @@ private CompletableFuture<LookupResult> historicalLookup(
                     new UnsupportedOperationException(
                             "Lookup with insertIfNotExists is not supported for historical partition lookup."));
         }
+        // TODO: Support historical lookup on rescaled tables. The lake data of the original
+        //  partition is laid out with the bucket count that partition was tiered with, but that
+        //  partition is already dropped from Fluss, so neither the table-level count nor the
+        //  historical partition's own count is guaranteed to match it. Only the lake metadata
+        //  still knows it (Paimon exposes it as DataSplit#totalBuckets), so the server has to
+        //  resolve it and route to the bucket itself. Until then, reject the lookup instead of
+        //  silently reading the wrong bucket.
+        if (tableInfo.getBucketLayoutEpoch() > 0) {
+            return completedExceptionally(
+                    new UnsupportedOperationException(
+                            String.format(
+                                    "Historical partition lookup is not supported on table %s because its "
+                                            + "'bucket.num' has been altered, so the bucket layout of already "
+                                            + "tiered partitions can no longer be determined by the client.",
+                                    tableInfo.getTablePath())));
+        }
         PhysicalTablePath historicalPartitionPath =
                 PhysicalTablePath.of(tableInfo.getTablePath(), HISTORICAL_PARTITION_VALUE);
         try {
@@ -192,7 +228,8 @@ private CompletableFuture<LookupResult> historicalLookup(
                     metadataUpdater.getPartitionIdOrElseThrow(historicalPartitionPath);
             TableBucket tableBucket =
                     new TableBucket(tableInfo.getTableId(), historicalPartitionId, bucketId);
-            return lookupBucket(tableBucket, keyBytes, false, true, originalPartitionName);
+            return lookupBucket(
+                    tableBucket, keyBytes, false, true, originalPartitionName, numBuckets);
         } catch (Throwable t) {
             return completedExceptionally(t);
         }
@@ -203,15 +240,17 @@ private CompletableFuture<LookupResult> lookupBucket(
             byte[] keyBytes,
             boolean insertIfNotExists,
             boolean historicalLookup,
-            @Nullable String originalPartitionName) {
+            @Nullable String originalPartitionName,
+            int bucketCount) {
         CompletableFuture<LookupResult> lookupFuture = new CompletableFuture<>();
         lookupClient
                 .lookup(
                         tableInfo.getTablePath(),
                         tableBucket,
                         keyBytes,
                         insertIfNotExists,
-                        historicalLookup ? originalPartitionName : null)
+                        historicalLookup ? originalPartitionName : null,
+                        bucketCount)
                 .whenComplete(
                         (result, error) -> {
                             if (error != null) {