diff --git a/fluss-client/src/main/java/org/apache/fluss/client/admin/Admin.java b/fluss-client/src/main/java/org/apache/fluss/client/admin/Admin.java index 5d749b6d432..8ec899a9985 100644 --- a/fluss-client/src/main/java/org/apache/fluss/client/admin/Admin.java +++ b/fluss-client/src/main/java/org/apache/fluss/client/admin/Admin.java @@ -324,6 +324,30 @@ CompletableFuture alterTable( */ CompletableFuture> listPartitionInfos(TablePath tablePath); + /** + * List partitions in the given table asynchronously, optionally including coordinator-managed + * system partitions. + * + *

System partitions are created and managed by the coordinator to support internal table + * functionality. For example, the {@code __historical__} partition is created when {@link + * ConfigOptions#TABLE_DATALAKE_HISTORICAL_PARTITION_ENABLED} is enabled. It provides a shared + * routing target for writes to expired partitions and, for primary-key tables, lookups of + * expired partition data in lake storage. + * + *

The following exceptions can be anticipated when calling {@code get()} on returned future. + * + *

+ * + * @param tablePath The path of the table. + * @param includeSystemPartitions If true, also include system partitions such as {@code + * __historical__}; otherwise, return only regular partitions. + */ + CompletableFuture> listPartitionInfos( + TablePath tablePath, boolean includeSystemPartitions); + /** * List all partitions in fluss cluster that are under the given table and the given partial * PartitionSpec asynchronously. diff --git a/fluss-client/src/main/java/org/apache/fluss/client/admin/FlussAdmin.java b/fluss-client/src/main/java/org/apache/fluss/client/admin/FlussAdmin.java index a0909ceec73..d9251f4fc0a 100644 --- a/fluss-client/src/main/java/org/apache/fluss/client/admin/FlussAdmin.java +++ b/fluss-client/src/main/java/org/apache/fluss/client/admin/FlussAdmin.java @@ -25,6 +25,7 @@ import org.apache.fluss.client.metadata.MetadataUpdater; import org.apache.fluss.client.metadata.RemoteLogManifestInfo; import org.apache.fluss.client.utils.ClientRpcMessageUtils; +import org.apache.fluss.client.utils.ClientUtils; import org.apache.fluss.cluster.Cluster; import org.apache.fluss.cluster.ServerNode; import org.apache.fluss.cluster.rebalance.GoalType; @@ -34,6 +35,7 @@ import org.apache.fluss.config.cluster.ConfigEntry; import org.apache.fluss.exception.FlussRuntimeException; import org.apache.fluss.exception.LeaderNotAvailableException; +import org.apache.fluss.exception.PartitionNotExistException; import org.apache.fluss.metadata.DatabaseChange; import org.apache.fluss.metadata.DatabaseDescriptor; import org.apache.fluss.metadata.DatabaseInfo; @@ -41,6 +43,7 @@ import org.apache.fluss.metadata.PartitionInfo; import org.apache.fluss.metadata.PartitionSpec; import org.apache.fluss.metadata.PhysicalTablePath; +import org.apache.fluss.metadata.ResolvedPartitionSpec; import org.apache.fluss.metadata.Schema; import org.apache.fluss.metadata.SchemaInfo; import org.apache.fluss.metadata.TableBucket; @@ -87,13 +90,14 @@ import org.apache.fluss.rpc.messages.ListOffsetsRequest; import org.apache.fluss.rpc.messages.ListOffsetsResponse; import org.apache.fluss.rpc.messages.ListPartitionInfosRequest; +import org.apache.fluss.rpc.messages.ListPartitionInfosResponse; import org.apache.fluss.rpc.messages.ListRebalanceProgressRequest; import org.apache.fluss.rpc.messages.ListRemoteLogManifestsRequest; import org.apache.fluss.rpc.messages.ListTablesRequest; 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.PbPartitionSpec; +import org.apache.fluss.rpc.messages.PbPartitionInfo; import org.apache.fluss.rpc.messages.PbTablePath; import org.apache.fluss.rpc.messages.PbTableStatsRespForBucket; import org.apache.fluss.rpc.messages.RebalanceRequest; @@ -104,6 +108,7 @@ import org.apache.fluss.rpc.protocol.ApiError; import org.apache.fluss.security.acl.AclBinding; import org.apache.fluss.security.acl.AclBindingFilter; +import org.apache.fluss.utils.ExceptionUtils; import org.apache.fluss.utils.concurrent.ExecutorThreadFactory; import org.apache.fluss.utils.concurrent.FutureUtils; @@ -135,6 +140,7 @@ import static org.apache.fluss.rpc.util.CommonRpcMessageUtils.toPbAclBindingFilters; import static org.apache.fluss.rpc.util.CommonRpcMessageUtils.toPbAclFilter; import static org.apache.fluss.rpc.util.CommonRpcMessageUtils.toPbAclInfos; +import static org.apache.fluss.utils.PartitionUtils.HISTORICAL_PARTITION_VALUE; import static org.apache.fluss.utils.Preconditions.checkNotNull; /** @@ -344,7 +350,8 @@ public CompletableFuture 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 @@ -375,25 +382,141 @@ public CompletableFuture> listTables(String databaseName) { @Override public CompletableFuture> listPartitionInfos(TablePath tablePath) { - return listPartitionInfos(tablePath, null); + return listPartitionInfos(tablePath, null, false); + } + + @Override + public CompletableFuture> listPartitionInfos( + TablePath tablePath, boolean includeSystemPartitions) { + return listPartitionInfos(tablePath, null, includeSystemPartitions); } @Override public CompletableFuture> listPartitionInfos( TablePath tablePath, PartitionSpec partitionSpec) { + return listPartitionInfos(tablePath, partitionSpec, false); + } + + private CompletableFuture> listPartitionInfos( + TablePath tablePath, + @Nullable PartitionSpec partitionSpec, + boolean includeSystemPartitions) { ListPartitionInfosRequest request = new ListPartitionInfosRequest(); request.setTablePath( new PbTablePath() .setDatabaseName(tablePath.getDatabaseName()) .setTableName(tablePath.getTableName())); - if (partitionSpec != null) { - PbPartitionSpec pbPartitionSpec = makePbPartitionSpec(partitionSpec); - request.setPartialPartitionSpec(pbPartitionSpec); + request.setPartialPartitionSpec(makePbPartitionSpec(partitionSpec)); + } + if (includeSystemPartitions) { + request.setIncludeSystemPartitions(true); } + return readOnlyGateway .listPartitionInfos(request) - .thenApply(ClientRpcMessageUtils::toPartitionInfos); + .thenCompose( + response -> + handleListPartitionInfosResponse( + tablePath, includeSystemPartitions, response)); + } + + @VisibleForTesting + CompletableFuture> handleListPartitionInfosResponse( + TablePath tablePath, + boolean includeSystemPartitions, + ListPartitionInfosResponse response) { + boolean allHaveBucketCount = + response.getPartitionsInfosList().stream() + .allMatch(PbPartitionInfo::hasBucketCount); + boolean systemPartitionsIncluded = + response.hasSystemPartitionsIncluded() && response.isSystemPartitionsIncluded(); + if (allHaveBucketCount && (!includeSystemPartitions || systemPartitionsIncluded)) { + return CompletableFuture.completedFuture( + ClientRpcMessageUtils.toPartitionInfos(response, -1)); + } + return getTableInfo(tablePath) + .thenCompose( + tableInfo -> { + int defaultBucketCount = + allHaveBucketCount + ? -1 + : ClientUtils.fallbackBucketCountOrFail( + tableInfo, tablePath); + List partitionInfos = + ClientRpcMessageUtils.toPartitionInfos( + response, defaultBucketCount); + if (includeSystemPartitions && !systemPartitionsIncluded) { + return appendLegacyHistoricalPartition(tableInfo, partitionInfos); + } + return CompletableFuture.completedFuture(partitionInfos); + }); + } + + private CompletableFuture> appendLegacyHistoricalPartition( + TableInfo tableInfo, List partitionInfos) { + if (!tableInfo.getTableConfig().isHistoricalPartitionEnabled() + || partitionInfos.stream() + .anyMatch( + partitionInfo -> + HISTORICAL_PARTITION_VALUE.equals( + partitionInfo.getPartitionName()))) { + return CompletableFuture.completedFuture(partitionInfos); + } + + TablePath tablePath = tableInfo.getTablePath(); + PhysicalTablePath historicalPartitionPath = + PhysicalTablePath.of(tablePath, HISTORICAL_PARTITION_VALUE); + Cluster cluster = metadataUpdater.getCluster(); + // A cached partition may belong to a previous table with the same name. + Optional historicalPartitionId = + cluster.getTableId(tablePath) + .filter(tableId -> tableId == tableInfo.getTableId()) + .flatMap(ignored -> cluster.getPartitionId(historicalPartitionPath)); + CompletableFuture> partitionIdFuture; + if (historicalPartitionId.isPresent()) { + partitionIdFuture = CompletableFuture.completedFuture(historicalPartitionId); + } else { + partitionIdFuture = + CompletableFuture.supplyAsync( + () -> { + try { + Cluster refreshedCluster = + sendMetadataRequestAndRebuildCluster( + readOnlyGateway, + true, + cluster, + Collections.singleton(tablePath), + Collections.singleton(historicalPartitionPath), + null); + return refreshedCluster.getPartitionId(historicalPartitionPath); + } catch (Exception e) { + Throwable cause = ExceptionUtils.stripExecutionException(e); + if (cause instanceof PartitionNotExistException) { + return Optional.empty(); + } + throw new FlussRuntimeException( + "Failed to resolve historical partition for " + + tablePath, + cause); + } + }, + refreshExecutor); + } + return partitionIdFuture.thenApply( + partitionId -> { + if (partitionId.isPresent()) { + partitionInfos.add( + new PartitionInfo( + partitionId.get(), + ResolvedPartitionSpec.fromPartitionName( + tableInfo.getPartitionKeys(), + HISTORICAL_PARTITION_VALUE), + null, + tableInfo.getNumBuckets())); + } + return partitionInfos; + }); } /** @@ -549,30 +672,30 @@ public CompletableFuture getTableStats(TablePath tablePath) { metadataUpdater.updateTableOrPartitionMetadata(tablePath, null); TableInfo tableInfo = getTableInfo(tablePath).join(); try { - int bucketCount = tableInfo.getNumBuckets(); + int tableBucketCount = tableInfo.getNumBuckets(); List partitionInfos; if (tableInfo.isPartitioned()) { partitionInfos = listPartitionInfos(tablePath).get(); } else { partitionInfos = Collections.singletonList(null); } - // create all TableBuckets for each partition and bucket combination + + long tableId = tableInfo.getTableId(); Map> bucketToRowCountMap = new HashMap<>(); for (PartitionInfo partitionInfo : partitionInfos) { + int bucketCount = + PartitionInfo.bucketCountOrDefault(partitionInfo, tableBucketCount); + Long partitionId = partitionInfo == null ? null : partitionInfo.getPartitionId(); for (int bucket = 0; bucket < bucketCount; bucket++) { - TableBucket tb = - new TableBucket( - tableInfo.getTableId(), - partitionInfo == null ? null : partitionInfo.getPartitionId(), - bucket); - bucketToRowCountMap.put(tb, new CompletableFuture<>()); + TableBucket tableBucket = new TableBucket(tableId, partitionId, bucket); + bucketToRowCountMap.put(tableBucket, new CompletableFuture<>()); } } + Map requestMap = prepareTableStatsRequests( metadataUpdater, bucketToRowCountMap.keySet(), tablePath); - sendTableStatsRequest( - metadataUpdater, tableInfo.getTableId(), requestMap, bucketToRowCountMap); + sendTableStatsRequest(metadataUpdater, tableId, requestMap, bucketToRowCountMap); return FutureUtils.combineAll(bucketToRowCountMap.values()) .thenApply( counts -> { @@ -606,13 +729,13 @@ private ListOffsetsResult listOffsets( buckets, offsetSpec, tableInfo.getTablePath()); - Map> bucketToOffsetMap = new ConcurrentHashMap<>(); + + Map> 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 @@ -818,7 +941,10 @@ private static Map prepareTableStatsRequests( Map requests = new HashMap<>(); nodeForBucketList.forEach( - (leader, tbs) -> requests.put(leader, makeGetTableStatsRequest(tbs))); + (leader, tbs) -> + requests.put( + leader, + makeGetTableStatsRequest(tbs, metadataUpdater.getCluster()))); return requests; } @@ -890,7 +1016,12 @@ private static Map prepareListOffsetsRequests( (leader, ids) -> listOffsetsRequests.put( leader, - makeListOffsetsRequest(tableId, partitionId, ids, offsetSpec))); + makeListOffsetsRequest( + tableId, + partitionId, + ids, + offsetSpec, + metadataUpdater.getCluster()))); return listOffsetsRequests; } diff --git a/fluss-client/src/main/java/org/apache/fluss/client/lookup/AbstractLookupQuery.java b/fluss-client/src/main/java/org/apache/fluss/client/lookup/AbstractLookupQuery.java index 19bea07bf20..6a0457d213d 100644 --- a/fluss-client/src/main/java/org/apache/fluss/client/lookup/AbstractLookupQuery.java +++ b/fluss-client/src/main/java/org/apache/fluss/client/lookup/AbstractLookupQuery.java @@ -39,22 +39,21 @@ public abstract class AbstractLookupQuery { */ 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); - } - public AbstractLookupQuery( TablePath tablePath, TableBucket tableBucket, byte[] key, - @Nullable String originalPartitionName) { + @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 +74,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; } diff --git a/fluss-client/src/main/java/org/apache/fluss/client/lookup/AbstractLookuper.java b/fluss-client/src/main/java/org/apache/fluss/client/lookup/AbstractLookuper.java index fd8ec59a4f4..d0de708739b 100644 --- a/fluss-client/src/main/java/org/apache/fluss/client/lookup/AbstractLookuper.java +++ b/fluss-client/src/main/java/org/apache/fluss/client/lookup/AbstractLookuper.java @@ -18,7 +18,9 @@ package org.apache.fluss.client.lookup; import org.apache.fluss.client.metadata.MetadataUpdater; +import org.apache.fluss.cluster.Cluster; import org.apache.fluss.memory.MemorySegment; +import org.apache.fluss.metadata.PhysicalTablePath; import org.apache.fluss.metadata.Schema; import org.apache.fluss.metadata.SchemaGetter; import org.apache.fluss.metadata.SchemaInfo; @@ -76,6 +78,17 @@ abstract class AbstractLookuper implements Lookuper { tableInfo.getTableConfig().getKvFormat(), tableInfo.getSchema())); } + protected PartitionRoutingInfo resolvePartitionRouting(String partitionName) { + PhysicalTablePath partitionPath = + PhysicalTablePath.of(tableInfo.getTablePath(), partitionName); + metadataUpdater.checkAndUpdatePartitionMetadata(partitionPath); + + Cluster cluster = metadataUpdater.getCluster(); + long partitionId = cluster.getPartitionIdOrElseThrow(partitionPath); + int bucketCount = cluster.getBucketCountOrFallback(tableInfo, partitionId); + return new PartitionRoutingInfo(partitionId, bucketCount); + } + protected void handleLookupResponse( List result, CompletableFuture lookupFuture) { List valueList = new ArrayList<>(result.size()); @@ -178,4 +191,22 @@ protected LookupResult processSchemaRequestedRows( } return new LookupResult(rowList); } + + static final class PartitionRoutingInfo { + private final long partitionId; + private final int bucketCount; + + private PartitionRoutingInfo(long partitionId, int bucketCount) { + this.partitionId = partitionId; + this.bucketCount = bucketCount; + } + + long getPartitionId() { + return partitionId; + } + + int getBucketCount() { + return bucketCount; + } + } } diff --git a/fluss-client/src/main/java/org/apache/fluss/client/lookup/LookupBatch.java b/fluss-client/src/main/java/org/apache/fluss/client/lookup/LookupBatch.java index 0d378ab4f21..a535b967b67 100644 --- a/fluss-client/src/main/java/org/apache/fluss/client/lookup/LookupBatch.java +++ b/fluss-client/src/main/java/org/apache/fluss/client/lookup/LookupBatch.java @@ -34,9 +34,12 @@ public class LookupBatch { private final List 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; } diff --git a/fluss-client/src/main/java/org/apache/fluss/client/lookup/LookupClient.java b/fluss-client/src/main/java/org/apache/fluss/client/lookup/LookupClient.java index f467b2b5f5c..21fda7fe5a5 100644 --- a/fluss-client/src/main/java/org/apache/fluss/client/lookup/LookupClient.java +++ b/fluss-client/src/main/java/org/apache/fluss/client/lookup/LookupClient.java @@ -113,17 +113,24 @@ public CompletableFuture 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> 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(); } diff --git a/fluss-client/src/main/java/org/apache/fluss/client/lookup/LookupQuery.java b/fluss-client/src/main/java/org/apache/fluss/client/lookup/LookupQuery.java index 3f0218d6d25..7dce9dedb9e 100644 --- a/fluss-client/src/main/java/org/apache/fluss/client/lookup/LookupQuery.java +++ b/fluss-client/src/main/java/org/apache/fluss/client/lookup/LookupQuery.java @@ -41,15 +41,26 @@ public class LookupQuery extends AbstractLookupQuery { 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; } + @VisibleForTesting + 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 diff --git a/fluss-client/src/main/java/org/apache/fluss/client/lookup/LookupSender.java b/fluss-client/src/main/java/org/apache/fluss/client/lookup/LookupSender.java index e608494e43b..7823388c4de 100644 --- a/fluss-client/src/main/java/org/apache/fluss/client/lookup/LookupSender.java +++ b/fluss-client/src/main/java/org/apache/fluss/client/lookup/LookupSender.java @@ -23,6 +23,7 @@ import org.apache.fluss.exception.ApiException; import org.apache.fluss.exception.FlussRuntimeException; import org.apache.fluss.exception.HistoricalPartitionThrottledException; +import org.apache.fluss.exception.InvalidBucketRoutingException; import org.apache.fluss.exception.InvalidMetadataException; import org.apache.fluss.exception.LeaderNotAvailableException; import org.apache.fluss.exception.PartitionNotExistException; @@ -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); } @@ -546,9 +547,11 @@ private void handleLookupError( destination, tableBucket, exception); - if (exception instanceof InvalidMetadataException) { + if (exception instanceof InvalidMetadataException + || exception instanceof InvalidBucketRoutingException) { LOG.warn( - "Invalid metadata error in {} request. Going to request metadata update.", + "Metadata or bucket routing error in {} request. Going to request metadata " + + "update.", lookupType, exception); long tableId = tableBucket.getTableId(); diff --git a/fluss-client/src/main/java/org/apache/fluss/client/lookup/PrefixKeyLookuper.java b/fluss-client/src/main/java/org/apache/fluss/client/lookup/PrefixKeyLookuper.java index 85b9aa5ecd7..fcb6e3d0cc2 100644 --- a/fluss-client/src/main/java/org/apache/fluss/client/lookup/PrefixKeyLookuper.java +++ b/fluss-client/src/main/java/org/apache/fluss/client/lookup/PrefixKeyLookuper.java @@ -20,6 +20,7 @@ import org.apache.fluss.bucketing.BucketingFunction; import org.apache.fluss.client.metadata.MetadataUpdater; import org.apache.fluss.client.table.getter.PartitionGetter; +import org.apache.fluss.exception.InvalidBucketRoutingException; import org.apache.fluss.exception.PartitionNotExistException; import org.apache.fluss.metadata.DataLakeFormat; import org.apache.fluss.metadata.SchemaGetter; @@ -39,8 +40,6 @@ import java.util.Set; import java.util.concurrent.CompletableFuture; -import static org.apache.fluss.client.utils.ClientUtils.getPartitionId; - /** * An implementation of {@link Lookuper} that lookups by prefix key. A prefix key is a prefix subset * of the primary key. @@ -165,26 +164,31 @@ public CompletableFuture lookup(InternalRow prefixKey) { prefixKeyEncoder == bucketKeyEncoder ? prefixKeyBytes : bucketKeyEncoder.encodeKey(prefixKey); - int bucketId = bucketingFunction.bucketing(bucketKeyBytes, numBuckets); Long partitionId = null; + int bucketCount = numBuckets; if (partitionGetter != null) { try { - partitionId = - getPartitionId( - prefixKey, - partitionGetter, - tableInfo.getTablePath(), - metadataUpdater); + PartitionRoutingInfo routing = + resolvePartitionRouting(partitionGetter.getPartition(prefixKey)); + partitionId = routing.getPartitionId(); + bucketCount = routing.getBucketCount(); } catch (PartitionNotExistException e) { return CompletableFuture.completedFuture(new LookupResult(Collections.emptyList())); + } catch (InvalidBucketRoutingException e) { + CompletableFuture failed = new CompletableFuture<>(); + failed.completeExceptionally(e); + return failed; } } + // Compute bucket ID after partition resolution — needs per-partition bucket count + int bucketId = bucketingFunction.bucketing(bucketKeyBytes, bucketCount); + CompletableFuture lookupFuture = new CompletableFuture<>(); TableBucket tableBucket = new TableBucket(tableInfo.getTableId(), partitionId, bucketId); lookupClient - .prefixLookup(tableInfo.getTablePath(), tableBucket, prefixKeyBytes) + .prefixLookup(tableInfo.getTablePath(), tableBucket, prefixKeyBytes, bucketCount) .whenComplete( (result, error) -> { if (error != null) { diff --git a/fluss-client/src/main/java/org/apache/fluss/client/lookup/PrefixLookupBatch.java b/fluss-client/src/main/java/org/apache/fluss/client/lookup/PrefixLookupBatch.java index 0932fcfbbfa..902dd9aa10f 100644 --- a/fluss-client/src/main/java/org/apache/fluss/client/lookup/PrefixLookupBatch.java +++ b/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 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> values) { if (values.size() != prefixLookups.size()) { completeExceptionally( diff --git a/fluss-client/src/main/java/org/apache/fluss/client/lookup/PrefixLookupQuery.java b/fluss-client/src/main/java/org/apache/fluss/client/lookup/PrefixLookupQuery.java index 5246aedca20..9d3c2d2ca99 100644 --- a/fluss-client/src/main/java/org/apache/fluss/client/lookup/PrefixLookupQuery.java +++ b/fluss-client/src/main/java/org/apache/fluss/client/lookup/PrefixLookupQuery.java @@ -32,11 +32,16 @@ public class PrefixLookupQuery extends AbstractLookupQuery> { private final CompletableFuture> 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> future() { return future; diff --git a/fluss-client/src/main/java/org/apache/fluss/client/lookup/PrimaryKeyLookuper.java b/fluss-client/src/main/java/org/apache/fluss/client/lookup/PrimaryKeyLookuper.java index d235649e036..dd5126a4834 100644 --- a/fluss-client/src/main/java/org/apache/fluss/client/lookup/PrimaryKeyLookuper.java +++ b/fluss-client/src/main/java/org/apache/fluss/client/lookup/PrimaryKeyLookuper.java @@ -20,6 +20,7 @@ import org.apache.fluss.bucketing.BucketingFunction; import org.apache.fluss.client.metadata.MetadataUpdater; import org.apache.fluss.client.table.getter.PartitionGetter; +import org.apache.fluss.exception.InvalidBucketRoutingException; import org.apache.fluss.exception.PartitionNotExistException; import org.apache.fluss.metadata.DataLakeFormat; import org.apache.fluss.metadata.PhysicalTablePath; @@ -39,7 +40,6 @@ import java.util.Set; import java.util.concurrent.CompletableFuture; -import static org.apache.fluss.client.utils.ClientUtils.getPartitionId; import static org.apache.fluss.utils.PartitionUtils.HISTORICAL_PARTITION_VALUE; import static org.apache.fluss.utils.PartitionUtils.isPastAutoPartition; import static org.apache.fluss.utils.Preconditions.checkArgument; @@ -124,35 +124,42 @@ public CompletableFuture lookup(InternalRow lookupKey) { bucketKeyEncoder == primaryKeyEncoder ? pkBytes : bucketKeyEncoder.encodeKey(lookupKey); - int bucketId = bucketingFunction.bucketing(bkBytes, numBuckets); Long partitionId = null; String originalPartitionName = null; + int bucketCount = numBuckets; if (partitionGetter != null) { originalPartitionName = partitionGetter.getPartition(lookupKey); if (confirmedHistoricalPartitions.contains(originalPartitionName)) { - return historicalLookup(bucketId, pkBytes, originalPartitionName); + return historicalLookup(bkBytes, pkBytes, originalPartitionName); } try { - partitionId = - getPartitionId( - lookupKey, - partitionGetter, - tableInfo.getTablePath(), - metadataUpdater); + PartitionRoutingInfo routing = resolvePartitionRouting(originalPartitionName); + partitionId = routing.getPartitionId(); + bucketCount = routing.getBucketCount(); } catch (PartitionNotExistException e) { - return mayFallbackToHistoricalLookup(bucketId, pkBytes, originalPartitionName); + return mayFallbackToHistoricalLookup(bkBytes, pkBytes, originalPartitionName); + } catch (InvalidBucketRoutingException e) { + return completedExceptionally(e); } } + int bucketId = bucketingFunction.bucketing(bkBytes, bucketCount); TableBucket tableBucket = new TableBucket(tableInfo.getTableId(), partitionId, bucketId); - return lookupBucket(tableBucket, pkBytes, insertIfNotExists, false, originalPartitionName); + return lookupBucket( + tableBucket, + bkBytes, + pkBytes, + insertIfNotExists, + false, + originalPartitionName, + bucketCount); } /** * Falls back to historical lookup when the normal partition is missing and fallback is enabled. */ private CompletableFuture mayFallbackToHistoricalLookup( - int bucketId, byte[] keyBytes, String originalPartitionName) { + byte[] bucketKeyBytes, byte[] keyBytes, String originalPartitionName) { // Clear the stale normal-partition route before deciding whether to fall back so that a // partition created later can be discovered by the next lookup. metadataUpdater.invalidPhysicalTableBucketAndPartitionMeta( @@ -171,28 +178,33 @@ private CompletableFuture mayFallbackToHistoricalLookup( return CompletableFuture.completedFuture(new LookupResult(Collections.emptyList())); } confirmedHistoricalPartitions.add(originalPartitionName); - return historicalLookup(bucketId, keyBytes, originalPartitionName); + return historicalLookup(bucketKeyBytes, keyBytes, originalPartitionName); } private CompletableFuture historicalLookup( - int bucketId, byte[] keyBytes, String originalPartitionName) { + byte[] bucketKeyBytes, byte[] keyBytes, String originalPartitionName) { if (insertIfNotExists) { return completedExceptionally( new UnsupportedOperationException( "Lookup with insertIfNotExists is not supported for historical partition lookup.")); } - PhysicalTablePath historicalPartitionPath = - PhysicalTablePath.of(tableInfo.getTablePath(), HISTORICAL_PARTITION_VALUE); try { - if (!metadataUpdater.checkAndUpdatePartitionMetadata(historicalPartitionPath)) { - throw new PartitionNotExistException( - "Historical partition " + historicalPartitionPath + " does not exist."); - } - Long historicalPartitionId = - metadataUpdater.getPartitionIdOrElseThrow(historicalPartitionPath); + PartitionRoutingInfo routing = resolvePartitionRouting(HISTORICAL_PARTITION_VALUE); + // Route by the historical partition's own count, which an ALTER bucket.num does not + // change. The bucket the lake data lives in is resolved on the server. + int routingBucketId = + bucketingFunction.bucketing(bucketKeyBytes, routing.getBucketCount()); TableBucket tableBucket = - new TableBucket(tableInfo.getTableId(), historicalPartitionId, bucketId); - return lookupBucket(tableBucket, keyBytes, false, true, originalPartitionName); + new TableBucket( + tableInfo.getTableId(), routing.getPartitionId(), routingBucketId); + return lookupBucket( + tableBucket, + bucketKeyBytes, + keyBytes, + false, + true, + originalPartitionName, + routing.getBucketCount()); } catch (Throwable t) { return completedExceptionally(t); } @@ -200,10 +212,12 @@ private CompletableFuture historicalLookup( private CompletableFuture lookupBucket( TableBucket tableBucket, + byte[] bucketKeyBytes, byte[] keyBytes, boolean insertIfNotExists, boolean historicalLookup, - @Nullable String originalPartitionName) { + @Nullable String originalPartitionName, + int bucketCount) { CompletableFuture lookupFuture = new CompletableFuture<>(); lookupClient .lookup( @@ -211,7 +225,8 @@ private CompletableFuture lookupBucket( tableBucket, keyBytes, insertIfNotExists, - historicalLookup ? originalPartitionName : null) + historicalLookup ? originalPartitionName : null, + bucketCount) .whenComplete( (result, error) -> { if (error != null) { @@ -227,9 +242,7 @@ private CompletableFuture lookupBucket( } mayFallbackToHistoricalLookup( - tableBucket.getBucket(), - keyBytes, - originalPartitionName) + bucketKeyBytes, keyBytes, originalPartitionName) .whenComplete( (historicalResult, historicalError) -> { if (historicalError != null) { diff --git a/fluss-client/src/main/java/org/apache/fluss/client/table/scanner/TableScan.java b/fluss-client/src/main/java/org/apache/fluss/client/table/scanner/TableScan.java index 796f6da3f33..3410ddc240e 100644 --- a/fluss-client/src/main/java/org/apache/fluss/client/table/scanner/TableScan.java +++ b/fluss-client/src/main/java/org/apache/fluss/client/table/scanner/TableScan.java @@ -257,7 +257,7 @@ public BatchScanner createBatchScanner() throws IOException { partitionInfos.stream() .flatMap( partitionInfo -> - IntStream.range(0, bucketCount) + IntStream.range(0, partitionInfo.getBucketCount()) .mapToObj( bucketId -> new TableBucket( diff --git a/fluss-client/src/main/java/org/apache/fluss/client/table/scanner/batch/KvBatchScanner.java b/fluss-client/src/main/java/org/apache/fluss/client/table/scanner/batch/KvBatchScanner.java index 2062f9c1fc5..3b200bc7306 100644 --- a/fluss-client/src/main/java/org/apache/fluss/client/table/scanner/batch/KvBatchScanner.java +++ b/fluss-client/src/main/java/org/apache/fluss/client/table/scanner/batch/KvBatchScanner.java @@ -20,10 +20,12 @@ import org.apache.fluss.annotation.Internal; import org.apache.fluss.annotation.VisibleForTesting; import org.apache.fluss.client.metadata.MetadataUpdater; +import org.apache.fluss.cluster.Cluster; import org.apache.fluss.exception.LeaderNotAvailableException; import org.apache.fluss.metadata.SchemaGetter; import org.apache.fluss.metadata.TableBucket; import org.apache.fluss.metadata.TableInfo; +import org.apache.fluss.metadata.TableOrPartition; import org.apache.fluss.metadata.TablePath; import org.apache.fluss.record.DefaultValueRecordBatch; import org.apache.fluss.record.ValueRecord; @@ -177,6 +179,7 @@ private void openScanner() { "Leader for bucket " + bucket + " is not available. Please retry the scan."); } + Cluster cluster = metadataUpdater.getCluster(); PbScanReqForBucket bucketReq = new PbScanReqForBucket() .setTableId(bucket.getTableId()) @@ -184,6 +187,8 @@ private void openScanner() { if (bucket.getPartitionId() != null) { bucketReq.setPartitionId(bucket.getPartitionId()); } + cluster.getBucketCount(TableOrPartition.of(bucket.getTableId(), bucket.getPartitionId())) + .ifPresent(bucketReq::setRoutingBucketCount); ScanKvRequest request = new ScanKvRequest() diff --git a/fluss-client/src/main/java/org/apache/fluss/client/table/scanner/batch/LimitBatchScanner.java b/fluss-client/src/main/java/org/apache/fluss/client/table/scanner/batch/LimitBatchScanner.java index b5a381ad156..58e39241dcc 100644 --- a/fluss-client/src/main/java/org/apache/fluss/client/table/scanner/batch/LimitBatchScanner.java +++ b/fluss-client/src/main/java/org/apache/fluss/client/table/scanner/batch/LimitBatchScanner.java @@ -18,11 +18,13 @@ package org.apache.fluss.client.table.scanner.batch; import org.apache.fluss.client.metadata.MetadataUpdater; +import org.apache.fluss.cluster.Cluster; import org.apache.fluss.exception.LeaderNotAvailableException; import org.apache.fluss.metadata.KvFormat; import org.apache.fluss.metadata.SchemaGetter; import org.apache.fluss.metadata.TableBucket; import org.apache.fluss.metadata.TableInfo; +import org.apache.fluss.metadata.TableOrPartition; import org.apache.fluss.record.DefaultValueRecordBatch; import org.apache.fluss.record.LogRecord; import org.apache.fluss.record.LogRecordBatch; @@ -103,10 +105,14 @@ public LimitBatchScanner( .setBucketId(tableBucket.getBucket()) .setLimit(limit); + metadataUpdater.checkAndUpdateMetadata(tableInfo.getTablePath(), tableBucket); + Cluster cluster = metadataUpdater.getCluster(); if (tableBucket.getPartitionId() != null) { limitScanRequest.setPartitionId(tableBucket.getPartitionId()); - metadataUpdater.checkAndUpdateMetadata(tableInfo.getTablePath(), tableBucket); } + cluster.getBucketCount( + TableOrPartition.of(tableBucket.getTableId(), tableBucket.getPartitionId())) + .ifPresent(limitScanRequest::setRoutingBucketCount); // because that rocksdb is not suitable to projection, thus do it in client. int leader = metadataUpdater.leaderFor(tableInfo.getTablePath(), tableBucket); diff --git a/fluss-client/src/main/java/org/apache/fluss/client/table/scanner/log/LogFetcher.java b/fluss-client/src/main/java/org/apache/fluss/client/table/scanner/log/LogFetcher.java index 77a8504ca33..2eec56b34cd 100644 --- a/fluss-client/src/main/java/org/apache/fluss/client/table/scanner/log/LogFetcher.java +++ b/fluss-client/src/main/java/org/apache/fluss/client/table/scanner/log/LogFetcher.java @@ -26,6 +26,7 @@ import org.apache.fluss.config.ConfigOptions; import org.apache.fluss.config.Configuration; import org.apache.fluss.exception.ApiException; +import org.apache.fluss.exception.InvalidBucketRoutingException; import org.apache.fluss.exception.InvalidMetadataException; import org.apache.fluss.exception.LeaderNotAvailableException; import org.apache.fluss.exception.PartitionNotExistException; @@ -34,6 +35,7 @@ import org.apache.fluss.metadata.SchemaGetter; import org.apache.fluss.metadata.TableBucket; import org.apache.fluss.metadata.TableInfo; +import org.apache.fluss.metadata.TableOrPartition; import org.apache.fluss.metadata.TablePartition; import org.apache.fluss.metadata.TablePath; import org.apache.fluss.predicate.Predicate; @@ -507,9 +509,10 @@ private synchronized void handleFetchLogResponse( private void handleFetchLogExceptionForBucket(TableBucket tb, int destination, ApiError error) { ApiException exception = error.error().exception(); LOG.error("Failed to fetch log from node {} for bucket {}", destination, tb, exception); - if (exception instanceof InvalidMetadataException) { + if (exception instanceof InvalidMetadataException + || exception instanceof InvalidBucketRoutingException) { LOG.warn( - "Invalid metadata error in fetch log request. " + "Metadata or bucket routing error in fetch log request. " + "Going to request metadata update.", exception); long tableId = tb.getTableId(); @@ -599,6 +602,10 @@ Map prepareFetchLogRequests(List fetchabl if (tb.getPartitionId() != null) { fetchLogReqForBucket.setPartitionId(tb.getPartitionId()); } + metadataUpdater + .getCluster() + .getBucketCount(TableOrPartition.of(tb.getTableId(), tb.getPartitionId())) + .ifPresent(fetchLogReqForBucket::setRoutingBucketCount); fetchReqsByLeaderAndTable .computeIfAbsent(leader, k -> new HashMap<>()) .computeIfAbsent(tb.getTableId(), k -> new ArrayList<>()) diff --git a/fluss-client/src/main/java/org/apache/fluss/client/table/writer/MultiTableWriterImpl.java b/fluss-client/src/main/java/org/apache/fluss/client/table/writer/MultiTableWriterImpl.java index 5702188d3a8..51337802873 100644 --- a/fluss-client/src/main/java/org/apache/fluss/client/table/writer/MultiTableWriterImpl.java +++ b/fluss-client/src/main/java/org/apache/fluss/client/table/writer/MultiTableWriterImpl.java @@ -243,7 +243,8 @@ private static TableInfo withSchema(TableInfo base, int schemaId, Schema schema) base.getRemoteDataDir(), base.getComment().orElse(null), base.getCreatedTime(), - base.getModifiedTime()); + base.getModifiedTime(), + base.getBucketCountEpoch()); } private TableInfo getTableInfo(TablePath tablePath) { diff --git a/fluss-client/src/main/java/org/apache/fluss/client/utils/ClientRpcMessageUtils.java b/fluss-client/src/main/java/org/apache/fluss/client/utils/ClientRpcMessageUtils.java index 45b867e89c1..ee535f2b0d7 100644 --- a/fluss-client/src/main/java/org/apache/fluss/client/utils/ClientRpcMessageUtils.java +++ b/fluss-client/src/main/java/org/apache/fluss/client/utils/ClientRpcMessageUtils.java @@ -31,6 +31,7 @@ import org.apache.fluss.client.metadata.RemoteLogManifestInfo; import org.apache.fluss.client.write.KvWriteBatch; import org.apache.fluss.client.write.ReadyWriteBatch; +import org.apache.fluss.cluster.Cluster; import org.apache.fluss.cluster.rebalance.RebalancePlanForBucket; import org.apache.fluss.cluster.rebalance.RebalanceProgress; import org.apache.fluss.cluster.rebalance.RebalanceResultForBucket; @@ -49,6 +50,7 @@ import org.apache.fluss.metadata.PhysicalTablePath; import org.apache.fluss.metadata.TableBucket; import org.apache.fluss.metadata.TableChange; +import org.apache.fluss.metadata.TableOrPartition; import org.apache.fluss.metadata.TablePath; import org.apache.fluss.rpc.messages.AcquireKvSnapshotLeaseRequest; import org.apache.fluss.rpc.messages.AcquireKvSnapshotLeaseResponse; @@ -83,6 +85,7 @@ import org.apache.fluss.rpc.messages.PbKvSnapshotLeaseForTable; import org.apache.fluss.rpc.messages.PbLakeSnapshotForBucket; import org.apache.fluss.rpc.messages.PbLookupReqForBucket; +import org.apache.fluss.rpc.messages.PbModifyBucketCount; import org.apache.fluss.rpc.messages.PbModifyColumn; import org.apache.fluss.rpc.messages.PbPartitionSpec; import org.apache.fluss.rpc.messages.PbPrefixLookupReqForBucket; @@ -143,6 +146,7 @@ public static ProduceLogRequest makeProduceLogRequest( PbProduceLogReqForBucket pbProduceLogReqForBucket = request.addBucketsReq() .setBucketId(tableBucket.getBucket()) + .setRoutingBucketCount(readyBatch.writeBatch().getBucketCount()) .setRecordsBytesView(readyBatch.writeBatch().build()); if (tableBucket.getPartitionId() != null) { pbProduceLogReqForBucket.setPartitionId(tableBucket.getPartitionId()); @@ -202,6 +206,7 @@ public static PutKvRequest makePutKvRequest( PbPutKvReqForBucket pbPutKvReqForBucket = request.addBucketsReq() .setBucketId(tableBucket.getBucket()) + .setRoutingBucketCount(readyBatch.writeBatch().getBucketCount()) .setRecordsBytesView(readyBatch.writeBatch().build()); if (tableBucket.getPartitionId() != null) { pbPutKvReqForBucket.setPartitionId(tableBucket.getPartitionId()); @@ -235,6 +240,11 @@ public static LookupRequest makeLookupRequest( if (tb.getPartitionId() != null) { pbLookupReqForBucket.setPartitionId(tb.getPartitionId()); } + // Carry the bucket count the bucketId was calculated with so the server can + // validate it; 0 means unknown (legacy) and leaves the field unset. + if (batch.getBucketCount() > 0) { + pbLookupReqForBucket.setRoutingBucketCount(batch.getBucketCount()); + } if (batch.originalPartitionName() != null) { pbLookupReqForBucket.setOriginalPartitionName( batch.originalPartitionName()); @@ -255,6 +265,11 @@ public static PrefixLookupRequest makePrefixLookupRequest( if (tb.getPartitionId() != null) { pbPrefixLookupReqForBucket.setPartitionId(tb.getPartitionId()); } + // Carry the bucket count the bucketId was calculated with so the server can + // validate it; 0 means unknown (legacy) and leaves the field unset. + if (batch.getBucketCount() > 0) { + pbPrefixLookupReqForBucket.setRoutingBucketCount(batch.getBucketCount()); + } batch.lookups().forEach(get -> pbPrefixLookupReqForBucket.addKey(get.key())); }); return request; @@ -357,7 +372,8 @@ public static ListOffsetsRequest makeListOffsetsRequest( long tableId, @Nullable Long partitionId, List bucketIdList, - OffsetSpec offsetSpec) { + OffsetSpec offsetSpec, + Cluster cluster) { ListOffsetsRequest listOffsetsRequest = new ListOffsetsRequest(); listOffsetsRequest .setFollowerServerId(-1) // -1 indicate the request from client. @@ -366,6 +382,8 @@ public static ListOffsetsRequest makeListOffsetsRequest( if (partitionId != null) { listOffsetsRequest.setPartitionId(partitionId); } + cluster.getBucketCount(TableOrPartition.of(tableId, partitionId)) + .ifPresent(listOffsetsRequest::setRoutingBucketCount); if (offsetSpec instanceof OffsetSpec.EarliestSpec) { listOffsetsRequest.setOffsetType(OffsetSpec.LIST_EARLIEST_OFFSET); @@ -420,6 +438,7 @@ public static AlterTableRequest makeAlterTableRequest( List renameColumns = new ArrayList<>(); List modifyColumns = new ArrayList<>(); List alterConfigs = new ArrayList<>(); + PbModifyBucketCount modifyBucketCount = null; for (TableChange tableChange : tableChanges) { if (tableChange instanceof TableChange.AddColumn) { addColumns.add(toPbAddColumn((TableChange.AddColumn) tableChange)); @@ -429,6 +448,16 @@ public static AlterTableRequest makeAlterTableRequest( renameColumns.add(toPbRenameColumn((TableChange.RenameColumn) tableChange)); } else if (tableChange instanceof TableChange.ModifyColumn) { modifyColumns.add(toPbModifyColumn((TableChange.ModifyColumn) tableChange)); + } else if (tableChange instanceof TableChange.ModifyBucketCount) { + if (modifyBucketCount != null) { + throw new IllegalArgumentException( + "Only one bucket count change is supported per ALTER TABLE request."); + } + modifyBucketCount = + new PbModifyBucketCount() + .setNewBucketCount( + ((TableChange.ModifyBucketCount) tableChange) + .getNewBucketCount()); } else if (tableChange instanceof TableChange.SetOption || tableChange instanceof TableChange.ResetOption) { alterConfigs.add(toPbAlterConfigs(tableChange)); @@ -442,6 +471,9 @@ public static AlterTableRequest makeAlterTableRequest( .addAllDropColumns(dropColumns) .addAllRenameColumns(renameColumns) .addAllModifyColumns(modifyColumns); + if (modifyBucketCount != null) { + request.setModifyBucketCount(modifyBucketCount); + } return request; } @@ -643,7 +675,8 @@ private static RebalancePlanForBucket toRebalancePlanForBucket( Arrays.stream(rebalancePlan.getNewReplicas()).boxed().collect(Collectors.toList())); } - public static List toPartitionInfos(ListPartitionInfosResponse response) { + public static List toPartitionInfos( + ListPartitionInfosResponse response, int defaultBucketCount) { return response.getPartitionsInfosList().stream() .map( pbPartitionInfo -> @@ -654,7 +687,12 @@ public static List toPartitionInfos(ListPartitionInfosResponse re // clusters do not include the remote data dir pbPartitionInfo.hasRemoteDataDir() ? pbPartitionInfo.getRemoteDataDir() - : null)) + : null, + // old clusters do not send the per-partition bucket count; + // resolve to the table-level count here + pbPartitionInfo.hasBucketCount() + ? pbPartitionInfo.getBucketCount() + : defaultBucketCount)) .collect(Collectors.toList()); } @@ -865,7 +903,8 @@ public static List toDatabaseSummaries(ListDatabasesResponse re return databaseSummaries; } - public static GetTableStatsRequest makeGetTableStatsRequest(List buckets) { + public static GetTableStatsRequest makeGetTableStatsRequest( + List buckets, Cluster cluster) { if (buckets.isEmpty()) { throw new IllegalArgumentException("Buckets list cannot be empty"); } @@ -884,6 +923,11 @@ public static GetTableStatsRequest makeGetTableStatsRequest(List bu if (bucket.getPartitionId() != null) { pbBucket.setPartitionId(bucket.getPartitionId()); } + cluster.getBucketCount( + TableOrPartition.of( + bucket.getTableId(), + bucket.getPartitionId())) + .ifPresent(pbBucket::setRoutingBucketCount); return pbBucket; }) .collect(Collectors.toList()); diff --git a/fluss-client/src/main/java/org/apache/fluss/client/utils/ClientUtils.java b/fluss-client/src/main/java/org/apache/fluss/client/utils/ClientUtils.java index 8629c6f95ab..6f491f63605 100644 --- a/fluss-client/src/main/java/org/apache/fluss/client/utils/ClientUtils.java +++ b/fluss-client/src/main/java/org/apache/fluss/client/utils/ClientUtils.java @@ -17,14 +17,9 @@ package org.apache.fluss.client.utils; -import org.apache.fluss.client.metadata.MetadataUpdater; -import org.apache.fluss.client.table.getter.PartitionGetter; import org.apache.fluss.config.ConfigOptions; import org.apache.fluss.exception.IllegalConfigurationException; -import org.apache.fluss.exception.PartitionNotExistException; -import org.apache.fluss.metadata.PhysicalTablePath; -import org.apache.fluss.metadata.TablePath; -import org.apache.fluss.row.InternalRow; +import org.apache.fluss.metadata.TableInfo; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -35,8 +30,6 @@ import java.util.regex.Matcher; import java.util.regex.Pattern; -import static org.apache.fluss.utils.Preconditions.checkNotNull; - /** Utils for Fluss Client. */ public final class ClientUtils { @@ -116,20 +109,20 @@ public static Integer getPort(String address) { } /** - * Return the id of the partition the row belongs to. It'll try to update the metadata if the - * partition doesn't exist. If the partition doesn't exist yet after update metadata, it'll - * throw {@link PartitionNotExistException}. + * Resolves the routing bucket count when it is unavailable in the metadata. Falling back to the + * table-level count is safe only when {@code bucketCountEpoch == 0}, which proves the table was + * never rescaled; otherwise this fails instead of silently returning a wrong answer. */ - public static Long getPartitionId( - InternalRow row, - PartitionGetter partitionGetter, - TablePath tablePath, - MetadataUpdater metadataUpdater) - throws PartitionNotExistException { - checkNotNull(partitionGetter, "partitionGetter shouldn't be null."); - String partitionName = partitionGetter.getPartition(row); - PhysicalTablePath physicalTablePath = PhysicalTablePath.of(tablePath, partitionName); - metadataUpdater.checkAndUpdatePartitionMetadata(physicalTablePath); - return metadataUpdater.getCluster().getPartitionIdOrElseThrow(physicalTablePath); + public static int fallbackBucketCountOrFail(TableInfo tableInfo, Object target) { + long epoch = tableInfo.getBucketCountEpoch(); + if (epoch > 0) { + throw new IllegalStateException( + "Routing bucket count is unavailable for " + + target + + " at bucketCountEpoch " + + epoch + + "; refusing to fall back to the table-level count."); + } + return tableInfo.getNumBuckets(); } } diff --git a/fluss-client/src/main/java/org/apache/fluss/client/utils/MetadataUtils.java b/fluss-client/src/main/java/org/apache/fluss/client/utils/MetadataUtils.java index 2990054999b..ca291f6a361 100644 --- a/fluss-client/src/main/java/org/apache/fluss/client/utils/MetadataUtils.java +++ b/fluss-client/src/main/java/org/apache/fluss/client/utils/MetadataUtils.java @@ -24,6 +24,7 @@ import org.apache.fluss.exception.StaleMetadataException; import org.apache.fluss.metadata.PhysicalTablePath; import org.apache.fluss.metadata.TableBucket; +import org.apache.fluss.metadata.TableOrPartition; import org.apache.fluss.metadata.TablePath; import org.apache.fluss.rpc.GatewayClientProxy; import org.apache.fluss.rpc.RpcClient; @@ -121,6 +122,7 @@ public static Cluster sendMetadataRequestAndRebuildCluster( Map newTablePathToTableId; Map> newBucketLocations; Map newPartitionIdByPath; + Map newBucketCountByTableOrPartition; NewTableMetadata newTableMetadata = getTableMetadataToUpdate(originCluster, response); @@ -134,10 +136,15 @@ public static Cluster sendMetadataRequestAndRebuildCluster( new HashMap<>(originCluster.getBucketLocationsByPath()); newPartitionIdByPath = new HashMap<>(originCluster.getPartitionIdByPath()); + newBucketCountByTableOrPartition = + new HashMap<>( + originCluster.getBucketCountByTableOrPartition()); newTablePathToTableId.putAll(newTableMetadata.tablePathToTableId); newBucketLocations.putAll(newTableMetadata.bucketLocations); newPartitionIdByPath.putAll(newTableMetadata.partitionIdByPath); + newBucketCountByTableOrPartition.putAll( + newTableMetadata.bucketCountByTableOrPartition); } else { // If full update, we will clear all tables info out ot the origin @@ -145,6 +152,8 @@ public static Cluster sendMetadataRequestAndRebuildCluster( newTablePathToTableId = newTableMetadata.tablePathToTableId; newBucketLocations = newTableMetadata.bucketLocations; newPartitionIdByPath = newTableMetadata.partitionIdByPath; + newBucketCountByTableOrPartition = + newTableMetadata.bucketCountByTableOrPartition; } return new Cluster( @@ -152,7 +161,8 @@ public static Cluster sendMetadataRequestAndRebuildCluster( coordinatorServer, newBucketLocations, newTablePathToTableId, - newPartitionIdByPath); + newPartitionIdByPath, + newBucketCountByTableOrPartition); }) .get(30, TimeUnit.SECONDS); // TODO currently, we don't have timeout logic in // RpcClient, it will let the get() block forever. So we @@ -162,8 +172,10 @@ public static Cluster sendMetadataRequestAndRebuildCluster( private static NewTableMetadata getTableMetadataToUpdate( Cluster cluster, MetadataResponse metadataResponse) { Map newTablePathToTableId = new HashMap<>(); + Map newTablePathByTableId = new HashMap<>(); Map> newBucketLocations = new HashMap<>(); Map newPartitionIdByPath = new HashMap<>(); + Map newBucketCountByTableOrPartition = new HashMap<>(); // iterate all table metadata List pbTableMetadataList = metadataResponse.getTableMetadatasList(); @@ -177,6 +189,7 @@ private static NewTableMetadata getTableMetadataToUpdate( protoTablePath.getDatabaseName(), protoTablePath.getTableName()); newTablePathToTableId.put(tablePath, tableId); + newTablePathByTableId.put(tableId, tablePath); // Get all buckets for the table. List pbBucketMetadataList = @@ -185,6 +198,12 @@ private static NewTableMetadata getTableMetadataToUpdate( PhysicalTablePath.of(tablePath), toBucketLocations( tablePath, tableId, null, null, pbBucketMetadataList)); + // An empty bucket list means the assignment is not generated yet; keeping the + // entry out lets callers fall back to the table-level count instead of 0. + if (!pbBucketMetadataList.isEmpty()) { + newBucketCountByTableOrPartition.put( + TableOrPartition.ofTable(tableId), pbBucketMetadataList.size()); + } }); List pbPartitionMetadataList = @@ -194,8 +213,10 @@ private static NewTableMetadata getTableMetadataToUpdate( pbPartitionMetadataList.forEach( pbPartitionMetadata -> { long tableId = pbPartitionMetadata.getTableId(); - // the table path should be initialized at begin - TablePath tablePath = cluster.getTablePathOrElseThrow(tableId); + TablePath tablePath = newTablePathByTableId.get(tableId); + if (tablePath == null) { + tablePath = cluster.getTablePathOrElseThrow(tableId); + } PhysicalTablePath physicalTablePath = PhysicalTablePath.of(tablePath, pbPartitionMetadata.getPartitionName()); newPartitionIdByPath.put( @@ -208,24 +229,43 @@ private static NewTableMetadata getTableMetadataToUpdate( pbPartitionMetadata.getPartitionId(), pbPartitionMetadata.getPartitionName(), pbPartitionMetadata.getBucketMetadatasList())); + // Old servers do not send the explicit count but still send the complete + // assignment, so normalize both protocol versions into the same Cluster map. + // An empty assignment means the layout is not available yet, not zero buckets. + int bucketCount = + pbPartitionMetadata.hasBucketCount() + && pbPartitionMetadata.getBucketCount() > 0 + ? pbPartitionMetadata.getBucketCount() + : pbPartitionMetadata.getBucketMetadatasList().size(); + if (bucketCount > 0) { + newBucketCountByTableOrPartition.put( + TableOrPartition.ofPartition(pbPartitionMetadata.getPartitionId()), + bucketCount); + } }); return new NewTableMetadata( - newTablePathToTableId, newBucketLocations, newPartitionIdByPath); + newTablePathToTableId, + newBucketLocations, + newPartitionIdByPath, + newBucketCountByTableOrPartition); } private static final class NewTableMetadata { private final Map tablePathToTableId; private final Map> bucketLocations; private final Map partitionIdByPath; + private final Map bucketCountByTableOrPartition; public NewTableMetadata( Map tablePathToTableId, Map> bucketLocations, - Map partitionIdByPath) { + Map partitionIdByPath, + Map bucketCountByTableOrPartition) { this.tablePathToTableId = tablePathToTableId; this.bucketLocations = bucketLocations; this.partitionIdByPath = partitionIdByPath; + this.bucketCountByTableOrPartition = bucketCountByTableOrPartition; } } diff --git a/fluss-client/src/main/java/org/apache/fluss/client/write/AbstractRowLogWriteBatch.java b/fluss-client/src/main/java/org/apache/fluss/client/write/AbstractRowLogWriteBatch.java index f0c3bb4f295..73ba23aebf9 100644 --- a/fluss-client/src/main/java/org/apache/fluss/client/write/AbstractRowLogWriteBatch.java +++ b/fluss-client/src/main/java/org/apache/fluss/client/write/AbstractRowLogWriteBatch.java @@ -46,6 +46,7 @@ abstract class AbstractRowLogWriteBatch extends WriteBatch { protected AbstractRowLogWriteBatch( long tableId, int bucketId, + int bucketCount, PhysicalTablePath physicalTablePath, int schemaId, WriteFormat writeFormat, @@ -57,6 +58,7 @@ protected AbstractRowLogWriteBatch( super( tableId, bucketId, + bucketCount, physicalTablePath, schemaId, writeFormat, diff --git a/fluss-client/src/main/java/org/apache/fluss/client/write/ArrowLogWriteBatch.java b/fluss-client/src/main/java/org/apache/fluss/client/write/ArrowLogWriteBatch.java index 54956bc4bac..f4e3b84f5d2 100644 --- a/fluss-client/src/main/java/org/apache/fluss/client/write/ArrowLogWriteBatch.java +++ b/fluss-client/src/main/java/org/apache/fluss/client/write/ArrowLogWriteBatch.java @@ -54,6 +54,7 @@ public class ArrowLogWriteBatch extends WriteBatch { public ArrowLogWriteBatch( long tableId, int bucketId, + int bucketCount, PhysicalTablePath physicalTablePath, int schemaId, ArrowWriter arrowWriter, @@ -64,6 +65,7 @@ public ArrowLogWriteBatch( super( tableId, bucketId, + bucketCount, physicalTablePath, schemaId, WriteFormat.ARROW_LOG, diff --git a/fluss-client/src/main/java/org/apache/fluss/client/write/CompactedLogWriteBatch.java b/fluss-client/src/main/java/org/apache/fluss/client/write/CompactedLogWriteBatch.java index 244295118e2..05f08d8bfbe 100644 --- a/fluss-client/src/main/java/org/apache/fluss/client/write/CompactedLogWriteBatch.java +++ b/fluss-client/src/main/java/org/apache/fluss/client/write/CompactedLogWriteBatch.java @@ -42,6 +42,7 @@ public final class CompactedLogWriteBatch extends AbstractRowLogWriteBatch new ArrayDeque<>()); synchronized (dq) { - RecordAppendResult appendResult = tryAppend(writeRecord, callback, dq); + RecordAppendResult appendResult = tryAppend(writeRecord, callback, bucketCount, dq); if (appendResult != null) { return appendResult; } @@ -245,7 +247,13 @@ public RecordAppendResult append( synchronized (dq) { RecordAppendResult appendResult = appendNewBatch( - writeRecord, callback, bucketId, tableInfo, dq, memorySegments); + writeRecord, + callback, + bucketId, + bucketCount, + tableInfo, + dq, + memorySegments); if (appendResult.newBatchCreated) { memorySegments = Collections.emptyList(); } @@ -281,6 +289,7 @@ public ReadyCheckResult ready(Cluster cluster) { Set readyNodes = new HashSet<>(); long nextReadyCheckDelayMs = batchTimeoutMs; Set unknownLeaderTables = new HashSet<>(); + Set invalidBucketRoutingTables = new HashSet<>(); // Go table by table so that we can get queue sizes for buckets in a table and calculate // cumulative frequency table (used in bucket assigner). @@ -290,13 +299,15 @@ public ReadyCheckResult ready(Cluster cluster) { bucketAndWriteBatches, readyNodes, unknownLeaderTables, + invalidBucketRoutingTables, cluster, nextReadyCheckDelayMs); } // TODO and the earliest time at which any non-send-able bucket will be ready; - return new ReadyCheckResult(readyNodes, nextReadyCheckDelayMs, unknownLeaderTables); + return new ReadyCheckResult( + readyNodes, nextReadyCheckDelayMs, unknownLeaderTables, invalidBucketRoutingTables); } /** @@ -440,11 +451,18 @@ boolean isHistoricalPartitionEnabled(TablePath tablePath) { return Boolean.TRUE.equals(historicalPartitionEnabledByTable.get(tablePath)); } - /** Reroutes queued batches for {@code originalPath} to the historical target. */ - void rerouteQueuedWritesToHistorical( + /** + * Reroutes queued batches for {@code originalPath} to the historical target. + * + * @return {@code true} if the batches were rerouted successfully or were already targeting the + * historical partition; {@code false} if any queued batch was routed with a bucket count + * that differs from the historical partition's bucket count + */ + boolean rerouteQueuedWritesToHistorical( PhysicalTablePath originalPath, PhysicalTablePath historicalPath, - long historicalPartitionId) { + long historicalPartitionId, + @Nullable Integer historicalBucketCount) { BucketAndWriteBatches writeTarget = checkNotNull( writeBatches.get(originalPath), @@ -454,7 +472,17 @@ void rerouteQueuedWritesToHistorical( synchronized (writeTarget) { if (writeTarget.isHistoricalWriteTarget()) { writeTarget.partitionId = historicalPartitionId; - return; + return true; + } + if (historicalBucketCount != null) { + for (Deque deque : writeTarget.batches.values()) { + for (WriteBatch batch : deque) { + if (batch.getBucketCount() > 0 + && batch.getBucketCount() != historicalBucketCount) { + return false; + } + } + } } // New appends observe the historical route and are marked as historical. Existing // queued batches are converted below before the Sender can drain again. @@ -482,6 +510,7 @@ void rerouteQueuedWritesToHistorical( } } } + return true; } /** Aborts incomplete batches whose current RPC target is {@code targetPath}. */ @@ -666,6 +695,7 @@ private long bucketReady( BucketAndWriteBatches bucketAndWriteBatches, Set readyNodes, Set unknownLeaderTables, + Set invalidBucketRoutingTables, Cluster cluster, long nextReadyCheckDelayMs) { // first check this table has partitionId. @@ -684,6 +714,13 @@ private long bucketReady( } } + Integer actualBucketCount = + bucketAndWriteBatches.isPartitionedTable + ? cluster.getBucketCount( + TableOrPartition.ofPartition( + bucketAndWriteBatches.partitionId)) + .orElse(null) + : null; Map> batches = bucketAndWriteBatches.batches; // Collect the queue sizes for available buckets to be used in adaptive bucket allocate. @@ -691,6 +728,7 @@ private long bucketReady( for (Map.Entry> entry : batches.entrySet()) { Deque deque = entry.getValue(); + final WriteBatch batch; final long waitedTimeMs; final int dequeSize; final boolean full; @@ -702,7 +740,7 @@ private long bucketReady( synchronized (deque) { // Deque are often empty in this path, esp with large bucket counts, // so we exit early if we can. - WriteBatch batch = deque.peekFirst(); + batch = deque.peekFirst(); if (batch == null) { continue; } @@ -712,6 +750,11 @@ private long bucketReady( full = dequeSize > 1 || batch.isClosed(); } + if (actualBucketCount != null && batch.getBucketCount() != actualBucketCount) { + invalidBucketRoutingTables.add(targetPath); + return nextReadyCheckDelayMs; + } + int bucketId = entry.getKey(); Optional tableIdOpt = cluster.getTableId(targetPath.getTablePath()); if (!tableIdOpt.isPresent()) { @@ -796,6 +839,7 @@ private RecordAppendResult appendNewBatch( WriteRecord writeRecord, WriteCallback callback, int bucketId, + int bucketCount, TableInfo tableInfo, Deque deque, List segments) @@ -813,7 +857,7 @@ private RecordAppendResult appendNewBatch( ? bucketAndWriteBatches : deque; synchronized (routeLock) { - RecordAppendResult appendResult = tryAppend(writeRecord, callback, deque); + RecordAppendResult appendResult = tryAppend(writeRecord, callback, bucketCount, deque); if (appendResult != null) { // Somebody else found us a batch, return the one we waited for! Hopefully this // doesn't happen often... @@ -828,6 +872,7 @@ private RecordAppendResult appendNewBatch( createWriteBatch( writeRecord, bucketId, + bucketCount, tableInfo, writeFormat, physicalTablePath, @@ -845,6 +890,7 @@ private RecordAppendResult appendNewBatch( private WriteBatch createWriteBatch( WriteRecord writeRecord, int bucketId, + int bucketCount, TableInfo tableInfo, WriteFormat writeFormat, PhysicalTablePath physicalTablePath, @@ -858,6 +904,7 @@ private WriteBatch createWriteBatch( return new KvWriteBatch( tableInfo.getTableId(), bucketId, + bucketCount, physicalTablePath, tableInfo.getSchemaId(), writeFormat.toKvFormat(), @@ -885,6 +932,7 @@ private WriteBatch createWriteBatch( return new ArrowLogWriteBatch( tableInfo.getTableId(), bucketId, + bucketCount, physicalTablePath, tableInfo.getSchemaId(), arrowWriter, @@ -897,6 +945,7 @@ private WriteBatch createWriteBatch( return new CompactedLogWriteBatch( tableInfo.getTableId(), bucketId, + bucketCount, physicalTablePath, schemaId, outputView.getPreAllocatedSize(), @@ -908,6 +957,7 @@ private WriteBatch createWriteBatch( return new IndexedLogWriteBatch( tableInfo.getTableId(), bucketId, + bucketCount, physicalTablePath, tableInfo.getSchemaId(), outputView.getPreAllocatedSize(), @@ -921,18 +971,22 @@ private WriteBatch createWriteBatch( } private RecordAppendResult tryAppend( - WriteRecord writeRecord, WriteCallback callback, Deque deque) + WriteRecord writeRecord, + WriteCallback callback, + int bucketCount, + Deque deque) throws Exception { if (closed) { throw new FlussRuntimeException("Writer closed while send in progress"); } WriteBatch last = deque.peekLast(); if (last != null) { - boolean success = last.tryAppend(writeRecord, callback); + boolean success = + last.getBucketCount() == bucketCount && last.tryAppend(writeRecord, callback); if (!success) { - // The last batch is either full/closed or belongs to a different table/write - // format/schema. Close it so the incoming record rolls over to a compatible new - // batch. + // The last batch is either full/closed or belongs to a different table, write + // format, schema, or bucket layout. Close it so the incoming record rolls over to + // a compatible new batch. // TODO For ArrowLogWriteBatch, close here is a heavy operation (including build // logic), we need to avoid do that in an lock which locked dq. However, why we not // remove build logic out of close for ArrowLogWriteBatch is that we want to release @@ -1357,14 +1411,17 @@ public static final class ReadyCheckResult { public final Set readyNodes; public final long nextReadyCheckDelayMs; public final Set unknownLeaderTables; + public final Set invalidBucketRoutingTables; public ReadyCheckResult( Set readyNodes, long nextReadyCheckDelayMs, - Set unknownLeaderTables) { + Set unknownLeaderTables, + Set invalidBucketRoutingTables) { this.readyNodes = readyNodes; this.nextReadyCheckDelayMs = nextReadyCheckDelayMs; this.unknownLeaderTables = unknownLeaderTables; + this.invalidBucketRoutingTables = invalidBucketRoutingTables; } } diff --git a/fluss-client/src/main/java/org/apache/fluss/client/write/Sender.java b/fluss-client/src/main/java/org/apache/fluss/client/write/Sender.java index f848cbede62..87ece38cab0 100644 --- a/fluss-client/src/main/java/org/apache/fluss/client/write/Sender.java +++ b/fluss-client/src/main/java/org/apache/fluss/client/write/Sender.java @@ -22,6 +22,7 @@ import org.apache.fluss.client.metrics.WriterMetricGroup; import org.apache.fluss.client.write.RecordAccumulator.ReadyCheckResult; import org.apache.fluss.cluster.Cluster; +import org.apache.fluss.exception.InvalidBucketRoutingException; import org.apache.fluss.exception.InvalidMetadataException; import org.apache.fluss.exception.LeaderNotAvailableException; import org.apache.fluss.exception.OutOfOrderSequenceException; @@ -31,6 +32,8 @@ import org.apache.fluss.exception.UnknownTableOrBucketException; import org.apache.fluss.metadata.PhysicalTablePath; import org.apache.fluss.metadata.TableBucket; +import org.apache.fluss.metadata.TableOrPartition; +import org.apache.fluss.metadata.TablePartition; import org.apache.fluss.rpc.gateway.TabletServerGateway; import org.apache.fluss.rpc.messages.PbProduceLogRespForBucket; import org.apache.fluss.rpc.messages.PbPutKvRespForBucket; @@ -56,6 +59,7 @@ import java.util.Objects; import java.util.Set; import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Consumer; import static org.apache.fluss.client.utils.ClientRpcMessageUtils.makeProduceLogRequest; import static org.apache.fluss.client.utils.ClientRpcMessageUtils.makePutKvRequest; @@ -118,6 +122,13 @@ public class Sender implements Runnable { private final WriterMetricGroup writerMetricGroup; + /** + * Called when a write batch is rejected for invalid bucket routing so the owning {@link + * WriterClient} can remove the stale {@link BucketAssigner}. The next {@code send} will refresh + * metadata and create a new assigner with the updated bucket count. + */ + private final Consumer bucketAssignerInvalidator; + public Sender( RecordAccumulator accumulator, int maxRequestTimeoutMs, @@ -126,7 +137,8 @@ public Sender( int retries, MetadataUpdater metadataUpdater, IdempotenceManager idempotenceManager, - WriterMetricGroup writerMetricGroup) { + WriterMetricGroup writerMetricGroup, + Consumer bucketAssignerInvalidator) { this.accumulator = accumulator; this.maxRequestSize = maxRequestSize; this.maxRequestTimeoutMs = maxRequestTimeoutMs; @@ -140,6 +152,7 @@ public Sender( this.idempotenceManager = idempotenceManager; this.writerMetricGroup = writerMetricGroup; + this.bucketAssignerInvalidator = bucketAssignerInvalidator; // TODO add retry logic while send failed. See FLUSS-56364375 } @@ -242,6 +255,26 @@ private void sendWriteData() throws Exception { // get the list of buckets with data ready to send. ReadyCheckResult readyCheckResult = accumulator.ready(clusterSnapshot); + if (!readyCheckResult.invalidBucketRoutingTables.isEmpty()) { + for (PhysicalTablePath physicalTablePath : + readyCheckResult.invalidBucketRoutingTables) { + accumulator.abortBatches( + physicalTablePath, + new InvalidBucketRoutingException( + "The bucket count changed before queued records for " + + physicalTablePath + + " could be sent. Retry the failed records.")); + Long tableId = + clusterSnapshot.getTableId(physicalTablePath.getTablePath()).orElse(null); + Long partitionId = clusterSnapshot.getPartitionId(physicalTablePath).orElse(null); + if (tableId != null) { + bucketAssignerInvalidator.accept(new TableBucket(tableId, partitionId, 0)); + } + } + metadataUpdater.invalidPhysicalTableBucketAndPartitionMeta( + readyCheckResult.invalidBucketRoutingTables); + } + // if there are any buckets whose leaders are not known yet, force metadata update if (!readyCheckResult.unknownLeaderTables.isEmpty()) { try { @@ -694,6 +727,26 @@ private Set handleWriteBatchException( // re-enqueues the batch. accumulator.updateThrottle(readyWriteBatch.tableBucket(), 1.0f); } + if (error.error() == Errors.INVALID_BUCKET_ROUTING) { + // The bucketId in this batch was computed with invalid routing information, and the + // server rejected it during pre-append validation, so it was provably never written. + // Reclaim its batch sequence (adjustBatchSequences=true): otherwise a permanent hole is + // left at this sequence, and the next batch that reaches the server on this bucket + // (created after the metadata refresh, carrying a valid routing count) would send the + // following sequence against a lower expected one, raising OUT_OF_ORDER_SEQUENCE and + // resetting the writer id — which discards idempotence for every bucket of this writer. + // Do not re-enqueue (the bucketId is fixed); invalidate metadata and drop the + // BucketAssigner so the next send re-routes with the updated count. + LOG.warn( + "Received {} in write request on table bucket {}. Failing batch and " + + "invalidating BucketAssigner.", + error.error(), + readyWriteBatch.tableBucket()); + failBatch(readyWriteBatch, error.exception(), true); + invalidMetadataTables.add(writeBatch.physicalTablePath()); + bucketAssignerInvalidator.accept(readyWriteBatch.tableBucket()); + return invalidMetadataTables; + } if (error.error() == Errors.DUPLICATE_SEQUENCE_EXCEPTION) { // If we have received a duplicate batch sequence error, it means that the batch // sequence has advanced beyond the sequence of the current batch. @@ -844,10 +897,44 @@ private void handleMissingPartition(PhysicalTablePath targetPath, Throwable caus @Nullable Throwable historicalTargetCause = null; try { if (metadataUpdater.checkAndUpdatePartitionMetadata(historicalPath)) { - accumulator.rerouteQueuedWritesToHistorical( - targetPath, - historicalPath, - metadataUpdater.getPartitionIdOrElseThrow(historicalPath)); + // The queued batches were routed by the original partition's bucket count; they can + // only land in the right buckets of the historical partition if its own count + // matches. Otherwise the bucket ids would be hashes against the wrong layout. + TablePartition historicalPartition = + metadataUpdater + .getCluster() + .getTablePartition(historicalPath) + .orElseThrow( + () -> + new PartitionNotExistException( + "Historical partition " + + historicalPath + + " does not exist.")); + Integer historicalBucketCount = + metadataUpdater + .getCluster() + .getBucketCount( + TableOrPartition.ofPartition( + historicalPartition.getPartitionId())) + .orElse(null); + boolean rerouted = + accumulator.rerouteQueuedWritesToHistorical( + targetPath, + historicalPath, + historicalPartition.getPartitionId(), + historicalBucketCount); + if (!rerouted) { + abortBatches( + targetPath, + newPartitionNotExistException( + "Cannot reroute writes from " + + targetPath + + " to the historical partition because their " + + "bucket ids were routed by a different bucket " + + "count than the historical partition's.", + historicalTargetCause)); + return; + } LOG.info( "Rerouted writes from partition {} to historical partition {}.", targetPath, diff --git a/fluss-client/src/main/java/org/apache/fluss/client/write/WriteBatch.java b/fluss-client/src/main/java/org/apache/fluss/client/write/WriteBatch.java index 3b6685dda05..2329dd3c297 100644 --- a/fluss-client/src/main/java/org/apache/fluss/client/write/WriteBatch.java +++ b/fluss-client/src/main/java/org/apache/fluss/client/write/WriteBatch.java @@ -52,6 +52,10 @@ public abstract class WriteBatch { private final WriteFormat writeFormat; private final int bucketId; + // The bucket count used to calculate this batch's bucketId; carried into the request so the + // TabletServer can validate it against the actual count (INVALID_BUCKET_ROUTING on mismatch). + private final int bucketCount; + protected final List callbacks = new ArrayList<>(); private final AtomicReference finalState = new AtomicReference<>(null); private final AtomicInteger attempts = new AtomicInteger(0); @@ -70,6 +74,7 @@ public abstract class WriteBatch { public WriteBatch( long tableId, int bucketId, + int bucketCount, PhysicalTablePath physicalTablePath, int schemaId, WriteFormat writeFormat, @@ -82,6 +87,7 @@ public WriteBatch( this.writeFormat = checkNotNull(writeFormat, "write format must be not null"); this.isHistoricalPartition = isHistoricalPartition; this.bucketId = bucketId; + this.bucketCount = bucketCount; this.requestFuture = new RequestFuture(); this.recordCount = 0; } @@ -205,6 +211,10 @@ public int bucketId() { return bucketId; } + public int getBucketCount() { + return bucketCount; + } + public long tableId() { return tableId; } diff --git a/fluss-client/src/main/java/org/apache/fluss/client/write/WriterClient.java b/fluss-client/src/main/java/org/apache/fluss/client/write/WriterClient.java index 2239b98d95e..cce7e10ad42 100644 --- a/fluss-client/src/main/java/org/apache/fluss/client/write/WriterClient.java +++ b/fluss-client/src/main/java/org/apache/fluss/client/write/WriterClient.java @@ -30,7 +30,9 @@ import org.apache.fluss.exception.IllegalConfigurationException; import org.apache.fluss.exception.PartitionNotExistException; import org.apache.fluss.metadata.PhysicalTablePath; +import org.apache.fluss.metadata.TableBucket; import org.apache.fluss.metadata.TableInfo; +import org.apache.fluss.metadata.TableOrPartition; import org.apache.fluss.rpc.gateway.TabletServerGateway; import org.apache.fluss.rpc.metrics.ClientMetricGroup; import org.apache.fluss.utils.AutoPartitionStrategy; @@ -96,7 +98,7 @@ public class WriterClient { private final Sender sender; private final ExecutorService ioThreadPool; private final MetadataUpdater metadataUpdater; - private final Map bucketAssignerMap = new CopyOnWriteMap<>(); + private final Map bucketAssigners = new CopyOnWriteMap<>(); private final IdempotenceManager idempotenceManager; private final WriterMetricGroup writerMetricGroup; private final DynamicPartitionCreator dynamicPartitionCreator; @@ -199,33 +201,51 @@ private void doSend(WriteRecord record, WriteCallback callback) { TableInfo tableInfo = record.getTableInfo(); PhysicalTablePath physicalTablePath = record.getPhysicalTablePath(); - // Skip the call entirely on non-partitioned tables; there is no partition to create. + // The path the record is physically written to. A retired partition's records land in + // the historical partition, whose own bucket count must drive the assignment. + PhysicalTablePath routingPath = physicalTablePath; if (tableInfo.isPartitioned()) { boolean historicalPartitionEnabled = accumulator.checkAndCacheHistoricalPartitionEnabled(tableInfo); if (historicalPartitionEnabled && mayBeExpiredHistoricalPartition( physicalTablePath, tableInfo, Instant.now())) { - resolveHistoricalWriteTarget(physicalTablePath); + routingPath = resolveHistoricalWriteTarget(physicalTablePath); } else { dynamicPartitionCreator.checkAndCreatePartitionAsync( physicalTablePath, tableInfo); } } - // maybe create bucket assigner. Cluster cluster = metadataUpdater.getCluster(); + long tableId = tableInfo.getTableId(); + Long partitionId = + tableInfo.isPartitioned() + ? cluster.getPartitionId(routingPath).orElse(null) + : null; + int bucketCount = + partitionId == null + ? tableInfo.getNumBuckets() + : cluster.getBucketCountOrFallback(tableInfo, partitionId); + final PhysicalTablePath finalRoutingPath = routingPath; BucketAssigner bucketAssigner = - bucketAssignerMap.computeIfAbsent( - physicalTablePath, - k -> createBucketAssigner(tableInfo, physicalTablePath, conf)); + bucketAssigners.computeIfAbsent( + TableOrPartition.of(tableId, partitionId), + k -> + createBucketAssigner( + tableInfo, finalRoutingPath, bucketCount, conf)); // Append the record to the accumulator. int bucketId = bucketAssigner.assignBucket(record.getBucketKey(), cluster); RecordAppendResult result = accumulator.append( - record, callback, cluster, bucketId, bucketAssigner.abortIfBatchFull()); + record, + callback, + cluster, + bucketId, + bucketCount, + bucketAssigner.abortIfBatchFull()); if (result.abortRecordForNewBatch) { int prevBucketId = bucketId; @@ -236,7 +256,8 @@ && mayBeExpiredHistoricalPartition( physicalTablePath, bucketId, prevBucketId); - result = accumulator.append(record, callback, cluster, bucketId, false); + result = + accumulator.append(record, callback, cluster, bucketId, bucketCount, false); } if (result.batchIsFull || result.newBatchCreated) { @@ -291,13 +312,14 @@ static boolean mayBeExpiredHistoricalPartition( return partitionName.compareTo(earliestRetainedPartition) < 0; } - private void resolveHistoricalWriteTarget(PhysicalTablePath originalPath) { + /** Returns the path the records of this original partition are physically written to. */ + private PhysicalTablePath resolveHistoricalWriteTarget(PhysicalTablePath originalPath) { // Keep refreshing while the target is still the original partition so its retirement can // be detected before more records are appended to the stale route. Ideally, the Client // should learn the server-authoritative partition status without synchronously refreshing // metadata on the per-record path; see https://github.com/apache/fluss/issues/4161. if (accumulator.hasHistoricalWriteTarget(originalPath)) { - return; + return PhysicalTablePath.of(originalPath.getTablePath(), HISTORICAL_PARTITION_VALUE); } PhysicalTablePath targetPath = originalPath; @@ -324,6 +346,7 @@ private void resolveHistoricalWriteTarget(PhysicalTablePath originalPath) { accumulator.routeWritesTo( originalPath, targetPath, metadataUpdater.getPartitionIdOrElseThrow(targetPath)); + return targetPath; } private void maybeAbortBatches(Throwable t) { @@ -408,7 +431,8 @@ private Sender newSender(short acks, int retries) { retries, metadataUpdater, idempotenceManager, - writerMetricGroup); + writerMetricGroup, + this::invalidateBucketAssigner); } public void close(Duration timeout) { @@ -461,22 +485,36 @@ private ExecutorService createThreadPool() { return Executors.newFixedThreadPool(1, new ExecutorThreadFactory(SENDER_THREAD_PREFIX)); } + /** + * Removes the {@link BucketAssigner} associated with the given table bucket. Called by {@link + * Sender} when a write batch is rejected for invalid bucket routing, so the next {@code send} + * creates a new assigner with the refreshed bucket count. + */ + private void invalidateBucketAssigner(TableBucket tableBucket) { + bucketAssigners.remove(TableOrPartition.ofTable(tableBucket.getTableId())); + if (tableBucket.getPartitionId() != null) { + bucketAssigners.remove(TableOrPartition.ofPartition(tableBucket.getPartitionId())); + } + } + private BucketAssigner createBucketAssigner( - TableInfo tableInfo, PhysicalTablePath physicalTablePath, Configuration conf) { - int bucketNumber = tableInfo.getNumBuckets(); + TableInfo tableInfo, + PhysicalTablePath physicalTablePath, + int bucketCount, + Configuration conf) { List bucketKeys = tableInfo.getBucketKeys(); if (!bucketKeys.isEmpty()) { BucketingFunction function = BucketingFunction.of( tableInfo.getTableConfig().getDataLakeFormat().orElse(null)); - return new HashBucketAssigner(bucketNumber, function); + return new HashBucketAssigner(bucketCount, function); } else { ConfigOptions.NoKeyAssigner noKeyAssigner = conf.get(ConfigOptions.CLIENT_WRITER_BUCKET_NO_KEY_ASSIGNER); if (noKeyAssigner == ROUND_ROBIN) { - return new RoundRobinBucketAssigner(physicalTablePath, bucketNumber); + return new RoundRobinBucketAssigner(physicalTablePath, bucketCount); } else if (noKeyAssigner == STICKY) { - return new StickyBucketAssigner(physicalTablePath, bucketNumber); + return new StickyBucketAssigner(physicalTablePath, bucketCount); } else { throw new IllegalArgumentException( "Unsupported append only row bucket assigner: " + noKeyAssigner); diff --git a/fluss-client/src/test/java/org/apache/fluss/client/admin/FlussAdmin2ITCase.java b/fluss-client/src/test/java/org/apache/fluss/client/admin/FlussAdmin2ITCase.java new file mode 100644 index 00000000000..3af429ba35c --- /dev/null +++ b/fluss-client/src/test/java/org/apache/fluss/client/admin/FlussAdmin2ITCase.java @@ -0,0 +1,180 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.fluss.client.admin; + +import org.apache.fluss.config.AutoPartitionTimeUnit; +import org.apache.fluss.config.ConfigOptions; +import org.apache.fluss.exception.TableNotPartitionedException; +import org.apache.fluss.metadata.PartitionInfo; +import org.apache.fluss.metadata.Schema; +import org.apache.fluss.metadata.TableDescriptor; +import org.apache.fluss.metadata.TablePath; +import org.apache.fluss.rpc.messages.ListPartitionInfosRequest; +import org.apache.fluss.rpc.messages.ListPartitionInfosResponse; +import org.apache.fluss.rpc.messages.PbTablePath; +import org.apache.fluss.types.DataTypes; + +import org.junit.jupiter.api.Test; + +import java.util.List; +import java.util.Map; + +import static org.apache.fluss.metadata.DataLakeFormat.PAIMON; +import static org.apache.fluss.record.TestData.DATA1_PARTITIONED_TABLE_DESCRIPTOR; +import static org.apache.fluss.record.TestData.DATA1_TABLE_DESCRIPTOR_PK; +import static org.apache.fluss.utils.PartitionUtils.HISTORICAL_PARTITION_VALUE; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** + * Additional integration tests for {@link FlussAdmin}. + * + *

This class contains additional tests because {@link FlussAdminITCase} is close to Checkstyle's + * 3000-line limit per file. Add new FlussAdmin integration tests here. + */ +class FlussAdmin2ITCase extends ClientToServerITCaseBase { + + @Test + void testListPartitionInfos() throws Exception { + String dbName = "test_db"; + TablePath nonPartitionedTablePath = TablePath.of(dbName, "test_non_partitioned_table"); + createTable(nonPartitionedTablePath, DATA1_TABLE_DESCRIPTOR_PK, false); + assertThatThrownBy(() -> admin.listPartitionInfos(nonPartitionedTablePath).get()) + .cause() + .isInstanceOf(TableNotPartitionedException.class) + .hasMessage("Table '%s' is not a partitioned table.", nonPartitionedTablePath); + + TableDescriptor partitionedTable = + TableDescriptor.builder() + .schema( + Schema.newBuilder() + .column("id", DataTypes.STRING()) + .column("name", DataTypes.STRING()) + .column("pt", DataTypes.STRING()) + .primaryKey("id", "pt") + .build()) + .distributedBy(3, "id") + .partitionedBy("pt") + .property(ConfigOptions.TABLE_AUTO_PARTITION_ENABLED, true) + .property(ConfigOptions.TABLE_AUTO_PARTITION_KEY, "pt") + .property( + ConfigOptions.TABLE_AUTO_PARTITION_TIME_UNIT, + AutoPartitionTimeUnit.YEAR) + .property(ConfigOptions.TABLE_DATALAKE_ENABLED, true) + .property(ConfigOptions.TABLE_DATALAKE_FORMAT, PAIMON) + .property(ConfigOptions.TABLE_DATALAKE_HISTORICAL_PARTITION_ENABLED, true) + .build(); + TablePath partitionedTablePath = TablePath.of(dbName, "test_partitioned_table"); + admin.createTable(partitionedTablePath, partitionedTable, false).get(); + Map partitionIdByNames = + FLUSS_CLUSTER_EXTENSION.waitUntilPartitionAllReady( + partitionedTablePath, + ConfigOptions.TABLE_AUTO_PARTITION_NUM_PRECREATE.defaultValue() + 1); + assertThat(partitionIdByNames).containsKey(HISTORICAL_PARTITION_VALUE); + + List partitionInfos = admin.listPartitionInfos(partitionedTablePath).get(); + assertThat(partitionInfos) + .hasSize(partitionIdByNames.size() - 1) + .extracting(PartitionInfo::getPartitionName) + .doesNotContain(HISTORICAL_PARTITION_VALUE); + + List allPartitionInfos = + admin.listPartitionInfos(partitionedTablePath, true).get(); + assertThat(allPartitionInfos) + .hasSize(partitionIdByNames.size()) + .extracting(PartitionInfo::getPartitionName) + .contains(HISTORICAL_PARTITION_VALUE); + PartitionInfo historicalPartitionInfo = + allPartitionInfos.stream() + .filter( + partitionInfo -> + HISTORICAL_PARTITION_VALUE.equals( + partitionInfo.getPartitionName())) + .findFirst() + .get(); + assertThat(historicalPartitionInfo.getPartitionId()) + .isEqualTo(partitionIdByNames.get(HISTORICAL_PARTITION_VALUE)); + assertThat(historicalPartitionInfo.getBucketCount()).isEqualTo(3); + + FlussAdmin flussAdmin = (FlussAdmin) admin; + ListPartitionInfosRequest legacyRequest = requestFor(partitionedTablePath, false); + ListPartitionInfosResponse legacyResponse = + flussAdmin.getAdminReadOnlyGateway().listPartitionInfos(legacyRequest).get(); + assertThat(legacyResponse.hasSystemPartitionsIncluded()).isFalse(); + List legacyCompatibleInfos = + flussAdmin + .handleListPartitionInfosResponse( + partitionedTablePath, true, legacyResponse) + .get(); + assertThat(legacyCompatibleInfos) + .hasSize(partitionIdByNames.size()) + .extracting(PartitionInfo::getPartitionName) + .contains(HISTORICAL_PARTITION_VALUE); + legacyResponse.setSystemPartitionsIncluded(false); + assertThat( + flussAdmin + .handleListPartitionInfosResponse( + partitionedTablePath, true, legacyResponse) + .get()) + .extracting(PartitionInfo::getPartitionName) + .contains(HISTORICAL_PARTITION_VALUE); + + TablePath noSystemPartitionTablePath = + TablePath.of(dbName, "test_partitioned_table_without_system_partition"); + admin.createTable(noSystemPartitionTablePath, DATA1_PARTITIONED_TABLE_DESCRIPTOR, false) + .get(); + ListPartitionInfosResponse responseWithoutSystemPartitions = + flussAdmin + .getAdminReadOnlyGateway() + .listPartitionInfos(requestFor(noSystemPartitionTablePath, true)) + .get(); + assertThat(responseWithoutSystemPartitions.hasSystemPartitionsIncluded()).isTrue(); + assertThat(responseWithoutSystemPartitions.isSystemPartitionsIncluded()).isTrue(); + + ListPartitionInfosResponse legacyResponseWithoutSystemPartitions = + flussAdmin + .getAdminReadOnlyGateway() + .listPartitionInfos(requestFor(noSystemPartitionTablePath, false)) + .get(); + assertThat( + flussAdmin + .handleListPartitionInfosResponse( + noSystemPartitionTablePath, + true, + legacyResponseWithoutSystemPartitions) + .get()) + .extracting(PartitionInfo::getPartitionName) + .doesNotContain(HISTORICAL_PARTITION_VALUE); + } + + private static ListPartitionInfosRequest requestFor( + TablePath tablePath, boolean includeSystemPartitions) { + ListPartitionInfosRequest request = + new ListPartitionInfosRequest() + .setTablePath( + new PbTablePath() + .setDatabaseName(tablePath.getDatabaseName()) + .setTableName(tablePath.getTableName())); + if (includeSystemPartitions) { + request.setIncludeSystemPartitions(true); + } + return request; + } +} diff --git a/fluss-client/src/test/java/org/apache/fluss/client/admin/FlussAdminITCase.java b/fluss-client/src/test/java/org/apache/fluss/client/admin/FlussAdminITCase.java index 70ad760f0e6..7a0a9c1e247 100644 --- a/fluss-client/src/test/java/org/apache/fluss/client/admin/FlussAdminITCase.java +++ b/fluss-client/src/test/java/org/apache/fluss/client/admin/FlussAdminITCase.java @@ -25,6 +25,7 @@ import org.apache.fluss.client.table.Table; import org.apache.fluss.client.table.writer.AppendWriter; import org.apache.fluss.client.table.writer.UpsertWriter; +import org.apache.fluss.cluster.Cluster; import org.apache.fluss.cluster.ServerNode; import org.apache.fluss.cluster.rebalance.ServerTag; import org.apache.fluss.config.AutoPartitionTimeUnit; @@ -53,7 +54,6 @@ import org.apache.fluss.exception.ServerTagAlreadyExistException; import org.apache.fluss.exception.ServerTagNotExistException; import org.apache.fluss.exception.TableNotExistException; -import org.apache.fluss.exception.TableNotPartitionedException; import org.apache.fluss.exception.TooManyBucketsException; import org.apache.fluss.exception.TooManyPartitionsException; import org.apache.fluss.fs.FsPath; @@ -127,11 +127,15 @@ import static org.apache.fluss.testutils.DataTestUtils.row; import static org.apache.fluss.testutils.InternalRowAssert.assertThatRow; import static org.apache.fluss.testutils.common.CommonTestUtils.waitUntil; -import static org.apache.fluss.utils.PartitionUtils.HISTORICAL_PARTITION_VALUE; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; -/** Test for {@link FlussAdmin}. */ +/** + * Integration tests for {@link FlussAdmin}. + * + *

Tests are split between this class and {@link FlussAdmin2ITCase} to stay within Checkstyle's + * 3000-line limit per file. Add new tests to {@link FlussAdmin2ITCase}. + */ class FlussAdminITCase extends ClientToServerITCaseBase { protected static final TablePath DEFAULT_TABLE_PATH = TablePath.of("test_db", "person"); @@ -166,11 +170,9 @@ void testMultiClient() throws Exception { Admin admin1 = conn.getAdmin(); Admin admin2 = conn.getAdmin(); assertThat(admin1).isEqualTo(admin2); - TableInfo t1 = admin1.getTableInfo(DEFAULT_TABLE_PATH).get(); TableInfo t2 = admin2.getTableInfo(DEFAULT_TABLE_PATH).get(); assertThat(t1).isEqualTo(t2); - admin1.close(); admin2.close(); } @@ -1161,56 +1163,6 @@ void testListDatabasesAndTables() throws Exception { .isInstanceOf(DatabaseNotExistException.class); } - @Test - void testListPartitionInfos() throws Exception { - String dbName = DEFAULT_TABLE_PATH.getDatabaseName(); - TablePath nonPartitionedTablePath = TablePath.of(dbName, "test_non_partitioned_table"); - admin.createTable(nonPartitionedTablePath, DEFAULT_TABLE_DESCRIPTOR, true).get(); - assertThatThrownBy(() -> admin.listPartitionInfos(nonPartitionedTablePath).get()) - .cause() - .isInstanceOf(TableNotPartitionedException.class) - .hasMessage("Table '%s' is not a partitioned table.", nonPartitionedTablePath); - - TableDescriptor partitionedTable = - TableDescriptor.builder() - .schema( - Schema.newBuilder() - .column("id", DataTypes.STRING()) - .column("name", DataTypes.STRING()) - .column("pt", DataTypes.STRING()) - .primaryKey("id", "pt") - .build()) - .comment("test table") - .distributedBy(3, "id") - .partitionedBy("pt") - .property(ConfigOptions.TABLE_AUTO_PARTITION_ENABLED, true) - .property(ConfigOptions.TABLE_AUTO_PARTITION_KEY, "pt") - .property( - ConfigOptions.TABLE_AUTO_PARTITION_TIME_UNIT, - AutoPartitionTimeUnit.YEAR) - .property(ConfigOptions.TABLE_DATALAKE_ENABLED, true) - .property(ConfigOptions.TABLE_DATALAKE_FORMAT, PAIMON) - .property(ConfigOptions.TABLE_DATALAKE_HISTORICAL_PARTITION_ENABLED, true) - .build(); - TablePath partitionedTablePath = TablePath.of(dbName, "test_partitioned_table"); - admin.createTable(partitionedTablePath, partitionedTable, true).get(); - Map partitionIdByNames = - FLUSS_CLUSTER_EXTENSION.waitUntilPartitionAllReady( - partitionedTablePath, - ConfigOptions.TABLE_AUTO_PARTITION_NUM_PRECREATE.defaultValue() + 1); - assertThat(partitionIdByNames).containsKey(HISTORICAL_PARTITION_VALUE); - - List partitionInfos = admin.listPartitionInfos(partitionedTablePath).get(); - assertThat(partitionInfos) - .hasSize(partitionIdByNames.size() - 1) - .extracting(PartitionInfo::getPartitionName) - .doesNotContain(HISTORICAL_PARTITION_VALUE); - for (PartitionInfo partitionInfo : partitionInfos) { - assertThat(partitionIdByNames.get(partitionInfo.getPartitionName())) - .isEqualTo(partitionInfo.getPartitionId()); - } - } - @Test void testListPartitionInfosAfterTabletServerRestart() throws Exception { String dbName = DEFAULT_TABLE_PATH.getDatabaseName(); @@ -1247,13 +1199,11 @@ void testKvSnapshotLeaseAfterCoordinatorServerRestart() throws Exception { // Restart the coordinator server so that the lease uses a stale cached address. restartCoordinatorServer(zkClient); - lease.acquireSnapshots(snapshots).get(); assertThat(zkClient.getKvSnapshotLeaseMetadata(lease.leaseId())).isPresent(); // Verify that release also refreshes metadata and retries against the new coordinator. restartCoordinatorServer(zkClient); - lease.releaseSnapshots(Collections.singleton(tableBucket)).get(); assertThat(zkClient.getKvSnapshotLeaseMetadata(lease.leaseId())).isNotPresent(); @@ -2925,7 +2875,11 @@ public CompletableFuture listOffsets( ListOffsetsRequest request = makeListOffsetsRequest( - 1L, null, Arrays.asList(0, 1, 2), new OffsetSpec.LatestSpec()); + 1L, + null, + Arrays.asList(0, 1, 2), + new OffsetSpec.LatestSpec(), + Cluster.empty()); Map leaderToRequestMap = new HashMap<>(); leaderToRequestMap.put(1, request); diff --git a/fluss-client/src/test/java/org/apache/fluss/client/lookup/LookupSenderTest.java b/fluss-client/src/test/java/org/apache/fluss/client/lookup/LookupSenderTest.java index 03b62d52e65..d6564769f89 100644 --- a/fluss-client/src/test/java/org/apache/fluss/client/lookup/LookupSenderTest.java +++ b/fluss-client/src/test/java/org/apache/fluss/client/lookup/LookupSenderTest.java @@ -22,6 +22,7 @@ import org.apache.fluss.config.ConfigOptions; import org.apache.fluss.config.Configuration; import org.apache.fluss.exception.HistoricalPartitionThrottledException; +import org.apache.fluss.exception.InvalidBucketRoutingException; import org.apache.fluss.exception.InvalidTableException; import org.apache.fluss.exception.NotLeaderOrFollowerException; import org.apache.fluss.exception.TableNotExistException; @@ -462,6 +463,27 @@ void testNonRetriableExceptionDoesNotRetry() { assertThat(query.retries()).isEqualTo(0); // no retries } + @Test + void testInvalidBucketRoutingFailsWithoutRetryAndInvalidatesMetadata() { + AtomicInteger attemptCount = new AtomicInteger(); + gateway.setLookupHandler( + request -> { + attemptCount.incrementAndGet(); + return createFailedResponse( + request, new InvalidBucketRoutingException("invalid bucket routing")); + }); + + LookupQuery query = new LookupQuery(DATA1_TABLE_PATH_PK, TABLE_BUCKET, new byte[0]); + lookupQueue.appendLookup(query); + + assertThatThrownBy(() -> query.future().get(5, TimeUnit.SECONDS)) + .isInstanceOf(ExecutionException.class) + .hasRootCauseInstanceOf(InvalidBucketRoutingException.class); + assertThat(attemptCount).hasValue(1); + assertThat(query.retries()).isZero(); + assertThat(metadataUpdater.getBucketLocation(TABLE_BUCKET)).isEmpty(); + } + @Test void testMaxRetriesEnforced() { // setup: always fail with retriable exception @@ -599,6 +621,68 @@ void testMultipleConcurrentLookupsWithRetries() throws Exception { .isGreaterThanOrEqualTo(2); // at least 1 failure + 1 success for the batch } + @Test + void testLookupRequestCarriesPinnedRoutingBucketCount() throws Exception { + // TOCTOU: the bucket count pinned at T1 (lookup time) must be carried to T2 (send time) + // as the request's routing_bucket_count, not re-read from cluster metadata at T2. + List receivedRequests = Collections.synchronizedList(new ArrayList<>()); + gateway.setLookupHandler( + request -> { + receivedRequests.add(request); + return createSuccessResponse(request, "value".getBytes()); + }); + + // T1: create query with bucketCount=4 (the partition's actual count at lookup time) + LookupQuery query = + new LookupQuery(DATA1_TABLE_PATH_PK, TABLE_BUCKET, bytes("key"), false, null, 4); + // The pinned value is visible on the query object + assertThat(query.bucketCount()).isEqualTo(4); + + lookupSender.sendLookups(1, LookupType.LOOKUP, Collections.singletonList(query)); + + // T2: the request must carry the T1-pinned count as routing_bucket_count + assertThat(receivedRequests).hasSize(1); + LookupRequest request = receivedRequests.get(0); + assertThat(request.getBucketsReqAt(0).hasRoutingBucketCount()).isTrue(); + assertThat(request.getBucketsReqAt(0).getRoutingBucketCount()).isEqualTo(4); + + // A legacy query (bucketCount=0) must not set routing_bucket_count at all, letting + // the server's epoch check decide. + receivedRequests.clear(); + LookupQuery legacyQuery = new LookupQuery(DATA1_TABLE_PATH_PK, TABLE_BUCKET, bytes("key")); + assertThat(legacyQuery.bucketCount()).isEqualTo(0); + + lookupSender.sendLookups(1, LookupType.LOOKUP, Collections.singletonList(legacyQuery)); + + assertThat(receivedRequests).hasSize(1); + assertThat(receivedRequests.get(0).getBucketsReqAt(0).hasRoutingBucketCount()).isFalse(); + } + + @Test + void testPrefixLookupRequestCarriesPinnedRoutingBucketCount() throws Exception { + // TOCTOU: same anchoring for prefix lookup path. + List receivedRequests = + Collections.synchronizedList(new ArrayList<>()); + gateway.setPrefixLookupHandler( + request -> { + receivedRequests.add(request); + return createSuccessPrefixLookupResponse(request); + }); + + // T1: create prefix query with bucketCount=4 + PrefixLookupQuery query = + new PrefixLookupQuery(DATA1_TABLE_PATH_PK, TABLE_BUCKET, bytes("prefix"), 4); + assertThat(query.bucketCount()).isEqualTo(4); + + lookupSender.sendLookups(1, LookupType.PREFIX_LOOKUP, Collections.singletonList(query)); + + // T2: the request must carry the T1-pinned count as routing_bucket_count + assertThat(receivedRequests).hasSize(1); + PrefixLookupRequest request = receivedRequests.get(0); + assertThat(request.getBucketsReqAt(0).hasRoutingBucketCount()).isTrue(); + assertThat(request.getBucketsReqAt(0).getRoutingBucketCount()).isEqualTo(4); + } + // Helper methods private CompletableFuture createPartitionNameEchoResponse( diff --git a/fluss-client/src/test/java/org/apache/fluss/client/lookup/PrimaryKeyLookuperTest.java b/fluss-client/src/test/java/org/apache/fluss/client/lookup/PrimaryKeyLookuperTest.java index 63c7af9d42e..25206eb3a3f 100644 --- a/fluss-client/src/test/java/org/apache/fluss/client/lookup/PrimaryKeyLookuperTest.java +++ b/fluss-client/src/test/java/org/apache/fluss/client/lookup/PrimaryKeyLookuperTest.java @@ -23,6 +23,7 @@ import org.apache.fluss.cluster.ServerNode; import org.apache.fluss.config.ConfigOptions; import org.apache.fluss.config.Configuration; +import org.apache.fluss.exception.InvalidBucketRoutingException; import org.apache.fluss.exception.PartitionNotExistException; import org.apache.fluss.metadata.DataLakeFormat; import org.apache.fluss.metadata.PhysicalTablePath; @@ -59,6 +60,7 @@ import static org.apache.fluss.client.metadata.TestingMetadataUpdater.NODE1; import static org.apache.fluss.utils.PartitionUtils.HISTORICAL_PARTITION_VALUE; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; /** Tests for {@link PrimaryKeyLookuper}. */ class PrimaryKeyLookuperTest { @@ -138,7 +140,51 @@ void testFallbackUsesOriginalPartitionWhenLookupRowIsReused() throws Exception { } } + @Test + void testRescaledPartitionMissingBucketCountFailsWithIllegalStateException() throws Exception { + // A rescaled table cannot safely route a partition whose actual bucket count is missing. + // Preserve the async API while exposing a diagnostic IllegalStateException as the cause. + TableInfo rescaledTableInfo = createTableInfo(1L); + ControllableLookupGateway gateway = new ControllableLookupGateway(); + TestingMetadataUpdater metadataUpdater = + TestingMetadataUpdater.builder( + Collections.singletonMap(TABLE_PATH, rescaledTableInfo)) + .withTabletServerGateway(NODE1.id(), gateway) + .build(); + // The cluster carries no per-partition bucket count for the active partition. + metadataUpdater.updateCluster(createCluster()); + + LookupClient lookupClient = new LookupClient(new Configuration(), metadataUpdater); + try { + PrimaryKeyLookuper lookuper = + new PrimaryKeyLookuper( + rescaledTableInfo, + new TestingSchemaGetter( + rescaledTableInfo.getSchemaId(), rescaledTableInfo.getSchema()), + metadataUpdater, + lookupClient, + false); + ProjectedRow lookupKey = + ProjectedRow.from(new int[] {0, 1}) + .replaceRow(GenericRow.of(1, BinaryString.fromString(PARTITION_A))); + + CompletableFuture resultFuture = lookuper.lookup(lookupKey); + + assertThat(resultFuture).isCompletedExceptionally(); + assertThatThrownBy(() -> resultFuture.get(5, TimeUnit.SECONDS)) + .hasCauseInstanceOf(InvalidBucketRoutingException.class) + .hasMessageContaining("Routing bucket count is unavailable") + .hasMessageContaining("bucketCountEpoch 1"); + } finally { + lookupClient.close(Duration.ofSeconds(5)); + } + } + private static TableInfo createTableInfo() { + return createTableInfo(0L); + } + + private static TableInfo createTableInfo(long bucketCountEpoch) { Schema schema = Schema.newBuilder() .column("id", DataTypes.INT()) @@ -156,7 +202,8 @@ private static TableInfo createTableInfo() { .property(ConfigOptions.TABLE_DATALAKE_FORMAT, DataLakeFormat.PAIMON) .property(ConfigOptions.TABLE_DATALAKE_HISTORICAL_PARTITION_ENABLED, true) .build(); - return TableInfo.of(TABLE_PATH, TABLE_ID, 1, tableDescriptor, null, 0L, 0L); + return TableInfo.of( + TABLE_PATH, TABLE_ID, 1, tableDescriptor, null, 0L, 0L, bucketCountEpoch); } private static Cluster createCluster() { @@ -185,7 +232,8 @@ private static Cluster createCluster() { COORDINATOR, bucketLocationsByPath, Collections.singletonMap(TABLE_PATH, TABLE_ID), - partitionIdsByPath); + partitionIdsByPath, + Collections.emptyMap()); } private static BucketLocation createBucketLocation( diff --git a/fluss-client/src/test/java/org/apache/fluss/client/metadata/MetadataUpdaterITCase.java b/fluss-client/src/test/java/org/apache/fluss/client/metadata/MetadataUpdaterITCase.java index 6d099a24e65..caecbe5452d 100644 --- a/fluss-client/src/test/java/org/apache/fluss/client/metadata/MetadataUpdaterITCase.java +++ b/fluss-client/src/test/java/org/apache/fluss/client/metadata/MetadataUpdaterITCase.java @@ -108,6 +108,7 @@ void testUpdateWithEmptyMetadataResponse() throws Exception { null, Collections.emptyMap(), Collections.emptyMap(), + Collections.emptyMap(), Collections.emptyMap()); metadataUpdater = new MetadataUpdater(rpcClient, new Configuration(), newCluster); diff --git a/fluss-client/src/test/java/org/apache/fluss/client/metadata/MetadataUpdaterTest.java b/fluss-client/src/test/java/org/apache/fluss/client/metadata/MetadataUpdaterTest.java index 1cda13b9309..d08d0b6a120 100644 --- a/fluss-client/src/test/java/org/apache/fluss/client/metadata/MetadataUpdaterTest.java +++ b/fluss-client/src/test/java/org/apache/fluss/client/metadata/MetadataUpdaterTest.java @@ -17,15 +17,20 @@ package org.apache.fluss.client.metadata; +import org.apache.fluss.client.utils.MetadataUtils; import org.apache.fluss.cluster.Cluster; import org.apache.fluss.cluster.ServerNode; import org.apache.fluss.cluster.ServerType; import org.apache.fluss.config.Configuration; import org.apache.fluss.exception.StaleMetadataException; +import org.apache.fluss.metadata.TableOrPartition; +import org.apache.fluss.metadata.TablePath; import org.apache.fluss.rpc.RpcClient; import org.apache.fluss.rpc.gateway.AdminReadOnlyGateway; import org.apache.fluss.rpc.messages.MetadataRequest; import org.apache.fluss.rpc.messages.MetadataResponse; +import org.apache.fluss.rpc.messages.PbPartitionMetadata; +import org.apache.fluss.rpc.messages.PbTableMetadata; import org.apache.fluss.rpc.metrics.TestingClientMetricGroup; import org.apache.fluss.server.coordinator.TestCoordinatorGateway; @@ -70,6 +75,83 @@ void testInitializeClusterWithRetries() throws Exception { .hasMessageContaining("The metadata is stale."); } + @Test + void testMetadataBucketCountCompatibility() throws Exception { + long tableId = 1L; + long legacyPartitionId = 2L; + long explicitPartitionId = 3L; + long unassignedPartitionId = 4L; + TablePath tablePath = TablePath.of("db", "table"); + + MetadataResponse response = new MetadataResponse(); + response.addTabletServer() + .setNodeId(TS_NODE.id()) + .setHost(TS_NODE.host()) + .setPort(TS_NODE.port()); + + PbTableMetadata tableMetadata = response.addTableMetadata().setTableId(tableId); + tableMetadata + .setTablePath() + .setDatabaseName(tablePath.getDatabaseName()) + .setTableName(tablePath.getTableName()); + for (int bucketId = 0; bucketId < 3; bucketId++) { + tableMetadata.addBucketMetadata().setBucketId(bucketId); + } + + PbPartitionMetadata legacyPartition = + response.addPartitionMetadata() + .setTableId(tableId) + .setPartitionId(legacyPartitionId) + .setPartitionName("legacy"); + legacyPartition.addBucketMetadata().setBucketId(0); + legacyPartition.addBucketMetadata().setBucketId(1); + assertThat(legacyPartition.hasBucketCount()).isFalse(); + + PbPartitionMetadata explicitPartition = + response.addPartitionMetadata() + .setTableId(tableId) + .setPartitionId(explicitPartitionId) + .setPartitionName("explicit") + .setBucketCount(4); + explicitPartition.addBucketMetadata().setBucketId(0); + explicitPartition.addBucketMetadata().setBucketId(1); + + response.addPartitionMetadata() + .setTableId(tableId) + .setPartitionId(unassignedPartitionId) + .setPartitionName("unassigned"); + + Cluster originCluster = + new Cluster( + Collections.singletonMap(TS_NODE.id(), TS_NODE), + CS_NODE, + Collections.emptyMap(), + Collections.emptyMap(), + Collections.emptyMap(), + Collections.emptyMap()); + AdminReadOnlyGateway gateway = + new TestCoordinatorGateway() { + @Override + public CompletableFuture metadata(MetadataRequest request) { + return CompletableFuture.completedFuture(response); + } + }; + + Cluster updatedCluster = + MetadataUtils.sendMetadataRequestAndRebuildCluster( + gateway, true, originCluster, null, null, null); + + assertThat(updatedCluster.getBucketCount(TableOrPartition.ofTable(tableId))).hasValue(3); + assertThat(updatedCluster.getBucketCount(TableOrPartition.ofPartition(legacyPartitionId))) + .hasValue(2); + assertThat(updatedCluster.getBucketCount(TableOrPartition.ofPartition(explicitPartitionId))) + .hasValue(4); + assertThat( + updatedCluster.getBucketCount( + TableOrPartition.ofPartition(unassignedPartitionId))) + .isEmpty(); + } + private static final class TestingAdminReadOnlyGateway extends TestCoordinatorGateway { private final int maxRetryCount; diff --git a/fluss-client/src/test/java/org/apache/fluss/client/metadata/TestingMetadataUpdater.java b/fluss-client/src/test/java/org/apache/fluss/client/metadata/TestingMetadataUpdater.java index 303e453180a..ccacf47b7be 100644 --- a/fluss-client/src/test/java/org/apache/fluss/client/metadata/TestingMetadataUpdater.java +++ b/fluss-client/src/test/java/org/apache/fluss/client/metadata/TestingMetadataUpdater.java @@ -223,6 +223,7 @@ private void initializeCluster( coordinatorServer, tablePathToBucketLocations, tableIdByPath, + Collections.emptyMap(), Collections.emptyMap()); } } diff --git a/fluss-client/src/test/java/org/apache/fluss/client/table/PartitionBucketCountActualRescaleITCase.java b/fluss-client/src/test/java/org/apache/fluss/client/table/PartitionBucketCountActualRescaleITCase.java new file mode 100644 index 00000000000..06574806924 --- /dev/null +++ b/fluss-client/src/test/java/org/apache/fluss/client/table/PartitionBucketCountActualRescaleITCase.java @@ -0,0 +1,566 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.fluss.client.table; + +import org.apache.fluss.client.ConnectionFactory; +import org.apache.fluss.client.admin.ClientToServerITCaseBase; +import org.apache.fluss.client.lookup.Lookuper; +import org.apache.fluss.client.table.scanner.ScanRecord; +import org.apache.fluss.client.table.scanner.batch.BatchScanner; +import org.apache.fluss.client.table.scanner.log.LogScanner; +import org.apache.fluss.client.table.scanner.log.ScanRecords; +import org.apache.fluss.client.table.writer.AppendWriter; +import org.apache.fluss.client.table.writer.UpsertResult; +import org.apache.fluss.client.table.writer.UpsertWriter; +import org.apache.fluss.config.ConfigOptions; +import org.apache.fluss.exception.InvalidAlterTableException; +import org.apache.fluss.exception.InvalidBucketRoutingException; +import org.apache.fluss.metadata.PartitionInfo; +import org.apache.fluss.metadata.Schema; +import org.apache.fluss.metadata.TableBucket; +import org.apache.fluss.metadata.TableChange; +import org.apache.fluss.metadata.TableDescriptor; +import org.apache.fluss.metadata.TableInfo; +import org.apache.fluss.metadata.TablePath; +import org.apache.fluss.row.InternalRow; +import org.apache.fluss.types.DataTypes; +import org.apache.fluss.types.RowType; +import org.apache.fluss.utils.CloseableIterator; + +import org.junit.jupiter.api.Test; + +import java.time.Duration; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutionException; + +import static org.apache.fluss.testutils.DataTestUtils.row; +import static org.apache.fluss.testutils.InternalRowAssert.assertThatRow; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** + * End-to-end IT case verifying that reads and writes route by the per-partition bucket count after + * an ALTER TABLE ... SET ('bucket.num' = N): old partitions keep their original bucket count while + * partitions created after the ALTER use the new count, and data written to each partition is read + * back correctly through its own bucket range. + */ +class PartitionBucketCountRescaleITCase extends ClientToServerITCaseBase { + + private static final int OLD_BUCKET_NUM = 2; + private static final int NEW_BUCKET_NUM = 4; + private static final int RECORDS_PER_PARTITION = 12; + private static final List OLD_NEW_PARTITIONS = Arrays.asList("old", "new"); + + @Test + void testLogTableReadWriteAcrossRescale() throws Exception { + // Write records to partitions with different bucket counts (old partition keeps + // OLD_BUCKET_NUM, new partition uses NEW_BUCKET_NUM). Routing by the wrong (table-level) + // count would miss rows when reading back each partition's own bucket range. + + TablePath tablePath = TablePath.of("test_db_1", "test_rescale_log_table"); + Schema schema = logSchema(); + createPartitionedTable(tablePath, schema); + List partitionInfos = setupOldNewPartitions(tablePath); + Map idByName = partitionIdByName(partitionInfos); + + // append RECORDS_PER_PARTITION rows to each partition + Table table = conn.getTable(tablePath); + AppendWriter appendWriter = table.newAppend().createWriter(); + Map> expectedByPartitionId = new HashMap<>(); + for (String partitionName : OLD_NEW_PARTITIONS) { + long partitionId = idByName.get(partitionName); + for (int j = 0; j < RECORDS_PER_PARTITION; j++) { + InternalRow r = row(j, "v" + j, partitionName); + appendWriter.append(r); + expectedByPartitionId.computeIfAbsent(partitionId, k -> new ArrayList<>()).add(r); + } + } + appendWriter.flush(); + + // read back by subscribing EACH partition's own bucket range [0, bucketCount) + Map> actualByPartitionId = + scanAllBucketsPerPartition(table, partitionInfos); + + assertRowsPerPartition(schema.getRowType(), actualByPartitionId, expectedByPartitionId); + } + + @Test + void testPkTableReadPathsAcrossRescale() throws Exception { + TablePath tablePath = TablePath.of("test_db_1", "test_rescale_pk_read_paths"); + Schema schema = pkSchema("a", "c"); + createPartitionedTable(tablePath, schema); + List partitionInfos = setupOldNewPartitions(tablePath); + Map idByName = partitionIdByName(partitionInfos); + Map bucketCountByName = bucketCountByName(partitionInfos); + + Table table = conn.getTable(tablePath); + long tableId = table.getTableInfo().getTableId(); + upsertRowsToOldAndNew(table); + + // 1. Lookup: write routing N must equal lookup routing N per partition, otherwise the + // lookup would hit the wrong bucket and miss the row. + Lookuper lookuper = table.newLookup().createLookuper(); + for (String partitionName : OLD_NEW_PARTITIONS) { + for (int j = 0; j < RECORDS_PER_PARTITION; j++) { + InternalRow expected = row(j, "v" + j, partitionName); + InternalRow looked = lookuper.lookup(row(j, partitionName)).get().getSingletonRow(); + assertThatRow(looked).withSchema(schema.getRowType()).isEqualTo(expected); + } + } + + // 2. PK stream read: LogScanner subscribes by per-partition bucket count. Each partition's + // total polled records must equal the writes; if routing used the wrong count, the + // subscribed bucket range would miss rows. + Map perBucketCount = + pollRecordCountPerBucket(table, partitionInfos, 2 * RECORDS_PER_PARTITION); + Map streamCountsPerPartition = new HashMap<>(); + perBucketCount.forEach( + (tb, c) -> streamCountsPerPartition.merge(tb.getPartitionId(), c, Integer::sum)); + for (String partitionName : OLD_NEW_PARTITIONS) { + assertThat(streamCountsPerPartition.get(idByName.get(partitionName))) + .as("PK stream count for partition %s", partitionName) + .isEqualTo(RECORDS_PER_PARTITION); + } + + // 3. Batch read: for each partition, trigger a KV snapshot on every bucket, then scan back + // per (partition, bucket) with BatchScanner using the partition's own bucket count. + for (String partitionName : OLD_NEW_PARTITIONS) { + long partitionId = idByName.get(partitionName); + int bucketCount = bucketCountByName.get(partitionName); + int partitionSum = 0; + for (int bucketId = 0; bucketId < bucketCount; bucketId++) { + TableBucket tb = new TableBucket(tableId, partitionId, bucketId); + long snapshotId = + FLUSS_CLUSTER_EXTENSION.triggerAndWaitSnapshot(tb).getSnapshotID(); + try (BatchScanner batchScanner = + table.newScan().createBatchScanner(tb, snapshotId)) { + while (true) { + CloseableIterator it = + batchScanner.pollBatch(Duration.ofSeconds(10)); + if (it == null) { + break; + } + try { + while (it.hasNext()) { + it.next(); + partitionSum++; + } + } finally { + it.close(); + } + } + } + } + assertThat(partitionSum) + .as("PK batch count for partition %s", partitionName) + .isEqualTo(RECORDS_PER_PARTITION); + } + + // 4. count(*) must use each partition's own bucket count, not the table-level count + // (which would enumerate out-of-range buckets for old partitions and skew the total). + assertThat(admin.getTableStats(tablePath).get().getRowCount()) + .isEqualTo(2L * RECORDS_PER_PARTITION); + } + + @Test + void testSameValueBucketNumAlterIsNoOp() throws Exception { + // SET ('bucket.num' = currentValue) does not change the bucket layout, so it must not + // advance bucketCountEpoch: legacy-client routing and historical lookup both read + // epoch > 0 as evidence that mixed bucket layouts may exist. + TablePath tablePath = TablePath.of("test_db_1", "test_same_value_bucket_num_alter"); + createPartitionedTable(tablePath, logSchema()); + + TableInfo before = admin.getTableInfo(tablePath).get(); + assertThat(before.getNumBuckets()).isEqualTo(OLD_BUCKET_NUM); + + alterBucketNum(tablePath, OLD_BUCKET_NUM); + + TableInfo sameValued = admin.getTableInfo(tablePath).get(); + assertThat(sameValued.getNumBuckets()).isEqualTo(OLD_BUCKET_NUM); + assertThat(sameValued.getBucketCountEpoch()).isEqualTo(before.getBucketCountEpoch()); + + // A real rescale does advance the epoch, which also proves the assertions above observe a + // value that actually moves. + alterBucketNum(tablePath, NEW_BUCKET_NUM); + + TableInfo rescaled = admin.getTableInfo(tablePath).get(); + assertThat(rescaled.getNumBuckets()).isEqualTo(NEW_BUCKET_NUM); + assertThat(rescaled.getBucketCountEpoch()).isGreaterThan(before.getBucketCountEpoch()); + + // Repeating the ALTER at the new count is a no-op too, so the comparison is against the + // current bucket count rather than the one the table was created with. + alterBucketNum(tablePath, NEW_BUCKET_NUM); + + TableInfo after = admin.getTableInfo(tablePath).get(); + assertThat(after.getNumBuckets()).isEqualTo(NEW_BUCKET_NUM); + assertThat(after.getBucketCountEpoch()).isEqualTo(rescaled.getBucketCountEpoch()); + } + + @Test + void testBucketCountChangeCannotBeMixedWithPropertyChange() throws Exception { + TablePath tablePath = TablePath.of("test_db_1", "test_mixed_bucket_and_property_alter"); + createPartitionedTable(tablePath, logSchema()); + + assertThatThrownBy( + () -> + admin.alterTable( + tablePath, + Arrays.asList( + TableChange.modifyBucketCount( + NEW_BUCKET_NUM), + TableChange.set("custom-key", "value")), + false) + .get()) + .cause() + .isInstanceOf(InvalidAlterTableException.class) + .hasMessageContaining("table properties, table schema, or table distribution"); + assertThat(admin.getTableInfo(tablePath).get().getNumBuckets()).isEqualTo(OLD_BUCKET_NUM); + } + + @Test + void testDynamicallyCreatedPartitionUsesPostAlterBucketCount() throws Exception { + // A partition created dynamically by the WRITER after an ALTER must use the new bucket + // count and be readable through that range. + clientConf.set(ConfigOptions.CLIENT_WRITER_DYNAMIC_CREATE_PARTITION_ENABLED, true); + conn.close(); + conn = ConnectionFactory.createConnection(clientConf); + admin = conn.getAdmin(); + + TablePath tablePath = TablePath.of("test_db_1", "test_rescale_dynamic_create"); + Schema schema = logSchema(); + createPartitionedTable(tablePath, schema); + alterBucketNum(tablePath, NEW_BUCKET_NUM); + + // write to a partition that does not exist yet; the writer creates it dynamically + Table table = conn.getTable(tablePath); + AppendWriter appendWriter = table.newAppend().createWriter(); + Map> expectedByPartitionId = new HashMap<>(); + List expectedRows = new ArrayList<>(); + for (int j = 0; j < RECORDS_PER_PARTITION; j++) { + InternalRow r = row(j, "v" + j, "auto"); + appendWriter.append(r); + expectedRows.add(r); + } + appendWriter.flush(); + + // the dynamically created partition carries the post-ALTER bucket count + List partitionInfos = admin.listPartitionInfos(tablePath).get(); + PartitionInfo autoPartition = + partitionInfos.stream() + .filter(p -> "auto".equals(p.getPartitionName())) + .findFirst() + .orElseThrow(() -> new AssertionError("dynamic partition was not created")); + assertThat(autoPartition.getBucketCount()).isEqualTo(NEW_BUCKET_NUM); + expectedByPartitionId.put(autoPartition.getPartitionId(), expectedRows); + + // all rows are readable through the partition's own bucket range + Map> actualByPartitionId = + scanAllBucketsPerPartition(table, Collections.singletonList(autoPartition)); + assertRowsPerPartition(schema.getRowType(), actualByPartitionId, expectedByPartitionId); + } + + @Test + void testStaleTableHandleWritesToDynamicallyCreatedPartitionAfterAlter() throws Exception { + // Old writers hold a stale table-level bucket count; new partitions must route by + // their actual (post-ALTER) count, otherwise lookups miss. + clientConf.set(ConfigOptions.CLIENT_WRITER_DYNAMIC_CREATE_PARTITION_ENABLED, true); + conn.close(); + conn = ConnectionFactory.createConnection(clientConf); + admin = conn.getAdmin(); + + TablePath tablePath = TablePath.of("test_db_1", "test_rescale_stale_handle"); + Schema schema = pkSchema("a", "c"); + createPartitionedTable(tablePath, schema); + + // open the handle BEFORE the ALTER, then rescale on the server side + Table staleTable = conn.getTable(tablePath); + UpsertWriter upsertWriter = staleTable.newUpsert().createWriter(); + alterBucketNum(tablePath, NEW_BUCKET_NUM); + + // The stale handle initially routes by the old table-level count. Dynamic creation stays + // asynchronous; once the new partition metadata arrives, affected batches fail instead of + // being sent with a bucket id computed from the wrong count. + List rows = new ArrayList<>(); + List> initialFutures = new ArrayList<>(); + for (int j = 0; j < RECORDS_PER_PARTITION; j++) { + InternalRow record = row(j, "v" + j, "auto"); + rows.add(record); + initialFutures.add(upsertWriter.upsert(record)); + } + upsertWriter.flush(); + + List failedRows = new ArrayList<>(); + for (int i = 0; i < initialFutures.size(); i++) { + try { + initialFutures.get(i).get(); + } catch (ExecutionException e) { + assertThat(e.getCause()).isInstanceOf(InvalidBucketRoutingException.class); + failedRows.add(rows.get(i)); + } + } + assertThat(failedRows).isNotEmpty(); + + // Retry only failed records. The first rejection invalidated stale routing metadata and the + // assigner, so the same stale table handle now resolves the partition's actual count. + List> retryFutures = new ArrayList<>(); + for (InternalRow failedRow : failedRows) { + retryFutures.add(upsertWriter.upsert(failedRow)); + } + upsertWriter.flush(); + for (CompletableFuture retryFuture : retryFutures) { + retryFuture.get(); + } + + // the dynamically created partition carries the post-ALTER bucket count + List partitionInfos = admin.listPartitionInfos(tablePath).get(); + assertThat(bucketCountByName(partitionInfos)).containsEntry("auto", NEW_BUCKET_NUM); + + // every key must be found after retrying the batches rejected during the rescale window + Lookuper lookuper = staleTable.newLookup().createLookuper(); + for (int j = 0; j < RECORDS_PER_PARTITION; j++) { + InternalRow expected = row(j, "v" + j, "auto"); + InternalRow looked = lookuper.lookup(row(j, "auto")).get().getSingletonRow(); + assertThatRow(looked).withSchema(schema.getRowType()).isEqualTo(expected); + } + } + + @Test + void testPrefixLookupAcrossPartitionsWithDifferentBucketCounts() throws Exception { + // Prefix lookup must resolve the bucket with the correct per-partition count; a mismatch + // would query the wrong bucket and miss rows. + TablePath tablePath = TablePath.of("test_db_1", "test_rescale_prefix_lookup"); + Schema schema = pkSchema("a", "b", "c"); + createPartitionedTable(tablePath, schema, "a"); + setupOldNewPartitions(tablePath); + + int aCardinality = 8; + int bPerA = 3; + Table table = conn.getTable(tablePath); + UpsertWriter upsertWriter = table.newUpsert().createWriter(); + for (String partitionName : OLD_NEW_PARTITIONS) { + for (int a = 0; a < aCardinality; a++) { + for (int k = 0; k < bPerA; k++) { + upsertWriter.upsert(row(a, "b" + k, partitionName)); + } + } + } + upsertWriter.flush(); + + Lookuper prefixLookuper = + table.newLookup().lookupBy(Arrays.asList("a", "c")).createLookuper(); + for (String partitionName : OLD_NEW_PARTITIONS) { + for (int a = 0; a < aCardinality; a++) { + List rows = + prefixLookuper.lookup(row(a, partitionName)).get().getRowList(); + assertThat(rows) + .as("prefix (a=%d, c=%s) should return %d rows", a, partitionName, bPerA) + .hasSize(bPerA); + } + } + } + + // ==================== helpers ==================== + + private static Schema logSchema() { + return Schema.newBuilder() + .column("a", DataTypes.INT()) + .column("b", DataTypes.STRING()) + .column("c", DataTypes.STRING()) + .build(); + } + + private static Schema pkSchema(String... primaryKeys) { + return Schema.newBuilder() + .column("a", DataTypes.INT()) + .column("b", DataTypes.STRING()) + .column("c", DataTypes.STRING()) + .primaryKey(primaryKeys) + .build(); + } + + /** Creates a table partitioned by "c" with OLD_BUCKET_NUM buckets and given bucket keys. */ + private void createPartitionedTable(TablePath tablePath, Schema schema, String... bucketKeys) + throws Exception { + createTable( + tablePath, + TableDescriptor.builder() + .schema(schema) + .distributedBy(OLD_BUCKET_NUM, bucketKeys) + .partitionedBy("c") + .build(), + true); + } + + /** + * Creates the "old" partition, ALTERs bucket.num to NEW_BUCKET_NUM, creates the "new" + * partition, and asserts the reported per-partition bucket counts. + */ + private List setupOldNewPartitions(TablePath tablePath) throws Exception { + // old partition (created before ALTER -> OLD_BUCKET_NUM buckets) + admin.createPartition(tablePath, newPartitionSpec("c", "old"), false).get(); + alterBucketNum(tablePath, NEW_BUCKET_NUM); + // new partition (created after ALTER -> NEW_BUCKET_NUM buckets) + admin.createPartition(tablePath, newPartitionSpec("c", "new"), false).get(); + + List partitionInfos = admin.listPartitionInfos(tablePath).get(); + assertThat(bucketCountByName(partitionInfos)) + .containsEntry("old", OLD_BUCKET_NUM) + .containsEntry("new", NEW_BUCKET_NUM); + return partitionInfos; + } + + /** Upserts RECORDS_PER_PARTITION rows (j, "v"+j, partition) into "old" and "new". */ + private static void upsertRowsToOldAndNew(Table table) throws Exception { + UpsertWriter upsertWriter = table.newUpsert().createWriter(); + for (String partitionName : OLD_NEW_PARTITIONS) { + for (int j = 0; j < RECORDS_PER_PARTITION; j++) { + upsertWriter.upsert(row(j, "v" + j, partitionName)); + } + } + upsertWriter.flush(); + } + + private void alterBucketNum(TablePath tablePath, int newBucketNum) throws Exception { + admin.alterTable( + tablePath, + Collections.singletonList(TableChange.modifyBucketCount(newBucketNum)), + false) + .get(); + } + + private static Map bucketCountByName(List partitionInfos) { + Map map = new HashMap<>(); + for (PartitionInfo p : partitionInfos) { + map.put(p.getPartitionName(), p.getBucketCount()); + } + return map; + } + + private static Map partitionIdByName(List partitionInfos) { + Map map = new HashMap<>(); + for (PartitionInfo p : partitionInfos) { + map.put(p.getPartitionName(), p.getPartitionId()); + } + return map; + } + + private static void subscribeAllBuckets( + LogScanner logScanner, List partitionInfos) { + for (PartitionInfo partitionInfo : partitionInfos) { + for (int bucketId = 0; bucketId < partitionInfo.getBucketCount(); bucketId++) { + logScanner.subscribeFromBeginning(partitionInfo.getPartitionId(), bucketId); + } + } + } + + /** + * Subscribes every bucket of every partition using that partition's own bucket count and polls + * until {@code expectedTotal} records arrive, returning the record count per bucket. If write + * routing used the wrong (table-level) bucket count for a partition, the scan of its real + * bucket range would miss rows and the expected count would never be reached. + */ + private static Map pollRecordCountPerBucket( + Table table, List partitionInfos, int expectedTotal) throws Exception { + Map perBucketCount = new HashMap<>(); + int scanned = 0; + try (LogScanner logScanner = table.newScan().createLogScanner()) { + subscribeAllBuckets(logScanner, partitionInfos); + long deadline = System.currentTimeMillis() + Duration.ofMinutes(1).toMillis(); + while (scanned < expectedTotal && System.currentTimeMillis() < deadline) { + ScanRecords scanRecords = logScanner.poll(Duration.ofSeconds(1)); + for (TableBucket scanBucket : scanRecords.buckets()) { + int c = 0; + for (ScanRecord ignored : scanRecords.records(scanBucket)) { + c++; + } + perBucketCount.merge(scanBucket, c, Integer::sum); + scanned += c; + } + } + } + assertThat(scanned).isEqualTo(expectedTotal); + return perBucketCount; + } + + /** + * Subscribes every bucket of every partition using that partition's own bucket count and + * collects all rows grouped by partition id. If write routing used the wrong (table-level) + * bucket count for a partition, the scan of its real bucket range would miss rows and the + * expected count would never be reached. + */ + private static Map> scanAllBucketsPerPartition( + Table table, List partitionInfos) throws Exception { + int totalExpected = partitionInfos.size() * RECORDS_PER_PARTITION; + Map> actual = new HashMap<>(); + int scanned = 0; + try (LogScanner logScanner = table.newScan().createLogScanner()) { + subscribeAllBuckets(logScanner, partitionInfos); + long deadline = System.currentTimeMillis() + Duration.ofMinutes(1).toMillis(); + while (scanned < totalExpected && System.currentTimeMillis() < deadline) { + ScanRecords scanRecords = logScanner.poll(Duration.ofSeconds(1)); + for (TableBucket scanBucket : scanRecords.buckets()) { + for (ScanRecord record : scanRecords.records(scanBucket)) { + actual.computeIfAbsent(scanBucket.getPartitionId(), k -> new ArrayList<>()) + .add(record.getRow()); + } + } + scanned += scanRecords.count(); + } + } + assertThat(scanned).isEqualTo(totalExpected); + return actual; + } + + private static void assertRowsPerPartition( + RowType rowType, + Map> actual, + Map> expected) { + assertThat(actual.keySet()).isEqualTo(expected.keySet()); + for (Map.Entry> entry : expected.entrySet()) { + List actualRows = actual.get(entry.getKey()); + List expectedRows = entry.getValue(); + // rows from different buckets of the same partition may interleave, so compare as a + // multiset: same size and same elements regardless of order. + assertThat(actualRows).hasSameSizeAs(expectedRows); + for (InternalRow expectedRow : expectedRows) { + boolean found = + actualRows.stream() + .anyMatch( + a -> { + try { + assertThatRow(a) + .withSchema(rowType) + .isEqualTo(expectedRow); + return true; + } catch (AssertionError e) { + return false; + } + }); + assertThat(found) + .as("expected row %s present in partition rows", expectedRow) + .isTrue(); + } + } + } +} diff --git a/fluss-client/src/test/java/org/apache/fluss/client/table/PartitionedTableITCase.java b/fluss-client/src/test/java/org/apache/fluss/client/table/PartitionedTableITCase.java index d204d087a4d..9154c83a68e 100644 --- a/fluss-client/src/test/java/org/apache/fluss/client/table/PartitionedTableITCase.java +++ b/fluss-client/src/test/java/org/apache/fluss/client/table/PartitionedTableITCase.java @@ -45,7 +45,6 @@ import static org.apache.fluss.record.TestData.DATA1_TABLE_PATH_PK; import static org.apache.fluss.testutils.DataTestUtils.row; import static org.apache.fluss.testutils.InternalRowAssert.assertThatRow; -import static org.apache.fluss.testutils.common.CommonTestUtils.retry; import static org.apache.fluss.testutils.common.CommonTestUtils.waitValue; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; @@ -221,19 +220,14 @@ void testCreatePartitionExceedMaxPartitionNumber() throws Exception { upsertWriter.upsert(row).get(); } - // add one row will not throw TooManyPartitionsException immediately. - upsertWriter.upsert(row(10, "a" + 10, "10")); - - // add another rows will throw TooManyPartitionsException final. - retry( - Duration.ofMinutes(1), - () -> - assertThatThrownBy(() -> upsertWriter.upsert(row(10, "a" + 10, "10")).get()) - .rootCause() - .isInstanceOf(TooManyPartitionsException.class) - .hasMessageContaining( - "Exceed the maximum number of partitions for table " - + "test_db_1.test_pk_table_1, only allow 10 partitions.")); + // Dynamic partition creation is synchronous (the bucket id needs the partition's own + // bucket count), so a partition that cannot be created fails the record right away. + assertThatThrownBy(() -> upsertWriter.upsert(row(10, "a" + 10, "10")).get()) + .rootCause() + .isInstanceOf(TooManyPartitionsException.class) + .hasMessageContaining( + "Exceed the maximum number of partitions for table " + + "test_db_1.test_pk_table_1, only allow 10 partitions."); } private Schema createPartitionedTable(TablePath tablePath, boolean isPrimaryTable) diff --git a/fluss-client/src/test/java/org/apache/fluss/client/table/scanner/log/LogFetcherITCase.java b/fluss-client/src/test/java/org/apache/fluss/client/table/scanner/log/LogFetcherITCase.java index d4572577655..ddbfaf462b7 100644 --- a/fluss-client/src/test/java/org/apache/fluss/client/table/scanner/log/LogFetcherITCase.java +++ b/fluss-client/src/test/java/org/apache/fluss/client/table/scanner/log/LogFetcherITCase.java @@ -280,7 +280,8 @@ void testFetchWhenDestinationIsNullInMetadata() throws Exception { oldCluster.getCoordinatorServer(), oldCluster.getBucketLocationsByPath(), oldCluster.getTableIdByPath(), - oldCluster.getPartitionIdByPath()); + oldCluster.getPartitionIdByPath(), + Collections.emptyMap()); metadataUpdater = new MetadataUpdater(rpcClient, clientConf, newCluster); LogScannerStatus logScannerStatus = new LogScannerStatus(); diff --git a/fluss-client/src/test/java/org/apache/fluss/client/utils/ClientRpcMessageUtilsTest.java b/fluss-client/src/test/java/org/apache/fluss/client/utils/ClientRpcMessageUtilsTest.java index 14f341fb7fd..07c63e464e4 100644 --- a/fluss-client/src/test/java/org/apache/fluss/client/utils/ClientRpcMessageUtilsTest.java +++ b/fluss-client/src/test/java/org/apache/fluss/client/utils/ClientRpcMessageUtilsTest.java @@ -22,13 +22,22 @@ import org.apache.fluss.memory.MemorySegment; import org.apache.fluss.memory.PreAllocatedPagedOutputView; import org.apache.fluss.metadata.KvFormat; +import org.apache.fluss.metadata.PartitionInfo; import org.apache.fluss.metadata.PhysicalTablePath; import org.apache.fluss.metadata.TableBucket; +import org.apache.fluss.metadata.TableChange; +import org.apache.fluss.rpc.messages.AlterTableRequest; +import org.apache.fluss.rpc.messages.ListPartitionInfosResponse; +import org.apache.fluss.rpc.messages.PbKeyValue; +import org.apache.fluss.rpc.messages.PbPartitionInfo; +import org.apache.fluss.rpc.messages.PbPartitionSpec; import org.apache.fluss.rpc.messages.PutKvRequest; import org.apache.fluss.rpc.protocol.MergeMode; import org.junit.jupiter.api.Test; +import javax.annotation.Nullable; + import java.util.Arrays; import java.util.Collections; import java.util.List; @@ -126,6 +135,69 @@ void testMakePutKvRequestWithSingleBatch() throws Exception { assertThat(request.getAggMode()).isEqualTo(MergeMode.OVERWRITE.getProtoValue()); } + @Test + void testMakeAlterTableRequestWithBucketCountChange() { + AlterTableRequest request = + ClientRpcMessageUtils.makeAlterTableRequest( + DATA1_TABLE_PATH_PK, + Collections.singletonList(TableChange.modifyBucketCount(8)), + false); + + assertThat(request.hasModifyBucketCount()).isTrue(); + assertThat(request.getModifyBucketCount().getNewBucketCount()).isEqualTo(8); + assertThat(request.getConfigChangesList()).isEmpty(); + } + + @Test + void testToPartitionInfosParsesBucketCount() { + // one partition with bucket_count set, one without (simulating an old cluster / old + // partition that did not persist per-partition bucket count) + ListPartitionInfosResponse response = + new ListPartitionInfosResponse() + .addAllPartitionsInfos( + Arrays.asList( + makePbPartitionInfo(1L, "20240101", "file://dir1", 8), + makePbPartitionInfo(2L, "20240102", null, null))); + + List partitionInfos = ClientRpcMessageUtils.toPartitionInfos(response, 4); + + assertThat(partitionInfos).hasSize(2); + + PartitionInfo withBucketCount = partitionInfos.get(0); + assertThat(withBucketCount.getPartitionId()).isEqualTo(1L); + assertThat(withBucketCount.getPartitionName()).isEqualTo("20240101"); + assertThat(withBucketCount.getRemoteDataDir()).isEqualTo("file://dir1"); + assertThat(withBucketCount.getBucketCount()).isEqualTo(8); + + // backward compatibility: missing bucket_count must resolve to the given table-level + // default, not the proto default 0 + PartitionInfo withoutBucketCount = partitionInfos.get(1); + assertThat(withoutBucketCount.getPartitionId()).isEqualTo(2L); + assertThat(withoutBucketCount.getPartitionName()).isEqualTo("20240102"); + assertThat(withoutBucketCount.getRemoteDataDir()).isNull(); + assertThat(withoutBucketCount.getBucketCount()).isEqualTo(4); + } + + private static PbPartitionInfo makePbPartitionInfo( + long partitionId, + String partitionValue, + @Nullable String remoteDataDir, + @Nullable Integer bucketCount) { + PbPartitionSpec partitionSpec = new PbPartitionSpec(); + PbKeyValue keyValue = new PbKeyValue().setKey("dt").setValue(partitionValue); + partitionSpec.addAllPartitionKeyValues(Collections.singletonList(keyValue)); + + PbPartitionInfo pbPartitionInfo = + new PbPartitionInfo().setPartitionId(partitionId).setPartitionSpec(partitionSpec); + if (remoteDataDir != null) { + pbPartitionInfo.setRemoteDataDir(remoteDataDir); + } + if (bucketCount != null) { + pbPartitionInfo.setBucketCount(bucketCount); + } + return pbPartitionInfo; + } + private KvWriteBatch createKvWriteBatch(int bucketId, MergeMode mergeMode) throws Exception { MemorySegment segment = MemorySegment.allocateHeapMemory(1024); PreAllocatedPagedOutputView outputView = @@ -133,6 +205,7 @@ private KvWriteBatch createKvWriteBatch(int bucketId, MergeMode mergeMode) throw return new KvWriteBatch( DATA1_TABLE_ID_PK, bucketId, + DATA1_TABLE_INFO_PK.getNumBuckets(), PhysicalTablePath.of(DATA1_TABLE_PATH_PK), DATA1_TABLE_INFO_PK.getSchemaId(), KvFormat.COMPACTED, diff --git a/fluss-client/src/test/java/org/apache/fluss/client/write/ArrowLogWriteBatchTest.java b/fluss-client/src/test/java/org/apache/fluss/client/write/ArrowLogWriteBatchTest.java index 54032b2f0e0..c00b1df29a7 100644 --- a/fluss-client/src/test/java/org/apache/fluss/client/write/ArrowLogWriteBatchTest.java +++ b/fluss-client/src/test/java/org/apache/fluss/client/write/ArrowLogWriteBatchTest.java @@ -128,6 +128,7 @@ void testAppendWithPreAllocatedMemorySegments() throws Exception { new ArrowLogWriteBatch( tb.getTableId(), tb.getBucket(), + DATA1_TABLE_INFO.getNumBuckets(), DATA1_PHYSICAL_TABLE_PATH, DATA1_TABLE_INFO.getSchemaId(), writerProvider.getOrCreateWriter( @@ -210,6 +211,7 @@ void testArrowCompressionRatioEstimated() throws Exception { new ArrowLogWriteBatch( tb.getTableId(), tb.getBucket(), + DATA1_TABLE_INFO.getNumBuckets(), DATA1_PHYSICAL_TABLE_PATH, DATA1_TABLE_INFO.getSchemaId(), arrowWriter, @@ -308,6 +310,7 @@ private ArrowLogWriteBatch createArrowLogWriteBatch(TableBucket tb, int maxSizeI return new ArrowLogWriteBatch( tb.getTableId(), tb.getBucket(), + DATA1_TABLE_INFO.getNumBuckets(), DATA1_PHYSICAL_TABLE_PATH, DATA1_TABLE_INFO.getSchemaId(), writerProvider.getOrCreateWriter( diff --git a/fluss-client/src/test/java/org/apache/fluss/client/write/CompactedLogWriteBatchTest.java b/fluss-client/src/test/java/org/apache/fluss/client/write/CompactedLogWriteBatchTest.java index c4eac5b766f..8f10e365ed9 100644 --- a/fluss-client/src/test/java/org/apache/fluss/client/write/CompactedLogWriteBatchTest.java +++ b/fluss-client/src/test/java/org/apache/fluss/client/write/CompactedLogWriteBatchTest.java @@ -250,6 +250,7 @@ private CompactedLogWriteBatch createLogWriteBatch( return new CompactedLogWriteBatch( tb.getTableId(), tb.getBucket(), + DATA1_TABLE_INFO.getNumBuckets(), DATA1_PHYSICAL_TABLE_PATH, DATA1_TABLE_INFO.getSchemaId(), writeLimit, diff --git a/fluss-client/src/test/java/org/apache/fluss/client/write/IndexedLogWriteBatchTest.java b/fluss-client/src/test/java/org/apache/fluss/client/write/IndexedLogWriteBatchTest.java index 347ef1d1ee8..ba7f2d03930 100644 --- a/fluss-client/src/test/java/org/apache/fluss/client/write/IndexedLogWriteBatchTest.java +++ b/fluss-client/src/test/java/org/apache/fluss/client/write/IndexedLogWriteBatchTest.java @@ -212,6 +212,7 @@ private IndexedLogWriteBatch createLogWriteBatch( return new IndexedLogWriteBatch( tb.getTableId(), tb.getBucket(), + DATA1_TABLE_INFO.getNumBuckets(), DATA1_PHYSICAL_TABLE_PATH, DATA1_TABLE_INFO.getSchemaId(), writeLimit, diff --git a/fluss-client/src/test/java/org/apache/fluss/client/write/KvWriteBatchTest.java b/fluss-client/src/test/java/org/apache/fluss/client/write/KvWriteBatchTest.java index 984c17741a2..f99149b50fc 100644 --- a/fluss-client/src/test/java/org/apache/fluss/client/write/KvWriteBatchTest.java +++ b/fluss-client/src/test/java/org/apache/fluss/client/write/KvWriteBatchTest.java @@ -223,6 +223,7 @@ private KvWriteBatch createKvWriteBatch( return new KvWriteBatch( tb.getTableId(), tb.getBucket(), + DATA1_TABLE_INFO_PK.getNumBuckets(), PhysicalTablePath.of(DATA1_TABLE_PATH_PK), DATA1_TABLE_INFO_PK.getSchemaId(), KvFormat.COMPACTED, @@ -320,6 +321,7 @@ private KvWriteBatch createKvWriteBatchWithMergeMode(TableBucket tb, MergeMode m return new KvWriteBatch( tb.getTableId(), tb.getBucket(), + DATA1_TABLE_INFO_PK.getNumBuckets(), PhysicalTablePath.of(DATA1_TABLE_PATH_PK), DATA1_TABLE_INFO_PK.getSchemaId(), KvFormat.COMPACTED, diff --git a/fluss-client/src/test/java/org/apache/fluss/client/write/RecordAccumulatorTest.java b/fluss-client/src/test/java/org/apache/fluss/client/write/RecordAccumulatorTest.java index 18fd9f81fee..866a35a0b32 100644 --- a/fluss-client/src/test/java/org/apache/fluss/client/write/RecordAccumulatorTest.java +++ b/fluss-client/src/test/java/org/apache/fluss/client/write/RecordAccumulatorTest.java @@ -102,6 +102,8 @@ class RecordAccumulatorTest { System.currentTimeMillis(), System.currentTimeMillis()); + private final int numBuckets = DATA1_TABLE_INFO.getNumBuckets(); + ServerNode node1 = new ServerNode(1, "localhost", 90, ServerType.TABLET_SERVER, "rack1"); ServerNode node2 = new ServerNode(2, "localhost", 91, ServerType.TABLET_SERVER, "rack2"); ServerNode node3 = new ServerNode(3, "localhost", 92, ServerType.TABLET_SERVER, "rack3"); @@ -157,7 +159,7 @@ void testDrainBatches() throws Exception { // initial data. for (int i = 0; i < 4; i++) { - accum.append(createRecord(row), writeCallback, cluster, i, false); + accum.append(createRecord(row), writeCallback, cluster, i, numBuckets, false); } // drain batches from 2 nodes: node1 => tb1, node2 => tb3, because the max request size is @@ -170,8 +172,8 @@ void testDrainBatches() throws Exception { verifyTableBucketInBatches(batches1, tb1, tb3); // add record for tb1, tb3 - accum.append(createRecord(row), writeCallback, cluster, 0, false); - accum.append(createRecord(row), writeCallback, cluster, 2, false); + accum.append(createRecord(row), writeCallback, cluster, 0, numBuckets, false); + accum.append(createRecord(row), writeCallback, cluster, 2, numBuckets, false); // drain batches from 2 nodes: node1 => tb2, node2 => tb4, because the max request size is // full after the first batch drained. The drain index should start from next table bucket, @@ -251,7 +253,14 @@ private void appendUntilBatchFull(RecordAccumulator accum, int bucketId) throws PhysicalTablePath tablePath = PhysicalTablePath.of(ZSTD_TABLE_INFO.getTablePath()); WriteRecord record = WriteRecord.forArrowAppend(ZSTD_TABLE_INFO, tablePath, row, null); // append until the batch is full - if (accum.append(record, writeCallback, cluster, bucketId, false).batchIsFull) { + if (accum.append( + record, + writeCallback, + cluster, + bucketId, + ZSTD_TABLE_INFO.getNumBuckets(), + false) + .batchIsFull) { break; } } @@ -267,7 +276,7 @@ void testFull() throws Exception { int appends = expectedNumAppends(row, batchSize); for (int i = 0; i < appends; i++) { // append to the first batch - accum.append(createRecord(row), writeCallback, cluster, 0, false); + accum.append(createRecord(row), writeCallback, cluster, 0, numBuckets, false); Deque writeBatches = accum.getReadyDeque(DATA1_PHYSICAL_TABLE_PATH, tb1.getBucket()); assertThat(writeBatches).hasSize(1); @@ -280,7 +289,7 @@ void testFull() throws Exception { // this appends doesn't fit in the first batch, so a new batch is created and the first // batch is closed. - accum.append(createRecord(row), writeCallback, cluster, 0, false); + accum.append(createRecord(row), writeCallback, cluster, 0, numBuckets, false); Deque writeBatches = accum.getReadyDeque(DATA1_PHYSICAL_TABLE_PATH, tb1.getBucket()); assertThat(writeBatches).hasSize(2); @@ -310,6 +319,27 @@ void testFull() throws Exception { } } + @Test + void testAppendRollsNewBatchWhenBucketCountChanges() throws Exception { + int batchSize = 1024; + IndexedRow row = indexedRow(DATA1_ROW_TYPE, new Object[] {1, "a"}); + RecordAccumulator accum = createTestRecordAccumulator(batchSize, 10L * batchSize); + + accum.append(createRecord(row), writeCallback, cluster, 0, numBuckets, false); + accum.append(createRecord(row), writeCallback, cluster, 0, numBuckets + 1, false); + + Deque writeBatches = + accum.getReadyDeque(DATA1_PHYSICAL_TABLE_PATH, tb1.getBucket()); + assertThat(writeBatches).hasSize(2); + Iterator batchIterator = writeBatches.iterator(); + WriteBatch oldBatch = batchIterator.next(); + assertThat(oldBatch.isClosed()).isTrue(); + assertThat(oldBatch.getBucketCount()).isEqualTo(numBuckets); + WriteBatch newBatch = batchIterator.next(); + assertThat(newBatch.isClosed()).isFalse(); + assertThat(newBatch.getBucketCount()).isEqualTo(numBuckets + 1); + } + @Test void testAppendRollsNewBatchWhenSchemaIdChanges() throws Exception { int batchSize = 1024; @@ -318,10 +348,15 @@ void testAppendRollsNewBatchWhenSchemaIdChanges() throws Exception { int oldSchemaId = DATA1_TABLE_INFO.getSchemaId(); int newSchemaId = oldSchemaId + 1; - accum.append(createRecord(row), writeCallback, cluster, 0, false); + accum.append(createRecord(row), writeCallback, cluster, 0, numBuckets, false); // a record with a bumped schema id closes the old-schema batch and rolls a new one. accum.append( - createRecord(row, withSchemaId(newSchemaId)), writeCallback, cluster, 0, false); + createRecord(row, withSchemaId(newSchemaId)), + writeCallback, + cluster, + 0, + numBuckets, + false); Deque writeBatches = accum.getReadyDeque(DATA1_PHYSICAL_TABLE_PATH, tb1.getBucket()); @@ -347,7 +382,7 @@ void testAppendLarge() throws Exception { IndexedRow row1 = indexedRow(DATA1_ROW_TYPE, new Object[] {100000000, new String(new char[2 * 100])}); // row size > 10; - accum.append(createRecord(row1), writeCallback, cluster, 0, false); + accum.append(createRecord(row1), writeCallback, cluster, 0, numBuckets, false); // bucket's leader should be ready for bucket0. assertThat(accum.ready(cluster).readyNodes).isEqualTo(Collections.singleton(node1.id())); @@ -389,7 +424,7 @@ void testAppendWithStickyBucketAssigner() throws Exception { // Create first batch. int bucketId = bucketAssigner.assignBucket(cluster); - accum.append(createRecord(row), writeCallback, cluster, bucketId, false); + accum.append(createRecord(row), writeCallback, cluster, bucketId, numBuckets, false); int appends = 1; boolean switchBucket = false; @@ -397,7 +432,8 @@ void testAppendWithStickyBucketAssigner() throws Exception { // Append to the first batch. bucketId = bucketAssigner.assignBucket(cluster); RecordAccumulator.RecordAppendResult result = - accum.append(createRecord(row), writeCallback, cluster, bucketId, true); + accum.append( + createRecord(row), writeCallback, cluster, bucketId, numBuckets, true); int numBatches = getBatchNumInAccum(accum); // Only one batch is created because the bucket is sticky. assertThat(numBatches).isEqualTo(1); @@ -418,7 +454,7 @@ void testAppendWithStickyBucketAssigner() throws Exception { // Writer would call this method in this case, make second batch. bucketAssigner.onNewBatch(cluster, bucketId); bucketId = bucketAssigner.assignBucket(cluster); - accum.append(createRecord(row), writeCallback, cluster, bucketId, false); + accum.append(createRecord(row), writeCallback, cluster, bucketId, numBuckets, false); appends++; // These append operations all go into the second batch. @@ -426,7 +462,8 @@ void testAppendWithStickyBucketAssigner() throws Exception { // Append to the first batch. bucketId = bucketAssigner.assignBucket(cluster); RecordAccumulator.RecordAppendResult result = - accum.append(createRecord(row), writeCallback, cluster, bucketId, true); + accum.append( + createRecord(row), writeCallback, cluster, bucketId, numBuckets, true); int numBatches = getBatchNumInAccum(accum); // Only one batch is created because the bucket is sticky. assertThat(numBatches).isEqualTo(2); @@ -450,7 +487,13 @@ void testPartialDrain() throws Exception { List buckets = Arrays.asList(tb1, tb2); for (TableBucket tb : buckets) { for (int i = 0; i < appends; i++) { - accum.append(createRecord(row), writeCallback, cluster, tb.getBucket(), false); + accum.append( + createRecord(row), + writeCallback, + cluster, + tb.getBucket(), + numBuckets, + false); } } @@ -467,7 +510,7 @@ void testFlush() throws Exception { RecordAccumulator accum = createTestRecordAccumulator(4 * 1024, 64 * 1024); for (int i = 0; i < 100; i++) { - accum.append(createRecord(row), writeCallback, cluster, i % 3, false); + accum.append(createRecord(row), writeCallback, cluster, i % 3, numBuckets, false); assertThat(accum.hasIncomplete()).isTrue(); } @@ -505,12 +548,14 @@ void testAbortAllBatchesHandlesConcurrentCompletion() throws Exception { (bucket, offset, exception) -> completedFuture.complete(exception), cluster, tb1.getBucket(), + numBuckets, false); accum.append( createRecord(row), (bucket, offset, exception) -> abortedFuture.complete(exception), cluster, tb2.getBucket(), + numBuckets, false); List batches = @@ -556,7 +601,7 @@ void testTableWithUnknownLeader() throws Exception { // add bucket1 which leader is unknown into cluster. cluster = updateCluster(Collections.singletonList(bucket1)); - accum.append(createRecord(row), writeCallback, cluster, 0, false); + accum.append(createRecord(row), writeCallback, cluster, 0, numBuckets, false); RecordAccumulator.ReadyCheckResult readyCheckResult = accum.ready(cluster); assertThat(readyCheckResult.unknownLeaderTables) .isEqualTo(Collections.singleton(DATA1_PHYSICAL_TABLE_PATH)); @@ -577,7 +622,7 @@ void testTableWithUnknownLeader() throws Exception { void testAwaitFlushComplete() throws Exception { IndexedRow row = indexedRow(DATA1_ROW_TYPE, new Object[] {1, "a"}); RecordAccumulator accum = createTestRecordAccumulator(4 * 1024, 64 * 1024); - accum.append(createRecord(row), writeCallback, cluster, 0, false); + accum.append(createRecord(row), writeCallback, cluster, 0, numBuckets, false); accum.beginFlush(); assertThat(accum.flushInProgress()).isTrue(); @@ -598,7 +643,13 @@ public void testNextReadyCheckDelay() throws Exception { // Add data for bucket 1 for (int i = 0; i < appends; i++) { - accum.append(createRecord(row), writeCallback, cluster, bucket1.getBucketId(), false); + accum.append( + createRecord(row), + writeCallback, + cluster, + bucket1.getBucketId(), + numBuckets, + false); } RecordAccumulator.ReadyCheckResult result = accum.ready(cluster); assertThat(result.readyNodes).isEmpty(); @@ -608,14 +659,26 @@ public void testNextReadyCheckDelay() throws Exception { // Add data for bucket 3 for (int i = 0; i < appends; i++) { - accum.append(createRecord(row), writeCallback, cluster, bucket3.getBucketId(), false); + accum.append( + createRecord(row), + writeCallback, + cluster, + bucket3.getBucketId(), + numBuckets, + false); } result = accum.ready(cluster); assertThat(result.readyNodes).hasSize(0); assertThat(result.nextReadyCheckDelayMs).isEqualTo(batchTimeout / 2); // Append one more data for bucket1 should make the batch full and sendable immediately - accum.append(createRecord(row), writeCallback, cluster, bucket1.getBucketId(), false); + accum.append( + createRecord(row), + writeCallback, + cluster, + bucket1.getBucketId(), + numBuckets, + false); result = accum.ready(cluster); // server for bucket1 should be ready now @@ -646,7 +709,7 @@ private TableInfo withSchemaId(int schemaId) { DATA1_TABLE_INFO.getSchema(), DATA1_TABLE_INFO.getBucketKeys(), DATA1_TABLE_INFO.getPartitionKeys(), - DATA1_TABLE_INFO.getNumBuckets(), + numBuckets, DATA1_TABLE_INFO.getProperties(), DATA1_TABLE_INFO.getCustomProperties(), DATA1_TABLE_INFO.getRemoteDataDir(), @@ -673,6 +736,7 @@ private Cluster updateCluster(List bucketLocations) { new ServerNode(0, "localhost", 89, ServerType.COORDINATOR), bucketsByPath, tableIdByPath, + Collections.emptyMap(), Collections.emptyMap()); } @@ -724,8 +788,8 @@ void testDrainContinuesWhenBucketAtMaxInflight() throws Exception { IndexedRow row = indexedRow(DATA1_ROW_TYPE, new Object[] {1, "a"}); // Drain both buckets so each has 1 in-flight batch. - accum.append(createRecord(row), writeCallback, cluster, tb1.getBucket(), false); - accum.append(createRecord(row), writeCallback, cluster, tb2.getBucket(), false); + accum.append(createRecord(row), writeCallback, cluster, tb1.getBucket(), numBuckets, false); + accum.append(createRecord(row), writeCallback, cluster, tb2.getBucket(), numBuckets, false); Map> firstDrain = accum.drain(cluster, Collections.singleton(node1.id()), Integer.MAX_VALUE); @@ -738,8 +802,8 @@ void testDrainContinuesWhenBucketAtMaxInflight() throws Exception { idempotenceManager.handleCompletedBatch(tb2Batch); // Append again to both. On drain, tb1 should be skipped but tb2 should still be drained. - accum.append(createRecord(row), writeCallback, cluster, tb1.getBucket(), false); - accum.append(createRecord(row), writeCallback, cluster, tb2.getBucket(), false); + accum.append(createRecord(row), writeCallback, cluster, tb1.getBucket(), numBuckets, false); + accum.append(createRecord(row), writeCallback, cluster, tb2.getBucket(), numBuckets, false); Map> secondDrain = accum.drain(cluster, Collections.singleton(node1.id()), Integer.MAX_VALUE); @@ -878,8 +942,8 @@ void testThrottledBucketSkippedInDrain() throws Exception { // Append records to tb1 and tb2. Both tb1 and tb2 lead on node1, so draining node1 // should only return tb2 when tb1 is throttled. - accum.append(createRecord(row), writeCallback, cluster, 0, false); - accum.append(createRecord(row), writeCallback, cluster, 1, false); + accum.append(createRecord(row), writeCallback, cluster, 0, numBuckets, false); + accum.append(createRecord(row), writeCallback, cluster, 1, numBuckets, false); // Throttle tb1 accum.updateThrottle(tb1, 0.5f); diff --git a/fluss-client/src/test/java/org/apache/fluss/client/write/SenderTest.java b/fluss-client/src/test/java/org/apache/fluss/client/write/SenderTest.java index 43ddafa3e5f..989059d8016 100644 --- a/fluss-client/src/test/java/org/apache/fluss/client/write/SenderTest.java +++ b/fluss-client/src/test/java/org/apache/fluss/client/write/SenderTest.java @@ -28,6 +28,7 @@ import org.apache.fluss.config.MemorySize; import org.apache.fluss.exception.AuthorizationException; import org.apache.fluss.exception.FlussRuntimeException; +import org.apache.fluss.exception.InvalidBucketRoutingException; import org.apache.fluss.exception.NetworkException; import org.apache.fluss.exception.OutOfOrderSequenceException; import org.apache.fluss.exception.PartitionNotExistException; @@ -39,6 +40,7 @@ import org.apache.fluss.metadata.TableBucket; import org.apache.fluss.metadata.TableDescriptor; import org.apache.fluss.metadata.TableInfo; +import org.apache.fluss.metadata.TableOrPartition; import org.apache.fluss.metadata.TablePath; import org.apache.fluss.record.MemoryLogRecords; import org.apache.fluss.row.BinaryRow; @@ -215,6 +217,75 @@ void testReroutesWriteAfterExplicitMissingPartitionResponse() throws Exception { assertThat(future.get()).isNull(); } + @Test + void testFailsQueuedBatchWhenPartitionBucketCountChanged() throws Exception { + sender.destroyResources(); + TableInfo tableInfo = createHistoricalTableInfo(); + PhysicalTablePath partitionPath = + PhysicalTablePath.of(tableInfo.getTablePath(), "20990101"); + TableBucket tableBucket = new TableBucket(tableInfo.getTableId(), 21L, 0); + Map bucketCounts = new HashMap<>(); + bucketCounts.put(TableOrPartition.ofPartition(tableBucket.getPartitionId()), 4); + metadataUpdater.updateCluster( + partitionedCluster( + tableInfo, + Collections.singletonMap(partitionPath, tableBucket), + bucketCounts)); + sender = setupWithIdempotenceState(); + + CompletableFuture future = + appendKvRecord(tableInfo, partitionPath, 1, metadataUpdater.getCluster(), 2); + sender.runOnce(); + + assertThat(future.get()) + .isInstanceOf(InvalidBucketRoutingException.class) + .hasMessageContaining("bucket count changed"); + assertThatThrownBy(() -> node1Gateway().getRequest(0)) + .isInstanceOf(IllegalStateException.class) + .hasMessageContaining("No requests pending"); + } + + @Test + void testAbortsRerouteWhenQueuedBatchBucketCountDiffers() throws Exception { + sender.destroyResources(); + TableInfo tableInfo = createHistoricalTableInfo(); + PhysicalTablePath originalPath = PhysicalTablePath.of(tableInfo.getTablePath(), "20990101"); + PhysicalTablePath historicalPath = + PhysicalTablePath.of(tableInfo.getTablePath(), HISTORICAL_PARTITION_VALUE); + TableBucket originalBucket = new TableBucket(tableInfo.getTableId(), 21L, 0); + TableBucket historicalBucket = new TableBucket(tableInfo.getTableId(), 22L, 0); + metadataUpdater = missingPartitionMetadataUpdater(tableInfo, originalPath); + Map tableBucketsByPath = new HashMap<>(); + tableBucketsByPath.put(originalPath, originalBucket); + tableBucketsByPath.put(historicalPath, historicalBucket); + // The historical partition keeps one bucket while the queued batch was routed by a + // rescaled partition's count of four: its bucket id cannot be moved to the historical + // layout, so the reroute must abort instead of silently misrouting the records. + Map bucketCountByTableOrPartition = new HashMap<>(); + bucketCountByTableOrPartition.put( + TableOrPartition.ofPartition(historicalBucket.getPartitionId()), 1); + metadataUpdater.updateCluster( + partitionedCluster(tableInfo, tableBucketsByPath, bucketCountByTableOrPartition)); + sender = setupWithIdempotenceState(); + + CompletableFuture future = + appendKvRecord(tableInfo, originalPath, 1, metadataUpdater.getCluster(), 4); + sender.runOnce(); + + TestTabletServerGateway gateway = node1Gateway(); + gateway.response( + 0, createPutKvResponse(originalBucket, Errors.UNKNOWN_TABLE_OR_BUCKET_EXCEPTION)); + sender.runOnce(); + + assertThat(future.get()) + .isInstanceOf(PartitionNotExistException.class) + .hasMessageContaining("different bucket count"); + // Nothing was sent to the historical partition: aborting is the whole point. + assertThatThrownBy(() -> gateway.getRequest(0)) + .isInstanceOf(IllegalStateException.class) + .hasMessageContaining("No requests pending"); + } + @Test void testAbortsOnlyMissingPartitionWhenRerouteIsUnsafe() throws Exception { sender.destroyResources(); @@ -1381,7 +1452,8 @@ void testSendWhenDestinationIsNullInMetadata() throws Exception { oldCluster.getCoordinatorServer(), oldCluster.getBucketLocationsByPath(), oldCluster.getTableIdByPath(), - oldCluster.getPartitionIdByPath()); + oldCluster.getPartitionIdByPath(), + Collections.emptyMap()); metadataUpdater.updateCluster(newCluster); @@ -1474,6 +1546,7 @@ void testRetryPutKeyWithSchemaNotExistException() throws Exception { (tb, leo, e) -> future.complete(e), metadataUpdater.getCluster(), 0, + DATA1_TABLE_INFO_PK.getNumBuckets(), false); sender.runOnce(); finishRequest(tableBucket, 0, createPutKvResponse(tableBucket, SCHEMA_NOT_EXIST)); @@ -1661,6 +1734,99 @@ void testSendWhenTableIdChanges() throws Exception { assertThat(future2.get()).isNull(); } + @Test + void testInvalidBucketRoutingFailsBatchAndInvalidatesBucketAssigner() throws Exception { + // Recreate sender with a tracking bucketAssignerInvalidator. + IdempotenceManager idempotenceManager = createIdempotenceManager(false); + Configuration conf = new Configuration(); + conf.set(ConfigOptions.CLIENT_WRITER_BUFFER_MEMORY_SIZE, new MemorySize(TOTAL_MEMORY_SIZE)); + conf.set(ConfigOptions.CLIENT_WRITER_BATCH_SIZE, new MemorySize(BATCH_SIZE)); + conf.set(ConfigOptions.CLIENT_WRITER_BUFFER_PAGE_SIZE, new MemorySize(PAGE_SIZE)); + conf.set(ConfigOptions.CLIENT_WRITER_BATCH_TIMEOUT, Duration.ofMillis(0)); + accumulator = + new RecordAccumulator( + conf, idempotenceManager, writerMetricGroup, SystemClock.getInstance()); + AtomicReference invalidatedBucket = new AtomicReference<>(); + Sender staleSender = + new Sender( + accumulator, + REQUEST_TIMEOUT, + MAX_REQUEST_SIZE, + ACKS_ALL, + Integer.MAX_VALUE, + metadataUpdater, + idempotenceManager, + writerMetricGroup, + invalidatedBucket::set); + + // Append one record and send it. + CompletableFuture future = new CompletableFuture<>(); + appendToAccumulator(tb1, row(1, "a"), (tb, leo, e) -> future.complete(e)); + staleSender.runOnce(); + assertThat(staleSender.numOfInFlightBatches(tb1)).isEqualTo(1); + + // Server rejects the bucketId computed with a stale count. + Cluster clusterBeforeError = metadataUpdater.getCluster(); + finishRequest(tb1, 0, createProduceLogResponse(tb1, Errors.INVALID_BUCKET_ROUTING)); + + // The batch is failed (not re-enqueued for retry — the bucketId is stale and must not be + // reused). + assertThat(staleSender.numOfInFlightBatches(tb1)).isEqualTo(0); + + // The BucketAssigner for this bucket was invalidated so the next send rebuilds it with + // the refreshed bucket count. + assertThat(invalidatedBucket.get()).isEqualTo(tb1); + + // The table's bucket metadata was invalidated so the next send requests it again. + assertThat(metadataUpdater.getCluster()).isNotSameAs(clusterBeforeError); + + // The write callback receives the non-retriable routing error. + Exception exception = future.get(); + assertThat(exception).isInstanceOf(InvalidBucketRoutingException.class); + } + + @Test + void testInvalidBucketRoutingReclaimsBatchSequenceWhenIdempotenceEnabled() throws Exception { + // Bucket routing is validated for hash-distributed tables, so exercise the client reclaim + // path on a primary-key table bucket rather than a keyless one. + TableBucket keyedBucket = new TableBucket(DATA1_TABLE_ID_PK, 0); + IdempotenceManager idempotenceManager = createIdempotenceManager(true); + Sender staleSender = setupWithIdempotenceState(idempotenceManager); + staleSender.runOnce(); + long writerId = idempotenceManager.writerId(); + assertThat(idempotenceManager.isWriterIdValid()).isTrue(); + assertThat(idempotenceManager.nextSequence(keyedBucket)).isEqualTo(0); + + // Drain and send one batch: it takes batch sequence 0 and nextSequence advances to 1. + CompletableFuture future = new CompletableFuture<>(); + appendKvToAccumulator( + keyedBucket, + compactedRow(DATA1_ROW_TYPE, new Object[] {1, "a"}), + (tb, leo, e) -> future.complete(e)); + staleSender.runOnce(); + assertThat(idempotenceManager.nextSequence(keyedBucket)).isEqualTo(1); + + // The server rejects the batch during pre-append routing validation, so it was provably + // never written. Its batch sequence (0) must be reclaimed. + finishRequest( + keyedBucket, 0, createPutKvResponse(keyedBucket, Errors.INVALID_BUCKET_ROUTING)); + staleSender.runOnce(); + + // The write callback receives the non-retriable routing error. + assertThat(future.get()).isInstanceOf(InvalidBucketRoutingException.class); + + // The writer id must survive: nothing was accepted, so there is no lost message to guard. + assertThat(idempotenceManager.isWriterIdValid()).isTrue(); + assertThat(idempotenceManager.writerId()).isEqualTo(writerId); + + // The reclaimed sequence must roll nextSequence back to 0. Otherwise a permanent hole at + // sequence 0 remains: the next batch that reaches the server on this bucket (created after + // the metadata refresh, carrying a valid routing count) would send sequence 1 against an + // expected 0, triggering OUT_OF_ORDER_SEQUENCE_EXCEPTION and resetWriterId, which wipes + // idempotence for every bucket of this writer. + assertThat(idempotenceManager.nextSequence(keyedBucket)).isEqualTo(0); + } + private TestingMetadataUpdater initializeMetadataUpdater() { Map tableInfos = new HashMap<>(); tableInfos.put(DATA1_TABLE_PATH, DATA1_TABLE_INFO); @@ -1732,6 +1898,13 @@ public boolean checkAndUpdatePartitionMetadata(PhysicalTablePath physicalTablePa private static Cluster partitionedCluster( TableInfo tableInfo, Map tableBucketsByPath) { + return partitionedCluster(tableInfo, tableBucketsByPath, Collections.emptyMap()); + } + + private static Cluster partitionedCluster( + TableInfo tableInfo, + Map tableBucketsByPath, + Map bucketCountByTableOrPartition) { int[] replicas = new int[] {TestingMetadataUpdater.NODE1.id()}; Map> bucketLocationsByPath = new HashMap<>(); Map partitionIdsByPath = new HashMap<>(); @@ -1753,12 +1926,23 @@ private static Cluster partitionedCluster( TestingMetadataUpdater.COORDINATOR, bucketLocationsByPath, Collections.singletonMap(tableInfo.getTablePath(), tableInfo.getTableId()), - partitionIdsByPath); + partitionIdsByPath, + bucketCountByTableOrPartition); } private CompletableFuture appendKvRecord( TableInfo tableInfo, PhysicalTablePath physicalTablePath, int id, Cluster cluster) throws Exception { + return appendKvRecord(tableInfo, physicalTablePath, id, cluster, 0); + } + + private CompletableFuture appendKvRecord( + TableInfo tableInfo, + PhysicalTablePath physicalTablePath, + int id, + Cluster cluster, + int bucketCount) + throws Exception { accumulator.checkAndCacheHistoricalPartitionEnabled(tableInfo); BinaryRow row = compactedRow( @@ -1782,6 +1966,7 @@ private CompletableFuture appendKvRecord( (tableBucket, logEndOffset, error) -> future.complete(error), cluster, 0, + bucketCount, false); return future; } @@ -1800,6 +1985,7 @@ private void appendToAccumulator( writeCallback, metadataUpdater.getCluster(), tb.getBucket(), + tableInfo.getNumBuckets(), false); } @@ -1828,6 +2014,7 @@ private void appendKvToAccumulator( writeCallback, metadataUpdater.getCluster(), tableBucket.getBucket(), + tableInfo.getNumBuckets(), false); } @@ -1964,7 +2151,8 @@ private Sender setupWithIdempotenceState( reties, metadataUpdater, idempotenceManager, - writerMetricGroup); + writerMetricGroup, + tb -> {}); } private IdempotenceManager createIdempotenceManager(boolean idempotenceEnabled) { diff --git a/fluss-client/src/test/java/org/apache/fluss/client/write/StickyStaticBucketAssignerTest.java b/fluss-client/src/test/java/org/apache/fluss/client/write/StickyStaticBucketAssignerTest.java index f8289063397..6009b96f51a 100644 --- a/fluss-client/src/test/java/org/apache/fluss/client/write/StickyStaticBucketAssignerTest.java +++ b/fluss-client/src/test/java/org/apache/fluss/client/write/StickyStaticBucketAssignerTest.java @@ -215,6 +215,7 @@ private Cluster updateCluster(List bucketLocations) { new ServerNode(0, "localhost", 89, ServerType.COORDINATOR), bucketsByPath, tableIdByPath, + Collections.emptyMap(), Collections.emptyMap()); } } diff --git a/fluss-common/src/main/java/org/apache/fluss/cluster/Cluster.java b/fluss-common/src/main/java/org/apache/fluss/cluster/Cluster.java index f78ade43b2d..41752f7dbc3 100644 --- a/fluss-common/src/main/java/org/apache/fluss/cluster/Cluster.java +++ b/fluss-common/src/main/java/org/apache/fluss/cluster/Cluster.java @@ -18,9 +18,13 @@ package org.apache.fluss.cluster; import org.apache.fluss.annotation.Internal; +import org.apache.fluss.exception.InvalidBucketRoutingException; import org.apache.fluss.exception.PartitionNotExistException; import org.apache.fluss.metadata.PhysicalTablePath; import org.apache.fluss.metadata.TableBucket; +import org.apache.fluss.metadata.TableInfo; +import org.apache.fluss.metadata.TableOrPartition; +import org.apache.fluss.metadata.TablePartition; import org.apache.fluss.metadata.TablePath; import javax.annotation.Nullable; @@ -28,6 +32,7 @@ import java.util.ArrayList; import java.util.Collections; import java.util.HashMap; +import java.util.HashSet; import java.util.List; import java.util.Map; import java.util.Optional; @@ -52,19 +57,23 @@ public final class Cluster { private final Map pathByTableId; private final Map partitionsIdByPath; private final Map partitionNameById; + private final Map bucketCountByTableOrPartition; public Cluster( Map aliveTabletServersById, @Nullable ServerNode coordinatorServer, Map> bucketLocationsByPath, Map tableIdByPath, - Map partitionsIdByPath) { + Map partitionsIdByPath, + Map bucketCountByTableOrPartition) { this.coordinatorServer = coordinatorServer; this.aliveTabletServersById = Collections.unmodifiableMap(aliveTabletServersById); this.aliveTabletServers = Collections.unmodifiableList(new ArrayList<>(aliveTabletServersById.values())); this.tableIdByPath = Collections.unmodifiableMap(tableIdByPath); this.partitionsIdByPath = Collections.unmodifiableMap(partitionsIdByPath); + this.bucketCountByTableOrPartition = + Collections.unmodifiableMap(bucketCountByTableOrPartition); // Index the bucket locations by table path, and index bucket location by bucket. // Note that this code is performance sensitive if there are a large number of buckets, @@ -127,12 +136,35 @@ public Cluster invalidPhysicalTableBucketMeta(Set physicalTab new ArrayList<>(tablePathAndBucketLocations.getValue())); } } + // resolve the invalid table or partition keys so the bucket count map can be filtered + Set invalidTableOrPartitions = new HashSet<>(); + for (PhysicalTablePath path : physicalTablesToInvalid) { + if (path.getPartitionName() == null) { + Long tableId = tableIdByPath.get(path.getTablePath()); + if (tableId != null) { + invalidTableOrPartitions.add(TableOrPartition.ofTable(tableId)); + } + } else { + Long partitionId = partitionsIdByPath.get(path); + if (partitionId != null) { + invalidTableOrPartitions.add(TableOrPartition.ofPartition(partitionId)); + } + } + } + Map newBucketCountByTableOrPartition = new HashMap<>(); + for (Map.Entry entry : + bucketCountByTableOrPartition.entrySet()) { + if (!invalidTableOrPartitions.contains(entry.getKey())) { + newBucketCountByTableOrPartition.put(entry.getKey(), entry.getValue()); + } + } return new Cluster( new HashMap<>(aliveTabletServersById), coordinatorServer, newBucketLocationsByPath, new HashMap<>(tableIdByPath), - new HashMap<>(partitionsIdByPath)); + new HashMap<>(partitionsIdByPath), + newBucketCountByTableOrPartition); } /** Invalidates bucket metadata and partition ID mappings for the given physical table paths. */ @@ -148,7 +180,8 @@ public Cluster invalidPhysicalTableBucketAndPartitionMeta( coordinatorServer, new HashMap<>(cluster.availableLocationsByPath), new HashMap<>(tableIdByPath), - newPartitionsIdByPath); + newPartitionsIdByPath, + new HashMap<>(cluster.bucketCountByTableOrPartition)); } @Nullable @@ -226,6 +259,24 @@ public Optional getPartitionId(PhysicalTablePath physicalTablePath) { return Optional.ofNullable(partitionsIdByPath.get(physicalTablePath)); } + /** + * Resolve a {@link PhysicalTablePath} to its current {@link TablePartition} (tableId + + * partitionId) from this snapshot. Retained for name resolution; the actual bucket-count lookup + * uses {@link #getBucketCount(TableOrPartition)}. Resolving both ids from the same snapshot + * avoids combining a stale tableId/partitionId with a newer one after a replacement. + */ + public Optional getTablePartition(PhysicalTablePath physicalTablePath) { + Long partitionId = partitionsIdByPath.get(physicalTablePath); + if (partitionId == null) { + return Optional.empty(); + } + Long tableId = tableIdByPath.get(physicalTablePath.getTablePath()); + if (tableId == null) { + return Optional.empty(); + } + return Optional.of(new TablePartition(tableId, partitionId)); + } + public TableBucket getTableBucket( long tableId, PhysicalTablePath physicalTablePath, int bucketId) { if (physicalTablePath.getPartitionName() != null) { @@ -274,6 +325,42 @@ public Map getPartitionIdByPath() { return partitionsIdByPath; } + /** + * Get the actual bucket count for the given table or partition. Returns empty if its bucket + * layout is not available yet. + */ + public Optional getBucketCount(TableOrPartition tableOrPartition) { + return Optional.ofNullable(bucketCountByTableOrPartition.get(tableOrPartition)); + } + + /** + * Gets the actual bucket count for the given table or partition, falling back to the + * table-level count when the table has never been rescaled. + */ + public int getBucketCountOrFallback(TableInfo tableInfo, @Nullable Long partitionId) { + TableOrPartition tableOrPartition = + TableOrPartition.of(tableInfo.getTableId(), partitionId); + Integer bucketCount = bucketCountByTableOrPartition.get(tableOrPartition); + if (bucketCount != null) { + return bucketCount; + } + long bucketCountEpoch = tableInfo.getBucketCountEpoch(); + if (bucketCountEpoch > 0) { + throw new InvalidBucketRoutingException( + "Routing bucket count is unavailable for " + + tableOrPartition + + " at bucketCountEpoch " + + bucketCountEpoch + + "; refusing to fall back to the table-level count."); + } + return tableInfo.getNumBuckets(); + } + + /** Get the table or partition to bucket count map. */ + public Map getBucketCountByTableOrPartition() { + return bucketCountByTableOrPartition; + } + /** Create an empty cluster instance with no nodes and no table-buckets. */ public static Cluster empty() { return new Cluster( @@ -281,6 +368,7 @@ public static Cluster empty() { null, Collections.emptyMap(), Collections.emptyMap(), + Collections.emptyMap(), Collections.emptyMap()); } diff --git a/fluss-common/src/main/java/org/apache/fluss/exception/InvalidBucketRoutingException.java b/fluss-common/src/main/java/org/apache/fluss/exception/InvalidBucketRoutingException.java new file mode 100644 index 00000000000..7688c18a87d --- /dev/null +++ b/fluss-common/src/main/java/org/apache/fluss/exception/InvalidBucketRoutingException.java @@ -0,0 +1,31 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to you under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.fluss.exception; + +import org.apache.fluss.annotation.PublicEvolving; + +/** Thrown when a request cannot be routed reliably with its supplied bucket information. */ +@PublicEvolving +public class InvalidBucketRoutingException extends ApiException { + + private static final long serialVersionUID = 1L; + + public InvalidBucketRoutingException(String message) { + super(message); + } +} diff --git a/fluss-common/src/main/java/org/apache/fluss/lake/lakestorage/LakeCatalog.java b/fluss-common/src/main/java/org/apache/fluss/lake/lakestorage/LakeCatalog.java index bc64cdf93fb..b7f05723efe 100644 --- a/fluss-common/src/main/java/org/apache/fluss/lake/lakestorage/LakeCatalog.java +++ b/fluss-common/src/main/java/org/apache/fluss/lake/lakestorage/LakeCatalog.java @@ -51,6 +51,12 @@ void createTable(TablePath tablePath, TableDescriptor tableDescriptor, Context c /** * Alter a table in lake. * + *

A {@link TableChange.ModifyBucketCount} is Fluss's coordinator-orchestrated bucket count + * rescale: implementations supporting rescale must apply it to their bucket layout option, + * others should throw {@link UnsupportedOperationException}. User-facing changes to the + * lake-native bucket option (e.g. Paimon {@code bucket}) keep being rejected to prevent the two + * systems from diverging. + * * @param tablePath path of the table to be altered * @param tableChanges The changes to be applied to the table * @param context contextual information needed for alter table diff --git a/fluss-common/src/main/java/org/apache/fluss/lake/lakestorage/LakeTableLookuper.java b/fluss-common/src/main/java/org/apache/fluss/lake/lakestorage/LakeTableLookuper.java index 6c5f6cd2849..6c93d9226e4 100644 --- a/fluss-common/src/main/java/org/apache/fluss/lake/lakestorage/LakeTableLookuper.java +++ b/fluss-common/src/main/java/org/apache/fluss/lake/lakestorage/LakeTableLookuper.java @@ -74,7 +74,7 @@ default void requestRefresh() { /** Context for a lake table point lookup. */ final class LookupContext { private final ResolvedPartitionSpec partitionSpec; - private final int bucketId; + private final @Nullable Integer bucketId; private final short schemaId; private final RowType valueRowType; private final LookupMetricRecorder lookupMetricRecorder; @@ -83,14 +83,15 @@ final class LookupContext { * Creates a lookup context. * * @param partitionSpec resolved Fluss partition spec for the lookup - * @param bucketId target bucket id in the lake table + * @param bucketId target bucket id in the lake table, or null when the caller cannot + * determine it and the implementation has to resolve it from the lake metadata * @param schemaId schema id to encode the returned Fluss value with * @param valueRowType row type to encode the returned Fluss value with * @param lookupMetricRecorder recorder for lake table point lookup metrics */ public LookupContext( ResolvedPartitionSpec partitionSpec, - int bucketId, + @Nullable Integer bucketId, short schemaId, RowType valueRowType, LookupMetricRecorder lookupMetricRecorder) { @@ -107,8 +108,11 @@ public ResolvedPartitionSpec partitionSpec() { return partitionSpec; } - /** Returns the target bucket id in the lake table. */ - public int bucketId() { + /** + * Returns the target bucket id in the lake table, or null when the implementation has to + * resolve it from the lake metadata. + */ + public @Nullable Integer bucketId() { return bucketId; } diff --git a/fluss-common/src/main/java/org/apache/fluss/lake/writer/WriterInitContext.java b/fluss-common/src/main/java/org/apache/fluss/lake/writer/WriterInitContext.java index f268a5c48af..0d031cc94f6 100644 --- a/fluss-common/src/main/java/org/apache/fluss/lake/writer/WriterInitContext.java +++ b/fluss-common/src/main/java/org/apache/fluss/lake/writer/WriterInitContext.java @@ -105,4 +105,12 @@ default long tieringRoundTimestamp() { default String[] ioTmpDirs() { return null; } + + /** + * Returns the actual bucket count of the target partition, or the table-level count for + * non-partitioned tables. After an ALTER bucket.num old partitions keep their original count, + * so lake writers must stamp bucket layouts with this value instead of the lake table's current + * schema-level bucket setting. + */ + int bucketCount(); } diff --git a/fluss-common/src/main/java/org/apache/fluss/metadata/PartitionInfo.java b/fluss-common/src/main/java/org/apache/fluss/metadata/PartitionInfo.java index d8845fc03f8..83b872b97ba 100644 --- a/fluss-common/src/main/java/org/apache/fluss/metadata/PartitionInfo.java +++ b/fluss-common/src/main/java/org/apache/fluss/metadata/PartitionInfo.java @@ -35,11 +35,22 @@ public class PartitionInfo { private final ResolvedPartitionSpec partitionSpec; private final @Nullable String remoteDataDir; + /** + * The bucket count of this partition. Always resolved: for partitions created by older versions + * that did not persist a per-partition bucket count, the table-level bucket count is filled in + * at construction time. + */ + private final int bucketCount; + public PartitionInfo( - long partitionId, ResolvedPartitionSpec partitionSpec, @Nullable String remoteDataDir) { + long partitionId, + ResolvedPartitionSpec partitionSpec, + @Nullable String remoteDataDir, + int bucketCount) { this.partitionId = partitionId; this.partitionSpec = partitionSpec; this.remoteDataDir = remoteDataDir; + this.bucketCount = bucketCount; } /** Get the partition id. The id is globally unique in the Fluss cluster. */ @@ -68,6 +79,25 @@ public String getRemoteDataDir() { return remoteDataDir; } + /** + * Get the bucket count of this partition. For partitions created by older versions without a + * persisted per-partition bucket count, this is the table-level bucket count. + */ + public int getBucketCount() { + return bucketCount; + } + + /** + * Resolves the effective bucket count for a (possibly absent) partition: returns the + * partition's bucket count when {@code partitionInfo} is non-null, otherwise the table-level + * bucket count. The null case represents a non-partitioned table or a partition whose + * PartitionInfo is not available. + */ + public static int bucketCountOrDefault( + @Nullable PartitionInfo partitionInfo, int tableBucketCount) { + return partitionInfo != null ? partitionInfo.getBucketCount() : tableBucketCount; + } + @Override public boolean equals(Object o) { if (this == o) { @@ -79,12 +109,13 @@ public boolean equals(Object o) { PartitionInfo that = (PartitionInfo) o; return partitionId == that.partitionId && Objects.equals(partitionSpec, that.partitionSpec) - && Objects.equals(remoteDataDir, that.remoteDataDir); + && Objects.equals(remoteDataDir, that.remoteDataDir) + && bucketCount == that.bucketCount; } @Override public int hashCode() { - return Objects.hash(partitionId, partitionSpec, remoteDataDir); + return Objects.hash(partitionId, partitionSpec, remoteDataDir, bucketCount); } @Override @@ -96,6 +127,8 @@ public String toString() { + partitionId + ", remoteDataDir=" + remoteDataDir + + ", bucketCount=" + + bucketCount + '}'; } } diff --git a/fluss-common/src/main/java/org/apache/fluss/metadata/TableChange.java b/fluss-common/src/main/java/org/apache/fluss/metadata/TableChange.java index 2890ca82e6e..41be2b85c46 100644 --- a/fluss-common/src/main/java/org/apache/fluss/metadata/TableChange.java +++ b/fluss-common/src/main/java/org/apache/fluss/metadata/TableChange.java @@ -24,6 +24,8 @@ import java.util.Objects; import java.util.Optional; +import static org.apache.fluss.utils.Preconditions.checkArgument; + /** {@link TableChange} represents the modification of the Fluss Table. */ public interface TableChange { @@ -132,6 +134,18 @@ static ResetOption reset(String key) { return new ResetOption(key); } + /** + * Changes the default bucket count for newly created partitions. + * + *

Existing partitions retain their bucket counts. + * + * @param newBucketCount the new default bucket count; must be positive + * @return the bucket count change + */ + static ModifyBucketCount modifyBucketCount(int newBucketCount) { + return new ModifyBucketCount(newBucketCount); + } + /** * A table change to set the table option. * @@ -229,6 +243,50 @@ public String toString() { } } + /** A change to the table's distribution. */ + interface DistributionChange extends TableChange {} + + /** Changes the default bucket count for newly created partitions. */ + final class ModifyBucketCount implements DistributionChange { + + private final int newBucketCount; + + private ModifyBucketCount(int newBucketCount) { + checkArgument( + newBucketCount > 0, + "Bucket count must be positive, but was %s.", + newBucketCount); + this.newBucketCount = newBucketCount; + } + + /** Returns the new default bucket count. */ + public int getNewBucketCount() { + return newBucketCount; + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (!(o instanceof ModifyBucketCount)) { + return false; + } + ModifyBucketCount that = (ModifyBucketCount) o; + return newBucketCount == that.newBucketCount; + } + + @Override + public int hashCode() { + return Objects.hash(newBucketCount); + } + + @Override + public String toString() { + return "ModifyBucketCount{" + "newBucketCount=" + newBucketCount + '}'; + } + } + /** A table change to modify the table schema. */ interface SchemaChange extends TableChange {} diff --git a/fluss-common/src/main/java/org/apache/fluss/metadata/TableInfo.java b/fluss-common/src/main/java/org/apache/fluss/metadata/TableInfo.java index 76ef538de41..b4be1f0049d 100644 --- a/fluss-common/src/main/java/org/apache/fluss/metadata/TableInfo.java +++ b/fluss-common/src/main/java/org/apache/fluss/metadata/TableInfo.java @@ -68,6 +68,7 @@ public final class TableInfo { private final long createdTime; private final long modifiedTime; + private final long bucketCountEpoch; private int[] cachedStatsIndexMapping = null; @@ -85,6 +86,38 @@ public TableInfo( @Nullable String comment, long createdTime, long modifiedTime) { + this( + tablePath, + tableId, + schemaId, + schema, + bucketKeys, + partitionKeys, + numBuckets, + properties, + customProperties, + remoteDataDir, + comment, + createdTime, + modifiedTime, + 0L); + } + + public TableInfo( + TablePath tablePath, + long tableId, + int schemaId, + Schema schema, + List bucketKeys, + List partitionKeys, + int numBuckets, + Configuration properties, + Configuration customProperties, + @Nullable String remoteDataDir, + @Nullable String comment, + long createdTime, + long modifiedTime, + long bucketCountEpoch) { this.tablePath = tablePath; this.tableId = tableId; this.schemaId = schemaId; @@ -102,6 +135,7 @@ public TableInfo( this.comment = comment; this.createdTime = createdTime; this.modifiedTime = modifiedTime; + this.bucketCountEpoch = bucketCountEpoch; } /** @@ -403,6 +437,14 @@ public long getModifiedTime() { return modifiedTime; } + /** + * Returns the bucket layout epoch of the table. New tables start at 0; every committed + * bucket.num change increments it (see {@code TableRegistration#withBucketCount(int)}). + */ + public long getBucketCountEpoch() { + return bucketCountEpoch; + } + /** * Converts this table info to a {@link TableDescriptor}. * @@ -431,6 +473,30 @@ public static TableInfo of( String remoteDataDir, long createdTime, long modifiedTime) { + return of( + tablePath, + tableId, + schemaId, + tableDescriptor, + remoteDataDir, + createdTime, + modifiedTime, + 0L); + } + + /** + * Creates a {@link TableInfo} from a {@link TableDescriptor} and other metadata, including the + * bucket layout epoch. + */ + public static TableInfo of( + TablePath tablePath, + long tableId, + int schemaId, + TableDescriptor tableDescriptor, + String remoteDataDir, + long createdTime, + long modifiedTime, + long bucketCountEpoch) { Schema schema = tableDescriptor.getSchema(); int numBuckets = tableDescriptor @@ -453,7 +519,8 @@ public static TableInfo of( remoteDataDir, tableDescriptor.getComment().orElse(null), createdTime, - modifiedTime); + modifiedTime, + bucketCountEpoch); } @Override @@ -466,6 +533,7 @@ public boolean equals(Object o) { return tableId == that.tableId && schemaId == that.schemaId && numBuckets == that.numBuckets + && bucketCountEpoch == that.bucketCountEpoch && Objects.equals(tablePath, that.tablePath) && Objects.equals(rowType, that.rowType) && Objects.equals(primaryKeys, that.primaryKeys) @@ -494,7 +562,8 @@ public int hashCode() { properties, customProperties, remoteDataDir, - comment); + comment, + bucketCountEpoch); } @Override @@ -529,6 +598,8 @@ public String toString() { + createdTime + ", modifiedTime=" + modifiedTime + + ", bucketCountEpoch=" + + bucketCountEpoch + '}'; } diff --git a/fluss-common/src/main/java/org/apache/fluss/metadata/TableOrPartition.java b/fluss-common/src/main/java/org/apache/fluss/metadata/TableOrPartition.java new file mode 100644 index 00000000000..29871fbd7bf --- /dev/null +++ b/fluss-common/src/main/java/org/apache/fluss/metadata/TableOrPartition.java @@ -0,0 +1,90 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.fluss.metadata; + +import org.apache.fluss.annotation.Internal; + +import javax.annotation.Nullable; + +import java.util.Objects; + +/** + * A class to identify a table or a partition, containing the table id and the optional partition + * id. + */ +@Internal +public class TableOrPartition { + + @Nullable private final Long tableId; + @Nullable private final Long partitionId; + + /** Create a {@link TableOrPartition} instance for a table. */ + public static TableOrPartition ofTable(long tableId) { + return new TableOrPartition(tableId, null); + } + + /** Create a {@link TableOrPartition} instance for a partition. */ + public static TableOrPartition ofPartition(long partitionId) { + return new TableOrPartition(null, partitionId); + } + + /** + * Create a {@link TableOrPartition} instance for the given table id and optional partition id: + * a partition when {@code partitionId} is non-null, otherwise the table itself. Note that a + * partition is identified by its (globally unique) partition id alone, so the table id is not + * retained in that case. + */ + public static TableOrPartition of(long tableId, @Nullable Long partitionId) { + return partitionId == null ? ofTable(tableId) : ofPartition(partitionId); + } + + private TableOrPartition(@Nullable Long tableId, @Nullable Long partitionId) { + this.tableId = tableId; + this.partitionId = partitionId; + } + + @Nullable + public Long getTableId() { + return tableId; + } + + @Nullable + public Long getPartitionId() { + return partitionId; + } + + @Override + public boolean equals(Object o) { + if (o == null || getClass() != o.getClass()) { + return false; + } + TableOrPartition that = (TableOrPartition) o; + return Objects.equals(tableId, that.tableId) + && Objects.equals(partitionId, that.partitionId); + } + + @Override + public int hashCode() { + return Objects.hash(tableId, partitionId); + } + + @Override + public String toString() { + return "TableOrPartition{" + "tableId=" + tableId + ", partitionId=" + partitionId + '}'; + } +} diff --git a/fluss-common/src/test/java/org/apache/fluss/cluster/ClusterTest.java b/fluss-common/src/test/java/org/apache/fluss/cluster/ClusterTest.java index f46ae574c0f..9116b874073 100644 --- a/fluss-common/src/test/java/org/apache/fluss/cluster/ClusterTest.java +++ b/fluss-common/src/test/java/org/apache/fluss/cluster/ClusterTest.java @@ -17,7 +17,10 @@ package org.apache.fluss.cluster; +import org.apache.fluss.exception.InvalidBucketRoutingException; import org.apache.fluss.metadata.PhysicalTablePath; +import org.apache.fluss.metadata.TableInfo; +import org.apache.fluss.metadata.TableOrPartition; import org.apache.fluss.metadata.TablePath; import org.junit.jupiter.api.BeforeEach; @@ -33,10 +36,13 @@ import static org.apache.fluss.record.TestData.DATA1_PHYSICAL_TABLE_PATH; import static org.apache.fluss.record.TestData.DATA1_PHYSICAL_TABLE_PATH_PA_2024; +import static org.apache.fluss.record.TestData.DATA1_TABLE_DESCRIPTOR; import static org.apache.fluss.record.TestData.DATA1_TABLE_ID; +import static org.apache.fluss.record.TestData.DATA1_TABLE_INFO; import static org.apache.fluss.record.TestData.DATA1_TABLE_PATH; import static org.apache.fluss.record.TestData.DATA2_TABLE_ID; import static org.apache.fluss.record.TestData.DATA2_TABLE_PATH; +import static org.apache.fluss.record.TestData.DEFAULT_REMOTE_DATA_DIR; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; @@ -97,6 +103,7 @@ void testInvalidMetaAndUpdate() { COORDINATOR_SERVER, new HashMap<>(cluster.getBucketLocationsByPath()), new HashMap<>(cluster.getTableIdByPath()), + Collections.emptyMap(), Collections.emptyMap()); } @@ -125,7 +132,8 @@ void testInvalidPartitionMeta() { COORDINATOR_SERVER, new HashMap<>(initialCluster.getBucketLocationsByPath()), new HashMap<>(initialCluster.getTableIdByPath()), - Collections.singletonMap(DATA1_PHYSICAL_TABLE_PATH_PA_2024, partitionId)); + Collections.singletonMap(DATA1_PHYSICAL_TABLE_PATH_PA_2024, partitionId), + Collections.emptyMap()); assertThat(cluster.getPartitionId(DATA1_PHYSICAL_TABLE_PATH_PA_2024)).hasValue(partitionId); assertThat(cluster.getPartitionName(partitionId)).hasValue("2024"); @@ -146,6 +154,53 @@ void testInvalidPartitionMeta() { assertThat(cluster.getPartitionName(partitionId)).isNotPresent(); } + @Test + void testGetBucketCountOrFallback() { + long partitionId = 42L; + Map bucketCounts = new HashMap<>(); + bucketCounts.put(TableOrPartition.ofTable(DATA1_TABLE_ID), 5); + bucketCounts.put(TableOrPartition.ofPartition(partitionId), 7); + Cluster cluster = + new Cluster( + aliveTabletServersById, + COORDINATOR_SERVER, + Collections.emptyMap(), + Collections.singletonMap(DATA1_TABLE_PATH, DATA1_TABLE_ID), + Collections.emptyMap(), + bucketCounts); + + assertThat(cluster.getBucketCountOrFallback(DATA1_TABLE_INFO, null)).isEqualTo(5); + assertThat(cluster.getBucketCountOrFallback(DATA1_TABLE_INFO, partitionId)).isEqualTo(7); + + Cluster clusterWithoutBucketCount = + new Cluster( + aliveTabletServersById, + COORDINATOR_SERVER, + Collections.emptyMap(), + Collections.singletonMap(DATA1_TABLE_PATH, DATA1_TABLE_ID), + Collections.emptyMap(), + Collections.emptyMap()); + assertThat(clusterWithoutBucketCount.getBucketCountOrFallback(DATA1_TABLE_INFO, null)) + .isEqualTo(DATA1_TABLE_INFO.getNumBuckets()); + + TableInfo rescaledTableInfo = + TableInfo.of( + DATA1_TABLE_PATH, + DATA1_TABLE_ID, + 1, + DATA1_TABLE_DESCRIPTOR, + DEFAULT_REMOTE_DATA_DIR, + 1L, + 1L, + 1L); + assertThatThrownBy( + () -> + clusterWithoutBucketCount.getBucketCountOrFallback( + rescaledTableInfo, partitionId)) + .isInstanceOf(InvalidBucketRoutingException.class) + .hasMessageContaining("bucketCountEpoch 1"); + } + @Test void testGetRandomTabletServer() { Map aliveTabletServersById = new HashMap<>(); @@ -209,6 +264,7 @@ private Cluster createCluster(Map aliveTabletServersById) { COORDINATOR_SERVER, tablePathToBucketLocations, tablePathToTableId, + Collections.emptyMap(), Collections.emptyMap()); } } diff --git a/fluss-common/src/test/java/org/apache/fluss/lake/source/TestingLakeSource.java b/fluss-common/src/test/java/org/apache/fluss/lake/source/TestingLakeSource.java index f9ee6756cc6..0ca98206792 100644 --- a/fluss-common/src/test/java/org/apache/fluss/lake/source/TestingLakeSource.java +++ b/fluss-common/src/test/java/org/apache/fluss/lake/source/TestingLakeSource.java @@ -28,6 +28,7 @@ import java.io.DataOutputStream; import java.io.IOException; import java.util.ArrayList; +import java.util.Collections; import java.util.List; /** A testing implementation of {@link LakeSource}. */ @@ -41,20 +42,41 @@ public class TestingLakeSource implements LakeSource { // partition infos of partitions contain lake splits private final List partitionInfos; + private final List explicitSplits; + private final boolean useExplicitSplits; + public TestingLakeSource() { this.bucketNum = 0; this.partitionInfos = null; + this.explicitSplits = Collections.emptyList(); + this.useExplicitSplits = false; } public TestingLakeSource(int bucketNum, List partitionInfos) { this.bucketNum = bucketNum; this.partitionInfos = partitionInfos; + this.explicitSplits = Collections.emptyList(); + this.useExplicitSplits = false; + } + + private TestingLakeSource(List explicitSplits) { + this.bucketNum = 0; + this.partitionInfos = null; + this.explicitSplits = new ArrayList<>(explicitSplits); + this.useExplicitSplits = true; } private TestingLakeSource(TestingLakeSource source) { this.bucketNum = source.bucketNum; this.partitionInfos = source.partitionInfos == null ? null : new ArrayList<>(source.partitionInfos); + this.explicitSplits = new ArrayList<>(source.explicitSplits); + this.useExplicitSplits = source.useExplicitSplits; + } + + /** Creates a source whose planner returns exactly the supplied splits. */ + public static TestingLakeSource fromSplits(List splits) { + return new TestingLakeSource(splits); } @Override @@ -75,7 +97,9 @@ public FilterPushDownResult withFilters(List predicates) { @Override public Planner createPlanner(PlannerContext context) throws IOException { - return new TestingPlanner(bucketNum, partitionInfos); + return useExplicitSplits + ? TestingPlanner.fromSplits(explicitSplits) + : new TestingPlanner(bucketNum, partitionInfos); } @Override diff --git a/fluss-common/src/test/java/org/apache/fluss/lake/source/TestingPlanner.java b/fluss-common/src/test/java/org/apache/fluss/lake/source/TestingPlanner.java index 2ae27221ef3..3c3d2a65f1c 100644 --- a/fluss-common/src/test/java/org/apache/fluss/lake/source/TestingPlanner.java +++ b/fluss-common/src/test/java/org/apache/fluss/lake/source/TestingPlanner.java @@ -21,6 +21,7 @@ import java.io.IOException; import java.util.ArrayList; +import java.util.Collections; import java.util.List; /** A testing implementation of {@link Planner}. */ @@ -28,14 +29,34 @@ public class TestingPlanner implements Planner { private final int bucketNum; private final List partitionInfos; + private final List explicitSplits; + private final boolean useExplicitSplits; public TestingPlanner(int bucketNum, List partitionInfos) { this.bucketNum = bucketNum; this.partitionInfos = partitionInfos; + this.explicitSplits = Collections.emptyList(); + this.useExplicitSplits = false; + } + + private TestingPlanner(List explicitSplits) { + this.bucketNum = 0; + this.partitionInfos = Collections.emptyList(); + this.explicitSplits = new ArrayList<>(explicitSplits); + this.useExplicitSplits = true; + } + + /** Creates a planner that returns exactly the supplied splits. */ + public static TestingPlanner fromSplits(List splits) { + return new TestingPlanner(splits); } @Override public List plan() throws IOException { + if (useExplicitSplits) { + return new ArrayList<>(explicitSplits); + } + List splits = new ArrayList<>(); for (PartitionInfo partitionInfo : partitionInfos) { diff --git a/fluss-common/src/test/java/org/apache/fluss/lake/writer/WriterInitContextTest.java b/fluss-common/src/test/java/org/apache/fluss/lake/writer/WriterInitContextTest.java index f6618907602..ee460e3f72c 100644 --- a/fluss-common/src/test/java/org/apache/fluss/lake/writer/WriterInitContextTest.java +++ b/fluss-common/src/test/java/org/apache/fluss/lake/writer/WriterInitContextTest.java @@ -54,6 +54,11 @@ public String partition() { public TableInfo tableInfo() { return null; } + + @Override + public int bucketCount() { + throw new UnsupportedOperationException("not used in this test"); + } }; assertThat(context.splitIndex()).isEqualTo(WriterInitContext.UNKNOWN_SPLIT_INDEX); diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/FlinkConnectorOptions.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/FlinkConnectorOptions.java index 8f2ce7cbea0..bc7529fe23c 100644 --- a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/FlinkConnectorOptions.java +++ b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/FlinkConnectorOptions.java @@ -54,7 +54,10 @@ public class FlinkConnectorOptions { ConfigOptions.key("bucket.num") .intType() .noDefaultValue() - .withDescription("The number of buckets of a Fluss table."); + .withDescription( + "The target number of buckets for a Fluss table. " + + "For partitioned tables, this value applies to newly created " + + "partitions; existing partitions retain their original bucket count."); public static final ConfigOption BUCKET_KEY = ConfigOptions.key("bucket.key") @@ -268,8 +271,8 @@ public class FlinkConnectorOptions { public static final List ALTER_DISALLOW_OPTIONS = Arrays.asList( AUTO_INCREMENT_FIELDS.key(), - BUCKET_NUMBER.key(), BUCKET_KEY.key(), + BUCKET_NUMBER.key(), BOOTSTRAP_SERVERS.key()); // ------------------------------------------------------------------------------------------- diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/action/orphan/OrphanCleanUtils.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/action/orphan/OrphanCleanUtils.java index 24381ab7527..624e62055cc 100644 --- a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/action/orphan/OrphanCleanUtils.java +++ b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/action/orphan/OrphanCleanUtils.java @@ -64,7 +64,7 @@ public static PhysicalTablePath physicalPath( */ public static List enumerateBuckets( TableInfo tableInfo, @Nullable PartitionInfo partitionInfo) { - int n = tableInfo.getNumBuckets(); + int n = PartitionInfo.bucketCountOrDefault(partitionInfo, tableInfo.getNumBuckets()); List buckets = new ArrayList(n); long tableId = tableInfo.getTableId(); for (int b = 0; b < n; b++) { diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/lake/LakeSplitGenerator.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/lake/LakeSplitGenerator.java index f26ebddfdd6..b13dc4e8f4e 100644 --- a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/lake/LakeSplitGenerator.java +++ b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/lake/LakeSplitGenerator.java @@ -109,14 +109,8 @@ public List generateHybridLakeFlussSplits() throws Exception { Map tableBucketsOffset = lakeSnapshotInfo.getTableBucketsOffset(); if (isPartitioned) { Set partitionInfos = listPartitionSupplier.get(); - Map partitionNameById = - partitionInfos.stream() - .collect( - Collectors.toMap( - PartitionInfo::getPartitionId, - PartitionInfo::getPartitionName)); return generatePartitionTableSplit( - lakeSplits, isLogTable, tableBucketsOffset, partitionNameById); + lakeSplits, isLogTable, tableBucketsOffset, partitionInfos); } else { Map> nonPartitionLakeSplits = lakeSplits.isEmpty() ? null : lakeSplits.values().iterator().next(); @@ -144,14 +138,14 @@ private List generatePartitionTableSplit( Map>> lakeSplits, boolean isLogTable, Map tableBucketSnapshotLogOffset, - Map partitionNameById) { + Set partitionInfos) { List splits = new ArrayList<>(); - Map flussPartitionIdByName = - partitionNameById.entrySet().stream() + Map flussPartitionByName = + partitionInfos.stream() .collect( Collectors.toMap( - Map.Entry::getValue, - Map.Entry::getKey, + PartitionInfo::getPartitionName, + partitionInfo -> partitionInfo, (existing, replacement) -> existing, LinkedHashMap::new)); long lakeSplitPartitionId = -1L; @@ -161,21 +155,23 @@ private List generatePartitionTableSplit( lakeSplits.entrySet()) { String partitionName = lakeSplitEntry.getKey(); Map> lakeSplitsOfPartition = lakeSplitEntry.getValue(); - Long partitionId = flussPartitionIdByName.remove(partitionName); - if (partitionId != null) { + PartitionInfo flussPartition = flussPartitionByName.remove(partitionName); + if (flussPartition != null) { // mean the partition also exist in fluss partition + int partitionBucketCount = flussPartition.getBucketCount(); Map bucketEndOffset = stoppingOffsetInitializer.getBucketOffsets( partitionName, - IntStream.range(0, bucketCount) + IntStream.range(0, partitionBucketCount) .boxed() .collect(Collectors.toList()), bucketOffsetsRetriever); splits.addAll( generateSplit( lakeSplitsOfPartition, - partitionId, + flussPartition.getPartitionId(), partitionName, + partitionBucketCount, isLogTable, tableBucketSnapshotLogOffset, bucketEndOffset)); @@ -196,19 +192,22 @@ private List generatePartitionTableSplit( } // iterate remain fluss splits - for (Map.Entry partitionIdByNameEntry : flussPartitionIdByName.entrySet()) { - String partitionName = partitionIdByNameEntry.getKey(); - Long partitionId = partitionIdByNameEntry.getValue(); + for (PartitionInfo flussPartition : flussPartitionByName.values()) { + String partitionName = flussPartition.getPartitionName(); + int partitionBucketCount = flussPartition.getBucketCount(); Map bucketEndOffset = stoppingOffsetInitializer.getBucketOffsets( partitionName, - IntStream.range(0, bucketCount).boxed().collect(Collectors.toList()), + IntStream.range(0, partitionBucketCount) + .boxed() + .collect(Collectors.toList()), bucketOffsetsRetriever); splits.addAll( generateSplit( null, - partitionId, + flussPartition.getPartitionId(), partitionName, + partitionBucketCount, isLogTable, // pass empty map since we won't read lake splits Collections.emptyMap(), @@ -221,6 +220,7 @@ private List generateSplit( @Nullable Map> lakeSplits, @Nullable Long partitionId, @Nullable String partitionName, + int numBuckets, boolean isLogTable, Map tableBucketSnapshotLogOffset, Map bucketEndOffset) { @@ -229,7 +229,7 @@ private List generateSplit( if (lakeSplits != null) { splits.addAll(toLakeSnapshotSplits(lakeSplits, partitionName, partitionId)); } - for (int bucket = 0; bucket < bucketCount; bucket++) { + for (int bucket = 0; bucket < numBuckets; bucket++) { TableBucket tableBucket = new TableBucket(tableInfo.getTableId(), partitionId, bucket); Long snapshotLogOffset = tableBucketSnapshotLogOffset.get(tableBucket); @@ -259,7 +259,28 @@ private List generateSplit( } } else { // it's primary key table - for (int bucket = 0; bucket < bucketCount; bucket++) { + if (lakeSplits != null) { + // Pairing below only visits buckets in [0, numBuckets), so any lake split with a + // bucket id outside that range would otherwise be lost from the union read. + List outOfRangeBuckets = + lakeSplits.keySet().stream() + .filter(bucket -> bucket >= numBuckets) + .sorted() + .collect(Collectors.toList()); + if (!outOfRangeBuckets.isEmpty()) { + throw new IllegalStateException( + String.format( + "Lake snapshot of table %s partition %s contains buckets %s " + + "outside the enumerated range [0, %d); refusing to " + + "generate union-read splits that would silently " + + "drop lake data.", + tableInfo.getTablePath(), + partitionName, + outOfRangeBuckets, + numBuckets)); + } + } + for (int bucket = 0; bucket < numBuckets; bucket++) { TableBucket tableBucket = new TableBucket(tableInfo.getTableId(), partitionId, bucket); Long snapshotLogOffset = tableBucketSnapshotLogOffset.get(tableBucket); @@ -323,6 +344,12 @@ private List generateNoPartitionedTableSplit( IntStream.range(0, bucketCount).boxed().collect(Collectors.toList()), bucketOffsetsRetriever); return generateSplit( - lakeSplits, null, null, isLogTable, tableBucketSnapshotLogOffset, bucketEndOffset); + lakeSplits, + null, + null, + bucketCount, + isLogTable, + tableBucketSnapshotLogOffset, + bucketEndOffset); } } diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/sink/undo/RecoveryOffsetManager.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/sink/undo/RecoveryOffsetManager.java index 38e22957c81..ab39682a0f3 100644 --- a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/sink/undo/RecoveryOffsetManager.java +++ b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/sink/undo/RecoveryOffsetManager.java @@ -223,11 +223,11 @@ public RecoveryDecision determineRecoveryStrategy( producerId); RecoveryStateKind stateKind = classifyRecoveredState(recoveredState); - Map partitionNames = getPartitionNameMap(); + Map partitionInfos = getPartitionInfoMap(); Map recoveryOffsets = stateKind == RecoveryStateKind.NO_FLINK_STATE ? getProducerOffsets() - : mergeCheckpointState(recoveredState, partitionNames); + : mergeCheckpointState(recoveredState, partitionInfos); LOG.info( "Recovery offsets for subtask {} (source={}): {}", @@ -236,7 +236,7 @@ public RecoveryDecision determineRecoveryStrategy( recoveryOffsets); Set allBuckets = getAllBuckets(); - Set filteredBuckets = filterBucketsBySharding(allBuckets, partitionNames); + Set filteredBuckets = filterBucketsBySharding(allBuckets, partitionInfos); LOG.info( "Subtask {}: filteredBuckets={}, recoveryOffsets={}", @@ -245,7 +245,7 @@ public RecoveryDecision determineRecoveryStrategy( recoveryOffsets); Map currentOffsets = - fetchCurrentOffsets(filteredBuckets, partitionNames); + fetchCurrentOffsets(filteredBuckets, partitionInfos); LOG.info("Subtask {}: currentOffsets={}", subtaskIndex, currentOffsets); @@ -361,7 +361,7 @@ private RecoveryStateKind classifyRecoveredState( } private Map mergeCheckpointState( - Collection states, Map partitionNames) { + Collection states, Map partitionInfos) { Map merged = new HashMap<>(); for (WriterState state : states) { if (state.getStateFormat() == WriterState.StateFormat.V2_COMPLETE @@ -375,7 +375,7 @@ private Map mergeCheckpointState( TableBucket bucket = entry.getKey(); validateTableId(bucket); validateBaselineOffset(bucket, entry.getValue()); - if (!isLiveStateBucket(bucket, partitionNames)) { + if (!isLiveStateBucket(bucket, partitionInfos)) { continue; } putMergedOffset(merged, bucket, entry.getValue()); @@ -404,14 +404,14 @@ private void validateTableId(TableBucket bucket) { } } - private boolean isLiveStateBucket(TableBucket bucket, Map partitionNames) { + private boolean isLiveStateBucket(TableBucket bucket, Map partitionInfos) { Long partitionId = bucket.getPartitionId(); if (isPartitioned) { if (partitionId == null) { throw new IllegalStateException( "State bucket " + bucket + " has no partition ID for a partitioned table."); } - return partitionNames.containsKey(partitionId); + return partitionInfos.containsKey(partitionId); } if (partitionId != null) { throw new IllegalStateException( @@ -489,7 +489,8 @@ private Set getAllBuckets() throws Exception { Set buckets = new HashSet<>(); if (isPartitioned) { for (PartitionInfo partition : getPartitionInfos()) { - for (int bucketId = 0; bucketId < numBuckets; bucketId++) { + int partitionBucketCount = partition.getBucketCount(); + for (int bucketId = 0; bucketId < partitionBucketCount; bucketId++) { buckets.add(new TableBucket(tableId, partition.getPartitionId(), bucketId)); } } @@ -504,10 +505,10 @@ private Set getAllBuckets() throws Exception { // ==================== Step 3: Filter by Sharding ==================== private Set filterBucketsBySharding( - Set buckets, Map partitionNames) { + Set buckets, Map partitionInfos) { Set filtered = new HashSet<>(); for (TableBucket bucket : buckets) { - if (isAssignedToSubtask(bucket, partitionNames)) { + if (isAssignedToSubtask(bucket, partitionInfos)) { filtered.add(bucket); } } @@ -518,24 +519,28 @@ private Set filterBucketsBySharding( * Determines if a bucket is assigned to the current subtask. * *

Uses {@link ChannelComputer#shouldCombinePartitionInSharding} and {@link - * ChannelComputer#select} to ensure consistent sharding logic with {@link - * org.apache.fluss.flink.sink.FlinkRowDataChannelComputer}. + * ChannelComputer#select} to keep the sharding logic aligned with {@link + * org.apache.fluss.flink.sink.FlinkRowDataChannelComputer}. A partition that kept its own + * bucket layout across an ALTER bucket.num is sharded by that partition's actual bucket count, + * not by the table-level one. * - *

For partitioned tables, if the partition has been deleted (partitionName not found in - * partitionNames map), the bucket is considered not assigned to any subtask and will be + *

For partitioned tables, if the partition has been deleted (partition not found in + * partitionInfos map), the bucket is considered not assigned to any subtask and will be * skipped. * * @param bucket the bucket to check - * @param partitionNames map of partition ID to partition name + * @param partitionInfos map of partition ID to partition info * @return true if the bucket is assigned to this subtask, false if not assigned or partition * deleted */ - private boolean isAssignedToSubtask(TableBucket bucket, Map partitionNames) { - // For partitioned table bucket, get partition name first + private boolean isAssignedToSubtask( + TableBucket bucket, Map partitionInfos) { + // For partitioned table bucket, get partition name and its own bucket count first String partitionName = null; + int shardingBucketCount = numBuckets; if (bucket.getPartitionId() != null) { - partitionName = partitionNames.get(bucket.getPartitionId()); - if (partitionName == null) { + PartitionInfo partitionInfo = partitionInfos.get(bucket.getPartitionId()); + if (partitionInfo == null) { // Partition has been deleted, skip this bucket LOG.debug( "Partition {} not found (deleted?), skipping bucket {}", @@ -543,12 +548,14 @@ private boolean isAssignedToSubtask(TableBucket bucket, Map partit bucket); return false; } + partitionName = partitionInfo.getPartitionName(); + shardingBucketCount = partitionInfo.getBucketCount(); } // Use shared logic to determine sharding strategy and compute channel int channel; if (ChannelComputer.shouldCombinePartitionInSharding( - isPartitioned, numBuckets, parallelism)) { + isPartitioned, shardingBucketCount, parallelism)) { // When shouldCombinePartitionInSharding is true, partitionName is guaranteed non-null // because: 1) isPartitioned=true means bucket has partitionId // 2) deleted partitions already returned false above @@ -562,7 +569,7 @@ private boolean isAssignedToSubtask(TableBucket bucket, Map partit // ==================== Step 4: Fetch Current Offsets ==================== private Map fetchCurrentOffsets( - Set buckets, Map partitionNames) throws Exception { + Set buckets, Map partitionInfos) throws Exception { Map offsets = new HashMap<>(); // Group buckets by partition @@ -586,12 +593,12 @@ private Map fetchCurrentOffsets( // Fetch partitioned buckets for (Map.Entry> entry : byPartition.entrySet()) { Long partitionId = entry.getKey(); - String partitionName = partitionNames.get(partitionId); - if (partitionName == null) { + PartitionInfo partitionInfo = partitionInfos.get(partitionId); + if (partitionInfo == null) { throw new IllegalStateException( "Partition " + partitionId + " not found in partition info cache"); } - fetchBucketOffsets(partitionName, entry.getValue(), offsets); + fetchBucketOffsets(partitionInfo.getPartitionName(), entry.getValue(), offsets); } return offsets; @@ -607,15 +614,15 @@ private List getPartitionInfos() throws Exception { return cachedPartitionInfos; } - private Map getPartitionNameMap() throws Exception { + private Map getPartitionInfoMap() throws Exception { if (!isPartitioned) { return new HashMap<>(); } - Map nameMap = new HashMap<>(); + Map infoMap = new HashMap<>(); for (PartitionInfo partition : getPartitionInfos()) { - nameMap.put(partition.getPartitionId(), partition.getPartitionName()); + infoMap.put(partition.getPartitionId(), partition); } - return nameMap; + return infoMap; } // ==================== Offset Fetching Helpers ==================== @@ -625,10 +632,13 @@ private Map fetchAllBucketOffsets() throws Exception { if (isPartitioned) { for (PartitionInfo partition : getPartitionInfos()) { fetchPartitionOffsets( - partition.getPartitionName(), partition.getPartitionId(), offsets); + partition.getPartitionName(), + partition.getPartitionId(), + partition.getBucketCount(), + offsets); } } else { - fetchPartitionOffsets(null, null, offsets); + fetchPartitionOffsets(null, null, numBuckets, offsets); } return offsets; } @@ -636,10 +646,11 @@ private Map fetchAllBucketOffsets() throws Exception { private void fetchPartitionOffsets( @Nullable String partitionName, @Nullable Long partitionId, + int bucketCount, Map offsets) throws Exception { - List bucketIds = new ArrayList<>(numBuckets); - for (int i = 0; i < numBuckets; i++) { + List bucketIds = new ArrayList<>(bucketCount); + for (int i = 0; i < bucketCount; i++) { bucketIds.add(i); } ListOffsetsResult result = listOffsets(partitionName, bucketIds); diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/enumerator/FlinkSourceEnumerator.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/enumerator/FlinkSourceEnumerator.java index 0dd5b1b45f1..54ae770c24a 100644 --- a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/enumerator/FlinkSourceEnumerator.java +++ b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/enumerator/FlinkSourceEnumerator.java @@ -772,22 +772,25 @@ private List generateFlussOnlyBatchSplits( Set partitionInfos = listPartitions(); List splits = new ArrayList<>(); for (PartitionInfo partitionInfo : partitionInfos) { - splits.addAll( - buildKvBatchSplits( - partitionInfo.getPartitionId(), - partitionInfo.getPartitionName())); + splits.addAll(buildKvBatchSplits(partitionInfo)); } return splits; } - return buildKvBatchSplits(null, null); + return buildKvBatchSplits(null); } return flussOnlyBatchSplitGenerator.generate(); } - private List buildKvBatchSplits( - @Nullable Long partitionId, @Nullable String partitionName) { + private List buildKvBatchSplits(@Nullable PartitionInfo partitionInfo) { + // A partition keeps the bucket count it was created with, so its buckets must be + // enumerated by that count; the table-level count only applies to a non-partitioned + // table, whose single bucket layout is the table's own. + int bucketCount = + partitionInfo != null ? partitionInfo.getBucketCount() : tableInfo.getNumBuckets(); + Long partitionId = partitionInfo != null ? partitionInfo.getPartitionId() : null; + String partitionName = partitionInfo != null ? partitionInfo.getPartitionName() : null; List splits = new ArrayList<>(); - for (int bucketId = 0; bucketId < tableInfo.getNumBuckets(); bucketId++) { + for (int bucketId = 0; bucketId < bucketCount; bucketId++) { TableBucket tb = new TableBucket(tableInfo.getTableId(), partitionId, bucketId); if (ignoreTableBucket(tb)) { continue; @@ -834,7 +837,7 @@ private List initNonPartitionedSplits() { if (hasPrimaryKey && startingOffsetsInitializer instanceof SnapshotOffsetsInitializer) { return getSnapshotAndLogSplits(getLatestKvSnapshotsAndRegister(null), null); } else { - return getLogSplit(null, null); + return getNonPartitionedLogSplit(); } } @@ -967,7 +970,12 @@ private PartitionChange getPartitionChange( Set fetchedPartitionInfos, boolean initialDiscovery) { final Set allNewPartitions = fetchedPartitionInfos.stream() - .map(p -> new Partition(p.getPartitionId(), p.getPartitionName())) + .map( + p -> + new Partition( + p.getPartitionId(), + p.getPartitionName(), + p.getBucketCount())) .collect(Collectors.toSet()); final Set removedPartitions = new HashSet<>(); @@ -1058,7 +1066,8 @@ private List initLogTablePartitionSplits( getLogSplit( partition.getPartitionId(), partition.getPartitionName(), - effectiveOffsetsInitializer)); + effectiveOffsetsInitializer, + partition.getBucketCount())); } return splits; } @@ -1206,19 +1215,19 @@ private List getSnapshotAndLogSplits( return splits; } - private List getLogSplit( - @Nullable Long partitionId, @Nullable String partitionName) { - return getLogSplit(partitionId, partitionName, startingOffsetsInitializer); + private List getNonPartitionedLogSplit() { + return getLogSplit(null, null, startingOffsetsInitializer, tableInfo.getNumBuckets()); } private List getLogSplit( @Nullable Long partitionId, @Nullable String partitionName, - OffsetsInitializer effectiveStartingOffsetsInitializer) { + OffsetsInitializer effectiveStartingOffsetsInitializer, + int bucketCount) { // always assume the bucket is from 0 to bucket num List splits = new ArrayList<>(); List bucketsNeedInitOffset = new ArrayList<>(); - for (int bucketId = 0; bucketId < tableInfo.getNumBuckets(); bucketId++) { + for (int bucketId = 0; bucketId < bucketCount; bucketId++) { TableBucket tableBucket = new TableBucket(tableInfo.getTableId(), partitionId, bucketId); if (ignoreTableBucket(tableBucket)) { @@ -1823,12 +1832,27 @@ public boolean isEmpty() { /** A container class to hold the partition id and partition name. */ private static class Partition { + /** Marks comparison-only instances that do not carry a bucket count. */ + private static final int NO_BUCKET_COUNT = -1; + final long partitionId; final String partitionName; + /** + * The actual bucket count of this partition, already resolved by {@link PartitionInfo}. It + * is {@link #NO_BUCKET_COUNT} only for instances created for diff comparison or removal + * handling, which never generate splits. + */ + final int bucketCount; + Partition(long partitionId, String partitionName) { + this(partitionId, partitionName, NO_BUCKET_COUNT); + } + + Partition(long partitionId, String partitionName, int bucketCount) { this.partitionId = partitionId; this.partitionName = partitionName; + this.bucketCount = bucketCount; } public long getPartitionId() { @@ -1839,6 +1863,16 @@ public String getPartitionName() { return partitionName; } + public int getBucketCount() { + checkState( + bucketCount != NO_BUCKET_COUNT, + "Partition %s (id %s) does not carry a bucket count; comparison-only " + + "instances must not be used to generate splits.", + partitionName, + partitionId); + return bucketCount; + } + @Override public boolean equals(Object o) { if (o == null || getClass() != o.getClass()) { diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/enumerator/FlussOnlyBatchSplitGenerator.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/enumerator/FlussOnlyBatchSplitGenerator.java index 4036f047a52..84992e5d31d 100644 --- a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/enumerator/FlussOnlyBatchSplitGenerator.java +++ b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/enumerator/FlussOnlyBatchSplitGenerator.java @@ -86,7 +86,7 @@ List generate() { } else { return hasPrimaryKey ? getBatchSnapshotAndLogSplits(kvSnapshotsRetriever.get(null), null) - : getLogSplits(null, null); + : getNonPartitionedLogSplits(); } } @@ -105,7 +105,11 @@ private List generatePrimaryKeyTableSplits( private List generateLogTableSplits(Collection partitions) { List splits = new ArrayList<>(); for (PartitionInfo partition : partitions) { - splits.addAll(getLogSplits(partition.getPartitionId(), partition.getPartitionName())); + splits.addAll( + getLogSplits( + partition.getPartitionId(), + partition.getPartitionName(), + partition.getBucketCount())); } return splits; } @@ -163,11 +167,15 @@ private List getBatchSnapshotAndLogSplits( return splits; } + private List getNonPartitionedLogSplits() { + return getLogSplits(null, null, tableInfo.getNumBuckets()); + } + private List getLogSplits( - @Nullable Long partitionId, @Nullable String partitionName) { + @Nullable Long partitionId, @Nullable String partitionName, int bucketCount) { List splits = new ArrayList<>(); List bucketsNeedInitOffset = new ArrayList<>(); - for (int bucketId = 0; bucketId < tableInfo.getNumBuckets(); bucketId++) { + for (int bucketId = 0; bucketId < bucketCount; bucketId++) { TableBucket tableBucket = new TableBucket(tableInfo.getTableId(), partitionId, bucketId); if (!tableBucketSkipper.test(tableBucket)) { diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/tiering/source/TieringSplitReader.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/tiering/source/TieringSplitReader.java index 60b911b6f74..fc30e0ff657 100644 --- a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/tiering/source/TieringSplitReader.java +++ b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/tiering/source/TieringSplitReader.java @@ -19,12 +19,14 @@ import org.apache.fluss.annotation.VisibleForTesting; import org.apache.fluss.client.Connection; +import org.apache.fluss.client.admin.Admin; import org.apache.fluss.client.table.Table; import org.apache.fluss.client.table.scanner.ScanRecord; import org.apache.fluss.client.table.scanner.log.ArrowScanRecords; import org.apache.fluss.client.table.scanner.log.LogScanner; import org.apache.fluss.client.table.scanner.log.LogScannerImpl; import org.apache.fluss.client.table.scanner.log.ScanRecords; +import org.apache.fluss.exception.FlussRuntimeException; import org.apache.fluss.flink.source.reader.BoundedSplitReader; import org.apache.fluss.flink.source.reader.RecordAndPos; import org.apache.fluss.flink.tiering.source.metrics.TieringMetrics; @@ -37,6 +39,7 @@ import org.apache.fluss.lake.writer.SupportsRecordBatchWrite; import org.apache.fluss.metadata.DataLakeFormat; import org.apache.fluss.metadata.LogFormat; +import org.apache.fluss.metadata.PartitionInfo; import org.apache.fluss.metadata.TableBucket; import org.apache.fluss.metadata.TableInfo; import org.apache.fluss.metadata.TablePath; @@ -116,6 +119,9 @@ public class TieringSplitReader private final Map currentTableSplitsByBucket; private final Map currentTableStoppingOffsets; + // partition id -> actual bucket count for the current table's partitions + private final Map currentTablePartitionBucketCounts = new HashMap<>(); + private final Map currentTableTieredOffsetAndTimestamp; private final Set currentEmptySplits; @@ -334,6 +340,22 @@ private Table getOrMoveToTable(TieringSplit split) { currentTableInfo.getTableId(), tablePath, split.getTableBucket().getTableId()); + // Snapshot each partition's actual bucket count so lake writers can stamp per-partition + // bucket layouts correctly after an ALTER bucket.num. + if (currentTableInfo.isPartitioned()) { + try { + // the admin is a shared per-connection instance, so it must not be closed here + Admin admin = connection.getAdmin(); + for (PartitionInfo partitionInfo : + admin.listPartitionInfos(tablePath, true).get()) { + currentTablePartitionBucketCounts.put( + partitionInfo.getPartitionId(), partitionInfo.getBucketCount()); + } + } catch (Exception e) { + throw new FlussRuntimeException( + "Failed to list partition infos for table " + tablePath, e); + } + } LOG.info("Start to tier table {} with table id {}.", currentTablePath, currentTableId); } return currentTable; @@ -620,6 +642,10 @@ private LakeWriter getOrCreateLakeWriter( throws IOException { LakeWriter lakeWriter = lakeWriters.get(bucket); if (lakeWriter == null) { + Integer partitionBucketCount = + bucket.getPartitionId() != null + ? currentTablePartitionBucketCounts.get(bucket.getPartitionId()) + : null; lakeWriter = lakeTieringFactory.createLakeWriter( new TieringWriterInitContext( @@ -629,6 +655,7 @@ private LakeWriter getOrCreateLakeWriter( currentTable.getTableInfo(), splitIndex, tieringRoundTimestamp, + partitionBucketCount, ioTmpDirs)); lakeWriters.put(bucket, lakeWriter); } @@ -782,6 +809,7 @@ private void finishCurrentTable() throws IOException { currentTableStoppingOffsets.clear(); currentTableTieredOffsetAndTimestamp.clear(); currentTableSplitsByBucket.clear(); + currentTablePartitionBucketCounts.clear(); } /** diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/tiering/source/TieringWriterInitContext.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/tiering/source/TieringWriterInitContext.java index f67b44176be..658467f5371 100644 --- a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/tiering/source/TieringWriterInitContext.java +++ b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/tiering/source/TieringWriterInitContext.java @@ -24,6 +24,8 @@ import javax.annotation.Nullable; +import static org.apache.fluss.utils.Preconditions.checkNotNull; + /** The implementation of {@link WriterInitContext}. */ public class TieringWriterInitContext implements WriterInitContext { @@ -33,40 +35,9 @@ public class TieringWriterInitContext implements WriterInitContext { private final TableInfo tableInfo; private final int splitIndex; private final long tieringRoundTimestamp; + private final int bucketCount; @Nullable private final String[] ioTmpDirs; - public TieringWriterInitContext( - TablePath tablePath, - TableBucket tableBucket, - @Nullable String partition, - TableInfo tableInfo) { - this( - tablePath, - tableBucket, - partition, - tableInfo, - UNKNOWN_SPLIT_INDEX, - UNKNOWN_TIERING_ROUND_TIMESTAMP, - (String[]) null); - } - - public TieringWriterInitContext( - TablePath tablePath, - TableBucket tableBucket, - @Nullable String partition, - TableInfo tableInfo, - int splitIndex, - long tieringRoundTimestamp) { - this( - tablePath, - tableBucket, - partition, - tableInfo, - splitIndex, - tieringRoundTimestamp, - (String[]) null); - } - public TieringWriterInitContext( TablePath tablePath, TableBucket tableBucket, @@ -74,6 +45,7 @@ public TieringWriterInitContext( TableInfo tableInfo, int splitIndex, long tieringRoundTimestamp, + @Nullable Integer bucketCount, @Nullable String[] ioTmpDirs) { this.tablePath = tablePath; this.tableBucket = tableBucket; @@ -82,6 +54,19 @@ public TieringWriterInitContext( this.splitIndex = splitIndex; this.tieringRoundTimestamp = tieringRoundTimestamp; this.ioTmpDirs = ioTmpDirs; + if (tableBucket.getPartitionId() == null) { + this.bucketCount = tableInfo.getNumBuckets(); + } else { + // Writing with a wrong bucket count would silently corrupt the lake table's bucket + // layout metadata, so a missing per-partition count must fail here. + this.bucketCount = + checkNotNull( + bucketCount, + "No actual bucket count known for partition %s (id %s) of table %s.", + partition, + tableBucket.getPartitionId(), + tablePath); + } } @Override @@ -120,4 +105,9 @@ public long tieringRoundTimestamp() { public String[] ioTmpDirs() { return ioTmpDirs; } + + @Override + public int bucketCount() { + return bucketCount; + } } diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/tiering/source/enumerator/TieringSourceEnumerator.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/tiering/source/enumerator/TieringSourceEnumerator.java index eb27397c9c6..f1bf98527af 100644 --- a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/tiering/source/enumerator/TieringSourceEnumerator.java +++ b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/tiering/source/enumerator/TieringSourceEnumerator.java @@ -158,7 +158,7 @@ public void start() { this.coordinatorGateway = GatewayClientProxy.createGatewayProxy( metadataUpdater::getCoordinatorServer, rpcClient, CoordinatorGateway.class); - this.splitGenerator = new TieringSplitGenerator(flussAdmin, metadataUpdater); + this.splitGenerator = new TieringSplitGenerator(flussAdmin); LOG.info("Starting register Tiering Service to Fluss Coordinator..."); try { diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/tiering/source/split/TieringSplitGenerator.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/tiering/source/split/TieringSplitGenerator.java index 5dd0bad2374..5eea8b6a377 100644 --- a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/tiering/source/split/TieringSplitGenerator.java +++ b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/tiering/source/split/TieringSplitGenerator.java @@ -22,10 +22,8 @@ import org.apache.fluss.client.initializer.OffsetsInitializer.BucketOffsetsRetriever; import org.apache.fluss.client.metadata.KvSnapshots; import org.apache.fluss.client.metadata.LakeSnapshot; -import org.apache.fluss.client.metadata.MetadataUpdater; import org.apache.fluss.exception.LakeTableSnapshotNotExistException; import org.apache.fluss.metadata.PartitionInfo; -import org.apache.fluss.metadata.PhysicalTablePath; import org.apache.fluss.metadata.TableBucket; import org.apache.fluss.metadata.TableInfo; import org.apache.fluss.metadata.TablePath; @@ -55,11 +53,9 @@ public class TieringSplitGenerator { private static final Logger LOG = LoggerFactory.getLogger(TieringSplitGenerator.class); private final Admin flussAdmin; - private final MetadataUpdater metadataUpdater; - public TieringSplitGenerator(Admin flussAdmin, MetadataUpdater metadataUpdater) { + public TieringSplitGenerator(Admin flussAdmin) { this.flussAdmin = flussAdmin; - this.metadataUpdater = metadataUpdater; } public List generateTableSplits(TableInfo tableInfo) throws Exception { @@ -91,31 +87,26 @@ public List generateTableSplits(TableInfo tableInfo) throws Except // partitioned table if (tableInfo.isPartitioned()) { List partitionInfos = - flussAdmin.listPartitionInfos(tableInfo.getTablePath()).get(); + flussAdmin.listPartitionInfos(tableInfo.getTablePath(), true).get(); Map partitionNameById = partitionInfos.stream() .collect( Collectors.toMap( PartitionInfo::getPartitionId, PartitionInfo::getPartitionName)); - if (tableInfo.getTableConfig().isHistoricalPartitionEnabled()) { - // The internal historical partition is intentionally omitted from - // listPartitionInfos(), but tiering must consume it to synchronize historical - // writes to the lake table. Resolve it explicitly and include it in the splits. - PhysicalTablePath historicalPath = - PhysicalTablePath.of(tablePath, HISTORICAL_PARTITION_VALUE); - // Partition metadata is decoded using the tableId-to-path mapping already present - // in the Cluster, so initialize the table metadata before requesting the internal - // partition directly. - metadataUpdater.checkAndUpdateTableMetadata(Collections.singleton(tablePath)); - metadataUpdater.checkAndUpdatePartitionMetadata(historicalPath); - partitionNameById.put( - metadataUpdater.getPartitionIdOrElseThrow(historicalPath), - HISTORICAL_PARTITION_VALUE); - } + Map bucketCountById = + partitionInfos.stream() + .collect( + Collectors.toMap( + PartitionInfo::getPartitionId, + PartitionInfo::getBucketCount)); return generatePartitionTableSplit( - tableInfo, partitionNameById, bucketOffsetsRetriever, lakeSnapshotInfo); + tableInfo, + partitionNameById, + bucketCountById, + bucketOffsetsRetriever, + lakeSnapshotInfo); } else { // non-partitioned table return generateNonPartitionedTableSplit( @@ -127,6 +118,7 @@ public List generateTableSplits(TableInfo tableInfo) throws Except private List generatePartitionTableSplit( TableInfo tableInfo, Map partitionNameById, + Map bucketCountById, BucketOffsetsRetriever bucketOffsetsRetriever, @Nullable LakeSnapshot lakeSnapshotInfo) { List splits = new ArrayList<>(); @@ -134,10 +126,11 @@ private List generatePartitionTableSplit( long partitionId = partitionNameByIdEntry.getKey(); String partitionName = partitionNameByIdEntry.getValue(); boolean historicalPartition = HISTORICAL_PARTITION_VALUE.equals(partitionName); + int partitionBucketCount = bucketCountById.get(partitionId); Map latestBucketsOffset = bucketOffsetsRetriever.latestOffsets( partitionName, - IntStream.range(0, tableInfo.getNumBuckets()) + IntStream.range(0, partitionBucketCount) .boxed() .collect(Collectors.toList())); KvSnapshots latestKvSnapshots = null; @@ -172,6 +165,7 @@ private List generatePartitionTableSplit( tableInfo, partitionId, partitionName, + partitionBucketCount, lakeSnapshotInfo, latestKvSnapshots, latestBucketsOffset)); @@ -204,13 +198,20 @@ private List generateNonPartitionedTableSplit( } return generateTableSplit( - tableInfo, null, null, lakeSnapshotInfo, latestKvSnapshots, latestBucketsOffset); + tableInfo, + null, + null, + tableInfo.getNumBuckets(), + lakeSnapshotInfo, + latestKvSnapshots, + latestBucketsOffset); } private List generateTableSplit( TableInfo tableInfo, @Nullable Long partitionId, @Nullable String partitionName, + int numBuckets, @Nullable LakeSnapshot lakeSnapshotInfo, @Nullable KvSnapshots latestKvSnapshots, Map latestBucketsOffset) { @@ -219,7 +220,7 @@ private List generateTableSplit( if (tableInfo.hasPrimaryKey()) { // it's primary key table checkState(latestKvSnapshots != null); - for (int bucket = 0; bucket < tableInfo.getNumBuckets(); bucket++) { + for (int bucket = 0; bucket < numBuckets; bucket++) { TableBucket tableBucket = new TableBucket(tableInfo.getTableId(), partitionId, bucket); Long lastCommittedBucketOffset = @@ -249,7 +250,7 @@ private List generateTableSplit( } else { // it's log table - for (int bucket = 0; bucket < tableInfo.getNumBuckets(); bucket++) { + for (int bucket = 0; bucket < numBuckets; bucket++) { TableBucket tableBucket = new TableBucket(tableInfo.getTableId(), partitionId, bucket); Long lastCommittedOffset = diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/utils/FlinkConversions.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/utils/FlinkConversions.java index 9a6fb038c78..ec949ab5310 100644 --- a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/utils/FlinkConversions.java +++ b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/utils/FlinkConversions.java @@ -502,8 +502,11 @@ private static TableChange.ColumnPosition toFlussColumnPosition( } } - private static TableChange.SetOption convertSetOption( + private static TableChange convertSetOption( org.apache.flink.table.catalog.TableChange.SetOption flinkSetOption) { + if (BUCKET_NUMBER.key().equals(flinkSetOption.getKey())) { + return TableChange.modifyBucketCount(Integer.parseInt(flinkSetOption.getValue())); + } return TableChange.set(flinkSetOption.getKey(), flinkSetOption.getValue()); } diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/utils/PushdownUtils.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/utils/PushdownUtils.java index 1dde6bf3da3..ea4b47959a5 100644 --- a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/utils/PushdownUtils.java +++ b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/utils/PushdownUtils.java @@ -389,8 +389,6 @@ private static long countLogTable(Admin flussAdmin, TablePath tablePath) throws "The Fluss cluster doesn't support count(*) on primary key table yet. Please upgrade to newer version (≥ 0.9)."); } int bucketCount = tableInfo.getNumBuckets(); - Collection buckets = - IntStream.range(0, bucketCount).boxed().collect(Collectors.toList()); List partitionInfos; if (tableInfo.isPartitioned()) { partitionInfos = flussAdmin.listPartitionInfos(tablePath).get(); @@ -399,7 +397,7 @@ private static long countLogTable(Admin flussAdmin, TablePath tablePath) throws } List> countFutureList = - offsetLengthes(flussAdmin, tablePath, partitionInfos, buckets); + offsetLengthes(flussAdmin, tablePath, partitionInfos, bucketCount); // wait for all the response CompletableFuture.allOf(countFutureList.toArray(new CompletableFuture[0])).join(); long count = 0; @@ -413,10 +411,13 @@ private static List> offsetLengthes( Admin flussAdmin, TablePath tablePath, List partitionInfos, - Collection buckets) { + int tableBucketCount) { List> list = new ArrayList<>(); for (@Nullable PartitionInfo info : partitionInfos) { String partitionName = info != null ? info.getPartitionName() : null; + int partitionBucketCount = PartitionInfo.bucketCountOrDefault(info, tableBucketCount); + Collection buckets = + IntStream.range(0, partitionBucketCount).boxed().collect(Collectors.toList()); ListOffsetsResult earliestOffsets = listOffsets( flussAdmin, diff --git a/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/catalog/FlinkCatalogITCase.java b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/catalog/FlinkCatalogITCase.java index 86bf210dc69..23bec8268c4 100644 --- a/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/catalog/FlinkCatalogITCase.java +++ b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/catalog/FlinkCatalogITCase.java @@ -30,6 +30,7 @@ import org.apache.fluss.exception.InvalidTableException; import org.apache.fluss.flink.FlinkConnectorOptions; import org.apache.fluss.metadata.DataLakeFormat; +import org.apache.fluss.metadata.PartitionInfo; import org.apache.fluss.metadata.TableInfo; import org.apache.fluss.metadata.TablePath; import org.apache.fluss.row.encode.KvValueLayout; @@ -272,10 +273,22 @@ void testAlterTableConfig() throws Exception { .hasMessage( "Currently, auto partition is only supported for partitioned table, please set table property 'table.auto-partition.enabled' to false."); + // altering bucket.num is no longer blocked at the catalog layer; it is rejected by the + // server. This table is non-partitioned, so it fails with the non-partitioned rescale + // message (partitioned-table rescale is supported; non-partitioned is not yet). String unSupportedDml2 = "alter table test_alter_table_append_only set ('bucket.num' = '1000')"; assertThatThrownBy(() -> tEnv.executeSql(unSupportedDml2)) .rootCause() + .isInstanceOf(InvalidAlterTableException.class) + .hasMessageContaining("Cannot alter 'bucket.num' on non-partitioned table") + .hasMessageContaining("not yet supported"); + + assertThatThrownBy( + () -> + tEnv.executeSql( + "alter table test_alter_table_append_only reset ('bucket.num')")) + .rootCause() .isInstanceOf(CatalogException.class) .hasMessage("The option 'bucket.num' is not supported to alter yet."); @@ -301,6 +314,47 @@ void testAlterTableConfig() throws Exception { .hasMessage("The option 'auto-increment.fields' is not supported to alter yet."); } + @Test + void testAlterPartitionedTableBucketCount() throws Exception { + String tableName = "test_alter_partitioned_table_bucket_count"; + ObjectPath objectPath = new ObjectPath(DEFAULT_DB, tableName); + TablePath tablePath = TablePath.of(DEFAULT_DB, tableName); + tEnv.executeSql( + "create table " + + tableName + + " (a int, pt string) partitioned by (pt) with ('bucket.num' = '2')"); + tEnv.executeSql("alter table " + tableName + " add partition (pt = 'old')"); + + CatalogTable table = (CatalogTable) catalog.getTable(objectPath); + assertThat(table.getOptions()).containsEntry(BUCKET_NUMBER.key(), "2"); + + try (Connection conn = + ConnectionFactory.createConnection(FLUSS_CLUSTER_EXTENSION.getClientConfig())) { + Admin admin = conn.getAdmin(); + assertThat(admin.listPartitionInfos(tablePath).get()) + .singleElement() + .satisfies(partition -> assertThat(partition.getBucketCount()).isEqualTo(2)); + + tEnv.executeSql("alter table " + tableName + " set ('bucket.num' = '4')"); + + table = (CatalogTable) catalog.getTable(objectPath); + assertThat(table.getOptions()).containsEntry(BUCKET_NUMBER.key(), "4"); + + tEnv.executeSql("alter table " + tableName + " add partition (pt = 'new')"); + + Map bucketCountByPartition = + admin.listPartitionInfos(tablePath).get().stream() + .collect( + Collectors.toMap( + PartitionInfo::getPartitionName, + PartitionInfo::getBucketCount)); + assertThat(bucketCountByPartition) + .hasSize(2) + .containsEntry("old", 2) + .containsEntry("new", 4); + } + } + @Test void testAlterTableSchema() throws Exception { ObjectPath objectPath = new ObjectPath(DEFAULT_DB, "append_only_table"); diff --git a/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/lake/LakeSplitGeneratorTest.java b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/lake/LakeSplitGeneratorTest.java new file mode 100644 index 00000000000..c95b5424522 --- /dev/null +++ b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/lake/LakeSplitGeneratorTest.java @@ -0,0 +1,173 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.fluss.flink.lake; + +import org.apache.fluss.client.initializer.OffsetsInitializer; +import org.apache.fluss.client.metadata.LakeSnapshot; +import org.apache.fluss.flink.lake.split.LakeSnapshotAndFlussLogSplit; +import org.apache.fluss.flink.sink.testutils.TestAdminAdapter; +import org.apache.fluss.flink.source.split.SourceSplitBase; +import org.apache.fluss.lake.source.LakeSource; +import org.apache.fluss.lake.source.LakeSplit; +import org.apache.fluss.lake.source.TestingLakeSource; +import org.apache.fluss.lake.source.TestingLakeSplit; +import org.apache.fluss.metadata.PartitionInfo; +import org.apache.fluss.metadata.ResolvedPartitionSpec; +import org.apache.fluss.metadata.Schema; +import org.apache.fluss.metadata.TableDescriptor; +import org.apache.fluss.metadata.TableInfo; +import org.apache.fluss.metadata.TablePath; +import org.apache.fluss.types.DataTypes; + +import org.junit.jupiter.api.Test; + +import java.util.Collection; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CompletableFuture; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** + * Unit test for the fail-loud guard in {@link LakeSplitGenerator}: for a primary-key table, if the + * lake snapshot of a partition contains a bucket id outside the partition's enumerated bucket range + * (which can only happen if the per-partition bucket count is inconsistent with the tiered data), + * union-read split generation must refuse rather than silently drop the out-of-range lake data. + */ +class LakeSplitGeneratorTest { + + /** Table-level bucket count, kept different from the per-partition counts used below. */ + private static final int TABLE_LEVEL_BUCKET_COUNT = 3; + + /** + * Builds a {@link LakeSplitGenerator} for a partitioned primary-key table (schema: a INT, b + * STRING, c STRING; PK a+c) whose single partition "p" has {@code partitionBucketCount} + * enumerated buckets and a single lake split landing in {@code lakeSplitBucket}. + */ + private static LakeSplitGenerator createGenerator(int partitionBucketCount, int lakeSplitBucket) + throws Exception { + TablePath tablePath = TablePath.of("db", "pk_table"); + TableDescriptor descriptor = + TableDescriptor.builder() + .schema( + Schema.newBuilder() + .column("a", DataTypes.INT()) + .column("b", DataTypes.STRING()) + .column("c", DataTypes.STRING()) + .primaryKey("a", "c") + .build()) + .distributedBy(TABLE_LEVEL_BUCKET_COUNT, "a") + .partitionedBy("c") + .build(); + TableInfo tableInfo = TableInfo.of(tablePath, 1L, 1, descriptor, null, 1L, 1L); + + LakeSplit lakeSplit = new TestingLakeSplit(lakeSplitBucket, Collections.singletonList("p")); + LakeSource lakeSource = + TestingLakeSource.fromSplits(Collections.singletonList(lakeSplit)); + TestAdminAdapter admin = + new TestAdminAdapter() { + @Override + public CompletableFuture getReadableLakeSnapshot( + TablePath ignored) { + return CompletableFuture.completedFuture( + new LakeSnapshot(1L, new HashMap<>())); + } + }; + OffsetsInitializer.BucketOffsetsRetriever retriever = new ZeroOffsetsRetriever(); + + // the partition "p" carries its own bucket count (so an out-of-range lake bucket can be + // detected against the enumerated range [0, partitionBucketCount)) + PartitionInfo partitionInfo = + new PartitionInfo( + 7L, + ResolvedPartitionSpec.fromPartitionName(tableInfo.getPartitionKeys(), "p"), + null, + partitionBucketCount); + + return new LakeSplitGenerator( + tableInfo, + admin, + lakeSource, + retriever, + OffsetsInitializer.latest(), + partitionBucketCount, + () -> Collections.singleton(partitionInfo)); + } + + private static final class ZeroOffsetsRetriever + implements OffsetsInitializer.BucketOffsetsRetriever { + + @Override + public Map latestOffsets(String partitionName, Collection buckets) { + return zeroOffsets(buckets); + } + + @Override + public Map earliestOffsets( + String partitionName, Collection buckets) { + return zeroOffsets(buckets); + } + + @Override + public Map offsetsFromTimestamp( + String partitionName, Collection buckets, long timestamp) { + return zeroOffsets(buckets); + } + + private static Map zeroOffsets(Collection buckets) { + Map offsets = new HashMap<>(); + for (Integer bucket : buckets) { + offsets.put(bucket, 0L); + } + return offsets; + } + } + + @Test + void testPrimaryKeyOutOfRangeLakeBucketFailsLoud() throws Exception { + // lake split lands in bucket 5, which is outside [0, 2) + LakeSplitGenerator generator = createGenerator(2, 5); + assertThatThrownBy(generator::generateHybridLakeFlussSplits) + .isInstanceOf(IllegalStateException.class) + .hasMessageContaining("outside the enumerated range") + .hasMessageContaining("refusing to generate union-read splits"); + } + + @Test + void testPrimaryKeyInRangeLakeBucketSucceeds() throws Exception { + // lake split lands in bucket 1, which is within [0, 4) + int partitionBucketCount = 4; + LakeSplitGenerator generator = createGenerator(partitionBucketCount, 1); + + // no out-of-range bucket: generation succeeds and produces one hybrid lake+log split per + // bucket of the partition's enumerated range [0, partitionBucketCount) + List splits = generator.generateHybridLakeFlussSplits(); + assertThat(partitionBucketCount).isNotEqualTo(TABLE_LEVEL_BUCKET_COUNT); + assertThat(splits).isNotNull().hasSize(partitionBucketCount); + for (int bucket = 0; bucket < partitionBucketCount; bucket++) { + SourceSplitBase split = splits.get(bucket); + assertThat(split).isInstanceOf(LakeSnapshotAndFlussLogSplit.class); + assertThat(split.getPartitionName()).isEqualTo("p"); + assertThat(split.getTableBucket().getPartitionId()).isEqualTo(7L); + assertThat(split.getTableBucket().getBucket()).isEqualTo(bucket); + } + } +} diff --git a/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/sink/testutils/TestAdminAdapter.java b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/sink/testutils/TestAdminAdapter.java index e8120c83d26..5e058194686 100644 --- a/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/sink/testutils/TestAdminAdapter.java +++ b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/sink/testutils/TestAdminAdapter.java @@ -173,6 +173,12 @@ public CompletableFuture> listPartitionInfos(TablePath table throw new UnsupportedOperationException("Not implemented in TestAdminAdapter"); } + @Override + public CompletableFuture> listPartitionInfos( + TablePath tablePath, boolean includeSystemPartitions) { + throw new UnsupportedOperationException("Not implemented in TestAdminAdapter"); + } + @Override public CompletableFuture> listPartitionInfos( TablePath tablePath, PartitionSpec partialPartitionSpec) { diff --git a/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/sink/undo/RecoveryOffsetManagerTest.java b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/sink/undo/RecoveryOffsetManagerTest.java index a9e3ec66f41..e15cbd79aa3 100644 --- a/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/sink/undo/RecoveryOffsetManagerTest.java +++ b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/sink/undo/RecoveryOffsetManagerTest.java @@ -101,7 +101,8 @@ private static RecoveryOffsetManager createManager( private static PartitionInfo createPartitionInfo(long partitionId, String partitionName) { ResolvedPartitionSpec spec = ResolvedPartitionSpec.fromPartitionValue("pt", partitionName); - return new PartitionInfo(partitionId, spec, DEFAULT_REMOTE_DATA_DIR); + // matches the 1-bucket table used by the tests calling this helper + return new PartitionInfo(partitionId, spec, DEFAULT_REMOTE_DATA_DIR, 1); } // ==================== FRESH_START Tests ==================== @@ -804,6 +805,116 @@ void testCleanupOffsetsNonTask0() { assertThat(admin.wasDeleteCalled()).isFalse(); } + @Test + void testCheckpointRecoveryEnumeratesPerPartitionBucketCount() throws Exception { + // Two partitions with different bucketCount: partition 1 has 2 buckets (old, + // pre-ALTER), partition 2 has 4 buckets (new, post-ALTER). If getAllBuckets used the + // table-level count for both, the old partition would spuriously enumerate buckets 2/3 or + // the new partition would be truncated. + long oldPartitionId = 1L; + long newPartitionId = 2L; + int oldBucketCount = 2; + int newBucketCount = 4; + + Map currentOffsets = new HashMap<>(); + for (int b = 0; b < oldBucketCount; b++) { + currentOffsets.put(new TableBucket(TABLE_ID, oldPartitionId, b), 200L); + } + for (int b = 0; b < newBucketCount; b++) { + currentOffsets.put(new TableBucket(TABLE_ID, newPartitionId, b), 200L); + } + + RecoveryTestAdmin admin = new RecoveryTestAdmin(currentOffsets); + admin.setPartitions( + Arrays.asList( + createPartitionInfoWithBucketCount(oldPartitionId, "old", oldBucketCount), + createPartitionInfoWithBucketCount(newPartitionId, "new", newBucketCount))); + // Table-level bucket count is 4 (post-ALTER). Old partition must be enumerated at 2. + TableInfo tableInfo = createTableInfo(4, true); + RecoveryOffsetManager manager = + new RecoveryOffsetManager( + admin, PRODUCER_ID, 0, 1, 10L, 5000L, TABLE_PATH, tableInfo); + + Map chkOffsets = new HashMap<>(); + for (Map.Entry e : currentOffsets.entrySet()) { + chkOffsets.put(e.getKey(), 100L); + } + WriterState state = new WriterState(chkOffsets); + + RecoveryOffsetManager.RecoveryDecision decision = + manager.determineRecoveryStrategy(Collections.singleton(state)); + + assertThat(decision.getStrategy()) + .isEqualTo(RecoveryOffsetManager.RecoveryStrategy.CHECKPOINT_RECOVERY); + assertThat(decision.getUndoOffsets()).hasSize(oldBucketCount + newBucketCount); + for (int b = 0; b < oldBucketCount; b++) { + assertThat(decision.getUndoOffsets()) + .as("old partition bucket %d must be in decision", b) + .containsKey(new TableBucket(TABLE_ID, oldPartitionId, b)); + } + for (int b = 0; b < newBucketCount; b++) { + assertThat(decision.getUndoOffsets()) + .as("new partition bucket %d must be in decision", b) + .containsKey(new TableBucket(TABLE_ID, newPartitionId, b)); + } + assertThat(decision.getUndoOffsets()) + .doesNotContainKey(new TableBucket(TABLE_ID, oldPartitionId, 2)) + .doesNotContainKey(new TableBucket(TABLE_ID, oldPartitionId, 3)); + } + + @Test + void testProducerOffsetRegistrationUsesPerPartitionBucketCount() throws Exception { + // Empty checkpoint on Task0 → registerCurrentOffsets writes ALL buckets it enumerates. + // fetchAllBucketOffsets must enumerate each partition using its own bucketCount so + // the registered set is exactly the union of per-partition [0, bucketCount) ranges. + long oldPartitionId = 1L; + long newPartitionId = 2L; + int oldBucketCount = 2; + int newBucketCount = 4; + + Map currentOffsets = new HashMap<>(); + long offset = 100L; + for (int b = 0; b < oldBucketCount; b++) { + currentOffsets.put(new TableBucket(TABLE_ID, oldPartitionId, b), offset++); + } + for (int b = 0; b < newBucketCount; b++) { + currentOffsets.put(new TableBucket(TABLE_ID, newPartitionId, b), offset++); + } + + RecoveryTestAdmin admin = new RecoveryTestAdmin(currentOffsets); + admin.setPartitions( + Arrays.asList( + createPartitionInfoWithBucketCount(oldPartitionId, "old", oldBucketCount), + createPartitionInfoWithBucketCount(newPartitionId, "new", newBucketCount))); + admin.setInitialOffsetsForRegistration(currentOffsets); + TableInfo tableInfo = createTableInfo(4, true); + RecoveryOffsetManager manager = + new RecoveryOffsetManager( + admin, PRODUCER_ID, 0, 1, 10L, 5000L, TABLE_PATH, tableInfo); + + // null recoveredState triggers producer-offset recovery on Task0, which internally calls + // fetchAllBucketOffsets to build the registration map. + manager.determineRecoveryStrategy(null); + + Map registered = admin.registeredOffsets; + assertThat(registered).hasSize(oldBucketCount + newBucketCount); + for (int b = 0; b < oldBucketCount; b++) { + assertThat(registered).containsKey(new TableBucket(TABLE_ID, oldPartitionId, b)); + } + for (int b = 0; b < newBucketCount; b++) { + assertThat(registered).containsKey(new TableBucket(TABLE_ID, newPartitionId, b)); + } + assertThat(registered) + .doesNotContainKey(new TableBucket(TABLE_ID, oldPartitionId, 2)) + .doesNotContainKey(new TableBucket(TABLE_ID, oldPartitionId, 3)); + } + + private static PartitionInfo createPartitionInfoWithBucketCount( + long partitionId, String partitionName, int bucketCount) { + ResolvedPartitionSpec spec = ResolvedPartitionSpec.fromPartitionValue("pt", partitionName); + return new PartitionInfo(partitionId, spec, DEFAULT_REMOTE_DATA_DIR, bucketCount); + } + // ==================== Test Admin Implementation ==================== /** diff --git a/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/source/enumerator/FlinkSourceEnumeratorTest.java b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/source/enumerator/FlinkSourceEnumeratorTest.java index 22646efabe9..35826c3239e 100644 --- a/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/source/enumerator/FlinkSourceEnumeratorTest.java +++ b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/source/enumerator/FlinkSourceEnumeratorTest.java @@ -39,16 +39,21 @@ import org.apache.fluss.flink.source.split.SnapshotSplit; import org.apache.fluss.flink.source.split.SourceSplitBase; import org.apache.fluss.flink.source.state.SourceEnumeratorState; +import org.apache.fluss.flink.tiering.source.split.TieringSplit; +import org.apache.fluss.flink.tiering.source.split.TieringSplitGenerator; import org.apache.fluss.flink.utils.FlinkTestBase; import org.apache.fluss.lake.source.LakeSource; import org.apache.fluss.lake.source.LakeSplit; import org.apache.fluss.lake.source.TestingLakeSource; import org.apache.fluss.lake.source.TestingLakeSplit; import org.apache.fluss.metadata.PartitionInfo; +import org.apache.fluss.metadata.PartitionSpec; import org.apache.fluss.metadata.ResolvedPartitionSpec; import org.apache.fluss.metadata.Schema; import org.apache.fluss.metadata.TableBucket; +import org.apache.fluss.metadata.TableChange; import org.apache.fluss.metadata.TableDescriptor; +import org.apache.fluss.metadata.TableInfo; import org.apache.fluss.metadata.TablePath; import org.apache.fluss.predicate.Predicate; import org.apache.fluss.predicate.PredicateBuilder; @@ -406,7 +411,10 @@ void testRestoreFlussOnlySourceWithLakeSourceDoesNotGenerateLakeSplits(@TempDir DEFAULT_BUCKET_NUM, Collections.singletonList( new PartitionInfo( - partitionId, partitionSpec, DEFAULT_REMOTE_DATA_DIR))); + partitionId, + partitionSpec, + DEFAULT_REMOTE_DATA_DIR, + DEFAULT_BUCKET_NUM))); SourceEnumeratorState checkpointState; try (MockSplitEnumeratorContext context = @@ -1567,14 +1575,22 @@ void testPartitionsExpiredInFlussButExistInLake( Collections.singletonList(isPrimaryKeyTable ? "date" : "name"), partitionName); lakePartitionInfos.add( - new PartitionInfo(partitionId, partitionSpec, DEFAULT_REMOTE_DATA_DIR)); + new PartitionInfo( + partitionId, + partitionSpec, + DEFAULT_REMOTE_DATA_DIR, + DEFAULT_BUCKET_NUM)); } ResolvedPartitionSpec partitionSpec = ResolvedPartitionSpec.fromPartitionName( Collections.singletonList(isPrimaryKeyTable ? "date" : "name"), hybridPartitionName); lakePartitionInfos.add( - new PartitionInfo(hybridPartitionId, partitionSpec, DEFAULT_REMOTE_DATA_DIR)); + new PartitionInfo( + hybridPartitionId, + partitionSpec, + DEFAULT_REMOTE_DATA_DIR, + DEFAULT_BUCKET_NUM)); LakeSource lakeSource = new TestingLakeSource(DEFAULT_BUCKET_NUM, lakePartitionInfos); @@ -2030,4 +2046,150 @@ private Map putRows(TablePath tablePath, int rowsNum) throws E } return bucketRows; } + + // ==================== Per-Partition Bucket Count Tests ==================== + + private static final int OLD_BUCKET_NUM = 2; + private static final int NEW_BUCKET_NUM = 4; + + private static final TableDescriptor RESCALE_LOG_TABLE = + TableDescriptor.builder() + .schema( + Schema.newBuilder() + .column("id", DataTypes.INT()) + .column("name", DataTypes.STRING()) + .build()) + .distributedBy(OLD_BUCKET_NUM) + .partitionedBy("name") + .property(ConfigOptions.TABLE_AUTO_PARTITION_ENABLED, false) + .build(); + + /** + * Creates a partitioned log table, creates "old" partition, ALTERs bucket.num to {@link + * #NEW_BUCKET_NUM}, creates "new" partition, writes rows to both. Returns [tablePath, + * tableInfo, oldPartitionId, newPartitionId]. + */ + private Object[] setupRescaledPartitionedTable() throws Exception { + TablePath tablePath = TablePath.of(DEFAULT_DB, "rescale_split_" + System.nanoTime()); + createTable(tablePath, RESCALE_LOG_TABLE); + PartitionSpec oldSpec = new PartitionSpec(Collections.singletonMap("name", "old")); + PartitionSpec newSpec = new PartitionSpec(Collections.singletonMap("name", "new")); + admin.createPartition(tablePath, oldSpec, false).get(); + writeRows(conn, tablePath, genRows(10, "old"), true); + admin.alterTable( + tablePath, + Collections.singletonList(TableChange.modifyBucketCount(NEW_BUCKET_NUM)), + false) + .get(); + admin.createPartition(tablePath, newSpec, false).get(); + writeRows(conn, tablePath, genRows(20, "new"), true); + List infos = admin.listPartitionInfos(tablePath).get(); + PartitionInfo oldInfo = + infos.stream().filter(i -> "old".equals(i.getPartitionName())).findFirst().get(); + PartitionInfo newInfo = + infos.stream().filter(i -> "new".equals(i.getPartitionName())).findFirst().get(); + assertThat(oldInfo.getBucketCount()).isEqualTo(OLD_BUCKET_NUM); + assertThat(newInfo.getBucketCount()).isEqualTo(NEW_BUCKET_NUM); + return new Object[] { + tablePath, + admin.getTableInfo(tablePath).get(), + oldInfo.getPartitionId(), + newInfo.getPartitionId() + }; + } + + private static List genRows(int count, String partition) { + List rows = new ArrayList<>(); + for (int i = 0; i < count; i++) { + rows.add(row(i, partition)); + } + return rows; + } + + /** + * Tests that {@link FlinkSourceEnumerator} generates splits using each partition's actual + * bucket count after an ALTER bucket.num, in both streaming and batch modes. + */ + @ParameterizedTest + @ValueSource(booleans = {true, false}) + void testFlussSplitsEnumeratePerPartitionBucketCount(boolean streaming) throws Throwable { + Object[] ctx = setupRescaledPartitionedTable(); + try (MockSplitEnumeratorContext context = + new MockSplitEnumeratorContext<>(3); + MockWorkExecutor workExecutor = new MockWorkExecutor(context); + FlinkSourceEnumerator enumerator = + new FlinkSourceEnumerator( + (TablePath) ctx[0], + flussConf, + false, + true, + context, + Collections.emptySet(), + Collections.emptyMap(), + null, + streaming + ? OffsetsInitializer.earliest() + : OffsetsInitializer.full(), + DEFAULT_SCAN_PARTITION_DISCOVERY_INTERVAL_MS, + streaming, + null, + null, + workExecutor, + LeaseContext.DEFAULT, + false)) { + enumerator.start(); + if (streaming) { + runPeriodicPartitionDiscovery(workExecutor); + } else { + workExecutor.runNextOneTimeCallable(); + } + for (int i = 0; i < 3; i++) { + registerReader(context, enumerator, i); + } + long oldId = (long) ctx[2]; + long newId = (long) ctx[3]; + List splits = + getReadersAssignments(context).values().stream() + .flatMap(List::stream) + .collect(Collectors.toList()); + assertThat(splits).allMatch(s -> s instanceof LogSplit); + assertThat( + splits.stream() + .filter(s -> s.getTableBucket().getPartitionId() == oldId) + .map(s -> s.getTableBucket().getBucket()) + .collect(Collectors.toList())) + .containsExactlyInAnyOrder(0, 1); + assertThat( + splits.stream() + .filter(s -> s.getTableBucket().getPartitionId() == newId) + .map(s -> s.getTableBucket().getBucket()) + .collect(Collectors.toList())) + .containsExactlyInAnyOrder(0, 1, 2, 3); + } + } + + /** + * Tests that {@link TieringSplitGenerator} generates tiering splits using each partition's + * actual bucket count after an ALTER bucket.num. + */ + @Test + void testTieringSplitsEnumeratePerPartitionBucketCount() throws Throwable { + Object[] ctx = setupRescaledPartitionedTable(); + TableInfo tableInfo = (TableInfo) ctx[1]; + long oldPartitionId = (long) ctx[2]; + long newPartitionId = (long) ctx[3]; + List splits = new TieringSplitGenerator(admin).generateTableSplits(tableInfo); + assertThat( + splits.stream() + .filter(s -> s.getTableBucket().getPartitionId() == oldPartitionId) + .map(s -> s.getTableBucket().getBucket()) + .collect(Collectors.toList())) + .containsExactlyInAnyOrder(0, 1); + assertThat( + splits.stream() + .filter(s -> s.getTableBucket().getPartitionId() == newPartitionId) + .map(s -> s.getTableBucket().getBucket()) + .collect(Collectors.toList())) + .containsExactlyInAnyOrder(0, 1, 2, 3); + } } diff --git a/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/tiering/source/TieringWriterInitContextTest.java b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/tiering/source/TieringWriterInitContextTest.java index 8002f476c22..640ce820431 100644 --- a/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/tiering/source/TieringWriterInitContextTest.java +++ b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/tiering/source/TieringWriterInitContextTest.java @@ -17,28 +17,108 @@ package org.apache.fluss.flink.tiering.source; +import org.apache.fluss.config.Configuration; +import org.apache.fluss.metadata.Schema; +import org.apache.fluss.metadata.TableBucket; +import org.apache.fluss.metadata.TableInfo; +import org.apache.fluss.metadata.TablePath; +import org.apache.fluss.types.DataTypes; + import org.junit.jupiter.api.Test; +import java.util.Collections; + +import static org.apache.fluss.record.TestData.DEFAULT_REMOTE_DATA_DIR; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; /** Tests for {@link TieringWriterInitContext}. */ class TieringWriterInitContextTest { + private static final long TABLE_ID = 1L; + private static final TablePath TABLE_PATH = TablePath.of("test_db", "test_table"); + private static final int TABLE_BUCKET_COUNT = 8; + @Test void testIoTmpDir() { TieringWriterInitContext defaultContext = - new TieringWriterInitContext(null, null, null, null); + newContext(new TableBucket(TABLE_ID, 0), null, null, null); TieringWriterInitContext context = - new TieringWriterInitContext( + newContext( + new TableBucket(TABLE_ID, 0), null, null, - null, - null, - 0, - 1L, new String[] {"/flink_tmp_0/fluss", "/flink_tmp_1/fluss"}); assertThat(defaultContext.ioTmpDirs()).isNull(); assertThat(context.ioTmpDirs()).containsExactly("/flink_tmp_0/fluss", "/flink_tmp_1/fluss"); } + + @Test + void testNonPartitionedFallsBackToTableLevelCount() { + // A non-partitioned bucket carries no per-partition count; the table-level count applies. + TieringWriterInitContext context = newContext(new TableBucket(TABLE_ID, 0), null, null); + assertThat(context.bucketCount()).isEqualTo(TABLE_BUCKET_COUNT); + } + + @Test + void testPartitionedUsesPerPartitionCount() { + // A partitioned bucket must use its own actual bucket count. + TieringWriterInitContext context = + newContext(new TableBucket(TABLE_ID, 1L, 0), "2024-01", 4); + assertThat(context.bucketCount()).isEqualTo(4); + } + + @Test + void testPartitionedWithoutCountFailsLoud() { + // A partitioned bucket with no resolved per-partition count must fail loudly rather than + // silently guessing (which would corrupt the lake table's bucket layout). + assertThatThrownBy(() -> newContext(new TableBucket(TABLE_ID, 1L, 0), "2024-01", null)) + .isInstanceOf(NullPointerException.class) + .hasMessageContaining("actual bucket count"); + } + + private static TieringWriterInitContext newContext( + TableBucket tableBucket, String partition, Integer partitionBucketCount) { + return newContext(tableBucket, partition, partitionBucketCount, null); + } + + private static TieringWriterInitContext newContext( + TableBucket tableBucket, + String partition, + Integer partitionBucketCount, + String[] ioTmpDirs) { + return new TieringWriterInitContext( + TABLE_PATH, + tableBucket, + partition, + createTableInfo(TABLE_BUCKET_COUNT), + 0, + 0L, + partitionBucketCount, + ioTmpDirs); + } + + private static TableInfo createTableInfo(int numBuckets) { + Schema schema = + Schema.newBuilder() + .column("id", DataTypes.INT()) + .column("value", DataTypes.STRING()) + .primaryKey("id") + .build(); + return new TableInfo( + TABLE_PATH, + TABLE_ID, + 0, + schema, + Collections.emptyList(), + Collections.emptyList(), + numBuckets, + new Configuration(), + new Configuration(), + DEFAULT_REMOTE_DATA_DIR, + null, + System.currentTimeMillis(), + System.currentTimeMillis()); + } } diff --git a/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/utils/FlinkConversionsTest.java b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/utils/FlinkConversionsTest.java index bbb4aa3b1a4..67f843a7276 100644 --- a/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/utils/FlinkConversionsTest.java +++ b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/utils/FlinkConversionsTest.java @@ -20,6 +20,7 @@ import org.apache.fluss.config.ConfigOptions; import org.apache.fluss.flink.catalog.TestSchemaResolver; import org.apache.fluss.metadata.KvFormat; +import org.apache.fluss.metadata.TableChange; import org.apache.fluss.metadata.TableDescriptor; import org.apache.fluss.metadata.TableInfo; import org.apache.fluss.metadata.TablePath; @@ -74,6 +75,26 @@ /** Test for {@link FlinkConversions}. */ public class FlinkConversionsTest { + @Test + void testConvertBucketCountChange() { + assertThat( + FlinkConversions.toFlussTableChanges( + org.apache.flink.table.catalog.TableChange.set( + BUCKET_NUMBER.key(), "8"))) + .containsExactly(TableChange.modifyBucketCount(8)); + assertThat( + FlinkConversions.toFlussTableChanges( + org.apache.flink.table.catalog.TableChange.set("key", "value"))) + .containsExactly(TableChange.set("key", "value")); + assertThatThrownBy( + () -> + FlinkConversions.toFlussTableChanges( + org.apache.flink.table.catalog.TableChange.set( + BUCKET_NUMBER.key(), "0"))) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("Bucket count must be positive"); + } + @Test void testTypeConversion() { // create a list with all fluss data types diff --git a/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/utils/FlinkTestBase.java b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/utils/FlinkTestBase.java index fc8371e27a7..968a71e99e2 100644 --- a/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/utils/FlinkTestBase.java +++ b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/utils/FlinkTestBase.java @@ -247,7 +247,8 @@ public static Map createPartitions( tableInfo.getTableId(), assignment.getBucketAssignments()), zkClient.getDefaultRemoteDataDir(), tablePath, - tableInfo.getTableId()); + tableInfo.getTableId(), + tableInfo.getNumBuckets()); } return newPartitionIds; } diff --git a/fluss-lake/fluss-lake-hudi/src/test/java/org/apache/fluss/lake/hudi/tiering/HudiTieringTest.java b/fluss-lake/fluss-lake-hudi/src/test/java/org/apache/fluss/lake/hudi/tiering/HudiTieringTest.java index fca678dcbe8..eccdf6e40b1 100644 --- a/fluss-lake/fluss-lake-hudi/src/test/java/org/apache/fluss/lake/hudi/tiering/HudiTieringTest.java +++ b/fluss-lake/fluss-lake-hudi/src/test/java/org/apache/fluss/lake/hudi/tiering/HudiTieringTest.java @@ -654,6 +654,11 @@ public int splitIndex() { public long tieringRoundTimestamp() { return tieringRoundTimestamp; } + + @Override + public int bucketCount() { + return tableInfo.getNumBuckets(); + } } private static class TestingCommitterInitContext implements CommitterInitContext { diff --git a/fluss-lake/fluss-lake-iceberg/src/main/java/org/apache/fluss/lake/iceberg/IcebergLakeCatalog.java b/fluss-lake/fluss-lake-iceberg/src/main/java/org/apache/fluss/lake/iceberg/IcebergLakeCatalog.java index 70886c9c839..8650dc48ce6 100644 --- a/fluss-lake/fluss-lake-iceberg/src/main/java/org/apache/fluss/lake/iceberg/IcebergLakeCatalog.java +++ b/fluss-lake/fluss-lake-iceberg/src/main/java/org/apache/fluss/lake/iceberg/IcebergLakeCatalog.java @@ -155,6 +155,12 @@ public void createTable(TablePath tablePath, TableDescriptor tableDescriptor, Co @Override public void alterTable(TablePath tablePath, List tableChanges, Context context) throws TableNotExistException { + for (TableChange change : tableChanges) { + if (change instanceof TableChange.ModifyBucketCount) { + throw new UnsupportedOperationException( + "Bucket count rescale is not supported by the Iceberg lake catalog yet."); + } + } try { Table table = icebergCatalog.loadTable(toIcebergTableIdentifier(tablePath)); diff --git a/fluss-lake/fluss-lake-iceberg/src/test/java/org/apache/fluss/lake/iceberg/tiering/IcebergTieringTest.java b/fluss-lake/fluss-lake-iceberg/src/test/java/org/apache/fluss/lake/iceberg/tiering/IcebergTieringTest.java index c3ecc0af390..d2a800935a5 100644 --- a/fluss-lake/fluss-lake-iceberg/src/test/java/org/apache/fluss/lake/iceberg/tiering/IcebergTieringTest.java +++ b/fluss-lake/fluss-lake-iceberg/src/test/java/org/apache/fluss/lake/iceberg/tiering/IcebergTieringTest.java @@ -316,6 +316,11 @@ public String partition() { public TableInfo tableInfo() { return tableInfo; } + + @Override + public int bucketCount() { + return tableInfo.getNumBuckets(); + } }); } diff --git a/fluss-lake/fluss-lake-lance/src/test/java/org/apache/fluss/lake/lance/tiering/LanceTieringTest.java b/fluss-lake/fluss-lake-lance/src/test/java/org/apache/fluss/lake/lance/tiering/LanceTieringTest.java index 2943d9195f5..82756c98a48 100644 --- a/fluss-lake/fluss-lake-lance/src/test/java/org/apache/fluss/lake/lance/tiering/LanceTieringTest.java +++ b/fluss-lake/fluss-lake-lance/src/test/java/org/apache/fluss/lake/lance/tiering/LanceTieringTest.java @@ -407,6 +407,11 @@ public String partition() { public TableInfo tableInfo() { return tableInfo; } + + @Override + public int bucketCount() { + return tableInfo.getNumBuckets(); + } }); } diff --git a/fluss-lake/fluss-lake-paimon/src/main/java/org/apache/fluss/lake/paimon/PaimonLakeCatalog.java b/fluss-lake/fluss-lake-paimon/src/main/java/org/apache/fluss/lake/paimon/PaimonLakeCatalog.java index 4cd127f441e..1457c0f6562 100644 --- a/fluss-lake/fluss-lake-paimon/src/main/java/org/apache/fluss/lake/paimon/PaimonLakeCatalog.java +++ b/fluss-lake/fluss-lake-paimon/src/main/java/org/apache/fluss/lake/paimon/PaimonLakeCatalog.java @@ -29,6 +29,7 @@ import org.apache.fluss.metadata.TablePath; import org.apache.fluss.utils.IOUtils; +import org.apache.paimon.CoreOptions; import org.apache.paimon.catalog.Catalog; import org.apache.paimon.catalog.CatalogContext; import org.apache.paimon.catalog.CatalogFactory; @@ -45,6 +46,7 @@ import org.slf4j.LoggerFactory; import java.util.ArrayList; +import java.util.Collections; import java.util.LinkedHashMap; import java.util.List; import java.util.stream.Collectors; @@ -63,6 +65,9 @@ public class PaimonLakeCatalog implements LakeCatalog { private static final Logger LOG = LoggerFactory.getLogger(PaimonLakeCatalog.class); private static final String PAIMON_PATH_KEY = "paimon.path"; + + private static final String BUCKET_NUM_PROPERTY = "bucket.num"; + public static final LinkedHashMap LEGACY_SYSTEM_COLUMNS = new LinkedHashMap<>(); @@ -119,11 +124,32 @@ public void createTable(TablePath tablePath, TableDescriptor tableDescriptor, Co @Override public void alterTable(TablePath tablePath, List tableChanges, Context context) throws TableNotExistException { + // Apply the bucket count rescale separately so the schema-compat branches below cannot + // swallow it. + Integer newBucketCount = null; + List remainingChanges = new ArrayList<>(tableChanges.size()); + for (TableChange tableChange : tableChanges) { + if (tableChange instanceof TableChange.ModifyBucketCount) { + newBucketCount = ((TableChange.ModifyBucketCount) tableChange).getNewBucketCount(); + if (newBucketCount <= 0) { + throw new InvalidAlterTableException( + "Invalid value for '" + BUCKET_NUM_PROPERTY + "': " + newBucketCount); + } + } else { + remainingChanges.add(tableChange); + } + } + if (newBucketCount != null) { + applyBucketCountChange(tablePath, newBucketCount); + } + if (remainingChanges.isEmpty()) { + return; + } try { Table table = paimonCatalog.getTable(toPaimon(tablePath)); FileStoreTable fileStoreTable = (FileStoreTable) table; List changesToApply = - validateAndFilterPaimonPathChanges(fileStoreTable.location(), tableChanges); + validateAndFilterPaimonPathChanges(fileStoreTable.location(), remainingChanges); // Avoid creating a new Paimon schema version for a path-only no-op. if (changesToApply.isEmpty()) { @@ -159,7 +185,7 @@ currentPaimonSchema, toPaimonSchema(context.getExpectedTable()))) { + "rather than applying other table changes: %s.", currentPaimonSchema, context.getCurrentTable().getSchema(), - tableChanges)); + remainingChanges)); } if (!paimonSchemaChanges.isEmpty()) { @@ -291,6 +317,23 @@ private void createDatabase(String databaseName) { } } + private void applyBucketCountChange(TablePath tablePath, int newBucketCount) + throws TableNotExistException { + // Bypass toPaimonSchemaChanges (which rejects Paimon's own BUCKET key via + // PAIMON_UNSETTABLE_OPTIONS) and set BUCKET directly. + List changes = + Collections.singletonList( + SchemaChange.setOption( + CoreOptions.BUCKET.key(), String.valueOf(newBucketCount))); + try { + paimonCatalog.alterTable(toPaimon(tablePath), changes, false); + } catch (Catalog.TableNotExistException e) { + throw new TableNotExistException("Table " + tablePath + " does not exist."); + } catch (Catalog.ColumnAlreadyExistException | Catalog.ColumnNotExistException e) { + throw new InvalidAlterTableException(e.getMessage(), e); + } + } + @Override public void close() { IOUtils.closeQuietly(paimonCatalog, "paimon catalog"); diff --git a/fluss-lake/fluss-lake-paimon/src/main/java/org/apache/fluss/lake/paimon/lookup/PaimonLakeTableLookuper.java b/fluss-lake/fluss-lake-paimon/src/main/java/org/apache/fluss/lake/paimon/lookup/PaimonLakeTableLookuper.java index caa296ae8a6..fbacfe40c26 100644 --- a/fluss-lake/fluss-lake-paimon/src/main/java/org/apache/fluss/lake/paimon/lookup/PaimonLakeTableLookuper.java +++ b/fluss-lake/fluss-lake-paimon/src/main/java/org/apache/fluss/lake/paimon/lookup/PaimonLakeTableLookuper.java @@ -17,6 +17,7 @@ package org.apache.fluss.lake.paimon.lookup; +import org.apache.fluss.bucketing.BucketingFunction; import org.apache.fluss.config.Configuration; import org.apache.fluss.config.MemorySize; import org.apache.fluss.config.TableConfig; @@ -25,6 +26,7 @@ import org.apache.fluss.lake.lakestorage.LakeTableLookuper; import org.apache.fluss.lake.paimon.utils.PaimonPartitionBucket; import org.apache.fluss.lake.paimon.utils.PaimonRowAsFlussRow; +import org.apache.fluss.metadata.DataLakeFormat; import org.apache.fluss.metadata.TablePath; import org.apache.fluss.row.BinaryRow; import org.apache.fluss.row.InternalRow; @@ -104,6 +106,10 @@ */ public class PaimonLakeTableLookuper implements LakeTableLookuper { + /** Bucketing function the lake data was written with, used to recompute a bucket. */ + private static final BucketingFunction BUCKETING_FUNCTION = + BucketingFunction.of(DataLakeFormat.PAIMON); + private final Configuration paimonConfig; private final TablePath tablePath; private final String ioTmpDir; @@ -119,6 +125,9 @@ public class PaimonLakeTableLookuper implements LakeTableLookuper { // Remains non-zero until a refresh completes without observing another request. private final AtomicLong pendingRefreshRequests; + /** Bucket count each partition was written with, resolved from the lake metadata. */ + private final Map totalBucketsByPartition; + private @Nullable Catalog catalog; private @Nullable FileStoreTable fileStoreTable; private @Nullable IOManager ioManager; @@ -153,6 +162,7 @@ public PaimonLakeTableLookuper( this.lookupStateLock = new Object(); this.registeredFiles = new ConcurrentHashMap<>(); this.pendingRefreshRequests = new AtomicLong(); + this.totalBucketsByPartition = new ConcurrentHashMap<>(); } @Override @@ -191,6 +201,7 @@ public void close() { IOUtils.closeQuietly(ioManager, "Paimon lookup IO manager"); IOUtils.closeQuietly(catalog, "Paimon catalog"); registeredFiles.clear(); + totalBucketsByPartition.clear(); localTableQuery = null; compactedKeyDecoder = null; trimmedPrimaryKeys = null; @@ -353,10 +364,14 @@ private org.apache.paimon.data.BinaryRow getKey(byte[] key, LookupContext contex } private @Nullable byte[] lookupInternal(byte[] key, LookupContext context) { + org.apache.paimon.data.BinaryRow partition = getPartition(context); + Integer bucket = resolveLakeBucket(key, partition, context); + if (bucket == null) { + return null; + } org.apache.paimon.data.InternalRow paimonRow; try { - paimonRow = - lookupPaimon(getPartition(context), context.bucketId(), getKey(key, context)); + paimonRow = lookupPaimon(partition, bucket, getKey(key, context)); } catch (IOException e) { // Historical Paimon point lookup is part of the Fluss KV lookup path. Expose a // persistent I/O failure as a retriable KV error so the existing KV RPC retry @@ -373,6 +388,88 @@ private org.apache.paimon.data.BinaryRow getKey(byte[] key, LookupContext contex return encodeValue(paimonRow, context.schemaId(), context.valueRowType()); } + /** + * Resolves the bucket the key lives in: the caller's bucket id when it matches the lake layout, + * otherwise recomputed from the bucket count the partition was written with. Returns null when + * the partition holds no data in the lake. + */ + private @Nullable Integer resolveLakeBucket( + byte[] key, org.apache.paimon.data.BinaryRow partition, LookupContext context) { + if (context.bucketId() != null) { + return context.bucketId(); + } + Integer totalBuckets = resolveTotalBuckets(partition); + if (totalBuckets == null) { + return null; + } + return BUCKETING_FUNCTION.bucketing(deriveBucketKey(key, context), totalBuckets); + } + + /** + * Reads the bucket count the given partition was written with from the lake metadata. Fluss + * writes each partition with a single bucket count, so more than one value means the lake data + * cannot be routed reliably. + */ + private @Nullable Integer resolveTotalBuckets(org.apache.paimon.data.BinaryRow partition) { + org.apache.paimon.data.BinaryRow partitionKey = partition.copy(); + Integer cachedTotalBuckets = totalBucketsByPartition.get(partitionKey); + if (cachedTotalBuckets != null) { + return cachedTotalBuckets; + } + + // Keep metadata I/O outside ConcurrentHashMap's bin lock. + Set totalBuckets = new HashSet<>(); + InnerTableScan tableScan = + fileStoreTable + .newScan() + .withPartitionFilter(Collections.singletonList(partitionKey)); + for (Split split : tableScan.plan().splits()) { + if (split instanceof DataSplit) { + totalBuckets.add(((DataSplit) split).totalBuckets()); + } + } + if (totalBuckets.isEmpty()) { + return null; + } + if (totalBuckets.size() > 1) { + throw new KvStorageException( + "Cannot look up historical data of table " + + tablePath + + " because its lake data reports multiple bucket counts " + + totalBuckets + + " for one partition, so the bucket a key was written to " + + "cannot be determined."); + } + + Integer resolvedTotalBuckets = totalBuckets.iterator().next(); + Integer previousTotalBuckets = + totalBucketsByPartition.putIfAbsent(partitionKey, resolvedTotalBuckets); + return previousTotalBuckets == null ? resolvedTotalBuckets : previousTotalBuckets; + } + + /** + * Returns the bucket key bytes the lake bucketing function expects. The lookup key already is + * the bucket key when the table uses the default bucket key; otherwise the primary key is + * decoded and the bucket key fields are re-encoded with the lake encoder. + */ + private byte[] deriveBucketKey(byte[] key, LookupContext context) { + List bucketKeys = fileStoreTable.schema().bucketKeys(); + if (bucketKeys.equals(trimmedPrimaryKeys)) { + return key; + } + RowType primaryKeyRowType = context.valueRowType().project(trimmedPrimaryKeys); + InternalRow primaryKeyRow; + if (compactedKeyDecoder != null) { + primaryKeyRow = compactedKeyDecoder.decodeKey(key); + } else { + org.apache.paimon.data.BinaryRow paimonKeyRow = + new org.apache.paimon.data.BinaryRow(trimmedPrimaryKeys.size()); + paimonKeyRow.pointTo(MemorySegment.wrap(key), 0, key.length); + primaryKeyRow = new PaimonRowAsFlussRow(paimonKeyRow); + } + return new PaimonKeyEncoder(primaryKeyRowType, bucketKeys).encodeKey(primaryKeyRow); + } + private @Nullable org.apache.paimon.data.InternalRow lookupPaimon( org.apache.paimon.data.BinaryRow partition, int bucket, diff --git a/fluss-lake/fluss-lake-paimon/src/main/java/org/apache/fluss/lake/paimon/tiering/PaimonLakeWriter.java b/fluss-lake/fluss-lake-paimon/src/main/java/org/apache/fluss/lake/paimon/tiering/PaimonLakeWriter.java index a93d61b2e01..05e7ca7eb75 100644 --- a/fluss-lake/fluss-lake-paimon/src/main/java/org/apache/fluss/lake/paimon/tiering/PaimonLakeWriter.java +++ b/fluss-lake/fluss-lake-paimon/src/main/java/org/apache/fluss/lake/paimon/tiering/PaimonLakeWriter.java @@ -33,8 +33,10 @@ import org.apache.paimon.catalog.Catalog; import org.apache.paimon.table.FileStoreTable; +import javax.annotation.Nullable; + import java.io.IOException; -import java.util.Collections; +import java.util.HashMap; import java.util.List; import java.util.Map; @@ -51,11 +53,19 @@ public PaimonLakeWriter( PaimonCatalogProvider paimonCatalogProvider, WriterInitContext writerInitContext) throws IOException { this.paimonCatalog = paimonCatalogProvider.get(); + // Only Fixed Bucket tables (bucket keys non-empty) carry a positive BUCKET in Paimon. + // Overriding on an Unaware Bucket table (BUCKET = -1) would change its bucket mode. + // The context always resolves the actual bucket count. + Integer bucketOverride = + !writerInitContext.tableInfo().getBucketKeys().isEmpty() + ? writerInitContext.bucketCount() + : null; TablePath lakeTablePath = writerInitContext.tableInfo().getLakeTablePath(); FileStoreTable fileStoreTable = getTable( lakeTablePath, - writerInitContext.tableInfo().getTableConfig().isDataLakeAutoCompaction()); + writerInitContext.tableInfo().getTableConfig().isDataLakeAutoCompaction(), + bucketOverride); List partitionKeys = fileStoreTable.partitionKeys(); RowType flussRowType = writerInitContext.tableInfo().getRowType(); @@ -139,15 +149,23 @@ public void close() throws IOException { } } - private FileStoreTable getTable(TablePath tablePath, boolean isAutoCompaction) + private FileStoreTable getTable( + TablePath tablePath, boolean isAutoCompaction, @Nullable Integer bucketOverride) throws IOException { try { FileStoreTable table = (FileStoreTable) paimonCatalog.getTable(toPaimon(tablePath)); - Map compactionOptions = - Collections.singletonMap( - CoreOptions.WRITE_ONLY.key(), - isAutoCompaction ? Boolean.FALSE.toString() : Boolean.TRUE.toString()); - return table.copy(compactionOptions); + if (bucketOverride != null) { + // copy(Map) rejects BUCKET as immutable, so swap it in via a schema copy, + // which only rebuilds the in-memory table view. + Map schemaOptions = new HashMap<>(table.schema().options()); + schemaOptions.put(CoreOptions.BUCKET.key(), String.valueOf(bucketOverride)); + table = table.copy(table.schema().copy(schemaOptions)); + } + Map dynamicOptions = new HashMap<>(); + dynamicOptions.put( + CoreOptions.WRITE_ONLY.key(), + isAutoCompaction ? Boolean.FALSE.toString() : Boolean.TRUE.toString()); + return table.copy(dynamicOptions); } catch (Exception e) { throw new IOException("Failed to get table " + tablePath + " in Paimon.", e); } diff --git a/fluss-lake/fluss-lake-paimon/src/main/java/org/apache/fluss/lake/paimon/tiering/append/AppendOnlyWriter.java b/fluss-lake/fluss-lake-paimon/src/main/java/org/apache/fluss/lake/paimon/tiering/append/AppendOnlyWriter.java index a6d4f4292fd..f6431cb124e 100644 --- a/fluss-lake/fluss-lake-paimon/src/main/java/org/apache/fluss/lake/paimon/tiering/append/AppendOnlyWriter.java +++ b/fluss-lake/fluss-lake-paimon/src/main/java/org/apache/fluss/lake/paimon/tiering/append/AppendOnlyWriter.java @@ -59,9 +59,7 @@ public AppendOnlyWriter( boolean historicalPartition) { //noinspection unchecked super( - (TableWriteImpl) - // todo: set ioManager to support write-buffer-spillable - fileStoreTable.newWrite(FLUSS_LAKE_TIERING_COMMIT_USER), + buildTableWrite(fileStoreTable), fileStoreTable.rowType(), tableBucket, partition, @@ -73,6 +71,15 @@ public AppendOnlyWriter( this.paimonIncludingSystemColumns = paimonIncludingSystemColumns; } + @SuppressWarnings("unchecked") + private static TableWriteImpl buildTableWrite(FileStoreTable fileStoreTable) { + TableWriteImpl tableWrite = + (TableWriteImpl) + // todo: set ioManager to support write-buffer-spillable + fileStoreTable.newWrite(FLUSS_LAKE_TIERING_COMMIT_USER); + return tableWrite; + } + @Override public void write(LogRecord record) throws Exception { BinaryRow targetPartition = prepareRecordAndGetPartition(record); diff --git a/fluss-lake/fluss-lake-paimon/src/test/java/org/apache/fluss/lake/paimon/PaimonLakeCatalogTest.java b/fluss-lake/fluss-lake-paimon/src/test/java/org/apache/fluss/lake/paimon/PaimonLakeCatalogTest.java index c704bef954d..4b0482557e0 100644 --- a/fluss-lake/fluss-lake-paimon/src/test/java/org/apache/fluss/lake/paimon/PaimonLakeCatalogTest.java +++ b/fluss-lake/fluss-lake-paimon/src/test/java/org/apache/fluss/lake/paimon/PaimonLakeCatalogTest.java @@ -580,6 +580,55 @@ private org.apache.paimon.schema.Schema.Builder createPaimonSchemaBuilder( .option(CoreOptions.BUCKET_KEY.key(), bucketKey); } + @Test + void testUserFacingAlterTableStillRejectsBucketChange() throws Exception { + // Fluss's typed bucket-count change is applied, while any attempt to set Paimon's own + // bucket option directly through a property change keeps being rejected. + String database = "test_user_bucket_reject_db"; + String tableName = "test_user_bucket_reject_table"; + TablePath tablePath = TablePath.of(database, tableName); + createFixedBucketTable(database, tableName, 4); + + TableDescriptor fixedBucketDescriptor = fixedBucketTableDescriptor(4); + TestingLakeCatalogContext matchingContext = + new TestingLakeCatalogContext(fixedBucketDescriptor, fixedBucketDescriptor); + + List userChanges = + Collections.singletonList( + TableChange.set( + "paimon." + org.apache.paimon.CoreOptions.BUCKET.key(), "8")); + assertThatThrownBy( + () -> + flussPaimonCatalog.alterTable( + tablePath, userChanges, matchingContext)) + .hasMessageContaining("bucket") + .hasMessageContaining("cannot be changed"); + + flussPaimonCatalog.alterTable( + tablePath, + Collections.singletonList(TableChange.modifyBucketCount(8)), + matchingContext); + Identifier identifier = Identifier.create(database, tableName); + Table after = flussPaimonCatalog.getPaimonCatalog().getTable(identifier); + assertThat(after.options().get(org.apache.paimon.CoreOptions.BUCKET.key())).isEqualTo("8"); + } + + private TableDescriptor fixedBucketTableDescriptor(int bucketCount) { + return TableDescriptor.builder() + .schema(FLUSS_SCHEMA) + .property(TABLE_DATALAKE_ENABLED.key(), "true") + .property(TABLE_DATALAKE_FORMAT.key(), "paimon") + .property("table.datalake.paimon.warehouse", tempWarehouseDir.toURI().toString()) + .distributedBy(bucketCount, "id") + .build(); + } + + private void createFixedBucketTable(String database, String tableName, int initialBucketCount) { + TableDescriptor td = fixedBucketTableDescriptor(initialBucketCount); + flussPaimonCatalog.createTable( + TablePath.of(database, tableName), td, new TestingLakeCatalogContext(td, td)); + } + private void createTable(String database, String tableName) { TableDescriptor td = getTableDescriptor(FLUSS_SCHEMA); TablePath tablePath = TablePath.of(database, tableName); diff --git a/fluss-lake/fluss-lake-paimon/src/test/java/org/apache/fluss/lake/paimon/flink/FlinkUnionReadRescaleBucketITCase.java b/fluss-lake/fluss-lake-paimon/src/test/java/org/apache/fluss/lake/paimon/flink/FlinkUnionReadRescaleBucketITCase.java new file mode 100644 index 00000000000..27d58264fdb --- /dev/null +++ b/fluss-lake/fluss-lake-paimon/src/test/java/org/apache/fluss/lake/paimon/flink/FlinkUnionReadRescaleBucketITCase.java @@ -0,0 +1,500 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.fluss.lake.paimon.flink; + +import org.apache.fluss.client.initializer.BucketOffsetsRetrieverImpl; +import org.apache.fluss.config.ConfigOptions; +import org.apache.fluss.metadata.PartitionInfo; +import org.apache.fluss.metadata.PartitionSpec; +import org.apache.fluss.metadata.Schema; +import org.apache.fluss.metadata.TableBucket; +import org.apache.fluss.metadata.TableChange; +import org.apache.fluss.metadata.TableDescriptor; +import org.apache.fluss.metadata.TablePath; +import org.apache.fluss.row.InternalRow; +import org.apache.fluss.types.DataTypes; + +import org.apache.flink.core.execution.JobClient; +import org.apache.flink.types.Row; +import org.apache.flink.types.RowKind; +import org.apache.flink.util.CloseableIterator; +import org.apache.flink.util.CollectionUtil; +import org.apache.paimon.CoreOptions; +import org.apache.paimon.table.FileStoreTable; +import org.apache.paimon.table.source.DataSplit; +import org.apache.paimon.table.source.Split; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; + +import java.time.Duration; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.stream.Collectors; + +import static org.apache.fluss.flink.source.testutils.FlinkRowAssertionsUtils.collectRowsWithTimeout; +import static org.apache.fluss.lake.paimon.utils.PaimonConversions.toPaimon; +import static org.apache.fluss.testutils.DataTestUtils.row; +import static org.assertj.core.api.Assertions.assertThat; + +/** + * The IT case for union read (lake + fluss log) on a partitioned table whose partitions carry + * different bucket counts after an ALTER TABLE ... SET ('bucket.num' = N): the partition created + * before the ALTER keeps its original bucket count while the partition created afterwards uses the + * new one. Verifies that tiering stamps each partition's files with the partition's actual bucket + * count and that union read enumerates buckets per partition. + */ +class FlinkUnionReadRescaleBucketITCase extends FlinkUnionReadTestBase { + + private static final int OLD_BUCKET_NUM = 2; + private static final int NEW_BUCKET_NUM = 4; + private static final int RECORDS_PER_ROUND = 16; + + @BeforeAll + protected static void beforeAll() { + FlinkUnionReadTestBase.beforeAll(); + } + + @Test + void testUnionReadAcrossPartitionsWithDifferentBucketCounts() throws Exception { + String tableName = "rescale_bucket_log_table"; + TablePath tablePath = TablePath.of(DEFAULT_DB, tableName); + createPartitionedLogTable(tablePath, OLD_BUCKET_NUM); + + // "old" partition is created before the ALTER and keeps OLD_BUCKET_NUM buckets + createPartition(tablePath, "old"); + List expectedRows = new ArrayList<>(writeRows(tablePath, "old", 0)); + + // ALTER bucket.num, which also propagates the new BUCKET to the Paimon table + alterBucketNum(tablePath); + + // "new" partition is created after the ALTER and uses NEW_BUCKET_NUM buckets + createPartition(tablePath, "new"); + expectedRows.addAll(writeRows(tablePath, "new", 0)); + + Map bucketCountByPartition = bucketCountByPartitionName(tablePath); + assertThat(bucketCountByPartition) + .containsEntry("old", OLD_BUCKET_NUM) + .containsEntry("new", NEW_BUCKET_NUM); + + // the ALTER must have propagated the new BUCKET to the Paimon schema + FileStoreTable paimonTable = (FileStoreTable) paimonCatalog.getTable(toPaimon(tablePath)); + assertThat(paimonTable.options().get(CoreOptions.BUCKET.key())) + .isEqualTo(String.valueOf(NEW_BUCKET_NUM)); + + // start tiering and wait until both partitions are fully synced by their own bucket range + JobClient jobClient = buildTieringJob(execEnv); + try { + long tableId = admin.getTableInfo(tablePath).get().getTableId(); + waitUntilPartitionBucketsSynced(tablePath, tableId); + + // files of each partition must be stamped with the partition's actual bucket count + assertThat(totalBucketsOfPartition(tablePath, "old")).containsExactly(OLD_BUCKET_NUM); + assertThat(totalBucketsOfPartition(tablePath, "new")).containsExactly(NEW_BUCKET_NUM); + + // union read with all data in lake + List actual = + CollectionUtil.iteratorToList( + batchTEnv.executeSql("select * from " + tableName).collect()); + assertThat(actual).containsExactlyInAnyOrderElementsOf(expectedRows); + } finally { + jobClient.cancel().get(); + } + + // write more rows after tiering stopped so union read mixes lake splits and fluss log + expectedRows.addAll(writeRows(tablePath, "old", RECORDS_PER_ROUND)); + expectedRows.addAll(writeRows(tablePath, "new", RECORDS_PER_ROUND)); + + List actual = + CollectionUtil.iteratorToList( + batchTEnv.executeSql("select * from " + tableName).collect()); + assertThat(actual).containsExactlyInAnyOrderElementsOf(expectedRows); + + // partition filter on the partition with the original bucket count + List actualOldPartition = + CollectionUtil.iteratorToList( + batchTEnv + .executeSql("select * from " + tableName + " where c = 'old'") + .collect()); + List expectedOldPartition = new ArrayList<>(); + for (Row r : expectedRows) { + if ("old".equals(r.getField(2))) { + expectedOldPartition.add(r); + } + } + assertThat(actualOldPartition).containsExactlyInAnyOrderElementsOf(expectedOldPartition); + } + + @Test + void testUnionReadPkTableAcrossPartitionsWithDifferentBucketCounts() throws Exception { + String tableName = "rescale_bucket_pk_table"; + TablePath tablePath = TablePath.of(DEFAULT_DB, tableName); + createPartitionedPkTable(tablePath, OLD_BUCKET_NUM); + + // "old" partition keeps OLD_BUCKET_NUM; "new" partition (post-ALTER) uses NEW_BUCKET_NUM + createPartition(tablePath, "old"); + writeUpsertRows(tablePath, "old", 0); + alterBucketNum(tablePath); + createPartition(tablePath, "new"); + writeUpsertRows(tablePath, "new", 0); + + assertThat(bucketCountByPartitionName(tablePath)) + .containsEntry("old", OLD_BUCKET_NUM) + .containsEntry("new", NEW_BUCKET_NUM); + + JobClient jobClient = buildTieringJob(execEnv); + try { + long tableId = admin.getTableInfo(tablePath).get().getTableId(); + waitUntilPartitionBucketsSynced(tablePath, tableId); + + assertThat(totalBucketsOfPartition(tablePath, "old")).containsExactly(OLD_BUCKET_NUM); + assertThat(totalBucketsOfPartition(tablePath, "new")).containsExactly(NEW_BUCKET_NUM); + + // lake-only read: latest value per key across both partitions, verified by content + List expectedSnapshot = new ArrayList<>(); + for (String partition : new String[] {"old", "new"}) { + for (int i = 0; i < RECORDS_PER_ROUND; i++) { + expectedSnapshot.add(Row.of(i, "v" + i, partition)); + } + } + List lakeOnly = + CollectionUtil.iteratorToList( + batchTEnv.executeSql("select * from " + tableName).collect()); + assertThat(lakeOnly).containsExactlyInAnyOrderElementsOf(expectedSnapshot); + } finally { + jobClient.cancel().get(); + } + + // update existing keys after tiering so the read must merge lake snapshot with the fluss + // log tail (dedup by primary key), on both the old and new bucket-count partitions + List updates = new ArrayList<>(); + updates.add(row(0, "old-updated", "old")); + updates.add(row(0, "new-updated", "new")); + writeRows(tablePath, updates, false); + + // full expected state after the updates: key 0 of each partition carries the new value, + // every other key keeps its tiered value + List expectedMerged = new ArrayList<>(); + expectedMerged.add(Row.of(0, "old-updated", "old")); + expectedMerged.add(Row.of(0, "new-updated", "new")); + for (String partition : new String[] {"old", "new"}) { + for (int i = 1; i < RECORDS_PER_ROUND; i++) { + expectedMerged.add(Row.of(i, "v" + i, partition)); + } + } + List merged = + CollectionUtil.iteratorToList( + batchTEnv.executeSql("select * from " + tableName).collect()); + assertThat(merged).containsExactlyInAnyOrderElementsOf(expectedMerged); + + // partition filter on the partition that kept the original bucket count + List expectedOldOnly = new ArrayList<>(); + for (Row r : expectedMerged) { + if ("old".equals(r.getField(2))) { + expectedOldOnly.add(r); + } + } + List oldOnly = + CollectionUtil.iteratorToList( + batchTEnv + .executeSql("select * from " + tableName + " where c = 'old'") + .collect()); + assertThat(oldOnly).containsExactlyInAnyOrderElementsOf(expectedOldOnly); + } + + @Test + void testStreamUnionReadAcrossPartitionsWithDifferentBucketCounts() throws Exception { + String tableName = "rescale_bucket_stream_log_table"; + TablePath tablePath = TablePath.of(DEFAULT_DB, tableName); + createPartitionedLogTable(tablePath, OLD_BUCKET_NUM); + + createPartition(tablePath, "old"); + List expectedRows = new ArrayList<>(writeRows(tablePath, "old", 0)); + alterBucketNum(tablePath); + createPartition(tablePath, "new"); + expectedRows.addAll(writeRows(tablePath, "new", 0)); + + assertThat(bucketCountByPartitionName(tablePath)) + .containsEntry("old", OLD_BUCKET_NUM) + .containsEntry("new", NEW_BUCKET_NUM); + + JobClient jobClient = buildTieringJob(execEnv); + try { + long tableId = admin.getTableInfo(tablePath).get().getTableId(); + waitUntilPartitionBucketsSynced(tablePath, tableId); + assertThat(totalBucketsOfPartition(tablePath, "old")).containsExactly(OLD_BUCKET_NUM); + assertThat(totalBucketsOfPartition(tablePath, "new")).containsExactly(NEW_BUCKET_NUM); + + // streaming union read: read lake snapshot then keep streaming the fluss log tail + CloseableIterator iterator = + streamTEnv + .executeSql( + "select * from " + + tableName + + " /*+ OPTIONS('scan.partition.discovery.interval'='100ms') */") + .collect(); + // append more rows to both partitions after starting the stream + expectedRows.addAll(writeRows(tablePath, "old", RECORDS_PER_ROUND)); + expectedRows.addAll(writeRows(tablePath, "new", RECORDS_PER_ROUND)); + + List actual = collectRowsWithTimeout(iterator, expectedRows.size(), true); + assertThat(actual) + .containsExactlyInAnyOrderElementsOf( + expectedRows.stream().map(Row::toString).collect(Collectors.toList())); + } finally { + jobClient.cancel().get(); + } + } + + @Test + void testStreamUnionReadPkTableAcrossPartitionsWithDifferentBucketCounts() throws Exception { + String tableName = "rescale_bucket_stream_pk_table"; + TablePath tablePath = TablePath.of(DEFAULT_DB, tableName); + createPartitionedPkTable(tablePath, OLD_BUCKET_NUM); + + createPartition(tablePath, "old"); + writeUpsertRows(tablePath, "old", 0); + alterBucketNum(tablePath); + createPartition(tablePath, "new"); + writeUpsertRows(tablePath, "new", 0); + + assertThat(bucketCountByPartitionName(tablePath)) + .containsEntry("old", OLD_BUCKET_NUM) + .containsEntry("new", NEW_BUCKET_NUM); + + JobClient jobClient = buildTieringJob(execEnv); + try { + long tableId = admin.getTableInfo(tablePath).get().getTableId(); + waitUntilPartitionBucketsSynced(tablePath, tableId); + assertThat(totalBucketsOfPartition(tablePath, "old")).containsExactly(OLD_BUCKET_NUM); + assertThat(totalBucketsOfPartition(tablePath, "new")).containsExactly(NEW_BUCKET_NUM); + + // streaming union read on the PK table: the snapshot phase emits +I for the latest + // value of every key across both partitions, then keeps consuming the changelog tail + CloseableIterator iterator = + streamTEnv + .executeSql( + "select * from " + + tableName + + " /*+ OPTIONS('scan.partition.discovery.interval'='100ms') */") + .collect(); + + List expectedEvents = new ArrayList<>(); + for (int i = 0; i < RECORDS_PER_ROUND; i++) { + expectedEvents.add(Row.ofKind(RowKind.INSERT, i, "v" + i, "old").toString()); + expectedEvents.add(Row.ofKind(RowKind.INSERT, i, "v" + i, "new").toString()); + } + + // update one key in each partition after the stream started: the changelog must + // arrive as -U/+U from the fluss log tail on both the old (2-bucket) and the new + // (4-bucket) partition, proving the tail is subscribed by per-partition bucket range + List updates = new ArrayList<>(); + updates.add(row(0, "old-updated", "old")); + updates.add(row(0, "new-updated", "new")); + writeRows(tablePath, updates, false); + expectedEvents.add(Row.ofKind(RowKind.UPDATE_BEFORE, 0, "v0", "old").toString()); + expectedEvents.add( + Row.ofKind(RowKind.UPDATE_AFTER, 0, "old-updated", "old").toString()); + expectedEvents.add(Row.ofKind(RowKind.UPDATE_BEFORE, 0, "v0", "new").toString()); + expectedEvents.add( + Row.ofKind(RowKind.UPDATE_AFTER, 0, "new-updated", "new").toString()); + + List actual = collectRowsWithTimeout(iterator, expectedEvents.size(), true); + assertThat(actual).containsExactlyInAnyOrderElementsOf(expectedEvents); + } finally { + jobClient.cancel().get(); + } + } + + @Test + void testUnionReadLakeOnlyExpiredPartitionAfterRescale() throws Exception { + String tableName = "rescale_bucket_expired_log_table"; + TablePath tablePath = TablePath.of(DEFAULT_DB, tableName); + createPartitionedLogTable(tablePath, OLD_BUCKET_NUM); + + createPartition(tablePath, "old"); + List oldRows = writeRows(tablePath, "old", 0); + alterBucketNum(tablePath); + createPartition(tablePath, "new"); + List newRows = writeRows(tablePath, "new", 0); + + assertThat(bucketCountByPartitionName(tablePath)) + .containsEntry("old", OLD_BUCKET_NUM) + .containsEntry("new", NEW_BUCKET_NUM); + + // tier everything to the lake, then drop the rescaled "old" partition in Fluss so that it + // survives only in the lake (its files are stamped with OLD_BUCKET_NUM) + JobClient jobClient = buildTieringJob(execEnv); + try { + long tableId = admin.getTableInfo(tablePath).get().getTableId(); + waitUntilPartitionBucketsSynced(tablePath, tableId); + assertThat(totalBucketsOfPartition(tablePath, "old")).containsExactly(OLD_BUCKET_NUM); + } finally { + jobClient.cancel().get(); + } + + admin.dropPartition( + tablePath, new PartitionSpec(Collections.singletonMap("c", "old")), false) + .get(); + assertThat(admin.listPartitionInfos(tablePath).get()).hasSize(1); + + // union read: the expired "old" partition is served entirely from the lake (with its + // original bucket count), the "new" partition from Fluss+lake + List expected = new ArrayList<>(oldRows); + expected.addAll(newRows); + List actual = + CollectionUtil.iteratorToList( + batchTEnv.executeSql("select * from " + tableName).collect()); + assertThat(actual).containsExactlyInAnyOrderElementsOf(expected); + + // partition filter on the lake-only expired partition still returns its data + List oldActual = + CollectionUtil.iteratorToList( + batchTEnv + .executeSql("select * from " + tableName + " where c = 'old'") + .collect()); + assertThat(oldActual).containsExactlyInAnyOrderElementsOf(oldRows); + } + + private void createPartitionedLogTable(TablePath tablePath, int bucketNum) throws Exception { + TableDescriptor descriptor = + TableDescriptor.builder() + .schema( + Schema.newBuilder() + .column("a", DataTypes.INT()) + .column("b", DataTypes.STRING()) + .column("c", DataTypes.STRING()) + .build()) + .distributedBy(bucketNum, "a") + .partitionedBy("c") + .property(ConfigOptions.TABLE_DATALAKE_ENABLED.key(), "true") + .property(ConfigOptions.TABLE_DATALAKE_FRESHNESS, Duration.ofMillis(500)) + .build(); + createTable(tablePath, descriptor); + } + + private void createPartitionedPkTable(TablePath tablePath, int bucketNum) throws Exception { + TableDescriptor descriptor = + TableDescriptor.builder() + .schema( + Schema.newBuilder() + .column("a", DataTypes.INT()) + .column("b", DataTypes.STRING()) + .column("c", DataTypes.STRING()) + .primaryKey("a", "c") + .build()) + .distributedBy(bucketNum, "a") + .partitionedBy("c") + .property(ConfigOptions.TABLE_DATALAKE_ENABLED.key(), "true") + .property(ConfigOptions.TABLE_DATALAKE_FRESHNESS, Duration.ofMillis(500)) + .build(); + createTable(tablePath, descriptor); + } + + private void createPartition(TablePath tablePath, String value) throws Exception { + admin.createPartition( + tablePath, new PartitionSpec(Collections.singletonMap("c", value)), false) + .get(); + } + + private void alterBucketNum(TablePath tablePath) throws Exception { + admin.alterTable( + tablePath, + Collections.singletonList(TableChange.modifyBucketCount(NEW_BUCKET_NUM)), + false) + .get(); + } + + private void writeUpsertRows(TablePath tablePath, String partition, int keyOffset) + throws Exception { + List rows = new ArrayList<>(); + for (int i = keyOffset; i < keyOffset + RECORDS_PER_ROUND; i++) { + rows.add(row(i, "v" + i, partition)); + } + writeRows(tablePath, rows, false); + } + + private List writeRows(TablePath tablePath, String partition, int keyOffset) + throws Exception { + List rows = new ArrayList<>(); + List flinkRows = new ArrayList<>(); + for (int i = keyOffset; i < keyOffset + RECORDS_PER_ROUND; i++) { + rows.add(row(i, "v" + i, partition)); + flinkRows.add(Row.of(i, "v" + i, partition)); + } + writeRows(tablePath, rows, true); + return flinkRows; + } + + private Map bucketCountByPartitionName(TablePath tablePath) throws Exception { + List partitionInfos = admin.listPartitionInfos(tablePath).get(); + Map bucketCountByName = new java.util.HashMap<>(); + for (PartitionInfo partitionInfo : partitionInfos) { + bucketCountByName.put(partitionInfo.getPartitionName(), partitionInfo.getBucketCount()); + } + return bucketCountByName; + } + + private void waitUntilPartitionBucketsSynced(TablePath tablePath, long tableId) + throws Exception { + // empty buckets never get a tiering split (nothing to tier), so only wait for the lake + // sync marker on buckets that actually contain data + BucketOffsetsRetrieverImpl bucketOffsetsRetriever = + new BucketOffsetsRetrieverImpl(admin, tablePath); + Set tableBuckets = new HashSet<>(); + for (PartitionInfo partitionInfo : admin.listPartitionInfos(tablePath).get()) { + int bucketCount = partitionInfo.getBucketCount(); + List buckets = new ArrayList<>(); + for (int bucket = 0; bucket < bucketCount; bucket++) { + buckets.add(bucket); + } + Map latestOffsets = + bucketOffsetsRetriever.latestOffsets(partitionInfo.getPartitionName(), buckets); + for (int bucket = 0; bucket < bucketCount; bucket++) { + Long latestOffset = latestOffsets.get(bucket); + if (latestOffset != null && latestOffset > 0) { + tableBuckets.add( + new TableBucket(tableId, partitionInfo.getPartitionId(), bucket)); + } + } + } + waitUntilBucketsSynced(tableBuckets); + } + + private Set totalBucketsOfPartition(TablePath tablePath, String partition) + throws Exception { + FileStoreTable fileStoreTable = + (FileStoreTable) paimonCatalog.getTable(toPaimon(tablePath)); + List splits = + fileStoreTable + .newReadBuilder() + .withPartitionFilter(Collections.singletonMap("c", partition)) + .newScan() + .plan() + .splits(); + assertThat(splits).isNotEmpty(); + Set totalBuckets = new HashSet<>(); + for (Split split : splits) { + totalBuckets.add(((DataSplit) split).totalBuckets()); + } + return totalBuckets; + } +} diff --git a/fluss-lake/fluss-lake-paimon/src/test/java/org/apache/fluss/lake/paimon/tiering/PaimonTieringTest.java b/fluss-lake/fluss-lake-paimon/src/test/java/org/apache/fluss/lake/paimon/tiering/PaimonTieringTest.java index 19c59600f92..2d3365ac4b7 100644 --- a/fluss-lake/fluss-lake-paimon/src/test/java/org/apache/fluss/lake/paimon/tiering/PaimonTieringTest.java +++ b/fluss-lake/fluss-lake-paimon/src/test/java/org/apache/fluss/lake/paimon/tiering/PaimonTieringTest.java @@ -75,9 +75,11 @@ import java.util.Arrays; import java.util.Collections; import java.util.HashMap; +import java.util.HashSet; import java.util.LinkedHashMap; import java.util.List; import java.util.Map; +import java.util.Set; import java.util.stream.Stream; import static org.apache.fluss.lake.committer.LakeCommitter.FLUSS_LAKE_SNAP_BUCKET_OFFSET_PROPERTY; @@ -478,6 +480,125 @@ void testThreePartitionTiering() throws Exception { } } + @Test + void testTieringStampsPartitionBucketCountAcrossRounds() throws Exception { + // After ALTER bucket.num=8: files tiered for the "old" partition are stamped with its + // actual count 4 (writer override) while the "new" partition inherits the schema value 8; + // a second tiering round passes Paimon's native bucket-count check (historical 4 == + // writer 4), and all rows of both partitions stay readable via bucket-aware reads. + int schemaBucketCount = 8; + int oldPartitionBucketCount = 4; + int recordsPerBucketPerRound = 2; + TablePath tablePath = TablePath.of("paimon", "test_partition_bucket_count_stamp"); + createTable( + tablePath, + false, + true, + schemaBucketCount, + Collections.singletonMap(CoreOptions.BUCKET_KEY.key(), "c1")); + + TableDescriptor descriptor = + TableDescriptor.builder() + .schema( + org.apache.fluss.metadata.Schema.newBuilder() + .column("c1", org.apache.fluss.types.DataTypes.INT()) + .column("c2", org.apache.fluss.types.DataTypes.STRING()) + .column("c3", org.apache.fluss.types.DataTypes.STRING()) + .build()) + .partitionedBy("c3") + .distributedBy(schemaBucketCount, "c1") + .property(ConfigOptions.TABLE_DATALAKE_ENABLED, true) + .build(); + TableInfo tableInfo = + TableInfo.of(tablePath, 0, 1, descriptor, DEFAULT_REMOTE_DATA_DIR, 1L, 1L); + + // two independent tiering rounds against the SAME partitions; each round creates fresh + // writers and its own committer, exactly as TieringCommitOperator does + for (int round = 0; round < 2; round++) { + List paimonWriteResults = new ArrayList<>(); + // "old" partition: created before the ALTER, still routes by its original bucket count + for (int bucket = 0; bucket < oldPartitionBucketCount; bucket++) { + try (LakeWriter lakeWriter = + createLakeWriter( + tablePath, bucket, "old", 1L, tableInfo, oldPartitionBucketCount)) { + for (LogRecord logRecord : + genLogTableRecords("old", bucket, recordsPerBucketPerRound).f0) { + lakeWriter.write(logRecord); + } + paimonWriteResults.add(lakeWriter.complete()); + } + } + // "new" partition: created after the ALTER, routes by the schema bucket count + for (int bucket = 0; bucket < schemaBucketCount; bucket++) { + try (LakeWriter lakeWriter = + createLakeWriter(tablePath, bucket, "new", 2L, tableInfo, null)) { + for (LogRecord logRecord : + genLogTableRecords("new", bucket, recordsPerBucketPerRound).f0) { + lakeWriter.write(logRecord); + } + paimonWriteResults.add(lakeWriter.complete()); + } + } + try (LakeCommitter lakeCommitter = + createLakeCommitter(tablePath, tableInfo, new Configuration())) { + PaimonCommittable committable = lakeCommitter.toCommittable(paimonWriteResults); + lakeCommitter.commit(committable, Collections.emptyMap()); + } + } + + // files of BOTH rounds carry each partition's actual bucket count + assertThat(totalBucketsOfPartition(tablePath, "old")) + .containsExactly(oldPartitionBucketCount); + assertThat(totalBucketsOfPartition(tablePath, "new")).containsExactly(schemaBucketCount); + + // both partitions must stay fully readable through the bucket-aware read path, each + // holding exactly its own (rounds * records * bucketCount) rows + assertThat(rowCountOfPartition(tablePath, "old")) + .isEqualTo(2 * recordsPerBucketPerRound * oldPartitionBucketCount); + assertThat(rowCountOfPartition(tablePath, "new")) + .isEqualTo(2 * recordsPerBucketPerRound * schemaBucketCount); + } + + private int rowCountOfPartition(TablePath tablePath, String partition) throws Exception { + FileStoreTable fileStoreTable = + (FileStoreTable) paimonCatalog.getTable(toPaimon(tablePath)); + ReadBuilder readBuilder = + fileStoreTable + .newReadBuilder() + .withPartitionFilter(Collections.singletonMap("c3", partition)); + int rowCount = 0; + try (CloseableIterator iterator = + readBuilder + .newRead() + .createReader(readBuilder.newScan().plan().splits()) + .toCloseableIterator()) { + while (iterator.hasNext()) { + iterator.next(); + rowCount++; + } + } + return rowCount; + } + + private Set totalBucketsOfPartition(TablePath tablePath, String partition) + throws Exception { + FileStoreTable fileStoreTable = + (FileStoreTable) paimonCatalog.getTable(toPaimon(tablePath)); + List splits = + fileStoreTable + .newReadBuilder() + .withPartitionFilter(Collections.singletonMap("c3", partition)) + .newScan() + .plan() + .splits(); + assertThat(splits).isNotEmpty(); + Set totalBuckets = new HashSet<>(); + for (Split split : splits) { + totalBuckets.add(((DataSplit) split).totalBuckets()); + } + return totalBuckets; + } + @ParameterizedTest @MethodSource("snapshotExpireArgs") void testSnapshotExpiration( @@ -929,6 +1050,17 @@ private LakeWriter createLakeWriter( @Nullable Long partitionId, TableInfo tableInfo) throws IOException { + return createLakeWriter(tablePath, bucket, partition, partitionId, tableInfo, null); + } + + private LakeWriter createLakeWriter( + TablePath tablePath, + int bucket, + @Nullable String partition, + @Nullable Long partitionId, + TableInfo tableInfo, + @Nullable Integer partitionBucketCount) + throws IOException { return paimonLakeTieringFactory.createLakeWriter( new WriterInitContext() { @Override @@ -952,6 +1084,13 @@ public String partition() { public TableInfo tableInfo() { return tableInfo; } + + @Override + public int bucketCount() { + return partitionBucketCount != null + ? partitionBucketCount + : tableInfo.getNumBuckets(); + } }); } diff --git a/fluss-rpc/src/main/java/org/apache/fluss/rpc/netty/client/ServerConnection.java b/fluss-rpc/src/main/java/org/apache/fluss/rpc/netty/client/ServerConnection.java index d165cb2bc1f..8eb9d805c2b 100644 --- a/fluss-rpc/src/main/java/org/apache/fluss/rpc/netty/client/ServerConnection.java +++ b/fluss-rpc/src/main/java/org/apache/fluss/rpc/netty/client/ServerConnection.java @@ -26,6 +26,7 @@ import org.apache.fluss.exception.NetworkException; import org.apache.fluss.exception.RetriableAuthenticationException; import org.apache.fluss.exception.UnsupportedVersionException; +import org.apache.fluss.rpc.messages.AlterTableRequest; import org.apache.fluss.rpc.messages.ApiMessage; import org.apache.fluss.rpc.messages.ApiVersionsRequest; import org.apache.fluss.rpc.messages.ApiVersionsResponse; @@ -75,6 +76,7 @@ final class ServerConnection { private static final Logger LOG = LoggerFactory.getLogger(ServerConnection.class); private static final short HISTORICAL_PRODUCE_LOG_MIN_VERSION = 1; private static final short HISTORICAL_PUT_KV_MIN_VERSION = 3; + private static final short ALTER_BUCKET_COUNT_MIN_VERSION = 1; private final ServerNode node; @@ -381,8 +383,8 @@ private CompletableFuture doSend( } } - private void validateVersionCompatibility( - ApiKeys apiKey, short version, ApiMessage rawRequest) { + @VisibleForTesting + void validateVersionCompatibility(ApiKeys apiKey, short version, ApiMessage rawRequest) { if (apiKey == ApiKeys.PRODUCE_LOG && version < HISTORICAL_PRODUCE_LOG_MIN_VERSION) { ProduceLogRequest produceLogRequest = (ProduceLogRequest) rawRequest; if (hasHistoricalProduce(produceLogRequest)) { @@ -410,6 +412,20 @@ private void validateVersionCompatibility( + '.'); } } + + if (apiKey == ApiKeys.ALTER_TABLE && version < ALTER_BUCKET_COUNT_MIN_VERSION) { + AlterTableRequest alterTableRequest = (AlterTableRequest) rawRequest; + if (alterTableRequest.hasModifyBucketCount()) { + throw new UnsupportedVersionException( + "Modifying the bucket count requires ALTER_TABLE version " + + ALTER_BUCKET_COUNT_MIN_VERSION + + " or newer, but server " + + node + + " negotiated version " + + version + + '.'); + } + } } private void handleApiVersionsResponse(ApiMessage response, Throwable cause) { diff --git a/fluss-rpc/src/main/java/org/apache/fluss/rpc/protocol/ApiKeys.java b/fluss-rpc/src/main/java/org/apache/fluss/rpc/protocol/ApiKeys.java index ffdd1978de9..37d9f907449 100644 --- a/fluss-rpc/src/main/java/org/apache/fluss/rpc/protocol/ApiKeys.java +++ b/fluss-rpc/src/main/java/org/apache/fluss/rpc/protocol/ApiKeys.java @@ -91,7 +91,8 @@ public enum ApiKeys { DROP_ACLS(1041, 0, 0, PUBLIC), LAKE_TIERING_HEARTBEAT(1042, 0, 0, PRIVATE), CONTROLLED_SHUTDOWN(1043, 0, 0, PRIVATE), - ALTER_TABLE(1044, 0, 0, PUBLIC), + // Version 1: supports modifying the table distribution's bucket count. + ALTER_TABLE(1044, 0, 1, PUBLIC), DESCRIBE_CLUSTER_CONFIGS(1045, 0, 0, PUBLIC), ALTER_CLUSTER_CONFIGS(1046, 0, 0, PUBLIC), ADD_SERVER_TAG(1047, 0, 0, PUBLIC), diff --git a/fluss-rpc/src/main/java/org/apache/fluss/rpc/protocol/Errors.java b/fluss-rpc/src/main/java/org/apache/fluss/rpc/protocol/Errors.java index c7d074ed113..33d21db7f50 100644 --- a/fluss-rpc/src/main/java/org/apache/fluss/rpc/protocol/Errors.java +++ b/fluss-rpc/src/main/java/org/apache/fluss/rpc/protocol/Errors.java @@ -35,6 +35,7 @@ import org.apache.fluss.exception.IneligibleReplicaException; import org.apache.fluss.exception.InsufficientKvLeaderReplicaCapacityException; import org.apache.fluss.exception.InvalidAlterTableException; +import org.apache.fluss.exception.InvalidBucketRoutingException; import org.apache.fluss.exception.InvalidColumnProjectionException; import org.apache.fluss.exception.InvalidConfigException; import org.apache.fluss.exception.InvalidCoordinatorException; @@ -285,7 +286,12 @@ public enum Errors { HISTORICAL_PARTITION_THROTTLED( 73, "Historical partition request is throttled because too many historical requests are in flight.", - HistoricalPartitionThrottledException::new); + HistoricalPartitionThrottledException::new), + INVALID_BUCKET_ROUTING( + 74, + "The request's bucket routing information is missing or invalid. The client should " + + "refresh partition metadata and rebuild the request.", + InvalidBucketRoutingException::new); private static final Logger LOG = LoggerFactory.getLogger(Errors.class); diff --git a/fluss-rpc/src/main/proto/FlussApi.proto b/fluss-rpc/src/main/proto/FlussApi.proto index e0fd955d568..b778db8c24a 100644 --- a/fluss-rpc/src/main/proto/FlussApi.proto +++ b/fluss-rpc/src/main/proto/FlussApi.proto @@ -138,6 +138,11 @@ message AlterTableRequest { repeated PbDropColumn drop_columns = 5; repeated PbRenameColumn rename_columns = 6; repeated PbModifyColumn modify_columns = 7; + optional PbModifyBucketCount modify_bucket_count = 8; +} + +message PbModifyBucketCount { + required int32 new_bucket_count = 1; } message AlterTableResponse { @@ -157,6 +162,7 @@ message GetTableInfoResponse { required int64 created_time = 4; required int64 modified_time = 5; optional string remote_data_dir = 6; + optional int64 bucket_count_epoch = 7; } // list tables request and response @@ -295,6 +301,7 @@ message LimitScanRequest { optional int64 partition_id = 3; required int32 bucket_id = 4; required int32 limit = 5; + optional int32 routing_bucket_count = 6; } message LimitScanResponse{ @@ -315,6 +322,7 @@ message PbScanReqForBucket { required int32 bucket_id = 3; // If set, stops returning rows after this many records. optional int64 limit = 4; + optional int32 routing_bucket_count = 5; } message ScanKvRequest { @@ -399,6 +407,7 @@ message ListOffsetsRequest { optional int64 partition_id = 4; repeated int32 bucket_id = 5 [packed = true]; // it is recommended to use packed for repeated numerics to get more efficient encoding optional int64 startTimestamp = 6; + optional int32 routing_bucket_count = 7; } message ListOffsetsResponse { repeated PbListOffsetsRespForBucket buckets_resp = 1; @@ -507,10 +516,12 @@ message InitWriterResponse { message ListPartitionInfosRequest { required PbTablePath table_path = 1; optional PbPartitionSpec partial_partition_spec = 2; + optional bool include_system_partitions = 3; } message ListPartitionInfosResponse { repeated PbPartitionInfo partitions_info = 1; + optional bool system_partitions_included = 2; } // list remote log manifest entries (one per bucket of a table or partition) @@ -860,6 +871,8 @@ message PbTableMetadata { required int64 created_time = 6; required int64 modified_time = 7; optional string remote_data_dir = 8; + // A table-level, monotonically increasing version for bucket.num changes. + optional int64 bucket_count_epoch = 9; // TODO add a new filed 'deleted_table' to indicate this table is deleted in UpdateMetadataRequest. // trace by: https://github.com/apache/fluss/issues/981 @@ -871,6 +884,8 @@ message PbPartitionMetadata { required string partition_name = 2; required int64 partition_id = 3; repeated PbBucketMetadata bucket_metadata = 4; + // the actual bucket count for this partition, used for per-partition bucket rescale + optional int32 bucket_count = 5; } message PbBucketMetadata { @@ -891,6 +906,7 @@ message PbProduceLogReqForBucket { required bytes records = 3; // The original partition name for a historical write; unset for a normal write. optional string original_partition_name = 4; + optional int32 routing_bucket_count = 5; } message PbProduceLogRespForBucket { @@ -921,6 +937,7 @@ message PbFetchLogReqForBucket { // TODO leader epoch required int64 fetch_offset = 3; required int32 max_fetch_bytes = 4; + optional int32 routing_bucket_count = 5; } message PbFetchLogRespForTable { @@ -951,6 +968,9 @@ message PbPutKvReqForBucket { required bytes records = 3; // The original partition name for historical PK writes. It is unset for normal writes. optional string original_partition_name = 4; + // the bucket count the sender used to form bucket_id; the server compares it against the + // actual count to detect stale routing. Not authoritative metadata. + optional int32 routing_bucket_count = 5; } message PbPutKvRespForBucket { @@ -975,6 +995,7 @@ message PbLookupReqForBucket { repeated bytes keys = 3; // The original partition name for historical lookup. It is unset for normal lookup. optional string original_partition_name = 4; + optional int32 routing_bucket_count = 5; } message PbLookupRespForBucket { @@ -1000,6 +1021,7 @@ message PbPrefixLookupReqForBucket { optional int64 partition_id = 1; required int32 bucket_id = 2; repeated bytes keys = 3; + optional int32 routing_bucket_count = 4; } message PbPrefixLookupRespForBucket { @@ -1069,6 +1091,8 @@ message PbNotifyLeaderAndIsrReqForBucket { repeated int32 isr = 6 [packed = true]; required int32 bucket_epoch = 7; repeated int32 standby_replicas = 8 [packed = true]; + optional int32 bucket_count = 9; + optional int64 bucket_count_epoch = 10; } message PbNotifyLeaderAndIsrRespForBucket { @@ -1139,6 +1163,8 @@ message PbPartitionInfo { required int64 partition_id = 1; required PbPartitionSpec partition_spec = 2; optional string remote_data_dir = 3; + // the actual bucket count for this partition, used for per-partition bucket rescale + optional int32 bucket_count = 4; } message PbPartitionSpec { @@ -1349,6 +1375,7 @@ message PbKvSnapshotLeaseForBucket { message PbTableStatsReqForBucket { optional int64 partition_id = 1; required int32 bucket_id = 2; + optional int32 routing_bucket_count = 3; } message PbTableStatsRespForBucket { diff --git a/fluss-rpc/src/test/java/org/apache/fluss/rpc/netty/client/ServerConnectionTest.java b/fluss-rpc/src/test/java/org/apache/fluss/rpc/netty/client/ServerConnectionTest.java index da2d2e6e6d2..c5e10f39706 100644 --- a/fluss-rpc/src/test/java/org/apache/fluss/rpc/netty/client/ServerConnectionTest.java +++ b/fluss-rpc/src/test/java/org/apache/fluss/rpc/netty/client/ServerConnectionTest.java @@ -33,6 +33,7 @@ import org.apache.fluss.metrics.util.NOPMetricsGroup; import org.apache.fluss.rpc.TestingGatewayService; import org.apache.fluss.rpc.TestingTabletGatewayService; +import org.apache.fluss.rpc.messages.AlterTableRequest; import org.apache.fluss.rpc.messages.ApiMessage; import org.apache.fluss.rpc.messages.ApiVersionsRequest; import org.apache.fluss.rpc.messages.ApiVersionsResponse; @@ -295,6 +296,33 @@ public ChannelFuture connect(String host, int port) { .isInstanceOf(DisconnectException.class); } + @Test + void testRejectBucketCountChangeForOldServer() throws Exception { + ServerConnection connection = + new ServerConnection( + bootstrap, + serverNode, + TestingClientMetricGroup.newInstance(), + clientAuthenticator, + (con, ignore) -> {}); + try { + connection.validateVersionCompatibility( + ApiKeys.ALTER_TABLE, (short) 0, new AlterTableRequest()); + + AlterTableRequest bucketCountRequest = new AlterTableRequest(); + bucketCountRequest.setModifyBucketCount().setNewBucketCount(8); + assertThatThrownBy( + () -> + connection.validateVersionCompatibility( + ApiKeys.ALTER_TABLE, (short) 0, bucketCountRequest)) + .isInstanceOf(UnsupportedVersionException.class) + .hasMessageContaining("requires ALTER_TABLE version 1 or newer") + .hasMessageContaining("negotiated version 0"); + } finally { + connection.close().get(); + } + } + @Test void testRejectHistoricalWritesForOldServer() throws Exception { nettyServer.close(); diff --git a/fluss-rpc/src/test/java/org/apache/fluss/rpc/protocol/ApiErrorTest.java b/fluss-rpc/src/test/java/org/apache/fluss/rpc/protocol/ApiErrorTest.java index 9bd9d40080a..2d7a5faae80 100644 --- a/fluss-rpc/src/test/java/org/apache/fluss/rpc/protocol/ApiErrorTest.java +++ b/fluss-rpc/src/test/java/org/apache/fluss/rpc/protocol/ApiErrorTest.java @@ -18,6 +18,7 @@ package org.apache.fluss.rpc.protocol; import org.apache.fluss.exception.HistoricalPartitionThrottledException; +import org.apache.fluss.exception.InvalidBucketRoutingException; import org.apache.fluss.exception.NotEnoughReplicasException; import org.apache.fluss.exception.TimeoutException; import org.apache.fluss.exception.UnknownTableOrBucketException; @@ -89,6 +90,13 @@ private static Collection parameters() { Errors.HISTORICAL_PARTITION_THROTTLED, historicalLookupThrottledErrorMsg)); + String invalidBucketRoutingErrorMsg = "invalid bucket routing"; + arguments.add( + Arguments.of( + new InvalidBucketRoutingException(invalidBucketRoutingErrorMsg), + Errors.INVALID_BUCKET_ROUTING, + invalidBucketRoutingErrorMsg)); + return arguments; } } diff --git a/fluss-rpc/src/test/java/org/apache/fluss/rpc/protocol/ApiKeysTest.java b/fluss-rpc/src/test/java/org/apache/fluss/rpc/protocol/ApiKeysTest.java index bcaf1ad0b76..1ae5e55ad23 100644 --- a/fluss-rpc/src/test/java/org/apache/fluss/rpc/protocol/ApiKeysTest.java +++ b/fluss-rpc/src/test/java/org/apache/fluss/rpc/protocol/ApiKeysTest.java @@ -43,4 +43,9 @@ void testAllApiKeys() { void testSnapshotMetadataSupportsLayoutAwareClient() { assertThat(ApiKeys.GET_KV_SNAPSHOT_METADATA.highestSupportedVersion).isEqualTo((short) 1); } + + @Test + void testAlterTableSupportsBucketCountChange() { + assertThat(ApiKeys.ALTER_TABLE.highestSupportedVersion).isEqualTo((short) 1); + } } diff --git a/fluss-rust/bindings/python/test/conftest.py b/fluss-rust/bindings/python/test/conftest.py index 22de4d43e35..8ad2b649933 100644 --- a/fluss-rust/bindings/python/test/conftest.py +++ b/fluss-rust/bindings/python/test/conftest.py @@ -175,18 +175,10 @@ async def _wait(table_path, timeout=15, partition_name=None): table_path, [0], fluss.OffsetSpec.earliest() ) return - except (fluss.FlussError, Exception) as e: - # Catch "No leader found" or other errors that indicate the table/partition is still initializing - err_msg = str(e) - if any( - msg in err_msg - for msg in [ - "No leader found", - "Table not ready", - "Metadata not ready", - "not leader or follower", - ] - ): + except fluss.FlussError as e: + # Retriable means the table or partition is still initializing. A missing leader + # is a client-side error, which is never flagged retriable, so match it too. + if e.is_retriable or "No leader found" in str(e): await asyncio.sleep(1) continue raise diff --git a/fluss-rust/crates/fluss/proto/FlussApi.proto b/fluss-rust/crates/fluss/proto/FlussApi.proto index e0fd955d568..b778db8c24a 100644 --- a/fluss-rust/crates/fluss/proto/FlussApi.proto +++ b/fluss-rust/crates/fluss/proto/FlussApi.proto @@ -138,6 +138,11 @@ message AlterTableRequest { repeated PbDropColumn drop_columns = 5; repeated PbRenameColumn rename_columns = 6; repeated PbModifyColumn modify_columns = 7; + optional PbModifyBucketCount modify_bucket_count = 8; +} + +message PbModifyBucketCount { + required int32 new_bucket_count = 1; } message AlterTableResponse { @@ -157,6 +162,7 @@ message GetTableInfoResponse { required int64 created_time = 4; required int64 modified_time = 5; optional string remote_data_dir = 6; + optional int64 bucket_count_epoch = 7; } // list tables request and response @@ -295,6 +301,7 @@ message LimitScanRequest { optional int64 partition_id = 3; required int32 bucket_id = 4; required int32 limit = 5; + optional int32 routing_bucket_count = 6; } message LimitScanResponse{ @@ -315,6 +322,7 @@ message PbScanReqForBucket { required int32 bucket_id = 3; // If set, stops returning rows after this many records. optional int64 limit = 4; + optional int32 routing_bucket_count = 5; } message ScanKvRequest { @@ -399,6 +407,7 @@ message ListOffsetsRequest { optional int64 partition_id = 4; repeated int32 bucket_id = 5 [packed = true]; // it is recommended to use packed for repeated numerics to get more efficient encoding optional int64 startTimestamp = 6; + optional int32 routing_bucket_count = 7; } message ListOffsetsResponse { repeated PbListOffsetsRespForBucket buckets_resp = 1; @@ -507,10 +516,12 @@ message InitWriterResponse { message ListPartitionInfosRequest { required PbTablePath table_path = 1; optional PbPartitionSpec partial_partition_spec = 2; + optional bool include_system_partitions = 3; } message ListPartitionInfosResponse { repeated PbPartitionInfo partitions_info = 1; + optional bool system_partitions_included = 2; } // list remote log manifest entries (one per bucket of a table or partition) @@ -860,6 +871,8 @@ message PbTableMetadata { required int64 created_time = 6; required int64 modified_time = 7; optional string remote_data_dir = 8; + // A table-level, monotonically increasing version for bucket.num changes. + optional int64 bucket_count_epoch = 9; // TODO add a new filed 'deleted_table' to indicate this table is deleted in UpdateMetadataRequest. // trace by: https://github.com/apache/fluss/issues/981 @@ -871,6 +884,8 @@ message PbPartitionMetadata { required string partition_name = 2; required int64 partition_id = 3; repeated PbBucketMetadata bucket_metadata = 4; + // the actual bucket count for this partition, used for per-partition bucket rescale + optional int32 bucket_count = 5; } message PbBucketMetadata { @@ -891,6 +906,7 @@ message PbProduceLogReqForBucket { required bytes records = 3; // The original partition name for a historical write; unset for a normal write. optional string original_partition_name = 4; + optional int32 routing_bucket_count = 5; } message PbProduceLogRespForBucket { @@ -921,6 +937,7 @@ message PbFetchLogReqForBucket { // TODO leader epoch required int64 fetch_offset = 3; required int32 max_fetch_bytes = 4; + optional int32 routing_bucket_count = 5; } message PbFetchLogRespForTable { @@ -951,6 +968,9 @@ message PbPutKvReqForBucket { required bytes records = 3; // The original partition name for historical PK writes. It is unset for normal writes. optional string original_partition_name = 4; + // the bucket count the sender used to form bucket_id; the server compares it against the + // actual count to detect stale routing. Not authoritative metadata. + optional int32 routing_bucket_count = 5; } message PbPutKvRespForBucket { @@ -975,6 +995,7 @@ message PbLookupReqForBucket { repeated bytes keys = 3; // The original partition name for historical lookup. It is unset for normal lookup. optional string original_partition_name = 4; + optional int32 routing_bucket_count = 5; } message PbLookupRespForBucket { @@ -1000,6 +1021,7 @@ message PbPrefixLookupReqForBucket { optional int64 partition_id = 1; required int32 bucket_id = 2; repeated bytes keys = 3; + optional int32 routing_bucket_count = 4; } message PbPrefixLookupRespForBucket { @@ -1069,6 +1091,8 @@ message PbNotifyLeaderAndIsrReqForBucket { repeated int32 isr = 6 [packed = true]; required int32 bucket_epoch = 7; repeated int32 standby_replicas = 8 [packed = true]; + optional int32 bucket_count = 9; + optional int64 bucket_count_epoch = 10; } message PbNotifyLeaderAndIsrRespForBucket { @@ -1139,6 +1163,8 @@ message PbPartitionInfo { required int64 partition_id = 1; required PbPartitionSpec partition_spec = 2; optional string remote_data_dir = 3; + // the actual bucket count for this partition, used for per-partition bucket rescale + optional int32 bucket_count = 4; } message PbPartitionSpec { @@ -1349,6 +1375,7 @@ message PbKvSnapshotLeaseForBucket { message PbTableStatsReqForBucket { optional int64 partition_id = 1; required int32 bucket_id = 2; + optional int32 routing_bucket_count = 3; } message PbTableStatsRespForBucket { diff --git a/fluss-rust/crates/fluss/src/client/table/scanner.rs b/fluss-rust/crates/fluss/src/client/table/scanner.rs index 14c587e57d9..a8720b5ef79 100644 --- a/fluss-rust/crates/fluss/src/client/table/scanner.rs +++ b/fluss-rust/crates/fluss/src/client/table/scanner.rs @@ -2357,6 +2357,7 @@ impl LogFetcher { bucket_id: bucket.bucket_id(), fetch_offset: offset, max_fetch_bytes: self.fetch_max_bytes_for_bucket, + routing_bucket_count: None, }; fetch_log_req_for_buckets diff --git a/fluss-rust/crates/fluss/src/client/write/sender.rs b/fluss-rust/crates/fluss/src/client/write/sender.rs index 6cf5e8613da..ebbd42c120d 100644 --- a/fluss-rust/crates/fluss/src/client/write/sender.rs +++ b/fluss-rust/crates/fluss/src/client/write/sender.rs @@ -2314,6 +2314,7 @@ mod tests { created_time: 0, modified_time: 0, remote_data_dir: None, + bucket_count_epoch: None, } .encode(&mut body) .expect("encode GetTableInfoResponse"); diff --git a/fluss-rust/crates/fluss/src/metadata/partition.rs b/fluss-rust/crates/fluss/src/metadata/partition.rs index c63fe296c5c..5e00fe5657c 100644 --- a/fluss-rust/crates/fluss/src/metadata/partition.rs +++ b/fluss-rust/crates/fluss/src/metadata/partition.rs @@ -301,6 +301,7 @@ impl PartitionInfo { partition_id: self.partition_id, partition_spec: self.partition_spec.to_pb(), remote_data_dir: None, + bucket_count: None, } } diff --git a/fluss-rust/crates/fluss/src/metadata/table_stats.rs b/fluss-rust/crates/fluss/src/metadata/table_stats.rs index 53c6f72f35f..f459f7c2221 100644 --- a/fluss-rust/crates/fluss/src/metadata/table_stats.rs +++ b/fluss-rust/crates/fluss/src/metadata/table_stats.rs @@ -38,6 +38,7 @@ impl BucketStatsRequest { PbTableStatsReqForBucket { partition_id: self.partition_id, bucket_id: self.bucket_id, + routing_bucket_count: None, } } diff --git a/fluss-rust/crates/fluss/src/proto/fluss.rs b/fluss-rust/crates/fluss/src/proto/fluss.rs index 12295277351..4e4e06d4124 100644 --- a/fluss-rust/crates/fluss/src/proto/fluss.rs +++ b/fluss-rust/crates/fluss/src/proto/fluss.rs @@ -162,6 +162,13 @@ pub struct AlterTableRequest { pub rename_columns: ::prost::alloc::vec::Vec, #[prost(message, repeated, tag = "7")] pub modify_columns: ::prost::alloc::vec::Vec, + #[prost(message, optional, tag = "8")] + pub modify_bucket_count: ::core::option::Option, +} +#[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)] +pub struct PbModifyBucketCount { + #[prost(int32, required, tag = "1")] + pub new_bucket_count: i32, } #[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)] pub struct AlterTableResponse {} @@ -185,6 +192,8 @@ pub struct GetTableInfoResponse { pub modified_time: i64, #[prost(string, optional, tag = "6")] pub remote_data_dir: ::core::option::Option<::prost::alloc::string::String>, + #[prost(int64, optional, tag = "7")] + pub bucket_count_epoch: ::core::option::Option, } /// list tables request and response #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] @@ -365,6 +374,8 @@ pub struct LimitScanRequest { pub bucket_id: i32, #[prost(int32, required, tag = "5")] pub limit: i32, + #[prost(int32, optional, tag = "6")] + pub routing_bucket_count: ::core::option::Option, } #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] pub struct LimitScanResponse { @@ -392,6 +403,8 @@ pub struct PbScanReqForBucket { /// If set, stops returning rows after this many records. #[prost(int64, optional, tag = "4")] pub limit: ::core::option::Option, + #[prost(int32, optional, tag = "5")] + pub routing_bucket_count: ::core::option::Option, } #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] pub struct ScanKvRequest { @@ -512,6 +525,8 @@ pub struct ListOffsetsRequest { pub bucket_id: ::prost::alloc::vec::Vec, #[prost(int64, optional, tag = "6")] pub start_timestamp: ::core::option::Option, + #[prost(int32, optional, tag = "7")] + pub routing_bucket_count: ::core::option::Option, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct ListOffsetsResponse { @@ -661,11 +676,15 @@ pub struct ListPartitionInfosRequest { pub table_path: PbTablePath, #[prost(message, optional, tag = "2")] pub partial_partition_spec: ::core::option::Option, + #[prost(bool, optional, tag = "3")] + pub include_system_partitions: ::core::option::Option, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct ListPartitionInfosResponse { #[prost(message, repeated, tag = "1")] pub partitions_info: ::prost::alloc::vec::Vec, + #[prost(bool, optional, tag = "2")] + pub system_partitions_included: ::core::option::Option, } /// list remote log manifest entries (one per bucket of a table or partition) #[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)] @@ -1131,6 +1150,9 @@ pub struct PbTableMetadata { pub modified_time: i64, #[prost(string, optional, tag = "8")] pub remote_data_dir: ::core::option::Option<::prost::alloc::string::String>, + /// A table-level, monotonically increasing version for bucket.num changes. + #[prost(int64, optional, tag = "9")] + pub bucket_count_epoch: ::core::option::Option, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct PbPartitionMetadata { @@ -1143,6 +1165,9 @@ pub struct PbPartitionMetadata { pub partition_id: i64, #[prost(message, repeated, tag = "4")] pub bucket_metadata: ::prost::alloc::vec::Vec, + /// the actual bucket count for this partition, used for per-partition bucket rescale + #[prost(int32, optional, tag = "5")] + pub bucket_count: ::core::option::Option, } #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] pub struct PbBucketMetadata { @@ -1173,6 +1198,8 @@ pub struct PbProduceLogReqForBucket { /// The original partition name for a historical write; unset for a normal write. #[prost(string, optional, tag = "4")] pub original_partition_name: ::core::option::Option<::prost::alloc::string::String>, + #[prost(int32, optional, tag = "5")] + pub routing_bucket_count: ::core::option::Option, } #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] pub struct PbProduceLogRespForBucket { @@ -1219,6 +1246,8 @@ pub struct PbFetchLogReqForBucket { pub fetch_offset: i64, #[prost(int32, required, tag = "4")] pub max_fetch_bytes: i32, + #[prost(int32, optional, tag = "5")] + pub routing_bucket_count: ::core::option::Option, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct PbFetchLogRespForTable { @@ -1267,6 +1296,10 @@ pub struct PbPutKvReqForBucket { /// The original partition name for historical PK writes. It is unset for normal writes. #[prost(string, optional, tag = "4")] pub original_partition_name: ::core::option::Option<::prost::alloc::string::String>, + /// the bucket count the sender used to form bucket_id; the server compares it against the + /// actual count to detect stale routing. Not authoritative metadata. + #[prost(int32, optional, tag = "5")] + pub routing_bucket_count: ::core::option::Option, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct PbPutKvRespForBucket { @@ -1302,6 +1335,8 @@ pub struct PbLookupReqForBucket { /// The original partition name for historical lookup. It is unset for normal lookup. #[prost(string, optional, tag = "4")] pub original_partition_name: ::core::option::Option<::prost::alloc::string::String>, + #[prost(int32, optional, tag = "5")] + pub routing_bucket_count: ::core::option::Option, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct PbLookupRespForBucket { @@ -1338,6 +1373,8 @@ pub struct PbPrefixLookupReqForBucket { pub bucket_id: i32, #[prost(bytes = "bytes", repeated, tag = "3")] pub keys: ::prost::alloc::vec::Vec<::prost::bytes::Bytes>, + #[prost(int32, optional, tag = "4")] + pub routing_bucket_count: ::core::option::Option, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct PbPrefixLookupRespForBucket { @@ -1447,6 +1484,10 @@ pub struct PbNotifyLeaderAndIsrReqForBucket { pub bucket_epoch: i32, #[prost(int32, repeated, tag = "8")] pub standby_replicas: ::prost::alloc::vec::Vec, + #[prost(int32, optional, tag = "9")] + pub bucket_count: ::core::option::Option, + #[prost(int64, optional, tag = "10")] + pub bucket_count_epoch: ::core::option::Option, } #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] pub struct PbNotifyLeaderAndIsrRespForBucket { @@ -1550,6 +1591,9 @@ pub struct PbPartitionInfo { pub partition_spec: PbPartitionSpec, #[prost(string, optional, tag = "3")] pub remote_data_dir: ::core::option::Option<::prost::alloc::string::String>, + /// the actual bucket count for this partition, used for per-partition bucket rescale + #[prost(int32, optional, tag = "4")] + pub bucket_count: ::core::option::Option, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct PbPartitionSpec { @@ -1865,6 +1909,8 @@ pub struct PbTableStatsReqForBucket { pub partition_id: ::core::option::Option, #[prost(int32, required, tag = "2")] pub bucket_id: i32, + #[prost(int32, optional, tag = "3")] + pub routing_bucket_count: ::core::option::Option, } #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] pub struct PbTableStatsRespForBucket { diff --git a/fluss-rust/crates/fluss/src/rpc/message/alter_table.rs b/fluss-rust/crates/fluss/src/rpc/message/alter_table.rs index 64e9ad09dc6..72c706236f7 100644 --- a/fluss-rust/crates/fluss/src/rpc/message/alter_table.rs +++ b/fluss-rust/crates/fluss/src/rpc/message/alter_table.rs @@ -48,6 +48,7 @@ impl AlterTableRequest { drop_columns: drop_columns.iter().map(DropColumn::to_pb).collect(), rename_columns: rename_columns.iter().map(RenameColumn::to_pb).collect(), modify_columns: modify_columns.iter().map(ModifyColumn::to_pb).collect(), + modify_bucket_count: None, }, } } diff --git a/fluss-rust/crates/fluss/src/rpc/message/limit_scan.rs b/fluss-rust/crates/fluss/src/rpc/message/limit_scan.rs index 9dfa408eef7..9528eff31c2 100644 --- a/fluss-rust/crates/fluss/src/rpc/message/limit_scan.rs +++ b/fluss-rust/crates/fluss/src/rpc/message/limit_scan.rs @@ -42,6 +42,7 @@ impl LimitScanRequest { partition_id, bucket_id, limit, + routing_bucket_count: None, }; Self { diff --git a/fluss-rust/crates/fluss/src/rpc/message/list_offsets.rs b/fluss-rust/crates/fluss/src/rpc/message/list_offsets.rs index a690b82efe9..9d71d180d24 100644 --- a/fluss-rust/crates/fluss/src/rpc/message/list_offsets.rs +++ b/fluss-rust/crates/fluss/src/rpc/message/list_offsets.rs @@ -86,6 +86,7 @@ impl ListOffsetsRequest { partition_id, bucket_id: bucket_ids, start_timestamp: offset_spec.start_timestamp(), + routing_bucket_count: None, }, } } diff --git a/fluss-rust/crates/fluss/src/rpc/message/list_partition_infos.rs b/fluss-rust/crates/fluss/src/rpc/message/list_partition_infos.rs index 091af6da13d..83671aa039b 100644 --- a/fluss-rust/crates/fluss/src/rpc/message/list_partition_infos.rs +++ b/fluss-rust/crates/fluss/src/rpc/message/list_partition_infos.rs @@ -36,6 +36,7 @@ impl ListPartitionInfosRequest { inner_request: proto::ListPartitionInfosRequest { table_path: to_table_path(table_path), partial_partition_spec: partial_partition_spec.map(|s| s.to_pb()), + include_system_partitions: None, }, } } diff --git a/fluss-rust/crates/fluss/src/rpc/message/lookup.rs b/fluss-rust/crates/fluss/src/rpc/message/lookup.rs index 6a31d3c6fca..26f9e0370ce 100644 --- a/fluss-rust/crates/fluss/src/rpc/message/lookup.rs +++ b/fluss-rust/crates/fluss/src/rpc/message/lookup.rs @@ -44,6 +44,7 @@ impl LookupRequest { bucket_id, keys, original_partition_name: None, + routing_bucket_count: None, }, ) .collect(); diff --git a/fluss-rust/crates/fluss/src/rpc/message/prefix_lookup.rs b/fluss-rust/crates/fluss/src/rpc/message/prefix_lookup.rs index afafbbf71e9..befe9bbcac2 100644 --- a/fluss-rust/crates/fluss/src/rpc/message/prefix_lookup.rs +++ b/fluss-rust/crates/fluss/src/rpc/message/prefix_lookup.rs @@ -43,6 +43,7 @@ impl PrefixLookupRequest { partition_id, bucket_id, keys, + routing_bucket_count: None, }, ) .collect(); diff --git a/fluss-rust/crates/fluss/src/rpc/message/produce_log.rs b/fluss-rust/crates/fluss/src/rpc/message/produce_log.rs index 041e0adde01..eb31682c504 100644 --- a/fluss-rust/crates/fluss/src/rpc/message/produce_log.rs +++ b/fluss-rust/crates/fluss/src/rpc/message/produce_log.rs @@ -50,6 +50,7 @@ impl ProduceLogRequest { bucket_id: ready_batch.table_bucket.bucket_id(), records: ready_batch.write_batch.build()?, original_partition_name: None, + routing_bucket_count: None, }) } diff --git a/fluss-rust/crates/fluss/src/rpc/message/put_kv.rs b/fluss-rust/crates/fluss/src/rpc/message/put_kv.rs index c6ce9eb5cb4..5dcdaa7b7b3 100644 --- a/fluss-rust/crates/fluss/src/rpc/message/put_kv.rs +++ b/fluss-rust/crates/fluss/src/rpc/message/put_kv.rs @@ -52,6 +52,7 @@ impl PutKvRequest { bucket_id: ready_batch.table_bucket.bucket_id(), records: ready_batch.write_batch.build()?, original_partition_name: None, + routing_bucket_count: None, }) } diff --git a/fluss-server/src/main/java/org/apache/fluss/server/RpcServiceBase.java b/fluss-server/src/main/java/org/apache/fluss/server/RpcServiceBase.java index e2ef1094ee1..782c0629c28 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/RpcServiceBase.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/RpcServiceBase.java @@ -319,7 +319,8 @@ public CompletableFuture getTableInfo(GetTableInfoRequest .setTableId(tableInfo.getTableId()) .setRemoteDataDir(tableInfo.getRemoteDataDir()) .setCreatedTime(tableInfo.getCreatedTime()) - .setModifiedTime(tableInfo.getModifiedTime()); + .setModifiedTime(tableInfo.getModifiedTime()) + .setBucketCountEpoch(tableInfo.getBucketCountEpoch()); return CompletableFuture.completedFuture(response); } @@ -391,8 +392,15 @@ public CompletableFuture getLatestKvSnapshots( // get table id long tableId = tableInfo.getTableId(); int numBuckets = tableInfo.getNumBuckets(); - Long partitionId = - hasPartitionName ? getPartitionId(tablePath, request.getPartitionName()) : null; + Long partitionId = null; + if (hasPartitionName) { + PartitionRegistration partition = + getPartition(tablePath, request.getPartitionName()); + partitionId = partition.getPartitionId(); + numBuckets = + partition.getBucketCountOrDefault( + numBuckets, tableInfo.getBucketCountEpoch()); + } Map> snapshots; if (partitionId != null) { snapshots = zkClient.getPartitionLatestBucketSnapshot(partitionId); @@ -406,7 +414,7 @@ public CompletableFuture getLatestKvSnapshots( } } - private long getPartitionId(TablePath tablePath, String partitionName) { + private PartitionRegistration getPartition(TablePath tablePath, String partitionName) { Optional optPartitionRegistration; try { optPartitionRegistration = zkClient.getPartition(tablePath, partitionName); @@ -421,7 +429,7 @@ private long getPartitionId(TablePath tablePath, String partitionName) { "The partition '%s' of table '%s' does not exist.", partitionName, tablePath)); } - return optPartitionRegistration.get().getPartitionId(); + return optPartitionRegistration.get(); } @Override @@ -495,6 +503,12 @@ public CompletableFuture listPartitionInfos( TablePath tablePath = toTablePath(request.getTablePath()); authorizeTable(OperationType.DESCRIBE, tablePath); + // Read table metadata before reading partitions. This prevents a read spanning ALTER from + // combining a pre-ALTER PartitionRegistration (without bucketCount) with a post-ALTER + // TableInfo. + TableInfo tableInfo = metadataManager.getTable(tablePath); + List partitionKeys = tableInfo.getPartitionKeys(); + Map partitionRegistrations; if (request.hasPartialPartitionSpec()) { ResolvedPartitionSpec partitionSpecFromRequest = @@ -504,12 +518,21 @@ public CompletableFuture listPartitionInfos( } else { partitionRegistrations = metadataManager.listPartitions(tablePath); } - // TODO: Return the actual lake partitions instead of the internal historical partition. - partitionRegistrations.remove(HISTORICAL_PARTITION_VALUE); - TableInfo tableInfo = metadataManager.getTable(tablePath); - List partitionKeys = tableInfo.getPartitionKeys(); - return CompletableFuture.completedFuture( - toListPartitionInfosResponse(partitionKeys, partitionRegistrations)); + boolean includeSystemPartitions = + request.hasIncludeSystemPartitions() && request.isIncludeSystemPartitions(); + if (!includeSystemPartitions) { + partitionRegistrations.remove(HISTORICAL_PARTITION_VALUE); + } + ListPartitionInfosResponse response = + toListPartitionInfosResponse( + partitionKeys, + partitionRegistrations, + tableInfo.getNumBuckets(), + tableInfo.getBucketCountEpoch()); + if (includeSystemPartitions) { + response.setSystemPartitionsIncluded(true); + } + return CompletableFuture.completedFuture(response); } @Override diff --git a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/AutoPartitionManager.java b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/AutoPartitionManager.java index c2f3ce32c45..277b72bf30e 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/AutoPartitionManager.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/AutoPartitionManager.java @@ -261,7 +261,8 @@ void createHistoricalPartition(TableInfo tableInfo) { new ResolvedPartitionSpec( tableInfo.getPartitionKeys(), Collections.singletonList(HISTORICAL_PARTITION_VALUE)), - currentPartitions); + currentPartitions, + tableInfo.getNumBuckets()); } }); } @@ -278,9 +279,10 @@ void dropHistoricalPartition(TableInfo tableInfo) { && !currentPartitions.containsKey(HISTORICAL_PARTITION_VALUE)) { return; } + TablePath tablePath = tableInfo.getTablePath(); try { metadataManager.dropPartition( - tableInfo.getTablePath(), + tablePath, new ResolvedPartitionSpec( tableInfo.getPartitionKeys(), Collections.singletonList(HISTORICAL_PARTITION_VALUE)), @@ -288,13 +290,11 @@ void dropHistoricalPartition(TableInfo tableInfo) { if (currentPartitions != null) { currentPartitions.remove(HISTORICAL_PARTITION_VALUE); } - LOG.info( - "Deleted historical partition for table [{}].", - tableInfo.getTablePath()); + LOG.info("Deleted historical partition for table [{}].", tablePath); } catch (Exception e) { LOG.warn( - "Failed to delete historical partition for table [{}].", - tableInfo.getTablePath(), + "Failed to delete historical partition for table [{}] .", + tablePath, e); } }); @@ -483,31 +483,37 @@ private void createPartitions( } for (ResolvedPartitionSpec partition : partitionsToPreCreate) { - createPartition(tableInfo, partition, currentPartitions); + createPartition(tableInfo, partition, currentPartitions, tableInfo.getNumBuckets()); } } private void createPartition( TableInfo tableInfo, ResolvedPartitionSpec partition, - TreeMap> currentPartitions) { + TreeMap> currentPartitions, + int bucketCount) { TablePath tablePath = tableInfo.getTablePath(); long tableId = tableInfo.getTableId(); int replicaFactor = tableInfo.getTableConfig().getReplicationFactor(); TabletServerInfo[] servers = metadataCache.getLiveServers(); - long newKvLeaderReplicaCount = tableInfo.hasPrimaryKey() ? tableInfo.getNumBuckets() : 0; + long newKvLeaderReplicaCount = tableInfo.hasPrimaryKey() ? bucketCount : 0; try { replicaCapacityController.checkCanCreateKvLeaderReplicas(newKvLeaderReplicaCount); Map bucketAssignments = - generateAssignment(tableInfo.getNumBuckets(), replicaFactor, servers) - .getBucketAssignments(); + generateAssignment(bucketCount, replicaFactor, servers).getBucketAssignments(); PartitionAssignment partitionAssignment = new PartitionAssignment(tableInfo.getTableId(), bucketAssignments); String remoteDataDir = remoteDirDynamicLoader.getRemoteDirSelector().nextDataDir(); metadataManager.createPartition( - tablePath, tableId, remoteDataDir, partitionAssignment, partition, false); + tablePath, + tableId, + remoteDataDir, + partitionAssignment, + partition, + false, + bucketCount); currentPartitions.put(partition.getPartitionName(), null); LOG.info( "Auto partitioning created partition {} for table [{}].", partition, tablePath); @@ -612,8 +618,8 @@ private void dropPartitions( currentPartitions.headMap(lastRetainPartitionTime).entrySet().iterator(); while (iterator.hasNext()) { Map.Entry> entry = iterator.next(); - // Historical system partitions are managed explicitly by table configuration changes - // and Coordinator recovery, never by normal retention cleanup. + // Historical system partitions are managed explicitly by table configuration + // changes and Coordinator recovery, never by normal retention cleanup. if (HISTORICAL_PARTITION_VALUE.equals(entry.getKey())) { continue; } diff --git a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessor.java b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessor.java index cb57a20abb6..1e3b589e19a 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessor.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessor.java @@ -373,6 +373,11 @@ public int getCoordinatorEpoch() { return coordinatorContext.getCoordinatorEpoch(); } + /** The ZK version of the coordinator epoch znode, used for epoch fencing of ZK mutations. */ + public int getCoordinatorZkVersion() { + return coordinatorContext.getCoordinatorZkVersion(); + } + private void initCoordinatorContext() throws Exception { long start = System.currentTimeMillis(); // get all coordinator servers @@ -927,7 +932,8 @@ private void processSchemaChange(SchemaChangeEvent schemaChangeEvent) { oldTableInfo.getRemoteDataDir(), oldTableInfo.getComment().orElse(null), oldTableInfo.getCreatedTime(), - System.currentTimeMillis())); + System.currentTimeMillis(), + oldTableInfo.getBucketCountEpoch())); updateTabletServerMetadataCache( new HashSet<>(coordinatorContext.getLiveTabletServers().values()), @@ -1005,6 +1011,19 @@ private void postAlterTableProperties(TableInfo oldTableInfo, TableInfo newTable newAutoPartitionStrategy); autoPartitionManager.handleAutoPartitionStrategyChange( newTableInfo, oldAutoPartitionStrategy, newAutoPartitionStrategy); + } else if (newAutoPartitionStrategy.isAutoPartitionEnabled() + && oldTableInfo.getNumBuckets() != newTableInfo.getNumBuckets()) { + // bucket.num changed (e.g. via ALTER TABLE) without any auto-partition strategy + // change. Refresh the cached TableInfo so that newly auto-created partitions use + // the updated table-level bucket count as their per-partition bucket count. + LOG.info( + "Updating auto-partition metadata for table {} (tableId={}) after " + + "bucket.num changed from {} to {}.", + newTableInfo.getTablePath(), + newTableInfo.getTableId(), + oldTableInfo.getNumBuckets(), + newTableInfo.getNumBuckets()); + autoPartitionManager.updateAutoPartitionTables(newTableInfo); } // If standby replica config changed, trigger re-election for all online buckets diff --git a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorRequestBatch.java b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorRequestBatch.java index 6b416957858..b4bf7e90d41 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorRequestBatch.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorRequestBatch.java @@ -212,6 +212,9 @@ public void addNotifyLeaderRequestForTabletServers( TableBucket tableBucket, List bucketReplicas, LeaderAndIsr leaderAndIsr) { + Integer bucketCount = getBucketCount(tableBucket); + Long bucketCountEpoch = getBucketCountEpoch(tableBucket.getTableId()); + tabletServers.stream() .filter(s -> s >= 0 && !coordinatorContext.shuttingDownTabletServers().contains(s)) .forEach( @@ -226,17 +229,45 @@ public void addNotifyLeaderRequestForTabletServers( tablePath, tableBucket, bucketReplicas, - leaderAndIsr)); + leaderAndIsr, + bucketCount, + bucketCountEpoch)); notifyBucketLeaderAndIsr.put(tableBucket, notifyLeaderAndIsrForBucket); }); // TODO for these cases, we can send NotifyLeaderAndIsrRequest instead of another // updateMetadata request, trace by: https://github.com/apache/fluss/issues/983 - addUpdateMetadataRequestForTabletServers( - coordinatorContext.getLiveTabletServers().keySet(), - null, - null, - Collections.singleton(tableBucket)); + // A missing bucket count means the assignment required to build BucketMetadata is absent. + if (bucketCount != null) { + addUpdateMetadataRequestForTabletServers( + coordinatorContext.getLiveTabletServers().keySet(), + null, + null, + Collections.singleton(tableBucket)); + } + } + + /** + * The actual bucket count of the bucket's owning table/partition, or null when no assignment is + * in the coordinator context. The count is immutable per bucket, so it is carried with the + * activation instead of waiting for the metadata push. + */ + private @Nullable Integer getBucketCount(TableBucket tableBucket) { + Map> assignment; + if (tableBucket.getPartitionId() != null) { + assignment = + coordinatorContext.getPartitionAssignment( + new TablePartition( + tableBucket.getTableId(), tableBucket.getPartitionId())); + } else { + assignment = coordinatorContext.getTableAssignment(tableBucket.getTableId()); + } + return assignment.isEmpty() ? null : assignment.size(); + } + + private @Nullable Long getBucketCountEpoch(long tableId) { + TableInfo tableInfo = coordinatorContext.getTableInfoById(tableId); + return tableInfo == null ? null : tableInfo.getBucketCountEpoch(); } public void addStopReplicaRequestForTabletServers( @@ -685,6 +716,13 @@ private UpdateMetadataRequest buildUpdateMetadataRequest() { coordinatorContext.isPartitionQueuedForDeletion( new TablePartition(tableId, partitionId)); String partitionName = coordinatorContext.getPartitionName(partitionId); + // the partition assignment size is the partition's actual bucket count; + // null when the assignment is not in context + Map> partitionAssignment = + coordinatorContext.getPartitionAssignment( + new TablePartition(tableId, partitionId)); + Integer bucketCount = + partitionAssignment.isEmpty() ? null : partitionAssignment.size(); PartitionMetadata partitionMetadata; if (partitionName == null) { if (partitionQueuedForDeletion) { @@ -693,7 +731,8 @@ private UpdateMetadataRequest buildUpdateMetadataRequest() { tableId, DELETED_PARTITION_NAME, partitionId, - kvEntry.getValue()); + kvEntry.getValue(), + bucketCount); } else { throw new IllegalStateException( "Partition name is null for partition " + partitionId); @@ -706,7 +745,8 @@ private UpdateMetadataRequest buildUpdateMetadataRequest() { partitionQueuedForDeletion ? DELETED_PARTITION_ID : partitionId, - kvEntry.getValue()); + kvEntry.getValue(), + bucketCount); } // table partitionMetadataList.add(partitionMetadata); diff --git a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorService.java b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorService.java index 6e0c6f3c99d..ab34d7b44c9 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorService.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorService.java @@ -58,6 +58,7 @@ import org.apache.fluss.metadata.TableChange; import org.apache.fluss.metadata.TableDescriptor; import org.apache.fluss.metadata.TableInfo; +import org.apache.fluss.metadata.TablePartition; import org.apache.fluss.metadata.TablePath; import org.apache.fluss.row.encode.KvValueLayout; import org.apache.fluss.rpc.gateway.CoordinatorGateway; @@ -229,6 +230,7 @@ import static org.apache.fluss.server.utils.ServerRpcMessageUtils.makeDropAclsResponse; import static org.apache.fluss.server.utils.ServerRpcMessageUtils.makeListRemoteLogManifestsResponse; import static org.apache.fluss.server.utils.ServerRpcMessageUtils.toAlterTableConfigChanges; +import static org.apache.fluss.server.utils.ServerRpcMessageUtils.toAlterTableDistributionChanges; import static org.apache.fluss.server.utils.ServerRpcMessageUtils.toAlterTableSchemaChanges; import static org.apache.fluss.server.utils.ServerRpcMessageUtils.toDatabaseChanges; import static org.apache.fluss.server.utils.ServerRpcMessageUtils.toTableBucketOffsets; @@ -251,6 +253,7 @@ public final class CoordinatorService extends RpcServiceBase implements Coordina private final boolean kvTableAllowCreation; private final Supplier eventManagerSupplier; private final Supplier coordinatorEpochSupplier; + private final Supplier coordinatorZkVersionSupplier; private final CoordinatorMetadataCache metadataCache; private final Supplier snapshotStoreManagerSupplier; @@ -296,6 +299,8 @@ public CoordinatorService( () -> coordinatorEventProcessorSupplier.get().getCoordinatorEventManager(); this.coordinatorEpochSupplier = () -> coordinatorEventProcessorSupplier.get().getCoordinatorEpoch(); + this.coordinatorZkVersionSupplier = + () -> coordinatorEventProcessorSupplier.get().getCoordinatorZkVersion(); this.snapshotStoreManagerSupplier = () -> coordinatorEventProcessorSupplier.get().completedSnapshotStoreManager(); this.lakeTableTieringManager = lakeTableTieringManager; @@ -585,15 +590,20 @@ public CompletableFuture alterTable(AlterTableRequest reques toAlterTableConfigChanges(request.getConfigChangesList()); TablePropertyChanges tablePropertyChanges = toTablePropertyChanges(alterTableConfigChanges); List alterSchemaChanges = toAlterTableSchemaChanges(request); - - if (!alterSchemaChanges.isEmpty() && !alterTableConfigChanges.isEmpty()) { - // Only support one of alterTableConfigChanges and alterSchemaChanges for atomic change. + List alterDistributionChanges = + toAlterTableDistributionChanges(request); + + boolean hasConfigChanges = !alterTableConfigChanges.isEmpty(); + boolean hasSchemaChanges = !alterSchemaChanges.isEmpty(); + boolean hasDistributionChanges = !alterDistributionChanges.isEmpty(); + if ((hasConfigChanges && (hasSchemaChanges || hasDistributionChanges)) + || (hasSchemaChanges && hasDistributionChanges)) { throw new InvalidAlterTableException( "Table alteration can only be applied to one of the following: " - + "table properties or table schema."); + + "table properties, table schema, or table distribution."); } - if (!alterSchemaChanges.isEmpty()) { + if (hasSchemaChanges) { metadataManager.alterTableSchema( tablePath, alterSchemaChanges, @@ -601,7 +611,7 @@ public CompletableFuture alterTable(AlterTableRequest reques currentSession().getPrincipal()); } - if (!alterTableConfigChanges.isEmpty()) { + if (hasConfigChanges) { metadataManager.alterTableProperties( tablePath, alterTableConfigChanges, @@ -609,7 +619,19 @@ public CompletableFuture alterTable(AlterTableRequest reques request.isIgnoreIfNotExists(), currentSession().getPrincipal(), this::beforeTablePropertiesUpdate, - this::afterTablePropertiesUpdate); + this::afterTablePropertiesUpdate, + coordinatorZkVersionSupplier.get()); + } + + if (hasDistributionChanges) { + TableChange.ModifyBucketCount modifyBucketCount = + (TableChange.ModifyBucketCount) alterDistributionChanges.get(0); + metadataManager.alterBucketCount( + tablePath, + modifyBucketCount.getNewBucketCount(), + request.isIgnoreIfNotExists(), + currentSession().getPrincipal(), + coordinatorZkVersionSupplier.get()); } return CompletableFuture.completedFuture(new AlterTableResponse()); @@ -618,13 +640,13 @@ public CompletableFuture alterTable(AlterTableRequest reques private void beforeTablePropertiesUpdate(TableInfo currentTable, TableDescriptor updatedTable) { if (!currentTable.getTableConfig().isHistoricalPartitionEnabled() && isHistoricalPartitionEnabled(updatedTable)) { + TablePath tablePath = currentTable.getTablePath(); try { replicaCapacityController.checkCanCreateKvLeaderReplicas( getBucketCount(updatedTable)); - createHistoricalPartition( - currentTable.getTablePath(), currentTable.getTableId(), updatedTable); + createHistoricalPartition(tablePath, currentTable.getTableId(), updatedTable); } catch (Exception e) { - throw historicalPartitionEnableException(currentTable.getTablePath(), e); + throw historicalPartitionEnableException(tablePath, e); } } } @@ -636,11 +658,11 @@ private void afterTablePropertiesUpdate(TableInfo currentTable, TableDescriptor return; } + TablePath tablePath = currentTable.getTablePath(); try { - metadataManager.dropPartition( - currentTable.getTablePath(), historicalPartitionSpec(updatedTable), true); + metadataManager.dropPartition(tablePath, historicalPartitionSpec(updatedTable), true); } catch (Exception e) { - throw historicalPartitionDisableException(currentTable.getTablePath(), e); + throw historicalPartitionDisableException(tablePath, e); } } @@ -648,9 +670,9 @@ private void createHistoricalPartition( TablePath tablePath, long tableId, TableDescriptor tableDescriptor) { int replicaFactor = tableDescriptor.getReplicationFactor(); TabletServerInfo[] servers = metadataCache.getLiveServers(); + int bucketCount = getBucketCount(tableDescriptor); Map bucketAssignments = - generateAssignment(getBucketCount(tableDescriptor), replicaFactor, servers) - .getBucketAssignments(); + generateAssignment(bucketCount, replicaFactor, servers).getBucketAssignments(); PartitionAssignment partitionAssignment = new PartitionAssignment(tableId, bucketAssignments); String remoteDataDir = remoteDirDynamicLoader.getRemoteDirSelector().nextDataDir(); @@ -661,7 +683,8 @@ private void createHistoricalPartition( remoteDataDir, partitionAssignment, historicalPartitionSpec(tableDescriptor), - true); + true, + bucketCount); } private static ResolvedPartitionSpec historicalPartitionSpec(TableDescriptor tableDescriptor) { @@ -871,6 +894,8 @@ public CompletableFuture createPartition( authorizeTable(OperationType.WRITE, tablePath); CreatePartitionResponse response = new CreatePartitionResponse(); + // The table metadata (including bucket.num) is read fresh here, and the partition's + // registration persists its assignment and bucket count atomically in one ZK transaction TableInfo tableInfo = metadataManager.getTable(tablePath); if (!tableInfo.isPartitioned()) { throw new TableNotPartitionedException( @@ -919,7 +944,8 @@ public CompletableFuture createPartition( remoteDataDir, partitionAssignment, partitionToCreate, - request.isIgnoreIfNotExists()); + request.isIgnoreIfNotExists(), + tableInfo.getNumBuckets()); return CompletableFuture.completedFuture(response); } @@ -1055,6 +1081,17 @@ private CompletableFuture resolveNumBuckets(long tableId, @Nullable Lon AccessContextEvent event = new AccessContextEvent<>( ctx -> { + if (partitionId != null) { + // for partitions, the table-level bucket count may differ from the + // partition's actual bucket count after ALTER bucket.num; use the + // partition assignment size instead + Map> partitionAssignment = + ctx.getPartitionAssignment( + new TablePartition(tableId, partitionId)); + return partitionAssignment.isEmpty() + ? null + : partitionAssignment.size(); + } TablePath tablePath = ctx.getTablePathById(tableId); if (tablePath != null) { TableInfo tableInfo = ctx.getTableInfoById(tableId); diff --git a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/MetadataManager.java b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/MetadataManager.java index eb3947f270c..e8dcd33e931 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/MetadataManager.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/MetadataManager.java @@ -68,6 +68,7 @@ import java.util.ArrayList; import java.util.Collection; +import java.util.Collections; import java.util.HashMap; import java.util.HashSet; import java.util.Iterator; @@ -86,6 +87,12 @@ public class MetadataManager { private static final Logger LOG = LoggerFactory.getLogger(MetadataManager.class); + /** + * Max internal retries for the side-effect-free ALTER read-modify-write when a CAS/epoch + * conflict (BadVersionException) indicates a concurrent metadata change. + */ + private static final int MAX_ALTER_TABLE_RETRIES = 3; + private final ZooKeeperClient zookeeperClient; private final int maxPartitionNum; private final int maxBucketNum; @@ -508,6 +515,102 @@ private void syncSchemaChangesToLake( } } + private void propagateBucketCountToLake( + TablePath tablePath, + TableInfo tableInfo, + int newBucketCount, + FlussPrincipal flussPrincipal) { + if (!tableInfo.getTableConfig().isDataLakeEnabled()) { + return; + } + // Paimon only tracks a bucket count for Fixed Bucket tables (bucket-key non-empty). + if (tableInfo.getBucketKeys().isEmpty()) { + return; + } + LakeCatalog lakeCatalog = + lakeCatalogDynamicLoader.getLakeCatalogContainer().getLakeCatalog(); + if (lakeCatalog == null) { + throw new FlussRuntimeException( + "Cannot propagate ALTER bucket.num to the lake side for table " + + tablePath + + " because the Fluss cluster does not have a lake catalog configured."); + } + TableDescriptor currentDescriptor = tableInfo.toTableDescriptor(); + List bucketCountChange = + Collections.singletonList(TableChange.modifyBucketCount(newBucketCount)); + LakeCatalog.Context lakeCatalogContext = + new CoordinatorService.DefaultLakeCatalogContext( + false, + tableInfo.getLakeTablePath(), + flussPrincipal, + currentDescriptor, + currentDescriptor); + // Lake First: this runs BEFORE the Fluss ZK commit, so a lake failure aborts the ALTER + // with the Fluss side unchanged. + try { + lakeCatalog.alterTable( + tableInfo.getLakeTablePath(), bucketCountChange, lakeCatalogContext); + } catch (TableNotExistException e) { + throw new FlussRuntimeException( + "Lake table doesn't exist for lake-enabled table " + + tablePath + + ", which shouldn't happen. Please check if the lake table was deleted manually.", + e); + } catch (Exception e) { + throw new FlussRuntimeException( + String.format( + "ALTER bucket.num for table %s was aborted: propagating the new " + + "bucket count (%d) to the lake schema failed. The Fluss " + + "side was NOT changed. Re-run the same ALTER once the " + + "lake is reachable.", + tablePath, newBucketCount), + e); + } + } + + /** + * Validates an ALTER bucket.num request: only partitioned tables are supported and the new + * value must fall within [1, maxBucketNum]. Runs before the lake-side propagation so an invalid + * ALTER never mutates lake metadata. + */ + private void validateBucketNumRescale( + TablePath tablePath, TableInfo tableInfo, int newBucketNum) { + // Non-partitioned tables require creating new bucket assignments and initializing + // LogTablets on TabletServers, which is not yet implemented. + if (tableInfo.getPartitionKeys().isEmpty()) { + throw new InvalidAlterTableException( + String.format( + "Cannot alter 'bucket.num' on non-partitioned table %s. " + + "Non-partitioned table rescale is not yet supported.", + tablePath)); + } + // A rescaled table routes late writes of retired partitions through the historical + // partition, whose own fixed layout can diverge from post-rescale partitions; supporting + // the combination is left to future work. + if (tableInfo.getTableConfig().isHistoricalPartitionEnabled()) { + throw new InvalidAlterTableException( + String.format( + "Cannot alter 'bucket.num' on table %s with historical partition " + + "enabled. Altering 'bucket.num' on such tables is not " + + "supported yet.", + tablePath)); + } + if (newBucketNum < 1) { + throw new InvalidAlterTableException( + String.format( + "Cannot alter 'bucket.num' to %d on table %s. " + + "The bucket count must be at least 1.", + newBucketNum, tablePath)); + } + if (newBucketNum > maxBucketNum) { + throw new TooManyBucketsException( + String.format( + "Cannot alter 'bucket.num' to %d on table %s, " + + "exceeding the maximum of %d buckets per partition.", + newBucketNum, tablePath, maxBucketNum)); + } + } + /** Alters table properties and invokes the callbacks around the metadata update. */ public void alterTableProperties( TablePath tablePath, @@ -516,74 +619,297 @@ public void alterTableProperties( boolean ignoreIfNotExists, FlussPrincipal flussPrincipal, BiConsumer beforeUpdate, - BiConsumer afterUpdate) { - try { - // it throws TableNotExistException if the table or database not exists - TableRegistration tableReg = getTableRegistration(tablePath); - SchemaInfo schemaInfo = getLatestSchema(tablePath); - // we can't use MetadataManager#getTable here, because it will add the default - // lake options to the table properties, which may cause the validation failure - TableInfo tableInfo = tableReg.toTableInfo(tablePath, schemaInfo); + BiConsumer afterUpdate, + int coordinatorZkVersion) { + int attempt = 0; + while (true) { + try { + doAlterTablePropertiesOnce( + tablePath, + tableChanges, + tablePropertyChanges, + flussPrincipal, + beforeUpdate, + afterUpdate, + coordinatorZkVersion); + return; + } catch (TableNotExistException e) { + if (ignoreIfNotExists) { + return; + } + throw e; + } catch (KeeperException.NoNodeException e) { + if (!isTablePresent(tablePath)) { + if (ignoreIfNotExists) { + return; + } + throw new TableNotExistException("Table " + tablePath + " does not exist.", e); + } + retryAlterOrThrow(tablePath, ++attempt, e); + } catch (KeeperException.BadVersionException e) { + retryAlterOrThrow(tablePath, ++attempt, e); + } catch (RuntimeException e) { + throw e; + } catch (Exception e) { + throw new FlussRuntimeException( + "Failed to alter table properties: " + tablePath, e); + } + } + } - // validate the changes - validateAlterTableProperties(tableInfo, tablePropertyChanges.tableKeysToChange()); + /** Alters the default bucket count for newly created partitions. */ + public void alterBucketCount( + TablePath tablePath, + int newBucketCount, + boolean ignoreIfNotExists, + FlussPrincipal flussPrincipal, + int coordinatorZkVersion) { + int attempt = 0; + while (true) { + try { + ZooKeeperClient.VersionedData versionedTableReg = + getTableRegistrationWithVersion(tablePath); + TableRegistration tableReg = versionedTableReg.data(); + Map defaultTableLakeOptions = + lakeCatalogDynamicLoader + .getLakeCatalogContainer() + .getDefaultTableLakeOptions(); + Map tableLakeOptions = + defaultTableLakeOptions == null + ? null + : new HashMap<>(defaultTableLakeOptions); + removeSensitiveTableOptions(tableLakeOptions); + TableInfo tableInfo = + tableReg.toTableInfo( + tablePath, getLatestSchema(tablePath), tableLakeOptions); + validateBucketNumRescale(tablePath, tableInfo, newBucketCount); + if (newBucketCount == tableInfo.getNumBuckets()) { + return; + } - TableDescriptor tableDescriptor = tableInfo.toTableDescriptor(); - TableDescriptor newDescriptor = - getUpdatedTableDescriptor(tableDescriptor, tablePropertyChanges); + // Lake First: a lake failure aborts the ALTER with Fluss unchanged. The + // propagation is idempotent, so it re-runs after each metadata conflict. + propagateBucketCountToLake(tablePath, tableInfo, newBucketCount, flussPrincipal); - if (newDescriptor != null) { - // is to enable datalake for the table - if (isDataLakeEnabled(newDescriptor) && !isDataLakeEnabled(tableDescriptor)) { - // The table was created before cluster-level datalake was enabled. - // Backfill `table.datalake.format` before enabling datalake on the table - // so the updated table metadata stays consistent with the cluster setting. - if (!tableInfo.getTableConfig().getDataLakeFormat().isPresent()) { - DataLakeFormat dataLakeFormat = - lakeCatalogDynamicLoader - .getLakeCatalogContainer() - .getDataLakeFormat(); - if (dataLakeFormat == null) { - throw new InvalidAlterTableException( - "Cannot alter table " - + tablePath - + " in data lake, because the Fluss cluster doesn't enable datalake tables."); - } - newDescriptor = newDescriptor.withDataLakeFormat(dataLakeFormat); + // TODO: bucket-layout ALTERs should be rejected during a rolling server upgrade. + Map> backfills = + computePartitionBucketCountBackfill( + tablePath, + tableInfo.getNumBuckets(), + tableInfo.getBucketCountEpoch()); + zookeeperClient.updateTableWithPartitionBucketCountBackfill( + tablePath, + tableReg.newBucketCount(newBucketCount), + versionedTableReg.zkVersion(), + backfills, + coordinatorZkVersion); + return; + } catch (TableNotExistException e) { + if (ignoreIfNotExists) { + return; + } + throw e; + } catch (KeeperException.NoNodeException e) { + if (!isTablePresent(tablePath)) { + if (ignoreIfNotExists) { + return; } + throw new TableNotExistException("Table " + tablePath + " does not exist.", e); } - - // reuse the same validate logic with the createTable() method - validateTableDescriptor(newDescriptor); - - beforeUpdate.accept(tableInfo, newDescriptor); - - // pre alter table properties, e.g. create lake table in lake storage if it's to - // enable datalake for the table - preAlterTableProperties( - tablePath, tableDescriptor, newDescriptor, tableChanges, flussPrincipal); - // update the table to zk - TableRegistration updatedTableRegistration = - tableReg.newProperties( - newDescriptor.getProperties(), newDescriptor.getCustomProperties()); - zookeeperClient.updateTable(tablePath, updatedTableRegistration); - afterUpdate.accept(tableInfo, newDescriptor); - } else { - LOG.info( - "No properties changed when alter table {}, skip update table.", tablePath); + retryAlterOrThrow(tablePath, ++attempt, e); + } catch (KeeperException.BadVersionException e) { + retryAlterOrThrow(tablePath, ++attempt, e); + } catch (RuntimeException e) { + throw e; + } catch (Exception e) { + throw new FlussRuntimeException( + "Failed to alter bucket count for table: " + tablePath, e); } + } + } + + private void retryAlterOrThrow(TablePath tablePath, int attempt, Exception cause) { + if (attempt >= MAX_ALTER_TABLE_RETRIES) { + throw new FlussRuntimeException( + String.format( + "Failed to alter table properties for %s after %d retries " + + "due to concurrent metadata changes; please retry.", + tablePath, attempt), + cause); + } + LOG.info( + "Retrying ALTER on table {} due to a concurrent metadata change (attempt {}).", + tablePath, + attempt); + } + + /** Returns whether the table registration node still exists in ZooKeeper. */ + private boolean isTablePresent(TablePath tablePath) { + try { + return zookeeperClient.tableExist(tablePath); } catch (Exception e) { - if (e instanceof TableNotExistException) { - if (ignoreIfNotExists) { - return; + return false; + } + } + + /** + * One attempt of the property ALTER read-modify-write. The ZK write is CAS-guarded by the table + * version read here plus the coordinator epoch version. Since {@link #preAlterTableProperties} + * may apply external lake side effects, a version conflict is surfaced instead of being retried + * automatically. + */ + private void doAlterTablePropertiesOnce( + TablePath tablePath, + List tableChanges, + TablePropertyChanges tablePropertyChanges, + FlussPrincipal flussPrincipal, + BiConsumer beforeUpdate, + BiConsumer afterUpdate, + int coordinatorZkVersion) + throws Exception { + // it throws TableNotExistException if the table or database not exists + ZooKeeperClient.VersionedData versionedTableReg = + getTableRegistrationWithVersion(tablePath); + TableRegistration tableReg = versionedTableReg.data(); + int tableZkVersion = versionedTableReg.zkVersion(); + SchemaInfo schemaInfo = getLatestSchema(tablePath); + // we can't use MetadataManager#getTable here, because it will add the default + // lake options to the table properties, which may cause the validation failure + TableInfo tableInfo = tableReg.toTableInfo(tablePath, schemaInfo); + + // validate the changes + validateAlterTableProperties(tableInfo, tablePropertyChanges.tableKeysToChange()); + + TableDescriptor tableDescriptor = tableInfo.toTableDescriptor(); + TableDescriptor newDescriptor = + getUpdatedTableDescriptor(tableDescriptor, tablePropertyChanges); + + if (newDescriptor != null) { + // is to enable datalake for the table + if (isDataLakeEnabled(newDescriptor) && !isDataLakeEnabled(tableDescriptor)) { + // The table was created before cluster-level datalake was enabled. + // Backfill `table.datalake.format` before enabling datalake on the table + // so the updated table metadata stays consistent with the cluster setting. + if (!tableInfo.getTableConfig().getDataLakeFormat().isPresent()) { + DataLakeFormat dataLakeFormat = + lakeCatalogDynamicLoader.getLakeCatalogContainer().getDataLakeFormat(); + if (dataLakeFormat == null) { + throw new InvalidAlterTableException( + "Cannot alter table " + + tablePath + + " in data lake, because the Fluss cluster doesn't enable datalake tables."); + } + newDescriptor = newDescriptor.withDataLakeFormat(dataLakeFormat); } - throw (TableNotExistException) e; - } else if (e instanceof RuntimeException) { - throw (RuntimeException) e; - } else { + } + + // Enabling the historical partition on a rescaled table is unsupported: the + // historical partition would be created with the new table-level count while retired + // lake data still uses the pre-rescale layout. bucketCountEpoch never decreases, so + // this also rejects a table that was rescaled while the feature was temporarily + // disabled. Checked before validateTableDescriptor so the rescale rejection is not + // masked by unrelated option-dependency errors. + if (!tableInfo.getTableConfig().isHistoricalPartitionEnabled() + && Configuration.fromMap(newDescriptor.getProperties()) + .get(ConfigOptions.TABLE_DATALAKE_HISTORICAL_PARTITION_ENABLED) + && tableInfo.getBucketCountEpoch() > 0) { + throw new InvalidAlterTableException( + String.format( + "Cannot enable historical partition on table %s after " + + "'bucket.num' has been altered. Enabling it on a " + + "rescaled table is not supported yet.", + tablePath)); + } + + // reuse the same validate logic with the createTable() method + validateTableDescriptor(newDescriptor); + + beforeUpdate.accept(tableInfo, newDescriptor); + + // pre alter table properties, e.g. create lake table in lake storage if it's to + // enable datalake for the table. NOTE: this may have external (lake catalog) side + // effects and is therefore NOT safe to auto-retry. + preAlterTableProperties( + tablePath, tableDescriptor, newDescriptor, tableChanges, flussPrincipal); + + TableRegistration updatedTableRegistration = + tableReg.newProperties( + newDescriptor.getProperties(), newDescriptor.getCustomProperties()); + try { + zookeeperClient.updateTableWithPartitionBucketCountBackfill( + tablePath, + updatedTableRegistration, + tableZkVersion, + Collections.emptyMap(), + coordinatorZkVersion); + } catch (KeeperException.BadVersionException e) { throw new FlussRuntimeException( - "Failed to alter table properties: " + tablePath, e); + String.format( + "Concurrent metadata change while altering table %s; the change was " + + "not committed, please retry the ALTER.", + tablePath), + e); } + afterUpdate.accept(tableInfo, newDescriptor); + } else { + LOG.info("No properties changed when alter table {}, skip update table.", tablePath); + } + } + + /** + * Computes the bucket-count backfill for existing legacy partitions that do not persist their + * own count. Before the first bucket rescale, every partition has the same layout as the table, + * so the pre-ALTER table count is authoritative and no per-partition assignment lookup is + * needed. Nothing is written here: the caller commits the returned registrations together with + * the table-level bucket.num update in a single ZK transaction, CAS-guarded by the versions + * captured here. + * + *

Idempotent: partitions that already have a persisted bucket count are skipped. A missing + * count after the epoch has advanced indicates inconsistent metadata and cannot be inferred + * from the current table-level count. + */ + private Map> + computePartitionBucketCountBackfill( + TablePath tablePath, int oldBucketCount, long bucketCountEpoch) { + try { + Map> backfills = + new HashMap<>(); + Map> registrations = + zookeeperClient.getPartitionRegistrationsWithVersion(tablePath); + for (Map.Entry> entry : + registrations.entrySet()) { + String partitionName = entry.getKey(); + ZooKeeperClient.VersionedData versionedRegistration = + entry.getValue(); + PartitionRegistration reg = versionedRegistration.data(); + if (reg.getBucketCount() != null) { + // Already has bucket count persisted, skip. Idempotent so retries are safe. + continue; + } + if (bucketCountEpoch > 0) { + throw new InvalidAlterTableException( + String.format( + "Cannot alter 'bucket.num' on table %s: partition '%s' has no " + + "persisted bucket count after bucket count epoch %d.", + tablePath, partitionName, bucketCountEpoch)); + } + PartitionRegistration updatedReg = + new PartitionRegistration( + reg.getTableId(), + reg.getPartitionId(), + reg.getRemoteDataDir(), + oldBucketCount); + backfills.put( + partitionName, + new ZooKeeperClient.VersionedData<>( + updatedReg, versionedRegistration.zkVersion())); + } + return backfills; + } catch (RuntimeException e) { + throw e; + } catch (Exception e) { + throw new FlussRuntimeException( + "Failed to compute partition bucket count backfill for table: " + tablePath, e); } } @@ -792,6 +1118,24 @@ public TableRegistration getTableRegistration(TablePath tablePath) { return optionalTable.get(); } + /** + * Reads the table registration together with the ZK version of its znode, for a subsequent + * compare-and-set write. Throws {@link TableNotExistException} when the table does not exist. + */ + private ZooKeeperClient.VersionedData getTableRegistrationWithVersion( + TablePath tablePath) { + Optional> optionalTable; + try { + optionalTable = zookeeperClient.getTableWithVersion(tablePath); + } catch (Exception e) { + throw new RuntimeException(e); + } + if (!optionalTable.isPresent()) { + throw new TableNotExistException("Table '" + tablePath + "' does not exist."); + } + return optionalTable.get(); + } + public SchemaInfo getLatestSchema(TablePath tablePath) throws SchemaNotExistException { final int currentSchemaId; try { @@ -842,13 +1186,18 @@ public Set getPartitions(TablePath tablePath) { "Fail to get partitions from zookeeper for table " + tablePath); } + /** + * Creates a partition. The {@code bucketCount} is the table-level count the partition is + * created under; it becomes the partition's own persisted bucket count from here on. + */ public void createPartition( TablePath tablePath, long tableId, String remoteDataDir, PartitionAssignment partitionAssignment, ResolvedPartitionSpec partition, - boolean ignoreIfExists) { + boolean ignoreIfExists, + int bucketCount) { String partitionName = partition.getPartitionName(); Optional optionalPartitionRegistration = getOptionalPartitionRegistration(tablePath, partitionName); @@ -880,12 +1229,15 @@ public void createPartition( e); } - int bucketCount = partitionAssignment.getBucketAssignments().size(); - if (bucketCount > maxBucketNum) { + int assignmentBucketCount = partitionAssignment.getBucketAssignments().size(); + if (assignmentBucketCount > maxBucketNum) { throw new TooManyBucketsException( String.format( "Partition '%s' has %d buckets for table %s, exceeding the maximum of %d buckets per partition.", - partition.getPartitionName(), bucketCount, tablePath, maxBucketNum)); + partition.getPartitionName(), + assignmentBucketCount, + tablePath, + maxBucketNum)); } try { @@ -897,7 +1249,8 @@ public void createPartition( partitionAssignment, remoteDataDir, tablePath, - tableId); + tableId, + bucketCount); LOG.info( "Register partition {} to zookeeper for table [{}].", partitionName, tablePath); } catch (KeeperException.NodeExistsException nodeExistsException) { diff --git a/fluss-server/src/main/java/org/apache/fluss/server/entity/NotifyLeaderAndIsrData.java b/fluss-server/src/main/java/org/apache/fluss/server/entity/NotifyLeaderAndIsrData.java index e6702cdcaf6..87baebd59b5 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/entity/NotifyLeaderAndIsrData.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/entity/NotifyLeaderAndIsrData.java @@ -22,6 +22,8 @@ import org.apache.fluss.rpc.messages.NotifyLeaderAndIsrRequest; import org.apache.fluss.server.zk.data.LeaderAndIsr; +import javax.annotation.Nullable; + import java.util.List; /** The table bucket data of {@link NotifyLeaderAndIsrRequest}. */ @@ -30,16 +32,31 @@ public final class NotifyLeaderAndIsrData { private final TableBucket tableBucket; private final List replicas; private final LeaderAndIsr leaderAndIsr; + // null when a legacy coordinator omits the fields + private final @Nullable Integer bucketCount; + private final @Nullable Long bucketCountEpoch; public NotifyLeaderAndIsrData( PhysicalTablePath physicalTablePath, TableBucket tableBucket, List replicas, LeaderAndIsr leaderAndIsr) { + this(physicalTablePath, tableBucket, replicas, leaderAndIsr, null, null); + } + + public NotifyLeaderAndIsrData( + PhysicalTablePath physicalTablePath, + TableBucket tableBucket, + List replicas, + LeaderAndIsr leaderAndIsr, + @Nullable Integer bucketCount, + @Nullable Long bucketCountEpoch) { this.physicalTablePath = physicalTablePath; this.tableBucket = tableBucket; this.replicas = replicas; this.leaderAndIsr = leaderAndIsr; + this.bucketCount = bucketCount; + this.bucketCountEpoch = bucketCountEpoch; } public PhysicalTablePath getPhysicalTablePath() { @@ -93,4 +110,14 @@ public List getStandbyReplicas() { public int[] getStandbyReplicasArray() { return leaderAndIsr.standbyReplicas().stream().mapToInt(Integer::intValue).toArray(); } + + /** The actual bucket count of the owning table/partition, or null if not carried. */ + public @Nullable Integer getBucketCount() { + return bucketCount; + } + + /** The bucket layout epoch of the owning table, or null if not carried. */ + public @Nullable Long getBucketCountEpoch() { + return bucketCountEpoch; + } } diff --git a/fluss-server/src/main/java/org/apache/fluss/server/metadata/CoordinatorMetadataProvider.java b/fluss-server/src/main/java/org/apache/fluss/server/metadata/CoordinatorMetadataProvider.java index 9cae3977f7b..6f04908667b 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/metadata/CoordinatorMetadataProvider.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/metadata/CoordinatorMetadataProvider.java @@ -105,7 +105,12 @@ public Optional getPartitionMetadataFromCache( partitionId, ctx.getPartitionAssignment(new TablePartition(tableId, partitionId))); return Optional.of( - new PartitionMetadata(tableId, partitionName, partitionId, bucketMetadataList)); + new PartitionMetadata( + tableId, + partitionName, + partitionId, + bucketMetadataList, + bucketMetadataList.size())); } @Override diff --git a/fluss-server/src/main/java/org/apache/fluss/server/metadata/PartitionMetadata.java b/fluss-server/src/main/java/org/apache/fluss/server/metadata/PartitionMetadata.java index 74c17b3606c..362518f616b 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/metadata/PartitionMetadata.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/metadata/PartitionMetadata.java @@ -17,6 +17,8 @@ package org.apache.fluss.server.metadata; +import javax.annotation.Nullable; + import java.util.List; /** This entity used to describe the table's partition metadata. */ @@ -40,16 +42,27 @@ public class PartitionMetadata { private final String partitionName; private final long partitionId; private final List bucketMetadataList; + @Nullable private final Integer bucketCount; public PartitionMetadata( long tableId, String partitionName, long partitionId, List bucketMetadataList) { + this(tableId, partitionName, partitionId, bucketMetadataList, null); + } + + public PartitionMetadata( + long tableId, + String partitionName, + long partitionId, + List bucketMetadataList, + @Nullable Integer bucketCount) { this.tableId = tableId; this.partitionName = partitionName; this.partitionId = partitionId; this.bucketMetadataList = bucketMetadataList; + this.bucketCount = bucketCount; } public long getTableId() { @@ -67,4 +80,10 @@ public long getPartitionId() { public List getBucketMetadataList() { return bucketMetadataList; } + + /** Returns the actual bucket count for this partition, or null if not set (old data). */ + @Nullable + public Integer getBucketCount() { + return bucketCount; + } } diff --git a/fluss-server/src/main/java/org/apache/fluss/server/metadata/ServerMetadataSnapshot.java b/fluss-server/src/main/java/org/apache/fluss/server/metadata/ServerMetadataSnapshot.java index a3faeb90e9d..2ead046a951 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/metadata/ServerMetadataSnapshot.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/metadata/ServerMetadataSnapshot.java @@ -23,6 +23,7 @@ import org.apache.fluss.cluster.TabletServerInfo; import org.apache.fluss.metadata.PhysicalTablePath; import org.apache.fluss.metadata.TableBucket; +import org.apache.fluss.metadata.TablePartition; import org.apache.fluss.metadata.TablePath; import javax.annotation.Nullable; @@ -60,6 +61,14 @@ public class ServerMetadataSnapshot { // bucketMetadata> private final Map> bucketMetadataMapForPartitions; + // TablePartition -> bucket count; absent only when a legacy Coordinator omits it; a new + // Coordinator always sends the field. + private final Map partitionBucketCounts; + + // tableId -> bucketCountEpoch; a TabletServer keeps the latest value and ignores a lower + // epoch to prevent an older bucket layout (ALTER bucket.num) from replacing a newer one. + private final Map bucketCountEpochByTableId; + public ServerMetadataSnapshot( @Nullable ServerInfo coordinatorServer, Map aliveTabletServers, @@ -67,7 +76,9 @@ public ServerMetadataSnapshot( Map pathByTableId, Map partitionIdByPath, Map> bucketMetadataMapForTables, - Map> bucketMetadataMapForPartitions) { + Map> bucketMetadataMapForPartitions, + Map partitionBucketCounts, + Map bucketCountEpochByTableId) { this.coordinatorServer = coordinatorServer; this.aliveTabletServers = Collections.unmodifiableMap(aliveTabletServers); @@ -84,6 +95,8 @@ public ServerMetadataSnapshot( this.bucketMetadataMapForTables = Collections.unmodifiableMap(bucketMetadataMapForTables); this.bucketMetadataMapForPartitions = Collections.unmodifiableMap(bucketMetadataMapForPartitions); + this.partitionBucketCounts = Collections.unmodifiableMap(partitionBucketCounts); + this.bucketCountEpochByTableId = Collections.unmodifiableMap(bucketCountEpochByTableId); } /** Create an empty cluster instance with no nodes and no table-buckets. */ @@ -95,6 +108,8 @@ public static ServerMetadataSnapshot empty() { Collections.emptyMap(), Collections.emptyMap(), Collections.emptyMap(), + Collections.emptyMap(), + Collections.emptyMap(), Collections.emptyMap()); } @@ -159,6 +174,31 @@ public Map getBucketMetadataForPartition(long partition return bucketMetadataMapForPartitions.getOrDefault(partitionId, Collections.emptyMap()); } + /** + * Returns the actual bucket count for the given table partition, or null when the coordinator + * didn't send an explicit bucket count for it. + */ + public @Nullable Integer getPartitionBucketCount(TablePartition tablePartition) { + return partitionBucketCounts.get(tablePartition); + } + + public Map getPartitionBucketCounts() { + return partitionBucketCounts; + } + + /** + * Returns the bucket layout epoch for the given tableId, or empty if not known (legacy table + * without the field, read as 0). + */ + public OptionalLong getBucketCountEpoch(long tableId) { + Long epoch = bucketCountEpochByTableId.get(tableId); + return epoch == null ? OptionalLong.empty() : OptionalLong.of(epoch); + } + + public Map getBucketCountEpochByTableId() { + return bucketCountEpochByTableId; + } + public Map getPartitionIdByPath() { return partitionIdByPath; } diff --git a/fluss-server/src/main/java/org/apache/fluss/server/metadata/TabletServerMetadataCache.java b/fluss-server/src/main/java/org/apache/fluss/server/metadata/TabletServerMetadataCache.java index 8f307ae1437..c314e8fa7fb 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/metadata/TabletServerMetadataCache.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/metadata/TabletServerMetadataCache.java @@ -26,6 +26,7 @@ import org.apache.fluss.metadata.SchemaInfo; import org.apache.fluss.metadata.TableBucket; import org.apache.fluss.metadata.TableInfo; +import org.apache.fluss.metadata.TablePartition; import org.apache.fluss.metadata.TablePath; import org.apache.fluss.server.coordinator.MetadataManager; import org.apache.fluss.server.tablet.TabletServer; @@ -151,6 +152,14 @@ public void updateLatestSchema(long tableId, SchemaInfo schemaInfo) { tableId, (short) schemaInfo.getSchemaId(), schemaInfo.getSchema()); } + /** + * Returns the bucket layout epoch for the table, or empty if not known (legacy table without + * the field, read as 0). + */ + public OptionalLong getBucketCountEpoch(long tableId) { + return serverMetadataSnapshot.getBucketCountEpoch(tableId); + } + public Optional getPartitionMetadata(PhysicalTablePath partitionPath) { TablePath tablePath = partitionPath.getTablePath(); String partitionName = partitionPath.getPartitionName(); @@ -161,13 +170,19 @@ public Optional getPartitionMetadata(PhysicalTablePath partit if (tableIdOpt.isPresent() && partitionIdOpt.isPresent()) { long tableId = tableIdOpt.getAsLong(); long partitionId = partitionIdOpt.get(); + List bucketMetadataList = + new ArrayList<>(snapshot.getBucketMetadataForPartition(partitionId).values()); + // prefer the explicit bucket count sent by the coordinator; the merged bucket + // metadata list may be transiently partial during incremental updates + Integer bucketCount = + snapshot.getPartitionBucketCount(new TablePartition(tableId, partitionId)); return Optional.of( new PartitionMetadata( tableId, partitionName, partitionId, - new ArrayList<>( - snapshot.getBucketMetadataForPartition(partitionId).values()))); + bucketMetadataList, + bucketCount != null ? bucketCount : bucketMetadataList.size())); } else { return Optional.empty(); @@ -207,22 +222,20 @@ public Set updateClusterMetadata(ClusterMetadata clusterMetadata) { new HashMap<>(serverMetadataSnapshot.getTableIdByPath()); Map> bucketMetadataMapForTables = new HashMap<>(serverMetadataSnapshot.getBucketMetadataMapForTables()); + Map bucketCountEpochByTableId = + new HashMap<>(serverMetadataSnapshot.getBucketCountEpochByTableId()); for (TableMetadata tableMetadata : clusterMetadata.getTableMetadataList()) { TableInfo tableInfo = tableMetadata.getTableInfo(); TablePath tablePath = tableInfo.getTablePath(); long tableId = tableInfo.getTableId(); - // Update schema metadata. - // todo: apply schema id and schema info if needs - int schemaId = tableInfo.getSchemaId(); - Schema schema = tableInfo.getSchema(); - serverSchemaCache.updateLatestSchema(tableId, (short) schemaId, schema); if (tableId == DELETED_TABLE_ID) { Long removedTableId = tableIdByPath.remove(tablePath); if (removedTableId != null) { bucketMetadataMapForTables.remove(removedTableId); deletedTableIds.add(removedTableId); + bucketCountEpochByTableId.remove(removedTableId); } } else if (tablePath == DELETED_TABLE_PATH) { serverMetadataSnapshot @@ -230,7 +243,23 @@ public Set updateClusterMetadata(ClusterMetadata clusterMetadata) { .ifPresent(tableIdByPath::remove); bucketMetadataMapForTables.remove(tableId); deletedTableIds.add(tableId); + bucketCountEpochByTableId.remove(tableId); } else { + // Ignore an older UpdateMetadata to prevent an older bucket + // layout (ALTER bucket.num) from replacing a newer one. + long newEpoch = tableInfo.getBucketCountEpoch(); + long currentEpoch = bucketCountEpochByTableId.getOrDefault(tableId, 0L); + if (newEpoch < currentEpoch) { + continue; + } + bucketCountEpochByTableId.put(tableId, newEpoch); + + // Update schema metadata. + // todo: apply schema id and schema info if needs + int schemaId = tableInfo.getSchemaId(); + Schema schema = tableInfo.getSchema(); + serverSchemaCache.updateLatestSchema(tableId, (short) schemaId, schema); + tableIdByPath.put(tablePath, tableId); tableMetadata .getBucketMetadataList() @@ -255,6 +284,8 @@ public Set updateClusterMetadata(ClusterMetadata clusterMetadata) { Map> bucketMetadataMapForPartitions = new HashMap<>( serverMetadataSnapshot.getBucketMetadataMapForPartitions()); + Map partitionBucketCounts = + new HashMap<>(serverMetadataSnapshot.getPartitionBucketCounts()); for (PartitionMetadata partitionMetadata : clusterMetadata.getPartitionMetadataList()) { @@ -268,14 +299,22 @@ public Set updateClusterMetadata(ClusterMetadata clusterMetadata) { Long removedPartitionId = partitionIdByPath.remove(physicalTablePath); if (removedPartitionId != null) { bucketMetadataMapForPartitions.remove(removedPartitionId); + partitionBucketCounts + .keySet() + .removeIf(k -> k.getPartitionId() == removedPartitionId); } } else if (partitionName.equals(DELETED_PARTITION_NAME)) { serverMetadataSnapshot .getPhysicalTablePath(partitionId) .ifPresent(partitionIdByPath::remove); bucketMetadataMapForPartitions.remove(partitionId); + partitionBucketCounts + .keySet() + .removeIf(k -> k.getPartitionId() == partitionId); } else { partitionIdByPath.put(physicalTablePath, partitionId); + mergePartitionBucketCount( + partitionBucketCounts, tableId, partitionId, partitionMetadata); partitionMetadata .getBucketMetadataList() .forEach( @@ -298,7 +337,9 @@ public Set updateClusterMetadata(ClusterMetadata clusterMetadata) { newPathByTableId, partitionIdByPath, bucketMetadataMapForTables, - bucketMetadataMapForPartitions); + bucketMetadataMapForPartitions, + partitionBucketCounts, + bucketCountEpochByTableId); return deletedTableIds; }); } @@ -319,6 +360,8 @@ public void clearTableMetadata() { Collections.emptyMap(), Collections.emptyMap(), Collections.emptyMap(), + Collections.emptyMap(), + Collections.emptyMap(), Collections.emptyMap()); }); } @@ -340,6 +383,17 @@ public void updateTableMetadata(TableMetadata tableMetadata) { // Get current snapshot ServerMetadataSnapshot currentSnapshot = serverMetadataSnapshot; + // Ignore an older UpdateMetadata for this table to prevent it from + // overwriting newer state when messages arrive out of order. + long newEpoch = tableInfo.getBucketCountEpoch(); + long currentEpoch = + currentSnapshot + .getBucketCountEpochByTableId() + .getOrDefault(tableId, 0L); + if (newEpoch < currentEpoch) { + return; + } + // Create new maps based on current state Map tableIdByPath = new HashMap<>(currentSnapshot.getTableIdByPath()); @@ -367,6 +421,11 @@ public void updateTableMetadata(TableMetadata tableMetadata) { // Build pathByTableId from tableIdByPath tableIdByPath.forEach((path, id) -> pathByTableId.put(id, path)); + // Update epoch for this table + Map bucketCountEpochByTableId = + new HashMap<>(currentSnapshot.getBucketCountEpochByTableId()); + bucketCountEpochByTableId.put(tableId, tableInfo.getBucketCountEpoch()); + // Create new snapshot serverMetadataSnapshot = new ServerMetadataSnapshot( @@ -376,7 +435,9 @@ public void updateTableMetadata(TableMetadata tableMetadata) { pathByTableId, partitionIdByPath, bucketMetadataMapForTables, - bucketMetadataMapForPartitions); + bucketMetadataMapForPartitions, + currentSnapshot.getPartitionBucketCounts(), + bucketCountEpochByTableId); }); } @@ -427,6 +488,11 @@ public void updatePartitionMetadata(PartitionMetadata partitionMetadata) { } bucketMetadataMapForPartitions.put(partitionId, partitionBucketMetadata); + Map partitionBucketCounts = + new HashMap<>(currentSnapshot.getPartitionBucketCounts()); + mergePartitionBucketCount( + partitionBucketCounts, tableId, partitionId, partitionMetadata); + // Copy other existing data Map> bucketMetadataMapForTables = new HashMap<>(currentSnapshot.getBucketMetadataMapForTables()); @@ -443,10 +509,28 @@ public void updatePartitionMetadata(PartitionMetadata partitionMetadata) { pathByTableId, partitionIdByPath, bucketMetadataMapForTables, - bucketMetadataMapForPartitions); + bucketMetadataMapForPartitions, + partitionBucketCounts, + currentSnapshot.getBucketCountEpochByTableId()); }); } + /** + * Merges the coordinator-sent per-partition bucket count into the cache map. Coordinators of + * older versions do not send it; in that case the cache keeps no entry and readers fall back to + * the merged bucket metadata size (see {@link #getPartitionMetadata}). + */ + private static void mergePartitionBucketCount( + Map partitionBucketCounts, + long tableId, + long partitionId, + PartitionMetadata partitionMetadata) { + if (partitionMetadata.getBucketCount() != null) { + partitionBucketCounts.put( + new TablePartition(tableId, partitionId), partitionMetadata.getBucketCount()); + } + } + @VisibleForTesting public ServerSchemaCache getServerSchemaCache() { return serverSchemaCache; diff --git a/fluss-server/src/main/java/org/apache/fluss/server/metadata/ZkBasedMetadataProvider.java b/fluss-server/src/main/java/org/apache/fluss/server/metadata/ZkBasedMetadataProvider.java index e83241b6596..88e0da30877 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/metadata/ZkBasedMetadataProvider.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/metadata/ZkBasedMetadataProvider.java @@ -110,7 +110,8 @@ public List getPartitionsMetadataFromZK( tableId, partitionName, partitionId, - bucketMetadataList); + bucketMetadataList, + bucketMetadataList.size()); result.add(partitionMetadata); }); return result; diff --git a/fluss-server/src/main/java/org/apache/fluss/server/replica/Replica.java b/fluss-server/src/main/java/org/apache/fluss/server/replica/Replica.java index 079e1f96502..a74056ebc95 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/replica/Replica.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/replica/Replica.java @@ -193,6 +193,12 @@ public final class Replica { private final SchemaGetter schemaGetter; private volatile TableInfo tableInfo; + + // Routing state carried with activation: both values are immutable per bucket, so they are + // set once and never change. Null until the coordinator notifies them. + private volatile @Nullable Integer routingBucketCount; + private volatile @Nullable Long bucketCountEpoch; + private final boolean historicalPartition; // logFormat and arrowCompressionInfo are immutable and used in hot-path, so cache them here. private final LogFormat logFormat; @@ -449,6 +455,29 @@ public LogFormat getLogFormat() { return logFormat; } + /** + * Adopts the routing state carried by the notification. Both values are immutable per bucket, + * so unset fields never overwrite known ones. + */ + public void updateRoutingState(NotifyLeaderAndIsrData data) { + if (data.getBucketCount() != null) { + this.routingBucketCount = data.getBucketCount(); + } + if (data.getBucketCountEpoch() != null) { + this.bucketCountEpoch = data.getBucketCountEpoch(); + } + } + + /** The actual bucket count of the owning table/partition, or null if not yet notified. */ + public @Nullable Integer getRoutingBucketCount() { + return routingBucketCount; + } + + /** The bucket layout epoch of the owning table, or null if not yet notified. */ + public @Nullable Long getBucketCountEpoch() { + return bucketCountEpoch; + } + public void makeLeader(NotifyLeaderAndIsrData data) throws IOException { boolean leaderHWIncremented = inWriteLock( @@ -458,6 +487,7 @@ public void makeLeader(NotifyLeaderAndIsrData data) throws IOException { validateBucketEpoch(requestBucketEpoch); coordinatorEpoch = data.getCoordinatorEpoch(); + updateRoutingState(data); long currentTimeMs = clock.milliseconds(); // Updating the assignment and ISR state is safe if the bucket epoch is @@ -530,6 +560,7 @@ public boolean makeFollower(NotifyLeaderAndIsrData data) { validateBucketEpoch(requestBucketEpoch); coordinatorEpoch = data.getCoordinatorEpoch(); + updateRoutingState(data); updateAssignmentAndIsr( Collections.emptyList(), @@ -1853,6 +1884,13 @@ public long getOffset(RemoteLogManager remoteLogManager, ListOffsetsParam listOf return inReadLock( leaderIsrUpdateLock, () -> { + if (!isLeader()) { + throw new NotLeaderOrFollowerException( + String.format( + "Leader not local for bucket %s on tabletServer %d", + tableBucket, localTabletServerId)); + } + int offsetType = listOffsetsParam.getOffsetType(); if (offsetType == ListOffsetsParam.TIMESTAMP_OFFSET_TYPE) { return getOffsetByTimestamp(remoteLogManager, listOffsetsParam); diff --git a/fluss-server/src/main/java/org/apache/fluss/server/replica/ReplicaManager.java b/fluss-server/src/main/java/org/apache/fluss/server/replica/ReplicaManager.java index 64c986fadd2..c918147e4f7 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/replica/ReplicaManager.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/replica/ReplicaManager.java @@ -25,6 +25,7 @@ import org.apache.fluss.exception.ConfigException; import org.apache.fluss.exception.FencedLeaderEpochException; import org.apache.fluss.exception.HistoricalPartitionThrottledException; +import org.apache.fluss.exception.InvalidBucketRoutingException; import org.apache.fluss.exception.InvalidColumnProjectionException; import org.apache.fluss.exception.InvalidCoordinatorException; import org.apache.fluss.exception.InvalidPartitionException; @@ -2516,6 +2517,9 @@ protected Optional maybeCreateReplica(NotifyLeaderAndIsrData data) { clock, remoteLogManager, scannerManager); + // Initialize the routing state before the replica becomes visible, so a + // ready leader always has its routing bucket count ready. + replica.updateRoutingState(data); if (!existingLogTabletOpt.isPresent()) { localDiskManager.recordReplicaLoad(dataDir, isKvTable); } @@ -2547,6 +2551,63 @@ public Replica getReplicaOrException(TableBucket tableBucket) { } } + /** + * Validates the routing bucket count of a client request against the replica-local routing + * state. The target bucket is resolved first through {@link + * #getReplicaOrException(TableBucket)}, so an unknown, non-local, or offline replica fails + * immediately with its standard API exception. Missing or mismatched routing information fails + * with INVALID_BUCKET_ROUTING. Online followers skip the leader-only routing validation. + * + *

Only client requests may be validated; follower-initiated requests carry bucket ids + * assigned authoritatively by NotifyLeaderAndIsr and must skip this check. + */ + public void validateRoutingBucketCount(TableBucket tableBucket, int routingBucketCount) { + Replica replica = getReplicaOrException(tableBucket); + if (!replica.isLeader()) { + return; + } + + Integer actual = replica.getRoutingBucketCount(); + if (routingBucketCount <= 0) { + if (resolveBucketCountEpoch(replica) > 0) { + throw new InvalidBucketRoutingException( + "Invalid bucket routing for " + + tableBucket + + ": the request did not include a routing bucket count; expected " + + actual + + ". Refresh partition metadata, recompute the bucket id, and " + + "rebuild the request."); + } + return; + } + + if (actual == null) { + return; + } + if (routingBucketCount != actual) { + throw new InvalidBucketRoutingException( + "Invalid bucket routing for " + + tableBucket + + ": requested bucket count " + + routingBucketCount + + ", expected " + + actual + + ". Refresh partition metadata, recompute the bucket id, and rebuild " + + "the request."); + } + } + + /** + * Resolves the effective bucket layout epoch as the maximum of the replica-local value and the + * metadata cache: ALTER bucket.num advances only the cache, and the epoch is monotonic. + */ + private long resolveBucketCountEpoch(Replica replica) { + Long replicaEpoch = replica.getBucketCountEpoch(); + long cachedEpoch = + metadataCache.getBucketCountEpoch(replica.getTableBucket().getTableId()).orElse(0L); + return Math.max(replicaEpoch == null ? 0L : replicaEpoch, cachedEpoch); + } + public HostedReplica getReplica(TableBucket tableBucket) { return allReplicas.getOrDefault(tableBucket, new NoneReplica()); } diff --git a/fluss-server/src/main/java/org/apache/fluss/server/replica/historical/HistoricalLakeLookupManager.java b/fluss-server/src/main/java/org/apache/fluss/server/replica/historical/HistoricalLakeLookupManager.java index 5e309481431..8210d9ce8e5 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/replica/historical/HistoricalLakeLookupManager.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/replica/historical/HistoricalLakeLookupManager.java @@ -361,10 +361,16 @@ private LookupContext createLookupContext( LakeTableLookuper.LookupMetricRecorder lookupMetricRecorder) { TableBucket tableBucket = lookupData.tableBucket(); TablePath tablePath = tableInfo.getLakeTablePath(); + + // The request's bucket id only routes the request. It matches the lake layout only while + // the table was never rescaled; otherwise the lake lookuper resolves the bucket itself. + Integer lakeBucketId = + tableInfo.getBucketCountEpoch() == 0 ? tableBucket.getBucket() : null; + LakeTableLookuper.LookupContext lookupContext = new LakeTableLookuper.LookupContext( originalPartitionSpec, - tableBucket.getBucket(), + lakeBucketId, (short) schemaInfo.getSchemaId(), schemaInfo.getSchema().getRowType(), lookupMetricRecorder); diff --git a/fluss-server/src/main/java/org/apache/fluss/server/tablet/TabletService.java b/fluss-server/src/main/java/org/apache/fluss/server/tablet/TabletService.java index c091913acd2..ae19ae90eb6 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/tablet/TabletService.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/tablet/TabletService.java @@ -18,6 +18,7 @@ package org.apache.fluss.server.tablet; import org.apache.fluss.cluster.ServerType; +import org.apache.fluss.exception.ApiException; import org.apache.fluss.exception.AuthorizationException; import org.apache.fluss.exception.InvalidScanRequestException; import org.apache.fluss.exception.InvalidTableException; @@ -37,7 +38,9 @@ import org.apache.fluss.rpc.entity.LookupResultForBucket; import org.apache.fluss.rpc.entity.PrefixLookupResultForBucket; import org.apache.fluss.rpc.entity.ProduceLogResultForBucket; +import org.apache.fluss.rpc.entity.PutKvResultForBucket; import org.apache.fluss.rpc.entity.ResultForBucket; +import org.apache.fluss.rpc.entity.TableStatsResultForBucket; import org.apache.fluss.rpc.gateway.CoordinatorGateway; import org.apache.fluss.rpc.gateway.TabletServerGateway; import org.apache.fluss.rpc.messages.FetchLogRequest; @@ -64,6 +67,7 @@ import org.apache.fluss.rpc.messages.NotifyLeaderAndIsrResponse; import org.apache.fluss.rpc.messages.NotifyRemoteLogOffsetsRequest; import org.apache.fluss.rpc.messages.NotifyRemoteLogOffsetsResponse; +import org.apache.fluss.rpc.messages.PbFetchLogReqForTable; import org.apache.fluss.rpc.messages.PbScanReqForBucket; import org.apache.fluss.rpc.messages.PrefixLookupRequest; import org.apache.fluss.rpc.messages.PrefixLookupResponse; @@ -114,6 +118,7 @@ import javax.annotation.Nullable; +import java.util.ArrayList; import java.util.Collection; import java.util.HashMap; import java.util.List; @@ -124,6 +129,8 @@ import java.util.concurrent.ExecutorService; import java.util.function.BiFunction; import java.util.function.Consumer; +import java.util.function.Function; +import java.util.function.ToIntFunction; import java.util.stream.Collectors; import static org.apache.fluss.rpc.util.CommonRpcMessageUtils.hasHistoricalLookup; @@ -218,11 +225,35 @@ public void shutdown() {} @Override public CompletableFuture produceLog(ProduceLogRequest request) { authorizeTable(WRITE, request.getTableId()); - CompletableFuture response = new CompletableFuture<>(); + long tableId = request.getTableId(); + Map routingErrors = new HashMap<>(); + collectRoutingErrors( + request.getBucketsReqsList(), + pbBucket -> + new TableBucket( + tableId, + pbBucket.hasPartitionId() ? pbBucket.getPartitionId() : null, + pbBucket.getBucketId()), + pbBucket -> pbBucket.hasRoutingBucketCount() ? pbBucket.getRoutingBucketCount() : 0, + ProduceLogResultForBucket::new, + routingErrors); + List produceLogData = toProduceLogDataForBuckets(request); + if (!routingErrors.isEmpty()) { + produceLogData.removeIf( + bucketData -> routingErrors.containsKey(bucketData.tableBucket())); + if (produceLogData.isEmpty()) { + return CompletableFuture.completedFuture( + makeProduceLogResponse(routingErrors.values())); + } + } + + CompletableFuture response = new CompletableFuture<>(); UserContext userContext = new UserContext(currentSession().getPrincipal()); Consumer> responseCallback = - results -> response.complete(makeProduceLogResponse(results)); + results -> + response.complete( + makeProduceLogResponse(withRoutingErrors(results, routingErrors))); if (hasHistoricalProduce(request)) { replicaManager.appendHistoricalRecordsToLog( request.getTimeoutMs(), @@ -245,10 +276,51 @@ public CompletableFuture produceLog(ProduceLogRequest reques return response; } + /** + * Validates one request-scoped bucket and propagates any standard replica or routing {@link + * ApiException} to fail the single-bucket request immediately. + */ + private void validateRoutingBucketCountOrThrow( + TableBucket tableBucket, int routingBucketCount) { + replicaManager.validateRoutingBucketCount(tableBucket, routingBucketCount); + } + + /** + * Bucket-count validation applies to client requests only ({@code followerServerId < 0}). + * Server-internal replication traffic (follower fetch and follower listOffsets) never carries a + * bucket count, and a follower's bucket ids come from {@code NotifyLeaderAndIsr}, which is + * authoritative. Validating them against the leader's metadata cache would stall replication + * whenever that cache lags behind or the table has been rescaled. + */ + private static boolean isFromClient(int followerServerId) { + return followerServerId < 0; + } + @Override public CompletableFuture fetchLog(FetchLogRequest request) { Map fetchLogData = getFetchLogData(request); Map errorResponseMap = new HashMap<>(); + if (isFromClient(request.getFollowerServerId())) { + for (PbFetchLogReqForTable pbTable : request.getTablesReqsList()) { + long tableId = pbTable.getTableId(); + collectRoutingErrors( + pbTable.getBucketsReqsList(), + pbBucket -> + new TableBucket( + tableId, + pbBucket.hasPartitionId() + ? pbBucket.getPartitionId() + : null, + pbBucket.getBucketId()), + pbBucket -> + pbBucket.hasRoutingBucketCount() + ? pbBucket.getRoutingBucketCount() + : 0, + FetchLogResultForBucket::error, + errorResponseMap); + } + fetchLogData.keySet().removeAll(errorResponseMap.keySet()); + } Map interesting = // TODO: we should also authorize for follower, otherwise, users can mock follower // to skip the authorization. @@ -307,14 +379,37 @@ private static FetchParams getFetchParams(FetchLogRequest request) { @Override public CompletableFuture putKv(PutKvRequest request) { authorizeTable(WRITE, request.getTableId()); + long tableId = request.getTableId(); + Map routingErrors = new HashMap<>(); + collectRoutingErrors( + request.getBucketsReqsList(), + pbBucket -> + new TableBucket( + tableId, + pbBucket.hasPartitionId() ? pbBucket.getPartitionId() : null, + pbBucket.getBucketId()), + pbBucket -> pbBucket.hasRoutingBucketCount() ? pbBucket.getRoutingBucketCount() : 0, + PutKvResultForBucket::new, + routingErrors); List putKvData = toPutKvDataForBuckets(request); + if (!routingErrors.isEmpty()) { + putKvData.removeIf(bucketData -> routingErrors.containsKey(bucketData.tableBucket())); + if (putKvData.isEmpty()) { + return CompletableFuture.completedFuture(makePutKvResponse(routingErrors.values())); + } + } + // Get mergeMode from request, default to DEFAULT if not set MergeMode mergeMode = request.hasAggMode() ? MergeMode.fromValue(request.getAggMode()) : MergeMode.DEFAULT; CompletableFuture response = new CompletableFuture<>(); + Consumer> responseCallback = + results -> + response.complete( + makePutKvResponse(withRoutingErrors(results, routingErrors))); if (hasHistoricalPut(request)) { replicaManager.putHistoricalRecordsToKv( request.getTimeoutMs(), @@ -323,7 +418,7 @@ public CompletableFuture putKv(PutKvRequest request) { getTargetColumns(request), mergeMode, currentSession().getApiVersion(), - bucketResponse -> response.complete(makePutKvResponse(bucketResponse))); + responseCallback); } else { Map recordsByBucket = new HashMap<>(); putKvData.forEach( @@ -335,14 +430,25 @@ public CompletableFuture putKv(PutKvRequest request) { getTargetColumns(request), mergeMode, currentSession().getApiVersion(), - bucketResponse -> response.complete(makePutKvResponse(bucketResponse))); + responseCallback); } return response; } @Override public CompletableFuture lookup(LookupRequest request) { + long tableId = request.getTableId(); Map errorResponseMap = new HashMap<>(); + collectRoutingErrors( + request.getBucketsReqsList(), + pbBucket -> + new TableBucket( + tableId, + pbBucket.hasPartitionId() ? pbBucket.getPartitionId() : null, + pbBucket.getBucketId()), + pbBucket -> pbBucket.hasRoutingBucketCount() ? pbBucket.getRoutingBucketCount() : 0, + LookupResultForBucket::new, + errorResponseMap); CompletableFuture response = new CompletableFuture<>(); if (request.hasInsertIfNotExists() && request.isInsertIfNotExists()) { @@ -353,6 +459,10 @@ public CompletableFuture lookup(LookupRequest request) { + "historical partition lookup."); } Map> normalLookupData = toLookupData(request); + normalLookupData.keySet().removeAll(errorResponseMap.keySet()); + if (normalLookupData.isEmpty()) { + return CompletableFuture.completedFuture(makeLookupResponse(errorResponseMap)); + } replicaManager.lookups( request.isInsertIfNotExists(), request.getTimeoutMs(), @@ -365,11 +475,20 @@ public CompletableFuture lookup(LookupRequest request) { if (historicalLookupRequest) { List historicalLookupData = toHistoricalLookupData(request); authorizeTable(READ, request.getTableId()); + historicalLookupData.removeIf( + bucketData -> errorResponseMap.containsKey(bucketData.tableBucket())); + if (historicalLookupData.isEmpty()) { + return CompletableFuture.completedFuture(makeLookupResponse(errorResponseMap)); + } replicaManager.historicalLookups( historicalLookupData, - value -> response.complete(makeLookupResponse(value))); + value -> + response.complete( + makeLookupResponse( + withRoutingErrors(value, errorResponseMap)))); } else { Map> normalLookupData = toLookupData(request); + normalLookupData.keySet().removeAll(errorResponseMap.keySet()); Map> interesting = authorizeRequestData( READ, @@ -390,8 +509,20 @@ public CompletableFuture lookup(LookupRequest request) { @Override public CompletableFuture prefixLookup(PrefixLookupRequest request) { + long tableId = request.getTableId(); Map> prefixLookupData = toPrefixLookupData(request); Map errorResponseMap = new HashMap<>(); + collectRoutingErrors( + request.getBucketsReqsList(), + pbBucket -> + new TableBucket( + tableId, + pbBucket.hasPartitionId() ? pbBucket.getPartitionId() : null, + pbBucket.getBucketId()), + pbBucket -> pbBucket.hasRoutingBucketCount() ? pbBucket.getRoutingBucketCount() : 0, + PrefixLookupResultForBucket::new, + errorResponseMap); + prefixLookupData.keySet().removeAll(errorResponseMap.keySet()); Map> interesting = authorizeRequestData( READ, prefixLookupData, errorResponseMap, PrefixLookupResultForBucket::new); @@ -410,6 +541,12 @@ public CompletableFuture prefixLookup(PrefixLookupRequest @Override public CompletableFuture limitScan(LimitScanRequest request) { authorizeTable(READ, request.getTableId()); + validateRoutingBucketCountOrThrow( + new TableBucket( + request.getTableId(), + request.hasPartitionId() ? request.getPartitionId() : null, + request.getBucketId()), + request.hasRoutingBucketCount() ? request.getRoutingBucketCount() : 0); CompletableFuture response = new CompletableFuture<>(); replicaManager.limitScan( @@ -425,11 +562,35 @@ public CompletableFuture limitScan(LimitScanRequest request) @Override public CompletableFuture getTableStats(GetTableStatsRequest request) { authorizeTable(READ, request.getTableId()); + long tableId = request.getTableId(); + Map routingErrors = new HashMap<>(); + collectRoutingErrors( + request.getBucketsReqsList(), + pbBucket -> + new TableBucket( + tableId, + pbBucket.hasPartitionId() ? pbBucket.getPartitionId() : null, + pbBucket.getBucketId()), + pbBucket -> pbBucket.hasRoutingBucketCount() ? pbBucket.getRoutingBucketCount() : 0, + TableStatsResultForBucket::new, + routingErrors); + + List requestedBuckets = getTableStatsRequestData(request); + if (!routingErrors.isEmpty()) { + requestedBuckets.removeAll(routingErrors.keySet()); + if (requestedBuckets.isEmpty()) { + return CompletableFuture.completedFuture( + makeGetTableStatsResponse(new ArrayList<>(routingErrors.values()))); + } + } CompletableFuture response = new CompletableFuture<>(); replicaManager.getTableStats( - getTableStatsRequestData(request), - result -> response.complete(makeGetTableStatsResponse(result))); + requestedBuckets, + result -> + response.complete( + makeGetTableStatsResponse( + withRoutingErrors(result, routingErrors)))); return response; } @@ -496,6 +657,19 @@ public CompletableFuture stopReplica( @Override public CompletableFuture listOffsets(ListOffsetsRequest request) { authorizeTable(DESCRIBE, request.getTableId()); + if (isFromClient(request.getFollowerServerId())) { + // Unlike the per-bucket requests, both the partition and the routing count are + // request-scoped here, so every requested bucket shares one bucket layout: a stale + // route invalidates the whole request at once and failing it as a whole is exact. + Long partitionId = request.hasPartitionId() ? request.getPartitionId() : null; + int routingBucketCount = + request.hasRoutingBucketCount() ? request.getRoutingBucketCount() : 0; + for (int bucketId : request.getBucketIds()) { + validateRoutingBucketCountOrThrow( + new TableBucket(request.getTableId(), partitionId, bucketId), + routingBucketCount); + } + } CompletableFuture response = new CompletableFuture<>(); Set tableBuckets = getListOffsetsData(request); replicaManager.listOffsets( @@ -629,13 +803,16 @@ public CompletableFuture scanKv(ScanKvRequest request) { if (request.hasBucketScanReq()) { PbScanReqForBucket bucketReq = request.getBucketScanReq(); long tableId = bucketReq.getTableId(); - authorizeTable(READ, tableId); - TableBucket tableBucket = new TableBucket( tableId, bucketReq.hasPartitionId() ? bucketReq.getPartitionId() : null, bucketReq.getBucketId()); + validateRoutingBucketCountOrThrow( + tableBucket, + bucketReq.hasRoutingBucketCount() ? bucketReq.getRoutingBucketCount() : 0); + authorizeTable(READ, tableId); + Long limit = bucketReq.hasLimit() ? bucketReq.getLimit() : null; Replica replica = replicaManager.getReplicaOrException(tableBucket); @@ -906,6 +1083,50 @@ private Map authorizeRequestData( return interesting; } + /** + * Records a per-bucket error for every request bucket whose routing validation fails. This + * includes stale bucket counts and the standard unknown, non-local, or offline replica errors + * raised while resolving the target bucket. + * + *

A failure on one bucket does not invalidate other buckets in the same batched request, so + * each {@link ApiException} is converted into the corresponding bucket result. Request-scoped + * validation (for example {@link #listOffsets}) instead propagates the exception and fails the + * whole request. + */ + private void collectRoutingErrors( + List

bucketReqs, + Function toTableBucket, + ToIntFunction

routingBucketCountOf, + BiFunction resultCreator, + Map errorsOut) { + for (P bucketReq : bucketReqs) { + TableBucket tableBucket = toTableBucket.apply(bucketReq); + try { + replicaManager.validateRoutingBucketCount( + tableBucket, routingBucketCountOf.applyAsInt(bucketReq)); + } catch (ApiException e) { + errorsOut.put( + tableBucket, resultCreator.apply(tableBucket, ApiError.fromThrowable(e))); + } + } + } + + /** + * Appends routing errors to the results produced for the accepted buckets, so that the response + * covers every bucket the client asked about. Returns {@code results} untouched when every + * bucket has valid routing information. + */ + private static List withRoutingErrors( + List results, Map routingErrors) { + if (routingErrors.isEmpty()) { + return results; + } + List merged = new ArrayList<>(results.size() + routingErrors.size()); + merged.addAll(results); + merged.addAll(routingErrors.values()); + return merged; + } + private Set filterAuthorizedTables( Collection tableBuckets, OperationType operationType) { return tableBuckets.stream() diff --git a/fluss-server/src/main/java/org/apache/fluss/server/utils/ServerRpcMessageUtils.java b/fluss-server/src/main/java/org/apache/fluss/server/utils/ServerRpcMessageUtils.java index 079bb4c9947..242f92a0ea4 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/utils/ServerRpcMessageUtils.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/utils/ServerRpcMessageUtils.java @@ -323,6 +323,15 @@ public static List toAlterTableConfigChanges(List al .collect(Collectors.toList()); } + public static List toAlterTableDistributionChanges( + AlterTableRequest request) { + if (!request.hasModifyBucketCount()) { + return Collections.emptyList(); + } + return Collections.singletonList( + TableChange.modifyBucketCount(request.getModifyBucketCount().getNewBucketCount())); + } + private static DatabaseChange toDatabaseChange(PbAlterConfig pbAlterConfig) { AlterConfigOpType opType = AlterConfigOpType.from(pbAlterConfig.getOpType()); String configKey = pbAlterConfig.getConfigKey(); @@ -609,7 +618,8 @@ private static PbTableMetadata toPbTableMetadata(TableMetadata tableMetadata) { .setTableJson(tableInfo.toTableDescriptor().toJsonBytes()) .setRemoteDataDir(tableInfo.getRemoteDataDir()) .setCreatedTime(tableInfo.getCreatedTime()) - .setModifiedTime(tableInfo.getModifiedTime()); + .setModifiedTime(tableInfo.getModifiedTime()) + .setBucketCountEpoch(tableInfo.getBucketCountEpoch()); TablePath tablePath = tableInfo.getTablePath(); pbTableMetadata .setTablePath() @@ -628,6 +638,16 @@ private static PbPartitionMetadata toPbPartitionMetadata(PartitionMetadata parti .setPartitionName(partitionMetadata.getPartitionName()); pbPartitionMetadata.addAllBucketMetadatas( toPbBucketMetadata(partitionMetadata.getBucketMetadataList())); + Integer bucketCount = partitionMetadata.getBucketCount(); + int effectiveBucketCount = + bucketCount != null + ? bucketCount + : partitionMetadata.getBucketMetadataList().size(); + // 0 means the partition assignment is not known yet, not a zero-bucket layout; + // omitting the field keeps the client on its table-level fallback instead of 0. + if (effectiveBucketCount > 0) { + pbPartitionMetadata.setBucketCount(effectiveBucketCount); + } return pbPartitionMetadata; } @@ -684,7 +704,11 @@ private static TableMetadata toTableMetaData(PbTableMetadata pbTableMetadata) { ? pbTableMetadata.getRemoteDataDir() : null, pbTableMetadata.getCreatedTime(), - pbTableMetadata.getModifiedTime()); + pbTableMetadata.getModifiedTime(), + // legacy Coordinators predate bucketCountEpoch; read as 0 (never ALTERed) + pbTableMetadata.hasBucketCountEpoch() + ? pbTableMetadata.getBucketCountEpoch() + : 0L); List bucketMetadata = new ArrayList<>(); for (PbBucketMetadata pbBucketMetadata : pbTableMetadata.getBucketMetadatasList()) { @@ -713,7 +737,8 @@ private static PartitionMetadata toPartitionMetadata(PbPartitionMetadata pbParti pbPartitionMetadata.getPartitionId(), pbPartitionMetadata.getBucketMetadatasList().stream() .map(ServerRpcMessageUtils::toBucketMetadata) - .collect(Collectors.toList())); + .collect(Collectors.toList()), + pbPartitionMetadata.hasBucketCount() ? pbPartitionMetadata.getBucketCount() : null); } public static NotifyLeaderAndIsrRequest makeNotifyLeaderAndIsrRequest( @@ -747,6 +772,12 @@ public static PbNotifyLeaderAndIsrReqForBucket makeNotifyBucketLeaderAndIsr( .setPhysicalTablePath(fromPhysicalTablePath(physicalTablePath)) .setReplicas(notifyLeaderAndIsrData.getReplicasArray()) .setIsrs(notifyLeaderAndIsrData.getIsrArray()); + if (notifyLeaderAndIsrData.getBucketCount() != null) { + reqForBucket.setBucketCount(notifyLeaderAndIsrData.getBucketCount()); + } + if (notifyLeaderAndIsrData.getBucketCountEpoch() != null) { + reqForBucket.setBucketCountEpoch(notifyLeaderAndIsrData.getBucketCountEpoch()); + } return reqForBucket; } @@ -783,7 +814,11 @@ public static List getNotifyLeaderAndIsrRequestData( isr, standbyReplicas, request.getCoordinatorEpoch(), - reqForBucket.getBucketEpoch()))); + reqForBucket.getBucketEpoch()), + reqForBucket.hasBucketCount() ? reqForBucket.getBucketCount() : null, + reqForBucket.hasBucketCountEpoch() + ? reqForBucket.getBucketCountEpoch() + : null)); } return notifyLeaderAndIsrDataList; } @@ -1852,18 +1887,24 @@ public static NotifyKvSnapshotOffsetRequest makeNotifyKvSnapshotOffsetRequest( } public static ListPartitionInfosResponse toListPartitionInfosResponse( - List partitionKeys, Map partitionRegistrations) { + List partitionKeys, + Map partitionRegistrations, + int tableBucketCount, + long bucketCountEpoch) { ListPartitionInfosResponse listPartitionsResponse = new ListPartitionInfosResponse(); for (Map.Entry partitionRegistration : partitionRegistrations.entrySet()) { ResolvedPartitionSpec spec = ResolvedPartitionSpec.fromPartitionName( partitionKeys, partitionRegistration.getKey()); + PartitionRegistration partition = partitionRegistration.getValue(); listPartitionsResponse .addPartitionsInfo() - .setPartitionId(partitionRegistration.getValue().getPartitionId()) + .setPartitionId(partition.getPartitionId()) .setPartitionSpec(makePbPartitionSpec(spec)) - .setRemoteDataDir(partitionRegistration.getValue().getRemoteDataDir()); + .setRemoteDataDir(partition.getRemoteDataDir()) + .setBucketCount( + partition.getBucketCountOrDefault(tableBucketCount, bucketCountEpoch)); } return listPartitionsResponse; } diff --git a/fluss-server/src/main/java/org/apache/fluss/server/zk/ZkAsyncResponse.java b/fluss-server/src/main/java/org/apache/fluss/server/zk/ZkAsyncResponse.java index 62752014683..06fe9bf9792 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/zk/ZkAsyncResponse.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/zk/ZkAsyncResponse.java @@ -22,6 +22,8 @@ import org.apache.fluss.shaded.zookeeper3.org.apache.zookeeper.KeeperException.Code; import org.apache.fluss.shaded.zookeeper3.org.apache.zookeeper.data.Stat; +import javax.annotation.Nullable; + import java.util.List; import java.util.Optional; @@ -77,19 +79,28 @@ public boolean hasError() { public static class ZkGetDataResponse extends ZkAsyncResponse { private final byte[] data; + private final @Nullable Stat stat; - public ZkGetDataResponse(String path, Code resultCode, byte[] data) { + public ZkGetDataResponse(String path, Code resultCode, byte[] data, @Nullable Stat stat) { super(path, resultCode); this.data = data; + this.stat = stat; } public byte[] getData() { return data; } + public @Nullable Stat getStat() { + return stat; + } + public static ZkGetDataResponse create(CuratorEvent event) { return new ZkGetDataResponse( - event.getPath(), Code.get(event.getResultCode()), event.getData()); + event.getPath(), + Code.get(event.getResultCode()), + event.getData(), + event.getStat()); } } diff --git a/fluss-server/src/main/java/org/apache/fluss/server/zk/ZooKeeperClient.java b/fluss-server/src/main/java/org/apache/fluss/server/zk/ZooKeeperClient.java index 1911196095f..010b2e840ee 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/zk/ZooKeeperClient.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/zk/ZooKeeperClient.java @@ -222,6 +222,18 @@ public Optional getOrEmpty(String path) throws Exception { } } + /** + * Reads the znode data and captures its {@link Stat} (hence the ZK version) atomically. Used by + * compare-and-set callers that must write back with the exact version they read. + */ + private Optional getDataWithStat(String path, Stat stat) throws Exception { + try { + return Optional.of(zkClient.getData().storingStatIn(stat).forPath(path)); + } catch (KeeperException.NoNodeException e) { + return Optional.empty(); + } + } + public String getDefaultRemoteDataDir() { return defaultRemoteDataDir; } @@ -801,6 +813,19 @@ public Optional getTable(TablePath tablePath) throws Exceptio t -> t.remoteDataDir == null ? t.newRemoteDataDir(defaultRemoteDataDir) : t); } + /** + * Get the table registration together with the ZK version of its znode, so callers can perform + * a compare-and-set write (see {@link #updateTableWithPartitionBucketCountBackfill}). + */ + public Optional> getTableWithVersion(TablePath tablePath) + throws Exception { + Stat stat = new Stat(); + Optional bytes = getDataWithStat(TableZNode.path(tablePath), stat); + return bytes.map(TableZNode::decode) + .map(t -> t.remoteDataDir == null ? t.newRemoteDataDir(defaultRemoteDataDir) : t) + .map(t -> new VersionedData<>(t, stat.getVersion())); + } + /** Get the tables in ZK. */ public Map getTables(Collection tablePaths) throws Exception { @@ -1033,6 +1058,134 @@ public Optional getPartition(TablePath tablePath, String p -> p.getRemoteDataDir() == null ? p.newRemoteDataDir(defaultRemoteDataDir) : p); } + /** + * Get a partition registration together with the ZK version of its znode, so callers can + * perform a compare-and-set backfill (see {@link + * #updateTableWithPartitionBucketCountBackfill}). + */ + public Optional> getPartitionWithVersion( + TablePath tablePath, String partitionName) throws Exception { + String path = PartitionZNode.path(tablePath, partitionName); + Stat stat = new Stat(); + return getDataWithStat(path, stat) + .map(PartitionZNode::decode) + .map( + p -> + p.getRemoteDataDir() == null + ? p.newRemoteDataDir(defaultRemoteDataDir) + : p) + .map(p -> new VersionedData<>(p, stat.getVersion())); + } + + /** + * Gets all partition registrations of a table together with their ZK versions. The partition + * znodes are fetched concurrently to avoid one synchronous ZooKeeper round trip per partition. + * A partition dropped after the children listing is omitted from the result. + */ + public Map> getPartitionRegistrationsWithVersion( + TablePath tablePath) throws Exception { + Set partitionNames = getPartitions(tablePath); + if (partitionNames.isEmpty()) { + return Collections.emptyMap(); + } + + Map pathToPartitionName = + partitionNames.stream() + .collect(toMap(name -> PartitionZNode.path(tablePath, name), name -> name)); + List responses = getDataInBackground(pathToPartitionName.keySet()); + Map> registrations = new HashMap<>(); + for (ZkGetDataResponse response : responses) { + if (response.getResultCode() == KeeperException.Code.NONODE) { + continue; + } + response.maybeThrow(); + byte[] data = + checkNotNull( + response.getData(), + "Partition registration data must not be null for %s", + response.getPath()); + Stat stat = + checkNotNull( + response.getStat(), + "Partition registration stat must not be null for %s", + response.getPath()); + PartitionRegistration registration = PartitionZNode.decode(data); + if (registration.getRemoteDataDir() == null) { + registration = registration.newRemoteDataDir(defaultRemoteDataDir); + } + registrations.put( + pathToPartitionName.get(response.getPath()), + new VersionedData<>(registration, stat.getVersion())); + } + return registrations; + } + + /** + * Overwrites a partition's registration znode without a version check. NOT used in production + * (the ALTER bucket.num backfill goes through {@link + * #updateTableWithPartitionBucketCountBackfill}); this is a test-only backdoor for constructing + * legacy partition znodes, e.g. one with a null per-partition bucket count (v1 data) or a stale + * znode version. + */ + @VisibleForTesting + public void updatePartitionRegistration( + TablePath tablePath, String partitionName, PartitionRegistration registration) + throws Exception { + String path = PartitionZNode.path(tablePath, partitionName); + byte[] data = PartitionZNode.encode(registration); + zkClient.setData().forPath(path, data); + } + + /** + * Updates the table registration and the given partition registrations in one atomic ZooKeeper + * transaction. Every {@code setData} is CAS-guarded by its expected ZK version and the whole + * transaction is fenced on the coordinator epoch znode ({@link ZkVersion#MATCH_ANY_VERSION} + * skips the fence), so a stale snapshot or a deposed coordinator fails with {@link + * KeeperException.BadVersionException} instead of committing. + * + * @param tablePath the table to update + * @param tableRegistration the new table-level registration + * @param expectedTableZkVersion the expected ZK version of the table znode + * @param partitionBackfills partition name -> (updated registration + expected ZK version) + * @param expectedCoordinatorEpochZkVersion the coordinator epoch znode version to fence on + */ + public void updateTableWithPartitionBucketCountBackfill( + TablePath tablePath, + TableRegistration tableRegistration, + int expectedTableZkVersion, + Map> partitionBackfills, + int expectedCoordinatorEpochZkVersion) + throws Exception { + List ops = new ArrayList<>(partitionBackfills.size() + 1); + for (Map.Entry> entry : + partitionBackfills.entrySet()) { + String partitionPath = PartitionZNode.path(tablePath, entry.getKey()); + byte[] partitionData = PartitionZNode.encode(entry.getValue().data()); + ops.add( + zkClient.transactionOp() + .setData() + .withVersion(entry.getValue().zkVersion()) + .forPath(partitionPath, partitionData)); + } + String tablePathStr = TableZNode.path(tablePath); + byte[] tableData = TableZNode.encode(tableRegistration); + ops.add( + zkClient.transactionOp() + .setData() + .withVersion(expectedTableZkVersion) + .forPath(tablePathStr, tableData)); + + List fencedOps = + wrapRequestsWithEpochCheck(ops, expectedCoordinatorEpochZkVersion); + zkClient.transaction().forOperations(fencedOps); + LOG.info( + "Atomically backfilled bucket count for {} partition(s) and updated table {} in one " + + "transaction (CAS + epoch fence {}).", + partitionBackfills.size(), + tablePath, + expectedCoordinatorEpochZkVersion); + } + /** Get partition id and table id for each partition in a batch async way. */ public Map getPartitionIds( Collection partitionPaths) throws Exception { @@ -1077,7 +1230,8 @@ public void registerPartitionAssignmentAndMetadata( PartitionAssignment partitionAssignment, String remoteDataDir, TablePath tablePath, - long tableId) + long tableId, + int bucketCount) throws Exception { // Merge "registerPartitionAssignment()" and "registerPartition()" // into one transaction. This is to avoid the case that the partition assignment is @@ -1123,7 +1277,7 @@ public void registerPartitionAssignmentAndMetadata( metadataPath, PartitionZNode.encode( new PartitionRegistration( - tableId, partitionId, remoteDataDir))); + tableId, partitionId, remoteDataDir, bucketCount))); ops.add(tabletServerPartitionNode); ops.add(metadataPartitionNode); @@ -1400,6 +1554,28 @@ public List listRemoteLogManifestHandles( return result; } + /** + * A decoded znode value together with the ZK version of its znode. Used to carry the version + * captured at read time so a later write can compare-and-set against it. + */ + public static final class VersionedData { + private final T data; + private final int zkVersion; + + public VersionedData(T data, int zkVersion) { + this.data = data; + this.zkVersion = zkVersion; + } + + public T data() { + return data; + } + + public int zkVersion() { + return zkVersion; + } + } + /** Tuple of a table bucket and its current remote log manifest handle. */ public static final class TableBucketAndManifest { private final TableBucket tableBucket; @@ -1717,6 +1893,15 @@ public CuratorFramework getCuratorClient() { return zkClient; } + /** + * Returns the Curator wrapper owned by this client for tests that need a decorating client over + * the same ZooKeeper connection. The returned wrapper must not be closed by the caller. + */ + @VisibleForTesting + public CuratorFrameworkWithUnhandledErrorListener getCuratorFrameworkWrapper() { + return curatorFrameworkWrapper; + } + // -------------------------------------------------------------------------------------------- // Table and Partition Metadata // -------------------------------------------------------------------------------------------- diff --git a/fluss-server/src/main/java/org/apache/fluss/server/zk/data/PartitionRegistration.java b/fluss-server/src/main/java/org/apache/fluss/server/zk/data/PartitionRegistration.java index 2aef5858f82..d342c686bbc 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/zk/data/PartitionRegistration.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/zk/data/PartitionRegistration.java @@ -45,10 +45,22 @@ public class PartitionRegistration { */ private final @Nullable String remoteDataDir; - public PartitionRegistration(long tableId, long partitionId, @Nullable String remoteDataDir) { + /** + * The bucket count of this partition. It is null when deserialized from an older version that + * does not persist per-partition bucket count. In that case, callers should fall back to the + * table-level bucket count. + */ + private final @Nullable Integer bucketCount; + + public PartitionRegistration( + long tableId, + long partitionId, + @Nullable String remoteDataDir, + @Nullable Integer bucketCount) { this.tableId = tableId; this.partitionId = partitionId; this.remoteDataDir = remoteDataDir; + this.bucketCount = bucketCount; } public long getTableId() { @@ -64,6 +76,41 @@ public String getRemoteDataDir() { return remoteDataDir; } + /** Returns the bucket count of this partition, or null if not persisted (old data). */ + @Nullable + public Integer getBucketCount() { + return bucketCount; + } + + /** + * Returns the bucket count of this partition, falling back to the given table-level bucket + * count when this partition was persisted by an older version that does not store the + * per-partition count. + * + *

The fallback is only valid at {@code bucketCountEpoch == 0} (legacy table or old server). + * At {@code bucketCountEpoch > 0}, the first ALTER bucket.num atomically backfills the + * per-partition count of every existing partition in the same ZooKeeper transaction that bumps + * the epoch (see {@code ZooKeeperClient#updateTableWithPartitionBucketCountBackfill}), so a + * partition observed at {@code bucketCountEpoch > 0} always carries a count. Reaching the throw + * below therefore means that invariant was broken (an internal bug, e.g. a partial backfill); + * it is theoretically unreachable and is surfaced as an {@link IllegalStateException} rather + * than a retriable error, since a client metadata refresh cannot repair server-side state. + */ + public int getBucketCountOrDefault(int tableBucketCount, long bucketCountEpoch) { + if (bucketCount != null) { + return bucketCount; + } + if (bucketCountEpoch == 0) { + return tableBucketCount; + } + throw new IllegalStateException( + "Partition " + + partitionId + + " is missing a per-partition bucket count at bucketCountEpoch " + + bucketCountEpoch + + "; the ALTER bucket.num backfill invariant was broken."); + } + public TablePartition toTablePartition() { return new TablePartition(tableId, partitionId); } @@ -77,7 +124,7 @@ public TablePartition toTablePartition() { * @return a new registration with the given remote data directory */ public PartitionRegistration newRemoteDataDir(String remoteDataDir) { - return new PartitionRegistration(tableId, partitionId, remoteDataDir); + return new PartitionRegistration(tableId, partitionId, remoteDataDir, bucketCount); } @Override @@ -88,12 +135,13 @@ public boolean equals(Object o) { PartitionRegistration that = (PartitionRegistration) o; return tableId == that.tableId && partitionId == that.partitionId - && Objects.equals(remoteDataDir, that.remoteDataDir); + && Objects.equals(remoteDataDir, that.remoteDataDir) + && Objects.equals(bucketCount, that.bucketCount); } @Override public int hashCode() { - return Objects.hash(tableId, partitionId, remoteDataDir); + return Objects.hash(tableId, partitionId, remoteDataDir, bucketCount); } @Override @@ -106,6 +154,8 @@ public String toString() { + ", remoteDataDir='" + remoteDataDir + '\'' + + ", bucketCount=" + + bucketCount + '}'; } } diff --git a/fluss-server/src/main/java/org/apache/fluss/server/zk/data/PartitionRegistrationJsonSerde.java b/fluss-server/src/main/java/org/apache/fluss/server/zk/data/PartitionRegistrationJsonSerde.java index ad83b18a079..37e4d5036b5 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/zk/data/PartitionRegistrationJsonSerde.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/zk/data/PartitionRegistrationJsonSerde.java @@ -37,7 +37,8 @@ public class PartitionRegistrationJsonSerde private static final String TABLE_ID_KEY = "table_id"; private static final String PARTITION_ID_KEY = "partition_id"; private static final String REMOTE_DATA_DIR_KEY = "remote_data_dir"; - private static final int VERSION = 1; + private static final String BUCKET_COUNT_KEY = "bucket_count"; + private static final int VERSION = 2; @Override public void serialize(PartitionRegistration registration, JsonGenerator generator) @@ -49,6 +50,9 @@ public void serialize(PartitionRegistration registration, JsonGenerator generato if (registration.getRemoteDataDir() != null) { generator.writeStringField(REMOTE_DATA_DIR_KEY, registration.getRemoteDataDir()); } + if (registration.getBucketCount() != null) { + generator.writeNumberField(BUCKET_COUNT_KEY, registration.getBucketCount()); + } generator.writeEndObject(); } @@ -62,6 +66,12 @@ public PartitionRegistration deserialize(JsonNode node) { if (node.has(REMOTE_DATA_DIR_KEY)) { remoteDataDir = node.get(REMOTE_DATA_DIR_KEY).asText(); } - return new PartitionRegistration(tableId, partitionId, remoteDataDir); + // When deserialize from an old version (v1), bucket_count may not exist. + // Callers should fall back to table-level bucket count when this is null. + Integer bucketCount = null; + if (node.has(BUCKET_COUNT_KEY)) { + bucketCount = node.get(BUCKET_COUNT_KEY).asInt(); + } + return new PartitionRegistration(tableId, partitionId, remoteDataDir, bucketCount); } } diff --git a/fluss-server/src/main/java/org/apache/fluss/server/zk/data/TableRegistration.java b/fluss-server/src/main/java/org/apache/fluss/server/zk/data/TableRegistration.java index 89e7716d01a..2ba3076a5d0 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/zk/data/TableRegistration.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/zk/data/TableRegistration.java @@ -66,6 +66,14 @@ public class TableRegistration { public final long createdTime; public final long modifiedTime; + /** + * A table-level, monotonically increasing version for bucket.num changes. New tables start at + * 0; a legacy JSON without the field is read as 0; every committed bucket.num change increments + * it. It is used to decide whether legacy clients without bucket count are still allowed (epoch + * 0) and to let TabletServers ignore older UpdateMetadata messages. + */ + public final long bucketCountEpoch; + public TableRegistration( long tableId, @Nullable String comment, @@ -76,6 +84,30 @@ public TableRegistration( @Nullable String remoteDataDir, long createdTime, long modifiedTime) { + this( + tableId, + comment, + partitionKeys, + tableDistribution, + properties, + customProperties, + remoteDataDir, + createdTime, + modifiedTime, + 0L); + } + + public TableRegistration( + long tableId, + @Nullable String comment, + List partitionKeys, + TableDistribution tableDistribution, + Map properties, + Map customProperties, + @Nullable String remoteDataDir, + long createdTime, + long modifiedTime, + long bucketCountEpoch) { checkArgument( tableDistribution.getBucketCount().isPresent(), "Bucket count is required for table registration."); @@ -89,6 +121,7 @@ public TableRegistration( this.remoteDataDir = remoteDataDir; this.createdTime = createdTime; this.modifiedTime = modifiedTime; + this.bucketCountEpoch = bucketCountEpoch; } public boolean isPartitioned() { @@ -127,7 +160,8 @@ public TableInfo toTableInfo( this.remoteDataDir, this.comment, this.createdTime, - this.modifiedTime); + this.modifiedTime, + this.bucketCountEpoch); } public static TableRegistration newTable( @@ -160,7 +194,29 @@ public TableRegistration newProperties( newCustomProperties, remoteDataDir, createdTime, - currentMillis); + currentMillis, + bucketCountEpoch); + } + + /** + * Returns a new registration with the given table-level bucket count and an incremented {@code + * bucketCountEpoch}. For a partitioned table, the new count applies to partitions created after + * this ALTER; existing partitions retain their actual bucket counts in their partition + * registrations. + */ + public TableRegistration newBucketCount(int newBucketCount) { + final long currentMillis = System.currentTimeMillis(); + return new TableRegistration( + tableId, + comment, + partitionKeys, + new TableDistribution(newBucketCount, bucketKeys), + properties, + customProperties, + remoteDataDir, + createdTime, + currentMillis, + bucketCountEpoch + 1); } /** @@ -181,7 +237,8 @@ public TableRegistration newRemoteDataDir(String remoteDataDir) { customProperties, remoteDataDir, createdTime, - modifiedTime); + modifiedTime, + bucketCountEpoch); } @Override @@ -197,6 +254,7 @@ public boolean equals(Object o) { return tableId == that.tableId && createdTime == that.createdTime && modifiedTime == that.modifiedTime + && bucketCountEpoch == that.bucketCountEpoch && Objects.equals(comment, that.comment) && Objects.equals(partitionKeys, that.partitionKeys) && Objects.equals(bucketCount, that.bucketCount) @@ -218,7 +276,8 @@ public int hashCode() { customProperties, remoteDataDir, createdTime, - modifiedTime); + modifiedTime, + bucketCountEpoch); } @Override @@ -245,6 +304,8 @@ public String toString() { + createdTime + ", modifiedTime=" + modifiedTime + + ", bucketCountEpoch=" + + bucketCountEpoch + '}'; } } diff --git a/fluss-server/src/main/java/org/apache/fluss/server/zk/data/TableRegistrationJsonSerde.java b/fluss-server/src/main/java/org/apache/fluss/server/zk/data/TableRegistrationJsonSerde.java index 6a9f93380ec..32313eed17e 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/zk/data/TableRegistrationJsonSerde.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/zk/data/TableRegistrationJsonSerde.java @@ -48,8 +48,9 @@ public class TableRegistrationJsonSerde static final String REMOTE_DATA_DIR = "remote_data_dir"; static final String CREATED_TIME = "created_time"; static final String MODIFIED_TIME = "modified_time"; + static final String BUCKET_COUNT_EPOCH = "bucket_count_epoch"; private static final String VERSION_KEY = "version"; - private static final int VERSION = 1; + private static final int VERSION = 2; @Override public void serialize(TableRegistration tableReg, JsonGenerator generator) throws IOException { @@ -112,6 +113,9 @@ public void serialize(TableRegistration tableReg, JsonGenerator generator) throw // serialize modifiedTime generator.writeNumberField(MODIFIED_TIME, tableReg.modifiedTime); + // serialize bucketCountEpoch + generator.writeNumberField(BUCKET_COUNT_EPOCH, tableReg.bucketCountEpoch); + generator.writeEndObject(); } @@ -157,6 +161,11 @@ public TableRegistration deserialize(JsonNode node) { long createdTime = node.get(CREATED_TIME).asLong(); long modifiedTime = node.get(MODIFIED_TIME).asLong(); + // When deserializing from a legacy version, the bucket layout epoch may not exist; + // read it as 0 (the table has never been ALTERed). + long bucketCountEpoch = + node.has(BUCKET_COUNT_EPOCH) ? node.get(BUCKET_COUNT_EPOCH).asLong() : 0L; + return new TableRegistration( tableId, comment, @@ -166,7 +175,8 @@ public TableRegistration deserialize(JsonNode node) { customProperties, remoteDataDir, createdTime, - modifiedTime); + modifiedTime, + bucketCountEpoch); } private Map deserializeProperties(JsonNode node) { diff --git a/fluss-server/src/test/java/org/apache/fluss/server/coordinator/AlterBucketNumTest.java b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/AlterBucketNumTest.java new file mode 100644 index 00000000000..70420a862cd --- /dev/null +++ b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/AlterBucketNumTest.java @@ -0,0 +1,1139 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.fluss.server.coordinator; + +import org.apache.fluss.cluster.Endpoint; +import org.apache.fluss.cluster.TabletServerInfo; +import org.apache.fluss.config.AutoPartitionTimeUnit; +import org.apache.fluss.config.ConfigOptions; +import org.apache.fluss.config.Configuration; +import org.apache.fluss.exception.FlussRuntimeException; +import org.apache.fluss.exception.InvalidAlterTableException; +import org.apache.fluss.exception.TableNotExistException; +import org.apache.fluss.exception.TooManyBucketsException; +import org.apache.fluss.lake.lakestorage.LakeCatalog; +import org.apache.fluss.lake.lakestorage.LakeStorage; +import org.apache.fluss.lake.lakestorage.LakeStoragePlugin; +import org.apache.fluss.metadata.DataLakeFormat; +import org.apache.fluss.metadata.DatabaseDescriptor; +import org.apache.fluss.metadata.ResolvedPartitionSpec; +import org.apache.fluss.metadata.Schema; +import org.apache.fluss.metadata.TableChange; +import org.apache.fluss.metadata.TableDescriptor; +import org.apache.fluss.metadata.TableInfo; +import org.apache.fluss.metadata.TablePath; +import org.apache.fluss.plugin.PluginManager; +import org.apache.fluss.server.entity.TablePropertyChanges; +import org.apache.fluss.server.lakehouse.TestingPaimonStoragePlugin; +import org.apache.fluss.server.lakehouse.TestingPaimonStoragePlugin.TestingPaimonLakeStorage; +import org.apache.fluss.server.zk.CuratorFrameworkWithUnhandledErrorListener; +import org.apache.fluss.server.zk.NOPErrorHandler; +import org.apache.fluss.server.zk.ZooKeeperClient; +import org.apache.fluss.server.zk.ZooKeeperExtension; +import org.apache.fluss.server.zk.data.CoordinatorAddress; +import org.apache.fluss.server.zk.data.PartitionAssignment; +import org.apache.fluss.server.zk.data.PartitionRegistration; +import org.apache.fluss.server.zk.data.TableAssignment; +import org.apache.fluss.server.zk.data.TableRegistration; +import org.apache.fluss.server.zk.data.TabletServerRegistration; +import org.apache.fluss.server.zk.data.ZkData; +import org.apache.fluss.server.zk.data.ZkVersion; +import org.apache.fluss.shaded.zookeeper3.org.apache.zookeeper.KeeperException; +import org.apache.fluss.testutils.common.AllCallbackWrapper; +import org.apache.fluss.types.DataTypes; +import org.apache.fluss.utils.function.RunnableWithException; + +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.RegisterExtension; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.CsvSource; +import org.junit.jupiter.params.provider.EnumSource; +import org.junit.jupiter.params.provider.MethodSource; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.stream.Stream; + +import static org.apache.fluss.config.ConfigOptions.DEFAULT_LISTENER_NAME; +import static org.apache.fluss.metadata.ResolvedPartitionSpec.fromPartitionName; +import static org.apache.fluss.server.utils.TableAssignmentUtils.generateAssignment; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** Tests for ALTER TABLE SET ('bucket.num' = 'N') per-partition bucket count rescale. */ +class AlterBucketNumTest { + + private static final String DEFAULT_DB = "db"; + + @RegisterExtension + public static final AllCallbackWrapper ZOO_KEEPER_EXTENSION_WRAPPER = + new AllCallbackWrapper<>(new ZooKeeperExtension()); + + private static ZooKeeperClient zookeeperClient; + private static MetadataManager metadataManager; + private static String remoteDataDir; + + @BeforeAll + static void beforeAll() throws Exception { + zookeeperClient = + ZOO_KEEPER_EXTENSION_WRAPPER + .getCustomExtension() + .getZooKeeperClient(NOPErrorHandler.INSTANCE); + metadataManager = + new MetadataManager( + zookeeperClient, + new Configuration(), + new LakeCatalogDynamicLoader(new Configuration(), null, true)); + + // register coordinator server + zookeeperClient.registerCoordinatorLeader( + new CoordinatorAddress( + "1", Endpoint.fromListenersString("CLIENT://localhost:10012"))); + zookeeperClient.fenceBecomeCoordinatorLeader("1"); + + // register 3 tablet servers + for (int i = 0; i < 3; i++) { + zookeeperClient.registerTabletServer( + i, + new TabletServerRegistration( + "rack" + i, + Collections.singletonList( + new Endpoint("host" + i, 1000, DEFAULT_LISTENER_NAME)), + System.currentTimeMillis())); + } + + // create database + metadataManager.createDatabase(DEFAULT_DB, DatabaseDescriptor.builder().build(), false); + remoteDataDir = zookeeperClient.getDefaultRemoteDataDir(); + } + + // ====================== Lake Propagation Tests ====================== + + @Test + void testAlterBucketNumOnLakeTablePassesValidationButAbortsWithoutLakeCatalog() + throws Exception { + // A lake table is no longer rejected by validation; the ALTER proceeds to the lake + // propagation, which aborts here because this harness wires no lake catalog. Covers the + // lakeCatalog == null branch (distinct from a propagation call that fails, tested below). + TablePath tablePath = TablePath.of(DEFAULT_DB, "test_lake_table_alter_allowed"); + int originalBucketCount = 4; + createSeededLakePartitionedTable(metadataManager, tablePath, originalBucketCount); + + // Lake First: the ALTER validates and then attempts to propagate the new bucket count to + // the lake side BEFORE the Fluss ZK commit. This unit-test harness has no real lake + // catalog wired in, so the propagation step fails with a FlussRuntimeException whose + // message clearly points at the propagation stage. + assertThatThrownBy(() -> alterBucketNum(metadataManager, tablePath, 8)) + .isInstanceOf(FlussRuntimeException.class) + .hasMessageContaining("propagate ALTER bucket.num to the lake side"); + + // The propagation failure aborts the ALTER BEFORE the Fluss ZK commit, so table-level and + // pre-existing partition state must both be unchanged. + assertThat(metadataManager.getTable(tablePath).getNumBuckets()) + .isEqualTo(originalBucketCount); + Optional pre = zookeeperClient.getPartition(tablePath, "2024-01"); + assertThat(pre).isPresent(); + assertThat(pre.get().getBucketCount()).isEqualTo(originalBucketCount); + } + + @Test + void testAlterBucketNumLakePropagationFailureAbortsAlter() throws Exception { + CountingLakeCatalog stub = new CountingLakeCatalog(true); + MetadataManager mm = buildMetadataManagerWithLakeCatalog(stub); + TablePath tablePath = TablePath.of(DEFAULT_DB, "test_lake_alter_persistent_fail"); + int originalBucketCount = 4; + createSeededLakePartitionedTable(mm, tablePath, originalBucketCount); + + // A lake failure fails loud: the ALTER aborts BEFORE the Fluss ZK commit with a clear + // error telling the operator nothing was changed on the Fluss side and to re-run the + // ALTER once the lake is reachable. + assertThatThrownBy(() -> alterBucketNum(mm, tablePath, 8)) + .isInstanceOf(FlussRuntimeException.class) + .hasMessageContaining("to the lake schema failed") + .hasMessageContaining("The Fluss side was NOT changed") + .hasMessageContaining("Re-run the same ALTER"); + + // Propagation is attempted exactly once. + assertThat(stub.attempts.get()).isEqualTo(1); + + // Lake First: the Fluss ZK commit never ran, so table-level bucket count and the + // pre-existing partition are both unchanged. + assertThat(mm.getTable(tablePath).getNumBuckets()).isEqualTo(originalBucketCount); + Optional pre = zookeeperClient.getPartition(tablePath, "2024-01"); + assertThat(pre).isPresent(); + assertThat(pre.get().getBucketCount()).isEqualTo(originalBucketCount); + } + + @Test + void testAlterBucketNumLakePropagationSucceedsFirstTry() throws Exception { + CountingLakeCatalog stub = new CountingLakeCatalog(false); + MetadataManager mm = buildMetadataManagerWithLakeCatalog(stub); + TablePath tablePath = TablePath.of(DEFAULT_DB, "test_lake_alter_success"); + int newBucketCount = 8; + createSeededLakePartitionedTable(mm, tablePath, 4); + + alterBucketNum(mm, tablePath, newBucketCount); + + // Propagation succeeded on the first attempt. + assertThat(stub.attempts.get()).isEqualTo(1); + assertThat(stub.lastBucketCount).isEqualTo(newBucketCount); + assertThat(mm.getTable(tablePath).getNumBuckets()).isEqualTo(newBucketCount); + } + + @Test + void testAlterBucketNumSkipsLakePropagationForUnawareBucketTable() throws Exception { + // A lake table WITHOUT bucket keys is an Unaware Bucket table in Paimon (BUCKET = -1 + // encodes the bucket MODE); propagating a positive BUCKET would flip its mode. The + // propagation must be skipped entirely while the Fluss-side rescale still succeeds. + CountingLakeCatalog stub = new CountingLakeCatalog(false); + MetadataManager mm = buildMetadataManagerWithLakeCatalog(stub); + TablePath tablePath = TablePath.of(DEFAULT_DB, "test_lake_alter_unaware_skip"); + int originalBucketCount = 4; + int newBucketCount = 8; + TableDescriptor unawareLakeTable = + TableDescriptor.builder() + .schema( + Schema.newBuilder() + .column("a", DataTypes.INT()) + .column("b", DataTypes.STRING()) + .build()) + .distributedBy(originalBucketCount) + .partitionedBy("b") + .property(ConfigOptions.TABLE_DATALAKE_ENABLED.key(), "true") + .property(ConfigOptions.TABLE_DATALAKE_FORMAT.key(), "paimon") + .build() + .withReplicationFactor(3); + mm.createTable( + tablePath, + remoteDataDir, + unawareLakeTable, + generateAssignment(originalBucketCount, 3, getTabletServers()), + false); + + alterBucketNum(mm, tablePath, newBucketCount); + + // no call ever reached the lake catalog, and the Fluss side still rescaled + assertThat(stub.attempts.get()).isEqualTo(0); + assertThat(mm.getTable(tablePath).getNumBuckets()).isEqualTo(newBucketCount); + } + + /** + * Builds a coordinator-side {@link MetadataManager} through the real plugin and lake-catalog + * initialization path, with the supplied catalog returned by the testing Paimon storage. + */ + private static MetadataManager buildMetadataManagerWithLakeCatalog(LakeCatalog catalog) { + LakeStoragePlugin plugin = + new TestingPaimonStoragePlugin() { + @Override + public LakeStorage createLakeStorage(Configuration configuration) { + return new TestingPaimonLakeStorage() { + @Override + public LakeCatalog createLakeCatalog() { + return catalog; + } + }; + } + }; + + PluginManager pluginManager = + new PluginManager() { + @Override + @SuppressWarnings("unchecked") + public

Iterator

load(Class

service) { + if (service == LakeStoragePlugin.class) { + return Collections.singletonList((P) plugin).iterator(); + } + return Collections.

emptyList().iterator(); + } + }; + + Configuration conf = new Configuration(); + conf.set(ConfigOptions.DATALAKE_ENABLED, true); + conf.set(ConfigOptions.DATALAKE_FORMAT, DataLakeFormat.PAIMON); + + LakeCatalogDynamicLoader loader = new LakeCatalogDynamicLoader(conf, pluginManager, true); + return new MetadataManager(zookeeperClient, conf, loader); + } + + /** + * Creates a lake-enabled, partitioned Fixed Bucket table with the given original bucket count + * and seeds one pre-existing partition "2024-01" carrying that bucket count. + */ + private static void createSeededLakePartitionedTable( + MetadataManager mm, TablePath tablePath, int originalBucketCount) throws Exception { + createSeededLakePartitionedTable(mm, tablePath, originalBucketCount, false); + } + + private static void createSeededLakePartitionedTable( + MetadataManager mm, + TablePath tablePath, + int originalBucketCount, + boolean historicalPartitionEnabled) + throws Exception { + TableDescriptor.Builder builder = + TableDescriptor.builder() + .schema( + Schema.newBuilder() + .column("a", DataTypes.INT()) + .column("b", DataTypes.STRING()) + .build()) + .distributedBy(originalBucketCount, "a") + .partitionedBy("b") + .property(ConfigOptions.TABLE_DATALAKE_ENABLED.key(), "true") + .property(ConfigOptions.TABLE_DATALAKE_FORMAT.key(), "paimon") + // the historical partition requires auto-partitioning + .property(ConfigOptions.TABLE_AUTO_PARTITION_ENABLED.key(), "true") + .property(ConfigOptions.TABLE_AUTO_PARTITION_KEY.key(), "b") + .property( + ConfigOptions.TABLE_AUTO_PARTITION_TIME_UNIT.key(), + AutoPartitionTimeUnit.DAY.toString()); + if (historicalPartitionEnabled) { + builder.property( + ConfigOptions.TABLE_DATALAKE_HISTORICAL_PARTITION_ENABLED.key(), "true"); + } + TableDescriptor lakeTable = builder.build().withReplicationFactor(3); + TableAssignment tableAssignment = + generateAssignment(originalBucketCount, 3, getTabletServers()); + mm.createTable(tablePath, remoteDataDir, lakeTable, tableAssignment, false); + TableInfo tableInfo = mm.getTable(tablePath); + mm.createPartition( + tablePath, + tableInfo.getTableId(), + remoteDataDir, + new PartitionAssignment( + tableInfo.getTableId(), tableAssignment.getBucketAssignments()), + fromPartitionName(tableInfo.getPartitionKeys(), "2024-01"), + false, + originalBucketCount); + } + + /** + * A stub lake catalog that counts bucket-count propagations and can simulate transient faults. + */ + private static final class CountingLakeCatalog implements LakeCatalog { + private final AtomicInteger attempts = new AtomicInteger(); + private final boolean failing; + private volatile Integer lastBucketCount; + + CountingLakeCatalog(boolean failing) { + this.failing = failing; + } + + @Override + public void createTable( + TablePath tablePath, TableDescriptor tableDescriptor, Context context) { + // not used by these tests + } + + @Override + public void alterTable( + TablePath tablePath, java.util.List tableChanges, Context context) { + for (TableChange change : tableChanges) { + if (change instanceof TableChange.ModifyBucketCount) { + attempts.incrementAndGet(); + if (failing) { + throw new RuntimeException("simulated transient lake failure"); + } + lastBucketCount = ((TableChange.ModifyBucketCount) change).getNewBucketCount(); + } + } + } + } + + @Test + void testAlterBucketNumRejectedOnNonPartitionedTable() throws Exception { + TablePath tablePath = TablePath.of(DEFAULT_DB, "test_non_partitioned_reject"); + TableDescriptor logTable = + TableDescriptor.builder() + .schema( + Schema.newBuilder() + .column("a", DataTypes.INT()) + .column("b", DataTypes.STRING()) + .build()) + .distributedBy(4) + .build() + .withReplicationFactor(3); + + TableAssignment tableAssignment = generateAssignment(4, 3, getTabletServers()); + metadataManager.createTable(tablePath, remoteDataDir, logTable, tableAssignment, false); + + // ALTER bucket.num on non-partitioned table should be rejected + assertThatThrownBy(() -> alterBucketNum(metadataManager, tablePath, 8)) + .isInstanceOf(InvalidAlterTableException.class) + .hasMessageContaining("Cannot alter 'bucket.num' on non-partitioned table"); + } + + @Test + void testAlterBucketNumRejectedOnHistoricalPartitionTable() throws Exception { + TablePath tablePath = TablePath.of(DEFAULT_DB, "test_reject_rescale_on_historical"); + int originalBucketCount = 4; + createSeededLakePartitionedTable(metadataManager, tablePath, originalBucketCount, true); + + // The rejection happens during validation, before the lake propagation, so the default + // manager without a lake catalog never reaches the propagation failure. + assertThatThrownBy(() -> alterBucketNum(metadataManager, tablePath, 8)) + .isInstanceOf(InvalidAlterTableException.class) + .hasMessageContaining("with historical partition enabled") + .hasMessageContaining("not supported yet"); + + // The bucket layout is untouched: neither the count nor the epoch moved. + TableInfo afterTableInfo = metadataManager.getTable(tablePath); + assertThat(afterTableInfo.getNumBuckets()).isEqualTo(originalBucketCount); + assertThat(afterTableInfo.getBucketCountEpoch()).isEqualTo(0L); + Optional partition = + zookeeperClient.getPartition(tablePath, "2024-01"); + assertThat(partition).isPresent(); + assertThat(partition.get().getBucketCount()).isEqualTo(originalBucketCount); + } + + @Test + void testEnableHistoricalPartitionRejectedAfterRescale() throws Exception { + MetadataManager mm = buildMetadataManagerWithLakeCatalog(new CountingLakeCatalog(false)); + TablePath tablePath = TablePath.of(DEFAULT_DB, "test_reject_historical_after_rescale"); + int originalBucketCount = 4; + createSeededLakePartitionedTable(mm, tablePath, originalBucketCount, false); + + // Rescale the table first, which advances the bucketCountEpoch. + alterBucketNum(mm, tablePath, 8); + assertThat(mm.getTable(tablePath).getBucketCountEpoch()).isEqualTo(1L); + + // Enabling the historical partition on the rescaled table must be rejected. + assertThatThrownBy(() -> alterHistoricalPartition(mm, tablePath, true)) + .isInstanceOf(InvalidAlterTableException.class) + .hasMessageContaining("Cannot enable historical partition") + .hasMessageContaining("not supported yet"); + + // The rejection left no side effects: no historical partition was created and the + // bucket layout stays at the rescaled state. + assertThat(zookeeperClient.getPartition(tablePath, "__historical__")).isEmpty(); + assertThat(mm.getTable(tablePath).getNumBuckets()).isEqualTo(8); + assertThat(mm.getTable(tablePath).getBucketCountEpoch()).isEqualTo(1L); + } + + @Test + void testEnableHistoricalPartitionRejectedAfterRescaleWhileDisabled() throws Exception { + MetadataManager mm = buildMetadataManagerWithLakeCatalog(new CountingLakeCatalog(false)); + TablePath tablePath = + TablePath.of(DEFAULT_DB, "test_reject_historical_after_disable_rescale"); + int originalBucketCount = 4; + createSeededLakePartitionedTable(mm, tablePath, originalBucketCount, true); + + // Disable the historical partition, then rescale: both steps succeed apart. + alterHistoricalPartition(mm, tablePath, false); + alterBucketNum(mm, tablePath, 8); + assertThat(mm.getTable(tablePath).getBucketCountEpoch()).isEqualTo(1L); + + // Re-enabling must still be rejected: the epoch never decreases, and the historical + // partition would be created with the new count while retired lake data keeps the + // pre-rescale layout. + assertThatThrownBy(() -> alterHistoricalPartition(mm, tablePath, true)) + .isInstanceOf(InvalidAlterTableException.class) + .hasMessageContaining("Cannot enable historical partition") + .hasMessageContaining("not supported yet"); + assertThat(mm.getTable(tablePath).getBucketCountEpoch()).isEqualTo(1L); + } + + @Test + void testEnableHistoricalPartitionOnNeverRescaledTableSucceeds() throws Exception { + MetadataManager mm = buildMetadataManagerWithLakeCatalog(new CountingLakeCatalog(false)); + TablePath tablePath = TablePath.of(DEFAULT_DB, "test_enable_historical_on_fresh_table"); + int originalBucketCount = 4; + createSeededLakePartitionedTable(mm, tablePath, originalBucketCount, false); + + // A never-rescaled table (epoch 0) can still enable the historical partition. + alterHistoricalPartition(mm, tablePath, true); + assertThat(mm.getTable(tablePath).getTableConfig().isHistoricalPartitionEnabled()).isTrue(); + } + + // ========================== Success Tests ========================== + + @ParameterizedTest(name = "bucketNum {0} -> {1}") + @CsvSource({"3, 6", "6, 3"}) + void testBackfillOnlyAffectsPartitionsWithoutBucketCount( + int originalBucketCount, int newBucketCount) throws Exception { + TablePath tablePath = + TablePath.of( + DEFAULT_DB, + "test_backfill_idempotent_" + originalBucketCount + "_" + newBucketCount); + TableAssignment tableAssignment = + generateAssignment(originalBucketCount, 3, getTabletServers()); + metadataManager.createTable( + tablePath, + remoteDataDir, + partitionedLogTable(originalBucketCount), + tableAssignment, + false); + TableInfo tableInfo = metadataManager.getTable(tablePath); + long tableId = tableInfo.getTableId(); + + // Create two partitions with bucketCount = originalBucketCount + PartitionAssignment partitionAssignment = + new PartitionAssignment(tableId, tableAssignment.getBucketAssignments()); + metadataManager.createPartition( + tablePath, + tableId, + remoteDataDir, + partitionAssignment, + fromPartitionName(tableInfo.getPartitionKeys(), "legacy"), + false, + originalBucketCount); + metadataManager.createPartition( + tablePath, + tableId, + remoteDataDir, + partitionAssignment, + fromPartitionName(tableInfo.getPartitionKeys(), "new"), + false, + originalBucketCount); + + // Simulate a legacy partition by overwriting its registration with null bucketCount. + // This models a partition created before per-partition bucket count was introduced. + Optional legacyReg = + zookeeperClient.getPartition(tablePath, "legacy"); + assertThat(legacyReg).isPresent(); + PartitionRegistration nullBucketCountReg = + new PartitionRegistration( + legacyReg.get().getTableId(), + legacyReg.get().getPartitionId(), + legacyReg.get().getRemoteDataDir(), + null); + zookeeperClient.updatePartitionRegistration(tablePath, "legacy", nullBucketCountReg); + + // Verify: "legacy" has null bucketCount, "new" has originalBucketCount + Optional beforeLegacy = + zookeeperClient.getPartition(tablePath, "legacy"); + assertThat(beforeLegacy).isPresent(); + assertThat(beforeLegacy.get().getBucketCount()).isNull(); + + Optional beforeNew = zookeeperClient.getPartition(tablePath, "new"); + assertThat(beforeNew).isPresent(); + assertThat(beforeNew.get().getBucketCount()).isEqualTo(originalBucketCount); + + // ALTER bucket.num in both directions (scale-up 3->6 and scale-down 6->3) + alterBucketNum(metadataManager, tablePath, newBucketCount); + + // Verify: table-level bucket count was updated and persisted in ZK + assertThat(metadataManager.getTable(tablePath).getNumBuckets()).isEqualTo(newBucketCount); + + // Verify: "legacy" was backfilled with actual bucket count (from old table bucket count) + Optional afterLegacy = + zookeeperClient.getPartition(tablePath, "legacy"); + assertThat(afterLegacy).isPresent(); + assertThat(afterLegacy.get().getBucketCount()).isEqualTo(originalBucketCount); + + // Verify: "new" still has the original bucketCount (not overwritten to the new value) + Optional afterNew = zookeeperClient.getPartition(tablePath, "new"); + assertThat(afterNew).isPresent(); + assertThat(afterNew.get().getBucketCount()).isEqualTo(originalBucketCount); + } + + @Test + void testPartitionCreatedInsideAlterWindowKeepsItsOwnBucketCount() throws Exception { + // A partition created between the backfill enumeration and the commit is not part of the + // backfill, so it must keep the bucket count of the assignment it was created with. + TablePath tablePath = TablePath.of(DEFAULT_DB, "test_alter_vs_create_occ"); + int originalBucketCount = 4; + TableAssignment tableAssignment = + generateAssignment(originalBucketCount, 3, getTabletServers()); + metadataManager.createTable( + tablePath, + remoteDataDir, + partitionedLogTable(originalBucketCount), + tableAssignment, + false); + TableInfo tableInfo = metadataManager.getTable(tablePath); + long tableId = tableInfo.getTableId(); + + // an existing partition created with the original bucket count (4) + PartitionAssignment partitionAssignment = + new PartitionAssignment(tableId, tableAssignment.getBucketAssignments()); + metadataManager.createPartition( + tablePath, + tableId, + remoteDataDir, + partitionAssignment, + fromPartitionName(tableInfo.getPartitionKeys(), "2024-01"), + false, + originalBucketCount); + + // The concurrent partition is created with the original count, mirroring a creation that + // read the table registration before the ALTER committed. + MetadataManager alterManager = + metadataManagerOver( + zkClientRunningBeforeFirstCommit( + () -> + metadataManager.createPartition( + tablePath, + tableId, + remoteDataDir, + new PartitionAssignment( + tableId, + tableAssignment.getBucketAssignments()), + fromPartitionName( + tableInfo.getPartitionKeys(), "2024-02"), + false, + originalBucketCount))); + + alterBucketNum(alterManager, tablePath, 8); + + // The ALTER committed the new table-level count. + assertThat(metadataManager.getTable(tablePath).getNumBuckets()).isEqualTo(8); + + // Both partitions keep a count consistent with the assignment they were created with. + assertPartitionCountMatchesAssignment(tablePath, "2024-01", originalBucketCount); + assertPartitionCountMatchesAssignment(tablePath, "2024-02", originalBucketCount); + } + + /** + * Asserts the partition's persisted bucket count equals both its assignment size and {@code + * expected}. + */ + private static void assertPartitionCountMatchesAssignment( + TablePath tablePath, String partitionName, int expected) throws Exception { + Optional partition = + zookeeperClient.getPartition(tablePath, partitionName); + assertThat(partition).isPresent(); + Optional assignment = + zookeeperClient.getPartitionAssignment(partition.get().getPartitionId()); + assertThat(assignment).isPresent(); + assertThat(partition.get().getBucketCount()) + .isEqualTo(assignment.get().getBucketAssignments().size()) + .isEqualTo(expected); + } + + private enum StaleFence { + TABLE_VERSION, + PARTITION_VERSION, + COORDINATOR_EPOCH + } + + @ParameterizedTest + @EnumSource(StaleFence.class) + void testBackfillCommitRejectsStaleFence(StaleFence fence) throws Exception { + TablePath tablePath = TablePath.of(DEFAULT_DB, "test_fence_" + fence.name().toLowerCase()); + int originalBucketCount = 4; + TableAssignment tableAssignment = + generateAssignment(originalBucketCount, 3, getTabletServers()); + metadataManager.createTable( + tablePath, + remoteDataDir, + partitionedLogTable(originalBucketCount), + tableAssignment, + false); + + String partitionName = "2024-01"; + if (fence == StaleFence.PARTITION_VERSION) { + TableInfo tableInfo = metadataManager.getTable(tablePath); + metadataManager.createPartition( + tablePath, + tableInfo.getTableId(), + remoteDataDir, + new PartitionAssignment( + tableInfo.getTableId(), tableAssignment.getBucketAssignments()), + fromPartitionName(tableInfo.getPartitionKeys(), partitionName), + false, + originalBucketCount); + } + + // Capture fresh versions, then make exactly the fenced dimension stale (simulating a + // concurrent read-modify-write or a deposed coordinator). + ZooKeeperClient.VersionedData table = + zookeeperClient.getTableWithVersion(tablePath).get(); + int tableVersion = table.zkVersion(); + Map> backfills = + new HashMap<>(); + int epochVersion = ZkVersion.MATCH_ANY_VERSION.getVersion(); + switch (fence) { + case TABLE_VERSION: + zookeeperClient.updateTable(tablePath, table.data()); + break; + case PARTITION_VERSION: + ZooKeeperClient.VersionedData partition = + zookeeperClient.getPartitionWithVersion(tablePath, partitionName).get(); + zookeeperClient.updatePartitionRegistration( + tablePath, partitionName, partition.data()); + backfills.put(partitionName, partition); + break; + case COORDINATOR_EPOCH: + epochVersion = zookeeperClient.getCurrentEpoch().getCoordinatorEpochZkVersion() + 1; + break; + } + + int staleEpochVersion = epochVersion; + assertThatThrownBy( + () -> + zookeeperClient.updateTableWithPartitionBucketCountBackfill( + tablePath, + table.data().newBucketCount(8), + tableVersion, + backfills, + staleEpochVersion)) + .isInstanceOf(KeeperException.BadVersionException.class); + // The atomic transaction rejected everything: table-level unchanged. + assertThat(metadataManager.getTable(tablePath).getNumBuckets()) + .isEqualTo(originalBucketCount); + + // With every dimension fresh the same commit succeeds. + ZooKeeperClient.VersionedData freshTable = + zookeeperClient.getTableWithVersion(tablePath).get(); + Map> freshBackfills = + new HashMap<>(); + if (fence == StaleFence.PARTITION_VERSION) { + freshBackfills.put( + partitionName, + zookeeperClient.getPartitionWithVersion(tablePath, partitionName).get()); + } + zookeeperClient.updateTableWithPartitionBucketCountBackfill( + tablePath, + freshTable.data().newBucketCount(8), + freshTable.zkVersion(), + freshBackfills, + zookeeperClient.getCurrentEpoch().getCoordinatorEpochZkVersion()); + assertThat(metadataManager.getTable(tablePath).getNumBuckets()).isEqualTo(8); + } + + @ParameterizedTest(name = "newBucketNum={0}") + @MethodSource("outOfRangeBucketNums") + void testAlterBucketNumRejectedOutOfRange( + int newBucketNum, Class expectedException, String message) + throws Exception { + TablePath tablePath = + TablePath.of(DEFAULT_DB, "test_alter_bucket_num_out_of_range_" + newBucketNum); + int originalBucketCount = 4; + metadataManager.createTable( + tablePath, + remoteDataDir, + partitionedLogTable(originalBucketCount), + generateAssignment(originalBucketCount, 3, getTabletServers()), + false); + + assertThatThrownBy(() -> alterBucketNum(metadataManager, tablePath, newBucketNum)) + .isInstanceOf(expectedException) + .hasMessageContaining(message); + // table-level bucket count unchanged + assertThat(metadataManager.getTable(tablePath).getNumBuckets()) + .isEqualTo(originalBucketCount); + } + + private static Stream outOfRangeBucketNums() { + return Stream.of( + Arguments.of(0, InvalidAlterTableException.class, "at least 1"), + Arguments.of( + ConfigOptions.MAX_BUCKET_NUM.defaultValue() + 1, + TooManyBucketsException.class, + "exceeding the maximum")); + } + + @Test + void testAlterBucketNumRetriesOnceThenSucceeds() throws Exception { + TablePath tablePath = TablePath.of(DEFAULT_DB, "test_alter_bucket_num_retry_success"); + int originalBucketCount = 4; + metadataManager.createTable( + tablePath, + remoteDataDir, + partitionedLogTable(originalBucketCount), + generateAssignment(originalBucketCount, 3, getTabletServers()), + false); + + // A ZK client that throws BadVersion on the first bucket-count commit, then delegates. + // This deterministically exercises the retry loop in alterTableProperties: attempt 1 + // hits BadVersion, attempt 2 re-reads a fresh version and commits successfully. + AtomicInteger commitCalls = new AtomicInteger(); + MetadataManager retryMetadataManager = + metadataManagerOver(zkClientFailingCommits(1, commitCalls)); + + alterBucketNum(retryMetadataManager, tablePath, 8); + + // The retry loop must have re-invoked the commit exactly once after the injected failure. + assertThat(commitCalls.get()).isEqualTo(2); + // Table-level bucket count is now the new value, confirming the retried attempt succeeded. + assertThat(metadataManager.getTable(tablePath).getNumBuckets()).isEqualTo(8); + } + + @Test + void testAlterBucketNumFailsAfterMaxRetries() throws Exception { + TablePath tablePath = TablePath.of(DEFAULT_DB, "test_alter_bucket_num_retry_exhaust"); + int originalBucketCount = 4; + metadataManager.createTable( + tablePath, + remoteDataDir, + partitionedLogTable(originalBucketCount), + generateAssignment(originalBucketCount, 3, getTabletServers()), + false); + + // Every commit throws BadVersion so the retry loop exhausts its budget and wraps the + // failure into FlussRuntimeException with the "after 3 retries" message. + AtomicInteger commitCalls = new AtomicInteger(); + MetadataManager exhaustRetryManager = + metadataManagerOver(zkClientFailingCommits(Integer.MAX_VALUE, commitCalls)); + + assertThatThrownBy(() -> alterBucketNum(exhaustRetryManager, tablePath, 8)) + .isInstanceOf(FlussRuntimeException.class) + .hasMessageContaining("after 3 retries") + .hasCauseInstanceOf(KeeperException.BadVersionException.class); + // Commit was attempted exactly MAX_ALTER_TABLE_RETRIES=3 times. + assertThat(commitCalls.get()).isEqualTo(3); + // Nothing was persisted. + assertThat(metadataManager.getTable(tablePath).getNumBuckets()) + .isEqualTo(originalBucketCount); + } + + @Test + void testAlterBackfillUsesOldTableCountWithoutAssignmentLookup() throws Exception { + TablePath tablePath = TablePath.of(DEFAULT_DB, "test_alter_bucket_num_no_assign_lookup"); + int originalBucketCount = 4; + String legacyPartition = "2024-02"; + createTableWithLegacyPartition(tablePath, legacyPartition); + + AtomicInteger assignmentReads = new AtomicInteger(); + Configuration wrapperConfig = new Configuration(); + wrapperConfig.set(ConfigOptions.REMOTE_DATA_DIR, remoteDataDir); + ZooKeeperClient noAssignmentReadClient = + new ZooKeeperClient(sharedZkWrapper(), wrapperConfig) { + @Override + public Optional getPartitionAssignment(long partitionId) { + assignmentReads.incrementAndGet(); + return Optional.empty(); + } + }; + MetadataManager manager = metadataManagerOver(noAssignmentReadClient); + + alterBucketNum(manager, tablePath, 8); + + assertThat(assignmentReads).hasValue(0); + assertThat(metadataManager.getTable(tablePath).getNumBuckets()).isEqualTo(8); + assertPartitionCountMatchesAssignment(tablePath, legacyPartition, originalBucketCount); + } + + @Test + void testAlterRejectsMissingPartitionCountAfterEpochAdvanced() throws Exception { + TablePath tablePath = TablePath.of(DEFAULT_DB, "test_missing_count_after_rescale"); + String partitionName = "2024-03"; + createTableWithLegacyPartition(tablePath, partitionName); + + alterBucketNum(metadataManager, tablePath, 8); + assertThat(metadataManager.getTable(tablePath).getBucketCountEpoch()).isEqualTo(1L); + + PartitionRegistration registration = + zookeeperClient.getPartition(tablePath, partitionName).get(); + zookeeperClient.updatePartitionRegistration( + tablePath, + partitionName, + new PartitionRegistration( + registration.getTableId(), + registration.getPartitionId(), + registration.getRemoteDataDir(), + null)); + + assertThatThrownBy(() -> alterBucketNum(metadataManager, tablePath, 16)) + .isInstanceOf(InvalidAlterTableException.class) + .hasMessageContaining("has no persisted bucket count after bucket count epoch 1"); + assertThat(metadataManager.getTable(tablePath).getNumBuckets()).isEqualTo(8); + } + + @Test + void testAlterRetriesOnConcurrentPartitionDeleteNoNode() throws Exception { + // A partition deleted after the backfill enumeration but before the transaction commit + // makes the commit fail with NoNode. The ALTER must re-read the metadata, skip the + // vanished partition on retry, and succeed — not fail permanently. + TablePath tablePath = TablePath.of(DEFAULT_DB, "test_alter_no_node_retry"); + int originalBucketCount = 4; + TableAssignment tableAssignment = + generateAssignment(originalBucketCount, 3, getTabletServers()); + metadataManager.createTable( + tablePath, + remoteDataDir, + partitionedLogTable(originalBucketCount), + tableAssignment, + false); + TableInfo tableInfo = metadataManager.getTable(tablePath); + long tableId = tableInfo.getTableId(); + + String victimPartition = "2024-01"; + metadataManager.createPartition( + tablePath, + tableId, + remoteDataDir, + new PartitionAssignment(tableId, tableAssignment.getBucketAssignments()), + fromPartitionName(tableInfo.getPartitionKeys(), victimPartition), + false, + originalBucketCount); + + MetadataManager noNodeManager = + metadataManagerOver( + zkClientNoNodeOnce( + () -> { + try { + // The concurrent delete wins the race just before commit. + metadataManager.dropPartition( + tablePath, + ResolvedPartitionSpec.fromPartitionName( + tableInfo.getPartitionKeys(), + victimPartition), + true); + } catch (Exception e) { + throw new RuntimeException(e); + } + })); + + alterBucketNum(noNodeManager, tablePath, 8); + + // The retry re-enumerated without the deleted partition and committed successfully. + assertThat(metadataManager.getTable(tablePath).getNumBuckets()).isEqualTo(8); + assertThat(zookeeperClient.getPartition(tablePath, victimPartition)).isEmpty(); + } + + @Test + void testAlterFailsCleanlyWhenTableDeletedMidAlter() throws Exception { + // When the table itself is dropped while an ALTER is in flight, the next retry must + // surface a clear TableNotExistException instead of an opaque ZK error. + TablePath tablePath = TablePath.of(DEFAULT_DB, "test_alter_table_gone"); + int originalBucketCount = 4; + metadataManager.createTable( + tablePath, + remoteDataDir, + partitionedLogTable(originalBucketCount), + generateAssignment(originalBucketCount, 3, getTabletServers()), + false); + + MetadataManager tableGoneManager = + metadataManagerOver( + zkClientNoNodeOnce( + () -> { + try { + metadataManager.dropTable(tablePath, true); + } catch (Exception e) { + throw new RuntimeException(e); + } + })); + + assertThatThrownBy(() -> alterBucketNum(tableGoneManager, tablePath, 8)) + .isInstanceOf(TableNotExistException.class); + } + + // ========================== Helpers ========================== + + private static TabletServerInfo[] getTabletServers() { + return new TabletServerInfo[] { + new TabletServerInfo(0, "rack0"), + new TabletServerInfo(1, "rack1"), + new TabletServerInfo(2, "rack2") + }; + } + + /** A partitioned log table: INT column "a", STRING partition key "dt". */ + private static TableDescriptor partitionedLogTable(int bucketCount) { + return TableDescriptor.builder() + .schema( + Schema.newBuilder() + .column("a", DataTypes.INT()) + .column("dt", DataTypes.STRING()) + .build()) + .distributedBy(bucketCount) + .partitionedBy("dt") + .build() + .withReplicationFactor(3); + } + + private static void alterBucketNum( + MetadataManager manager, TablePath tablePath, int newBucketCount) { + manager.alterBucketCount( + tablePath, newBucketCount, false, null, ZkVersion.MATCH_ANY_VERSION.getVersion()); + } + + private static void alterHistoricalPartition( + MetadataManager manager, TablePath tablePath, boolean enable) { + String key = ConfigOptions.TABLE_DATALAKE_HISTORICAL_PARTITION_ENABLED.key(); + TablePropertyChanges.Builder builder = TablePropertyChanges.builder(); + List changes = new ArrayList<>(); + if (enable) { + builder.setTableProperty(key, "true"); + changes.add(TableChange.set(key, "true")); + } else { + builder.resetTableProperty(key); + changes.add(TableChange.reset(key)); + } + manager.alterTableProperties( + tablePath, + changes, + builder.build(), + false, + null, + (currentTable, updatedTable) -> {}, + (currentTable, updatedTable) -> {}, + ZkVersion.MATCH_ANY_VERSION.getVersion()); + } + + /** + * Creates a partitioned log table with one partition whose persisted bucket count is then + * cleared, so an ALTER bucket.num must enter the backfill path for it. Returns the partition + * id. + */ + private static long createTableWithLegacyPartition(TablePath tablePath, String partitionName) + throws Exception { + int originalBucketCount = 4; + TableAssignment tableAssignment = + generateAssignment(originalBucketCount, 3, getTabletServers()); + metadataManager.createTable( + tablePath, + remoteDataDir, + partitionedLogTable(originalBucketCount), + tableAssignment, + false); + TableInfo tableInfo = metadataManager.getTable(tablePath); + metadataManager.createPartition( + tablePath, + tableInfo.getTableId(), + remoteDataDir, + new PartitionAssignment( + tableInfo.getTableId(), tableAssignment.getBucketAssignments()), + fromPartitionName(tableInfo.getPartitionKeys(), partitionName), + false, + originalBucketCount); + ZooKeeperClient.VersionedData versioned = + zookeeperClient.getPartitionWithVersion(tablePath, partitionName).get(); + zookeeperClient.updatePartitionRegistration( + tablePath, + partitionName, + new PartitionRegistration( + versioned.data().getTableId(), + versioned.data().getPartitionId(), + versioned.data().getRemoteDataDir(), + null)); + return versioned.data().getPartitionId(); + } + + /** Builds a MetadataManager over a decorated ZK client sharing the test cluster. */ + private static MetadataManager metadataManagerOver(ZooKeeperClient decoratedClient) { + return new MetadataManager( + decoratedClient, + new Configuration(), + new LakeCatalogDynamicLoader(new Configuration(), null, true)); + } + + /** Shares the test ZK connection so decorating subclasses can override single methods. */ + private static CuratorFrameworkWithUnhandledErrorListener sharedZkWrapper() { + return zookeeperClient.getCuratorFrameworkWrapper(); + } + + /** + * A ZK client sharing the test connection whose bucket-count commit throws BadVersion for the + * first {@code failures} calls (counted in {@code commitCalls}) and delegates afterwards. + */ + private static ZooKeeperClient zkClientFailingCommits(int failures, AtomicInteger commitCalls) + throws Exception { + Configuration wrapperConfig = new Configuration(); + wrapperConfig.set(ConfigOptions.REMOTE_DATA_DIR, remoteDataDir); + return new ZooKeeperClient(sharedZkWrapper(), wrapperConfig) { + @Override + public void updateTableWithPartitionBucketCountBackfill( + TablePath tp, + TableRegistration reg, + int expectedTableZkVersion, + Map> backfills, + int expectedCoordinatorEpochZkVersion) + throws Exception { + if (commitCalls.getAndIncrement() < failures) { + throw new KeeperException.BadVersionException(); + } + super.updateTableWithPartitionBucketCountBackfill( + tp, + reg, + expectedTableZkVersion, + backfills, + expectedCoordinatorEpochZkVersion); + } + }; + } + + /** + * A ZK client sharing the test connection whose first bucket-count commit performs {@code + * beforeFirstCommit} and then throws NoNode, simulating a concurrent delete winning the race + * just before the commit; later calls delegate. + */ + private static ZooKeeperClient zkClientNoNodeOnce(Runnable beforeFirstCommit) throws Exception { + Configuration wrapperConfig = new Configuration(); + wrapperConfig.set(ConfigOptions.REMOTE_DATA_DIR, remoteDataDir); + return new ZooKeeperClient(sharedZkWrapper(), wrapperConfig) { + private boolean failed; + + @Override + public void updateTableWithPartitionBucketCountBackfill( + TablePath tp, + TableRegistration reg, + int expectedTableZkVersion, + Map> backfills, + int expectedCoordinatorEpochZkVersion) + throws Exception { + if (!failed) { + failed = true; + beforeFirstCommit.run(); + throw new KeeperException.NoNodeException(ZkData.TableZNode.path(tp)); + } + super.updateTableWithPartitionBucketCountBackfill( + tp, + reg, + expectedTableZkVersion, + backfills, + expectedCoordinatorEpochZkVersion); + } + }; + } + + /** + * A ZK client sharing the test connection that runs {@code beforeFirstCommit} right before the + * first bucket-count commit and then delegates, giving a deterministic interleaving without + * injecting a failure. + */ + private static ZooKeeperClient zkClientRunningBeforeFirstCommit(RunnableWithException action) + throws Exception { + Configuration wrapperConfig = new Configuration(); + wrapperConfig.set(ConfigOptions.REMOTE_DATA_DIR, remoteDataDir); + return new ZooKeeperClient(sharedZkWrapper(), wrapperConfig) { + private boolean ran; + + @Override + public void updateTableWithPartitionBucketCountBackfill( + TablePath tp, + TableRegistration reg, + int expectedTableZkVersion, + Map> backfills, + int expectedCoordinatorEpochZkVersion) + throws Exception { + if (!ran) { + ran = true; + action.run(); + } + super.updateTableWithPartitionBucketCountBackfill( + tp, + reg, + expectedTableZkVersion, + backfills, + expectedCoordinatorEpochZkVersion); + } + }; + } +} diff --git a/fluss-server/src/test/java/org/apache/fluss/server/coordinator/AutoPartitionManagerTest.java b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/AutoPartitionManagerTest.java index 72a71d6373a..6f470255ae8 100644 --- a/fluss-server/src/test/java/org/apache/fluss/server/coordinator/AutoPartitionManagerTest.java +++ b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/AutoPartitionManagerTest.java @@ -370,7 +370,8 @@ void testAddPartitionedTable(TestParams params) throws Exception { remoteDataDir, partitionAssignment, fromPartitionName(table.getPartitionKeys(), partitionName), - false); + false, + table.getNumBuckets()); // mock the partition is created in zk. autoPartitionManager.addPartition(tableId, partitionName); } @@ -454,7 +455,8 @@ void testDayFormatWithDashes() throws Exception { remoteDataDir, partitionAssignment, fromPartitionName(table.getPartitionKeys(), "2024-09-15"), - false); + false, + table.getNumBuckets()); autoPartitionManager.addPartition(table.getTableId(), "2024-09-15"); metadataManager.dropPartition( @@ -551,7 +553,8 @@ void testMaxPartitions() throws Exception { remoteDataDir, partitionAssignment, fromPartitionName(table.getPartitionKeys(), i + ""), - false); + false, + table.getNumBuckets()); // mock the partition is created in zk. autoPartitionManager.addPartition(tableId, i + ""); } @@ -734,6 +737,69 @@ void testMaxBucketNumPerPartition() throws Exception { assertThat(partitionsNum).isEqualTo(5); } + /** + * Verifies that the cached {@link TableInfo} update is the eventual-consistency boundary for + * auto-created partition bucket counts: partitions created before the refresh keep the old + * count, while later partitions use the new count. + */ + @Test + void testAutoCreatedPartitionUsesCachedBucketCount() throws Exception { + ZonedDateTime startTime = + LocalDateTime.parse("2024-09-10T00:00:00").atZone(ZoneId.systemDefault()); + long startMs = startTime.toInstant().toEpochMilli(); + ManualClock clock = new ManualClock(startMs); + ManuallyTriggeredScheduledExecutorService periodicExecutor = + new ManuallyTriggeredScheduledExecutorService(); + + AutoPartitionManager autoPartitionManager = + new AutoPartitionManager( + new TestingServerMetadataCache(3), + metadataManager, + remoteDirDynamicLoader, + new Configuration(), + disabledCapacityController(), + clock, + periodicExecutor); + autoPartitionManager.start(); + + // DAY-partitioned table with 4 buckets per partition, never auto-drop, pre-create 4 + TableInfo table = createPartitionedTableWithBuckets(-1, 4, AutoPartitionTimeUnit.DAY, 4); + TablePath tablePath = table.getTablePath(); + autoPartitionManager.addAutoPartitionTable(table, true); + periodicExecutor.triggerNonPeriodicScheduledTask(); + + Map partitions = + zookeeperClient.getPartitionRegistrations(tablePath); + assertThat(partitions.keySet()) + .containsExactlyInAnyOrder("20240910", "20240911", "20240912", "20240913"); + // all pre-created partitions carry the original bucket count 4 + for (PartitionRegistration reg : partitions.values()) { + assertThat(reg.getBucketCount()).isEqualTo(4); + } + + // Simulate the first half of ALTER bucket.num 4 -> 8: ZK is updated, but the coordinator + // event has not refreshed AutoPartitionManager's cached TableInfo yet. + TableRegistration reg = zookeeperClient.getTable(tablePath).get(); + zookeeperClient.updateTable(tablePath, reg.newBucketCount(8)); + clock.advanceTime(Duration.ofDays(1).plusHours(23)); + periodicExecutor.triggerPeriodicScheduledTasks(); + + partitions = zookeeperClient.getPartitionRegistrations(tablePath); + assertThat(partitions.get("20240914").getBucketCount()).isEqualTo(4); + + // Complete event propagation and create the next partition from the refreshed cache. + TableInfo updatedTable = createUpdatedBucketCountTableInfo(table, 8); + autoPartitionManager.updateAutoPartitionTables(updatedTable); + periodicExecutor.triggerNonPeriodicScheduledTask(); + clock.advanceTime(Duration.ofDays(1)); + periodicExecutor.triggerPeriodicScheduledTasks(); + + partitions = zookeeperClient.getPartitionRegistrations(tablePath); + assertThat(partitions.get("20240910").getBucketCount()).isEqualTo(4); + assertThat(partitions.get("20240914").getBucketCount()).isEqualTo(4); + assertThat(partitions.get("20240915").getBucketCount()).isEqualTo(8); + } + @Test void testAutoCreatePartitionChecksCapacityWithoutReservation() throws Exception { ZonedDateTime startTime = @@ -795,14 +861,16 @@ void testAutoDropPartitionDoesNotMutateObservedKvLeaderReplicaCount() throws Exc remoteDataDir, partitionAssignment, fromPartitionName(table.getPartitionKeys(), "2025042600"), - false); + false, + table.getNumBuckets()); metadataManager.createPartition( tablePath, table.getTableId(), remoteDataDir, partitionAssignment, fromPartitionName(table.getPartitionKeys(), "2025042601"), - false); + false, + table.getNumBuckets()); autoPartitionManager.addPartition(table.getTableId(), "2025042600"); autoPartitionManager.addPartition(table.getTableId(), "2025042601"); capacityController.updateObservedKvLeaderReplicaCount((long) table.getNumBuckets() * 2); @@ -1218,7 +1286,8 @@ private void createPartition( remoteDataDir, partitionAssignment, fromPartitionName(tableInfo.getPartitionKeys(), partitionName), - false); + false, + bucketAssignments.size()); autoPartitionManager.addPartition(tableInfo.getTableId(), partitionName); } @@ -1420,6 +1489,24 @@ private TableInfo createUpdatedHistoricalPartitionEnabledTableInfo( return createUpdatedTableInfo(original, newProperties); } + /** Creates a new TableInfo with an updated table-level bucket count, reusing the original. */ + private TableInfo createUpdatedBucketCountTableInfo(TableInfo original, int newNumBuckets) { + return new TableInfo( + original.getTablePath(), + original.getTableId(), + original.getSchemaId(), + original.getSchema(), + original.getBucketKeys(), + original.getPartitionKeys(), + newNumBuckets, + original.getProperties(), + original.getCustomProperties(), + original.getRemoteDataDir(), + original.getComment().orElse(null), + original.getCreatedTime(), + System.currentTimeMillis()); + } + private TableInfo createUpdatedTableInfo(TableInfo original, Configuration newProperties) { return new TableInfo( original.getTablePath(), diff --git a/fluss-server/src/test/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessorTest.java b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessorTest.java index a39b3c935db..b26d8cabca2 100644 --- a/fluss-server/src/test/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessorTest.java +++ b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessorTest.java @@ -27,7 +27,8 @@ import org.apache.fluss.exception.InvalidAlterTableException; import org.apache.fluss.exception.InvalidCoordinatorException; import org.apache.fluss.fs.FsPath; -import org.apache.fluss.metadata.DatabaseDescriptor; +import org.apache.fluss.metadata.PartitionSpec; +import org.apache.fluss.metadata.ResolvedPartitionSpec; import org.apache.fluss.metadata.Schema; import org.apache.fluss.metadata.TableBucket; import org.apache.fluss.metadata.TableBucketReplica; @@ -60,7 +61,6 @@ import org.apache.fluss.server.coordinator.event.CoordinatorEventManager; import org.apache.fluss.server.coordinator.event.NotifyLeaderAndIsrResponseReceivedEvent; import org.apache.fluss.server.coordinator.event.RetryOfflineLeaderEvent; -import org.apache.fluss.server.coordinator.lease.KvSnapshotLeaseManager; import org.apache.fluss.server.coordinator.remote.RemoteDirDynamicLoader; import org.apache.fluss.server.coordinator.statemachine.BucketState; import org.apache.fluss.server.coordinator.statemachine.ReplicaState; @@ -79,33 +79,21 @@ import org.apache.fluss.server.metrics.group.TestingMetricGroups; import org.apache.fluss.server.tablet.TestTabletServerGateway; import org.apache.fluss.server.zk.NOPErrorHandler; -import org.apache.fluss.server.zk.ZkEpoch; import org.apache.fluss.server.zk.ZooKeeperClient; -import org.apache.fluss.server.zk.ZooKeeperExtension; import org.apache.fluss.server.zk.data.BucketAssignment; -import org.apache.fluss.server.zk.data.CoordinatorAddress; import org.apache.fluss.server.zk.data.LeaderAndIsr; import org.apache.fluss.server.zk.data.PartitionAssignment; import org.apache.fluss.server.zk.data.TableAssignment; import org.apache.fluss.server.zk.data.TabletServerRegistration; import org.apache.fluss.server.zk.data.ZkData; -import org.apache.fluss.server.zk.data.ZkData.PartitionIdsZNode; -import org.apache.fluss.server.zk.data.ZkData.TableIdsZNode; -import org.apache.fluss.testutils.common.AllCallbackWrapper; +import org.apache.fluss.server.zk.data.ZkVersion; import org.apache.fluss.testutils.common.ManuallyTriggeredScheduledExecutorService; import org.apache.fluss.types.DataTypes; import org.apache.fluss.utils.ExceptionUtils; import org.apache.fluss.utils.clock.SystemClock; -import org.apache.fluss.utils.concurrent.ExecutorThreadFactory; -import org.apache.fluss.utils.concurrent.FlussScheduler; -import org.apache.fluss.utils.concurrent.Scheduler; import org.apache.fluss.utils.types.Tuple2; -import org.junit.jupiter.api.AfterEach; -import org.junit.jupiter.api.BeforeAll; -import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.extension.RegisterExtension; import org.junit.jupiter.api.io.TempDir; import java.nio.file.Path; @@ -121,7 +109,6 @@ import java.util.concurrent.CompletableFuture; import java.util.concurrent.ConcurrentLinkedDeque; import java.util.concurrent.CountDownLatch; -import java.util.concurrent.Executors; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; import java.util.function.BiConsumer; @@ -149,7 +136,7 @@ import static org.assertj.core.api.Assertions.assertThatThrownBy; /** Test for {@link CoordinatorEventProcessor}. */ -class CoordinatorEventProcessorTest { +class CoordinatorEventProcessorTest extends CoordinatorEventProcessorTestBase { private static final int N_BUCKETS = 3; private static final int REPLICATION_FACTOR = 3; @@ -166,107 +153,6 @@ class CoordinatorEventProcessorTest { .build() .withReplicationFactor(REPLICATION_FACTOR); - @RegisterExtension - public static final AllCallbackWrapper ZOO_KEEPER_EXTENSION_WRAPPER = - new AllCallbackWrapper<>(new ZooKeeperExtension()); - - private static ZooKeeperClient zookeeperClient; - private static MetadataManager metadataManager; - private static ZkEpoch zkEpoch; - - private CoordinatorEventProcessor eventProcessor; - private final String defaultDatabase = "db"; - private TestCoordinatorChannelManager testCoordinatorChannelManager; - private AutoPartitionManager autoPartitionManager; - private LakeTableTieringManager lakeTableTieringManager; - private CompletedSnapshotStoreManager completedSnapshotStoreManager; - private CoordinatorMetadataCache serverMetadataCache; - private ReplicaCapacityController replicaCapacityController; - private KvSnapshotLeaseManager kvSnapshotLeaseManager; - private Scheduler scheduler; - private String remoteDataDir; - - @BeforeAll - static void baseBeforeAll() throws Exception { - zookeeperClient = - ZOO_KEEPER_EXTENSION_WRAPPER - .getCustomExtension() - .getZooKeeperClient(NOPErrorHandler.INSTANCE); - metadataManager = - new MetadataManager( - zookeeperClient, - new Configuration(), - new LakeCatalogDynamicLoader(new Configuration(), null, true)); - - // register coordinator server - zookeeperClient.registerCoordinatorLeader( - new CoordinatorAddress( - "2", Endpoint.fromListenersString("CLIENT://localhost:10012"))); - - zkEpoch = zookeeperClient.fenceBecomeCoordinatorLeader("2"); - // register 3 tablet servers - for (int i = 0; i < 3; i++) { - zookeeperClient.registerTabletServer( - i, - new TabletServerRegistration( - "rack" + i, - Collections.singletonList( - new Endpoint("host" + i, 1000, DEFAULT_LISTENER_NAME)), - System.currentTimeMillis())); - } - } - - @BeforeEach - void beforeEach() { - serverMetadataCache = new CoordinatorMetadataCache(); - // set a test channel manager for the context - testCoordinatorChannelManager = new TestCoordinatorChannelManager(); - lakeTableTieringManager = - new LakeTableTieringManager(TestingMetricGroups.LAKE_TIERING_METRICS); - remoteDataDir = zookeeperClient.getDefaultRemoteDataDir(); - Configuration conf = new Configuration(); - conf.setString(ConfigOptions.REMOTE_DATA_DIR, remoteDataDir); - replicaCapacityController = new ReplicaCapacityController(conf, serverMetadataCache); - autoPartitionManager = - new AutoPartitionManager( - serverMetadataCache, - metadataManager, - new RemoteDirDynamicLoader(conf), - conf, - replicaCapacityController); - kvSnapshotLeaseManager = - new KvSnapshotLeaseManager( - Duration.ofMinutes(10).toMillis(), - zookeeperClient, - remoteDataDir, - SystemClock.getInstance(), - TestingMetricGroups.COORDINATOR_METRICS); - kvSnapshotLeaseManager.start(); - - scheduler = new FlussScheduler(1); - scheduler.startup(); - - eventProcessor = buildCoordinatorEventProcessor(); - eventProcessor.startup(); - metadataManager.createDatabase( - defaultDatabase, DatabaseDescriptor.builder().build(), false); - completedSnapshotStoreManager = eventProcessor.completedSnapshotStoreManager(); - } - - @AfterEach - void afterEach() throws Exception { - if (eventProcessor != null) { - eventProcessor.shutdown(); - } - if (scheduler != null) { - scheduler.shutdown(); - } - metadataManager.dropDatabase(defaultDatabase, false, true); - // clear the assignment info for all tables; - ZOO_KEEPER_EXTENSION_WRAPPER.getCustomExtension().cleanupPath(TableIdsZNode.path()); - ZOO_KEEPER_EXTENSION_WRAPPER.getCustomExtension().cleanupPath(PartitionIdsZNode.path()); - } - @Test void testLoadedAssignmentsTrackKnownKvAndUnknownTablesConservatively() throws Exception { long kvTableId = 10001L; @@ -1772,6 +1658,96 @@ void testSchemaChange() throws Exception { 3, new TableMetadata(tableInfo2, Collections.emptyList()))); } + @Test + void testSchemaChangeKeepsBucketCountEpochAfterRescale() throws Exception { + initCoordinatorChannel(); + TablePath t1 = TablePath.of(defaultDatabase, "schema_change_keeps_epoch"); + int originalBucketCount = 3; + TableAssignment tableAssignment = + generateAssignment( + originalBucketCount, + REPLICATION_FACTOR, + new TabletServerInfo[] { + new TabletServerInfo(0, "rack0"), + new TabletServerInfo(1, "rack1"), + new TabletServerInfo(2, "rack2") + }); + TableDescriptor partitionedTable = + TableDescriptor.builder() + .schema( + Schema.newBuilder() + .column("a", DataTypes.INT()) + .column("b", DataTypes.STRING()) + .primaryKey("a", "b") + .build()) + .distributedBy(originalBucketCount) + .partitionedBy("b") + .build() + .withReplicationFactor(REPLICATION_FACTOR); + long tableId = + metadataManager.createTable( + t1, remoteDataDir, partitionedTable, tableAssignment, false); + // create one partition so the ALTER bucket.num rescale can commit + metadataManager.createPartition( + t1, + tableId, + remoteDataDir, + new PartitionAssignment( + tableId, + generateAssignment( + originalBucketCount, + REPLICATION_FACTOR, + new TabletServerInfo[] { + new TabletServerInfo(0, "rack0"), + new TabletServerInfo(1, "rack1"), + new TabletServerInfo(2, "rack2") + }) + .getBucketAssignments()), + ResolvedPartitionSpec.fromPartitionSpec( + Collections.singletonList("b"), + new PartitionSpec(Collections.singletonMap("b", "2024-01-01"))), + false, + originalBucketCount); + + // ALTER bucket.num advances the bucket count epoch (persisted in ZK TableRegistration) + metadataManager.alterBucketCount( + t1, 8, false, null, ZkVersion.MATCH_ANY_VERSION.getVersion()); + + long epochAfterAlter = metadataManager.getTable(t1).getBucketCountEpoch(); + assertThat(epochAfterAlter).isGreaterThan(0L); + + // A later schema change rebuilds the context TableInfo; the epoch must survive it + alterTable( + t1, + Collections.singletonList( + TableChange.addColumn( + "add_column", + DataTypes.INT(), + null, + TableChange.ColumnPosition.last()))); + + retryVerifyContext( + ctx -> { + TableInfo tableInfoInCtx = ctx.getTableInfoById(tableId); + assertThat(tableInfoInCtx).isNotNull(); + // the schema change took effect + assertThat(tableInfoInCtx.getSchema().getColumnNames()).contains("add_column"); + // and the bucket count epoch did NOT roll back to 0 + assertThat(tableInfoInCtx.getBucketCountEpoch()).isEqualTo(epochAfterAlter); + }); + + // the UpdateMetadata pushed for the schema change carries the same epoch + TableInfo tableInfoAfterSchemaChange = metadataManager.getTable(t1); + assertThat(tableInfoAfterSchemaChange.getBucketCountEpoch()).isEqualTo(epochAfterAlter); + retry( + Duration.ofMinutes(1), + () -> + verifyMetadataUpdateRequest( + 3, + new TableMetadata( + tableInfoAfterSchemaChange, Collections.emptyList()))); + } + @Test void testTableRegistrationChange() throws Exception { // make sure all request to gateway should be successful @@ -1822,7 +1798,8 @@ void testTableRegistrationChange() throws Exception { false, null, (currentTable, updatedTable) -> {}, - (currentTable, updatedTable) -> {}); + (currentTable, updatedTable) -> {}, + ZkVersion.MATCH_ANY_VERSION.getVersion()); // get updated table info and verify metadata update request is sent TableInfo updatedTableInfo = metadataManager.getTable(t1); @@ -1889,7 +1866,8 @@ void testAlterStandbyReplicaEnabled() throws Exception { false, null, (currentTable, updatedTable) -> {}, - (currentTable, updatedTable) -> {}); + (currentTable, updatedTable) -> {}, + ZkVersion.MATCH_ANY_VERSION.getVersion()); // verify standby replicas are removed after re-election retryVerifyContext( @@ -1963,7 +1941,8 @@ void testAlterEnableStandbyReplicaForExistingTable() throws Exception { false, null, (currentTable, updatedTable) -> {}, - (currentTable, updatedTable) -> {}); + (currentTable, updatedTable) -> {}, + ZkVersion.MATCH_ANY_VERSION.getVersion()); // Verify re-election happened: standby assigned and leaderEpoch incremented retryVerifyContext( @@ -2026,7 +2005,8 @@ void testAlterStandbyReplicaEnabledForLogTable() throws Exception { false, null, (currentTable, updatedTable) -> {}, - (currentTable, updatedTable) -> {})) + (currentTable, updatedTable) -> {}, + ZkVersion.MATCH_ANY_VERSION.getVersion())) .isInstanceOf(InvalidAlterTableException.class) .hasMessageContaining("can only be altered on primary key tables"); @@ -2043,7 +2023,8 @@ void testAlterStandbyReplicaEnabledForLogTable() throws Exception { false, null, (currentTable, updatedTable) -> {}, - (currentTable, updatedTable) -> {})) + (currentTable, updatedTable) -> {}, + ZkVersion.MATCH_ANY_VERSION.getVersion())) .isInstanceOf(InvalidAlterTableException.class) .hasMessageContaining("can only be altered on primary key tables"); } @@ -2325,27 +2306,6 @@ private void verifyIsr(TableBucket tb, int expectedLeader, List expecte .hasSameElementsAs(expectedIsr); } - private CoordinatorEventProcessor buildCoordinatorEventProcessor() { - Configuration conf = new Configuration(); - conf.set(ConfigOptions.REMOTE_DATA_DIR, remoteDataDir); - conf.set(ConfigOptions.COORDINATOR_OFFLINE_LEADER_RETRY_DELAY, Duration.ofDays(1)); - return new CoordinatorEventProcessor( - zookeeperClient, - serverMetadataCache, - testCoordinatorChannelManager, - new CoordinatorContext(zkEpoch), - replicaCapacityController, - autoPartitionManager, - lakeTableTieringManager, - TestingMetricGroups.COORDINATOR_METRICS, - conf, - Executors.newFixedThreadPool(1, new ExecutorThreadFactory("test-coordinator-io")), - metadataManager, - kvSnapshotLeaseManager, - scheduler, - SystemClock.getInstance()); - } - private static class FailingUpdateMetadataChannelManager extends TestCoordinatorChannelManager { private final CountDownLatch failureObserved = new CountDownLatch(1); @@ -2468,14 +2428,16 @@ private Tuple2 preparePartitionAssignment( partitionAssignment, remoteDataDir, tablePath, - tableId); + tableId, + partitionAssignment.getBucketAssignments().size()); zookeeperClient.registerPartitionAssignmentAndMetadata( partition2Id, partition2Name, partitionAssignment, remoteDataDir, tablePath, - tableId); + tableId, + partitionAssignment.getBucketAssignments().size()); return Tuple2.of( new PartitionIdName(partition1Id, partition1Name), diff --git a/fluss-server/src/test/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessorTestBase.java b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessorTestBase.java new file mode 100644 index 00000000000..372abf19f8e --- /dev/null +++ b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessorTestBase.java @@ -0,0 +1,183 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.fluss.server.coordinator; + +import org.apache.fluss.cluster.Endpoint; +import org.apache.fluss.config.ConfigOptions; +import org.apache.fluss.config.Configuration; +import org.apache.fluss.metadata.DatabaseDescriptor; +import org.apache.fluss.server.coordinator.lease.KvSnapshotLeaseManager; +import org.apache.fluss.server.coordinator.remote.RemoteDirDynamicLoader; +import org.apache.fluss.server.metadata.CoordinatorMetadataCache; +import org.apache.fluss.server.metrics.group.TestingMetricGroups; +import org.apache.fluss.server.zk.NOPErrorHandler; +import org.apache.fluss.server.zk.ZkEpoch; +import org.apache.fluss.server.zk.ZooKeeperClient; +import org.apache.fluss.server.zk.ZooKeeperExtension; +import org.apache.fluss.server.zk.data.CoordinatorAddress; +import org.apache.fluss.server.zk.data.TabletServerRegistration; +import org.apache.fluss.server.zk.data.ZkData.PartitionIdsZNode; +import org.apache.fluss.server.zk.data.ZkData.TableIdsZNode; +import org.apache.fluss.testutils.common.AllCallbackWrapper; +import org.apache.fluss.utils.clock.SystemClock; +import org.apache.fluss.utils.concurrent.ExecutorThreadFactory; +import org.apache.fluss.utils.concurrent.FlussScheduler; +import org.apache.fluss.utils.concurrent.Scheduler; + +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.extension.RegisterExtension; + +import java.time.Duration; +import java.util.Collections; +import java.util.concurrent.Executors; + +import static org.apache.fluss.config.ConfigOptions.DEFAULT_LISTENER_NAME; + +/** + * The shared lifecycle harness for {@link CoordinatorEventProcessor} unit tests: a ZooKeeper test + * cluster, a coordinator event processor rebuilt per test, and the cleanup between tests. + * + *

Extracted from {@code CoordinatorEventProcessorTest} to deduplicate the harness and keep the + * test class within the checkstyle file-length limit. + */ +class CoordinatorEventProcessorTestBase { + + @RegisterExtension + public static final AllCallbackWrapper ZOO_KEEPER_EXTENSION_WRAPPER = + new AllCallbackWrapper<>(new ZooKeeperExtension()); + + protected static ZooKeeperClient zookeeperClient; + protected static MetadataManager metadataManager; + protected static ZkEpoch zkEpoch; + + protected CoordinatorEventProcessor eventProcessor; + protected final String defaultDatabase = "db"; + protected TestCoordinatorChannelManager testCoordinatorChannelManager; + protected AutoPartitionManager autoPartitionManager; + protected LakeTableTieringManager lakeTableTieringManager; + protected CompletedSnapshotStoreManager completedSnapshotStoreManager; + protected CoordinatorMetadataCache serverMetadataCache; + protected ReplicaCapacityController replicaCapacityController; + protected KvSnapshotLeaseManager kvSnapshotLeaseManager; + protected Scheduler scheduler; + protected String remoteDataDir; + + @BeforeAll + static void baseBeforeAll() throws Exception { + zookeeperClient = + ZOO_KEEPER_EXTENSION_WRAPPER + .getCustomExtension() + .getZooKeeperClient(NOPErrorHandler.INSTANCE); + metadataManager = + new MetadataManager( + zookeeperClient, + new Configuration(), + new LakeCatalogDynamicLoader(new Configuration(), null, true)); + + // register coordinator server + zookeeperClient.registerCoordinatorLeader( + new CoordinatorAddress( + "2", Endpoint.fromListenersString("CLIENT://localhost:10012"))); + + zkEpoch = zookeeperClient.fenceBecomeCoordinatorLeader("2"); + // register 3 tablet servers + for (int i = 0; i < 3; i++) { + zookeeperClient.registerTabletServer( + i, + new TabletServerRegistration( + "rack" + i, + Collections.singletonList( + new Endpoint("host" + i, 1000, DEFAULT_LISTENER_NAME)), + System.currentTimeMillis())); + } + } + + @BeforeEach + void beforeEach() { + serverMetadataCache = new CoordinatorMetadataCache(); + // set a test channel manager for the context + testCoordinatorChannelManager = new TestCoordinatorChannelManager(); + lakeTableTieringManager = + new LakeTableTieringManager(TestingMetricGroups.LAKE_TIERING_METRICS); + remoteDataDir = zookeeperClient.getDefaultRemoteDataDir(); + Configuration conf = new Configuration(); + conf.setString(ConfigOptions.REMOTE_DATA_DIR, remoteDataDir); + replicaCapacityController = new ReplicaCapacityController(conf, serverMetadataCache); + autoPartitionManager = + new AutoPartitionManager( + serverMetadataCache, + metadataManager, + new RemoteDirDynamicLoader(conf), + conf, + replicaCapacityController); + kvSnapshotLeaseManager = + new KvSnapshotLeaseManager( + Duration.ofMinutes(10).toMillis(), + zookeeperClient, + remoteDataDir, + SystemClock.getInstance(), + TestingMetricGroups.COORDINATOR_METRICS); + kvSnapshotLeaseManager.start(); + + scheduler = new FlussScheduler(1); + scheduler.startup(); + + eventProcessor = buildCoordinatorEventProcessor(); + eventProcessor.startup(); + metadataManager.createDatabase( + defaultDatabase, DatabaseDescriptor.builder().build(), false); + completedSnapshotStoreManager = eventProcessor.completedSnapshotStoreManager(); + } + + @AfterEach + void afterEach() throws Exception { + if (eventProcessor != null) { + eventProcessor.shutdown(); + } + if (scheduler != null) { + scheduler.shutdown(); + } + metadataManager.dropDatabase(defaultDatabase, false, true); + // clear the assignment info for all tables; + ZOO_KEEPER_EXTENSION_WRAPPER.getCustomExtension().cleanupPath(TableIdsZNode.path()); + ZOO_KEEPER_EXTENSION_WRAPPER.getCustomExtension().cleanupPath(PartitionIdsZNode.path()); + } + + protected CoordinatorEventProcessor buildCoordinatorEventProcessor() { + Configuration conf = new Configuration(); + conf.set(ConfigOptions.REMOTE_DATA_DIR, remoteDataDir); + conf.set(ConfigOptions.COORDINATOR_OFFLINE_LEADER_RETRY_DELAY, Duration.ofDays(1)); + return new CoordinatorEventProcessor( + zookeeperClient, + serverMetadataCache, + testCoordinatorChannelManager, + new CoordinatorContext(zkEpoch), + replicaCapacityController, + autoPartitionManager, + lakeTableTieringManager, + TestingMetricGroups.COORDINATOR_METRICS, + conf, + Executors.newFixedThreadPool(1, new ExecutorThreadFactory("test-coordinator-io")), + metadataManager, + kvSnapshotLeaseManager, + scheduler, + SystemClock.getInstance()); + } +} diff --git a/fluss-server/src/test/java/org/apache/fluss/server/coordinator/CoordinatorRequestBatchTest.java b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/CoordinatorRequestBatchTest.java index 0b3892a1cc2..3c925914261 100644 --- a/fluss-server/src/test/java/org/apache/fluss/server/coordinator/CoordinatorRequestBatchTest.java +++ b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/CoordinatorRequestBatchTest.java @@ -20,9 +20,11 @@ import org.apache.fluss.exception.NetworkException; import org.apache.fluss.metadata.PhysicalTablePath; import org.apache.fluss.metadata.TableBucket; +import org.apache.fluss.metadata.TableInfo; import org.apache.fluss.metadata.TablePath; import org.apache.fluss.rpc.messages.NotifyLeaderAndIsrRequest; import org.apache.fluss.rpc.messages.NotifyLeaderAndIsrResponse; +import org.apache.fluss.rpc.messages.PbNotifyLeaderAndIsrReqForBucket; import org.apache.fluss.server.coordinator.event.AccessContextEvent; import org.apache.fluss.server.coordinator.event.EventManager; import org.apache.fluss.server.zk.ZkEpoch; @@ -33,8 +35,11 @@ import java.util.Arrays; import java.util.Collections; +import java.util.concurrent.atomic.AtomicReference; import java.util.function.BiConsumer; +import static org.apache.fluss.record.TestData.DATA1_TABLE_DESCRIPTOR; +import static org.apache.fluss.record.TestData.DEFAULT_REMOTE_DATA_DIR; import static org.assertj.core.api.Assertions.assertThat; /** Test for the {@link CoordinatorRequestBatch}. */ @@ -59,6 +64,7 @@ void testNotifyLeaderAndIsrSendFailureClearsLeaderPending() { TablePath tablePath = TablePath.of("db1", "t1"); coordinatorContext.putTablePath(tableId, tablePath); + putTableInfo(tableId, tablePath); coordinatorContext.setLiveTabletServers( CoordinatorTestUtils.createServers(Collections.singletonList(0))); coordinatorContext.updateBucketReplicaAssignment(tb, Collections.singletonList(0)); @@ -98,6 +104,7 @@ void testNotifyLeaderAndIsrSendFailureToFollowerDoesNotClearOtherPending() { TablePath tablePath = TablePath.of("db1", "t2"); coordinatorContext.putTablePath(tableId, tablePath); + putTableInfo(tableId, tablePath); coordinatorContext.setLiveTabletServers( CoordinatorTestUtils.createServers(Arrays.asList(0, 1))); coordinatorContext.updateBucketReplicaAssignment(followerTb, Arrays.asList(0, 1)); @@ -132,6 +139,70 @@ void testNotifyLeaderAndIsrSendFailureToFollowerDoesNotClearOtherPending() { .containsExactly(otherLeaderTb); } + @Test + void testNotifyLeaderAndIsrAllowsMissingRoutingState() { + long tableId = 300L; + TableBucket tb = new TableBucket(tableId, 0); + TablePath tablePath = TablePath.of("db1", "t3"); + coordinatorContext.putTablePath(tableId, tablePath); + coordinatorContext.setLiveTabletServers( + CoordinatorTestUtils.createServers(Collections.singletonList(0))); + LeaderAndIsr leaderAndIsr = + new LeaderAndIsr(0, 0, Collections.singletonList(0), Collections.emptyList(), 0, 0); + coordinatorContext.putBucketLeaderAndIsr(tb, leaderAndIsr); + + AtomicReference sentRequest = new AtomicReference<>(); + TestCoordinatorChannelManager channelManager = + new TestCoordinatorChannelManager() { + @Override + public void sendBucketLeaderAndIsrRequest( + int receiveServerId, + NotifyLeaderAndIsrRequest request, + BiConsumer + responseConsumer) { + sentRequest.set(request); + responseConsumer.accept( + null, new NetworkException("simulated send failure for test")); + } + }; + CoordinatorRequestBatch batch = + new CoordinatorRequestBatch( + channelManager, + newSynchronousAccessContextEventManager(), + coordinatorContext); + + batch.addNotifyLeaderRequestForTabletServers( + Collections.singleton(0), + PhysicalTablePath.of(tablePath), + tb, + Collections.singletonList(0), + leaderAndIsr); + + assertThat(coordinatorContext.getPendingLeaderActivationBuckets()).isEmpty(); + batch.sendRequestToTabletServers(0); + + assertThat(sentRequest.get()).isNotNull(); + assertThat(sentRequest.get().getNotifyBucketsLeaderReqsList()).hasSize(1); + PbNotifyLeaderAndIsrReqForBucket bucketRequest = + sentRequest.get().getNotifyBucketsLeaderReqsList().get(0); + assertThat(bucketRequest.hasBucketCount()).isFalse(); + assertThat(bucketRequest.hasBucketCountEpoch()).isFalse(); + assertThat(coordinatorContext.getPendingLeaderActivationBuckets()).isEmpty(); + } + + /** Registers table metadata so normal notifications carry the bucket layout epoch. */ + private void putTableInfo(long tableId, TablePath tablePath) { + coordinatorContext.putTableInfo( + TableInfo.of( + tablePath, + tableId, + 0, + DATA1_TABLE_DESCRIPTOR, + DEFAULT_REMOTE_DATA_DIR, + System.currentTimeMillis(), + System.currentTimeMillis())); + } + private static TestCoordinatorChannelManager newAlwaysFailingChannelManager() { return new TestCoordinatorChannelManager() { @Override diff --git a/fluss-server/src/test/java/org/apache/fluss/server/coordinator/TableManagerTest.java b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/TableManagerTest.java index 8116cfcb35b..235d09af841 100644 --- a/fluss-server/src/test/java/org/apache/fluss/server/coordinator/TableManagerTest.java +++ b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/TableManagerTest.java @@ -318,7 +318,8 @@ void testCreateAndDropPartition() throws Exception { partitionAssignment, DEFAULT_REMOTE_DATA_DIR, DATA1_TABLE_PATH, - tableId); + tableId, + partitionAssignment.getBucketAssignments().size()); // create partition tableManager.onCreateNewPartition( diff --git a/fluss-server/src/test/java/org/apache/fluss/server/coordinator/event/watcher/TableChangeWatcherTest.java b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/event/watcher/TableChangeWatcherTest.java index 0b0a1b1daa5..28a4131ecd8 100644 --- a/fluss-server/src/test/java/org/apache/fluss/server/coordinator/event/watcher/TableChangeWatcherTest.java +++ b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/event/watcher/TableChangeWatcherTest.java @@ -46,6 +46,7 @@ import org.apache.fluss.server.zk.data.PartitionAssignment; import org.apache.fluss.server.zk.data.TableAssignment; import org.apache.fluss.server.zk.data.TableRegistration; +import org.apache.fluss.server.zk.data.ZkVersion; import org.apache.fluss.testutils.common.AllCallbackWrapper; import org.apache.fluss.types.DataTypes; import org.apache.fluss.utils.clock.SystemClock; @@ -252,9 +253,21 @@ void testPartitionedTable() throws Exception { .getBucketAssignments()); // register assignment and metadata zookeeperClient.registerPartitionAssignmentAndMetadata( - 1L, "2011", partitionAssignment, remoteDataDir, tablePath, tableId); + 1L, + "2011", + partitionAssignment, + remoteDataDir, + tablePath, + tableId, + partitionAssignment.getBucketAssignments().size()); zookeeperClient.registerPartitionAssignmentAndMetadata( - 2L, "2022", partitionAssignment, remoteDataDir, tablePath, tableId); + 2L, + "2022", + partitionAssignment, + remoteDataDir, + tablePath, + tableId, + partitionAssignment.getBucketAssignments().size()); // create partitions events expectedEvents.add( @@ -415,7 +428,8 @@ void testTableRegistrationChange() { false, null, (currentTable, updatedTable) -> {}, - (currentTable, updatedTable) -> {}); + (currentTable, updatedTable) -> {}, + ZkVersion.MATCH_ANY_VERSION.getVersion()); // get the updated table registration TableRegistration updatedTableRegistration = diff --git a/fluss-server/src/test/java/org/apache/fluss/server/log/remote/RemoteLogManagerTest.java b/fluss-server/src/test/java/org/apache/fluss/server/log/remote/RemoteLogManagerTest.java index 5728dbf6f4a..9af86ea5ecf 100644 --- a/fluss-server/src/test/java/org/apache/fluss/server/log/remote/RemoteLogManagerTest.java +++ b/fluss-server/src/test/java/org/apache/fluss/server/log/remote/RemoteLogManagerTest.java @@ -487,7 +487,9 @@ void testRemoteFirstFetchRejectsNonLeader(boolean partitionTable) throws Excepti Arrays.asList(TABLET_SERVER_ID, newLeaderId), Collections.emptyList(), INITIAL_COORDINATOR_EPOCH, - INITIAL_BUCKET_EPOCH + 1))); + INITIAL_BUCKET_EPOCH + 1), + 3, + 0L)); CompletableFuture> fetchFuture = new CompletableFuture<>(); diff --git a/fluss-server/src/test/java/org/apache/fluss/server/log/remote/RemoteLogTestBase.java b/fluss-server/src/test/java/org/apache/fluss/server/log/remote/RemoteLogTestBase.java index a533c3ecda7..d0ed729eb8b 100644 --- a/fluss-server/src/test/java/org/apache/fluss/server/log/remote/RemoteLogTestBase.java +++ b/fluss-server/src/test/java/org/apache/fluss/server/log/remote/RemoteLogTestBase.java @@ -88,12 +88,9 @@ private Replica makeReplicaAndAddSegments( tb, Collections.singletonList(0), new LeaderAndIsr( - 0, - 0, - Collections.singletonList(0), - Collections.emptyList(), - 0, - 0))); + 0, 0, Collections.singletonList(0), Collections.emptyList(), 0, 0), + 3, + 0L)); addMultiSegmentsToLogTablet(replica.getLogTablet(), segmentSize); return replica; } diff --git a/fluss-server/src/test/java/org/apache/fluss/server/metadata/TabletServerMetadataCacheTest.java b/fluss-server/src/test/java/org/apache/fluss/server/metadata/TabletServerMetadataCacheTest.java index 213d60b0ebe..e0dcdee3e83 100644 --- a/fluss-server/src/test/java/org/apache/fluss/server/metadata/TabletServerMetadataCacheTest.java +++ b/fluss-server/src/test/java/org/apache/fluss/server/metadata/TabletServerMetadataCacheTest.java @@ -47,12 +47,14 @@ import static org.apache.fluss.record.TestData.DATA1_TABLE_PATH; import static org.apache.fluss.record.TestData.DEFAULT_REMOTE_DATA_DIR; import static org.apache.fluss.server.metadata.PartitionMetadata.DELETED_PARTITION_ID; +import static org.apache.fluss.server.metadata.PartitionMetadata.DELETED_PARTITION_NAME; import static org.apache.fluss.server.metadata.TableMetadata.DELETED_TABLE_ID; import static org.apache.fluss.server.zk.data.LeaderAndIsr.NO_LEADER; import static org.assertj.core.api.Assertions.assertThat; /** Test for {@link TabletServerMetadataCache}. */ public class TabletServerMetadataCacheTest { + private TabletServerMetadataCache serverMetadataCache; private ServerInfo coordinatorServer; private Set aliveTableServers; @@ -319,6 +321,145 @@ private void assertTableMetadataEquals( .hasSameElementsAs(expectedBucketMetadataList); } + @Test + void testPartitionBucketCountRemovedOnDelete() { + // Seed both partitions with explicit per-partition bucket counts. + int explicitBucketCount = 8; + serverMetadataCache.updateClusterMetadata( + new ClusterMetadata( + coordinatorServer, + aliveTableServers, + tableMetadataList, + Arrays.asList( + new PartitionMetadata( + partitionTableId, + partitionName1, + partitionId1, + initialBucketMetadata, + explicitBucketCount), + new PartitionMetadata( + partitionTableId, + partitionName2, + partitionId2, + initialBucketMetadata, + explicitBucketCount)))); + // The explicit count must be exposed via the updateClusterMetadata path (not the + // merged bucket-metadata-list size); the delete/re-add asserts below build on this. + assertThat( + serverMetadataCache + .getPartitionMetadata( + PhysicalTablePath.of(partitionedTablePath, partitionName1)) + .get() + .getBucketCount()) + .isEqualTo(explicitBucketCount); + + // Delete partition1 via DELETED_PARTITION_ID (partitionId marks deletion). + serverMetadataCache.updateClusterMetadata( + new ClusterMetadata( + coordinatorServer, + aliveTableServers, + Collections.emptyList(), + Collections.singletonList( + new PartitionMetadata( + partitionTableId, + partitionName1, + DELETED_PARTITION_ID, + Collections.emptyList())))); + assertThat( + serverMetadataCache.getPartitionMetadata( + PhysicalTablePath.of(partitionedTablePath, partitionName1))) + .isEmpty(); + + // Re-add partition1 WITHOUT an explicit bucketCount. The cache must NOT return the + // stale 8; the DELETED_PARTITION_ID path must have removed the prior entry from the + // partitionBucketCounts map so the fallback (bucketMetadataList.size()) applies. + serverMetadataCache.updateClusterMetadata( + new ClusterMetadata( + coordinatorServer, + aliveTableServers, + Collections.emptyList(), + Collections.singletonList( + new PartitionMetadata( + partitionTableId, + partitionName1, + partitionId1, + initialBucketMetadata)))); + assertThat( + serverMetadataCache + .getPartitionMetadata( + PhysicalTablePath.of(partitionedTablePath, partitionName1)) + .get() + .getBucketCount()) + .isEqualTo(initialBucketMetadata.size()); + + // Delete partition2 via DELETED_PARTITION_NAME (partitionName marks deletion). + serverMetadataCache.updateClusterMetadata( + new ClusterMetadata( + coordinatorServer, + aliveTableServers, + Collections.emptyList(), + Collections.singletonList( + new PartitionMetadata( + partitionTableId, + DELETED_PARTITION_NAME, + partitionId2, + Collections.emptyList())))); + assertThat( + serverMetadataCache.getPartitionMetadata( + PhysicalTablePath.of(partitionedTablePath, partitionName2))) + .isEmpty(); + + // Re-add partition2 WITHOUT explicit; the DELETED_PARTITION_NAME path must have also + // cleared partitionBucketCounts, so fallback (list size) applies rather than stale 8. + serverMetadataCache.updateClusterMetadata( + new ClusterMetadata( + coordinatorServer, + aliveTableServers, + Collections.emptyList(), + Collections.singletonList( + new PartitionMetadata( + partitionTableId, + partitionName2, + partitionId2, + initialBucketMetadata)))); + assertThat( + serverMetadataCache + .getPartitionMetadata( + PhysicalTablePath.of(partitionedTablePath, partitionName2)) + .get() + .getBucketCount()) + .isEqualTo(initialBucketMetadata.size()); + } + + @Test + void testUpdatePartitionMetadataPropagatesExplicitBucketCount() { + // Seed table metadata: updatePartitionMetadata bails out if the tableId is unknown. + serverMetadataCache.updateClusterMetadata( + new ClusterMetadata( + coordinatorServer, + aliveTableServers, + tableMetadataList, + Collections.emptyList())); + + // Route via the single-partition updatePartitionMetadata path (distinct from + // updateClusterMetadata). The explicit bucketCount must be applied to the cache. + int explicitBucketCount = 8; + serverMetadataCache.updatePartitionMetadata( + new PartitionMetadata( + partitionTableId, + partitionName1, + partitionId1, + initialBucketMetadata, + explicitBucketCount)); + assertThat( + serverMetadataCache + .getPartitionMetadata( + PhysicalTablePath.of(partitionedTablePath, partitionName1)) + .get() + .getBucketCount()) + .isEqualTo(explicitBucketCount); + } + private void assertPartitionMetadataEquals( long partitionId, long expectedTableId, diff --git a/fluss-server/src/test/java/org/apache/fluss/server/metadata/ZkBasedMetadataProviderTest.java b/fluss-server/src/test/java/org/apache/fluss/server/metadata/ZkBasedMetadataProviderTest.java index 5e74623c740..6200cf886d7 100644 --- a/fluss-server/src/test/java/org/apache/fluss/server/metadata/ZkBasedMetadataProviderTest.java +++ b/fluss-server/src/test/java/org/apache/fluss/server/metadata/ZkBasedMetadataProviderTest.java @@ -187,7 +187,8 @@ void testGetPartitionMetadataFromZk() throws Exception { partitionAssignment, DEFAULT_REMOTE_DATA_DIR, tablePath, - tableId); + tableId, + partitionAssignment.getBucketAssignments().size()); // Create leader and isr for partition buckets TableBucket partitionBucket0 = new TableBucket(tableId, partitionId, 0); @@ -272,14 +273,16 @@ void testBatchGetPartitionMetadataFromZkAsync() throws Exception { partitionAssignment1, DEFAULT_REMOTE_DATA_DIR, tablePath1, - tableId1); + tableId1, + partitionAssignment1.getBucketAssignments().size()); zookeeperClient.registerPartitionAssignmentAndMetadata( partitionId2, partitionName2, partitionAssignment2, DEFAULT_REMOTE_DATA_DIR, tablePath1, - tableId1); + tableId1, + partitionAssignment2.getBucketAssignments().size()); // Create partition for table2 long partitionId3 = 21L; @@ -295,7 +298,8 @@ void testBatchGetPartitionMetadataFromZkAsync() throws Exception { partitionAssignment3, DEFAULT_REMOTE_DATA_DIR, tablePath2, - tableId2); + tableId2, + partitionAssignment3.getBucketAssignments().size()); // Create leader and isr for all partition buckets TableBucket bucket1 = new TableBucket(tableId1, partitionId1, 0); diff --git a/fluss-server/src/test/java/org/apache/fluss/server/replica/AdjustIsrTest.java b/fluss-server/src/test/java/org/apache/fluss/server/replica/AdjustIsrTest.java index 77e425bbcc8..b481b2f855e 100644 --- a/fluss-server/src/test/java/org/apache/fluss/server/replica/AdjustIsrTest.java +++ b/fluss-server/src/test/java/org/apache/fluss/server/replica/AdjustIsrTest.java @@ -140,7 +140,9 @@ void testShrinkIsrUsesLatestStateAfterLeaderChange() throws Exception { DATA1_PHYSICAL_TABLE_PATH, tb, Arrays.asList(1, 2), - new LeaderAndIsr(1, 0, Arrays.asList(1, 2), Collections.emptyList(), 0, 0)); + new LeaderAndIsr(1, 0, Arrays.asList(1, 2), Collections.emptyList(), 0, 0), + 3, + 0L); ReentrantReadWriteLock leaderIsrUpdateLock = replica.getLeaderIsrUpdateLock(); @@ -349,6 +351,8 @@ private void notifyLeaderAndIsr( replicas, Collections.emptyList(), replica.getCoordinatorEpoch(), - bucketEpoch)))); + bucketEpoch), + 3, + 0L))); } } diff --git a/fluss-server/src/test/java/org/apache/fluss/server/replica/ReplicaManagerTest.java b/fluss-server/src/test/java/org/apache/fluss/server/replica/ReplicaManagerTest.java index 3364aa61b08..5d95f19c9a5 100644 --- a/fluss-server/src/test/java/org/apache/fluss/server/replica/ReplicaManagerTest.java +++ b/fluss-server/src/test/java/org/apache/fluss/server/replica/ReplicaManagerTest.java @@ -617,7 +617,9 @@ void testNewKvLeaderRejectedWhenDiskLocked() throws Exception { Collections.singletonList(TABLET_SERVER_ID), Collections.emptyList(), INITIAL_COORDINATOR_EPOCH, - INITIAL_BUCKET_EPOCH))), + INITIAL_BUCKET_EPOCH), + 3, + 0L)), future::complete); List results = future.get(); @@ -643,7 +645,9 @@ void testNewKvLeaderRejectedWhenDiskLocked() throws Exception { Collections.singletonList(TABLET_SERVER_ID), Collections.emptyList(), INITIAL_COORDINATOR_EPOCH, - INITIAL_BUCKET_EPOCH))), + INITIAL_BUCKET_EPOCH), + 3, + 0L)), future::complete); assertThat(future.get()).containsOnly(new NotifyLeaderAndIsrResultForBucket(logTb)); @@ -1787,7 +1791,9 @@ void becomeLeaderOrFollower() throws Exception { Arrays.asList(1, 2, 3), Collections.emptyList(), INITIAL_COORDINATOR_EPOCH, - INITIAL_BUCKET_EPOCH))), + INITIAL_BUCKET_EPOCH), + 3, + 0L)), future::complete); assertThat(future.get()).containsOnly(new NotifyLeaderAndIsrResultForBucket(tb)); assertReplicaEpochEquals( @@ -1808,7 +1814,9 @@ void becomeLeaderOrFollower() throws Exception { Arrays.asList(1, 2, 3), Collections.emptyList(), INITIAL_COORDINATOR_EPOCH, - INITIAL_BUCKET_EPOCH))), + INITIAL_BUCKET_EPOCH), + 3, + 0L)), future::complete); assertThat(future.get()) .containsOnly( @@ -1857,7 +1865,9 @@ void testLakeSnapshotReadFailureDoesNotFailLeaderTransition() throws Exception { Collections.singletonList(TABLET_SERVER_ID), Collections.emptyList(), INITIAL_COORDINATOR_EPOCH, - INITIAL_BUCKET_EPOCH))), + INITIAL_BUCKET_EPOCH), + 3, + 0L)), future::complete); assertThat(future.get()).containsOnly(new NotifyLeaderAndIsrResultForBucket(tableBucket)); @@ -1884,7 +1894,9 @@ void testStopReplica() throws Exception { Arrays.asList(1, 2, 3), Collections.emptyList(), INITIAL_COORDINATOR_EPOCH, - INITIAL_BUCKET_EPOCH))), + INITIAL_BUCKET_EPOCH), + 3, + 0L)), future::complete); assertThat(future.get()).containsOnly(new NotifyLeaderAndIsrResultForBucket(tb)); assertReplicaEpochEquals( @@ -1916,7 +1928,9 @@ void testStopReplica() throws Exception { Arrays.asList(1, 2, 3), Collections.emptyList(), INITIAL_COORDINATOR_EPOCH, - INITIAL_BUCKET_EPOCH))), + INITIAL_BUCKET_EPOCH), + 3, + 0L)), future::complete); assertThat(future.get()).containsOnly(new NotifyLeaderAndIsrResultForBucket(tb)); assertReplicaEpochEquals( diff --git a/fluss-server/src/test/java/org/apache/fluss/server/replica/ReplicaRoutingStateTest.java b/fluss-server/src/test/java/org/apache/fluss/server/replica/ReplicaRoutingStateTest.java new file mode 100644 index 00000000000..9871cc32de6 --- /dev/null +++ b/fluss-server/src/test/java/org/apache/fluss/server/replica/ReplicaRoutingStateTest.java @@ -0,0 +1,148 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.fluss.server.replica; + +import org.apache.fluss.exception.InvalidBucketRoutingException; +import org.apache.fluss.exception.RetriableException; +import org.apache.fluss.exception.UnknownTableOrBucketException; +import org.apache.fluss.metadata.PhysicalTablePath; +import org.apache.fluss.metadata.TableBucket; +import org.apache.fluss.metadata.TableInfo; +import org.apache.fluss.server.entity.NotifyLeaderAndIsrData; +import org.apache.fluss.server.entity.NotifyLeaderAndIsrResultForBucket; +import org.apache.fluss.server.metadata.ClusterMetadata; +import org.apache.fluss.server.metadata.TableMetadata; +import org.apache.fluss.server.zk.data.LeaderAndIsr; + +import org.junit.jupiter.api.Test; + +import java.util.Collections; +import java.util.List; +import java.util.concurrent.CompletableFuture; + +import static org.apache.fluss.record.TestData.DATA1_TABLE_DESCRIPTOR; +import static org.apache.fluss.record.TestData.DATA1_TABLE_ID; +import static org.apache.fluss.record.TestData.DATA1_TABLE_PATH; +import static org.apache.fluss.record.TestData.DEFAULT_REMOTE_DATA_DIR; +import static org.apache.fluss.server.coordinator.CoordinatorContext.INITIAL_COORDINATOR_EPOCH; +import static org.apache.fluss.server.zk.data.LeaderAndIsr.INITIAL_BUCKET_EPOCH; +import static org.apache.fluss.server.zk.data.LeaderAndIsr.INITIAL_LEADER_EPOCH; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** + * Test for the routing state a {@link Replica} is armed with on leader activation and the request + * validation driven by it. + */ +final class ReplicaRoutingStateTest extends ReplicaTestBase { + + /** + * A bucket no other test class in this package uses. {@code TestingMetricGroups} caches bucket + * metric groups per (table, bucket) in a static registry, so sharing a bucket coordinate hands + * stale gauges across test classes. + */ + private static final int TEST_BUCKET = 5; + + @Test + void testLeaderActivationAllowsMissingRoutingState() throws Exception { + TableBucket tb = new TableBucket(DATA1_TABLE_ID, TEST_BUCKET); + + // A legacy coordinator's notification carries no routing fields, but leader activation must + // still succeed because server upgrade order is not guaranteed. + CompletableFuture> legacyFuture = + new CompletableFuture<>(); + replicaManager.becomeLeaderOrFollower( + INITIAL_COORDINATOR_EPOCH, + Collections.singletonList(notifyDataWithRoutingState(tb, null, null)), + legacyFuture::complete); + assertThat(legacyFuture.get()).containsOnly(new NotifyLeaderAndIsrResultForBucket(tb)); + assertThat(replicaManager.getReplicaOrException(tb).isLeader()).isTrue(); + assertThat(replicaManager.getReplicaOrException(tb).getRoutingBucketCount()).isNull(); + assertThat(replicaManager.getReplicaOrException(tb).getBucketCountEpoch()).isNull(); + + // A later notification with routing state updates the already active leader. + makeLeaderWithRoutingState(tb, 3, 0L); + assertThat(replicaManager.getReplicaOrException(tb).isLeader()).isTrue(); + replicaManager.validateRoutingBucketCount(tb, 3); + assertThatThrownBy(() -> replicaManager.validateRoutingBucketCount(tb, 4)) + .isInstanceOf(InvalidBucketRoutingException.class) + .isNotInstanceOf(RetriableException.class); + + // A legacy client (no count) passes on a non-rescaled table... + replicaManager.validateRoutingBucketCount(tb, 0); + // ...but is rejected once an ALTER advances the metadata cache to epoch 1 through + // UpdateMetadata, which does not re-notify the already active replica. + replicaManager.maybeUpdateMetadataCache( + INITIAL_COORDINATOR_EPOCH, + new ClusterMetadata( + null, + Collections.emptySet(), + Collections.singletonList( + new TableMetadata( + TableInfo.of( + DATA1_TABLE_PATH, + DATA1_TABLE_ID, + 1, + DATA1_TABLE_DESCRIPTOR, + DEFAULT_REMOTE_DATA_DIR, + 1L, + 1L, + 1L), + Collections.emptyList())), + Collections.emptyList())); + assertThat(replicaManager.getReplicaOrException(tb).getBucketCountEpoch()).isEqualTo(0L); + assertThatThrownBy(() -> replicaManager.validateRoutingBucketCount(tb, 0)) + .isInstanceOf(InvalidBucketRoutingException.class); + + // Unknown or out-of-range buckets are rejected during routing validation instead of being + // silently deferred to the downstream replica lookup. + TableBucket unknownBucket = new TableBucket(DATA1_TABLE_ID, 99); + assertThatThrownBy(() -> replicaManager.validateRoutingBucketCount(unknownBucket, 3)) + .isInstanceOf(UnknownTableOrBucketException.class) + .hasMessageContaining(unknownBucket.toString()); + } + + private void makeLeaderWithRoutingState( + TableBucket tb, Integer bucketCount, Long bucketCountEpoch) throws Exception { + CompletableFuture> future = + new CompletableFuture<>(); + replicaManager.becomeLeaderOrFollower( + INITIAL_COORDINATOR_EPOCH, + Collections.singletonList( + notifyDataWithRoutingState(tb, bucketCount, bucketCountEpoch)), + future::complete); + assertThat(future.get()).containsOnly(new NotifyLeaderAndIsrResultForBucket(tb)); + } + + private static NotifyLeaderAndIsrData notifyDataWithRoutingState( + TableBucket tb, Integer bucketCount, Long bucketCountEpoch) { + return new NotifyLeaderAndIsrData( + PhysicalTablePath.of(DATA1_TABLE_PATH), + tb, + Collections.singletonList(TABLET_SERVER_ID), + new LeaderAndIsr( + TABLET_SERVER_ID, + INITIAL_LEADER_EPOCH, + Collections.singletonList(TABLET_SERVER_ID), + Collections.emptyList(), + INITIAL_COORDINATOR_EPOCH, + INITIAL_BUCKET_EPOCH), + bucketCount, + bucketCountEpoch); + } +} diff --git a/fluss-server/src/test/java/org/apache/fluss/server/replica/ReplicaTest.java b/fluss-server/src/test/java/org/apache/fluss/server/replica/ReplicaTest.java index b1e98b9d30f..b9e8cb7ce4b 100644 --- a/fluss-server/src/test/java/org/apache/fluss/server/replica/ReplicaTest.java +++ b/fluss-server/src/test/java/org/apache/fluss/server/replica/ReplicaTest.java @@ -19,6 +19,7 @@ import org.apache.fluss.config.ConfigOptions; import org.apache.fluss.config.Configuration; +import org.apache.fluss.exception.NotLeaderOrFollowerException; import org.apache.fluss.exception.OutOfOrderSequenceException; import org.apache.fluss.fs.FsPath; import org.apache.fluss.metadata.LogFormat; @@ -48,6 +49,7 @@ import org.apache.fluss.server.kv.snapshot.KvSnapshotDownloadSpec; import org.apache.fluss.server.kv.snapshot.TestingCompletedKvSnapshotCommitter; import org.apache.fluss.server.log.FetchParams; +import org.apache.fluss.server.log.ListOffsetsParam; import org.apache.fluss.server.log.LogAppendInfo; import org.apache.fluss.server.log.LogReadInfo; import org.apache.fluss.server.testutils.KvTestUtils; @@ -149,6 +151,21 @@ void testMakeLeader() throws Exception { assertThat(kvReplica.getKvTablet()).isNotNull(); } + @Test + void testGetOffsetRequiresLeader() throws Exception { + Replica replica = + makeLogReplica(DATA1_PHYSICAL_TABLE_PATH, new TableBucket(DATA1_TABLE_ID, 1)); + + assertThat(replica.isLeader()).isFalse(); + assertThatThrownBy( + () -> + replica.getOffset( + remoteLogManager, + new ListOffsetsParam( + -1, ListOffsetsParam.LATEST_OFFSET_TYPE, null))) + .isInstanceOf(NotLeaderOrFollowerException.class); + } + @Test void testAppendRecordsToLeader() throws Exception { Replica logReplica = @@ -237,7 +254,9 @@ void testBucketPhysicalStorageLocalLogSizeIncludesFollower() throws Exception { replicas, Collections.emptyList(), INITIAL_COORDINATOR_EPOCH, - followerLeaderEpoch))); + followerLeaderEpoch), + 3, + 0L)); assertThat(logReplica.isLeader()).isFalse(); assertThat(localLogSizeGauge.getValue()).isEqualTo(localLogSize); @@ -1147,7 +1166,9 @@ private void makeKvReplicaAsFollower(Replica replica, int leaderEpoch) { Collections.emptyList(), INITIAL_COORDINATOR_EPOCH, // we also use the leader epoch as bucket epoch - leaderEpoch))); + leaderEpoch), + 3, + 0L)); } private void makeLeaderReplica( @@ -1165,7 +1186,9 @@ private void makeLeaderReplica( Collections.emptyList(), INITIAL_COORDINATOR_EPOCH, // we also use the leader epoch as bucket epoch - leaderEpoch))); + leaderEpoch), + 3, + 0L)); } private static LogRecords fetchRecords(Replica replica) throws IOException { diff --git a/fluss-server/src/test/java/org/apache/fluss/server/replica/ReplicaTestBase.java b/fluss-server/src/test/java/org/apache/fluss/server/replica/ReplicaTestBase.java index 98d714edc1e..fbc9b25c512 100644 --- a/fluss-server/src/test/java/org/apache/fluss/server/replica/ReplicaTestBase.java +++ b/fluss-server/src/test/java/org/apache/fluss/server/replica/ReplicaTestBase.java @@ -457,7 +457,10 @@ protected void makeLogTableAsLeader( isr, Collections.emptyList(), INITIAL_COORDINATOR_EPOCH, - INITIAL_BUCKET_EPOCH)))); + INITIAL_BUCKET_EPOCH), + // all test tables are created distributedBy(3) + TEST_ROUTING_BUCKET_COUNT, + 0L))); } // TODO this is only for single tablet server unit test. @@ -500,9 +503,15 @@ protected void makeKvTableAsLeader( Collections.emptyList(), INITIAL_COORDINATOR_EPOCH, // use leader epoch as bucket epoch - leaderEpoch)))); + leaderEpoch), + // all test tables are created distributedBy(3) + TEST_ROUTING_BUCKET_COUNT, + 0L))); } + /** The routing bucket count carried by test notifications; test tables are distributedBy(3). */ + protected static final Integer TEST_ROUTING_BUCKET_COUNT = 3; + protected void makeLeaderAndFollower(List notifyLeaderAndIsrDataList) { replicaManager.becomeLeaderOrFollower(0, notifyLeaderAndIsrDataList, result -> {}); } diff --git a/fluss-server/src/test/java/org/apache/fluss/server/replica/fetcher/ReplicaFetcherManagerTest.java b/fluss-server/src/test/java/org/apache/fluss/server/replica/fetcher/ReplicaFetcherManagerTest.java index d61b6319c10..5c059dea1d2 100644 --- a/fluss-server/src/test/java/org/apache/fluss/server/replica/fetcher/ReplicaFetcherManagerTest.java +++ b/fluss-server/src/test/java/org/apache/fluss/server/replica/fetcher/ReplicaFetcherManagerTest.java @@ -96,7 +96,9 @@ void testAddAndRemoveBucket() { Arrays.asList(leader.id(), TABLET_SERVER_ID), Collections.emptyList(), INITIAL_COORDINATOR_EPOCH, - LeaderAndIsr.INITIAL_BUCKET_EPOCH))), + LeaderAndIsr.INITIAL_BUCKET_EPOCH), + 3, + 0L)), result -> {}); InitialFetchStatus initialFetchStatus = @@ -143,7 +145,9 @@ void testDoesNotAddFetcherWhenFollowerHasNoLeader() { Collections.emptyList(), Collections.emptyList(), INITIAL_COORDINATOR_EPOCH, - LeaderAndIsr.INITIAL_BUCKET_EPOCH))), + LeaderAndIsr.INITIAL_BUCKET_EPOCH), + 3, + 0L)), result::set); assertThat(result.get()).hasSize(1); diff --git a/fluss-server/src/test/java/org/apache/fluss/server/replica/fetcher/ReplicaFetcherThreadTest.java b/fluss-server/src/test/java/org/apache/fluss/server/replica/fetcher/ReplicaFetcherThreadTest.java index 8ff2e20a096..8262d080f1f 100644 --- a/fluss-server/src/test/java/org/apache/fluss/server/replica/fetcher/ReplicaFetcherThreadTest.java +++ b/fluss-server/src/test/java/org/apache/fluss/server/replica/fetcher/ReplicaFetcherThreadTest.java @@ -570,7 +570,9 @@ private void makeLeaderAndFollower( Arrays.asList(leaderServerId, followerServerId), Collections.emptyList(), INITIAL_COORDINATOR_EPOCH, - INITIAL_BUCKET_EPOCH))), + INITIAL_BUCKET_EPOCH), + 3, + 0L)), result -> {}); followerRM.becomeLeaderOrFollower( INITIAL_COORDINATOR_EPOCH, @@ -585,7 +587,9 @@ private void makeLeaderAndFollower( Arrays.asList(leaderServerId, followerServerId), Collections.emptyList(), INITIAL_COORDINATOR_EPOCH, - INITIAL_BUCKET_EPOCH))), + INITIAL_BUCKET_EPOCH), + 3, + 0L)), result -> {}); } diff --git a/fluss-server/src/test/java/org/apache/fluss/server/replica/historical/HistoricalPartitionManagerTest.java b/fluss-server/src/test/java/org/apache/fluss/server/replica/historical/HistoricalPartitionManagerTest.java index b18f96b3a7a..ff7d174358a 100644 --- a/fluss-server/src/test/java/org/apache/fluss/server/replica/historical/HistoricalPartitionManagerTest.java +++ b/fluss-server/src/test/java/org/apache/fluss/server/replica/historical/HistoricalPartitionManagerTest.java @@ -889,7 +889,9 @@ private TableInfo registerHistoricalTableAndBecomeLeader( Collections.singletonList(TABLET_SERVER_ID), Collections.emptyList(), INITIAL_COORDINATOR_EPOCH, - INITIAL_BUCKET_EPOCH))), + INITIAL_BUCKET_EPOCH), + 3, + 0L)), leaderFuture::complete); assertThat(leaderFuture.get(10, TimeUnit.SECONDS)) .containsOnly(new NotifyLeaderAndIsrResultForBucket(TABLE_BUCKET)); @@ -909,7 +911,9 @@ private static NotifyLeaderAndIsrData followerState() { replicas, Collections.emptyList(), INITIAL_COORDINATOR_EPOCH, - INITIAL_BUCKET_EPOCH + 1)); + INITIAL_BUCKET_EPOCH + 1), + 3, + 0L); } private static NotifyLeaderAndIsrData leaderStateAfterFollower() { @@ -924,7 +928,9 @@ private static NotifyLeaderAndIsrData leaderStateAfterFollower() { replicas, Collections.emptyList(), INITIAL_COORDINATOR_EPOCH, - INITIAL_BUCKET_EPOCH + 2)); + INITIAL_BUCKET_EPOCH + 2), + 3, + 0L); } private static void await(CountDownLatch latch) { diff --git a/fluss-server/src/test/java/org/apache/fluss/server/tablet/TabletServiceITCase.java b/fluss-server/src/test/java/org/apache/fluss/server/tablet/TabletServiceITCase.java index 84708d0b6c1..5ca0cd51225 100644 --- a/fluss-server/src/test/java/org/apache/fluss/server/tablet/TabletServiceITCase.java +++ b/fluss-server/src/test/java/org/apache/fluss/server/tablet/TabletServiceITCase.java @@ -19,7 +19,9 @@ import org.apache.fluss.config.ConfigOptions; import org.apache.fluss.exception.FlussRuntimeException; +import org.apache.fluss.exception.InvalidBucketRoutingException; import org.apache.fluss.exception.InvalidRequiredAcksException; +import org.apache.fluss.exception.UnknownTableOrBucketException; import org.apache.fluss.metadata.KvFormat; import org.apache.fluss.metadata.LogFormat; import org.apache.fluss.metadata.PhysicalTablePath; @@ -37,8 +39,11 @@ import org.apache.fluss.row.encode.ValueEncoder; import org.apache.fluss.rpc.gateway.TabletServerGateway; import org.apache.fluss.rpc.messages.FetchLogResponse; +import org.apache.fluss.rpc.messages.GetTableStatsRequest; +import org.apache.fluss.rpc.messages.GetTableStatsResponse; import org.apache.fluss.rpc.messages.InitWriterRequest; import org.apache.fluss.rpc.messages.InitWriterResponse; +import org.apache.fluss.rpc.messages.ListOffsetsRequest; import org.apache.fluss.rpc.messages.ListOffsetsResponse; import org.apache.fluss.rpc.messages.NotifyLeaderAndIsrRequest; import org.apache.fluss.rpc.messages.NotifyLeaderAndIsrResponse; @@ -49,6 +54,7 @@ import org.apache.fluss.rpc.messages.PbNotifyLeaderAndIsrReqForBucket; import org.apache.fluss.rpc.messages.PbPrefixLookupRespForBucket; import org.apache.fluss.rpc.messages.PbPutKvRespForBucket; +import org.apache.fluss.rpc.messages.PbTableStatsRespForBucket; import org.apache.fluss.rpc.messages.ProduceLogResponse; import org.apache.fluss.rpc.messages.PutKvResponse; import org.apache.fluss.rpc.messages.ScanKvRequest; @@ -766,6 +772,138 @@ void testLimitScanLogTable() throws Exception { expected2); } + @Test + void testRoutingBucketCountValidationAppliesToClientRequestsOnly() throws Exception { + // Routing validation only applies to hash-distributed tables: a keyless table may place a + // record in any bucket, so a stale count is harmless there (see + // ReplicaManager#validateRoutingBucketCount). + long tableId = + createTable( + FLUSS_CLUSTER_EXTENSION, DATA1_TABLE_PATH_PK, DATA1_TABLE_DESCRIPTOR_PK); + TableBucket tb = new TableBucket(tableId, 0); + + FLUSS_CLUSTER_EXTENSION.waitUntilAllReplicaReady(tb); + + int leader = FLUSS_CLUSTER_EXTENSION.waitAndGetLeader(tb); + TabletServerGateway leaderGateWay = + FLUSS_CLUSTER_EXTENSION.newTabletServerClientForNode(leader); + + // a client whose bucket count doesn't match the actual one computed its bucketId from a + // count the server has confirmed to be stale. + assertThatThrownBy( + () -> + leaderGateWay + .listOffsets( + newListOffsetsRequestWithRoutingBucketCount( + -1, + ListOffsetsParam.LATEST_OFFSET_TYPE, + tableId, + 0, + 999)) + .get()) + .cause() + .isInstanceOf(InvalidBucketRoutingException.class); + + // the very same count coming from a follower is not validated: a follower's bucket ids come + // from NotifyLeaderAndIsr, so replication must not depend on the leader's metadata cache. + assertListOffsetsResponse( + leaderGateWay + .listOffsets( + newListOffsetsRequestWithRoutingBucketCount( + 1, ListOffsetsParam.LATEST_OFFSET_TYPE, tableId, 0, 999)) + .get(), + 0L, + Errors.NONE.code(), + null); + + // Request-scoped validation resolves the target immediately, so an unknown table/bucket + // fails the whole RPC with the standard replica lookup exception. + assertThatThrownBy( + () -> + leaderGateWay + .listOffsets( + newListOffsetsRequestWithRoutingBucketCount( + -1, + ListOffsetsParam.LATEST_OFFSET_TYPE, + 10005L, + 0, + 3)) + .get()) + .cause() + .isInstanceOf(UnknownTableOrBucketException.class) + .hasMessageContaining("Unknown table or bucket"); + } + + private static ListOffsetsRequest newListOffsetsRequestWithRoutingBucketCount( + int followerServerId, + int offsetType, + long tableId, + int bucketId, + int routingBucketCount) { + return newListOffsetsRequest(followerServerId, offsetType, tableId, bucketId) + .setRoutingBucketCount(routingBucketCount); + } + + @Test + void testInvalidRoutingBucketCountOnlyFailsTheOffendingBucket() throws Exception { + // 9 buckets over 3 tablet servers, so at least one server necessarily leads two of them + // and a single request can carry two buckets hosted by the same leader. + int bucketCount = 9; + TablePath tablePath = TablePath.of("test_db_1", "test_stale_routing_per_bucket"); + long tableId = + createTable( + FLUSS_CLUSTER_EXTENSION, + tablePath, + TableDescriptor.builder() + .schema(DATA1_SCHEMA) + // hash-distributed: routing validation is skipped for keyless + // tables + .distributedBy(bucketCount, "a") + .build()); + + Map> bucketsByLeader = new HashMap<>(); + for (int bucketId = 0; bucketId < bucketCount; bucketId++) { + TableBucket tb = new TableBucket(tableId, bucketId); + FLUSS_CLUSTER_EXTENSION.waitUntilAllReplicaReady(tb); + bucketsByLeader + .computeIfAbsent( + FLUSS_CLUSTER_EXTENSION.waitAndGetLeader(tb), k -> new ArrayList<>()) + .add(bucketId); + } + Map.Entry> coLocated = + bucketsByLeader.entrySet().stream() + .filter(entry -> entry.getValue().size() >= 2) + .findFirst() + .orElseThrow( + () -> + new AssertionError( + "9 buckets over 3 servers must co-locate two leaders")); + int healthyBucket = coLocated.getValue().get(0); + int staleBucket = coLocated.getValue().get(1); + + GetTableStatsRequest request = new GetTableStatsRequest().setTableId(tableId); + request.addBucketsReq().setBucketId(healthyBucket).setRoutingBucketCount(bucketCount); + request.addBucketsReq().setBucketId(staleBucket).setRoutingBucketCount(bucketCount + 1); + + GetTableStatsResponse response = + FLUSS_CLUSTER_EXTENSION + .newTabletServerClientForNode(coLocated.getKey()) + .getTableStats(request) + .get(); + + assertThat(response.getBucketsRespsCount()).isEqualTo(2); + Map respByBucket = new HashMap<>(); + for (PbTableStatsRespForBucket bucketResp : response.getBucketsRespsList()) { + respByBucket.put(bucketResp.getBucketId(), bucketResp); + } + + // the co-batched bucket whose routing is still valid is served as usual + assertThat(respByBucket.get(healthyBucket).hasErrorCode()).isFalse(); + // Only the bucket routed by an invalid count is rejected without retrying the fixed route. + assertThat(respByBucket.get(staleBucket).getErrorCode()) + .isEqualTo(Errors.INVALID_BUCKET_ROUTING.code()); + } + @Test void testListOffsets() throws Exception { long tableId = @@ -988,7 +1126,12 @@ private NotifyLeaderAndIsrRequest makeNotifyLeaderAndIsrRequest( PbNotifyLeaderAndIsrReqForBucket reqForBucket = makeNotifyBucketLeaderAndIsr( new NotifyLeaderAndIsrData( - physicalTablePath, tableBucket, leaderAndIsr.isr(), leaderAndIsr)); + physicalTablePath, + tableBucket, + leaderAndIsr.isr(), + leaderAndIsr, + 3, + 0L)); return ServerRpcMessageUtils.makeNotifyLeaderAndIsrRequest( 0, Collections.singletonList(reqForBucket)); } diff --git a/fluss-server/src/test/java/org/apache/fluss/server/testutils/FlussClusterExtension.java b/fluss-server/src/test/java/org/apache/fluss/server/testutils/FlussClusterExtension.java index 0cc615482ab..cff3f788c99 100644 --- a/fluss-server/src/test/java/org/apache/fluss/server/testutils/FlussClusterExtension.java +++ b/fluss-server/src/test/java/org/apache/fluss/server/testutils/FlussClusterExtension.java @@ -738,7 +738,11 @@ public void triggerAndWaitSnapshot(TablePath tablePath) throws Exception { Map partitions = zooKeeperClient.getPartitionRegistrations(tablePath); for (PartitionRegistration partition : partitions.values()) { - for (int bucketId = 0; bucketId < bucketCount; bucketId++) { + // partitions diverge from the table-level count after ALTER bucket.num + int partitionBucketCount = + partition.getBucketCountOrDefault( + bucketCount, tableRegistration.bucketCountEpoch); + for (int bucketId = 0; bucketId < partitionBucketCount; bucketId++) { tableBuckets.add( new TableBucket(tableId, partition.getPartitionId(), bucketId)); } @@ -866,7 +870,9 @@ public void notifyLeaderAndIsr( PhysicalTablePath.of(tablePath), tableBucket, replicas, - leaderAndIsr)); + leaderAndIsr, + 3, + 0L)); NotifyLeaderAndIsrRequest notifyLeaderAndIsrRequest = ServerRpcMessageUtils.makeNotifyLeaderAndIsrRequest( 0, Collections.singletonList(reqForBucket)); diff --git a/fluss-server/src/test/java/org/apache/fluss/server/testutils/PartitionMetadataAssert.java b/fluss-server/src/test/java/org/apache/fluss/server/testutils/PartitionMetadataAssert.java index 349fe9a6fd6..89aabbb1ac6 100644 --- a/fluss-server/src/test/java/org/apache/fluss/server/testutils/PartitionMetadataAssert.java +++ b/fluss-server/src/test/java/org/apache/fluss/server/testutils/PartitionMetadataAssert.java @@ -41,6 +41,12 @@ private PartitionMetadataAssert(PartitionMetadata actual) { public PartitionMetadataAssert isEqualTo(PartitionMetadata expected) { assertThat(expected.getPartitionName()).isEqualTo(actual.getPartitionName()); + // actual bucketCount is always non-null (falls back to bucketMetadataList size), so + // only compare when expected sets it — otherwise legacy callers passing null fail + // spuriously. + if (expected.getBucketCount() != null) { + assertThat(actual.getBucketCount()).isEqualTo(expected.getBucketCount()); + } List bucketMetadataList = expected.getBucketMetadataList(); List actualBucketMetadataList = actual.getBucketMetadataList(); assertThat(bucketMetadataList) diff --git a/fluss-server/src/test/java/org/apache/fluss/server/utils/ServerRpcMessageUtilsTest.java b/fluss-server/src/test/java/org/apache/fluss/server/utils/ServerRpcMessageUtilsTest.java index 5459d66a2e0..8ae2c1494a2 100644 --- a/fluss-server/src/test/java/org/apache/fluss/server/utils/ServerRpcMessageUtilsTest.java +++ b/fluss-server/src/test/java/org/apache/fluss/server/utils/ServerRpcMessageUtilsTest.java @@ -18,6 +18,7 @@ package org.apache.fluss.server.utils; import org.apache.fluss.metadata.TableBucket; +import org.apache.fluss.metadata.TableChange; import org.apache.fluss.metadata.TablePath; import org.apache.fluss.record.KvRecordBatch; import org.apache.fluss.record.MemoryLogRecords; @@ -26,6 +27,7 @@ import org.apache.fluss.rpc.entity.LookupResultForBucket; import org.apache.fluss.rpc.entity.ProduceLogResultForBucket; import org.apache.fluss.rpc.entity.PutKvResultForBucket; +import org.apache.fluss.rpc.messages.AlterTableRequest; import org.apache.fluss.rpc.messages.FetchLogResponse; import org.apache.fluss.rpc.messages.LookupResponse; import org.apache.fluss.rpc.messages.PbBucketMetadata; @@ -68,6 +70,16 @@ /** Tests for {@link ServerRpcMessageUtils}. */ class ServerRpcMessageUtilsTest { + @Test + void testAlterTableDistributionChanges() { + AlterTableRequest request = new AlterTableRequest(); + assertThat(ServerRpcMessageUtils.toAlterTableDistributionChanges(request)).isEmpty(); + + request.setModifyBucketCount().setNewBucketCount(8); + assertThat(ServerRpcMessageUtils.toAlterTableDistributionChanges(request)) + .containsExactly(TableChange.modifyBucketCount(8)); + } + @Test void testFetchLogResponseContainsMinRetainOffset() { TableBucket tableBucket = new TableBucket(1L, 0); diff --git a/fluss-server/src/test/java/org/apache/fluss/server/zk/ZooKeeperClientTest.java b/fluss-server/src/test/java/org/apache/fluss/server/zk/ZooKeeperClientTest.java index cbc0b85c6a2..60020824a53 100644 --- a/fluss-server/src/test/java/org/apache/fluss/server/zk/ZooKeeperClientTest.java +++ b/fluss-server/src/test/java/org/apache/fluss/server/zk/ZooKeeperClientTest.java @@ -81,6 +81,7 @@ import static org.mockito.ArgumentMatchers.anyCollection; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.verify; /** Test for {@link ZooKeeperClient}. */ class ZooKeeperClientTest { @@ -594,7 +595,8 @@ void testGetLatestBucketSnapshotsInBatch() throws Exception { new PartitionAssignment(tableId, partitionBucketAssignments), remoteDataDir, TablePath.of("db", "partitioned_table"), - tableId); + tableId, + partitionBucketAssignments.size()); TableBucket partitionBucket = new TableBucket(tableId, partitionId, 0); BucketSnapshot partitionSnapshot = new BucketSnapshot(5L, 50L, "oss://test/partition-cp5"); zookeeperClient.registerTableBucketSnapshot(partitionBucket, partitionSnapshot); @@ -734,9 +736,21 @@ void testPartition() throws Exception { }) .getBucketAssignments()); zookeeperClient.registerPartitionAssignmentAndMetadata( - 1L, "p1", partitionAssignment, remoteDataDir, tablePath, tableId); + 1L, + "p1", + partitionAssignment, + remoteDataDir, + tablePath, + tableId, + partitionAssignment.getBucketAssignments().size()); zookeeperClient.registerPartitionAssignmentAndMetadata( - 2L, "p2", partitionAssignment, remoteDataDir, tablePath, tableId); + 2L, + "p2", + partitionAssignment, + remoteDataDir, + tablePath, + tableId, + partitionAssignment.getBucketAssignments().size()); // check created partitions partitions = zookeeperClient.getPartitions(tablePath); @@ -748,8 +762,30 @@ void testPartition() throws Exception { assertThat(zookeeperClient.getPartitionsForTables(Arrays.asList(tablePath))) .containsValues(new ArrayList<>(partitions)); - // test delete partition - zookeeperClient.deletePartition(tablePath, "p1"); + // A batch read returns every registration and preserves the version needed by CAS updates. + PartitionRegistration p1Registration = zookeeperClient.getPartition(tablePath, "p1").get(); + zookeeperClient.updatePartitionRegistration(tablePath, "p1", p1Registration); + ZooKeeperClient batchReadClient = spy(zookeeperClient); + Map> registrations = + batchReadClient.getPartitionRegistrationsWithVersion(tablePath); + verify(batchReadClient).getDataInBackground(anyCollection()); + assertThat(registrations).containsOnlyKeys("p1", "p2"); + assertThat(registrations.get("p1").data().getPartitionId()).isEqualTo(1L); + assertThat(registrations.get("p1").zkVersion()).isEqualTo(1); + assertThat(registrations.get("p2").data().getPartitionId()).isEqualTo(2L); + assertThat(registrations.get("p2").zkVersion()).isZero(); + + // A partition dropped after the children listing is omitted from the batch result. + ZooKeeperClient raceTestingClient = spy(zookeeperClient); + doAnswer( + invocation -> { + zookeeperClient.deletePartition(tablePath, "p1"); + return invocation.callRealMethod(); + }) + .when(raceTestingClient) + .getDataInBackground(anyCollection()); + assertThat(raceTestingClient.getPartitionRegistrationsWithVersion(tablePath)) + .containsOnlyKeys("p2"); partitions = zookeeperClient.getPartitions(tablePath); assertThat(partitions).containsExactly("p2"); } diff --git a/fluss-server/src/test/java/org/apache/fluss/server/zk/data/PartitionRegistrationJsonSerdeTest.java b/fluss-server/src/test/java/org/apache/fluss/server/zk/data/PartitionRegistrationJsonSerdeTest.java index 3676688f5c0..cd3dca553bc 100644 --- a/fluss-server/src/test/java/org/apache/fluss/server/zk/data/PartitionRegistrationJsonSerdeTest.java +++ b/fluss-server/src/test/java/org/apache/fluss/server/zk/data/PartitionRegistrationJsonSerdeTest.java @@ -36,10 +36,14 @@ class PartitionRegistrationJsonSerdeTest extends JsonSerdeTestBase + if (info.getBucketCount != tableBucketCount) { + throw new UnsupportedOperationException( + s"Spark does not yet support per-partition bucket count rescale. " + + s"Table $tablePath partition ${info.getPartitionName} has bucket count " + + s"${info.getBucketCount} but the table-level count is $tableBucketCount.") + } + } + infos + } private var allDataForTriggerAvailableNow: Option[TableBucketOffsets] = None diff --git a/fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/read/SplitPlanner.scala b/fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/read/SplitPlanner.scala index e8386c94ffa..e178ffd6891 100644 --- a/fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/read/SplitPlanner.scala +++ b/fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/read/SplitPlanner.scala @@ -130,8 +130,22 @@ abstract class AbstractSplitPlanner( admin0 } - protected lazy val partitionInfos: util.List[PartitionInfo] = - admin.listPartitionInfos(tablePath).get() + protected lazy val partitionInfos: util.List[PartitionInfo] = { + val infos = admin.listPartitionInfos(tablePath).get() + // Fail fast if any partition's bucket count differs from the table-level count. + // Per-partition bucket count rescale (ALTER bucket.num) is not yet supported in Spark. + val tableBucketCount = tableInfo.getNumBuckets + infos.asScala.foreach { + info => + if (info.getBucketCount != tableBucketCount) { + throw new UnsupportedOperationException( + s"Spark does not yet support per-partition bucket count rescale. " + + s"Table $tablePath partition ${info.getPartitionName} has bucket count " + + s"${info.getBucketCount} but the table-level count is $tableBucketCount.") + } + } + infos + } protected def stoppingOffsetsInitializer: OffsetsInitializer diff --git a/website/docs/engine-flink/ddl.md b/website/docs/engine-flink/ddl.md index 37ca0090a9a..63b000675be 100644 --- a/website/docs/engine-flink/ddl.md +++ b/website/docs/engine-flink/ddl.md @@ -290,6 +290,7 @@ When using SET to modify [Storage Options](engine-flink/options.md#storage-optio **Supported Options to modify** - All [Read Options](engine-flink/options.md#read-options), [Write Options](engine-flink/options.md#write-options), [Lookup Options](engine-flink/options.md#lookup-options) and [Other Options](engine-flink/options.md#other-options) except `bootstrap.servers`. +- `bucket.num`: Set the target number of buckets. For partitioned tables, the new value applies to newly created partitions; existing partitions retain their original bucket count. Not supported on non-partitioned tables, and among lake-enabled tables only Paimon is supported. - The following [Storage Options](engine-flink/options.md#storage-options): - `table.datalake.enabled`: Enable or disable lakehouse storage for the table. - `table.datalake.historical-partition.enabled`: Enable or disable historical partition lookup. @@ -299,6 +300,9 @@ When using SET to modify [Storage Options](engine-flink/options.md#storage-optio - `table.auto-partition.num-precreate`: Set the number of future partitions to pre-create for auto partitioning. ```sql title="Flink SQL" +-- Change the bucket count for a partitioned table (applies to new partitions only) +ALTER TABLE my_table SET ('bucket.num' = '8'); + -- Enable lakehouse storage for the table ALTER TABLE my_table SET ('table.datalake.enabled' = 'true'); diff --git a/website/docs/engine-flink/options.md b/website/docs/engine-flink/options.md index 895051e7d9e..a30f93dcb85 100644 --- a/website/docs/engine-flink/options.md +++ b/website/docs/engine-flink/options.md @@ -64,7 +64,7 @@ See more details about [ALTER TABLE ... SET](engine-flink/ddl.md#set-properties) | Option | Type | Default | Description | |-----------------------------------------|----------|-------------------------------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| | auto-increment.fields | String | (None) | Defines the auto increment columns. The auto increment column can only be used in primary-key table. With an auto increment column in the table, whenever a new row is inserted into the table, the new row will be assigned with the next available value from the auto-increment sequence. The data type of the auto increment column must be INT or BIGINT. Currently a table can have only one auto-increment column. Adding an auto increment column to an existing table is not supported. | -| bucket.num | int | The bucket number of Fluss cluster. | The number of buckets of a Fluss table. | +| bucket.num | int | The bucket number of Fluss cluster. | The target number of buckets for a Fluss table. For partitioned tables, this value applies to newly created partitions; existing partitions retain their original bucket count. | | bucket.key | String | (None) | Specific the distribution policy of the Fluss table. Data will be distributed to each bucket according to the hash value of bucket-key (It must be a subset of the primary keys excluding partition keys of the primary key table). If you specify multiple fields, delimiter is `,`. If the table has a primary key and a bucket key is not specified, the bucket key will be used as primary key(excluding the partition key). If the table has no primary key and the bucket key is not specified, the data will be distributed to each bucket randomly. | | table.log.ttl | Duration | 7 days | The time to live for log segments. The configuration controls the maximum time log segments are retained before they become eligible for deletion. When remote log tiering is enabled, this value controls the retention of remote log segments. Setting the value to '0ms' disables TTL-based deletion. The default value is 7 days. | | table.log.local-ttl | Duration | (None) | The time to live for local log segments. The configuration controls the maximum time local log segments are retained before they become eligible for deletion. When remote log tiering is enabled, an expired local segment is deleted only after it has been copied to remote storage. Setting the value to '0ms' disables TTL-based deletion. If not configured, the value inherits `table.log.ttl`. When both values are positive, it must be less than or equal to `table.log.ttl`. | diff --git a/website/docs/table-design/data-distribution/bucketing.md b/website/docs/table-design/data-distribution/bucketing.md index 615c1cb0d0b..89f6829c800 100644 --- a/website/docs/table-design/data-distribution/bucketing.md +++ b/website/docs/table-design/data-distribution/bucketing.md @@ -8,7 +8,7 @@ sidebar_position: 1 A bucketing strategy is a data distribution technique that divides table data into small pieces and distributes the data to multiple hosts and services. -When creating a Fluss table, you can specify the number of buckets by setting `'bucket.num' = ''` property for the table, see more details in [DDL](engine-flink/ddl.md). +When creating a Fluss table, you can specify the number of buckets by setting `'bucket.num' = ''` property for the table, see more details in [DDL](engine-flink/ddl.md). For partitioned tables, `bucket.num` can be altered via `ALTER TABLE SET ('bucket.num' = '')` — the new value applies to newly created partitions while existing partitions retain their original bucket count. Currently, Fluss supports 3 bucketing strategies: **Hash Bucketing**, **Sticky Bucketing** and **Round-Robin Bucketing**. Primary-Key Tables only allow to use **Hash Bucketing**. Log Tables use **Sticky Bucketing** by default but can use other two bucketing strategies.