diff --git a/docs/docs/pypaimon/multimodal-api.mdx b/docs/docs/pypaimon/multimodal-api.mdx index e8732875b176..b987c3eb5f2f 100644 --- a/docs/docs/pypaimon/multimodal-api.mdx +++ b/docs/docs/pypaimon/multimodal-api.mdx @@ -709,6 +709,49 @@ result = ( ) ``` +### Temporal alignment + +Use `align` to assemble a bounded training timeline from independently sampled +tables. The anchor scan defines one output row per step. Named secondary scans +match within the same episode (or other `by` keys) by exact, backward, forward, +or nearest timestamp. + +```python +from datetime import timedelta +from pypaimon.multimodal import align, backward, nearest + +steps = align( + actions.scan().select(["episode_id", "event_time", "action"]), + on="event_time", + by="episode_id", + sources={ + "camera_left": nearest( + images.scan().where("camera = 'left'").select("image"), + tolerance=timedelta(milliseconds=20), + ), + "robot_state": backward( + topics.scan().where("topic = '/robot/state'").select("value"), + tolerance=timedelta(milliseconds=50), + ), + }, +) + +for batch in steps.to_arrow_batch_reader(batch_size=128): + train(batch) +``` + +Temporal keys are injected into the internal metadata scan, so secondary +queries may select only their payload columns. Output payload columns are named +`__`. Every source also emits `__valid`, +`__matched_time`, and signed `__time_delta` audit columns. +Tolerance is inclusive, and nearest ties choose the earlier row. + +Planning reads only grouping keys, timestamps, and row IDs. Selected payload +rows are then fetched in batches from the same pinned snapshots. BLOB values +remain descriptors, allowing the existing BLOB and video readers to fetch or +decode only selected samples. This bounded API does not interpolate values or +build training windows; those operations remain explicit downstream steps. + ### Reading BLOB columns `scan().read_blobs(column)` bulk-fetches a BLOB column's bytes for the filtered diff --git a/paimon-python/pypaimon/multimodal/__init__.py b/paimon-python/pypaimon/multimodal/__init__.py index b669267b41da..18c68f42e9a2 100644 --- a/paimon-python/pypaimon/multimodal/__init__.py +++ b/paimon-python/pypaimon/multimodal/__init__.py @@ -36,6 +36,14 @@ text_route, vector_route, ) +from pypaimon.multimodal.temporal import ( + AlignedScan, + align, + backward, + exact, + forward, + nearest, +) from pypaimon.multimodal.video import VideoFrameCollator from pypaimon.table.row.blob import Blob, BlobDescriptor, VideoFrameDescriptor from pypaimon.table.data_evolution_merge_into import ( @@ -49,6 +57,7 @@ "BlobDescriptor", "BlobObject", "BlobStore", + "AlignedScan", "Hdf5File", "Hdf5LoadResult", "MultimodalConnection", @@ -61,7 +70,12 @@ "VideoFrameCollator", "VideoFrameDescriptor", "connect", + "align", + "backward", + "exact", + "forward", "lit", + "nearest", "source_col", "target_col", "text_route", diff --git a/paimon-python/pypaimon/multimodal/temporal.py b/paimon-python/pypaimon/multimodal/temporal.py new file mode 100644 index 000000000000..82fd5f8f23c7 --- /dev/null +++ b/paimon-python/pypaimon/multimodal/temporal.py @@ -0,0 +1,457 @@ +# 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. + +"""Bounded temporal alignment for multimodal table scans.""" + +from bisect import bisect_left +from collections.abc import Mapping +from dataclasses import dataclass +from datetime import timedelta +from typing import Optional + +import pyarrow as pa +import pyarrow.compute as pc + +from pypaimon.common.options.core_options import CoreOptions +from pypaimon.common.predicate_builder import PredicateBuilder +from pypaimon.multimodal.query import ScanQuery +from pypaimon.schema.data_types import PyarrowFieldParser +from pypaimon.table.special_fields import SpecialFields + + +_ROW_ID = SpecialFields.ROW_ID.name + + +@dataclass(frozen=True) +class _MatchSpec: + query: ScanQuery + method: str + tolerance: object = None + on: Optional[str] = None + + +def exact(query, *, on=None): + """Match a secondary scan at exactly the anchor timestamp.""" + return _match_spec(query, "exact", None, on) + + +def backward(query, *, tolerance, on=None): + """Match the latest secondary row at or before the anchor timestamp.""" + return _match_spec(query, "backward", tolerance, on) + + +def forward(query, *, tolerance, on=None): + """Match the earliest secondary row at or after the anchor timestamp.""" + return _match_spec(query, "forward", tolerance, on) + + +def nearest(query, *, tolerance, on=None): + """Match the closest secondary row; equal-distance ties choose earlier.""" + return _match_spec(query, "nearest", tolerance, on) + + +def align(anchor, *, on, by, sources): + """Build a batch-streaming, episode-local temporal alignment. + + The anchor scan defines output rows. Each named secondary source + contributes at most one row according to its match policy. All scans are + pinned to their current snapshots when this object is constructed. + """ + if not isinstance(on, str) or not on: + raise ValueError("on must be a non-empty column name.") + if isinstance(by, str): + by = (by,) + else: + try: + by = tuple(by) + except TypeError as error: + raise ValueError("by must be a column name or sequence.") from error + if not by: + raise ValueError("align requires at least one grouping column in by.") + if (any(not isinstance(name, str) or not name for name in by) + or len(set(by)) != len(by)): + raise ValueError("by must contain unique, non-empty column names.") + if not isinstance(sources, Mapping): + raise TypeError("sources must be a mapping of names to match specs.") + if not sources: + raise ValueError("align requires at least one named secondary source.") + return AlignedScan(anchor, on, by, sources) + + +class AlignedScan: + """Executable result of :func:`align`.""" + + def __init__(self, anchor, on, by, sources): + self._anchor = _pin_scan(_require_scan(anchor, "anchor")) + self._on = on + self._by = by + self._sources = { + name: _PinnedSource(name, _pin_match_spec(spec, name), on, by) + for name, spec in sources.items() + } + self._anchor_schema = _query_schema(self._anchor) + self._anchor_table_schema = _table_schema(self._anchor) + self._validate_anchor() + self.schema = self._output_schema() + + def to_arrow_batch_reader(self, *, batch_size=1024): + """Plan scalar timestamps once, then fetch selected rows in batches.""" + if not isinstance(batch_size, int) or batch_size <= 0: + raise ValueError("batch_size must be a positive integer.") + + anchor_rows = _metadata_rows(self._anchor, self._on, self._by) + anchor_rows.sort(key=_metadata_sort_key(self._on, self._by)) + for source in self._sources.values(): + source.plan() + + def batches(): + for start in range(0, len(anchor_rows), batch_size): + rows = anchor_rows[start:start + batch_size] + yield self._build_batch(rows) + + return pa.RecordBatchReader.from_batches(self.schema, batches()) + + def to_arrow(self): + reader = self.to_arrow_batch_reader() + try: + return reader.read_all() + finally: + reader.close() + + def to_pandas(self): + return self.to_arrow().to_pandas() + + def to_list(self): + return self.to_arrow().to_pylist() + + def _validate_anchor(self): + _require_columns( + self._anchor_table_schema, self._by + (self._on,), "anchor") + anchor_type = self._anchor_table_schema.field(self._on).type + _delta_type(anchor_type) + for source in self._sources.values(): + if source.time_type != anchor_type: + raise TypeError( + "Anchor and source %r temporal columns must have the same " + "type; got %s and %s." + % (source.name, anchor_type, source.time_type) + ) + + def _output_schema(self): + fields = list(self._anchor_schema) + names = set(self._anchor_schema.names) + anchor_time_type = self._anchor_table_schema.field(self._on).type + for source in self._sources.values(): + for field in source.payload_schema: + output = pa.field( + source.output_name(field.name), field.type, nullable=True) + if output.name in names: + raise ValueError( + "Duplicate aligned column %r." % output.name) + fields.append(output) + names.add(output.name) + audit = [ + pa.field( + source.output_name("valid"), pa.bool_(), nullable=False), + pa.field( + source.output_name("matched_time"), + source.time_type, + nullable=True, + ), + pa.field( + source.output_name("time_delta"), + _delta_type(anchor_time_type), + nullable=True, + ), + ] + for field in audit: + if field.name in names: + raise ValueError( + "Duplicate aligned column %r." % field.name) + fields.append(field) + names.add(field.name) + return pa.schema(fields) + + def _build_batch(self, anchor_rows): + anchor_ids = [row[_ROW_ID] for row in anchor_rows] + anchor = _fetch_rows(self._anchor, anchor_ids) + arrays = [anchor[name] for name in self._anchor_schema.names] + + for source in self._sources.values(): + matches = [source.match(row) for row in anchor_rows] + matched_ids = [match[0] for match in matches if match is not None] + unique_ids = list(dict.fromkeys(matched_ids)) + values = _fetch_rows(source.spec.query, unique_ids) + positions = { + row_id: index for index, row_id in enumerate(unique_ids) + } + take = pa.array([ + None if match is None else positions[match[0]] + for match in matches + ], type=pa.int64()) + for field in source.payload_schema: + arrays.append(pc.take(values[field.name], take)) + + arrays.extend([ + pa.array( + [match is not None for match in matches], type=pa.bool_()), + pa.array( + [None if match is None else match[1] for match in matches], + type=source.time_type, + ), + pa.array( + [None if match is None else match[2] for match in matches], + type=_delta_type( + self._anchor_table_schema.field(self._on).type), + ), + ]) + + table = pa.Table.from_arrays( + arrays, schema=self.schema).combine_chunks() + return table.to_batches(max_chunksize=len(anchor_rows))[0] + + +class _PinnedSource: + + def __init__(self, name, spec, anchor_on, by): + if not name or "__" in name: + raise ValueError( + "Aligned source names must be non-empty and cannot contain " + "'__'.") + self.name = name + self.spec = spec + self.anchor_on = anchor_on + self.on = spec.on or anchor_on + self.by = by + table_schema = _table_schema(spec.query) + _require_columns(table_schema, by + (self.on,), "source %r" % name) + self.time_type = table_schema.field(self.on).type + _delta_type(self.time_type) + schema = _query_schema(spec.query) + self.payload_schema = pa.schema([ + field for field in schema + if field.name not in set(by + (self.on,)) + ]) + self._index = None + + def output_name(self, field): + return "%s__%s" % (self.name, field) + + def plan(self): + rows = _metadata_rows(self.spec.query, self.on, self.by) + groups = {} + for row in rows: + key = tuple(row[name] for name in self.by) + groups.setdefault(key, []).append((row[self.on], row[_ROW_ID])) + + self._index = {} + for key, values in groups.items(): + values.sort() + times = [value[0] for value in values] + if len(times) != len(set(times)): + raise ValueError( + "Source %r has duplicate timestamps within group %r." + % (self.name, key) + ) + self._index[key] = (times, [value[1] for value in values]) + + def match(self, anchor_row): + key = tuple(anchor_row[name] for name in self.by) + values = self._index.get(key) + if values is None: + return None + times, row_ids = values + target = anchor_row[self.anchor_on] + index = _match_index(times, target, self.spec.method) + if index is None: + return None + delta = times[index] - target + if (self.spec.tolerance is not None + and abs(delta) > self.spec.tolerance): + return None + return row_ids[index], times[index], delta + + +def _match_spec(query, method, tolerance, on): + _require_scan(query, method) + if method != "exact" and tolerance is None: + raise ValueError("%s requires a non-null tolerance." % method) + if tolerance is not None: + try: + zero = timedelta(0) if isinstance(tolerance, timedelta) else 0 + negative = tolerance < zero + except TypeError: + negative = False + if negative: + raise ValueError("tolerance must be non-negative.") + return _MatchSpec(query, method, tolerance, on) + + +def _pin_match_spec(spec, name): + if not isinstance(spec, _MatchSpec): + raise TypeError( + "Source %r must use exact(), backward(), forward(), or nearest()." + % name + ) + return _MatchSpec( + _pin_scan(spec.query), spec.method, spec.tolerance, spec.on) + + +def _require_scan(query, label): + if type(query) is not ScanQuery: + raise TypeError("%s must be a MultimodalTable.scan() query." % label) + return query + + +def _pin_scan(query): + table = query._table + options = table.options + if not options.row_tracking_enabled(False): + raise ValueError("align requires 'row-tracking.enabled' = 'true'.") + empty = False + if (options.scan_snapshot_id() is None + and options.scan_tag_name() is None): + snapshot = table.snapshot_manager().get_latest_snapshot() + if snapshot is None: + empty = True + else: + table = table.copy({ + CoreOptions.SCAN_SNAPSHOT_ID.key(): str(snapshot.id), + }) + pinned = ScanQuery(table) + pinned._predicate = query._predicate + pinned._projection = query._projection + pinned._limit = query._limit + pinned._include_row_id = query._include_row_id + pinned._temporal_empty = empty + return pinned + + +def _query_schema(query): + table = query._table.copy({CoreOptions.BLOB_AS_DESCRIPTOR.key(): "true"}) + builder = query._configured_read_builder(table) + return PyarrowFieldParser.from_paimon_schema(builder.read_type()) + + +def _table_schema(query): + return PyarrowFieldParser.from_paimon_schema(query._table.fields) + + +def _metadata_rows(query, on, by): + columns = list(dict.fromkeys(by + (on, _ROW_ID))) + if getattr(query, "_temporal_empty", False): + return [] + table = query._table.copy({CoreOptions.BLOB_AS_DESCRIPTOR.key(): "true"}) + builder = table.new_read_builder() + if query._predicate is not None: + builder = builder.with_filter(query._predicate) + for name in query._predicate_fields(): + if name not in columns: + columns.append(name) + builder = builder.with_projection(columns) + if query._limit is not None: + builder = builder.with_limit(query._limit) + arrow = builder.new_read().to_arrow(builder.new_scan().plan().splits()) + rows = arrow.select(list(by) + [on, _ROW_ID]).to_pylist() + for row in rows: + if row[on] is None or any(row[name] is None for name in by): + raise ValueError("Temporal keys cannot be null: %r." % row) + return rows + + +def _fetch_rows(query, row_ids): + schema = _query_schema(query) + if not row_ids: + return pa.Table.from_arrays( + [pa.array([], type=field.type) for field in schema], schema=schema) + + table = query._table.copy({CoreOptions.BLOB_AS_DESCRIPTOR.key(): "true"}) + visible = schema.names + projection = list(dict.fromkeys(visible + [_ROW_ID])) + for name in query._predicate_fields(): + if name not in projection: + projection.append(name) + builder = table.new_read_builder().with_projection(projection) + row_id_predicate = builder.new_predicate_builder().is_in(_ROW_ID, row_ids) + predicates = [row_id_predicate] + if query._predicate is not None: + predicates.insert(0, query._predicate) + builder = builder.with_filter(PredicateBuilder.and_predicates(predicates)) + arrow = builder.new_read().to_arrow(builder.new_scan().plan().splits()) + + found = arrow[_ROW_ID].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: + raise RuntimeError( + "Aligned row ids disappeared from pinned snapshot: %r." % missing) + take = pa.array([positions[row_id] for row_id in row_ids], type=pa.int64()) + return arrow.select(visible).take(take) + + +def _metadata_sort_key(on, by): + def key(row): + return tuple(row[name] for name in by) + (row[on], row[_ROW_ID]) + return key + + +def _match_index(times, target, method): + position = bisect_left(times, target) + if method == "exact": + if position < len(times) and times[position] == target: + return position + return None + if method == "backward": + if position < len(times) and times[position] == target: + return position + return position - 1 if position else None + if method == "forward": + return position if position < len(times) else None + if method == "nearest": + if position == 0: + return 0 + if position == len(times): + return len(times) - 1 + before = position - 1 + if target - times[before] <= times[position] - target: + return before + return position + raise ValueError("Unknown temporal match method %r." % method) + + +def _delta_type(data_type): + if pa.types.is_timestamp(data_type): + return pa.duration(data_type.unit) + if pa.types.is_integer(data_type): + return pa.int64() + if pa.types.is_floating(data_type): + return pa.float64() + raise TypeError( + "Temporal columns must be integer, floating point, or timestamp; " + "got %s." + % data_type + ) + + +def _require_columns(schema, columns, label): + missing = [name for name in columns if name not in schema.names] + if missing: + raise ValueError( + "%s is missing temporal columns %r." % (label, missing)) diff --git a/paimon-python/pypaimon/tests/multimodal_temporal_test.py b/paimon-python/pypaimon/tests/multimodal_temporal_test.py new file mode 100644 index 000000000000..a11e861ae679 --- /dev/null +++ b/paimon-python/pypaimon/tests/multimodal_temporal_test.py @@ -0,0 +1,262 @@ +# 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 os +import shutil +import tempfile +import unittest +from datetime import datetime, timedelta + +import pyarrow as pa +import pypaimon.multimodal as pmm + + +class MultimodalTemporalTest(unittest.TestCase): + + def setUp(self): + self.temp_dir = tempfile.mkdtemp(prefix="pypaimon_temporal_") + self.conn = pmm.connect(options={ + "warehouse": os.path.join(self.temp_dir, "warehouse"), + }) + + def tearDown(self): + shutil.rmtree(self.temp_dir, ignore_errors=True) + + def test_alignment_requires_an_explicit_group_boundary(self): + table = self._table("missing_group", { + "event_time": pa.int64(), + "value": pa.int32(), + }) + with self.assertRaisesRegex(ValueError, "grouping column"): + pmm.align( + table.scan(), + on="event_time", + by=(), + sources={"value": pmm.exact(table.scan())}, + ) + + def test_aligns_named_sources_in_episode_local_batches(self): + actions = self._table("actions", { + "episode_id": pa.string(), + "event_time": pa.int64(), + "action": pa.int32(), + }) + images = self._table("images", { + "episode_id": pa.string(), + "event_time": pa.int64(), + "camera": pa.string(), + "image": pa.string(), + }) + states = self._table("states", { + "episode_id": pa.string(), + "event_time": pa.int64(), + "state": pa.int32(), + }) + commands = self._table("commands", { + "episode_id": pa.string(), + "event_time": pa.int64(), + "command": pa.string(), + }) + actions.add([ + {"episode_id": "ep-2", "event_time": 100, "action": 4}, + {"episode_id": "ep-1", "event_time": 300, "action": 3}, + {"episode_id": "ep-1", "event_time": 100, "action": 1}, + {"episode_id": "ep-1", "event_time": 200, "action": 2}, + ]) + images.add([ + {"episode_id": "ep-1", "event_time": 90, + "camera": "left", "image": "early"}, + {"episode_id": "ep-1", "event_time": 90, + "camera": "right", "image": "ignored"}, + {"episode_id": "ep-1", "event_time": 110, + "camera": "left", "image": "late"}, + {"episode_id": "ep-1", "event_time": 215, + "camera": "left", "image": "middle"}, + {"episode_id": "ep-2", "event_time": 99, + "camera": "left", "image": "other"}, + ]) + states.add([ + {"episode_id": "ep-1", "event_time": 80, "state": 8}, + {"episode_id": "ep-1", "event_time": 190, "state": 19}, + {"episode_id": "ep-2", "event_time": 95, "state": 95}, + ]) + commands.add([ + {"episode_id": "ep-1", "event_time": 100, "command": "open"}, + {"episode_id": "ep-1", "event_time": 220, "command": "close"}, + {"episode_id": "ep-2", "event_time": 100, "command": "hold"}, + ]) + + aligned = pmm.align( + actions.scan().select(["episode_id", "event_time", "action"]), + on="event_time", + by="episode_id", + sources={ + "camera": pmm.nearest( + images.scan().where("camera = 'left'").select("image"), + tolerance=20, + ), + "state": pmm.backward( + states.scan().select("state"), tolerance=25), + "command": pmm.exact(commands.scan().select("command")), + "next_command": pmm.forward( + commands.scan().select("command"), tolerance=25), + }, + ) + reader = aligned.to_arrow_batch_reader(batch_size=2) + batches = list(reader) + rows = pa.Table.from_batches(batches).to_pylist() + + self.assertEqual([2, 2], [batch.num_rows for batch in batches]) + self.assertEqual( + [("ep-1", 100), ("ep-1", 200), ("ep-1", 300), ("ep-2", 100)], + [(row["episode_id"], row["event_time"]) for row in rows], + ) + # Equal-distance nearest ties select the earlier row. + self.assertEqual( + ["early", "middle", None, "other"], + [row["camera__image"] for row in rows], + ) + self.assertEqual([-10, 15, None, -1], [ + row["camera__time_delta"] for row in rows + ]) + self.assertEqual([8, 19, None, 95], [ + row["state__state"] for row in rows + ]) + self.assertEqual(["open", None, None, "hold"], [ + row["command__command"] for row in rows + ]) + self.assertEqual(["open", "close", None, "hold"], [ + row["next_command__command"] for row in rows + ]) + self.assertEqual([True, True, False, True], [ + row["camera__valid"] for row in rows + ]) + + def test_alignment_pins_each_scan_snapshot(self): + anchors = self._table("pinned_anchors", { + "episode_id": pa.int32(), + "event_time": pa.int64(), + "value": pa.string(), + }) + secondary = self._table("pinned_secondary", { + "episode_id": pa.int32(), + "event_time": pa.int64(), + "value": pa.string(), + }) + anchors.add([{"episode_id": 1, "event_time": 100, "value": "old"}]) + secondary.add([ + {"episode_id": 1, "event_time": 90, "value": "old-match"} + ]) + aligned = pmm.align( + anchors.scan(), + on="event_time", + by="episode_id", + sources={ + "secondary": pmm.nearest(secondary.scan(), tolerance=20), + }, + ) + + anchors.add([{"episode_id": 1, "event_time": 200, "value": "new"}]) + secondary.add([ + {"episode_id": 1, "event_time": 100, "value": "new-match"} + ]) + + self.assertEqual([{ + "episode_id": 1, + "event_time": 100, + "value": "old", + "secondary__value": "old-match", + "secondary__valid": True, + "secondary__matched_time": 90, + "secondary__time_delta": -10, + }], aligned.to_list()) + + def test_alignment_keeps_blob_payloads_as_descriptors(self): + anchors = self._table("blob_anchors", { + "episode_id": pa.int32(), + "event_time": pa.int64(), + }) + images = self._table("blob_images", { + "episode_id": pa.int32(), + "event_time": pa.int64(), + "image": pa.large_binary(), + }) + anchors.add([{"episode_id": 1, "event_time": 100}]) + images.add([{ + "episode_id": 1, + "event_time": 100, + "image": b"encoded-image", + }]) + + row = pmm.align( + anchors.scan(), + on="event_time", + by="episode_id", + sources={ + "camera": pmm.exact(images.scan().select("image")), + }, + ).to_list()[0] + + descriptor = pmm.BlobDescriptor.deserialize(row["camera__image"]) + self.assertTrue(descriptor.uri.endswith(".blob")) + self.assertEqual(len(b"encoded-image"), descriptor.length) + + def test_alignment_supports_timestamp_columns_with_different_names(self): + anchors = self._table("timestamp_anchors", { + "episode_id": pa.int32(), + "event_time": pa.timestamp("ms"), + }) + samples = self._table("timestamp_samples", { + "episode_id": pa.int32(), + "captured_at": pa.timestamp("ms"), + "value": pa.int32(), + }) + anchor_time = datetime(2026, 9, 1, 12, 0, 0, 100000) + sample_time = anchor_time - timedelta(milliseconds=5) + anchors.add([{"episode_id": 1, "event_time": anchor_time}]) + samples.add([{ + "episode_id": 1, + "captured_at": sample_time, + "value": 7, + }]) + + row = pmm.align( + anchors.scan(), + on="event_time", + by="episode_id", + sources={ + "sample": pmm.nearest( + samples.scan().select("value"), + on="captured_at", + tolerance=timedelta(milliseconds=10), + ), + }, + ).to_list()[0] + + self.assertEqual(7, row["sample__value"]) + self.assertEqual(sample_time, row["sample__matched_time"]) + self.assertEqual( + timedelta(milliseconds=-5), row["sample__time_delta"]) + + def _table(self, name, fields): + return self.conn.create_table(name, schema=pa.schema([ + pa.field(field_name, field_type) + for field_name, field_type in fields.items() + ])) + + +if __name__ == "__main__": + unittest.main()