Merge f80b373 into 1f915f6

Author: vlad-gogov

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

@@ -0,0 +1,113 @@
+import logging
+from .base import ColumnFamilyTestBase, ColumnTableHelper
+from typing import Callable
+from ydb.tests.library.common.helpers import plain_or_under_sanitizer
+from ydb.tests.olap.scenario.helpers.thread_helper import TestThread
+
+logger = logging.getLogger(__name__)
+
+
+class TestAlterCompression(ColumnFamilyTestBase):
+    class_name = "alter_compression"
+    single_upsert_row_count: int = 10**5
+    count_upsert: int = 10
+
+    @classmethod
+    def setup_class(cls):
+        super(TestAlterCompression, cls).setup_class()
+
+    def test_all_supported_compression(self):
+        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: f"FAMILY default ({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 i in range(len(tables_path)):
+            self.ydb_client.query(
+                f"""
+                CREATE TABLE `{tables_path[i]}` (
+                    value Uint64 NOT NULL,
+                    value1 Uint64,
+                    PRIMARY KEY(value),
+                    {tables_family[i]}
+                )
+                WITH (STORE = COLUMN)
+                """
+            )
+            logger.info(f"Table {tables_path[i]} created")
+            tables.append(ColumnTableHelper(self.ydb_client, tables_path[i]))
+
+        def upsert_and_wait_portions(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'] == self.single_upsert_row_count * self.count_upsert, plain_or_under_sanitizer(70, 140)
+            ):
+                raise Exception("not all portions have been updated")
+
+        assert len(tables) == len(tables_path)
+
+        tasks: list[TestThread] = []
+        for table in tables:
+            tasks.append(TestThread(target=upsert_and_wait_portions, args=[table, self.single_upsert_row_count, self.count_upsert]))
+
+        for task in tasks:
+            task.start()
+
+        for task in tasks:
+            task.join()
+
+        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'] == self.single_upsert_row_count * self.count_upsert
+            assert self.count_upsert * self.single_upsert_row_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

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

@@ -0,0 +1,127 @@
+import yatest.common
+import os
+import time
+import ydb
+import logging
+
+from ydb.tests.library.harness.kikimr_runner import KiKiMR
+from ydb.tests.library.harness.kikimr_config import KikimrConfigGenerator
+from ydb.tests.library.harness.util import LogLevels
+
+
+logger = logging.getLogger(__name__)
+
+
+class YdbClient:
+    def __init__(self, endpoint, database):
+        self.driver = ydb.Driver(endpoint=endpoint, database=database, oauth=None)
+        self.database = database
+        self.session_pool = ydb.QuerySessionPool(self.driver)
+
+    def stop(self):
+        self.session_pool.stop()
+        self.driver.stop()
+
+    def wait_connection(self, timeout=5):
+        self.driver.wait(timeout, fail_fast=True)
+
+    def query(self, statement):
+        return self.session_pool.execute_with_retries(statement)
+
+
+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_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 _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
+
+
+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,23 @@
+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
+)
+
+DEPENDS(
+    ydb/apps/ydbd
+)
+
+END()
+

ydb/tests/olap/column_family/ya.make

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

ydb/tests/olap/ya.make

@@ -31,4 +31,5 @@ RECURSE(
     docs
     load
     ttl_tiering
+    column_family
 )