JingsongLi commented on code in PR #9831:
URL: https://github.com/apache/paimon/pull/9831#discussion_r4068931132
##########
paimon-python/pypaimon/multimodal/lerobot/dataset.py:
##########
@@ -1579,6 +1608,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[anchor][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=False, any_frame=False, stream=stream)
Review Comment:
[P1] Seek backward to the persisted keyframe anchor
`backward=False` can skip the indexed keyframe, but the loop below requires
an exact match for `anchor_pts`. With PyAV 18.1.0, a valid 150-frame fragmented
H.264 MP4 (`movflags=frag_keyframe+empty_moov+default_base_moof`) is
automatically indexed successfully, yet all 150 indexed reads fail: seeking to
the first anchor at PTS 1024 starts at the next GOP, PTS 8704. Ordinary HEVC
MP4 also fails for several frame ranges. The same forward seek skips the anchor
on the complete file, so fetching more bytes or retrying earlier anchors does
not fix it; the unindexed reader returns the correct frames.
Please seek backward and decode until the stored anchor is found, and add
real fMP4/HEVC regression coverage. Changing only this flag to `backward=True`
made all 750 frame comparisons pass across closed/open-GOP H.264, fragmented
H.264, and HEVC, without changing the stored fields.
##########
paimon-python/pypaimon/table/row/video_keyframe_index.py:
##########
@@ -0,0 +1,292 @@
+# 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('<BQII')
+ METADATA_RANGE = struct.Struct('<qq')
+ ENTRY = struct.Struct('<qqq')
+ _CHUNK_SIZE = 64 * 1024
+
+ def __init__(self, metadata_ranges, keyframes):
+ self.metadata_ranges = tuple(
+ (operator.index(offset), operator.index(length))
+ for offset, length in metadata_ranges
+ )
+ self.keyframes = tuple(
+ (operator.index(ordinal), operator.index(pts),
+ operator.index(position))
+ for ordinal, pts, position in keyframes
+ )
+ previous_end = 0
+ for offset, length in self.metadata_ranges:
+ if offset < previous_end or length <= 0:
+ raise ValueError(
+ "Video metadata ranges must be ordered, non-overlapping, "
+ "and non-empty."
+ )
+ previous_end = offset + length
+ if not self.keyframes or self.keyframes[0][0] != 0:
+ raise ValueError("Video keyframe index requires an initial
keyframe.")
+ previous_ordinal, previous_pts, previous_position = -1, None, -1
+ for ordinal, pts, position in self.keyframes:
+ 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:
+ raise ValueError(
+ "Keyframe packet positions must be non-negative and "
+ "strictly increasing."
+ )
+ previous_ordinal, previous_pts, previous_position = (
+ ordinal, pts, position)
+
+ def serialize(self):
+ entries = b''.join(self.ENTRY.pack(*entry) for entry in self.keyframes)
+ return self.HEADER.pack(
+ self.VERSION,
+ self.MAGIC,
+ len(self.metadata_ranges),
+ len(self.keyframes),
+ ) + b''.join(
+ self.METADATA_RANGE.pack(*value)
+ for value in self.metadata_ranges
+ ) + zlib.compress(entries)
+
+ @classmethod
+ def inspect(cls, source, payload_length):
+ """Build an index for the first video stream of an ISO BMFF file."""
+ try:
+ import av
+ except ImportError as error:
+ raise ImportError(
+ "Video seek-index generation requires PyAV."
+ ) from error
+
+ metadata_ranges = cls._iso_bmff_metadata_ranges(
+ source, payload_length)
+ source.seek(0)
+ with av.open(source) as container:
+ if not container.streams.video:
+ raise ValueError("Video has no video stream.")
+ stream = container.streams.video[0]
+ packet_positions = {}
+ for packet in container.demux(stream):
+ if (not packet.is_keyframe or packet.pts is None
+ or packet.pos is None or packet.pos < 0
+ or getattr(packet, "is_discard", False)):
+ continue
+ pts = cls._stream_pts(
+ packet.pts,
+ packet.time_base or stream.time_base,
+ stream.time_base,
+ )
+ if pts in packet_positions:
+ raise ValueError("Video has duplicate keyframe
timestamps.")
+ packet_positions[pts] = int(packet.pos)
+
+ source.seek(0)
+ with av.open(source) as container:
+ stream = container.streams.video[0]
+ keyframes = []
+ previous_pts = None
+ for ordinal, frame in enumerate(container.decode(stream)):
+ if frame.pts is None:
+ raise ValueError("Video frame has no presentation
timestamp.")
+ pts = cls._stream_pts(
+ frame.pts,
+ frame.time_base or stream.time_base,
+ stream.time_base,
+ )
+ if previous_pts is not None and pts <= previous_pts:
+ raise ValueError(
+ "Video frame timestamps must be strictly increasing."
+ )
+ if frame.key_frame:
+ position = packet_positions.get(pts)
+ if position is None:
+ raise ValueError(
+ "Video keyframe has no packet byte position."
+ )
+ keyframes.append((ordinal, pts, position))
+ previous_pts = pts
+ return cls(metadata_ranges, keyframes)
+
+ @staticmethod
+ def _stream_pts(pts, time_base, stream_time_base):
+ value = pts * time_base / stream_time_base
+ if value.denominator != 1:
+ raise ValueError(
+ "Video timestamp does not fit stream time base."
+ )
+ return int(value)
+
+ @staticmethod
+ def _iso_bmff_metadata_ranges(source, payload_length):
+ ranges = []
+ has_moov = False
+ offset = 0
+ while offset < payload_length:
+ source.seek(offset)
+ header = source.read(8)
+ if len(header) != 8:
+ raise ValueError("Invalid ISO BMFF box header.")
+ size, box_type = struct.unpack(">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 == b"mdat" else size))
Review Comment:
[P2] Exclude padding bodies from initialization ranges
Treating every non-`mdat` box as initialization data includes the entire
body of `free`/`skip` boxes. These ranges are fetched for every uncached
indexed read, so MP4 files with reserved space can retain almost full-file read
amplification despite the sparse index. I reproduced this with a valid MP4
remuxed using FFmpeg's `-moov_size 16777216`: reading frame 75 requests
16,779,069 bytes out of a 16,785,845-byte file. Keeping only the padding box
headers reduces that request to 4,329 bytes, and all 150 decoded frames still
match full decoding.
Please preserve the headers needed for box traversal while excluding
`free`/`skip` bodies from the persisted initialization ranges. Other boxes
should be handled according to their decoding requirements rather than dropped
indiscriminately.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
