[WIP][core] Cache manifest entry in TM

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<Boolean> 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<Integer> 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<Boolean> 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);
     }

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> statistics() {
     public Optional<WriteSelector> 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());

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<InternalRow> newWrite(String commitUser) {
 
     @Override
     public TableWriteImpl<InternalRow> newWrite(String commitUser, @Nullable Integer writeId) {
+        return newWrite(commitUser, writeId, createRowKeyExtractor());
+    }
+
+    @Override
+    public TableWriteImpl<InternalRow> 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(),

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) {

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) {

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<WriteSelector> 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<String, String> dynamicOptions) {
+        return new OverwriteFileStoreTable(wrapped().copy(dynamicOptions));
+    }
+
+    @Override
+    public FileStoreTable copy(TableSchema newTableSchema) {
+        return new OverwriteFileStoreTable(wrapped().copy(newTableSchema));
+    }
+
+    @Override
+    public FileStoreTable copyWithoutTimeTravel(Map<String, String> 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));
+    }
+}

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<KeyValue> newWrite(String commitUser) {
 
     @Override
     public TableWriteImpl<KeyValue> newWrite(String commitUser, @Nullable Integer writeId) {
-        return newWrite(store().newWrite(commitUser, writeId));
+        return newWrite(commitUser, writeId, createRowKeyExtractor());
+    }
+
+    @Override
+    public TableWriteImpl<KeyValue> newWrite(
+            String commitUser, @Nullable Integer writeId, RowKeyExtractor rowKeyExtractor) {
+        return newWrite(store().newWrite(commitUser, writeId), rowKeyExtractor);
     }
 
     @Override
     public TableWriteImpl<KeyValue> newPostponeFixedBucketWrite(
             String commitUser, @Nullable Integer writeId) {
-        return newWrite(store().newPostponeFixedBucketWrite(commitUser));
+        return newWrite(store().newPostponeFixedBucketWrite(commitUser), createRowKeyExtractor());
     }
 
-    private TableWriteImpl<KeyValue> newWrite(AbstractFileStoreWrite<KeyValue> storeWrite) {
+    private TableWriteImpl<KeyValue> newWrite(
+            AbstractFileStoreWrite<KeyValue> storeWrite, RowKeyExtractor rowKeyExtractor) {
         KeyValue kv = new KeyValue();
         return new TableWriteImpl<>(
                 rowType(),
                 storeWrite,
-                createRowKeyExtractor(),
+                rowKeyExtractor,
                 (record, rowKind) ->
                         kv.replace(
                                 record.primaryKey(),

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,20 +64,29 @@ private BinaryRow bucketKey() {
         }
 
         if (reuseBucketKey == null) {
-            reuseBucketKey = bucketKeyProjection.apply(record);
+            reuseBucketKey = bucketKeyProjection().apply(record);
         }
         return reuseBucketKey;
     }
 
     @Override
     public int bucket() {
         if (reuseBucket == null) {
-            reuseBucket = bucket(numBuckets);
+            reuseBucket = bucket(partitionBucketMapping.resolveNumBuckets(partition()));
         }
         return reuseBucket;
     }
 
     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;
+    }
 }

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<InternalRow> 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);