Merge 812eddf into adbecd5

Author: vlad-gogov

ydb/tests/olap/column_family/compression/alter_compression.py

@@ -0,0 +1,177 @@
+import logging
+from .base import ColumnFamilyTestBase
+from typing import Callable
+from ydb.tests.library.common.helpers import plain_or_under_sanitizer
+from ydb.tests.olap.common.thread_helper import TestThread, TestThreads
+from ydb.tests.olap.common.column_table_helper import ColumnTableHelper
+
+logger = logging.getLogger(__name__)
+
+
+class TestAlterCompression(ColumnFamilyTestBase):
+    class_name = "alter_compression"
+
+    @classmethod
+    def setup_class(cls):
+        super(TestAlterCompression, cls).setup_class()
+
+    def upsert_and_wait_portions(self, table: ColumnTableHelper, number_rows_for_insert: int, count_upsert: int):
+        prev_number_rows: int = table.get_row_count()
+        for _ in range(count_upsert):
+            self.ydb_client.query(
+                """
+                $row_count = %i;
+                $prev_count = %i;
+                $rows= ListMap(ListFromRange(0, $row_count), ($i) -> {
+                    return <|
+                        value: $i + $prev_count,
+                        value1: $i + $prev_count,
+                    |>;
+                });
+                UPSERT INTO `%s`
+                SELECT * FROM AS_TABLE($rows);
+            """
+                % (number_rows_for_insert, prev_number_rows, table.path)
+            )
+            prev_number_rows += number_rows_for_insert
+            logger.info(
+                f"{prev_number_rows} rows inserted in {table.path}. portions: {table.get_portion_stat_by_tier()}, blobs: {table.get_blob_stat_by_tier()}"
+            )
+        assert table.get_row_count() == prev_number_rows
+
+        if not self.wait_for(
+            lambda: len(table.get_portion_stat_by_tier()) != 0, plain_or_under_sanitizer(70, 140)
+        ):
+            raise Exception("not all portions have been updated")
+
+        if not self.wait_for(
+            lambda: table.get_portion_stat_by_tier()['__DEFAULT']['Rows'] == number_rows_for_insert * count_upsert, plain_or_under_sanitizer(70, 140)
+        ):
+            raise Exception("not all portions have been updated")
+
+    def add_family_in_create(self, name: str, settings: str):
+        return f"FAMILY {name} ({settings})"
+
+    def test_all_supported_compression(self):
+        ''' Implements https://github.com/ydb-platform/ydb/issues/13640 '''
+
+        single_upsert_rows_count: int = 10**5
+        upsert_count: int = 10
+        test_name: str = "all_supported_compression"
+        test_dir: str = f"{self.ydb_client.database}/{self.class_name}/{test_name}"
+        tables_path: list[str] = [
+            f"{test_dir}/off_compression",
+            f"{test_dir}/lz4_compression",
+            f"{test_dir}/zstd_compression",
+        ]
+        add_defaut_family: Callable[[str], str] = lambda settings: self.add_family_in_create(name='default', settings=settings)
+        tables_family: list[str] = [
+            add_defaut_family('COMPRESSION = "off"'),
+            add_defaut_family('COMPRESSION = "lz4"'),
+            add_defaut_family('COMPRESSION = "zstd"'),
+        ]
+
+        for i in range(2, 22):
+            tables_path.append(f"{test_dir}/zstd_{i}_compression")
+            tables_family.append(add_defaut_family(f'COMPRESSION = "zstd", COMPRESSION_LEVEL = {i}'))
+
+        assert len(tables_path) == len(tables_family)
+
+        tables: list[ColumnTableHelper] = []
+        for table_path, table_family in zip(tables_path, tables_family):
+            self.ydb_client.query(
+                f"""
+                CREATE TABLE `{table_path}` (
+                    value Uint64 NOT NULL,
+                    value1 Uint64,
+                    PRIMARY KEY(value),
+                    {table_family}
+                )
+                WITH (STORE = COLUMN)
+                """
+            )
+            logger.info(f"Table {table_path} created")
+            tables.append(ColumnTableHelper(self.ydb_client, table_path))
+
+        assert len(tables) == len(tables_path)
+
+        tasks: TestThreads = TestThreads()
+        for table in tables:
+            tasks.append(TestThread(target=self.upsert_and_wait_portions, args=[table, single_upsert_rows_count, upsert_count]))
+
+        tasks.start_and_wait_all()
+
+        volumes_without_compression: tuple[int, int] = tables[0].get_volumes_column("value")
+        for table in tables:
+            assert table.get_portion_stat_by_tier()['__DEFAULT']['Rows'] == single_upsert_rows_count * upsert_count
+            assert upsert_count * single_upsert_rows_count * 8 == volumes_without_compression[0]
+
+        for i in range(1, len(tables_path)):
+            volumes: tuple[int, int] = tables[i].get_volumes_column("value")
+            koef: float = volumes_without_compression[1] / volumes[1]
+            logging.info(
+                f"compression in `{tables[i].path}` {volumes_without_compression[1]} / {volumes[1]}: {koef}"
+            )
+            assert koef > 1
+
+    def test_availability_data(self):
+        ''' Implements https://github.com/ydb-platform/ydb/issues/13643 '''
+
+        single_upsert_rows_count: int = 10 ** 2
+        upsert_rows_count: int = 10
+        rows_count: int = 0
+        test_name: str = "availability_data"
+        test_dir: str = f"{self.ydb_client.database}/{self.class_name}/{test_name}"
+        tables_path: str = f"{test_dir}/test_table"
+
+        tables_family: list[str] = [
+            self.add_family_in_create(name = 'default', settings = 'COMPRESSION = "off"'),
+            self.add_family_in_create(name = 'family_lz4', settings = 'COMPRESSION = "lz4"'),
+            self.add_family_in_create(name = 'family_zstd', settings = 'COMPRESSION = "zstd"'),
+            self.add_family_in_create(name = 'family_zstd_10', settings = 'COMPRESSION = "zstd", COMPRESSION_LEVEL = 10'),
+        ]
+
+        self.ydb_client.query(
+            f"""
+            CREATE TABLE `{tables_path}` (
+                value Uint64 NOT NULL,
+                value1 Uint64,
+                PRIMARY KEY(value),
+                {','.join(tables_family)}
+            )
+            WITH (STORE = COLUMN)
+            """
+        )
+        logger.info(f"Table {tables_path} created")
+        test_table: ColumnTableHelper = ColumnTableHelper(self.ydb_client, tables_path)
+
+        def check_data(table: ColumnTableHelper, rows_cont: int):
+            assert table.get_row_count() == rows_cont
+            count_row: int = 0
+            result_set = self.ydb_client.query(f"SELECT * FROM `{table.path}` ORDER BY `value`")
+            for result in result_set:
+                logging.info(result)
+                for row in result.rows:
+                    assert row['value'] == count_row and row['value1'] == count_row
+                    count_row += 1
+            assert count_row == rows_cont
+        
+        self.upsert_and_wait_portions(test_table, single_upsert_rows_count, upsert_rows_count)
+        rows_count += single_upsert_rows_count * upsert_rows_count        
+        check_data(table=test_table, rows_cont=rows_count)
+
+        self.ydb_client.query(f"ALTER TABLE `{tables_path}` ALTER COLUMN `value` SET FAMILY family_lz4")
+        self.upsert_and_wait_portions(test_table, single_upsert_rows_count, upsert_rows_count)
+        rows_count += single_upsert_rows_count * upsert_rows_count
+        check_data(table=test_table, rows_cont=rows_count)
+
+        self.ydb_client.query(f"ALTER TABLE `{tables_path}` ALTER COLUMN `value1` SET FAMILY family_zstd")
+        self.upsert_and_wait_portions(test_table, single_upsert_rows_count, upsert_rows_count)
+        rows_count += single_upsert_rows_count * upsert_rows_count
+        check_data(table=test_table, rows_cont=rows_count)
+
+        self.ydb_client.query(f"ALTER TABLE `{tables_path}` ALTER COLUMN `value` SET FAMILY family_zstd_10")
+        self.ydb_client.query(f"ALTER TABLE `{tables_path}` ALTER COLUMN `value1` SET FAMILY family_zstd_10")
+        self.upsert_and_wait_portions(test_table, single_upsert_rows_count, upsert_rows_count)
+        rows_count += single_upsert_rows_count * upsert_rows_count
+        check_data(table=test_table, rows_cont=rows_count)

ydb/tests/olap/column_family/compression/base.py

@@ -0,0 +1,51 @@
+import yatest.common
+import os
+import time
+import logging
+
+from ydb.tests.library.harness.kikimr_runner import KiKiMR
+from ydb.tests.library.harness.kikimr_config import KikimrConfigGenerator
+from ydb.tests.olap.common.ydb_client import YdbClient
+
+
+logger = logging.getLogger(__name__)
+
+
+class ColumnFamilyTestBase(object):
+    @classmethod
+    def setup_class(cls):
+        cls._setup_ydb()
+
+    @classmethod
+    def teardown_class(cls):
+        cls.ydb_client.stop()
+        cls.cluster.stop()
+
+    @classmethod
+    def _setup_ydb(cls):
+        ydb_path = yatest.common.build_path(os.environ.get("YDB_DRIVER_BINARY"))
+        logger.info(yatest.common.execute([ydb_path, "-V"], wait=True).stdout.decode("utf-8"))
+        config = KikimrConfigGenerator(
+            extra_feature_flags={"enable_olap_compression": True},
+            column_shard_config={
+                "lag_for_compaction_before_tierings_ms": 0,
+                "compaction_actualization_lag_ms": 0,
+                "optimizer_freshness_check_duration_ms": 0,
+                "small_portion_detect_size_limit": 0,
+                "max_read_staleness_ms": 5000,
+            },
+        )
+        cls.cluster = KiKiMR(config)
+        cls.cluster.start()
+        node = cls.cluster.nodes[1]
+        cls.ydb_client = YdbClient(database=f"/{config.domain_name}", endpoint=f"grpc://{node.host}:{node.port}")
+        cls.ydb_client.wait_connection()
+
+    @staticmethod
+    def wait_for(condition_func, timeout_seconds):
+        t0 = time.time()
+        while time.time() - t0 < timeout_seconds:
+            if condition_func():
+                return True
+            time.sleep(1)
+        return False

ydb/tests/olap/column_family/compression/ya.make

@@ -0,0 +1,24 @@
+PY3TEST()
+ENV(YDB_DRIVER_BINARY="ydb/apps/ydbd/ydbd")
+
+TEST_SRCS(
+    base.py
+    alter_compression.py
+)
+
+SIZE(MEDIUM)
+
+PEERDIR(
+    ydb/tests/library
+    ydb/public/sdk/python
+    ydb/public/sdk/python/enable_v3_new_behavior
+    ydb/tests/olap/scenario/helpers
+    ydb/tests/olap/common
+)
+
+DEPENDS(
+    ydb/apps/ydbd
+)
+
+END()
+

ydb/tests/olap/column_family/ya.make

@@ -0,0 +1,3 @@
+RECURSE(
+    compression
+)

ydb/tests/olap/common/column_table_helper.py

@@ -0,0 +1,76 @@
+import logging
+import time
+from .ydb_client import YdbClient
+
+
+logger = logging.getLogger(__name__)
+
+
+class ColumnTableHelper:
+    def __init__(self, ydb_client: YdbClient, path: str):
+        self.ydb_client = ydb_client
+        self.path = path
+
+    def get_row_count(self) -> int:
+        count_row: int = 0
+        result_set = self.ydb_client.query(f"SELECT COUNT(*) AS Rows FROM `{self.path}`")
+        for result in result_set:
+            for row in result.rows:
+                count_row += row["Rows"]
+        return count_row
+    
+    def get_portion_count(self) -> int:
+        return self.ydb_client.query(f"select count(*) as Rows from `{self.path}/.sys/primary_index_portion_stats`")[0].rows[0]["Rows"]
+
+    def get_portion_stat_by_tier(self) -> dict[str, dict[str, int]]:
+        results = self.ydb_client.query(
+            f"select TierName, sum(Rows) as Rows, count(*) as Portions from `{self.path}/.sys/primary_index_portion_stats` group by TierName"
+        )
+        return {
+            row["TierName"]: {"Rows": row["Rows"], "Portions": row["Portions"]}
+            for result_set in results
+            for row in result_set.rows
+        }
+
+    def get_blob_stat_by_tier(self) -> dict[str, (int, int)]:
+        stmt = f"""
+            select TierName, count(*) as Portions, sum(BlobSize) as BlobSize, sum(BlobCount) as BlobCount from (
+                select TabletId, PortionId, TierName, sum(BlobRangeSize) as BlobSize, count(*) as BlobCount from `{self.path}/.sys/primary_index_stats` group by TabletId, PortionId, TierName
+            ) group by TierName
+        """
+        results = self.ydb_client.query(stmt)
+        return {
+            row["TierName"]: {"Portions": row["Portions"], "BlobSize": row["BlobSize"], "BlobCount": row["BlobCount"]}
+            for result_set in results
+            for row in result_set.rows
+        }
+    
+    def set_fast_compaction(self):
+        self.ydb_client.query(
+            f"""
+            ALTER OBJECT `{self.path}` (TYPE TABLE) SET (ACTION=UPSERT_OPTIONS, `COMPACTION_PLANNER.CLASS_NAME`=`lc-buckets`, `COMPACTION_PLANNER.FEATURES`=`
+                  {{"levels" : [{{"class_name" : "Zero", "portions_live_duration" : "5s", "expected_blobs_size" : 1000000000000, "portions_count_available" : 2}},
+                                {{"class_name" : "Zero"}}]}}`);
+            """
+        )
+
+    def _coollect_volumes_column(self, column_name: str) -> tuple[int, int]:
+        query = f'SELECT * FROM `{self.path}/.sys/primary_index_stats` WHERE Activity == 1 AND EntityName = \"{column_name}\"'
+        result_set = self.ydb_client.query(query)
+        raw_bytes, bytes = 0, 0
+        for result in result_set:
+            for rows in result.rows:
+                raw_bytes += rows["RawBytes"]
+                bytes += rows["BlobRangeSize"]
+        return raw_bytes, bytes
+
+    def get_volumes_column(self, column_name: str) -> tuple[int, int]:
+        pred_raw_bytes, pred_bytes = 0, 0
+        raw_bytes, bytes = self._coollect_volumes_column(column_name)
+        while pred_raw_bytes != raw_bytes and pred_bytes != bytes:
+            pred_raw_bytes = raw_bytes
+            pred_bytes = bytes
+            time.sleep(10)
+            raw_bytes, bytes = self._coollect_volumes_column(column_name)
+        logging.info(f"Table `{self.path}`, volumes `{column_name}` ({raw_bytes}, {bytes})")
+        return raw_bytes, bytes

ydb/tests/olap/common/thread_helper.py

@@ -0,0 +1,42 @@
+import threading
+
+
+class TestThread(threading.Thread):
+    def run(self) -> None:
+        self.exc = None
+        try:
+            self.ret = self._target(*self._args, **self._kwargs)
+        except BaseException as e:
+            self.exc = e
+
+    def join(self, timeout=None):
+        super().join(timeout)
+        if self.exc:
+            raise self.exc
+        return self.ret
+
+class TestThreads():
+    def __init__(self):
+        self.threads: list[TestThread] = list()
+
+    def append(self, thread: TestThread) -> int:
+        self.threads.append(thread)
+        return len(self.threads) - 1
+
+    def start_thread(self, thread_index: int):
+        self.threads[thread_index].start()
+
+    def start_all(self) -> None:
+        for thread in self.threads:
+            thread.start()
+
+    def join_thread(self, thread_index: int):
+        self.threads[thread_index].join()
+
+    def join_all(self, timeout=None) -> None:
+        for thread in self.threads:
+            thread.join(timeout=timeout)
+
+    def start_and_wait_all(self):
+        self.start_all()
+        self.join_all()

ydb/tests/olap/common/ya.make

@@ -0,0 +1,13 @@
+PY3_LIBRARY()
+
+    PY_SRCS (
+        column_table_helper.py
+        thread_helper.py
+        ydb_client.py
+    )
+
+    PEERDIR(
+        ydb/public/sdk/python
+    )
+
+END()