apache · sakshichitnis27 · Sep 6, 2026
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
@@ -661,6 +661,9 @@ ListOffsetsResult listOffsets(
      * <p>More details, Fluss collects the cluster's load information and optimizes to perform load
      * balancing according to the user-defined {@code priorityGoals}.
      *
+     * <p>{@link GoalType#PREFERRED_LEADER_ELECTION} must be requested as a standalone goal. It
+     * changes only bucket leadership and leaves replica assignments unchanged.
+     *
      * <p>Currently, Fluss only supports one active rebalance task in the cluster. If an uncompleted
      * rebalance task exists, Fluss will return the uncompleted rebalance task's progress.
      *

diff --git a/fluss-client/src/test/java/org/apache/fluss/client/admin/RebalanceITCase.java b/fluss-client/src/test/java/org/apache/fluss/client/admin/RebalanceITCase.java
@@ -29,19 +29,23 @@
 import org.apache.fluss.exception.RebalanceFailureException;
 import org.apache.fluss.metadata.DatabaseDescriptor;
 import org.apache.fluss.metadata.PartitionSpec;
+import org.apache.fluss.metadata.TableBucket;
 import org.apache.fluss.metadata.TableDescriptor;
 import org.apache.fluss.metadata.TablePath;
 import org.apache.fluss.server.replica.ReplicaManager;
 import org.apache.fluss.server.testutils.FlussClusterExtension;
+import org.apache.fluss.server.zk.ZooKeeperClient;
 
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
 
 import java.time.Duration;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.List;
 import java.util.Optional;
 
 import static org.apache.fluss.record.TestData.DATA1_SCHEMA;
@@ -209,6 +213,128 @@ void testRebalanceForLogTable() throws Exception {
         admin.removeServerTag(Collections.singletonList(0), ServerTag.PERMANENT_OFFLINE).get();
     }
 
+    @Test
+    void testPreferredLeaderElectionAfterRecovery() throws Exception {
+        String dbName = "db-preferred-leader";
+        admin.createDatabase(dbName, DatabaseDescriptor.EMPTY, false).get();
+        long tableId =
+                createTable(
+                        new TablePath(dbName, "preferred-leader-table"), DATA1_TABLE_DESCRIPTOR);
+        FLUSS_CLUSTER_EXTENSION.waitUntilTableReady(tableId);
+
+        TableBucket tableBucket = new TableBucket(tableId, 0);
+        ZooKeeperClient zkClient = FLUSS_CLUSTER_EXTENSION.getZooKeeperClient();
+        List<Integer> assignment =
+                new ArrayList<>(
+                        zkClient.getTableAssignment(tableId)
+                                .get()
+                                .getBucketAssignment(tableBucket.getBucket())
+                                .getReplicas());
+        int preferredLeader = assignment.get(0);
+        assertThat(FLUSS_CLUSTER_EXTENSION.waitAndGetLeader(tableBucket))
+                .isEqualTo(preferredLeader);
+
+        boolean preferredLeaderStopped = false;
+        boolean preferredLeaderTagged = false;
+        try {
+            FLUSS_CLUSTER_EXTENSION.stopTabletServer(preferredLeader);
+            preferredLeaderStopped = true;
+            FLUSS_CLUSTER_EXTENSION.waitUntilReplicaShrinkFromIsr(tableBucket, preferredLeader);
+            retry(
+                    Duration.ofMinutes(1),
+                    () ->
+                            assertThat(FLUSS_CLUSTER_EXTENSION.waitAndGetLeader(tableBucket))
+                                    .isNotEqualTo(preferredLeader));
+            int failoverLeader = FLUSS_CLUSTER_EXTENSION.waitAndGetLeader(tableBucket);
+
+            String unavailableRebalanceId =
+                    admin.rebalance(Collections.singletonList(GoalType.PREFERRED_LEADER_ELECTION))
+                            .get();
+            waitUntilRebalanceCompletes(unavailableRebalanceId);
+            assertThat(FLUSS_CLUSTER_EXTENSION.waitAndGetLeader(tableBucket))
+                    .isEqualTo(failoverLeader);
+            assertThat(
+                            zkClient.getTableAssignment(tableId)
+                                    .get()
+                                    .getBucketAssignment(tableBucket.getBucket())
+                                    .getReplicas())
+                    .containsExactlyElementsOf(assignment);
+
+            FLUSS_CLUSTER_EXTENSION.startTabletServer(preferredLeader);
+            preferredLeaderStopped = false;
+            FLUSS_CLUSTER_EXTENSION.waitUntilReplicaExpandToIsr(tableBucket, preferredLeader);
+
+            admin.addServerTag(
+                            Collections.singletonList(preferredLeader), ServerTag.TEMPORARY_OFFLINE)
+                    .get();
+            preferredLeaderTagged = true;
+            String taggedRebalanceId =
+                    admin.rebalance(Collections.singletonList(GoalType.PREFERRED_LEADER_ELECTION))
+                            .get();
+            waitUntilRebalanceCompletes(taggedRebalanceId);
+            assertThat(FLUSS_CLUSTER_EXTENSION.waitAndGetLeader(tableBucket))
+                    .isEqualTo(failoverLeader);
+
+            admin.removeServerTag(
+                            Collections.singletonList(preferredLeader), ServerTag.TEMPORARY_OFFLINE)
+                    .get();
+            preferredLeaderTagged = false;
+
+            String rebalanceId =
+                    admin.rebalance(Collections.singletonList(GoalType.PREFERRED_LEADER_ELECTION))
+                            .get();
+            waitUntilRebalanceCompletes(rebalanceId);
+            assertThat(FLUSS_CLUSTER_EXTENSION.waitAndGetLeader(tableBucket))
+                    .isEqualTo(preferredLeader);
+            assertThat(
+                            zkClient.getTableAssignment(tableId)
+                                    .get()
+                                    .getBucketAssignment(tableBucket.getBucket())
+                                    .getReplicas())
+                    .containsExactlyElementsOf(assignment);
+
+            int leaderEpoch = zkClient.getLeaderAndIsr(tableBucket).get().leaderEpoch();
+            String idempotentRebalanceId =
+                    admin.rebalance(Collections.singletonList(GoalType.PREFERRED_LEADER_ELECTION))
+                            .get();
+            waitUntilRebalanceCompletes(idempotentRebalanceId);
+            assertThat(zkClient.getLeaderAndIsr(tableBucket).get().leaderEpoch())
+                    .isEqualTo(leaderEpoch);
+
+            assertThatThrownBy(
+                            () ->
+                                    admin.rebalance(
+                                                    Arrays.asList(
+                                                            GoalType.PREFERRED_LEADER_ELECTION,
+                                                            GoalType.LEADER_DISTRIBUTION))
+                                            .get())
+                    .rootCause()
+                    .isInstanceOf(RebalanceFailureException.class)
+                    .hasMessageContaining("must be used as a standalone rebalance goal");
+        } finally {
+            if (preferredLeaderTagged) {
+                admin.removeServerTag(
+                                Collections.singletonList(preferredLeader),
+                                ServerTag.TEMPORARY_OFFLINE)
+                        .get();
+            }
+            if (preferredLeaderStopped) {
+                FLUSS_CLUSTER_EXTENSION.startTabletServer(preferredLeader);
+            }
+        }
+    }
+
+    private void waitUntilRebalanceCompletes(String rebalanceId) {
+        retry(
+                Duration.ofMinutes(2),
+                () -> {
+                    Optional<RebalanceProgress> progress =
+                            admin.listRebalanceProgress(rebalanceId).get();
+                    assertThat(progress).isPresent();
+                    assertThat(progress.get().status()).isEqualTo(RebalanceStatus.COMPLETED);
+                });
+    }
+
     @Test
     void testListRebalanceProgress() throws Exception {
         String dbName = "db-rebalance-list";

diff --git a/fluss-common/src/main/java/org/apache/fluss/cluster/rebalance/GoalType.java b/fluss-common/src/main/java/org/apache/fluss/cluster/rebalance/GoalType.java
@@ -44,7 +44,13 @@ public enum GoalType {
      * Goal to generate replica movement tasks to ensure that the number of replicas on each
      * tabletServer is near balanced and the replicas are distributed across racks.
      */
-    RACK_AWARE(2);
+    RACK_AWARE(2),
+
+    /**
+     * Goal to move leadership to the first replica in each persisted bucket assignment without
+     * changing replica assignments.
+     */
+    PREFERRED_LEADER_ELECTION(3);
 
     public final int value;
 
@@ -59,6 +65,8 @@ public static GoalType valueOf(int value) {
             return LEADER_DISTRIBUTION;
         } else if (value == RACK_AWARE.value) {
             return RACK_AWARE;
+        } else if (value == PREFERRED_LEADER_ELECTION.value) {
+            return PREFERRED_LEADER_ELECTION;
         } else {
             throw new IllegalArgumentException(
                     String.format(

diff --git a/fluss-common/src/test/java/org/apache/fluss/cluster/rebalance/GoalTypeTest.java b/fluss-common/src/test/java/org/apache/fluss/cluster/rebalance/GoalTypeTest.java
@@ -0,0 +1,33 @@
+/*
+ * 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.cluster.rebalance;
+
+import org.junit.jupiter.api.Test;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/** Test for {@link GoalType}. */
+class GoalTypeTest {
+
+    @Test
+    void testPreferredLeaderElectionFromName() {
+        assertThat(GoalType.fromName("preferred_leader_election"))
+                .isEqualTo(GoalType.PREFERRED_LEADER_ELECTION);
+        assertThat(GoalType.valueOf(3)).isEqualTo(GoalType.PREFERRED_LEADER_ELECTION);
+    }
+}
diff --git a/...fluss-flink-common/src/main/java/org/apache/fluss/flink/procedure/RebalanceProcedure.java b/...fluss-flink-common/src/main/java/org/apache/fluss/flink/procedure/RebalanceProcedure.java
@@ -41,6 +41,8 @@
  * CALL sys.rebalance('REPLICA_DISTRIBUTION');
  * -- Trigger rebalance with REPLICA_DISTRIBUTION and LEADER_DISTRIBUTION goals
  * CALL sys.rebalance('REPLICA_DISTRIBUTION,LEADER_DISTRIBUTION');
+ * -- Restore leadership to the first replica in each assignment
+ * CALL sys.rebalance('PREFERRED_LEADER_ELECTION');
  * </pre>
  */
 public class RebalanceProcedure extends ProcedureBase {

diff --git a/...uss-flink-common/src/test/java/org/apache/fluss/flink/procedure/FlinkProcedureITCase.java b/...uss-flink-common/src/test/java/org/apache/fluss/flink/procedure/FlinkProcedureITCase.java
@@ -27,6 +27,7 @@
 import org.apache.fluss.config.Configuration;
 import org.apache.fluss.config.MemorySize;
 import org.apache.fluss.exception.NoRebalanceInProgressException;
+import org.apache.fluss.exception.RebalanceFailureException;
 import org.apache.fluss.exception.SecurityDisabledException;
 import org.apache.fluss.metadata.DataLakeFormat;
 import org.apache.fluss.metadata.TablePath;
@@ -743,6 +744,40 @@ void testRebalance(boolean upperCase) throws Exception {
         FLUSS_CLUSTER_EXTENSION.getZooKeeperClient().deleteRebalanceTask();
     }
 
+    @Test
+    void testPreferredLeaderElectionProcedure() throws Exception {
+        try {
+            String rebalance =
+                    String.format(
+                            "Call %s.sys.rebalance('PREFERRED_LEADER_ELECTION')", CATALOG_NAME);
+            try (CloseableIterator<Row> rows = tEnv.executeSql(rebalance).collect()) {
+                assertThat(CollectionUtil.iteratorToList(rows)).hasSize(1);
+            }
+
+            retry(
+                    Duration.ofMinutes(2),
+                    () -> {
+                        Optional<RebalanceProgress> progress =
+                                admin.listRebalanceProgress(null).get();
+                        assertThat(progress).isPresent();
+                        assertThat(progress.get().status()).isEqualTo(RebalanceStatus.COMPLETED);
+                    });
+
+            assertThatThrownBy(
+                            () ->
+                                    tEnv.executeSql(
+                                                    String.format(
+                                                            "Call %s.sys.rebalance('PREFERRED_LEADER_ELECTION,LEADER_DISTRIBUTION')",
+                                                            CATALOG_NAME))
+                                            .await())
+                    .rootCause()
+                    .isInstanceOf(RebalanceFailureException.class)
+                    .hasMessageContaining("must be used as a standalone rebalance goal");
+        } finally {
+            FLUSS_CLUSTER_EXTENSION.getZooKeeperClient().deleteRebalanceTask();
+        }
+    }
+
     @Test
     void testListRebalanceProgress() throws Exception {
         // add server tag PERMANENT_OFFLINE for server 3, this will avoid to generate bucket

diff --git a/fluss-rust/crates/fluss/src/metadata/goal_type.rs b/fluss-rust/crates/fluss/src/metadata/goal_type.rs
@@ -23,6 +23,7 @@ pub enum GoalType {
     ReplicaDistribution,
     LeaderDistribution,
     RackAware,
+    PreferredLeaderElection,
 }
 
 impl GoalType {
@@ -31,6 +32,7 @@ impl GoalType {
             Self::ReplicaDistribution => 0,
             Self::LeaderDistribution => 1,
             Self::RackAware => 2,
+            Self::PreferredLeaderElection => 3,
         }
     }
 
@@ -39,6 +41,7 @@ impl GoalType {
             0 => Ok(Self::ReplicaDistribution),
             1 => Ok(Self::LeaderDistribution),
             2 => Ok(Self::RackAware),
+            3 => Ok(Self::PreferredLeaderElection),
             _ => Err(Error::IllegalArgument {
                 message: format!("Unsupported GoalType: {value}"),
             }),
@@ -56,6 +59,7 @@ mod tests {
             GoalType::ReplicaDistribution,
             GoalType::LeaderDistribution,
             GoalType::RackAware,
+            GoalType::PreferredLeaderElection,
         ] {
             assert_eq!(GoalType::try_from_i32(goal.to_i32()).unwrap(), goal);
         }

diff --git a/...s-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessor.java b/...s-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessor.java
@@ -1605,6 +1605,23 @@ public void tryToExecuteRebalanceTask(RebalancePlanForBucket planForBucket) {
                         coordinatorContext.getAssignment(tableBucket), newReplicas);
 
         if (planForBucket.isLeaderChanged() && !reassignment.isBeingReassigned()) {
+            int targetLeader = planForBucket.getNewLeader();
+            Optional<LeaderAndIsr> leaderAndIsr =
+                    coordinatorContext.getBucketLeaderAndIsr(tableBucket);
+            Optional<ServerTag> targetServerTag = coordinatorContext.getServerTag(targetLeader);
+            if (!leaderAndIsr.isPresent()
+                    || !leaderAndIsr.get().isr().contains(targetLeader)
+                    || !coordinatorContext.isReplicaOnline(targetLeader, tableBucket)
+                    || (targetServerTag.isPresent()
+                            && (targetServerTag.get() == ServerTag.TEMPORARY_OFFLINE
+                                    || targetServerTag.get() == ServerTag.PERMANENT_OFFLINE))) {
+                LOG.warn(
+                        "Skipping leader-only rebalance for tableBucket {} because target leader {} is no longer eligible.",
+                        tableBucket,
+                        targetLeader);
+                rebalanceManager.finishRebalanceTask(tableBucket, RebalanceStatus.FAILED);
+                return;
+            }
             // buckets only need to change leader like leader replica rebalance.
             // Don't finish the task immediately; wait for the NotifyLeaderAndIsr response
             // from the tablet server to confirm the leader change has been applied.
@@ -1614,7 +1631,7 @@ public void tryToExecuteRebalanceTask(RebalancePlanForBucket planForBucket) {
             tableBucketStateMachine.handleStateChange(
                     Collections.singleton(tableBucket),
                     OnlineBucket,
-                    new ReassignmentLeaderElection(newReplicas));
+                    new ReassignmentLeaderElection(newReplicas, false));
         } else {
             try {
                 LOG.info(