Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Kaixuan-Duan wants to merge 16 commits into
base: main
Choose a base branch
from
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
16 commits
Select commit Hold shift + click to select a range
917013c
[server][client] Support per-partition bucket count for partitioned t…
Kaixuan-Duan Aug 27, 2026
bf61c55
Guard the historical reroute against mismatched bucket counts
Kaixuan-Duan Sep 3, 2026
c223a7e
Rename bucket_count_actual to bucket_count and bucket_layout_epoch to…
Kaixuan-Duan Sep 3, 2026
a26a06b
Resolve the historical partition's bucket count in the tiering reader
Kaixuan-Duan Sep 4, 2026
8ad6c43
[server] Report a stale routing bucket count per bucket, not per request
platinumhamburg Sep 4, 2026
687d036
[server] Skip routing bucket count validation for keyless tables
platinumhamburg Sep 4, 2026
dc2028d
[client] Reclaim the batch sequence when a write is rejected with STA…
platinumhamburg Sep 4, 2026
ae15934
[client] Fall back to table-level bucket count for never-rescaled par…
platinumhamburg Sep 4, 2026
8f1f1fd
Use hash-distributed tables in the routing validation IT cases
Kaixuan-Duan Sep 4, 2026
9c79417
[flink] Enumerate KV batch splits by the partition's own bucket count
Kaixuan-Duan Sep 4, 2026
3316d3b
[client] Centralize the missing-bucket-count fallback into one policy
platinumhamburg Sep 4, 2026
d7ae991
[server] Preserve bucketCountEpoch across schema changes
Kaixuan-Duan Sep 4, 2026
bec161c
Revert "[server] Skip routing bucket count validation for keyless tab…
Kaixuan-Duan Sep 8, 2026
3a3a0c3
[server] Make bucket rescale and historical partition mutually exclusive
Kaixuan-Duan Sep 8, 2026
0e99fff
[test] Extract CoordinatorEventProcessorTest harness into base class
Kaixuan-Duan Sep 8, 2026
c60e5b4
[test] Remove the historical-partition rescale lookup ITCase
Kaixuan-Duan Sep 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
82 changes: 65 additions & 17 deletions fluss-client/src/main/java/org/apache/fluss/client/admin/FlussAdmin.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -93,6 +93,7 @@
import org.apache.fluss.rpc.messages.ListTablesResponse;
import org.apache.fluss.rpc.messages.PbAlterConfig;
import org.apache.fluss.rpc.messages.PbListOffsetsRespForBucket;
import org.apache.fluss.rpc.messages.PbPartitionInfo;
import org.apache.fluss.rpc.messages.PbPartitionSpec;
import org.apache.fluss.rpc.messages.PbTablePath;
import org.apache.fluss.rpc.messages.PbTableStatsRespForBucket;
Expand Down Expand Up @@ -344,7 +345,8 @@ public CompletableFuture<TableInfo> getTableInfo(TablePath tablePath) {
// clusters do not include the remote data dir
r.hasRemoteDataDir() ? r.getRemoteDataDir() : null,
r.getCreatedTime(),
r.getModifiedTime()));
r.getModifiedTime(),
r.hasBucketCountEpoch() ? r.getBucketCountEpoch() : 0L));
}

@Override
Expand Down Expand Up @@ -393,7 +395,41 @@ public CompletableFuture<List<PartitionInfo>> listPartitionInfos(
}
return readOnlyGateway
.listPartitionInfos(request)
.thenApply(ClientRpcMessageUtils::toPartitionInfos);
.thenCompose(
response -> {
boolean allHaveBucketCount =
response.getPartitionsInfosList().stream()
.allMatch(PbPartitionInfo::hasBucketCount);
if (allHaveBucketCount) {
// Every partition already carries its own bucket count, so skip
// the extra getTableInfo RPC (the -1 default is never used).
return CompletableFuture.completedFuture(
ClientRpcMessageUtils.toPartitionInfos(response, -1));
}
// Upgrade contract for the per-partition bucket count fallback:
// 1) upgrade clients first;
// 2) upgrade the CoordinatorServer before (or together with) the
// TabletServers; a TabletServer newer than its Coordinator
// fails leader activation with UnsupportedVersionException;
// 3) prohibit ALTER bucket.num during the server rolling upgrade;
// 4) a fully old cluster omits the partition count, while the
// table-level count is still safe (bucketCountEpoch == 0);
// 5) after every server is upgraded, ListPartitionInfosResponse
// must return an explicit partitionId and bucketCount;
// 6) if a new server still returns a missing count, fail loud
// instead of calling getTableInfo to guess it.
// TODO: a future FIP should enforce server-side rejection of
// bucket-layout ALTERs during a rolling upgrade so clients never
// observe a half-upgraded cluster.
return getTableInfo(tablePath)
.thenApply(
tableInfo ->
ClientRpcMessageUtils.toPartitionInfos(
response,
ClientRpcMessageUtils
.fallbackBucketCountOrFail(
tableInfo, tablePath)));
});
}

/**
Expand Down Expand Up @@ -549,30 +585,34 @@ public CompletableFuture<TableStats> getTableStats(TablePath tablePath) {
metadataUpdater.updateTableOrPartitionMetadata(tablePath, null);
TableInfo tableInfo = getTableInfo(tablePath).join();
try {
int bucketCount = tableInfo.getNumBuckets();
int tableBucketCount = tableInfo.getNumBuckets();
List<PartitionInfo> partitionInfos;
if (tableInfo.isPartitioned()) {
partitionInfos = listPartitionInfos(tablePath).get();
} else {
partitionInfos = Collections.singletonList(null);
}
// create all TableBuckets for each partition and bucket combination
Map<TableBucket, CompletableFuture<Long>> bucketToRowCountMap = new HashMap<>();
List<TableBucket> tableBuckets = new ArrayList<>();
for (PartitionInfo partitionInfo : partitionInfos) {
int bucketCount =
PartitionInfo.bucketCountOrDefault(partitionInfo, tableBucketCount);
for (int bucket = 0; bucket < bucketCount; bucket++) {
TableBucket tb =
tableBuckets.add(
new TableBucket(
tableInfo.getTableId(),
partitionInfo == null ? null : partitionInfo.getPartitionId(),
bucket);
bucketToRowCountMap.put(tb, new CompletableFuture<>());
bucket));
}
}
long tableId = tableInfo.getTableId();
Map<TableBucket, CompletableFuture<Long>> bucketToRowCountMap = new HashMap<>();
for (TableBucket tb : tableBuckets) {
bucketToRowCountMap.put(tb, new CompletableFuture<>());
}
Comment on lines +595 to +611

@wuchong wuchong Sep 9, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
List<TableBucket> tableBuckets = new ArrayList<>();
for (PartitionInfo partitionInfo : partitionInfos) {
int bucketCount =
PartitionInfo.bucketCountOrDefault(partitionInfo, tableBucketCount);
for (int bucket = 0; bucket < bucketCount; bucket++) {
TableBucket tb =
tableBuckets.add(
new TableBucket(
tableInfo.getTableId(),
partitionInfo == null ? null : partitionInfo.getPartitionId(),
bucket);
bucketToRowCountMap.put(tb, new CompletableFuture<>());
bucket));
}
}
long tableId = tableInfo.getTableId();
Map<TableBucket, CompletableFuture<Long>> bucketToRowCountMap = new HashMap<>();
for (TableBucket tb : tableBuckets) {
bucketToRowCountMap.put(tb, new CompletableFuture<>());
}
long tableId = tableInfo.getTableId();
Map<TableBucket, CompletableFuture<Long>> bucketToRowCountMap = new HashMap<>();
for (PartitionInfo partitionInfo : partitionInfos) {
int bucketCount =
PartitionInfo.bucketCountOrDefault(partitionInfo, tableBucketCount);
for (int bucket = 0; bucket < bucketCount; bucket++) {
TableBucket tb =
new TableBucket(
tableId,
partitionInfo == null ? null : partitionInfo.getPartitionId(),
bucket);
bucketToRowCountMap.put(tb, new CompletableFuture<>());
}
}

Can simplify this 2 foreach into 1 foreach and simplify the logic.

Map<Integer, GetTableStatsRequest> requestMap =
prepareTableStatsRequests(
metadataUpdater, bucketToRowCountMap.keySet(), tablePath);
sendTableStatsRequest(
metadataUpdater, tableInfo.getTableId(), requestMap, bucketToRowCountMap);
sendTableStatsRequest(metadataUpdater, tableId, requestMap, bucketToRowCountMap);
return FutureUtils.combineAll(bucketToRowCountMap.values())
.thenApply(
counts -> {
Expand Down Expand Up @@ -606,13 +646,13 @@ private ListOffsetsResult listOffsets(
buckets,
offsetSpec,
tableInfo.getTablePath());
Map<Integer, CompletableFuture<Long>> bucketToOffsetMap = new ConcurrentHashMap<>();

Map<Integer, CompletableFuture<Long>> resultMap = new ConcurrentHashMap<>();
for (int bucket : buckets) {
bucketToOffsetMap.put(bucket, new CompletableFuture<>());
resultMap.put(bucket, new CompletableFuture<>());
}

sendListOffsetsRequest(metadataUpdater, requestMap, bucketToOffsetMap);
return new ListOffsetsResult(bucketToOffsetMap);
sendListOffsetsRequest(metadataUpdater, requestMap, resultMap);
return new ListOffsetsResult(resultMap);
}

@Override
Expand Down Expand Up @@ -818,7 +858,10 @@ private static Map<Integer, GetTableStatsRequest> prepareTableStatsRequests(

Map<Integer, GetTableStatsRequest> requests = new HashMap<>();
nodeForBucketList.forEach(
(leader, tbs) -> requests.put(leader, makeGetTableStatsRequest(tbs)));
(leader, tbs) ->
requests.put(
leader,
makeGetTableStatsRequest(tbs, metadataUpdater.getCluster())));
return requests;
}

Expand Down Expand Up @@ -890,7 +933,12 @@ private static Map<Integer, ListOffsetsRequest> prepareListOffsetsRequests(
(leader, ids) ->
listOffsetsRequests.put(
leader,
makeListOffsetsRequest(tableId, partitionId, ids, offsetSpec)));
makeListOffsetsRequest(
tableId,
partitionId,
ids,
offsetSpec,
metadataUpdater.getCluster())));
return listOffsetsRequests;
}

Expand Down
18 changes: 17 additions & 1 deletion fluss-client/src/main/java/org/apache/fluss/client/lookup/AbstractLookupQuery.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Comment thread
Kaixuan-Duan marked this conversation as resolved.
this(tablePath, tableBucket, key, null);
this(tablePath, tableBucket, key, null, 0);
}

public AbstractLookupQuery(
Comment thread
Kaixuan-Duan marked this conversation as resolved.
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;
}
Expand All @@ -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;
}
Expand Down
17 changes: 17 additions & 0 deletions fluss-client/src/main/java/org/apache/fluss/client/lookup/AbstractLookuper.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,11 +18,13 @@
package org.apache.fluss.client.lookup;

import org.apache.fluss.client.metadata.MetadataUpdater;
import org.apache.fluss.client.utils.ClientRpcMessageUtils;
import org.apache.fluss.memory.MemorySegment;
import org.apache.fluss.metadata.Schema;
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.row.encode.KvValueLayout;
Expand Down Expand Up @@ -76,6 +78,21 @@ abstract class AbstractLookuper implements Lookuper {
tableInfo.getTableConfig().getKvFormat(), tableInfo.getSchema()));
}

/**
* Resolves the effective bucket count for the target partition: the per-partition bucket count
* when the cluster metadata has it, otherwise the shared fallback policy (safe table-level
* count only when the table was never rescaled; fail loud otherwise).
*/
protected int resolvePartitionBucketCount(TablePartition tablePartition) {
return metadataUpdater
.getCluster()
.getBucketCount(tablePartition)
.orElseGet(
() ->
ClientRpcMessageUtils.fallbackBucketCountOrFail(
tableInfo, tablePartition));
}

protected void handleLookupResponse(
List<byte[]> result, CompletableFuture<LookupResult> lookupFuture) {
List<MemorySegment> valueList = new ArrayList<>(result.size());
Expand Down
10 changes: 9 additions & 1 deletion fluss-client/src/main/java/org/apache/fluss/client/lookup/LookupBatch.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand All @@ -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;
}
Expand Down
15 changes: 11 additions & 4 deletions fluss-client/src/main/java/org/apache/fluss/client/lookup/LookupClient.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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();
}
Expand Down
16 changes: 13 additions & 3 deletions fluss-client/src/main/java/org/apache/fluss/client/lookup/LookupQuery.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(

@wuchong wuchong Sep 9, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Add @VisibleForTesting

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
Expand Down
12 changes: 10 additions & 2 deletions fluss-client/src/main/java/org/apache/fluss/client/lookup/LookupSender.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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);
}

Expand Down Expand Up @@ -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);
}

Expand Down Expand Up @@ -565,6 +566,13 @@ private void handleLookupError(
invalidTableOrPartitions(tableOrPartitions);
}

if (error.error() == Errors.STALE_METADATA) {
for (AbstractLookupQuery<?> lookup : lookups) {
lookup.future().completeExceptionally(exception);
}
return;
}
Comment on lines +569 to +574

@wuchong wuchong Sep 9, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we need to specially handling STALE_METADATA error? If it is not a RetriableException, I think existing code can already correctly completeExceptionally the futures.


for (AbstractLookupQuery<?> lookup : lookups) {
String originalPartitionNameMsg =
lookup.originalPartitionName() == null
Expand Down
Loading
Loading