From 0a3a923db421af4df43d7fa010d398e442da34f2 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Mon, 21 Sep 2026 23:08:56 -0700 Subject: [PATCH 1/8] [format][python] Persist keyframe indexes for efficient random video frame reads --- docs/docs/concepts/spec/fileformat.md | 22 +- docs/docs/multimodal-table/video.mdx | 37 +- docs/docs/pypaimon/lerobot.md | 10 +- docs/docs/pypaimon/video.md | 3 + .../paimon/data/VideoFrameDescriptor.java | 54 +- .../paimon/data/VideoFrameDescriptorTest.java | 48 +- .../paimon/data/video-frame-descriptor-v1.hex | 2 +- .../apache/paimon/append/BlobTableTest.java | 71 ++- .../paimon/format/blob/VideoFileMeta.java | 54 +- .../paimon/format/blob/VideoFormatReader.java | 6 +- .../paimon/format/blob/VideoFormatWriter.java | 57 +- .../format/blob/VideoKeyframeIndex.java | 130 +++++ .../format/blob/VideoFileFormatTest.java | 223 +++++++- .../apache/paimon/format/blob/video-v1.hex | 2 +- .../pypaimon/multimodal/lerobot/dataset.py | 209 +++++++- .../pypaimon/multimodal/lerobot/loader.py | 2 + paimon-python/pypaimon/multimodal/table.py | 19 +- paimon-python/pypaimon/multimodal/video.py | 52 +- .../read/reader/video_format_reader.py | 67 ++- paimon-python/pypaimon/table/row/blob.py | 59 +- .../table/row/video_keyframe_index.py | 295 ++++++++++ paimon-python/pypaimon/tests/blob_test.py | 8 +- .../tests/data_evolution_row_rolling_test.py | 24 +- .../pypaimon/tests/multimodal_lerobot_test.py | 166 ++++-- .../pypaimon/tests/multimodal_table_test.py | 124 ++++- .../pypaimon/tests/multimodal_video_test.py | 6 +- .../pypaimon/tests/video_format_test.py | 507 ++++++++++++++++-- .../tests/video_keyframe_index_test.py | 142 +++++ .../pypaimon/write/video_format_writer.py | 90 +++- 29 files changed, 2235 insertions(+), 254 deletions(-) create mode 100644 paimon-format/src/main/java/org/apache/paimon/format/blob/VideoKeyframeIndex.java create mode 100644 paimon-python/pypaimon/table/row/video_keyframe_index.py create mode 100644 paimon-python/pypaimon/tests/video_keyframe_index_test.py diff --git a/docs/docs/concepts/spec/fileformat.md b/docs/docs/concepts/spec/fileformat.md index 5ece594d1e85..4d409a78d308 100644 --- a/docs/docs/concepts/spec/fileformat.md +++ b/docs/docs/concepts/spec/fileformat.md @@ -413,8 +413,14 @@ the ordinary BLOB entry header, length trailer, or per-entry CRC: +----------------------------+ | ... | +----------------------------+ +| Keyframe Index 1 | Video metadata ranges and compressed keyframe entries ++----------------------------+ +| Keyframe Index 2 | ++----------------------------+ | Physical Length Index | Delta-Varint video lengths +----------------------------+ +| Keyframe-Index Length Index | Delta-Varint keyframe-index lengths ++----------------------------+ | Run Length Index | Delta-Varint logical row counts +----------------------------+ | Run Reference Index | Delta-Varint physical video ordinals @@ -422,6 +428,7 @@ the ordinary BLOB entry header, length trailer, or per-entry CRC: | Run First-Frame Index | Delta-Varint frame ordinals +----------------------------+ | Physical Index Length | 4 bytes (Little Endian) +| Keyframe Length-Index Size | 4 bytes (Little Endian) | Run-Length Index Length | 4 bytes (Little Endian) | Run-Reference Index Length | 4 bytes (Little Endian) | First-Frame Index Length | 4 bytes (Little Endian) @@ -434,7 +441,18 @@ The run arrays have equal element counts. A non-negative run reference is an ord physical length index. For logical row `r` in a run beginning at logical row `s`, the returned `VideoFrameDescriptor` identifies the referenced raw video range and frame ordinal `run_first_frame + (r - s)`. `-1` is a NULL run and `-2` is a data-evolution placeholder run. -Non-negative runs have fixed frame stride one in version 1; a discontinuity starts another run. +Non-negative runs have fixed frame stride one; a discontinuity starts another run. Each physical +video may store one sparse keyframe index; a zero length selects the scan fallback. + +A keyframe-index block has a 17-byte header: version (`1`, uint8), magic (`0x564944454F4B4649`, uint64), +metadata-range count (uint32), and keyframe count (uint32). It then stores metadata `(offset, +length)` pairs (two int64 values) and zlib-compressed `(frame ordinal, PTS, packet position)` +keyframe entries (three int64 values). Offsets are relative to the encoded video; writers reject +out-of-range values. One block is limited to 16 MiB and all blocks in one file to 64 MiB. + +The index covers the first video stream. Its time base remains in the video. A reader fetches the +metadata and target GOP, seeks by PTS, and decodes forward in presentation order. It may include the +following GOP for reordered frames. The serialized `VideoFrameDescriptor` stored in an Arrow/data-file cell has its own versioned wire layout. All numeric values are little endian: @@ -448,6 +466,8 @@ wire layout. All numeric values are little endian: | Offset | 8 bytes | Start of the complete encoded-video payload | | Length | 8 bytes | Encoded-video payload length | | Frame index | 8 bytes | Zero-based presentation-order frame ordinal | +| Keyframe-index offset | 8 bytes | Index offset in the `.video` file, or `-1` | +| Keyframe-index length | 8 bytes | Index length, or `0` | Descriptor bytes are independently versioned from the `.video` container. Java and Python share canonical descriptor and container fixtures to keep both implementations byte-compatible. diff --git a/docs/docs/multimodal-table/video.mdx b/docs/docs/multimodal-table/video.mdx index 417ba15a2977..d8833ee48fb0 100644 --- a/docs/docs/multimodal-table/video.mdx +++ b/docs/docs/multimodal-table/video.mdx @@ -81,12 +81,28 @@ to write complete videos and their logical frame rows. ## Descriptors and Physical Layout -Each logical value is a `VideoFrameDescriptor`: its URI range identifies one complete encoded -video and its frame ordinal selects a frame inside that video. On write, a `.video` data region -concatenates raw video payloads without ordinary BLOB entry wrappers. Four embedded delta-varint -indexes record physical video lengths, logical run lengths, run-to-video references, and the first -frame ordinal of each run. Consecutive frames therefore need one run entry rather than one index -entry per row. NULL and data-evolution placeholders use negative run references. +Each `VideoFrameDescriptor` identifies an encoded video range and a frame ordinal. A `.video` file +packs raw videos, optional keyframe indexes, and delta-varint row mappings. Consecutive frames share +one run; negative references represent NULL or data-evolution placeholders. + +
+.video (version 1, in file order)
+├─ Complete encoded videos A, B, ...
+├─ Keyframe index blocks A, B, ...
+│  ├─ Video metadata ranges: offset and length within the encoded video
+│  └─ Compressed keyframe entries: frame ordinal, PTS, and packet byte position
+├─ Metadata indexes
+│  ├─ Video payload lengths
+│  ├─ Keyframe index block lengths
+│  └─ Run lengths, video references, and starting frame indices
+└─ Footer: five 4-byte index sizes, 4-byte magic, and 1-byte version
+
+ +Each video index stores initialization ranges such as MP4 `moov` and keyframe positions. A cold +reader fetches these ranges and a bounded GOP window around the target. Zero length selects the +scan fallback. PyPaimon generates indexes for supported ISO BMFF videos and preserves them during +rewrites. See the +[file format specification](../concepts/spec/fileformat#video). ## Reuse, Rolling, and Compaction @@ -101,10 +117,9 @@ BLOB, and vector rolling is deferred until the next payload boundary and all act together. The resulting file group remains row-aligned and a single large episode may exceed the configured target. -When BLOB compaction is enabled, it byte-copies the complete encoded-video ranges into a new self-contained `.video` pack -and rebuilds the embedded indexes; it does not decode or re-encode frames. The aligned normal data -file contains only application columns such as `episode_id`, state, and action. Paimon stores the -frame mapping in the `.video` descriptor/index path. +When BLOB compaction is enabled, it copies videos and keyframe indexes into a new self-contained +`.video` pack without decoding frames. The aligned normal data file contains only application +columns such as `episode_id`, state, and action. `blob-compaction.enabled` defaults to `false`; ordinary normal-file compaction can leave the video packs unchanged. See [Data Evolution Maintenance](./data-evolution-maintenance#ordinary-compaction). @@ -115,8 +130,6 @@ leave the video packs unchanged. See [Data Evolution Maintenance](./data-evoluti `.blob`. - Non-null writes must be exact descriptor-backed `BlobRef` values containing a `VideoFrameDescriptor`. Inline bytes and ordinary `BlobDescriptor` values are rejected. -- Version 1 addresses frames by zero-based presentation-order ordinal with stride one. It does not - store PTS values or parse codec/container metadata. - A `BlobConsumer` callback is not supported for the video field. ## Read Frames diff --git a/docs/docs/pypaimon/lerobot.md b/docs/docs/pypaimon/lerobot.md index 99110cdb5d7a..34db12bf63d6 100644 --- a/docs/docs/pypaimon/lerobot.md +++ b/docs/docs/pypaimon/lerobot.md @@ -107,9 +107,8 @@ together, and keep writers paused until the call returns. Scalars map to scalar types, vectors to `VECTOR`, higher-rank tensors to nested `ARRAY`, and images to `BLOB`. Images keep their compressed bytes. -Video features map to `BLOB`. Frame rows reference MP4 payloads copied once per -aligned file group. Video imports use the video grouping policy and check -rolling before each Episode. They require a bucket-unaware table. Use +Video features map to `BLOB`. Each MP4 is copied once per aligned file group and indexed for range +reads. Imports require a bucket-unaware table and check rolling before each Episode. Use `VideoFrameCollator` for scans or `PaimonLeRobotDataset` for training. ## Capture LeRobot frames directly into Paimon @@ -223,9 +222,8 @@ loader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=4) ``` Without `tag_name`, the latest snapshots are used. Frame lookups use the BTree -on `index`; payloads remain lazy. Video decoding prefers TorchCodec, falls back -to PyAV, and reuses a bounded decoder cache. Set `video_backend` to force -either decoder. +on `index`; payloads remain lazy. Indexed videos use PyAV range reads for metadata and target GOPs; +unindexed videos use the TorchCodec/PyAV scan path. Set `video_backend` to force either decoder. Subclass `PaimonDatasetReader` for a custom logical frame layout: diff --git a/docs/docs/pypaimon/video.md b/docs/docs/pypaimon/video.md index 881ce62fdef7..7cbd3762dd74 100644 --- a/docs/docs/pypaimon/video.md +++ b/docs/docs/pypaimon/video.md @@ -88,6 +88,9 @@ video = pm.BlobDescriptor( frames.add_video(video, episode_43_rows) ``` +With PyAV, PyPaimon indexes supported MP4 files for range reads; other videos +use the scan fallback. + The writer deduplicates exact payload descriptor identity inside each `.video` file. Its video grouping policy coordinates normal, BLOB, and vector rolling at payload boundaries. A file may exceed its target before the next boundary. diff --git a/paimon-common/src/main/java/org/apache/paimon/data/VideoFrameDescriptor.java b/paimon-common/src/main/java/org/apache/paimon/data/VideoFrameDescriptor.java index 45b995db97a5..4f15d6d1c1c4 100644 --- a/paimon-common/src/main/java/org/apache/paimon/data/VideoFrameDescriptor.java +++ b/paimon-common/src/main/java/org/apache/paimon/data/VideoFrameDescriptor.java @@ -37,15 +37,31 @@ public class VideoFrameDescriptor extends BlobDescriptor { private static final long MAGIC = 0x564944454F46524DL; // "VIDEOFRM" private static final byte CURRENT_VERSION = 1; private static final int FIXED_LENGTH = - Byte.BYTES + Long.BYTES + Integer.BYTES + 3 * Long.BYTES; + Byte.BYTES + Long.BYTES + Integer.BYTES + 5 * Long.BYTES; private final long frameIndex; - - public VideoFrameDescriptor(String uri, long offset, long length, long frameIndex) { + private final long keyframeIndexOffset; + private final long keyframeIndexLength; + + public VideoFrameDescriptor( + String uri, + long offset, + long length, + long frameIndex, + long keyframeIndexOffset, + long keyframeIndexLength) { super(uri, offset, length); checkArgument( frameIndex >= 0, "Video frame index must be non-negative, but was %s.", frameIndex); + checkArgument( + keyframeIndexLength >= 0, "Video keyframe index length must be non-negative."); + checkArgument( + (keyframeIndexLength == 0 && keyframeIndexOffset == -1) + || (keyframeIndexLength > 0 && keyframeIndexOffset >= 0), + "Invalid video keyframe index range."); this.frameIndex = frameIndex; + this.keyframeIndexOffset = keyframeIndexOffset; + this.keyframeIndexLength = keyframeIndexLength; } public long frameIndex() { @@ -57,6 +73,22 @@ public BlobDescriptor payloadDescriptor() { return new BlobDescriptor(uri(), offset(), length()); } + public @Nullable BlobDescriptor keyframeIndexDescriptor() { + return keyframeIndexLength == 0 + ? null + : new BlobDescriptor(uri(), keyframeIndexOffset, keyframeIndexLength); + } + + /** Returns the persisted keyframe index carried by an exact frame reference. */ + public static @Nullable Blob keyframeIndexBlob(Blob blob) { + VideoFrameDescriptor frame = fromBlob(blob); + BlobDescriptor mapping = frame == null ? null : frame.keyframeIndexDescriptor(); + if (mapping == null) { + return null; + } + return Blob.fromDescriptor(((BlobRef) blob).uriReader(), mapping); + } + /** Returns the video frame carried by an exact lazy blob reference, or {@code null}. */ public static @Nullable VideoFrameDescriptor fromBlob(@Nullable Blob blob) { if (blob == null || blob.getClass() != BlobRef.class) { @@ -86,6 +118,8 @@ public byte[] serialize() { buffer.putLong(offset()); buffer.putLong(length()); buffer.putLong(frameIndex); + buffer.putLong(keyframeIndexOffset); + buffer.putLong(keyframeIndexLength); return buffer.array(); } @@ -109,7 +143,7 @@ public static VideoFrameDescriptor deserialize(byte[] bytes) { throw invalidPayload("missing magic header"); } int uriLength = buffer.getInt(); - // checked by comparison and subtraction: uriLength + 3 * Long.BYTES wraps negative + // checked by comparison and subtraction: uriLength + 5 * Long.BYTES wraps negative // for a uriLength near Integer.MAX_VALUE if (uriLength < 0) { throw invalidPayload("negative URI length: " + uriLength); @@ -117,7 +151,7 @@ public static VideoFrameDescriptor deserialize(byte[] bytes) { if (uriLength > buffer.remaining()) { throw invalidPayload("URI length exceeds data size"); } - if (buffer.remaining() - uriLength < 3 * Long.BYTES) { + if (buffer.remaining() - uriLength < 5 * Long.BYTES) { throw invalidPayload("missing offset/length/frame index"); } @@ -127,13 +161,16 @@ public static VideoFrameDescriptor deserialize(byte[] bytes) { long offset = buffer.getLong(); long length = buffer.getLong(); long frameIndex = buffer.getLong(); + long keyframeIndexOffset = buffer.getLong(); + long keyframeIndexLength = buffer.getLong(); if (buffer.hasRemaining()) { throw invalidPayload("trailing bytes"); } if (frameIndex < 0) { throw invalidPayload("negative frame index: " + frameIndex); } - return new VideoFrameDescriptor(uri, offset, length, frameIndex); + return new VideoFrameDescriptor( + uri, offset, length, frameIndex, keyframeIndexOffset, keyframeIndexLength); } public static boolean isVideoFrameDescriptor(byte[] bytes) { @@ -154,12 +191,13 @@ public boolean equals(Object o) { } VideoFrameDescriptor that = (VideoFrameDescriptor) o; return frameIndex == that.frameIndex - && payloadDescriptor().equals(that.payloadDescriptor()); + && payloadDescriptor().equals(that.payloadDescriptor()) + && Objects.equals(keyframeIndexDescriptor(), that.keyframeIndexDescriptor()); } @Override public int hashCode() { - return Objects.hash(payloadDescriptor(), frameIndex); + return Objects.hash(payloadDescriptor(), frameIndex, keyframeIndexDescriptor()); } @Override diff --git a/paimon-common/src/test/java/org/apache/paimon/data/VideoFrameDescriptorTest.java b/paimon-common/src/test/java/org/apache/paimon/data/VideoFrameDescriptorTest.java index ea452559f14f..07d9be6e3f92 100644 --- a/paimon-common/src/test/java/org/apache/paimon/data/VideoFrameDescriptorTest.java +++ b/paimon-common/src/test/java/org/apache/paimon/data/VideoFrameDescriptorTest.java @@ -36,7 +36,7 @@ public class VideoFrameDescriptorTest { @Test public void testRoundTripAndPayloadIdentity() { VideoFrameDescriptor frame = - new VideoFrameDescriptor("oss://bucket/source.mp4", 17, 103, 42); + new VideoFrameDescriptor("oss://bucket/source.mp4", 17, 103, 42, -1, 0); assertThat(VideoFrameDescriptor.isVideoFrameDescriptor(frame.serialize())).isTrue(); assertThat(BlobDescriptor.isBlobDescriptor(frame.serialize())).isFalse(); @@ -46,26 +46,22 @@ public void testRoundTripAndPayloadIdentity() { .isEqualTo(new BlobDescriptor("oss://bucket/source.mp4", 17, 103)); VideoFrameDescriptor next = - new VideoFrameDescriptor("oss://bucket/source.mp4", 17, 103, 43); + new VideoFrameDescriptor("oss://bucket/source.mp4", 17, 103, 43, -1, 0); assertThat(next).isNotEqualTo(frame); assertThat(next.payloadDescriptor()).isEqualTo(frame.payloadDescriptor()); + + VideoFrameDescriptor indexed = + new VideoFrameDescriptor("oss://bucket/source.mp4", 17, 103, 42, 120, 8); + assertThat(VideoFrameDescriptor.deserialize(indexed.serialize())).isEqualTo(indexed); + assertThat(indexed.keyframeIndexDescriptor()) + .isEqualTo(new BlobDescriptor("oss://bucket/source.mp4", 120, 8)); } @Test public void testCrossLanguageWireFixture() throws Exception { - VideoFrameDescriptor expected = new VideoFrameDescriptor("s3://bucket/视频.mp4", 7, 99, 42); - byte[] fixture = - fromHex( - new String( - IOUtils.readFully( - VideoFrameDescriptorTest.class - .getClassLoader() - .getResourceAsStream( - "org/apache/paimon/data/video-frame-descriptor-v1.hex"), - true), - StandardCharsets.UTF_8) - .trim()); - + VideoFrameDescriptor expected = + new VideoFrameDescriptor("s3://bucket/视频.mp4", 7, 99, 42, 106, 8); + byte[] fixture = fixture("video-frame-descriptor-v1.hex"); assertThat(expected.serialize()).isEqualTo(fixture); assertThat(BlobDescriptor.deserialize(fixture)).isEqualTo(expected); assertThat(BlobDescriptor.isSerializedDescriptor(fixture)).isTrue(); @@ -73,7 +69,7 @@ public void testCrossLanguageWireFixture() throws Exception { @Test public void testBlobFromBytesPreservesFrameDescriptor() { - VideoFrameDescriptor expected = new VideoFrameDescriptor("file:/video.mp4", 0, 9, 7); + VideoFrameDescriptor expected = new VideoFrameDescriptor("file:/video.mp4", 0, 9, 7, -1, 0); Blob blob = Blob.fromBytes(expected.serialize(), null, null); assertThat(blob).isInstanceOf(BlobRef.class); @@ -82,17 +78,18 @@ public void testBlobFromBytesPreservesFrameDescriptor() { @Test public void testRejectInvalidPayload() { - VideoFrameDescriptor descriptor = new VideoFrameDescriptor("file:/video.mp4", 0, 9, 7); + VideoFrameDescriptor descriptor = + new VideoFrameDescriptor("file:/video.mp4", 0, 9, 7, -1, 0); byte[] trailing = Arrays.copyOf(descriptor.serialize(), descriptor.serialize().length + 1); assertThatThrownBy(() -> VideoFrameDescriptor.deserialize(trailing)) .isInstanceOf(IllegalArgumentException.class) .hasMessageContaining("trailing bytes"); - assertThatThrownBy(() -> new VideoFrameDescriptor("file:/video.mp4", 0, 9, -1)) + assertThatThrownBy(() -> new VideoFrameDescriptor("file:/video.mp4", 0, 9, -1, -1, 0)) .isInstanceOf(IllegalArgumentException.class) .hasMessageContaining("non-negative"); - // The old check let this through because uriLength + 3 * Long.BYTES wrapped + // The old check let this through because uriLength + 5 * Long.BYTES wrapped // negative, and deserialize went on to allocate ~2GB. byte[] hostileUriLength = descriptor.serialize(); ByteBuffer.wrap(hostileUriLength) @@ -129,4 +126,17 @@ private static byte[] fromHex(String hex) { } return bytes; } + + private static byte[] fixture(String name) throws Exception { + return fromHex( + new String( + IOUtils.readFully( + VideoFrameDescriptorTest.class + .getClassLoader() + .getResourceAsStream( + "org/apache/paimon/data/" + name), + true), + StandardCharsets.UTF_8) + .trim()); + } } diff --git a/paimon-common/src/test/resources/org/apache/paimon/data/video-frame-descriptor-v1.hex b/paimon-common/src/test/resources/org/apache/paimon/data/video-frame-descriptor-v1.hex index 9b45e36167b4..db91bd4e0e8a 100644 --- a/paimon-common/src/test/resources/org/apache/paimon/data/video-frame-descriptor-v1.hex +++ b/paimon-common/src/test/resources/org/apache/paimon/data/video-frame-descriptor-v1.hex @@ -15,4 +15,4 @@ # specific language governing permissions and limitations # under the License. -014d52464f454449561600000073333a2f2f6275636b65742fe8a786e9a2912e6d7034070000000000000063000000000000002a00000000000000 +014d52464f454449561600000073333a2f2f6275636b65742fe8a786e9a2912e6d7034070000000000000063000000000000002a000000000000006a000000000000000800000000000000 diff --git a/paimon-core/src/test/java/org/apache/paimon/append/BlobTableTest.java b/paimon-core/src/test/java/org/apache/paimon/append/BlobTableTest.java index ed69827b9311..4660b0f998aa 100644 --- a/paimon-core/src/test/java/org/apache/paimon/append/BlobTableTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/append/BlobTableTest.java @@ -1561,31 +1561,31 @@ public void testVideoRollingAndCompaction() throws Exception { Blob.fromDescriptor( sourceReader, new VideoFrameDescriptor( - firstUri, 0, firstBytes.length, 0))), + firstUri, 0, firstBytes.length, 0, -1, 0))), GenericRow.of( 1, Blob.fromDescriptor( sourceReader, new VideoFrameDescriptor( - firstUri, 0, firstBytes.length, 1))), + firstUri, 0, firstBytes.length, 1, -1, 0))), GenericRow.of( 2, Blob.fromDescriptor( sourceReader, new VideoFrameDescriptor( - firstUri, 0, firstBytes.length, 2))), + firstUri, 0, firstBytes.length, 2, -1, 0))), GenericRow.of( 3, Blob.fromDescriptor( sourceReader, new VideoFrameDescriptor( - secondUri, 0, secondBytes.length, 0))), + secondUri, 0, secondBytes.length, 0, -1, 0))), GenericRow.of( 4, Blob.fromDescriptor( sourceReader, new VideoFrameDescriptor( - secondUri, 0, secondBytes.length, 1))))); + secondUri, 0, secondBytes.length, 1, -1, 0))))); FileStoreTable table = getTableDefault(); List videoFiles = liveVideoFiles(table); @@ -1651,7 +1651,7 @@ public void testVideoRollingByBlobTargetSize() throws Exception { schemaBuilder.column("id", DataTypes.INT()); schemaBuilder.column("video", DataTypes.BLOB()); schemaBuilder.option(CoreOptions.TARGET_FILE_SIZE.key(), "1 GB"); - schemaBuilder.option(CoreOptions.BLOB_TARGET_FILE_SIZE.key(), "1 b"); + schemaBuilder.option(CoreOptions.BLOB_TARGET_FILE_SIZE.key(), "20 b"); schemaBuilder.option(CoreOptions.TARGET_FILE_ROW_NUM.key(), "1000"); schemaBuilder.option(CoreOptions.ROW_TRACKING_ENABLED.key(), "true"); schemaBuilder.option(CoreOptions.DATA_EVOLUTION_ENABLED.key(), "true"); @@ -1660,10 +1660,14 @@ public void testVideoRollingByBlobTargetSize() throws Exception { byte[] firstBytes = "first-video".getBytes(); byte[] secondBytes = "second-video".getBytes(); + byte[] keyframeIndex = + fromHex( + "0149464b4f45444956010000000200000000000000000000000100000000000000" + + "789c6360c00e78a0f4821e08cd04e503001394013b"); java.nio.file.Path firstSource = tempPath.resolve("size-first-source.mp4"); java.nio.file.Path secondSource = tempPath.resolve("size-second-source.mp4"); - java.nio.file.Files.write(firstSource, firstBytes); - java.nio.file.Files.write(secondSource, secondBytes); + java.nio.file.Files.write(firstSource, concat("first-video".getBytes(), keyframeIndex)); + java.nio.file.Files.write(secondSource, concat("second-video".getBytes(), keyframeIndex)); UriReader sourceReader = UriReader.fromFile(LocalFileIO.create()); String firstUri = new Path(firstSource.toUri()).toString(); String secondUri = new Path(secondSource.toUri()).toString(); @@ -1676,7 +1680,12 @@ public void testVideoRollingByBlobTargetSize() throws Exception { Blob.fromDescriptor( sourceReader, new VideoFrameDescriptor( - firstUri, 0, firstBytes.length, frame)))); + firstUri, + 0, + firstBytes.length, + frame, + firstBytes.length, + keyframeIndex.length)))); } for (int frame = 0; frame < 2; frame++) { rows.add( @@ -1685,7 +1694,12 @@ public void testVideoRollingByBlobTargetSize() throws Exception { Blob.fromDescriptor( sourceReader, new VideoFrameDescriptor( - secondUri, 0, secondBytes.length, frame)))); + secondUri, + 0, + secondBytes.length, + frame, + secondBytes.length, + keyframeIndex.length)))); } writeRows(getTableDefault(), rows); @@ -1739,11 +1753,21 @@ public void testMultipleVideoFieldsRollAtAlignedEpisodeBoundaries() throws Excep Blob.fromDescriptor( sourceReader, new VideoFrameDescriptor( - uris[cameraA], 0, payloads[cameraA].length, frame)), + uris[cameraA], + 0, + payloads[cameraA].length, + frame, + -1, + 0)), Blob.fromDescriptor( sourceReader, new VideoFrameDescriptor( - uris[cameraB], 0, payloads[cameraB].length, frame)))); + uris[cameraB], + 0, + payloads[cameraB].length, + frame, + -1, + 0)))); } writeRows(getTableDefault(), input); @@ -1807,11 +1831,12 @@ public void testMultipleVideoFieldsAllowNestedEpisodeBoundaries() throws Excepti row, Blob.fromDescriptor( sourceReader, - new VideoFrameDescriptor(uris[0], 0, payload.length, row)), + new VideoFrameDescriptor( + uris[0], 0, payload.length, row, -1, 0)), Blob.fromDescriptor( sourceReader, new VideoFrameDescriptor( - uris[cameraB], 0, payload.length, row % 2)))); + uris[cameraB], 0, payload.length, row % 2, -1, 0)))); } writeRows(getTableDefault(), rows); @@ -2453,6 +2478,24 @@ private static void writeFile(FileIO fileIO, Path path, byte[] bytes) throws IOE } } + private static byte[] concat(byte[] first, byte[] second) { + byte[] result = Arrays.copyOf(first, first.length + second.length); + System.arraycopy(second, 0, result, first.length, second.length); + return result; + } + + private static byte[] fromHex(String hex) { + byte[] bytes = new byte[hex.length() / 2]; + for (int i = 0; i < bytes.length; i++) { + int offset = i * 2; + bytes[i] = + (byte) + ((Character.digit(hex.charAt(offset), 16) << 4) + + Character.digit(hex.charAt(offset + 1), 16)); + } + return bytes; + } + private static long countFilesWithSuffix(FileIO fileIO, Path root, String suffix) throws IOException { long count = 0; diff --git a/paimon-format/src/main/java/org/apache/paimon/format/blob/VideoFileMeta.java b/paimon-format/src/main/java/org/apache/paimon/format/blob/VideoFileMeta.java index 4aca43304290..6fceaa83f245 100644 --- a/paimon-format/src/main/java/org/apache/paimon/format/blob/VideoFileMeta.java +++ b/paimon-format/src/main/java/org/apache/paimon/format/blob/VideoFileMeta.java @@ -35,6 +35,8 @@ public class VideoFileMeta { private final long[] physicalVideoLengths; private final long[] physicalVideoOffsets; + private final long[] keyframeIndexLengths; + private final long[] keyframeIndexOffsets; private final long[] runEnds; private final long[] runReferences; private final long[] runFirstFrames; @@ -54,11 +56,12 @@ public VideoFileMeta(SeekableInputStream in, long fileSize, @Nullable RoaringBit byte[] footer = new byte[VideoFormatWriter.FILE_FOOTER_LENGTH]; IOUtils.readFully(in, footer); int physicalIndexLength = BytesUtils.getInt(footer, 0); - int runLengthIndexLength = BytesUtils.getInt(footer, Integer.BYTES); - int runReferenceIndexLength = BytesUtils.getInt(footer, Integer.BYTES * 2); - int firstFrameIndexLength = BytesUtils.getInt(footer, Integer.BYTES * 3); - int magic = BytesUtils.getInt(footer, Integer.BYTES * 4); - byte version = footer[Integer.BYTES * 5]; + int keyframeLengthIndexLength = BytesUtils.getInt(footer, Integer.BYTES); + int runLengthIndexLength = BytesUtils.getInt(footer, Integer.BYTES * 2); + int runReferenceIndexLength = BytesUtils.getInt(footer, Integer.BYTES * 3); + int firstFrameIndexLength = BytesUtils.getInt(footer, Integer.BYTES * 4); + int magic = BytesUtils.getInt(footer, Integer.BYTES * 5); + byte version = footer[Integer.BYTES * 6]; if (magic != VideoFormatWriter.MAGIC_NUMBER) { throw corrupt("invalid footer magic %s.", magic); } @@ -68,6 +71,7 @@ public VideoFileMeta(SeekableInputStream in, long fileSize, @Nullable RoaringBit int[] indexLengths = { physicalIndexLength, + keyframeLengthIndexLength, runLengthIndexLength, runReferenceIndexLength, firstFrameIndexLength @@ -87,26 +91,50 @@ public VideoFileMeta(SeekableInputStream in, long fileSize, @Nullable RoaringBit long offset = indexStart; long[] physicalVideoLengths = readIndex(in, offset, physicalIndexLength, "physical video"); offset += physicalIndexLength; + long[] keyframeIndexLengths = + readIndex(in, offset, keyframeLengthIndexLength, "keyframe index"); + offset += keyframeLengthIndexLength; + if (keyframeIndexLengths.length != physicalVideoLengths.length) { + throw corrupt("physical video and keyframe index indexes have different counts."); + } long[] runLengths = readIndex(in, offset, runLengthIndexLength, "run length"); offset += runLengthIndexLength; long[] runReferences = readIndex(in, offset, runReferenceIndexLength, "run reference"); offset += runReferenceIndexLength; long[] runFirstFrames = readIndex(in, offset, firstFrameIndexLength, "run first-frame"); + long keyframeIndexSize = 0; + for (long length : keyframeIndexLengths) { + if (length < 0 || keyframeIndexSize > Long.MAX_VALUE - length) { + throw corrupt("invalid keyframe index length %s.", length); + } + keyframeIndexSize += length; + } + long keyframeIndexStart = indexStart - keyframeIndexSize; + if (keyframeIndexStart < 0) { + throw corrupt("keyframe indexes exceed the file size."); + } + long[] physicalVideoOffsets = new long[physicalVideoLengths.length]; long payloadOffset = 0; for (int i = 0; i < physicalVideoLengths.length; i++) { long length = physicalVideoLengths[i]; - if (length <= 0 || length > indexStart - payloadOffset) { + if (length <= 0 || length > keyframeIndexStart - payloadOffset) { throw corrupt("invalid physical video length %s at ordinal %s.", length, i); } physicalVideoOffsets[i] = payloadOffset; payloadOffset += length; } - if (payloadOffset != indexStart) { + if (payloadOffset != keyframeIndexStart) { throw corrupt( "indexed videos use %s bytes, but payload region contains %s bytes.", - payloadOffset, indexStart); + payloadOffset, keyframeIndexStart); + } + long[] keyframeIndexOffsets = new long[keyframeIndexLengths.length]; + long keyframeIndexOffset = keyframeIndexStart; + for (int i = 0; i < keyframeIndexLengths.length; i++) { + keyframeIndexOffsets[i] = keyframeIndexOffset; + keyframeIndexOffset += keyframeIndexLengths[i]; } if (runLengths.length != runReferences.length @@ -162,6 +190,8 @@ public VideoFileMeta(SeekableInputStream in, long fileSize, @Nullable RoaringBit this.physicalVideoLengths = physicalVideoLengths; this.physicalVideoOffsets = physicalVideoOffsets; + this.keyframeIndexLengths = keyframeIndexLengths; + this.keyframeIndexOffsets = keyframeIndexOffsets; this.runEnds = runEnds; this.runReferences = runReferences; this.runFirstFrames = runFirstFrames; @@ -192,6 +222,14 @@ public long frameIndex(int returnedRow) { return runFirstFrames[run] + row - runStart; } + public long keyframeIndexOffset(int returnedRow) { + return keyframeIndexOffsets[physicalOrdinal(returnedRow)]; + } + + public long keyframeIndexLength(int returnedRow) { + return keyframeIndexLengths[physicalOrdinal(returnedRow)]; + } + public int returnedPosition(int currentPosition) { return logicalPosition(currentPosition - 1); } diff --git a/paimon-format/src/main/java/org/apache/paimon/format/blob/VideoFormatReader.java b/paimon-format/src/main/java/org/apache/paimon/format/blob/VideoFormatReader.java index 91c3511ef08d..f6f3463a9bd7 100644 --- a/paimon-format/src/main/java/org/apache/paimon/format/blob/VideoFormatReader.java +++ b/paimon-format/src/main/java/org/apache/paimon/format/blob/VideoFormatReader.java @@ -89,7 +89,11 @@ public InternalRow next() { filePath.toString(), fileMeta.videoOffset(currentPosition), fileMeta.videoLength(currentPosition), - fileMeta.frameIndex(currentPosition)); + fileMeta.frameIndex(currentPosition), + fileMeta.keyframeIndexLength(currentPosition) == 0 + ? -1 + : fileMeta.keyframeIndexOffset(currentPosition), + fileMeta.keyframeIndexLength(currentPosition)); field = Blob.fromDescriptor(uriReader, descriptor); } currentPosition++; diff --git a/paimon-format/src/main/java/org/apache/paimon/format/blob/VideoFormatWriter.java b/paimon-format/src/main/java/org/apache/paimon/format/blob/VideoFormatWriter.java index 4bc3602b69c6..d868e808a2da 100644 --- a/paimon-format/src/main/java/org/apache/paimon/format/blob/VideoFormatWriter.java +++ b/paimon-format/src/main/java/org/apache/paimon/format/blob/VideoFormatWriter.java @@ -33,8 +33,11 @@ import org.apache.paimon.utils.LongArrayList; import java.io.IOException; +import java.util.ArrayList; import java.util.HashMap; +import java.util.List; import java.util.Map; +import java.util.Objects; import static org.apache.paimon.utils.Preconditions.checkArgument; import static org.apache.paimon.utils.StreamUtils.intToLittleEndian; @@ -52,20 +55,25 @@ public class VideoFormatWriter implements FileAwareFormatWriter { public static final int MAGIC_NUMBER = 0x4F454449; // "IDEO" in little endian public static final long NULL_REFERENCE = -1L; public static final long PLACEHOLDER_REFERENCE = -2L; - public static final int FILE_FOOTER_LENGTH = Integer.BYTES * 5 + Byte.BYTES; + public static final int FILE_FOOTER_LENGTH = Integer.BYTES * 6 + Byte.BYTES; + static final long MAX_KEYFRAME_INDEX_BYTES = 16L * 1024 * 1024; + static final long MAX_TOTAL_KEYFRAME_INDEX_BYTES = 64L * 1024 * 1024; private final PositionOutputStream out; private final RawVideoPayloadWriter payloadWriter; private final LongArrayList physicalVideoLengths; + private final List keyframeIndexes; private final LongArrayList runLengths; private final LongArrayList runReferences; private final LongArrayList runFirstFrames; private final Map physicalVideos; + private final Map physicalVideoKeyframeIndexes; private long currentRunLength; private long currentRunReference; private long currentRunFirstFrame; private long currentRunLastFrame; + private long keyframeIndexBytes; private boolean closed; public VideoFormatWriter( @@ -86,10 +94,12 @@ public VideoFormatWriter( blobFetchMetricReporter, copyBufferSize); this.physicalVideoLengths = new LongArrayList(16); + this.keyframeIndexes = new ArrayList<>(); this.runLengths = new LongArrayList(16); this.runReferences = new LongArrayList(16); this.runFirstFrames = new LongArrayList(16); this.physicalVideos = new HashMap<>(); + this.physicalVideoKeyframeIndexes = new HashMap<>(); } @Override @@ -121,8 +131,12 @@ public void addElement(InternalRow element) throws IOException { "Video fields require an exact BlobRef containing a VideoFrameDescriptor."); BlobDescriptor payload = frame.payloadDescriptor(); + BlobDescriptor keyframeIndexDescriptor = frame.keyframeIndexDescriptor(); Integer ordinal = physicalVideos.get(payload); if (ordinal == null) { + long keyframeIndexLength = + keyframeIndexDescriptor == null ? 0 : keyframeIndexDescriptor.length(); + checkKeyframeIndexSize(keyframeIndexLength, keyframeIndexBytes); long length = payloadWriter.write(element); if (length == BlobFormatWriter.NULL_LENGTH) { append(NULL_REFERENCE, 0); @@ -130,14 +144,32 @@ public void addElement(InternalRow element) throws IOException { } ordinal = physicalVideoLengths.size(); physicalVideoLengths.add(length); + Blob keyframeIndex = VideoFrameDescriptor.keyframeIndexBlob(blob); + byte[] mapping = keyframeIndex == null ? new byte[0] : keyframeIndex.toData(); + checkArgument( + mapping.length == keyframeIndexLength, + "Video keyframe index length changed while reading: expected %s, read %s.", + keyframeIndexLength, + mapping.length); + if (mapping.length > 0) { + VideoKeyframeIndex.validate(mapping, length); + } + keyframeIndexes.add(mapping); + keyframeIndexBytes += mapping.length; physicalVideos.put(payload, ordinal); + physicalVideoKeyframeIndexes.put(payload, keyframeIndexDescriptor); + } else { + checkArgument( + Objects.equals( + physicalVideoKeyframeIndexes.get(payload), keyframeIndexDescriptor), + "Video frames for the same payload must use the same keyframe index."); } append(ordinal, frame.frameIndex()); } @Override public boolean reachTargetSize(boolean suggestedCheck, long targetSize) throws IOException { - return out.getPos() >= targetSize; + return keyframeIndexBytes >= targetSize || out.getPos() >= targetSize - keyframeIndexBytes; } @Override @@ -148,15 +180,24 @@ public void close() throws IOException { flushRun(); payloadWriter.close(); + for (byte[] mapping : keyframeIndexes) { + out.write(mapping); + } byte[] physicalIndex = DeltaVarintCompressor.compressLongArrayList(physicalVideoLengths); + LongArrayList keyframeIndexLengths = new LongArrayList(keyframeIndexes.size()); + keyframeIndexes.forEach(mapping -> keyframeIndexLengths.add(mapping.length)); + byte[] keyframeLengthIndex = + DeltaVarintCompressor.compressLongArrayList(keyframeIndexLengths); byte[] runLengthIndex = DeltaVarintCompressor.compressLongArrayList(runLengths); byte[] runReferenceIndex = DeltaVarintCompressor.compressLongArrayList(runReferences); byte[] firstFrameIndex = DeltaVarintCompressor.compressLongArrayList(runFirstFrames); out.write(physicalIndex); + out.write(keyframeLengthIndex); out.write(runLengthIndex); out.write(runReferenceIndex); out.write(firstFrameIndex); out.write(intToLittleEndian(physicalIndex.length)); + out.write(intToLittleEndian(keyframeLengthIndex.length)); out.write(intToLittleEndian(runLengthIndex.length)); out.write(intToLittleEndian(runReferenceIndex.length)); out.write(intToLittleEndian(firstFrameIndex.length)); @@ -173,6 +214,18 @@ int runCount() { return runLengths.size() + (currentRunLength == 0 ? 0 : 1); } + static void checkKeyframeIndexSize(long length, long bufferedBytes) { + checkArgument( + length >= 0 && length <= MAX_KEYFRAME_INDEX_BYTES, + "Video keyframe index length %s exceeds the %s-byte limit.", + length, + MAX_KEYFRAME_INDEX_BYTES); + checkArgument( + bufferedBytes <= MAX_TOTAL_KEYFRAME_INDEX_BYTES - length, + "Buffered video keyframe indexes exceed the %s-byte limit.", + MAX_TOTAL_KEYFRAME_INDEX_BYTES); + } + private void append(long reference, long frameIndex) { if (canExtend(reference, frameIndex)) { currentRunLength++; diff --git a/paimon-format/src/main/java/org/apache/paimon/format/blob/VideoKeyframeIndex.java b/paimon-format/src/main/java/org/apache/paimon/format/blob/VideoKeyframeIndex.java new file mode 100644 index 000000000000..15c01adb9b0a --- /dev/null +++ b/paimon-format/src/main/java/org/apache/paimon/format/blob/VideoKeyframeIndex.java @@ -0,0 +1,130 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.paimon.format.blob; + +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.util.zip.DataFormatException; +import java.util.zip.Inflater; + +/** Validates persisted video metadata ranges and sparse keyframe entries. */ +final class VideoKeyframeIndex { + + private static final int VERSION = 1; + private static final long MAGIC = 0x564944454F4B4649L; // "VIDEOKFI" + private static final int HEADER_LENGTH = Byte.BYTES + Long.BYTES + Integer.BYTES * 2; + private static final int METADATA_RANGE_LENGTH = Long.BYTES * 2; + private static final int ENTRY_LENGTH = Long.BYTES * 3; + + private VideoKeyframeIndex() {} + + static void validate(byte[] data, long payloadLength) { + if (data == null || data.length <= HEADER_LENGTH) { + throw invalid("too short"); + } + if (payloadLength <= 0) { + throw invalid("empty video payload"); + } + + ByteBuffer header = ByteBuffer.wrap(data, 0, HEADER_LENGTH).order(ByteOrder.LITTLE_ENDIAN); + int version = Byte.toUnsignedInt(header.get()); + long magic = header.getLong(); + long metadataRangeCount = Integer.toUnsignedLong(header.getInt()); + long keyframeCount = Integer.toUnsignedLong(header.getInt()); + if (version != VERSION || magic != MAGIC) { + throw invalid("version or magic"); + } + if (keyframeCount == 0) { + throw invalid("empty"); + } + + long entriesOffset = HEADER_LENGTH + metadataRangeCount * METADATA_RANGE_LENGTH; + if (entriesOffset >= data.length) { + throw invalid("metadata ranges"); + } + ByteBuffer ranges = ByteBuffer.wrap(data).order(ByteOrder.LITTLE_ENDIAN); + ranges.position(HEADER_LENGTH); + long previousEnd = 0; + for (long i = 0; i < metadataRangeCount; i++) { + long offset = ranges.getLong(); + long length = ranges.getLong(); + if (offset < previousEnd || length <= 0 || offset > Long.MAX_VALUE - length) { + throw invalid("metadata ranges"); + } + previousEnd = offset + length; + if (previousEnd > payloadLength) { + throw invalid("metadata range outside video payload"); + } + } + + Inflater inflater = new Inflater(); + inflater.setInput(data, (int) entriesOffset, data.length - (int) entriesOffset); + byte[] entry = new byte[ENTRY_LENGTH]; + long previousOrdinal = -1; + long previousPts = 0; + long previousPacketPosition = -1; + try { + for (long i = 0; i < keyframeCount; i++) { + inflateEntry(inflater, entry); + ByteBuffer values = ByteBuffer.wrap(entry).order(ByteOrder.LITTLE_ENDIAN); + long ordinal = values.getLong(); + long pts = values.getLong(); + long packetPosition = values.getLong(); + if ((i == 0 && ordinal != 0) || ordinal <= previousOrdinal) { + throw invalid("keyframe ordinals must be strictly increasing from zero"); + } + if (i > 0 && pts <= previousPts) { + throw invalid("keyframe timestamps must be strictly increasing"); + } + if (packetPosition <= previousPacketPosition || packetPosition >= payloadLength) { + throw invalid( + "keyframe packet positions must be within the video payload and strictly increasing"); + } + previousOrdinal = ordinal; + previousPts = pts; + previousPacketPosition = packetPosition; + } + + if (inflater.inflate(entry, 0, 1) != 0 + || !inflater.finished() + || inflater.getRemaining() != 0) { + throw invalid("entry count"); + } + } catch (DataFormatException e) { + throw new IllegalArgumentException("Invalid video keyframe index payload.", e); + } finally { + inflater.end(); + } + } + + private static void inflateEntry(Inflater inflater, byte[] entry) throws DataFormatException { + int offset = 0; + while (offset < entry.length) { + int length = inflater.inflate(entry, offset, entry.length - offset); + if (length == 0) { + throw invalid("entries"); + } + offset += length; + } + } + + private static IllegalArgumentException invalid(String reason) { + return new IllegalArgumentException("Invalid video keyframe index: " + reason + '.'); + } +} diff --git a/paimon-format/src/test/java/org/apache/paimon/format/blob/VideoFileFormatTest.java b/paimon-format/src/test/java/org/apache/paimon/format/blob/VideoFileFormatTest.java index b2ece4a8b2ac..4d3f98e906e8 100644 --- a/paimon-format/src/test/java/org/apache/paimon/format/blob/VideoFileFormatTest.java +++ b/paimon-format/src/test/java/org/apache/paimon/format/blob/VideoFileFormatTest.java @@ -60,6 +60,10 @@ /** Tests for {@link VideoFileFormat}. */ public class VideoFileFormatTest { + private static final String KEYFRAME_INDEX_HEX = + "0149464b4f45444956010000000200000000000000000000000100000000000000" + + "789c6360c00e78a0f4821e08cd04e503001394013b"; + @TempDir java.nio.file.Path tempPath; private FileIO fileIO; @@ -122,23 +126,47 @@ public void testPackRawVideosAndMapFrameRuns() throws IOException { @Test public void testCrossLanguageV1Fixture() throws IOException { - byte[] fixture = - fromHex( - new String( - IOUtils.readFully( - VideoFileFormatTest.class - .getClassLoader() - .getResourceAsStream( - "org/apache/paimon/format/blob/video-v1.hex"), - true), - StandardCharsets.UTF_8) - .trim()); - - Blob a2 = sourceFrame("a.mp4", "abc".getBytes(StandardCharsets.UTF_8), 2); - Blob a3 = sourceFrame("a.mp4", "abc".getBytes(StandardCharsets.UTF_8), 3); + byte[] fixture = fixture("video-v1.hex"); + byte[] video = "abc".getBytes(StandardCharsets.UTF_8); + byte[] mapping = fromHex(KEYFRAME_INDEX_HEX); + java.nio.file.Path source = tempPath.resolve("indexed.mp4"); + byte[] sourceBytes = new byte[video.length + mapping.length]; + System.arraycopy(video, 0, sourceBytes, 0, video.length); + System.arraycopy(mapping, 0, sourceBytes, video.length, mapping.length); + Files.write(source, sourceBytes); + VideoFrameDescriptor descriptor = + new VideoFrameDescriptor( + new Path(source.toUri()).toString(), + 0, + video.length, + 2, + video.length, + mapping.length); + Blob a2 = + Blob.fromDescriptor(org.apache.paimon.utils.UriReader.fromFile(fileIO), descriptor); + Blob a3 = + Blob.fromDescriptor( + org.apache.paimon.utils.UriReader.fromFile(fileIO), + new VideoFrameDescriptor( + descriptor.uri(), + 0, + video.length, + 3, + video.length, + mapping.length)); Blob b7 = sourceFrame("b.mp4", "WXYZ".getBytes(StandardCharsets.UTF_8), 7); Blob b8 = sourceFrame("b.mp4", "WXYZ".getBytes(StandardCharsets.UTF_8), 8); - Blob a10 = sourceFrame("a.mp4", "abc".getBytes(StandardCharsets.UTF_8), 10); + Blob a10 = + Blob.fromDescriptor( + org.apache.paimon.utils.UriReader.fromFile(fileIO), + new VideoFrameDescriptor( + descriptor.uri(), + 0, + video.length, + 10, + video.length, + mapping.length)); + write(a2, a3, null, BlobPlaceholder.INSTANCE, b7, b8, a10); assertThat(Files.readAllBytes(java.nio.file.Paths.get(file.toUri()))).isEqualTo(fixture); @@ -146,12 +174,148 @@ public void testCrossLanguageV1Fixture() throws IOException { VideoFileMeta meta = new VideoFileMeta(in, fixture.length, null); assertThat(meta.recordNumber()).isEqualTo(7); assertThat(meta.physicalVideoNumber()).isEqualTo(2); - assertThat(meta.frameIndex(0)).isEqualTo(2); - assertThat(meta.frameIndex(1)).isEqualTo(3); - assertThat(meta.frameIndex(4)).isEqualTo(7); - assertThat(meta.frameIndex(5)).isEqualTo(8); - assertThat(meta.frameIndex(6)).isEqualTo(10); + assertThat(meta.videoLength(0)).isEqualTo(video.length); + assertThat(meta.keyframeIndexOffset(0)).isEqualTo(7); + assertThat(meta.keyframeIndexLength(0)).isEqualTo(mapping.length); + assertThat(meta.keyframeIndexLength(4)).isZero(); } + VideoFrameDescriptor restored = descriptor(read(null).get(0)); + assertThat( + VideoFrameDescriptor.keyframeIndexBlob( + Blob.fromDescriptor( + org.apache.paimon.utils.UriReader.fromFile(fileIO), + restored)) + .toData()) + .isEqualTo(mapping); + } + + @Test + public void testRejectInvalidKeyframeIndex() throws IOException { + byte[] video = "video".getBytes(StandardCharsets.UTF_8); + byte[] mapping = "mapping".getBytes(StandardCharsets.UTF_8); + java.nio.file.Path source = tempPath.resolve("invalid-index.mp4"); + byte[] sourceBytes = new byte[video.length + mapping.length]; + System.arraycopy(video, 0, sourceBytes, 0, video.length); + System.arraycopy(mapping, 0, sourceBytes, video.length, mapping.length); + Files.write(source, sourceBytes); + Blob frame = + Blob.fromDescriptor( + org.apache.paimon.utils.UriReader.fromFile(fileIO), + new VideoFrameDescriptor( + new Path(source.toUri()).toString(), + 0, + video.length, + 0, + video.length, + mapping.length)); + + assertThatThrownBy(() -> write(frame)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("Invalid video keyframe index"); + } + + @Test + public void testRejectOversizedKeyframeIndexesBeforeFetch() throws IOException { + String missing = new Path(tempPath.resolve("missing.mp4").toUri()).toString(); + Blob frame = + Blob.fromDescriptor( + org.apache.paimon.utils.UriReader.fromFile(fileIO), + new VideoFrameDescriptor( + missing, + 0, + 1, + 0, + 1, + VideoFormatWriter.MAX_KEYFRAME_INDEX_BYTES + 1)); + + assertThatThrownBy(() -> write(frame)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("keyframe index length") + .hasMessageContaining("limit"); + assertThatThrownBy( + () -> + VideoFormatWriter.checkKeyframeIndexSize( + 1, VideoFormatWriter.MAX_TOTAL_KEYFRAME_INDEX_BYTES)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("Buffered video keyframe indexes") + .hasMessageContaining("limit"); + } + + @Test + public void testRejectSeekOffsetsOutsideVideoPayload() throws IOException { + byte[] video = "video".getBytes(StandardCharsets.UTF_8); + byte[] mapping = + fromHex( + "0149464b4f45444956010000000100000000000000000000000600000000000000" + + "789c6360c00e0000180001"); + java.nio.file.Path source = tempPath.resolve("out-of-range-index.mp4"); + byte[] sourceBytes = new byte[video.length + mapping.length]; + int offset = put(sourceBytes, 0, video); + put(sourceBytes, offset, mapping); + Files.write(source, sourceBytes); + Blob frame = + Blob.fromDescriptor( + org.apache.paimon.utils.UriReader.fromFile(fileIO), + new VideoFrameDescriptor( + new Path(source.toUri()).toString(), + 0, + video.length, + 0, + video.length, + mapping.length)); + + assertThatThrownBy(() -> write(frame)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("video payload"); + } + + @Test + public void testRejectInconsistentKeyframeIndexesForSamePayload() throws IOException { + byte[] video = "video".getBytes(StandardCharsets.UTF_8); + byte[] firstIndex = fromHex(KEYFRAME_INDEX_HEX); + byte[] secondIndex = + fromHex( + "0149464b4f454449560000000002000000789c6360c00ed8a074801b846682f201" + + "09ea009f"); + java.nio.file.Path source = tempPath.resolve("inconsistent-index.mp4"); + byte[] sourceBytes = new byte[video.length + firstIndex.length + secondIndex.length]; + int offset = put(sourceBytes, 0, video); + offset = put(sourceBytes, offset, firstIndex); + put(sourceBytes, offset, secondIndex); + Files.write(source, sourceBytes); + String uri = new Path(source.toUri()).toString(); + org.apache.paimon.utils.UriReader reader = + org.apache.paimon.utils.UriReader.fromFile(fileIO); + Blob unindexed = + Blob.fromDescriptor( + reader, new VideoFrameDescriptor(uri, 0, video.length, 0, -1, 0)); + Blob first = + Blob.fromDescriptor( + reader, + new VideoFrameDescriptor( + uri, 0, video.length, 1, video.length, firstIndex.length)); + Blob second = + Blob.fromDescriptor( + reader, + new VideoFrameDescriptor( + uri, + 0, + video.length, + 2, + video.length + firstIndex.length, + secondIndex.length)); + + assertThatThrownBy(() -> write(unindexed, first)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("same payload"); + fileIO.delete(file, false); + assertThatThrownBy(() -> write(first, unindexed)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("same payload"); + fileIO.delete(file, false); + assertThatThrownBy(() -> write(first, second)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("same payload"); } @Test @@ -224,22 +388,26 @@ public void testIndependentFormatRegistrationAndClassification() { @Test public void testRejectCorruptRunReference() throws IOException { byte[] physicalIndex = DeltaVarintCompressor.compress(new long[0]); + byte[] keyframeLengthIndex = DeltaVarintCompressor.compress(new long[0]); byte[] runLengthIndex = DeltaVarintCompressor.compress(new long[] {1}); byte[] runReferenceIndex = DeltaVarintCompressor.compress(new long[] {0}); byte[] firstFrameIndex = DeltaVarintCompressor.compress(new long[] {0}); byte[] bytes = new byte [physicalIndex.length + + keyframeLengthIndex.length + runLengthIndex.length + runReferenceIndex.length + firstFrameIndex.length + VideoFormatWriter.FILE_FOOTER_LENGTH]; int position = 0; position = put(bytes, position, physicalIndex); + position = put(bytes, position, keyframeLengthIndex); position = put(bytes, position, runLengthIndex); position = put(bytes, position, runReferenceIndex); position = put(bytes, position, firstFrameIndex); position = putInt(bytes, position, physicalIndex.length); + position = putInt(bytes, position, keyframeLengthIndex.length); position = putInt(bytes, position, runLengthIndex.length); position = putInt(bytes, position, runReferenceIndex.length); position = putInt(bytes, position, firstFrameIndex.length); @@ -265,7 +433,7 @@ private Blob sourceFrame(String name, byte[] bytes, long frameIndex) throws IOEx } VideoFrameDescriptor descriptor = new VideoFrameDescriptor( - new Path(source.toUri()).toString(), 0, bytes.length, frameIndex); + new Path(source.toUri()).toString(), 0, bytes.length, frameIndex, -1, 0); return Blob.fromDescriptor(org.apache.paimon.utils.UriReader.fromFile(fileIO), descriptor); } @@ -318,4 +486,17 @@ private static byte[] fromHex(String hex) { } return bytes; } + + private static byte[] fixture(String name) throws IOException { + return fromHex( + new String( + IOUtils.readFully( + VideoFileFormatTest.class + .getClassLoader() + .getResourceAsStream( + "org/apache/paimon/format/blob/" + name), + true), + StandardCharsets.UTF_8) + .trim()); + } } diff --git a/paimon-format/src/test/resources/org/apache/paimon/format/blob/video-v1.hex b/paimon-format/src/test/resources/org/apache/paimon/format/blob/video-v1.hex index 1d9b2533aabb..93f047531e7b 100644 --- a/paimon-format/src/test/resources/org/apache/paimon/format/blob/video-v1.hex +++ b/paimon-format/src/test/resources/org/apache/paimon/format/blob/video-v1.hex @@ -15,4 +15,4 @@ # specific language governing permissions and limitations # under the License. -6162635758595a0602040100020100010106010403000e06020000000500000005000000050000004944454f01 +6162635758595a0149464b4f45444956010000000200000000000000000000000100000000000000789c6360c00e78a0f4821e08cd04e503001394013b06026c6b040100020100010106010403000e0602000000020000000500000005000000050000004944454f01 diff --git a/paimon-python/pypaimon/multimodal/lerobot/dataset.py b/paimon-python/pypaimon/multimodal/lerobot/dataset.py index 494c111597f5..fbb32ccb5cd9 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/dataset.py +++ b/paimon-python/pypaimon/multimodal/lerobot/dataset.py @@ -23,12 +23,14 @@ import operator import pickle import sys +import tempfile import zlib from abc import ABC, abstractmethod from collections import OrderedDict from collections.abc import Mapping from concurrent.futures import ThreadPoolExecutor from functools import partial +from types import SimpleNamespace import pyarrow as pa @@ -48,6 +50,7 @@ ) from pypaimon.multimodal.table import _target_schema, _time_travel_table from pypaimon.multimodal.video import VideoFrameCollator +from pypaimon.table.row.video_keyframe_index import VideoKeyframeIndex _TORCH_DTYPE_NAMES = { @@ -273,6 +276,7 @@ def _set_frame_rows( decode_batch_fn=_decode_video_frames, output_column=key, collate_fn=_identity, + range_parallelism=self.blob_parallelism, ) for key in self._video_keys ] @@ -1368,14 +1372,20 @@ def decode(collator): # Cross-file windows use the regular frame assembly path. if any(d.payload_descriptor != payload for d in window): return key, None - requests.setdefault(payload, []).append(( + keyframe_index = window[0].keyframe_index_descriptor + if any(d.keyframe_index_descriptor != keyframe_index + for d in window): + raise ValueError( + "One video payload references different keyframe indexes." + ) + requests.setdefault((payload, keyframe_index), []).append(( offset, [d.frame_index for d in window])) output = [None] * len(plans) with torch.inference_mode(inference_enabled), \ torch.set_grad_enabled(grad_enabled): - for payload, windows in requests.items(): - decoder = collator._decoder(payload) + for (payload, keyframe_index), windows in requests.items(): + decoder = collator._decoder(payload, keyframe_index) get_frames = getattr(decoder, "get_frames_at", None) if not callable(get_frames): return key, None @@ -1500,6 +1510,11 @@ def _video_tensor(frame, feature, return_uint8=False): def _open_video_decoder(stream, backend=None): + keyframe_index = getattr(stream, "video_keyframe_index", None) + if not isinstance(keyframe_index, VideoKeyframeIndex): + keyframe_index = None + if backend is None and keyframe_index is not None: + return _PyAVVideoDecoder(stream, keyframe_index) if backend in (None, "torchcodec"): try: return _open_torchcodec_decoder(stream) @@ -1507,7 +1522,8 @@ def _open_video_decoder(stream, backend=None): if backend == "torchcodec": raise stream.seek(0) - return _PyAVVideoDecoder(stream) + return _PyAVVideoDecoder(stream) if keyframe_index is None \ + else _PyAVVideoDecoder(stream, keyframe_index) def _open_torchcodec_decoder(stream): @@ -1531,7 +1547,7 @@ class _PyAVVideoDecoder: # Reuse common overlapping delta windows without retaining a whole video. _FRAME_CACHE_SIZE = 8 - def __init__(self, stream): + def __init__(self, stream, keyframe_index=None): try: import av except ImportError as error: @@ -1539,13 +1555,22 @@ def __init__(self, stream): "Video-backed PaimonLeRobotDataset requires PyAV from " "'pypaimon[lerobot]'." ) from error - self._container = av.open(stream) - self._stream = self._container.streams.video[0] + self._source = stream + self._keyframe_index = keyframe_index + self._anchor_ordinals = ( + [entry[0] for entry in keyframe_index.keyframes] + if keyframe_index is not None else []) + self._container = None + self._stream = None self._next_index = 0 self._timestamps = [] self._keyframes = [] self._cache = OrderedDict() - self._frames = iter(self._container.decode(self._stream)) + self._frames = None + if keyframe_index is None: + self._container = av.open(stream) + self._stream = self._container.streams.video[0] + self._frames = iter(self._container.decode(self._stream)) def __getitem__(self, index): index = operator.index(index) @@ -1554,7 +1579,11 @@ def __getitem__(self, index): frame = self._cache.pop(index, None) if frame is not None: self._cache[index] = frame - return self._tensor(frame) + return frame if self._keyframe_index is not None \ + else self._tensor(frame) + + if self._keyframe_index is not None: + return self._read_indexed([index])[0] indexed = ( index > 0 and not self._timestamps @@ -1598,6 +1627,150 @@ def __getitem__(self, index): ) from error raise IndexError("Video frame index %d is out of range." % index) + def get_frames_at(self, *, indices): + import torch + frames = ( + self._read_indexed(indices) + if self._keyframe_index is not None + else [self[index] for index in indices] + ) + return SimpleNamespace(data=torch.stack(frames)) + + def _read_indexed(self, indices): + requested = [operator.index(index) for index in indices] + if any(index < 0 for index in requested): + raise IndexError("Video frame index is out of range.") + missing = sorted(set(requested) - set(self._cache)) + decoded = self._decode_sparse(missing) if missing else {} + try: + return [ + self._cache[index] if index in self._cache else decoded[index] + for index in requested + ] + except KeyError as error: + raise IndexError( + "Video frame index %d is out of range." % error.args[0] + ) from error + + def _decode_sparse(self, indices): + decoded = {} + import av + pending = list(indices) + last_error = None + for retry in range(3): + groups, ranges = self._sparse_plan(pending, retry) + bodies = self._read_video_ranges(ranges) + failed = [] + with tempfile.TemporaryFile() as sparse: + sparse.truncate(self._video_length()) + for (offset, unused_length), body in zip(ranges, bodies): + sparse.seek(offset) + sparse.write(body) + sparse.seek(0) + with av.open(sparse) as container: + stream = container.streams.video[0] + for anchor, targets in sorted(groups.items()): + try: + self._decode_sparse_group( + container, + stream, + anchor, + sorted(targets), + decoded, + ) + except (av.error.FFmpegError, ValueError) as error: + last_error = error + failed.extend(targets) + if not failed: + return decoded + pending = failed + raise ValueError( + "Cannot decode video frames from persisted seek index." + ) from last_error + + def _sparse_plan(self, indices, retry): + mapping = self._keyframe_index + groups = {} + ranges = list(mapping.metadata_ranges) + for index in indices: + preceding = bisect.bisect_right( + self._anchor_ordinals, index) - 1 + if preceding < 0: + raise IndexError( + "Video frame index %d is out of range." % index) + anchor = max(0, preceding - retry) + groups.setdefault(anchor, []).append(index) + start = mapping.keyframes[max(0, anchor - 1)][2] + range_end = preceding + 2 + end = ( + mapping.keyframes[range_end][2] + if range_end < len(mapping.keyframes) + else self._video_length() + ) + ranges.append((start, end - start)) + return groups, _merge_video_ranges(ranges) + + def _decode_sparse_group( + self, container, stream, anchor, targets, decoded): + anchor_ordinal, anchor_pts, unused_position = ( + self._keyframe_index.keyframes[anchor]) + container.seek( + anchor_pts, backward=True, any_frame=False, stream=stream) + ordinal = anchor_ordinal + found_anchor = False + remaining = set(targets) + found = {} + for frame in container.decode(stream): + if not found_anchor: + if frame.pts != anchor_pts or not frame.key_frame: + continue + found_anchor = True + if ordinal in remaining: + found[ordinal] = self._tensor(frame) + remaining.remove(ordinal) + if not remaining: + decoded.update(found) + for index, tensor in found.items(): + self._remember(index, tensor) + return + ordinal += 1 + if ordinal > targets[-1]: + break + missing = min(remaining) if remaining else targets[-1] + raise ValueError( + "Cannot decode indexed video frame %d from keyframe %d." + % (missing, anchor_ordinal) + ) + + def _video_length(self): + length = getattr(self._source, "video_length", None) + if length is not None: + return length + position = self._source.tell() + try: + return self._source.seek(0, 2) + finally: + self._source.seek(position) + + def _read_video_ranges(self, ranges): + reader = getattr(self._source, "video_read_ranges", None) + if callable(reader): + bodies = reader(ranges) + else: + bodies = [] + for offset, length in ranges: + self._source.seek(offset) + bodies.append(self._source.read(length)) + if len(bodies) != len(ranges): + raise IOError("Video range reader returned an invalid result count.") + for (unused_offset, length), body in zip(ranges, bodies): + if len(body) != length: + raise IOError( + "Video byte range changed while reading: expected %d, " + "read %d." % (length, len(body)) + ) + return bodies + def _index_packets(self): entries = [] for packet in self._container.demux(self._stream): @@ -1647,7 +1820,23 @@ def _tensor(frame): return torch.from_numpy(array).permute(2, 0, 1) def close(self): - self._container.close() + if self._container is not None: + self._container.close() + + +def _merge_video_ranges(ranges): + merged = [] + for offset, length in sorted(ranges): + end = offset + length + if merged and offset <= merged[-1][0] + merged[-1][1]: + previous_offset, previous_length = merged[-1] + merged[-1] = ( + previous_offset, + max(previous_offset + previous_length, end) - previous_offset, + ) + else: + merged.append((offset, length)) + return merged def _decode_video_frames(decoder, frame_indices, unused_rows): diff --git a/paimon-python/pypaimon/multimodal/lerobot/loader.py b/paimon-python/pypaimon/multimodal/lerobot/loader.py index 6925886b97de..c64ef56cfcef 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/loader.py +++ b/paimon-python/pypaimon/multimodal/lerobot/loader.py @@ -424,6 +424,8 @@ def _video_frame_descriptors( timestamp_type, name, ), + -1, + 0, ).serialize() for unused_frame_index, timestamp in video_rows ] diff --git a/paimon-python/pypaimon/multimodal/table.py b/paimon-python/pypaimon/multimodal/table.py index 81f302c49948..610350c2c543 100644 --- a/paimon-python/pypaimon/multimodal/table.py +++ b/paimon-python/pypaimon/multimodal/table.py @@ -298,7 +298,8 @@ def replace_video( """ column = self._resolve_video_frame_column(video_column) target_schema = _target_schema(self.raw_table) - payload, first_frame = _video_payload(video, first_frame) + payload, keyframe_index, first_frame = _video_payload( + video, first_frame) row_ids = ( self.scan() @@ -313,7 +314,7 @@ def replace_video( return self descriptors = _video_frame_descriptors( - payload, len(row_ids), first_frame) + payload, keyframe_index, len(row_ids), first_frame) update_data = pa.Table.from_arrays( [ pa.array(row_ids, type=pa.int64()), @@ -761,35 +762,41 @@ def _video_payload(video, first_frame): "video must be a path, Blob, or BlobDescriptor, got %r." % type(video) ) + keyframe_index = None if isinstance(payload, VideoFrameDescriptor): + keyframe_index = payload.keyframe_index_descriptor payload = payload.payload_descriptor - return payload, first_frame + return payload, keyframe_index, first_frame -def _video_frame_descriptors(payload, count, first_frame): +def _video_frame_descriptors(payload, keyframe_index, count, first_frame): from pypaimon.table.row.blob import VideoFrameDescriptor + index_offset = -1 if keyframe_index is None else keyframe_index.offset + index_length = 0 if keyframe_index is None else keyframe_index.length return [ VideoFrameDescriptor( payload.uri, payload.offset, payload.length, first_frame + index, + index_offset, + index_length, ).serialize() for index in range(count) ] def _video_frame_table(video, frames, video_column, first_frame, target_schema): - payload, first_frame = _video_payload(video, first_frame) + payload, keyframe_index, first_frame = _video_payload(video, first_frame) non_video_schema = pa.schema([ field for field in target_schema if field.name != video_column ]) frame_table = _to_arrow_table(frames, non_video_schema) descriptor_values = _video_frame_descriptors( - payload, frame_table.num_rows, first_frame) + payload, keyframe_index, frame_table.num_rows, first_frame) arrays = [] for field in target_schema: if field.name == video_column: diff --git a/paimon-python/pypaimon/multimodal/video.py b/paimon-python/pypaimon/multimodal/video.py index ac2e7039e83c..4d627239f111 100644 --- a/paimon-python/pypaimon/multimodal/video.py +++ b/paimon-python/pypaimon/multimodal/video.py @@ -22,6 +22,7 @@ from collections.abc import Mapping from pypaimon.table.row.blob import Blob, VideoFrameDescriptor +from pypaimon.table.row.video_keyframe_index import VideoKeyframeIndex class VideoFrameCollator: @@ -57,7 +58,8 @@ def __init__( output_column="frame", max_open_videos=8, collate_fn=None, - decode_batch_fn=None): + decode_batch_fn=None, + range_parallelism=4): if not video_column: raise ValueError("video_column is required.") if not callable(decoder_factory): @@ -77,6 +79,10 @@ def __init__( raise ValueError("max_open_videos must be a positive int.") if collate_fn is not None and not callable(collate_fn): raise ValueError("collate_fn must be callable or None.") + if (isinstance(range_parallelism, bool) + or not isinstance(range_parallelism, int) + or range_parallelism <= 0): + raise ValueError("range_parallelism must be a positive int.") raw_table = getattr(table, "raw_table", table) file_io = getattr(raw_table, "file_io", None) @@ -91,6 +97,7 @@ def __init__( self.output_column = output_column self.max_open_videos = max_open_videos self.collate_fn = collate_fn + self.range_parallelism = range_parallelism self._decoders = OrderedDict() self._owner_pid = os.getpid() @@ -132,12 +139,17 @@ def _decode_rows(self, rows): if descriptor is None: decoded[position] = output continue - grouped.setdefault(descriptor.payload_descriptor, []).append( - (descriptor.frame_index, position, output) - ) + payload = descriptor.payload_descriptor + index = descriptor.keyframe_index_descriptor + group = grouped.setdefault(payload, [index, []]) + if group[0] != index: + raise ValueError( + "One video payload references different keyframe indexes." + ) + group[1].append((descriptor.frame_index, position, output)) - for payload, frames in grouped.items(): - decoder = self._decoder(payload) + for payload, (index, frames) in grouped.items(): + decoder = self._decoder(payload, index) frames.sort(key=lambda frame: frame[0]) if self.decode_batch_fn is not None: values = self.decode_batch_fn( @@ -190,10 +202,11 @@ def _prepare_row(self, row): ) return output, descriptor - def _decoder(self, descriptor): - resource = self._decoders.pop(descriptor, None) + def _decoder(self, descriptor, keyframe_index_descriptor=None): + cache_key = (descriptor, keyframe_index_descriptor) + resource = self._decoders.pop(cache_key, None) if resource is not None: - self._decoders[descriptor] = resource + self._decoders[cache_key] = resource return resource[0] # Reuse the table's resolved FileIO. Rebuilding a reader from raw URI @@ -205,12 +218,31 @@ def _decoder(self, descriptor): descriptor.length, ).new_input_stream() try: + if keyframe_index_descriptor is not None: + mapping = self.file_io.read_file_range( + keyframe_index_descriptor.uri, + keyframe_index_descriptor.offset, + keyframe_index_descriptor.length, + ) + stream.video_keyframe_index = ( + VideoKeyframeIndex.deserialize(mapping, descriptor.length)) + stream.video_length = descriptor.length + stream.video_read_ranges = lambda ranges: ( + self.file_io.read_ranges_coalesced( + [( + descriptor.uri, + descriptor.offset + offset, + length, + ) for offset, length in ranges], + self.range_parallelism, + ) + ) decoder = self.decoder_factory(stream) except Exception: stream.close() raise resource = (decoder, stream) - self._decoders[descriptor] = resource + self._decoders[cache_key] = resource if len(self._decoders) > self.max_open_videos: _, evicted = self._decoders.popitem(last=False) self._close_resource(evicted) diff --git a/paimon-python/pypaimon/read/reader/video_format_reader.py b/paimon-python/pypaimon/read/reader/video_format_reader.py index 3c2b91efdf14..04ac0be4b545 100644 --- a/paimon-python/pypaimon/read/reader/video_format_reader.py +++ b/paimon-python/pypaimon/read/reader/video_format_reader.py @@ -30,7 +30,7 @@ class VideoFileMeta: VERSION = 1 MAGIC_NUMBER = 0x4F454449 - FOOTER_SIZE = 21 + FOOTER_SIZE = 25 NULL_REFERENCE = -1 PLACE_HOLDER_REFERENCE = -2 @@ -44,9 +44,9 @@ def __init__(self, stream, file_size: int): footer = stream.read(self.FOOTER_SIZE) if len(footer) != self.FOOTER_SIZE: raise IOError("Corrupt video file: cannot read footer.") - lengths = struct.unpack(' footer_start: raise IOError("Corrupt video file: indexes exceed the file size.") index_start = footer_start - total_index_length indexes = [] offset = index_start - for name, length in zip( - ("physical video", "run length", "run reference", "first frame"), - index_lengths): + for name, length in zip(index_names, index_lengths): stream.seek(offset) raw = stream.read(length) if len(raw) != length: @@ -72,22 +75,49 @@ def __init__(self, stream, file_size: int): indexes.append(DeltaVarintCompressor.decompress(raw)) offset += length - physical_lengths, run_lengths, references, first_frames = indexes + ( + physical_lengths, + keyframe_index_lengths, + run_lengths, + references, + first_frames, + ) = indexes + if len(keyframe_index_lengths) != len(physical_lengths): + raise IOError( + "Corrupt video file: physical video and keyframe index " + "indexes have different counts." + ) + if any(length < 0 for length in keyframe_index_lengths): + raise IOError( + "Corrupt video file: negative keyframe index length." + ) + + keyframe_index_size = sum(keyframe_index_lengths) + mapping_start = index_start - keyframe_index_size + if mapping_start < 0: + raise IOError( + "Corrupt video file: keyframe indexes exceed the file size." + ) physical_offsets = [] payload_offset = 0 for ordinal, length in enumerate(physical_lengths): - if length <= 0 or length > index_start - payload_offset: + if length <= 0 or length > mapping_start - payload_offset: raise IOError( "Corrupt video file: invalid physical video length %s " "at ordinal %s." % (length, ordinal) ) physical_offsets.append(payload_offset) payload_offset += length - if payload_offset != index_start: + if payload_offset != mapping_start: raise IOError( "Corrupt video file: indexed videos use %s bytes, but payload " - "region contains %s bytes." % (payload_offset, index_start) + "region contains %s bytes." % (payload_offset, mapping_start) ) + keyframe_index_offsets = [] + mapping_offset = mapping_start + for length in keyframe_index_lengths: + keyframe_index_offsets.append(mapping_offset) + mapping_offset += length if not (len(run_lengths) == len(references) == len(first_frames)): raise IOError( @@ -123,6 +153,8 @@ def __init__(self, stream, file_size: int): self.physical_lengths = physical_lengths self.physical_offsets = physical_offsets + self.keyframe_index_lengths = keyframe_index_lengths + self.keyframe_index_offsets = keyframe_index_offsets self.run_ends = run_ends self.references = references self.first_frames = first_frames @@ -163,11 +195,17 @@ def frame(self, returned_row: int): if reference == self.PLACE_HOLDER_REFERENCE: return Blob.PLACE_HOLDER run_start = 0 if run == 0 else self.run_ends[run - 1] - return ( + frame = ( self.physical_offsets[reference], self.physical_lengths[reference], self.first_frames[run] + logical - run_start, ) + if self.keyframe_index_lengths[reference] == 0: + return frame + (-1, 0) + return frame + ( + self.keyframe_index_offsets[reference], + self.keyframe_index_lengths[reference], + ) class VideoFrameRecordIterator: @@ -188,10 +226,7 @@ def __next__(self): raise StopIteration value = self.meta.frame(self.current_position) if isinstance(value, tuple): - offset, length, frame_index = value - descriptor = VideoFrameDescriptor( - self.file_path, offset, length, frame_index - ) + descriptor = VideoFrameDescriptor(self.file_path, *value) value = Blob.from_descriptor(self._uri_reader, descriptor) self.current_position += 1 return GenericRow([value], [self.field], RowKind.INSERT) diff --git a/paimon-python/pypaimon/table/row/blob.py b/paimon-python/pypaimon/table/row/blob.py index f2d666b64db2..317763ab12f4 100644 --- a/paimon-python/pypaimon/table/row/blob.py +++ b/paimon-python/pypaimon/table/row/blob.py @@ -229,9 +229,11 @@ class VideoFrameDescriptor(BlobDescriptor): CURRENT_VERSION = 1 MAGIC = 0x564944454F46524D # "VIDEOFRM" - _FIXED_LENGTH = 1 + 8 + 4 + 8 + 8 + 8 + _FIXED_LENGTH = 1 + 8 + 4 + 8 + 8 + 8 + 8 + 8 - def __init__(self, uri: str, offset: int, length: int, frame_index: int): + def __init__( + self, uri: str, offset: int, length: int, frame_index: int, + keyframe_index_offset: int, keyframe_index_length: int): if isinstance(frame_index, bool) or not isinstance(frame_index, int): raise TypeError("Video frame index must be an int.") if frame_index < 0: @@ -239,8 +241,15 @@ def __init__(self, uri: str, offset: int, length: int, frame_index: int): "Video frame index must be non-negative, but was %s." % frame_index ) + if keyframe_index_length < 0: + raise ValueError("Video keyframe index length must be non-negative.") + if ((keyframe_index_length == 0 and keyframe_index_offset != -1) + or (keyframe_index_length > 0 and keyframe_index_offset < 0)): + raise ValueError("Invalid video keyframe index range.") super().__init__(uri, offset, length) self._frame_index = frame_index + self._keyframe_index_offset = keyframe_index_offset + self._keyframe_index_length = keyframe_index_length @property def frame_index(self) -> int: @@ -251,13 +260,28 @@ def payload_descriptor(self) -> BlobDescriptor: """Physical video identity without the logical frame locator.""" return BlobDescriptor(self.uri, self.offset, self.length) + @property + def keyframe_index_descriptor(self) -> Optional[BlobDescriptor]: + if self._keyframe_index_length == 0: + return None + return BlobDescriptor( + self.uri, + self._keyframe_index_offset, + self._keyframe_index_length, + ) + def serialize(self) -> bytes: uri_bytes = self.uri.encode('utf-8') return ( struct.pack(' 'VideoFrameDescriptor': uri_end = 13 + uri_length uri = raw[13:uri_end].decode('utf-8') - offset, length, frame_index = struct.unpack( - ' bool: isinstance(other, VideoFrameDescriptor) and self.payload_descriptor == other.payload_descriptor and self.frame_index == other.frame_index + and self.keyframe_index_descriptor == other.keyframe_index_descriptor ) def __hash__(self) -> int: - return hash((self.payload_descriptor, self.frame_index)) + return hash(( + self.payload_descriptor, + self.frame_index, + self.keyframe_index_descriptor, + )) def __str__(self) -> str: return ( diff --git a/paimon-python/pypaimon/table/row/video_keyframe_index.py b/paimon-python/pypaimon/table/row/video_keyframe_index.py new file mode 100644 index 000000000000..1a265f80271e --- /dev/null +++ b/paimon-python/pypaimon/table/row/video_keyframe_index.py @@ -0,0 +1,295 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Sparse seek metadata for encoded video payloads.""" + +import operator +import struct +import zlib + + +class VideoKeyframeIndex: + """Persist video metadata ranges and presentation-order keyframes.""" + + VERSION = 1 + MAGIC = 0x564944454F4B4649 # "VIDEOKFI" + HEADER = struct.Struct('I4s", header) + header_size = 8 + if size == 1: + extended = source.read(8) + if len(extended) != 8: + raise ValueError("Invalid ISO BMFF extended box header.") + size = struct.unpack(">Q", extended)[0] + header_size = 16 + elif size == 0: + size = payload_length - offset + if size < header_size or offset + size > payload_length: + raise ValueError("Invalid ISO BMFF box size.") + ranges.append(( + offset, + header_size if box_type in (b"mdat", b"free", b"skip") + else size, + )) + has_moov |= box_type == b"moov" + offset += size + if offset != payload_length or not has_moov: + raise ValueError("Video is not a supported ISO BMFF file.") + return ranges + + @classmethod + def validate(cls, data, payload_length=None): + metadata_count, keyframe_count, entries_offset = cls._read_header(data) + for _ in cls._iter_metadata_ranges( + data, metadata_count, payload_length): + pass + for _ in cls._iter_keyframes( + data, entries_offset, keyframe_count, payload_length): + pass + + @classmethod + def deserialize(cls, data, payload_length=None): + metadata_count, keyframe_count, entries_offset = cls._read_header(data) + return cls( + cls._iter_metadata_ranges(data, metadata_count, payload_length), + cls._iter_keyframes( + data, entries_offset, keyframe_count, payload_length), + ) + + @classmethod + def _read_header(cls, data): + if not isinstance(data, (bytes, bytearray)): + raise TypeError("VideoKeyframeIndex expects bytes.") + if len(data) <= cls.HEADER.size: + raise ValueError("Invalid video keyframe index: too short.") + version, magic, metadata_count, keyframe_count = ( + cls.HEADER.unpack(data[:cls.HEADER.size])) + if version != cls.VERSION or magic != cls.MAGIC: + raise ValueError("Invalid video keyframe index version or magic.") + if keyframe_count == 0: + raise ValueError("Invalid video keyframe index header.") + entries_offset = ( + cls.HEADER.size + metadata_count * cls.METADATA_RANGE.size) + if entries_offset >= len(data): + raise ValueError("Invalid video keyframe index metadata ranges.") + return metadata_count, keyframe_count, entries_offset + + @classmethod + def _iter_metadata_ranges(cls, data, count, payload_length): + previous_end = 0 + offset = cls.HEADER.size + for _ in range(count): + start, length = cls.METADATA_RANGE.unpack_from(data, offset) + end = start + length + if start < previous_end or length <= 0: + raise ValueError( + "Invalid video keyframe index metadata ranges." + ) + if payload_length is not None and end > payload_length: + raise ValueError( + "Video metadata range is outside the video payload." + ) + previous_end = end + offset += cls.METADATA_RANGE.size + yield start, length + + @classmethod + def _iter_keyframes(cls, data, input_offset, count, payload_length): + try: + decoder = zlib.decompressobj() + pending = b'' + remainder = b'' + decoded_count = 0 + previous_ordinal, previous_pts, previous_position = -1, None, -1 + while not decoder.eof: + if not pending and input_offset < len(data): + end = min(input_offset + cls._CHUNK_SIZE, len(data)) + pending = data[input_offset:end] + input_offset = end + output = decoder.decompress(pending, cls._CHUNK_SIZE) + pending = decoder.unconsumed_tail + entries = remainder + output + complete = len(entries) - len(entries) % cls.ENTRY.size + remainder = entries[complete:] + for ordinal, pts, position in cls.ENTRY.iter_unpack( + entries[:complete]): + if decoded_count >= count: + raise ValueError( + "Invalid video keyframe index entries." + ) + if decoded_count == 0 and ordinal != 0: + raise ValueError( + "Video keyframe index requires an initial " + "keyframe." + ) + if ordinal <= previous_ordinal: + raise ValueError( + "Keyframe ordinals must be strictly increasing." + ) + if previous_pts is not None and pts <= previous_pts: + raise ValueError( + "Keyframe timestamps must be strictly increasing." + ) + if (position <= previous_position + or (payload_length is not None + and position >= payload_length)): + raise ValueError( + "Keyframe packet positions must be within the " + "video payload and strictly increasing." + ) + decoded_count += 1 + previous_ordinal, previous_pts, previous_position = ( + ordinal, pts, position) + yield ordinal, pts, position + if not output and not pending and input_offset == len(data): + break + if (not decoder.eof or decoder.unused_data + or pending or input_offset != len(data) + or remainder or decoded_count != count): + raise ValueError("Invalid video keyframe index entries.") + except (struct.error, zlib.error) as error: + raise ValueError( + "Invalid video keyframe index payload." + ) from error diff --git a/paimon-python/pypaimon/tests/blob_test.py b/paimon-python/pypaimon/tests/blob_test.py index 7e25b720a371..fd61c0073dfa 100644 --- a/paimon-python/pypaimon/tests/blob_test.py +++ b/paimon-python/pypaimon/tests/blob_test.py @@ -1479,7 +1479,7 @@ def test_blob_descriptor_detection(self): self.assertIsNone(BlobDescriptor.parse_if_serialized(v1_shaped_inline)) self.assertIsNone(BlobDescriptor.parse_if_serialized(b"tiny")) - video = VideoFrameDescriptor("file:///v.mp4", 0, 10, 2) + video = VideoFrameDescriptor("file:///v.mp4", 0, 10, 2, -1, 0) video_bytes = video.serialize() self.assertEqual(video_bytes, BlobDescriptor.deserialize(video_bytes).serialize()) self.assertEqual(video, BlobDescriptor.parse_if_serialized(video_bytes)) @@ -1666,7 +1666,7 @@ def test_dedicated_writer_accepts_exact_v1_descriptor_bytes(self): pa.RecordBatch.from_arrays( [pa.array([v1], type=pa.large_binary())], names=["payload"])) - video_bytes = VideoFrameDescriptor("file:///v.mp4", 0, 10, 2).serialize() + video_bytes = VideoFrameDescriptor("file:///v.mp4", 0, 10, 2, -1, 0).serialize() writer._validate_inline_stored_fields_input( pa.RecordBatch.from_arrays( [pa.array([video_bytes], type=pa.large_binary())], names=["payload"])) @@ -1760,7 +1760,7 @@ def test_offset_row_get_blob_video_frame_descriptor_bytes(self): data = b"video-frame-payload" descriptor = VideoFrameDescriptor( - "file-backed/video.mp4", 0, len(data), 2) + "file-backed/video.mp4", 0, len(data), 2, -1, 0) file_io = self._token_aware_file_io(data) row = OffsetRow( (descriptor.serialize(),), 0, 1, @@ -1788,7 +1788,7 @@ def test_blob_inline_convert_reader_resolves_video_frame_descriptor(self): data = b"convert video blob" descriptor = VideoFrameDescriptor( - "file-backed/convert.mp4", 0, len(data), 2) + "file-backed/convert.mp4", 0, len(data), 2, -1, 0) file_io = self._token_aware_file_io(data) batch = RecordBatch.from_arrays( [pa.array([descriptor.serialize()], type=pa.large_binary())], 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 bc797398e35e..00e9e9019ded 100644 --- a/paimon-python/pypaimon/tests/data_evolution_row_rolling_test.py +++ b/paimon-python/pypaimon/tests/data_evolution_row_rolling_test.py @@ -244,7 +244,7 @@ def test_video_writer_rolls_between_payload_groups(self): first_descriptor.uri, first_descriptor.offset, first_descriptor.length, - frame, + frame, -1, 0, ).serialize() for frame in range(3) ] + [ @@ -252,7 +252,7 @@ def test_video_writer_rolls_between_payload_groups(self): second_descriptor.uri, second_descriptor.offset, second_descriptor.length, - frame, + frame, -1, 0, ).serialize() for frame in range(2) ], @@ -280,12 +280,12 @@ def test_video_batches_are_preserved_at_payload_boundaries(self): 'id': list(range(5)), 'payload': [ VideoFrameDescriptor( - first.uri, first.offset, first.length, frame + first.uri, first.offset, first.length, frame, -1, 0 ).serialize() for frame in range(3) ] + [ VideoFrameDescriptor( - second.uri, second.offset, second.length, frame + second.uri, second.offset, second.length, frame, -1, 0 ).serialize() for frame in range(2) ], @@ -340,7 +340,7 @@ def test_multiple_video_fields_roll_as_episode_aligned_groups(self): descriptors[0 if row < 3 else 2].uri, 0, descriptors[0 if row < 3 else 2].length, - row if row < 3 else row - 3, + row if row < 3 else row - 3, -1, 0, ).serialize() for row in range(5) ], @@ -349,7 +349,7 @@ def test_multiple_video_fields_roll_as_episode_aligned_groups(self): descriptors[1 if row < 3 else 3].uri, 0, descriptors[1 if row < 3 else 3].length, - row if row < 3 else row - 3, + row if row < 3 else row - 3, -1, 0, ).serialize() for row in range(5) ], @@ -405,13 +405,13 @@ def test_multiple_video_fields_allow_nested_episode_boundaries(self): 'id': list(range(4)), 'camera_a': [ VideoFrameDescriptor( - camera_a.uri, 0, camera_a.length, frame + camera_a.uri, 0, camera_a.length, frame, -1, 0 ).serialize() for frame in range(4) ], 'camera_b': [ VideoFrameDescriptor( - descriptor.uri, 0, descriptor.length, frame % 2 + descriptor.uri, 0, descriptor.length, frame % 2, -1, 0 ).serialize() for frame, descriptor in enumerate( [camera_b_0, camera_b_0, camera_b_1, camera_b_1] @@ -456,7 +456,7 @@ def test_video_episodes_roll_before_shared_payload_group(self): descriptor.uri, descriptor.offset, descriptor.length, - frame, + frame, -1, 0, ).serialize() for frame in range(6) ], @@ -505,13 +505,13 @@ def test_video_columns_roll_together_at_episode_boundaries(self): 'id': list(range(4)), 'camera_a': [ VideoFrameDescriptor( - camera_a.uri, 0, camera_a.length, frame + camera_a.uri, 0, camera_a.length, frame, -1, 0 ).serialize() for frame in range(4) ], 'camera_b': [ VideoFrameDescriptor( - camera_b.uri, 0, camera_b.length, frame + camera_b.uri, 0, camera_b.length, frame, -1, 0 ).serialize() for frame in range(4) ], @@ -559,7 +559,7 @@ def test_vector_rolling_waits_for_video_episode_boundary(self): descriptor.uri, descriptor.offset, descriptor.length, - frame, + frame, -1, 0, ).serialize() for frame in range(4) ], diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index 649538185020..056e7b641fdf 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -91,6 +91,7 @@ _validate_info_paths, ) from pypaimon.multimodal.table import _target_schema +from pypaimon.table.row.video_keyframe_index import VideoKeyframeIndex try: from lerobot.datasets.lerobot_dataset import LeRobotDataset @@ -390,40 +391,68 @@ def close(self): av is not None and importlib.util.find_spec("torch") is not None, "PyAV and Torch are required for video decoding", ) - def test_pyav_decoder_seeks_before_b_frames(self): - output = io.BytesIO() - with av.open(output, mode="w", format="mp4") as container: - stream = container.add_stream("mpeg4", rate=30) - stream.width = 16 - stream.height = 16 - stream.pix_fmt = "yuv420p" - stream.gop_size = 12 - stream.codec_context.max_b_frames = 2 - for index in range(70): - image = np.full( - (16, 16, 3), index + 24, dtype=np.uint8) - frame = av.VideoFrame.from_ndarray(image, format="rgb24") - frame.pts = index - frame.time_base = Fraction(1, 30) - for packet in stream.encode(frame): - container.mux(packet) - for packet in stream.encode(): - container.mux(packet) - - payload = output.getvalue() - with av.open(io.BytesIO(payload)) as container: - expected = [ - np.array(frame.to_ndarray(format="rgb24"), copy=True) - for frame in container.decode(video=0) - ] + def test_indexed_pyav_decoder_handles_b_frames_and_fragmented_video(self): + cases = ( + ("mpeg4", None, None), + ( + "libx264", + {"movflags": "frag_keyframe+empty_moov+default_base_moof"}, + None, + ), + ( + "libx265", + None, + {"x265-params": + "keyint=12:min-keyint=12:scenecut=0:log-level=error"}, + ), + ) + for codec, container_options, stream_options in cases: + with self.subTest(codec=codec): + output = io.BytesIO() + with av.open( + output, + mode="w", + format="mp4", + options=container_options) as container: + stream = container.add_stream(codec, rate=30) + stream.width = 16 + stream.height = 16 + stream.pix_fmt = "yuv420p" + stream.gop_size = 12 + stream.codec_context.max_b_frames = 2 + if stream_options is not None: + stream.options = stream_options + for index in range(60): + image = np.full( + (16, 16, 3), index + 24, dtype=np.uint8) + frame = av.VideoFrame.from_ndarray( + image, format="rgb24") + frame.pts = index + frame.time_base = Fraction(1, 30) + for packet in stream.encode(frame): + container.mux(packet) + for packet in stream.encode(): + container.mux(packet) - decoder = _PyAVVideoDecoder(io.BytesIO(payload)) - try: - for index in (69, 20, 35, 1, 68): - actual = decoder[index].permute(1, 2, 0).numpy() - np.testing.assert_array_equal(expected[index], actual) - finally: - decoder.close() + payload = output.getvalue() + index = VideoKeyframeIndex.inspect( + io.BytesIO(payload), len(payload)) + with av.open(io.BytesIO(payload)) as container: + expected = np.stack([ + np.array(frame.to_ndarray(format="rgb24"), copy=True) + for frame in container.decode(video=0) + ]) + + source = io.BytesIO(payload) + source.video_length = len(payload) + decoder = _PyAVVideoDecoder(source, index) + try: + actual = decoder.get_frames_at( + indices=list(range(len(expected))) + ).data.permute(0, 2, 3, 1).numpy() + np.testing.assert_array_equal(expected, actual) + finally: + decoder.close() def test_default_video_backend_falls_back_on_os_error(self): stream = Mock() @@ -480,7 +509,8 @@ def read_indices(self, indices, columns): "frame_index": index, "timestamp": index / 10, "task_index": 0, "camera": pmm.VideoFrameDescriptor( - "file:///shared.video", 0, 5, index + 3).serialize(), + "file:///shared.video", 0, 5, index + 3, -1, 0 + ).serialize(), } for index in indices], schema=self.schema).select(columns) for backend, batch in (("torchcodec", True), (None, True), @@ -597,7 +627,8 @@ def test_torchcodec_batch_matches_single_frame_decoding(self): ) as decode_batch: result = collator([{ "video": pmm.VideoFrameDescriptor( - "file:///episode.mp4", 0, len(payload), index).serialize(), + "file:///episode.mp4", 0, len(payload), index, -1, 0 + ).serialize(), } for index in indices]) decode_batch.assert_called_once_with(indices=[1, 5, 9, 9]) actual = torch.stack([row["frame"] for row in result]) @@ -1251,7 +1282,7 @@ def test_video_windows_decode_directly_and_fall_back(self): key = "camera" feature = {key: {"dtype": "video", "shape": [4, 5, 3]}} - rows = {i: {key: VideoFrameDescriptor("a.video", 0, 1, i).serialize()} + rows = {i: {key: VideoFrameDescriptor("a.video", 0, 1, i, -1, 0).serialize()} for i in range(2)} plans = [{"windows": {key: [1, 0, 0]}}] * 2 pixels = torch.arange(120).reshape(2, 4, 5, 3).to(torch.uint8) @@ -1295,8 +1326,8 @@ def test_video_windows_decode_directly_and_fall_back(self): decoder.get_frames_at.reset_mock() separate = dict(rows) - separate[2] = {key: VideoFrameDescriptor("b.video", 0, 1, 0).serialize()} - separate[3] = {key: VideoFrameDescriptor("b.video", 0, 1, 1).serialize()} + separate[2] = {key: VideoFrameDescriptor("b.video", 0, 1, 0, -1, 0).serialize()} + separate[3] = {key: VideoFrameDescriptor("b.video", 0, 1, 1, -1, 0).serialize()} interleaved = [{"windows": {key: window}} for window in ([1, 0, 0], [2, 3], [0, 1], [3, 2, 2])] result = _decode_video_windows( @@ -1314,7 +1345,7 @@ def test_video_windows_decode_directly_and_fall_back(self): open_decoder.reset_mock() mixed = dict(rows) - mixed[1] = {key: VideoFrameDescriptor("b.video", 0, 1, 0).serialize()} + mixed[1] = {key: VideoFrameDescriptor("b.video", 0, 1, 0, -1, 0).serialize()} self.assertEqual({}, _decode_video_windows( plans, mixed, [collator], feature, True)) open_decoder.assert_not_called() @@ -2483,7 +2514,8 @@ def sample_timestamps(unused_dataset, uri): path for path in source_file_io.opened_paths if path.endswith(".mp4") ] - self.assertEqual(3, len(opened_videos)) + # PyAV-enabled writers inspect each source once before copying it. + self.assertEqual(6 if av is not None else 3, len(opened_videos)) self.assertEqual(1, source_file_io.close_count) _, remote_bodies = connection.get_table( "remote_frames").scan().select([ @@ -2635,7 +2667,10 @@ def _assert_imported_video_payload_can_be_decoded(self, training_reads): "videos/camera_b/to_timestamp": 0.4, }, ] - physical_frame_values = [24, 56, 88, 120, 168, 216] + physical_frame_values = [24] * 60 + for index, value in { + 1: 56, 2: 88, 3: 120, 5: 168, 6: 216}.items(): + physical_frame_values[index] = value expected_frame_values = [168, 216, 56, 88, 120] info_dir = temp_dir / "meta" @@ -2659,8 +2694,9 @@ def _assert_imported_video_payload_can_be_decoded(self, training_reads): 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): + stream.gop_size = 10 + stream.codec_context.max_b_frames = 2 + for pts, value in enumerate(physical_frame_values): image = np.full((16, 16, 3), value, dtype=np.uint8) frame = av.VideoFrame.from_ndarray(image, format="rgb24") frame.pts = pts @@ -2702,6 +2738,7 @@ def read_batch(self, begin, end): connection = pmm.connect(options={ "warehouse": str(temp_dir / "warehouse"), + "file-io.read-coalesce.max-gap": "0 b", }) with patch( "pypaimon.multimodal.lerobot.api." @@ -2717,14 +2754,22 @@ def read_batch(self, begin, end): table = connection.get_table("frames") rows = table.scan().select([ - "index", "camera" + "index", "camera", "camera_b" ]).to_arrow().sort_by("index").to_pylist() descriptors = [ pmm.VideoFrameDescriptor.deserialize(row["camera"]) for row in rows ] + camera_b_descriptors = [ + pmm.VideoFrameDescriptor.deserialize(row["camera_b"]) + for row in rows + ] + self.assertTrue(all( + descriptor.keyframe_index_descriptor is not None + for descriptor in descriptors + camera_b_descriptors + )) self.assertEqual( - [4, 5, 1, 2, 3], + [5, 6, 1, 2, 3], [descriptor.frame_index for descriptor in descriptors], ) @@ -2774,7 +2819,19 @@ def close(self): delta_timestamps={"camera": [0.0, 0.1]}, ) try: - last, first = dataset.__getitems__([4, 0]) + video_ranges = [] + file_io = table.raw_table.file_io + read_ranges = file_io.read_ranges_coalesced + + def track_ranges(ranges, parallelism): + video_ranges.extend(ranges) + return read_ranges(ranges, parallelism) + + with patch.object( + file_io, + "read_ranges_coalesced", + side_effect=track_ranges): + last, first = dataset.__getitems__([4, 0]) self.assertEqual( [2, 3, 16, 16], list(last["camera"].shape)) self.assertEqual( @@ -2795,6 +2852,23 @@ def close(self): [False, True], last["camera_is_pad"].tolist()) self.assertEqual( 1, len(dataset._video_collators[0]._decoders)) + for descriptor in descriptors + camera_b_descriptors: + loaded = bytearray(descriptor.length) + for path, offset, length in video_ranges: + if path != descriptor.uri: + continue + begin = max(offset, descriptor.offset) + end = min( + offset + length, + descriptor.offset + descriptor.length, + ) + if begin < end: + loaded[ + begin - descriptor.offset: + end - descriptor.offset + ] = b"\1" * (end - begin) + self.assertGreater(sum(loaded), 0) + self.assertLess(sum(loaded), descriptor.length) from torch.utils.data import DataLoader worker_indices = [] diff --git a/paimon-python/pypaimon/tests/multimodal_table_test.py b/paimon-python/pypaimon/tests/multimodal_table_test.py index 2e7e680f723b..09043f8e9222 100644 --- a/paimon-python/pypaimon/tests/multimodal_table_test.py +++ b/paimon-python/pypaimon/tests/multimodal_table_test.py @@ -200,6 +200,128 @@ def test_add_videos_packs_multiple_videos_in_one_commit(self): self.assertTrue(descriptors[0].uri.endswith(".video")) self.assertEqual(2, len({d.payload_descriptor for d in descriptors})) + def test_add_video_indexes_path_and_local_blob_with_unknown_length(self): + av = pytest.importorskip("av") + np = pytest.importorskip("numpy") + from fractions import Fraction + from pypaimon.table.row.blob import Blob, VideoFrameDescriptor + from pypaimon.table.row.video_keyframe_index import VideoKeyframeIndex + + table = self.conn.create_table( + "local_video_indexes", + schema=_schema({ + "frame_id": pa.int32(), + "video": pa.large_binary(), + }), + options=dict(_PARQUET_OPTIONS, **{ + "video-frame-field": "video", + "blob-as-descriptor": "true", + }), + ) + video_path = os.path.join(self.temp_dir, "local-indexed.mp4") + with av.open(video_path, mode="w", format="mp4") as container: + stream = container.add_stream("libx264", rate=10) + stream.width = stream.height = 16 + stream.pix_fmt = "yuv420p" + stream.gop_size = 4 + for index in range(12): + frame = av.VideoFrame.from_ndarray( + np.full((16, 16, 3), index * 16, dtype=np.uint8), + format="rgb24", + ) + frame.pts = index + frame.time_base = Fraction(1, 10) + for packet in stream.encode(frame): + container.mux(packet) + for packet in stream.encode(): + container.mux(packet) + + table.add_video(video_path, [{"frame_id": 0}]) + table.add_video( + Blob.from_local(video_path), [{"frame_id": 1}], first_frame=1) + + descriptors = [ + VideoFrameDescriptor.deserialize(row["video"]) + for row in sorted( + table.scan().select(["frame_id", "video"]).to_list(), + key=lambda row: row["frame_id"], + ) + ] + self.assertEqual([0, 1], [value.frame_index for value in descriptors]) + for descriptor in descriptors: + stored_index = descriptor.keyframe_index_descriptor + self.assertIsNotNone(stored_index) + VideoKeyframeIndex.deserialize( + Blob.from_file( + table.raw_table.file_io, + stored_index.uri, + stored_index.offset, + stored_index.length, + ).to_data(), + descriptor.length, + ) + + def test_video_write_apis_preserve_existing_keyframe_index(self): + from pypaimon.table.row.blob import Blob, VideoFrameDescriptor + from pypaimon.table.row.video_keyframe_index import VideoKeyframeIndex + + table = self.conn.create_table( + "indexed_video_write_apis", + schema=_schema({ + "frame_id": pa.int32(), + "video": pa.large_binary(), + }), + options=dict(_PARQUET_OPTIONS, **{ + "video-frame-field": "video", + "blob-as-descriptor": "true", + }), + ) + video = b"indexed-video" + index = VideoKeyframeIndex( + [(0, 1)], [(0, 0, 1)] + ).serialize() + source_path = os.path.join(self.temp_dir, "indexed-source.mp4") + with open(source_path, "wb") as output: + output.write(video + index) + source = VideoFrameDescriptor( + "file://" + source_path, + 0, + len(video), + 0, + len(video), + len(index), + ) + + table.add_video(source, [{"frame_id": 1}], first_frame=1) + table.add_videos([(source, [{"frame_id": 2}], 2)]) + plain_path = os.path.join(self.temp_dir, "plain-source.mp4") + with open(plain_path, "wb") as output: + output.write(b"plain-video") + table.add_video(Blob.from_local(plain_path), [{"frame_id": 3}]) + table.replace_video("frame_id = 3", source, first_frame=3) + + rows = sorted( + table.scan().select(["frame_id", "video"]).to_list(), + key=lambda row: row["frame_id"], + ) + descriptors = [ + VideoFrameDescriptor.deserialize(row["video"]) + for row in rows + ] + self.assertEqual([1, 2, 3], [value.frame_index for value in descriptors]) + for descriptor in descriptors: + stored_index = descriptor.keyframe_index_descriptor + self.assertIsNotNone(stored_index) + self.assertEqual( + index, + Blob.from_file( + table.raw_table.file_io, + stored_index.uri, + stored_index.offset, + stored_index.length, + ).to_data(), + ) + def test_normal_update_preserves_video_descriptors(self): from pypaimon.table.row.blob import Blob @@ -342,7 +464,7 @@ def test_add_batches_aborts_before_commit_on_invalid_video_frame(self): payload = Blob.from_local(video_path).to_descriptor() valid = VideoFrameDescriptor( - payload.uri, payload.offset, payload.length, 0 + payload.uri, payload.offset, payload.length, 0, -1, 0 ).serialize() with self.assertRaisesRegex(ValueError, "VideoFrameDescriptor"): table.add_batches([ diff --git a/paimon-python/pypaimon/tests/multimodal_video_test.py b/paimon-python/pypaimon/tests/multimodal_video_test.py index 918e917a32b3..d5f14d92d761 100644 --- a/paimon-python/pypaimon/tests/multimodal_video_test.py +++ b/paimon-python/pypaimon/tests/multimodal_video_test.py @@ -244,7 +244,7 @@ def decode_batch(decoder, indices, rows): def row(request, offset, length, index): return {"request": request, "video": VideoFrameDescriptor( - path, offset, length, index).serialize()} + path, offset, length, index, -1, 0).serialize()} collator = VideoFrameCollator( self.table, @@ -403,7 +403,7 @@ def new_input_stream(self, path): file_io = ResolvedFileIO() table = SimpleNamespace(raw_table=SimpleNamespace(file_io=file_io)) descriptor = VideoFrameDescriptor( - "oss://bucket/internal.video", 0, 14, 2 + "oss://bucket/internal.video", 0, 14, 2, -1, 0 ).serialize() collator = VideoFrameCollator( table, @@ -425,7 +425,7 @@ def _descriptor(self, name, data, frame_index): with open(path, "wb") as output: output.write(data) return VideoFrameDescriptor( - path, 0, len(data), frame_index + path, 0, len(data), frame_index, -1, 0 ).serialize() diff --git a/paimon-python/pypaimon/tests/video_format_test.py b/paimon-python/pypaimon/tests/video_format_test.py index 33d26e325d4c..e4c7ad174758 100644 --- a/paimon-python/pypaimon/tests/video_format_test.py +++ b/paimon-python/pypaimon/tests/video_format_test.py @@ -15,10 +15,16 @@ # specific language governing permissions and limitations # under the License. +import bisect +import io import struct import tempfile +import tracemalloc import unittest +import zlib +from fractions import Fraction from pathlib import Path +from unittest import mock from pypaimon.common.delta_varint_compressor import DeltaVarintCompressor from pypaimon.common.options import Options @@ -34,8 +40,16 @@ ) from pypaimon.table.row.generic_row import GenericRow from pypaimon.table.row.row_kind import RowKind +from pypaimon.table.row.video_keyframe_index import VideoKeyframeIndex from pypaimon.write.video_format_writer import VideoFormatWriter +try: + import av + import numpy as np +except ImportError: + av = None + np = None + class VideoFormatTest(unittest.TestCase): @@ -61,7 +75,7 @@ def tearDown(self): self.temp_dir.cleanup() def test_descriptor_round_trip_preserves_payload_and_frame(self): - descriptor = VideoFrameDescriptor("s3://bucket/a.video", 7, 99, 42) + descriptor = VideoFrameDescriptor("s3://bucket/a.video", 7, 99, 42, -1, 0) serialized = descriptor.serialize() self.assertTrue( @@ -80,12 +94,20 @@ def test_descriptor_round_trip_preserves_payload_and_frame(self): with self.assertRaisesRegex(ValueError, "trailing bytes"): VideoFrameDescriptor.deserialize(serialized + b"x") with self.assertRaisesRegex(ValueError, "non-negative"): - VideoFrameDescriptor("x", 0, 1, -1) + VideoFrameDescriptor("x", 0, 1, -1, -1, 0) + + indexed = VideoFrameDescriptor("s3://bucket/a.video", 7, 99, 42, 106, 8) + restored = VideoFrameDescriptor.deserialize(indexed.serialize()) + self.assertEqual(indexed, restored) + self.assertEqual( + BlobDescriptor("s3://bucket/a.video", 106, 8), + restored.keyframe_index_descriptor, + ) def test_cross_language_descriptor_fixture(self): fixture = self._fixture_bytes(self.DESCRIPTOR_FIXTURE) - expected = VideoFrameDescriptor("s3://bucket/视频.mp4", 7, 99, 42) - + expected = VideoFrameDescriptor( + "s3://bucket/视频.mp4", 7, 99, 42, 106, 8) self.assertEqual(fixture, expected.serialize()) self.assertEqual(expected, BlobDescriptor.deserialize(fixture)) @@ -114,12 +136,12 @@ def test_pack_raw_videos_and_map_frame_runs(self): with self.file_io.new_input_stream(target) as stream: meta = VideoFileMeta(stream, len(stored)) self.assertEqual(6, meta.record_count) - self.assertEqual((0, len(first_bytes), 0), meta.frame(0)) - self.assertEqual((0, len(first_bytes), 1), meta.frame(1)) + self.assertEqual((0, len(first_bytes), 0, -1, 0), meta.frame(0)) + self.assertEqual((0, len(first_bytes), 1, -1, 0), meta.frame(1)) self.assertEqual( - (len(first_bytes), len(second_bytes), 7), meta.frame(2) + (len(first_bytes), len(second_bytes), 7, -1, 0), meta.frame(2) ) - self.assertEqual((0, len(first_bytes), 4), meta.frame(3)) + self.assertEqual((0, len(first_bytes), 4, -1, 0), meta.frame(3)) self.assertIsNone(meta.frame(4)) self.assertIs(Blob.PLACE_HOLDER, meta.frame(5)) @@ -133,39 +155,335 @@ def test_pack_raw_videos_and_map_frame_runs(self): def test_cross_language_video_v1_fixture(self): fixture = self._fixture_bytes(self.VIDEO_FIXTURE) - fixture_path = self.root / "fixture.video" - fixture_path.write_bytes(fixture) - target = fixture_path.as_uri() + encoded = VideoKeyframeIndex( + [(0, 1)], [(0, 0, 0), (12, 36000, 2)] + ).serialize() + video = b"abc" + source_path = self.root / "indexed.mp4" + source_path.write_bytes(video + encoded) - with self.file_io.new_input_stream(target) as stream: - meta = VideoFileMeta(stream, len(fixture)) - self.assertEqual(7, meta.record_count) - self.assertEqual((0, 3, 2), meta.frame(0)) - self.assertEqual((0, 3, 3), meta.frame(1)) - self.assertIsNone(meta.frame(2)) - self.assertIs(Blob.PLACE_HOLDER, meta.frame(3)) - self.assertEqual((3, 4, 7), meta.frame(4)) - self.assertEqual((3, 4, 8), meta.frame(5)) - self.assertEqual((0, 3, 10), meta.frame(6)) - - written_target = (self.root / "written.video").as_uri() - writer = VideoFormatWriter( - self.file_io.new_output_stream(written_target), - file_path=written_target, - ) + def indexed_frame(frame_index): + descriptor = VideoFrameDescriptor( + source_path.as_uri(), 0, len(video), frame_index, + len(video), len(encoded) + ) + return Blob.from_descriptor( + self.file_io.uri_reader_factory.create(descriptor.uri), + descriptor, + ) + target = (self.root / "indexed.video").as_uri() + + writer = VideoFormatWriter(self.file_io.new_output_stream(target)) values = ( - self._source_frame("a.mp4", b"abc", 2), - self._source_frame("a.mp4", b"abc", 3), + indexed_frame(2), + indexed_frame(3), None, Blob.PLACE_HOLDER, self._source_frame("b.mp4", b"WXYZ", 7), self._source_frame("b.mp4", b"WXYZ", 8), - self._source_frame("a.mp4", b"abc", 10), + indexed_frame(10), ) for value in values: writer.add_element(GenericRow([value], [self.field], RowKind.INSERT)) writer.close() - self.assertEqual(fixture, (self.root / "written.video").read_bytes()) + + stored = (self.root / "indexed.video").read_bytes() + self.assertEqual(fixture, stored) + with self.file_io.new_input_stream(target) as stream: + meta = VideoFileMeta(stream, len(stored)) + self.assertEqual( + (0, len(video), 2, 7, len(encoded)), meta.frame(0) + ) + self.assertEqual((3, 4, 7, -1, 0), meta.frame(4)) + serialized = self._read(target, row_indices=[0])[0] + value = VideoFrameDescriptor.deserialize(serialized) + self.assertEqual(2, value.frame_index) + mapping_descriptor = value.keyframe_index_descriptor + self.assertEqual( + encoded, + Blob.from_file( + self.file_io, + mapping_descriptor.uri, + mapping_descriptor.offset, + mapping_descriptor.length, + ).to_data(), + ) + + rewritten = (self.root / "rewritten.video").as_uri() + writer = VideoFormatWriter(self.file_io.new_output_stream(rewritten)) + writer.add_element(GenericRow([ + Blob.from_bytes(serialized, file_io=self.file_io) + ], [self.field], RowKind.INSERT)) + writer.close() + rewritten_value = VideoFrameDescriptor.deserialize( + self._read(rewritten)[0] + ) + rewritten_index = rewritten_value.keyframe_index_descriptor + self.assertEqual( + encoded, + Blob.from_file( + self.file_io, + rewritten_index.uri, + rewritten_index.offset, + rewritten_index.length, + ).to_data(), + ) + + @unittest.skipUnless( + av is not None and np is not None, + "PyAV and NumPy are required for sparse video seek validation", + ) + def test_real_mp4_index_supports_sparse_random_reads(self): + payload = self._real_mp4() + metadata_ranges, keyframes, has_b_frames, is_vfr = ( + self._mp4_seek_index(payload)) + index = VideoKeyframeIndex(metadata_ranges, keyframes) + source = self.root / "real-indexed.mp4" + source.write_bytes(payload) + target_frames = [44, 1, 30, 14, 59, 15] + target = (self.root / "real-indexed.video").as_uri() + + writer = VideoFormatWriter(self.file_io.new_output_stream(target)) + for frame_index in target_frames: + descriptor = VideoFrameDescriptor( + source.as_uri(), 0, len(payload), frame_index, + -1, 0) + value = Blob.from_descriptor( + self.file_io.uri_reader_factory.create(descriptor.uri), + descriptor, + ) + writer.add_element( + GenericRow([value], [self.field], RowKind.INSERT)) + writer.close() + + descriptors = [ + VideoFrameDescriptor.deserialize(value) + for value in self._read(target) + ] + payload_descriptor = descriptors[0].payload_descriptor + index_descriptor = descriptors[0].keyframe_index_descriptor + stored_payload = Blob.from_file( + self.file_io, + payload_descriptor.uri, + payload_descriptor.offset, + payload_descriptor.length, + ).to_data() + stored_index = VideoKeyframeIndex.deserialize( + Blob.from_file( + self.file_io, + index_descriptor.uri, + index_descriptor.offset, + index_descriptor.length, + ).to_data()) + + self.assertEqual(target_frames, [value.frame_index for value in descriptors]) + self.assertEqual(index.metadata_ranges, stored_index.metadata_ranges) + self.assertEqual(index.keyframes, stored_index.keyframes) + self.assertTrue(has_b_frames) + self.assertTrue(is_vfr) + + with av.open(io.BytesIO(stored_payload)) as container: + expected = [ + frame.to_ndarray(format="rgb24") + for frame in container.decode(video=0) + ] + ordinals = [entry[0] for entry in stored_index.keyframes] + sparse_read_sizes = [] + for target_frame in target_frames: + keyframe = bisect.bisect_right(ordinals, target_frame) - 1 + range_end = ( + stored_index.keyframes[keyframe + 2][2] + if keyframe + 2 < len(stored_index.keyframes) + else len(stored_payload) + ) + for retry in range(3): + anchor = max(0, keyframe - retry) + anchor_ordinal, anchor_pts, unused_position = ( + stored_index.keyframes[anchor]) + range_start = stored_index.keyframes[ + max(0, anchor - 1)][2] + sparse, available_bytes = self._sparse_video( + stored_payload, + list(stored_index.metadata_ranges) + + [(range_start, range_end - range_start)], + ) + try: + actual = self._decode_sparse_frame( + sparse, anchor_ordinal, anchor_pts, target_frame) + break + except (av.error.FFmpegError, ValueError): + continue + else: + self.fail("Sparse video retries did not reach the target frame.") + np.testing.assert_array_equal(expected[target_frame], actual) + sparse_read_sizes.append(available_bytes) + self.assertLess(min(sparse_read_sizes), len(stored_payload)) + + def test_target_size_counts_buffered_keyframe_index(self): + mapping = VideoKeyframeIndex([], [(0, 0, 0)]).serialize() + video = b"video" + source = self.root / "target-size.mp4" + source.write_bytes(video + mapping) + descriptor = VideoFrameDescriptor( + source.as_uri(), 0, len(video), 0, len(video), len(mapping) + ) + blob = Blob.from_descriptor( + self.file_io.uri_reader_factory.create(descriptor.uri), descriptor + ) + writer = VideoFormatWriter(io.BytesIO()) + + writer.add_element(GenericRow([blob], [self.field], RowKind.INSERT)) + + self.assertFalse(writer.reach_target_size(len(video) + len(mapping) + 1)) + self.assertTrue(writer.reach_target_size(len(video) + len(mapping))) + + def test_rejects_invalid_keyframe_index(self): + video = b"video" + mapping = b"mapping" + source = self.root / "invalid-index.mp4" + source.write_bytes(video + mapping) + descriptor = VideoFrameDescriptor( + source.as_uri(), 0, len(video), 0, len(video), len(mapping) + ) + blob = Blob.from_descriptor( + self.file_io.uri_reader_factory.create(descriptor.uri), descriptor + ) + writer = VideoFormatWriter(io.BytesIO()) + + with self.assertRaisesRegex(ValueError, "Invalid video keyframe index"): + writer.add_element( + GenericRow([blob], [self.field], RowKind.INSERT) + ) + + def test_rejects_oversized_keyframe_indexes_before_fetch(self): + descriptor = VideoFrameDescriptor( + (self.root / "missing.mp4").as_uri(), + 0, + 1, + 0, + 1, + VideoFormatWriter.MAX_KEYFRAME_INDEX_BYTES + 1, + ) + blob = Blob.from_descriptor( + self.file_io.uri_reader_factory.create(descriptor.uri), descriptor + ) + + with self.assertRaisesRegex(ValueError, "keyframe index length.*limit"): + VideoFormatWriter(io.BytesIO()).add_element( + GenericRow([blob], [self.field], RowKind.INSERT) + ) + + def test_rejects_cumulative_keyframe_indexes_before_fetch(self): + mapping = VideoKeyframeIndex([], [(0, 0, 0)]).serialize() + first = self._indexed_frame("first-indexed.mp4", b"a", mapping) + second = self._indexed_frame("second-indexed.mp4", b"b", mapping) + writer = VideoFormatWriter(io.BytesIO()) + + with mock.patch.object( + VideoFormatWriter, + 'MAX_TOTAL_KEYFRAME_INDEX_BYTES', + len(mapping)): + writer.add_element( + GenericRow([first], [self.field], RowKind.INSERT) + ) + with self.assertRaisesRegex( + ValueError, "Buffered video keyframe indexes.*limit"): + writer.add_element( + GenericRow([second], [self.field], RowKind.INSERT) + ) + + def test_rejects_seek_offsets_outside_video_payload(self): + video = b"video" + mapping = VideoKeyframeIndex( + [(0, len(video) + 1)], [(0, 0, 0)] + ).serialize() + source = self.root / "out-of-range-index.mp4" + source.write_bytes(video + mapping) + descriptor = VideoFrameDescriptor( + source.as_uri(), 0, len(video), 0, len(video), len(mapping) + ) + blob = Blob.from_descriptor( + self.file_io.uri_reader_factory.create(descriptor.uri), descriptor + ) + + with self.assertRaisesRegex(ValueError, "video payload"): + VideoFormatWriter(io.BytesIO()).add_element( + GenericRow([blob], [self.field], RowKind.INSERT) + ) + + def test_rejects_compressed_invalid_index_without_expanding_it(self): + count = 500_000 + entry = VideoKeyframeIndex.ENTRY.pack(0, 0, 0) + mapping = VideoKeyframeIndex.HEADER.pack( + VideoKeyframeIndex.VERSION, VideoKeyframeIndex.MAGIC, 0, count + ) + zlib.compress(entry * count) + video = b"video" + source = self.root / "compressed-invalid-index.mp4" + source.write_bytes(video + mapping) + descriptor = VideoFrameDescriptor( + source.as_uri(), 0, len(video), 0, len(video), len(mapping) + ) + blob = Blob.from_descriptor( + self.file_io.uri_reader_factory.create(descriptor.uri), descriptor + ) + writer = VideoFormatWriter(io.BytesIO()) + + self.assertLess(len(mapping), 20_000) + tracemalloc.start() + try: + with self.assertRaisesRegex( + ValueError, + "Keyframe ordinals must be strictly increasing"): + writer.add_element( + GenericRow([blob], [self.field], RowKind.INSERT) + ) + _, peak = tracemalloc.get_traced_memory() + finally: + tracemalloc.stop() + self.assertLess(peak, 1_000_000) + + def test_rejects_inconsistent_keyframe_indexes_for_same_payload(self): + video = b"video" + first_index = VideoKeyframeIndex( + [(0, 1)], [(0, 0, 0), (12, 36000, 4)] + ).serialize() + second_index = VideoKeyframeIndex( + [], [(0, 0, 0), (6, 18000, 3)] + ).serialize() + source = self.root / "inconsistent-index.mp4" + source.write_bytes(video + first_index + second_index) + + def frame(frame_index, index_offset=-1, index_length=0): + descriptor = VideoFrameDescriptor( + source.as_uri(), 0, len(video), frame_index, + index_offset, index_length + ) + return Blob.from_descriptor( + self.file_io.uri_reader_factory.create(descriptor.uri), + descriptor, + ) + + unindexed = frame(0) + first = frame(1, len(video), len(first_index)) + second = frame( + 2, len(video) + len(first_index), len(second_index) + ) + for values in ( + (unindexed, first), + (first, unindexed), + (first, second), + ): + with self.subTest(values=values): + writer = VideoFormatWriter(io.BytesIO()) + writer.add_element( + GenericRow([values[0]], [self.field], RowKind.INSERT) + ) + with self.assertRaisesRegex(ValueError, "same payload"): + writer.add_element( + GenericRow( + [values[1]], [self.field], RowKind.INSERT + ) + ) def test_selection_keeps_logical_frame_positions(self): target = (self.root / "selection.video").as_uri() @@ -232,6 +550,7 @@ def test_rejects_non_video_frame_input(self): def test_rejects_out_of_range_run_reference(self): target_path = self.root / "corrupt.video" indexes = [ + DeltaVarintCompressor.compress([]), DeltaVarintCompressor.compress([]), DeltaVarintCompressor.compress([1]), DeltaVarintCompressor.compress([0]), @@ -240,7 +559,7 @@ def test_rejects_out_of_range_run_reference(self): target_path.write_bytes( b"".join(indexes) + struct.pack( - 'I4s", payload, offset) + header_size = 8 + if size == 1: + size = struct.unpack_from(">Q", payload, offset + 8)[0] + header_size = 16 + elif size == 0: + size = len(payload) - offset + self.assertGreaterEqual(size, header_size) + self.assertLessEqual(offset + size, len(payload)) + metadata_ranges.append( + (offset, header_size if box_type == b"mdat" else size)) + box_types.append(box_type) + offset += size + self.assertEqual(b"moov", box_types[-1]) + + with av.open(io.BytesIO(payload)) as container: + self.assertEqual(2, len(container.streams.video)) + stream = container.streams.video[0] + packet_positions = {} + has_b_frames = False + for packet in container.demux(stream): + if packet.pts is not None and packet.dts is not None: + has_b_frames |= packet.pts != packet.dts + if (packet.is_keyframe and packet.pts is not None + and packet.pos is not None and packet.pos >= 0): + packet_positions[int(packet.pts)] = int(packet.pos) + + with av.open(io.BytesIO(payload)) as container: + stream = container.streams.video[0] + keyframes = [] + frame_pts = [] + for ordinal, frame in enumerate(container.decode(stream)): + frame_pts.append(int(frame.pts)) + if frame.key_frame: + keyframes.append(( + ordinal, + int(frame.pts), + packet_positions[int(frame.pts)], + )) + self.assertEqual(60, len(frame_pts)) + is_vfr = len({ + right - left for left, right in zip(frame_pts, frame_pts[1:]) + }) > 1 + return metadata_ranges, keyframes, has_b_frames, is_vfr + + def _decode_sparse_frame( + self, source, anchor_ordinal, anchor_pts, target_ordinal): + with av.open(source) as container: + stream = container.streams.video[0] + container.seek( + anchor_pts, backward=True, any_frame=False, stream=stream) + ordinal = anchor_ordinal + found_anchor = False + for frame in container.decode(stream): + if not found_anchor: + if frame.pts != anchor_pts or not frame.key_frame: + continue + found_anchor = True + if ordinal == target_ordinal: + return frame.to_ndarray(format="rgb24") + ordinal += 1 + raise ValueError("Sparse video read did not reach the target frame.") + + @staticmethod + def _sparse_video(payload, ranges): + sparse = bytearray(len(payload)) + loaded = bytearray(len(payload)) + for offset, length in ranges: + sparse[offset:offset + length] = payload[offset:offset + length] + loaded[offset:offset + length] = b'\1' * length + return io.BytesIO(sparse), sum(loaded) + @staticmethod def _fixture_bytes(path): hex_value = "".join( diff --git a/paimon-python/pypaimon/tests/video_keyframe_index_test.py b/paimon-python/pypaimon/tests/video_keyframe_index_test.py new file mode 100644 index 000000000000..8ead98efe199 --- /dev/null +++ b/paimon-python/pypaimon/tests/video_keyframe_index_test.py @@ -0,0 +1,142 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import io +import struct +import unittest +import zlib +from unittest import mock + +import pypaimon.table.row.video_keyframe_index as keyframe_index_module +from pypaimon.table.row.video_keyframe_index import VideoKeyframeIndex + + +class VideoKeyframeIndexTest(unittest.TestCase): + + def test_iso_bmff_metadata_ranges_exclude_padding_bodies(self): + def box(box_type, body): + return struct.pack(">I4s", len(body) + 8, box_type) + body + + boxes = [ + box(b"ftyp", b"isom"), + box(b"free", b"x" * 1024), + box(b"skip", b"y" * 512), + box(b"mdat", b"video"), + box(b"moov", b"metadata"), + ] + payload = b''.join(boxes) + offsets = [] + offset = 0 + for value in boxes: + offsets.append(offset) + offset += len(value) + + self.assertEqual( + [ + (offsets[0], len(boxes[0])), + (offsets[1], 8), + (offsets[2], 8), + (offsets[3], 8), + (offsets[4], len(boxes[4])), + ], + VideoKeyframeIndex._iso_bmff_metadata_ranges( + io.BytesIO(payload), len(payload)), + ) + + def test_round_trip_and_validation(self): + index = VideoKeyframeIndex( + [(0, 128), (256, 64)], + [(0, 500, 512), (60, 9000, 4096)], + ) + restored = VideoKeyframeIndex.deserialize(index.serialize()) + self.assertEqual(index.metadata_ranges, restored.metadata_ranges) + self.assertEqual(index.keyframes, restored.keyframes) + self.assertEqual(17, index.HEADER.size) + self.assertFalse(hasattr(restored, 'time_base')) + self.assertFalse(hasattr(restored, 'frame_count')) + self.assertFalse(hasattr(restored, 'stream_index')) + for ranges in ([(-1, 1)], [(0, 0)], [(4, 2), (3, 1)]): + with self.subTest(ranges=ranges): + with self.assertRaises(ValueError): + VideoKeyframeIndex(ranges, [(0, 0, 0)]) + for anchors in ( + [], [(1, 0, 0)], [(0, 5, 0), (0, 9, 1)], + [(0, 5, 0), (2, 5, 1)], + [(0, 5, 0), (2, 9, 0)], + [(0, 5, 0), (2, 9, 2), (1, 10, 3)], + ): + with self.subTest(anchors=anchors): + with self.assertRaises(ValueError): + VideoKeyframeIndex([], anchors) + for data in (b'', index.serialize()[:-1], index.serialize() + b'junk'): + with self.assertRaises(ValueError): + VideoKeyframeIndex.deserialize(data) + header = index.serialize()[:index.HEADER.size] + with self.assertRaises(ValueError): + VideoKeyframeIndex.deserialize( + header + zlib.compress(struct.pack(' None: @@ -86,14 +92,29 @@ def add_element(self, row) -> None: "VideoFrameDescriptor." ) payload = frame.payload_descriptor + keyframe_index_descriptor = frame.keyframe_index_descriptor ordinal = self._physical_videos.get(payload) if ordinal is None: + keyframe_index = self._keyframe_index(value, frame) length = self._write_video_payload(value) ordinal = len(self._physical_lengths) self._physical_lengths.append(length) + self._keyframe_indexes.append(keyframe_index) + self._keyframe_index_bytes += len(keyframe_index) self._physical_videos[payload] = ordinal + self._physical_video_keyframe_indexes[ + payload] = keyframe_index_descriptor + elif (self._physical_video_keyframe_indexes[payload] + != keyframe_index_descriptor): + raise ValueError( + "Video frames for the same payload must use the same " + "keyframe index." + ) self._append(ordinal, frame.frame_index) + def reach_target_size(self, target_size: int) -> bool: + return self.position + self._keyframe_index_bytes >= target_size + @property def physical_video_count(self) -> int: return len(self._physical_lengths) @@ -106,7 +127,12 @@ def close(self) -> None: if self._closed: return self._flush_run() + for keyframe_index in self._keyframe_indexes: + self.output_stream.write(keyframe_index) physical_index = DeltaVarintCompressor.compress(self._physical_lengths) + keyframe_length_index = DeltaVarintCompressor.compress( + [len(mapping) for mapping in self._keyframe_indexes] + ) run_length_index = DeltaVarintCompressor.compress(self._run_lengths) run_reference_index = DeltaVarintCompressor.compress( self._run_references @@ -115,27 +141,73 @@ def close(self) -> None: self._run_first_frames ) for index in ( - physical_index, - run_length_index, - run_reference_index, - first_frame_index, - ): + physical_index, + keyframe_length_index, + run_length_index, + run_reference_index, + first_frame_index): self.output_stream.write(index) - self.output_stream.write(struct.pack( - ' self.MAX_KEYFRAME_INDEX_BYTES: + raise ValueError( + "Video keyframe index length %s exceeds the %s-byte limit." + % (length, self.MAX_KEYFRAME_INDEX_BYTES) + ) + if (self._keyframe_index_bytes + > self.MAX_TOTAL_KEYFRAME_INDEX_BYTES - length): + raise ValueError( + "Buffered video keyframe indexes exceed the %s-byte limit." + % self.MAX_TOTAL_KEYFRAME_INDEX_BYTES + ) + def _write_video_payload(self, blob: BlobRef) -> int: start = self.position stream = blob.new_input_stream() From 31275f2ab0bb1e8a6c380732aa2bd442cfdb7e72 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Tue, 22 Sep 2026 00:22:40 -0700 Subject: [PATCH 2/8] [python] Read indexed video ranges on demand --- docs/docs/multimodal-table/video.mdx | 6 +- docs/docs/pypaimon/lerobot.md | 5 +- .../pypaimon/multimodal/lerobot/dataset.py | 189 ++++++++++++++---- .../pypaimon/tests/multimodal_lerobot_test.py | 64 +++++- 4 files changed, 213 insertions(+), 51 deletions(-) diff --git a/docs/docs/multimodal-table/video.mdx b/docs/docs/multimodal-table/video.mdx index d8833ee48fb0..aaab226bbaaa 100644 --- a/docs/docs/multimodal-table/video.mdx +++ b/docs/docs/multimodal-table/video.mdx @@ -99,9 +99,9 @@ one run; negative references represent NULL or data-evolution placeholders. Each video index stores initialization ranges such as MP4 `moov` and keyframe positions. A cold -reader fetches these ranges and a bounded GOP window around the target. Zero length selects the -scan fallback. PyPaimon generates indexes for supported ISO BMFF videos and preserves them during -rewrites. See the +reader prefetches these ranges and a bounded GOP window, then fetches uncached decoder reads on +demand. Zero length selects the scan fallback. PyPaimon generates indexes for supported ISO BMFF +videos and preserves them during rewrites. See the [file format specification](../concepts/spec/fileformat#video). ## Reuse, Rolling, and Compaction diff --git a/docs/docs/pypaimon/lerobot.md b/docs/docs/pypaimon/lerobot.md index 34db12bf63d6..e97ee770e7cd 100644 --- a/docs/docs/pypaimon/lerobot.md +++ b/docs/docs/pypaimon/lerobot.md @@ -222,8 +222,9 @@ loader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=4) ``` Without `tag_name`, the latest snapshots are used. Frame lookups use the BTree -on `index`; payloads remain lazy. Indexed videos use PyAV range reads for metadata and target GOPs; -unindexed videos use the TorchCodec/PyAV scan path. Set `video_backend` to force either decoder. +on `index`; payloads remain lazy. Indexed videos prefetch metadata and target GOPs, then fetch +uncached PyAV reads on demand. Unindexed videos use the TorchCodec/PyAV scan path. Set +`video_backend` to force either decoder. Subclass `PaimonDatasetReader` for a custom logical frame layout: diff --git a/paimon-python/pypaimon/multimodal/lerobot/dataset.py b/paimon-python/pypaimon/multimodal/lerobot/dataset.py index fbb32ccb5cd9..7a56919049d2 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/dataset.py +++ b/paimon-python/pypaimon/multimodal/lerobot/dataset.py @@ -18,12 +18,12 @@ """LeRobot-compatible map-style reads from a multimodal Paimon table.""" import bisect +import io import json import math import operator import pickle import sys -import tempfile import zlib from abc import ABC, abstractmethod from collections import OrderedDict @@ -1542,6 +1542,118 @@ def _open_torchcodec_decoder(stream): return VideoDecoder(stream.read(), seek_mode="exact") +class _RangeBackedVideo(io.RawIOBase): + + def __init__(self, length, read_ranges): + if length < 0: + raise ValueError("Video length must be non-negative.") + self._length = length + self._read_ranges = read_ranges + self._position = 0 + self._segments = [] + + def readable(self): + return True + + def seekable(self): + return True + + def tell(self): + return self._position + + def seek(self, offset, whence=io.SEEK_SET): + offset = operator.index(offset) + if whence == io.SEEK_SET: + position = offset + elif whence == io.SEEK_CUR: + position = self._position + offset + elif whence == io.SEEK_END: + position = self._length + offset + else: + raise ValueError("Invalid whence: %s" % whence) + if position < 0: + raise ValueError("Negative seek position: %s" % position) + self._position = position + return position + + def read(self, size=-1): + read_all = size is None or size < 0 + if read_all: + end = self._length + else: + end = min(self._position + operator.index(size), self._length) + if self._position >= end: + return b'' + start = self._position + self._ensure([(start, end - start)]) + self._position = end + return self._cached(start, end) + + def readinto(self, value): + data = self.read(len(value)) + value[:len(data)] = data + return len(data) + + def prefetch(self, ranges): + self._ensure(ranges) + + def _ensure(self, ranges): + requested = [] + for offset, length in _merge_video_ranges(ranges): + if offset < 0 or length < 0 or offset + length > self._length: + raise ValueError("Video byte range is outside the payload.") + requested.extend(self._missing(offset, offset + length)) + missing = _merge_video_ranges(requested) + if not missing: + return + bodies = self._read_ranges(missing) + if len(bodies) != len(missing): + raise IOError("Video range reader returned an invalid result count.") + for (offset, length), body in zip(missing, bodies): + body = bytes(body) + if len(body) != length: + raise IOError( + "Video byte range changed while reading: expected %d, " + "read %d." % (length, len(body)) + ) + self._segments.append((offset, offset + length, body)) + self._segments.sort(key=lambda value: value[0]) + + def _missing(self, start, end): + missing = [] + position = start + for segment_start, segment_end, unused_body in self._segments: + if segment_end <= position: + continue + if segment_start >= end: + break + if segment_start > position: + missing.append(( + position, min(segment_start, end) - position)) + position = max(position, segment_end) + if position >= end: + break + if position < end: + missing.append((position, end - position)) + return missing + + def _cached(self, start, end): + result = [] + position = start + for segment_start, segment_end, body in self._segments: + if segment_end <= position: + continue + if segment_start > position: + break + limit = min(segment_end, end) + result.append(body[ + position - segment_start:limit - segment_start]) + position = limit + if position == end: + return b''.join(result) + raise IOError("Video range cache is incomplete.") + + class _PyAVVideoDecoder: # Reuse common overlapping delta windows without retaining a whole video. @@ -1641,7 +1753,7 @@ def _read_indexed(self, indices): if any(index < 0 for index in requested): raise IndexError("Video frame index is out of range.") missing = sorted(set(requested) - set(self._cache)) - decoded = self._decode_sparse(missing) if missing else {} + decoded = self._decode_indexed(missing) if missing else {} try: return [ self._cache[index] if index in self._cache else decoded[index] @@ -1652,55 +1764,41 @@ def _read_indexed(self, indices): "Video frame index %d is out of range." % error.args[0] ) from error - def _decode_sparse(self, indices): - decoded = {} + def _decode_indexed(self, indices): import av - pending = list(indices) - last_error = None - for retry in range(3): - groups, ranges = self._sparse_plan(pending, retry) - bodies = self._read_video_ranges(ranges) - failed = [] - with tempfile.TemporaryFile() as sparse: - sparse.truncate(self._video_length()) - for (offset, unused_length), body in zip(ranges, bodies): - sparse.seek(offset) - sparse.write(body) - sparse.seek(0) - with av.open(sparse) as container: - stream = container.streams.video[0] - for anchor, targets in sorted(groups.items()): - try: - self._decode_sparse_group( - container, - stream, - anchor, - sorted(targets), - decoded, - ) - except (av.error.FFmpegError, ValueError) as error: - last_error = error - failed.extend(targets) - if not failed: - return decoded - pending = failed - raise ValueError( - "Cannot decode video frames from persisted seek index." - ) from last_error + groups, ranges = self._indexed_plan(indices) + decoded = {} + source = _RangeBackedVideo( + self._video_length(), self._read_video_ranges) + try: + source.prefetch(_merge_video_ranges( + list(self._keyframe_index.metadata_ranges) + ranges)) + with av.open(source) as container: + stream = container.streams.video[0] + for anchor, targets in sorted(groups.items()): + self._decode_indexed_group( + container, + stream, + anchor, + sorted(targets), + decoded, + ) + return decoded + finally: + source.close() - def _sparse_plan(self, indices, retry): + def _indexed_plan(self, indices): mapping = self._keyframe_index groups = {} - ranges = list(mapping.metadata_ranges) + ranges = [] for index in indices: preceding = bisect.bisect_right( self._anchor_ordinals, index) - 1 if preceding < 0: raise IndexError( "Video frame index %d is out of range." % index) - anchor = max(0, preceding - retry) - groups.setdefault(anchor, []).append(index) - start = mapping.keyframes[max(0, anchor - 1)][2] + groups.setdefault(preceding, []).append(index) + start = mapping.keyframes[max(0, preceding - 1)][2] range_end = preceding + 2 end = ( mapping.keyframes[range_end][2] @@ -1710,8 +1808,13 @@ def _sparse_plan(self, indices, retry): ranges.append((start, end - start)) return groups, _merge_video_ranges(ranges) - def _decode_sparse_group( - self, container, stream, anchor, targets, decoded): + def _decode_indexed_group( + self, + container, + stream, + anchor, + targets, + decoded): anchor_ordinal, anchor_pts, unused_position = ( self._keyframe_index.keyframes[anchor]) container.seek( diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index 056e7b641fdf..093e9b655b08 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -48,6 +48,7 @@ from pypaimon.multimodal.lerobot.dataset import ( _PaimonLeRobotMetadata, _PyAVVideoDecoder, + _RangeBackedVideo, _arrow_rows, _decode_video_frames, _decode_video_rows, @@ -387,6 +388,43 @@ def close(self): finally: decoder.close() + def test_range_backed_video_fetches_only_missing_bytes(self): + payload = bytes(range(32)) + calls = [] + + def read_ranges(ranges): + calls.append(list(ranges)) + return [ + payload[offset:offset + length] + for offset, length in ranges + ] + + reader = _RangeBackedVideo(len(payload), read_ranges) + reader.prefetch([(0, 4), (12, 4)]) + self.assertEqual([[(0, 4), (12, 4)]], calls) + + calls.clear() + reader.seek(2) + self.assertEqual(payload[2:4], reader.read(2)) + self.assertEqual([], calls) + + reader.seek(6) + self.assertEqual(payload[6:14], reader.read(8)) + self.assertEqual([[(6, 6)]], calls) + + calls.clear() + reader.seek(18) + output = bytearray(4) + self.assertEqual(4, reader.readinto(output)) + self.assertEqual(payload[18:22], bytes(output)) + self.assertEqual([[(18, 4)]], calls) + + self.assertEqual(len(payload) - 3, reader.seek(-3, io.SEEK_END)) + self.assertEqual(payload[-3:], reader.read()) + with self.assertRaisesRegex(ValueError, "Negative seek"): + reader.seek(-1) + reader.close() + @unittest.skipUnless( av is not None and importlib.util.find_spec("torch") is not None, "PyAV and Torch are required for video decoding", @@ -445,11 +483,31 @@ def test_indexed_pyav_decoder_handles_b_frames_and_fragmented_video(self): source = io.BytesIO(payload) source.video_length = len(payload) + range_calls = [] + + def read_ranges(ranges): + range_calls.append(list(ranges)) + return [ + payload[offset:offset + length] + for offset, length in ranges + ] + + source.video_read_ranges = read_ranges decoder = _PyAVVideoDecoder(source, index) try: - actual = decoder.get_frames_at( - indices=list(range(len(expected))) - ).data.permute(0, 2, 3, 1).numpy() + indices = list(range(len(expected))) + if codec == "mpeg4": + plan = ({0: indices}, []) + with patch.object( + decoder, "_indexed_plan", return_value=plan): + actual = decoder.get_frames_at( + indices=indices + ).data.permute(0, 2, 3, 1).numpy() + self.assertGreater(len(range_calls), 1) + else: + actual = decoder.get_frames_at( + indices=indices + ).data.permute(0, 2, 3, 1).numpy() np.testing.assert_array_equal(expected, actual) finally: decoder.close() From 55993488663516f6227ff47826ad352ea4307c1e Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Tue, 22 Sep 2026 01:10:24 -0700 Subject: [PATCH 3/8] [python] Bound indexed video decoder resources --- docs/docs/concepts/spec/fileformat.md | 3 +- .../format/blob/VideoKeyframeIndex.java | 4 ++ .../format/blob/VideoFileFormatTest.java | 19 +++++++ .../pypaimon/multimodal/lerobot/dataset.py | 5 +- .../table/row/video_keyframe_index.py | 33 ++++++++++-- .../pypaimon/tests/multimodal_lerobot_test.py | 26 +++++++++ .../pypaimon/tests/video_format_test.py | 2 +- .../tests/video_keyframe_index_test.py | 54 ++++++++++++++++++- 8 files changed, 137 insertions(+), 9 deletions(-) diff --git a/docs/docs/concepts/spec/fileformat.md b/docs/docs/concepts/spec/fileformat.md index 4d409a78d308..c3920db7d6e4 100644 --- a/docs/docs/concepts/spec/fileformat.md +++ b/docs/docs/concepts/spec/fileformat.md @@ -448,7 +448,8 @@ A keyframe-index block has a 17-byte header: version (`1`, uint8), magic (`0x564 metadata-range count (uint32), and keyframe count (uint32). It then stores metadata `(offset, length)` pairs (two int64 values) and zlib-compressed `(frame ordinal, PTS, packet position)` keyframe entries (three int64 values). Offsets are relative to the encoded video; writers reject -out-of-range values. One block is limited to 16 MiB and all blocks in one file to 64 MiB. +out-of-range values. All numeric fields are little endian. One block is limited to 65,536 +keyframes and 16 MiB; all blocks in one file are limited to 64 MiB. The index covers the first video stream. Its time base remains in the video. A reader fetches the metadata and target GOP, seeks by PTS, and decodes forward in presentation order. It may include the diff --git a/paimon-format/src/main/java/org/apache/paimon/format/blob/VideoKeyframeIndex.java b/paimon-format/src/main/java/org/apache/paimon/format/blob/VideoKeyframeIndex.java index 15c01adb9b0a..8892f0f4a85f 100644 --- a/paimon-format/src/main/java/org/apache/paimon/format/blob/VideoKeyframeIndex.java +++ b/paimon-format/src/main/java/org/apache/paimon/format/blob/VideoKeyframeIndex.java @@ -31,6 +31,7 @@ final class VideoKeyframeIndex { private static final int HEADER_LENGTH = Byte.BYTES + Long.BYTES + Integer.BYTES * 2; private static final int METADATA_RANGE_LENGTH = Long.BYTES * 2; private static final int ENTRY_LENGTH = Long.BYTES * 3; + static final long MAX_KEYFRAME_COUNT = 64 * 1024; private VideoKeyframeIndex() {} @@ -53,6 +54,9 @@ static void validate(byte[] data, long payloadLength) { if (keyframeCount == 0) { throw invalid("empty"); } + if (keyframeCount > MAX_KEYFRAME_COUNT) { + throw invalid("keyframe count exceeds limit"); + } long entriesOffset = HEADER_LENGTH + metadataRangeCount * METADATA_RANGE_LENGTH; if (entriesOffset >= data.length) { diff --git a/paimon-format/src/test/java/org/apache/paimon/format/blob/VideoFileFormatTest.java b/paimon-format/src/test/java/org/apache/paimon/format/blob/VideoFileFormatTest.java index 4d3f98e906e8..adf9f9ac500d 100644 --- a/paimon-format/src/test/java/org/apache/paimon/format/blob/VideoFileFormatTest.java +++ b/paimon-format/src/test/java/org/apache/paimon/format/blob/VideoFileFormatTest.java @@ -48,6 +48,8 @@ import org.junit.jupiter.api.io.TempDir; import java.io.IOException; +import java.nio.ByteBuffer; +import java.nio.ByteOrder; import java.nio.charset.StandardCharsets; import java.nio.file.Files; import java.util.ArrayList; @@ -214,6 +216,23 @@ public void testRejectInvalidKeyframeIndex() throws IOException { .hasMessageContaining("Invalid video keyframe index"); } + @Test + public void testRejectTooManyKeyframes() { + byte[] mapping = + ByteBuffer.allocate(18) + .order(ByteOrder.LITTLE_ENDIAN) + .put((byte) 1) + .putLong(0x564944454F4B4649L) + .putInt(0) + .putInt((int) VideoKeyframeIndex.MAX_KEYFRAME_COUNT + 1) + .put((byte) 0) + .array(); + + assertThatThrownBy(() -> VideoKeyframeIndex.validate(mapping, 1)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("keyframe count exceeds limit"); + } + @Test public void testRejectOversizedKeyframeIndexesBeforeFetch() throws IOException { String missing = new Path(tempPath.resolve("missing.mp4").toUri()).toString(); diff --git a/paimon-python/pypaimon/multimodal/lerobot/dataset.py b/paimon-python/pypaimon/multimodal/lerobot/dataset.py index 7a56919049d2..fbe3c26a9e54 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/dataset.py +++ b/paimon-python/pypaimon/multimodal/lerobot/dataset.py @@ -1514,7 +1514,10 @@ def _open_video_decoder(stream, backend=None): if not isinstance(keyframe_index, VideoKeyframeIndex): keyframe_index = None if backend is None and keyframe_index is not None: - return _PyAVVideoDecoder(stream, keyframe_index) + try: + return _PyAVVideoDecoder(stream, keyframe_index) + except (ImportError, OSError, RuntimeError): + stream.seek(0) if backend in (None, "torchcodec"): try: return _open_torchcodec_decoder(stream) diff --git a/paimon-python/pypaimon/table/row/video_keyframe_index.py b/paimon-python/pypaimon/table/row/video_keyframe_index.py index 1a265f80271e..6f7ecf925d9d 100644 --- a/paimon-python/pypaimon/table/row/video_keyframe_index.py +++ b/paimon-python/pypaimon/table/row/video_keyframe_index.py @@ -30,6 +30,7 @@ class VideoKeyframeIndex: HEADER = struct.Struct('= self.MAX_KEYFRAME_COUNT: + raise ValueError( + "Video keyframe index exceeds the %s-entry limit." + % self.MAX_KEYFRAME_COUNT + ) + normalized_keyframes.append(( + operator.index(ordinal), operator.index(pts), + operator.index(position), + )) + self.keyframes = tuple(normalized_keyframes) previous_end = 0 for offset, length in self.metadata_ranges: if offset < previous_end or length <= 0: @@ -108,6 +116,11 @@ def inspect(cls, source, payload_length): ) if pts in packet_positions: raise ValueError("Video has duplicate keyframe timestamps.") + if len(packet_positions) >= cls.MAX_KEYFRAME_COUNT: + raise ValueError( + "Video keyframe index exceeds the %s-entry limit." + % cls.MAX_KEYFRAME_COUNT + ) packet_positions[pts] = int(packet.pos) source.seek(0) @@ -133,6 +146,11 @@ def inspect(cls, source, payload_length): raise ValueError( "Video keyframe has no packet byte position." ) + if len(keyframes) >= cls.MAX_KEYFRAME_COUNT: + raise ValueError( + "Video keyframe index exceeds the %s-entry limit." + % cls.MAX_KEYFRAME_COUNT + ) keyframes.append((ordinal, pts, position)) previous_pts = pts return cls(metadata_ranges, keyframes) @@ -210,6 +228,11 @@ def _read_header(cls, data): raise ValueError("Invalid video keyframe index version or magic.") if keyframe_count == 0: raise ValueError("Invalid video keyframe index header.") + if keyframe_count > cls.MAX_KEYFRAME_COUNT: + raise ValueError( + "Video keyframe index exceeds the %s-entry limit." + % cls.MAX_KEYFRAME_COUNT + ) entries_offset = ( cls.HEADER.size + metadata_count * cls.METADATA_RANGE.size) if entries_offset >= len(data): diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index 093e9b655b08..763231375946 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -532,6 +532,32 @@ def test_default_video_backend_falls_back_on_os_error(self): with self.assertRaises(OSError): _open_video_decoder(stream, backend="torchcodec") + def test_indexed_default_backend_falls_back_to_torchcodec(self): + stream = Mock() + stream.video_keyframe_index = VideoKeyframeIndex( + [], [(0, 0, 0)]) + decoder = object() + module = "pypaimon.multimodal.lerobot.dataset." + with patch( + module + "_PyAVVideoDecoder", + side_effect=ImportError("no av")) as pyav, patch( + module + "_open_torchcodec_decoder", + return_value=decoder) as torchcodec: + self.assertIs(decoder, _open_video_decoder(stream)) + stream.seek.assert_called_once_with(0) + pyav.assert_called_once_with( + stream, stream.video_keyframe_index) + torchcodec.assert_called_once_with(stream) + + stream.reset_mock() + with patch( + module + "_PyAVVideoDecoder", + side_effect=ImportError("no av")), patch( + module + "_open_torchcodec_decoder") as torchcodec: + with self.assertRaisesRegex(ImportError, "no av"): + _open_video_decoder(stream, backend="pyav") + torchcodec.assert_not_called() + def test_video_batches_include_delta_frames_and_preserve_backends(self): try: import torch diff --git a/paimon-python/pypaimon/tests/video_format_test.py b/paimon-python/pypaimon/tests/video_format_test.py index e4c7ad174758..ec9c070cee18 100644 --- a/paimon-python/pypaimon/tests/video_format_test.py +++ b/paimon-python/pypaimon/tests/video_format_test.py @@ -433,7 +433,7 @@ def test_rejects_compressed_invalid_index_without_expanding_it(self): try: with self.assertRaisesRegex( ValueError, - "Keyframe ordinals must be strictly increasing"): + "entry limit"): writer.add_element( GenericRow([blob], [self.field], RowKind.INSERT) ) diff --git a/paimon-python/pypaimon/tests/video_keyframe_index_test.py b/paimon-python/pypaimon/tests/video_keyframe_index_test.py index 8ead98efe199..4b10f54628e5 100644 --- a/paimon-python/pypaimon/tests/video_keyframe_index_test.py +++ b/paimon-python/pypaimon/tests/video_keyframe_index_test.py @@ -16,7 +16,10 @@ # under the License. import io +from fractions import Fraction import struct +import sys +from types import SimpleNamespace import unittest import zlib from unittest import mock @@ -104,7 +107,7 @@ def test_validate_rejects_ranges_outside_video_payload(self): VideoKeyframeIndex.deserialize(data, payload_length=10) def test_valid_large_index_uses_bounded_input_chunks(self): - count = 100_000 + count = 60_000 entries = b''.join( VideoKeyframeIndex.ENTRY.pack(value, value, value) for value in range(count) @@ -138,5 +141,54 @@ def __getattr__(self, name): max(input_sizes), VideoKeyframeIndex._CHUNK_SIZE ) + def test_rejects_too_many_keyframes_before_decompression(self): + count = VideoKeyframeIndex.MAX_KEYFRAME_COUNT + 1 + data = VideoKeyframeIndex.HEADER.pack( + VideoKeyframeIndex.VERSION, VideoKeyframeIndex.MAGIC, 0, count + ) + zlib.compress(VideoKeyframeIndex.ENTRY.pack(0, 0, 0)) + + with mock.patch.object( + keyframe_index_module.zlib, + 'decompressobj') as decompress: + with self.assertRaisesRegex(ValueError, "entry limit"): + VideoKeyframeIndex.deserialize(data) + decompress.assert_not_called() + + with self.assertRaisesRegex(ValueError, "entry limit"): + VideoKeyframeIndex( + [], + ((value, value, value) for value in range(count)), + ) + + def test_inspect_bounds_keyframe_packets(self): + stream = SimpleNamespace(time_base=Fraction(1, 30)) + packets = [ + SimpleNamespace( + is_keyframe=True, + pts=value, + pos=value, + time_base=stream.time_base, + is_discard=False, + ) + for value in range(2) + ] + container = mock.MagicMock() + container.__enter__.return_value = container + container.streams.video = [stream] + container.demux.return_value = packets + av = SimpleNamespace(open=mock.Mock(return_value=container)) + + with mock.patch.dict(sys.modules, {'av': av}), mock.patch.object( + VideoKeyframeIndex, + '_iso_bmff_metadata_ranges', + return_value=[]), mock.patch.object( + VideoKeyframeIndex, + 'MAX_KEYFRAME_COUNT', + 1): + with self.assertRaisesRegex(ValueError, "entry limit"): + VideoKeyframeIndex.inspect(mock.Mock(), 10) + self.assertEqual(1, av.open.call_count) + + if __name__ == '__main__': unittest.main() From 7fa244ccf1aeb7610211e1a7a1be953d45e782c0 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Tue, 22 Sep 2026 01:24:42 -0700 Subject: [PATCH 4/8] [python] Benchmark cold random video reads --- .../pypaimon/benchmark/video_random_read.py | 432 ++++++++++++++++++ 1 file changed, 432 insertions(+) create mode 100644 paimon-python/pypaimon/benchmark/video_random_read.py diff --git a/paimon-python/pypaimon/benchmark/video_random_read.py b/paimon-python/pypaimon/benchmark/video_random_read.py new file mode 100644 index 000000000000..ddb713e6a7dd --- /dev/null +++ b/paimon-python/pypaimon/benchmark/video_random_read.py @@ -0,0 +1,432 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Benchmark cold indexed and unindexed random video reads. + +The run command expects short-lived object URLs in environment variables so +credentials never appear in arguments or results. Each sample uses a fresh +decoder, HTTP session, and byte cache. Object-store service caches cannot be +evicted by a client and remain outside this benchmark contract. +""" + +import argparse +import hashlib +import io +import json +import os +import platform +import random +import statistics +import threading +import time +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from urllib.parse import urlsplit + +import numpy as np +import requests + +from pypaimon.multimodal.lerobot.dataset import _PyAVVideoDecoder +from pypaimon.table.row.video_keyframe_index import VideoKeyframeIndex + + +VIDEO_URL_ENV = "PYPAIMON_BENCHMARK_VIDEO_URL" +INDEX_URL_ENV = "PYPAIMON_BENCHMARK_INDEX_URL" + + +def _percentile(values, fraction): + ordered = sorted(values) + return ordered[ + min(round((len(ordered) - 1) * fraction), len(ordered) - 1) + ] + + +def _summary(runs): + latencies = [run["seconds"] * 1000 for run in runs] + gets = [run["gets"] for run in runs] + transferred = [run["bytes"] for run in runs] + return { + "runs": len(runs), + "latency_ms_p50": statistics.median(latencies), + "latency_ms_p95": _percentile(latencies, 0.95), + "gets_p50": statistics.median(gets), + "gets_p95": _percentile(gets, 0.95), + "gets_total": sum(gets), + "bytes_p50": statistics.median(transferred), + "bytes_p95": _percentile(transferred, 0.95), + "bytes_total": sum(transferred), + } + + +def _runtime(source_commit): + import av + import pyarrow + model = "" + try: + for line in Path("/proc/cpuinfo").read_text().splitlines(): + if line.startswith("model name"): + model = line.split(":", 1)[1].strip() + break + except OSError: + pass + return { + "source_commit": source_commit, + "python": platform.python_version(), + "pyav": av.__version__, + "pyarrow": pyarrow.__version__, + "numpy": np.__version__, + "requests": requests.__version__, + "machine": platform.machine(), + "cpu_count": os.cpu_count(), + "cpu_model": model, + } + + +def _file_sha256(path): + digest = hashlib.sha256() + with path.open("rb") as source: + while True: + chunk = source.read(1024 * 1024) + if not chunk: + return digest.hexdigest() + digest.update(chunk) + + +class _HttpObjects: + + def __init__(self, video_url, index_url): + self.video_url = video_url + self.index_url = index_url + self._local = threading.local() + self._sessions = [] + self._lock = threading.Lock() + self.gets = 0 + self.bytes = 0 + self.index_gets = 0 + self.index_bytes = 0 + + def _session(self): + session = getattr(self._local, "session", None) + if session is None: + session = requests.Session() + session.trust_env = False + self._local.session = session + with self._lock: + self._sessions.append(session) + return session + + def read_index(self): + response = self._session().get(self.index_url, timeout=120) + response.raise_for_status() + data = response.content + self._record(len(data), True) + return data + + def read_video(self, offset, length): + if length <= 0: + return b"" + response = self._session().get( + self.video_url, + headers={"Range": "bytes=%d-%d" % ( + offset, offset + length - 1)}, + timeout=120, + ) + response.raise_for_status() + data = response.content + if response.status_code != 206 or len(data) != length: + raise IOError( + "Object store did not return the requested video range." + ) + self._record(len(data), False) + return data + + def _record(self, length, index): + with self._lock: + self.gets += 1 + self.bytes += length + if index: + self.index_gets += 1 + self.index_bytes += length + + def close(self): + for session in self._sessions: + session.close() + + +class _HttpRangeStream(io.RawIOBase): + + def __init__(self, objects, length, parallelism): + self._objects = objects + self._length = length + self._parallelism = parallelism + self._position = 0 + + def readable(self): + return True + + def seekable(self): + return True + + def tell(self): + return self._position + + def seek(self, offset, whence=io.SEEK_SET): + if whence == io.SEEK_SET: + position = offset + elif whence == io.SEEK_CUR: + position = self._position + offset + elif whence == io.SEEK_END: + position = self._length + offset + else: + raise ValueError("Invalid whence: %s" % whence) + if position < 0: + raise ValueError("Negative seek position: %s" % position) + self._position = position + return position + + def read(self, size=-1): + length = ( + self._length - self._position + if size is None or size < 0 + else min(size, self._length - self._position) + ) + if length <= 0: + return b"" + data = self._objects.read_video(self._position, length) + self._position += len(data) + return data + + def readinto(self, value): + data = self.read(len(value)) + value[:len(data)] = data + return len(data) + + def video_read_ranges(self, ranges): + if len(ranges) < 2 or self._parallelism == 1: + return [ + self._objects.read_video(offset, length) + for offset, length in ranges + ] + with ThreadPoolExecutor( + max_workers=min(self._parallelism, len(ranges))) as pool: + return list(pool.map( + lambda value: self._objects.read_video(*value), ranges)) + + +def _frame_count(path): + import av + with av.open(str(path)) as container: + stream = container.streams.video[0] + if stream.frames: + return stream.frames + return sum(1 for unused in container.decode(stream)) + + +def _generate_index(video, repetitions): + payload_length = video.stat().st_size + values = [] + serialized = None + for unused in range(repetitions): + started = time.perf_counter() + with video.open("rb") as source: + current = VideoKeyframeIndex.inspect(source, payload_length) + values.append(time.perf_counter() - started) + current_bytes = current.serialize() + if serialized is not None and current_bytes != serialized: + raise RuntimeError("Video index generation is not deterministic.") + serialized = current_bytes + return serialized, values + + +def prepare(args): + video = args.video.resolve() + serialized, durations = _generate_index(video, args.repetitions) + args.index.write_bytes(serialized) + result = { + "contract": "pypaimon-video-index-ingestion-v1", + "runtime": _runtime(args.source_commit), + "video_bytes": video.stat().st_size, + "video_sha256": _file_sha256(video), + "frames": _frame_count(video), + "index_bytes": len(serialized), + "index_ratio": len(serialized) / video.stat().st_size, + "generation_runs": len(durations), + "generation_ms_p50": statistics.median(durations) * 1000, + "generation_ms_p95": _percentile(durations, 0.95) * 1000, + "generation_seconds": durations, + } + args.output.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n") + print(json.dumps(result, indent=2, sort_keys=True)) + + +def _decode(mode, targets, video_length, expected_index, parallelism): + objects = _HttpObjects( + os.environ[VIDEO_URL_ENV], os.environ[INDEX_URL_ENV]) + stream = _HttpRangeStream(objects, video_length, parallelism) + started = time.perf_counter() + decoder = None + try: + mapping = None + if mode == "indexed": + index_bytes = objects.read_index() + if index_bytes != expected_index: + raise RuntimeError("Remote video index differs from local index.") + mapping = VideoKeyframeIndex.deserialize( + index_bytes, video_length) + stream.video_length = video_length + decoder = _PyAVVideoDecoder(stream, mapping) + frames = np.stack( + decoder._read_indexed(targets) + if mapping is not None else [decoder[index] for index in targets] + ) + digest = hashlib.sha256(frames.tobytes()).hexdigest() + elapsed = time.perf_counter() - started + return { + "targets": targets, + "seconds": elapsed, + "gets": objects.gets, + "bytes": objects.bytes, + "index_gets": objects.index_gets, + "index_bytes": objects.index_bytes, + "video_gets": objects.gets - objects.index_gets, + "video_bytes": objects.bytes - objects.index_bytes, + "sha256": digest, + } + finally: + if decoder is not None: + decoder.close() + stream.close() + objects.close() + + +def _paired_runs(cases, video_length, expected_index, parallelism): + runs = {"indexed": [], "unindexed": []} + for position, targets in enumerate(cases): + order = ( + ("unindexed", "indexed") + if position % 2 == 0 else ("indexed", "unindexed") + ) + pair = {} + for mode in order: + pair[mode] = _decode( + mode, targets, video_length, expected_index, parallelism) + runs[mode].append(pair[mode]) + if pair["indexed"]["sha256"] != pair["unindexed"]["sha256"]: + raise RuntimeError("Indexed and unindexed frames differ.") + return runs + + +def run(args): + video_url = os.environ.get(VIDEO_URL_ENV) + index_url = os.environ.get(INDEX_URL_ENV) + if not video_url or not index_url: + raise ValueError( + "%s and %s must contain object URLs." + % (VIDEO_URL_ENV, INDEX_URL_ENV) + ) + expected_index = args.index.read_bytes() + video_length = args.video.stat().st_size + frames = _frame_count(args.video) + if frames < 3: + raise ValueError("Video must contain at least three frames.") + if args.batch_size > frames - 2: + raise ValueError("--batch-size exceeds the available target frames.") + randomizer = random.Random(args.seed) + first_cases = [ + [randomizer.randrange(1, frames - 1)] + for unused in range(args.first_frame_runs) + ] + batch_cases = [ + randomizer.sample(range(1, frames - 1), args.batch_size) + for unused in range(args.batch_runs) + ] + + # Keep decoder and RGB conversion costs while avoiding a Torch dependency. + _PyAVVideoDecoder._tensor = staticmethod( + lambda frame: np.array( + frame.to_ndarray(format="rgb24"), copy=True + ).transpose(2, 0, 1) + ) + first = _paired_runs( + first_cases, video_length, expected_index, args.parallelism) + batches = _paired_runs( + batch_cases, video_length, expected_index, args.parallelism) + result = { + "contract": "pypaimon-cold-object-video-random-read-v1", + "runtime": _runtime(args.source_commit), + "storage_host": urlsplit(video_url).hostname, + "service_cache": "not controlled", + "fresh_state": "decoder, HTTP sessions, and client byte cache", + "video_bytes": video_length, + "video_sha256": _file_sha256(args.video), + "frames": frames, + "index_bytes": len(expected_index), + "seed": args.seed, + "parallelism": args.parallelism, + "first_frame": { + "indexed": _summary(first["indexed"]), + "unindexed": _summary(first["unindexed"]), + "raw": first, + }, + "random_batch": { + "batch_size": args.batch_size, + "indexed": _summary(batches["indexed"]), + "unindexed": _summary(batches["unindexed"]), + "raw": batches, + }, + "decoded_frames_equal": True, + } + args.output.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n") + print(json.dumps(result, indent=2, sort_keys=True)) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + commands = parser.add_subparsers(dest="command") + + prepare_parser = commands.add_parser("prepare") + prepare_parser.add_argument("--video", type=Path, required=True) + prepare_parser.add_argument("--index", type=Path, required=True) + prepare_parser.add_argument("--output", type=Path, required=True) + prepare_parser.add_argument("--source-commit", required=True) + prepare_parser.add_argument("--repetitions", type=int, default=3) + prepare_parser.set_defaults(function=prepare) + + run_parser = commands.add_parser("run") + run_parser.add_argument("--video", type=Path, required=True) + run_parser.add_argument("--index", type=Path, required=True) + run_parser.add_argument("--output", type=Path, required=True) + run_parser.add_argument("--source-commit", required=True) + run_parser.add_argument("--first-frame-runs", type=int, default=20) + run_parser.add_argument("--batch-runs", type=int, default=10) + run_parser.add_argument("--batch-size", type=int, default=16) + run_parser.add_argument("--parallelism", type=int, default=8) + run_parser.add_argument("--seed", type=int, default=9831) + run_parser.set_defaults(function=run) + + args = parser.parse_args() + if not hasattr(args, "function"): + parser.error("a command is required") + for name in ( + "repetitions", "first_frame_runs", "batch_runs", "batch_size", + "parallelism"): + if hasattr(args, name) and getattr(args, name) <= 0: + parser.error("--%s must be positive" % name.replace("_", "-")) + args.function(args) + + +if __name__ == "__main__": + main() From f3bb9da803a3adf131c696bbc5520cf854ccce67 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Tue, 22 Sep 2026 02:02:16 -0700 Subject: [PATCH 5/8] [python] Bound video metadata ranges --- docs/docs/concepts/spec/fileformat.md | 2 +- .../format/blob/VideoKeyframeIndex.java | 4 + .../format/blob/VideoFileFormatTest.java | 17 +++ .../table/row/video_keyframe_index.py | 50 +++++-- .../pypaimon/tests/multimodal_lerobot_test.py | 128 ++++++++++-------- .../tests/video_keyframe_index_test.py | 41 +++++- 6 files changed, 175 insertions(+), 67 deletions(-) diff --git a/docs/docs/concepts/spec/fileformat.md b/docs/docs/concepts/spec/fileformat.md index c3920db7d6e4..37f5b6ec73eb 100644 --- a/docs/docs/concepts/spec/fileformat.md +++ b/docs/docs/concepts/spec/fileformat.md @@ -449,7 +449,7 @@ metadata-range count (uint32), and keyframe count (uint32). It then stores metad length)` pairs (two int64 values) and zlib-compressed `(frame ordinal, PTS, packet position)` keyframe entries (three int64 values). Offsets are relative to the encoded video; writers reject out-of-range values. All numeric fields are little endian. One block is limited to 65,536 -keyframes and 16 MiB; all blocks in one file are limited to 64 MiB. +metadata ranges, 65,536 keyframes, and 16 MiB; all blocks in one file are limited to 64 MiB. The index covers the first video stream. Its time base remains in the video. A reader fetches the metadata and target GOP, seeks by PTS, and decodes forward in presentation order. It may include the diff --git a/paimon-format/src/main/java/org/apache/paimon/format/blob/VideoKeyframeIndex.java b/paimon-format/src/main/java/org/apache/paimon/format/blob/VideoKeyframeIndex.java index 8892f0f4a85f..34cfc0e59010 100644 --- a/paimon-format/src/main/java/org/apache/paimon/format/blob/VideoKeyframeIndex.java +++ b/paimon-format/src/main/java/org/apache/paimon/format/blob/VideoKeyframeIndex.java @@ -31,6 +31,7 @@ final class VideoKeyframeIndex { private static final int HEADER_LENGTH = Byte.BYTES + Long.BYTES + Integer.BYTES * 2; private static final int METADATA_RANGE_LENGTH = Long.BYTES * 2; private static final int ENTRY_LENGTH = Long.BYTES * 3; + static final long MAX_METADATA_RANGE_COUNT = 64 * 1024; static final long MAX_KEYFRAME_COUNT = 64 * 1024; private VideoKeyframeIndex() {} @@ -54,6 +55,9 @@ static void validate(byte[] data, long payloadLength) { if (keyframeCount == 0) { throw invalid("empty"); } + if (metadataRangeCount > MAX_METADATA_RANGE_COUNT) { + throw invalid("metadata range count exceeds limit"); + } if (keyframeCount > MAX_KEYFRAME_COUNT) { throw invalid("keyframe count exceeds limit"); } diff --git a/paimon-format/src/test/java/org/apache/paimon/format/blob/VideoFileFormatTest.java b/paimon-format/src/test/java/org/apache/paimon/format/blob/VideoFileFormatTest.java index adf9f9ac500d..ddc971d440ef 100644 --- a/paimon-format/src/test/java/org/apache/paimon/format/blob/VideoFileFormatTest.java +++ b/paimon-format/src/test/java/org/apache/paimon/format/blob/VideoFileFormatTest.java @@ -233,6 +233,23 @@ public void testRejectTooManyKeyframes() { .hasMessageContaining("keyframe count exceeds limit"); } + @Test + public void testRejectTooManyMetadataRanges() { + byte[] mapping = + ByteBuffer.allocate(18) + .order(ByteOrder.LITTLE_ENDIAN) + .put((byte) 1) + .putLong(0x564944454F4B4649L) + .putInt((int) VideoKeyframeIndex.MAX_METADATA_RANGE_COUNT + 1) + .putInt(1) + .put((byte) 0) + .array(); + + assertThatThrownBy(() -> VideoKeyframeIndex.validate(mapping, 1)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("metadata range count exceeds limit"); + } + @Test public void testRejectOversizedKeyframeIndexesBeforeFetch() throws IOException { String missing = new Path(tempPath.resolve("missing.mp4").toUri()).toString(); diff --git a/paimon-python/pypaimon/table/row/video_keyframe_index.py b/paimon-python/pypaimon/table/row/video_keyframe_index.py index 6f7ecf925d9d..9d8f34e158ca 100644 --- a/paimon-python/pypaimon/table/row/video_keyframe_index.py +++ b/paimon-python/pypaimon/table/row/video_keyframe_index.py @@ -30,14 +30,30 @@ class VideoKeyframeIndex: HEADER = struct.Struct('= self.MAX_METADATA_RANGE_COUNT: + raise ValueError( + "Video keyframe index exceeds the %s-metadata-range " + "limit." % self.MAX_METADATA_RANGE_COUNT + ) + metadata_range_count += 1 + offset, length = operator.index(offset), operator.index(length) + if normalized_metadata_ranges: + previous_offset, previous_length = ( + normalized_metadata_ranges[-1]) + if previous_offset + previous_length == offset: + normalized_metadata_ranges[-1] = ( + previous_offset, previous_length + length) + continue + normalized_metadata_ranges.append((offset, length)) + self.metadata_ranges = tuple(normalized_metadata_ranges) normalized_keyframes = [] for ordinal, pts, position in keyframes: if len(normalized_keyframes) >= self.MAX_KEYFRAME_COUNT: @@ -164,8 +180,8 @@ def _stream_pts(pts, time_base, stream_time_base): ) return int(value) - @staticmethod - def _iso_bmff_metadata_ranges(source, payload_length): + @classmethod + def _iso_bmff_metadata_ranges(cls, source, payload_length): ranges = [] has_moov = False offset = 0 @@ -186,11 +202,20 @@ def _iso_bmff_metadata_ranges(source, payload_length): size = payload_length - offset if size < header_size or offset + size > payload_length: raise ValueError("Invalid ISO BMFF box size.") - ranges.append(( - offset, + length = ( header_size if box_type in (b"mdat", b"free", b"skip") - else size, - )) + else size + ) + if ranges and ranges[-1][0] + ranges[-1][1] == offset: + ranges[-1] = (ranges[-1][0], ranges[-1][1] + length) + else: + if len(ranges) >= cls.MAX_METADATA_RANGE_COUNT: + raise ValueError( + "Video keyframe index exceeds the " + "%s-metadata-range limit." + % cls.MAX_METADATA_RANGE_COUNT + ) + ranges.append((offset, length)) has_moov |= box_type == b"moov" offset += size if offset != payload_length or not has_moov: @@ -228,6 +253,11 @@ def _read_header(cls, data): raise ValueError("Invalid video keyframe index version or magic.") if keyframe_count == 0: raise ValueError("Invalid video keyframe index header.") + if metadata_count > cls.MAX_METADATA_RANGE_COUNT: + raise ValueError( + "Video keyframe index exceeds the %s-metadata-range limit." + % cls.MAX_METADATA_RANGE_COUNT + ) if keyframe_count > cls.MAX_KEYFRAME_COUNT: raise ValueError( "Video keyframe index exceeds the %s-entry limit." diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index 763231375946..929032615d07 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -2705,12 +2705,12 @@ def _assert_imported_video_payload_can_be_decoded(self, training_reads): "action": {"dtype": "float32", "shape": [1]}, "camera": { "dtype": "video", - "shape": [16, 16, 3], + "shape": [128, 128, 3], "video_info": {"video.fps": 10.0}, }, "camera_b": { "dtype": "video", - "shape": [16, 16, 3], + "shape": [128, 128, 3], "video_info": {"video.fps": 10.0}, }, }, @@ -2726,12 +2726,12 @@ def _assert_imported_video_payload_can_be_decoded(self, training_reads): "tasks": ["pick"], "videos/camera/chunk_index": 0, "videos/camera/file_index": 0, - "videos/camera/from_timestamp": 0.5, - "videos/camera/to_timestamp": 0.7, + "videos/camera/from_timestamp": 5.5, + "videos/camera/to_timestamp": 5.7, "videos/camera_b/chunk_index": 0, "videos/camera_b/file_index": 0, - "videos/camera_b/from_timestamp": 0.5, - "videos/camera_b/to_timestamp": 0.7, + "videos/camera_b/from_timestamp": 5.5, + "videos/camera_b/to_timestamp": 5.7, }, { "episode_index": 1, @@ -2743,18 +2743,17 @@ def _assert_imported_video_payload_can_be_decoded(self, training_reads): "tasks": ["pick"], "videos/camera/chunk_index": 0, "videos/camera/file_index": 0, - "videos/camera/from_timestamp": 0.1, - "videos/camera/to_timestamp": 0.4, + "videos/camera/from_timestamp": 6.1, + "videos/camera/to_timestamp": 6.4, "videos/camera_b/chunk_index": 0, "videos/camera_b/file_index": 0, - "videos/camera_b/from_timestamp": 0.1, - "videos/camera_b/to_timestamp": 0.4, + "videos/camera_b/from_timestamp": 6.1, + "videos/camera_b/to_timestamp": 6.4, }, ] - physical_frame_values = [24] * 60 - for index, value in { - 1: 56, 2: 88, 3: 120, 5: 168, 6: 216}.items(): - physical_frame_values[index] = value + target_frame_values = { + 55: 168, 56: 216, 61: 56, 62: 88, 63: 120, + } expected_frame_values = [168, 216, 56, 88, 120] info_dir = temp_dir / "meta" @@ -2774,14 +2773,21 @@ def _assert_imported_video_payload_can_be_decoded(self, training_reads): 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.width = 128 + stream.height = 128 stream.pix_fmt = "yuv420p" stream.time_base = Fraction(1, 10) - stream.gop_size = 10 + stream.gop_size = 12 stream.codec_context.max_b_frames = 2 - for pts, value in enumerate(physical_frame_values): - image = np.full((16, 16, 3), value, dtype=np.uint8) + for pts in range(120): + value = target_frame_values.get(pts) + if value is None: + noise = np.random.RandomState(pts).randint( + 0, 256, size=(128, 128, 3)) + image = noise.astype(np.uint8) + else: + image = np.full( + (128, 128, 3), value, dtype=np.uint8) frame = av.VideoFrame.from_ndarray(image, format="rgb24") frame.pts = pts frame.time_base = Fraction(1, 10) @@ -2853,7 +2859,7 @@ def read_batch(self, begin, end): for descriptor in descriptors + camera_b_descriptors )) self.assertEqual( - [5, 6, 1, 2, 3], + [55, 56, 61, 62, 63], [descriptor.frame_index for descriptor in descriptors], ) @@ -2892,7 +2898,7 @@ def close(self): ) np.testing.assert_allclose( [row["decoded"][1] for row in decoded_rows], - [0.5, 0.6, 0.1, 0.2, 0.3], + [5.5, 5.6, 6.1, 6.2, 6.3], atol=1e-6, ) if not training_reads: @@ -2903,56 +2909,70 @@ def close(self): delta_timestamps={"camera": [0.0, 0.1]}, ) try: - video_ranges = [] + transferred = [] file_io = table.raw_table.file_io read_ranges = file_io.read_ranges_coalesced + read_file_range = file_io.read_file_range + selected = [descriptors[2], camera_b_descriptors[2]] + tracked = [] + for descriptor in selected: + tracked.append(( + descriptor.uri, + descriptor.offset, + descriptor.offset + descriptor.length, + )) + index = descriptor.keyframe_index_descriptor + tracked.append(( + index.uri, index.offset, index.offset + index.length)) + + def track(path, offset, length): + end = offset + length + if any( + path == tracked_path + and offset < tracked_end and end > tracked_begin + for tracked_path, tracked_begin, tracked_end + in tracked): + transferred.append(length) def track_ranges(ranges, parallelism): - video_ranges.extend(ranges) + for path, offset, length in ranges: + track(path, offset, length) return read_ranges(ranges, parallelism) + def track_file_range(path, offset, length): + track(path, offset, length) + return read_file_range(path, offset, length) + with patch.object( file_io, "read_ranges_coalesced", - side_effect=track_ranges): - last, first = dataset.__getitems__([4, 0]) - self.assertEqual( - [2, 3, 16, 16], list(last["camera"].shape)) + side_effect=track_ranges), patch.object( + file_io, + "read_file_range", + side_effect=track_file_range): + middle, = dataset.__getitems__([2]) self.assertEqual( - [2, 3, 16, 16], list(first["camera"].shape)) + [2, 3, 128, 128], list(middle["camera"].shape)) self.assertEqual( - [3, 16, 16], list(first["camera_b"].shape)) - self.assertEqual("torch.float32", str(last["camera"].dtype)) + [3, 128, 128], list(middle["camera_b"].shape)) + self.assertEqual("torch.float32", str(middle["camera"].dtype)) np.testing.assert_allclose( [ - float(last["camera"][0].mean()) * 255, - float(first["camera"][0].mean()) * 255, - float(first["camera"][1].mean()) * 255, + float(middle["camera"][0].mean()) * 255, + float(middle["camera"][1].mean()) * 255, ], - [120, 168, 216], + [56, 88], atol=5, ) self.assertEqual( - [False, True], last["camera_is_pad"].tolist()) + [False, False], middle["camera_is_pad"].tolist()) self.assertEqual( 1, len(dataset._video_collators[0]._decoders)) - for descriptor in descriptors + camera_b_descriptors: - loaded = bytearray(descriptor.length) - for path, offset, length in video_ranges: - if path != descriptor.uri: - continue - begin = max(offset, descriptor.offset) - end = min( - offset + length, - descriptor.offset + descriptor.length, - ) - if begin < end: - loaded[ - begin - descriptor.offset: - end - descriptor.offset - ] = b"\1" * (end - begin) - self.assertGreater(sum(loaded), 0) - self.assertLess(sum(loaded), descriptor.length) + self.assertGreater(sum(transferred), 0) + self.assertLess( + sum(transferred), + sum(descriptor.length for descriptor in selected), + ) from torch.utils.data import DataLoader worker_indices = [] @@ -2977,9 +2997,9 @@ def track_ranges(ranges, parallelism): np.testing.assert_allclose( [0.0, 1.0], item["action"].tolist()) self.assertEqual( - [3, 16, 16], list(item["camera"].shape)) + [3, 128, 128], list(item["camera"].shape)) self.assertEqual( - [3, 16, 16], list(item["camera_b"].shape)) + [3, 128, 128], list(item["camera_b"].shape)) finally: action_dataset.close() finally: diff --git a/paimon-python/pypaimon/tests/video_keyframe_index_test.py b/paimon-python/pypaimon/tests/video_keyframe_index_test.py index 4b10f54628e5..94884326329c 100644 --- a/paimon-python/pypaimon/tests/video_keyframe_index_test.py +++ b/paimon-python/pypaimon/tests/video_keyframe_index_test.py @@ -50,8 +50,7 @@ def box(box_type, body): self.assertEqual( [ - (offsets[0], len(boxes[0])), - (offsets[1], 8), + (offsets[0], len(boxes[0]) + 8), (offsets[2], 8), (offsets[3], 8), (offsets[4], len(boxes[4])), @@ -68,6 +67,11 @@ def test_round_trip_and_validation(self): restored = VideoKeyframeIndex.deserialize(index.serialize()) self.assertEqual(index.metadata_ranges, restored.metadata_ranges) self.assertEqual(index.keyframes, restored.keyframes) + self.assertEqual( + ((0, 2),), + VideoKeyframeIndex( + [(0, 1), (1, 1)], [(0, 0, 0)]).metadata_ranges, + ) self.assertEqual(17, index.HEADER.size) self.assertFalse(hasattr(restored, 'time_base')) self.assertFalse(hasattr(restored, 'frame_count')) @@ -160,6 +164,39 @@ def test_rejects_too_many_keyframes_before_decompression(self): ((value, value, value) for value in range(count)), ) + def test_rejects_too_many_metadata_ranges_before_materialization(self): + count = VideoKeyframeIndex.MAX_METADATA_RANGE_COUNT + 1 + data = VideoKeyframeIndex.HEADER.pack( + VideoKeyframeIndex.VERSION, VideoKeyframeIndex.MAGIC, count, 1 + ) + zlib.compress(VideoKeyframeIndex.ENTRY.pack(0, 0, 0)) + + with mock.patch.object( + VideoKeyframeIndex, '_iter_metadata_ranges') as ranges: + with self.assertRaisesRegex(ValueError, "metadata-range limit"): + VideoKeyframeIndex.deserialize(data) + ranges.assert_not_called() + + with self.assertRaisesRegex(ValueError, "metadata-range limit"): + VideoKeyframeIndex( + ((value * 2, 1) for value in range(count)), + [(0, 0, 0)], + ) + + def test_iso_bmff_metadata_range_count_is_bounded_after_merging(self): + def box(box_type, body): + return struct.pack(">I4s", len(body) + 8, box_type) + body + + payload = b''.join([ + box(b"free", b"x"), + box(b"skip", b"y"), + box(b"moov", b"metadata"), + ]) + with mock.patch.object( + VideoKeyframeIndex, 'MAX_METADATA_RANGE_COUNT', 2): + with self.assertRaisesRegex(ValueError, "metadata-range limit"): + VideoKeyframeIndex._iso_bmff_metadata_ranges( + io.BytesIO(payload), len(payload)) + def test_inspect_bounds_keyframe_packets(self): stream = SimpleNamespace(time_base=Fraction(1, 30)) packets = [ From d3808f0e6f2794c8a762a4eb1aae2c5064a3ea7c Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Tue, 22 Sep 2026 03:25:42 -0700 Subject: [PATCH 6/8] [python] Fix video stream fallbacks --- .../filesystem/jindo_file_system_handler.py | 1 + .../pypaimon/multimodal/lerobot/dataset.py | 42 ++++++++++++++++- .../pypaimon/tests/jindo_file_system_test.py | 22 ++++++++- .../pypaimon/tests/multimodal_lerobot_test.py | 35 +++++++++++++++ .../pypaimon/tests/video_format_test.py | 45 +++++++++++++++++++ 5 files changed, 143 insertions(+), 2 deletions(-) diff --git a/paimon-python/pypaimon/filesystem/jindo_file_system_handler.py b/paimon-python/pypaimon/filesystem/jindo_file_system_handler.py index 73ad09077fa1..7521ab115258 100644 --- a/paimon-python/pypaimon/filesystem/jindo_file_system_handler.py +++ b/paimon-python/pypaimon/filesystem/jindo_file_system_handler.py @@ -139,6 +139,7 @@ def seek(self, position: int, whence: int = 0): if self.closed: raise ValueError("I/O operation on closed file") self._stream.seek(position, whence) + return self._stream.tell() def tell(self) -> int: if self.closed: diff --git a/paimon-python/pypaimon/multimodal/lerobot/dataset.py b/paimon-python/pypaimon/multimodal/lerobot/dataset.py index fbe3c26a9e54..6d9bc1d62414 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/dataset.py +++ b/paimon-python/pypaimon/multimodal/lerobot/dataset.py @@ -1515,9 +1515,11 @@ def _open_video_decoder(stream, backend=None): keyframe_index = None if backend is None and keyframe_index is not None: try: - return _PyAVVideoDecoder(stream, keyframe_index) + decoder = _PyAVVideoDecoder(stream, keyframe_index) except (ImportError, OSError, RuntimeError): stream.seek(0) + else: + return _FallbackVideoDecoder(stream, decoder) if backend in (None, "torchcodec"): try: return _open_torchcodec_decoder(stream) @@ -1545,6 +1547,44 @@ def _open_torchcodec_decoder(stream): return VideoDecoder(stream.read(), seek_mode="exact") +class _FallbackVideoDecoder: + + def __init__(self, stream, decoder): + self._stream = stream + self._decoder = decoder + self._pending = True + + def __getitem__(self, index): + return self._call("__getitem__", index) + + def get_frames_at(self, *, indices): + return self._call("get_frames_at", indices=indices) + + def _call(self, method, *args, **kwargs): + try: + result = getattr(self._decoder, method)(*args, **kwargs) + except (ImportError, OSError, RuntimeError) as error: + if not self._pending: + raise + self._pending = False + close = getattr(self._decoder, "close", None) + if close is not None: + close() + self._stream.seek(0) + try: + self._decoder = _open_torchcodec_decoder(self._stream) + except (ImportError, OSError, RuntimeError): + raise error + return getattr(self._decoder, method)(*args, **kwargs) + self._pending = False + return result + + def close(self): + close = getattr(self._decoder, "close", None) + if close is not None: + close() + + class _RangeBackedVideo(io.RawIOBase): def __init__(self, length, read_ranges): diff --git a/paimon-python/pypaimon/tests/jindo_file_system_test.py b/paimon-python/pypaimon/tests/jindo_file_system_test.py index b828fd8b8307..d43717319689 100644 --- a/paimon-python/pypaimon/tests/jindo_file_system_test.py +++ b/paimon-python/pypaimon/tests/jindo_file_system_test.py @@ -15,6 +15,7 @@ # specific language governing permissions and limitations # under the License. +import io import os import types import unittest @@ -27,7 +28,11 @@ from pypaimon.common.options import Options from pypaimon.common.options.config import OssOptions from pypaimon.filesystem import jindo_file_system_handler as jindo_module -from pypaimon.filesystem.jindo_file_system_handler import JindoFileSystemHandler, JINDO_AVAILABLE +from pypaimon.filesystem.jindo_file_system_handler import ( + JindoFileSystemHandler, + JindoInputFile, + JINDO_AVAILABLE, +) class _RecordingConfig: @@ -40,6 +45,21 @@ def set(self, key, value): class JindoConfigTest(unittest.TestCase): + def test_input_file_seek_returns_position(self): + class JindoStream: + + def __init__(self): + self._stream = io.BytesIO(b"video") + + def seek(self, offset, whence=io.SEEK_SET): + self._stream.seek(offset, whence) + + def tell(self): + return self._stream.tell() + + stream = JindoInputFile(JindoStream()) + self.assertEqual(5, stream.seek(0, io.SEEK_END)) + def test_forwards_native_options_to_connect(self): created_config = _RecordingConfig() config_factory = mock.Mock(return_value=created_config) diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index 929032615d07..4606ad39aae6 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -558,6 +558,41 @@ def test_indexed_default_backend_falls_back_to_torchcodec(self): _open_video_decoder(stream, backend="pyav") torchcodec.assert_not_called() + def test_indexed_default_backend_falls_back_after_lazy_failure(self): + stream = Mock() + stream.video_keyframe_index = VideoKeyframeIndex( + [], [(0, 0, 0)]) + pyav_decoder = Mock() + pyav_decoder.get_frames_at.side_effect = OSError("cannot open") + torchcodec_decoder = Mock() + expected = object() + torchcodec_decoder.get_frames_at.return_value = expected + module = "pypaimon.multimodal.lerobot.dataset." + + with patch( + module + "_PyAVVideoDecoder", + return_value=pyav_decoder), patch( + module + "_open_torchcodec_decoder", + return_value=torchcodec_decoder) as torchcodec: + decoder = _open_video_decoder(stream) + self.assertIs(expected, decoder.get_frames_at(indices=[3])) + pyav_decoder.close.assert_called_once_with() + stream.seek.assert_called_once_with(0) + torchcodec.assert_called_once_with(stream) + decoder.close() + torchcodec_decoder.close.assert_called_once_with() + + stream.reset_mock() + with patch( + module + "_PyAVVideoDecoder", + return_value=pyav_decoder), patch( + module + "_open_torchcodec_decoder") as torchcodec: + decoder = _open_video_decoder(stream, backend="pyav") + with self.assertRaisesRegex(OSError, "cannot open"): + decoder.get_frames_at(indices=[3]) + stream.seek.assert_not_called() + torchcodec.assert_not_called() + def test_video_batches_include_delta_frames_and_preserve_backends(self): try: import torch diff --git a/paimon-python/pypaimon/tests/video_format_test.py b/paimon-python/pypaimon/tests/video_format_test.py index ec9c070cee18..fac7960b0d86 100644 --- a/paimon-python/pypaimon/tests/video_format_test.py +++ b/paimon-python/pypaimon/tests/video_format_test.py @@ -28,6 +28,7 @@ from pypaimon.common.delta_varint_compressor import DeltaVarintCompressor from pypaimon.common.options import Options +from pypaimon.filesystem.jindo_file_system_handler import JindoInputFile from pypaimon.filesystem.local_file_io import LocalFileIO from pypaimon.read.reader.format_blob_reader import FormatBlobReader from pypaimon.read.reader.video_format_reader import VideoFileMeta @@ -319,6 +320,50 @@ def test_real_mp4_index_supports_sparse_random_reads(self): sparse_read_sizes.append(available_bytes) self.assertLess(min(sparse_read_sizes), len(stored_payload)) + @unittest.skipUnless( + av is not None and np is not None, + "PyAV and NumPy are required for video index generation", + ) + def test_unknown_length_jindo_video_generates_index(self): + payload = self._real_mp4() + + class JindoStream: + + def __init__(self): + self._stream = io.BytesIO(payload) + + @property + def closed(self): + return self._stream.closed + + def read(self, size=-1): + return self._stream.read(size) + + def seek(self, offset, whence=io.SEEK_SET): + self._stream.seek(offset, whence) + + def tell(self): + return self._stream.tell() + + def close(self): + self._stream.close() + + reader = mock.Mock() + reader.new_input_stream.side_effect = lambda unused: JindoInputFile( + JindoStream()) + descriptor = VideoFrameDescriptor( + "oss://bucket/video.mp4", 0, -1, 0, -1, 0) + value = Blob.from_descriptor(reader, descriptor) + target = (self.root / "jindo.video").as_uri() + + writer = VideoFormatWriter(self.file_io.new_output_stream(target)) + writer.add_element(GenericRow([value], [self.field], RowKind.INSERT)) + writer.close() + + stored = VideoFrameDescriptor.deserialize(self._read(target)[0]) + self.assertEqual(len(payload), stored.length) + self.assertGreater(stored.keyframe_index_descriptor.length, 0) + def test_target_size_counts_buffered_keyframe_index(self): mapping = VideoKeyframeIndex([], [(0, 0, 0)]).serialize() video = b"video" From 85b6e621553173b361b7e966ea9ecb0ae031d538 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Tue, 22 Sep 2026 03:41:02 -0700 Subject: [PATCH 7/8] [python] Handle PyAV fallback errors --- .../pypaimon/multimodal/lerobot/dataset.py | 13 +++++-- .../pypaimon/tests/multimodal_lerobot_test.py | 36 +++++++++++++++++-- 2 files changed, 44 insertions(+), 5 deletions(-) diff --git a/paimon-python/pypaimon/multimodal/lerobot/dataset.py b/paimon-python/pypaimon/multimodal/lerobot/dataset.py index 6d9bc1d62414..cc5558c99a31 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/dataset.py +++ b/paimon-python/pypaimon/multimodal/lerobot/dataset.py @@ -1509,6 +1509,13 @@ def _video_tensor(frame, feature, return_uint8=False): return frame +def _video_decoder_fallback_errors(): + errors = (ImportError, OSError, RuntimeError, ValueError) + av = sys.modules.get("av") + ffmpeg_error = getattr(getattr(av, "error", None), "FFmpegError", None) + return errors if ffmpeg_error is None else errors + (ffmpeg_error,) + + def _open_video_decoder(stream, backend=None): keyframe_index = getattr(stream, "video_keyframe_index", None) if not isinstance(keyframe_index, VideoKeyframeIndex): @@ -1516,7 +1523,7 @@ def _open_video_decoder(stream, backend=None): if backend is None and keyframe_index is not None: try: decoder = _PyAVVideoDecoder(stream, keyframe_index) - except (ImportError, OSError, RuntimeError): + except _video_decoder_fallback_errors(): stream.seek(0) else: return _FallbackVideoDecoder(stream, decoder) @@ -1563,7 +1570,7 @@ def get_frames_at(self, *, indices): def _call(self, method, *args, **kwargs): try: result = getattr(self._decoder, method)(*args, **kwargs) - except (ImportError, OSError, RuntimeError) as error: + except _video_decoder_fallback_errors() as error: if not self._pending: raise self._pending = False @@ -1573,7 +1580,7 @@ def _call(self, method, *args, **kwargs): self._stream.seek(0) try: self._decoder = _open_torchcodec_decoder(self._stream) - except (ImportError, OSError, RuntimeError): + except _video_decoder_fallback_errors(): raise error return getattr(self._decoder, method)(*args, **kwargs) self._pending = False diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index 4606ad39aae6..8fd7a90628f0 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -559,11 +559,20 @@ def test_indexed_default_backend_falls_back_to_torchcodec(self): torchcodec.assert_not_called() def test_indexed_default_backend_falls_back_after_lazy_failure(self): + if av is None: + self.skipTest("PyAV is required for its real exception types") + try: + av.open(io.BytesIO(b"invalid video")) + except av.error.InvalidDataError as error: + invalid_data_error = error + else: + self.fail("Invalid video did not raise InvalidDataError") + stream = Mock() stream.video_keyframe_index = VideoKeyframeIndex( [], [(0, 0, 0)]) pyav_decoder = Mock() - pyav_decoder.get_frames_at.side_effect = OSError("cannot open") + pyav_decoder.get_frames_at.side_effect = invalid_data_error torchcodec_decoder = Mock() expected = object() torchcodec_decoder.get_frames_at.return_value = expected @@ -588,11 +597,34 @@ def test_indexed_default_backend_falls_back_after_lazy_failure(self): return_value=pyav_decoder), patch( module + "_open_torchcodec_decoder") as torchcodec: decoder = _open_video_decoder(stream, backend="pyav") - with self.assertRaisesRegex(OSError, "cannot open"): + with self.assertRaises(av.error.InvalidDataError): decoder.get_frames_at(indices=[3]) stream.seek.assert_not_called() torchcodec.assert_not_called() + def test_indexed_default_backend_falls_back_on_ffmpeg_error(self): + if av is None: + self.skipTest("PyAV is required for its real exception types") + stream = Mock() + stream.video_keyframe_index = VideoKeyframeIndex( + [], [(0, 0, 0)]) + pyav_decoder = Mock() + pyav_decoder.get_frames_at.side_effect = av.error.FFmpegError( + 1, "cannot open", "video") + torchcodec_decoder = Mock() + expected = object() + torchcodec_decoder.get_frames_at.return_value = expected + module = "pypaimon.multimodal.lerobot.dataset." + + with patch( + module + "_PyAVVideoDecoder", + return_value=pyav_decoder), patch( + module + "_open_torchcodec_decoder", + return_value=torchcodec_decoder): + decoder = _open_video_decoder(stream) + self.assertIs(expected, decoder.get_frames_at(indices=[3])) + stream.seek.assert_called_once_with(0) + def test_video_batches_include_delta_frames_and_preserve_backends(self): try: import torch From cf35baeebfb591b3d9be805d865bbbabbbde18f4 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Tue, 22 Sep 2026 20:22:09 -0700 Subject: [PATCH 8/8] [python] Stabilize remote video import test --- paimon-python/pypaimon/tests/multimodal_lerobot_test.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index 8fd7a90628f0..6fdfe17415a6 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -2665,8 +2665,13 @@ def sample_timestamps(unused_dataset, uri): path for path in source_file_io.opened_paths if path.endswith(".mp4") ] - # PyAV-enabled writers inspect each source once before copying it. - self.assertEqual(6 if av is not None else 3, len(opened_videos)) + self.assertEqual( + { + remote + "/videos/" + relative + for relative in payloads + }, + set(opened_videos), + ) self.assertEqual(1, source_file_io.close_count) _, remote_bodies = connection.get_table( "remote_frames").scan().select([