This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/paimon-rust.git
The following commit(s) were added to refs/heads/main by this push:
new c056ef0 Add video frame SQL UDF (#465)
c056ef0 is described below
commit c056ef07d6722c7b01b14fc696eb4cfd44b471c3
Author: Jingsong Lee <[email protected]>
AuthorDate: Tue Jul 7 13:08:37 2026 +0800
Add video frame SQL UDF (#465)
---
bindings/python/python/pypaimon_rust/functions.py | 88 ++++++++++++++++++++---
bindings/python/src/context.rs | 44 ++++++++----
bindings/python/tests/test_datafusion.py | 63 ++++++++++++++--
3 files changed, 170 insertions(+), 25 deletions(-)
diff --git a/bindings/python/python/pypaimon_rust/functions.py
b/bindings/python/python/pypaimon_rust/functions.py
index 9a3b4cf..2460a36 100644
--- a/bindings/python/python/pypaimon_rust/functions.py
+++ b/bindings/python/python/pypaimon_rust/functions.py
@@ -117,15 +117,44 @@ def _decode_video_snapshot(
break
if candidate is not None:
- try:
- image = candidate.to_image()
- except ImportError as e:
- raise ImportError(
- "Pillow is required to encode video_snapshot images"
- ) from e
- output = io.BytesIO()
- image.save(output, format=image_format)
- return output.getvalue()
+ return _encode_video_frame(candidate, image_format)
+ return None
+
+
+def _encode_video_frame(frame, image_format: str) -> bytes:
+ try:
+ image = frame.to_image()
+ except ImportError as e:
+ raise ImportError("Pillow is required to encode video frame images")
from e
+ output = io.BytesIO()
+ image.save(output, format=image_format)
+ return output.getvalue()
+
+
+def _decode_video_frame(
+ stream: BinaryIO,
+ image_format: str,
+ frame_index: int,
+) -> bytes | None:
+ try:
+ import av
+ except ImportError as e:
+ raise ImportError("PyAV is required to decode video frames") from e
+
+ with av.open(stream, mode="r") as container:
+ format_names = set((container.format.name or "").split(","))
+ if format_names & _STILL_IMAGE_FORMATS:
+ logger.debug(
+ "video_frame input is a still image format: %s",
+ container.format.name,
+ )
+ return None
+ if not container.streams.video:
+ return None
+
+ for index, frame in enumerate(container.decode(video=0)):
+ if index == frame_index:
+ return _encode_video_frame(frame, image_format)
return None
@@ -177,3 +206,44 @@ def _make_video_snapshot(image_format: str = "PNG",
blob_reader_registry=None):
return pa.array(frames, type=pa.binary())
return video_snapshot
+
+
+def _make_video_frame(image_format: str = "PNG", blob_reader_registry=None):
+ image_format = image_format.upper()
+
+ def video_frame(values, frame_indices):
+ try:
+ import pyarrow as pa
+ except ImportError as e:
+ raise ImportError("pyarrow is required to return video_frame
results") from e
+
+ frames = []
+ raw_values = values.to_pylist()
+ frame_index_values = frame_indices.to_pylist()
+ if len(frame_index_values) != len(raw_values):
+ raise ValueError("video_frame index argument must have the same
row count")
+
+ for raw_value, frame_index in zip(raw_values, frame_index_values):
+ if raw_value is None or frame_index is None:
+ frames.append(None)
+ continue
+
+ try:
+ frame_index = int(frame_index)
+ if frame_index < 0:
+ frames.append(None)
+ continue
+ stream = open_blob_descriptor_stream(raw_value,
blob_reader_registry)
+ try:
+ frames.append(_decode_video_frame(stream, image_format,
frame_index))
+ finally:
+ stream.close()
+ except ImportError:
+ raise
+ except Exception as e:
+ logger.warning("Failed to decode video frame: %s", e)
+ frames.append(None)
+
+ return pa.array(frames, type=pa.binary())
+
+ return video_frame
diff --git a/bindings/python/src/context.rs b/bindings/python/src/context.rs
index 43b1fdb..f6e2a6e 100644
--- a/bindings/python/src/context.rs
+++ b/bindings/python/src/context.rs
@@ -129,17 +129,18 @@ pub struct PySQLContext {
}
impl PySQLContext {
- fn register_video_snapshot_builtin(&self, py: Python<'_>) -> PyResult<()> {
+ fn register_video_builtins(&self, py: Python<'_>) -> PyResult<()> {
let functions = py.import("pypaimon_rust.functions")?;
let blob_reader_registry = Py::new(
py,
PyBlobReaderRegistry::new(self.inner.blob_reader_registry()),
)?;
- let func = functions
+ let snapshot_blob_reader_registry = blob_reader_registry.clone_ref(py);
+ let snapshot_func = functions
.getattr("_make_video_snapshot")?
- .call1(("PNG", blob_reader_registry))?
+ .call1(("PNG", snapshot_blob_reader_registry))?
.unbind();
- let signature = Signature::one_of(
+ let snapshot_signature = Signature::one_of(
vec![
TypeSignature::Exact(vec![ArrowDataType::Binary]),
TypeSignature::Exact(vec![ArrowDataType::Binary,
ArrowDataType::Int32]),
@@ -147,23 +148,42 @@ impl PySQLContext {
],
Volatility::Volatile,
);
- let udf = build_python_scalar_udf(
+ let snapshot_udf = build_python_scalar_udf(
"video_snapshot".to_string(),
- func,
+ snapshot_func,
+ ArrowDataType::Binary,
+ snapshot_signature,
+ );
+ self.inner.ctx().register_udf(snapshot_udf);
+
+ let frame_func = functions
+ .getattr("_make_video_frame")?
+ .call1(("PNG", blob_reader_registry))?
+ .unbind();
+ let frame_signature = Signature::one_of(
+ vec![
+ TypeSignature::Exact(vec![ArrowDataType::Binary,
ArrowDataType::Int32]),
+ TypeSignature::Exact(vec![ArrowDataType::Binary,
ArrowDataType::Int64]),
+ ],
+ Volatility::Volatile,
+ );
+ let frame_udf = build_python_scalar_udf(
+ "video_frame".to_string(),
+ frame_func,
ArrowDataType::Binary,
- signature,
+ frame_signature,
);
- self.inner.ctx().register_udf(udf);
+ self.inner.ctx().register_udf(frame_udf);
Ok(())
}
- fn warn_video_snapshot_registration_failure(py: Python<'_>, err: PyErr) {
+ fn warn_video_builtin_registration_failure(py: Python<'_>, err: PyErr) {
if let Ok(warnings) = py.import("warnings") {
let category = py.get_type::<PyRuntimeWarning>();
let _ = warnings.call_method1(
"warn",
(
- format!("video_snapshot built-in could not be registered:
{err}"),
+ format!("video built-ins could not be registered: {err}"),
category,
),
);
@@ -178,8 +198,8 @@ impl PySQLContext {
let ctx = Self {
inner: SQLContext::new(),
};
- if let Err(err) = ctx.register_video_snapshot_builtin(py) {
- Self::warn_video_snapshot_registration_failure(py, err);
+ if let Err(err) = ctx.register_video_builtins(py) {
+ Self::warn_video_builtin_registration_failure(py, err);
}
Ok(ctx)
}
diff --git a/bindings/python/tests/test_datafusion.py
b/bindings/python/tests/test_datafusion.py
index 236b130..5430b2a 100644
--- a/bindings/python/tests/test_datafusion.py
+++ b/bindings/python/tests/test_datafusion.py
@@ -82,13 +82,20 @@ def extract_rows(batches):
def test_video_snapshot_builtin_registered_on_context_init():
ctx = SQLContext()
- batches = ctx.sql("SELECT video_snapshot(CAST(NULL AS BYTEA)) AS
cover_png")
+ batches = ctx.sql(
+ """
+ SELECT
+ video_snapshot(CAST(NULL AS BYTEA)) AS cover_png,
+ video_frame(CAST(NULL AS BYTEA), 0) AS frame_png
+ """
+ )
table = pa.Table.from_batches(batches)
assert table["cover_png"].to_pylist() == [None]
+ assert table["frame_png"].to_pylist() == [None]
-def test_sql_context_survives_video_snapshot_registration_failure(monkeypatch):
+def test_sql_context_survives_video_builtins_registration_failure(monkeypatch):
monkeypatch.setitem(
sys.modules,
"pypaimon_rust.functions",
@@ -97,7 +104,7 @@ def
test_sql_context_survives_video_snapshot_registration_failure(monkeypatch):
with pytest.warns(
RuntimeWarning,
- match="video_snapshot built-in could not be registered",
+ match="video built-ins could not be registered",
):
ctx = SQLContext()
@@ -280,6 +287,55 @@ def test_video_snapshot_accepts_timestamp_ms():
ctx.sql("DROP TEMPORARY TABLE paimon.default.videos")
+def test_video_frame_accepts_frame_index():
+ image_module = pytest.importorskip("PIL.Image")
+
+ with tempfile.TemporaryDirectory() as warehouse:
+ video_path = Path(warehouse) / "sample.mp4"
+ write_sample_video(
+ video_path,
+ colors=((240, 40, 80), (40, 220, 80), (40, 80, 240)),
+ )
+ video_bytes = video_path.read_bytes()
+
+ ctx = SQLContext()
+ ctx.register_catalog("paimon", {"warehouse": warehouse})
+ ctx.register_batch(
+ "paimon.default.videos",
+ pa.record_batch(
+ [[1], pa.array([video_bytes], type=pa.binary())],
+ names=["id", "video"],
+ ),
+ )
+
+ batches = ctx.sql(
+ """
+ SELECT
+ video_frame(video, 0) AS first_png,
+ video_frame(video, CAST(1 AS BIGINT)) AS second_png,
+ video_frame(video, CAST(2 AS INT)) AS third_png,
+ video_frame(video, 3) AS beyond_duration_png,
+ video_frame(video, -1) AS negative_png
+ FROM paimon.default.videos
+ """
+ )
+ row = pa.Table.from_batches(batches).to_pylist()[0]
+
+ assert row["first_png"].startswith(PNG_SIGNATURE)
+ assert row["second_png"].startswith(PNG_SIGNATURE)
+ assert row["third_png"].startswith(PNG_SIGNATURE)
+ assert row["beyond_duration_png"] is None
+ assert row["negative_png"] is None
+
+ first_image =
image_module.open(io.BytesIO(row["first_png"])).convert("RGB")
+ second_image =
image_module.open(io.BytesIO(row["second_png"])).convert("RGB")
+ third_image =
image_module.open(io.BytesIO(row["third_png"])).convert("RGB")
+ assert first_image.getpixel((16, 16)) != second_image.getpixel((16,
16))
+ assert second_image.getpixel((16, 16)) != third_image.getpixel((16,
16))
+
+ ctx.sql("DROP TEMPORARY TABLE paimon.default.videos")
+
+
def test_query_simple_table_via_catalog_provider():
catalog = PaimonCatalog({"warehouse": WAREHOUSE})
ctx = SessionContext()
@@ -684,4 +740,3 @@ def test_list_databases_and_tables():
assert "name" in field_names
# simple_log_table is non-partitioned, so partition keys are empty.
assert schema.partition_keys() == []
-
