mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-26 19:26:16 +00:00
Pepijn
111 lines
3.9 KiB
Python
111 lines
3.9 KiB
Python
#!/usr/bin/env python
|
|
|
|
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
|
|
#
|
|
# Licensed 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
|
|
|
|
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
|
|
import fsspec
|
|
import pytest
|
|
|
|
pytest.importorskip("pyarrow", reason="pyarrow is required (install lerobot[dataset])")
|
|
|
|
import pyarrow as pa
|
|
import pyarrow.parquet as pq
|
|
|
|
from lerobot.datasets.episode_parquet import EpisodeParquetReader
|
|
|
|
|
|
def _table(episodes: list[int]) -> pa.Table:
|
|
frame_counts: dict[int, int] = {}
|
|
frame_indices = []
|
|
values = []
|
|
ignored = []
|
|
for episode in episodes:
|
|
frame_index = frame_counts.get(episode, 0)
|
|
frame_counts[episode] = frame_index + 1
|
|
frame_indices.append(frame_index)
|
|
values.append(episode * 10 + frame_index)
|
|
ignored.append(f"ignored-{episode}-{frame_index}")
|
|
return pa.table(
|
|
{
|
|
"episode_index": episodes,
|
|
"frame_index": frame_indices,
|
|
"value": values,
|
|
"ignored": ignored,
|
|
}
|
|
)
|
|
|
|
|
|
def _write_episode_row_groups(path: Path) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
writer = pq.ParquetWriter(path, _table([0]).schema)
|
|
try:
|
|
writer.write_table(_table([0, 0]))
|
|
writer.write_table(_table([1, 1, 1]))
|
|
finally:
|
|
writer.close()
|
|
|
|
|
|
def test_reader_projects_columns_and_reads_matching_row_group(tmp_path: Path) -> None:
|
|
path = tmp_path / "data/chunk-000/file-000.parquet"
|
|
_write_episode_row_groups(path)
|
|
reader = EpisodeParquetReader(tmp_path, columns=("episode_index", "frame_index", "value"))
|
|
|
|
table = reader.read_episode(path.relative_to(tmp_path), episode_index=1, expected_rows=3)
|
|
|
|
assert table.column_names == ["episode_index", "frame_index", "value"]
|
|
assert table.column("value").to_pylist() == [10, 11, 12]
|
|
|
|
|
|
def test_reader_filters_legacy_mixed_row_group(tmp_path: Path) -> None:
|
|
path = tmp_path / "data/chunk-000/file-000.parquet"
|
|
path.parent.mkdir(parents=True)
|
|
pq.write_table(_table([0, 0, 1, 1, 1]), path)
|
|
reader = EpisodeParquetReader(tmp_path, columns=("episode_index", "frame_index", "value"))
|
|
|
|
table = reader.read_episode(path.relative_to(tmp_path), episode_index=1, expected_rows=3)
|
|
|
|
assert table.column("episode_index").to_pylist() == [1, 1, 1]
|
|
assert table.column("frame_index").to_pylist() == [0, 1, 2]
|
|
|
|
|
|
def test_reader_rejects_partial_episode(tmp_path: Path) -> None:
|
|
path = tmp_path / "data/chunk-000/file-000.parquet"
|
|
path.parent.mkdir(parents=True)
|
|
pq.write_table(_table([2, 2]), path)
|
|
reader = EpisodeParquetReader(tmp_path, columns=("episode_index", "frame_index"))
|
|
|
|
with pytest.raises(ValueError, match="expected 3 rows, found 2"):
|
|
reader.read_episode(path.relative_to(tmp_path), episode_index=2, expected_rows=3)
|
|
|
|
|
|
def test_reader_rejects_missing_episode(tmp_path: Path) -> None:
|
|
path = tmp_path / "data/chunk-000/file-000.parquet"
|
|
path.parent.mkdir(parents=True)
|
|
pq.write_table(_table([0, 0]), path)
|
|
reader = EpisodeParquetReader(tmp_path, columns=("episode_index", "frame_index"))
|
|
|
|
with pytest.raises(ValueError, match="episode 4"):
|
|
reader.read_episode(path.relative_to(tmp_path), episode_index=4, expected_rows=1)
|
|
|
|
|
|
def test_reader_supports_fsspec_remote_root() -> None:
|
|
filesystem = fsspec.filesystem("memory")
|
|
root = "memory://episode-reader"
|
|
path = "episode-reader/data/chunk-000/file-000.parquet"
|
|
with filesystem.open(path, "wb") as output:
|
|
pq.write_table(_table([0, 0, 0]), output)
|
|
reader = EpisodeParquetReader(root, columns=("episode_index", "value"))
|
|
|
|
table = reader.read_episode("data/chunk-000/file-000.parquet", episode_index=0, expected_rows=3)
|
|
|
|
assert table.column("value").to_pylist() == [0, 1, 2]
|