This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new 8a0b4ef768 [python] Import LeRobot video datasets (#9494)
8a0b4ef768 is described below
commit 8a0b4ef768bf1b9c3e76ee639ab15f194cc2a077
Author: XiaoHongbo <[email protected]>
AuthorDate: Thu Sep 10 23:33:54 2026 +0800
[python] Import LeRobot video datasets (#9494)
---
docs/docs/pypaimon/lerobot.md | 6 +
docs/docs/pypaimon/video.md | 8 +-
paimon-python/pypaimon/multimodal/lerobot/api.py | 49 +-
.../pypaimon/multimodal/lerobot/loader.py | 363 +++++++++-
.../pypaimon/multimodal/lerobot/metadata.py | 88 ++-
.../pypaimon/multimodal/lerobot/schema.py | 14 +-
.../pypaimon/multimodal/lerobot/source.py | 133 +++-
.../pypaimon/multimodal/lerobot/writer.py | 6 +
.../tests/data_evolution_row_rolling_test.py | 187 ++++++
.../pypaimon/tests/multimodal_lerobot_test.py | 733 ++++++++++++++++++++-
.../tests/multimodal_lerobot_writer_test.py | 16 +
paimon-python/pypaimon/write/file_store_write.py | 7 +
paimon-python/pypaimon/write/table_write.py | 16 +
.../pypaimon/write/writer/blob_file_writer.py | 6 +-
paimon-python/pypaimon/write/writer/blob_writer.py | 15 +-
.../write/writer/dedicated_format_writer.py | 44 +-
.../pypaimon/write/writer/vector_writer.py | 27 +-
17 files changed, 1643 insertions(+), 75 deletions(-)
diff --git a/docs/docs/pypaimon/lerobot.md b/docs/docs/pypaimon/lerobot.md
index 9554f0e234..722d63f3f2 100644
--- a/docs/docs/pypaimon/lerobot.md
+++ b/docs/docs/pypaimon/lerobot.md
@@ -107,6 +107,12 @@ together, and keep writers paused until the call returns.
Scalars map to scalar types, vectors to `VECTOR`, higher-rank tensors to nested
`ARRAY`, and images to `BLOB`. Images keep their compressed bytes.
+Video features map to `BLOB`. Frame rows reference MP4 payloads copied once per
+aligned file group. Video imports use the video grouping policy and check
+rolling before each Episode. They require a bucket-unaware table. Read them
+with a Paimon scan and `VideoFrameCollator`; `PaimonLeRobotDataset` currently
+supports image features only.
+
## Capture LeRobot frames directly into Paimon
`PaimonLeRobotWriter` implements the write-side surface used by LeRobot's
diff --git a/docs/docs/pypaimon/video.md b/docs/docs/pypaimon/video.md
index 987cb54322..881ce62fde 100644
--- a/docs/docs/pypaimon/video.md
+++ b/docs/docs/pypaimon/video.md
@@ -89,11 +89,9 @@ frames.add_video(video, episode_43_rows)
```
The writer deduplicates exact payload descriptor identity inside each `.video`
-file. Payload boundaries may be nested across configured video columns, but
-every change must occur at an episode boundary. When a file target is reached
-within an episode, normal, BLOB, and vector rolling waits for the next payload
-boundary and closes the aligned file group together. A single large episode
-may therefore exceed the target. The normal `.blob` format is unchanged.
+file. Its video grouping policy coordinates normal, BLOB, and vector rolling
+at payload boundaries. A file may exceed its target before the next boundary.
+The normal `.blob` format is unchanged.
### Update frame rows and replace a video
diff --git a/paimon-python/pypaimon/multimodal/lerobot/api.py
b/paimon-python/pypaimon/multimodal/lerobot/api.py
index 84e3d2eed5..20a054f90f 100644
--- a/paimon-python/pypaimon/multimodal/lerobot/api.py
+++ b/paimon-python/pypaimon/multimodal/lerobot/api.py
@@ -37,6 +37,7 @@ from pypaimon.multimodal.lerobot.schema import (
_require_v3,
_schema_from_info,
_validate_v3_required_features,
+ _video_feature_names,
)
from pypaimon.multimodal.lerobot.source import (
_close_quietly,
@@ -52,6 +53,12 @@ from pypaimon.multimodal.source_utils import (
)
+_VIDEO_LAYOUT_ERROR = (
+ "LeRobot video import requires a bucket-unaware target table so each "
+ "Episode is written by one writer."
+)
+
+
def load_from_lerobot(
connection,
table_name: str,
@@ -94,14 +101,20 @@ def load_from_lerobot(
_positive_integer(local_info.get("fps"), "fps")
_validated_counts(local_info, resolved_source.path)
_validate_v3_required_features(local_info)
+ video_fields = _video_feature_names(local_info)
LeRobotDataset = _import_lerobot_dataset()
dataset = _open_resolved_dataset(
- LeRobotDataset, resolved_source, local_info)
+ LeRobotDataset,
+ resolved_source,
+ local_info,
+ download_videos=bool(video_fields),
+ )
try:
info = dict(dataset.meta.info)
_require_v3(info, resolved_source.path)
_validated_counts(info, resolved_source.path)
_validate_v3_required_features(info)
+ video_fields = _video_feature_names(info)
lerobot_schema = _schema_from_info(info)
metadata = _load_dataset_metadata(
@@ -117,6 +130,7 @@ def load_from_lerobot(
options,
metadata,
tag_name,
+ video_fields,
)
finally:
close = getattr(dataset, "close", None)
@@ -134,14 +148,21 @@ def _import_dataset(
batch_size,
options,
metadata,
- tag_name):
+ tag_name,
+ video_fields):
table = _create_target_table(
- connection, table_name, source_schema, options, metadata)
+ connection,
+ table_name,
+ source_schema,
+ options,
+ metadata,
+ video_fields,
+ )
tables = _prepare_metadata_tables(
connection, table.raw_table, metadata)
episodes_snapshot_id = _append_arrow_tables(
tables["episodes"],
- _validated_episode_tables(metadata),
+ _validated_episode_tables(metadata, video_fields),
)
frames_snapshot_id = None
if int(info["total_frames"]) > 0:
@@ -153,6 +174,7 @@ def _import_dataset(
source_schema,
batch_size,
metadata,
+ video_fields,
)
_commit_metadata(
connection,
@@ -195,8 +217,25 @@ def _required_count(info, name, source):
def _create_target_table(
- connection, table_name, source_schema, options, metadata):
+ connection, table_name, source_schema, options, metadata,
+ video_fields=()):
create_options = dict(options or {})
+ configured = create_options.get("video-frame-field")
+ if configured is not None:
+ requested = {
+ name.strip() for name in str(configured).split(",")
+ if name.strip()
+ }
+ if requested != set(video_fields):
+ raise ValueError(
+ "LeRobot video features %s do not match "
+ "'video-frame-field'=%r."
+ % (list(video_fields), configured)
+ )
+ if video_fields:
+ if str(create_options.get("bucket", "-1")).strip() != "-1":
+ raise ValueError(_VIDEO_LAYOUT_ERROR)
+ create_options["video-frame-field"] = ",".join(video_fields)
managed_options = _managed_table_options(
connection._identifier(table_name), metadata)
reserved_options = set(_COMPANION_OPTION_KEYS.values()).intersection(
diff --git a/paimon-python/pypaimon/multimodal/lerobot/loader.py
b/paimon-python/pypaimon/multimodal/lerobot/loader.py
index 71719869c3..6925886b97 100644
--- a/paimon-python/pypaimon/multimodal/lerobot/loader.py
+++ b/paimon-python/pypaimon/multimodal/lerobot/loader.py
@@ -19,7 +19,10 @@
import io
import math
import numbers
+from array import array
+from bisect import bisect_left
from pathlib import Path
+from urllib.parse import unquote, urlparse
import pyarrow as pa
@@ -27,6 +30,7 @@ from pypaimon.multimodal.arrow_utils import strict_arrow_table
from pypaimon.multimodal.hdf5 import _SnapshotRecorder
from pypaimon.multimodal.lerobot.schema import _feature_shape
from pypaimon.multimodal.table import _target_schema
+from pypaimon.table.row.blob import VideoFrameDescriptor
_DECLARED_NUMERIC_RANGES = {
@@ -49,6 +53,7 @@ _NUMERIC_DTYPES = {
"float64",
}
_BOOLEAN_DTYPES = {"bool", "boolean"}
+_VIDEO_TIMESTAMP_TOLERANCE = 1e-4
def _strict_lerobot_table(data, target_schema, source, batch_index):
@@ -68,7 +73,8 @@ def _write_dataset(
source,
source_schema,
batch_size,
- metadata):
+ metadata,
+ video_fields=()):
target_schema = _target_schema(table.raw_table)
write_builder = table.raw_table.new_batch_write_builder()
table_write = None
@@ -81,13 +87,23 @@ def _write_dataset(
expected_tasks = set()
observed_tasks = set()
snapshot_recorder = _SnapshotRecorder()
+ video_sources = {}
try:
table_write = write_builder.new_write()
+ reader_factory = getattr(
+ dataset, "video_uri_reader_factory", None)
+ if video_fields and reader_factory is not None:
+ table_write.with_blob_uri_reader_factory(reader_factory)
table_commit = write_builder.new_commit()
table_commit.add_commit_callback(snapshot_recorder)
- for episode_index, episode_begin, task_indices, begin, end in \
- _episode_batches(dataset, info, batch_size, episodes):
+ for source_episode, episode_index, episode_begin, task_indices, \
+ begin, end in _episode_batches(
+ info,
+ batch_size,
+ episodes,
+ video_fields,
+ ):
if episode_index != current_episode:
if current_episode is not None:
_validate_episode_tasks(
@@ -95,8 +111,18 @@ def _write_dataset(
current_episode = episode_index
expected_tasks = set(task_indices)
observed_tasks = set()
+ if video_fields:
+ table_write.roll_before_group_if_needed(
+ int(source_episode["length"]))
batch = _read_batch(
- dataset, info, begin, end, source_schema)
+ dataset,
+ info,
+ begin,
+ end,
+ source_schema,
+ episode=source_episode,
+ video_sources=video_sources,
+ )
seen_tasks = _validate_frame_controls(
batch,
int(info["fps"]),
@@ -146,9 +172,9 @@ def _write_dataset(
table_commit.close()
-def _episode_batches(dataset, info, batch_size, episodes):
+def _episode_batches(info, batch_size, episodes, video_fields=()):
episode_count = int(info.get("total_episodes", 0))
- total_frames = int(info.get("total_frames", len(dataset)))
+ total_frames = int(info["total_frames"])
expected_begin = 0
for ordinal in range(episode_count):
episode = episodes.iloc[ordinal] if hasattr(episodes, "iloc") \
@@ -164,9 +190,11 @@ def _episode_batches(dataset, info, batch_size, episodes):
% ordinal)
episode_begin = begin
task_indices = episode.get("task_indices", ())
+ source_episode = episode if video_fields else None
while begin < end:
batch_end = min(begin + batch_size, end)
yield (
+ source_episode,
episode_index,
episode_begin,
task_indices,
@@ -280,7 +308,9 @@ def _control_integer(value, name, frame_index):
return int(value)
-def _read_batch(dataset, info, begin, end, schema):
+def _read_batch(
+ dataset, info, begin, end, schema, episode=None,
+ video_sources=None):
read_batch = getattr(dataset, "read_batch", None)
if callable(read_batch):
raw = read_batch(begin, end)
@@ -291,32 +321,331 @@ def _read_batch(dataset, info, begin, end, schema):
elif not isinstance(raw, pa.Table):
raw = pa.Table.from_pydict(raw)
features = info["features"]
+ video_rows = None
+ video_timestamp_type = None
+ if any(feature.get("dtype") == "video"
+ for feature in features.values()):
+ video_rows = _validate_video_rows(
+ raw, info, episode, begin, end)
+ video_timestamp_type = raw.schema.field("timestamp").type
arrays = []
fields = []
for name, feature in features.items():
field = schema.field(name)
dtype = feature["dtype"]
- if name not in raw.column_names:
+ if dtype == "video":
+ values = _video_frame_descriptors(
+ dataset,
+ info,
+ episode,
+ video_rows,
+ name,
+ feature,
+ begin,
+ end,
+ video_sources if video_sources is not None else {},
+ video_timestamp_type,
+ )
+ elif name not in raw.column_names:
raise ValueError(
"LeRobot data is missing metadata feature %s." % name)
- values = raw.column(name).to_pylist()
- if dtype == "image":
- image_reader = getattr(dataset, "image_bytes", None)
- if callable(image_reader):
- values = [image_reader(value) for value in values]
+ else:
+ values = raw.column(name).to_pylist()
+ if dtype == "image":
+ image_reader = getattr(dataset, "image_bytes", None)
+ if callable(image_reader):
+ values = [image_reader(value) for value in values]
+ else:
+ values = [_image_bytes(value, dataset.root)
+ for value in values]
else:
- values = [_image_bytes(value, dataset.root)
+ values = [_normalize_value(value, feature, name)
for value in values]
- else:
- values = [_normalize_value(value, feature, name)
- for value in values]
arrays.append(_safe_array(values, field, name, dtype))
fields.append(field)
return pa.Table.from_arrays(arrays, schema=pa.schema(fields))
+def _video_frame_descriptors(
+ dataset, info, episode, video_rows, name, feature, begin, end, cache,
+ timestamp_type):
+ if episode is None:
+ raise ValueError("LeRobot video import requires Episode metadata.")
+ episode_begin = _nonnegative_integer(
+ episode["dataset_from_index"], "dataset_from_index")
+ episode_end = _nonnegative_integer(
+ episode["dataset_to_index"], "dataset_to_index")
+ if begin < episode_begin or end > episode_end:
+ raise ValueError(
+ "LeRobot video batch [%d, %d) crosses Episode range [%d, %d)."
+ % (begin, end, episode_begin, episode_end)
+ )
+
+ fps = _video_fps(info, feature, name)
+ prefix = "videos/%s/" % name
+ try:
+ chunk_index = _nonnegative_integer(
+ episode[prefix + "chunk_index"], prefix + "chunk_index")
+ file_index = _nonnegative_integer(
+ episode[prefix + "file_index"], prefix + "file_index")
+ from_timestamp = float(_python_scalar(
+ episode[prefix + "from_timestamp"]))
+ to_timestamp = float(_python_scalar(
+ episode[prefix + "to_timestamp"]))
+ except (KeyError, TypeError, ValueError) as error:
+ raise ValueError(
+ "LeRobot Episode metadata is missing video mapping for %s."
+ % name
+ ) from error
+
+ episode_length = episode_end - episode_begin
+ _validate_video_timestamp_range(
+ from_timestamp, to_timestamp, fps, episode_length, name)
+
+ source_key = (name, chunk_index, file_index)
+ source = cache.get(source_key)
+ if source is None:
+ uri, length = _video_source(
+ dataset, info, episode, name, chunk_index, file_index)
+ source = (uri, length, _video_sample_timestamps(dataset, uri))
+ cache[source_key] = source
+ uri, length, sample_timestamps = source
+ return [
+ VideoFrameDescriptor(
+ uri,
+ 0,
+ length,
+ _video_frame_ordinal(
+ sample_timestamps,
+ timestamp,
+ from_timestamp,
+ timestamp_type,
+ name,
+ ),
+ ).serialize()
+ for unused_frame_index, timestamp in video_rows
+ ]
+
+
+def _video_sample_timestamps(dataset, uri):
+ resolver = getattr(dataset, "video_sample_timestamps", None)
+ if callable(resolver):
+ values = resolver(uri)
+ else:
+ try:
+ import av
+ except ImportError as error:
+ raise ImportError(
+ "LeRobot video import requires PyAV. Install it with "
+ "`pip install 'pypaimon[lerobot]'`.") from error
+
+ input_stream = None
+ parsed = urlparse(uri)
+ if parsed.scheme in ("", "file"):
+ source = unquote(parsed.path) if parsed.scheme else uri
+ else:
+ factory = getattr(dataset, "video_uri_reader_factory", None)
+ if factory is None:
+ raise ValueError(
+ "LeRobot video source %s cannot be inspected." % uri)
+ input_stream = factory.create(uri).new_input_stream(uri)
+ source = input_stream
+ try:
+ with av.open(source) as container:
+ stream = container.streams.video[0]
+ values = [
+ float(packet.pts * (packet.time_base or stream.time_base))
+ for packet in container.demux(stream)
+ if packet.pts is not None and not packet.is_discard
+ ]
+ finally:
+ if input_stream is not None:
+ input_stream.close()
+
+ timestamps = array("d", sorted(float(value) for value in values))
+ if not timestamps or any(not math.isfinite(value) for value in timestamps):
+ raise ValueError(
+ "LeRobot video source %s has no valid frame timestamps." % uri)
+ return timestamps
+
+
+def _video_frame_ordinal(
+ timestamps, frame_timestamp, from_timestamp, timestamp_type, name):
+ timestamp = from_timestamp + frame_timestamp
+ position = bisect_left(timestamps, timestamp)
+ candidates = []
+ if position:
+ candidates.append(position - 1)
+ if position < len(timestamps):
+ candidates.append(position)
+ ordinal = min(
+ candidates,
+ key=lambda index: (abs(timestamps[index] - timestamp), index),
+ )
+ expected = pa.scalar(
+ timestamps[ordinal] - from_timestamp, type=timestamp_type).as_py()
+ if not math.isclose(
+ frame_timestamp,
+ float(expected),
+ rel_tol=0.0,
+ abs_tol=_VIDEO_TIMESTAMP_TOLERANCE):
+ raise ValueError(
+ "LeRobot video feature %s has no frame matching timestamp %s."
+ % (name, timestamp))
+ return ordinal
+
+
+def _validate_video_rows(raw, info, episode, begin, end):
+ if episode is None:
+ raise ValueError("LeRobot video import requires Episode metadata.")
+ required = ("episode_index", "frame_index", "timestamp")
+ missing = [name for name in required if name not in raw.column_names]
+ if missing:
+ raise ValueError(
+ "LeRobot video import requires frame columns %s."
+ % ", ".join(missing)
+ )
+ episode_index = _nonnegative_integer(
+ episode["episode_index"], "episode_index")
+ episode_begin = _nonnegative_integer(
+ episode["dataset_from_index"], "dataset_from_index")
+ expected_frames = range(begin - episode_begin, end - episode_begin)
+ actual_episodes = raw.column("episode_index").to_pylist()
+ actual_frames = raw.column("frame_index").to_pylist()
+ timestamps = raw.column("timestamp").to_pylist()
+ timestamp_type = raw.schema.field("timestamp").type
+ fps = _positive_fps(info.get("fps"), "dataset")
+ result = []
+ for offset, expected_frame in enumerate(expected_frames):
+ actual_episode = _nonnegative_integer(
+ actual_episodes[offset], "episode_index")
+ actual_frame = _nonnegative_integer(
+ actual_frames[offset], "frame_index")
+ if actual_episode != episode_index or actual_frame != expected_frame:
+ raise ValueError(
+ "LeRobot frame rows do not match Episode %d range [%d, %d)."
+ % (episode_index, begin, end)
+ )
+ timestamp = timestamps[offset]
+ expected_timestamp = pa.scalar(
+ expected_frame / fps, type=timestamp_type).as_py()
+ if (isinstance(timestamp, bool)
+ or not isinstance(timestamp, numbers.Real)
+ or not math.isclose(
+ float(timestamp), float(expected_timestamp),
+ rel_tol=0.0,
+ abs_tol=_VIDEO_TIMESTAMP_TOLERANCE)):
+ raise ValueError(
+ "LeRobot frame %d has timestamp %r; expected %r."
+ % (actual_frame, timestamp, expected_timestamp)
+ )
+ result.append((actual_frame, float(timestamp)))
+ return result
+
+
+def _validate_video_timestamp_range(
+ from_timestamp, to_timestamp, fps, episode_length, name):
+ if (not math.isfinite(from_timestamp)
+ or not math.isfinite(to_timestamp)
+ or from_timestamp < 0
+ or to_timestamp <= from_timestamp):
+ raise ValueError(
+ "LeRobot video feature %s has invalid timestamp range [%s, %s)."
+ % (name, from_timestamp, to_timestamp)
+ )
+ duration = to_timestamp - from_timestamp
+ expected_duration = episode_length / fps
+ if not math.isclose(
+ duration,
+ expected_duration,
+ rel_tol=0.0,
+ abs_tol=_VIDEO_TIMESTAMP_TOLERANCE):
+ raise ValueError(
+ "LeRobot video feature %s has duration %s, but Episode with %d "
+ "frames at FPS %s requires %s."
+ % (name, duration, episode_length, fps, expected_duration)
+ )
+
+
+def _video_fps(info, feature, name):
+ global_fps = _positive_fps(info.get("fps"), "dataset")
+ values = []
+ for key in ("info", "video_info"):
+ details = feature.get(key)
+ if isinstance(details, dict) and "video.fps" in details:
+ values.append(details["video.fps"])
+ if "fps" in feature:
+ values.append(feature["fps"])
+ value = values[0] if values else global_fps
+ fps = _positive_fps(value, "video feature %s" % name)
+ if any(not math.isclose(
+ fps, _positive_fps(other, "video feature %s" % name),
+ rel_tol=1e-6, abs_tol=1e-6) for other in values[1:]) \
+ or not math.isclose(
+ fps, global_fps, rel_tol=1e-6, abs_tol=1e-6):
+ raise ValueError(
+ "LeRobot video feature %s FPS does not match dataset FPS."
+ % name
+ )
+ return fps
+
+
+def _positive_fps(value, owner):
+ try:
+ fps = float(value)
+ except (TypeError, ValueError) as error:
+ raise ValueError(
+ "LeRobot %s is missing a valid FPS." % owner
+ ) from error
+ if not math.isfinite(fps) or fps <= 0:
+ raise ValueError(
+ "LeRobot %s is missing a valid FPS." % owner)
+ return fps
+
+
+def _nonnegative_integer(value, name):
+ value = _python_scalar(value)
+ if isinstance(value, bool) \
+ or not isinstance(value, numbers.Integral) or value < 0:
+ raise ValueError(
+ "LeRobot %s must be a non-negative integer; found %r."
+ % (name, value))
+ return int(value)
+
+
+def _video_source(
+ dataset, info, episode, name, chunk_index, file_index):
+ resolver = getattr(dataset, "video_source", None)
+ if callable(resolver):
+ return resolver(name, episode)
+
+ template = info.get("video_path")
+ if not isinstance(template, str) or not template:
+ raise ValueError("LeRobot v3 metadata is missing info.video_path.")
+ relative = template.format(
+ video_key=name,
+ chunk_index=chunk_index,
+ file_index=file_index,
+ )
+ root = Path(dataset.root).resolve()
+ path = (root / relative).resolve()
+ try:
+ path.relative_to(root)
+ except ValueError as error:
+ raise ValueError(
+ "LeRobot video path must stay within the source directory: %s"
+ % relative
+ ) from error
+ if not path.is_file():
+ raise FileNotFoundError("LeRobot video file does not exist: %s" % path)
+ length = path.stat().st_size
+ if length <= 0:
+ raise ValueError("LeRobot video file is empty: %s" % path)
+ return path.as_uri(), length
+
+
def _safe_array(values, field, name, dtype):
_validate_declared_range(values, field.type, name, dtype)
try:
diff --git a/paimon-python/pypaimon/multimodal/lerobot/metadata.py
b/paimon-python/pypaimon/multimodal/lerobot/metadata.py
index e9daa02066..df75ce9ab0 100644
--- a/paimon-python/pypaimon/multimodal/lerobot/metadata.py
+++ b/paimon-python/pypaimon/multimodal/lerobot/metadata.py
@@ -73,15 +73,33 @@ _EPISODE_CONTROL_COLUMNS = [
class _EpisodeIndex:
- def __init__(self):
+ def __init__(self, video_fields=()):
self._ranges = array("q")
self._task_offsets = array("q", [0])
self._task_indices = array("q")
+ self._video_fields = tuple(video_fields)
+ self._video_indices = {
+ field: (array("q"), array("q"))
+ for field in self._video_fields
+ }
+ self._video_timestamps = {
+ field: (array("d"), array("d"))
+ for field in self._video_fields
+ }
- def append(self, begin, end, task_indices):
+ def append(self, begin, end, task_indices, video_values=None):
self._ranges.extend((begin, end))
self._task_indices.extend(task_indices)
self._task_offsets.append(len(self._task_indices))
+ video_values = video_values or {}
+ for field in self._video_fields:
+ values = video_values[field]
+ chunk_indices, file_indices = self._video_indices[field]
+ from_timestamps, to_timestamps = self._video_timestamps[field]
+ chunk_indices.append(values[0])
+ file_indices.append(values[1])
+ from_timestamps.append(values[2])
+ to_timestamps.append(values[3])
def __len__(self):
return len(self._ranges) // 2
@@ -95,13 +113,22 @@ class _EpisodeIndex:
task_end = self._task_offsets[index + 1]
begin = self._ranges[index * 2]
end = self._ranges[index * 2 + 1]
- return {
+ result = {
"episode_index": index,
"dataset_from_index": begin,
"dataset_to_index": end,
"length": end - begin,
"task_indices": self._task_indices[task_begin:task_end],
}
+ for field in self._video_fields:
+ chunk_indices, file_indices = self._video_indices[field]
+ from_timestamps, to_timestamps = self._video_timestamps[field]
+ prefix = "videos/%s/" % field
+ result[prefix + "chunk_index"] = chunk_indices[index]
+ result[prefix + "file_index"] = file_indices[index]
+ result[prefix + "from_timestamp"] = from_timestamps[index]
+ result[prefix + "to_timestamp"] = to_timestamps[index]
+ return result
def _load_dataset_metadata(dataset, info, source):
@@ -535,14 +562,29 @@ def _source_episodes(dataset, source):
return {"paths": paths, "schema": schema}
-def _validated_episode_tables(metadata):
- episodes = _EpisodeIndex()
+def _validated_episode_tables(metadata, video_fields=()):
+ episodes = _EpisodeIndex(video_fields)
expected_begin = 0
+ video_columns = [
+ "videos/%s/%s" % (field, suffix)
+ for field in video_fields
+ for suffix in (
+ "chunk_index", "file_index", "from_timestamp", "to_timestamp")
+ ]
+ required_columns = _EPISODE_CONTROL_COLUMNS + video_columns
for table in _source_episode_tables(metadata):
- controls = table.select(_EPISODE_CONTROL_COLUMNS)
+ missing = [
+ name for name in required_columns
+ if name not in table.column_names
+ ]
+ if missing:
+ raise ValueError(
+ "LeRobot Episode metadata is missing columns: %s."
+ % ", ".join(missing))
+ controls = table.select(required_columns)
columns = {
name: controls.column(name)
- for name in _EPISODE_CONTROL_COLUMNS
+ for name in required_columns
}
for offset in range(controls.num_rows):
index = _integer(
@@ -581,7 +623,37 @@ def _validated_episode_tables(metadata):
if len(set(task_indices)) != len(task_indices):
raise ValueError(
"LeRobot Episode %d repeats a task." % index)
- episodes.append(begin, end, task_indices)
+ video_values = {}
+ for field in video_fields:
+ prefix = "videos/%s/" % field
+ chunk_index = _integer(
+ columns[prefix + "chunk_index"][offset].as_py(),
+ prefix + "chunk_index",
+ )
+ file_index = _integer(
+ columns[prefix + "file_index"][offset].as_py(),
+ prefix + "file_index",
+ )
+ if chunk_index < 0 or file_index < 0:
+ raise ValueError(
+ "LeRobot Episode video indices must be non-negative.")
+ from_timestamp = columns[
+ prefix + "from_timestamp"][offset].as_py()
+ to_timestamp = columns[
+ prefix + "to_timestamp"][offset].as_py()
+ if isinstance(from_timestamp, bool) \
+ or not isinstance(from_timestamp, numbers.Real) \
+ or isinstance(to_timestamp, bool) \
+ or not isinstance(to_timestamp, numbers.Real):
+ raise ValueError(
+ "LeRobot Episode video timestamps must be numeric.")
+ video_values[field] = (
+ chunk_index,
+ file_index,
+ float(from_timestamp),
+ float(to_timestamp),
+ )
+ episodes.append(begin, end, task_indices, video_values)
expected_begin = end
yield table
del table
diff --git a/paimon-python/pypaimon/multimodal/lerobot/schema.py
b/paimon-python/pypaimon/multimodal/lerobot/schema.py
index de16c95a64..188127a14d 100644
--- a/paimon-python/pypaimon/multimodal/lerobot/schema.py
+++ b/paimon-python/pypaimon/multimodal/lerobot/schema.py
@@ -83,6 +83,14 @@ def _validate_v3_required_features(info):
"shape=[1]." % (name, dtype))
+def _video_feature_names(info):
+ return [
+ name
+ for name, feature in info.get("features", {}).items()
+ if isinstance(feature, dict) and feature.get("dtype") == "video"
+ ]
+
+
def _validate_lerobot_schema(source_schema, target_schema, source):
"""Require an existing table to preserve the LeRobot feature contract."""
for source_field in source_schema:
@@ -122,11 +130,7 @@ def _feature_field(name, feature):
"LeRobot feature %s metadata must be an object." % name)
dtype = str(feature.get("dtype", ""))
shape = _feature_shape(feature, name)
- if dtype == "video":
- raise ValueError(
- "LeRobot video feature %s is not supported yet; use an "
- "image-based dataset." % name)
- if dtype == "image":
+ if dtype in ("image", "video"):
arrow_type = pa.large_binary()
else:
scalar_type = _SCALAR_DTYPES.get(dtype)
diff --git a/paimon-python/pypaimon/multimodal/lerobot/source.py
b/paimon-python/pypaimon/multimodal/lerobot/source.py
index 04547db804..995c14a438 100644
--- a/paimon-python/pypaimon/multimodal/lerobot/source.py
+++ b/paimon-python/pypaimon/multimodal/lerobot/source.py
@@ -32,7 +32,9 @@ import pyarrow.fs as pafs
import pyarrow.parquet as pq
from pypaimon.common.options import Options
+from pypaimon.common.uri_reader import FileUriReader
from pypaimon.filesystem.pyarrow_file_io import LegacyOssDirectoryListingError
+from pypaimon.multimodal.lerobot.schema import _video_feature_names
from pypaimon.multimodal.source_utils import (
_SourceFileIO,
_normalize_source_path,
@@ -61,17 +63,37 @@ class _RemoteLeRobotMeta:
class _RemoteEpisodeIndex:
- def __init__(self):
+ def __init__(self, video_fields=()):
self.starts = array("q")
self._ends = array("q")
self._chunk_indices = array("q")
self._file_indices = array("q")
+ self._video_fields = tuple(video_fields)
+ self._video_indices = {
+ field: (array("q"), array("q"))
+ for field in self._video_fields
+ }
+ self._video_timestamps = {
+ field: (array("d"), array("d"))
+ for field in self._video_fields
+ }
- def append(self, begin, end, chunk_index, file_index):
+ def append(
+ self, begin, end, chunk_index, file_index,
+ video_values=None):
self.starts.append(begin)
self._ends.append(end)
self._chunk_indices.append(chunk_index)
self._file_indices.append(file_index)
+ video_values = video_values or {}
+ for field in self._video_fields:
+ values = video_values[field]
+ chunk_indices, file_indices = self._video_indices[field]
+ from_timestamps, to_timestamps = self._video_timestamps[field]
+ chunk_indices.append(values[0])
+ file_indices.append(values[1])
+ from_timestamps.append(values[2])
+ to_timestamps.append(values[3])
def __len__(self):
return len(self.starts)
@@ -81,13 +103,23 @@ class _RemoteEpisodeIndex:
index += len(self)
if index < 0 or index >= len(self):
raise IndexError(index)
- return {
+ result = {
"episode_index": index,
"dataset_from_index": self.starts[index],
"dataset_to_index": self._ends[index],
+ "length": self._ends[index] - self.starts[index],
"data/chunk_index": self._chunk_indices[index],
"data/file_index": self._file_indices[index],
}
+ for field in self._video_fields:
+ chunk_indices, file_indices = self._video_indices[field]
+ from_timestamps, to_timestamps = self._video_timestamps[field]
+ prefix = "videos/%s/" % field
+ result[prefix + "chunk_index"] = chunk_indices[index]
+ result[prefix + "file_index"] = file_indices[index]
+ result[prefix + "from_timestamp"] = from_timestamps[index]
+ result[prefix + "to_timestamp"] = to_timestamps[index]
+ return result
@contextmanager
@@ -197,17 +229,17 @@ def _load_hub_info(source):
% (source.path, error)) from error
-def _open_dataset(LeRobotDataset, source):
+def _open_dataset(LeRobotDataset, source, download_videos=False):
try:
if source.root is not None:
return LeRobotDataset(
repo_id=source.repo_id,
root=source.root,
- download_videos=False,
+ download_videos=download_videos,
)
return LeRobotDataset(
repo_id=source.repo_id,
- download_videos=False,
+ download_videos=download_videos,
)
except Exception as error:
raise ValueError(
@@ -215,10 +247,11 @@ def _open_dataset(LeRobotDataset, source):
% (source.path, error)) from error
-def _open_resolved_dataset(LeRobotDataset, source, info):
+def _open_resolved_dataset(
+ LeRobotDataset, source, info, download_videos=False):
if source.file_io is not None:
return _RemoteLeRobotDataset(source, info)
- return _open_dataset(LeRobotDataset, source)
+ return _open_dataset(LeRobotDataset, source, download_videos)
class _RemoteLeRobotDataset:
@@ -291,35 +324,94 @@ class _RemoteLeRobotDataset:
return _read_remote_bytes(self._file_io, source_path)
return _encode_media_frame(value)
+ def video_source(self, video_key, episode):
+ relative_path = self.meta.info["video_path"].format(
+ video_key=video_key,
+ chunk_index=int(episode[
+ "videos/%s/chunk_index" % video_key]),
+ file_index=int(episode[
+ "videos/%s/file_index" % video_key]),
+ )
+ relative_path = _relative_dataset_path(
+ relative_path, "info.video_path")
+ path = _remote_source_path(
+ self.source.path,
+ relative_path,
+ "info.video_path",
+ self._file_io,
+ )
+ status = self._file_io.get_file_status(path)
+ if status.type != pafs.FileType.File \
+ or status.size is None or status.size <= 0:
+ raise ValueError("LeRobot video file is empty: %s" % path)
+ return path, int(status.size)
+
+ @property
+ def video_uri_reader_factory(self):
+ return _SourceUriReaderFactory(self._file_io)
+
def _load_episodes(self, info):
episode_count = int(info.get("total_episodes", 0))
if episode_count == 0:
return []
directory = _remote_path(self.source.path, "meta/episodes")
paths = _remote_parquet_files(self._file_io, directory)
- episodes = _RemoteEpisodeIndex()
+ video_fields = _video_feature_names(info)
+ selected_columns = list(self._EPISODE_COLUMNS)
+ for field in video_fields:
+ selected_columns.extend([
+ "videos/%s/%s" % (field, suffix)
+ for suffix in (
+ "chunk_index",
+ "file_index",
+ "from_timestamp",
+ "to_timestamp",
+ )
+ ])
+ episodes = _RemoteEpisodeIndex(video_fields)
for path in paths:
table = _read_remote_parquet(
self._file_io,
path,
- columns=self._EPISODE_COLUMNS,
+ columns=selected_columns,
)
- columns = {
+ control_columns = {
name: table.column(name)
for name in self._EPISODE_COLUMNS
}
for offset in range(table.num_rows):
episode_index = int(
- columns["episode_index"][offset].as_py())
+ control_columns["episode_index"][offset].as_py())
if episode_index != len(episodes):
raise ValueError(
"LeRobot Episode metadata must be ordered by "
"episode_index.")
episodes.append(
- int(columns["dataset_from_index"][offset].as_py()),
- int(columns["dataset_to_index"][offset].as_py()),
- int(columns["data/chunk_index"][offset].as_py()),
- int(columns["data/file_index"][offset].as_py()),
+ int(control_columns[
+ "dataset_from_index"][offset].as_py()),
+ int(control_columns[
+ "dataset_to_index"][offset].as_py()),
+ int(control_columns[
+ "data/chunk_index"][offset].as_py()),
+ int(control_columns[
+ "data/file_index"][offset].as_py()),
+ {
+ field: (
+ int(table.column(
+ "videos/%s/chunk_index" % field
+ )[offset].as_py()),
+ int(table.column(
+ "videos/%s/file_index" % field
+ )[offset].as_py()),
+ float(table.column(
+ "videos/%s/from_timestamp" % field
+ )[offset].as_py()),
+ float(table.column(
+ "videos/%s/to_timestamp" % field
+ )[offset].as_py()),
+ )
+ for field in video_fields
+ },
)
if len(episodes) != episode_count:
raise ValueError(
@@ -389,6 +481,15 @@ def _remote_path(root, relative_path):
return "%s/%s" % (root.rstrip("/"), relative_path.lstrip("/"))
+class _SourceUriReaderFactory:
+
+ def __init__(self, file_io):
+ self._reader = FileUriReader(file_io)
+
+ def create(self, unused_uri):
+ return self._reader
+
+
def _pandas_index_column(schema, component):
encoded = (schema.metadata or {}).get(b"pandas")
try:
diff --git a/paimon-python/pypaimon/multimodal/lerobot/writer.py
b/paimon-python/pypaimon/multimodal/lerobot/writer.py
index f73fadf110..8fc152e6f1 100644
--- a/paimon-python/pypaimon/multimodal/lerobot/writer.py
+++ b/paimon-python/pypaimon/multimodal/lerobot/writer.py
@@ -52,6 +52,7 @@ from pypaimon.multimodal.lerobot.schema import (
_schema_from_info,
_validate_lerobot_schema,
_validate_v3_required_features,
+ _video_feature_names,
)
from pypaimon.multimodal.table import _target_schema
@@ -288,6 +289,11 @@ class PaimonLeRobotWriter:
raise ValueError("features must be a non-empty mapping.")
if "task" in features:
raise ValueError("task is managed by PaimonLeRobotWriter.")
+ video_fields = _video_feature_names({"features": features})
+ if video_fields:
+ raise ValueError(
+ "PaimonLeRobotWriter does not support video features: %s."
+ % ", ".join(video_fields))
requested_subtasks = _validate_subtasks(
subtasks, "subtask_index" in features)
_lerobot_stats_functions()
diff --git a/paimon-python/pypaimon/tests/data_evolution_row_rolling_test.py
b/paimon-python/pypaimon/tests/data_evolution_row_rolling_test.py
index eb5b61786c..bc797398e3 100644
--- a/paimon-python/pypaimon/tests/data_evolution_row_rolling_test.py
+++ b/paimon-python/pypaimon/tests/data_evolution_row_rolling_test.py
@@ -134,6 +134,20 @@ class DataEvolutionRowRollingTest(unittest.TestCase):
tw.close()
return files
+ def _write_episode_files(self, table, data, episode_lengths):
+ wb = table.new_batch_write_builder()
+ writer = wb.new_write()
+ offset = 0
+ for length in episode_lengths:
+ writer.roll_before_group_if_needed(length)
+ writer.write_arrow(data.slice(offset, length))
+ offset += length
+ messages = writer.prepare_commit()
+ files = [file for message in messages for file in message.new_files]
+ wb.new_commit().commit(messages)
+ writer.close()
+ return files
+
def _read_ids(self, table):
rb = table.new_read_builder().with_projection(['id'])
return sorted(
@@ -419,6 +433,179 @@ class DataEvolutionRowRollingTest(unittest.TestCase):
self.assertEqual([4, 4], sorted(file.row_count for file in
video_files))
self.assertEqual(list(range(4)), self._read_ids(table))
+ def test_video_episodes_roll_before_shared_payload_group(self):
+ path = os.path.join(self.tempdir, 'shared-episodes.mp4')
+ payload = b'shared-video'
+ with open(path, 'wb') as output:
+ output.write(payload)
+ descriptor = BlobDescriptor(path, 0, len(payload))
+ table = self._create_with_schema(
+ self.blob_schema,
+ {
+ **self.de_options,
+ 'target-file-row-num': '3',
+ 'video-frame-field': 'payload',
+ 'blob-as-descriptor': 'true',
+ },
+ )
+ rows = pa.Table.from_pydict(
+ {
+ 'id': list(range(6)),
+ 'payload': [
+ VideoFrameDescriptor(
+ descriptor.uri,
+ descriptor.offset,
+ descriptor.length,
+ frame,
+ ).serialize()
+ for frame in range(6)
+ ],
+ },
+ schema=self.blob_schema,
+ )
+
+ files = self._write_episode_files(table, rows, [2, 4])
+
+ normal_rows = sorted(
+ file.row_count for file in files
+ if not file.file_name.endswith('.video')
+ )
+ video_rows = sorted(
+ file.row_count for file in files
+ if file.file_name.endswith('.video')
+ )
+ self.assertEqual([2, 4], normal_rows)
+ self.assertEqual([2, 4], video_rows)
+ self.assertEqual(list(range(6)), self._read_ids(table))
+
+ def test_video_columns_roll_together_at_episode_boundaries(self):
+ paths = [
+ os.path.join(self.tempdir, name)
+ for name in ('small-camera.mp4', 'large-camera.mp4')
+ ]
+ payloads = [b'a', b'b' * 100]
+ for path, payload in zip(paths, payloads):
+ with open(path, 'wb') as output:
+ output.write(payload)
+ camera_a, camera_b = [
+ BlobDescriptor(path, 0, len(payload))
+ for path, payload in zip(paths, payloads)
+ ]
+ table = self._create_with_schema(
+ self.multi_video_schema,
+ {
+ **self.de_options,
+ 'video-frame-field': 'camera_a,camera_b',
+ 'blob-as-descriptor': 'true',
+ 'blob.target-file-size': '50 b',
+ },
+ )
+ rows = pa.Table.from_pydict(
+ {
+ 'id': list(range(4)),
+ 'camera_a': [
+ VideoFrameDescriptor(
+ camera_a.uri, 0, camera_a.length, frame
+ ).serialize()
+ for frame in range(4)
+ ],
+ 'camera_b': [
+ VideoFrameDescriptor(
+ camera_b.uri, 0, camera_b.length, frame
+ ).serialize()
+ for frame in range(4)
+ ],
+ },
+ schema=self.multi_video_schema,
+ )
+
+ files = self._write_episode_files(table, rows, [2, 2])
+
+ files_by_column = {}
+ for file in files:
+ if file.file_name.endswith('.video'):
+ files_by_column.setdefault(file.write_cols[0], []).append(
+ file.row_count)
+ self.assertEqual([2, 2], files_by_column['camera_a'])
+ self.assertEqual([2, 2], files_by_column['camera_b'])
+ self.assertEqual([2, 2], sorted(
+ file.row_count for file in files
+ if not file.file_name.endswith('.video')
+ ))
+ result = table.new_read_builder().new_read().to_arrow(
+ table.new_read_builder().new_scan().plan().splits()
+ ).sort_by('id')
+ for column in ('camera_a', 'camera_b'):
+ self.assertEqual(
+ list(range(4)),
+ [
+ VideoFrameDescriptor.deserialize(value.as_py()).frame_index
+ for value in result[column]
+ ],
+ )
+ self.assertEqual(list(range(4)), self._read_ids(table))
+
+ def test_vector_rolling_waits_for_video_episode_boundary(self):
+ path = os.path.join(self.tempdir, 'vector-episodes.mp4')
+ payload = b'shared-video'
+ with open(path, 'wb') as output:
+ output.write(payload)
+ descriptor = BlobDescriptor(path, 0, len(payload))
+ rows = pa.Table.from_pydict(
+ {
+ 'id': list(range(4)),
+ 'payload': [
+ VideoFrameDescriptor(
+ descriptor.uri,
+ descriptor.offset,
+ descriptor.length,
+ frame,
+ ).serialize()
+ for frame in range(4)
+ ],
+ 'embedding': [
+ [float(frame), float(frame + 1), float(frame + 2)]
+ for frame in range(4)
+ ],
+ },
+ schema=self.blob_vector_schema,
+ )
+
+ for episode_lengths, expected_rows in (
+ ([2, 2], [2, 2]), ([4], [4])):
+ with self.subTest(episode_lengths=episode_lengths):
+ table = self._create_with_schema(
+ self.blob_vector_schema,
+ {
+ **self.de_options,
+ 'target-file-row-num': '100',
+ 'video-frame-field': 'payload',
+ 'blob-as-descriptor': 'true',
+ 'vector.file.format': 'parquet',
+ 'vector.target-file-size': '1 b',
+ },
+ )
+ files = self._write_episode_files(
+ table, rows, episode_lengths)
+
+ normal_rows = sorted(
+ file.row_count for file in files
+ if not file.file_name.endswith('.video')
+ and '.vector.' not in file.file_name
+ )
+ video_rows = sorted(
+ file.row_count for file in files
+ if file.file_name.endswith('.video')
+ )
+ vector_rows = sorted(
+ file.row_count for file in files
+ if '.vector.' in file.file_name
+ )
+ self.assertEqual(expected_rows, normal_rows)
+ self.assertEqual(expected_rows, video_rows)
+ self.assertEqual(expected_rows, vector_rows)
+ self.assertEqual(list(range(4)), self._read_ids(table))
+
def test_blob_consumer_descriptors_survive_abort_after_rolling(self):
table = self._create_with_schema(
self.blob_schema,
diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py
b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py
index 5be17dc7b5..82e2b7631c 100644
--- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py
+++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py
@@ -16,6 +16,7 @@
import builtins
from array import array
+from fractions import Fraction
import io
import json
import pickle
@@ -26,6 +27,7 @@ import threading
import unittest
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
+from types import SimpleNamespace
from unittest.mock import Mock, patch
import numpy as np
@@ -48,6 +50,7 @@ from pypaimon.multimodal.lerobot.dataset import (
_selected_episodes,
_torch_row,
)
+from pypaimon.multimodal.lerobot.api import _create_target_table
from pypaimon.multimodal.lerobot.metadata import (
_append_arrow_tables,
_companion_identifier,
@@ -61,6 +64,8 @@ from pypaimon.multimodal.lerobot.metadata import (
from pypaimon.multimodal.lerobot.loader import (
_image_bytes,
_read_batch,
+ _video_frame_ordinal,
+ _video_sample_timestamps,
_validate_frame_controls,
)
from pypaimon.multimodal.lerobot.schema import (
@@ -83,6 +88,11 @@ try:
except ImportError:
LeRobotDataset = None
+try:
+ import av
+except ImportError:
+ av = None
+
def _replaced_contract(field, old, new):
description = field.metadata[b"description"].decode("utf-8")
@@ -882,8 +892,10 @@ class LeRobotValidationTest(unittest.TestCase):
"shape": [8, 10, 3],
}
}
- with self.assertRaisesRegex(ValueError, "video feature camera.*not
supported"):
- _schema_from_info(info)
+ self.assertEqual(
+ pa.large_binary(),
+ _schema_from_info(info).field("camera").type,
+ )
def test_existing_schema_preserves_lerobot_feature_contract(self):
source = _schema_from_info({
@@ -993,6 +1005,7 @@ class LeRobotValidationTest(unittest.TestCase):
"episode_index": [0],
"dataset_from_index": [0],
"dataset_to_index": [1],
+ "length": [1],
"data/chunk_index": [0],
"data/file_index": [0],
})
@@ -1026,6 +1039,16 @@ class LeRobotValidationTest(unittest.TestCase):
"fps": 30,
"features": {
"index": {"dtype": "int64", "shape": [1]},
+ "timestamp": {
+ "dtype": "float32",
+ "shape": [1],
+ "fps": 10.0,
+ },
+ "camera": {
+ "dtype": "video",
+ "shape": [8, 10, 3],
+ "video_info": {"video.fps": 10.0},
+ },
},
}))
connection = pmm.connect(options={
@@ -1042,6 +1065,72 @@ class LeRobotValidationTest(unittest.TestCase):
finally:
shutil.rmtree(temp_dir, ignore_errors=True)
+ def test_video_options_must_match_metadata(self):
+ temp_dir = Path(tempfile.mkdtemp(prefix="pypaimon_lerobot_options_"))
+ try:
+ info = {
+ "features": {
+ "index": {"dtype": "int64", "shape": [1]},
+ "camera_a": {"dtype": "video", "shape": [8, 10, 3]},
+ "camera_b": {"dtype": "video", "shape": [8, 10, 3]},
+ },
+ }
+ connection = pmm.connect(options={
+ "warehouse": str(temp_dir / "warehouse"),
+ })
+ schema = _schema_from_info(info)
+ metadata = {"stats_table": None, "subtasks_table": None}
+
+ with self.assertRaisesRegex(ValueError, "do not match"):
+ _create_target_table(
+ connection,
+ "conflict",
+ schema,
+ options={"video-frame-field": "camera_a"},
+ metadata=metadata,
+ video_fields=("camera_a", "camera_b"),
+ )
+
+ table = _create_target_table(
+ connection,
+ "reordered",
+ schema,
+ options={"video-frame-field": "camera_b,camera_a"},
+ metadata=metadata,
+ video_fields=("camera_a", "camera_b"),
+ )
+ self.assertEqual(
+ {"camera_a", "camera_b"},
+ table.raw_table.options.video_frame_fields(),
+ )
+ finally:
+ shutil.rmtree(temp_dir, ignore_errors=True)
+
+ def test_video_import_requires_single_writer_layout(self):
+ temp_dir = Path(tempfile.mkdtemp(prefix="pypaimon_lerobot_layout_"))
+ try:
+ info = {
+ "features": {
+ "episode_index": {"dtype": "int64", "shape": [1]},
+ "camera": {"dtype": "video", "shape": [8, 10, 3]},
+ },
+ }
+ schema = _schema_from_info(info)
+ connection = pmm.connect(options={
+ "warehouse": str(temp_dir / "warehouse"),
+ })
+ with self.assertRaisesRegex(ValueError, "bucket-unaware"):
+ _create_target_table(
+ connection,
+ "bucketed",
+ schema,
+ options={"bucket": "1"},
+ metadata={"stats_table": None, "subtasks_table": None},
+ video_fields=("camera",),
+ )
+ finally:
+ shutil.rmtree(temp_dir, ignore_errors=True)
+
def test_empty_fast_path_validates_required_counts(self):
temp_dir = Path(tempfile.mkdtemp(prefix="pypaimon_lerobot_counts_"))
try:
@@ -1339,26 +1428,644 @@ class LeRobotValidationTest(unittest.TestCase):
finally:
shutil.rmtree(temp_dir, ignore_errors=True)
- def test_local_video_is_rejected_before_opening(self):
+ def test_episode_aware_multi_video_import(self):
+ import pandas as pd
+
temp_dir = Path(tempfile.mkdtemp(prefix="pypaimon_lerobot_video_"))
try:
info_dir = temp_dir / "meta"
info_dir.mkdir()
- (info_dir / "info.json").write_text(json.dumps({
+ info = {
"codebase_version": "v3.0",
+ "fps": 10,
+ "total_frames": 5,
+ "total_episodes": 2,
+ "total_tasks": 1,
+ "data_path": (
+ "data/chunk-{chunk_index:03d}/"
+ "file-{file_index:03d}.parquet"
+ ),
+ "video_path": (
+ "videos/{video_key}/chunk-{chunk_index:03d}/"
+ "file-{file_index:03d}.mp4"
+ ),
"features": {
- "camera": {"dtype": "video", "shape": [8, 10, 3]},
+ "index": {"dtype": "int64", "shape": [1]},
+ "episode_index": {"dtype": "int64", "shape": [1]},
+ "frame_index": {"dtype": "int64", "shape": [1]},
+ "timestamp": {
+ "dtype": "float32",
+ "shape": [1],
+ "fps": 10.0,
+ },
+ "task_index": {"dtype": "int64", "shape": [1]},
+ "observation.state": {
+ "dtype": "float32",
+ "shape": [3],
+ },
+ "camera_a": {
+ "dtype": "video",
+ "shape": [8, 10, 3],
+ "video_info": {"video.fps": 10.0},
+ },
+ "camera_b": {
+ "dtype": "video",
+ "shape": [8, 10, 3],
+ "video_info": {"video.fps": 10.0},
+ },
},
- }))
+ }
+ (info_dir / "info.json").write_text(json.dumps(info))
+ payloads = {
+ "camera_a/chunk-000/file-000.mp4": b"camera-a",
+ "camera_b/chunk-000/file-000.mp4": b"camera-b-0",
+ "camera_b/chunk-000/file-001.mp4": b"camera-b-1",
+ }
+ for relative, payload in payloads.items():
+ path = temp_dir / "videos" / relative
+ path.parent.mkdir(parents=True, exist_ok=True)
+ path.write_bytes(payload)
+
+ episodes = [
+ {
+ "episode_index": 0,
+ "dataset_from_index": 0,
+ "dataset_to_index": 2,
+ "length": 2,
+ "data/chunk_index": 0,
+ "data/file_index": 0,
+ "tasks": ["pick"],
+ "videos/camera_a/chunk_index": 0,
+ "videos/camera_a/file_index": 0,
+ "videos/camera_a/from_timestamp": 0.5,
+ "videos/camera_a/to_timestamp": 0.7,
+ "videos/camera_b/chunk_index": 0,
+ "videos/camera_b/file_index": 0,
+ "videos/camera_b/from_timestamp": 0.0,
+ "videos/camera_b/to_timestamp": 0.2,
+ },
+ {
+ "episode_index": 1,
+ "dataset_from_index": 2,
+ "dataset_to_index": 5,
+ "length": 3,
+ "data/chunk_index": 0,
+ "data/file_index": 0,
+ "tasks": ["pick"],
+ "videos/camera_a/chunk_index": 0,
+ "videos/camera_a/file_index": 0,
+ "videos/camera_a/from_timestamp": 0.1,
+ "videos/camera_a/to_timestamp": 0.4,
+ "videos/camera_b/chunk_index": 0,
+ "videos/camera_b/file_index": 1,
+ "videos/camera_b/from_timestamp": 0.0,
+ "videos/camera_b/to_timestamp": 0.3,
+ },
+ ]
+ # The published Episode Parquet, not the dataset object's stale
+ # metadata cache, must determine the imported video descriptors.
+ cached_episodes = [dict(episode) for episode in episodes]
+ cached_episodes[1].update({
+ "videos/camera_b/file_index": 0,
+ "videos/camera_b/from_timestamp": 0.2,
+ "videos/camera_b/to_timestamp": 0.5,
+ })
+
+ class Dataset:
+
+ root = temp_dir
+ meta = SimpleNamespace(
+ info=info, episodes=cached_episodes, tasks=["pick"])
+ rows = pa.table({
+ "index": pa.array(range(5), type=pa.int64()),
+ "episode_index": pa.array(
+ [0, 0, 1, 1, 1], type=pa.int64()),
+ "frame_index": pa.array(
+ [0, 1, 0, 1, 2], type=pa.int64()),
+ "timestamp": pa.array(
+ [0.0, 0.1, 0.0, 0.1, 0.2],
+ type=pa.float32(),
+ ),
+ "task_index": pa.array([0] * 5, type=pa.int64()),
+ "observation.state": pa.array(
+ [[float(index), 0.0, 1.0] for index in range(5)],
+ type=pa.list_(pa.float32(), 3),
+ ),
+ })
+
+ def __len__(self):
+ return 5
+
+ def read_batch(self, begin, end):
+ return self.rows.slice(begin, end - begin)
+
connection = pmm.connect(options={
"warehouse": str(temp_dir / "warehouse"),
})
- with self.assertRaisesRegex(
- ValueError, "video feature camera.*not supported"):
- connection.load_from_lerobot("frames", temp_dir)
+ episodes_path = (
+ temp_dir / "meta/episodes/chunk-000/file-000.parquet")
+ episodes_path.parent.mkdir(parents=True)
+ pq.write_table(pa.Table.from_pylist(episodes), episodes_path)
+ pq.write_table(pa.Table.from_pandas(pd.DataFrame(
+ {"task_index": [0]},
+ index=pd.Index(["pick"], name="task"),
+ )), temp_dir / "meta/tasks.parquet")
+
+ def sample_timestamps(unused_dataset, uri):
+ return (
+ [0.0, 0.1, 0.2, 0.3, 0.5, 0.6]
+ if "camera_a" in uri else
+ [0.0, 0.1, 0.2, 0.3, 0.4]
+ )
+
+ with patch(
+ "pypaimon.multimodal.lerobot.api."
+ "_import_lerobot_dataset",
+ return_value=object,
+ ), patch(
+ "pypaimon.multimodal.lerobot.api."
+ "_open_resolved_dataset",
+ return_value=Dataset(),
+ ), patch(
+ "pypaimon.multimodal.lerobot.loader."
+ "_video_sample_timestamps",
+ side_effect=sample_timestamps,
+ ):
+ result = connection.load_from_lerobot(
+ "frames", temp_dir, batch_size=1)
+
+ self.assertIsNone(result)
+ table = connection.get_table("frames")
+ self.assertEqual(
+ {"camera_a", "camera_b"},
+ table.raw_table.options.video_frame_fields(),
+ )
+ field_types = {
+ field.name: str(field.type)
+ for field in table.raw_table.fields
+ }
+ self.assertEqual(
+ "VECTOR<FLOAT, 3> NOT NULL",
+ field_types["observation.state"],
+ )
+ rows = table.scan().select([
+ "index", "camera_a", "camera_b"
+ ]).to_arrow().sort_by("index").to_pylist()
+ camera_a = [
+ pmm.VideoFrameDescriptor.deserialize(row["camera_a"])
+ for row in rows
+ ]
+ camera_b = [
+ pmm.VideoFrameDescriptor.deserialize(row["camera_b"])
+ for row in rows
+ ]
+ self.assertEqual(
+ [4, 5, 1, 2, 3],
+ [descriptor.frame_index for descriptor in camera_a],
+ )
+ self.assertEqual(
+ [0, 1, 0, 1, 2],
+ [descriptor.frame_index for descriptor in camera_b],
+ )
+ _, bodies = table.scan().select([
+ "index", "camera_a", "camera_b"
+ ]).read_blobs()
+ self.assertEqual(
+ [payloads["camera_a/chunk-000/file-000.mp4"]] * 5,
+ bodies["camera_a"],
+ )
+ self.assertEqual(
+ [payloads["camera_b/chunk-000/file-000.mp4"]] * 2
+ + [payloads["camera_b/chunk-000/file-001.mp4"]] * 3,
+ bodies["camera_b"],
+ )
+
+ data_path = temp_dir / "data/chunk-000/file-000.parquet"
+ data_path.parent.mkdir(parents=True)
+ pq.write_table(Dataset.rows, data_path)
+ remote = "oss://source-bucket/robot-videos"
+ source_file_io = _RemoteLeRobotFileIO(temp_dir, remote)
+ with patch(
+ "pypaimon.multimodal.lerobot.source._SourceFileIO",
+ return_value=source_file_io,
+ ), patch(
+ "pypaimon.multimodal.lerobot.api."
+ "_import_lerobot_dataset",
+ return_value=object,
+ ), patch(
+ "pypaimon.multimodal.lerobot.loader."
+ "_video_sample_timestamps",
+ side_effect=sample_timestamps,
+ ):
+ remote_result = connection.load_from_lerobot(
+ "remote_frames", remote, batch_size=1)
+
+ self.assertIsNone(remote_result)
+ opened_videos = [
+ path for path in source_file_io.opened_paths
+ if path.endswith(".mp4")
+ ]
+ self.assertEqual(3, len(opened_videos))
+ self.assertEqual(1, source_file_io.close_count)
+ _, remote_bodies = connection.get_table(
+ "remote_frames").scan().select([
+ "index", "camera_a", "camera_b"
+ ]).read_blobs()
+ self.assertEqual(bodies, remote_bodies)
+
+ # Both cameras now share one physical MP4 across Episodes. The
+ # logical Episode boundary must still control normal-file rolling.
+ episodes[1].update({
+ "videos/camera_b/file_index": 0,
+ "videos/camera_b/from_timestamp": 0.2,
+ "videos/camera_b/to_timestamp": 0.5,
+ })
+ pq.write_table(pa.Table.from_pylist(episodes), episodes_path)
+ with patch(
+ "pypaimon.multimodal.lerobot.api."
+ "_import_lerobot_dataset",
+ return_value=object,
+ ), patch(
+ "pypaimon.multimodal.lerobot.api."
+ "_open_resolved_dataset",
+ return_value=Dataset(),
+ ), patch(
+ "pypaimon.multimodal.lerobot.loader."
+ "_video_sample_timestamps",
+ side_effect=sample_timestamps,
+ ):
+ connection.load_from_lerobot(
+ "shared_video_frames",
+ temp_dir,
+ batch_size=1,
+ options={"target-file-row-num": "1"},
+ )
+ raw_table = connection.get_table(
+ "shared_video_frames").raw_table
+ files = {
+ file.file_name: file
+ for split in
raw_table.new_read_builder().new_scan().plan().splits()
+ for file in split.files
+ }.values()
+ self.assertEqual(
+ [2, 3],
+ sorted(
+ file.row_count for file in files
+ if not file.file_name.endswith(".video")
+ and ".vector." not in file.file_name
+ ),
+ )
+ self.assertEqual(
+ [2, 3],
+ sorted(
+ file.row_count for file in files
+ if ".vector." in file.file_name
+ ),
+ )
+ finally:
+ shutil.rmtree(temp_dir, ignore_errors=True)
+
+ @unittest.skipUnless(av is not None, "PyAV is required for MP4 decoding")
+ def test_imported_video_payload_can_be_decoded(self):
+ import pandas as pd
+
+ temp_dir = Path(tempfile.mkdtemp(prefix="pypaimon_lerobot_mp4_"))
+ try:
+ info = {
+ "codebase_version": "v3.0",
+ "fps": 10,
+ "total_frames": 5,
+ "total_episodes": 2,
+ "total_tasks": 1,
+ "data_path": (
+ "data/chunk-{chunk_index:03d}/"
+ "file-{file_index:03d}.parquet"
+ ),
+ "video_path": (
+ "videos/{video_key}/chunk-{chunk_index:03d}/"
+ "file-{file_index:03d}.mp4"
+ ),
+ "features": {
+ "index": {"dtype": "int64", "shape": [1]},
+ "episode_index": {"dtype": "int64", "shape": [1]},
+ "frame_index": {"dtype": "int64", "shape": [1]},
+ "timestamp": {
+ "dtype": "float32",
+ "shape": [1],
+ "fps": 10.0,
+ },
+ "task_index": {"dtype": "int64", "shape": [1]},
+ "camera": {
+ "dtype": "video",
+ "shape": [16, 16, 3],
+ "video_info": {"video.fps": 10.0},
+ },
+ },
+ }
+ episodes = [
+ {
+ "episode_index": 0,
+ "dataset_from_index": 0,
+ "dataset_to_index": 2,
+ "length": 2,
+ "data/chunk_index": 0,
+ "data/file_index": 0,
+ "tasks": ["pick"],
+ "videos/camera/chunk_index": 0,
+ "videos/camera/file_index": 0,
+ "videos/camera/from_timestamp": 0.5,
+ "videos/camera/to_timestamp": 0.7,
+ },
+ {
+ "episode_index": 1,
+ "dataset_from_index": 2,
+ "dataset_to_index": 5,
+ "length": 3,
+ "data/chunk_index": 0,
+ "data/file_index": 0,
+ "tasks": ["pick"],
+ "videos/camera/chunk_index": 0,
+ "videos/camera/file_index": 0,
+ "videos/camera/from_timestamp": 0.1,
+ "videos/camera/to_timestamp": 0.4,
+ },
+ ]
+ physical_frame_values = [24, 56, 88, 120, 168, 216]
+ expected_frame_values = [168, 216, 56, 88, 120]
+
+ info_dir = temp_dir / "meta"
+ info_dir.mkdir()
+ (info_dir / "info.json").write_text(json.dumps(info))
+ episodes_path = (
+ info_dir / "episodes/chunk-000/file-000.parquet")
+ episodes_path.parent.mkdir(parents=True)
+ pq.write_table(pa.Table.from_pylist(episodes), episodes_path)
+ pq.write_table(pa.Table.from_pandas(pd.DataFrame(
+ {"task_index": [0]},
+ index=pd.Index(["pick"], name="task"),
+ )), info_dir / "tasks.parquet")
+
+ video_path = (
+ temp_dir / "videos/camera/chunk-000/file-000.mp4")
+ video_path.parent.mkdir(parents=True)
+ with av.open(str(video_path), mode="w") as container:
+ stream = container.add_stream("mpeg4", rate=10)
+ stream.width = 16
+ stream.height = 16
+ stream.pix_fmt = "yuv420p"
+ stream.time_base = Fraction(1, 10)
+ for pts, value in zip(
+ [0, 1, 2, 3, 5, 6], physical_frame_values):
+ image = np.full((16, 16, 3), value, dtype=np.uint8)
+ frame = av.VideoFrame.from_ndarray(image, format="rgb24")
+ frame.pts = pts
+ frame.time_base = Fraction(1, 10)
+ for packet in stream.encode(frame):
+ container.mux(packet)
+ for packet in stream.encode():
+ container.mux(packet)
+
+ class Dataset:
+
+ root = temp_dir
+ meta = SimpleNamespace(
+ info=info, episodes=episodes, tasks=["pick"])
+ rows = pa.table({
+ "index": pa.array(range(5), type=pa.int64()),
+ "episode_index": pa.array(
+ [0, 0, 1, 1, 1], type=pa.int64()),
+ "frame_index": pa.array(
+ [0, 1, 0, 1, 2], type=pa.int64()),
+ "timestamp": pa.array(
+ [0.0, 0.1, 0.0, 0.1, 0.2],
+ type=pa.float32(),
+ ),
+ "task_index": pa.array([0] * 5, type=pa.int64()),
+ })
+
+ def __len__(self):
+ return 5
+
+ def read_batch(self, begin, end):
+ return self.rows.slice(begin, end - begin)
+
+ connection = pmm.connect(options={
+ "warehouse": str(temp_dir / "warehouse"),
+ })
+ with patch(
+ "pypaimon.multimodal.lerobot.api."
+ "_import_lerobot_dataset",
+ return_value=object,
+ ), patch(
+ "pypaimon.multimodal.lerobot.api."
+ "_open_resolved_dataset",
+ return_value=Dataset(),
+ ):
+ connection.load_from_lerobot(
+ "frames", temp_dir, batch_size=1)
+
+ table = connection.get_table("frames")
+ rows = table.scan().select([
+ "index", "camera"
+ ]).to_arrow().sort_by("index").to_pylist()
+ descriptors = [
+ pmm.VideoFrameDescriptor.deserialize(row["camera"])
+ for row in rows
+ ]
+ self.assertEqual(
+ [4, 5, 1, 2, 3],
+ [descriptor.frame_index for descriptor in descriptors],
+ )
+
+ class Decoder:
+
+ def __init__(self, source):
+ self.container = av.open(source)
+ self.frames = list(self.container.decode(video=0))
+
+ def value(self, frame_index):
+ frame = self.frames[frame_index]
+ return float(frame.to_ndarray(
+ format="rgb24").mean()), float(frame.time)
+
+ def close(self):
+ self.container.close()
+
+ collator = pmm.VideoFrameCollator(
+ table,
+ video_column="camera",
+ decoder_factory=Decoder,
+ decode_fn=lambda decoder, frame_index, row: decoder.value(
+ frame_index),
+ output_column="decoded",
+ collate_fn=lambda decoded_rows: decoded_rows,
+ )
+ try:
+ decoded_rows = collator(rows)
+ self.assertEqual(1, len(collator._decoders))
+ finally:
+ collator.close()
+ np.testing.assert_allclose(
+ [row["decoded"][0] for row in decoded_rows],
+ expected_frame_values,
+ atol=5,
+ )
+ np.testing.assert_allclose(
+ [row["decoded"][1] for row in decoded_rows],
+ [0.5, 0.6, 0.1, 0.2, 0.3],
+ atol=1e-6,
+ )
finally:
shutil.rmtree(temp_dir, ignore_errors=True)
+ def test_video_ordinals_follow_source_timestamps(self):
+ info = {
+ "fps": 10,
+ "features": {
+ "episode_index": {"dtype": "int64", "shape": [1]},
+ "frame_index": {"dtype": "int64", "shape": [1]},
+ "timestamp": {"dtype": "float32", "shape": [1]},
+ "camera": {
+ "dtype": "video",
+ "shape": [8, 10, 3],
+ "video_info": {"video.fps": 10.0},
+ },
+ },
+ }
+ rows = pa.table({
+ "episode_index": pa.array([0, 0, 1, 1], type=pa.int64()),
+ "frame_index": pa.array([0, 1, 0, 1], type=pa.int64()),
+ "timestamp": pa.array(
+ [0.0, 0.1, 0.0, 0.1], type=pa.float32()),
+ })
+
+ class Dataset:
+
+ root = Path("/")
+
+ def __init__(self, rows):
+ self.rows = rows
+
+ def read_batch(self, begin, end):
+ return self.rows.slice(begin, end - begin)
+
+ def video_sample_timestamps(self, unused_uri):
+ return [0.0, 0.1, 0.2, 0.3, 0.5, 0.6]
+
+ schema = _schema_from_info(info)
+ episodes = [
+ {
+ "episode_index": 0,
+ "length": 2,
+ "dataset_from_index": 0,
+ "dataset_to_index": 2,
+ "videos/camera/chunk_index": 0,
+ "videos/camera/file_index": 0,
+ "videos/camera/from_timestamp": 0.5,
+ "videos/camera/to_timestamp": 0.7,
+ },
+ {
+ "episode_index": 1,
+ "length": 2,
+ "dataset_from_index": 2,
+ "dataset_to_index": 4,
+ "videos/camera/chunk_index": 0,
+ "videos/camera/file_index": 0,
+ "videos/camera/from_timestamp": 0.1,
+ "videos/camera/to_timestamp": 0.3,
+ },
+ ]
+ video_sources = {}
+ with patch(
+ "pypaimon.multimodal.lerobot.loader._video_source",
+ return_value=("file:/video.mp4", 10),
+ ):
+ results = [
+ _read_batch(
+ Dataset(rows), info, begin, begin + 2, schema,
+ episode=episode, video_sources=video_sources,
+ )
+ for begin, episode in zip((0, 2), episodes)
+ ]
+ self.assertEqual(
+ [4, 5, 1, 2],
+ [
+ pmm.VideoFrameDescriptor.deserialize(value.as_py()).frame_index
+ for result in results
+ for value in result["camera"]
+ ],
+ )
+
+ missing = dict(episodes[0])
+ missing.update({
+ "videos/camera/from_timestamp": 0.8,
+ "videos/camera/to_timestamp": 1.0,
+ })
+ with patch(
+ "pypaimon.multimodal.lerobot.loader._video_source",
+ return_value=("file:/video.mp4", 10),
+ ), self.assertRaisesRegex(ValueError, "has no frame"):
+ _read_batch(
+ Dataset(rows), info, 0, 2, schema,
+ episode=missing, video_sources={})
+
+ for timestamp_range in (
+ (-0.1, 0.1), (0.0, float("nan")), (0.0, 0.3)):
+ invalid_episode = dict(episodes[0])
+ invalid_episode.update({
+ "videos/camera/from_timestamp": timestamp_range[0],
+ "videos/camera/to_timestamp": timestamp_range[1],
+ })
+ with self.subTest(timestamp_range=timestamp_range), \
+ self.assertRaisesRegex(ValueError, "timestamp|duration"):
+ _read_batch(
+ Dataset(rows), info, 0, 2, schema,
+ episode=invalid_episode, video_sources={})
+
+ def test_video_sample_timestamps_skip_discard_packets(self):
+ class Container:
+
+ streams = SimpleNamespace(video=[SimpleNamespace(
+ time_base=Fraction(1, 10))])
+
+ def __enter__(self):
+ return self
+
+ def __exit__(self, unused_type, unused_value, unused_traceback):
+ pass
+
+ def demux(self, unused_stream):
+ return [
+ SimpleNamespace(
+ pts=-1, time_base=Fraction(1, 10),
+ is_discard=True),
+ SimpleNamespace(
+ pts=0, time_base=Fraction(1, 10),
+ is_discard=False),
+ SimpleNamespace(
+ pts=1, time_base=Fraction(1, 10),
+ is_discard=False),
+ ]
+
+ fake_av = SimpleNamespace(open=lambda unused_source: Container())
+ with patch.dict(sys.modules, {"av": fake_av}):
+ timestamps = _video_sample_timestamps(
+ SimpleNamespace(), "file:/video.mp4")
+ self.assertEqual([0.0, 0.1], list(timestamps))
+
+ def test_video_frame_matching_quantizes_float32_timestamp(self):
+ frame_index = 61441
+ timestamps = array(
+ "d", (index / 30 for index in range(frame_index + 1)))
+ frame_timestamp = pa.scalar(
+ frame_index / 30, type=pa.float32()).as_py()
+
+ self.assertEqual(
+ frame_index,
+ _video_frame_ordinal(
+ timestamps, frame_timestamp, 0.0, pa.float32(), "camera"),
+ )
+
class _RemoteLeRobotFileIO:
@@ -1384,7 +2091,9 @@ class _RemoteLeRobotFileIO:
native_path = remote_path.split("://", 1)[1]
file_type = pafs.FileType.Directory if local_path.is_dir() \
else pafs.FileType.File
- return pafs.FileInfo(native_path, file_type)
+ size = local_path.stat().st_size \
+ if file_type == pafs.FileType.File else None
+ return pafs.FileInfo(native_path, file_type, size=size)
def get_file_status(self, remote_path):
local_path = self._local_path(remote_path)
@@ -2387,7 +3096,8 @@ class LeRobotImportTest(unittest.TestCase):
source,
source_schema,
batch_size,
- metadata):
+ metadata,
+ video_fields=()):
table.add(_read_batch(
dataset,
info,
@@ -2403,6 +3113,7 @@ class LeRobotImportTest(unittest.TestCase):
source_schema,
batch_size,
metadata,
+ video_fields,
)
with patch.object(
diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_writer_test.py
b/paimon-python/pypaimon/tests/multimodal_lerobot_writer_test.py
index 26fb37a248..b796e04061 100644
--- a/paimon-python/pypaimon/tests/multimodal_lerobot_writer_test.py
+++ b/paimon-python/pypaimon/tests/multimodal_lerobot_writer_test.py
@@ -437,6 +437,22 @@ class PaimonLeRobotWriterTest(unittest.TestCase):
},
)
+ def test_video_features_are_rejected(self):
+ with self.assertRaisesRegex(
+ ValueError, "does not support video features: camera"):
+ PaimonLeRobotWriter(
+ self.connection,
+ "video",
+ fps=10,
+ features={
+ "camera": {
+ "dtype": "video",
+ "shape": (3, 4, 5),
+ "names": ["channels", "height", "width"],
+ },
+ },
+ )
+
def
test_missing_lerobot_stats_dependency_fails_before_table_creation(self):
self.connection.catalog.create_database(
"default", ignore_if_exists=True)
diff --git a/paimon-python/pypaimon/write/file_store_write.py
b/paimon-python/pypaimon/write/file_store_write.py
index df3022b37d..a73b9df77d 100644
--- a/paimon-python/pypaimon/write/file_store_write.py
+++ b/paimon-python/pypaimon/write/file_store_write.py
@@ -48,6 +48,7 @@ class FileStoreWrite:
self.max_seq_numbers: dict = {}
self.write_cols = None
self.blob_consumer = None
+ self.blob_uri_reader_factory = None
self.commit_identifier = 0
self.options = CoreOptions.copy(table.options)
self.changelog_producer = self.options.changelog_producer()
@@ -110,6 +111,11 @@ class FileStoreWrite:
)
writer.write(data.to_batches()[0])
+ def roll_before_group_if_needed(self, row_count: int):
+ for writer in self.data_writers.values():
+ if isinstance(writer, DedicatedFormatWriter):
+ writer.roll_before_group_if_needed(row_count)
+
def _check_runtime_bucket(self, partition, bucket, total_buckets):
if total_buckets is None:
return
@@ -168,6 +174,7 @@ class FileStoreWrite:
write_cols=self.write_cols,
blob_consumer=self.blob_consumer,
changelog_producer=self.changelog_producer,
+ blob_uri_reader_factory=self.blob_uri_reader_factory,
)
elif self._has_vector_columns() and options.with_vector_format():
return DataVectorWriter(
diff --git a/paimon-python/pypaimon/write/table_write.py
b/paimon-python/pypaimon/write/table_write.py
index f0a68bcf6f..fd09282c1d 100644
--- a/paimon-python/pypaimon/write/table_write.py
+++ b/paimon-python/pypaimon/write/table_write.py
@@ -85,6 +85,13 @@ class TableWrite:
sub_table = pa.compute.take(data, row_indices)
self._write_partition_bucket_batch(partition, bucket, sub_table)
+ def roll_before_group_if_needed(self, row_count: int):
+ """Roll current files before the next logical write group if needed."""
+ if isinstance(row_count, bool) or not isinstance(row_count, int) \
+ or row_count <= 0:
+ raise ValueError("Group row count must be a positive integer.")
+ self.file_store_write.roll_before_group_if_needed(row_count)
+
def _write_partition_bucket_batch(self, partition, bucket, data):
self.file_store_write.write(partition, bucket, data)
@@ -226,6 +233,15 @@ class TableWrite:
self.file_store_write.blob_consumer = blob_consumer
return self
+ def with_blob_uri_reader_factory(self, uri_reader_factory):
+ if self.file_store_write.data_writers:
+ raise RuntimeError(
+ "with_blob_uri_reader_factory must be called before any "
+ "write operation."
+ )
+ self.file_store_write.blob_uri_reader_factory = uri_reader_factory
+ return self
+
def write_ray(
self,
dataset: "Dataset",
diff --git a/paimon-python/pypaimon/write/writer/blob_file_writer.py
b/paimon-python/pypaimon/write/writer/blob_file_writer.py
index 887d8cd792..77605ac16c 100644
--- a/paimon-python/pypaimon/write/writer/blob_file_writer.py
+++ b/paimon-python/pypaimon/write/writer/blob_file_writer.py
@@ -45,9 +45,10 @@ class BlobFileWriter:
def __init__(self, file_io, file_path: Path, blob_consumer:
Optional[BlobConsumer] = None,
copy_buffer_size: int = BlobFormatWriter.BUFFER_SIZE,
- video: bool = False):
+ video: bool = False, uri_reader_factory=None):
self.file_io = file_io
self.file_path = file_path
+ self._uri_reader_factory = uri_reader_factory
self._blob_consumer = blob_consumer
if video:
if blob_consumer is not None:
@@ -122,7 +123,8 @@ class BlobFileWriter:
if isinstance(col_data, bytes):
if BlobDescriptorSerde.is_descriptor(col_data):
descriptor = BlobDescriptorSerde.deserialize(col_data)
- uri_reader =
self.file_io.uri_reader_factory.create(descriptor.uri)
+ factory = self._uri_reader_factory or
self.file_io.uri_reader_factory
+ uri_reader = factory.create(descriptor.uri)
return Blob.from_descriptor(uri_reader, descriptor)
return BlobData(col_data)
diff --git a/paimon-python/pypaimon/write/writer/blob_writer.py
b/paimon-python/pypaimon/write/writer/blob_writer.py
index c32205b2c9..166144e5c6 100644
--- a/paimon-python/pypaimon/write/writer/blob_writer.py
+++ b/paimon-python/pypaimon/write/writer/blob_writer.py
@@ -37,7 +37,7 @@ class BlobWriter(AppendOnlyDataWriter):
def __init__(self, table, partition: Tuple, bucket: int, max_seq_number:
int, blob_column: str,
options: Dict[str, str] = None, blob_consumer:
Optional[BlobConsumer] = None,
- video: bool = False):
+ video: bool = False, uri_reader_factory=None):
super().__init__(table, partition, bucket, max_seq_number,
options, write_cols=[blob_column])
@@ -61,6 +61,7 @@ class BlobWriter(AppendOnlyDataWriter):
self.blob_copy_buffer_size = self.options.blob_copy_buffer_size()
self._blob_consumer = blob_consumer
+ self._uri_reader_factory = uri_reader_factory
self.current_writer: Optional[BlobFileWriter] = None
self.current_file_path: Optional[str] = None
self.record_count = 0
@@ -140,6 +141,7 @@ class BlobWriter(AppendOnlyDataWriter):
blob_consumer=self._blob_consumer,
copy_buffer_size=self.blob_copy_buffer_size,
video=self.video,
+ uri_reader_factory=self._uri_reader_factory,
)
def rolling_file(self) -> bool:
@@ -151,6 +153,17 @@ class BlobWriter(AppendOnlyDataWriter):
or
self.current_writer.reach_target_size(self.blob_target_file_size)
)
+ def should_roll_before_group(self, row_count: int) -> bool:
+ return (
+ self._video_group_policy is not None
+ and self.current_writer is not None
+ and (
+ self._video_group_policy.pending_roll
+ or self.current_writer.row_count + row_count
+ > self.target_file_row_num
+ )
+ )
+
def close_current_writer(self):
"""Close current writer and create metadata."""
if self.current_writer is None:
diff --git a/paimon-python/pypaimon/write/writer/dedicated_format_writer.py
b/paimon-python/pypaimon/write/writer/dedicated_format_writer.py
index 3df6db795d..c894182fbe 100644
--- a/paimon-python/pypaimon/write/writer/dedicated_format_writer.py
+++ b/paimon-python/pypaimon/write/writer/dedicated_format_writer.py
@@ -68,7 +68,8 @@ class DedicatedFormatWriter(DataWriter):
def __init__(self, table, partition: Tuple, bucket: int, max_seq_number:
int, options: CoreOptions = None,
write_cols: Optional[List[str]] = None, blob_consumer:
Optional[BlobConsumer] = None,
- changelog_producer: ChangelogProducer =
ChangelogProducer.NONE):
+ changelog_producer: ChangelogProducer =
ChangelogProducer.NONE,
+ blob_uri_reader_factory=None):
super().__init__(table, partition, bucket, max_seq_number, options,
write_cols=write_cols,
changelog_producer=changelog_producer)
@@ -181,6 +182,7 @@ class DedicatedFormatWriter(DataWriter):
options=options,
blob_consumer=blob_consumer,
video=blob_column in configured_video_fields,
+ uri_reader_factory=blob_uri_reader_factory,
)
# Initialize vector writer when vector.file.format is configured.
@@ -195,6 +197,7 @@ class DedicatedFormatWriter(DataWriter):
vector_columns=self.vector_write_columns,
vector_file_format=options.vector_file_format(),
options=options,
+ rolling_managed_by_parent=bool(self.video_frame_columns),
)
logger.info(
@@ -270,8 +273,8 @@ class DedicatedFormatWriter(DataWriter):
self.record_count += data.num_rows
- # Check if normal data rolling is needed
- if self._should_roll_normal():
+ # Defer any active video-group roll to its Episode boundary.
+ if self._should_roll_active_group():
self._roll_or_defer_for_video_group()
def write_row(self, row):
@@ -325,7 +328,7 @@ class DedicatedFormatWriter(DataWriter):
self.vector_writer.write(vector_data)
self.record_count += 1
- if self._should_roll_normal():
+ if self._should_roll_active_group():
self._roll_or_defer_for_video_group()
except Exception as e:
@@ -520,6 +523,39 @@ class DedicatedFormatWriter(DataWriter):
# Check if normal data exceeds target size
return self._normal_buffer.nbytes > self.target_file_size
+ def roll_before_group_if_needed(self, row_count: int):
+ """Roll current files before the next logical write group if needed."""
+ self._require_finished_flush()
+ if self._video_group_policy is None:
+ return
+
+ pending_rows = self.pending_row_count
+ should_roll = pending_rows > 0 and (
+ self._video_group_policy.pending_roll
+ or pending_rows + row_count > self.target_file_row_num
+ or (
+ not self._normal_buffer.is_empty
+ and self._normal_buffer.nbytes > self.target_file_size
+ )
+ )
+ should_roll = should_roll or any(
+ self.blob_writers[column].should_roll_before_group(row_count)
+ for column in self.video_frame_columns
+ )
+ should_roll = should_roll or (
+ self.vector_writer is not None
+ and self.vector_writer.should_roll_before_group(row_count)
+ )
+ if should_roll:
+ self._close_current_writers()
+
+ def _should_roll_active_group(self) -> bool:
+ return self._should_roll_normal() or (
+ self._video_group_policy is not None
+ and self.vector_writer is not None
+ and self.vector_writer.rolling_file()
+ )
+
def _roll_or_defer_for_video_group(self):
if self._video_group_policy is not None and
self._video_group_policy.defer_roll():
return
diff --git a/paimon-python/pypaimon/write/writer/vector_writer.py
b/paimon-python/pypaimon/write/writer/vector_writer.py
index b51e2e49ae..a7cd6bfc58 100644
--- a/paimon-python/pypaimon/write/writer/vector_writer.py
+++ b/paimon-python/pypaimon/write/writer/vector_writer.py
@@ -36,16 +36,41 @@ class VectorWriter(AppendOnlyDataWriter):
"""
def __init__(self, table, partition: Tuple, bucket: int, max_seq_number:
int,
- vector_columns: List[str], vector_file_format: str, options:
CoreOptions = None):
+ vector_columns: List[str], vector_file_format: str,
+ options: CoreOptions = None,
+ rolling_managed_by_parent: bool = False):
super().__init__(table, partition, bucket, max_seq_number,
options, write_cols=vector_columns)
self.vector_columns = vector_columns
self.vector_file_format = vector_file_format
self.file_format = vector_file_format
self.target_file_size = options.vector_target_file_size()
+ # The parent coordinates rolling at logical group boundaries.
+ self.rolling_managed_by_parent = rolling_managed_by_parent
self.file_uuid = str(uuid.uuid4())
self.file_count = 0
+ def _check_and_roll_if_needed(self):
+ if not self.rolling_managed_by_parent:
+ super()._check_and_roll_if_needed()
+
+ def rolling_file(self) -> bool:
+ return (
+ self._buffer.num_rows >= self.target_file_row_num
+ or self._buffer.nbytes > self.target_file_size
+ )
+
+ def should_roll_before_group(self, row_count: int) -> bool:
+ return (
+ self.rolling_managed_by_parent
+ and self._buffer.num_rows > 0
+ and (
+ self.rolling_file()
+ or self._buffer.num_rows + row_count
+ > self.target_file_row_num
+ )
+ )
+
def _write_data_to_file(self, data: pa.Table):
if data.num_rows == 0:
return
