support single shard tables for query generator

src/test/regress/citus_tests/query_generator/README.md

@@ -172,7 +172,7 @@ Tool supports following citus table types:
 targetTables:
   - Table:
     ...
-    citusType: <one of (DISTRIBUTED || REFERENCE || POSTGRES)>
+    citusType: <one of (HASH_DISTRIBUTED || SINGLE_SHARD_DISTRIBUTED || REFERENCE || POSTGRES)>
     ...
 ```
 

src/test/regress/citus_tests/query_generator/config/config.py

@@ -111,8 +111,12 @@ def getMaxAllowedCountForTable(tableName):
     return filtered[0].maxAllowedUseOnQuery
 
 
-def isTableDistributed(table):
+def isTableHashDistributed(table):
-    return table.citusType == CitusType.DISTRIBUTED
+    return table.citusType == CitusType.HASH_DISTRIBUTED
+
+
+def isTableSingleShardDistributed(table):
+    return table.citusType == CitusType.SINGLE_SHARD_DISTRIBUTED
 
 
 def isTableReference(table):

src/test/regress/citus_tests/query_generator/config/config.yaml

@@ -26,7 +26,7 @@ commonColName: id
 targetTables:
   - Table:
       name: dist
-      citusType: DISTRIBUTED
+      citusType: HASH_DISTRIBUTED
       maxAllowedUseOnQuery: 10
       rowCount: 10
       nullRate: 0.1
@@ -36,6 +36,18 @@ targetTables:
           name: id
           type: int
       distinctCopyCount: 2
+  # - Table:
+  #     name: single_dist
+  #     citusType: SINGLE_SHARD_DISTRIBUTED
+  #     maxAllowedUseOnQuery: 10
+  #     rowCount: 10
+  #     nullRate: 0.1
+  #     duplicateRate: 0.1
+  #     columns:
+  #       - Column:
+  #         name: id
+  #         type: int
+  #     distinctCopyCount: 2
   - Table:
       name: ref
       citusType: REFERENCE

src/test/regress/citus_tests/query_generator/data_gen.py

@@ -17,7 +17,10 @@ def getTableData():
         dataGenerationSql += "\n"
 
         # generate null rows
-        if not table.citusType == CitusType.DISTRIBUTED:
+        if table.citusType not in (
+            CitusType.HASH_DISTRIBUTED,
+            CitusType.SINGLE_SHARD_DISTRIBUTED,
+        ):
             targetNullRows = int(table.rowCount * table.nullRate)
             dataGenerationSql += _genNullData(table.name, targetNullRows)
             dataGenerationSql += "\n"

src/test/regress/citus_tests/query_generator/ddl_gen.py

@@ -1,4 +1,9 @@
-from config.config import getConfig, isTableDistributed, isTableReference
+from config.config import (
+    getConfig,
+    isTableHashDistributed,
+    isTableReference,
+    isTableSingleShardDistributed,
+)
 
 
 def getTableDDLs():
@@ -23,7 +28,7 @@ def _genTableDDL(table):
         ddl += _genColumnDDL(table.columns[-1])
     ddl += ");\n"
 
-    if isTableDistributed(table):
+    if isTableHashDistributed(table):
         ddl += (
             "SELECT create_distributed_table("
             + "'"
@@ -34,6 +39,9 @@ def _genTableDDL(table):
             + ");"
         )
         ddl += "\n"
+    if isTableSingleShardDistributed(table):
+        ddl += "SELECT create_distributed_table(" + "'" + table.name + "'" ",NULL);"
+        ddl += "\n"
     elif isTableReference(table):
         ddl += "SELECT create_reference_table(" + "'" + table.name + "'" + ");"
         ddl += "\n"

src/test/regress/citus_tests/query_generator/node_defs.py

@@ -22,9 +22,10 @@ class RestrictOp(Enum):
 
 
 class CitusType(Enum):
-    DISTRIBUTED = 1
+    HASH_DISTRIBUTED = 1
-    REFERENCE = 2
+    SINGLE_SHARD_DISTRIBUTED = 2
-    POSTGRES = 3
+    REFERENCE = 3
+    POSTGRES = 4
 
 
 class Table: