| # 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 struct | |
| from pypaimon.common.mp4_moov import MoovLocator | |
| from pypaimon.common.delta_varint_compressor import DeltaVarintCompressor | |
| from pypaimon.schema.data_types import is_blob_type | |
| from pypaimon.table.row.blob import ( | |
| Blob, | |
| BlobRef, | |
| VideoFrameDescriptor, | |
| ) | |
| from pypaimon.execution.python.write.blob_format_writer import BlobFormatWriter | |
| class VideoFormatWriter(BlobFormatWriter): | |
| """Pack complete encoded videos and an embedded logical frame-run index.""" | |
| VERSION = 1 | |
| FOOTER_MAGIC_NUMBER = 0x4F454449 | |
| FOOTER_SIZE = 29 | |
| NULL_REFERENCE = -1 | |
| PLACE_HOLDER_REFERENCE = -2 | |
| def __init__( | |
| self, | |
| output_stream, | |
| file_path=None, | |
| copy_buffer_size=BlobFormatWriter.BUFFER_SIZE): | |
| super().__init__( | |
| output_stream, | |
| blob_consumer=None, | |
| file_path=file_path, | |
| copy_buffer_size=copy_buffer_size, | |
| ) | |
| self._physical_lengths = [] | |
| self._moov_offsets = [] | |
| self._moov_lengths = [] | |
| self._physical_videos = {} | |
| self._run_lengths = [] | |
| self._run_references = [] | |
| self._run_first_frames = [] | |
| self._current_run_length = 0 | |
| self._current_run_reference = None | |
| self._current_run_first_frame = 0 | |
| self._current_run_last_frame = 0 | |
| self._closed = False | |
| def add_element(self, row) -> None: | |
| if not hasattr(row, 'values') or len(row.values) != 1: | |
| raise ValueError("VideoFormatWriter only supports one field") | |
| if not is_blob_type(row.fields[0].type): | |
| raise ValueError( | |
| "VideoFormatWriter only supports one scalar BLOB field" | |
| ) | |
| value = row.values[0] | |
| if value is None: | |
| self._append(self.NULL_REFERENCE, 0) | |
| return | |
| if value is Blob.PLACE_HOLDER: | |
| self._append(self.PLACE_HOLDER_REFERENCE, 0) | |
| return | |
| if type(value) is not BlobRef: | |
| raise ValueError( | |
| "Video fields require an exact BlobRef containing a " | |
| "VideoFrameDescriptor." | |
| ) | |
| frame = value.to_descriptor() | |
| if not isinstance(frame, VideoFrameDescriptor): | |
| raise ValueError( | |
| "Video fields require an exact BlobRef containing a " | |
| "VideoFrameDescriptor." | |
| ) | |
| payload = frame.payload_descriptor | |
| ordinal = self._physical_videos.get(payload) | |
| if ordinal is None: | |
| length = self._write_video_payload(value) | |
| ordinal = len(self._physical_lengths) | |
| self._physical_lengths.append(length) | |
| self._physical_videos[payload] = ordinal | |
| self._append(ordinal, frame.frame_index) | |
| @property | |
| def physical_video_count(self) -> int: | |
| return len(self._physical_lengths) | |
| @property | |
| def run_count(self) -> int: | |
| return len(self._run_lengths) + (1 if self._current_run_length else 0) | |
| def close(self) -> None: | |
| if self._closed: | |
| return | |
| try: | |
| self._flush_run() | |
| physical_index = DeltaVarintCompressor.compress(self._physical_lengths) | |
| moov_offset_index = DeltaVarintCompressor.compress(self._moov_offsets) | |
| moov_length_index = DeltaVarintCompressor.compress(self._moov_lengths) | |
| run_length_index = DeltaVarintCompressor.compress(self._run_lengths) | |
| run_reference_index = DeltaVarintCompressor.compress( | |
| self._run_references | |
| ) | |
| first_frame_index = DeltaVarintCompressor.compress( | |
| self._run_first_frames | |
| ) | |
| for index in ( | |
| physical_index, | |
| moov_offset_index, | |
| moov_length_index, | |
| run_length_index, | |
| run_reference_index, | |
| first_frame_index, | |
| ): | |
| self.output_stream.write(index) | |
| self.output_stream.write(struct.pack( | |
| '<IIIIIIIB', | |
| len(physical_index), | |
| len(moov_offset_index), | |
| len(moov_length_index), | |
| len(run_length_index), | |
| len(run_reference_index), | |
| len(first_frame_index), | |
| self.FOOTER_MAGIC_NUMBER, | |
| self.VERSION, | |
| )) | |
| if hasattr(self.output_stream, 'flush'): | |
| self.output_stream.flush() | |
| finally: | |
| self._closed = True | |
| if hasattr(self.output_stream, 'close'): | |
| self.output_stream.close() | |
| def _write_video_payload(self, blob: BlobRef) -> int: | |
| start = self.position | |
| locator = MoovLocator() | |
| stream = blob.new_input_stream() | |
| try: | |
| while True: | |
| chunk = stream.read(self.copy_buffer_size) | |
| if not chunk: | |
| break | |
| locator.feed(chunk) | |
| self.output_stream.write(chunk) | |
| self.position += len(chunk) | |
| finally: | |
| stream.close() | |
| length = self.position - start | |
| if length <= 0: | |
| raise ValueError("Encoded video payload must not be empty.") | |
| moov_offset, moov_length = locator.finish() | |
| self._moov_offsets.append(moov_offset) | |
| self._moov_lengths.append(moov_length) | |
| return length | |
| def _append(self, reference: int, frame_index: int) -> None: | |
| if self._can_extend(reference, frame_index): | |
| self._current_run_length += 1 | |
| self._current_run_last_frame = frame_index | |
| return | |
| self._flush_run() | |
| self._current_run_reference = reference | |
| self._current_run_first_frame = frame_index | |
| self._current_run_last_frame = frame_index | |
| self._current_run_length = 1 | |
| def _can_extend(self, reference: int, frame_index: int) -> bool: | |
| if ( | |
| self._current_run_length == 0 | |
| or self._current_run_reference != reference | |
| ): | |
| return False | |
| return ( | |
| reference < 0 | |
| or frame_index == self._current_run_last_frame + 1 | |
| ) | |
| def _flush_run(self) -> None: | |
| if self._current_run_length == 0: | |
| return | |
| self._run_lengths.append(self._current_run_length) | |
| self._run_references.append(self._current_run_reference) | |
| self._run_first_frames.append(self._current_run_first_frame) | |
| self._current_run_length = 0 |