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 0a6de28448 [python] Vectorize temporal alignment row ID fetching
(#10843)
0a6de28448 is described below
commit 0a6de2844838f85d6dc655ccf1b4a63c4d2d9fbf
Author: XiaoHongbo <[email protected]>
AuthorDate: Wed Oct 28 15:15:44 2026 +0800
[python] Vectorize temporal alignment row ID fetching (#10843)
---
paimon-python/pypaimon/multimodal/temporal.py | 42 ++++++++++++++--------
.../pypaimon/tests/multimodal_temporal_test.py | 40 +++++++++++++++++++++
2 files changed, 68 insertions(+), 14 deletions(-)
diff --git a/paimon-python/pypaimon/multimodal/temporal.py
b/paimon-python/pypaimon/multimodal/temporal.py
index ca2d299dbe..1affc09c5c 100644
--- a/paimon-python/pypaimon/multimodal/temporal.py
+++ b/paimon-python/pypaimon/multimodal/temporal.py
@@ -23,6 +23,7 @@ import json
import math
from numbers import Integral, Real
+import numpy as np
import pyarrow as pa
import pyarrow.compute as pc
@@ -1238,14 +1239,21 @@ class _RowIdFetcher:
return Range.sort_and_merge_overlap(ranges, True)
def fetch(self, row_ids):
- if not row_ids:
+ if len(row_ids) == 0:
return pa.Table.from_arrays(
[pa.array([], type=field.type) for field in self._schema],
schema=self._schema,
)
- wanted = Range.sort_and_merge_overlap(
- [Range(row_id, row_id) for row_id in set(row_ids)], True)
+ requested = np.asarray(row_ids, dtype=np.int64)
+ unique = np.unique(requested)
+ boundaries = np.flatnonzero(unique[1:] - unique[:-1] != 1) + 1
+ starts = np.concatenate(([0], boundaries))
+ ends = np.concatenate((boundaries - 1, [len(unique) - 1]))
+ wanted = [
+ Range(int(unique[start]), int(unique[end]))
+ for start, end in zip(starts, ends)
+ ]
selected_splits = []
for split_index in self._find_splits(wanted):
original = self._splits[split_index]
@@ -1265,21 +1273,27 @@ class _RowIdFetcher:
selected_splits.append(indexed)
arrow = self._project_fetch(self._read.to_arrow(selected_splits))
- found = arrow[self._row_id_name].to_pylist()
- positions = {}
- for index, row_id in enumerate(found):
- if row_id in positions:
- raise RuntimeError(
- "Duplicate row id %r in aligned scan." % row_id)
- positions[row_id] = index
- missing = [row_id for row_id in row_ids if row_id not in positions]
- if missing:
+ found = arrow[self._row_id_name]
+ if found.null_count:
+ raise RuntimeError("Null row id in aligned scan.")
+ found = found.to_numpy()
+ order = None
+ if len(found) > 1 and not np.all(found[1:] > found[:-1]):
+ order = np.argsort(found)
+ found = found[order]
+ if np.any(found[1:] == found[:-1]):
+ raise RuntimeError("Duplicate row id in aligned scan.")
+ positions = np.searchsorted(found, requested)
+ present = positions < len(found)
+ if len(found):
+ present &= found[np.minimum(positions, len(found) - 1)] ==
requested
+ if not np.all(present):
raise RuntimeError(
"Aligned row ids disappeared from pinned snapshot: %r."
- % missing
+ % requested[~present].tolist()
)
take = pa.array(
- [positions[row_id] for row_id in row_ids], type=pa.int64())
+ positions if order is None else order[positions], type=pa.int64())
visible = pa.Table.from_arrays(
[arrow.column(index) for index in range(len(self._schema))],
schema=self._schema,
diff --git a/paimon-python/pypaimon/tests/multimodal_temporal_test.py
b/paimon-python/pypaimon/tests/multimodal_temporal_test.py
index 8c5199eeff..b65f5f7309 100644
--- a/paimon-python/pypaimon/tests/multimodal_temporal_test.py
+++ b/paimon-python/pypaimon/tests/multimodal_temporal_test.py
@@ -36,6 +36,46 @@ from pypaimon.read.table_scan import TableScan
class MultimodalTemporalTest(unittest.TestCase):
+ def test_row_id_fetch_preserves_request_order(self):
+ for found in ([1, 2, 3, 9], [9, 2, 1, 3], [2**63 - 1, 0]):
+ for requested in (list(reversed(found)), [found[0]] * 3):
+ with self.subTest(found=found, requested=requested):
+ fetcher = self._row_id_fetcher(found)
+ result = fetcher.fetch(np.asarray(requested,
dtype=np.int64))
+ self.assertEqual(result["value"].to_pylist(), requested)
+
+ def test_row_id_fetch_keeps_exact_ranges(self):
+ fetcher = self._row_id_fetcher([1, 2, 4, 5, 9])
+ fetcher.fetch([9, 1, 2, 4, 5, 1])
+ fetcher._find_splits.assert_called_once_with([
+ temporal.Range(1, 2), temporal.Range(4, 5), temporal.Range(9, 9)])
+
+ def test_row_id_fetch_rejects_missing_and_duplicate_ids(self):
+ for found, requested, error in (
+ ([1, 1], [1], "Duplicate row id"),
+ ([1, 2], [3], "disappeared"),
+ ([], [1], "disappeared"),
+ ([None], [1], "Null row id")):
+ with self.subTest(found=found, requested=requested):
+ with self.assertRaisesRegex(RuntimeError, error):
+ self._row_id_fetcher(found).fetch(requested)
+ self.assertEqual(
+ self._row_id_fetcher([]).fetch(np.array([],
dtype=np.int64)).num_rows, 0)
+
+ @staticmethod
+ def _row_id_fetcher(ids):
+ fetcher = temporal._RowIdFetcher.__new__(temporal._RowIdFetcher)
+ fetcher._schema = pa.schema([("value", pa.int64())])
+ fetcher._row_id_name = temporal._ROW_ID
+ fetcher._name_paths = None
+ fetcher._find_splits = mock.Mock(return_value=[])
+ fetcher._read = mock.Mock()
+ fetcher._read.to_arrow.return_value = pa.table({
+ "value": pa.array(ids, type=pa.int64()),
+ temporal._ROW_ID: pa.array(ids, type=pa.int64()),
+ })
+ return fetcher
+
def setUp(self):
self.temp_dir = tempfile.mkdtemp(prefix="pypaimon_temporal_")
self.conn = pmm.connect(options={
