diff --git a/docs/docs/pypaimon/multimodal-api.mdx b/docs/docs/pypaimon/multimodal-api.mdx index e8732875b176..8b4c204bd516 100644 --- a/docs/docs/pypaimon/multimodal-api.mdx +++ b/docs/docs/pypaimon/multimodal-api.mdx @@ -556,11 +556,15 @@ snapshot_id = conn.load_from_lerobot( Missing tables are created from metadata; existing tables use strict schema validation and append semantics. Scalars map to scalar types, vectors to -`VECTOR`, higher-rank tensors to nested `ARRAY`, and images to `BLOB`. Images -keep their compressed bytes. - -Only v3 is supported. Video features, `uint64`, and language event structures -are rejected. +`VECTOR`, higher-rank tensors to nested `ARRAY`, and images or video frames to +`BLOB`. Images keep their compressed bytes; videos pack encoded MP4 payloads +instead of repeating them per frame, with frame ordinals in the video field. +Video imports keep each Episode in one aligned file group. When the normal +file or a camera sidecar needs to roll, all active writers close together at +the next Episode boundary. The target table must be unpartitioned and +bucket-unaware. + +Only v3 is supported. `uint64` and language event structures are rejected. ## Overwrite diff --git a/paimon-python/pypaimon/multimodal/lerobot/api.py b/paimon-python/pypaimon/multimodal/lerobot/api.py index 0b7f76d78a46..9015d133160d 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/api.py +++ b/paimon-python/pypaimon/multimodal/lerobot/api.py @@ -34,6 +34,7 @@ _require_v3, _schema_from_info, _validate_lerobot_schema, + _video_feature_names, ) from pypaimon.multimodal.lerobot.source import ( _has_tasks, @@ -48,6 +49,13 @@ _validate_source_kerberos, ) from pypaimon.multimodal.table import _target_schema +from pypaimon.table.bucket_mode import BucketMode + + +_VIDEO_LAYOUT_ERROR = ( + "LeRobot video import requires an unpartitioned, bucket-unaware " + "target table so each Episode is written by one writer." +) def load_from_lerobot( @@ -83,6 +91,7 @@ def load_from_lerobot( _require_v3(local_info, resolved_source.path) _validate_info_paths(local_info) _schema_from_info(local_info, include_task=False) + video_fields = _video_feature_names(local_info) total_frames, _, total_tasks = \ _validated_counts(local_info, resolved_source.path) if total_frames == 0: @@ -96,11 +105,16 @@ def load_from_lerobot( source_schema, options, resolved_source, + video_fields, ) return None 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) @@ -109,12 +123,14 @@ def load_from_lerobot( source_schema = _schema_from_info( info, include_task=_has_tasks(dataset, info)) + video_fields = _video_feature_names(info) table = _validated_table( connection, table_name, source_schema, options, resolved_source, + video_fields, ) if row_count == 0: @@ -126,6 +142,7 @@ def load_from_lerobot( resolved_source, source_schema, batch_size, + video_fields, ) finally: close = getattr(dataset, "close", None) @@ -160,9 +177,10 @@ def _required_count(info, name, source): def _validated_table( - connection, table_name, source_schema, options, source): + connection, table_name, source_schema, options, source, + video_fields=()): table = _get_or_create_table( - connection, table_name, source_schema, options) + connection, table_name, source_schema, options, video_fields) target_schema = _target_schema(table.raw_table) _validate_lerobot_schema( source_schema, target_schema, source.path) @@ -172,13 +190,42 @@ def _validated_table( source, 0, ) + configured = table.raw_table.options.video_frame_fields() + if configured != set(video_fields): + raise ValueError( + "LeRobot video features %s require table option " + "'video-frame-field'=%r; found %s." + % (list(video_fields), ",".join(video_fields), sorted(configured)) + ) + if video_fields and ( + table.raw_table.partition_keys + or table.raw_table.bucket_mode() != BucketMode.BUCKET_UNAWARE): + raise ValueError(_VIDEO_LAYOUT_ERROR) return table -def _get_or_create_table(connection, table_name, schema, options): +def _get_or_create_table( + connection, table_name, schema, options, video_fields=()): try: return connection.get_table(table_name) except (DatabaseNotExistException, TableNotExistException): + options = dict(options or {}) + if video_fields and str(options.get("bucket", "-1")).strip() != "-1": + raise ValueError(_VIDEO_LAYOUT_ERROR) + configured = 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: + options["video-frame-field"] = ",".join(video_fields) return connection.create_table( table_name, schema=schema, diff --git a/paimon-python/pypaimon/multimodal/lerobot/loader.py b/paimon-python/pypaimon/multimodal/lerobot/loader.py index 1f46e047d2b4..f6df046b6671 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/loader.py +++ b/paimon-python/pypaimon/multimodal/lerobot/loader.py @@ -27,6 +27,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 +50,7 @@ "float64", } _BOOLEAN_DTYPES = {"bool", "boolean"} +_VIDEO_TIMESTAMP_TOLERANCE = 1e-4 def _strict_lerobot_table(data, target_schema, source, batch_index): @@ -67,7 +69,8 @@ def _write_dataset( info, source, source_schema, - batch_size): + batch_size, + video_fields=()): target_schema = _target_schema(table.raw_table) write_builder = table.raw_table.new_batch_write_builder() table_write = None @@ -76,23 +79,39 @@ def _write_dataset( batch_count = 0 row_count = 0 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 begin, end in _episode_batches(dataset, info, batch_size): - batch = _read_batch( - dataset, info, begin, end, source_schema) - batch = _strict_lerobot_table( - batch, - target_schema, - source, - batch_count, - ) - table_write.write_arrow(batch) - batch_count += 1 - row_count += batch.num_rows + for episode, episode_begin, episode_end in _episodes(dataset, info): + if video_fields: + table_write.begin_video_episode(episode_end - episode_begin) + for begin in range(episode_begin, episode_end, batch_size): + end = min(begin + batch_size, episode_end) + batch = _read_batch( + dataset, + info, + begin, + end, + source_schema, + episode=episode, + video_sources=video_sources, + ) + batch = _strict_lerobot_table( + batch, + target_schema, + source, + batch_count, + ) + table_write.write_arrow(batch) + batch_count += 1 + row_count += batch.num_rows expected_rows = int(info.get("total_frames", len(dataset))) if row_count != expected_rows: @@ -119,7 +138,7 @@ def _write_dataset( table_commit.close() -def _episode_batches(dataset, info, batch_size): +def _episodes(dataset, info): episodes = getattr(dataset.meta, "episodes", None) episode_count = int(info.get("total_episodes", 0)) total_frames = int(info.get("total_frames", len(dataset))) @@ -129,17 +148,26 @@ def _episode_batches(dataset, info, batch_size): for ordinal in range(episode_count): episode = episodes.iloc[ordinal] if hasattr(episodes, "iloc") \ else episodes[ordinal] - begin = int(_python_scalar(episode["dataset_from_index"])) - end = int(_python_scalar(episode["dataset_to_index"])) - if begin != expected_begin or end <= begin: + try: + episode_index = _nonnegative_integer( + episode["episode_index"], "episode_index") + length = _nonnegative_integer(episode["length"], "length") + begin = _nonnegative_integer( + episode["dataset_from_index"], "dataset_from_index") + end = _nonnegative_integer( + episode["dataset_to_index"], "dataset_to_index") + except (KeyError, TypeError) as error: + raise ValueError( + "LeRobot episode %d is missing required boundary metadata." + % ordinal + ) from error + if episode_index != ordinal or begin != expected_begin \ + or end <= begin or length != end - begin: raise ValueError( - "LeRobot episode %d has invalid frame range [%d, %d); " - "expected it to start at %d." + "LeRobot episode %d has invalid index, length, or frame " + "range [%d, %d); expected it to start at %d." % (ordinal, begin, end, expected_begin)) - while begin < end: - batch_end = min(begin + batch_size, end) - yield begin, batch_end - begin = batch_end + yield episode, begin, end expected_begin = end if expected_begin != total_frames: raise ValueError( @@ -147,7 +175,9 @@ def _episode_batches(dataset, info, batch_size): % (expected_begin, total_frames)) -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) @@ -158,26 +188,47 @@ 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": + if episode is None: + raise ValueError( + "LeRobot video import requires Episode metadata.") + 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) @@ -191,6 +242,212 @@ def _read_batch(dataset, info, begin, end, schema): return pa.Table.from_arrays(arrays, schema=pa.schema(fields)) +def _video_frame_descriptors( + dataset, info, episode, video_rows, name, feature, begin, end, cache): + 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 + + first_frame = _aligned_frame_ordinal( + from_timestamp, fps, "video feature %s from_timestamp" % name) + to_frame = _aligned_frame_ordinal( + to_timestamp, fps, "video feature %s to_timestamp" % name) + if first_frame < 0 or to_frame - first_frame != episode_end - episode_begin: + raise ValueError( + "LeRobot video feature %s has frame range [%d, %d), but " + "Episode [%d, %d) contains %d frames." + % ( + name, + first_frame, + to_frame, + episode_begin, + episode_end, + episode_end - episode_begin, + ) + ) + + source_key = (name, chunk_index, file_index) + source = cache.get(source_key) + if source is None: + source = _video_source( + dataset, info, episode, name, chunk_index, file_index) + cache[source_key] = source + uri, length = source + descriptors = [] + for episode_frame_index, timestamp in video_rows: + frame_index = first_frame + episode_frame_index + shifted_timestamp = from_timestamp + timestamp + if not math.isclose( + shifted_timestamp, + frame_index / fps, + rel_tol=0.0, + abs_tol=_VIDEO_TIMESTAMP_TOLERANCE): + raise ValueError( + "LeRobot video feature %s frame %d has shifted timestamp " + "%s, expected %s." + % (name, episode_frame_index, shifted_timestamp, + frame_index / fps) + ) + descriptors.append(VideoFrameDescriptor( + uri, 0, length, frame_index).serialize()) + return descriptors + + +def _validate_video_rows(raw, info, episode, begin, end): + 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 = list(range( + begin - episode_begin, end - episode_begin)) + actual_episodes = [ + _nonnegative_integer(value, "episode_index") + for value in raw.column("episode_index").to_pylist() + ] + actual_frames = [ + _nonnegative_integer(value, "frame_index") + for value in raw.column("frame_index").to_pylist() + ] + if actual_episodes != [episode_index] * (end - begin) \ + or actual_frames != expected_frames: + raise ValueError( + "LeRobot frame rows do not match Episode %d range [%d, %d)." + % (episode_index, begin, end) + ) + timestamps = raw.column("timestamp").to_pylist() + global_fps = _positive_fps(info.get("fps"), "dataset") + timestamp_values = [] + for frame_index, timestamp in zip(actual_frames, timestamps): + try: + value = float(_python_scalar(timestamp)) + except (TypeError, ValueError) as error: + raise ValueError("LeRobot frame timestamp is invalid.") from error + if not math.isfinite(value) or not math.isclose( + value, + frame_index / global_fps, + rel_tol=0.0, + abs_tol=_VIDEO_TIMESTAMP_TOLERANCE): + raise ValueError( + "LeRobot frame %d has timestamp %r, expected %s." + % (frame_index, timestamp, frame_index / global_fps) + ) + timestamp_values.append(value) + return list(zip(actual_frames, timestamp_values)) + + +def _aligned_frame_ordinal(timestamp, fps, owner): + if not math.isfinite(timestamp): + raise ValueError("LeRobot %s is invalid." % owner) + frame = int(round(timestamp * fps)) + if not math.isclose( + timestamp, + frame / fps, + rel_tol=0.0, + abs_tol=_VIDEO_TIMESTAMP_TOLERANCE): + raise ValueError( + "LeRobot %s %s is not aligned to FPS %s." + % (owner, timestamp, fps) + ) + return frame + + +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 _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: @@ -298,6 +555,17 @@ def _python_scalar(value): return value +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 metadata field %s must be a non-negative integer; " + "found %r." % (name, value) + ) + return int(value) + + def _image_bytes(value, root): if value is None: raise ValueError("LeRobot image feature contains a null frame.") diff --git a/paimon-python/pypaimon/multimodal/lerobot/schema.py b/paimon-python/pypaimon/multimodal/lerobot/schema.py index 6b8ccfb8b537..8946eec097ff 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/schema.py +++ b/paimon-python/pypaimon/multimodal/lerobot/schema.py @@ -65,6 +65,14 @@ def _schema_from_info(info, include_task): return pa.schema(fields) +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: @@ -104,11 +112,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 eaf995d69eda..0e379a146a2b 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/source.py +++ b/paimon-python/pypaimon/multimodal/lerobot/source.py @@ -30,6 +30,7 @@ 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.hdf5 import ( _Hdf5SourceFileIO, @@ -37,6 +38,7 @@ _qualified_status_path, ) from pypaimon.multimodal.lerobot.loader import _encode_media_frame +from pypaimon.multimodal.lerobot.schema import _video_feature_names @dataclass(frozen=True) @@ -154,17 +156,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( @@ -172,10 +174,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: @@ -184,6 +187,7 @@ class _RemoteLeRobotDataset: "episode_index", "dataset_from_index", "dataset_to_index", + "length", "data/chunk_index", "data/file_index", ] @@ -251,18 +255,55 @@ 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) + columns = list(self._EPISODE_COLUMNS) + for video_key in _video_feature_names(info): + columns.extend([ + "videos/%s/%s" % (video_key, suffix) + for suffix in ( + "chunk_index", + "file_index", + "from_timestamp", + "to_timestamp", + ) + ]) rows = [] for path in paths: rows.extend(_read_remote_parquet( self._file_io, path, - columns=self._EPISODE_COLUMNS, + columns=columns, ).to_pylist()) rows.sort(key=lambda row: int(row["episode_index"])) if len(rows) != episode_count: @@ -332,6 +373,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 _relative_dataset_path(path, name): if not isinstance(path, str) or not path: raise ValueError("LeRobot %s must be a relative path." % name) 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..aba837311143 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,184 @@ 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)) + 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', + }, + ) + 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, + ) + + wb = table.new_batch_write_builder() + writer = wb.new_write() + for episode_start in (0, 2): + writer.begin_video_episode(2) + for row in range(episode_start, episode_start + 2): + writer.write_arrow(rows.slice(row, 1)) + messages = writer.prepare_commit() + files = [file for message in messages for file in message.new_files] + wb.new_commit().commit(messages) + writer.close() + + 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([2, 2], normal_rows) + self.assertEqual([2, 2], video_rows) + self.assertEqual([2, 2], 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 7f1f8d10a8ce..b23c7afdd616 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -21,17 +21,20 @@ import tempfile import unittest from pathlib import Path +from types import SimpleNamespace from unittest.mock import Mock, patch import numpy as np import pyarrow as pa import pyarrow.fs as pafs +import pyarrow.parquet as pq import pypaimon.multimodal as pmm from pypaimon.common.options import Options from pypaimon.multimodal.hdf5 import _Hdf5SourceFileIO from pypaimon.multimodal.lerobot import load_from_lerobot from pypaimon.multimodal.lerobot.loader import ( + _episodes, _image_bytes, _read_batch, _task_name, @@ -66,6 +69,36 @@ def _replaced_contract(field, old, new): class LeRobotValidationTest(unittest.TestCase): + def test_episode_boundaries_require_integer_metadata(self): + base = { + "episode_index": 0, + "length": 2, + "dataset_from_index": 0, + "dataset_to_index": 2, + } + info = {"total_episodes": 1, "total_frames": 2} + + class Dataset: + + def __init__(self, episode): + self.meta = SimpleNamespace(episodes=[episode]) + + def __len__(self): + return 2 + + for field, value in ( + ("episode_index", 0.9), + ("length", 2.9), + ("dataset_from_index", 0.9), + ("dataset_to_index", 2.9), + ("length", True), + ("length", "2")): + episode = dict(base) + episode[field] = value + with self.subTest(field=field, value=value), \ + self.assertRaisesRegex(ValueError, field): + list(_episodes(Dataset(episode), info)) + def test_dataset_open_never_downloads_videos(self): calls = [] @@ -228,8 +261,11 @@ 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, include_task=False) + self.assertEqual( + pa.large_binary(), + _schema_from_info( + info, include_task=False).field("camera").type, + ) def test_existing_schema_preserves_lerobot_feature_contract(self): source = _schema_from_info({ @@ -313,6 +349,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], }) @@ -343,6 +380,16 @@ def test_empty_local_dataset_returns_before_opening_lerobot(self): "total_tasks": 0, "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={ @@ -355,11 +402,102 @@ def test_empty_local_dataset_returns_before_opening_lerobot(self): "empty_frames", source)) import_lerobot.assert_not_called() table = connection.get_table("empty_frames") + self.assertEqual( + {"camera"}, table.raw_table.options.video_frame_fields()) self.assertIsNone( table.raw_table.snapshot_manager().get_latest_snapshot()) 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: + source = temp_dir / "source" + (source / "meta").mkdir(parents=True) + info = { + "codebase_version": "v3.0", + "fps": 10, + "total_frames": 0, + "total_episodes": 0, + "total_tasks": 0, + "features": { + "index": {"dtype": "int64", "shape": [1]}, + "camera_a": {"dtype": "video", "shape": [8, 10, 3]}, + "camera_b": {"dtype": "video", "shape": [8, 10, 3]}, + }, + } + (source / "meta" / "info.json").write_text(json.dumps(info)) + connection = pmm.connect(options={ + "warehouse": str(temp_dir / "warehouse"), + }) + + with self.assertRaisesRegex(ValueError, "do not match"): + connection.load_from_lerobot( + "conflict", + source, + options={"video-frame-field": "camera_a"}, + ) + + schema = _schema_from_info(info, include_task=False) + connection.create_table("missing_option", schema=schema) + with self.assertRaisesRegex(ValueError, "require table option"): + connection.load_from_lerobot("missing_option", source) + + connection.create_table( + "reordered", + schema=schema, + options={"video-frame-field": "camera_b,camera_a"}, + ) + self.assertIsNone(connection.load_from_lerobot( + "reordered", source)) + 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: + source = temp_dir / "source" + (source / "meta").mkdir(parents=True) + info = { + "codebase_version": "v3.0", + "fps": 10, + "total_frames": 0, + "total_episodes": 0, + "total_tasks": 0, + "features": { + "episode_index": {"dtype": "int64", "shape": [1]}, + "camera": {"dtype": "video", "shape": [8, 10, 3]}, + }, + } + (source / "meta" / "info.json").write_text(json.dumps(info)) + schema = _schema_from_info(info, include_task=False) + connection = pmm.connect(options={ + "warehouse": str(temp_dir / "warehouse"), + }) + connection.create_table( + "partitioned", + schema=schema, + options={"video-frame-field": "camera"}, + partitioned=["episode_index"], + ) + connection.create_table( + "bucketed", + schema=schema, + options={"video-frame-field": "camera", "bucket": "1"}, + ) + + for table_name in ("partitioned", "bucketed"): + with self.subTest(table_name=table_name): + with self.assertRaisesRegex( + ValueError, "unpartitioned, bucket-unaware"): + connection.load_from_lerobot(table_name, source) + with self.assertRaisesRegex( + ValueError, "unpartitioned, bucket-unaware"): + connection.load_from_lerobot( + "new_bucketed", source, options={"bucket": "1"}) + 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: @@ -464,26 +602,307 @@ 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): 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": 0, + "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, + }, + "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, + "videos/camera_a/chunk_index": 0, + "videos/camera_a/file_index": 0, + "videos/camera_a/from_timestamp": 0.0, + "videos/camera_a/to_timestamp": 0.2, + "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, + "videos/camera_a/chunk_index": 0, + "videos/camera_a/file_index": 0, + "videos/camera_a/from_timestamp": 0.2, + "videos/camera_a/to_timestamp": 0.5, + "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, + }, + ] + + class Dataset: + + root = temp_dir + meta = SimpleNamespace( + info=info, episodes=episodes, tasks=None) + 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(), + ), + }) + + 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) + with patch( + "pypaimon.multimodal.lerobot.api." + "_import_lerobot_dataset", + return_value=object, + ), patch( + "pypaimon.multimodal.lerobot.api." + "_open_resolved_dataset", + return_value=Dataset(), + ): + snapshot_id = connection.load_from_lerobot( + "frames", temp_dir, batch_size=1) + + self.assertEqual(1, snapshot_id) + table = connection.get_table("frames") + self.assertEqual( + {"camera_a", "camera_b"}, + table.raw_table.options.video_frame_fields(), + ) + 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( + [0, 1, 2, 3, 4], + [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) + 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) + remote = "oss://source-bucket/robot-videos" + source_file_io = _RemoteLeRobotFileIO(temp_dir, remote) + with patch( + "pypaimon.multimodal.lerobot.source._Hdf5SourceFileIO", + return_value=source_file_io, + ), patch( + "pypaimon.multimodal.lerobot.api." + "_import_lerobot_dataset", + return_value=object, + ): + remote_snapshot = connection.load_from_lerobot( + "remote_frames", remote, batch_size=1) + + self.assertEqual(1, remote_snapshot) + 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, + }) + 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( + "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") + ), + ) finally: shutil.rmtree(temp_dir, ignore_errors=True) + def test_video_timestamp_phase_must_match_frame_ordinals(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], type=pa.int64()), + "frame_index": pa.array([0, 1], type=pa.int64()), + "timestamp": pa.array([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) + + schema = _schema_from_info(info, include_task=False) + for from_timestamp in (0.05, 1000000.0005): + episode = { + "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": from_timestamp, + "videos/camera/to_timestamp": from_timestamp + 0.2, + } + with self.subTest(from_timestamp=from_timestamp), patch( + "pypaimon.multimodal.lerobot.loader._video_source", + return_value=("file:/video.mp4", 10), + ), self.assertRaisesRegex(ValueError, "not aligned"): + _read_batch( + Dataset(rows), info, 0, 2, schema, episode=episode, + video_sources={}) + + shifted_rows = rows.set_column( + rows.schema.get_field_index("timestamp"), + "timestamp", + pa.array([0.00009, 0.10009], type=pa.float32()), + ) + episode.update({ + "videos/camera/from_timestamp": 0.00009, + "videos/camera/to_timestamp": 0.20009, + }) + with patch( + "pypaimon.multimodal.lerobot.loader._video_source", + return_value=("file:/video.mp4", 10), + ), self.assertRaisesRegex(ValueError, "shifted timestamp"): + _read_batch( + Dataset(shifted_rows), info, 0, 2, schema, + episode=episode, video_sources={}) + class _RemoteLeRobotFileIO: @@ -509,7 +928,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) 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 8b2097741f21..f8da5c5b36ab 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) @@ -172,6 +173,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. @@ -186,6 +188,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( @@ -261,8 +264,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): @@ -315,7 +318,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: @@ -500,6 +503,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