[python] Read independent temporal sources concurrently

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

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,32 +195,37 @@ 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
 
         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:

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)),