[Data] read_clickhouse materializes an entire ReadTask result in one block

What happened + What you expected to happen

ray.data.read_clickhouse() can materialize the complete result of one ReadTask in a worker before Ray Data receives any output block. Although the datasource opens an Arrow stream, the read function does not return an output block until the full query result has been materialized.

With a filtered query, or without order_by, the datasource falls back to one task. A large result can therefore exceed worker memory even though the source uses an Arrow stream.

Ray's ReadTask contract allows a task to yield multiple blocks to avoid running out of memory during the read. However, the current ClickHouse read_fn() does not produce an iterable until _execute_block_query() has materialized the entire result into one Table. The fused ReadClickHouse->SplitBlocks(6) stage can only split that already-materialized block, so it cannot bound the memory used before the first block is emitted.

Expected behavior: the ReadTask should emit Arrow blocks incrementally so memory does not scale with the complete result of one task. Existing ordering and filter semantics should remain unchanged.

Versions / Dependencies

• Ray 2.58.0 (latest stable release; source-verified as still affected)
• Python 3.12
• PyArrow 19.0.1
• clickhouse-connect 1.5.0
• ClickHouse Server 25.3.2.39
• Docker-based Linux Ray cluster with a 2 GiB worker container memory limit, Dashboard disabled, and the driver running inside the cluster

Reproduction script

Use an existing MergeTree table containing 10 million rows and an approximately 512-byte payload column.

Run this from a Ray driver in the cluster configured above. Omitting order_by intentionally selects the single-ReadTask path:

import ray

ray.init(address="auto")
dataset = ray.data.read_clickhouse(
    table="default.oom_events",
    dsn="clickhouse+http://default@clickhouse:8123/default",
    columns=["id", "payload"],
    client_settings={"max_block_size": 8192},
    # Avoid retrying the OOM task so the reproduction has one terminal failure.
    ray_remote_args={"max_retries": 0},
)
print(dataset.take(1))

The runtime log below was captured on Ray 2.55.1. Ray 2.58.0 was checked separately at source level and retains the same materialization path.

Observed result:

ray.exceptions.OutOfMemoryError: 1 worker(s) were killed due to the node running low on memory
Memory on the node ... was 2.00GB / 2.00GB
task name=ReadClickHouse->SplitBlocks(6), actual memory used=1.57GB
Object store memory usage ... bytes in use: 10766

The 6 is the read operator's additional output split factor. It is derived during Ray's read planning from read-task metadata estimates and the target block-size policy; it does not make the ClickHouse query itself incremental.

Relevant implementation

The ClickHouse datasource currently wraps each query as one block-producing function:

def _create_read_fn(self, query):
    def read_fn():
        return [self._execute_block_query(query)]
    return read_fn

The query function collects the complete Arrow stream before returning. The following excerpt is from the Ray 2.58.0 source; comments and the local pyarrow import are omitted:

client = self._init_client()
try:
    with client.query_arrow_stream(query) as stream:
        record_batches = list(stream)
    return pa.Table.from_batches(record_batches)
except Exception as e:
    raise RuntimeError(f"Failed to execute block query: {e}")
finally:
    client.close()

The planner samples the in-memory bytes per row and uses it for task metadata:

estimated_size_bytes_per_row, sample_block_schema = self._get_sampled_estimates()
size_bytes=estimated_size_bytes_per_row * block_rows

Ray's ReadTask applies per_task_row_limit only after calling read_fn(); the unlimited branch also forwards the result directly:

result = self._read_fn()
if self._per_task_row_limit is None:
    yield from result
    return

yield from _iter_sliced_blocks(result, self._per_task_row_limit)

Therefore, the ReadClickHouse->SplitBlocks(...) operator cannot protect the worker from the earlier full-result materialization.

Suggested fix direction

Reuse the existing total-size and sampled bytes-per-row estimates to inform Ray's read planning and to bound each block emitted by the datasource. The datasource may reduce the requested parallelism under its existing ordering and filtering rules, but should not claim ownership of the global target block-size policy. Estimates should be treated as planning hints rather than hard RSS limits; if an estimate is unavailable or invalid, use a conservative fallback.

Change the worker read function to iterate query_arrow_stream() and yield each Arrow batch, or a bounded row/byte slice of each batch, before requesting the next batch. The producer should not fetch the next batch until the current bounded block has been yielded, so peak worker memory is bounded by the emitted block plus normal Ray transport overhead. The stream and client must still close on normal EOF, errors, cancellation, and early consumer close. This removes the list(stream) materialization while preserving the existing query, ordering, and filter behavior.

Relevant source: https://github.com/ray-project/ray/blob/ray-2.58.0/python/ray/data/_internal/datasource/clickhouse_datasource.py#L182-L256

The ReadTask multi-block contract and post-read slicing are documented here: https://github.com/ray-project/ray/blob/ray-2.58.0/python/ray/data/datasource/datasource.py#L304-L321 and https://github.com/ray-project/ray/blob/ray-2.58.0/python/ray/data/datasource/datasource.py#L351-L362

The read-operator split planning is implemented here: https://github.com/ray-project/ray/blob/ray-2.58.0/python/ray/data/_internal/logical/rules/set_read_parallelism.py#L23-L146

Issue Severity

High: large ClickHouse reads can terminate Ray workers even when the downstream consumer requests only a small result.