diff --git a/docs/docs/pypaimon/video.md b/docs/docs/pypaimon/video.md index 30b347755014..589f1647d731 100644 --- a/docs/docs/pypaimon/video.md +++ b/docs/docs/pypaimon/video.md @@ -24,7 +24,48 @@ under the License. # Video Frames -Store complete encoded videos and read their logical frame rows for training. First create a multimodal `conn` as shown in [Multimodal Tables](./multimodal-tables#connect). Provide your own video files and frame metadata for the examples below. +Store encoded videos in `.video` and read individual frames for training. + +## Import a Video + +With Python 3.10 or newer, install `pypaimon[video]`: + +```shell +pip install 'pypaimon[video]' +``` + +```python +import pyarrow as pa +import pypaimon.multimodal as pm + +conn = pm.connect(options={"warehouse": "file:///tmp/video_lake"}) +frames = conn.create_table( + "auto_frames", + schema=pa.schema([ + pa.field("frame_index", pa.int64()), + pa.field("timestamp", pa.float64()), + pa.field("episode_id", pa.int64()), + pa.field("video", pa.large_binary()), + ]), + options={"video-frame-field": "video"}, +) +frames.add_video("/data/example.mp4") +frames.add_video("oss://source-bucket/episode-42.mp4", metadata={"episode_id": 42}) +# Raw HEVC: explicitly supply a fixed frame rate. +frames.add_video("oss://source-bucket/episode-42.h265", source_fps=30) +``` + +Generates zero-based `frame_index` and `timestamp` in seconds from the first +frame. `metadata` supplies per-video constants. Containerized videos retain +original bytes and variable frame timing; require one decodable video stream +with strictly increasing PTS. + +Raw HEVC is remuxed to MP4 without re-encoding. `source_fps` generates fixed-rate +timestamps, not original capture times. Temporary raw/MP4 files are cleaned up +after import, including failures. Omit `source_fps` for containerized videos. + +Configure remote credentials on the connection and keep the source unchanged +during import. To supply frame metadata yourself, pass `frames` as shown below. ## Video Frame Storage diff --git a/paimon-python/pypaimon/common/file_io.py b/paimon-python/pypaimon/common/file_io.py index 178802f3f23f..0049b9410421 100644 --- a/paimon-python/pypaimon/common/file_io.py +++ b/paimon-python/pypaimon/common/file_io.py @@ -31,6 +31,15 @@ _LOG = logging.getLogger(__name__) +def _is_local_path(path): + from urllib.parse import urlparse + + uri = urlparse(path) + # urlparse treats a Windows drive letter as a URI scheme. + return not uri.scheme or uri.scheme == "file" or ( + os.name == "nt" and len(uri.scheme) == 1 and not uri.netloc) + + def normalize_naive_datetimes(value): """Recursively attach UTC to naive datetimes, including ones nested inside ROW (dict) and ARRAY/MAP (list/tuple) values. @@ -670,9 +679,7 @@ def get(path: str, catalog_options: Optional[Options] = None) -> 'FileIO': uri = urlparse(path) scheme = uri.scheme - # urlparse treats a Windows drive letter as a URI scheme. - if not scheme or scheme == "file" or ( - _os.name == "nt" and len(scheme) == 1 and not uri.netloc): + if _is_local_path(path): from pypaimon.filesystem.local_file_io import LocalFileIO return LocalFileIO(path, catalog_options) diff --git a/paimon-python/pypaimon/filesystem/local_file_io.py b/paimon-python/pypaimon/filesystem/local_file_io.py index ea0acbe46020..5cb7b4946073 100644 --- a/paimon-python/pypaimon/filesystem/local_file_io.py +++ b/paimon-python/pypaimon/filesystem/local_file_io.py @@ -69,7 +69,8 @@ def __init__(self, path: str = None, catalog_options: Optional[Options] = None): def create(): return LocalFileIO() - def _to_file(self, path: str) -> Path: + @staticmethod + def _to_file(path: str) -> Path: parsed = urlparse(path) if parsed.scheme and len(parsed.scheme) == 1 and not parsed.netloc: diff --git a/paimon-python/pypaimon/multimodal/table.py b/paimon-python/pypaimon/multimodal/table.py index 859a9036dd80..4c9a36b95c91 100644 --- a/paimon-python/pypaimon/multimodal/table.py +++ b/paimon-python/pypaimon/multimodal/table.py @@ -124,17 +124,116 @@ def add(self, data): def add_video( self, video, - frames, + frames=None, *, video_column=None, - first_frame=0): - """Append logical frame rows backed by one complete encoded video. - - ``frames`` supplies every table column except the configured - ``video-frame-field``. Frame ordinals are generated from - ``first_frame`` and stored in the video descriptor, not in the normal - data file. + first_frame=0, + metadata=None, + source_fps=None): + """Append frame rows sharing one encoded video. + + ``frames`` supplies non-video columns; descriptor ordinals start at + ``first_frame``. Omit it to generate ``frame_index`` (int64) and + ``timestamp`` (float64 seconds relative to the first frame). + ``metadata`` supplies constant fields. Auto import requires PyAV + and one video stream. For raw HEVC, ``source_fps`` sets fixed-rate + timing and remuxes to MP4 without re-encoding. """ + if frames is None: + import os + import stat + import tempfile + from contextlib import ExitStack + from urllib.parse import urlparse + + from pypaimon.table.row.blob import Blob, BlobDescriptor + from pypaimon.common.file_io import _is_local_path + from pypaimon.filesystem.local_file_io import LocalFileIO + from pypaimon.multimodal.video import ( + _inspect_video_frames, _require_video_av, _raw_video_time_base, _remux_raw_hevc, + ) + + time_base = _raw_video_time_base(source_fps) if source_fps is not None else None + + if isinstance(first_frame, bool) or not isinstance(first_frame, int) or first_frame != 0: + raise ValueError("first_frame must be 0 when frames is omitted.") + column = self._resolve_video_frame_column(video_column) + schema = _target_schema(self.raw_table) + generated = {"frame_index": pa.int64(), "timestamp": pa.float64()} + for name, dtype in generated.items(): + if name == column or name not in schema.names or schema.field(name).type != dtype: + raise ValueError( + "Automatic video import requires column %r of type %s." + % (name, dtype)) + if metadata is not None and not isinstance(metadata, Mapping): + raise ValueError("metadata must be a mapping or None.") + metadata = dict(metadata or {}) + allowed = set(schema.names) - set(generated) - set(self.raw_table.options.video_frame_fields()) + if set(metadata) - allowed: + raise ValueError("metadata must contain only non-video, non-generated columns.") + fields = [field for field in schema if field.name not in generated and field.name != column] + constants = _to_arrow_table([metadata], pa.schema(fields)) if fields else None + for field in fields: + if not field.nullable and constants[field.name].null_count: + raise ValueError("Missing required metadata: %s." % field.name) + if isinstance(video, os.PathLike): + video = os.fspath(video) + if isinstance(video, str): + if not urlparse(video).scheme: + video = os.path.abspath(video) + video = BlobDescriptor(video, 0, -1) + if not isinstance(video, BlobDescriptor): + raise ValueError("Automatic video import requires a path, URI, or BlobDescriptor.") + _require_video_av() + local_path = LocalFileIO._to_file(video.uri) if _is_local_path(video.uri) else None + + def source_state(): + if local_path is None: + return None + info = os.stat(local_path) + if not stat.S_ISREG(info.st_mode): + raise ValueError("Automatic video import requires a regular local file.") + return (info.st_dev, info.st_ino, info.st_size, info.st_mtime_ns, info.st_ctime_ns) + + with ExitStack() as stack: + original = source_state() + # Reuse the writer's source credentials. + reader = self.raw_table.file_io.uri_reader_factory.create(video.uri) + with Blob.from_descriptor(reader, video).new_input_stream() as stream: + if time_base is not None: + directory = stack.enter_context(tempfile.TemporaryDirectory(prefix="paimon_hevc_")) + path = os.path.join(directory, "video.mp4") + count = _remux_raw_hevc(stream, path, time_base) + video = BlobDescriptor(path, 0, os.path.getsize(path)) + timestamps = _inspect_video_frames(path) + if len(timestamps) != count: + raise ValueError("HEVC frame count changed during remux.") + else: + timestamps = _inspect_video_frames(stream) + + def check_source(): + if source_state() != original: + raise ValueError("Video source changed during import.") + + check_source() + + def batches(): + for start in range(0, len(timestamps), 1024): + end = min(start + 1024, len(timestamps)) + values = { + "frame_index": pa.array(range(start, end), type=pa.int64()), + "timestamp": pa.array(timestamps[start:end], type=pa.float64()), + } + if constants is not None: + repeated = constants.take(pa.array([0] * (end - start), type=pa.int32())) + values.update({name: repeated[name] for name in repeated.column_names}) + yield _video_frame_table(video, pa.table(values), column, start, schema) + # Check again before commit. + check_source() + + return self.add_batches(batches()) + elif metadata is not None or source_fps is not None: + raise ValueError("metadata and source_fps are only supported when frames is omitted.") return self.add_videos( [(video, frames, first_frame)], video_column=video_column ) diff --git a/paimon-python/pypaimon/multimodal/video.py b/paimon-python/pypaimon/multimodal/video.py index a45b01b0a0cf..f70ea6395a39 100644 --- a/paimon-python/pypaimon/multimodal/video.py +++ b/paimon-python/pypaimon/multimodal/video.py @@ -30,6 +30,131 @@ from pypaimon.table.row.video_keyframe_index import VideoKeyframeIndex +def _require_video_av(): + """Check the optional decoder before opening a video source.""" + try: + import av + except ImportError as error: + raise ImportError( + "Automatic video import requires PyAV; install pypaimon[video]." + ) from error + return av + + +def _raw_video_time_base(fps): + from numbers import Real + + if isinstance(fps, bool) or not isinstance(fps, Real): + raise ValueError("source_fps must be a positive finite number.") + try: + rate = Fraction(str(fps)) + if rate <= 0 or max(rate.numerator, rate.denominator) > 2147483647: + raise ValueError() + except (ValueError, OverflowError): + raise ValueError("source_fps must be a positive finite representable frame rate.") + return 1 / rate + + +def _remux_raw_hevc(source, output_path, time_base): + """Preserve HEVC packets and assign constant-rate presentation timestamps.""" + from array import array + import shutil + import tempfile + + av = _require_video_av() + # Snapshot to disk for both passes. + with tempfile.TemporaryFile() as raw: + shutil.copyfileobj(source, raw, length=1024 * 1024) + raw.seek(0) + ordinals = array("q") + frame_count = 0 + delay = 0 + with av.open(raw) as container: + if container.format.name != "hevc" or len(container.streams) != 1: + raise ValueError("source_fps is only supported for a raw H.265/HEVC stream.") + for packet in container.demux(video=0): + if packet.size: + # Track packet IDs through B-frame reordering. + packet.pts = len(ordinals) + packet.dts = None + packet.time_base = time_base + ordinals.append(-1) + for frame in packet.decode(): + identity = frame.pts + if frame.is_corrupt or identity is None or not 0 <= identity < len(ordinals) \ + or ordinals[identity] != -1: + raise ValueError("Cannot determine raw HEVC frame order.") + ordinals[identity] = frame_count + delay = max(delay, identity - frame_count) + frame_count += 1 + if not frame_count or frame_count != len(ordinals): + raise ValueError("Raw HEVC must contain one decodable frame per packet.") + raw.seek(0) + with av.open(raw) as container, av.open(output_path, "w", format="mp4") as output: + if hasattr(output, "add_stream_from_template"): + stream = output.add_stream_from_template(container.streams.video[0], opaque=True) + else: + stream = output.add_stream(template=container.streams.video[0]) + stream.time_base = time_base + output.start_encoding() + if (time_base / stream.time_base).denominator != 1: + raise ValueError("MP4 cannot represent source_fps exactly.") + packet_time_bases = {} + count = 0 + for packet in container.demux(video=0): + if not packet.size: + continue + # Keep PyAV 12's read-only duration; bound the AVRational denominator. + duration = packet.duration + if duration <= 0: + raise ValueError("Raw HEVC packet has no duration.") + if duration not in packet_time_bases: + exact = time_base / duration + bounded = exact.limit_denominator(2147483647) + error = abs(bounded - exact) * duration * (frame_count + delay + 1) + # Every PTS/DTS and duration must round to the same MP4 tick. + if bounded <= 0 or bounded.numerator > 2147483647 or error >= stream.time_base / 2: + raise ValueError("Cannot represent raw HEVC timestamps exactly in MP4.") + packet_time_bases[duration] = bounded + packet.pts = ordinals[count] * duration + packet.dts = (count - delay) * duration + packet.time_base = packet_time_bases[duration] + packet.stream = stream + output.mux(packet) + count += 1 + if count != frame_count: + raise ValueError("Raw HEVC packet count changed during remux.") + return frame_count + + +def _inspect_video_frames(source): + """Read presentation timestamps without retaining decoded images.""" + av = _require_video_av() + from array import array + + timestamps = array("d") + first_time = None + previous_time = None + with av.open(source) as container: + if len(container.streams.video) != 1: + raise ValueError("Automatic video import requires one video stream.") + for frame in container.decode(container.streams.video[0]): + if frame.is_corrupt: + raise ValueError("Video contains a corrupt frame.") + if frame.pts is None or frame.time_base is None: + raise ValueError("Video frame is missing a presentation timestamp.") + current_time = frame.pts * frame.time_base + if previous_time is not None and current_time <= previous_time: + raise ValueError("Video presentation timestamps must be strictly increasing.") + if first_time is None: + first_time = current_time + timestamps.append(float(current_time - first_time)) + previous_time = current_time + if not timestamps: + raise ValueError("Video contains no decodable frames.") + return timestamps + + class _VideoDecodePool: """One lazily created decode pool per dataset and worker process.""" diff --git a/paimon-python/pypaimon/tests/multimodal_table_test.py b/paimon-python/pypaimon/tests/multimodal_table_test.py index e1db3625de58..be94381585a6 100644 --- a/paimon-python/pypaimon/tests/multimodal_table_test.py +++ b/paimon-python/pypaimon/tests/multimodal_table_test.py @@ -253,6 +253,557 @@ def alter(changes): with self.assertRaisesRegex(RuntimeError, "not supported yet"): alter([C.set_option("video-frame-field", "video,second,later")]) + def _auto_video_table(self, name="auto_frames", fields=None, video_fields="video"): + return self.conn.create_table( + name, + schema=pa.schema(fields or [ + pa.field("frame_index", pa.int64(), nullable=False), + pa.field("timestamp", pa.float64(), nullable=False), + pa.field("episode_id", pa.int64()), + pa.field("video", pa.large_binary()), + ]), + options=dict(_PARQUET_OPTIONS, **{ + "video-frame-field": video_fields, "blob-as-descriptor": "true", + }), + ) + + def _auto_video_file(self, pts=(0, 1, 2, 3, 4), streams=1): + av = pytest.importorskip("av") + import numpy as np + from fractions import Fraction + path = os.path.join(self.temp_dir, "auto.mp4") + with av.open(path, "w", format="mp4") as container: + videos = [container.add_stream("libx264", rate=10) for _ in range(streams)] + for stream in videos: + stream.width = stream.height = 16 + stream.pix_fmt = "yuv420p" + stream.gop_size = 4 + stream.codec_context.max_b_frames = 2 + for i, pts_value in enumerate(pts): + for stream in videos: + frame = av.VideoFrame.from_ndarray( + np.full((16, 16, 3), i % 200, dtype=np.uint8), format="rgb24") + frame.pts, frame.time_base = pts_value, Fraction(1, 10) + for packet in stream.encode(frame): + container.mux(packet) + for stream in videos: + for packet in stream.encode(): + container.mux(packet) + return path + + def test_auto_video_real_timestamps_and_payload(self): + av = pytest.importorskip("av") + import numpy as np + from pathlib import Path + for number, pts in enumerate([(0, 1, 2, 3, 4), (50, 51, 54, 60, 62)]): + with self.subTest(pts=pts): + path = self._auto_video_file(pts) + original = Path(path).read_bytes() + with av.open(path) as container: + decoded = list(container.decode(video=0)) + times = [f.pts * f.time_base for f in decoded] + pixels = [f.to_ndarray(format="rgb24") for f in decoded] + self.assertEqual(len(pts), len(times)) + if number: + self.assertGreater(times[0], 0) + self.assertGreater(len(set(b - a for a, b in zip(times, times[1:]))), 1) + table = self._auto_video_table("auto_frames_%d" % number) + self.assertIs(table, table.add_video(Path(path), metadata={"episode_id": 42})) + rows = sorted(table.scan().to_list(), key=lambda r: r["frame_index"]) + self.assertEqual(list(range(len(pts))), [r["frame_index"] for r in rows]) + self.assertEqual([float(t - times[0]) for t in times], [r["timestamp"] for r in rows]) + self.assertEqual([42] * len(pts), [r["episode_id"] for r in rows]) + descriptors = [pmm.VideoFrameDescriptor.deserialize(r["video"]) for r in rows] + self.assertEqual(list(range(len(pts))), [d.frame_index for d in descriptors]) + self.assertTrue(descriptors[0].uri.endswith(".video")) + _, bodies = table.scan().read_blobs("video") + self.assertTrue(all(body == original for body in bodies["video"])) + with av.open(io.BytesIO(bodies["video"][0])) as container: + for expected, frame in zip(pixels, container.decode(video=0)): + np.testing.assert_array_equal(expected, frame.to_ndarray(format="rgb24")) + + def _raw_hevc_file(self, b_frames=3): + av = pytest.importorskip("av") + import numpy as np + if "libx265" not in av.codecs_available: + self.skipTest("libx265 is unavailable") + path = os.path.join(self.temp_dir, "source.h265") + with av.open(path, "w", format="hevc") as container: + stream = container.add_stream("libx265", rate=24) + stream.width = stream.height = 64 + stream.pix_fmt = "yuv420p" + stream.codec_context.gop_size = 8 + stream.codec_context.max_b_frames = b_frames + stream.codec_context.options = { + "preset": "ultrafast", + "x265-params": "log-level=error:pools=none:frame-threads=1:scenecut=0", + } + for i in range(24): + frame = av.VideoFrame.from_ndarray( + np.full((64, 64, 3), i * 9, dtype=np.uint8), format="rgb24") + for packet in stream.encode(frame): + container.mux(packet) + for packet in stream.encode(): + container.mux(packet) + return path + + def test_auto_video_raw_hevc_remux_and_indexed_read(self): + av = pytest.importorskip("av") + import numpy as np + from fractions import Fraction + from pathlib import Path + from pypaimon.multimodal.lerobot.dataset import _PyAVVideoDecoder + from pypaimon.table.row.blob import Blob + from pypaimon.table.row.video_keyframe_index import VideoKeyframeIndex + + cases = [(0, 24), (3, Fraction(30000, 1001)), (3, 59.94), (3, Fraction(60000, 1001))] + for number, (b_frames, fps) in enumerate(cases): + with self.subTest(b_frames=b_frames, fps=fps): + path = self._raw_hevc_file(b_frames) + original = Path(path).read_bytes() + with av.open(path) as container: + decoded = list(container.decode(video=0)) + pixels = [f.to_ndarray(format="rgb24") for f in decoded] + self.assertTrue(all(f.pts is None for f in decoded)) + self.assertEqual(bool(b_frames), any(f.pict_type == 3 for f in decoded)) + table = self._auto_video_table("hevc_%d" % number) + self.assertIs(table, table.add_video(path, source_fps=fps, metadata={"episode_id": 42})) + self.assertEqual(original, Path(path).read_bytes()) + rows = sorted(table.scan().to_list(), key=lambda r: r["frame_index"]) + self.assertEqual(list(range(24)), [r["frame_index"] for r in rows]) + np.testing.assert_allclose([i / float(fps) for i in range(24)], + [r["timestamp"] for r in rows], rtol=0, atol=1e-9) + self.assertTrue(all(r["episode_id"] == 42 for r in rows)) + descriptor = pmm.VideoFrameDescriptor.deserialize(rows[0]["video"]) + self.assertTrue(descriptor.uri.endswith(".video")) + index_descriptor = descriptor.keyframe_index_descriptor + self.assertIsNotNone(index_descriptor) + index = VideoKeyframeIndex.deserialize(Blob.from_file( + table.raw_table.file_io, index_descriptor.uri, + index_descriptor.offset, index_descriptor.length).to_data(), descriptor.length) + self.assertGreater(len(index.keyframes), 1) + _, bodies = table.scan().read_blobs("video") + payload = bodies["video"][0] + self.assertTrue(all(b == payload for b in bodies["video"])) + with av.open(io.BytesIO(payload)) as container: + stream = container.streams.video[0] + self.assertEqual("hevc", stream.codec_context.name) + self.assertEqual(24 / Fraction(str(fps)), stream.duration * stream.time_base) + actual = list(container.decode(video=0)) + self.assertEqual([i / Fraction(str(fps)) for i in range(24)], + [frame.pts * frame.time_base for frame in actual]) + self.assertEqual(24, len(actual)) + for expected, frame in zip(pixels, actual): + np.testing.assert_array_equal(expected, frame.to_ndarray(format="rgb24")) + decoder = _PyAVVideoDecoder(io.BytesIO(payload), keyframe_index=index) + try: + for i in [23, 0, 12, 7, 18, 1]: + np.testing.assert_array_equal(pixels[i], decoder[i].permute(1, 2, 0).numpy()) + finally: + decoder.close() + + def test_auto_video_raw_hevc_remote_descriptor(self): + from pathlib import Path + from unittest.mock import Mock + original = Path(self._raw_hevc_file()).read_bytes() + uri = "oss://source-bucket/video.h265" + table = self._auto_video_table() + factory = table.raw_table.file_io.uri_reader_factory + create = factory.create + stream = io.BytesIO(b"prefix" + original + b"suffix") + reader = Mock() + reader.new_input_stream.return_value = stream + with patch.object(factory, "create", side_effect=lambda path: reader if path == uri else create(path)): + table.add_video(pmm.BlobDescriptor(uri, 6, len(original)), source_fps=24) + reader.new_input_stream.assert_called_once_with(uri) + self.assertTrue(stream.closed) + self.assertEqual(24, len(table.scan().to_list())) + + def test_auto_video_raw_hevc_http(self): + from http.server import BaseHTTPRequestHandler, HTTPServer + from pathlib import Path + from threading import Thread + original = Path(self._raw_hevc_file()).read_bytes() + + class Handler(BaseHTTPRequestHandler): + def do_GET(self): + self.send_response(200) + self.send_header("Content-Length", str(len(original))) + self.end_headers() + self.wfile.write(original) + + def log_message(self, *args): + pass + + server = HTTPServer(("127.0.0.1", 0), Handler) + thread = Thread(target=server.serve_forever) + thread.start() + try: + table = self._auto_video_table() + with tempfile.TemporaryDirectory() as directory: + with patch("tempfile.tempdir", directory): + table.add_video("http://127.0.0.1:%d/video.h265" % server.server_port, source_fps=30) + self.assertEqual([], os.listdir(directory)) + rows = table.scan().select(["frame_index", "timestamp"]).to_list() + self.assertEqual(24, len(rows)) + self.assertTrue(all(abs(r["timestamp"] - r["frame_index"] / 30) < 1e-9 for r in rows)) + finally: + server.shutdown() + server.server_close() + thread.join() + + def test_auto_video_raw_hevc_source_change_before_write(self): + from pathlib import Path + from pypaimon.multimodal.video import _remux_raw_hevc + path = self._raw_hevc_file() + table = self._auto_video_table() + + def remux(*args): + count = _remux_raw_hevc(*args) + Path(path).write_bytes(b"replaced") + return count + + with patch("pypaimon.multimodal.video._remux_raw_hevc", side_effect=remux): + with self.assertRaisesRegex(ValueError, "source changed"): + table.add_video(path, source_fps=30) + self.assertIsNone(table.raw_table.snapshot_manager().get_latest_snapshot()) + + def test_auto_video_raw_hevc_frame_count_mismatch(self): + table = self._auto_video_table() + path = self._raw_hevc_file() + with patch("pypaimon.multimodal.video._inspect_video_frames", return_value=[0.0]): + with self.assertRaisesRegex(ValueError, "frame count changed"): + table.add_video(path, source_fps=30) + self.assertIsNone(table.raw_table.snapshot_manager().get_latest_snapshot()) + + def test_auto_video_raw_hevc_invalid_input_and_cleanup(self): + from pathlib import Path + table = self._auto_video_table() + factory = table.raw_table.file_io.uri_reader_factory + import sys + with patch.dict(sys.modules, {"av": None}), patch.object(factory, "create") as create: + with self.assertRaises(ImportError): + table.add_video("oss://bucket/video.h265", source_fps=30) + create.assert_not_called() + with patch.object(factory, "create") as create: + for fps in [0, -1, True, "30", float("nan"), float("inf"), 1e20]: + with self.subTest(fps=fps), self.assertRaisesRegex(ValueError, "source_fps"): + table.add_video("oss://bucket/video.h265", source_fps=fps) + with self.assertRaisesRegex(ValueError, "frames is omitted"): + table.add_video("oss://bucket/video.h265", [], source_fps=30) + create.assert_not_called() + path = self._raw_hevc_file() + with self.assertRaisesRegex(ValueError, "missing a presentation timestamp"): + table.add_video(path) + with self.assertRaisesRegex(ValueError, "raw H.265"): + table.add_video(self._auto_video_file(), source_fps=30) + with tempfile.TemporaryDirectory() as directory: + with patch("tempfile.tempdir", directory): + with patch.object(table, "add_batches", side_effect=RuntimeError("write failed")): + with self.assertRaisesRegex(RuntimeError, "write failed"): + table.add_video(path, source_fps=30) + self.assertEqual([], os.listdir(directory)) + Path(path).write_bytes(b"not a video") + with self.assertRaises(Exception): + table.add_video(path, source_fps=30) + self.assertEqual([], os.listdir(directory)) + self.assertIsNone(table.raw_table.snapshot_manager().get_latest_snapshot()) + + def test_auto_video_missing_pyav_does_not_access_source(self): + import sys + table = self._auto_video_table() + factory = table.raw_table.file_io.uri_reader_factory + for source in ["http://example.invalid/video.mp4", "oss://bucket/video.mp4", "/missing.mp4"]: + with self.subTest(source=source), patch.dict(sys.modules, {"av": None}), \ + patch.object(factory, "create") as create, patch("os.stat") as stat: + with self.assertRaisesRegex(ImportError, r"pypaimon\[video\]"): + table.add_video(source) + create.assert_not_called() + stat.assert_not_called() + self.assertIsNone(table.raw_table.snapshot_manager().get_latest_snapshot()) + + def test_auto_video_local_uri_changes_use_fileio_path(self): + from pathlib import PurePosixPath, PureWindowsPath + from types import SimpleNamespace + from unittest.mock import Mock + pytest.importorskip("av") + table = self._auto_video_table() + factory = table.raw_table.file_io.uri_reader_factory + cases = [ + ("nt", r"C:\video.mp4", r"C:\video.mp4"), + ("nt", "file://server/share/video%20one.mp4", r"\\server\share\video one.mp4"), + ("posix", "file://server/share/video%20one.mp4", "//server/share/video one.mp4"), + ("nt", "file://C:/video.mp4", r"C:\video.mp4"), + ("nt", r"\\server\share\video.mp4", r"\\server\share\video.mp4"), + ] + for platform, uri, expected in cases: + with self.subTest(platform=platform, uri=uri): + source = io.BytesIO(b"video") + reader = Mock() + reader.new_input_stream.return_value = source + states = [SimpleNamespace(st_mode=0o100644, st_dev=1, st_ino=2, + st_size=5, st_mtime_ns=i, st_ctime_ns=i) for i in (1, 2)] + with patch.object(os, "name", platform), \ + patch("pypaimon.filesystem.local_file_io.Path", + PureWindowsPath if platform == "nt" else PurePosixPath), \ + patch("os.stat", side_effect=states) as stat, \ + patch.object(factory, "create", return_value=reader), \ + patch("pypaimon.multimodal.video._inspect_video_frames", return_value=[0.0]): + with self.assertRaisesRegex(ValueError, "source changed"): + value = pmm.BlobDescriptor(uri, 0, -1) if uri.startswith("\\\\") else uri + table.add_video(value) + self.assertEqual([expected, expected], [str(call[0][0]) for call in stat.call_args_list]) + self.assertTrue(source.closed) + self.assertIsNone(table.raw_table.snapshot_manager().get_latest_snapshot()) + + def test_auto_video_remote_uri_and_descriptor(self): + from pathlib import Path + from unittest.mock import Mock + original = Path(self._auto_video_file((50, 51, 54, 60, 62))).read_bytes() + uri = "oss://source-bucket/episode.mp4" + for number, source in enumerate([uri, pmm.BlobDescriptor(uri, 6, len(original))]): + with self.subTest(source=source): + table = self._auto_video_table("remote_%d" % number) + factory = table.raw_table.file_io.uri_reader_factory + create = factory.create + streams = [] + + def open_source(path): + self.assertEqual(uri, path) + stream = io.BytesIO(original if number == 0 else b"prefix" + original + b"suffix") + streams.append(stream) + return stream + + reader = Mock() + reader.new_input_stream.side_effect = open_source + with patch.object(factory, "create", side_effect=lambda path: reader if path == uri else create(path)): + table.add_video(source, metadata={"episode_id": 42}) + self.assertGreaterEqual(len(streams), 2) + self.assertTrue(all(stream.closed for stream in streams)) + rows = sorted(table.scan().select(["frame_index", "timestamp"]).to_list(), + key=lambda row: row["frame_index"]) + self.assertEqual(list(range(5)), [row["frame_index"] for row in rows]) + self.assertEqual([0.0, 0.1, 0.4, 1.0, 1.2], [row["timestamp"] for row in rows]) + _, bodies = table.scan().read_blobs("video") + self.assertEqual([original] * 5, bodies["video"]) + + def test_auto_video_http(self): + from http.server import BaseHTTPRequestHandler, HTTPServer + from pathlib import Path + from threading import Thread + original = Path(self._auto_video_file()).read_bytes() + + class Handler(BaseHTTPRequestHandler): + def do_GET(self): + self.send_response(200) + self.send_header("Content-Length", str(len(original))) + self.end_headers() + self.wfile.write(original) + + def log_message(self, *args): + pass + + server = HTTPServer(("127.0.0.1", 0), Handler) + thread = Thread(target=server.serve_forever) + thread.start() + try: + table = self._auto_video_table() + table.add_video("http://127.0.0.1:%d/video.mp4" % server.server_port) + _, bodies = table.scan().read_blobs("video") + self.assertEqual([original] * 5, bodies["video"]) + finally: + server.shutdown() + server.server_close() + thread.join() + + def test_auto_video_remote_failures_close_streams(self): + from pathlib import Path + from unittest.mock import Mock + original = Path(self._auto_video_file()).read_bytes() + uri = "oss://source-bucket/episode.mp4" + for number, (data, fail_after) in enumerate([(b"corrupt", 100), (original, 0), (original, 1)]): + with self.subTest(fail_after=fail_after): + table = self._auto_video_table("remote_failure_%d" % number) + factory = table.raw_table.file_io.uri_reader_factory + create = factory.create + streams = [] + + def open_source(path): + if len(streams) >= fail_after: + raise IOError("Source unavailable") + stream = io.BytesIO(data) + streams.append(stream) + return stream + + reader = Mock() + reader.new_input_stream.side_effect = open_source + with patch.object(factory, "create", side_effect=lambda path: reader if path == uri else create(path)): + with self.assertRaises(Exception): + table.add_video(uri) + self.assertTrue(all(stream.closed for stream in streams)) + self.assertIsNone(table.raw_table.snapshot_manager().get_latest_snapshot()) + + def test_auto_video_documented_schema(self): + table = self.conn.create_table( + "documented_frames", + schema=pa.schema([ + pa.field("frame_index", pa.int64()), + pa.field("timestamp", pa.float64()), + pa.field("episode_id", pa.int64()), + pa.field("video", pa.large_binary()), + ]), + options={"video-frame-field": "video"}, + ) + table.add_video(self._auto_video_file()) + self.assertEqual(5, len(table.scan().select(["frame_index", "timestamp"]).to_list())) + + def test_auto_video_batches_and_optional_metadata(self): + path = self._auto_video_file(range(1030)) + table = self._auto_video_table() + table.add_video(path) + rows = sorted(table.scan().select(["frame_index", "timestamp", "episode_id", "video"]).to_list(), + key=lambda r: r["frame_index"]) + self.assertEqual(1030, len(rows)) + self.assertEqual(list(range(1030)), [r["frame_index"] for r in rows]) + ordinals = [pmm.VideoFrameDescriptor.deserialize(r["video"]).frame_index for r in rows] + self.assertEqual(list(range(1030)), ordinals) + self.assertTrue(all(r["episode_id"] is None for r in rows)) + self.assertEqual(1, table.raw_table.snapshot_manager().get_latest_snapshot().id) + + def test_auto_video_validation_before_decode(self): + table = self._auto_video_table() + module = "pypaimon.multimodal.video._inspect_video_frames" + with patch(module) as inspect, patch("pypaimon.multimodal.video._require_video_av"): + for kwargs in [{"metadata": []}, {"metadata": {"timestamp": 1}}, + {"metadata": {"video": b"x"}}, {"metadata": {"unknown": 1}}, + {"metadata": {"episode_id": "wrong"}}, {"first_frame": 1}, + {"first_frame": False}, {"first_frame": 0.0}]: + with self.subTest(kwargs=kwargs), self.assertRaises((ValueError, TypeError, pa.ArrowException)): + table.add_video("missing.mp4", **kwargs) + for source in [b"bytes", io.BytesIO(), self.temp_dir]: + with self.subTest(source=source), self.assertRaises(ValueError): + table.add_video(source) + inspect.assert_not_called() + self.assertIsNone(table.raw_table.snapshot_manager().get_latest_snapshot()) + required = self._auto_video_table("required", fields=[ + pa.field("frame_index", pa.int64()), pa.field("timestamp", pa.float64()), + pa.field("episode_id", pa.int64(), nullable=False), pa.field("video", pa.large_binary()), + ]) + with self.assertRaisesRegex(ValueError, "required metadata"): + required.add_video("missing.mp4") + wrong = self._auto_video_table("wrong", fields=[ + pa.field("frame_index", pa.int32()), pa.field("timestamp", pa.float64()), + pa.field("video", pa.large_binary()), + ]) + with self.assertRaisesRegex(ValueError, "frame_index"): + wrong.add_video("missing.mp4") + with self.assertRaisesRegex(ValueError, "only supported"): + table.add_video("missing.mp4", [{"frame_index": 0}], metadata={}) + + def test_auto_video_multiple_columns(self): + table = self._auto_video_table(fields=[ + pa.field("frame_index", pa.int64()), pa.field("timestamp", pa.float64()), + pa.field("video", pa.large_binary()), pa.field("other", pa.large_binary()), + ], video_fields="video,other") + path = self._auto_video_file() + with self.assertRaisesRegex(ValueError, "video_column is required"): + table.add_video(path) + with self.assertRaisesRegex(ValueError, "non-video"): + table.add_video(path, video_column="video", metadata={"other": b"x"}) + table.add_video(path, video_column="video") + self.assertTrue(all(r["other"] is None for r in table.scan().to_list())) + + def test_auto_video_failure_preserves_snapshot(self): + from pathlib import Path + table = self._auto_video_table() + path = self._auto_video_file() + table.add_video(path) + snapshot = table.raw_table.snapshot_manager().get_latest_snapshot().id + Path(path).write_bytes(b"not a video") + with self.assertRaises(Exception): + table.add_video(path) + import wave + audio_path = os.path.join(self.temp_dir, "audio.wav") + with wave.open(audio_path, "wb") as audio: + audio.setnchannels(1) + audio.setsampwidth(2) + audio.setframerate(8000) + audio.writeframes(b"\x00\x00" * 800) + with self.assertRaisesRegex(ValueError, "one video stream"): + table.add_video(audio_path) + path = self._auto_video_file(streams=2) + with self.assertRaisesRegex(ValueError, "one video stream"): + table.add_video(path) + self.assertEqual(snapshot, table.raw_table.snapshot_manager().get_latest_snapshot().id) + + def test_auto_video_timestamp_failures_close_container(self): + from fractions import Fraction + from unittest.mock import MagicMock, Mock + pytest.importorskip("av") + table = self._auto_video_table() + path = self._auto_video_file() + for frames, message in [([], "no decodable"), + ([Mock(is_corrupt=True)], "corrupt"), + ([Mock(is_corrupt=False, pts=None, time_base=Fraction(1, 10))], "missing"), + ([Mock(is_corrupt=False, pts=0, time_base=None)], "missing"), + ([Mock(is_corrupt=False, pts=2, time_base=Fraction(1, 10)), + Mock(is_corrupt=False, pts=1, time_base=Fraction(1, 10))], "strictly increasing"), + ([Mock(is_corrupt=False, pts=1, time_base=Fraction(1, 10)), + Mock(is_corrupt=False, pts=1, time_base=Fraction(1, 10))], "strictly increasing")]: + container = MagicMock() + container.__enter__.return_value = container + container.streams.video = [object()] + container.decode.return_value = iter(frames) + with patch("av.open", return_value=container), self.assertRaisesRegex(ValueError, message): + table.add_video(path) + container.__exit__.assert_called_once() + self.assertIsNone(table.raw_table.snapshot_manager().get_latest_snapshot()) + container = MagicMock() + container.__enter__.return_value = container + container.streams.video = [] + with patch("av.open", return_value=container), self.assertRaisesRegex(ValueError, "one video stream"): + table.add_video(path) + container.__exit__.assert_called_once() + + def broken_decode(): + yield Mock(is_corrupt=False, pts=0, time_base=Fraction(1, 10)) + raise RuntimeError("decode failed") + + container.streams.video = [object()] + container.decode.return_value = broken_decode() + with patch("av.open", return_value=container), self.assertRaisesRegex(RuntimeError, "decode failed"): + table.add_video(path) + self.assertIsNone(table.raw_table.snapshot_manager().get_latest_snapshot()) + + def test_auto_video_changed_source_aborts_before_commit(self): + from pypaimon.multimodal.video import _inspect_video_frames + table = self._auto_video_table() + path = self._auto_video_file() + + def inspect_changed(source): + result = _inspect_video_frames(source) + with open(path, "ab") as output: + output.write(b"changed") + return result + + with patch("pypaimon.multimodal.video._inspect_video_frames", side_effect=inspect_changed): + with self.assertRaisesRegex(ValueError, "changed"): + table.add_video(path) + self.assertIsNone(table.raw_table.snapshot_manager().get_latest_snapshot()) + path = self._auto_video_file() + from pypaimon.multimodal.table import _video_frame_table + + def frame_table(*args): + result = _video_frame_table(*args) + with open(path, "ab") as output: + output.write(b"changed during write") + return result + + with patch("pypaimon.multimodal.table._video_frame_table", side_effect=frame_table): + with self.assertRaisesRegex(ValueError, "changed"): + table.add_video(path) + self.assertIsNone(table.raw_table.snapshot_manager().get_latest_snapshot()) + def test_add_video_embeds_frame_ordinals_outside_normal_data(self): table = self.conn.create_table( "video_frames", diff --git a/paimon-python/pypaimon/tests/multimodal_video_test.py b/paimon-python/pypaimon/tests/multimodal_video_test.py index 8e4a8fa1b682..8c7ff6fad948 100644 --- a/paimon-python/pypaimon/tests/multimodal_video_test.py +++ b/paimon-python/pypaimon/tests/multimodal_video_test.py @@ -50,6 +50,58 @@ def close(self): raise RuntimeError("decoder close failed") +class RawVideoRemuxTest(unittest.TestCase): + + def test_read_only_packet_duration(self): + from fractions import Fraction + from unittest.mock import MagicMock, patch + from pypaimon.multimodal.video import _remux_raw_hevc + + class Packet: + size = 1 + + def __init__(self, duration): + self._duration = duration + + @property + def duration(self): + return self._duration + + def decode(self): + return [SimpleNamespace(pts=self.pts, is_corrupt=False)] + + class Streams(list): + @property + def video(self): + return self + + for duration in (48000, 0): + with self.subTest(duration=duration), tempfile.TemporaryDirectory() as directory: + packets = [Packet(24000), Packet(duration)] + first, second, output = MagicMock(), MagicMock(), MagicMock() + for container in (first, second, output): + container.__enter__.return_value = container + container.streams = Streams([SimpleNamespace()]) + first.format.name = "hevc" + first.demux.return_value = [Packet(24000), Packet(48000)] + second.demux.return_value = packets + av = MagicMock() + av.open.side_effect = [first, second, output] + with patch("pypaimon.multimodal.video._require_video_av", return_value=av): + if duration == 0: + with self.assertRaisesRegex(ValueError, "no duration"): + _remux_raw_hevc(io.BytesIO(b"source"), os.path.join(directory, "out.mp4"), Fraction(1, 30)) + else: + self.assertEqual(2, _remux_raw_hevc( + io.BytesIO(b"source"), os.path.join(directory, "out.mp4"), Fraction(1, 30))) + self.assertEqual(2, output.mux.call_count) + for i, packet in enumerate(packets): + self.assertEqual(Fraction(i, 30), packet.pts * packet.time_base) + self.assertEqual(Fraction(i, 30), packet.dts * packet.time_base) + self.assertEqual(Fraction(1, 30), packet.duration * packet.time_base) + output.__exit__.assert_called_once() + + class ImageVideoEncoderTest(unittest.TestCase): def setUp(self):