diff --git a/.github/workflows/ci-python.yml b/.github/workflows/ci-python.yml index 580a6a2a78a8..20dde453aa2c 100644 --- a/.github/workflows/ci-python.yml +++ b/.github/workflows/ci-python.yml @@ -228,12 +228,16 @@ jobs: import importlib.metadata as metadata import json + import pypaimon_rust import pypaimon_rust.datafusion distribution = metadata.distribution('pypaimon-rust') source = json.loads(distribution.read_text('direct_url.json')) print('pypaimon-rust', distribution.version, 'commit', source['vcs_info']['commit_id']) + assert source['vcs_info']['requested_revision'] == 'main' + assert callable(getattr(pypaimon_rust, 'temporal_match', None)), ( + 'Native CI requires the real temporal_match kernel') PY - name: Run Python tests with native planning, reading, writing, updating and committing diff --git a/paimon-python/pypaimon/multimodal/temporal.py b/paimon-python/pypaimon/multimodal/temporal.py index ca2d299dbe6f..3f4d9b601dd3 100644 --- a/paimon-python/pypaimon/multimodal/temporal.py +++ b/paimon-python/pypaimon/multimodal/temporal.py @@ -19,10 +19,12 @@ from bisect import bisect_left, bisect_right from datetime import timedelta from fractions import Fraction +from functools import lru_cache import json import math from numbers import Integral, Real +import numpy as np import pyarrow as pa import pyarrow.compute as pc @@ -215,11 +217,8 @@ def to_arrow_batch_reader(self, *, batch_size=1024): def batches(): try: for metadata in anchor_metadata: - rows = _metadata_rows( - metadata, self._on, - self._anchor_table_schema.field(self._on).type) yield from self._build_batches( - rows, anchor_fetcher, source_fetchers, schema) + metadata, anchor_fetcher, source_fetchers, schema) finally: anchor_metadata.close() @@ -310,13 +309,21 @@ def _output_schema(self, anchor_schema=None, source_fetchers=None): def _build_batch( self, anchor_rows, anchor_fetcher, source_fetchers, schema): - anchor_ids = [row[_ROW_ID] for row in anchor_rows] + native_keys = None + if isinstance(anchor_rows, pa.Table): + time_type = self._anchor_table_schema.field(self._on).type + native_keys = _native_anchor_keys(anchor_rows, self._on, self._by, time_type) + anchor_ids = anchor_rows[_ROW_ID].to_pylist() + if native_keys is None: + 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) 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)) + arrays.extend(source.build_arrays(anchor_rows, fetcher, native_keys)) if arrays: table = pa.Table.from_arrays( @@ -446,7 +453,21 @@ def output_fields(self, payload_schema, effective=True): for name in self.payload_schema.names ] - def build_arrays(self, anchor_rows, fetcher): + def build_arrays(self, anchor_rows, fetcher, native_keys=None): + native = _native_matches(self, native_keys, self.direction) + if native is not None: + first, _ = native + missing = first < 0 + ids = self._row_ids.to_numpy(zero_copy_only=False) + requested = (ids[np.maximum(first, 0)] if len(ids) + else np.zeros(len(first), dtype=np.int64)) + unique = np.unique(requested[~missing]) + values = fetcher.fetch(unique.tolist()) + take = pa.array(np.searchsorted(unique, requested), + mask=missing, type=pa.int64()) + return [pc.take(values[field.name], take) + for field in self.payload_schema] + anchor_rows = _reference_rows(anchor_rows, self) matches = [self.match(row) for row in anchor_rows] matched_ids = [match for match in matches if match is not None] unique_ids = list(dict.fromkeys(matched_ids)) @@ -474,6 +495,12 @@ def __init__(self, label, query, anchor_on, by, tolerance, label, query, anchor_on, by, "nearest", tolerance, right_on, suffix) + def plan(self): + super().plan() + times = self._time_keys + self._time_min = times.min() if len(times) else None + self._time_max = times.max() if len(times) else None + @staticmethod def output_field(field, effective=True): try: @@ -516,8 +543,28 @@ def match(self, anchor_row): *_linear_weight(target, before_time, after_time, self.time_type), ) - def build_arrays(self, anchor_rows, fetcher): - matches = [self.match(row) for row in anchor_rows] + def build_arrays(self, anchor_rows, fetcher, native_keys=None): + native = _native_matches(self, native_keys, "interpolate") + if native is None: + anchor_rows = _reference_rows(anchor_rows, self) + matches = [self.match(row) for row in anchor_rows] + else: + arrays = self._native_float_arrays(native_keys[0], fetcher, native) + if arrays is not None: + return arrays + first, second = native + matches = [] + for index, (before, after) in enumerate(zip(first, second)): + if before < 0: + matches.append(None) + continue + before_id = self._row_ids[int(before)].as_py() + after_id = self._row_ids[int(after)].as_py() + ratio = ((0.0, 0, 1) if before == after else _linear_weight( + _python_scalar(native_keys[0][index]), + _python_scalar(self._time_keys[before]), + _python_scalar(self._time_keys[after]), self.time_type)) + matches.append((before_id, after_id, *ratio)) matched_ids = [] for match in matches: if match is not None: @@ -556,6 +603,34 @@ def build_arrays(self, anchor_rows, fetcher): arrays.append(array) return arrays + def _native_float_arrays(self, targets, fetcher, indices): + if (not len(targets) or not len(self._row_ids) + or not all(pa.types.is_floating(fetcher.schema.field(field.name).type) + for field in self.payload_schema)): + return None + times = self._time_keys + # Keep exact-ratio interpolation for wide spans and integer payloads. + if (max(int(self._time_max), int(targets.max())) + - min(int(self._time_min), int(targets.min())) > 2 ** 53): + return None + first, second = indices + missing = first < 0 + first, second = np.maximum(first, 0), np.maximum(second, 0) + ids = self._row_ids.to_numpy(zero_copy_only=False) + before_ids, after_ids = ids[first], ids[second] + unique = np.unique(np.concatenate((before_ids[~missing], after_ids[~missing]))) + values = fetcher.fetch(unique.tolist()) + before = pa.array(np.searchsorted(unique, before_ids), mask=missing, type=pa.int64()) + after = pa.array(np.searchsorted(unique, after_ids), mask=missing, type=pa.int64()) + numerator = targets - times[first] + denominator = times[second] - times[first] + weights = np.zeros(len(targets), dtype=np.float64) + np.divide(numerator, denominator, out=weights, where=denominator != 0) + weights = pa.array(weights, mask=missing) + return [_interpolate_array(pc.take(values[field.name], before), + pc.take(values[field.name], after), weights, None) + for field in self.payload_schema] + class _WindowJoinRight(_AsOfJoinRight): @@ -652,8 +727,13 @@ def _match_range(self, anchor_row): ) return first, last - def build_arrays(self, anchor_rows, fetcher): - ranges = [self._match_range(row) for row in anchor_rows] + def build_arrays(self, anchor_rows, fetcher, native_keys=None): + native = _native_matches(self, native_keys, "window") + ranges = ([self._match_range(row) for row in + _reference_rows(anchor_rows, self)] + if native is None else [ + (0, 0) if start < 0 else (int(start), int(end)) + for start, end in zip(*native)]) slices = [(0, 0)] * len(ranges) unique_ids = [] covered_end = 0 @@ -1025,11 +1105,23 @@ def _metadata_batches(query, on, by, batch_size): return reader = read_builder.new_read()._to_managed_arrow_batch_reader(splits) try: + pending = [] + pending_rows = 0 for batch in reader: metadata = pa.Table.from_batches([batch]).select(output_columns) _validate_metadata(query, metadata, key_columns) - for start in range(0, len(metadata), batch_size): - yield metadata.slice(start, batch_size) + start = 0 + while start < len(metadata): + size = min(batch_size - pending_rows, len(metadata) - start) + pending.append(metadata.slice(start, size)) + pending_rows += size + start += size + if pending_rows == batch_size: + yield pa.concat_tables(pending) + pending = [] + pending_rows = 0 + if pending: + yield pa.concat_tables(pending) finally: reader.close() @@ -1323,6 +1415,75 @@ def _arrow_rows(table): ] +def _reference_rows(metadata, source): + return (_metadata_rows(metadata, source.anchor_on, source.time_type) + if isinstance(metadata, pa.Table) else metadata) + + +def _native_anchor_keys(metadata, on, by, time_type): + targets = _time_search_keys(metadata[on].combine_chunks(), time_type) + if targets.dtype != np.dtype("int64") or _temporal_match_kernel() is None: + return None + if len(by) == 1: + column = metadata[by[0]].combine_chunks() + data_type = column.type + if (pa.types.is_integer(data_type) + or pa.types.is_string(data_type) + or pa.types.is_large_string(data_type) + or pa.types.is_binary(data_type) + or pa.types.is_large_binary(data_type) + or pa.types.is_boolean(data_type)): + encoded = pc.dictionary_encode(column) + return (targets, + [(key,) for key in encoded.dictionary.to_pylist()], + encoded.indices.to_numpy(zero_copy_only=False)) + keys, positions, codes = [], {}, [] + for row in metadata.select(list(by)).to_pylist(): + key = tuple(row[name] for name in by) + if key not in positions: + positions[key] = len(keys) + keys.append(key) + codes.append(positions[key]) + return targets, keys, np.asarray(codes, dtype=np.int64) + + +@lru_cache(maxsize=1) +def _temporal_match_kernel(): + try: + import pypaimon_rust + except ImportError: + return None + return getattr(pypaimon_rust, "temporal_match", None) + + +def _native_matches(source, native_keys, mode): + if native_keys is None or source._time_keys.dtype != np.dtype("int64"): + return None + tolerance = source._tolerance_key + preceding = getattr(source, "_preceding_key", 0) + following = getattr(source, "_following_key", 0) + if any(value is not None and ( + not isinstance(value, Integral) or not 0 <= value <= _MAX_INT64) + for value in (tolerance, preceding, following)): + return None + kernel = _temporal_match_kernel() + if kernel is None: + return None + targets, keys, codes = native_keys + bounds = np.asarray([source._index.get(key, (0, 0)) for key in keys], + dtype=np.int64).reshape((-1, 2)) + closed = getattr(source, "closed", "both") + result = kernel( + pa.array(source._time_keys), pa.array(targets), + pa.array(bounds[:, 0][codes]), pa.array(bounds[:, 1][codes]), + mode, exact=source.allow_exact_matches, tolerance=tolerance, + preceding=preceding, following=following, + left_closed=closed in ("both", "left"), + right_closed=closed in ("both", "right")) + return tuple(column.to_numpy(zero_copy_only=False) + for column in result.columns) + + def _metadata_rows(table, on, time_type): rows = _arrow_rows(table) if pa.types.is_timestamp(time_type): diff --git a/paimon-python/pypaimon/tests/multimodal_temporal_test.py b/paimon-python/pypaimon/tests/multimodal_temporal_test.py index 8c5199eefff8..c08a507119c5 100644 --- a/paimon-python/pypaimon/tests/multimodal_temporal_test.py +++ b/paimon-python/pypaimon/tests/multimodal_temporal_test.py @@ -45,6 +45,136 @@ def setUp(self): def tearDown(self): shutil.rmtree(self.temp_dir, ignore_errors=True) + def test_metadata_batches_coalesce_without_reading_ahead(self): + metadata = pa.table({"episode_id": [1] * 11, + "event_time": list(range(11)), + temporal._ROW_ID: list(range(11))}) + consumed = [] + closed = [] + + class Reader: + def __iter__(self): + for start, size in ((0, 2), (2, 0), (2, 3), (5, 6)): + consumed.append(start) + if size: + yield metadata.slice(start, size).to_batches()[0] + else: + yield pa.RecordBatch.from_arrays( + [pa.array([], type=field.type) for field in metadata.schema], + schema=metadata.schema) + + def close(self): + closed.append(True) + + builder = mock.Mock() + builder.new_read.return_value._to_managed_arrow_batch_reader.return_value = Reader() + query = mock.Mock(_temporal_empty=False) + with mock.patch.object(temporal, "_metadata_builders", return_value=( + builder, [], ["episode_id", "event_time"], metadata.column_names)): + batches = temporal._metadata_batches(query, "event_time", ("episode_id",), 4) + first = next(batches) + self.assertEqual([0, 1, 2, 3], first["event_time"].to_pylist()) + self.assertEqual([0, 2, 2], consumed) + remaining = list(batches) + self.assertEqual([4, 3], [len(batch) for batch in remaining]) + self.assertEqual(metadata, pa.concat_tables([first] + remaining)) + self.assertEqual([True], closed) + + def test_native_interpolation_scans_right_time_bounds_only_during_plan(self): + reductions = [] + + class TrackedTimes(np.ndarray): + def min(self, *args, **kwargs): + reductions.append("min") + return super().min(*args, **kwargs) + + def max(self, *args, **kwargs): + reductions.append("max") + return super().max(*args, **kwargs) + + right = temporal._LinearInterpolationRight.__new__( + temporal._LinearInterpolationRight) + right.query = mock.Mock() + right.on = "event_time" + right.by = ("episode_id",) + right.time_type = pa.int64() + right.payload_schema = pa.schema([("value", pa.float64())]) + row_ids = np.arange(1_000_000, dtype=np.int64) + times = (row_ids * 2).view(TrackedTimes) + metadata = pa.table({"episode_id": np.zeros(len(times), dtype=np.int32), + "event_time": times, temporal._ROW_ID: row_ids}) + fetcher = mock.Mock(schema=right.payload_schema) + fetcher.fetch.side_effect = lambda ids: pa.table({ + "value": pa.array(ids, type=pa.float64())}) + with mock.patch.object(temporal, "_metadata_table", return_value=metadata), \ + mock.patch.object(temporal, "_time_search_keys", return_value=times): + right.plan() + self.assertEqual(["min", "max"], reductions) + for batch_size in (32, 1024): + for start in (0, 1024, 2048): + before = np.arange(start, start + batch_size, dtype=np.int64) + targets = before * 2 + 1 + arrays = right._native_float_arrays( + targets, fetcher, (before, before + 1)) + self.assertEqual((before + 0.5).tolist(), arrays[0].to_pylist()) + # Wide target spans still fall back without scanning the right side. + self.assertIsNone(right._native_float_arrays( + np.array([2 ** 54], dtype=np.int64), fetcher, + (np.array([0]), np.array([1])))) + self.assertEqual(["min", "max"], reductions) + + def test_native_anchor_keys_preserve_timestamp_precision(self): + times = pa.array([1, 2, 3], type=pa.timestamp("ns")) + metadata = pa.table({"episode_id": [1, 1, 2], "event_time": times, + temporal._ROW_ID: [8, 9, 10]}) + with mock.patch.object(temporal, "_temporal_match_kernel", return_value=object()): + targets, keys, codes = temporal._native_anchor_keys( + metadata, "event_time", ("episode_id",), times.type) + self.assertEqual([1, 2, 3], targets.tolist()) + self.assertEqual([(1,), (2,)], keys) + self.assertEqual([0, 0, 1], codes.tolist()) + + def test_native_temporal_kernel_is_optional(self): + for module in (None, mock.Mock(spec=[])): + with self.subTest(module=module): + temporal._temporal_match_kernel.cache_clear() + with mock.patch.dict(sys.modules, {"pypaimon_rust": module}): + self.assertIsNone(temporal._temporal_match_kernel()) + temporal._temporal_match_kernel.cache_clear() + + def test_native_ci_rejects_missing_temporal_kernel(self): + with mock.patch.dict(os.environ, {"PYPAIMON_TEST_NATIVE_READ": "1"}), \ + mock.patch.object(temporal, "_temporal_match_kernel", return_value=None): + with self.assertRaisesRegex(AssertionError, "Native CI requires temporal_match"): + self.test_native_temporal_matching_matches_reference() + + def test_native_temporal_matching_matches_reference(self): + kernel = temporal._temporal_match_kernel() + if os.environ.get("PYPAIMON_TEST_NATIVE_READ") == "1": + self.assertIsNotNone(kernel, "Native CI requires temporal_match") + if kernel is None: + self.skipTest("Requires pypaimon_rust temporal_match") + rng = np.random.default_rng(123) + for direction in ("backward", "forward", "nearest"): + for exact in (False, True): + for tolerance in (None, 0, 3): + source = object.__new__(temporal._AsOfJoinRight) + source.by = ("episode_id",) + source._time_keys = np.sort(rng.integers(-20, 20, size=40)) + source._row_ids = pa.array(range(40), type=pa.int64()) + source._index = {(1,): (0, 40)} + source.direction = direction + source.allow_exact_matches = exact + source._tolerance_key = tolerance + metadata = pa.table({"episode_id": [1] * 50 + [2], + "event_time": rng.integers(-30, 30, size=51)}) + native_keys = temporal._native_anchor_keys( + metadata, "event_time", ("episode_id",), pa.int64()) + first, _ = temporal._native_matches(source, native_keys, direction) + expected = [source.match(row) for row in + temporal._metadata_rows(metadata, "event_time", pa.int64())] + self.assertEqual(expected, [None if value < 0 else int(value) for value in first]) + def test_alignment_requires_an_explicit_group_boundary(self): table = self._table("missing_group", { "event_time": pa.int64(),