diff --git a/docs/docs/pypaimon/multimodal-reading.md b/docs/docs/pypaimon/multimodal-reading.md index e00aef81718a..9db933334b66 100644 --- a/docs/docs/pypaimon/multimodal-reading.md +++ b/docs/docs/pypaimon/multimodal-reading.md @@ -175,6 +175,10 @@ for batch in aligned.to_arrow_batch_reader(batch_size=128): train(batch) ``` +Set `parallelism=4` in `to_arrow_batch_reader()` to plan and read sources concurrently +(default: 1). Output order is preserved; per-table reader concurrency still applies. +Custom catalog and authorization callbacks must be thread-safe. + `direction` is `backward`, `forward`, or `nearest`; tolerance is inclusive and zero permits only exact matches. Nearest ties use the earlier time. For duplicate timestamps, backward uses the last row and forward uses the first. Nearest uses the last diff --git a/paimon-python/pypaimon/multimodal/temporal.py b/paimon-python/pypaimon/multimodal/temporal.py index 4cedf35238bc..4ae151d71a69 100644 --- a/paimon-python/pypaimon/multimodal/temporal.py +++ b/paimon-python/pypaimon/multimodal/temporal.py @@ -17,12 +17,14 @@ """Temporal alignment for multimodal table scans.""" from bisect import bisect_left, bisect_right +from concurrent.futures import FIRST_EXCEPTION, ThreadPoolExecutor, wait from datetime import timedelta from fractions import Fraction -from functools import lru_cache +from functools import lru_cache, partial import json import math from numbers import Integral, Real +from threading import Event import numpy as np import pyarrow as pa @@ -193,24 +195,29 @@ def _append(self, source): result.schema = result._output_schema() return result - def to_arrow_batch_reader(self, *, batch_size=1024): - """Index right-side timestamps, then stream aligned rows in batches.""" + def to_arrow_batch_reader(self, *, batch_size=1024, parallelism=1): + """Stream aligned batches; ``parallelism`` limits concurrent sources (default 1).""" if (isinstance(batch_size, bool) or not isinstance(batch_size, int) or batch_size <= 0): raise ValueError("batch_size must be a positive integer.") + if (isinstance(parallelism, bool) + or not isinstance(parallelism, int) or parallelism <= 0): + raise ValueError("parallelism must be a positive integer.") + anchor_metadata = _metadata_batches( self._anchor, self._on, self._by, batch_size) row_group_cache = _ParquetRowGroupCacheRequest( _TEMPORAL_ROW_GROUP_CACHE_MAX_SIZE) - anchor_fetcher = _RowIdFetcher(self._anchor, row_group_cache) - source_fetchers = [] - for source in self._sources: + + def prepare_source(source): source.plan() - source_fetchers.append( - _RowIdFetcher( - source.query, row_group_cache, source._fetch_names)) + return _RowIdFetcher(source.query, row_group_cache, source._fetch_names) + + tasks = [partial(_RowIdFetcher, self._anchor, row_group_cache)] + tasks.extend(partial(prepare_source, source) for source in self._sources) + anchor_fetcher, *source_fetchers = _run_source_tasks(tasks, parallelism) schema = self._output_schema(anchor_fetcher.schema, source_fetchers) self.schema = schema @@ -218,7 +225,7 @@ def batches(): try: for metadata in anchor_metadata: yield from self._build_batches( - metadata, anchor_fetcher, source_fetchers, schema) + metadata, anchor_fetcher, source_fetchers, schema, parallelism) finally: anchor_metadata.close() @@ -308,7 +315,7 @@ def _output_schema(self, anchor_schema=None, source_fetchers=None): return pa.schema(fields, metadata=anchor_schema.metadata) def _build_batch( - self, anchor_rows, anchor_fetcher, source_fetchers, schema): + self, anchor_rows, anchor_fetcher, source_fetchers, schema, parallelism=1): native_keys = None if isinstance(anchor_rows, pa.Table): time_type = self._anchor_table_schema.field(self._on).type @@ -318,12 +325,14 @@ def _build_batch( anchor_rows = _metadata_rows(anchor_rows, self._on, time_type) else: anchor_ids = [row[_ROW_ID] for row in anchor_rows] - anchor = anchor_fetcher.fetch(anchor_ids) + tasks = [partial(anchor_fetcher.fetch, anchor_ids)] + tasks.extend(partial(source.build_arrays, anchor_rows, fetcher, native_keys) + for source, fetcher in zip(self._sources, source_fetchers)) + anchor, *source_arrays = _run_source_tasks(tasks, parallelism) anchor.validate() arrays = [anchor[name] for name in self._anchor_schema.names] - - for source, fetcher in zip(self._sources, source_fetchers): - arrays.extend(source.build_arrays(anchor_rows, fetcher, native_keys)) + for columns in source_arrays: + arrays.extend(columns) if arrays: table = pa.Table.from_arrays( @@ -334,20 +343,52 @@ def _build_batch( return batch.replace_schema_metadata(schema.metadata) def _build_batches( - self, anchor_rows, anchor_fetcher, source_fetchers, schema): + self, anchor_rows, anchor_fetcher, source_fetchers, schema, parallelism=1): try: yield self._build_batch( - anchor_rows, anchor_fetcher, source_fetchers, schema) + anchor_rows, anchor_fetcher, source_fetchers, schema, parallelism) except pa.ArrowInvalid as error: if len(anchor_rows) < 2 or "offset" not in str(error).lower(): raise middle = len(anchor_rows) // 2 yield from self._build_batches( anchor_rows[:middle], anchor_fetcher, - source_fetchers, schema) + source_fetchers, schema, parallelism) yield from self._build_batches( anchor_rows[middle:], anchor_fetcher, - source_fetchers, schema) + source_fetchers, schema, parallelism) + + +def _run_source_tasks(tasks, parallelism): + if parallelism == 1 or len(tasks) < 2: + return [task() for task in tasks] + failed = Event() + + def run(task): + # A worker may dequeue another task before the caller observes failure. + if failed.is_set(): + return None + try: + return task() + except BaseException: + failed.set() + raise + + with ThreadPoolExecutor(max_workers=min(parallelism, len(tasks))) as executor: + futures = [] + try: + for task in tasks: + futures.append(executor.submit(run, task)) + done, _ = wait(futures, return_when=FIRST_EXCEPTION) + for future in done: + if future.exception() is not None: + future.result() + return [future.result() for future in futures] + except BaseException: + failed.set() + for future in futures: + future.cancel() + raise class _AsOfJoinRight: diff --git a/paimon-python/pypaimon/tests/multimodal_temporal_test.py b/paimon-python/pypaimon/tests/multimodal_temporal_test.py index 8e95268f6dfa..318cf63b3e3f 100644 --- a/paimon-python/pypaimon/tests/multimodal_temporal_test.py +++ b/paimon-python/pypaimon/tests/multimodal_temporal_test.py @@ -19,6 +19,7 @@ import shutil import sys import tempfile +import threading import tracemalloc import unittest from datetime import datetime, timedelta @@ -85,6 +86,171 @@ def setUp(self): def tearDown(self): shutil.rmtree(self.temp_dir, ignore_errors=True) + def _parallel_alignment(self): + tables = [] + for name in ("a", "b", "c", "d"): + table = self._table("parallel_" + name, { + "episode_id": pa.int32(), "event_time": pa.int64(), name: pa.int64()}) + table.add([{"episode_id": 1, "event_time": i, name: i} for i in range(3)]) + tables.append(table) + aligned = pmm.join_asof( + tables[0].scan(), tables[1].scan(), + on="event_time", by="episode_id") + return aligned.join_asof(tables[2].scan()).join_asof(tables[3].scan()) + + def test_parallel_sources_preserve_results_and_bound_reads(self): + aligned = self._parallel_alignment() + reader = aligned.to_arrow_batch_reader(parallelism=1) + try: + expected = reader.read_all() + finally: + reader.close() + barrier = threading.Barrier(2) + lock = threading.Lock() + active = [0, 0] + original = temporal._RowIdFetcher.fetch + + def fetch(fetcher, row_ids): + with lock: + active[0] += 1 + active[1] = max(active) + try: + barrier.wait(timeout=10) + return original(fetcher, row_ids) + finally: + with lock: + active[0] -= 1 + + with mock.patch.object(temporal._RowIdFetcher, "fetch", fetch): + reader = aligned.to_arrow_batch_reader(batch_size=2, parallelism=2) + try: + actual = reader.read_all() + finally: + reader.close() + self.assertEqual(expected, actual) + self.assertEqual([0, 2], active) + + def test_parallel_sources_plan_independently(self): + aligned = self._parallel_alignment() + barrier = threading.Barrier(3) + original = temporal._AsOfJoinRight.plan + + def plan(source): + barrier.wait(timeout=10) + return original(source) + + with mock.patch.object(temporal._AsOfJoinRight, "plan", plan): + reader = aligned.to_arrow_batch_reader(parallelism=4) + try: + self.assertEqual(3, len(reader.read_all())) + finally: + reader.close() + + def test_parallel_payload_failure_joins_workers_and_closes_input(self): + aligned = self._parallel_alignment() + barrier = threading.Barrier(4) + lock = threading.Lock() + finished = [] + closed = [] + original_fetch = temporal._RowIdFetcher.fetch + original_metadata = temporal._metadata_batches + + def metadata(*args): + try: + yield from original_metadata(*args) + finally: + closed.append(True) + + def fetch(fetcher, row_ids): + try: + barrier.wait(timeout=10) + if "b" in fetcher.schema.names: + raise RuntimeError("payload failed") + return original_fetch(fetcher, row_ids) + finally: + with lock: + finished.append(True) + + with mock.patch.object(temporal, "_metadata_batches", metadata), \ + mock.patch.object(temporal._RowIdFetcher, "fetch", fetch): + reader = aligned.to_arrow_batch_reader(parallelism=4) + try: + with self.assertRaisesRegex(Exception, "payload failed"): + reader.read_all() + self.assertEqual(4, len(finished)) + self.assertEqual([True], closed) + finally: + reader.close() + + def test_parallel_planning_failure_joins_other_sources(self): + aligned = self._parallel_alignment() + barrier = threading.Barrier(3) + finished = [] + original = temporal._AsOfJoinRight.plan + + def plan(source): + try: + barrier.wait(timeout=10) + if "b" in source.payload_schema.names: + raise RuntimeError("planning failed") + return original(source) + finally: + finished.append(source.label) + + with mock.patch.object(temporal._AsOfJoinRight, "plan", plan): + with self.assertRaisesRegex(RuntimeError, "planning failed"): + aligned.to_arrow_batch_reader(parallelism=4) + self.assertEqual(3, len(finished)) + + def test_source_failure_skips_queued_work_before_slow_source_finishes(self): + slow_started = threading.Event() + release_slow = threading.Event() + slow_finished = threading.Event() + queued_started = threading.Event() + original_wait = temporal.wait + failure = RuntimeError("source failed") + + def slow(): + slow_started.set() + try: + if not release_slow.wait(timeout=10): + raise AssertionError("Slow source was not released") + finally: + slow_finished.set() + + def fail(): + if not slow_started.wait(timeout=10): + raise AssertionError("Slow source did not start") + raise failure + + def observe(futures, **kwargs): + try: + # Let the worker dequeue the third wrapper before cancellation. + self.assertIsNone(futures[2].result(timeout=10)) + self.assertFalse(queued_started.is_set()) + self.assertFalse(slow_finished.is_set()) + done, pending = original_wait(futures, **kwargs) + self.assertIn(futures[1], done) + self.assertIn(futures[0], pending) + return done, pending + finally: + release_slow.set() + + with mock.patch.object(temporal, "wait", side_effect=observe) as waiting: + with self.assertRaises(RuntimeError) as error: + temporal._run_source_tasks([slow, fail, queued_started.set], 2) + self.assertIs(failure, error.exception) + waiting.assert_called_once() + self.assertTrue(slow_finished.is_set()) + self.assertFalse(queued_started.is_set()) + + def test_temporal_parallelism_validation(self): + aligned = self._parallel_alignment() + for value in (True, 0, -1, 1.5, None): + with self.subTest(value=value): + with self.assertRaisesRegex(ValueError, "parallelism"): + aligned.to_arrow_batch_reader(parallelism=value) + def test_metadata_batches_coalesce_without_reading_ahead(self): metadata = pa.table({"episode_id": [1] * 11, "event_time": list(range(11)),