blob: 6a520f7083d7d2a1c04eff4908f7510e70c274ae [file]
# 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 io
import json
import os
import shutil
import tempfile
import unittest
from dataclasses import replace
from datetime import timedelta
from unittest.mock import patch
import pyarrow as pa
import pytest
import pypaimon.multimodal as pmm
from pypaimon.execution.python.read.table_read import PythonTableRead
from pypaimon.multimodal import source_col
from pypaimon.common.json_util import JSON
from pypaimon.common.predicate_builder import PredicateBuilder
from pypaimon import Schema as PaimonSchema
from pypaimon.globalindex.global_index_result import GlobalIndexResult
from pypaimon.utils.range import Range
try:
import ray
except ImportError:
ray = None
_PARQUET_OPTIONS = {
"row-tracking.enabled": "true",
"data-evolution.enabled": "true",
"deletion-vectors.enabled": "true",
"file.format": "parquet",
"vector.file.format": "parquet",
}
def _schema(fields):
return pa.schema([
pa.field(name, field_type)
for name, field_type in fields.items()
])
def _vector(dim):
return pa.list_(pa.float32(), dim)
def _raw_schema(options=None, primary_keys=None):
return PaimonSchema.from_pyarrow_schema(
_schema({"id": pa.int32(), "name": pa.string()}),
primary_keys=list(primary_keys or []),
options=dict(options or {"file.format": "parquet"}),
)
class MultimodalTableTest(unittest.TestCase):
def setUp(self):
self.temp_dir = tempfile.mkdtemp(prefix="pypaimon_mm_")
self.warehouse = os.path.join(self.temp_dir, "warehouse")
self.conn = pmm.connect(options={"warehouse": self.warehouse})
def tearDown(self):
shutil.rmtree(self.temp_dir, ignore_errors=True)
def test_connect_accepts_options(self):
conn = pmm.connect(
database="analytics",
options={
"warehouse": os.path.join(self.temp_dir, "warehouse_options"),
},
)
table = conn.create_table(
"docs",
schema=_schema({"id": pa.int32()}),
options=_PARQUET_OPTIONS,
)
self.assertEqual("analytics.docs", table.identifier)
def test_connect_rejects_positional_warehouse(self):
with self.assertRaises(TypeError):
pmm.connect(self.warehouse)
def test_copy_applies_dynamic_table_options(self):
table = self.conn.create_table(
"frames", schema=_schema({"id": pa.int32()}),
options=_PARQUET_OPTIONS)
copied = table.copy({"read.video.max-open-decoders": "8"})
self.assertEqual(table.identifier, copied.identifier)
self.assertIs(table.catalog, copied.catalog)
self.assertNotIn(
"read.video.max-open-decoders",
table.raw_table.table_schema.options)
self.assertEqual(
"8", copied.raw_table.table_schema.options[
"read.video.max-open-decoders"])
table.add([{"id": 1}])
snapshot_id = table.raw_table.snapshot_manager().get_latest_snapshot().id
table.raw_table.create_tag("training", snapshot_id=snapshot_id)
tagged = copied.raw_table.copy({"scan.tag-name": "training"})
self.assertEqual(
"8", tagged.table_schema.options[
"read.video.max-open-decoders"])
def test_create_table_defaults_data_evolution_options(self):
table = self.conn.create_table(
"docs",
schema=_schema({
"id": pa.int32(),
"content": pa.string(),
"embedding": _vector(3),
"payload": pa.large_binary(),
}),
)
options = table.raw_table.table_schema.options
self.assertEqual("true", options["row-tracking.enabled"])
self.assertEqual("true", options["data-evolution.enabled"])
self.assertNotIn("deletion-vectors.enabled", options)
self.assertEqual("true", options["blob-as-descriptor"])
self.assertNotIn("data-evolution.row-sidecar.enabled", options)
self.assertEqual("full", options["global-index.search-mode"])
self.assertNotIn("vector.file.format", options)
self.assertEqual("default.docs", self.conn.get_table("docs").identifier)
self.assertEqual(["id", "content", "embedding", "payload"],
[field.name for field in table.raw_table.fields])
def test_add_video_column_to_existing_images(self):
from pypaimon.schema.data_types import AtomicType
from pypaimon.schema.schema_change import SchemaChange
table = self.conn.create_table(
"evolved_video", schema=_schema({"id": pa.int32(), "image": pa.large_binary()}),
options=_PARQUET_OPTIONS)
table.add([{"id": 0, "image": b"image-0"}, {"id": 1, "image": b"image-1"}])
table.raw_table.create_tag("before_video")
old_table = table
self.conn.catalog.alter_table("default.evolved_video", [
SchemaChange.add_column("video", AtomicType("BYTES"), "__VIDEO_FRAME_FIELD"),
])
table = self.conn.get_table("evolved_video")
self.assertEqual([None, None], [r["video"] for r in table.scan().to_list()])
old_table.add([{"id": 2, "image": b"image-2"}])
video_path = os.path.join(self.temp_dir, "backfill.mp4")
with open(video_path, "wb") as out:
out.write(b"complete-video-payload")
table.replace_video("id < 2", pmm.Blob.from_local(video_path), first_frame=7)
# A retained writer must not stamp its old schema onto a new snapshot.
old_table.add([{"id": 3, "image": b"image-3"}])
snapshot = table.raw_table.snapshot_manager().get_latest_snapshot()
self.assertEqual(table.raw_table.table_schema.id, snapshot.schema_id)
table.raw_table.create_tag("after_old_append")
rows = sorted(table.scan().select(["id", "image", "video", "_ROW_ID"]).to_list(),
key=lambda row: row["id"])
self.assertEqual([0, 1, 2, 3], [row["_ROW_ID"] for row in rows])
self.assertEqual([7, 8], [pmm.VideoFrameDescriptor.deserialize(row["video"]).frame_index
for row in rows[:2]])
self.assertTrue(all(row["video"] is None for row in rows[2:]))
_, images = table.scan().read_blobs("image")
self.assertEqual([b"image-0", b"image-1", b"image-2", b"image-3"], images["image"])
_, videos = table.scan().where("id < 2").read_blobs("video")
self.assertEqual([b"complete-video-payload"] * 2, videos["video"])
for options in ({"scan.snapshot-id": str(snapshot.id)}, {"scan.tag-name": "after_old_append"}):
with self.subTest(options=options):
historical = table.raw_table.copy(options)
self.assertEqual(["id", "image", "video"], [field.name for field in historical.fields])
read = historical.new_read_builder()
restored = sorted(read.new_read().to_arrow(read.new_scan().plan().splits()).to_pylist(),
key=lambda row: row["id"])
self.assertEqual([0, 1, 2, 3], [row["id"] for row in restored])
self.assertEqual([row["video"] for row in rows], [row["video"] for row in restored])
historical = self.conn.catalog.get_table("default.evolved_video").copy({"scan.tag-name": "before_video"})
self.assertNotIn("video-frame-field", historical.table_schema.options)
# The retained old schema still reads the original rows after the backfill.
read = historical.new_read_builder()
historical_rows = read.new_read().to_arrow(read.new_scan().plan().splits()).to_pylist()
self.assertEqual([0, 1], sorted(row["id"] for row in historical_rows))
self.assertTrue(all(row.get("video") is None for row in historical_rows))
self.conn.catalog.alter_table("default.evolved_video", [
SchemaChange.add_column("second", AtomicType("BYTES"), "__VIDEO_FRAME_FIELD")])
for column, remaining in (("video", "second"), ("second", None)):
self.conn.catalog.alter_table("default.evolved_video", [SchemaChange.drop_column(column)])
current = self.conn.get_table("evolved_video")
self.assertEqual(remaining, current.raw_table.table_schema.options.get("video-frame-field"))
self.assertEqual(4, len(current.scan().to_list()))
self.conn.catalog.alter_table("default.evolved_video", [
SchemaChange.add_column("video", AtomicType("BYTES"), "__VIDEO_FRAME_FIELD")])
current = self.conn.get_table("evolved_video")
self.assertEqual([None] * 4, [row["video"] for row in current.scan().to_list()])
def test_add_video_column_validation(self):
from pypaimon.schema.data_types import AtomicType
from pypaimon.schema.schema_change import SchemaChange as C
table = self.conn.create_table(
"video_evolution_validation",
schema=_schema({"id": pa.int32(), "image": pa.large_binary()}),
options=_PARQUET_OPTIONS)
table.add([{"id": 0, "image": b"original"}])
identifier = "default.video_evolution_validation"
with self.assertRaisesRegex(ValueError, "immutable option"):
table.raw_table.copy({"video-frame-field": "image"})
def alter(changes):
self.conn.catalog.alter_table(identifier, changes)
for changes in [
[C.set_option("video-frame-field", "image")],
[C.set_option("video-frame-field", "missing")],
[C.add_column("bad", AtomicType("INT"), "__VIDEO_FRAME_FIELD")],
[C.add_column("bad", AtomicType("BLOB", False), "__VIDEO_FRAME_FIELD")],
[C.add_column("bad", AtomicType("BLOB")), C.set_option("video-frame-field", "image,bad")],
]:
with self.subTest(changes=changes):
with self.assertRaises(RuntimeError):
alter(changes)
self.assertEqual(0, self.conn.catalog.get_table(identifier).table_schema.id)
alter([C.update_column_comment("image", "__VIDEO_FRAME_FIELD")])
self.assertNotIn("video-frame-field", self.conn.catalog.get_table(identifier).table_schema.options)
alter([C.add_column("video", AtomicType("BYTES"), "__VIDEO_FRAME_FIELD")])
alter([C.add_column("second", AtomicType("BLOB"), "__VIDEO_FRAME_FIELD")])
self.assertEqual(
"video,second", self.conn.catalog.get_table(identifier).table_schema.options["video-frame-field"])
for change in [C.set_option("video-frame-field", "second"),
C.set_option("video-frame-field", ""), C.remove_option("video-frame-field")]:
with self.assertRaises(RuntimeError):
alter([change])
# A column added by a previous commit may already contain ordinary BLOB data.
alter([C.add_column("later", AtomicType("BLOB"))])
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.multimodal.mp4 import ContainerFrameIndex
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"))
_, bodies = table.scan().read_blobs("video")
payload = bodies["video"][0]
index = ContainerFrameIndex(io.BytesIO(payload), len(payload))
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",
schema=_schema({
"episode_id": pa.int64(),
"video": pa.large_binary(),
}),
options=dict(_PARQUET_OPTIONS, **{
"video-frame-field": "video",
"blob-as-descriptor": "true",
}),
)
video_path = os.path.join(self.temp_dir, "episode-42.mp4")
video_bytes = b"fake-mp4-payload"
with open(video_path, "wb") as output:
output.write(video_bytes)
video = pmm.Blob.from_local(video_path)
table.add_video(
video,
[{"episode_id": 42} for _ in range(3)],
first_frame=7,
)
rows = table.scan().select(["episode_id", "video"]).to_list()
descriptors = [
pmm.VideoFrameDescriptor.deserialize(row["video"])
for row in rows
]
self.assertEqual([7, 8, 9], [value.frame_index for value in descriptors])
self.assertTrue(all(
value.payload_descriptor == descriptors[0].payload_descriptor
for value in descriptors
))
self.assertTrue(descriptors[0].uri.endswith(".video"))
self.assertEqual(len(video_bytes), descriptors[0].length)
_, bodies = table.scan().read_blobs("video", parallelism=2)
self.assertEqual([video_bytes] * 3, bodies["video"])
def test_add_images_as_video_encodes_and_stores_frame_rows(self):
table = self.conn.create_table(
"image_video_frames",
schema=_schema({
"episode_id": pa.int64(),
"video": pa.large_binary(),
}),
options=dict(_PARQUET_OPTIONS, **{
"video-frame-field": "video",
"blob-as-descriptor": "true",
}),
)
calls = []
def encode(images, output_path, **options):
images = list(images)
calls.append((images, options))
with open(output_path, "wb") as output:
output.write(b"encoded-video")
return len(images)
with patch(
"pypaimon.multimodal.video._encode_images_to_video",
encode):
table.add_images_as_video(
[b"frame-0", b"frame-1"],
[{"episode_id": 42}, {"episode_id": 42}],
fps=30,
codec="mpeg4",
pixel_format="yuv420p",
gop_size=2,
codec_options={"qscale": "3"},
)
self.assertEqual(
[
(
[b"frame-0", b"frame-1"],
{
"fps": 30,
"codec": "mpeg4",
"pixel_format": "yuv420p",
"gop_size": 2,
"codec_options": {"qscale": "3"},
},
)
],
calls,
)
rows = table.scan().select(["episode_id", "video"]).to_list()
descriptors = [
pmm.VideoFrameDescriptor.deserialize(row["video"])
for row in rows
]
self.assertEqual([0, 1], [value.frame_index for value in descriptors])
_, bodies = table.scan().read_blobs("video")
self.assertEqual([b"encoded-video"] * 2, bodies["video"])
def test_add_images_as_video_rejects_frame_count_mismatch(self):
table = self.conn.create_table(
"mismatched_image_video_frames",
schema=_schema({
"episode_id": pa.int64(),
"video": pa.large_binary(),
}),
options=dict(_PARQUET_OPTIONS, **{
"video-frame-field": "video",
"blob-as-descriptor": "true",
}),
)
output_paths = []
def encode(unused_images, output_path, **unused_options):
output_paths.append(output_path)
with open(output_path, "wb") as output:
output.write(b"one-frame-video")
return 1
with patch(
"pypaimon.multimodal.video._encode_images_to_video",
encode):
with self.assertRaisesRegex(
ValueError,
"Image count 1 does not match frame row count 2"):
table.add_images_as_video(
[b"frame-0"],
[{"episode_id": 42}, {"episode_id": 42}],
fps=30,
codec="mpeg4",
pixel_format="yuv420p",
gop_size=2,
)
self.assertFalse(os.path.exists(output_paths[0]))
self.assertIsNone(
table.raw_table.snapshot_manager().get_latest_snapshot()
)
def test_add_videos_packs_multiple_videos_in_one_commit(self):
from pypaimon.table.row.blob import Blob, VideoFrameDescriptor
table = self.conn.create_table(
"batched_video_frames",
schema=_schema({
"episode_id": pa.int64(),
"video": pa.large_binary(),
}),
options=dict(_PARQUET_OPTIONS, **{
"video-frame-field": "video",
"blob-as-descriptor": "true",
}),
)
paths = [os.path.join(self.temp_dir, "episode-%s.mp4" % index)
for index in (1, 2)]
payloads = [b"first-video", b"second-video"]
for path, payload in zip(paths, payloads):
with open(path, "wb") as output:
output.write(payload)
table.add_videos([
(Blob.from_local(paths[0]), [{"episode_id": 1}] * 2, 3),
(Blob.from_local(paths[1]), [{"episode_id": 2}] * 3, 10),
])
snapshot = table.raw_table.snapshot_manager().get_latest_snapshot()
self.assertEqual(1, snapshot.id)
rows = table.scan().select(["episode_id", "video"]).to_list()
rows.sort(key=lambda row: (row["episode_id"], row["video"]))
descriptors = [
VideoFrameDescriptor.deserialize(row["video"])
for row in rows
]
self.assertEqual([3, 4, 10, 11, 12], [d.frame_index for d in descriptors])
self.assertEqual(1, len({d.uri for d in descriptors}))
self.assertTrue(descriptors[0].uri.endswith(".video"))
self.assertEqual(2, len({d.payload_descriptor for d in descriptors}))
def test_add_video_indexes_path_and_local_blob_with_unknown_length(self):
av = pytest.importorskip("av")
np = pytest.importorskip("numpy")
from fractions import Fraction
from pypaimon.table.row.blob import Blob, VideoFrameDescriptor
from pypaimon.multimodal.mp4 import ContainerFrameIndex
table = self.conn.create_table(
"local_video_indexes",
schema=_schema({
"frame_id": pa.int32(),
"video": pa.large_binary(),
}),
options=dict(_PARQUET_OPTIONS, **{
"video-frame-field": "video",
"blob-as-descriptor": "true",
}),
)
video_path = os.path.join(self.temp_dir, "local-indexed.mp4")
with av.open(video_path, mode="w", format="mp4") as container:
stream = container.add_stream("libx264", rate=10)
stream.width = stream.height = 16
stream.pix_fmt = "yuv420p"
stream.gop_size = 4
for index in range(12):
frame = av.VideoFrame.from_ndarray(
np.full((16, 16, 3), index * 16, dtype=np.uint8),
format="rgb24",
)
frame.pts = index
frame.time_base = Fraction(1, 10)
for packet in stream.encode(frame):
container.mux(packet)
for packet in stream.encode():
container.mux(packet)
table.add_video(video_path, [{"frame_id": 0}])
table.add_video(
Blob.from_local(video_path), [{"frame_id": 1}], first_frame=1)
descriptors = [
VideoFrameDescriptor.deserialize(row["video"])
for row in sorted(
table.scan().select(["frame_id", "video"]).to_list(),
key=lambda row: row["frame_id"],
)
]
self.assertEqual([0, 1], [value.frame_index for value in descriptors])
for descriptor in descriptors:
payload = Blob.from_file(table.raw_table.file_io, descriptor.uri,
descriptor.offset, descriptor.length).to_data()
self.assertEqual(open(video_path, "rb").read(), payload)
self.assertGreater(len(ContainerFrameIndex(io.BytesIO(payload), len(payload)).keyframes), 1)
def test_video_write_apis_preserve_payload(self):
from pypaimon.table.row.blob import Blob, VideoFrameDescriptor
table = self.conn.create_table(
"indexed_video_write_apis",
schema=_schema({
"frame_id": pa.int32(),
"video": pa.large_binary(),
}),
options=dict(_PARQUET_OPTIONS, **{
"video-frame-field": "video",
"blob-as-descriptor": "true",
}),
)
video = b"indexed-video"
source_path = os.path.join(self.temp_dir, "indexed-source.mp4")
with open(source_path, "wb") as output:
output.write(video)
source = VideoFrameDescriptor(
"file://" + source_path,
0,
len(video),
0,
)
table.add_video(source, [{"frame_id": 1}], first_frame=1)
table.add_videos([(source, [{"frame_id": 2}], 2)])
plain_path = os.path.join(self.temp_dir, "plain-source.mp4")
with open(plain_path, "wb") as output:
output.write(b"plain-video")
table.add_video(Blob.from_local(plain_path), [{"frame_id": 3}])
table.replace_video("frame_id = 3", source, first_frame=3)
rows = sorted(
table.scan().select(["frame_id", "video"]).to_list(),
key=lambda row: row["frame_id"],
)
descriptors = [
VideoFrameDescriptor.deserialize(row["video"])
for row in rows
]
self.assertEqual([1, 2, 3], [value.frame_index for value in descriptors])
for descriptor in descriptors:
self.assertEqual(video, Blob.from_file(
table.raw_table.file_io, descriptor.uri, descriptor.offset, descriptor.length).to_data())
def test_normal_update_preserves_video_descriptors(self):
from pypaimon.table.row.blob import Blob
table = self.conn.create_table(
"update_video_frame_metadata",
schema=_schema({
"episode_id": pa.int64(),
"frame_id": pa.int32(),
"label": pa.string(),
"video": pa.large_binary(),
}),
options=dict(_PARQUET_OPTIONS, **{
"video-frame-field": "video",
"blob-as-descriptor": "false",
}),
)
video_path = os.path.join(self.temp_dir, "update-metadata.mp4")
with open(video_path, "wb") as output:
output.write(b"video-kept-by-normal-update")
table.add_video(
Blob.from_local(video_path),
[
{"episode_id": 42, "frame_id": index, "label": "old"}
for index in range(3)
],
)
before = sorted(
table.scan().select(["frame_id", "video"]).to_list(),
key=lambda row: row["frame_id"],
)
table.update("episode_id = 42", {"label": "new"})
after = sorted(
table.scan().select(["frame_id", "label", "video"]).to_list(),
key=lambda row: row["frame_id"],
)
self.assertEqual(["new"] * 3, [row["label"] for row in after])
self.assertEqual(
[row["video"] for row in before],
[row["video"] for row in after],
)
def test_update_rejects_video_field(self):
table = self.conn.create_table(
"reject_generic_video_update",
schema=_schema({
"frame_id": pa.int32(),
"video": pa.large_binary(),
}),
options=dict(_PARQUET_OPTIONS, **{
"video-frame-field": "video",
}),
)
with self.assertRaisesRegex(ValueError, "replace_video"):
table.update("frame_id = 0", {"video": b"not-an-mp4"})
self.assertIsNone(
table.raw_table.snapshot_manager().get_latest_snapshot()
)
def test_replace_video_updates_only_matching_frame_rows(self):
from pypaimon.table.row.blob import Blob, VideoFrameDescriptor
table = self.conn.create_table(
"replace_video_frames",
schema=_schema({
"episode_id": pa.int64(),
"frame_id": pa.int32(),
"video": pa.large_binary(),
}),
options=dict(_PARQUET_OPTIONS, **{
"video-frame-field": "video",
"blob-as-descriptor": "false",
}),
)
old_path = os.path.join(self.temp_dir, "old-episode.mp4")
new_path = os.path.join(self.temp_dir, "new-episode.mp4")
old_payload = b"old-complete-video"
new_payload = b"new-complete-video"
with open(old_path, "wb") as output:
output.write(old_payload)
with open(new_path, "wb") as output:
output.write(new_payload)
table.add_video(
Blob.from_local(old_path),
[
{"episode_id": 42, "frame_id": index}
for index in range(3)
],
)
table.replace_video(
"frame_id >= 1",
Blob.from_local(new_path),
first_frame=20,
)
rows = sorted(
table.scan().select(["frame_id", "video"]).to_list(),
key=lambda row: row["frame_id"],
)
descriptors = [
VideoFrameDescriptor.deserialize(row["video"])
for row in rows
]
self.assertEqual([0, 20, 21], [value.frame_index for value in descriptors])
self.assertNotEqual(descriptors[0].uri, descriptors[1].uri)
self.assertEqual(descriptors[1].payload_descriptor,
descriptors[2].payload_descriptor)
self.assertTrue(descriptors[1].uri.endswith(".video"))
scalar, bodies = table.scan().read_blobs("video", parallelism=2)
body_by_frame = dict(zip(
scalar["frame_id"].to_pylist(), bodies["video"]
))
self.assertEqual({
0: old_payload,
1: new_payload,
2: new_payload,
}, body_by_frame)
def test_add_batches_aborts_before_commit_on_invalid_video_frame(self):
from pypaimon.table.row.blob import Blob, VideoFrameDescriptor
table = self.conn.create_table(
"failed_batched_video_frames",
schema=_schema({
"episode_id": pa.int32(),
"video": pa.large_binary(),
}),
options=dict(_PARQUET_OPTIONS, **{
"video-frame-field": "video",
}),
)
video_path = os.path.join(self.temp_dir, "valid-before-failure.mp4")
with open(video_path, "wb") as output:
output.write(b"valid-video")
payload = Blob.from_local(video_path).to_descriptor()
valid = VideoFrameDescriptor(
payload.uri, payload.offset, payload.length, 0
).serialize()
with self.assertRaisesRegex(ValueError, "VideoFrameDescriptor"):
table.add_batches([
[{"episode_id": 1, "video": valid}],
[{"episode_id": 1, "video": b"inline-is-invalid"}],
])
self.assertIsNone(
table.raw_table.snapshot_manager().get_latest_snapshot()
)
def test_add_batches_empty_iterable_is_noop(self):
table = self.conn.create_table(
"empty_batched_frames",
schema=_schema({"frame_index": pa.int32()}),
options=_PARQUET_OPTIONS,
)
self.assertIs(table, table.add_batches(iter(())))
self.assertIsNone(
table.raw_table.snapshot_manager().get_latest_snapshot()
)
def test_scan_to_torch_keeps_blob_descriptors(self):
table = self.conn.create_table(
"torch_video_frames",
schema=_schema({
"episode_id": pa.int32(),
"video": pa.large_binary(),
}),
options=dict(_PARQUET_OPTIONS, **{
"video-frame-field": "video",
"blob-as-descriptor": "false",
}),
)
from pypaimon.table.row.blob import Blob
video_path = os.path.join(self.temp_dir, "torch-episode.mp4")
with open(video_path, "wb") as output:
output.write(b"fake-video")
table.add_video(
Blob.from_local(video_path), [{"episode_id": 1}]
)
sentinel = object()
captured = {}
def fake_to_torch(table_read, splits, **kwargs):
captured["descriptor"] = table_read.table.options.blob_as_descriptor()
captured["splits"] = splits
captured["kwargs"] = kwargs
return sentinel
with patch(
"pypaimon.read.table_read.TableRead.to_torch",
autospec=True,
side_effect=fake_to_torch,
):
result = table.scan().select(
["episode_id", "video"]
).to_torch(
streaming=True,
prefetch_concurrency=2,
shuffle=False,
)
self.assertIs(sentinel, result)
self.assertTrue(captured["descriptor"])
self.assertTrue(captured["splits"])
self.assertEqual(True, captured["kwargs"]["streaming"])
self.assertEqual(2, captured["kwargs"]["prefetch_concurrency"])
def test_create_table_uses_options_and_partitioned(self):
table = self.conn.create_table(
"docs",
schema=_schema({
"id": pa.int32(),
"embedding": _vector(3),
"dt": pa.string(),
}),
options=dict(_PARQUET_OPTIONS, **{
"deletion-vectors.enabled": "false",
}),
partitioned=["dt"],
)
options = table.raw_table.table_schema.options
self.assertEqual(["dt"], table.raw_table.table_schema.partition_keys)
self.assertEqual("false", options["deletion-vectors.enabled"])
self.assertEqual("parquet", options["file.format"])
self.assertEqual("parquet", options["vector.file.format"])
def test_create_table_rejects_non_arrow_schema(self):
with self.assertRaisesRegex(ValueError, "pyarrow.Schema"):
self.conn.create_table("docs", schema={"id": pa.int32()})
def test_create_table_accepts_pyarrow_schema_types(self):
table = self.conn.create_table(
"typed",
schema=pa.schema([
pa.field("flag", pa.bool_(), nullable=False),
pa.field("tiny", pa.int8()),
pa.field("small", pa.int16()),
pa.field("id", pa.int64()),
pa.field("score", pa.float32()),
pa.field("ratio", pa.float64()),
pa.field("title", pa.string()),
pa.field("payload", pa.binary()),
pa.field("hash", pa.binary(16)),
pa.field("blob", pa.large_binary()),
pa.field("amount", pa.decimal128(12, 2)),
pa.field("dt", pa.date32()),
pa.field("created_at", pa.timestamp("us")),
pa.field("created_ltz", pa.timestamp("us", tz="UTC")),
pa.field("tags", pa.list_(pa.string())),
pa.field("attrs", pa.map_(pa.string(), pa.int32())),
pa.field("meta", pa.struct([pa.field("rank", pa.int32())])),
pa.field("embedding", _vector(3)),
]),
)
types_by_name = {
field.name: str(field.type) for field in table.raw_table.fields
}
self.assertEqual("BOOLEAN NOT NULL", types_by_name["flag"])
self.assertEqual("TINYINT", types_by_name["tiny"])
self.assertEqual("SMALLINT", types_by_name["small"])
self.assertEqual("BIGINT", types_by_name["id"])
self.assertEqual("FLOAT", types_by_name["score"])
self.assertEqual("DOUBLE", types_by_name["ratio"])
self.assertEqual("STRING", types_by_name["title"])
self.assertEqual("BYTES", types_by_name["payload"])
self.assertEqual("BINARY(16)", types_by_name["hash"])
self.assertEqual("BLOB", types_by_name["blob"])
self.assertEqual("DECIMAL(12, 2)", types_by_name["amount"])
self.assertEqual("DATE", types_by_name["dt"])
self.assertEqual("TIMESTAMP(6)", types_by_name["created_at"])
self.assertEqual("TIMESTAMP_LTZ(6)", types_by_name["created_ltz"])
self.assertEqual("ARRAY<STRING>", types_by_name["tags"])
self.assertEqual("MAP<STRING NOT NULL, INT>", types_by_name["attrs"])
self.assertEqual("ROW<rank: INT>", types_by_name["meta"])
self.assertEqual("VECTOR<FLOAT, 3>", types_by_name["embedding"])
def test_blob_store_put_objects_get_list_and_delete(self):
table = self.conn.create_table(
"objects",
schema=_schema({
"key": pa.string(),
"image": pa.large_binary(),
"content_type": pa.string(),
"owner": pa.string(),
}),
options=_PARQUET_OPTIONS,
)
store = table.blobs(column="image")
results = store.put_objects([
{
"key": "images/cat.jpg",
"body": b"cat-image-old",
"columns": {"content_type": "image/gif", "owner": "ignored"},
},
{
"key": "images/cat.jpg",
"body": b"cat-image-v1",
"columns": {"content_type": "image/jpeg", "owner": "alice"},
},
{
"key": "images/dog.jpg",
"body": bytearray(b"dog-image"),
"columns": {"content_type": "image/jpeg", "owner": "bob"},
},
])
self.assertEqual(["images/cat.jpg", "images/dog.jpg"],
[result.key for result in results])
self.assertEqual([12, 9], [result.size for result in results])
cat = store.get_object("images/cat.jpg")
self.assertEqual(b"cat-image-v1", cat.read())
self.assertEqual(table.raw_table.table_path, cat.table_root)
with patch.object(cat.file_io, 'create_blob_presigned_url',
return_value="https://signed-url") as presign:
validity = timedelta(minutes=30)
self.assertEqual("https://signed-url", cat.to_presigned_url(validity))
presign.assert_called_once_with(cat.table_root, cat.descriptor, validity)
presign.reset_mock()
ranged = store.get_object("images/cat.jpg", range="bytes=1-4")
self.assertEqual("https://signed-url", ranged.to_presigned_url(validity))
root, descriptor, expiry = presign.call_args[0]
self.assertEqual(cat.table_root, root)
self.assertEqual(cat.descriptor.uri, descriptor.uri)
self.assertEqual(cat.descriptor.offset + 1, descriptor.offset)
self.assertEqual(4, descriptor.length)
self.assertEqual(validity, expiry)
presign.reset_mock()
with self.assertRaisesRegex(ValueError, "bound to a table root"):
replace(cat, table_root=None).to_presigned_url(validity)
presign.assert_not_called()
self.assertEqual(b"at-i", store.get_object(
"images/cat.jpg", range="bytes=1-4").read())
clipped = store.get_object("images/cat.jpg", range="bytes=10-999")
self.assertEqual(2, clipped.content_length)
self.assertEqual(b"v1", clipped.read())
with self.assertRaisesRegex(ValueError, "Range start"):
store.get_object("images/cat.jpg", range="bytes=12-13")
self.assertEqual("image/jpeg", cat.columns["content_type"])
self.assertEqual("alice", cat.columns["owner"])
owner_only = store.get_object(
"images/cat.jpg",
columns=["owner"],
)
self.assertEqual({"owner": "alice"}, owner_only.columns)
listed = store.list_objects(prefix="images/")
self.assertEqual(["images/cat.jpg", "images/dog.jpg"],
sorted(obj.key for obj in listed))
self.assertEqual(
{
"images/cat.jpg": "image/jpeg",
"images/dog.jpg": "image/jpeg",
},
{obj.key: obj.columns["content_type"] for obj in listed},
)
listed_without_columns = store.list_objects(
prefix="images/",
columns=[],
)
self.assertEqual(
{"images/cat.jpg": {}, "images/dog.jpg": {}},
{obj.key: obj.columns for obj in listed_without_columns},
)
self.assertEqual([], store.list_objects(prefix="images/", limit=0))
with self.assertRaisesRegex(ValueError, "limit"):
store.list_objects(prefix="images/", limit=-1)
store.put_object(
"images/cat.jpg",
b"cat-image-v2",
columns={"content_type": "image/png", "owner": "alice"},
)
self.assertEqual(b"cat-image-v2", store.get_object("images/cat.jpg").read())
info = store.head_object("images/cat.jpg")
self.assertEqual("image/png", info.columns["content_type"])
self.assertEqual(
{"content_type": "image/png"},
store.head_object(
"images/cat.jpg",
columns="content_type",
).columns,
)
self.assertEqual(2, table.scan().to_arrow().num_rows)
previous_descriptor = info.descriptor
updated = store.update_object_columns(
"images/cat.jpg",
{"content_type": "image/webp"},
)
self.assertEqual(previous_descriptor, updated.descriptor)
self.assertEqual("image/webp", updated.columns["content_type"])
self.assertEqual("alice", updated.columns["owner"])
self.assertEqual(b"cat-image-v2", store.get_object("images/cat.jpg").read())
batch_updates = store.update_objects_columns([
{"key": "images/cat.jpg", "columns": {"owner": "carol"}},
{"key": "images/dog.jpg", "columns": {"owner": "dave"}},
])
self.assertEqual(
["images/cat.jpg", "images/dog.jpg"],
[obj.key for obj in batch_updates],
)
self.assertEqual("carol", store.head_object("images/cat.jpg").columns["owner"])
self.assertEqual("dave", store.head_object("images/dog.jpg").columns["owner"])
self.assertEqual(2, table.scan().to_arrow().num_rows)
with self.assertRaisesRegex(ValueError, "columns must not be empty"):
store.update_object_columns("images/cat.jpg", {})
with self.assertRaises(pmm.NoSuchKey):
store.update_object_columns("images/missing.jpg", {"owner": "nobody"})
with self.assertRaisesRegex(ValueError, "columns must not include"):
store.update_object_columns("images/cat.jpg", {"image": b"new"})
store.delete_objects(["images/dog.jpg", "images/dog.jpg"])
with self.assertRaises(pmm.NoSuchKey):
store.head_object("images/dog.jpg")
self.assertEqual(["images/cat.jpg"],
[obj.key for obj in store.list_objects(prefix="images/")])
store.delete_object("images/cat.jpg")
self.assertEqual([], store.list_objects(prefix="images/"))
@pytest.mark.python_read
def test_blob_store_list_reads_batches_and_stops_at_limit(self):
from pypaimon.read.table_read import TableRead
table = self.conn.create_table(
"list_batches", schema=_schema({
"key": pa.string(), "image": pa.large_binary(), "owner": pa.string(),
}), options=dict(_PARQUET_OPTIONS, **{"read.batch-size": "2"}))
table.add([
{"key": key, "image": b"body", "owner": "alice"}
for key in ("other/a", "other/b", "images/a", "images/b", "images/c", "images/d")
])
store = table.blobs(column="image")
original = PythonTableRead._arrow_batch_generator
batches, closed = [], []
def tracked_read(read, *args, **kwargs):
reader = original(read, *args, **kwargs)
try:
for batch in reader:
batches.append(batch.num_rows)
yield batch
finally:
reader.close()
closed.append(True)
with patch.object(TableRead, "to_arrow", side_effect=AssertionError("full read")), \
patch.object(PythonTableRead, "_arrow_batch_generator", tracked_read):
self.assertEqual([], store.list_objects(limit=0))
self.assertEqual([], batches)
listed = store.list_objects(prefix="images/", limit=1, columns=["owner"])
self.assertEqual(["images/a"], [obj.key for obj in listed])
self.assertEqual({"owner": "alice"}, listed[0].columns)
self.assertEqual(4, listed[0].size)
self.assertEqual([2, 2], batches)
self.assertEqual([True], closed)
batches.clear()
listed = store.list_objects(limit=1, columns=[])
self.assertEqual(1, len(listed))
self.assertEqual({}, listed[0].columns)
self.assertEqual([1], batches)
self.assertEqual(2, len(closed))
with patch.object(store, "_row_to_info", side_effect=ValueError("invalid descriptor")):
try:
store.list_objects()
except ValueError as error:
self.assertEqual("invalid descriptor", str(error))
# Check while the traceback still holds the reader's frame.
self.assertEqual(3, len(closed))
else:
self.fail("Expected invalid descriptor")
self.assertEqual([], store.list_objects(prefix="absent/"))
self.assertEqual(4, len(closed))
def test_blob_store_list_preserves_non_string_prefix_semantics(self):
table = self.conn.create_table(
"numeric_keys", schema=_schema({"key": pa.int32(), "image": pa.large_binary()}),
options=_PARQUET_OPTIONS)
table.add([
{"key": value, "image": b"body"} for value in (None, 10, 20, 21)
])
store = table.blobs(column="image")
self.assertEqual([None], [obj.key for obj in store.list_objects(prefix="N", limit=1)])
self.assertEqual([20], [obj.key for obj in store.list_objects(prefix="2", limit=1)])
self.assertEqual(4, len(store.list_objects(prefix="")))
def test_blob_store_put_object_accepts_blob_without_materializing(self):
from pypaimon.table.row.blob import Blob, BlobDescriptor
table = self.conn.create_table(
"streamed_objects",
schema=_schema({
"key": pa.string(),
"payload": pa.large_binary(),
}),
options=_PARQUET_OPTIONS,
)
data = b"streamed-managed-payload"
stream_uri = "stream://payloads/1"
stream_calls = []
class StreamOnlyReader:
def new_input_stream(self, uri):
stream_calls.append(("new_input_stream", uri))
return io.BytesIO(data)
class StreamOnlyReaderFactory:
def create(self, uri):
stream_calls.append(("create", uri))
return StreamOnlyReader()
class DescriptorOnlyBlob(Blob):
def to_data(self):
raise AssertionError("put_object should not materialize Blob data")
def to_descriptor(self):
return BlobDescriptor(stream_uri, 0, len(data))
def new_input_stream(self):
raise AssertionError("put_object should use the URI stream")
store = table.blobs(column="payload")
original_factory = table.raw_table.file_io.uri_reader_factory
table.raw_table.file_io.uri_reader_factory = StreamOnlyReaderFactory()
try:
result = store.put_object("payloads/1", DescriptorOnlyBlob())
finally:
table.raw_table.file_io.uri_reader_factory = original_factory
self.assertEqual(len(data), result.size)
self.assertEqual([
("create", stream_uri),
("new_input_stream", stream_uri),
], stream_calls)
stored = store.get_object("payloads/1")
self.assertNotEqual(stream_uri, stored.descriptor.uri)
self.assertEqual(data, stored.read())
def test_blob_store_put_object_reference_preserves_descriptor_uri(self):
table = self.conn.create_table(
"referenced_objects",
schema=_schema({
"object_key": pa.string(),
"payload": pa.large_binary(),
"media_type": pa.string(),
}),
options=dict(_PARQUET_OPTIONS, **{
"blob-descriptor-field": "payload",
}),
)
data = b"external-video-payload"
external_path = os.path.join(self.temp_dir, "video.bin")
with open(external_path, "wb") as f:
f.write(data)
store = table.blobs(column="payload")
result = store.put_object(
"videos/1",
uri=external_path,
length=len(data),
columns={"media_type": "video/mp4"},
)
descriptor = store.head_object("videos/1").descriptor
more = store.put_objects([
{
"key": "videos/2",
"descriptor": descriptor,
"columns": {"media_type": "video/mp4"},
}
])
self.assertEqual("videos/1", result.key)
self.assertEqual(len(data), result.size)
self.assertEqual("videos/2", more[0].key)
obj = store.get_object("videos/1")
with patch.object(obj.file_io, "new_input_stream", side_effect=AssertionError(
"external object must use its URI reader")):
self.assertEqual(data, obj.read())
self.assertEqual("video/mp4", obj.columns["media_type"])
self.assertEqual(external_path, store.head_object("videos/1").descriptor.uri)
self.assertEqual(data, store.get_object("videos/2").read())
def test_blob_store_managed_object_reads_through_fuse_file_io(self):
from pypaimon.common.file_io import FileIO
from pypaimon.filesystem.local_file_io import FuseLocalFileIO
from pypaimon.table.row.blob import BlobDescriptor
table = self.conn.create_table(
"mounted_objects",
schema=_schema({"key": pa.string(), "payload": pa.large_binary()}),
options=_PARQUET_OPTIONS,
)
store = table.blobs(column="payload")
store.put_object("example.hdf5", body=b"hdf5-payload")
obj = store.get_object("example.hdf5")
remote_root = "oss://test-bucket/table"
remote_uri = remote_root + obj.descriptor.uri[len(obj.table_root):]
mounted = replace(
obj,
descriptor=BlobDescriptor(
remote_uri, obj.descriptor.offset, obj.descriptor.length),
file_io=FuseLocalFileIO(remote_root, obj.table_root),
table_root=remote_root,
)
with patch.object(FileIO, "get", side_effect=AssertionError(
"managed object must use the table FileIO")):
with mounted.open() as stream:
self.assertEqual(b"hdf5-payload", stream.read())
self.assertEqual(b"df5", replace(
mounted, range_header="bytes=1-3").read())
suffix = remote_uri[len(remote_root):]
for variant in (
"OSS://test-bucket/table" + suffix,
remote_root + "//./" + suffix.lstrip("/"),
remote_root + "/sub/../" + suffix.lstrip("/")):
with self.subTest(uri=variant), patch.object(
mounted.file_io.uri_reader_factory, "create",
side_effect=AssertionError("managed object must not use a URI reader")):
normalized = replace(
mounted,
descriptor=BlobDescriptor(
variant, obj.descriptor.offset, obj.descriptor.length),
)
self.assertEqual(b"hdf5-payload", normalized.read())
def test_blob_store_external_uris_do_not_use_fuse_file_io(self):
from pypaimon.filesystem.local_file_io import FuseLocalFileIO
from pypaimon.multimodal.blob_store import BlobObject
from pypaimon.table.row.blob import BlobDescriptor
class ExternalReader:
def new_input_stream(self, uri):
return io.BytesIO(b"external")
root = "oss://test-bucket/table"
file_io = FuseLocalFileIO(root, self.temp_dir)
for uri in (
root + "/../other/payload.bin",
"oss://other-bucket/table/payload.bin",
"oss://test-bucket/table-sibling/payload.bin",
"s3://test-bucket/table/payload.bin",
root + "/payload.bin?token=external"):
with self.subTest(uri=uri), patch.object(
file_io, "new_input_stream",
side_effect=AssertionError("external object must not use table FileIO")), \
patch.object(file_io.uri_reader_factory, "create",
return_value=ExternalReader()) as create_reader:
obj = BlobObject(
key="external",
descriptor=BlobDescriptor(uri, 0, len(b"external")),
columns={},
file_io=file_io,
table_root=root,
)
self.assertEqual(b"external", obj.read())
create_reader.assert_called_once_with(uri)
def test_blob_store_put_object_uri_streams_into_managed_blob(self):
table = self.conn.create_table(
"managed_only_objects",
schema=_schema({
"key": pa.string(),
"payload": pa.large_binary(),
}),
options=_PARQUET_OPTIONS,
)
data = b"external-managed-payload"
external_path = os.path.join(self.temp_dir, "managed.bin")
with open(external_path, "wb") as f:
f.write(data)
store = table.blobs(column="payload")
result = store.put_object(
"payloads/1",
uri=external_path,
length=len(data),
)
self.assertEqual(len(data), result.size)
stored = store.get_object("payloads/1")
self.assertNotEqual(external_path, stored.descriptor.uri)
self.assertEqual(data, stored.read())
def test_drop_table_can_ignore_missing_table(self):
self.conn.drop_table("missing", ignore_if_not_exists=True)
def test_get_table_rejects_non_data_evolution_table(self):
self.conn.catalog.create_database("default", ignore_if_exists=True)
self.conn.catalog.create_table(
"default.raw",
_raw_schema(),
False,
)
with self.assertRaisesRegex(ValueError, "data-evolution.enabled"):
self.conn.get_table("raw")
def test_get_table_rejects_primary_key_table(self):
self.conn.catalog.create_database("default", ignore_if_exists=True)
self.conn.catalog.create_table(
"default.pk", _raw_schema(options=_PARQUET_OPTIONS), False)
# Row tracking tables cannot define primary keys, so such a table can
# only come from writers that skipped schema validation. Simulate it
# by rewriting the committed schema file.
table = self.conn.catalog.get_table("default.pk")
legacy_schema = replace(table.table_schema, primary_keys=["id"])
table.file_io.overwrite_file_utf8(
table.schema_manager._to_schema_path(legacy_schema.id),
JSON.to_json(legacy_schema, indent=2))
with self.assertRaisesRegex(ValueError, "primary keys"):
self.conn.get_table("pk")
def test_create_table_ignores_invalid_options_when_table_exists(self):
schema = _schema({"id": pa.int32()})
expected = self.conn.create_table("existing", schema=schema)
actual = self.conn.create_table(
"existing",
schema=_schema({
"id": pa.int32(),
"embedding": _vector(3),
}),
options={"data-evolution.enabled": "false"},
ignore_if_exists=True,
)
self.assertEqual(expected.identifier, actual.identifier)
with patch(
"pypaimon.multimodal.connection._table_exists",
return_value=False):
raced = self.conn.create_table(
"existing",
schema=_schema({
"id": pa.int32(),
"embedding": _vector(3),
}),
options={"data-evolution.enabled": "false"},
ignore_if_exists=True,
)
self.assertEqual(expected.identifier, raced.identifier)
def test_create_table_handles_concurrent_delete_when_ignoring(self):
schema = _schema({"id": pa.int32()})
self.conn.create_table("deleted", schema=schema)
original_get = self.conn.get_table
deleted = [False]
def delete_once(name):
if not deleted[0]:
deleted[0] = True
self.conn.catalog.drop_table("default.deleted", False)
return original_get(name)
with patch.object(
self.conn, "get_table", side_effect=delete_once):
table = self.conn.create_table(
"deleted",
data=pa.table({"id": [1, 2, 3]}),
schema=schema,
ignore_if_exists=True,
)
self.assertEqual("default.deleted", table.identifier)
self.assertEqual(
[1, 2, 3], table.scan().to_arrow()["id"].to_pylist())
self.conn.create_table("fallback_deleted", schema=schema)
deleted[0] = False
def delete_fallback_once(name):
if not deleted[0]:
deleted[0] = True
self.conn.catalog.drop_table(
"default.fallback_deleted", False)
return original_get(name)
with patch(
"pypaimon.multimodal.connection._table_exists",
return_value=False):
with patch.object(
self.conn,
"get_table",
side_effect=delete_fallback_once):
with self.assertRaisesRegex(
ValueError, "data-evolution.enabled"):
self.conn.create_table(
"fallback_deleted",
schema=schema,
options={"data-evolution.enabled": "false"},
ignore_if_exists=True,
)
def test_create_table_does_not_add_data_to_concurrent_winner(self):
schema = _schema({"id": pa.int32()})
winner = self.conn.create_table("winner", schema=schema)
with patch(
"pypaimon.multimodal.connection._table_exists",
return_value=False):
actual = self.conn.create_table(
"winner",
data=pa.table({"id": [1, 2, 3]}),
schema=schema,
ignore_if_exists=True,
)
self.assertEqual(winner.identifier, actual.identifier)
self.assertEqual([], actual.scan().to_arrow().to_pylist())
def test_create_table_preserves_unknown_create_error(self):
schema = _schema({"id": pa.int32()})
create = self.conn.catalog.create_table
create_error = TimeoutError("create response lost")
def create_then_lose_response(*args, **kwargs):
create(*args, **kwargs)
raise create_error
with patch.object(
self.conn.catalog,
"create_table",
side_effect=create_then_lose_response):
with self.assertRaises(TimeoutError) as context:
self.conn.create_table(
"failed_create",
data=pa.table({"id": [1, 2, 3]}),
schema=schema,
ignore_if_exists=True,
)
self.assertIs(create_error, context.exception)
self.assertEqual(
[],
self.conn.get_table("failed_create")
.scan().to_arrow().to_pylist(),
)
def test_create_table_can_add_initial_data_and_get_by_short_name(self):
self.conn.create_table(
"users",
data=[
{"id": 1, "name": "Alice", "age": 30},
{"id": 2, "name": "Bob", "age": 25},
],
schema=_schema({
"id": pa.int32(),
"name": pa.string(),
"age": pa.int32(),
}),
options=_PARQUET_OPTIONS,
)
users = self.conn.get_table("users")
result = users.scan().select(["id", "name"]).to_arrow()
self.assertEqual(["id", "name"], result.column_names)
self.assertEqual([1, 2], result["id"].to_pylist())
def test_add_preserves_fields_missing_from_first_row(self):
table = self.conn.create_table(
"sparse_rows", schema=_schema({
"id": pa.int32(), "caption": pa.string(), "quality": pa.float64(),
"image": pa.large_binary(), "missing": pa.string(),
}), options=_PARQUET_OPTIONS)
table.add([
{"id": "1"},
{"id": "2", "caption": "keep-me", "quality": 0.9,
"image": b"payload", "extra": "ignored"},
])
scalar, blobs = table.scan().read_blobs("image")
rows = sorted(zip(scalar.to_pylist(), blobs["image"]), key=lambda row: row[0]["id"])
self.assertEqual([
({"id": 1, "caption": None, "quality": None, "missing": None}, None),
({"id": 2, "caption": "keep-me", "quality": 0.9, "missing": None}, b"payload"),
], rows)
def test_add_python_map_values(self):
schema = _schema({
"id": pa.int32(), "assets": pa.map_(pa.string(), pa.large_binary()),
"scores": pa.map_(pa.string(), pa.int64()),
})
rows = [
{"id": 1, "assets": {"image": pmm.Blob.from_data(b"body"), "missing": None},
"scores": {"quality": 9}},
{"id": 2, "assets": [("empty", b""), ("image", b"second")], "scores": []},
{"id": 3, "assets": [], "scores": None},
{"id": 4, "assets": None, "scores": None},
]
for columnar in (False, True):
with self.subTest(columnar=columnar):
table = self.conn.create_table(
"map_input_%s" % columnar, schema=schema, options=_PARQUET_OPTIONS)
data = {name: [row[name] for row in rows] for name in schema.names} if columnar else rows
table.add(data)
scalar, blobs = table.scan().read_blobs("assets")
self.assertEqual({
1: [("image", b"body"), ("missing", None)],
2: [("empty", b""), ("image", b"second")], 3: [], 4: None,
}, dict(zip(scalar["id"].to_pylist(), blobs["assets"])))
self.assertEqual({1: [("quality", 9)], 2: [], 3: None, 4: None},
dict(zip(scalar["id"].to_pylist(), scalar["scores"].to_pylist())))
def test_python_input_preserves_safe_casts_and_empty_rows(self):
from pypaimon.multimodal.table import _to_arrow_table
schema = _schema({"id": pa.int32(), "caption": pa.string()})
for data in ([{"id": 1.5}], {"id": [1.5]}):
with self.assertRaises(pa.ArrowInvalid):
_to_arrow_table(data, schema)
self.assertEqual([], _to_arrow_table([], schema).to_pylist())
self.assertEqual([{"id": None, "caption": None}] * 2,
_to_arrow_table([{}, {}], schema).to_pylist())
self.assertEqual([{"id": 2, "caption": None}],
_to_arrow_table({"id": ["2"]}, schema).to_pylist())
def test_add_scan_where_select_limit(self):
users = self.conn.create_table(
"users",
schema=_schema({
"id": pa.int32(),
"name": pa.string(),
"age": pa.int32(),
}),
options=_PARQUET_OPTIONS,
)
users.add([
{"id": 1, "name": "Alice", "age": 30},
{"id": 2, "name": "Bob", "age": 25},
{"id": 3, "name": "Carol", "age": 40},
])
result = (
users.scan()
.where("age >= 30")
.select(["id", "name"])
.limit(1)
.to_arrow()
)
self.assertEqual(["id", "name"], result.column_names)
self.assertEqual(1, result.num_rows)
self.assertEqual([1], result["id"].to_pylist())
@pytest.mark.python_read
def test_scan_to_arrow_batch_reader(self):
users = self.conn.create_table(
"batch_users",
data=[
{"id": 1, "age": 20},
{"id": 2, "age": 30},
{"id": 3, "age": 40},
],
schema=_schema({"id": pa.int32(), "age": pa.int32()}),
options=_PARQUET_OPTIONS,
)
reader = (
users.scan()
.where("age >= 30")
.select("id")
.to_arrow_batch_reader()
)
self.assertEqual([{"id": 2}, {"id": 3}], reader.read_all().to_pylist())
closed = []
def batches(*args):
try:
yield pa.record_batch(
[pa.array([1], type=pa.int32())], names=["id"])
yield pa.record_batch(
[pa.array([2], type=pa.int32())], names=["id"])
finally:
closed.append(True)
with patch(
"pypaimon.execution.python.read.table_read.PythonTableRead._arrow_batch_generator",
new=batches), patch(
"pypaimon.read.read_support."
"_RECORD_BATCH_READER_FROM_STREAM", None):
reader = users.scan().select("id").to_arrow_batch_reader()
reader.read_next_batch()
reader.close()
self.assertEqual([True], closed)
closed.clear()
with patch(
"pypaimon.execution.python.read.table_read.PythonTableRead._arrow_batch_generator",
new=batches):
reader = users.scan().select("id").to_arrow_batch_reader()
reader.read_next_batch()
reader.close()
self.assertEqual([True], closed)
read_builder = users.raw_table.new_read_builder()
splits = read_builder.new_scan().plan().splits()
with patch(
"pypaimon.read.read_support."
"_RECORD_BATCH_READER_FROM_STREAM", None):
reader = read_builder.new_read().to_arrow_batch_reader(splits)
self.assertIsInstance(reader, pa.RecordBatchReader)
reader.close()
with self.assertRaisesRegex(TypeError, "only supported on scan"):
users.search("thirty", column="id").to_arrow_batch_reader()
def test_scan_with_row_id_returns_system_column(self):
users = self.conn.create_table(
"users",
data=[
{"id": 1, "name": "Alice"},
{"id": 2, "name": "Bob"},
],
schema=_schema({
"id": pa.int32(),
"name": pa.string(),
}),
options=_PARQUET_OPTIONS,
)
result = (
users.scan()
.with_row_id()
.select(["id"])
.to_arrow()
)
self.assertEqual(["id", "_ROW_ID"], result.column_names)
self.assertEqual([1, 2], result["id"].to_pylist())
self.assertEqual([0, 1], result["_ROW_ID"].to_pylist())
def test_take_row_ids_reads_projected_rows(self):
docs = self.conn.create_table(
"docs",
data=[
{"id": 1, "content": "alpha"},
{"id": 2, "content": "beta"},
{"id": 3, "content": "gamma"},
],
schema=_schema({
"id": pa.int32(),
"content": pa.string(),
}),
options=_PARQUET_OPTIONS,
)
manifest = {
row["id"]: row["_ROW_ID"]
for row in docs.scan().select(["id"]).with_row_id().to_list()
}
rows = sorted(
docs.take_row_ids([manifest[3], manifest[1]])
.select(["id", "content"])
.with_row_id()
.to_list(),
key=lambda row: row["id"],
)
self.assertEqual(
[
{"id": 1, "content": "alpha", "_ROW_ID": manifest[1]},
{"id": 3, "content": "gamma", "_ROW_ID": manifest[3]},
],
rows,
)
def test_take_row_ids_accepts_empty_manifest(self):
docs = self.conn.create_table(
"docs",
data=[{"id": 1}],
schema=_schema({"id": pa.int32()}),
options=_PARQUET_OPTIONS,
)
self.assertEqual([], docs.take_row_ids([]).select(["id"]).to_list())
def test_row_id_filter_and_ordered_batch_have_distinct_contracts(self):
docs = self.conn.create_table(
"docs",
data=[{"id": i, "content": str(i)} for i in range(3)],
schema=_schema({"id": pa.int32(), "content": pa.string()}),
options=_PARQUET_OPTIONS,
)
requested = [2, 0, 2]
filtered = docs.take_row_ids(requested).select("id").with_row_id().to_list()
self.assertEqual(
[{"id": 0, "_ROW_ID": 0}, {"id": 2, "_ROW_ID": 2}],
sorted(filtered, key=lambda row: row["_ROW_ID"]),
)
batch = docs.take_batch(requested, columns=["content", "_ROW_ID"])
self.assertIsInstance(batch, pa.Table)
self.assertEqual(["content", "_ROW_ID"], batch.column_names)
self.assertEqual(
[{"content": str(i), "_ROW_ID": i} for i in requested],
batch.to_pylist(),
)
self.assertEqual(
[{"id": i, "content": str(i)} for i in requested],
docs.take_batch(iter(requested)).to_pylist(),
)
self.assertEqual(
[{"id": i} for i in requested],
docs.take_batch(pa.chunked_array([[2], [0, 2]]), columns="id").to_pylist(),
)
def test_take_batch_empty_and_missing_rows(self):
docs = self.conn.create_table(
"docs", data=[{"id": i} for i in range(3)],
schema=_schema({"id": pa.int32()}), options=_PARQUET_OPTIONS,
)
empty = docs.take_batch([], columns=["id", "_ROW_ID"])
self.assertEqual(
pa.schema([pa.field("id", pa.int32()), pa.field("_ROW_ID", pa.int64())]),
empty.schema,
)
self.assertEqual([], empty.to_pylist())
self.assertEqual(3, docs.take_batch([2, 0, 2], columns=[]).num_rows)
docs.delete(where="id = 1")
for row_ids in ([2, 1, 2], [99], [99, 1, 99]):
with self.subTest(row_ids=row_ids):
with self.assertRaisesRegex(KeyError, "Row IDs not found"):
docs.take_batch(row_ids, columns="id")
self.assertEqual(
[2, 0, 2], docs.take_batch([2, 0, 2], columns="id")["id"].to_pylist())
def test_take_batch_reads_one_snapshot_with_order_and_duplicates(self):
docs = self.conn.create_table(
"docs", data=[{"id": i} for i in range(3)],
schema=_schema({"id": pa.int32()}), options=_PARQUET_OPTIONS,
)
snapshot_id = docs.raw_table.snapshot_manager().get_latest_snapshot().id
docs.raw_table.create_tag("training", snapshot_id=snapshot_id)
docs.delete(where="id = 2")
for selection in ({"snapshot_id": snapshot_id}, {"tag_name": "training"}):
with self.subTest(selection=selection):
self.assertEqual(
[{"id": i, "_ROW_ID": i} for i in [2, 0, 2]],
docs.take_batch(
[2, 0, 2], columns=["id", "_ROW_ID"], **selection).to_pylist(),
)
with self.assertRaisesRegex(ValueError, "cannot be set at the same time"):
docs.take_batch([0], snapshot_id=snapshot_id, tag_name="training")
def test_row_id_reads_require_row_tracking(self):
docs = self.conn.create_table(
"docs", schema=_schema({"id": pa.int32()}), options=_PARQUET_OPTIONS)
untracked = docs.copy({"row-tracking.enabled": "false"})
for read in (untracked.take_row_ids, untracked.take_batch):
with self.subTest(read=read.__name__):
with self.assertRaisesRegex(ValueError, "require row-tracking.enabled=true"):
read([0])
def test_take_row_ids_supports_snapshot_and_tag_time_travel(self):
docs = self.conn.create_table(
"docs",
schema=_schema({
"id": pa.int32(),
"content": pa.string(),
"embedding": _vector(3),
}),
options=_PARQUET_OPTIONS,
)
docs.add([
{"id": 1, "content": "first", "embedding": [1.0, 0.0, 0.0]},
])
snapshot_id = docs.raw_table.snapshot_manager().get_latest_snapshot().id
docs.raw_table.create_tag("v1", snapshot_id=snapshot_id)
row_id = (
docs.scan(snapshot_id=snapshot_id)
.select(["id"])
.with_row_id()
.to_list()[0]["_ROW_ID"]
)
docs.delete(where="id = 1")
self.assertEqual(
[],
docs.take_row_ids([row_id]).select(["id"]).to_list(),
)
self.assertEqual(
[{"id": 1, "_ROW_ID": row_id}],
docs.take_row_ids([row_id], snapshot_id=snapshot_id)
.select(["id"])
.with_row_id()
.to_list(),
)
self.assertEqual(
[{"id": 1}],
docs.take_row_ids([row_id], tag_name="v1")
.select(["id"])
.to_list(),
)
self.assertEqual(
snapshot_id,
docs.search(
[1.0, 0.0, 0.0],
column="embedding",
snapshot_id=snapshot_id,
)._table.options.scan_snapshot_id(),
)
def test_time_travel_rejects_snapshot_id_and_tag_name_together(self):
docs = self.conn.create_table(
"docs",
data=[{"id": 1}],
schema=_schema({"id": pa.int32()}),
options=_PARQUET_OPTIONS,
)
with self.assertRaisesRegex(ValueError, "cannot be set at the same time"):
docs.take_row_ids([0], snapshot_id=1, tag_name="v1")
def test_overwrite_replaces_unpartitioned_table(self):
users = self.conn.create_table(
"users",
data=[
{"id": 1, "name": "Alice", "age": 30},
{"id": 2, "name": "Bob", "age": 25},
],
schema=_schema({
"id": pa.int32(),
"name": pa.string(),
"age": pa.int32(),
}),
options=_PARQUET_OPTIONS,
)
result = users.overwrite([
{"id": 3, "name": "Carol", "age": 40},
])
self.assertIs(users, result)
self.assertEqual(
[
{"id": 3, "name": "Carol", "age": 40},
],
users.scan().to_list(),
)
def test_empty_overwrite_clears_unpartitioned_table(self):
users = self.conn.create_table(
"users",
data=[
{"id": 1, "name": "Alice"},
{"id": 2, "name": "Bob"},
],
schema=_schema({
"id": pa.int32(),
"name": pa.string(),
}),
options=_PARQUET_OPTIONS,
)
users.overwrite([])
self.assertEqual([], users.scan().to_list())
def test_overwrite_replaces_dynamic_partitions(self):
users = self.conn.create_table(
"users",
data=[
{"id": 1, "name": "Alice", "dt": "2024-01-01"},
{"id": 2, "name": "Bob", "dt": "2024-01-02"},
],
schema=_schema({
"id": pa.int32(),
"name": pa.string(),
"dt": pa.string(),
}),
options=_PARQUET_OPTIONS,
partitioned=["dt"],
)
users.overwrite([
{"id": 3, "name": "Carol", "dt": "2024-01-01"},
])
rows = sorted(users.scan().to_list(), key=lambda r: r["id"])
self.assertEqual(
[
{"id": 2, "name": "Bob", "dt": "2024-01-02"},
{"id": 3, "name": "Carol", "dt": "2024-01-01"},
],
rows,
)
def test_overwrite_replaces_static_partition(self):
users = self.conn.create_table(
"users",
data=[
{"id": 1, "name": "Alice", "dt": "2024-01-01"},
{"id": 2, "name": "Bob", "dt": "2024-01-02"},
],
schema=_schema({
"id": pa.int32(),
"name": pa.string(),
"dt": pa.string(),
}),
options=dict(_PARQUET_OPTIONS, **{
"dynamic-partition-overwrite": "false",
}),
partitioned=["dt"],
)
users.overwrite([
{"id": 3, "name": "Carol", "dt": "2024-01-01"},
], partition={"dt": "2024-01-01"})
rows = sorted(users.scan().to_list(), key=lambda r: r["id"])
self.assertEqual(
[
{"id": 2, "name": "Bob", "dt": "2024-01-02"},
{"id": 3, "name": "Carol", "dt": "2024-01-01"},
],
rows,
)
def test_scan_read_blobs(self):
obs = self.conn.create_table(
"obs",
schema=_schema({
"clip": pa.string(),
"idx": pa.int32(),
"image": pa.large_binary(),
}),
options=_PARQUET_OPTIONS,
partitioned=["clip"],
)
payloads = [("blob-%d-" % i).encode() * (i + 1) for i in range(6)]
obs.add([
{"clip": "c1", "idx": 0, "image": payloads[0]},
{"clip": "c1", "idx": 1, "image": payloads[1]},
{"clip": "c1", "idx": 2, "image": payloads[2]},
{"clip": "c2", "idx": 0, "image": payloads[3]},
{"clip": "c2", "idx": 1, "image": payloads[4]},
{"clip": "c3", "idx": 0, "image": None},
])
# scalar table is row-aligned with the blob list and drops the blob column
scalar, blobs = obs.scan().where("clip = 'c1'").read_blobs("image")
images = blobs["image"]
self.assertEqual(3, len(images))
self.assertNotIn("image", scalar.column_names)
got = dict(zip(scalar.column("idx").to_pylist(), images))
self.assertEqual({0: payloads[0], 1: payloads[1], 2: payloads[2]}, got)
# blob column auto-detected when columns=None
_, blobs2 = obs.scan().where("clip = 'c2'").read_blobs()
self.assertEqual({payloads[3], payloads[4]}, set(blobs2["image"]))
# null blob -> None
_, blobs3 = obs.scan().where("clip = 'c3'").read_blobs("image")
self.assertEqual([None], blobs3["image"])
# non-BLOB column is rejected
with self.assertRaisesRegex(ValueError, "not a BLOB column"):
obs.scan().read_blobs("idx")
# empty result: 0 matching rows -> empty blob list, schema preserved
scalar_e, blobs_e = obs.scan().where("clip = 'none'").read_blobs("image")
self.assertEqual(0, scalar_e.num_rows)
self.assertEqual([], blobs_e["image"])
def test_scan_read_blobs_multi_column(self):
obs = self.conn.create_table(
"multi",
schema=_schema({
"clip": pa.string(),
"idx": pa.int32(),
"img": pa.large_binary(),
"aud": pa.large_binary(),
}),
options=_PARQUET_OPTIONS,
partitioned=["clip"],
)
obs.add([
{"clip": "c1", "idx": i,
"img": ("img-%d" % i).encode(), "aud": ("aud-%d" % i).encode()}
for i in range(4)
])
# both BLOB columns fetched in one call, each row-aligned with the scalars
scalar, blobs = obs.scan().where("clip = 'c1'").read_blobs(["img", "aud"])
self.assertEqual({"img", "aud"}, set(blobs))
self.assertNotIn("img", scalar.column_names)
self.assertNotIn("aud", scalar.column_names)
idx = scalar.column("idx").to_pylist()
self.assertEqual({i: ("img-%d" % i).encode() for i in range(4)},
dict(zip(idx, blobs["img"])))
self.assertEqual({i: ("aud-%d" % i).encode() for i in range(4)},
dict(zip(idx, blobs["aud"])))
# reading only one BLOB column must not leak the other into scalar as
# descriptor bytes
scalar1, blobs1 = obs.scan().where("clip = 'c1'").read_blobs("img")
self.assertEqual({"img"}, set(blobs1))
self.assertEqual(["clip", "idx"], scalar1.column_names)
# columns=None intersects all BLOBs with select() -> only projected BLOB
_, blobs2 = obs.scan().where("clip = 'c1'").select(["clip", "idx", "img"]).read_blobs()
self.assertEqual({"img"}, set(blobs2))
# duplicate columns are de-duplicated
_, blobs3 = obs.scan().where("clip = 'c1'").read_blobs(["img", "img"])
self.assertEqual({"img"}, set(blobs3))
def test_scan_read_and_stream_map_blobs(self):
map_blob_type = pa.map_(pa.string(), pa.large_binary())
schema = _schema({
"id": pa.int32(),
"preview": pa.large_binary(),
"assets": map_blob_type,
})
obs = self.conn.create_table(
"map_blobs",
schema=schema,
options=_PARQUET_OPTIONS,
)
obs.add(pa.Table.from_pydict({
"id": [1, 2, 3, 4],
"preview": [b"preview-1", None, b"preview-3", b"preview-4"],
"assets": pa.array(
[
[
("thumb", b"thumb-1"),
("empty", b""),
("missing", None),
],
None,
[],
[("original", b"original-4")],
],
type=map_blob_type,
),
}, schema=schema))
scalar, blobs = obs.scan().read_blobs(
["preview", "assets"], parallelism=2
)
ids = scalar.column("id").to_pylist()
self.assertNotIn("assets", scalar.column_names)
self.assertEqual(
{
1: b"preview-1",
2: None,
3: b"preview-3",
4: b"preview-4",
},
dict(zip(ids, blobs["preview"])),
)
self.assertEqual(
{
1: {"thumb": b"thumb-1", "empty": b"", "missing": None},
2: None,
3: {},
4: {"original": b"original-4"},
},
{
row_id: None if values is None else dict(values)
for row_id, values in zip(ids, blobs["assets"])
},
)
streamed = {}
for scalar_batch, blob_batch in obs.scan().stream_blobs("assets"):
for row_id, values in zip(
scalar_batch.column("id").to_pylist(),
blob_batch["assets"]):
streamed[row_id] = None if values is None else dict(values)
self.assertEqual(
{
1: {"thumb": b"thumb-1", "empty": b"", "missing": None},
2: None,
3: {},
4: {"original": b"original-4"},
},
streamed,
)
_, selected = obs.scan().select(["id", "assets"]).read_blobs()
self.assertEqual({"assets"}, set(selected))
def test_scan_read_and_stream_array_blobs(self):
schema = _schema({
"id": pa.int32(),
"preview": pa.large_binary(),
"pages": pa.list_(pa.large_binary()),
"assets": pa.map_(pa.string(), pa.large_binary()),
})
obs = self.conn.create_table(
"array_blobs", schema=schema,
options=dict(_PARQUET_OPTIONS, **{"read.batch-size": "2"}))
data = {
"id": [1, 2, 3, 4],
"preview": [b"p1", None, b"p3", b"p4"],
"pages": [[b"first", None, b"", b"last", b"first"], None, [], [b"fourth"]],
"assets": [[("cover", b"c1")], None, [], [("missing", None)]],
}
obs.add(pa.Table.from_pydict(data, schema=schema))
expected = dict(zip(data["id"], data["pages"]))
scalar, blobs = obs.scan().read_blobs()
self.assertEqual(["id"], scalar.column_names)
self.assertEqual({"preview", "pages", "assets"}, set(blobs))
for name in blobs:
self.assertEqual(dict(zip(data["id"], data[name])),
dict(zip(scalar.column("id").to_pylist(), blobs[name])))
streamed = {}
for scalar_batch, blob_batch in obs.scan().stream_blobs("pages"):
self.assertEqual(["id"], scalar_batch.schema.names)
self.assertLessEqual(scalar_batch.num_rows, 2)
streamed.update(zip(scalar_batch.column("id").to_pylist(), blob_batch["pages"]))
self.assertEqual(expected, streamed)
# Projection hides the filter column and other BLOBs, but keeps row IDs.
query = obs.scan().select(["pages"]).with_row_id().where("id = 4").limit(1)
for scalar, blobs in [query.read_blobs(), *list(query.stream_blobs())]:
self.assertEqual(["_ROW_ID"], scalar.schema.names)
self.assertEqual(1, scalar.num_rows)
self.assertEqual({"pages": [[b"fourth"]]}, blobs)
scalar, blobs = obs.scan().where("id < 0").read_blobs("pages")
self.assertEqual(0, scalar.num_rows)
self.assertEqual({"pages": []}, blobs)
self.assertEqual([], list(obs.scan().where("id < 0").stream_blobs("pages")))
_, blobs = obs.scan().read_blobs(["pages", "pages"])
self.assertEqual({"pages"}, set(blobs))
scalar, _ = obs.scan().read_blobs("preview")
self.assertEqual(["id"], scalar.column_names)
def test_scan_read_blobs_filter_column_not_selected(self):
# The row filter must apply even when its column is not in select().
obs = self.conn.create_table(
"obs",
schema=_schema({
"clip": pa.string(),
"idx": pa.int32(),
"image": pa.large_binary(),
}),
options=_PARQUET_OPTIONS,
partitioned=["clip"],
)
obs.add([
{"clip": "c1", "idx": i, "image": ("v-%d" % i).encode()}
for i in range(4)
])
scalar, blobs = (
obs.scan().select(["image"]).where("idx = 1").read_blobs("image"))
self.assertEqual([b"v-1"], blobs["image"])
self.assertNotIn("idx", scalar.schema.names)
scalar2, blobs2 = (
obs.scan().select(["idx", "image"]).where("idx = 2").read_blobs("image"))
self.assertEqual([b"v-2"], blobs2["image"])
self.assertEqual([2], scalar2.column("idx").to_pylist())
streamed = []
for _, body in (
obs.scan().select(["image"]).where("idx = 1").stream_blobs("image")):
streamed.extend(body["image"])
self.assertEqual([b"v-1"], streamed)
scalar3, _ = obs.scan().select(["missing", "image"]).read_blobs("image")
self.assertNotIn("missing", scalar3.schema.names)
def test_scan_read_blobs_with_row_id(self):
# with_row_id() must expose _ROW_ID in the scalar table, like to_arrow().
obs = self.conn.create_table(
"obs",
schema=_schema({
"clip": pa.string(),
"idx": pa.int32(),
"image": pa.large_binary(),
}),
options=_PARQUET_OPTIONS,
partitioned=["clip"],
)
obs.add([
{"clip": "c1", "idx": i, "image": ("v-%d" % i).encode()}
for i in range(3)
])
row_id = "_ROW_ID"
expected = {r["idx"]: r[row_id]
for r in obs.scan().with_row_id().to_arrow().to_pylist()}
scalar, blobs = (
obs.scan().with_row_id().select(["idx", "image"]).read_blobs("image"))
self.assertIn(row_id, scalar.schema.names)
self.assertEqual(expected, dict(zip(scalar.column("idx").to_pylist(),
scalar.column(row_id).to_pylist())))
got = {}
for sb, _ in obs.scan().with_row_id().select(["idx", "image"]).stream_blobs("image"):
self.assertIn(row_id, sb.schema.names)
got.update(zip(sb.column("idx").to_pylist(), sb.column(row_id).to_pylist()))
self.assertEqual(expected, got)
def test_scan_read_blobs_where_on_blob_column(self):
# A where() on one BLOB column must still filter when a different BLOB is
# read (the predicate BLOB column is read as descriptor for filtering).
obs = self.conn.create_table(
"obs",
schema=_schema({
"clip": pa.string(),
"idx": pa.int32(),
"imga": pa.large_binary(),
"imgb": pa.large_binary(),
}),
options=_PARQUET_OPTIONS,
partitioned=["clip"],
)
obs.add([
{"clip": "c1", "idx": 0, "imga": b"a0", "imgb": b"b0"},
{"clip": "c1", "idx": 1, "imga": None, "imgb": b"b1"},
{"clip": "c1", "idx": 2, "imga": b"a2", "imgb": b"b2"},
])
scalar, blobs = obs.scan().where("imga IS NOT NULL").read_blobs("imgb")
self.assertEqual([b"b0", b"b2"], blobs["imgb"])
self.assertNotIn("imga", scalar.schema.names)
def test_search_query_rejects_blob_reads(self):
t = self.conn.create_table(
"srch",
schema=_schema({
"id": pa.int32(),
"emb": _vector(3),
"img": pa.large_binary(),
}),
options=_PARQUET_OPTIONS,
)
with self.assertRaisesRegex(TypeError, "only supported on scan"):
t.search([1.0, 0.0, 0.0], column="emb").read_blobs("img")
with self.assertRaisesRegex(TypeError, "only supported on scan"):
t.search([1.0, 0.0, 0.0], column="emb").stream_blobs("img")
def test_scan_stream_blobs(self):
obs = self.conn.create_table(
"obs",
schema=_schema({
"clip": pa.string(),
"idx": pa.int32(),
"image": pa.large_binary(),
}),
options=_PARQUET_OPTIONS,
partitioned=["clip"],
)
payloads = {i: ("s-%d-" % i).encode() * (i + 1) for i in range(5)}
obs.add([
{"clip": "c1", "idx": i, "image": payloads[i]} for i in range(5)
])
# collecting every streamed batch must reproduce the full, row-aligned result
got = {}
batches = 0
for scalar, blobs in obs.scan().where("clip = 'c1'").stream_blobs("image"):
batches += 1
self.assertNotIn("image", scalar.schema.names)
for idx, img in zip(scalar.column("idx").to_pylist(), blobs["image"]):
got[idx] = img
self.assertGreaterEqual(batches, 1)
self.assertEqual({i: payloads[i] for i in range(5)}, got)
# bad column is rejected eagerly (not deferred to first iteration)
with self.assertRaisesRegex(ValueError, "not a BLOB column"):
obs.scan().stream_blobs("idx")
# breaking out early closes the reader without error
it = obs.scan().where("clip = 'c1'").stream_blobs("image")
next(it)
it.close()
# empty result streams nothing
self.assertEqual(
[], list(obs.scan().where("clip = 'none'").stream_blobs("image")))
def test_scan_stream_blobs_closes_underlying_iterator(self):
table = self.conn.create_table(
"stream_cleanup", schema=_schema({
"id": pa.int32(), "image": pa.large_binary(),
}), options=dict(_PARQUET_OPTIONS, **{"read.batch-size": "1"}))
table.add([{"id": i, "image": b"body"} for i in range(3)])
original = PythonTableRead._arrow_batch_generator
closed = []
def tracked_read(read, *args, **kwargs):
reader = original(read, *args, **kwargs)
try:
yield from reader
finally:
reader.close()
closed.append(True)
with patch("pypaimon.read.read_dispatch._native_read_execution", return_value=None), \
patch.object(PythonTableRead, "_arrow_batch_generator", tracked_read):
stream = table.scan().stream_blobs("image")
next(stream)
stream.close()
self.assertEqual([True], closed)
for method in ("_fetch_bodies", "_scalar_columns"):
with self.subTest(method=method):
closed.clear()
query = table.scan()
with patch.object(query, method, side_effect=ValueError("batch failed")):
stream = query.stream_blobs("image")
try:
next(stream)
except ValueError as error:
self.assertEqual("batch failed", str(error))
# Assert while the traceback still retains the reader's frame.
self.assertEqual([True], closed)
else:
self.fail("Expected batch failure")
@unittest.skipIf(ray is None, "ray is not installed")
def test_scan_to_ray_map_with_blobs(self):
started_ray = False
if not ray.is_initialized():
ray.init(ignore_reinit_error=True, num_cpus=2)
started_ray = True
obs = self.conn.create_table(
"ray_obs",
schema=_schema({
"clip": pa.string(),
"idx": pa.int32(),
"image": pa.large_binary(),
"audio": pa.large_binary(),
}),
options=_PARQUET_OPTIONS,
partitioned=["clip"],
)
obs.add([
{"clip": "c1", "idx": 0, "image": b"img-0", "audio": b"aud-0"},
{"clip": "c1", "idx": 1, "image": b"img-1", "audio": b"aud-1"},
{"clip": "c2", "idx": 0, "image": b"img-x", "audio": b"aud-x"},
])
def collect_batch(scalar, blobs, prefix):
assert isinstance(scalar, pa.Table)
assert ["idx"] == scalar.column_names
idxs = scalar.column("idx").to_pylist()
rows = []
for idx, image in zip(idxs, blobs["image"]):
rows.append({"idx": idx, "image": prefix + image})
return pa.Table.from_pylist(rows)
try:
from pypaimon.ray import map_with_blobs
ds = (
obs.scan()
.where("clip = 'c1'")
.select(["idx", "image", "audio"])
.to_ray(concurrency=1, override_num_blocks=1)
)
ds = ds.filter(lambda row: row["idx"] >= 0)
with self.assertRaisesRegex(ValueError, "not a BLOB column"):
obs.map_with_blobs(ds, ["idx"], collect_batch)
with self.assertRaisesRegex(ValueError, "FileIO"):
map_with_blobs(ds, ["image"], collect_batch)
from pypaimon.table.row.blob import BlobDescriptor
descriptor = BlobDescriptor("oss://bucket/blob", 0, 1).serialize()
foreign_ds = ray.data.from_arrow(pa.table({
"idx": [0],
"image": [descriptor],
"foreign_blob": [descriptor],
}))
with self.assertRaisesRegex(Exception, "does not own"):
map_with_blobs(
foreign_ds,
["image"],
collect_batch,
file_io=obs.raw_table.file_io,
all_blob_columns=["image"],
batch_size=1,
).take_all()
result = obs.map_with_blobs(
ds,
["image"],
collect_batch,
parallelism=2,
batch_size=1,
blob_uri_affinity=True,
prefetch_bytes=32,
fn_kwargs={"prefix": b"got-"},
ray_remote_args={"num_cpus": 1},
)
rows = sorted(result.take_all(), key=lambda row: row["idx"])
self.assertEqual(
[
{"idx": 0, "image": b"got-img-0"},
{"idx": 1, "image": b"got-img-1"},
],
rows,
)
finally:
if started_ray:
ray.shutdown()
@unittest.skipIf(ray is None, "ray is not installed")
def test_scan_to_ray_map_with_nested_blobs(self):
from pypaimon.ray import map_with_blobs
schema = _schema({
"id": pa.int32(), "image": pa.large_binary(),
"pages": pa.list_(pa.large_binary()),
"assets": pa.map_(pa.string(), pa.large_binary()),
})
table = self.conn.create_table(
"ray_nested_blobs", schema=schema, options=_PARQUET_OPTIONS)
rows = [
{"id": 1, "image": b"preview", "pages": [b"a", None, b"", b"b"],
"assets": [("thumb", b"thumb"), ("missing", None), ("empty", b"")]},
{"id": 2, "image": None, "pages": None, "assets": None},
{"id": 3, "image": b"preview", "pages": [], "assets": []},
{"id": 4, "image": b"preview", "pages": [b"d"], "assets": [("thumb", b"last")]},
]
table.add(pa.Table.from_pylist(rows, schema=schema))
output_schema = pa.schema([schema.field(name) for name in ("id", "pages", "assets")])
def collect(scalar, blobs):
assert scalar.column_names == ["id"]
return pa.Table.from_pydict(dict(id=scalar["id"], **blobs), schema=output_schema)
started_ray = not ray.is_initialized()
if started_ray:
ray.init(ignore_reinit_error=True, num_cpus=2)
try:
dataset = table.scan().to_ray(concurrency=1, override_num_blocks=1)
for maps, arrays in ((["assets"], ["assets"]), (["unknown"], [])):
with self.assertRaisesRegex(ValueError, "disjoint subsets"):
map_with_blobs(dataset, ["pages", "assets"], collect,
map_blob_columns=maps, array_blob_columns=arrays)
expected = [{name: row[name] for name in output_schema.names} for row in rows]
for use_table in (False, True):
with self.subTest(use_table=use_table):
if use_table:
# Ray transformations do not preserve the source metadata.
result = table.map_with_blobs(
dataset.filter(lambda row: row["id"] > 0),
["pages", "assets"], collect, batch_size=2)
else:
result = map_with_blobs(
dataset, ["pages", "assets"], collect, batch_size=2)
actual = [row for batch in result.iter_batches(batch_format="pyarrow")
for row in batch.to_pylist()]
self.assertEqual(expected, sorted(actual, key=lambda row: row["id"]))
projected = table.scan().select(["id", "pages"]).to_ray()
result = table.map_with_blobs(
projected, "pages", lambda scalar, blobs: scalar, batch_size=2)
self.assertEqual([1, 2, 3, 4], sorted(row["id"] for row in result.take_all()))
empty = table.scan().where("id < 0").to_ray()
self.assertEqual([], map_with_blobs(empty, ["pages", "assets"], collect).take_all())
finally:
if started_ray:
ray.shutdown()
def test_map_with_blobs_rejects_foreign_nested_descriptors(self):
from pypaimon.filesystem.local_file_io import LocalFileIO
from pypaimon.ray.ray_paimon import _map_blob_batch
from pypaimon.table.row.blob import BlobDescriptor
descriptor = BlobDescriptor("oss://other-table/blob", 0, 1).serialize()
for arrow_type, values in (
(pa.list_(pa.large_binary()), [None, [], [b"inline", descriptor]]),
(pa.map_(pa.string(), pa.large_binary()),
[None, [], [("inline", b"inline"), ("foreign", descriptor)]]),
(pa.list_(pa.list_(pa.string())),
[None, [], [("inline", "inline"), ("foreign", descriptor.decode("utf-8"))]])):
with self.subTest(arrow_type=arrow_type):
batch = pa.table({
"image": [b"inline"] * 3,
"text": [descriptor.decode("utf-8")] * 3,
"foreign": pa.array(values, type=arrow_type),
})
with self.assertRaisesRegex(ValueError, "does not own"):
_map_blob_batch(batch, None, ["image"], ["image"], 1, None, {})
# A descriptor outside the current slice is not a foreign column value.
_map_blob_batch(batch.slice(0, 2), LocalFileIO(), ["image"], ["image"], 1,
lambda scalar, blobs: scalar, {})
empty = pa.table({
"image": pa.chunked_array([], type=pa.large_binary()),
"foreign": pa.chunked_array([], type=arrow_type),
})
result = _map_blob_batch(empty, LocalFileIO(), ["image"], ["image"], 1,
lambda scalar, blobs: scalar, {})
self.assertEqual(0, result.num_rows)
def test_map_with_blobs_restores_utf8_map_values(self):
from unittest.mock import Mock
from pypaimon.ray.ray_paimon import _map_blob_batch
from pypaimon.table.row.blob import BlobDescriptor
descriptor = BlobDescriptor("test://blob", 0, 4).serialize()
entries = [("blob", descriptor.decode("utf-8")), ("text", "\u4f60\u597d"),
("empty", ""), ("missing", None), ("blob", descriptor.decode("utf-8"))]
batch = pa.table({"assets": pa.array([None, [], entries],
type=pa.list_(pa.list_(pa.string())))})
file_io = Mock()
file_io.read_ranges_coalesced.return_value = [b"body", None, None, None, b"body"]
actual = _map_blob_batch(
batch, file_io, ["assets"], ["assets"], 2, lambda scalar, blobs: blobs, {},
map_blob_cols=["assets"])
self.assertEqual({"assets": [None, [], [
("blob", b"body"), ("text", "\u4f60\u597d".encode("utf-8")),
("empty", b""), ("missing", None), ("blob", b"body"),
]]}, actual)
file_io.read_ranges_coalesced.assert_called_once_with(
[("test://blob", 0, 4), None, None, None, ("test://blob", 0, 4)], 2)
@unittest.skipIf(ray is None, "ray is not installed")
def test_ray_filter_preserves_map_blob_bytes(self):
from pypaimon.filesystem.local_file_io import LocalFileIO
from pypaimon.ray import map_with_blobs
schema = _schema({"id": pa.int32(), "assets": pa.map_(pa.string(), pa.large_binary())})
def collect(scalar, blobs):
return pa.Table.from_pydict(dict(id=scalar["id"], **blobs), schema=schema)
started_ray = not ray.is_initialized()
if started_ray:
ray.init(ignore_reinit_error=True, num_cpus=2)
try:
# Exercise string inference and object fallback independently of descriptor URI length.
for payload in ("\u4f60\u597d\0".encode("utf-8"), b"\xff\0"):
with self.subTest(payload=payload):
rows = [
{"id": 1, "assets": None}, {"id": 2, "assets": []},
{"id": 3, "assets": [("body", payload), ("empty", b""), ("missing", None)]},
]
dataset = ray.data.from_arrow(pa.Table.from_pylist(rows, schema=schema))
result = map_with_blobs(
dataset.filter(lambda row: row["id"] > 0), "assets", collect,
file_io=LocalFileIO(), all_blob_columns=["assets"], map_blob_columns=["assets"],
batch_size=2)
actual = [row for batch in result.iter_batches(batch_format="pyarrow")
for row in batch.to_pylist()]
self.assertEqual(rows, sorted(actual, key=lambda row: row["id"]))
finally:
if started_ray:
ray.shutdown()
@unittest.skipIf(ray is None, "ray is not installed")
def test_map_with_blobs_python_object_columns(self):
from pypaimon.filesystem.local_file_io import LocalFileIO
from pypaimon.ray.ray_paimon import _map_blob_batch
from pypaimon.table.row.blob import BlobDescriptor
try:
from ray.data.extensions import ArrowPythonObjectArray
except ImportError:
self.skipTest("Ray has no Python object extension type")
expected = {
"assets": [None, [], [("image", b"body"), ("missing", None)]],
"pages": [None, [], [b"page", None]],
}
batch = pa.table({name: ArrowPythonObjectArray.from_objects(values)
for name, values in expected.items()})
actual = _map_blob_batch(
batch, LocalFileIO(), ["assets", "pages"], ["assets", "pages"], 1,
lambda scalar, blobs: blobs, {}, map_blob_cols=["assets"], array_blob_cols=["pages"])
self.assertEqual(expected, actual)
foreign = ArrowPythonObjectArray.from_objects([
None, [], [("inline", b"body"), ("foreign", BlobDescriptor("oss://other/blob", 0, 1).serialize())],
])
with self.assertRaisesRegex(ValueError, "does not own"):
_map_blob_batch(batch.append_column("foreign", foreign), None,
["assets"], ["assets", "pages"], 1, None, {}, map_blob_cols=["assets"])
@unittest.skipIf(ray is None, "ray is not installed")
def test_scan_to_ray_map_with_blobs_guards(self):
started_ray = False
if not ray.is_initialized():
ray.init(ignore_reinit_error=True, num_cpus=2)
started_ray = True
obs = self.conn.create_table(
"ray_obs_guards",
schema=_schema({
"clip": pa.string(),
"idx": pa.int32(),
"image": pa.large_binary(),
}),
options=_PARQUET_OPTIONS,
partitioned=["clip"],
)
obs.add([{"clip": "c1", "idx": 0, "image": b"img-0"}])
def return_none(scalar, blobs):
return None
try:
from pypaimon.ray import map_with_blobs
ds = (
obs.scan()
.where("clip = 'c1'")
.select(["idx", "image"])
.to_ray(concurrency=1, override_num_blocks=1)
)
with self.assertRaisesRegex(ValueError, "all_blob_columns"):
map_with_blobs(
ray.data.from_arrow(pa.table({"image": [b"inline"]})),
["image"],
lambda scalar, blobs: pa.table({"rows": [scalar.num_rows]}),
file_io=obs.raw_table.file_io,
)
with self.assertRaisesRegex(ValueError, "requires batch_size"):
obs.map_with_blobs(
ds,
["image"],
return_none,
batch_size=None,
blob_uri_affinity=True,
)
with self.assertRaisesRegex(ValueError, "prefetch_bytes"):
obs.map_with_blobs(
ds,
["image"],
return_none,
prefetch_bytes=0,
)
with self.assertRaisesRegex(Exception, "must return"):
obs.map_with_blobs(
ds,
["image"],
return_none,
batch_size=1,
).take_all()
empty_ds = (
obs.scan()
.where("clip = 'none'")
.select(["idx", "image"])
.to_ray(concurrency=1, override_num_blocks=1)
)
result = obs.map_with_blobs(
empty_ds,
["image"],
lambda scalar, blobs: pa.table({"rows": [scalar.num_rows]}),
batch_size=1,
)
self.assertEqual([], result.take_all())
finally:
if started_ray:
ray.shutdown()
def test_scan_to_ray_nested_projection_output_names(self):
obs = self.conn.create_table(
"ray_nested_obs",
schema=_schema({
"id": pa.int32(),
"tag": pa.string(),
"payload": pa.struct([("a", pa.int64()), ("b", pa.string())]),
"image": pa.large_binary(),
}),
options=_PARQUET_OPTIONS,
)
_, _, visible_columns = (
obs.scan()
.where("tag = 'keep'")
.select(["id", "payload.a"])
._blob_descriptor_query_read_builder()
)
self.assertEqual(["id", "payload_a"], visible_columns)
def test_fetch_bodies_decodes_descriptor_inline_and_null(self):
# Cells may be descriptor bytes (incl. -1 read-to-EOF), inline bytes, or null.
from pypaimon.multimodal.query import ScanQuery
from pypaimon.common.file_io import FileIO
from pypaimon.table.row.blob import BlobDescriptor
data = bytes(range(64))
path = os.path.join(self.temp_dir, "blob.bin")
with open(path, "wb") as f:
f.write(data)
file_io = FileIO.get("file://" + self.temp_dir, {})
cells = [
BlobDescriptor(path, 0, -1).serialize(),
BlobDescriptor(path, 10, 8).serialize(),
b"raw-inline-bytes",
None,
]
bodies = ScanQuery._fetch_bodies(
file_io, {"image": cells}, ["image"], 4)
self.assertEqual(
[data, data[10:18], b"raw-inline-bytes", None], bodies["image"])
def test_fetch_bodies_reads_via_file_io_not_uri_reader(self):
# Blobs must be read through file_io.read_ranges_coalesced (which carries the
# resolved DLF/OSS token), never via uri_reader_factory -- that would rebuild a
# FileIO from the raw catalog options and fail in DLF data-token mode.
from pypaimon.multimodal.query import ScanQuery
from pypaimon.table.row.blob import BlobDescriptor
class _FakeIO:
@property
def uri_reader_factory(self):
raise AssertionError("_fetch_bodies must not use uri_reader_factory")
def read_ranges_coalesced(self, ranges, parallelism):
return [None if r is None else b"BODY:" + r[0].encode() for r in ranges]
cells = [BlobDescriptor("oss://bucket/x", 4, 10).serialize(), b"inline-blob", None]
bodies = ScanQuery._fetch_bodies(_FakeIO(), {"img": cells}, ["img"], 8)
self.assertEqual([b"BODY:oss://bucket/x", b"inline-blob", None], bodies["img"])
def test_fetch_bodies_coalesces_mixed_array_map_and_scalar_blobs(self):
from pypaimon.multimodal.blob_read import fetch_blob_bodies
from pypaimon.common.file_io import FileIO
from pypaimon.table.row.blob import BlobDescriptor
path = os.path.join(self.temp_dir, "array-blob.bin")
with open(path, "wb") as stream:
stream.write(b"0123456789")
file_io = FileIO.get("file://" + self.temp_dir, {})
descriptor = BlobDescriptor(path, 2, 3).serialize()
cells = {
"image": [descriptor],
"pages": [[descriptor, None, b"", b"inline", descriptor], [], None],
"assets": [[("cover", descriptor), ("missing", None)]],
}
with patch.object(file_io, "read_ranges_coalesced",
wraps=file_io.read_ranges_coalesced) as read:
bodies = fetch_blob_bodies(
file_io, cells, list(cells), 2,
map_blob_cols=["assets"], array_blob_cols=["pages"])
self.assertEqual(1, read.call_count)
self.assertEqual({
"image": [b"234"],
"pages": [[b"234", None, b"", b"inline", b"234"], [], None],
"assets": [[("cover", b"234"), ("missing", None)]],
}, bodies)
def test_fetch_bodies_rejects_unresolved_blob_view(self):
from pypaimon.multimodal.query import ScanQuery
from pypaimon.table.row.blob import BlobViewStruct
view_bytes = BlobViewStruct("db.tbl", 1, 2).serialize()
with self.assertRaisesRegex(ValueError, "blob-view"):
ScanQuery._fetch_bodies(None, {"image": [view_bytes]}, ["image"], 4)
def test_scan_does_not_expose_pre_filter(self):
users = self.conn.create_table(
"users",
schema=_schema({"id": pa.int32()}),
options=_PARQUET_OPTIONS,
)
self.assertFalse(hasattr(users.scan(), "pre_filter"))
def test_where_rejects_predicate_object(self):
users = self.conn.create_table(
"users",
schema=_schema({"id": pa.int32()}),
options=_PARQUET_OPTIONS,
)
predicate = PredicateBuilder(users.raw_table.fields).equal("id", 1)
with self.assertRaisesRegex(ValueError, "SQL-like string"):
users.scan().where(predicate)
def test_update_by_filter(self):
users = self.conn.create_table(
"users",
data=[
{"id": 1, "name": "Alice", "age": 30},
{"id": 2, "name": "Bob", "age": 25},
],
schema=_schema({
"id": pa.int32(),
"name": pa.string(),
"age": pa.int32(),
}),
options=_PARQUET_OPTIONS,
)
users.update(where="id = 2", values={"age": 26})
rows = sorted(users.scan().to_list(), key=lambda r: r["id"])
self.assertEqual(
[
{"id": 1, "name": "Alice", "age": 30},
{"id": 2, "name": "Bob", "age": 26},
],
rows,
)
def test_update_by_row_id_preserves_duplicate_keys_and_media(self):
users = self.conn.create_table(
"users", schema=_schema({"id": pa.int32(), "label": pa.string(), "image": pa.large_binary()}),
data=[{"id": 7, "label": "first", "image": b"first"},
{"id": 7, "label": "second", "image": b"second"}],
options=_PARQUET_OPTIONS)
row_ids = {row["label"]: row["_ROW_ID"]
for row in users.scan().select(["label"]).with_row_id().to_list()}
files = {file.file_name for split in users.raw_table.new_read_builder().new_scan().plan_all().splits()
for file in split.files}
with patch.object(users.raw_table.file_io, "read_ranges_coalesced",
side_effect=AssertionError("must not fetch media")):
result = users.update(where="_ROW_ID = %s" % row_ids["second"], values={"label": "car"})
self.assertIs(users, result)
self.assertEqual(
{row_ids["first"]: "first", row_ids["second"]: "car"},
{row["_ROW_ID"]: row["label"]
for row in users.scan().select(["label"]).with_row_id().to_list()})
current_files = {file.file_name: file
for split in users.raw_table.new_read_builder().new_scan().plan_all().splits()
for file in split.files}
self.assertTrue(files.issubset(current_files))
self.assertTrue(all(file.write_cols == ["label"] for name, file in current_files.items() if name not in files))
def test_update_by_row_id_conditions_skip_missing_and_deleted_rows(self):
users = self.conn.create_table(
"users", schema=_schema({"part": pa.string(), "id": pa.int32(), "label": pa.string()}),
partitioned=["part"],
data=[{"part": "a" if i < 3 else "b", "id": i, "label": "old"} for i in [3, 4, 5, 0, 1, 2]],
options=_PARQUET_OPTIONS)
row_ids = {row["id"]: row["_ROW_ID"] for row in users.scan().select(["id"]).with_row_id().to_list()}
users.delete("id = 1")
users.update(
where="_ROW_ID IN (%s, %s, 999999) AND id = 0" % (row_ids[0], row_ids[1]),
values={"label": "selected"})
users.update(where="_ROW_ID = %s OR id = 2" % row_ids[5], values={"label": "either"})
users.update(
where="_ROW_ID >= %s AND _ROW_ID <= %s" % (row_ids[3], row_ids[3]), values={"label": None})
users.update(where="_ROW_ID BETWEEN %s AND %s" % (row_ids[4], row_ids[4]), values={"label": "range"})
expected = [
{"id": 0, "label": "selected"}, {"id": 2, "label": "either"},
{"id": 3, "label": None}, {"id": 4, "label": "range"}, {"id": 5, "label": "either"}]
self.assertEqual(expected, sorted(users.scan().select(["id", "label"]).to_list(), key=lambda row: row["id"]))
for row_id in (row_ids[1], 999999):
with self.subTest(row_id=row_id):
self.assertIs(users.update(where="_ROW_ID = %s" % row_id, values={"label": "absent"}), users)
self.assertEqual(expected, sorted(users.scan().select(["id", "label"]).to_list(), key=lambda row: row["id"]))
users.update(
where="_ROW_ID = %s OR _ROW_ID = %s AND id = 999" % (row_ids[0], row_ids[3]),
values={"label": "precedence"})
expected[0]["label"] = "precedence"
users.update(
where="(_ROW_ID = %s OR _ROW_ID = %s) AND id = 999" % (row_ids[0], row_ids[3]),
values={"label": "parenthesized"})
self.assertEqual(expected, sorted(users.scan().select(["id", "label"]).to_list(), key=lambda row: row["id"]))
def test_update_by_row_id_on_empty_and_fully_deleted_table(self):
users = self.conn.create_table(
"users", schema=_schema({"id": pa.int32(), "label": pa.string()}), options=_PARQUET_OPTIONS)
for deleted in (False, True):
with self.subTest(deleted=deleted):
if deleted:
users.add([{"id": 7, "label": "old"}])
users.delete("id = 7")
self.assertIs(users.update(where="_ROW_ID >= 0", values={"label": "new"}), users)
self.assertEqual([], users.scan().to_list())
def test_update_by_row_id_with_inverted_between(self):
users = self.conn.create_table(
"users", schema=_schema({"id": pa.int32(), "label": pa.string()}),
data=[{"id": i, "label": "old"} for i in range(3)], options=_PARQUET_OPTIONS)
row_ids = {row["id"]: row["_ROW_ID"] for row in users.scan().select(["id"]).with_row_id().to_list()}
lower, upper = max(row_ids.values()), min(row_ids.values())
users.update(
where="_ROW_ID BETWEEN %s AND %s OR _ROW_ID = %s" % (lower, upper, row_ids[1]),
values={"label": "new"})
expected = [{"id": i, "label": "new" if i == 1 else "old"} for i in range(3)]
self.assertEqual(expected, users.scan().to_list())
users.update(where="_ROW_ID BETWEEN %s AND %s" % (lower, upper), values={"label": "empty"})
self.assertEqual(expected, users.scan().to_list())
def test_update_row_id_predicates_keep_system_columns_read_only(self):
users = self.conn.create_table(
"users", schema=_schema({"id": pa.int32(), "label": pa.string()}),
data=[{"id": 7, "label": "old"}], options=_PARQUET_OPTIONS)
row_id = users.scan().select(["id"]).with_row_id().to_list()[0]["_ROW_ID"]
for field in ("_SEQUENCE_NUMBER", "_VALUE_KIND"):
with self.subTest(field=field), self.assertRaisesRegex(ValueError, "Unknown field"):
users.update(where="%s = 0" % field, values={"label": "new"})
with self.subTest(assignment=field), self.assertRaisesRegex(ValueError, "not in table schema"):
users.update(where="_ROW_ID = %s" % row_id, values={field: 0})
with self.assertRaisesRegex(ValueError, "not in table schema"):
users.update(where="_ROW_ID = %s" % row_id, values={"_ROW_ID": 999})
with self.assertRaises(ValueError):
users.update(where="_ROW_ID = 'invalid'", values={"label": "new"})
self.assertEqual([{"id": 7, "label": "old"}], users.scan().to_list())
self.assertIs(users.update(where=None, values={"label": "all"}), users)
self.assertEqual([{"id": 7, "label": "all"}], users.scan().to_list())
with self.assertRaisesRegex(ValueError, "Unknown field '_ROW_ID'"):
users.copy({"row-tracking.enabled": "false"}).update(
where="_ROW_ID = %s" % row_id, values={"label": "new"})
@pytest.mark.native_plan
def test_update_by_row_id_uses_native_predicate_update(self):
from pypaimon.execution.python.write.table_update import PythonTableUpdate
users = self.conn.create_table(
"users", schema=_schema({"id": pa.int32(), "label": pa.string()}),
data=[{"id": 7, "label": "first"}, {"id": 7, "label": "second"}],
options=dict(_PARQUET_OPTIONS, **{"write.native.enabled": "true"}))
row_ids = {row["label"]: row["_ROW_ID"]
for row in users.scan().select(["label"]).with_row_id().to_list()}
with patch.object(PythonTableUpdate, "_execute_predicate_update",
side_effect=AssertionError("must use native predicate update")):
users.update(where="_ROW_ID = %s" % row_ids["second"], values={"label": "car"})
self.assertEqual(
{row_ids["first"]: "first", row_ids["second"]: "car"},
{row["_ROW_ID"]: row["label"]
for row in users.scan().select(["label"]).with_row_id().to_list()})
def test_update_by_row_id_keeps_the_existing_backend_fallback_boundary(self):
from unittest.mock import Mock
users = self.conn.create_table(
"users", schema=_schema({"id": pa.int32(), "label": pa.string()}),
data=[{"id": 7, "label": "first"}, {"id": 7, "label": "second"}], options=_PARQUET_OPTIONS)
row_ids = {row["label"]: row["_ROW_ID"]
for row in users.scan().select(["label"]).with_row_id().to_list()}
users = users.copy({"write.native.enabled": "true"})
prepare = Mock(return_value=None)
with patch("pypaimon.write.update_dispatch.prepare_native_write_operation", prepare), patch(
"pypaimon.write.backend_dispatch.prepare_native_write_operation", prepare):
users.update(where="_ROW_ID = %s" % row_ids["second"], values={"label": "fallback"})
prepare.assert_called_once()
self.assertEqual("predicate_update", prepare.call_args[0][0])
predicate = prepare.call_args[0][3]
self.assertEqual("_ROW_ID", predicate.field)
self.assertEqual([row_ids["second"]], predicate.literals)
expected = {row_ids["first"]: "first", row_ids["second"]: "fallback"}
self.assertEqual(expected, {row["_ROW_ID"]: row["label"]
for row in users.scan().select(["label"]).with_row_id().to_list()})
snapshot = users.raw_table.snapshot_manager().get_latest_snapshot().id
backend = Mock()
error = OSError("native execution failed")
backend.update.side_effect = error
with patch("pypaimon.write.update_dispatch.prepare_native_write_operation", return_value=backend), patch(
"pypaimon.write.update_dispatch.create_python_update_operation",
side_effect=AssertionError("must not replay a failed Native update")):
with self.assertRaises(OSError) as failure:
users.update(where="_ROW_ID = %s" % row_ids["second"], values={"label": "bad"})
backend.update.assert_called_once()
self.assertIs(failure.exception, error)
self.assertEqual(snapshot, users.raw_table.snapshot_manager().get_latest_snapshot().id)
self.assertEqual(expected, {row["_ROW_ID"]: row["label"]
for row in users.scan().select(["label"]).with_row_id().to_list()})
def test_delete_by_filter(self):
users = self.conn.create_table(
"users",
data=[
{"id": 1, "name": "Alice", "age": 30},
{"id": 2, "name": "Bob", "age": 25},
{"id": 3, "name": "Carol", "age": 40},
],
schema=_schema({
"id": pa.int32(),
"name": pa.string(),
"age": pa.int32(),
}),
options=_PARQUET_OPTIONS,
)
users.delete(where="id = 2")
rows = sorted(users.scan().to_list(), key=lambda r: r["id"])
self.assertEqual(
[
{"id": 1, "name": "Alice", "age": 30},
{"id": 3, "name": "Carol", "age": 40},
],
rows,
)
def test_merge_updates_matches_and_inserts_new_rows(self):
users = self.conn.create_table(
"users",
data=[
{"id": 1, "name": "Alice", "age": 30},
{"id": 2, "name": "Bob", "age": 25},
],
schema=_schema({
"id": pa.int32(),
"name": pa.string(),
"age": pa.int32(),
}),
options=_PARQUET_OPTIONS,
)
users.merge("id") \
.when_matched_update() \
.when_not_matched_insert() \
.execute([
{"id": 2, "name": "Bob_v2", "age": 26},
{"id": 3, "name": "Carol", "age": 40},
])
rows = sorted(users.scan().to_list(), key=lambda r: r["id"])
self.assertEqual(
[
{"id": 1, "name": "Alice", "age": 30},
{"id": 2, "name": "Bob_v2", "age": 26},
{"id": 3, "name": "Carol", "age": 40},
],
rows,
)
def test_merge_deletes_matched_rows(self):
users = self.conn.create_table(
"users",
data=[
{"id": 1, "name": "Alice", "age": 30},
{"id": 2, "name": "Bob", "age": 25},
{"id": 3, "name": "Carol", "age": 40},
],
schema=_schema({
"id": pa.int32(),
"name": pa.string(),
"age": pa.int32(),
}),
options=_PARQUET_OPTIONS,
)
users.merge("id") \
.when_matched_delete() \
.execute([
{"id": 2},
{"id": 3},
])
rows = sorted(users.scan().to_list(), key=lambda r: r["id"])
self.assertEqual(
[
{"id": 1, "name": "Alice", "age": 30},
],
rows,
)
def test_merge_all_uses_only_source_columns(self):
users = self.conn.create_table(
"users",
data=[
{"id": 1, "name": "Alice", "age": 30},
{"id": 2, "name": "Bob", "age": 25},
],
schema=_schema({
"id": pa.int32(),
"name": pa.string(),
"age": pa.int32(),
}),
options=_PARQUET_OPTIONS,
)
users.merge("id") \
.when_matched_update() \
.when_not_matched_insert() \
.execute([
{"id": 2, "age": 26},
{"id": 3, "age": 40},
])
rows = sorted(users.scan().to_list(), key=lambda r: r["id"])
self.assertEqual(
[
{"id": 1, "name": "Alice", "age": 30},
{"id": 2, "name": "Bob", "age": 26},
{"id": 3, "name": None, "age": 40},
],
rows,
)
def test_merge_preserves_fields_first_present_in_later_source_rows(self):
users = self.conn.create_table(
"sparse_merge",
data=[
{"id": 1, "name": "Alice", "age": 30},
{"id": 2, "name": "Bob", "age": 25},
],
schema=_schema({
"id": pa.int32(), "name": pa.string(), "age": pa.int32(),
}),
options=_PARQUET_OPTIONS,
)
source = [{"id": 1}, {"id": 2, "name": "Bobby"},
{"id": 3, "name": "Carol"}]
users.merge("id").when_matched_update() \
.when_not_matched_insert().execute(source)
self.assertEqual([
{"id": 1, "name": None, "age": 30},
{"id": 2, "name": "Bobby", "age": 25},
{"id": 3, "name": "Carol", "age": None},
], sorted(users.scan().to_list(), key=lambda row: row["id"]))
self.assertEqual({"id": 1}, source[0])
def test_merge_preserves_later_source_only_fields_with_mapped_keys(self):
users = self.conn.create_table(
"sparse_mapped_merge",
data=[{"id": 1, "name": "Alice"}],
schema=_schema({"id": pa.int32(), "name": pa.string()}),
options=_PARQUET_OPTIONS,
)
values = {"name": source_col("new_name")}
users.merge({"id": "source_id"}).when_matched_update(values) \
.when_not_matched_insert({
"id": source_col("source_id"), **values,
}).execute([
{"source_id": 1},
{"source_id": 2, "new_name": "Bob"},
])
self.assertEqual([
{"id": 1, "name": None}, {"id": 2, "name": "Bob"},
], sorted(users.scan().to_list(), key=lambda row: row["id"]))
def test_merge_supports_source_key_mapping(self):
users = self.conn.create_table(
"users",
data=[
{"id": 1, "name": "Alice", "age": 30},
{"id": 2, "name": "Bob", "age": 25},
],
schema=_schema({
"id": pa.int32(),
"name": pa.string(),
"age": pa.int32(),
}),
options=_PARQUET_OPTIONS,
)
users.merge({"id": "source_id"}) \
.when_matched_update({"age": source_col("age")}) \
.when_not_matched_insert() \
.execute([
{"source_id": 2, "name": "Bob_v2", "age": 26},
{"source_id": 3, "name": "Carol", "age": 40},
])
rows = sorted(users.scan().to_list(), key=lambda r: r["id"])
self.assertEqual(
[
{"id": 1, "name": "Alice", "age": 30},
{"id": 2, "name": "Bob", "age": 26},
{"id": 3, "name": "Carol", "age": 40},
],
rows,
)
def test_merge_where_uses_source_and_target_aliases(self):
users = self.conn.create_table(
"users",
schema=_schema({
"id": pa.int32(),
"name": pa.string(),
"age": pa.int32(),
}),
options=_PARQUET_OPTIONS,
)
calls = {}
class FakeUpdate:
def merge_into(
self,
source,
on,
when_matched=None,
when_not_matched=None):
calls["on"] = on
calls["matched"] = list(when_matched or [])
calls["not_matched"] = list(when_not_matched or [])
return []
class FakeCommit:
def commit(self, messages):
calls["messages"] = messages
def close(self):
pass
class FakeWriteBuilder:
def new_update(self):
return FakeUpdate()
def new_commit(self):
return FakeCommit()
users.raw_table.new_batch_write_builder = lambda: FakeWriteBuilder()
(
users.merge("id")
.when_matched_update(
{"age": source_col("age")},
where="source.age > target.age and source.name != 'target.name'",
)
.when_not_matched_insert(where="source.age > 0")
.execute([
{"id": 1, "name": "Alice", "age": 31},
])
)
self.assertEqual({"id": "id"}, calls["on"])
self.assertEqual(
"s.age > t.age and s.name != 'target.name'",
calls["matched"][0].condition,
)
self.assertEqual("s.age > 0", calls["not_matched"][0].condition)
self.assertEqual([], calls["messages"])
def test_merge_delete_where_uses_source_and_target_aliases(self):
users = self.conn.create_table(
"users",
schema=_schema({
"id": pa.int32(),
"name": pa.string(),
"age": pa.int32(),
}),
options=_PARQUET_OPTIONS,
)
calls = {}
class FakeUpdate:
def merge_into(
self,
source,
on,
when_matched=None,
when_not_matched=None):
calls["matched"] = list(when_matched or [])
return []
class FakeCommit:
def commit(self, messages):
calls["messages"] = messages
def close(self):
pass
class FakeWriteBuilder:
def new_update(self):
return FakeUpdate()
def new_commit(self):
return FakeCommit()
users.raw_table.new_batch_write_builder = lambda: FakeWriteBuilder()
(
users.merge("id")
.when_matched_delete(
where="source.age < target.age and source.name != 'target.name'",
)
.execute([
{"id": 1, "name": "Alice", "age": 29},
])
)
self.assertTrue(calls["matched"][0].delete)
self.assertEqual(
"s.age < t.age and s.name != 'target.name'",
calls["matched"][0].condition,
)
self.assertEqual([], calls["messages"])
def test_merge_validates_on_columns(self):
users = self.conn.create_table(
"users",
schema=_schema({"id": pa.int32(), "name": pa.string()}),
options=_PARQUET_OPTIONS,
)
with self.assertRaisesRegex(ValueError, "source columns"):
users.merge({"id": "missing"}) \
.when_matched_update() \
.execute([{"id": 1, "name": "Alice"}])
def test_create_index_normalizes_full_text_alias(self):
docs = self.conn.create_table(
"docs",
schema=_schema({
"id": pa.int32(),
"content": pa.string(),
}),
options=_PARQUET_OPTIONS,
)
calls = []
def create_global_index(column, index_type, options=None, **kwargs):
calls.append((column, index_type, options))
return index_type
docs.raw_table.create_global_index = create_global_index
options = {"tokenizer": "default"}
self.assertEqual(
"full-text",
docs.create_index("content", index_type="full-text",
options=options),
)
self.assertEqual(
[
("content", "full-text", options),
],
calls,
)
def test_create_index_requires_index_type(self):
docs = self.conn.create_table(
"docs",
schema=_schema({
"id": pa.int32(),
"content": pa.string(),
}),
options=_PARQUET_OPTIONS,
)
def create_global_index(column, index_type, options=None, **kwargs):
return index_type
docs.raw_table.create_global_index = create_global_index
with self.assertRaises(TypeError):
docs.create_index("content")
def test_table_does_not_expose_predicate_builder(self):
docs = self.conn.create_table(
"docs",
schema=_schema({"id": pa.int32()}),
options=_PARQUET_OPTIONS,
)
self.assertFalse(hasattr(docs, "predicate_builder"))
def test_search_reads_vector_matching_rows(self):
docs = self.conn.create_table(
"docs",
schema=_schema({
"id": pa.int32(),
"content": pa.string(),
"embedding": _vector(3),
}),
options=_PARQUET_OPTIONS,
)
docs.add([
{"id": 1, "content": "a", "embedding": [1.0, 0.0, 0.0]},
{"id": 2, "content": "b", "embedding": [0.0, 1.0, 0.0]},
{"id": 3, "content": "c", "embedding": [0.0, 0.0, 1.0]},
])
calls = {}
class FakeVectorBuilder:
def with_vector_column(self, column):
calls["column"] = column
return self
def with_query_vector(self, vector):
calls["vector"] = vector
return self
def with_limit(self, limit):
calls["limit"] = limit
return self
def with_options(self, options):
calls["options"] = options
return self
def with_filter(self, predicate):
calls["filter"] = predicate
return self
def execute_local(self):
return GlobalIndexResult.from_range(Range(1, 1))
patcher = patch.object(type(docs.raw_table), "new_vector_search_builder", return_value=FakeVectorBuilder())
patcher.start()
self.addCleanup(patcher.stop)
result = (
docs.search([0.0, 1.0, 0.0])
.where("id >= 1")
.limit(5)
.to_arrow()
)
self.assertEqual("embedding", calls["column"])
self.assertEqual([0.0, 1.0, 0.0], calls["vector"])
self.assertEqual(5, calls["limit"])
self.assertEqual([2], result["id"].to_pylist())
def test_search_applies_pre_filter_to_vector_builder(self):
docs = self.conn.create_table(
"docs",
schema=_schema({
"id": pa.int32(),
"category": pa.string(),
"embedding": _vector(3),
}),
options=_PARQUET_OPTIONS,
)
docs.add([
{"id": 1, "category": "lake", "embedding": [1.0, 0.0, 0.0]},
{"id": 2, "category": "city", "embedding": [0.0, 1.0, 0.0]},
])
calls = {}
class FakeVectorBuilder:
def with_vector_column(self, column):
return self
def with_query_vector(self, vector):
return self
def with_limit(self, limit):
return self
def with_options(self, options):
return self
def with_filter(self, predicate):
calls["pre_filter"] = predicate
return self
def execute_local(self):
return GlobalIndexResult.from_range(Range(0, 0))
patcher = patch.object(type(docs.raw_table), "new_vector_search_builder", return_value=FakeVectorBuilder())
patcher.start()
self.addCleanup(patcher.stop)
docs.search(
[1.0, 0.0, 0.0],
pre_filter="category = 'lake'",
).limit(1).to_list()
self.assertEqual("equal", calls["pre_filter"].method)
self.assertEqual("category", calls["pre_filter"].field)
self.assertEqual(["lake"], calls["pre_filter"].literals)
def test_searches_infer_vector_column_by_query_dimension(self):
docs = self.conn.create_table(
"docs",
schema=_schema({
"id": pa.int32(),
"image_embedding": _vector(2),
"text_embedding": _vector(3),
}),
options=_PARQUET_OPTIONS,
)
docs.add([
{
"id": 1,
"image_embedding": [1.0, 0.0],
"text_embedding": [1.0, 0.0, 0.0],
},
{
"id": 2,
"image_embedding": [0.0, 1.0],
"text_embedding": [0.0, 1.0, 0.0],
},
])
result = (
docs.search([0.0, 1.0, 0.0])
.select(["id"])
.limit(1)
.to_list()
)
batch_result = (
docs.search_vectors([[0.0, 1.0, 0.0]])
.select(["id"])
.limit(1)
.to_list()
)
self.assertEqual([{"id": 2}], result)
self.assertEqual([[{"id": 2}]], batch_result)
def test_search_reports_vector_column_dimension_errors(self):
docs = self.conn.create_table(
"docs",
schema=_schema({
"id": pa.int32(),
"content": pa.string(),
"image_embedding": _vector(2),
"text_embedding": _vector(3),
}),
options=_PARQUET_OPTIONS,
)
with self.assertRaisesRegex(
ValueError,
r"No vector column found with dimension 4;.*"
r"image_embedding\(2\).*text_embedding\(3\)"):
docs.search([1.0, 0.0, 0.0, 0.0])
with self.assertRaisesRegex(ValueError, "not found in table schema"):
docs.search([1.0, 0.0], column="missing")
with self.assertRaisesRegex(ValueError, "not a fixed-size vector"):
docs.search([1.0, 0.0], column="content")
with self.assertRaisesRegex(
ValueError,
"Vector dimension 3 does not match column "
"'image_embedding' dimension 2"):
docs.search([1.0, 0.0, 0.0], column="image_embedding")
with self.assertRaisesRegex(ValueError, "same dimension"):
docs.search_vectors([
[1.0, 0.0, 0.0],
[1.0, 0.0],
])
def test_search_requires_column_for_same_dimension_vectors(self):
docs = self.conn.create_table(
"docs",
schema=_schema({
"id": pa.int32(),
"title_embedding": _vector(2),
"body_embedding": _vector(2),
}),
options=_PARQUET_OPTIONS,
)
with self.assertRaisesRegex(
ValueError,
r"Multiple vector columns found with dimension 2: "
r"title_embedding, body_embedding; pass column\."):
docs.search([1.0, 0.0])
def test_search_pre_filter_rejects_predicate_object(self):
docs = self.conn.create_table(
"docs",
schema=_schema({
"id": pa.int32(),
"embedding": _vector(3),
}),
options=_PARQUET_OPTIONS,
)
predicate = PredicateBuilder(docs.raw_table.fields).equal("id", 1)
with self.assertRaisesRegex(ValueError, "SQL-like string"):
docs.search([1.0, 0.0, 0.0]).pre_filter(predicate)
def test_search_accepts_generator_vector(self):
docs = self.conn.create_table(
"docs",
schema=_schema({
"id": pa.int32(),
"embedding": _vector(3),
}),
options=_PARQUET_OPTIONS,
)
docs.add([{"id": 1, "embedding": [1.0, 0.0, 0.0]}])
calls = {}
class FakeVectorBuilder:
def with_vector_column(self, column):
calls["column"] = column
return self
def with_query_vector(self, vector):
calls["vector"] = vector
return self
def with_limit(self, limit):
calls["limit"] = limit
return self
def with_options(self, options):
calls["options"] = options
return self
def execute_local(self):
return GlobalIndexResult.from_range(Range(0, 0))
patcher = patch.object(type(docs.raw_table), "new_vector_search_builder", return_value=FakeVectorBuilder())
patcher.start()
self.addCleanup(patcher.stop)
result = docs.search((v for v in [1.0, 0.0, 0.0])).limit(1).to_list()
self.assertEqual("embedding", calls["column"])
self.assertEqual([1.0, 0.0, 0.0], calls["vector"])
self.assertEqual([{"id": 1, "embedding": [1.0, 0.0, 0.0]}], result)
def test_search_with_row_id_returns_system_column(self):
docs = self.conn.create_table(
"docs",
schema=_schema({
"id": pa.int32(),
"embedding": _vector(3),
}),
options=_PARQUET_OPTIONS,
)
docs.add([{"id": 1, "embedding": [1.0, 0.0, 0.0]}])
class FakeVectorBuilder:
def with_vector_column(self, column):
return self
def with_query_vector(self, vector):
return self
def with_limit(self, limit):
return self
def with_options(self, options):
return self
def execute_local(self):
return GlobalIndexResult.from_range(Range(0, 0))
patcher = patch.object(type(docs.raw_table), "new_vector_search_builder", return_value=FakeVectorBuilder())
patcher.start()
self.addCleanup(patcher.stop)
result = (
docs.search([1.0, 0.0, 0.0], column="embedding")
.select(["id"])
.with_row_id()
.limit(1)
.to_list()
)
self.assertEqual([{"id": 1, "_ROW_ID": 0}], result)
def test_search_rejects_batch_vectors(self):
docs = self.conn.create_table(
"docs",
schema=_schema({
"id": pa.int32(),
"embedding": _vector(3),
}),
options=_PARQUET_OPTIONS,
)
with self.assertRaisesRegex(ValueError, "use search_vectors"):
docs.search([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]])
def test_search_vectors_reads_one_result_set_per_query_vector(self):
docs = self.conn.create_table(
"docs",
schema=_schema({
"id": pa.int32(),
"content": pa.string(),
"embedding": _vector(3),
}),
options=_PARQUET_OPTIONS,
)
docs.add([
{"id": 1, "content": "a", "embedding": [1.0, 0.0, 0.0]},
{"id": 2, "content": "b", "embedding": [0.0, 1.0, 0.0]},
{"id": 3, "content": "c", "embedding": [0.0, 0.0, 1.0]},
])
calls = {}
class FakeBatchVectorBuilder:
def with_vector_column(self, column):
calls["column"] = column
return self
def with_query_vectors(self, vectors):
calls["vectors"] = vectors
return self
def with_limit(self, limit):
calls["limit"] = limit
return self
def with_options(self, options):
calls["options"] = options
return self
def with_filter(self, predicate):
calls["filter"] = predicate
return self
def execute_batch_local(self):
return [
GlobalIndexResult.from_range(Range(0, 0)),
GlobalIndexResult.from_range(Range(2, 2)),
]
patcher = patch.object(type(docs.raw_table), "new_batch_vector_search_builder",
return_value=FakeBatchVectorBuilder())
patcher.start()
self.addCleanup(patcher.stop)
result = (
docs.search_vectors(
[[1.0, 0.0, 0.0], [0.0, 0.0, 1.0]],
options={"nprobe": "8"},
pre_filter="content = 'a'",
)
.where("id >= 1")
.select(["id"])
.limit(2)
.to_list()
)
self.assertEqual("embedding", calls["column"])
self.assertEqual(
[[1.0, 0.0, 0.0], [0.0, 0.0, 1.0]],
calls["vectors"],
)
self.assertEqual(2, calls["limit"])
self.assertEqual({"nprobe": "8"}, calls["options"])
self.assertEqual("content", calls["filter"].field)
self.assertEqual(["a"], calls["filter"].literals)
self.assertEqual([[{"id": 1}], [{"id": 3}]], result)
def test_search_vectors_rejects_text_parameter(self):
docs = self.conn.create_table(
"docs",
schema=_schema({
"id": pa.int32(),
"content": pa.string(),
"embedding": _vector(3),
}),
options=_PARQUET_OPTIONS,
)
with self.assertRaises(TypeError):
docs.search_vectors([[1.0, 0.0, 0.0]], text="paimon")
def test_search_reads_text_query(self):
docs = self.conn.create_table(
"docs",
schema=_schema({
"id": pa.int32(),
"content": pa.string(),
}),
options=_PARQUET_OPTIONS,
)
docs.add([
{"id": 1, "content": "paimon vector"},
{"id": 2, "content": "lakehouse"},
])
calls = {}
class FakeFullTextBuilder:
def with_query(self, field_name, query):
calls["query"] = (field_name, json.loads(query))
return self
def with_limit(self, limit):
calls["limit"] = limit
return self
def execute_local(self):
return GlobalIndexResult.from_range(Range(0, 0))
patcher = patch.object(type(docs.raw_table), "new_full_text_search_builder", return_value=FakeFullTextBuilder())
patcher.start()
self.addCleanup(patcher.stop)
result = (
docs.search("paimon vector")
.limit(1)
.to_arrow()
)
self.assertEqual(1, calls["limit"])
self.assertEqual("content", calls["query"][0])
self.assertEqual("paimon vector", calls["query"][1]["match"]["query"])
self.assertEqual([1], result["id"].to_pylist())
def test_search_string_requires_unambiguous_text_column(self):
docs = self.conn.create_table(
"docs",
schema=_schema({
"id": pa.int32(),
"title": pa.string(),
"content": pa.string(),
}),
options=_PARQUET_OPTIONS,
)
with self.assertRaisesRegex(ValueError, "Multiple text columns"):
docs.search("paimon")
def test_search_hybrid_rejects_shorthand_arguments(self):
docs = self.conn.create_table(
"docs",
schema=_schema({
"id": pa.int32(),
"embedding": _vector(3),
}),
options=_PARQUET_OPTIONS,
)
with self.assertRaises(TypeError):
docs.search_hybrid(vector=[1.0, 0.0, 0.0])
def test_hybrid_route_helpers_use_route_specific_arguments(self):
with self.assertRaises(TypeError):
pmm.vector_route([1.0, 0.0, 0.0])
route = pmm.text_route("paimon")
self.assertIsNone(route.column)
self.assertEqual("paimon", route.query)
def test_search_can_build_hybrid_routes(self):
docs = self.conn.create_table(
"docs",
schema=_schema({
"id": pa.int32(),
"content": pa.string(),
"embedding": _vector(3),
}),
options=_PARQUET_OPTIONS,
)
docs.add([
{"id": 1, "content": "paimon", "embedding": [1.0, 0.0, 0.0]},
{"id": 2, "content": "vector", "embedding": [0.0, 1.0, 0.0]},
])
calls = {}
class FakeHybridBuilder:
def with_limit(self, limit):
calls["limit"] = limit
return self
def with_ranker(self, ranker):
calls["ranker"] = ranker
return self
def add_vector_route(
self, column, vector, limit, weight=1.0, options=None):
calls["vector"] = (column, vector, limit, weight, options)
return self
def add_full_text_route(
self, field_name, query, limit, weight=1.0, options=None):
calls["text"] = (
field_name, json.loads(query), limit, weight, options)
return self
def with_filter(self, predicate):
calls["filter"] = predicate
return self
def execute_local(self):
return GlobalIndexResult.from_range(Range(0, 1))
patcher = patch.object(type(docs.raw_table), "new_hybrid_search_builder", return_value=FakeHybridBuilder())
patcher.start()
self.addCleanup(patcher.stop)
result = (
docs.search_hybrid(
[
pmm.vector_route("embedding", [1.0, 0.0, 0.0]),
pmm.text_route("paimon"),
],
)
.rerank("rrf")
.limit(2)
.to_arrow()
)
self.assertEqual(2, calls["limit"])
self.assertEqual("rrf", calls["ranker"])
self.assertEqual(
("embedding", [1.0, 0.0, 0.0], 2, 1.0, {}),
calls["vector"],
)
self.assertEqual("content", calls["text"][0])
self.assertEqual("paimon", calls["text"][1]["match"]["query"])
self.assertEqual([1, 2], result["id"].to_pylist())
def test_search_hybrid_applies_pre_filter_to_vector_routes(self):
docs = self.conn.create_table(
"docs",
schema=_schema({
"id": pa.int32(),
"category": pa.string(),
"embedding": _vector(3),
}),
options=_PARQUET_OPTIONS,
)
docs.add([
{"id": 1, "category": "lake", "embedding": [1.0, 0.0, 0.0]},
{"id": 2, "category": "city", "embedding": [0.0, 1.0, 0.0]},
])
calls = {}
class FakeHybridBuilder:
def with_limit(self, limit):
calls["limit"] = limit
return self
def with_ranker(self, ranker):
calls["ranker"] = ranker
return self
def add_vector_route(
self, column, vector, limit, weight=1.0, options=None):
calls["vector"] = (column, vector, limit, weight, options)
return self
def with_filter(self, predicate):
calls["pre_filter"] = predicate
return self
def execute_local(self):
return GlobalIndexResult.from_range(Range(0, 0))
patcher = patch.object(type(docs.raw_table), "new_hybrid_search_builder", return_value=FakeHybridBuilder())
patcher.start()
self.addCleanup(patcher.stop)
result = (
docs.search_hybrid(
[pmm.vector_route("embedding", [1.0, 0.0, 0.0])],
pre_filter="category = 'lake'",
)
.limit(1)
.to_list()
)
self.assertEqual(1, calls["limit"])
self.assertEqual(
("embedding", [1.0, 0.0, 0.0], 1, 1.0, {}),
calls["vector"],
)
self.assertEqual("category", calls["pre_filter"].field)
self.assertEqual(["lake"], calls["pre_filter"].literals)
self.assertEqual(
[{"id": 1, "category": "lake", "embedding": [1.0, 0.0, 0.0]}],
result,
)
def test_module_does_not_export_vectors_route(self):
self.assertFalse(hasattr(pmm, "vectors_route"))
self.assertFalse(hasattr(pmm, "VectorsRoute"))
def test_search_hybrid_can_build_multiple_vector_routes(self):
docs = self.conn.create_table(
"docs",
schema=_schema({
"id": pa.int32(),
"content": pa.string(),
"image_embedding": _vector(3),
"text_embedding": _vector(3),
}),
options=_PARQUET_OPTIONS,
)
docs.add([
{
"id": 1,
"content": "paimon",
"image_embedding": [1.0, 0.0, 0.0],
"text_embedding": [0.0, 1.0, 0.0],
},
{
"id": 2,
"content": "vector",
"image_embedding": [0.0, 1.0, 0.0],
"text_embedding": [1.0, 0.0, 0.0],
},
])
calls = {"vector_routes": [], "text_routes": []}
class FakeHybridBuilder:
def with_limit(self, limit):
calls["limit"] = limit
return self
def with_ranker(self, ranker):
calls["ranker"] = ranker
return self
def add_vector_route(
self, column, vector, limit, weight=1.0, options=None):
calls["vector_routes"].append(
(column, vector, limit, weight, options))
return self
def add_full_text_route(
self, field_name, query, limit, weight=1.0, options=None):
calls["text_routes"].append(
(field_name, json.loads(query), limit, weight, options))
return self
def execute_local(self):
return GlobalIndexResult.from_range(Range(0, 0))
patcher = patch.object(type(docs.raw_table), "new_hybrid_search_builder", return_value=FakeHybridBuilder())
patcher.start()
self.addCleanup(patcher.stop)
result = (
docs.search_hybrid(
[
pmm.vector_route(
"image_embedding",
[1.0, 0.0, 0.0],
weight=0.7,
limit=6,
options={"nprobe": "8"},
),
pmm.vector_route(
"text_embedding",
[0.0, 1.0, 0.0],
weight=0.3,
limit=4,
options={"nprobe": "4"},
),
pmm.text_route("paimon", weight=0.2),
],
ranker="weighted_score",
route_limit=4,
)
.limit(2)
.to_list()
)
self.assertEqual(2, calls["limit"])
self.assertEqual("weighted_score", calls["ranker"])
self.assertEqual(
[
(
"image_embedding",
[1.0, 0.0, 0.0],
6,
0.7,
{"nprobe": "8"},
),
(
"text_embedding",
[0.0, 1.0, 0.0],
4,
0.3,
{"nprobe": "4"},
),
],
calls["vector_routes"],
)
self.assertEqual("content", calls["text_routes"][0][0])
self.assertEqual("paimon", calls["text_routes"][0][1]["match"]["query"])
self.assertEqual(4, calls["text_routes"][0][2])
self.assertEqual(0.2, calls["text_routes"][0][3])
self.assertEqual(
[{
"id": 1,
"content": "paimon",
"image_embedding": [1.0, 0.0, 0.0],
"text_embedding": [0.0, 1.0, 0.0],
}],
result,
)
if __name__ == "__main__":
unittest.main()