YT-25975 - Task count for parallel reader

* Changelog entry
Type: feature
Component: python-sdk

Tune task count for parallel reader
commit_hash:4dd27a2d54e4478f7d086723de1beeffa9e855b6

Author: denvr

yt/python/yt/wrapper/default_config.py

@@ -236,7 +236,7 @@ class DefaultConfigQueryTrackerType(TypedDict, total=False):
 
     class DefaultConfigReadParallelType(TypedDict, total=False):
         max_thread_count: int
-        data_size_per_thread: int
+        data_size_per_thread: Union[int, None]
         enable: bool
 
     read_parallel: DefaultConfigReadParallelType
@@ -844,8 +844,8 @@ def get_dynamic_table_retries() -> DefaultConfigRetriesBackoffType:
     "read_parallel": {
         # Number of threads for reading table.
         "max_thread_count": 10,
-        # Approximate data size per one thread.
-        "data_size_per_thread": 8 * 1024 * 1024,
+        # Approximate data size per one task.
+        "data_size_per_thread": None,
         # Always run read parallel if it is possible.
         "enable": False
     },

yt/python/yt/wrapper/file_commands.py

@@ -9,7 +9,7 @@
 from .cypress_commands import (remove, exists, set_attribute, mkdir, find_free_subpath,
                                create, link, get, set)
 from .default_config import DEFAULT_WRITE_CHUNK_SIZE
-from .parallel_reader import make_read_parallel_request
+from .parallel_reader import make_read_parallel_request, _prepare_ranges_for_parallel_read
 from .parallel_writer import make_parallel_write_request
 from .retries import Retrier, default_chaos_monkey
 from .transaction import Transaction
@@ -90,23 +90,6 @@ def attributes(self):
         return self._attributes
 
 
-def _prepare_ranges_for_parallel_read(offset, length, data_size, data_size_per_thread):
-    offset = get_value(offset, 0)
-    offset = min(offset, data_size)
-
-    length = get_value(length, data_size)
-    length = min(length, data_size - offset)
-
-    result = []
-    while offset < data_size and length > 0:
-        range_size = min(data_size_per_thread, length)
-        result.append({"range" : (offset, range_size)})
-        offset += range_size
-        length -= range_size
-
-    return result
-
-
 def _prepare_params_for_parallel_read(params, range):
     params["offset"], params["length"] = range["range"][0], range["range"][1]
     return params

yt/python/yt/wrapper/parallel_reader.py

@@ -1,4 +1,4 @@
-from .common import YtError
+from .common import YtError, require, get_value
 from .config import get_config, get_option
 from .errors import YtChunkUnavailable
 from .format import YtFormatReadError
@@ -11,10 +11,18 @@
 from .thread_pool import ThreadPool
 from .ypath import TablePath
 
+import yt.logger as logger
+
 from yt.common import join_exceptions
 
+import builtins
 import copy
 import threading
+import typing
+
+
+DEFAULT_DATA_SIZE_PER_THREAD = 8 * 1024 * 1024
+DEFAULT_SINGLE_CHUNK_SPLIT = 3
 
 
 class ParallelReadRetrier(Retrier):
@@ -138,3 +146,90 @@ def make_read_parallel_request(command_name, path, ranges, params, prepare_param
         if transaction:
             transaction.abort()
         raise
+
+
+def _prepare_ranges_for_parallel_read(
+    offset: int,
+    length: int,
+    data_size: int,
+    data_size_per_thread: int,
+) -> typing.List[typing.Dict[str, typing.Tuple[int, int]]]:
+    if not data_size_per_thread:
+        data_size_per_thread = DEFAULT_DATA_SIZE_PER_THREAD
+
+    offset = get_value(offset, 0)
+    offset = min(offset, data_size)
+
+    length = get_value(length, data_size)
+    length = min(length, data_size - offset)
+
+    result = []
+    while offset < data_size and length > 0:
+        range_size = min(data_size_per_thread, length)
+        result.append({"range" : (offset, range_size)})
+        offset += range_size
+        length -= range_size
+
+    return result
+
+
+def _slice_row_ranges_for_parallel_read(
+    ranges: typing.List[typing.Mapping],
+    row_count: int,
+    chunk_count: int,
+    data_size: int,
+    replication_factor: int,
+    data_size_per_thread: typing.Union[int, None],
+) -> typing.Tuple[typing.List[typing.Dict[str, typing.Tuple[int, int]]], int]:
+    def _get_ranges(ranges, rows_per_task):
+        result = []
+        for range in ranges:
+            if "exact" in range:
+                require("row_index" in range["exact"], lambda: YtError('Invalid YPath: "row_index" not found'))
+                lower_limit = range["exact"]["row_index"]
+                upper_limit = lower_limit + 1
+            else:
+                if "lower_limit" in range:
+                    require("row_index" in range["lower_limit"], lambda: YtError('Invalid YPath: "row_index" not found'))
+                if "upper_limit" in range:
+                    require("row_index" in range["upper_limit"], lambda: YtError('Invalid YPath: "row_index" not found'))
+
+                lower_limit = 0 if "lower_limit" not in range else range["lower_limit"]["row_index"]
+                upper_limit = row_count if "upper_limit" not in range else range["upper_limit"]["row_index"]
+
+            for start in builtins.range(lower_limit, upper_limit, rows_per_task):
+                end = min(start + rows_per_task, upper_limit)
+                result.append({"range" : (start, end)})
+        return result
+
+    if row_count > 0:
+        row_size = data_size / float(row_count)
+    else:
+        row_size = 1
+
+    if data_size_per_thread:
+        rows_per_task = max(
+            int(data_size_per_thread / row_size),
+            1,
+        )
+        result = _get_ranges(ranges, rows_per_task)
+        logger.debug(f"Parallel read tasks count: {len(result)}, {row_size=}, {rows_per_task=}")
+        return result, rows_per_task * row_size
+    else:
+        if chunk_count > 0 and replication_factor >= 0:
+            data_size_per_thread = int(data_size / (chunk_count * replication_factor))
+        else:
+            data_size_per_thread = data_size
+        rows_per_task_by_table = int(data_size_per_thread / row_size)
+        rows_per_task_by_default = int(DEFAULT_DATA_SIZE_PER_THREAD / row_size)
+        rows_per_task = max(
+            rows_per_task_by_default,
+            rows_per_task_by_table,
+            1,
+        )
+        result = _get_ranges(ranges, rows_per_task)
+        if len(result) == 1 and rows_per_task > DEFAULT_SINGLE_CHUNK_SPLIT:
+            rows_per_task //= DEFAULT_SINGLE_CHUNK_SPLIT
+            result = _get_ranges(ranges, rows_per_task)
+        logger.debug(f"Parallel read tasks count: {len(result)}, {row_size=}, {rows_per_task=} ({rows_per_task_by_default}, {rows_per_task_by_table})")
+        return result, rows_per_task * row_size

yt/python/yt/wrapper/table_commands.py

@@ -25,7 +25,7 @@
 from .table_helpers import (_prepare_source_tables, _are_default_empty_table, _prepare_table_writer,
                             _remove_tables, DEFAULT_EMPTY_TABLE, _to_chunk_stream, _prepare_command_format)
 from .file_commands import _get_remote_temp_files_directory, _append_default_path_with_user_level
-from .parallel_reader import make_read_parallel_request
+from .parallel_reader import make_read_parallel_request, _slice_row_ranges_for_parallel_read
 from .schema import _SchemaRuntimeCtx, TableSchema, make_dataclass_from_table_schema
 from .stream import ItemStream, _ChunkStream
 from .ypath import TablePath, YPath, ypath_join
@@ -34,7 +34,6 @@
 import yt.yson as yson
 import yt.logger as logger
 
-import builtins
 from copy import deepcopy
 from datetime import timedelta
 import enum
@@ -449,35 +448,6 @@ def read_blob_table(table, part_index_column_name=None, data_column_name=None,
     return response
 
 
-def _slice_row_ranges_for_parallel_read(ranges, row_count, data_size, data_size_per_thread):
-    result = []
-    if row_count > 0:
-        row_size = data_size / float(row_count)
-    else:
-        row_size = 1
-
-    rows_per_thread = max(int(data_size_per_thread / row_size), 1)
-    for range in ranges:
-        if "exact" in range:
-            require("row_index" in range["exact"], lambda: YtError('Invalid YPath: "row_index" not found'))
-            lower_limit = range["exact"]["row_index"]
-            upper_limit = lower_limit + 1
-        else:
-            if "lower_limit" in range:
-                require("row_index" in range["lower_limit"], lambda: YtError('Invalid YPath: "row_index" not found'))
-            if "upper_limit" in range:
-                require("row_index" in range["upper_limit"], lambda: YtError('Invalid YPath: "row_index" not found'))
-
-            lower_limit = 0 if "lower_limit" not in range else range["lower_limit"]["row_index"]
-            upper_limit = row_count if "upper_limit" not in range else range["upper_limit"]["row_index"]
-
-        for start in builtins.range(lower_limit, upper_limit, rows_per_thread):
-            end = min(start + rows_per_thread, upper_limit)
-            result.append({"range" : (start, end)})
-
-    return result
-
-
 def _prepare_params_for_parallel_read(params, range):
     params["path"].attributes["ranges"] = [{"lower_limit": {"row_index": range["range"][0]},
                                             "upper_limit": {"row_index": range["range"][1]}}]
@@ -808,7 +778,7 @@ def _check_attributes_for_read_table(attributes, table, client):
 def _get_table_attributes(table, client):
     attributes = get(
         table + "/@",
-        attributes=["type", "chunk_count", "compressed_data_size", "dynamic", "row_count", "uncompressed_data_size"],
+        attributes=["type", "chunk_count", "compressed_data_size", "dynamic", "row_count", "replication_factor", "chunk_count", "uncompressed_data_size"],
         client=client)
     return attributes
 
@@ -861,11 +831,14 @@ def read_table(table, format=None, table_reader=None, control_attributes=None, u
                 table.attributes["ranges"] = [
                     {"lower_limit": {"row_index": 0},
                      "upper_limit": {"row_index": attributes["row_count"]}}]
-            ranges = _slice_row_ranges_for_parallel_read(
-                table.attributes["ranges"],
-                attributes["row_count"],
-                attributes["uncompressed_data_size"],
-                get_config(client)["read_parallel"]["data_size_per_thread"])
+            ranges, _ = _slice_row_ranges_for_parallel_read(
+                ranges=table.attributes["ranges"],
+                row_count=attributes["row_count"],
+                chunk_count=attributes["chunk_count"],
+                data_size=attributes["uncompressed_data_size"],
+                replication_factor=attributes["replication_factor"],
+                data_size_per_thread=get_config(client)["read_parallel"]["data_size_per_thread"],
+            )
             response_parameters = get_value(response_parameters, {})
             if not ranges:
                 response_parameters["start_row_index"] = 0
@@ -1229,12 +1202,14 @@ def _dump_file(table, output_file, output_path, enable_several_files, unordered,
                 "upper_limit": {"row_index": attributes["row_count"]},
             }]
 
-        data_size_per_thread = get_config(client)["read_parallel"]["data_size_per_thread"]
-        ranges = _slice_row_ranges_for_parallel_read(
-            table.attributes["ranges"],
-            attributes["row_count"],
-            attributes["uncompressed_data_size"],
-            data_size_per_thread)
+        ranges, data_size_per_thread = _slice_row_ranges_for_parallel_read(
+            ranges=table.attributes["ranges"],
+            row_count=attributes["row_count"],
+            chunk_count=attributes["chunk_count"],
+            data_size=attributes["uncompressed_data_size"],
+            replication_factor=attributes["replication_factor"],
+            data_size_per_thread=get_config(client)["read_parallel"]["data_size_per_thread"],
+        )
 
         range_count = len(ranges)
         result_ranges = []