[core] Allow custom primary-key compaction rewriters

docs/docs/program-api/java-writing.md

@@ -180,3 +180,32 @@ selector API: they require dedicated bucket assignment and `write(row, bucket)`
 
 For a Flink job, use [FlinkSinkBuilder](flink-api#write-to-table) to integrate routing, checkpoints,
 and commits with the engine.
+
+## Custom Primary-Key Compaction Rewriters
+
+Applications can install a `CompactRewriterFactory` on a table writer to replace or wrap the
+file-rewrite work for each primary-key partition and bucket. Paimon continues to select compaction
+inputs, schedule work, and collect results for checkpoint commits. The factory receives the normal
+rewriter selected for the table's merge engine, changelog producer, and deletion-vector options,
+so an implementation can delegate unsupported operations to it.
+
+```java
+write.withCompactRewriterFactory((partition, bucket, defaultRewriter) -> {
+    // Return a custom CompactRewriter here, or retain Paimon's implementation.
+    return defaultRewriter;
+});
+```
+
+Configure the factory before writing, restoring, or compacting any bucket. Install it again on each
+recovered writer; the factory and rewriters are not checkpoint state. The callback receives an
+independent partition copy and is invoked for each newly opened or restored bucket writer.
+
+A custom rewriter implements `rewrite(outputLevel, dropDelete, sections)` and
+`upgrade(outputLevel, file)`, returning `CompactResult` file changes. It must preserve Paimon's
+merge, sequence, changelog, deletion-vector, record-expiration, and metadata contracts. The returned rewriter owns the
+default rewriter and must close it when closed, even if it handles every operation itself. Paimon
+closes the default rewriter if factory creation fails or returns null.
+
+This hook supports primary-key merge-tree writers. Append, postpone, and primary-key clustering
+writers reject it. With `write-only = true`, the factory is never invoked. Installing a factory
+after a bucket writer has been created is rejected.

paimon-core/src/main/java/org/apache/paimon/mergetree/compact/CompactRewriterFactory.java

@@ -0,0 +1,46 @@
+/*
+ * 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.mergetree.compact;
+
+import org.apache.paimon.data.BinaryRow;
+
+/**
+ * Creates a compaction rewriter for one primary-key partition and bucket.
+ *
+ * <p>The supplied rewriter is Paimon's implementation selected for the table's merge engine,
+ * changelog producer, and deletion-vector options. A factory may return it unchanged, or wrap it to
+ * delegate compactions that its implementation does not support. A replacement must preserve the
+ * same records, sequence numbers, changelogs, deletion vectors, record expiration, and file
+ * metadata contracts.
+ *
+ * <p>After a successful call, the returned rewriter owns the supplied rewriter and must close it
+ * when closed. If creation fails or returns null, Paimon closes the supplied rewriter. The factory
+ * must release any other resources it allocated before failing. Each call must return a rewriter
+ * owned exclusively by that bucket; it is closed by Paimon's compaction manager.
+ */
+@FunctionalInterface
+public interface CompactRewriterFactory {
+
+    /**
+     * Creates a rewriter before the bucket starts compacting. The partition is an independent copy
+     * that may be retained. Capture table schema, options, and file access in the factory as
+     * needed.
+     */
+    CompactRewriter create(BinaryRow partition, int bucket, CompactRewriter defaultRewriter);
+}

paimon-core/src/main/java/org/apache/paimon/mergetree/compact/KvCompactionManagerFactory.java

@@ -97,6 +97,10 @@ static KvCompactionManagerFactory create(
 
     void withCompactionMetrics(@Nullable CompactionMetrics compactionMetrics);
 
+    default void withCompactRewriterFactory(CompactRewriterFactory factory) {
+        throw new UnsupportedOperationException("Custom compaction rewriters are not supported.");
+    }
+
     /** Create a {@link CompactManager} for the given partition and bucket. */
     CompactManager create(
             BinaryRow partition,

paimon-core/src/main/java/org/apache/paimon/mergetree/compact/MergeTreeCompactManagerFactory.java

@@ -22,6 +22,7 @@
 import org.apache.paimon.CoreOptions.ChangelogProducer;
 import org.apache.paimon.CoreOptions.MergeEngine;
 import org.apache.paimon.KeyValue;
+import org.apache.paimon.annotation.VisibleForTesting;
 import org.apache.paimon.codegen.RecordEqualiser;
 import org.apache.paimon.compact.CompactManager;
 import org.apache.paimon.compact.NoopCompactManager;
@@ -74,6 +75,8 @@
 import static org.apache.paimon.CoreOptions.MergeEngine.DEDUPLICATE;
 import static org.apache.paimon.lookup.LookupStoreFactory.bloomFilterBuilderFactory;
 import static org.apache.paimon.mergetree.LookupFile.localFilePrefix;
+import static org.apache.paimon.utils.Preconditions.checkNotNull;
+import static org.apache.paimon.utils.Preconditions.checkState;
 
 /** Factory to create {@link MergeTreeCompactManager}. */
 public class MergeTreeCompactManagerFactory implements KvCompactionManagerFactory {
@@ -97,6 +100,8 @@ public class MergeTreeCompactManagerFactory implements KvCompactionManagerFactor
     @Nullable private IOManager ioManager;
     @Nullable private CompactionMetrics compactionMetrics;
     @Nullable private Cache<String, LookupFile> lookupFileCache;
+    @Nullable private CompactRewriterFactory compactRewriterFactory;
+    private boolean initialized;
 
     public MergeTreeCompactManagerFactory(
             KeyValueFileReaderFactory.Builder readerFactoryBuilder,
@@ -141,6 +146,14 @@ public void withCompactionMetrics(@Nullable CompactionMetrics compactionMetrics)
         this.compactionMetrics = compactionMetrics;
     }
 
+    @Override
+    public void withCompactRewriterFactory(CompactRewriterFactory factory) {
+        checkState(
+                !initialized,
+                "Configure the compaction rewriter factory before creating bucket writers.");
+        this.compactRewriterFactory = checkNotNull(factory);
+    }
+
     @Override
     public CompactManager create(
             BinaryRow partition,
@@ -149,6 +162,7 @@ public CompactManager create(
             List<DataFileMeta> restoreFiles,
             @Nullable BucketedDvMaintainer dvMaintainer,
             boolean ignorePreviousFiles) {
+        initialized = true;
         if (options.writeOnly()) {
             return new NoopCompactManager();
         }
@@ -157,7 +171,7 @@ public CompactManager create(
         Comparator<InternalRow> keyComparator = keyComparatorSupplier.get();
         Levels levels = new Levels(keyComparator, restoreFiles, options.numLevels());
         @Nullable FieldsComparator userDefinedSeqComparator = udsComparatorSupplier.get();
-        MergeTreeCompactRewriter rewriter =
+        MergeTreeCompactRewriter defaultRewriter =
                 createRewriter(
                         partition,
                         bucket,
@@ -166,12 +180,14 @@ public CompactManager create(
                         levels,
                         dvMaintainer,
                         ignorePreviousFiles);
+        CompactRewriter rewriter =
+                wrapRewriter(compactRewriterFactory, partition, bucket, defaultRewriter);
         CompactionMetrics.Reporter metricsReporter =
                 compactionMetrics == null
                         ? null
                         : compactionMetrics.createReporter(partition, bucket);
         if (metricsReporter != null) {
-            rewriter.setMetricsReporter(metricsReporter);
+            defaultRewriter.setMetricsReporter(metricsReporter);
         }
         String bucketInfo = "bucket=" + bucket;
         if (partition.getFieldCount() > 0) {
@@ -250,6 +266,31 @@ private static Long estimateLastFullCompactionTime(
         return max < 0 ? null : max;
     }
 
+    @VisibleForTesting
+    static CompactRewriter wrapRewriter(
+            @Nullable CompactRewriterFactory compactRewriterFactory,
+            BinaryRow partition,
+            int bucket,
+            CompactRewriter defaultRewriter) {
+        if (compactRewriterFactory == null) {
+            return defaultRewriter;
+        }
+        try {
+            return checkNotNull(
+                    compactRewriterFactory.create(partition.copy(), bucket, defaultRewriter),
+                    "The compaction rewriter factory must return a rewriter.");
+        } catch (RuntimeException | Error failure) {
+            try {
+                defaultRewriter.close();
+            } catch (Exception closeFailure) {
+                if (closeFailure != failure) {
+                    failure.addSuppressed(closeFailure);
+                }
+            }
+            throw failure;
+        }
+    }
+
     private MergeTreeCompactRewriter createRewriter(
             BinaryRow partition,
             int bucket,

paimon-core/src/main/java/org/apache/paimon/operation/FileStoreWrite.java

@@ -27,6 +27,7 @@
 import org.apache.paimon.index.pk.BucketedPrimaryKeyIndexMaintainer;
 import org.apache.paimon.io.DataFileMeta;
 import org.apache.paimon.memory.MemoryPoolFactory;
+import org.apache.paimon.mergetree.compact.CompactRewriterFactory;
 import org.apache.paimon.metrics.MetricRegistry;
 import org.apache.paimon.table.sink.CommitMessage;
 import org.apache.paimon.table.sink.SinkRecord;
@@ -86,6 +87,11 @@ default void withWriteType(RowType writeType) {
 
     void withCompactExecutor(ExecutorService compactExecutor);
 
+    /** Installs a compaction rewriter factory before any bucket writer is created. */
+    default FileStoreWrite<T> withCompactRewriterFactory(CompactRewriterFactory factory) {
+        throw new UnsupportedOperationException("Custom compaction rewriters are not supported.");
+    }
+
     /**
      * Write the data to the store according to the partition and bucket.
      *

paimon-core/src/main/java/org/apache/paimon/operation/KeyValueFileStoreWrite.java

@@ -37,6 +37,7 @@
 import org.apache.paimon.io.KeyValueFileWriterFactory;
 import org.apache.paimon.io.RecordLevelExpire;
 import org.apache.paimon.mergetree.MergeTreeWriter;
+import org.apache.paimon.mergetree.compact.CompactRewriterFactory;
 import org.apache.paimon.mergetree.compact.KvCompactionManagerFactory;
 import org.apache.paimon.mergetree.compact.LookupMergeFunction;
 import org.apache.paimon.mergetree.compact.MergeFunctionFactory;
@@ -180,6 +181,12 @@ protected boolean ignorePreviousFilesForWriter(
         return ignorePreviousFiles;
     }
 
+    @Override
+    public KeyValueFileStoreWrite withCompactRewriterFactory(CompactRewriterFactory factory) {
+        compactManagerFactory.withCompactRewriterFactory(factory);
+        return this;
+    }
+
     @Override
     public KeyValueFileStoreWrite withIOManager(IOManager ioManager) {
         super.withIOManager(ioManager);

paimon-core/src/main/java/org/apache/paimon/table/sink/TableWrite.java

@@ -25,6 +25,7 @@
 import org.apache.paimon.disk.IOManager;
 import org.apache.paimon.io.BundleRecords;
 import org.apache.paimon.memory.MemoryPoolFactory;
+import org.apache.paimon.mergetree.compact.CompactRewriterFactory;
 import org.apache.paimon.metrics.MetricRegistry;
 import org.apache.paimon.table.Table;
 import org.apache.paimon.types.RowType;
@@ -53,6 +54,16 @@ public interface TableWrite extends AutoCloseable {
      */
     TableWrite withBlobConsumer(BlobConsumer blobConsumer);
 
+    /**
+     * Installs a rewriter factory for primary-key merge-tree compaction. Configure this before
+     * writing, restoring, or compacting any bucket, and configure it again on each recovered
+     * writer. Paimon retains compaction scheduling and commit coordination. Append and clustering
+     * writers do not support this hook; write-only writers never invoke it.
+     */
+    default TableWrite withCompactRewriterFactory(CompactRewriterFactory factory) {
+        throw new UnsupportedOperationException("Custom compaction rewriters are not supported.");
+    }
+
     /** Calculate which partition {@code row} belongs to. */
     BinaryRow getPartition(InternalRow row);
 

paimon-core/src/main/java/org/apache/paimon/table/sink/TableWriteImpl.java

@@ -27,6 +27,7 @@
 import org.apache.paimon.io.BundleRecords;
 import org.apache.paimon.io.DataFileMeta;
 import org.apache.paimon.memory.MemoryPoolFactory;
+import org.apache.paimon.mergetree.compact.CompactRewriterFactory;
 import org.apache.paimon.metrics.MetricRegistry;
 import org.apache.paimon.operation.BundleFileStoreWriter;
 import org.apache.paimon.operation.FileStoreWrite;
@@ -135,6 +136,12 @@ public TableWrite withBlobConsumer(BlobConsumer blobConsumer) {
         return this;
     }
 
+    @Override
+    public TableWriteImpl<T> withCompactRewriterFactory(CompactRewriterFactory factory) {
+        write.withCompactRewriterFactory(factory);
+        return this;
+    }
+
     public TableWriteImpl<T> withCompactExecutor(ExecutorService compactExecutor) {
         write.withCompactExecutor(compactExecutor);
         return this;

paimon-core/src/test/java/org/apache/paimon/mergetree/compact/MergeTreeCompactManagerFactoryTest.java

@@ -36,15 +36,20 @@
 import org.apache.paimon.types.RowType;
 
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
 
+import java.io.IOException;
 import java.util.Collections;
 import java.util.Comparator;
 import java.util.concurrent.ExecutorService;
 
 import static org.apache.paimon.CoreOptions.DELETION_VECTORS_ENABLED;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
 import static org.mockito.Answers.RETURNS_SELF;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
@@ -61,6 +66,46 @@ public class MergeTreeCompactManagerFactoryTest {
                     DataTypes.FIELD(0, "key", DataTypes.INT()),
                     DataTypes.FIELD(1, "value", DataTypes.INT()));
 
+    @ParameterizedTest
+    @ValueSource(booleans = {false, true})
+    public void testFailedFactoryClosesDefaultRewriter(boolean returnNull) throws Exception {
+        CompactRewriter delegate = mock(CompactRewriter.class);
+        CompactRewriterFactory factory =
+                (partition, bucket, rewriter) -> {
+                    if (returnNull) {
+                        return null;
+                    }
+                    throw new IllegalStateException("factory failed");
+                };
+        assertThatThrownBy(
+                        () ->
+                                MergeTreeCompactManagerFactory.wrapRewriter(
+                                        factory, BinaryRow.EMPTY_ROW, 0, delegate))
+                .isInstanceOf(returnNull ? NullPointerException.class : IllegalStateException.class)
+                .hasMessageContaining(returnNull ? "must return a rewriter" : "factory failed");
+        verify(delegate).close();
+    }
+
+    @Test
+    public void testCloseFailureDoesNotReplaceFactoryFailure() throws Exception {
+        CompactRewriter delegate = mock(CompactRewriter.class);
+        IOException closeFailure = new IOException("close failed");
+        doThrow(closeFailure).when(delegate).close();
+        IllegalStateException failure = new IllegalStateException("factory failed");
+        assertThatThrownBy(
+                        () ->
+                                MergeTreeCompactManagerFactory.wrapRewriter(
+                                        (partition, bucket, rewriter) -> {
+                                            throw failure;
+                                        },
+                                        BinaryRow.EMPTY_ROW,
+                                        0,
+                                        delegate))
+                .isSameAs(failure)
+                .hasSuppressedException(closeFailure);
+        verify(delegate).close();
+    }
+
     @Test
     public void testLookupValueProjection() throws Exception {
         Options options = new Options();