feat(sharding): add shard creation for new nodes during rebalance

- Log distinct messages for rebalancing due to inactive vs new nodes
- Detect new nodes joining cluster and create missing shard tables
- Set up replication from existing nodes to new nodes for new shards
- Improve shard map initialization to avoid undefined keys

Author: donhardman

src/Plugin/Sharding/Operator.php

@@ -113,7 +113,12 @@ public function checkBalance(): static {
 		if ($clusterHash === $currentHash) {
 			return $this;
 		}
-		Buddy::info("Rebalancing due to inactive nodes: {$inactiveNodes->join(', ')}");
+
+		if ($inactiveNodes->count() > 0) {
+			Buddy::info("Rebalancing due to inactive nodes: {$inactiveNodes->join(', ')}");
+		} else {
+			Buddy::info('Rebalancing due to new nodes joined');
+		}
 
 		// Get all tables from the state we have
 		$list = $this->state->listRegex('table:.+');

src/Plugin/Sharding/Table.php

@@ -443,6 +443,16 @@ public function rebalance(Queue $queue): void {
 				$cluster->processPendingTables($queue);
 			}
 
+			// Handle new nodes that need shard creation
+			$originalNodes = new Set($schema->map(fn($row) => $row['node']));
+			$newNodes = $activeNodes->diff($originalNodes);
+
+
+
+			if ($newNodes->count() > 0) {
+				$this->handleShardCreationForRebalancing($queue, $schema, $newSchema, $clusterMap);
+			}
+
 			// At this case we update schema
 			// before creating distributed table
 			$this->updateScheme($newSchema);
@@ -664,10 +674,12 @@ protected function getCreateShardedTableSQL(Set $shards): string {
 		$map = new Map;
 		foreach ($nodes as $row) {
 			foreach ($row['shards'] as $shard) {
-				$map[$shard] ??= new Set;
+				if (!$map->hasKey($shard)) {
+					$map[$shard] = new Set();
+				}
 				$shardName = $this->getShardName($shard);
-				// @phpstan-ignore-next-line
-				$map[$shard]->add("{$row['node']}:{$shardName}");
+
+				$map[$shard]?->add("{$row['node']}:{$shardName}");
 			}
 		}
 
@@ -748,4 +760,61 @@ protected static function parseShards(string $shards): Set {
 			: new Set
 		;
 	}
+
+	/**
+	 * Handle shard creation for rebalancing (all nodes that need new shards)
+	 * @param Queue $queue
+	 * @param Vector<array{node:string,shards:Set<int>,connections:Set<string>}> $oldSchema
+	 * @param Vector<array{node:string,shards:Set<int>,connections:Set<string>}> $newSchema
+	 * @param Map<string,Cluster> $clusterMap
+	 * @return void
+	 */
+	protected function handleShardCreationForRebalancing(Queue $queue, Vector $oldSchema, Vector $newSchema, Map $clusterMap): void {
+		// Create map of old schema for comparison
+		/** @var Map<string,Set<int>> */
+		$oldShardMap = new Map();
+		foreach ($oldSchema as $row) {
+			$oldShardMap[$row['node']] = $row['shards'];
+		}
+
+		foreach ($newSchema as $row) {
+			$oldShards = $oldShardMap->get($row['node'], new Set());
+			$newShards = $row['shards'];
+			$shardsToCreate = $newShards->diff($oldShards);
+
+			if ($shardsToCreate->isEmpty()) {
+				continue;
+			}
+
+			// Create missing shard tables on this node
+			foreach ($shardsToCreate as $shard) {
+				$sql = $this->getCreateTableShardSQL($shard);
+				$queue->add($row['node'], $sql);
+
+				// Find nodes that already have this shard in old schema for replication
+				$existingNodesWithShard = $oldSchema->filter(
+					fn($existingRow) => $existingRow['shards']->contains($shard)
+				);
+
+				// If there are existing nodes with this shard, set up replication
+				if ($existingNodesWithShard->count() <= 0) {
+					continue;
+				}
+
+				$sourceNode = $existingNodesWithShard->first()['node'];
+				$connectedNodes = new Set([$row['node'], $sourceNode]);
+
+
+
+				// Set up cluster replication for this shard
+				$this->handleReplication(
+					$sourceNode,
+					$queue,
+					$connectedNodes,
+					$clusterMap,
+					$shard
+				);
+			}
+		}
+	}
 }

src/Plugin/Sharding/Util.php

@@ -7,6 +7,7 @@
 use Ds\Vector;
 use RuntimeException;
 
+/** @package Manticoresearch\Buddy\Base\Plugin\Sharding */
 final class Util {
   /**
    * Generate sharding schema by using input nodes, shards and replication
@@ -137,6 +138,25 @@ public static function rebalanceShardingScheme(Vector $schema, Set $nodes): Vect
 		$newSchema = self::addNodesToSchema($newSchema, $nodes);
 		$inactiveShards = self::findInactiveShards($schema, $nodes);
 
+		// Check if we have new nodes (nodes with no shards assigned)
+		$hasNewNodes = $newSchema->filter(fn($row) => $row['shards']->isEmpty())->count() > 0;
+
+		if ($hasNewNodes) {
+			// For new nodes, use the same balanced logic as createShardingSchema
+			$totalShards = 0;
+			foreach ($schema as $row) {
+				$totalShards += $row['shards']->count();
+			}
+
+			if ($totalShards > 0) {
+				// Reuse the balanced assignment logic from createShardingSchema
+				$balancedSchema = self::initializeSchema($nodes);
+				$nodeMap = self::initializeNodeMap($nodes->count());
+				return self::assignNodesToSchema($balancedSchema, $nodeMap, $nodes, $totalShards, 2);
+			}
+		}
+
+		// For node failures only (no new nodes), use the original logic
 		return self::assignShardsToNodes($newSchema, $inactiveShards);
 	}