diff --git a/docs/docs/pypaimon/multimodal-api.mdx b/docs/docs/pypaimon/multimodal-api.mdx index 6edbeeb7980e..8d0287030d48 100644 --- a/docs/docs/pypaimon/multimodal-api.mdx +++ b/docs/docs/pypaimon/multimodal-api.mdx @@ -694,6 +694,10 @@ 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. When rolling is needed, writers roll together between +Episodes. Video imports require a bucket-unaware table. + ## Capture LeRobot frames directly into Paimon `PaimonLeRobotWriter` implements the write-side surface used by LeRobot's diff --git a/paimon-python/pypaimon/multimodal/lerobot/api.py b/paimon-python/pypaimon/multimodal/lerobot/api.py index 84e3d2eed5c1..20a054f90fa0 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/api.py +++ b/paimon-python/pypaimon/multimodal/lerobot/api.py @@ -37,6 +37,7 @@ _require_v3, _schema_from_info, _validate_v3_required_features, + _video_feature_names, ) from pypaimon.multimodal.lerobot.source import ( _close_quietly, @@ -52,6 +53,12 @@ ) +_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 71719869c39e..01708e866663 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.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 @@ "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.begin_video_episode( + 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,316 @@ 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 + if any(feature.get("dtype") == "video" + for feature in features.values()): + video_rows = _validate_video_rows( + raw, info, episode, begin, end) 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 {}, + ) + 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): + 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, from_timestamp + timestamp, 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, timestamp, name): + 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), + ) + distance = abs(timestamps[ordinal] - timestamp) + if distance >= _VIDEO_TIMESTAMP_TOLERANCE: + raise ValueError( + "LeRobot video feature %s has no frame within %s seconds of " + "timestamp %s." + % (name, _VIDEO_TIMESTAMP_TOLERANCE, 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 b1130bbbb1e4..80201b5a9e18 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/metadata.py +++ b/paimon-python/pypaimon/multimodal/lerobot/metadata.py @@ -69,15 +69,33 @@ 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 @@ -91,13 +109,22 @@ def __getitem__(self, index): 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): @@ -505,14 +532,29 @@ def read_schema(path): 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( @@ -551,7 +593,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 de16c95a6478..188127a14dbe 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 04547db80403..995c14a43864 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/source.py +++ b/paimon-python/pypaimon/multimodal/lerobot/source.py @@ -32,7 +32,9 @@ 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 @@ def __getitem__(self, index): 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 @@ def image_bytes(self, value): 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 faa3c6b5bdb4..a72a4a4b1ce0 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/writer.py +++ b/paimon-python/pypaimon/multimodal/lerobot/writer.py @@ -36,6 +36,7 @@ _feature_shape, _schema_from_info, _validate_lerobot_schema, + _video_feature_names, ) from pypaimon.multimodal.table import _target_schema @@ -79,6 +80,11 @@ def __init__( 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)) self.fps = fps self.episodes_per_commit = episodes_per_commit 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 eb5b61786cef..5efe37047ce4 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 @@ def _write_files(self, table, data): 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.begin_video_episode(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 @@ def test_multiple_video_fields_allow_nested_episode_boundaries(self): 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 5be17dc7b53e..8f7c7d5eac15 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 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 @@ _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,7 @@ from pypaimon.multimodal.lerobot.loader import ( _image_bytes, _read_batch, + _video_sample_timestamps, _validate_frame_controls, ) from pypaimon.multimodal.lerobot.schema import ( @@ -83,6 +87,11 @@ 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 +891,10 @@ def test_schema_comes_from_metadata_and_rejects_unsupported_types(self): "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 +1004,7 @@ def test_remote_episode_metadata_projects_stats_columns(self): "episode_index": [0], "dataset_from_index": [0], "dataset_to_index": [1], + "length": [1], "data/chunk_index": [0], "data/file_index": [0], }) @@ -1026,6 +1038,16 @@ def test_empty_local_dataset_is_rejected_before_opening_lerobot(self): "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 +1064,72 @@ def test_empty_local_dataset_is_rejected_before_opening_lerobot(self): 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 +1427,631 @@ def test_local_v2_is_rejected_before_opening(self): 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 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)) + class _RemoteLeRobotFileIO: @@ -1384,7 +2077,9 @@ def _status(self, local_path): 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 +3082,8 @@ def append_then_write( source, source_schema, batch_size, - metadata): + metadata, + video_fields=()): table.add(_read_batch( dataset, info, @@ -2403,6 +3099,7 @@ def append_then_write( 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 f446de75818f..821464c3f80d 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_writer_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_writer_test.py @@ -150,6 +150,22 @@ def test_existing_table_requires_matching_feature_schema(self): }, ) + 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_commits_multiple_completed_episodes_as_one_batch(self): writer = PaimonLeRobotWriter( self.connection, diff --git a/paimon-python/pypaimon/write/file_store_write.py b/paimon-python/pypaimon/write/file_store_write.py index df3022b37d2d..01683bdda9d3 100644 --- a/paimon-python/pypaimon/write/file_store_write.py +++ b/paimon-python/pypaimon/write/file_store_write.py @@ -48,6 +48,7 @@ def __init__(self, table, commit_user): 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 @@ def write_row( ) writer.write(data.to_batches()[0]) + def begin_video_episode(self, row_count: int): + for writer in self.data_writers.values(): + if isinstance(writer, DedicatedFormatWriter): + writer.begin_video_episode(row_count) + def _check_runtime_bucket(self, partition, bucket, total_buckets): if total_buckets is None: return @@ -168,6 +174,7 @@ def max_seq_number(): 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 f0a68bcf6fc1..c27de20eeb07 100644 --- a/paimon-python/pypaimon/write/table_write.py +++ b/paimon-python/pypaimon/write/table_write.py @@ -85,6 +85,13 @@ def write_arrow_batch(self, data: pa.RecordBatch): sub_table = pa.compute.take(data, row_indices) self._write_partition_bucket_batch(partition, bucket, sub_table) + def begin_video_episode(self, row_count: int): + """Keep the next video Episode within one aligned normal file.""" + if isinstance(row_count, bool) or not isinstance(row_count, int) \ + or row_count <= 0: + raise ValueError("Video Episode row count must be a positive integer.") + self.file_store_write.begin_video_episode(row_count) + def _write_partition_bucket_batch(self, partition, bucket, data): self.file_store_write.write(partition, bucket, data) @@ -226,6 +233,15 @@ def with_blob_consumer(self, blob_consumer: BlobConsumer): 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 887d8cd792b8..77605ac16c3f 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 @@ def _to_blob(self, col_data) -> Optional[Blob]: 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 c32205b2c985..bc3bfba8484f 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 @@ def __init__(self, table, partition: Tuple, bucket: int, max_seq_number: int, bl 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 @@ def open_current_writer(self): 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 @@ def rolling_file(self) -> bool: or self.current_writer.reach_target_size(self.blob_target_file_size) ) + def should_roll_before_video_episode(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 3df6db795d00..7ca0da784045 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 @@ def __init__(self, table, partition: Tuple, bucket: int, max_seq_number: int, op 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 @@ def __init__(self, table, partition: Tuple, bucket: int, max_seq_number: int, op 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 @@ def _write_batch(self, data: pa.RecordBatch): 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 @@ def write_row(self, row): 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,40 @@ def _should_roll_normal(self) -> bool: # Check if normal data exceeds target size return self._normal_buffer.nbytes > self.target_file_size + def begin_video_episode(self, row_count: int): + """Roll only between complete Episodes, before writing the next one.""" + 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_video_episode( + 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_video_episode(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 b51e2e49aecd..f492e355e1c3 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() + # Video tables close normal and sidecar files at one Episode boundary. + 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_video_episode(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