From 45804c0e5d58a30745494fb24faac0226e912127 Mon Sep 17 00:00:00 2001
From: Vinod Sridharan <14185211+visridha@users.noreply.github.com>
Date: Thu, 4 Apr 2024 00:33:06 +0000
Subject: [PATCH] more

---
 src/backend/distributed/operations/create_shards.c |  6 +++++-
 src/backend/distributed/shared_library_init.c      | 12 ++++++++++++
 src/include/distributed/pg_dist_shard.h            |  1 +
 src/test/regress/expected/multi_create_table.out   |  1 +
 src/test/regress/sql/multi_create_table.sql        |  1 +
 5 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/src/backend/distributed/operations/create_shards.c b/src/backend/distributed/operations/create_shards.c
index fe26d9ed4..59312d08b 100644
--- a/src/backend/distributed/operations/create_shards.c
+++ b/src/backend/distributed/operations/create_shards.c
@@ -54,6 +54,9 @@
 #include "distributed/transaction_management.h"
 #include "distributed/worker_manager.h"
 
+/* Config variables managed via guc */
+bool EnableSingleShardTableMultiNodePlacement = false;
+
 
 /* declarations for dynamic loading */
 PG_FUNCTION_INFO_V1(master_create_worker_shards);
@@ -165,7 +168,8 @@ CreateShardsWithRoundRobinPolicy(Oid distributedTableId, int32 shardCount,
 	char shardStorageType = ShardStorageType(distributedTableId);
 
 	int64 shardOffset = 0;
-	if (shardCount == 1 && shardStorageType == SHARD_STORAGE_TABLE)
+	if (EnableSingleShardTableMultiNodePlacement && shardCount == 1
+		&& shardStorageType == SHARD_STORAGE_TABLE)
 	{
 		/* For single shard distributed tables, use the colocationId to offset
 		 * where the shard is placed.
diff --git a/src/backend/distributed/shared_library_init.c b/src/backend/distributed/shared_library_init.c
index 45e212e8b..f004e8c43 100644
--- a/src/backend/distributed/shared_library_init.c
+++ b/src/backend/distributed/shared_library_init.c
@@ -86,6 +86,7 @@
 #include "distributed/multi_router_planner.h"
 #include "distributed/multi_server_executor.h"
 #include "distributed/pg_dist_partition.h"
+#include "distributed/pg_dist_shard.h"
 #include "distributed/placement_connection.h"
 #include "distributed/priority.h"
 #include "distributed/query_pushdown_planning.h"
@@ -1462,6 +1463,17 @@ RegisterCitusConfigVariables(void)
 		GUC_NO_SHOW_ALL | GUC_NOT_IN_SAMPLE,
 		NULL, NULL, NULL);
 
+	DefineCustomBoolVariable(
+		"citus.enable_single_shard_table_multi_node_placement",
+		gettext_noop("Enables placement of single shard distributed tables in"
+					 " all nodes of the cluster"),
+		NULL,
+		&EnableSingleShardTableMultiNodePlacement,
+		false,
+		PGC_USERSET,
+		GUC_NO_SHOW_ALL | GUC_NOT_IN_SAMPLE,
+		NULL, NULL, NULL);
+
 	DefineCustomBoolVariable(
 		"citus.enable_statistics_collection",
 		gettext_noop("Enables sending basic usage statistics to Citus."),
diff --git a/src/include/distributed/pg_dist_shard.h b/src/include/distributed/pg_dist_shard.h
index 5c98b755f..502e9e180 100644
--- a/src/include/distributed/pg_dist_shard.h
+++ b/src/include/distributed/pg_dist_shard.h
@@ -59,5 +59,6 @@ typedef FormData_pg_dist_shard *Form_pg_dist_shard;
 #define SHARD_STORAGE_TABLE 't'
 #define SHARD_STORAGE_VIRTUAL 'v'
 
+extern bool EnableSingleShardTableMultiNodePlacement;
 
 #endif   /* PG_DIST_SHARD_H */
diff --git a/src/test/regress/expected/multi_create_table.out b/src/test/regress/expected/multi_create_table.out
index 8d1ca713f..3e89fd8d0 100644
--- a/src/test/regress/expected/multi_create_table.out
+++ b/src/test/regress/expected/multi_create_table.out
@@ -468,6 +468,7 @@ DROP TABLE temp_table;
 DROP TABLE shard_count_table_3;
 -- test shard count 1 placement with colocate none.
 -- create a base table instance
+set citus.enable_single_hash_repartition_joins to on;
 CREATE TABLE shard_count_table_1_inst_1 (a int);
 SELECT create_distributed_table('shard_count_table_1_inst_1', 'a', shard_count:=1, colocate_with:='none');
  create_distributed_table
diff --git a/src/test/regress/sql/multi_create_table.sql b/src/test/regress/sql/multi_create_table.sql
index be7df50ea..e2165fc78 100644
--- a/src/test/regress/sql/multi_create_table.sql
+++ b/src/test/regress/sql/multi_create_table.sql
@@ -294,6 +294,7 @@ DROP TABLE shard_count_table_3;
 
 -- test shard count 1 placement with colocate none.
 -- create a base table instance
+set citus.enable_single_hash_repartition_joins to on;
 CREATE TABLE shard_count_table_1_inst_1 (a int);
 SELECT create_distributed_table('shard_count_table_1_inst_1', 'a', shard_count:=1, colocate_with:='none');