| # 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 json | |
| import os | |
| import pickle | |
| import shutil | |
| import tempfile | |
| import unittest | |
| from types import SimpleNamespace | |
| from unittest.mock import patch | |
| import pyarrow as pa | |
| import pytest | |
| import torch | |
| import pypaimon.multimodal as pmm | |
| from pypaimon.catalog.table_query_auth import TableQueryAuthResult | |
| from pypaimon.multimodal import window_dataset | |
| from pypaimon.multimodal.query import ScanQuery | |
| from pypaimon.multimodal.window_dataset import ContiguousWindowDataset | |
| from pypaimon.execution.python.read.reader.filter_record_batch_reader import (FilterRecordBatchReader) | |
| from pypaimon.execution.python.read.reader.format_pyarrow_reader import (FormatPyArrowReader) | |
| from pypaimon.read.table_scan import TableScan | |
| _TABLE_OPTIONS = { | |
| "row-tracking.enabled": "true", | |
| "data-evolution.enabled": "true", | |
| "deletion-vectors.enabled": "true", | |
| "file.format": "parquet", | |
| "vector.file.format": "parquet", | |
| } | |
| class _TensorColumnTransform: | |
| def __call__(self, values): | |
| return torch.tensor(values, dtype=torch.int64) | |
| class _WindowAdapter: | |
| def __call__(self, sample): | |
| return { | |
| "episode": sample["episode"], | |
| "start": sample["step"], | |
| "values": sample["value"], | |
| "padding_mask": sample["is_pad"], | |
| } | |
| class ContiguousWindowDatasetTest(unittest.TestCase): | |
| def setUp(self): | |
| self.temp_dir = tempfile.mkdtemp(prefix="pypaimon_windows_") | |
| self.conn = pmm.connect(options={ | |
| "warehouse": os.path.join(self.temp_dir, "warehouse"), | |
| }) | |
| def tearDown(self): | |
| shutil.rmtree(self.temp_dir, ignore_errors=True) | |
| @staticmethod | |
| def _schema(): | |
| return pa.schema([ | |
| pa.field("episode", pa.string(), nullable=False), | |
| pa.field("step", pa.int32(), nullable=False), | |
| pa.field("value", pa.int32(), nullable=False), | |
| pa.field("payload", pa.large_binary(), nullable=False), | |
| ]) | |
| @staticmethod | |
| def _row(episode, step): | |
| return { | |
| "episode": episode, | |
| "step": step, | |
| "value": step + (100 if episode == "episode-b" else 0), | |
| "payload": ("%s-%d" % (episode, step)).encode(), | |
| } | |
| def _table(self, name="frames"): | |
| table = self.conn.create_table( | |
| name, schema=self._schema(), options=_TABLE_OPTIONS) | |
| table.add([ | |
| self._row("episode-b", 2), | |
| self._row("episode-a", 1), | |
| self._row("episode-b", 0), | |
| self._row("episode-a", 0), | |
| self._row("episode-b", 3), | |
| self._row("episode-b", 1), | |
| ]) | |
| return table | |
| @staticmethod | |
| def _dataset(table, **kwargs): | |
| return ( | |
| table.scan() | |
| .to_contiguous_window_dataset( | |
| window_size=3, | |
| columns=["value", "payload"], | |
| group_key="episode", | |
| order_key="step", | |
| **kwargs, | |
| ) | |
| ) | |
| def test_per_column_offsets_combine_history_future_and_anchor(self): | |
| dataset = self._table().scan().to_contiguous_window_dataset( | |
| columns=["value", "payload"], group_key="episode", order_key="step", | |
| frame_offsets={"value": [-1, 0, 2]}, boundary="pad") | |
| self.assertEqual(6, len(dataset)) | |
| sample = dataset[0] | |
| self.assertEqual(("episode-a", 0), (sample["episode"], sample["step"])) | |
| self.assertEqual([0, 0, 1], sample["value"]) | |
| self.assertEqual([True, False, True], sample["value_is_pad"].tolist()) | |
| self.assertEqual([b"episode-a-0"], sample["payload"]) | |
| self.assertEqual([False], sample["payload_is_pad"].tolist()) | |
| self.assertEqual(torch.bool, sample["value_is_pad"].dtype) | |
| self.assertNotIn("is_pad", sample) | |
| self.assertEqual([101, 102, 103], dataset[4]["value"]) | |
| self.assertEqual([False, False, True], dataset[4]["value_is_pad"].tolist()) | |
| def test_frame_mode_defaults_to_single_frame_with_time_dimension(self): | |
| dataset = self._table().scan().select(["value"]).to_contiguous_window_dataset( | |
| group_key="episode", order_key="step") | |
| self.assertEqual(6, len(dataset)) | |
| self.assertEqual({"episode", "step", "value", "value_is_pad"}, set(dataset[0])) | |
| self.assertEqual([0], dataset[0]["value"]) | |
| self.assertEqual([False], dataset[0]["value_is_pad"].tolist()) | |
| self.assertEqual(103, dataset[-1]["value"][0]) | |
| def test_anchor_selection_preserves_context_across_files(self): | |
| table = self.conn.create_table( | |
| "anchor_context", schema=self._schema(), options=_TABLE_OPTIONS) | |
| for start in (0, 3, 6): | |
| table.add([self._row("episode-a", step) | |
| for step in range(start, min(start + 3, 8))]) | |
| row_ids = {row["step"]: row["_ROW_ID"] for row in | |
| table.scan().select(["step"]).with_row_id().to_list()} | |
| dataset = table.scan().to_contiguous_window_dataset( | |
| columns=["value", "payload"], group_key="episode", order_key="step", | |
| frame_offsets={"value": [-1, 0, 1], "payload": [0, 1]}, | |
| anchor_row_ids=[row_ids[5], row_ids[2], row_ids[5]]) | |
| self.assertEqual(2, len(dataset)) | |
| samples = dataset.__getitems__([1, 0, 1]) | |
| self.assertEqual([5, 2, 5], [sample["step"] for sample in samples]) | |
| self.assertEqual([[4, 5, 6], [1, 2, 3], [4, 5, 6]], | |
| [sample["value"] for sample in samples]) | |
| self.assertEqual([[b"episode-a-5", b"episode-a-6"], | |
| [b"episode-a-2", b"episode-a-3"], | |
| [b"episode-a-5", b"episode-a-6"]], | |
| [sample["payload"] for sample in samples]) | |
| self.assertTrue(all(not sample["value_is_pad"].any() for sample in samples)) | |
| def test_anchor_selection_keeps_stride_and_checks_only_selected_boundaries(self): | |
| table = self._table() | |
| row_ids = {row["step"]: row["_ROW_ID"] for row in | |
| table.scan().where("episode = 'episode-b'") | |
| .select(["step"]).with_row_id().to_list()} | |
| kwargs = dict( | |
| columns=["value"], group_key="episode", order_key="step", | |
| frame_offsets={"value": [-1, 0, 1]}, stride=2, boundary="error") | |
| dataset = table.scan().to_contiguous_window_dataset( | |
| anchor_row_ids=(row_ids[step] for step in (1, 2)), **kwargs) | |
| self.assertEqual(1, len(dataset)) | |
| self.assertEqual(2, dataset[0]["step"]) | |
| self.assertEqual([101, 102, 103], dataset[0]["value"]) | |
| with self.assertRaisesRegex(ValueError, "episode-b.*incomplete window at 0"): | |
| table.scan().to_contiguous_window_dataset( | |
| anchor_row_ids=[row_ids[0]], **kwargs) | |
| def test_empty_anchor_selection_has_no_windows_or_payload_reads(self): | |
| with patch("pypaimon.multimodal.window_dataset.fetch_blob_bodies") as fetch: | |
| dataset = self._table().scan().to_contiguous_window_dataset( | |
| columns=["value", "payload"], group_key="episode", order_key="step", | |
| frame_offsets={"value": [-2, 0, 2]}, boundary="error", | |
| anchor_row_ids=[]) | |
| self.assertEqual(0, len(dataset)) | |
| self.assertEqual([], dataset[:]) | |
| with self.assertRaises(IndexError): | |
| dataset[0] | |
| fetch.assert_not_called() | |
| def test_anchor_selection_rejects_invalid_or_out_of_context_ids(self): | |
| table = self._table() | |
| for selected, error, message in ( | |
| ([True], TypeError, "booleans"), | |
| ([1.5], TypeError, "integer row IDs"), | |
| (["1"], TypeError, "integer row IDs"), | |
| ([-1], ValueError, "nonnegative int64"), | |
| ([2 ** 63], ValueError, "nonnegative int64"), | |
| ([999], ValueError, "outside the context query")): | |
| with self.subTest(selected=selected): | |
| with self.assertRaisesRegex(error, message): | |
| self._dataset(table, anchor_row_ids=selected) | |
| excluded_id = (table.scan().where("episode = 'episode-b'") | |
| .select(["step"]).with_row_id().to_list()[0]["_ROW_ID"]) | |
| with self.assertRaisesRegex(ValueError, "outside the context query"): | |
| table.scan().where("episode = 'episode-a'").to_contiguous_window_dataset( | |
| columns=["value"], group_key="episode", order_key="step", | |
| anchor_row_ids=[excluded_id]) | |
| def test_drop_offsets_keep_stride_anchored_at_group_start(self): | |
| table = self._table() | |
| dataset = table.scan().to_contiguous_window_dataset( | |
| columns=["value", "payload"], group_key="episode", order_key="step", | |
| frame_offsets={"value": [-1, 0], "payload": [1]}, stride=2) | |
| self.assertEqual(1, len(dataset)) | |
| self.assertEqual(("episode-b", 2), (dataset[0]["episode"], dataset[0]["step"])) | |
| self.assertEqual([101, 102], dataset[0]["value"]) | |
| self.assertEqual([b"episode-b-3"], dataset[0]["payload"]) | |
| self.assertFalse(dataset[0]["value_is_pad"].any()) | |
| with self.assertRaisesRegex( | |
| ValueError, r"episode-a.*at 0: offset bounds=\[-1, 0\]"): | |
| table.scan().to_contiguous_window_dataset( | |
| columns=["value"], group_key="episode", order_key="step", | |
| frame_offsets={"value": [-1, 0]}, boundary="error") | |
| def test_sparse_repeated_offsets_preserve_batch_and_slice_order(self): | |
| dataset = self._table().scan().to_contiguous_window_dataset( | |
| columns=["value"], group_key="episode", order_key="step", | |
| frame_offsets={"value": [2, -1, 2, 0]}, boundary="pad") | |
| samples = dataset.__getitems__([4, 0, 4, -1]) | |
| self.assertEqual([[103, 101, 103, 102], [1, 0, 1, 0], | |
| [103, 101, 103, 102], [103, 102, 103, 103]], | |
| [sample["value"] for sample in samples]) | |
| self.assertEqual([True, False, True, False], samples[0]["value_is_pad"].tolist()) | |
| self.assertEqual([dataset[i]["value"] for i in (5, 3, 1)], | |
| [sample["value"] for sample in dataset[::-2]]) | |
| self.assertEqual([], dataset.__getitems__([])) | |
| with self.assertRaises(IndexError): | |
| dataset.__getitems__([0, len(dataset)]) | |
| def test_custom_padding_precedes_transforms_even_for_all_padding(self): | |
| dataset = self._table().scan().to_contiguous_window_dataset( | |
| columns=["value", "payload"], group_key="episode", order_key="step", | |
| frame_offsets={"value": [-10, 10], "payload": [0, 10]}, | |
| boundary="pad", pad_values={"value": -7, "payload": b"missing"}, | |
| column_transforms={"value": _TensorColumnTransform()}) | |
| with patch.object(dataset, "_read_rows", wraps=dataset._read_rows) as read: | |
| sample = dataset[0] | |
| self.assertEqual(1, read.call_count) | |
| self.assertEqual(["payload"], read.call_args.args[1]) | |
| self.assertTrue(torch.equal(torch.tensor([-7, -7]), sample["value"])) | |
| self.assertEqual([True, True], sample["value_is_pad"].tolist()) | |
| self.assertEqual([b"episode-a-0", b"missing"], sample["payload"]) | |
| self.assertEqual([False, True], sample["payload_is_pad"].tolist()) | |
| def test_seconds_offsets_match_frame_offsets_with_alignment_tolerance(self): | |
| table = self._table() | |
| kwargs = dict(columns=["value", "payload"], group_key="episode", | |
| order_key="step", boundary="pad") | |
| frames = table.scan().to_contiguous_window_dataset( | |
| frame_offsets={"value": [-1, 0, 2]}, **kwargs) | |
| seconds = table.scan().to_contiguous_window_dataset( | |
| delta_timestamps={"value": [-0.1, 0, 0.20001]}, fps=10, | |
| tolerance_s=0.0001, **kwargs) | |
| self.assertEqual(len(frames), len(seconds)) | |
| for expected, actual in zip(frames[:], seconds[:]): | |
| self.assertEqual(expected["value"], actual["value"]) | |
| self.assertEqual(expected["payload"], actual["payload"]) | |
| self.assertTrue(torch.equal(expected["value_is_pad"], actual["value_is_pad"])) | |
| def test_offset_projections_do_not_fetch_anchor_images_for_action_rows(self): | |
| dataset = self._table().scan().to_contiguous_window_dataset( | |
| columns=["value", "payload"], group_key="episode", order_key="step", | |
| frame_offsets={"value": [-1, 0, 1], "payload": [0]}, boundary="pad") | |
| with patch.object(dataset, "_read_rows", wraps=dataset._read_rows) as read, \ | |
| patch("pypaimon.multimodal.window_dataset.fetch_blob_bodies", | |
| side_effect=window_dataset.fetch_blob_bodies) as fetch: | |
| samples = dataset.__getitems__([3, 4, 3]) | |
| projections = {tuple(call.args[1]): len(call.args[0]) for call in read.call_args_list} | |
| self.assertEqual({("value",): 4, ("payload",): 2}, projections) | |
| self.assertEqual(2, sum(len(call.args[1]["payload"]) for call in fetch.call_args_list)) | |
| self.assertEqual([[100, 101, 102], [101, 102, 103], [100, 101, 102]], | |
| [sample["value"] for sample in samples]) | |
| def test_unused_pad_values_do_not_split_reads(self): | |
| table = self._table() | |
| for boundary in ("drop", "error"): | |
| with self.subTest(boundary=boundary): | |
| dataset = table.scan().to_contiguous_window_dataset( | |
| columns=["value", "payload"], group_key="episode", order_key="step", | |
| frame_offsets={"value": [0, 1], "payload": [0, 1]}, | |
| stride=2, boundary=boundary, pad_values={"value": -1}) | |
| with patch.object(dataset, "_read_rows", wraps=dataset._read_rows) as read: | |
| samples = dataset.__getitems__([0, 1, 0]) | |
| self.assertEqual(1, read.call_count) | |
| self.assertEqual(["value", "payload"], read.call_args.args[1]) | |
| self.assertEqual(4, len(read.call_args.args[0])) | |
| self.assertEqual([[0, 1], [100, 101], [0, 1]], | |
| [sample["value"] for sample in samples]) | |
| self.assertEqual([b"episode-a-0", b"episode-a-1"], samples[0]["payload"]) | |
| def test_offset_padding_and_repeated_cells_are_mutably_isolated(self): | |
| table = self.conn.create_table( | |
| "offset_mutable", schema=pa.schema([ | |
| pa.field("episode", pa.string(), nullable=False), | |
| pa.field("step", pa.int32(), nullable=False), | |
| pa.field("values", pa.list_(pa.int32()), nullable=False), | |
| ]), options=_TABLE_OPTIONS) | |
| table.add([{"episode": "a", "step": 0, "values": [0]}]) | |
| pad = [-1] | |
| dataset = table.scan().to_contiguous_window_dataset( | |
| columns=["values"], group_key="episode", order_key="step", | |
| frame_offsets={"values": [-1, 0, 0, 1]}, boundary="pad", | |
| pad_values={"values": pad}) | |
| first, second = dataset.__getitems__([0, 0]) | |
| first["values"][0].append(99) | |
| first["values"][1].append(99) | |
| self.assertEqual([0], first["values"][2]) | |
| self.assertEqual([-1], first["values"][3]) | |
| self.assertEqual([[-1], [0], [0], [-1]], second["values"]) | |
| self.assertEqual([[-1], [0], [0], [-1]], dataset[0]["values"]) | |
| self.assertEqual([-1], pad) | |
| def test_offset_dataloader_workers_and_distributed_sampler(self): | |
| dataset = self._table().scan().to_contiguous_window_dataset( | |
| columns=["value"], group_key="episode", order_key="step", | |
| frame_offsets={"value": [-1, 0, 1]}, boundary="pad", | |
| column_transforms={"value": _TensorColumnTransform()}) | |
| restored = pickle.loads(pickle.dumps(dataset)) | |
| self.assertEqual(dataset.snapshot_id, restored.snapshot_id) | |
| batches = list(torch.utils.data.DataLoader( | |
| restored, batch_size=2, num_workers=2, shuffle=False)) | |
| self.assertEqual([[0, 0, 1], [0, 1, 1], [100, 100, 101], | |
| [100, 101, 102], [101, 102, 103], [102, 103, 103]], | |
| [row for batch in batches for row in batch["value"].tolist()]) | |
| self.assertEqual((2, 3), tuple(batches[0]["value_is_pad"].shape)) | |
| self.assertEqual(torch.bool, batches[0]["value_is_pad"].dtype) | |
| partitions = [list(torch.utils.data.DistributedSampler( | |
| dataset, num_replicas=2, rank=rank, shuffle=False)) for rank in range(2)] | |
| self.assertEqual([[0, 2, 4], [1, 3, 5]], partitions) | |
| self.assertEqual([0, 100, 102], [dataset[i]["value"][1].item() for i in partitions[0]]) | |
| def test_rejects_invalid_offset_configuration_at_construction(self): | |
| table = self._table() | |
| cases = [ | |
| {"frame_offsets": {"missing": [0]}}, | |
| {"frame_offsets": {"value": []}}, | |
| {"frame_offsets": {"value": [1.0]}}, | |
| {"frame_offsets": {"value": [True]}}, | |
| {"frame_offsets": {"value": [0]}, "window_size": 2}, | |
| {"delta_timestamps": {"value": [0]}, "fps": 10, "window_size": 2}, | |
| {"frame_offsets": {"value": [0]}, "delta_timestamps": {"value": [0]}, "fps": 10}, | |
| {"delta_timestamps": {"value": [0]}}, | |
| {"delta_timestamps": {"value": []}, "fps": 10}, | |
| {"delta_timestamps": {"value": [float("nan")]}, "fps": 10}, | |
| {"delta_timestamps": {"value": [float("inf")]}, "fps": 10}, | |
| {"delta_timestamps": {"value": [0]}, "fps": 0}, | |
| {"delta_timestamps": {"value": [0]}, "fps": -1}, | |
| {"delta_timestamps": {"value": [0]}, "fps": float("inf")}, | |
| {"delta_timestamps": {"value": [0]}, "fps": float("nan")}, | |
| {"delta_timestamps": {"value": [0.05]}, "fps": 10}, | |
| {"frame_offsets": {"value": [0]}, "anchor_columns": ["value"]}, | |
| {"frame_offsets": {"value": [0]}, "tail": "pad"}, | |
| {"frame_offsets": {"value": [0]}, "tolerance_s": 0.01}, | |
| {"window_size": 2, "tolerance_s": 0.01}, | |
| {"window_size": 2, "tail": "pad", "boundary": "pad"}, | |
| {"boundary": "unknown"}, | |
| ] | |
| for kwargs in cases: | |
| with self.subTest(kwargs=kwargs), self.assertRaises((TypeError, ValueError)): | |
| table.scan().to_contiguous_window_dataset( | |
| columns=["value"], group_key="episode", order_key="step", **kwargs) | |
| def test_rejects_generated_padding_mask_collisions(self): | |
| table = self.conn.create_table( | |
| "mask_collision", schema=pa.schema([ | |
| pa.field("episode", pa.string(), nullable=False), | |
| pa.field("step", pa.int32(), nullable=False), | |
| pa.field("value", pa.int32(), nullable=False), | |
| pa.field("value_is_pad", pa.int32(), nullable=False), | |
| ]), options=_TABLE_OPTIONS) | |
| table.add([{"episode": "a", "step": 0, "value": 1, "value_is_pad": 2}]) | |
| with self.assertRaises(ValueError): | |
| table.scan().to_contiguous_window_dataset( | |
| columns=["value", "value_is_pad"], group_key="episode", order_key="step") | |
| with self.assertRaises(ValueError): | |
| table.scan().to_contiguous_window_dataset( | |
| columns=["value"], group_key="episode", order_key="value_is_pad") | |
| def test_offset_mode_allows_is_pad_as_a_stored_column(self): | |
| table = self.conn.create_table( | |
| "stored_is_pad", schema=pa.schema([ | |
| pa.field("episode", pa.string(), nullable=False), | |
| pa.field("step", pa.int32(), nullable=False), | |
| pa.field("is_pad", pa.int32(), nullable=False), | |
| pa.field("value", pa.int32(), nullable=False), | |
| ]), options=_TABLE_OPTIONS) | |
| table.add([{"episode": "a", "step": i, "is_pad": i, "value": 100 + i} | |
| for i in range(3)]) | |
| for group, order, column in (("episode", "step", "is_pad"), | |
| ("is_pad", "step", "value"), | |
| ("episode", "is_pad", "value")): | |
| with self.subTest(group=group, order=order, column=column): | |
| kwargs = dict(group_key=group, order_key=order, columns=[column]) | |
| dataset = table.scan().to_contiguous_window_dataset(**kwargs) | |
| self.assertEqual(3, len(dataset)) | |
| self.assertEqual([0] if column == "is_pad" else [100], dataset[0][column]) | |
| self.assertEqual([False], dataset[0][column + "_is_pad"].tolist()) | |
| with self.assertRaises(ValueError): | |
| table.scan().to_contiguous_window_dataset(window_size=1, **kwargs) | |
| def test_sorts_rows_and_never_crosses_episode_boundaries(self): | |
| dataset = self._dataset(self._table()) | |
| self.assertIsInstance(dataset, torch.utils.data.Dataset) | |
| self.assertEqual(2, len(dataset)) | |
| self.assertIsInstance(dataset.snapshot_id, int) | |
| self.assertNotIn("_episodes", vars(dataset)) | |
| first = dataset[0] | |
| second = dataset[1] | |
| self.assertEqual("episode-b", first["episode"]) | |
| self.assertEqual(0, first["step"]) | |
| self.assertEqual([100, 101, 102], first["value"]) | |
| self.assertEqual([101, 102, 103], second["value"]) | |
| self.assertFalse(first["is_pad"].any()) | |
| self.assertEqual({"episode-b"}, { | |
| window["episode"] for window in (first, second) | |
| }) | |
| def test_reads_blob_payloads_only_when_a_window_is_requested(self): | |
| table = self._table() | |
| with patch( | |
| "pypaimon.multimodal.window_dataset.fetch_blob_bodies", | |
| side_effect=window_dataset.fetch_blob_bodies) as fetch: | |
| dataset = self._dataset(table) | |
| self.assertEqual(0, fetch.call_count) | |
| sample = dataset[0] | |
| self.assertEqual(1, fetch.call_count) | |
| self.assertEqual(3, len(fetch.call_args.args[1]["payload"])) | |
| self.assertEqual( | |
| [b"episode-b-0", b"episode-b-1", b"episode-b-2"], | |
| sample["payload"], | |
| ) | |
| def test_reads_map_blob_payloads_for_map_only_and_mixed_windows(self): | |
| schema = pa.schema([ | |
| pa.field("episode", pa.string(), nullable=False), | |
| pa.field("step", pa.int32(), nullable=False), | |
| pa.field("payload", pa.large_binary()), | |
| pa.field( | |
| "attachments", | |
| pa.map_(pa.string(), pa.large_binary()), | |
| ), | |
| ]) | |
| table = self.conn.create_table( | |
| "map_blobs", | |
| schema=schema, | |
| options=_TABLE_OPTIONS, | |
| ) | |
| table.add(pa.Table.from_pylist([ | |
| { | |
| "episode": "episode-a", | |
| "step": 0, | |
| "payload": b"scalar-0", | |
| "attachments": { | |
| "body": b"map-0", "empty": b"", "null": None}, | |
| }, | |
| { | |
| "episode": "episode-a", | |
| "step": 1, | |
| "payload": b"scalar-1", | |
| "attachments": None, | |
| }, | |
| ], schema=schema)) | |
| def window(columns): | |
| return table.scan().to_contiguous_window_dataset( | |
| window_size=2, | |
| columns=columns, | |
| group_key="episode", | |
| order_key="step", | |
| )[0] | |
| map_only = window(["attachments"]) | |
| mixed = window(["payload", "attachments"]) | |
| self.assertEqual( | |
| {"body": b"map-0", "empty": b"", "null": None}, | |
| dict(map_only["attachments"][0]), | |
| ) | |
| self.assertIsNone(map_only["attachments"][1]) | |
| self.assertEqual([b"scalar-0", b"scalar-1"], mixed["payload"]) | |
| self.assertEqual(map_only["attachments"], mixed["attachments"]) | |
| def test_reads_array_blob_payloads_in_single_and_batched_windows(self): | |
| pages = [[b"first", None, b"", b"first"], None, [], [b"last"]] | |
| schema = pa.schema([ | |
| ("episode", pa.string()), ("step", pa.int32()), | |
| ("pages", pa.list_(pa.large_binary())), | |
| ("payload", pa.large_binary()), | |
| ("attachments", pa.map_(pa.string(), pa.large_binary())), | |
| ]) | |
| table = self.conn.create_table( | |
| "array_blobs", schema=schema, options=_TABLE_OPTIONS) | |
| table.add(pa.Table.from_pydict({ | |
| "episode": ["episode-a"] * 4, "step": [0, 1, 2, 3], | |
| "pages": pages, "payload": [b"cover"] * 4, | |
| "attachments": [[("thumb", b"thumbnail")]] * 4, | |
| }, schema=schema)) | |
| for columns in (["pages"], ["pages", "payload", "attachments"]): | |
| with self.subTest(columns=columns): | |
| dataset = table.scan().to_contiguous_window_dataset( | |
| window_size=4, columns=columns, | |
| group_key="episode", order_key="step") | |
| self.assertEqual(pages, dataset[0]["pages"]) | |
| self.assertEqual([pages, pages], [ | |
| row["pages"] for row in dataset.__getitems__([0, 0])]) | |
| if "payload" in columns: | |
| self.assertEqual([b"cover"] * 4, dataset[0]["payload"]) | |
| self.assertEqual( | |
| [[("thumb", b"thumbnail")]] * 4, | |
| dataset[0]["attachments"]) | |
| # Reordered/repeated offsets and endpoint padding retain cell structure. | |
| dataset = table.scan().to_contiguous_window_dataset( | |
| columns=["pages", "payload"], group_key="episode", order_key="step", | |
| frame_offsets={"pages": [-1, 0, 2, 0]}, boundary="pad") | |
| samples = dataset.__getitems__([3, 0, 3]) | |
| self.assertEqual([ | |
| [pages[2], pages[3], pages[3], pages[3]], | |
| [pages[0], pages[0], pages[2], pages[0]], | |
| [pages[2], pages[3], pages[3], pages[3]], | |
| ], [sample["pages"] for sample in samples]) | |
| self.assertEqual([False, False, True, False], | |
| samples[0]["pages_is_pad"].tolist()) | |
| samples[0]["pages"][1].append(b"changed") | |
| self.assertEqual([b"last"], samples[2]["pages"][1]) | |
| # Constant padding must not turn an empty ARRAY into a null cell. | |
| padded = table.scan().to_contiguous_window_dataset( | |
| columns=["pages"], group_key="episode", order_key="step", | |
| frame_offsets={"pages": [-1, 0]}, boundary="pad", | |
| pad_values={"pages": []}) | |
| self.assertEqual([[], pages[0]], padded[0]["pages"]) | |
| def test_anchor_columns_read_only_the_window_anchor(self): | |
| table = self._table() | |
| with patch( | |
| "pypaimon.multimodal.window_dataset.fetch_blob_bodies", | |
| side_effect=window_dataset.fetch_blob_bodies) as fetch: | |
| dataset = self._dataset(table, anchor_columns=["payload"]) | |
| sample = dataset[0] | |
| self.assertEqual([100, 101, 102], sample["value"]) | |
| self.assertEqual([b"episode-b-0"], sample["payload"]) | |
| self.assertEqual(1, fetch.call_count) | |
| self.assertEqual(1, len(fetch.call_args.args[1]["payload"])) | |
| def test_plural_access_coalesces_overlapping_window_reads(self): | |
| dataset = self._dataset( | |
| self._table(), anchor_columns=["payload"]) | |
| with patch.object( | |
| dataset, "_read_rows", wraps=dataset._read_rows) as read: | |
| actual = dataset.__getitems__([1, 0, 1]) | |
| self.assertEqual(2, read.call_count) | |
| self.assertEqual(4, len(read.call_args_list[0].args[0])) | |
| self.assertEqual(["value"], read.call_args_list[0].args[1]) | |
| self.assertEqual(2, len(read.call_args_list[1].args[0])) | |
| self.assertEqual(["payload"], read.call_args_list[1].args[1]) | |
| self.assertEqual( | |
| [("episode-b", 1), ("episode-b", 0), ("episode-b", 1)], | |
| [(sample["episode"], sample["step"]) for sample in actual], | |
| ) | |
| self.assertEqual( | |
| [[101, 102, 103], [100, 101, 102], [101, 102, 103]], | |
| [sample["value"] for sample in actual], | |
| ) | |
| self.assertEqual( | |
| [[b"episode-b-1"], [b"episode-b-0"], [b"episode-b-1"]], | |
| [sample["payload"] for sample in actual], | |
| ) | |
| def test_plural_access_isolates_mutable_cells_between_samples(self): | |
| table = self.conn.create_table( | |
| "mutable_cells", | |
| schema=pa.schema([ | |
| pa.field("episode", pa.string(), nullable=False), | |
| pa.field("step", pa.int32(), nullable=False), | |
| pa.field("values", pa.list_(pa.int32()), nullable=False), | |
| ]), | |
| options=_TABLE_OPTIONS, | |
| ) | |
| table.add([ | |
| {"episode": "episode-a", "step": step, "values": [step]} | |
| for step in range(3) | |
| ]) | |
| def mutate(values): | |
| for value in values: | |
| value.append(99) | |
| return values | |
| dataset = table.scan().to_contiguous_window_dataset( | |
| window_size=2, | |
| columns=["values"], | |
| group_key="episode", | |
| order_key="step", | |
| column_transforms={"values": mutate}, | |
| ) | |
| batched = dataset.__getitems__([0, 1, 0]) | |
| singles = [dataset[index] for index in (0, 1, 0)] | |
| self.assertEqual( | |
| [sample["values"] for sample in singles], | |
| [sample["values"] for sample in batched], | |
| ) | |
| def test_pad_tail_repeats_last_row_and_marks_real_padding(self): | |
| dataset = self._dataset( | |
| self._table(), tail="pad", pad_values={"value": -1}) | |
| self.assertEqual(6, len(dataset)) | |
| short_tail = dataset[1] | |
| long_tail = dataset[-1] | |
| self.assertEqual("episode-a", short_tail["episode"]) | |
| self.assertEqual([1, -1, -1], short_tail["value"]) | |
| self.assertEqual( | |
| [b"episode-a-1"] * 3, short_tail["payload"]) | |
| self.assertEqual([False, True, True], short_tail["is_pad"].tolist()) | |
| self.assertEqual("episode-b", long_tail["episode"]) | |
| self.assertEqual([103, -1, -1], long_tail["value"]) | |
| self.assertEqual([False, True, True], long_tail["is_pad"].tolist()) | |
| def test_error_tail_rejects_an_incomplete_scheduled_window(self): | |
| with self.assertRaisesRegex( | |
| ValueError, "episode-a.*incomplete.*window_size=3"): | |
| self._dataset(self._table(), tail="error") | |
| def test_stride_controls_scheduled_window_anchors(self): | |
| dataset = self._dataset(self._table(), stride=2, tail="pad") | |
| self.assertEqual( | |
| [("episode-a", 0), ("episode-b", 0), ("episode-b", 2)], | |
| [(dataset[index]["episode"], dataset[index]["step"]) | |
| for index in range(len(dataset))], | |
| ) | |
| self.assertEqual( | |
| [False, False, True], dataset[-1]["is_pad"].tolist()) | |
| def test_rejects_missing_and_duplicate_order_keys_within_a_group(self): | |
| gapped = self.conn.create_table( | |
| "gapped", schema=self._schema(), options=_TABLE_OPTIONS) | |
| gapped.add([ | |
| self._row("episode-a", 0), | |
| self._row("episode-a", 2), | |
| ]) | |
| with self.assertRaisesRegex( | |
| ValueError, "episode-a.*not contiguous.*0.*2"): | |
| self._dataset(gapped) | |
| table = self.conn.create_table( | |
| "duplicates", schema=self._schema(), options=_TABLE_OPTIONS) | |
| table.add([ | |
| self._row("episode-a", 0), | |
| self._row("episode-a", 0), | |
| self._row("episode-a", 1), | |
| ]) | |
| with self.assertRaisesRegex( | |
| ValueError, "episode-a.*duplicate.*order.*0"): | |
| self._dataset(table) | |
| def test_pins_snapshot_for_later_on_demand_reads(self): | |
| table = self._table() | |
| dataset = self._dataset(table) | |
| snapshot_id = dataset.snapshot_id | |
| table.add([self._row("episode-b", 4)]) | |
| self.assertEqual(snapshot_id, dataset.snapshot_id) | |
| self.assertNotEqual( | |
| snapshot_id, table.raw_table.snapshot_manager().get_latest_snapshot().id) | |
| self.assertEqual(2, len(dataset)) | |
| self.assertEqual([101, 102, 103], dataset[-1]["value"]) | |
| def test_snapshot_pin_clears_scan_mode_before_on_demand_reads(self): | |
| table = self._table() | |
| query = ScanQuery(table.raw_table.copy({"scan.mode": "latest-full"})) | |
| dataset = query.to_contiguous_window_dataset( | |
| window_size=3, | |
| columns=["value", "payload"], | |
| group_key="episode", | |
| order_key="step", | |
| ) | |
| self.assertEqual([100, 101, 102], dataset[0]["value"]) | |
| def test_pickle_round_trip_preserves_snapshot_and_window(self): | |
| dataset = self._dataset( | |
| self._table(), anchor_columns=["payload"]) | |
| expected = dataset[-1] | |
| restored = pickle.loads(pickle.dumps(dataset)) | |
| self.assertEqual(dataset.snapshot_id, restored.snapshot_id) | |
| self.assertEqual( | |
| dataset.snapshot_id, | |
| restored._table.options.scan_snapshot_id(), | |
| ) | |
| actual = restored[-1] | |
| self.assertEqual(expected["episode"], actual["episode"]) | |
| self.assertEqual(expected["step"], actual["step"]) | |
| self.assertEqual(expected["value"], actual["value"]) | |
| self.assertEqual(expected["payload"], actual["payload"]) | |
| self.assertTrue(torch.equal(expected["is_pad"], actual["is_pad"])) | |
| def test_projection_filter_transform_and_dataloader_workers(self): | |
| table = self._table() | |
| dataset = ( | |
| table.scan() | |
| .where("episode = 'episode-b'") | |
| .select(["value"]) | |
| .to_contiguous_window_dataset( | |
| window_size=2, | |
| group_key="episode", | |
| order_key="step", | |
| column_transforms={"value": _TensorColumnTransform()}, | |
| adapter=_WindowAdapter(), | |
| ) | |
| ) | |
| loader = torch.utils.data.DataLoader( | |
| dataset, batch_size=2, shuffle=False, num_workers=2) | |
| batches = list(loader) | |
| self.assertEqual(2, len(batches)) | |
| self.assertEqual(torch.int64, batches[0]["values"].dtype) | |
| self.assertEqual((2, 2), tuple(batches[0]["values"].shape)) | |
| self.assertEqual(torch.bool, batches[0]["padding_mask"].dtype) | |
| self.assertEqual([0, 1, 2], [ | |
| start for batch in batches for start in batch["start"].tolist() | |
| ]) | |
| self.assertEqual( | |
| [[100, 101], [101, 102], [102, 103]], | |
| [values for batch in batches for values in batch["values"].tolist()], | |
| ) | |
| self.assertTrue(all( | |
| episode == "episode-b" | |
| for batch in batches for episode in batch["episode"] | |
| )) | |
| def test_default_keys_and_public_from_query_entry_point(self): | |
| table = self.conn.create_table( | |
| "default_keys", | |
| schema=pa.schema([ | |
| pa.field("episode_index", pa.string(), nullable=False), | |
| pa.field("frame_index", pa.int32(), nullable=False), | |
| pa.field("value", pa.int32(), nullable=False), | |
| ]), | |
| options=_TABLE_OPTIONS, | |
| ) | |
| table.add([ | |
| {"episode_index": "episode-a", "frame_index": 0, "value": 10}, | |
| {"episode_index": "episode-a", "frame_index": 1, "value": 11}, | |
| ]) | |
| dataset = ContiguousWindowDataset.from_query( | |
| table.scan().select(["value"]), window_size=2) | |
| self.assertEqual(1, len(dataset)) | |
| self.assertEqual("episode-a", dataset[0]["episode_index"]) | |
| self.assertEqual(0, dataset[0]["frame_index"]) | |
| self.assertEqual([10, 11], dataset[0]["value"]) | |
| def test_rejects_scan_and_batch_vector_search_queries(self): | |
| table = self.conn.create_table( | |
| "vectors", | |
| schema=pa.schema([ | |
| pa.field("episode", pa.string(), nullable=False), | |
| pa.field("step", pa.int32(), nullable=False), | |
| pa.field("embedding", pa.list_(pa.float32(), 2)), | |
| ]), | |
| options=_TABLE_OPTIONS, | |
| ) | |
| table.add([ | |
| {"episode": "episode-a", "step": 0, "embedding": [1.0, 0.0]}, | |
| {"episode": "episode-a", "step": 1, "embedding": [0.0, 1.0]}, | |
| ]) | |
| kwargs = { | |
| "window_size": 2, | |
| "columns": ["embedding"], | |
| "group_key": "episode", | |
| "order_key": "step", | |
| } | |
| for query in (table.search([1.0, 0.0]), | |
| table.search_vectors([[1.0, 0.0]])): | |
| with self.subTest(query=type(query).__name__): | |
| with self.assertRaisesRegex(TypeError, "only supported on scan"): | |
| query.to_contiguous_window_dataset(**kwargs) | |
| with self.assertRaisesRegex(TypeError, "only supported on scan"): | |
| ContiguousWindowDataset.from_query(query, **kwargs) | |
| def test_reads_a_pinned_tag_after_its_snapshot_file_is_removed(self): | |
| table = self._table() | |
| table.raw_table.create_tag("v1") | |
| dataset = table.scan(tag_name="v1").to_contiguous_window_dataset( | |
| window_size=3, columns=["value"], | |
| group_key="episode", order_key="step") | |
| raw_table = table.raw_table | |
| raw_table.file_io.delete_quietly( | |
| raw_table.snapshot_manager().get_snapshot_path(dataset.snapshot_id)) | |
| self.assertEqual("v1", dataset._table.options.scan_tag_name()) | |
| self.assertIsNone(dataset._table.options.scan_snapshot_id()) | |
| self.assertEqual([100, 101, 102], dataset[0]["value"]) | |
| def test_reuses_context_plan_across_projections_and_pickle(self): | |
| dataset = self._dataset(self._table(), anchor_columns=["payload"]) | |
| with patch.object( | |
| TableScan, "plan", autospec=True, | |
| side_effect=TableScan.plan) as plan: | |
| samples = [dataset[index] for index in range(len(dataset))] | |
| restored = pickle.loads(pickle.dumps(dataset)) | |
| samples.extend(restored.__getitems__(list(range(len(dataset))))) | |
| self.assertEqual(0, plan.call_count) | |
| self.assertEqual( | |
| [[100, 101, 102], [101, 102, 103]] * 2, | |
| [sample["value"] for sample in samples]) | |
| self.assertEqual( | |
| [[b"episode-b-0"], [b"episode-b-1"]] * 2, | |
| [sample["payload"] for sample in samples]) | |
| @pytest.mark.python_read | |
| def test_window_ranges_do_not_need_a_second_predicate_scan(self): | |
| table = self.conn.create_table( | |
| "sparse_reads", schema=self._schema(), options=_TABLE_OPTIONS) | |
| table.add([self._row("episode-a", step) for step in range(12)]) | |
| dataset = table.scan().to_contiguous_window_dataset( | |
| window_size=2, columns=["value"], | |
| group_key="episode", order_key="step") | |
| filtered_rows = [] | |
| original_filter = FilterRecordBatchReader._filter_batch | |
| def record_filter(reader, batch): | |
| filtered_rows.append(batch.num_rows) | |
| return original_filter(reader, batch) | |
| with patch.object(FilterRecordBatchReader, "_filter_batch", record_filter): | |
| samples = dataset.__getitems__([5, 1, 5]) | |
| self.assertEqual([[5, 6], [1, 2], [5, 6]], | |
| [sample["value"] for sample in samples]) | |
| # IndexedSplit already selects these sparse rows; a predicate scanner | |
| # would evaluate the same membership again for every fetched batch. | |
| self.assertEqual(0, sum(filtered_rows)) | |
| def test_rejects_query_authorization_that_masks_row_ids(self): | |
| table = self._table() | |
| auth = TableQueryAuthResult( | |
| filter=None, | |
| column_masking={"_ROW_ID": json.dumps({"name": "NULL"})}, | |
| ) | |
| table.raw_table.catalog_environment.table_query_auth = ( | |
| lambda options, table_identifier: lambda select: auth) | |
| with self.assertRaisesRegex(ValueError, "masks _ROW_ID"): | |
| self._dataset(table) | |
| def test_rejects_reads_after_the_pinned_tag_moves_to_another_snapshot(self): | |
| table = self._table() | |
| raw_table = table.raw_table | |
| raw_table.create_tag("v1") | |
| dataset = table.scan(tag_name="v1").to_contiguous_window_dataset( | |
| window_size=3, columns=["value"], | |
| group_key="episode", order_key="step") | |
| restored = pickle.loads(pickle.dumps(dataset)) | |
| table.add([self._row("episode-b", 4)]) | |
| raw_table.replace_tag("v1") | |
| for pinned in (dataset, restored): | |
| with self.assertRaisesRegex(RuntimeError, "pinned snapshot moved"): | |
| pinned[0] | |
| def test_rejects_query_authorization_that_masks_the_group_key(self): | |
| table = self._table() | |
| auth = TableQueryAuthResult( | |
| filter=None, | |
| column_masking={"episode": json.dumps({ | |
| "name": "SUBSTRING", | |
| "inputs": [ | |
| {"name": "episode", "type": "STRING NOT NULL"}, 1, 1, | |
| ], | |
| })}, | |
| ) | |
| table.raw_table.catalog_environment.table_query_auth = ( | |
| lambda options, table_identifier: lambda select: auth) | |
| with self.assertRaisesRegex(ValueError, r"masks \['episode'\]"): | |
| self._dataset(table) | |
| @pytest.mark.python_read | |
| def test_reads_only_the_files_and_row_ranges_a_window_touches(self): | |
| table = self.conn.create_table( | |
| "many_files", schema=self._schema(), options=_TABLE_OPTIONS) | |
| for batch in range(4): | |
| table.add([ | |
| self._row("episode-a", batch * 8 + offset) | |
| for offset in range(8) | |
| ]) | |
| dataset = table.scan().to_contiguous_window_dataset( | |
| window_size=3, columns=["value"], | |
| group_key="episode", order_key="step") | |
| opened = [] | |
| decoded = [] | |
| original_init = FormatPyArrowReader.__init__ | |
| original_read = FormatPyArrowReader.read_arrow_batch | |
| def recording_init(reader, *args, **kwargs): | |
| original_init(reader, *args, **kwargs) | |
| opened.append((args[2], kwargs.get("row_ranges"))) | |
| def recording_read(reader): | |
| batch = original_read(reader) | |
| decoded.append(0 if batch is None else batch.num_rows) | |
| return batch | |
| with patch.object(FormatPyArrowReader, "__init__", recording_init), \ | |
| patch.object( | |
| FormatPyArrowReader, "read_arrow_batch", recording_read): | |
| sample = dataset[len(dataset) - 1] | |
| self.assertEqual([29, 30, 31], sample["value"]) | |
| self.assertEqual(1, len({file_path for file_path, _ in opened})) | |
| # The reader receives positions inside the file, so global rows 29-31 | |
| # of the fourth eight-row file arrive as 5-7. | |
| self.assertEqual([[(5, 7)]], [ranges for _, ranges in opened]) | |
| self.assertEqual(3, sum(decoded)) | |
| def test_rejects_nan_group_keys_that_never_compare_equal(self): | |
| table = self.conn.create_table( | |
| "nan_groups", | |
| schema=pa.schema([ | |
| pa.field("episode", pa.float64()), | |
| pa.field("step", pa.int32(), nullable=False), | |
| pa.field("value", pa.int32(), nullable=False), | |
| ]), | |
| options=_TABLE_OPTIONS, | |
| ) | |
| table.add([ | |
| {"episode": float("nan"), "step": 0, "value": 0}, | |
| {"episode": float("nan"), "step": 1, "value": 1}, | |
| ]) | |
| with self.assertRaisesRegex(ValueError, "episode must not contain NaN"): | |
| table.scan().to_contiguous_window_dataset( | |
| window_size=2, columns=["value"], | |
| group_key="episode", order_key="step") | |
| def test_decodes_video_windows_and_reuses_overlapping_frames(self): | |
| table = self.conn.create_table( | |
| "videos", | |
| schema=pa.schema([ | |
| pa.field("episode", pa.string(), nullable=False), | |
| pa.field("step", pa.int32(), nullable=False), | |
| pa.field("video", pa.large_binary()), | |
| ]), | |
| options=dict(_TABLE_OPTIONS, **{ | |
| "video-frame-field": "video", | |
| "blob-as-descriptor": "true", | |
| }), | |
| ) | |
| video_path = os.path.join(self.temp_dir, "episode-a.mp4") | |
| with open(video_path, "wb") as output: | |
| output.write(b"encoded-video") | |
| table.add_video( | |
| pmm.Blob.from_local(video_path), | |
| [ | |
| {"episode": "episode-a", "step": step} | |
| for step in range(3) | |
| ], | |
| first_frame=7, | |
| ) | |
| decoded = [] | |
| class Decoder: | |
| def get_frames_at(self, *, indices): | |
| decoded.append(tuple(indices)) | |
| return SimpleNamespace(data=torch.tensor(indices)) | |
| with patch( | |
| "pypaimon.multimodal.lerobot.dataset._open_video_decoder", | |
| return_value=Decoder()): | |
| dataset = table.scan().to_contiguous_window_dataset( | |
| window_size=2, columns=["video"], | |
| group_key="episode", order_key="step") | |
| try: | |
| samples = dataset.__getitems__([0, 1]) | |
| finally: | |
| close = getattr(dataset, "close", None) | |
| if close is not None: | |
| close() | |
| self.assertEqual([(7, 8, 9)], decoded) | |
| self.assertEqual( | |
| [[7, 8], [8, 9]], | |
| [[int(frame) for frame in sample["video"]] for sample in samples], | |
| ) | |
| for sample in samples: | |
| for frame in sample["video"]: | |
| self.assertEqual( | |
| frame.numel() * frame.element_size(), | |
| frame.untyped_storage().nbytes(), | |
| ) | |
| def test_decodes_encoded_video_windows_with_multiple_workers(self): | |
| try: | |
| import av # noqa: F401 | |
| from PIL import Image | |
| except ImportError: | |
| self.skipTest("PyAV and Pillow are required for video decoding") | |
| table = self.conn.create_table( | |
| "encoded_videos", | |
| schema=pa.schema([ | |
| pa.field("episode", pa.string(), nullable=False), | |
| pa.field("step", pa.int32(), nullable=False), | |
| pa.field("video", pa.large_binary()), | |
| ]), | |
| options=dict(_TABLE_OPTIONS, **{ | |
| "video-frame-field": "video", | |
| "blob-as-descriptor": "true", | |
| }), | |
| ) | |
| table.add_images_as_video( | |
| [ | |
| Image.new("RGB", (8, 6), color) | |
| for color in ("red", "green", "blue") | |
| ], | |
| [ | |
| {"episode": "episode-a", "step": step} | |
| for step in range(3) | |
| ], | |
| fps=10, | |
| codec="mpeg4", | |
| pixel_format="yuv420p", | |
| gop_size=1, | |
| codec_options={"qscale": "3"}, | |
| ) | |
| dataset = table.scan().to_contiguous_window_dataset( | |
| window_size=2, | |
| columns=["video"], | |
| group_key="episode", | |
| order_key="step", | |
| column_transforms={"video": torch.stack}, | |
| video_backend="pyav", | |
| ) | |
| try: | |
| batches = list(torch.utils.data.DataLoader( | |
| dataset, batch_size=2, num_workers=2, shuffle=False)) | |
| finally: | |
| dataset.close() | |
| self.assertEqual(1, len(batches)) | |
| self.assertEqual((2, 2, 3, 6, 8), tuple(batches[0]["video"].shape)) | |
| first = batches[0]["video"][0, 0].float().mean(dim=(1, 2)) | |
| last = batches[0]["video"][1, 1].float().mean(dim=(1, 2)) | |
| self.assertGreater(first[0], first[1] + 100) | |
| self.assertGreater(first[0], first[2] + 100) | |
| self.assertGreater(last[2], last[0] + 100) | |
| self.assertGreater(last[2], last[1] + 100) | |
| def test_validates_configuration_and_scan_only_contract(self): | |
| table = self._table() | |
| query = table.scan() | |
| for name, value in ( | |
| ("window_size", 0), | |
| ("stride", 0), | |
| ("tail", "unknown"), | |
| ("group_key", "missing"), | |
| ("order_key", "missing")): | |
| kwargs = { | |
| "window_size": 2, | |
| "columns": ["value"], | |
| "stride": 1, | |
| "tail": "drop", | |
| "group_key": "episode", | |
| "order_key": "step", | |
| } | |
| kwargs[name] = value | |
| with self.subTest(name=name), self.assertRaises((TypeError, ValueError)): | |
| query.to_contiguous_window_dataset(**kwargs) | |
| reserved_table = self.conn.create_table( | |
| "reserved", schema=pa.schema([ | |
| pa.field("is_pad", pa.string(), nullable=False), | |
| pa.field("step", pa.int32(), nullable=False), | |
| pa.field("value", pa.int32(), nullable=False), | |
| ]), options=_TABLE_OPTIONS) | |
| with self.assertRaisesRegex(ValueError, "must not be is_pad"): | |
| reserved_table.scan().to_contiguous_window_dataset( | |
| window_size=2, columns=["value"], | |
| group_key="is_pad", order_key="step") | |
| with self.assertRaisesRegex(TypeError, "only supported on scan"): | |
| table.search("anything", column="episode").to_contiguous_window_dataset( | |
| window_size=2, columns=["value"], | |
| group_key="episode", order_key="step") | |
| with self.assertRaisesRegex(ValueError, "video_backend must be"): | |
| table.scan().to_contiguous_window_dataset( | |
| window_size=2, columns=["value"], | |
| group_key="episode", order_key="step", | |
| video_backend="unsupported") | |
| if __name__ == "__main__": | |
| unittest.main() |