################################################################################
# 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.
################################################################################
"""Audio decode, encode, split, and waveform transform operators.
The transform layer separates encoded boundary inputs from decoded waveform
processing. Operators that modify samples accept ``AUDIO_WAVEFORM`` only, so a
pipeline should usually call ``audio_decode`` once, transform waveforms, and
then call ``audio_encode`` or a sink-specific writer at the boundary.
"""
from pyflink.common.types import Row
from pyflink.dataframe import DataType, udf
from pyflink.table.udf import ScalarFunction, TableFunction, udtf
from pyflink.multimodal.codec.audio import (
DEFAULT_MAX_DECODED_AUDIO_BYTES,
MAX_AUDIO_WAVEFORM_DATA_BYTES,
_normalize_timestamp_ranges,
concat_waveforms,
decode_audio_input,
encode_audio_input,
ensure_supported_audio_input,
is_audio_data_error,
normalize_audio_format,
resample_waveform,
iter_audio_clip_ref_by_duration,
iter_audio_clip_ref_by_speech,
iter_audio_clip_ref_by_timestamps,
iter_split_waveform_by_duration,
iter_split_waveform_by_speech,
iter_split_waveform_by_timestamps,
standardize_waveform,
)
from pyflink.multimodal.types.audio import (
AUDIO_CLIP_REF_TABLE_TYPE,
AUDIO_WAVEFORM_TABLE_TYPE,
AUDIO_WAVEFORM_TYPE,
)
from pyflink.multimodal.utils import _build_or_apply_udf, _udf_runtime_kwargs
__all__ = [
"audio_decode",
"audio_encode",
"audio_convert_format",
"audio_standardize",
"audio_resample",
"audio_split_by_duration",
"audio_split_by_timestamp",
"audio_split_by_speech",
"audio_concat",
]
_VALID_AUDIO_DECODE_ON_ERROR = {"raise", "null"}
_VALID_SEGMENT_TYPES = {"audio", "ref"}
def _bridge_client_or_none(function_context):
getter = getattr(function_context, "_get_file_system_bridge_client", None)
if getter is None:
return None
return getter()
def _validate_positive_int(name, value):
if not isinstance(value, int) or isinstance(value, bool) or value <= 0:
raise ValueError(f"{name} must be a positive integer, got {value!r}")
return value
def _validate_max_decoded_bytes(value):
value = _validate_positive_int("max_decoded_bytes", value)
if value > MAX_AUDIO_WAVEFORM_DATA_BYTES:
raise ValueError(
f"max_decoded_bytes must be <= {MAX_AUDIO_WAVEFORM_DATA_BYTES}, "
f"got {value!r}"
)
return value
def _validate_non_negative_int(name, value):
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 _validate_optional_positive_int(name, value):
if value is None:
return None
return _validate_positive_int(name, value)
def _normalize_audio_decode_on_error(on_error):
if on_error not in _VALID_AUDIO_DECODE_ON_ERROR:
raise ValueError(
f"on_error must be one of {sorted(_VALID_AUDIO_DECODE_ON_ERROR)}, "
f"got {on_error!r}"
)
return on_error
def _normalize_segment_type(segment_type):
if segment_type not in _VALID_SEGMENT_TYPES:
raise ValueError(
"segment_type must be one of "
f"{sorted(_VALID_SEGMENT_TYPES)}, got {segment_type!r}"
)
return segment_type
class _AudioDecode(ScalarFunction):
"""Decode BYTES, URI, or AUDIO_CLIP_REF into AUDIO_WAVEFORM."""
def __init__(self,
on_error="raise",
max_decoded_bytes=DEFAULT_MAX_DECODED_AUDIO_BYTES):
super().__init__()
self.on_error = _normalize_audio_decode_on_error(on_error)
self.max_decoded_bytes = _validate_max_decoded_bytes(max_decoded_bytes)
self._bridge_client = None
def open(self, function_context):
self._bridge_client = _bridge_client_or_none(function_context)
def eval(self, audio_input):
if audio_input is None:
return None
ensure_supported_audio_input(
"audio_decode",
("BYTES", "URI", "AUDIO_CLIP_REF"),
audio_input,
)
if self.on_error == "raise":
return decode_audio_input(
audio_input,
bridge_client=self._bridge_client,
max_decoded_bytes=self.max_decoded_bytes,
)
try:
return decode_audio_input(
audio_input,
bridge_client=self._bridge_client,
max_decoded_bytes=self.max_decoded_bytes,
)
except Exception as e:
if is_audio_data_error(e):
return None
raise
[docs]def audio_decode(*columns,
on_error="raise",
max_decoded_bytes=DEFAULT_MAX_DECODED_AUDIO_BYTES,
concurrency=None):
"""
Decode audio into a waveform row.
The waveform row contains ``data`` plus audio layout fields:
``sample_rate``, ``channels``, ``frames``, ``sample_format``, ``layout``,
``duration_ms``, and optional source/segment fields. ``data`` stores a
variable-shape float32 tensor exposed to Python as a ``numpy.ndarray`` with
shape ``(frames, channels)``. Downstream operators should consume the row
through audio operators rather than parsing ``data`` manually.
For remote URI inputs, large encoded files may be copied to a worker-local
temporary file after metadata probing confirms that the decoded waveform
fits ``max_decoded_bytes``. The temporary file is an internal optimization
to avoid moving large byte payloads across the JVM/Python boundary; it is
cleaned up by the operator and is never returned to downstream operators.
Args:
*columns: Optional audio columns. Supported inputs are ``BYTES``, URI
``STRING``, and ``AUDIO_CLIP_REF``.
on_error: Handling strategy for corrupt supported audio
inputs:
* ``"raise"`` (default): propagate decode failures.
* ``"null"``: return ``NULL`` for rows that cannot be decoded.
URI access errors, unsupported input kinds, and invalid
``AUDIO_CLIP_REF`` values still raise.
max_decoded_bytes: Maximum decoded PCM payload allowed in one
``AUDIO_WAVEFORM`` row. The default is 1 GiB. Values above the
PyFlink tensor payload limit are rejected. This guard is checked
before materializing decoded samples and prevents compressed audio
from expanding into an unexpectedly large per-row waveform.
concurrency: UDF concurrency. ``None`` uses the framework default.
Returns:
A UDF producing ``ROW<data TENSOR('float32'), sample_rate INT, channels INT,
frames BIGINT, sample_format STRING, layout STRING, duration_ms BIGINT,
source_uri STRING, start_time_ms BIGINT, end_time_ms BIGINT>``.
Raises:
ValueError: If ``on_error`` is unsupported or a row value is not one of
the supported input forms. ``on_error="raise"`` also propagates
row-level media failures, including decoded-size guard failures and
short reads from truncated inputs.
Examples::
>>> # Usage 1: create a reusable UDF and apply it to a column.
>>> decode = audio_decode(
... on_error="null",
... max_decoded_bytes=1024 * 1024 * 1024,
... )
>>> df = df.with_column("waveform", decode(col("audio_bytes")))
>>>
>>> # Usage 2: pass the column directly when building the expression.
>>> df = df.with_column(
... "waveform",
... audio_decode(
... col("audio_bytes"),
... on_error="null",
... max_decoded_bytes=1024 * 1024 * 1024,
... ),
... )
"""
wrapper = udf(
_AudioDecode(on_error=on_error, max_decoded_bytes=max_decoded_bytes),
return_dtype=AUDIO_WAVEFORM_TYPE,
**_udf_runtime_kwargs(concurrency=concurrency),
)
return _build_or_apply_udf(wrapper, *columns)
class _AudioEncode(ScalarFunction):
"""Encode AUDIO_WAVEFORM to audio bytes."""
def __init__(self, format="wav"):
super().__init__()
self.format = normalize_audio_format(format)
def eval(self, audio_input):
if audio_input is None:
return None
return encode_audio_input(audio_input, format=self.format)
[docs]def audio_encode(*columns, format="wav", concurrency=None):
"""
Encode ``AUDIO_WAVEFORM`` to bytes.
Args:
*columns: Optional audio columns. Supported input is
``AUDIO_WAVEFORM`` only.
format: Encoded output format. Common values are ``"wav"``,
``"flac"``, ``"ogg"``, ``"aiff"``, and ``"mp3"``. Actual
availability depends on the worker's ``soundfile``/libsndfile
runtime. PyFlink writes a stable default subtype for each supported
format; if the current libsndfile runtime does not support that
subtype, the operator fails instead of silently choosing an
environment-dependent subtype.
concurrency: UDF concurrency. ``None`` uses the framework default.
Returns:
A UDF producing encoded ``BYTES``.
Raises:
ValueError: If the input is not ``AUDIO_WAVEFORM`` or the format is
unsupported, if the default output subtype is unavailable in the
current libsndfile runtime, or if waveform samples contain
``NaN``/``Inf``.
Examples::
>>> # Usage 1: create a reusable UDF and apply it to a column.
>>> encode = audio_encode(format="wav")
>>> df = df.with_column("wav", encode(col("waveform")))
>>>
>>> # Usage 2: pass the column directly when building the expression.
>>> df = df.with_column("wav", audio_encode(col("waveform"), format="wav"))
"""
wrapper = udf(
_AudioEncode(format=format),
return_dtype=DataType.binary(),
**_udf_runtime_kwargs(concurrency=concurrency),
)
return _build_or_apply_udf(wrapper, *columns)
class _AudioConvertFormat(ScalarFunction):
"""Convert BYTES or URI audio to encoded bytes in the requested format."""
def __init__(self,
format="wav",
on_error="raise",
max_decoded_bytes=DEFAULT_MAX_DECODED_AUDIO_BYTES):
super().__init__()
self.format = normalize_audio_format(format)
self.on_error = _normalize_audio_decode_on_error(on_error)
self.max_decoded_bytes = _validate_max_decoded_bytes(max_decoded_bytes)
self._bridge_client = None
def open(self, function_context):
self._bridge_client = _bridge_client_or_none(function_context)
def eval(self, audio_input):
if audio_input is None:
return None
ensure_supported_audio_input(
"audio_convert_format",
("BYTES", "URI"),
audio_input,
)
if self.on_error == "raise":
waveform = decode_audio_input(
audio_input,
bridge_client=self._bridge_client,
max_decoded_bytes=self.max_decoded_bytes,
)
return encode_audio_input(waveform, format=self.format)
try:
waveform = decode_audio_input(
audio_input,
bridge_client=self._bridge_client,
max_decoded_bytes=self.max_decoded_bytes,
)
except Exception as e:
if is_audio_data_error(e):
return None
raise
return encode_audio_input(waveform, format=self.format)
class _AudioStandardize(ScalarFunction):
"""Standardize AUDIO_WAVEFORM to a sample rate and channel count."""
def __init__(self, sample_rate=16000, channels=1):
super().__init__()
self.sample_rate = _validate_positive_int("sample_rate", sample_rate)
self.channels = _validate_positive_int("channels", channels)
def eval(self, audio_input):
if audio_input is None:
return None
return standardize_waveform(
audio_input,
sample_rate=self.sample_rate,
channels=self.channels,
)
[docs]def audio_standardize(*columns, sample_rate=16000, channels=1, concurrency=None):
"""
Standardize an audio waveform.
The operator resamples when needed and converts channel count when needed.
The default ``16000`` Hz mono target matches Whisper-style speech model
input contracts. Channel conversion downmixes multi-channel input by
averaging channels when the target is mono. Mono-to-multi conversion
duplicates the mono channel. Multi-to-different-multi conversion downmixes
to mono first, then duplicates to the requested channel count.
Args:
*columns: Optional audio columns. Supported input is
``AUDIO_WAVEFORM`` only; call ``audio_decode`` first for encoded
inputs.
sample_rate: Target sample rate in samples per second.
channels: Target channel count.
concurrency: UDF concurrency. ``None`` uses the framework default.
Returns:
A UDF producing ``ROW<data TENSOR('float32'), sample_rate INT, channels INT,
frames BIGINT, sample_format STRING, layout STRING, duration_ms BIGINT,
source_uri STRING, start_time_ms BIGINT, end_time_ms BIGINT>``.
The output keeps source and segment fields from the input when present.
Raises:
ValueError: If the input is not an audio waveform row, or if
``sample_rate`` or ``channels`` is invalid.
Examples::
>>> decode = audio_decode()
>>> waveform = decode(col("audio_bytes"))
>>> # Usage 1: create a reusable UDF and apply it to a column.
>>> standardize = audio_standardize(sample_rate=16000, channels=1)
>>> df = df.with_column(
... "speech_ready",
... standardize(waveform),
... )
>>>
>>> # Usage 2: pass the column directly when building the expression.
>>> df = df.with_column(
... "speech_ready",
... audio_standardize(waveform, sample_rate=16000, channels=1),
... )
"""
wrapper = udf(
_AudioStandardize(sample_rate=sample_rate, channels=channels),
return_dtype=AUDIO_WAVEFORM_TYPE,
**_udf_runtime_kwargs(concurrency=concurrency),
)
return _build_or_apply_udf(wrapper, *columns)
class _AudioResample(ScalarFunction):
"""Resample AUDIO_WAVEFORM to a target sample rate."""
def __init__(self, sample_rate):
super().__init__()
self.sample_rate = _validate_positive_int("sample_rate", sample_rate)
def eval(self, audio_input):
if audio_input is None:
return None
return resample_waveform(audio_input, self.sample_rate)
[docs]def audio_resample(*columns, sample_rate, concurrency=None):
"""
Resample an audio waveform.
This changes the sample rate while preserving the channel count. Use
``audio_standardize`` when both sample rate and channel layout should be
normalized.
Args:
*columns: Optional audio columns. Supported input is
``AUDIO_WAVEFORM`` only; call ``audio_decode`` first for encoded
inputs.
sample_rate: Target sample rate in samples per second.
concurrency: UDF concurrency. ``None`` uses the framework default.
Returns:
A UDF producing ``ROW<data TENSOR('float32'), sample_rate INT, channels INT,
frames BIGINT, sample_format STRING, layout STRING, duration_ms BIGINT,
source_uri STRING, start_time_ms BIGINT, end_time_ms BIGINT>``.
Raises:
ValueError: If the input is not an audio waveform row, or if
``sample_rate`` is invalid.
Examples::
>>> # Usage 1: create a reusable UDF and apply it to a column.
>>> resample = audio_resample(sample_rate=8000)
>>> df = df.with_column("resampled", resample(col("waveform")))
>>>
>>> # Usage 2: pass the column directly when building the expression.
>>> df = df.with_column(
... "resampled",
... audio_resample(col("waveform"), sample_rate=8000),
... )
"""
wrapper = udf(
_AudioResample(sample_rate=sample_rate),
return_dtype=AUDIO_WAVEFORM_TYPE,
**_udf_runtime_kwargs(concurrency=concurrency),
)
return _build_or_apply_udf(wrapper, *columns)
class _AudioSplitByDuration(TableFunction):
"""Split audio by fixed duration into one segment row per output."""
def __init__(self, segment_duration_ms=None, segment_type="audio", max_segments=1024):
super().__init__()
self.segment_duration_ms = _validate_optional_positive_int(
"segment_duration_ms",
segment_duration_ms,
)
self.segment_type = _normalize_segment_type(segment_type)
self.max_segments = _validate_positive_int("max_segments", max_segments)
self._bridge_client = None
def open(self, function_context):
self._bridge_client = _bridge_client_or_none(function_context)
def eval(self,
audio_input,
segment_duration_ms=None,
segment_type=None,
max_segments=None):
if audio_input is None:
return
segment_duration_ms = _validate_positive_int(
"segment_duration_ms",
self.segment_duration_ms
if segment_duration_ms is None
else segment_duration_ms,
)
segment_type = _normalize_segment_type(
self.segment_type if segment_type is None else segment_type
)
max_segments = _validate_positive_int(
"max_segments",
self.max_segments if max_segments is None else max_segments,
)
if segment_type == "audio":
segments = iter_split_waveform_by_duration(
audio_input,
segment_duration_ms,
max_segments=max_segments,
)
else:
segments = iter_audio_clip_ref_by_duration(
audio_input,
segment_duration_ms,
bridge_client=self._bridge_client,
max_segments=max_segments,
)
for segment in segments:
yield Row(segment)
[docs]def audio_split_by_duration(
*columns,
segment_duration_ms,
segment_type="audio",
max_segments=1024,
concurrency=None,
):
"""
Split audio by fixed duration into one output row per segment.
``segment_type="audio"`` materializes decoded waveform segments.
``segment_type="ref"`` returns lazy reference segments with ``uri``,
``start_time_ms``, and ``end_time_ms`` fields; downstream operators can
decode only the segments they need.
SQL module calls put ``segment_type`` before ``max_segments``. For example,
use ``audio_split_by_duration(uri, 30000, 'ref')`` for lazy references, or
``audio_split_by_duration(waveform, 30000, 'audio', 1024)`` when overriding
``max_segments``.
Args:
*columns: Optional audio columns. ``segment_type="audio"`` supports
waveform rows. ``segment_type="ref"`` supports URI ``STRING`` and
reference rows.
segment_duration_ms: Segment duration in milliseconds.
segment_type: ``"audio"`` for materialized waveform rows or ``"ref"``
for reference rows.
max_segments: Maximum allowed number of returned segments. Exceeding
the limit raises ``ValueError`` instead of silently truncating data.
concurrency: UDF concurrency. ``None`` uses the framework default.
Returns:
A UDTF producing one ``segment`` column per emitted segment. For
``segment_type="audio"``, ``segment`` is a waveform row with
``data TENSOR('float32'), sample_rate INT, channels INT, frames BIGINT,
sample_format STRING, layout STRING, duration_ms BIGINT,
source_uri STRING, start_time_ms BIGINT, end_time_ms BIGINT``.
For ``segment_type="ref"``, ``segment`` has
``uri STRING, start_time_ms BIGINT, end_time_ms BIGINT``.
Raises:
ValueError: If the input does not match ``segment_type``, if
``segment_duration_ms`` or ``max_segments`` is invalid, or if the
split would produce more than ``max_segments`` items.
Examples::
>>> # Usage 1: create a reusable UDTF and join it laterally.
>>> split = audio_split_by_duration(segment_duration_ms=30000)
>>> segments = df.join_lateral(split(col("waveform")).alias("segment"))
>>>
>>> # Usage 2: pass the column directly when joining laterally.
>>> segments = df.join_lateral(
... audio_split_by_duration(
... col("waveform"),
... segment_duration_ms=30000,
... ).alias("segment")
... )
>>>
>>> # Usage 1: create a reusable UDTF for lazy reference segments.
>>> ref_split = audio_split_by_duration(
... segment_duration_ms=30000,
... segment_type="ref",
... )
>>> segments = df.join_lateral(ref_split(col("uri")).alias("segment"))
>>>
>>> # Usage 2: pass the URI column directly for lazy reference segments.
>>> segments = df.join_lateral(
... audio_split_by_duration(
... col("uri"),
... segment_duration_ms=30000,
... segment_type="ref",
... ).alias("segment")
... )
"""
segment_type = _normalize_segment_type(segment_type)
wrapper = udtf(
_AudioSplitByDuration(
segment_duration_ms=segment_duration_ms,
segment_type=segment_type,
max_segments=max_segments,
),
result_types=(
[AUDIO_WAVEFORM_TABLE_TYPE]
if segment_type == "audio"
else [AUDIO_CLIP_REF_TABLE_TYPE]
),
**_udf_runtime_kwargs(concurrency=concurrency),
)
return _build_or_apply_udf(wrapper, *columns)
class _AudioSplitByTimestamp(TableFunction):
"""Split audio by explicit timestamp ranges."""
def __init__(self, timestamps, segment_type="audio"):
super().__init__()
self.timestamps = _normalize_timestamp_ranges(
timestamps,
"audio_split_by_timestamp",
)
self.segment_type = _normalize_segment_type(segment_type)
def eval(self, audio_input):
if audio_input is None:
return
if self.segment_type == "audio":
segments = iter_split_waveform_by_timestamps(audio_input, self.timestamps)
else:
segments = iter_audio_clip_ref_by_timestamps(audio_input, self.timestamps)
for segment in segments:
yield Row(segment)
class _AudioSplitByTimestampSql(TableFunction):
"""SQL module variant: timestamps are runtime data."""
def __init__(self, segment_type="audio"):
super().__init__()
self.segment_type = _normalize_segment_type(segment_type)
def eval(self, audio_input, timestamps):
if audio_input is None or timestamps is None:
return
if self.segment_type == "audio":
segments = iter_split_waveform_by_timestamps(audio_input, timestamps)
else:
segments = iter_audio_clip_ref_by_timestamps(audio_input, timestamps)
for segment in segments:
yield Row(segment)
[docs]def audio_split_by_timestamp(
*columns,
timestamps=None,
segment_type="audio",
concurrency=None,
):
"""
Split audio by explicit timestamp ranges into one output row per segment.
Timestamp ranges are interpreted as millisecond ``start_ms``/``end_ms``
pairs. They may be provided as a literal Python sequence via
``timestamps=[...]`` or as a second DataFrame column expression.
``segment_type="audio"`` materializes waveform slices.
``segment_type="ref"`` returns lazy reference slices without opening the
encoded audio object. For URI input, timestamp ranges are not clipped to the
actual audio duration; for bounded ``AUDIO_CLIP_REF`` input, ranges are clipped
only to the existing reference boundary. That boundary is trusted as caller
metadata and is not validated against the real audio duration. SQL module
calls use ``segment_type`` as the third argument, for example
``audio_split_by_timestamp(uri, ranges, 'ref')``.
Args:
*columns: Optional audio columns. ``segment_type="audio"`` supports
waveform rows. ``segment_type="ref"`` supports URI ``STRING`` and
reference rows.
timestamps: Optional literal sequence of ranges. Each range may be a
Row/object/dict with ``start_ms`` and ``end_ms`` fields, or a
two-item sequence. Leave this unset when passing ranges as the
second input column.
segment_type: ``"audio"`` for materialized waveform rows or ``"ref"``
for reference rows.
concurrency: UDF concurrency. ``None`` uses the framework default.
Returns:
A UDTF producing one ``segment`` column per emitted segment. For
``segment_type="audio"``, ``segment`` is a waveform row with
``data TENSOR('float32'), sample_rate INT, channels INT, frames BIGINT,
sample_format STRING, layout STRING, duration_ms BIGINT,
source_uri STRING, start_time_ms BIGINT, end_time_ms BIGINT``.
For ``segment_type="ref"``, ``segment`` has
``uri STRING, start_time_ms BIGINT, end_time_ms BIGINT``.
Raises:
ValueError: If the input does not match ``segment_type``, if a timestamp
range is invalid, or if ``segment_type`` is unsupported.
Examples::
>>> ranges = [{"start_ms": 1000, "end_ms": 3500}]
>>> # Usage 1: create a reusable UDTF and join it laterally.
>>> split = audio_split_by_timestamp(timestamps=ranges)
>>> segments = df.join_lateral(split(col("waveform")).alias("segment"))
>>>
>>> # Usage 2: pass the column directly when joining laterally.
>>> segments = df.join_lateral(
... audio_split_by_timestamp(
... col("waveform"),
... timestamps=ranges,
... ).alias("segment")
... )
>>>
>>> # Dynamic timestamp ranges can be read from a column.
>>> segments = df.join_lateral(
... audio_split_by_timestamp(
... col("waveform"),
... col("ranges"),
... ).alias("segment")
... )
"""
segment_type = _normalize_segment_type(segment_type)
if timestamps is not None and len(columns) >= 2:
raise ValueError(
"audio_split_by_timestamp timestamps must be either a literal "
"keyword argument or the second input column, not both"
)
if timestamps is None:
if columns and len(columns) != 2:
raise ValueError(
"audio_split_by_timestamp requires an audio input column and "
"a timestamps column when timestamps is not provided as a "
"literal keyword argument"
)
wrapper = udtf(
_AudioSplitByTimestampSql(segment_type=segment_type),
result_types=(
[AUDIO_WAVEFORM_TABLE_TYPE]
if segment_type == "audio"
else [AUDIO_CLIP_REF_TABLE_TYPE]
),
**_udf_runtime_kwargs(concurrency=concurrency),
)
return _build_or_apply_udf(wrapper, *columns)
wrapper = udtf(
_AudioSplitByTimestamp(timestamps=timestamps, segment_type=segment_type),
result_types=(
[AUDIO_WAVEFORM_TABLE_TYPE]
if segment_type == "audio"
else [AUDIO_CLIP_REF_TABLE_TYPE]
),
**_udf_runtime_kwargs(concurrency=concurrency),
)
return _build_or_apply_udf(wrapper, *columns)
class _AudioSplitBySpeech(TableFunction):
"""Split audio by speech activity ranges."""
def __init__(self,
segment_type="audio",
pre_padding_ms=0,
post_padding_ms=0,
merge_gap_ms=0,
min_segment_ms=0,
max_segment_ms=None,
max_segments=1024):
super().__init__()
self.segment_type = _normalize_segment_type(segment_type)
self.pre_padding_ms = _validate_non_negative_int(
"audio_split_by_speech.pre_padding_ms",
pre_padding_ms,
)
self.post_padding_ms = _validate_non_negative_int(
"audio_split_by_speech.post_padding_ms",
post_padding_ms,
)
self.merge_gap_ms = _validate_non_negative_int(
"audio_split_by_speech.merge_gap_ms",
merge_gap_ms,
)
self.min_segment_ms = _validate_non_negative_int(
"audio_split_by_speech.min_segment_ms",
min_segment_ms,
)
self.max_segment_ms = _validate_optional_positive_int(
"audio_split_by_speech.max_segment_ms",
max_segment_ms,
)
self.max_segments = _validate_positive_int(
"audio_split_by_speech.max_segments",
max_segments,
)
self._bridge_client = None
def open(self, function_context):
self._bridge_client = _bridge_client_or_none(function_context)
def eval(self,
audio_input,
speech_activity,
segment_type=None,
pre_padding_ms=None,
post_padding_ms=None,
merge_gap_ms=None,
min_segment_ms=None,
max_segment_ms=None,
max_segments=None):
if audio_input is None or speech_activity is None:
return
segment_type = _normalize_segment_type(
self.segment_type if segment_type is None else segment_type
)
kwargs = {
"pre_padding_ms": _validate_non_negative_int(
"audio_split_by_speech.pre_padding_ms",
self.pre_padding_ms if pre_padding_ms is None else pre_padding_ms,
),
"post_padding_ms": _validate_non_negative_int(
"audio_split_by_speech.post_padding_ms",
self.post_padding_ms if post_padding_ms is None else post_padding_ms,
),
"merge_gap_ms": _validate_non_negative_int(
"audio_split_by_speech.merge_gap_ms",
self.merge_gap_ms if merge_gap_ms is None else merge_gap_ms,
),
"min_segment_ms": _validate_non_negative_int(
"audio_split_by_speech.min_segment_ms",
self.min_segment_ms if min_segment_ms is None else min_segment_ms,
),
"max_segment_ms": _validate_optional_positive_int(
"audio_split_by_speech.max_segment_ms",
self.max_segment_ms if max_segment_ms is None else max_segment_ms,
),
"max_segments": _validate_positive_int(
"audio_split_by_speech.max_segments",
self.max_segments if max_segments is None else max_segments,
),
}
if segment_type == "audio":
segments = iter_split_waveform_by_speech(
audio_input,
speech_activity,
**kwargs,
)
else:
segments = iter_audio_clip_ref_by_speech(
audio_input,
speech_activity,
bridge_client=self._bridge_client,
**kwargs,
)
for segment in segments:
yield Row(segment)
[docs]def audio_split_by_speech(
*columns,
segment_type="audio",
pre_padding_ms=0,
post_padding_ms=0,
merge_gap_ms=0,
min_segment_ms=0,
max_segment_ms=None,
max_segments=1024,
concurrency=None,
):
"""
Split audio by speech activity ranges into one output row per segment.
The second input column must contain speech ranges with ``start_ms`` and
``end_ms`` fields, such as output from a VAD or upstream speech activity
model. This operator does not detect speech by itself; it only applies
padding, merging, filtering, and splitting to caller-provided ranges.
Ranges outside the actual audio duration are clipped, and fully
out-of-range ranges are dropped. SQL module calls put ``segment_type`` before
the optional tuning arguments. Use ``audio_split_by_speech(uri, ranges,
'ref')`` for default lazy references, or provide all tuning literals as
``audio_split_by_speech(uri, ranges, 'ref', 200, 200, 30, 30, 1000, 1024)``.
Args:
*columns: Audio input followed by speech activity ranges.
``segment_type="audio"`` supports waveform rows.
``segment_type="ref"`` supports URI ``STRING`` and reference rows.
segment_type: ``"audio"`` for materialized waveform rows or ``"ref"``
for reference rows.
pre_padding_ms: Milliseconds to extend before each speech range.
post_padding_ms: Milliseconds to extend after each speech range.
merge_gap_ms: Merge adjacent ranges separated by this gap or less.
min_segment_ms: Drop segments shorter than this duration after padding
and merging.
max_segment_ms: Optionally split long speech ranges into chunks no
longer than this value.
max_segments: Maximum allowed number of returned segments.
concurrency: UDF concurrency. ``None`` uses the framework default.
Returns:
A UDTF producing one ``segment`` column per emitted segment. For
``segment_type="audio"``, ``segment`` is a waveform row with
``data TENSOR('float32'), sample_rate INT, channels INT, frames BIGINT,
sample_format STRING, layout STRING, duration_ms BIGINT,
source_uri STRING, start_time_ms BIGINT, end_time_ms BIGINT``.
For ``segment_type="ref"``, ``segment`` has
``uri STRING, start_time_ms BIGINT, end_time_ms BIGINT``.
Raises:
ValueError: If the input does not match ``segment_type``, if speech
ranges or split parameters are invalid, or if the split would produce
more than ``max_segments`` items.
Examples::
>>> # Usage 1: create a reusable UDTF and join it laterally.
>>> split = audio_split_by_speech(
... pre_padding_ms=200,
... post_padding_ms=200,
... )
>>> segments = df.join_lateral(
... split(col("waveform"), col("speech_ranges")).alias("segment")
... )
>>>
>>> # Usage 2: pass columns directly when joining laterally.
>>> segments = df.join_lateral(
... audio_split_by_speech(
... col("waveform"),
... col("speech_ranges"),
... pre_padding_ms=200,
... post_padding_ms=200,
... ).alias("segment")
... )
"""
segment_type = _normalize_segment_type(segment_type)
wrapper = udtf(
_AudioSplitBySpeech(
segment_type=segment_type,
pre_padding_ms=pre_padding_ms,
post_padding_ms=post_padding_ms,
merge_gap_ms=merge_gap_ms,
min_segment_ms=min_segment_ms,
max_segment_ms=max_segment_ms,
max_segments=max_segments,
),
result_types=(
[AUDIO_WAVEFORM_TABLE_TYPE]
if segment_type == "audio"
else [AUDIO_CLIP_REF_TABLE_TYPE]
),
**_udf_runtime_kwargs(concurrency=concurrency),
)
return _build_or_apply_udf(wrapper, *columns)
class _AudioConcat(ScalarFunction):
"""Concatenate ARRAY<AUDIO_WAVEFORM> into a single AUDIO_WAVEFORM."""
def eval(self, audio_inputs):
if audio_inputs is None:
return None
return concat_waveforms(audio_inputs)
[docs]def audio_concat(*columns, concurrency=None):
"""
Concatenate waveforms.
All input waveforms must have compatible sample rate, channel count, and
sample format. Use ``audio_standardize`` on individual segments before
concatenation when the source layout may differ.
Args:
*columns: Optional audio columns. Supported input is
``ARRAY<AUDIO_WAVEFORM>`` only.
concurrency: UDF concurrency. ``None`` uses the framework default.
Returns:
A UDF producing one ``ROW<data TENSOR('float32'), sample_rate INT,
channels INT, frames BIGINT, sample_format STRING, layout STRING,
duration_ms BIGINT, source_uri STRING, start_time_ms BIGINT,
end_time_ms BIGINT>``.
Raises:
ValueError: If the input is not an array of waveform rows, or if the
waveforms have incompatible sample rate, channel count, or sample
format.
Examples::
>>> # Usage 1: create a reusable UDF and apply it to a column.
>>> concat = audio_concat()
>>> df = df.with_column("joined", concat(col("clips")))
>>>
>>> # Usage 2: pass the column directly when building the expression.
>>> df = df.with_column("joined", audio_concat(col("clips")))
"""
wrapper = udf(
_AudioConcat(),
return_dtype=AUDIO_WAVEFORM_TYPE,
**_udf_runtime_kwargs(concurrency=concurrency),
)
return _build_or_apply_udf(wrapper, *columns)