Merge 6670f58 into d4e572d

Author: zverevgeny

ydb/tests/library/harness/kikimr_config.py

@@ -157,6 +157,7 @@ def __init__(
             generic_connector_config=None,  # typing.Optional[TGenericConnectorConfig]
             kafka_api_port=None,
             metadata_section=None,
+            column_shard_config=None,
     ):
         if extra_feature_flags is None:
             extra_feature_flags = []
@@ -282,11 +283,12 @@ def __init__(
             self.yaml_config['pqconfig']['require_credentials_in_new_protocol'] = False
             self.yaml_config['pqconfig']['root'] = '/Root/PQ'
             self.yaml_config['pqconfig']['quoting_config']['enable_quoting'] = False
-
         if pq_client_service_types:
             self.yaml_config['pqconfig']['client_service_type'] = []
             for service_type in pq_client_service_types:
                 self.yaml_config['pqconfig']['client_service_type'].append({'name': service_type})
+        if column_shard_config:
+            self.yaml_config["column_shard_config"] = column_shard_config
 
         self.yaml_config['grpc_config']['services'].extend(extra_grpc_services)
 

ydb/tests/workloads/olap_workload/__main__.py

@@ -137,8 +137,9 @@ def join(self):
 
 
 class WorkloadTablesCreateDrop(WorkloadBase):
-    def __init__(self, client, prefix, stop):
+    def __init__(self, client, prefix, stop, allow_nullables_in_pk):
         super().__init__(client, prefix, "create_drop", stop)
+        self.allow_nullables_in_pk = allow_nullables_in_pk
         self.created = 0
         self.deleted = 0
         self.tables = set()
@@ -166,11 +167,21 @@ def create_table(self, table):
         column_n = random.randint(1, 10000)
         primary_key_column_n = random.randint(1, column_n)
         partition_key_column_n = random.randint(1, primary_key_column_n)
-        columns = [random.choice(supported_pk_types) for _ in range(primary_key_column_n)] + [random.choice(supported_types) for _ in range(column_n - primary_key_column_n)]
+        column_defs = []
+        for i in range(column_n):
+            if i < primary_key_column_n:
+                c = random.choice(supported_pk_types)
+                if not self.allow_nullables_in_pk or random.choice([False, True]):
+                    c += " NOT NULL"
+            else:
+                c = random.choice(supported_types)
+                if random.choice([False, True]):
+                    c += " NOT NULL"
+            column_defs.append(c)
 
         stmt = f"""
                 CREATE TABLE `{path}` (
-                    {", ".join(["c" + str(i) + " " + columns[i] + " NOT NULL" for i in range(column_n)])},
+                    {", ".join(["c" + str(i) + " " + column_defs[i] for i in range(column_n)])},
                     PRIMARY KEY({", ".join(["c" + str(i) for i in range(primary_key_column_n)])})
                 )
                 PARTITION BY HASH({", ".join(["c" + str(i) for i in range(partition_key_column_n)])})
@@ -273,11 +284,12 @@ def get_workload_thread_funcs(self):
 
 
 class WorkloadRunner:
-    def __init__(self, client, name, duration):
+    def __init__(self, client, name, duration, allow_nullables_in_pk):
         self.client = client
-        self.name = name
+        self.name = args.path
         self.tables_prefix = "/".join([self.client.database, self.name])
-        self.duration = duration
+        self.duration = args.duration
+        self.allow_nullables_in_pk = allow_nullables_in_pk
 
     def __enter__(self):
         self._cleanup()
@@ -294,7 +306,7 @@ def _cleanup(self):
     def run(self):
         stop = threading.Event()
         workloads = [
-            WorkloadTablesCreateDrop(self.client, self.name, stop),
+            WorkloadTablesCreateDrop(self.client, self.name, stop, self.allow_nullables_in_pk),
             # WorkloadInsertDelete(self.client, self.name, stop), TODO fix https://github.com/ydb-platform/ydb/issues/12871
         ]
         for w in workloads:
@@ -320,8 +332,9 @@ def run(self):
     parser.add_argument("--database", default="Root/test", help="A database to connect")
     parser.add_argument("--path", default="olap_workload", help="A path prefix for tables")
     parser.add_argument("--duration", default=10 ** 9, type=lambda x: int(x), help="A duration of workload in seconds.")
+    parser.add_argument("--allow-nullables-in-pk", default=False, help="Allow nullable types for columns in a Primary Key.")
     args = parser.parse_args()
     client = YdbClient(args.endpoint, args.database, True)
     client.wait_connection()
-    with WorkloadRunner(client, args.path, args.duration) as runner:
+    with WorkloadRunner(client, args.path, args.duration, args.allow_nullables_in_pk) as runner:
         runner.run()

ydb/tests/workloads/olap_workload/tests/test_workload.py

@@ -9,7 +9,12 @@
 class TestYdbWorkload(object):
     @classmethod
     def setup_class(cls):
-        cls.cluster = KiKiMR(KikimrConfigGenerator(erasure=Erasure.MIRROR_3_DC))
+        cls.cluster = KiKiMR(KikimrConfigGenerator(
+            erasure=Erasure.MIRROR_3_DC,
+            column_shard_config={
+                "allow_nullable_columns_in_pk": True,
+            }
+        ))
         cls.cluster.start()
 
     @classmethod
@@ -24,6 +29,7 @@ def test(self):
                 "--endpoint", f"grpc://localhost:{self.cluster.nodes[1].grpc_port}",
                 "--database=/Root",
                 "--duration", "120",
+                "--allow-nullables-in-pk", "1",
             ],
             wait=True
         )