diff --git a/server/src/main/java/org/opensearch/cluster/routing/allocation/allocator/BalancedShardsAllocator.java b/server/src/main/java/org/opensearch/cluster/routing/allocation/allocator/BalancedShardsAllocator.java
index 785636fa7ff2a..641e9b8cdd4e2 100644
--- a/server/src/main/java/org/opensearch/cluster/routing/allocation/allocator/BalancedShardsAllocator.java
+++ b/server/src/main/java/org/opensearch/cluster/routing/allocation/allocator/BalancedShardsAllocator.java
@@ -294,7 +294,12 @@ private void updatePreferPrimaryShardBalanceBuffer(float preferPrimaryShardBalan
     }
 
     private void updateWeightFunction() {
-        weightFunction = new WeightFunction(this.indexBalanceFactor, this.shardBalanceFactor, this.preferPrimaryShardRebalanceBuffer);
+        weightFunction = new WeightFunction(
+            this.indexBalanceFactor,
+            this.shardBalanceFactor,
+            this.preferPrimaryShardRebalanceBuffer,
+            this.preferPrimaryShardBalance
+        );
     }
 
     /**
@@ -303,9 +308,7 @@ private void updateWeightFunction() {
      */
     private void setPreferPrimaryShardBalance(boolean preferPrimaryShardBalance) {
         this.preferPrimaryShardBalance = preferPrimaryShardBalance;
-        this.weightFunction.updateAllocationConstraint(INDEX_PRIMARY_SHARD_BALANCE_CONSTRAINT_ID, preferPrimaryShardBalance);
-        this.weightFunction.updateAllocationConstraint(CLUSTER_PRIMARY_SHARD_BALANCE_CONSTRAINT_ID, preferPrimaryShardBalance);
-        this.weightFunction.updateRebalanceConstraint(INDEX_PRIMARY_SHARD_BALANCE_CONSTRAINT_ID, preferPrimaryShardBalance);
+        this.weightFunction.updatePrimaryShardBalanceConstraints(preferPrimaryShardBalance);
     }
 
     private void setPreferPrimaryShardRebalance(boolean preferPrimaryShardRebalance) {
@@ -492,7 +495,7 @@ static class WeightFunction {
         private AllocationConstraints constraints;
         private RebalanceConstraints rebalanceConstraints;
 
-        WeightFunction(float indexBalance, float shardBalance, float preferPrimaryBalanceBuffer) {
+        WeightFunction(float indexBalance, float shardBalance, float preferPrimaryBalanceBuffer, boolean preferPrimaryShardBalance) {
             float sum = indexBalance + shardBalance;
             if (sum <= 0.0f) {
                 throw new IllegalArgumentException("Balance factors must sum to a value > 0 but was: " + sum);
@@ -506,6 +509,8 @@ static class WeightFunction {
             this.rebalanceConstraints = new RebalanceConstraints(rebalanceParameter);
             // Enable index shard per node breach constraint
             updateAllocationConstraint(INDEX_SHARD_PER_NODE_BREACH_CONSTRAINT_ID, true);
+            // Set primary shard balance constraints
+            updatePrimaryShardBalanceConstraints(preferPrimaryShardBalance);
         }
 
         public float weightWithAllocationConstraints(ShardsBalancer balancer, ModelNode node, String index) {
@@ -531,6 +536,12 @@ void updateAllocationConstraint(String constraint, boolean enable) {
         void updateRebalanceConstraint(String constraint, boolean add) {
             this.rebalanceConstraints.updateRebalanceConstraint(constraint, add);
         }
+
+        void updatePrimaryShardBalanceConstraints(boolean preferPrimaryShardBalance) {
+            updateAllocationConstraint(INDEX_PRIMARY_SHARD_BALANCE_CONSTRAINT_ID, preferPrimaryShardBalance);
+            updateAllocationConstraint(CLUSTER_PRIMARY_SHARD_BALANCE_CONSTRAINT_ID, preferPrimaryShardBalance);
+            updateRebalanceConstraint(INDEX_PRIMARY_SHARD_BALANCE_CONSTRAINT_ID, preferPrimaryShardBalance);
+        }
     }
 
     /**