This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new 872f7e8bae [python][ray] Reduce Blob request QPS with URI affinity
(#9130)
872f7e8bae is described below
commit 872f7e8bae83e8cb400303ee703269a4d125053b
Author: XiaoHongbo <[email protected]>
AuthorDate: Wed Sep 23 22:46:07 2026 +0800
[python][ray] Reduce Blob request QPS with URI affinity (#9130)
---
docs/docs/pypaimon/multimodal-reading.md | 11 +
paimon-python/pypaimon/ray/ray_paimon.py | 205 ++++++++++++++++-
.../pypaimon/tests/multimodal_table_test.py | 21 +-
.../pypaimon/tests/ray_blob_affinity_test.py | 250 +++++++++++++++++++++
4 files changed, 475 insertions(+), 12 deletions(-)
diff --git a/docs/docs/pypaimon/multimodal-reading.md
b/docs/docs/pypaimon/multimodal-reading.md
index fd3685ba27..e4b02f70dc 100644
--- a/docs/docs/pypaimon/multimodal-reading.md
+++ b/docs/docs/pypaimon/multimodal-reading.md
@@ -497,3 +497,14 @@ Best practices:
- Use business keys or application columns when writing inference or training
outputs back to a table.
- Reorder `take_row_ids` results client-side when input order matters.
+
+### BLOB URI affinity
+
+For small inference batches, `map_with_blobs(..., batch_size=32,
+blob_uri_affinity=True)` sorts scalar BLOB descriptors by URI and offset to
+coalesce reads across inference batches. This distributed sort may reorder
rows;
+use it when fewer storage requests outweigh the shuffle cost. `prefetch_bytes`
+defaults to 64 MiB per payload window, except for an oversized inference batch.
+Descriptor batches contain up to `max(1024, batch_size)` rows; the UDF still
+receives at most `batch_size` rows. MAP and ARRAY BLOB columns use the default
+mapping path without URI affinity.
diff --git a/paimon-python/pypaimon/ray/ray_paimon.py
b/paimon-python/pypaimon/ray/ray_paimon.py
index 58358fe81c..f14bc5a0ae 100644
--- a/paimon-python/pypaimon/ray/ray_paimon.py
+++ b/paimon-python/pypaimon/ray/ray_paimon.py
@@ -26,7 +26,9 @@ Usage::
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
@@ -151,6 +153,8 @@ def map_with_blobs(
array_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,
@@ -167,7 +171,10 @@ def map_with_blobs(
Supply ``map_blob_columns`` and ``array_blob_columns`` when transforms
erase the source metadata and Arrow nested types. The table method
supplies this information automatically.
- 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()
@@ -183,6 +190,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:
@@ -198,7 +213,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)
@@ -225,19 +240,90 @@ def map_with_blobs(
or (map_blob_columns | array_blob_columns) - all_blob):
raise ValueError("Nested BLOB columns must be disjoint subsets of
all_blob_columns.")
+ mapper = _map_blob_batch
+ affinity_cols = []
+ if blob_uri_affinity:
+ if (map_blob_columns | array_blob_columns).intersection(blob_cols):
+ raise ValueError("blob_uri_affinity supports scalar BLOB columns
only")
+ dataset, affinity_cols = _cluster_by_blob_uri(dataset, blob_cols)
+ mapper = _map_blob_affinity_block
+ # Batch descriptors separately from inference; Ray requires an explicit
+ # outer batch size when GPU resources are requested.
+ kwargs["batch_size"] = max(1024, batch_size)
+
+ mapper_kwargs = {
+ "file_io": resolved_file_io,
+ "blob_cols": blob_cols,
+ "all_blob_cols": list(all_blob_cols),
+ "map_blob_cols": list(map_blob_columns),
+ "array_blob_cols": list(array_blob_columns),
+ "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),
- "map_blob_cols": list(map_blob_columns),
- "array_blob_cols": list(array_blob_columns),
- "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 BlobDescriptorSerde
+
+ empty_key = b"\0" * 16
+ uri_keys = {}
+ keys = []
+ offsets = []
+ columns = [batch.column(name) for name in blob_cols]
+ if any(not (pa.types.is_binary(column.type) or
pa.types.is_large_binary(column.type)
+ or pa.types.is_null(column.type)) for column in columns):
+ raise ValueError("blob_uri_affinity supports scalar BLOB columns only")
+ 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 BlobDescriptorSerde.is_descriptor(raw):
+ descriptor = BlobDescriptorSerde.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):
@@ -303,6 +389,103 @@ def _map_blob_batch(
return result
+def _map_blob_affinity_block(
+ batch, file_io, blob_cols, all_blob_cols, parallelism, fn, fn_kwargs,
+ fn_batch_size, prefetch_bytes, affinity_cols,
+ map_blob_cols=(), array_blob_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)
+ del fn_bodies, bodies
+
+
+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)
+ 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(
+ "Column {!r} holds BLOB descriptors this table does not own "
+ "(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
+
+
+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 "
+ "pyarrow.Table. For side-effect-only processing, return an empty "
+ "pyarrow.Table instead of None.")
+ 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 BlobDescriptorSerde
+
+ total = 0
+ for name in blob_cols:
+ for value in batch.column(name):
+ if not value.is_valid:
+ continue
+ raw = value.as_py()
+ if BlobDescriptorSerde.is_descriptor(raw):
+ length = BlobDescriptorSerde.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 2e7e680f72..a61477f100 100644
--- a/paimon-python/pypaimon/tests/multimodal_table_test.py
+++ b/paimon-python/pypaimon/tests/multimodal_table_test.py
@@ -1856,10 +1856,12 @@ class MultimodalTableTest(unittest.TestCase):
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(
[
@@ -2086,6 +2088,23 @@ class MultimodalTableTest(unittest.TestCase):
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 0000000000..a63f2e628c
--- /dev/null
+++ b/paimon-python/pypaimon/tests/ray_blob_affinity_test.py
@@ -0,0 +1,250 @@
+# 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 weakref
+from unittest.mock import patch
+
+import pyarrow as pa
+
+try:
+ import ray
+except ImportError:
+ ray = None
+
+from pypaimon.ray.ray_paimon import (
+ _append_blob_affinity_keys,
+ _blob_prefetch_windows,
+ _map_blob_affinity_block,
+)
+from pypaimon.table.row.blob import BlobDescriptor, VideoFrameDescriptor
+
+
+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_video_frames_share_payload_uri_key_and_preserve_offsets(self):
+ batch = pa.table({"image": [
+ _descriptor("oss://bucket/a", 10, 100),
+ VideoFrameDescriptor("oss://bucket/a", 20, 100, 0).serialize(),
+ VideoFrameDescriptor("oss://bucket/a", 20, 100, 1).serialize(),
+ VideoFrameDescriptor("oss://bucket/b", 30, 100, 0).serialize(),
+ ]})
+ result = _append_blob_affinity_keys(batch, ["image"], "key", "offset")
+ keys = result.column("key").to_pylist()
+ self.assertEqual(keys[:3], [keys[0]] * 3)
+ self.assertNotEqual(keys[3], keys[0])
+ self.assertTrue(all(key != b"\0" * 16 for key in keys))
+ self.assertEqual(result.column("offset").to_pylist(), [10, 20, 20, 30])
+
+ def test_video_prefetch_windows_use_payload_length(self):
+ batch = pa.table({"image": [
+ VideoFrameDescriptor("oss://bucket/video-{}".format(i),
+ 0, 10 * 1024 * 1024, 0).serialize()
+ for i in range(3)
+ ]})
+ windows = list(_blob_prefetch_windows(
+ batch, ["image"], fn_batch_size=1, max_bytes=1024 * 1024))
+ self.assertEqual(windows, [(0, 1), (1, 2), (2, 3)])
+
+ def test_previous_payload_window_is_released_before_next_read(self):
+ class Payload(bytearray):
+ pass
+
+ refs = []
+ reads = []
+ test = self
+
+ class FileIO:
+ def read_ranges_coalesced(self, ranges, parallelism):
+ test.assertTrue(all(ref() is None for ref in refs))
+ payloads = [Payload(b"x") for _ in ranges]
+ refs.extend(weakref.ref(value) for value in payloads)
+ reads.append(len(ranges))
+ return payloads
+
+ batch = pa.table({
+ "id": list(range(4)),
+ "image": [_descriptor("file:///video", i, 1) for i in range(4)],
+ })
+ outputs = list(_map_blob_affinity_block(
+ batch, FileIO(), ["image"], ["image"], 1,
+ lambda scalar, blobs: scalar, {},
+ fn_batch_size=1, prefetch_bytes=2, affinity_cols=[]))
+ self.assertEqual(reads, [2, 2])
+ self.assertEqual([result.num_rows for result in outputs], [1] * 4)
+ self.assertTrue(all(ref() is None for ref in refs))
+
+ def test_affinity_rejects_nested_columns(self):
+ for data_type, value in [
+ (pa.list_(pa.binary()), [b"image"]),
+ (pa.map_(pa.string(), pa.binary()), [("key", b"image")])]:
+ batch = pa.table({"image": pa.array([value], type=data_type)})
+ with self.assertRaisesRegex(ValueError, "scalar BLOB"):
+ _append_blob_affinity_keys(batch, ["image"], "key", "offset")
+
+ 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)])
+
+
[email protected](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_gpu_arguments_keep_explicit_outer_and_inner_batch_sizes(self):
+ from pypaimon.ray import map_with_blobs
+
+ source = ray.data.from_arrow(pa.table({"image": [b"inline"]}))
+ original = ray.data.Dataset.map_batches
+ calls = []
+
+ def capture(dataset, fn, **kwargs):
+ calls.append(kwargs)
+ return original(dataset, fn, **kwargs)
+
+ for resources in ({"num_gpus": 1}, {"ray_remote_args": {"num_gpus":
1}}):
+ with self.subTest(resources=resources), \
+ patch.object(ray.data.Dataset, "map_batches", capture):
+ # Construct a real Ray plan without scheduling GPU work.
+ map_with_blobs(
+ source, ["image"], lambda scalar, blobs: scalar,
+ file_io=object(), all_blob_columns=["image"],
+ batch_size=32, blob_uri_affinity=True, **resources)
+ self.assertGreater(calls[-1]["batch_size"], 32)
+ self.assertEqual(calls[-1]["fn_kwargs"]["fn_batch_size"], 32)
+ self.assertEqual(calls[-1]["num_gpus"], 1)
+
+ 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()
