From 7064f798a7266713e3cc26715fb4848f7fafeec7 Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Sun, 9 Aug 2026 01:48:59 -0700 Subject: [PATCH] [python][ray] Reduce Blob request QPS with URI affinity --- docs/docs/pypaimon/multimodal-api.mdx | 19 ++ paimon-python/pypaimon/ray/ray_paimon.py | 180 ++++++++++++++++-- .../pypaimon/tests/multimodal_table_test.py | 21 +- .../pypaimon/tests/ray_blob_affinity_test.py | 164 ++++++++++++++++ 4 files changed, 370 insertions(+), 14 deletions(-) create mode 100644 paimon-python/pypaimon/tests/ray_blob_affinity_test.py diff --git a/docs/docs/pypaimon/multimodal-api.mdx b/docs/docs/pypaimon/multimodal-api.mdx index 2bb3d2562cc8..2f8d174e1503 100644 --- a/docs/docs/pypaimon/multimodal-api.mdx +++ b/docs/docs/pypaimon/multimodal-api.mdx @@ -423,12 +423,31 @@ ds = ( result_ds = docs.map_with_blobs(ds, blob_cols, process_batch) ``` +For small inference batches over scattered descriptors, enable URI affinity to +place adjacent ranges from the same Blob file together. Blob reads are then +coalesced across several inference batches without exposing the prefetched +bytes to Ray's object store: + +```python +result_ds = docs.map_with_blobs( + ds, + blob_cols, + process_batch, + batch_size=32, + blob_uri_affinity=True, +) +``` + Notes: - `process_batch` must return a small Ray-compatible batch; return an empty `pyarrow.Table` for side-effect-only jobs. - Avoid returning raw BLOB bytes, which would materialize payloads in Ray's object store. +- URI affinity performs a distributed sort and may reorder rows. Enable it when + reduced object-store request QPS outweighs the shuffle cost. `prefetch_bytes` + bounds each worker's payload window (64 MiB by default, except for one + oversized inference batch). - Tune `to_ray(...)` and `map_with_blobs(...)` parameters only when needed. ## Row IDs diff --git a/paimon-python/pypaimon/ray/ray_paimon.py b/paimon-python/pypaimon/ray/ray_paimon.py index 78309200ac8d..6e1cfc9098a7 100644 --- a/paimon-python/pypaimon/ray/ray_paimon.py +++ b/paimon-python/pypaimon/ray/ray_paimon.py @@ -26,7 +26,9 @@ write_paimon(ds, "db.table", catalog_options={"warehouse": "/path"}) """ +import hashlib import importlib +import uuid from typing import Any, Callable, Dict, List, Optional, TYPE_CHECKING from pypaimon.common.predicate import Predicate @@ -146,6 +148,8 @@ def map_with_blobs( all_blob_columns=None, parallelism: int = 64, batch_size: Optional[int] = 1024, + blob_uri_affinity: bool = False, + prefetch_bytes: int = 64 * 1024 * 1024, fn_kwargs: Optional[Dict[str, Any]] = None, ray_remote_args: Optional[Dict[str, Any]] = None, **map_args, @@ -157,7 +161,10 @@ def map_with_blobs( Ray-compatible batch; for side-effect-only work, return an empty ``pyarrow.Table`` instead of ``None``. Call this directly on ``scan().to_ray()`` output, or pass ``file_io`` and ``all_blob_columns``. - Tune ``batch_size`` for BLOB size and worker memory. + Tune ``batch_size`` for BLOB size and worker memory. Set + ``blob_uri_affinity=True`` to shuffle descriptors by URI and offset before + reading. This lets each worker coalesce adjacent ranges across multiple + ``fn`` batches, bounded by ``prefetch_bytes``. """ _require_ray_data() @@ -173,6 +180,14 @@ def map_with_blobs( raise ValueError("parallelism must be at least 1, got {}".format(parallelism)) if batch_size is not None and batch_size < 1: raise ValueError("batch_size must be at least 1, got {}".format(batch_size)) + if not isinstance(blob_uri_affinity, bool): + raise ValueError("blob_uri_affinity must be a boolean") + if blob_uri_affinity and batch_size is None: + raise ValueError("blob_uri_affinity requires batch_size") + if (isinstance(prefetch_bytes, bool) + or not isinstance(prefetch_bytes, int) + or prefetch_bytes < 1): + raise ValueError("prefetch_bytes must be a positive integer") resolved_file_io = file_io if resolved_file_io is None: @@ -188,7 +203,7 @@ def map_with_blobs( kwargs = dict(map_args) kwargs["batch_format"] = "pyarrow" - if batch_size is not None: + if batch_size is not None and not blob_uri_affinity: kwargs.setdefault("batch_size", batch_size) if ray_remote_args is not None: _set_map_batches_remote_args(dataset, kwargs, ray_remote_args) @@ -206,17 +221,81 @@ def map_with_blobs( if invalid: raise ValueError("Column {!r} is not a BLOB column.".format(invalid[0])) + mapper = _map_blob_batch + affinity_cols = [] + if blob_uri_affinity: + dataset, affinity_cols = _cluster_by_blob_uri(dataset, blob_cols) + mapper = _map_blob_affinity_block + kwargs["batch_size"] = None + + mapper_kwargs = { + "file_io": resolved_file_io, + "blob_cols": blob_cols, + "all_blob_cols": list(all_blob_cols), + "parallelism": parallelism, + "fn": fn, + "fn_kwargs": dict(fn_kwargs or {}), + } + if blob_uri_affinity: + mapper_kwargs.update({ + "fn_batch_size": batch_size, + "prefetch_bytes": prefetch_bytes, + "affinity_cols": affinity_cols, + }) return dataset.map_batches( - _map_blob_batch, + mapper, fn_kwargs=mapper_kwargs, **kwargs) + + +def _cluster_by_blob_uri(dataset, blob_cols): + token = uuid.uuid4().hex + key_col = "__paimon_blob_key_{}".format(token) + offset_col = "__paimon_blob_offset_{}".format(token) + with_keys = dataset.map_batches( + _append_blob_affinity_keys, fn_kwargs={ - "file_io": resolved_file_io, "blob_cols": blob_cols, - "all_blob_cols": list(all_blob_cols), - "parallelism": parallelism, - "fn": fn, - "fn_kwargs": dict(fn_kwargs or {}), + "key_col": key_col, + "offset_col": offset_col, }, - **kwargs) + batch_format="pyarrow", + zero_copy_batch=True, + ) + return with_keys.sort([key_col, offset_col]), [key_col, offset_col] + + +def _append_blob_affinity_keys(batch, blob_cols, key_col, offset_col): + import pyarrow as pa + from pypaimon.table.row.blob import BlobDescriptor + + empty_key = b"\0" * 16 + uri_keys = {} + keys = [] + offsets = [] + columns = [batch.column(name) for name in blob_cols] + for row in range(batch.num_rows): + descriptor = None + for column in columns: + value = column[row] + raw = value.as_py() if value.is_valid else None + if raw is not None and BlobDescriptor.is_blob_descriptor(raw): + descriptor = BlobDescriptor.deserialize(raw) + break + if descriptor is not None: + key = uri_keys.get(descriptor.uri) + if key is None: + key = hashlib.blake2b( + descriptor.uri.encode("utf-8"), digest_size=16).digest() + uri_keys[descriptor.uri] = key + keys.append(key) + offsets.append(descriptor.offset) + else: + keys.append(empty_key) + offsets.append(-1) + return batch.append_column( + key_col, pa.array(keys, type=pa.binary(16)) + ).append_column( + offset_col, pa.array(offsets, type=pa.int64()) + ) def _set_map_batches_remote_args(dataset, kwargs, ray_remote_args): @@ -233,13 +312,52 @@ def _map_blob_batch( batch, file_io, blob_cols, all_blob_cols, parallelism, fn, fn_kwargs): from pypaimon.multimodal.blob_read import fetch_blob_bodies + scalar_cols = _blob_scalar_columns(batch, blob_cols, all_blob_cols) + bodies = fetch_blob_bodies( + file_io, batch.select(blob_cols).to_pydict(), blob_cols, parallelism) + return _call_blob_fn(fn, batch.select(scalar_cols), bodies, fn_kwargs) + + +def _map_blob_affinity_block( + batch, file_io, blob_cols, all_blob_cols, parallelism, fn, fn_kwargs, + fn_batch_size, prefetch_bytes, affinity_cols): + from pypaimon.multimodal.blob_read import fetch_blob_bodies + + if batch.num_rows == 0: + return + + scalar_cols = _blob_scalar_columns( + batch, blob_cols, all_blob_cols, affinity_cols) + + for start, end in _blob_prefetch_windows( + batch, blob_cols, fn_batch_size, prefetch_bytes): + window = batch.slice(start, end - start) + bodies = fetch_blob_bodies( + file_io, + window.select(blob_cols).to_pydict(), + blob_cols, + parallelism, + ) + scalar = window.select(scalar_cols) + for batch_start in range(0, window.num_rows, fn_batch_size): + size = min(fn_batch_size, window.num_rows - batch_start) + fn_bodies = { + name: values[batch_start:batch_start + size] + for name, values in bodies.items() + } + yield _call_blob_fn( + fn, scalar.slice(batch_start, size), fn_bodies, fn_kwargs) + + +def _blob_scalar_columns(batch, blob_cols, all_blob_cols, internal_cols=()): missing = [name for name in blob_cols if name not in batch.schema.names] if missing: raise ValueError("BLOB column(s) not found in Ray Dataset: {}".format( ", ".join(missing))) all_blob = set(all_blob_cols) - scalar_cols = [name for name in batch.schema.names if name not in all_blob] + excluded = all_blob | set(internal_cols) + scalar_cols = [name for name in batch.schema.names if name not in excluded] unknown = _unknown_blob_descriptor_columns(batch, scalar_cols) if unknown: raise ValueError( @@ -247,10 +365,11 @@ def _map_blob_batch( "(likely from a joined BLOB table). Fetch it with its own " "table.map_with_blobs() in a separate pass, or drop it before " "mapping.".format(unknown[0])) + return scalar_cols - bodies = fetch_blob_bodies( - file_io, batch.select(blob_cols).to_pydict(), blob_cols, parallelism) - result = fn(batch.select(scalar_cols), bodies, **fn_kwargs) + +def _call_blob_fn(fn, scalar, bodies, fn_kwargs): + result = fn(scalar, bodies, **fn_kwargs) if result is None: raise ValueError( "map_with_blobs UDF must return a Ray-compatible batch, such as a " @@ -259,6 +378,41 @@ def _map_blob_batch( return result +def _blob_prefetch_windows(batch, blob_cols, fn_batch_size, max_bytes): + start = 0 + end = 0 + size = 0 + while end < batch.num_rows: + next_end = min(end + fn_batch_size, batch.num_rows) + next_size = _blob_payload_size( + batch.slice(end, next_end - end), blob_cols, max_bytes) + if end > start and size + next_size > max_bytes: + yield start, end + start = end + size = 0 + size += next_size + end = next_end + if end > start: + yield start, end + + +def _blob_payload_size(batch, blob_cols, unknown_size): + from pypaimon.table.row.blob import BlobDescriptor + + total = 0 + for name in blob_cols: + for value in batch.column(name): + if not value.is_valid: + continue + raw = value.as_py() + if BlobDescriptor.is_blob_descriptor(raw): + length = BlobDescriptor.deserialize(raw).length + total += length if length >= 0 else unknown_size + else: + total += len(raw) + return total + + def _unknown_blob_descriptor_columns(batch, scalar_cols): return [ name for name in scalar_cols diff --git a/paimon-python/pypaimon/tests/multimodal_table_test.py b/paimon-python/pypaimon/tests/multimodal_table_test.py index 2cfa1ad0ba9b..2bd7ec1a9915 100644 --- a/paimon-python/pypaimon/tests/multimodal_table_test.py +++ b/paimon-python/pypaimon/tests/multimodal_table_test.py @@ -1062,10 +1062,12 @@ def collect_batch(scalar, blobs, prefix): collect_batch, parallelism=2, batch_size=1, + blob_uri_affinity=True, + prefetch_bytes=32, fn_kwargs={"prefix": b"got-"}, ray_remote_args={"num_cpus": 1}, ) - rows = sorted(result.to_pandas().to_dict("records"), key=lambda row: row["idx"]) + rows = sorted(result.take_all(), key=lambda row: row["idx"]) self.assertEqual( [ @@ -1117,6 +1119,23 @@ def return_none(scalar, blobs): file_io=obs.raw_table.file_io, ) + with self.assertRaisesRegex(ValueError, "requires batch_size"): + obs.map_with_blobs( + ds, + ["image"], + return_none, + batch_size=None, + blob_uri_affinity=True, + ) + + with self.assertRaisesRegex(ValueError, "prefetch_bytes"): + obs.map_with_blobs( + ds, + ["image"], + return_none, + prefetch_bytes=0, + ) + with self.assertRaisesRegex(Exception, "must return"): obs.map_with_blobs( ds, diff --git a/paimon-python/pypaimon/tests/ray_blob_affinity_test.py b/paimon-python/pypaimon/tests/ray_blob_affinity_test.py new file mode 100644 index 000000000000..835f96591aa5 --- /dev/null +++ b/paimon-python/pypaimon/tests/ray_blob_affinity_test.py @@ -0,0 +1,164 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest + +import pyarrow as pa + +try: + import ray +except ImportError: + ray = None + +from pypaimon.ray.ray_paimon import ( + _append_blob_affinity_keys, + _blob_prefetch_windows, +) +from pypaimon.table.row.blob import BlobDescriptor + + +def _descriptor(uri, offset, length): + return BlobDescriptor(uri, offset, length).serialize() + + +class _ReadCounter: + def __init__(self): + self.reads = 0 + + def add(self, count): + self.reads += count + + def get(self): + return self.reads + + def reset(self): + self.reads = 0 + + +class _CountingFileIO: + def __init__(self, counter): + self.counter = counter + + def read_ranges_coalesced(self, ranges, parallelism): + paths = {value[0] for value in ranges if value is not None} + ray.get(self.counter.add.remote(len(paths))) + return [ + None if value is None else bytes([value[1] + 1]) * value[2] + for value in ranges + ] + + +class BlobAffinityHelperTest(unittest.TestCase): + def test_appends_uri_and_offset(self): + batch = pa.table({ + "id": [1, 2, 3], + "thumbnail": [None, None, None], + "image": [ + _descriptor("oss://bucket/a", 20, 2), + None, + b"inline", + ], + }) + + result = _append_blob_affinity_keys( + batch, ["thumbnail", "image"], "key", "offset") + + keys = result.column("key").to_pylist() + self.assertEqual(len(keys[0]), 16) + self.assertEqual(keys[1], b"\0" * 16) + self.assertEqual(keys[2], b"\0" * 16) + self.assertEqual(result.column("offset").to_pylist(), [20, -1, -1]) + + def test_prefetch_windows_end_on_function_batch_boundaries(self): + batch = pa.table({ + "image": [ + _descriptor("oss://bucket/a", i * 4, 4) + for i in range(5) + ], + }) + + windows = list(_blob_prefetch_windows( + batch, ["image"], fn_batch_size=2, max_bytes=8)) + + self.assertEqual(windows, [(0, 2), (2, 4), (4, 5)]) + + +@unittest.skipIf(ray is None, "ray is not installed") +class BlobAffinityRayTest(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.started_ray = not ray.is_initialized() + if cls.started_ray: + ray.init(ignore_reinit_error=True, num_cpus=2) + + @classmethod + def tearDownClass(cls): + if cls.started_ray: + ray.shutdown() + + def test_uri_affinity_coalesces_across_function_batches(self): + from pypaimon.ray import map_with_blobs + + counter = ray.remote(num_cpus=0)(_ReadCounter).remote() + file_io = _CountingFileIO(counter) + source = pa.table({ + "id": [1, 2, 3, 4], + "image": [ + _descriptor("oss://bucket/a", 0, 1), + _descriptor("oss://bucket/b", 0, 1), + _descriptor("oss://bucket/a", 1, 1), + _descriptor("oss://bucket/b", 1, 1), + ], + }) + + def consume(scalar, blobs): + return pa.table({ + "id": scalar.column("id"), + "image_size": [len(value) for value in blobs["image"]], + "fn_batch_size": [scalar.num_rows] * scalar.num_rows, + }) + + baseline = map_with_blobs( + ray.data.from_arrow(source), + ["image"], + consume, + file_io=file_io, + all_blob_columns=["image"], + batch_size=1, + ) + self.assertEqual(len(baseline.take_all()), 4) + self.assertEqual(ray.get(counter.get.remote()), 4) + + ray.get(counter.reset.remote()) + clustered = map_with_blobs( + ray.data.from_arrow(source), + ["image"], + consume, + file_io=file_io, + all_blob_columns=["image"], + batch_size=1, + blob_uri_affinity=True, + prefetch_bytes=16, + ) + rows = sorted(clustered.take_all(), key=lambda row: row["id"]) + + self.assertEqual([row["id"] for row in rows], [1, 2, 3, 4]) + self.assertEqual([row["fn_batch_size"] for row in rows], [1, 1, 1, 1]) + self.assertEqual(ray.get(counter.get.remote()), 2) + + +if __name__ == "__main__": + unittest.main()