diff --git a/paimon-api/src/main/java/org/apache/paimon/CoreOptions.java b/paimon-api/src/main/java/org/apache/paimon/CoreOptions.java index b30e01544267..2e345a694bab 100644 --- a/paimon-api/src/main/java/org/apache/paimon/CoreOptions.java +++ b/paimon-api/src/main/java/org/apache/paimon/CoreOptions.java @@ -2282,6 +2282,23 @@ public String toString() { "For DELETE manifest entry in manifest file, drop stats to reduce memory and storage." + " Default value is false only for compatibility of old reader."); + public static final ConfigOption PARTITION_BUCKET_MAPPING_CACHE_ENABLED = + key("partition-bucket-mapping.cache-enabled") + .booleanType() + .defaultValue(false) + .withDescription( + "If true, cache partition bucket mappings in the current JVM when initializing writers." + + " This avoids repeated manifest scans by multiple writers in the same TaskManager," + + " but the cached mapping is shared until the table snapshot changes."); + + public static final ConfigOption PARTITION_BUCKET_MAPPING_CACHE_MAX_ENTRIES = + key("partition-bucket-mapping.cache-max-entries") + .intType() + .defaultValue(128) + .withDescription( + "Maximum number of partition bucket mappings to cache in the current JVM." + + " Older snapshots of the same table are invalidated when a newer snapshot mapping is loaded."); + public static final ConfigOption DATA_FILE_THIN_MODE = key("data-file.thin-mode") .booleanType() @@ -3681,6 +3698,14 @@ public boolean manifestDeleteFileDropStats() { return options.get(MANIFEST_DELETE_FILE_DROP_STATS); } + public boolean partitionBucketMappingCacheEnabled() { + return options.get(PARTITION_BUCKET_MAPPING_CACHE_ENABLED); + } + + public int partitionBucketMappingCacheMaxEntries() { + return options.get(PARTITION_BUCKET_MAPPING_CACHE_MAX_ENTRIES); + } + public boolean disableNullToNotNull() { return options.get(DISABLE_ALTER_COLUMN_NULL_TO_NOT_NULL); } diff --git a/paimon-core/src/main/java/org/apache/paimon/table/AbstractFileStoreTable.java b/paimon-core/src/main/java/org/apache/paimon/table/AbstractFileStoreTable.java index 2f893d4e3b84..8585f807af71 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/AbstractFileStoreTable.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/AbstractFileStoreTable.java @@ -40,6 +40,7 @@ import org.apache.paimon.table.sink.DynamicBucketRowKeyExtractor; import org.apache.paimon.table.sink.FixedBucketRowKeyExtractor; import org.apache.paimon.table.sink.FixedBucketWriteSelector; +import org.apache.paimon.table.sink.PartitionBucketMapping; import org.apache.paimon.table.sink.PostponeBucketRowKeyExtractor; import org.apache.paimon.table.sink.RowKeyExtractor; import org.apache.paimon.table.sink.RowKindGenerator; @@ -228,7 +229,9 @@ public Optional statistics() { public Optional newWriteSelector() { switch (bucketMode()) { case HASH_FIXED: - return Optional.of(new FixedBucketWriteSelector(schema())); + return Optional.of( + new FixedBucketWriteSelector( + schema(), PartitionBucketMapping.loadFromTable(this))); case BUCKET_UNAWARE: case POSTPONE_MODE: return Optional.empty(); @@ -256,7 +259,8 @@ protected CatalogEnvironment newCatalogEnvironment(String branch) { public RowKeyExtractor createRowKeyExtractor() { switch (bucketMode()) { case HASH_FIXED: - return new FixedBucketRowKeyExtractor(schema()); + return new FixedBucketRowKeyExtractor( + schema(), PartitionBucketMapping.loadFromTable(this)); case HASH_DYNAMIC: case KEY_DYNAMIC: return new DynamicBucketRowKeyExtractor(schema()); diff --git a/paimon-core/src/main/java/org/apache/paimon/table/AppendOnlyFileStoreTable.java b/paimon-core/src/main/java/org/apache/paimon/table/AppendOnlyFileStoreTable.java index 63710698ca3b..31ff9bd4f458 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/AppendOnlyFileStoreTable.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/AppendOnlyFileStoreTable.java @@ -30,6 +30,7 @@ import org.apache.paimon.predicate.Predicate; import org.apache.paimon.schema.TableSchema; import org.apache.paimon.table.query.LocalTableQuery; +import org.apache.paimon.table.sink.RowKeyExtractor; import org.apache.paimon.table.sink.TableWriteImpl; import org.apache.paimon.table.source.AppendBatchTableScan; import org.apache.paimon.table.source.AppendOnlySplitGenerator; @@ -162,11 +163,17 @@ public TableWriteImpl newWrite(String commitUser) { @Override public TableWriteImpl newWrite(String commitUser, @Nullable Integer writeId) { + return newWrite(commitUser, writeId, createRowKeyExtractor()); + } + + @Override + public TableWriteImpl newWrite( + String commitUser, @Nullable Integer writeId, RowKeyExtractor rowKeyExtractor) { BaseAppendFileStoreWrite writer = store().newWrite(commitUser, writeId); return new TableWriteImpl<>( rowType(), writer, - createRowKeyExtractor(), + rowKeyExtractor, (record, rowKind) -> { Preconditions.checkState( rowKind.isAdd(), diff --git a/paimon-core/src/main/java/org/apache/paimon/table/DelegatedFileStoreTable.java b/paimon-core/src/main/java/org/apache/paimon/table/DelegatedFileStoreTable.java index 9a5e81bcb636..289aaf6517c1 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/DelegatedFileStoreTable.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/DelegatedFileStoreTable.java @@ -345,6 +345,12 @@ public TableWriteImpl newWrite(String commitUser, @Nullable Integer writeId) return wrapped.newWrite(commitUser, writeId); } + @Override + public TableWriteImpl newWrite( + String commitUser, @Nullable Integer writeId, RowKeyExtractor rowKeyExtractor) { + return wrapped.newWrite(commitUser, writeId, rowKeyExtractor); + } + @Override public TableWriteImpl newPostponeFixedBucketWrite( String commitUser, @Nullable Integer writeId) { diff --git a/paimon-core/src/main/java/org/apache/paimon/table/FileStoreTable.java b/paimon-core/src/main/java/org/apache/paimon/table/FileStoreTable.java index 970493430589..e1a6a14652d5 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/FileStoreTable.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/FileStoreTable.java @@ -139,6 +139,10 @@ default PostponeFixedBucketWriteBuilder newPostponeFixedBucketWriteBuilder() { TableWriteImpl newWrite(String commitUser, @Nullable Integer writeId); + /** Creates a new write using the supplied bucket assignment logic. */ + TableWriteImpl newWrite( + String commitUser, @Nullable Integer writeId, RowKeyExtractor rowKeyExtractor); + /** Creates a fixed-bucket merge-tree write for a postpone-bucket batch write. */ default TableWriteImpl newPostponeFixedBucketWrite( String commitUser, @Nullable Integer writeId) { diff --git a/paimon-core/src/main/java/org/apache/paimon/table/OverwriteFileStoreTable.java b/paimon-core/src/main/java/org/apache/paimon/table/OverwriteFileStoreTable.java new file mode 100644 index 000000000000..b96e3cf3a392 --- /dev/null +++ b/paimon-core/src/main/java/org/apache/paimon/table/OverwriteFileStoreTable.java @@ -0,0 +1,99 @@ +/* + * 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.paimon.table; + +import org.apache.paimon.schema.TableSchema; +import org.apache.paimon.table.sink.FixedBucketRowKeyExtractor; +import org.apache.paimon.table.sink.FixedBucketWriteSelector; +import org.apache.paimon.table.sink.PartitionBucketMapping; +import org.apache.paimon.table.sink.RowKeyExtractor; +import org.apache.paimon.table.sink.TableWriteImpl; +import org.apache.paimon.table.sink.WriteSelector; + +import javax.annotation.Nullable; + +import java.util.Map; +import java.util.Optional; + +/** + * A table wrapper for overwrite operations which routes rows using the target schema bucket count. + * Existing per-partition bucket mappings must not be used while rewriting a partition to a new + * bucket count. + */ +public class OverwriteFileStoreTable extends DelegatedFileStoreTable { + + public OverwriteFileStoreTable(FileStoreTable wrapped) { + super(wrapped); + } + + private PartitionBucketMapping targetBucketMapping() { + return new PartitionBucketMapping(schema().numBuckets()); + } + + @Override + public Optional newWriteSelector() { + return Optional.of(new FixedBucketWriteSelector(schema(), targetBucketMapping())); + } + + @Override + public RowKeyExtractor createRowKeyExtractor() { + return new FixedBucketRowKeyExtractor(schema(), targetBucketMapping()); + } + + @Override + public TableWriteImpl newWrite(String commitUser) { + return newWrite(commitUser, null); + } + + @Override + public TableWriteImpl newWrite(String commitUser, @Nullable Integer writeId) { + return wrapped().newWrite(commitUser, writeId, createRowKeyExtractor()); + } + + @Override + public TableWriteImpl newWrite( + String commitUser, @Nullable Integer writeId, RowKeyExtractor rowKeyExtractor) { + return wrapped().newWrite(commitUser, writeId, rowKeyExtractor); + } + + @Override + public FileStoreTable copy(Map dynamicOptions) { + return new OverwriteFileStoreTable(wrapped().copy(dynamicOptions)); + } + + @Override + public FileStoreTable copy(TableSchema newTableSchema) { + return new OverwriteFileStoreTable(wrapped().copy(newTableSchema)); + } + + @Override + public FileStoreTable copyWithoutTimeTravel(Map dynamicOptions) { + return new OverwriteFileStoreTable(wrapped().copyWithoutTimeTravel(dynamicOptions)); + } + + @Override + public FileStoreTable copyWithLatestSchema() { + return new OverwriteFileStoreTable(wrapped().copyWithLatestSchema()); + } + + @Override + public FileStoreTable switchToBranch(String branchName) { + return new OverwriteFileStoreTable(wrapped().switchToBranch(branchName)); + } +} diff --git a/paimon-core/src/main/java/org/apache/paimon/table/PrimaryKeyFileStoreTable.java b/paimon-core/src/main/java/org/apache/paimon/table/PrimaryKeyFileStoreTable.java index f521f2e9b17f..b8e41ccdddd6 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/PrimaryKeyFileStoreTable.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/PrimaryKeyFileStoreTable.java @@ -33,6 +33,7 @@ import org.apache.paimon.schema.KeyValueFieldsExtractor; import org.apache.paimon.schema.TableSchema; import org.apache.paimon.table.query.LocalTableQuery; +import org.apache.paimon.table.sink.RowKeyExtractor; import org.apache.paimon.table.sink.TableWriteImpl; import org.apache.paimon.table.source.DataTableScan; import org.apache.paimon.table.source.InnerTableRead; @@ -178,21 +179,28 @@ public TableWriteImpl newWrite(String commitUser) { @Override public TableWriteImpl newWrite(String commitUser, @Nullable Integer writeId) { - return newWrite(store().newWrite(commitUser, writeId)); + return newWrite(commitUser, writeId, createRowKeyExtractor()); + } + + @Override + public TableWriteImpl newWrite( + String commitUser, @Nullable Integer writeId, RowKeyExtractor rowKeyExtractor) { + return newWrite(store().newWrite(commitUser, writeId), rowKeyExtractor); } @Override public TableWriteImpl newPostponeFixedBucketWrite( String commitUser, @Nullable Integer writeId) { - return newWrite(store().newPostponeFixedBucketWrite(commitUser)); + return newWrite(store().newPostponeFixedBucketWrite(commitUser), createRowKeyExtractor()); } - private TableWriteImpl newWrite(AbstractFileStoreWrite storeWrite) { + private TableWriteImpl newWrite( + AbstractFileStoreWrite storeWrite, RowKeyExtractor rowKeyExtractor) { KeyValue kv = new KeyValue(); return new TableWriteImpl<>( rowType(), storeWrite, - createRowKeyExtractor(), + rowKeyExtractor, (record, rowKind) -> kv.replace( record.primaryKey(), diff --git a/paimon-core/src/main/java/org/apache/paimon/table/sink/FixedBucketRowKeyExtractor.java b/paimon-core/src/main/java/org/apache/paimon/table/sink/FixedBucketRowKeyExtractor.java index 146a45b43713..3d8c930bce96 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/sink/FixedBucketRowKeyExtractor.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/sink/FixedBucketRowKeyExtractor.java @@ -29,24 +29,26 @@ /** {@link KeyAndBucketExtractor} for {@link InternalRow}. */ public class FixedBucketRowKeyExtractor extends RowKeyExtractor { - private final int numBuckets; + private final PartitionBucketMapping partitionBucketMapping; private final boolean sameBucketKeyAndTrimmedPrimaryKey; - private final Projection bucketKeyProjection; + private transient Projection bucketKeyProjection; private BinaryRow reuseBucketKey; private Integer reuseBucket; private final BucketFunction bucketFunction; public FixedBucketRowKeyExtractor(TableSchema schema) { + this(schema, new PartitionBucketMapping(new CoreOptions(schema.options()).bucket())); + } + + public FixedBucketRowKeyExtractor( + TableSchema schema, PartitionBucketMapping partitionBucketMapping) { super(schema); - numBuckets = new CoreOptions(schema.options()).bucket(); bucketFunction = BucketFunction.create( new CoreOptions(schema.options()), schema.logicalBucketKeyType()); sameBucketKeyAndTrimmedPrimaryKey = schema.bucketKeys().equals(schema.trimmedPrimaryKeys()); - bucketKeyProjection = - CodeGenUtils.newProjection( - schema.logicalRowType(), schema.projection(schema.bucketKeys())); + this.partitionBucketMapping = partitionBucketMapping; } @Override @@ -62,7 +64,7 @@ private BinaryRow bucketKey() { } if (reuseBucketKey == null) { - reuseBucketKey = bucketKeyProjection.apply(record); + reuseBucketKey = bucketKeyProjection().apply(record); } return reuseBucketKey; } @@ -70,7 +72,7 @@ private BinaryRow bucketKey() { @Override public int bucket() { if (reuseBucket == null) { - reuseBucket = bucket(numBuckets); + reuseBucket = bucket(partitionBucketMapping.resolveNumBuckets(partition())); } return reuseBucket; } @@ -78,4 +80,13 @@ public int bucket() { public int bucket(int numBuckets) { return bucketFunction.bucket(bucketKey(), numBuckets); } + + private Projection bucketKeyProjection() { + if (bucketKeyProjection == null) { + bucketKeyProjection = + CodeGenUtils.newProjection( + schema.logicalRowType(), schema.projection(schema.bucketKeys())); + } + return bucketKeyProjection; + } } diff --git a/paimon-core/src/main/java/org/apache/paimon/table/sink/FixedBucketWriteSelector.java b/paimon-core/src/main/java/org/apache/paimon/table/sink/FixedBucketWriteSelector.java index e08841dd8cd3..245886e2db27 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/sink/FixedBucketWriteSelector.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/sink/FixedBucketWriteSelector.java @@ -28,17 +28,24 @@ public class FixedBucketWriteSelector implements WriteSelector { private static final long serialVersionUID = 1L; private final TableSchema schema; + private final PartitionBucketMapping partitionBucketMapping; private transient KeyAndBucketExtractor extractor; public FixedBucketWriteSelector(TableSchema schema) { + this(schema, new PartitionBucketMapping(schema.numBuckets())); + } + + public FixedBucketWriteSelector( + TableSchema schema, PartitionBucketMapping partitionBucketMapping) { this.schema = schema; + this.partitionBucketMapping = partitionBucketMapping; } @Override public int select(InternalRow row, int numWriters) { if (extractor == null) { - extractor = new FixedBucketRowKeyExtractor(schema); + extractor = new FixedBucketRowKeyExtractor(schema, partitionBucketMapping); } extractor.setRecord(row); return ChannelComputer.select(extractor.partition(), extractor.bucket(), numWriters); diff --git a/paimon-core/src/main/java/org/apache/paimon/table/sink/PartitionBucketMapping.java b/paimon-core/src/main/java/org/apache/paimon/table/sink/PartitionBucketMapping.java new file mode 100644 index 000000000000..002b862d3b49 --- /dev/null +++ b/paimon-core/src/main/java/org/apache/paimon/table/sink/PartitionBucketMapping.java @@ -0,0 +1,245 @@ +/* + * 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.paimon.table.sink; + +import org.apache.paimon.data.BinaryRow; +import org.apache.paimon.manifest.FileEntry; +import org.apache.paimon.manifest.SimpleFileEntry; +import org.apache.paimon.schema.TableSchema; +import org.apache.paimon.table.FileStoreTable; + +import org.apache.paimon.shade.caffeine2.com.github.benmanes.caffeine.cache.Cache; +import org.apache.paimon.shade.caffeine2.com.github.benmanes.caffeine.cache.Caffeine; + +import java.io.Serializable; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.function.Supplier; + +/** + * A mapping that resolves the number of buckets for each partition in a table. + * + *

Different partitions may have different bucket counts (e.g., after a rescale operation). This + * class maintains a per-partition bucket count mapping and falls back to a default bucket count for + * partitions that are not explicitly mapped. + * + *

This is used by components such as {@link FixedBucketRowKeyExtractor} and {@link + * FixedBucketWriteSelector} to correctly determine the bucket assignment for rows in tables where + * partitions may have been rescaled independently. + * + * @see #loadFromTable(FileStoreTable) + * @see #resolveNumBuckets(BinaryRow) + */ +public class PartitionBucketMapping implements Serializable { + + private static final long serialVersionUID = 1L; + + private static final Cache CACHE = + Caffeine.newBuilder().maximumSize(128).executor(Runnable::run).build(); + + /** The default number of buckets, used when a partition has no explicit mapping. */ + private final int defaultBucketCount; + + /** A map from partition to its specific bucket count. May be empty but never {@code null}. */ + private final Map partitionBucketMap; + + /** + * Creates a mapping with only a default bucket count and no per-partition overrides. + * + * @param defaultBucketCount the default number of buckets for all partitions + */ + public PartitionBucketMapping(int defaultBucketCount) { + this(defaultBucketCount, Collections.emptyMap()); + } + + /** + * Creates a mapping with a default bucket count and an explicit per-partition bucket map. + * + * @param defaultBucketCount the default number of buckets, used as a fallback + * @param partitionBucketMap a map from partition (as {@link BinaryRow}) to its bucket count + */ + public PartitionBucketMapping( + int defaultBucketCount, Map partitionBucketMap) { + this.defaultBucketCount = defaultBucketCount; + this.partitionBucketMap = Collections.unmodifiableMap(new HashMap<>(partitionBucketMap)); + } + + /** + * Loads a {@link PartitionBucketMapping} by scanning the manifest entries of the given table. + * + *

For non-partitioned tables, this returns a mapping with only the schema-defined default + * bucket count and an empty partition map. + * + *

For partitioned tables, the method scans all manifest entries and records the {@code + * totalBuckets} value for each partition. If the scan fails for any reason, a fallback mapping + * with only the default bucket count is returned. + * + * @param table the {@link FileStoreTable} to load the mapping from + * @return a {@link PartitionBucketMapping} reflecting the current bucket layout of the table + */ + public static PartitionBucketMapping loadFromTable(FileStoreTable table) { + if (!table.coreOptions().partitionBucketMappingCacheEnabled()) { + return loadFromTableWithoutCache(table); + } + + CacheKey key = CacheKey.from(table); + return getOrLoad( + key, + table.coreOptions().partitionBucketMappingCacheMaxEntries(), + () -> loadFromTableWithoutCache(table)); + } + + static PartitionBucketMapping getOrLoad( + CacheKey key, int maxEntries, Supplier mappingLoader) { + ensureMaximumSize(maxEntries); + return CACHE.get( + key, + ignored -> { + invalidateOlderSnapshots(key); + return mappingLoader.get(); + }); + } + + private static void ensureMaximumSize(int maxEntries) { + CACHE.policy().eviction().ifPresent(eviction -> eviction.setMaximum(maxEntries)); + } + + private static void invalidateOlderSnapshots(CacheKey key) { + CACHE.asMap().keySet().stream() + .filter( + existingKey -> + existingKey.isSameTable(key) + && existingKey.latestSnapshotId < key.latestSnapshotId) + .forEach(CACHE::invalidate); + } + + static PartitionBucketMapping loadFromTableWithoutCache(FileStoreTable table) { + int defaultBuckets = table.schema().numBuckets(); + if (table.partitionKeys().isEmpty()) { + return new PartitionBucketMapping(defaultBuckets, Collections.emptyMap()); + } + + List entries = table.store().newScan().readSimpleEntries(); + return loadFromEntries(entries, table.schema()); + } + + static void clearCache() { + CACHE.invalidateAll(); + } + + static PartitionBucketMapping getCachedMapping(FileStoreTable table) { + return CACHE.getIfPresent(CacheKey.from(table)); + } + + static PartitionBucketMapping getCachedMapping(CacheKey key) { + return CACHE.getIfPresent(key); + } + + public static PartitionBucketMapping loadFromEntries( + List entries, TableSchema tableSchema) { + int defaultBuckets = tableSchema.numBuckets(); + if (tableSchema.partitionKeys().isEmpty()) { + return new PartitionBucketMapping(defaultBuckets, Collections.emptyMap()); + } + + Map partitionBucketMap = new HashMap<>(); + for (FileEntry entry : entries) { + int totalBuckets = entry.totalBuckets(); + // Only store partitions whose bucket count differs from the default. + // This keeps the map empty for partitions that have never been rescaled, + // avoiding per-partition BinaryRow copies and Integer allocations entirely. + if (totalBuckets > 0 && totalBuckets != defaultBuckets) { + BinaryRow partition = entry.partition(); + partitionBucketMap.putIfAbsent(partition.copy(), totalBuckets); + } + } + return new PartitionBucketMapping(defaultBuckets, partitionBucketMap); + } + + /** + * Resolves the number of buckets for the given partition. + * + *

If the partition has an explicit entry in the partition-to-bucket map, that value is + * returned. Otherwise, the default bucket count is returned. + * + * @param partition the partition key as a {@link BinaryRow} + * @return the number of buckets for the given partition + */ + public int resolveNumBuckets(BinaryRow partition) { + if (partitionBucketMap != null) { + Integer partitionBucketCount = partitionBucketMap.get(partition); + if (partitionBucketCount != null) { + return partitionBucketCount; + } + } + return defaultBucketCount; + } + + static class CacheKey { + private final String tablePath; + private final long latestSnapshotId; + private final long schemaId; + private final int defaultBucketCount; + + CacheKey(String tablePath, long latestSnapshotId, long schemaId, int defaultBucketCount) { + this.tablePath = tablePath; + this.latestSnapshotId = latestSnapshotId; + this.schemaId = schemaId; + this.defaultBucketCount = defaultBucketCount; + } + + private static CacheKey from(FileStoreTable table) { + Long latestSnapshotId = table.snapshotManager().latestSnapshotId(); + return new CacheKey( + table.location().toString(), + latestSnapshotId == null ? -1L : latestSnapshotId, + table.schema().id(), + table.schema().numBuckets()); + } + + private boolean isSameTable(CacheKey other) { + return schemaId == other.schemaId + && defaultBucketCount == other.defaultBucketCount + && Objects.equals(tablePath, other.tablePath); + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (!(o instanceof CacheKey)) { + return false; + } + CacheKey cacheKey = (CacheKey) o; + return latestSnapshotId == cacheKey.latestSnapshotId + && schemaId == cacheKey.schemaId + && defaultBucketCount == cacheKey.defaultBucketCount + && Objects.equals(tablePath, cacheKey.tablePath); + } + + @Override + public int hashCode() { + return Objects.hash(tablePath, latestSnapshotId, schemaId, defaultBucketCount); + } + } +} diff --git a/paimon-core/src/main/java/org/apache/paimon/table/sink/RowKeyExtractor.java b/paimon-core/src/main/java/org/apache/paimon/table/sink/RowKeyExtractor.java index 455aaa4aa5e9..f94d3d9cd7db 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/sink/RowKeyExtractor.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/sink/RowKeyExtractor.java @@ -22,10 +22,16 @@ import org.apache.paimon.data.InternalRow; import org.apache.paimon.schema.TableSchema; +import java.io.Serializable; + /** {@link KeyAndBucketExtractor} for {@link InternalRow}. */ -public abstract class RowKeyExtractor implements KeyAndBucketExtractor { +public abstract class RowKeyExtractor implements KeyAndBucketExtractor, Serializable { + + private static final long serialVersionUID = 1L; + + private transient RowPartitionKeyExtractor partitionKeyExtractor; - private final RowPartitionKeyExtractor partitionKeyExtractor; + protected final TableSchema schema; protected InternalRow record; @@ -33,7 +39,7 @@ public abstract class RowKeyExtractor implements KeyAndBucketExtractor partitionMap = new HashMap<>(); + partitionMap.put(mappedPartition, 4); + PartitionBucketMapping mapping = new PartitionBucketMapping(defaultBuckets, partitionMap); + + FixedBucketRowKeyExtractor extractor = extractor("a", "b", "a,b", defaultBuckets, mapping); + + assertThat(bucket(extractor, GenericRow.of(1, 456, 7))).isEqualTo(3); + assertThat(bucket(extractor, GenericRow.of(99, 456, 7))).isEqualTo(47); + } + private int bucket(FixedBucketRowKeyExtractor extractor, InternalRow row) { extractor.setRecord(row); return extractor.bucket(); @@ -125,8 +139,29 @@ private FixedBucketRowKeyExtractor extractor( return extractor(rowType, partK, bk, pk, numBucket); } + private FixedBucketRowKeyExtractor extractor( + String partK, String bk, String pk, int numBucket, PartitionBucketMapping mapping) { + RowType rowType = + new RowType( + Arrays.asList( + new DataField(0, "a", new IntType()), + new DataField(1, "b", new IntType()), + new DataField(2, "c", new IntType()))); + return extractor(rowType, partK, bk, pk, numBucket, mapping); + } + private FixedBucketRowKeyExtractor extractor( RowType rowType, String partK, String bk, String pk, int numBucket) { + return extractor(rowType, partK, bk, pk, numBucket, new PartitionBucketMapping(numBucket)); + } + + private FixedBucketRowKeyExtractor extractor( + RowType rowType, + String partK, + String bk, + String pk, + int numBucket, + PartitionBucketMapping mapping) { List fields = TableSchema.newFields(rowType); Map options = new HashMap<>(); options.put(BUCKET_KEY.key(), bk); @@ -142,6 +177,6 @@ private FixedBucketRowKeyExtractor extractor( "".equals(pk) ? Collections.emptyList() : Arrays.asList(pk.split(",")), options, ""); - return new FixedBucketRowKeyExtractor(schema); + return new FixedBucketRowKeyExtractor(schema, mapping); } } diff --git a/paimon-core/src/test/java/org/apache/paimon/table/sink/PartitionBucketMappingTest.java b/paimon-core/src/test/java/org/apache/paimon/table/sink/PartitionBucketMappingTest.java new file mode 100644 index 000000000000..e9a0b5f8ee2e --- /dev/null +++ b/paimon-core/src/test/java/org/apache/paimon/table/sink/PartitionBucketMappingTest.java @@ -0,0 +1,341 @@ +/* + * 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.paimon.table.sink; + +import org.apache.paimon.CoreOptions; +import org.apache.paimon.data.BinaryRow; +import org.apache.paimon.manifest.FileKind; +import org.apache.paimon.manifest.SimpleFileEntry; +import org.apache.paimon.schema.TableSchema; +import org.apache.paimon.types.DataField; +import org.apache.paimon.types.IntType; +import org.apache.paimon.types.RowType; + +import org.junit.jupiter.api.Test; + +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.atomic.AtomicInteger; + +import static org.assertj.core.api.Assertions.assertThat; + +/** Tests for {@link PartitionBucketMapping}. */ +public class PartitionBucketMappingTest { + + @Test + public void testDefaultBucketCount() { + PartitionBucketMapping mapping = new PartitionBucketMapping(16); + + // Any partition should resolve to the default + assertThat(mapping.resolveNumBuckets(BinaryRow.EMPTY_ROW)).isEqualTo(16); + assertThat(mapping.resolveNumBuckets(partition(1))).isEqualTo(16); + assertThat(mapping.resolveNumBuckets(partition(42))).isEqualTo(16); + } + + @Test + public void testExplicitPartitionMapping() { + BinaryRow partA = partition(1); + BinaryRow partB = partition(2); + BinaryRow partC = partition(3); + + Map partitionMap = new HashMap<>(); + partitionMap.put(partA, 32); + partitionMap.put(partB, 64); + + PartitionBucketMapping mapping = new PartitionBucketMapping(16, partitionMap); + + // Mapped partitions return their specific bucket counts + assertThat(mapping.resolveNumBuckets(partA)).isEqualTo(32); + assertThat(mapping.resolveNumBuckets(partB)).isEqualTo(64); + + // Unmapped partition falls back to the default + assertThat(mapping.resolveNumBuckets(partC)).isEqualTo(16); + } + + @Test + public void testLoadFromEntries_emptyList() { + PartitionBucketMapping mapping = + PartitionBucketMapping.loadFromEntries( + Collections.emptyList(), partitionedSchema(16)); + + // With no entries, no per-partition mapping exists; everything returns the default. + assertThat(mapping.resolveNumBuckets(partition(1))).isEqualTo(16); + assertThat(mapping.resolveNumBuckets(BinaryRow.EMPTY_ROW)).isEqualTo(16); + } + + @Test + public void testLoadFromEntries_nonPartitionedTableSkipsScan() { + // For non-partitioned tables, loadFromEntries must short-circuit and never + // populate the per-partition map, even if entries report a different + // totalBuckets value (which can happen mid-rescale or with stale snapshots). + // Otherwise it would trigger spurious "bucket changed without overwrite" + // errors at commit time on non-partitioned tables. + List entries = + Arrays.asList(entry(BinaryRow.EMPTY_ROW, 0, 1), entry(BinaryRow.EMPTY_ROW, 0, 4)); + + PartitionBucketMapping mapping = + PartitionBucketMapping.loadFromEntries(entries, nonPartitionedSchema(2)); + + // Should resolve to the schema default, not anything from the entries. + assertThat(mapping.resolveNumBuckets(BinaryRow.EMPTY_ROW)).isEqualTo(2); + } + + @Test + public void testLoadFromEntries_allDefaultBuckets() { + // Entries whose totalBuckets matches the default are intentionally not stored + // in the per-partition map (memory optimisation), but resolveNumBuckets must + // still return the default for those partitions. + List entries = + Arrays.asList( + entry(partition(1), 0, 16), + entry(partition(2), 0, 16), + entry(partition(3), 1, 16)); + + PartitionBucketMapping mapping = + PartitionBucketMapping.loadFromEntries(entries, partitionedSchema(16)); + + assertThat(mapping.resolveNumBuckets(partition(1))).isEqualTo(16); + assertThat(mapping.resolveNumBuckets(partition(2))).isEqualTo(16); + assertThat(mapping.resolveNumBuckets(partition(3))).isEqualTo(16); + // Unseen partition still returns default. + assertThat(mapping.resolveNumBuckets(partition(99))).isEqualTo(16); + } + + @Test + public void testLoadFromEntries_heterogeneousBuckets() { + // Reproduces the scenario from the FileSystemWriteRestore bug fix: + // table default = 32, but partition A has been rescaled to 2 buckets + // and partition B to 64. Partition C uses default (no entry needed). + BinaryRow partA = partition(1); + BinaryRow partB = partition(2); + BinaryRow partC = partition(3); + + List entries = + Arrays.asList( + entry(partA, 0, 2), + entry(partA, 1, 2), + entry(partB, 0, 64), + entry(partC, 0, 32)); + + PartitionBucketMapping mapping = + PartitionBucketMapping.loadFromEntries(entries, partitionedSchema(32)); + + assertThat(mapping.resolveNumBuckets(partA)).isEqualTo(2); + assertThat(mapping.resolveNumBuckets(partB)).isEqualTo(64); + // partC matches default and was skipped from the map; resolves to default. + assertThat(mapping.resolveNumBuckets(partC)).isEqualTo(32); + // unseen partition resolves to default. + assertThat(mapping.resolveNumBuckets(partition(99))).isEqualTo(32); + } + + @Test + public void testLoadFromEntries_zeroOrNegativeTotalBucketsIgnored() { + // Entries with totalBuckets <= 0 represent metadata/legacy entries that + // should not influence the mapping. + BinaryRow partA = partition(1); + + List entries = Arrays.asList(entry(partA, 0, 0), entry(partA, 1, -1)); + + PartitionBucketMapping mapping = + PartitionBucketMapping.loadFromEntries(entries, partitionedSchema(32)); + + // Nothing was stored; partA resolves to the default. + assertThat(mapping.resolveNumBuckets(partA)).isEqualTo(32); + } + + @Test + public void testLoadFromEntries_putIfAbsentSemantics() { + // If multiple entries for the same partition somehow report different + // totalBuckets values, the first observed value is kept (putIfAbsent + // semantics in loadFromEntries). This is a defensive contract test. + BinaryRow partA = partition(1); + + List entries = Arrays.asList(entry(partA, 0, 2), entry(partA, 1, 4)); + + PartitionBucketMapping mapping = + PartitionBucketMapping.loadFromEntries(entries, partitionedSchema(32)); + + assertThat(mapping.resolveNumBuckets(partA)).isEqualTo(2); + } + + @Test + public void testGetOrLoadCachesMappingByKey() { + PartitionBucketMapping.clearCache(); + PartitionBucketMapping.CacheKey key = + new PartitionBucketMapping.CacheKey("table", 1, 0, 32); + AtomicInteger loadCount = new AtomicInteger(); + + PartitionBucketMapping first = + PartitionBucketMapping.getOrLoad( + key, + 128, + () -> { + loadCount.incrementAndGet(); + return new PartitionBucketMapping(32); + }); + PartitionBucketMapping second = + PartitionBucketMapping.getOrLoad( + key, + 128, + () -> { + loadCount.incrementAndGet(); + return new PartitionBucketMapping(64); + }); + + assertThat(second).isSameAs(first); + assertThat(second.resolveNumBuckets(partition(1))).isEqualTo(32); + assertThat(loadCount).hasValue(1); + PartitionBucketMapping.clearCache(); + } + + @Test + public void testGetOrLoadInvalidatesOlderSnapshotsForSameTable() { + PartitionBucketMapping.clearCache(); + PartitionBucketMapping.CacheKey oldSnapshot = + new PartitionBucketMapping.CacheKey("table", 1, 0, 32); + PartitionBucketMapping.CacheKey newSnapshot = + new PartitionBucketMapping.CacheKey("table", 2, 0, 32); + + PartitionBucketMapping oldMapping = + PartitionBucketMapping.getOrLoad( + oldSnapshot, 128, () -> new PartitionBucketMapping(32)); + PartitionBucketMapping newMapping = + PartitionBucketMapping.getOrLoad( + newSnapshot, 128, () -> new PartitionBucketMapping(64)); + + assertThat(newMapping).isNotSameAs(oldMapping); + assertThat(PartitionBucketMapping.getCachedMapping(oldSnapshot)).isNull(); + assertThat(PartitionBucketMapping.getCachedMapping(newSnapshot)).isSameAs(newMapping); + PartitionBucketMapping.clearCache(); + } + + @Test + public void testGetOrLoadSingleFlightForConcurrentWriters() throws Exception { + PartitionBucketMapping.clearCache(); + PartitionBucketMapping.CacheKey key = + new PartitionBucketMapping.CacheKey("table", 1, 0, 32); + AtomicInteger loadCount = new AtomicInteger(); + CountDownLatch loaderStarted = new CountDownLatch(1); + CountDownLatch releaseLoader = new CountDownLatch(1); + ExecutorService executor = Executors.newFixedThreadPool(2); + try { + Future first = + executor.submit( + () -> + PartitionBucketMapping.getOrLoad( + key, + 128, + () -> { + loadCount.incrementAndGet(); + loaderStarted.countDown(); + await(releaseLoader); + return new PartitionBucketMapping(32); + })); + + loaderStarted.await(); + + Future second = + executor.submit( + () -> + PartitionBucketMapping.getOrLoad( + key, + 128, + () -> { + loadCount.incrementAndGet(); + return new PartitionBucketMapping(64); + })); + + releaseLoader.countDown(); + + PartitionBucketMapping firstMapping = first.get(); + PartitionBucketMapping secondMapping = second.get(); + assertThat(secondMapping).isSameAs(firstMapping); + assertThat(loadCount).hasValue(1); + } finally { + executor.shutdownNow(); + PartitionBucketMapping.clearCache(); + } + } + + private static void await(CountDownLatch latch) { + try { + latch.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException(e); + } + } + + private static BinaryRow partition(int value) { + return BinaryRow.singleColumn(value); + } + + private static SimpleFileEntry entry(BinaryRow partition, int bucket, int totalBuckets) { + return new SimpleFileEntry( + FileKind.ADD, + partition, + bucket, + totalBuckets, + 0, + "data-" + partition.hashCode() + "-" + bucket + ".parquet", + Collections.emptyList(), + null, + BinaryRow.EMPTY_ROW, + BinaryRow.EMPTY_ROW, + null, + 0L, + null); + } + + /** + * Builds a minimal {@link TableSchema} for a partitioned table with the given default bucket + * count. The schema declares a single partition column ("p") and a single value column ("v"). + */ + private static TableSchema partitionedSchema(int defaultBuckets) { + return schema(Collections.singletonList("p"), defaultBuckets); + } + + /** Builds a minimal {@link TableSchema} for a non-partitioned table. */ + private static TableSchema nonPartitionedSchema(int defaultBuckets) { + return schema(Collections.emptyList(), defaultBuckets); + } + + private static TableSchema schema(List partitionKeys, int defaultBuckets) { + List fields = + Arrays.asList( + new DataField(0, "p", new IntType()), new DataField(1, "v", new IntType())); + Map options = new HashMap<>(); + options.put(CoreOptions.BUCKET.key(), String.valueOf(defaultBuckets)); + return new TableSchema( + 0, + fields, + RowType.currentHighestFieldId(fields), + partitionKeys, + Collections.emptyList(), + options, + ""); + } +} diff --git a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/sink/FlinkSinkBuilder.java b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/sink/FlinkSinkBuilder.java index 8bea0b5acfbe..dbf8c552200f 100644 --- a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/sink/FlinkSinkBuilder.java +++ b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/sink/FlinkSinkBuilder.java @@ -36,6 +36,7 @@ import org.apache.paimon.table.BlobDescriptorReaderFactory; import org.apache.paimon.table.BucketMode; import org.apache.paimon.table.FileStoreTable; +import org.apache.paimon.table.OverwriteFileStoreTable; import org.apache.paimon.table.PostponeUtils; import org.apache.paimon.table.Table; import org.apache.paimon.table.sink.ChannelComputer; @@ -360,9 +361,15 @@ protected DataStreamSink buildForFixedBucket(DataStream input) { + " then the parallelism of writerOperator will be set to bucketNums."); parallelism = bucketNums; } + FileStoreTable sinkTable = + overwritePartition == null ? table : new OverwriteFileStoreTable(table); DataStream partitioned = - partition(input, new RowDataChannelComputer(table.schema()), parallelism); - return configureBlobDescriptorReaderFactory(new FixedBucketSink(table, overwritePartition)) + partition( + input, + new RowDataChannelComputer(sinkTable.createRowKeyExtractor()), + parallelism); + return configureBlobDescriptorReaderFactory( + new FixedBucketSink(sinkTable, overwritePartition)) .sinkFrom(partitioned); } diff --git a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/sink/RowDataChannelComputer.java b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/sink/RowDataChannelComputer.java index 1df93c82bcb1..5234333f87ff 100644 --- a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/sink/RowDataChannelComputer.java +++ b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/sink/RowDataChannelComputer.java @@ -24,25 +24,28 @@ import org.apache.paimon.table.sink.ChannelComputer; import org.apache.paimon.table.sink.FixedBucketRowKeyExtractor; import org.apache.paimon.table.sink.KeyAndBucketExtractor; +import org.apache.paimon.table.sink.RowKeyExtractor; /** {@link ChannelComputer} for {@link InternalRow}. */ public class RowDataChannelComputer implements ChannelComputer { private static final long serialVersionUID = 1L; - private final TableSchema schema; + private final KeyAndBucketExtractor extractor; private transient int numChannels; - private transient KeyAndBucketExtractor extractor; public RowDataChannelComputer(TableSchema schema) { - this.schema = schema; + this(new FixedBucketRowKeyExtractor(schema)); + } + + public RowDataChannelComputer(RowKeyExtractor extractor) { + this.extractor = extractor; } @Override public void setup(int numChannels) { this.numChannels = numChannels; - this.extractor = new FixedBucketRowKeyExtractor(schema); } @Override