[DOP-22297] Parallel reading from JDBC sources

Author: Ilyas Gasanov

docs/changelog/next_release/219.improvement.rst

@@ -0,0 +1 @@
+Enable parallel reading from JDBC sources

syncmaster/dto/transfers.py

@@ -6,6 +6,7 @@
 
 from onetl.file.format import CSV, JSON, ORC, XML, Excel, JSONLine, Parquet
 
+from syncmaster.dto.transfers_resources import Resources
 from syncmaster.dto.transfers_strategy import FullStrategy, IncrementalStrategy
 
 
@@ -19,6 +20,7 @@ class DBTransferDTO(TransferDTO):
     id: int
     table_name: str
     strategy: FullStrategy | IncrementalStrategy
+    resources: Resources
     transformations: list[dict] | None = None
     options: dict | None = None
 
@@ -34,6 +36,7 @@ class FileTransferDTO(TransferDTO):
     directory_path: str
     file_format: CSV | JSONLine | JSON | Excel | XML | ORC | Parquet
     strategy: FullStrategy | IncrementalStrategy
+    resources: Resources
     options: dict
     file_name_template: str | None = None
     df_schema: dict | None = None

syncmaster/dto/transfers_resources.py

@@ -0,0 +1,10 @@
+# SPDX-FileCopyrightText: 2023-2024 MTS PJSC
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass
+
+
+@dataclass
+class Resources:
+    max_parallel_tasks: int
+    cpu_cores_per_task: int
+    ram_bytes_per_task: int

syncmaster/worker/controller.py

@@ -41,6 +41,7 @@
     SFTPTransferDTO,
     WebDAVTransferDTO,
 )
+from syncmaster.dto.transfers_resources import Resources
 from syncmaster.dto.transfers_strategy import Strategy
 from syncmaster.exceptions.connection import ConnectionTypeNotRecognizedError
 from syncmaster.schemas.v1.connection_types import FILE_CONNECTION_TYPES
@@ -169,6 +170,7 @@ def __init__(
             transfer_id=run.transfer.id,
             transfer_params=run.transfer.source_params,
             strategy_params=run.transfer.strategy_params,
+            resources=run.transfer.resources,
             transformations=run.transfer.transformations,
             connection_auth_data=source_auth_data,
             temp_dir=TemporaryDirectory(dir=self.temp_dir.name, prefix="downloaded_"),
@@ -179,6 +181,7 @@ def __init__(
             transfer_id=run.transfer.id,
             transfer_params=run.transfer.target_params,
             strategy_params=run.transfer.strategy_params,
+            resources=run.transfer.resources,
             transformations=run.transfer.transformations,
             connection_auth_data=target_auth_data,
             temp_dir=TemporaryDirectory(dir=self.temp_dir.name, prefix="written_"),
@@ -215,6 +218,7 @@ def get_handler(
         transfer_id: int,
         transfer_params: dict[str, Any],
         strategy_params: dict[str, Any],
+        resources: dict[str, Any],
         transformations: list[dict],
         temp_dir: TemporaryDirectory,
     ) -> Handler:
@@ -232,6 +236,7 @@ def get_handler(
             transfer_dto=transfer_dto(
                 id=transfer_id,
                 strategy=Strategy.from_dict(strategy_params),
+                resources=Resources(**resources),
                 transformations=transformations,
                 **transfer_params,
             ),

syncmaster/worker/handlers/db/base.py

@@ -45,6 +45,8 @@ def read(self) -> DataFrame:
                 expression=self.transfer_dto.strategy.increment_by,
             )
 
+        reader_params.update(self._get_reading_options())
+
         reader = DBReader(
             connection=self.connection,
             table=self.transfer_dto.table_name,
@@ -104,6 +106,35 @@ def _get_columns_filter_expressions(self) -> list[str] | None:
 
         return self._make_columns_filter_expressions(expressions)
 
+    def _get_reading_options(self) -> dict:
+        options = {}
+
+        if self.transfer_dto.resources.max_parallel_tasks < 2:
+            return options
+
+        if self.transfer_dto.strategy.type == "incremental":
+            # using "range" partitioning if HWM is available for efficient min/max lookup,
+            # since the incremental column is usually an indexed primary key
+            options["options"] = {
+                "partition_column": self.transfer_dto.strategy.increment_by,
+                "num_partitions": self.transfer_dto.resources.max_parallel_tasks,
+                "partitioning_mode": "range" if self.hwm and self.hwm.value else "hash",
+            }
+
+        elif self.transfer_dto.strategy.type == "full":
+            schema = self.connection.get_df_schema(
+                source=self.transfer_dto.table_name,
+                columns=self._get_columns_filter_expressions(),
+            )
+            # using "hash" partitioning as it works with any column type and improves full scan performance
+            options["options"] = {
+                "partition_column": self._quote_field(schema[0].name),
+                "num_partitions": self.transfer_dto.resources.max_parallel_tasks,
+                "partitioning_mode": "hash",
+            }
+
+        return options
+
     @staticmethod
     def _quote_field(field: str) -> str:
         return f'"{field}"'

syncmaster/worker/handlers/db/clickhouse.py

@@ -46,9 +46,9 @@ def write(self, df: DataFrame) -> None:
             (col for col in normalized_df.columns if col.lower().endswith("id")),
             normalized_df.columns[0],  # if there is no column with "id", take the first column
         )
-        quoted_sort_column = f'"{sort_column}"'
-
-        self.transfer_dto.options["createTableOptions"] = f"ENGINE = MergeTree() ORDER BY {quoted_sort_column}"
+        self.transfer_dto.options["createTableOptions"] = (
+            f"ENGINE = MergeTree() ORDER BY {self._quote_field(sort_column)}"
+        )
 
         if self.transfer_dto.strategy.type == "incremental" and self.hwm and self.hwm.value:
             self.transfer_dto.options["if_exists"] = "append"

syncmaster/worker/handlers/db/hive.py

@@ -60,6 +60,9 @@ def _make_rows_filter_expression(self, filters: list[dict]) -> str | None:
 
         return " AND ".join(expressions) or None
 
+    def _get_reading_options(self) -> dict:
+        return {}
+
     @staticmethod
     def _quote_field(field):
         return f"`{field}`"

tests/test_integration/test_run_transfer/test_clickhouse.py

@@ -47,6 +47,13 @@ async def postgres_to_clickhouse(
             "table_name": f"{clickhouse_for_conftest.user}.target_table",
         },
         strategy_params=strategy,
+        # this covers setting JDBC parallel reading options for full and incremental strategies
+        # TODO: add fixtures if usage increases
+        resources={
+            "max_parallel_tasks": 2,
+            "cpu_cores_per_task": 1,
+            "ram_bytes_per_task": 1024**3,
+        },
         transformations=transformations,
         queue_id=queue.id,
     )