################################################################################
# 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.
################################################################################
"""Video frame operators for extraction pipelines."""
import math
from pyflink.common import Row
from pyflink.dataframe import DataType, udf
from pyflink.table import DataTypes
from pyflink.table.udf import ScalarFunction, TableFunction, udtf
from pyflink.multimodal.codec.video import (
_normalize_frame_selector,
_normalize_video_ref,
_validate_positive_int,
extract_video_frames,
probe_video_metadata,
)
from pyflink.multimodal.types.video import (
VIDEO_FRAME_METADATA_TABLE_TYPE,
VIDEO_METADATA_FIELDS,
VIDEO_METADATA_TYPE,
VIDEO_STRUCT_TABLE_TYPE,
)
from pyflink.multimodal.utils import _build_or_apply_udf
__all__ = [
"video_metadata",
"video_split",
"video_explode_frames",
"video_extract_frames",
]
_VALID_VIDEO_ON_ERROR = {"raise", "null"}
def _normalize_on_error(on_error):
if on_error not in _VALID_VIDEO_ON_ERROR:
raise ValueError(
f"on_error must be one of {sorted(_VALID_VIDEO_ON_ERROR)}, got {on_error!r}"
)
return on_error
def _optional_time_ms(name, value):
if value is None:
return None
if not isinstance(value, int) or isinstance(value, bool) or value < 0:
raise ValueError(f"{name} must be a non-negative integer, got {value!r}")
return value
def _optional_positive_int(name, value):
if value is None:
return None
return _validate_positive_int(name, value)
def _optional_non_negative_int(name, value):
if value is None:
return None
if not isinstance(value, int) or isinstance(value, bool) or value < 0:
raise ValueError(f"{name} must be a non-negative integer, got {value!r}")
return value
def _validated_file_read_options(read_chunk_size, max_cached_blocks, read_ahead_blocks):
"""Validate public file-read options before passing them to codec helpers."""
options = {}
read_chunk_size = _optional_positive_int("read_chunk_size", read_chunk_size)
max_cached_blocks = _optional_non_negative_int(
"max_cached_blocks", max_cached_blocks)
read_ahead_blocks = _optional_non_negative_int(
"read_ahead_blocks", read_ahead_blocks)
if read_chunk_size is not None:
options["read_chunk_size"] = read_chunk_size
if max_cached_blocks is not None:
options["max_cached_blocks"] = max_cached_blocks
if read_ahead_blocks is not None:
options["read_ahead_blocks"] = read_ahead_blocks
return options
def _bridge_client(function_context, operator_name):
if not hasattr(function_context, "_get_file_system_bridge_client"):
raise RuntimeError(
f"{operator_name} requires a FunctionContext with FileSystem bridge access."
)
return function_context._get_file_system_bridge_client()
class _VideoMetadata(ScalarFunction):
"""Probe lightweight video container metadata."""
def __init__(self, on_error="raise", container_options=None,
read_chunk_size=None, max_cached_blocks=None,
read_ahead_blocks=None):
super().__init__()
self.on_error = _normalize_on_error(on_error)
self.container_options = container_options
self.file_read_options = _validated_file_read_options(
read_chunk_size, max_cached_blocks, read_ahead_blocks)
def open(self, function_context):
self._bridge_client = _bridge_client(function_context, "video_metadata")
def eval(self, uri):
if uri is None:
return None
try:
return probe_video_metadata(
uri,
self._bridge_client,
container_options=self.container_options,
**self.file_read_options,
)
except Exception:
if self.on_error == "null":
return None
raise
class _VideoSplit(TableFunction):
"""Split a video URI or segment ref into metadata-only segment refs."""
def __init__(
self,
segment_duration_ms=None,
num_segments=None,
video_duration_ms=None,
max_segments=1024,
on_error="raise",
container_options=None,
read_chunk_size=None,
max_cached_blocks=None,
read_ahead_blocks=None,
):
super().__init__()
if (segment_duration_ms is None) == (num_segments is None):
raise ValueError(
"exactly one of segment_duration_ms or num_segments is required")
self.segment_duration_ms = (
None if segment_duration_ms is None
else _validate_positive_int("segment_duration_ms", segment_duration_ms)
)
self.num_segments = (
None if num_segments is None
else _validate_positive_int("num_segments", num_segments)
)
self.video_duration_ms = _optional_time_ms(
"video_duration_ms", video_duration_ms)
self.max_segments = _validate_positive_int("max_segments", max_segments)
self.on_error = _normalize_on_error(on_error)
self.container_options = container_options
self.file_read_options = _validated_file_read_options(
read_chunk_size, max_cached_blocks, read_ahead_blocks)
def open(self, function_context):
self._bridge_client = _bridge_client(function_context, "video_split")
def eval(self, video_input, metadata=None):
try:
video_ref = _normalize_video_ref(video_input)
if video_ref is None:
return
duration_ms = _duration_ms_from_metadata(metadata)
if (
duration_ms is None
and video_ref.end_time_ms is None
and self.video_duration_ms is None
):
duration_ms = self._probe_duration_ms(video_ref.uri)
start_time_ms = video_ref.start_time_ms or 0
end_time_ms = (
video_ref.end_time_ms
if video_ref.end_time_ms is not None
else duration_ms - 1
if duration_ms is not None
else self.video_duration_ms - 1
if self.video_duration_ms is not None
else None
)
if end_time_ms is None:
raise ValueError(
"video_split requires video end_time_ms, video metadata "
"duration_ms, or video_duration_ms"
)
if end_time_ms < start_time_ms:
raise ValueError("segment end_time_ms must be >= start_time_ms")
if self.segment_duration_ms is not None:
yield from self._split_by_duration(
video_ref.uri, start_time_ms, end_time_ms)
else:
yield from self._split_by_count(
video_ref.uri, start_time_ms, end_time_ms)
except Exception:
if self.on_error == "null":
return
raise
def _split_by_duration(self, uri, start_time_ms, end_time_ms):
segment_start = start_time_ms
emitted = 0
while segment_start <= end_time_ms and emitted < self.max_segments:
segment_end = min(
segment_start + self.segment_duration_ms - 1, end_time_ms)
yield Row(_video_segment_ref(uri, segment_start, segment_end))
segment_start = segment_end + 1
emitted += 1
def _split_by_count(self, uri, start_time_ms, end_time_ms):
total_ms = end_time_ms - start_time_ms + 1
segment_duration_ms = int(math.ceil(float(total_ms) / self.num_segments))
emitted = 0
segment_start = start_time_ms
while (
segment_start <= end_time_ms
and emitted < self.num_segments
and emitted < self.max_segments
):
segment_end = min(segment_start + segment_duration_ms - 1,
end_time_ms)
yield Row(_video_segment_ref(uri, segment_start, segment_end))
segment_start = segment_end + 1
emitted += 1
def _probe_duration_ms(self, uri):
metadata = probe_video_metadata(
uri,
self._bridge_client,
container_options=self.container_options,
**self.file_read_options,
)
return _duration_ms_from_metadata(metadata)
[docs]def video_split(
*columns,
segment_duration_ms=None,
num_segments=None,
video_duration_ms=None,
max_segments=1024,
on_error="raise",
container_options=None,
read_chunk_size=None,
max_cached_blocks=None,
read_ahead_blocks=None,
concurrency=None,
):
"""
Split a video URI or segment ref into metadata-only segment refs.
For full-video URI inputs, the operator probes video metadata when no
metadata row, duration_ms value, explicit ``video_duration_ms``, or segment
``end_time_ms`` is provided. Passing metadata from ``video_metadata`` keeps
this as a metadata-only split and avoids the extra probe.
Args:
*columns: Optional ``uri``/``VIDEO_STRUCT`` column, plus an optional
video metadata row or duration_ms value. When provided, the
UDTF is applied directly instead of returning a factory.
segment_duration_ms: Fixed segment duration in milliseconds. Use this
when each output segment should cover the same time window, such as
one segment per second. Exactly one of ``segment_duration_ms`` or
``num_segments`` is required.
num_segments: Target number of approximately equal segments. Use this
when each input should produce a bounded number of segments
regardless of duration. Exactly one of ``segment_duration_ms`` or
``num_segments`` is required.
video_duration_ms: Optional constant duration fallback in milliseconds.
max_segments: Safety limit for emitted segments per input row.
on_error: Error handling policy. ``"raise"`` propagates failures;
``"null"`` drops unreadable inputs.
container_options: Optional PyAV ``av.open`` options used only when
this operator probes metadata itself.
read_chunk_size: Optional Java-backed file read chunk size in bytes
used only when this operator probes metadata itself. Must not
exceed the 16 MiB Java bridge read range limit.
max_cached_blocks: Optional per-file cache block count. ``0`` disables
block caching.
read_ahead_blocks: Optional number of blocks to prefetch after a cache
miss. ``0`` disables read-ahead.
concurrency: Optional UDTF operator parallelism.
Returns:
A UDTF producing one VIDEO_STRUCT column:
``ROW<uri STRING, start_time_ms BIGINT, end_time_ms BIGINT>``.
Example::
>>> segments = table.join_lateral(
... video_split(segment_duration_ms=1000)(
... col("uri")).alias("segment"))
"""
wrapper = udtf(
_VideoSplit(
segment_duration_ms=segment_duration_ms,
num_segments=num_segments,
video_duration_ms=video_duration_ms,
max_segments=max_segments,
on_error=on_error,
container_options=container_options,
read_chunk_size=read_chunk_size,
max_cached_blocks=max_cached_blocks,
read_ahead_blocks=read_ahead_blocks,
),
result_types=[VIDEO_STRUCT_TABLE_TYPE],
concurrency=concurrency,
)
return _build_or_apply_udf(wrapper, *columns)
class _VideoExplodeFrames(TableFunction):
"""Extract selected frames from a video or metadata-only segment."""
def __init__(
self,
frame_selector="keyframe",
sample_interval_ms=None,
max_frames=None,
on_error="raise",
container_options=None,
read_chunk_size=None,
max_cached_blocks=None,
read_ahead_blocks=None,
):
super().__init__()
self.frame_selector = _normalize_frame_selector(frame_selector)
self.sample_interval_ms = (
None if sample_interval_ms is None
else _validate_positive_int("sample_interval_ms", sample_interval_ms)
)
self.max_frames = _optional_positive_int("max_frames", max_frames)
self.on_error = _normalize_on_error(on_error)
self.container_options = container_options
self.file_read_options = _validated_file_read_options(
read_chunk_size, max_cached_blocks, read_ahead_blocks)
def open(self, function_context):
self._bridge_client = _bridge_client(
function_context, "video_explode_frames")
def eval(self, video_input):
try:
yield from extract_video_frames(
video_input,
self._bridge_client,
frame_selector=self.frame_selector,
sample_interval_ms=self.sample_interval_ms,
max_frames=self.max_frames,
container_options=self.container_options,
**self.file_read_options,
)
except Exception:
if self.on_error == "null":
return
raise
[docs]def video_explode_frames(
*columns,
frame_selector="keyframe",
sample_interval_ms=None,
max_frames=None,
on_error="raise",
container_options=None,
read_chunk_size=None,
max_cached_blocks=None,
read_ahead_blocks=None,
concurrency=None,
):
"""
Expand selected video frames into one output row per frame.
Use the returned UDTF with ``Table.join_lateral`` / ``flat_map``. URI
inputs read the full video; segment refs only emit frames within
``start_time_ms`` and ``end_time_ms``.
Args:
*columns: Optional URI or segment column. When provided, the UDTF is
applied directly instead of returning a factory.
frame_selector: Frame selection mode, ``"keyframe"``, ``"sample"``,
or ``"all_frames"``.
sample_interval_ms: Sampling interval for ``frame_selector="sample"``.
max_frames: Optional safety limit for emitted frames per input row.
``None`` emits all selected frames.
on_error: Error handling policy. ``"raise"`` propagates failures;
``"null"`` drops unreadable inputs.
container_options: Optional PyAV ``av.open`` options.
read_chunk_size: Optional Java-backed file read chunk size in bytes.
Must not exceed the 16 MiB Java bridge read range limit.
max_cached_blocks: Optional per-file cache block count. ``0`` disables
block caching.
read_ahead_blocks: Optional number of blocks to prefetch after a cache
miss. ``0`` disables read-ahead.
concurrency: Optional UDTF operator parallelism.
Returns:
A UDTF producing two columns: ``IMAGE`` and ``VIDEO_FRAME_METADATA``.
Example::
>>> frames = table.join_lateral(
... video_explode_frames(frame_selector="keyframe")(
... col("uri")).alias("frame", "metadata"))
"""
wrapper = udtf(
_VideoExplodeFrames(
frame_selector=frame_selector,
sample_interval_ms=sample_interval_ms,
max_frames=max_frames,
on_error=on_error,
container_options=container_options,
read_chunk_size=read_chunk_size,
max_cached_blocks=max_cached_blocks,
read_ahead_blocks=read_ahead_blocks,
),
result_types=[DataTypes.IMAGE(), VIDEO_FRAME_METADATA_TABLE_TYPE],
concurrency=concurrency,
)
return _build_or_apply_udf(wrapper, *columns)
class _VideoExtractFrames(ScalarFunction):
"""Extract selected frames as one list per input video."""
def __init__(
self,
frame_selector="keyframe",
sample_interval_ms=None,
max_frames=None,
on_error="raise",
container_options=None,
read_chunk_size=None,
max_cached_blocks=None,
read_ahead_blocks=None,
):
super().__init__()
self.frame_selector = _normalize_frame_selector(frame_selector)
self.sample_interval_ms = (
None if sample_interval_ms is None
else _validate_positive_int("sample_interval_ms", sample_interval_ms)
)
self.max_frames = _optional_positive_int("max_frames", max_frames)
self.on_error = _normalize_on_error(on_error)
self.container_options = container_options
self.file_read_options = _validated_file_read_options(
read_chunk_size, max_cached_blocks, read_ahead_blocks)
def open(self, function_context):
self._bridge_client = _bridge_client(
function_context, "video_extract_frames")
def eval(self, video_input):
if video_input is None:
return None
try:
return [
frame
for frame, _ in extract_video_frames(
video_input,
self._bridge_client,
frame_selector=self.frame_selector,
sample_interval_ms=self.sample_interval_ms,
max_frames=self.max_frames,
container_options=self.container_options,
**self.file_read_options,
)
]
except Exception:
if self.on_error == "null":
return None
raise
def _duration_ms_from_metadata(metadata):
if metadata is None:
return None
if isinstance(metadata, int) and not isinstance(metadata, bool):
return _optional_time_ms("duration_ms", metadata)
duration_ms = getattr(metadata, "duration_ms", None)
if duration_ms is not None:
return _optional_time_ms("duration_ms", duration_ms)
try:
duration_ms = metadata[VIDEO_METADATA_FIELDS.index("duration_ms")]
except (TypeError, IndexError) as e:
raise ValueError(
"video_metadata must be a duration_ms value or video_metadata row"
) from e
return _optional_time_ms("duration_ms", duration_ms)
def _video_segment_ref(uri, start_time_ms, end_time_ms):
return Row(uri=uri, start_time_ms=start_time_ms, end_time_ms=end_time_ms)