[flink] Add default parallelism for source and sink in rescale procedure, also make them configurable

Author: tsreaper

docs/content/flink/procedures.md

@@ -757,13 +757,15 @@ All available procedures are listed below.
    <tr>
       <td>rescale</td>
       <td>
-         CALL [catalog.]sys.rescale(`table` => 'identifier', `bucket_num` => bucket_num, `partition` => 'partition')
+         CALL [catalog.]sys.rescale(`table` => 'identifier', `bucket_num` => bucket_num, `partition` => 'partition', `source_parallelism` => 'source_parallelism', `sink_parallelism` => 'sink_parallelism')
       </td>
       <td>
          Rescale one partition of a table. Arguments:
          <li>identifier: The target table identifier. Cannot be empty.</li>
          <li>bucket_num: Resulting bucket number after rescale. The default value of argument bucket_num is the current bucket number of the table. Cannot be empty for postpone bucket tables.</li>
          <li>partition: What partition to rescale. For partitioned table this argument cannot be empty.</li>
+         <li>source_parallelism: Parallelism of source operator. The default value is the current bucket number of the partition.</li>
+         <li>partition: Parallelism of sink operator. The default value is equal to bucket_num.</li>
       </td>
       <td>
          CALL sys.rescale(`table` => 'default.T', `bucket_num` => 16, `partition` => 'dt=20250217,hh=08')

paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/action/RescaleAction.java

@@ -21,6 +21,7 @@
 import org.apache.paimon.CoreOptions;
 import org.apache.paimon.flink.sink.FlinkSinkBuilder;
 import org.apache.paimon.flink.source.FlinkSourceBuilder;
+import org.apache.paimon.manifest.ManifestEntry;
 import org.apache.paimon.partition.PartitionPredicate;
 import org.apache.paimon.predicate.Predicate;
 import org.apache.paimon.table.BucketMode;
@@ -38,13 +39,16 @@
 import javax.annotation.Nullable;
 
 import java.util.HashMap;
+import java.util.Iterator;
 import java.util.Map;
 
 /** Action to rescale one partition of a table. */
 public class RescaleAction extends TableActionBase {
 
     private @Nullable Integer bucketNum;
     private Map<String, String> partition = new HashMap<>();
+    private @Nullable Integer sourceParallelism;
+    private @Nullable Integer sinkParallelism;
 
     public RescaleAction(String databaseName, String tableName, Map<String, String> catalogConfig) {
         super(databaseName, tableName, catalogConfig);
@@ -60,6 +64,16 @@ public RescaleAction withPartition(Map<String, String> partition) {
         return this;
     }
 
+    public RescaleAction withSourceParallelism(int sourceParallelism) {
+        this.sourceParallelism = sourceParallelism;
+        return this;
+    }
+
+    public RescaleAction withSinkParallelism(int sinkParallelism) {
+        this.sinkParallelism = sinkParallelism;
+        return this;
+    }
+
     @Override
     public void build() throws Exception {
         Configuration flinkConf = new Configuration();
@@ -79,6 +93,8 @@ public void build() throws Exception {
                 new FlinkSourceBuilder(fileStoreTable)
                         .env(env)
                         .sourceBounded(true)
+                        .sourceParallelism(
+                                sourceParallelism == null ? currentBucketNum() : sourceParallelism)
                         .predicate(partitionPredicate)
                         .build();
 
@@ -92,12 +108,29 @@ public void build() throws Exception {
         }
         FileStoreTable rescaledTable =
                 fileStoreTable.copy(fileStoreTable.schema().copy(bucketOptions));
-        new FlinkSinkBuilder(rescaledTable).overwrite(partition).forRowData(source).build();
+        new FlinkSinkBuilder(rescaledTable)
+                .overwrite(partition)
+                .parallelism(sinkParallelism == null ? bucketNum : sinkParallelism)
+                .forRowData(source)
+                .build();
     }
 
     @Override
     public void run() throws Exception {
         build();
         env.execute("Rescale Postpone Bucket : " + table.fullName());
     }
+
+    private int currentBucketNum() {
+        FileStoreTable fileStoreTable = (FileStoreTable) table;
+        Iterator<ManifestEntry> it =
+                fileStoreTable
+                        .newSnapshotReader()
+                        .withPartitionFilter(partition)
+                        .readFileIterator();
+        Preconditions.checkArgument(
+                it.hasNext(),
+                "The specified partition does not have any data files. No need to rescale.");
+        return it.next().totalBuckets();
+    }
 }

paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/action/RescaleActionFactory.java

@@ -26,6 +26,8 @@ public class RescaleActionFactory implements ActionFactory {
     public static final String IDENTIFIER = "rescale";
     private static final String BUCKET_NUM = "bucket_num";
     private static final String PARTITION = "partition";
+    private static final String SOURCE_PARALLELISM = "source_parallelism";
+    private static final String SINK_PARALLELISM = "sink_parallelism";
 
     @Override
     public String identifier() {
@@ -43,10 +45,15 @@ public Optional<Action> create(MultipleParameterToolAdapter params) {
         if (params.has(BUCKET_NUM)) {
             action.withBucketNum(Integer.parseInt(params.get(BUCKET_NUM)));
         }
-
         if (params.has(PARTITION)) {
             action.withPartition(getPartitions(params).get(0));
         }
+        if (params.has(SOURCE_PARALLELISM)) {
+            action.withSourceParallelism(Integer.parseInt(params.get(SOURCE_PARALLELISM)));
+        }
+        if (params.has(SINK_PARALLELISM)) {
+            action.withSinkParallelism(Integer.parseInt(params.get(SINK_PARALLELISM)));
+        }
 
         return Optional.of(action);
     }
@@ -60,11 +67,16 @@ public void printHelp() {
         System.out.println(
                 "  rescale --warehouse <warehouse_path> --database <database_name> "
                         + "--table <table_name> [--bucket_num <bucket_num>] "
-                        + "[--partition <partition>]");
+                        + "[--partition <partition>] "
+                        + "[--source_parallelism <source_parallelism>] [--sink_parallelism <sink_parallelism>]");
         System.out.println(
-                "The default value of argument bucket_num is the current bucket number of the table. "
+                "The default value of argument bucket_num is the value of 'bucket' option of the table. "
                         + "For postpone bucket tables, this argument must be specified.");
         System.out.println(
                 "Argument partition must be specified if the table is a partitioned table.");
+        System.out.println(
+                "The default value of argument source_parallelism is the current bucket number of the partition.");
+        System.out.println(
+                "The default value of argument sink_parallelism is equal to bucket_num.");
     }
 }

paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/procedure/RescaleProcedure.java

@@ -38,13 +38,26 @@ public class RescaleProcedure extends ProcedureBase {
             argument = {
                 @ArgumentHint(name = "table", type = @DataTypeHint("STRING")),
                 @ArgumentHint(name = "bucket_num", type = @DataTypeHint("INT"), isOptional = true),
-                @ArgumentHint(name = "partition", type = @DataTypeHint("STRING"), isOptional = true)
+                @ArgumentHint(
+                        name = "partition",
+                        type = @DataTypeHint("STRING"),
+                        isOptional = true),
+                @ArgumentHint(
+                        name = "source_parallelism",
+                        type = @DataTypeHint("INT"),
+                        isOptional = true),
+                @ArgumentHint(
+                        name = "sink_parallelism",
+                        type = @DataTypeHint("INT"),
+                        isOptional = true)
             })
     public String[] call(
             ProcedureContext procedureContext,
             String tableId,
             @Nullable Integer bucketNum,
-            @Nullable String partition)
+            @Nullable String partition,
+            @Nullable Integer sourceParallelism,
+            @Nullable Integer sinkParallelism)
             throws Exception {
         Identifier identifier = Identifier.fromString(tableId);
         String databaseName = identifier.getDatabaseName();
@@ -57,6 +70,12 @@ public String[] call(
         if (partition != null) {
             action.withPartition(ParameterUtils.getPartitions(partition).get(0));
         }
+        if (sourceParallelism != null) {
+            action.withSourceParallelism(sourceParallelism);
+        }
+        if (sinkParallelism != null) {
+            action.withSinkParallelism(sinkParallelism);
+        }
 
         return execute(
                 procedureContext, action, "Rescale Postpone Bucket : " + identifier.getFullName());