| # 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 json | |
| import re | |
| from dataclasses import dataclass | |
| from typing import Dict, Mapping, Optional, Sequence | |
| import pyarrow as pa | |
| import pyarrow.compute as pc | |
| from pypaimon.common.options.core_options import CoreOptions | |
| from pypaimon.common.where_parser import parse_where_clause | |
| from pypaimon.multimodal.query import ( | |
| BatchVectorQuery, | |
| HybridQuery, | |
| ScanQuery, | |
| TextQuery, | |
| VectorQuery, | |
| ) | |
| from pypaimon.schema.data_types import ( | |
| PyarrowFieldParser, is_array_blob_type, is_blob_file_type, is_map_blob_type, | |
| ) | |
| from pypaimon.table.data_evolution_merge_into import ( | |
| WhenMatched, | |
| WhenNotMatched, | |
| source_col as _source_col, | |
| ) | |
| from pypaimon.table.special_fields import SpecialFields | |
| from pypaimon.write.update_validation import predicate_fields | |
| _ALL_SOURCE_COLUMNS = object() | |
| _MERGE_REF_PATTERN = re.compile(r'\b(source|target)\.(\w+)\b') | |
| _STRING_LITERAL_PATTERN = re.compile(r"'(?:[^']|'')*'") | |
| @dataclass(frozen=True) | |
| class VectorRoute: | |
| """Vector route spec for hybrid search.""" | |
| column: Optional[str] | |
| vector: object | |
| weight: float = 1.0 | |
| limit: Optional[int] = None | |
| options: Optional[Dict[str, str]] = None | |
| @dataclass(frozen=True) | |
| class TextRoute: | |
| """Full-text route spec for hybrid search.""" | |
| query: object | |
| column: Optional[str] = None | |
| weight: float = 1.0 | |
| limit: Optional[int] = None | |
| options: Optional[Dict[str, str]] = None | |
| def vector_route(column, vector, *, weight: float = 1.0, | |
| limit: Optional[int] = None, | |
| options: Optional[Dict[str, str]] = None) -> VectorRoute: | |
| return VectorRoute( | |
| column=column, | |
| vector=vector, | |
| weight=weight, | |
| limit=limit, | |
| options=dict(options or {}), | |
| ) | |
| def text_route(query, *, column: Optional[str] = None, weight: float = 1.0, | |
| limit: Optional[int] = None, | |
| options: Optional[Dict[str, str]] = None) -> TextRoute: | |
| return TextRoute( | |
| query=query, | |
| column=column, | |
| weight=weight, | |
| limit=limit, | |
| options=dict(options or {}), | |
| ) | |
| class MultimodalTable: | |
| """High-level table facade for mutable multimodal application data.""" | |
| def __init__(self, catalog, identifier: str, raw_table): | |
| self.catalog = catalog | |
| self.identifier = identifier | |
| self.name = identifier | |
| self.raw_table = raw_table | |
| self.table = raw_table | |
| def copy(self, options: dict) -> 'MultimodalTable': | |
| """Return a table view with dynamic Paimon options.""" | |
| raw_table = self.raw_table.copy(options) | |
| return MultimodalTable( | |
| self.catalog, raw_table.identifier.get_full_name(), raw_table) | |
| def add(self, data): | |
| arrow_table = _to_arrow_table(data, _target_schema(self.raw_table)) | |
| write_builder = self.raw_table.new_batch_write_builder() | |
| table_write = write_builder.new_write() | |
| table_commit = write_builder.new_commit() | |
| try: | |
| table_write.write_arrow(arrow_table) | |
| table_commit.commit(table_write.prepare_commit()) | |
| finally: | |
| table_write.close() | |
| table_commit.close() | |
| return self | |
| def add_video( | |
| self, | |
| video, | |
| frames=None, | |
| *, | |
| video_column=None, | |
| first_frame=0, | |
| metadata=None, | |
| source_fps=None): | |
| """Append frame rows sharing one encoded video. | |
| ``frames`` supplies non-video columns; descriptor ordinals start at | |
| ``first_frame``. Omit it to generate ``frame_index`` (int64) and | |
| ``timestamp`` (float64 seconds relative to the first frame). | |
| ``metadata`` supplies constant fields. Auto import requires PyAV | |
| and one video stream. For raw HEVC, ``source_fps`` sets fixed-rate | |
| timing and remuxes to MP4 without re-encoding. | |
| """ | |
| if frames is None: | |
| import os | |
| import stat | |
| import tempfile | |
| from contextlib import ExitStack | |
| from urllib.parse import urlparse | |
| from pypaimon.table.row.blob import Blob, BlobDescriptor | |
| from pypaimon.common.file_io import _is_local_path | |
| from pypaimon.filesystem.local_file_io import LocalFileIO | |
| from pypaimon.multimodal.video import ( | |
| _inspect_video_frames, _require_video_av, _raw_video_time_base, _remux_raw_hevc, | |
| ) | |
| time_base = _raw_video_time_base(source_fps) if source_fps is not None else None | |
| if isinstance(first_frame, bool) or not isinstance(first_frame, int) or first_frame != 0: | |
| raise ValueError("first_frame must be 0 when frames is omitted.") | |
| column = self._resolve_video_frame_column(video_column) | |
| schema = _target_schema(self.raw_table) | |
| generated = {"frame_index": pa.int64(), "timestamp": pa.float64()} | |
| for name, dtype in generated.items(): | |
| if name == column or name not in schema.names or schema.field(name).type != dtype: | |
| raise ValueError( | |
| "Automatic video import requires column %r of type %s." | |
| % (name, dtype)) | |
| if metadata is not None and not isinstance(metadata, Mapping): | |
| raise ValueError("metadata must be a mapping or None.") | |
| metadata = dict(metadata or {}) | |
| allowed = set(schema.names) - set(generated) - set(self.raw_table.options.video_frame_fields()) | |
| if set(metadata) - allowed: | |
| raise ValueError("metadata must contain only non-video, non-generated columns.") | |
| fields = [field for field in schema if field.name not in generated and field.name != column] | |
| constants = _to_arrow_table([metadata], pa.schema(fields)) if fields else None | |
| for field in fields: | |
| if not field.nullable and constants[field.name].null_count: | |
| raise ValueError("Missing required metadata: %s." % field.name) | |
| if isinstance(video, os.PathLike): | |
| video = os.fspath(video) | |
| if isinstance(video, str): | |
| if not urlparse(video).scheme: | |
| video = os.path.abspath(video) | |
| video = BlobDescriptor(video, 0, -1) | |
| if not isinstance(video, BlobDescriptor): | |
| raise ValueError("Automatic video import requires a path, URI, or BlobDescriptor.") | |
| _require_video_av() | |
| local_path = LocalFileIO._to_file(video.uri) if _is_local_path(video.uri) else None | |
| def source_state(): | |
| if local_path is None: | |
| return None | |
| info = os.stat(local_path) | |
| if not stat.S_ISREG(info.st_mode): | |
| raise ValueError("Automatic video import requires a regular local file.") | |
| return (info.st_dev, info.st_ino, info.st_size, info.st_mtime_ns, info.st_ctime_ns) | |
| with ExitStack() as stack: | |
| original = source_state() | |
| # Reuse the writer's source credentials. | |
| reader = self.raw_table.file_io.uri_reader_factory.create(video.uri) | |
| with Blob.from_descriptor(reader, video).new_input_stream() as stream: | |
| if time_base is not None: | |
| directory = stack.enter_context(tempfile.TemporaryDirectory(prefix="paimon_hevc_")) | |
| path = os.path.join(directory, "video.mp4") | |
| count = _remux_raw_hevc(stream, path, time_base) | |
| video = BlobDescriptor(path, 0, os.path.getsize(path)) | |
| timestamps = _inspect_video_frames(path) | |
| if len(timestamps) != count: | |
| raise ValueError("HEVC frame count changed during remux.") | |
| else: | |
| timestamps = _inspect_video_frames(stream) | |
| def check_source(): | |
| if source_state() != original: | |
| raise ValueError("Video source changed during import.") | |
| check_source() | |
| def batches(): | |
| for start in range(0, len(timestamps), 1024): | |
| end = min(start + 1024, len(timestamps)) | |
| values = { | |
| "frame_index": pa.array(range(start, end), type=pa.int64()), | |
| "timestamp": pa.array(timestamps[start:end], type=pa.float64()), | |
| } | |
| if constants is not None: | |
| repeated = constants.take(pa.array([0] * (end - start), type=pa.int32())) | |
| values.update({name: repeated[name] for name in repeated.column_names}) | |
| yield _video_frame_table(video, pa.table(values), column, start, schema) | |
| # Check again before commit. | |
| check_source() | |
| return self.add_batches(batches()) | |
| elif metadata is not None or source_fps is not None: | |
| raise ValueError("metadata and source_fps are only supported when frames is omitted.") | |
| return self.add_videos( | |
| [(video, frames, first_frame)], video_column=video_column | |
| ) | |
| def add_images_as_video( | |
| self, | |
| images, | |
| frames, | |
| *, | |
| fps, | |
| codec, | |
| pixel_format, | |
| gop_size, | |
| codec_options=None, | |
| video_column=None): | |
| """Encode ordered images as one MP4 and append its logical frame rows. | |
| Images may be encoded image bytes or Pillow images. ``frames`` supplies | |
| every table column except ``video_column`` and must contain one row per | |
| image. Encoding finishes before the Paimon write starts. | |
| """ | |
| import os | |
| import tempfile | |
| from pypaimon.multimodal.video import _encode_images_to_video | |
| column = self._resolve_video_frame_column(video_column) | |
| target_schema = _target_schema(self.raw_table) | |
| non_video_schema = pa.schema([ | |
| field for field in target_schema if field.name != column | |
| ]) | |
| frame_table = _to_arrow_table(frames, non_video_schema) | |
| with tempfile.TemporaryDirectory(prefix="pypaimon_video_") as directory: | |
| video_path = os.path.join(directory, "video.mp4") | |
| image_count = _encode_images_to_video( | |
| images, | |
| video_path, | |
| fps=fps, | |
| codec=codec, | |
| pixel_format=pixel_format, | |
| gop_size=gop_size, | |
| codec_options=codec_options, | |
| ) | |
| if image_count != frame_table.num_rows: | |
| raise ValueError( | |
| "Image count %d does not match frame row count %d." | |
| % (image_count, frame_table.num_rows) | |
| ) | |
| return self.add_video( | |
| video_path, frame_table, video_column=column | |
| ) | |
| def add_videos(self, videos, *, video_column=None): | |
| """Append several encoded videos with one writer and one commit. | |
| Each item is ``(video, frames)`` or ``(video, frames, first_frame)``. | |
| Keeping one writer open lets a single ``.video`` object pack multiple | |
| complete encoded videos up to the configured rolling target. | |
| """ | |
| column = self._resolve_video_frame_column(video_column) | |
| target_schema = _target_schema(self.raw_table) | |
| def frame_batches(): | |
| for item in videos: | |
| try: | |
| item = tuple(item) | |
| except TypeError as error: | |
| raise ValueError( | |
| "Each videos item must be (video, frames) or " | |
| "(video, frames, first_frame)." | |
| ) from error | |
| if len(item) == 2: | |
| video, frames = item | |
| first_frame = 0 | |
| elif len(item) == 3: | |
| video, frames, first_frame = item | |
| else: | |
| raise ValueError( | |
| "Each videos item must be (video, frames) or " | |
| "(video, frames, first_frame)." | |
| ) | |
| yield _video_frame_table( | |
| video, | |
| frames, | |
| column, | |
| first_frame, | |
| target_schema, | |
| ) | |
| return self.add_batches(frame_batches()) | |
| def _resolve_video_frame_column(self, requested): | |
| configured = self.raw_table.options.video_frame_fields() | |
| if not configured: | |
| raise ValueError( | |
| "add_video requires table option 'video-frame-field'." | |
| ) | |
| if requested is None and len(configured) > 1: | |
| raise ValueError( | |
| "video_column is required when 'video-frame-field' configures " | |
| "multiple fields." | |
| ) | |
| column = requested or next(iter(configured)) | |
| if column not in configured: | |
| raise ValueError( | |
| "Video column %r is not configured by 'video-frame-field'." | |
| % column | |
| ) | |
| return column | |
| def add_batches(self, batches): | |
| """Append an iterable of batches with one writer and one commit. | |
| Keeping the writer open across batches also preserves video payload | |
| groups and lets one ``.video`` file pack several encoded videos. | |
| """ | |
| try: | |
| iterator = iter(batches) | |
| except TypeError as error: | |
| raise ValueError("batches must be an iterable of input batches.") from error | |
| target_schema = _target_schema(self.raw_table) | |
| table_write = None | |
| table_commit = None | |
| commit_started = False | |
| try: | |
| for data in iterator: | |
| arrow_table = _to_arrow_table(data, target_schema) | |
| if arrow_table.num_rows == 0: | |
| continue | |
| if table_write is None: | |
| write_builder = self.raw_table.new_batch_write_builder() | |
| table_write = write_builder.new_write() | |
| table_commit = write_builder.new_commit() | |
| table_write.write_arrow(arrow_table) | |
| close_iterator = getattr(iterator, "close", None) | |
| iterator = None | |
| if close_iterator is not None: | |
| close_iterator() | |
| if table_write is None: | |
| return self | |
| commit_messages = table_write.prepare_commit() | |
| commit_started = True | |
| table_commit.commit(commit_messages) | |
| return self | |
| except BaseException: | |
| if table_write is not None and not commit_started: | |
| table_write.abort() | |
| raise | |
| finally: | |
| if iterator is not None: | |
| close_iterator = getattr(iterator, "close", None) | |
| if close_iterator is not None: | |
| close_iterator() | |
| try: | |
| if table_write is not None: | |
| table_write.close() | |
| finally: | |
| if table_commit is not None: | |
| table_commit.close() | |
| def overwrite(self, data, partition: Optional[Mapping[str, object]] = None): | |
| arrow_table = _to_arrow_table(data, _target_schema(self.raw_table)) | |
| overwrite_partition = dict(partition) if partition is not None else None | |
| write_builder = ( | |
| self.raw_table.new_batch_write_builder() | |
| .overwrite(overwrite_partition) | |
| ) | |
| table_write = write_builder.new_write() | |
| table_commit = write_builder.new_commit() | |
| try: | |
| table_write.write_arrow(arrow_table) | |
| table_commit.commit(table_write.prepare_commit()) | |
| finally: | |
| table_write.close() | |
| table_commit.close() | |
| return self | |
| def update(self, where, values): | |
| """Update visible rows matching a SQL-like predicate and return this table. | |
| Row-tracking tables also accept ``_ROW_ID`` predicates. Missing or | |
| deleted rows do not match. ``_ROW_ID`` is read-only. Use | |
| ``update_columns_by_row_id`` for external per-row values with strict | |
| ID matching. Row-ID predicates require a current table schema; reopen | |
| the table after external schema changes or a branch copy with a | |
| different schema. | |
| """ | |
| video_columns = self.raw_table.options.video_frame_fields() | |
| if isinstance(values, Mapping): | |
| assigned_video_columns = video_columns.intersection(values) | |
| if assigned_video_columns: | |
| raise ValueError( | |
| "update() cannot write video-frame-field %r; use " | |
| "replace_video() with a complete encoded video." | |
| % sorted(assigned_video_columns)[0] | |
| ) | |
| if isinstance(where, str) and self.raw_table.options.row_tracking_enabled(): | |
| predicate = parse_where_clause(where, SpecialFields.row_type_with_row_id(self.raw_table.fields)) | |
| else: | |
| predicate = self.scan().where(where)._predicate | |
| if predicate is not None and SpecialFields.ROW_ID.name in predicate_fields(predicate): | |
| _validate_row_id_update_schema(self) | |
| write_builder = self.raw_table.new_batch_write_builder() | |
| table_update = write_builder.new_update() | |
| table_commit = write_builder.new_commit() | |
| try: | |
| messages = table_update.update_by_predicate(predicate, values) | |
| table_commit.commit(messages) | |
| finally: | |
| table_commit.close() | |
| return self | |
| def replace_video( | |
| self, | |
| where, | |
| video, | |
| *, | |
| video_column=None, | |
| first_frame=0): | |
| """Replace the video backing the logical frame rows matching ``where``. | |
| Matching rows are ordered by ``_ROW_ID`` and assigned consecutive | |
| frame ordinals starting at ``first_frame``. Only the configured video | |
| column is updated; ordinary columns and the normal data files remain | |
| untouched. | |
| """ | |
| column = self._resolve_video_frame_column(video_column) | |
| target_schema = _target_schema(self.raw_table) | |
| payload, first_frame = _video_payload( | |
| video, first_frame) | |
| row_ids = ( | |
| self.scan() | |
| .where(where) | |
| .select([]) | |
| .with_row_id() | |
| .to_arrow()[SpecialFields.ROW_ID.name] | |
| .to_pylist() | |
| ) | |
| row_ids.sort() | |
| if not row_ids: | |
| return self | |
| descriptors = _video_frame_descriptors( | |
| payload, len(row_ids), first_frame) | |
| update_data = pa.Table.from_arrays( | |
| [ | |
| pa.array(row_ids, type=pa.int64()), | |
| pa.array(descriptors, type=target_schema.field(column).type), | |
| ], | |
| names=[SpecialFields.ROW_ID.name, column], | |
| ) | |
| write_builder = self.raw_table.new_batch_write_builder() | |
| table_update = write_builder.new_update().with_update_type([column]) | |
| table_commit = write_builder.new_commit() | |
| try: | |
| messages = table_update.update_by_arrow_with_row_id(update_data) | |
| table_commit.commit(messages) | |
| finally: | |
| table_commit.close() | |
| return self | |
| def delete(self, where): | |
| query = self.scan().where(where) | |
| predicate = query._predicate | |
| write_builder = self.raw_table.new_batch_write_builder() | |
| table_update = write_builder.new_update() | |
| table_commit = write_builder.new_commit() | |
| try: | |
| messages = table_update.delete_by_predicate(predicate) | |
| table_commit.commit(messages) | |
| finally: | |
| table_commit.close() | |
| return self | |
| def merge(self, on): | |
| return _MergeBuilder(self, on) | |
| def merge_columns(self, data, *, on): | |
| """Match external column values without inserting rows or rewriting other columns. | |
| Missing output columns are added as nullable columns using the input | |
| Arrow types. Matching follows an update-only MERGE: unmatched source | |
| rows are ignored, one source row may update multiple target rows, and | |
| multiple source rows matching the same target row are rejected. | |
| Existing output types must match the table types. Execution reuses | |
| the standard MERGE backend, including native execution when enabled. | |
| Schema publication and data commit are separate operations. Requires | |
| a current full-table view and a single writer for the operation. | |
| Returns :class:`pypaimon.multimodal.ColumnUpdateResult` with | |
| ``updated_rows=None`` because MERGE does not return logical row counts. | |
| """ | |
| from pypaimon.multimodal.columns import merge_columns | |
| return merge_columns(self, data, on=on) | |
| def add_columns(self, schema, *, inputs=None, fn=None, batch_size=1024): | |
| """Add nullable columns, optionally computing their values in batches. | |
| ``schema`` declares new columns. If supplied, ``fn`` receives an | |
| Arrow RecordBatch containing only ``inputs`` and returns a Table or | |
| RecordBatch of the declared columns in the same row order and count. | |
| Ordinary BLOB inputs are payload bytes. Video-frame inputs retain | |
| their descriptor semantics; decoding belongs in the function. | |
| The function may run repeatedly and must be side-effect-free. | |
| Outputs are grouped by the original file row-ID ranges before being | |
| written. Untouched columns, including media, retain their files. | |
| Schema publication precedes computation; failures leave new columns | |
| visible as NULL. Requires a current full-table view and a single | |
| writer. Returns :class:`pypaimon.multimodal.ColumnUpdateResult`. | |
| """ | |
| from pypaimon.multimodal.columns import add_columns | |
| return add_columns(self, schema, inputs=inputs, fn=fn, batch_size=batch_size) | |
| def update_columns_by_row_id(self, data, *, columns): | |
| """Update existing columns from one local Arrow table containing ``_ROW_ID``. | |
| ``columns`` is a required column name or sequence of output names. | |
| Only those columns are written; other input columns are ignored. | |
| IDs must be unique, non-NULL, non-negative integers representable as | |
| int64. Every ID must identify a visible row in the captured source | |
| snapshot. Input order does not affect matching. Outputs must be existing, | |
| non-partition columns with the table's types and nullability; use ``replace_video`` | |
| for video-frame columns. Omitted rows keep their values, and untouched | |
| columns retain their files. No columns or rows are added. | |
| Validation and writing use one snapshot and the standard update backend. | |
| All prepared updates are committed together. Requires a current | |
| full-table Data Evolution view and a single writer for the entire call. | |
| A schema mismatch with the current target branch is rejected before | |
| writing; reopen the table after external schema changes or a branch | |
| copy with a different schema. | |
| IDs from historical results are not remapped or checked for staleness. | |
| Rollback followed by an append can reuse row IDs and numeric snapshot | |
| IDs. Discard pending annotations from the rolled-back history before | |
| writing them back; a matching numeric snapshot ID does not prove validity. | |
| Returns :class:`pypaimon.multimodal.ColumnUpdateResult` with the input | |
| row count in ``updated_rows`` and ``added_columns=()``. Empty input | |
| creates no data snapshot. There are no automatic retries. | |
| """ | |
| from pypaimon.multimodal.columns import update_columns_by_row_id | |
| return update_columns_by_row_id(self, data, columns=columns) | |
| def compute_columns(self, columns, *, inputs, fn, batch_size=1024): | |
| """Compute existing columns for all visible rows without changing the schema. | |
| ``columns`` is a column name or a sequence of names. Outputs must be | |
| existing, non-partition columns; use ``replace_video`` for video-frame | |
| outputs. Types and nullability are taken from the table schema. | |
| ``fn`` follows the same input and output contract as ``add_columns``. | |
| Inputs may also be outputs; all reads use the captured source snapshot. | |
| Execution is local and requires a current full-table view and a single | |
| writer. Every call recomputes all visible rows, including non-NULL | |
| outputs. Untouched columns retain their files. There are no automatic | |
| retries or checkpoints. Returns :class:`pypaimon.multimodal.ColumnUpdateResult` | |
| with ``added_columns=()`` and the computed row count in ``updated_rows``. | |
| """ | |
| from pypaimon.multimodal.columns import compute_columns | |
| return compute_columns(self, columns, inputs=inputs, fn=fn, batch_size=batch_size) | |
| def map_with_blobs(self, dataset, columns, fn, **kwargs): | |
| from pypaimon.ray import map_with_blobs | |
| return map_with_blobs( | |
| dataset, | |
| columns, | |
| fn, | |
| file_io=self.raw_table.file_io, | |
| all_blob_columns=_blob_columns(self.raw_table), | |
| map_blob_columns=[field.name for field in self.raw_table.fields | |
| if is_map_blob_type(field.type)], | |
| array_blob_columns=[field.name for field in self.raw_table.fields | |
| if is_array_blob_type(field.type)], | |
| **kwargs, | |
| ) | |
| def scan( | |
| self, | |
| *, | |
| snapshot_id: Optional[int] = None, | |
| tag_name: Optional[str] = None): | |
| return ScanQuery(_time_travel_table( | |
| self.raw_table, snapshot_id=snapshot_id, tag_name=tag_name)) | |
| def take_row_ids( | |
| self, | |
| row_ids, | |
| *, | |
| snapshot_id: Optional[int] = None, | |
| tag_name: Optional[str] = None): | |
| """Return a scan filtering rows by Paimon row ID; requires row tracking. | |
| Input order and duplicates do not affect this membership query. | |
| Missing or deleted IDs are omitted. Use ``take_batch`` for ordered | |
| sampling with replacement. | |
| """ | |
| row_ids = _coerce_row_ids(row_ids) | |
| read_table = _time_travel_table( | |
| self.raw_table, snapshot_id=snapshot_id, tag_name=tag_name) | |
| if not read_table.options.row_tracking_enabled(): | |
| raise ValueError("Row ID reads require row-tracking.enabled=true.") | |
| read_builder = read_table.new_read_builder().with_projection( | |
| [field.name for field in read_table.fields] | |
| + [SpecialFields.ROW_ID.name] | |
| ) | |
| predicate = read_builder.new_predicate_builder().is_in( | |
| SpecialFields.ROW_ID.name, row_ids) | |
| query = ScanQuery(read_table) | |
| query._predicate = predicate | |
| return query | |
| def take_batch( | |
| self, | |
| row_ids, | |
| *, | |
| columns=None, | |
| snapshot_id: Optional[int] = None, | |
| tag_name: Optional[str] = None) -> pa.Table: | |
| """Fetch an ordered batch by Paimon row ID, preserving duplicates. | |
| Returns an Arrow table with one row per input ID, in input order. | |
| ``columns`` defaults to all table columns; include ``_ROW_ID`` | |
| explicitly to return it. Missing, deleted, or filtered-out IDs raise | |
| KeyError instead of shortening the batch. An empty input returns an | |
| empty table with the projected schema. | |
| Requires row tracking. IDs are table row handles, not Dataset position | |
| indices. Pass the source snapshot or tag when sampling a saved manifest. | |
| BLOB columns follow the table's read options, as in ``scan()``. | |
| """ | |
| row_ids = _coerce_row_ids(row_ids) | |
| query = self.take_row_ids( | |
| row_ids, snapshot_id=snapshot_id, tag_name=tag_name) | |
| if columns is None: | |
| columns = [field.name for field in query._table.fields] | |
| elif isinstance(columns, str): | |
| columns = [columns] | |
| else: | |
| columns = list(columns) | |
| batch = query.select(columns).with_row_id().to_arrow() | |
| requested = pa.array(row_ids, type=pa.int64()) | |
| positions = pc.index_in( | |
| requested, value_set=batch.column(SpecialFields.ROW_ID.name)) | |
| if positions.null_count: | |
| missing = requested.filter(pc.is_null(positions)).to_pylist() | |
| raise KeyError("Row IDs not found: %s" % list(dict.fromkeys(missing))) | |
| return batch.take(positions).select(columns) | |
| def blobs(self, *, column: Optional[str] = None, key_column: Optional[str] = None): | |
| from pypaimon.multimodal.blob_store import BlobStore | |
| return BlobStore(self, column=column, key_column=key_column) | |
| def search( | |
| self, | |
| query, | |
| *, | |
| column: Optional[str] = None, | |
| options: Optional[Dict[str, str]] = None, | |
| pre_filter=None, | |
| snapshot_id: Optional[int] = None, | |
| tag_name: Optional[str] = None): | |
| read_table = _time_travel_table( | |
| self.raw_table, snapshot_id=snapshot_id, tag_name=tag_name) | |
| schema = _target_schema(read_table) | |
| if isinstance(query, str): | |
| return TextQuery( | |
| read_table, | |
| text_query=_coerce_full_text_query( | |
| query, "search", schema, column=column), | |
| pre_filter=pre_filter, | |
| ) | |
| vector = _coerce_vector(query, "search") | |
| if _looks_like_batch_vectors(vector): | |
| raise ValueError( | |
| "search() accepts a single query; use search_vectors() for " | |
| "multiple vectors.") | |
| return VectorQuery( | |
| read_table, | |
| vector=vector, | |
| vector_column=_resolve_vector_column( | |
| schema, column, len(vector)), | |
| vector_options=options, | |
| pre_filter=pre_filter, | |
| ) | |
| def search_vectors( | |
| self, | |
| vectors, | |
| *, | |
| column: Optional[str] = None, | |
| options: Optional[Dict[str, str]] = None, | |
| pre_filter=None, | |
| snapshot_id: Optional[int] = None, | |
| tag_name: Optional[str] = None): | |
| read_table = _time_travel_table( | |
| self.raw_table, snapshot_id=snapshot_id, tag_name=tag_name) | |
| schema = _target_schema(read_table) | |
| vectors = _coerce_vectors(vectors) | |
| dimension = len(vectors[0]) | |
| if any(len(vector) != dimension for vector in vectors): | |
| raise ValueError( | |
| "search_vectors requires all query vectors to have the same " | |
| "dimension.") | |
| vector_column = _resolve_vector_column( | |
| schema, column, dimension) | |
| return BatchVectorQuery( | |
| read_table, | |
| vectors=vectors, | |
| vector_column=vector_column, | |
| vector_options=options, | |
| pre_filter=pre_filter, | |
| ) | |
| def search_hybrid( | |
| self, | |
| routes, | |
| *, | |
| ranker: str = "rrf", | |
| route_limit: Optional[int] = None, | |
| pre_filter=None, | |
| snapshot_id: Optional[int] = None, | |
| tag_name: Optional[str] = None): | |
| read_table = _time_travel_table( | |
| self.raw_table, snapshot_id=snapshot_id, tag_name=tag_name) | |
| schema = _target_schema(read_table) | |
| vector_routes, text_routes = _normalize_hybrid_routes( | |
| schema, | |
| routes=routes, | |
| method="search_hybrid", | |
| ) | |
| if not vector_routes and not text_routes: | |
| raise ValueError( | |
| "search_hybrid requires at least one route.") | |
| return HybridQuery( | |
| read_table, | |
| vector_routes=vector_routes, | |
| text_routes=text_routes, | |
| ranker=ranker, | |
| route_limit=route_limit, | |
| pre_filter=pre_filter, | |
| ) | |
| def create_index(self, column, index_type, options=None, *, execution="local", | |
| concurrency=None, ray_remote_args=None): | |
| return self.raw_table.create_global_index( | |
| column, index_type=_normalize_index_type(index_type), options=options, | |
| execution=execution, concurrency=concurrency, ray_remote_args=ray_remote_args) | |
| def _merge( | |
| self, | |
| source, | |
| on, | |
| when_matched: Sequence[WhenMatched], | |
| when_not_matched: Sequence[WhenNotMatched], | |
| operation: str): | |
| target_cols, source_cols = _normalize_merge_on(on, operation) | |
| target_schema = _target_schema(self.raw_table) | |
| source_table = _to_arrow_table(source) | |
| _validate_merge_on( | |
| target_schema, source_table, target_cols, source_cols, operation) | |
| source_table = _cast_merge_on_columns( | |
| source_table, target_schema, target_cols, source_cols) | |
| when_matched, when_not_matched = _resolve_merge_clauses( | |
| when_matched, | |
| when_not_matched, | |
| target_schema, | |
| source_table, | |
| dict(zip(target_cols, source_cols)), | |
| ) | |
| write_builder = self.raw_table.new_batch_write_builder() | |
| table_update = write_builder.new_update() | |
| table_commit = write_builder.new_commit() | |
| try: | |
| messages = table_update.merge_into( | |
| source_table, | |
| on=dict(zip(target_cols, source_cols)), | |
| when_matched=list(when_matched), | |
| when_not_matched=list(when_not_matched), | |
| ) | |
| table_commit.commit(messages) | |
| finally: | |
| table_commit.close() | |
| return self | |
| class _MergeBuilder: | |
| """Builder for idempotent merge writes.""" | |
| def __init__(self, table: MultimodalTable, on): | |
| self._table = table | |
| self._on = on | |
| self._when_matched = [] | |
| self._when_not_matched = [] | |
| def when_matched_update(self, values=None, where: Optional[str] = None): | |
| self._when_matched.append( | |
| WhenMatched.update( | |
| _ALL_SOURCE_COLUMNS if values is None else values, | |
| condition=_normalize_merge_where(where), | |
| )) | |
| return self | |
| def when_matched_delete(self, where: Optional[str] = None): | |
| self._when_matched.append( | |
| WhenMatched.delete( | |
| condition=_normalize_merge_where(where), | |
| )) | |
| return self | |
| def when_not_matched_insert(self, values=None, where: Optional[str] = None): | |
| self._when_not_matched.append( | |
| WhenNotMatched( | |
| insert=_ALL_SOURCE_COLUMNS if values is None else values, | |
| condition=_normalize_merge_where(where), | |
| )) | |
| return self | |
| def execute(self, data): | |
| if not self._when_matched and not self._when_not_matched: | |
| raise ValueError( | |
| "merge requires at least one matched or not-matched clause.") | |
| return self._table._merge( | |
| data, | |
| self._on, | |
| when_matched=self._when_matched, | |
| when_not_matched=self._when_not_matched, | |
| operation="merge", | |
| ) | |
| def _blob_columns(table): | |
| return tuple( | |
| field.name for field in table.fields | |
| if is_blob_file_type(field.type) | |
| ) | |
| def _to_arrow_table(data, target_schema=None): | |
| if target_schema is not None: | |
| data = _serialize_blob_values(data, target_schema) | |
| if isinstance(data, list): | |
| # Untyped from_pylist only discovers columns present in the first row. | |
| data = { | |
| field.name: [row.get(field.name) for row in data] | |
| for field in target_schema | |
| } | |
| if isinstance(data, dict): | |
| data = dict(data) | |
| for field in target_schema: | |
| if field.name in data and pa.types.is_map(field.type): | |
| # Inferred STRUCT/LIST arrays cannot be cast back to MAP. | |
| data[field.name] = pa.array(data[field.name], type=field.type) | |
| if isinstance(data, pa.Table): | |
| table = data | |
| elif isinstance(data, pa.RecordBatch): | |
| table = pa.Table.from_batches([data]) | |
| elif isinstance(data, list): | |
| # Inference from the first row alone drops fields in later source rows. | |
| names = dict.fromkeys(name for row in data for name in row) | |
| table = pa.Table.from_pydict({ | |
| name: [row.get(name) for row in data] for name in names | |
| }) | |
| elif isinstance(data, dict): | |
| table = pa.Table.from_pydict(data) | |
| elif hasattr(data, "__dataframe__") or data.__class__.__module__.startswith("pandas"): | |
| table = pa.Table.from_pandas(data, preserve_index=False) | |
| else: | |
| raise ValueError("Unsupported multimodal data type: %r" % type(data)) | |
| if target_schema is None: | |
| return table | |
| return _align_to_schema(table, target_schema) | |
| def _serialize_blob_values(data, target_schema): | |
| if isinstance(data, (pa.Table, pa.RecordBatch)): | |
| return data | |
| binary_fields = { | |
| field.name: field.type | |
| for field in target_schema | |
| if _contains_binary(field.type) | |
| } | |
| if not binary_fields: | |
| return data | |
| if isinstance(data, list): | |
| return [ | |
| { | |
| name: _serialize_blob_value(value, binary_fields.get(name)) | |
| for name, value in row.items() | |
| } | |
| if isinstance(row, Mapping) | |
| else row | |
| for row in data | |
| ] | |
| if isinstance(data, dict): | |
| converted = dict(data) | |
| for name, arrow_type in binary_fields.items(): | |
| if name not in converted: | |
| continue | |
| column = converted[name] | |
| if isinstance(column, (pa.Array, pa.ChunkedArray)): | |
| column = column.to_pylist() | |
| converted[name] = [ | |
| _serialize_blob_value(value, arrow_type) | |
| for value in column | |
| ] | |
| return converted | |
| if ( | |
| hasattr(data, "__dataframe__") | |
| or data.__class__.__module__.startswith("pandas") | |
| ): | |
| converted = data.copy() | |
| for name, arrow_type in binary_fields.items(): | |
| if name in converted.columns: | |
| converted[name] = converted[name].map( | |
| lambda value: _serialize_blob_value(value, arrow_type) | |
| ) | |
| return converted | |
| return data | |
| def _contains_binary(arrow_type): | |
| if pa.types.is_binary(arrow_type) or pa.types.is_large_binary(arrow_type): | |
| return True | |
| if pa.types.is_list(arrow_type) or pa.types.is_large_list(arrow_type): | |
| return _contains_binary(arrow_type.value_type) | |
| if pa.types.is_map(arrow_type): | |
| return _contains_binary(arrow_type.item_type) | |
| return False | |
| def _serialize_blob_value(value, arrow_type): | |
| if value is None or arrow_type is None: | |
| return value | |
| if pa.types.is_binary(arrow_type) or pa.types.is_large_binary(arrow_type): | |
| from pypaimon.table.row.blob import Blob, BlobDescriptor | |
| if isinstance(value, BlobDescriptor): | |
| return value.serialize() | |
| if isinstance(value, Blob): | |
| try: | |
| return value.to_descriptor().serialize() | |
| except RuntimeError: | |
| return value.to_data() | |
| return value | |
| if pa.types.is_list(arrow_type) or pa.types.is_large_list(arrow_type): | |
| return [ | |
| _serialize_blob_value(element, arrow_type.value_type) | |
| for element in value | |
| ] | |
| if pa.types.is_map(arrow_type): | |
| entries = value.items() if isinstance(value, Mapping) else value | |
| return [ | |
| (key, _serialize_blob_value(element, arrow_type.item_type)) | |
| for key, element in entries | |
| ] | |
| return value | |
| def _coerce_row_ids(row_ids): | |
| if row_ids is None or isinstance(row_ids, (str, bytes)): | |
| raise ValueError("row_ids must be an iterable of row id integers.") | |
| try: | |
| iterator = iter(row_ids) | |
| except TypeError: | |
| raise ValueError("row_ids must be an iterable of row id integers.") | |
| coerced = [] | |
| for row_id in iterator: | |
| if hasattr(row_id, "as_py"): | |
| row_id = row_id.as_py() | |
| coerced.append(int(row_id)) | |
| return coerced | |
| def _time_travel_table(table, snapshot_id=None, tag_name=None): | |
| if snapshot_id is not None and tag_name is not None: | |
| raise ValueError( | |
| "snapshot_id and tag_name cannot be set at the same time") | |
| if snapshot_id is not None: | |
| return table.copy({ | |
| CoreOptions.SCAN_SNAPSHOT_ID.key(): str(snapshot_id), | |
| }) | |
| if tag_name is not None: | |
| return table.copy({ | |
| CoreOptions.SCAN_TAG_NAME.key(): tag_name, | |
| }) | |
| return table | |
| def _validate_row_id_update_schema(owner): | |
| raw = owner.raw_table | |
| current = owner.catalog.get_table(raw.identifier) | |
| if current.table_schema.id != raw.table_schema.id or current.fields != raw.fields: | |
| raise ValueError("Table schema has changed. Reopen the table before updating columns by row ID.") | |
| def _target_schema(table): | |
| return PyarrowFieldParser.from_paimon_schema(table.table_schema.fields) | |
| def _video_payload(video, first_frame): | |
| from pypaimon.table.row.blob import ( | |
| Blob, | |
| BlobDescriptor, | |
| VideoFrameDescriptor, | |
| ) | |
| if isinstance(first_frame, bool) or not isinstance(first_frame, int): | |
| raise ValueError("first_frame must be a non-negative int.") | |
| if first_frame < 0: | |
| raise ValueError("first_frame must be a non-negative int.") | |
| if isinstance(video, str): | |
| video = Blob.from_local(video) | |
| if isinstance(video, Blob): | |
| try: | |
| payload = video.to_descriptor() | |
| except RuntimeError as error: | |
| raise ValueError( | |
| "video must be descriptor-backed; inline video bytes are not " | |
| "accepted by video write APIs." | |
| ) from error | |
| elif isinstance(video, BlobDescriptor): | |
| payload = video | |
| else: | |
| raise ValueError( | |
| "video must be a path, Blob, or BlobDescriptor, got %r." | |
| % type(video) | |
| ) | |
| if isinstance(payload, VideoFrameDescriptor): | |
| payload = payload.payload_descriptor | |
| return payload, first_frame | |
| def _video_frame_descriptors(payload, count, first_frame): | |
| from pypaimon.table.row.blob import VideoFrameDescriptor | |
| return [ | |
| VideoFrameDescriptor( | |
| payload.uri, | |
| payload.offset, | |
| payload.length, | |
| first_frame + index, | |
| ).serialize() | |
| for index in range(count) | |
| ] | |
| def _video_frame_table(video, frames, video_column, first_frame, target_schema): | |
| payload, first_frame = _video_payload(video, first_frame) | |
| non_video_schema = pa.schema([ | |
| field for field in target_schema if field.name != video_column | |
| ]) | |
| frame_table = _to_arrow_table(frames, non_video_schema) | |
| descriptor_values = _video_frame_descriptors( | |
| payload, frame_table.num_rows, first_frame) | |
| arrays = [] | |
| for field in target_schema: | |
| if field.name == video_column: | |
| arrays.append(pa.array(descriptor_values, type=field.type)) | |
| else: | |
| arrays.append(frame_table[field.name]) | |
| return pa.Table.from_arrays(arrays, schema=target_schema) | |
| def _align_to_schema( | |
| table: pa.Table, | |
| schema: pa.Schema, | |
| column_mapping: Optional[Dict[str, str]] = None) -> pa.Table: | |
| column_mapping = column_mapping or {} | |
| arrays = [] | |
| for field in schema: | |
| source_name = column_mapping.get(field.name, field.name) | |
| if source_name in table.column_names: | |
| array = table[source_name] | |
| if array.type != field.type: | |
| array = array.cast(field.type) | |
| else: | |
| array = pa.nulls(table.num_rows, type=field.type) | |
| arrays.append(array) | |
| return pa.Table.from_arrays(arrays, schema=schema) | |
| def _normalize_merge_on(on, operation: str): | |
| if isinstance(on, Mapping): | |
| target_cols = list(on.keys()) | |
| source_cols = list(on.values()) | |
| elif isinstance(on, str): | |
| target_cols = [on] | |
| source_cols = [on] | |
| else: | |
| target_cols = list(on) | |
| source_cols = list(on) | |
| if not target_cols: | |
| raise ValueError("%s requires at least one ON column." % operation) | |
| if len(target_cols) != len(source_cols): | |
| raise ValueError( | |
| "%s ON target/source column counts do not match." % operation) | |
| return target_cols, source_cols | |
| def _validate_merge_on( | |
| target_schema, | |
| source_table, | |
| target_cols, | |
| source_cols, | |
| operation: str): | |
| target_names = set(target_schema.names) | |
| source_names = set(source_table.column_names) | |
| missing_targets = [col for col in target_cols if col not in target_names] | |
| if missing_targets: | |
| raise ValueError( | |
| "%s ON target columns are not in table schema: %s" | |
| % (operation, missing_targets)) | |
| missing_sources = [col for col in source_cols if col not in source_names] | |
| if missing_sources: | |
| raise ValueError( | |
| "%s ON source columns are not in source data: %s" | |
| % (operation, missing_sources)) | |
| def _cast_merge_on_columns(source_table, target_schema, target_cols, source_cols): | |
| result = source_table | |
| for target_col, source_col_name in zip(target_cols, source_cols): | |
| target_type = target_schema.field(target_col).type | |
| index = result.schema.get_field_index(source_col_name) | |
| if index < 0: | |
| continue | |
| array = result[source_col_name] | |
| if array.type == target_type: | |
| continue | |
| source_field = result.schema.field(index) | |
| result = result.set_column( | |
| index, | |
| pa.field( | |
| source_col_name, | |
| target_type, | |
| nullable=source_field.nullable, | |
| metadata=source_field.metadata, | |
| ), | |
| array.cast(target_type), | |
| ) | |
| return result | |
| def _resolve_merge_clauses( | |
| when_matched, | |
| when_not_matched, | |
| target_schema, | |
| source_table, | |
| on_map): | |
| target_names = list(target_schema.names) | |
| source_names = set(source_table.column_names) | |
| return ( | |
| [ | |
| ( | |
| WhenMatched.delete(condition=clause.condition) | |
| if clause.delete | |
| else WhenMatched.update( | |
| _resolve_merge_set_spec( | |
| clause.update, target_names, source_names, on_map), | |
| condition=clause.condition, | |
| ) | |
| ) | |
| for clause in when_matched | |
| ], | |
| [ | |
| WhenNotMatched( | |
| insert=_resolve_merge_set_spec( | |
| clause.insert, target_names, source_names, on_map), | |
| condition=clause.condition, | |
| ) | |
| for clause in when_not_matched | |
| ], | |
| ) | |
| def _resolve_merge_set_spec(spec, target_names, source_names, on_map): | |
| if spec is not _ALL_SOURCE_COLUMNS: | |
| return spec | |
| return { | |
| target_name: _source_col(source_name) | |
| for target_name in target_names | |
| for source_name in [on_map.get(target_name, target_name)] | |
| if source_name in source_names | |
| } | |
| def _normalize_merge_where(where): | |
| if where is None: | |
| return None | |
| if not isinstance(where, str): | |
| raise ValueError("merge where must be a string.") | |
| parts, last = [], 0 | |
| for match in _STRING_LITERAL_PATTERN.finditer(where): | |
| parts.append(_rewrite_merge_refs(where[last:match.start()])) | |
| parts.append(match.group()) | |
| last = match.end() | |
| parts.append(_rewrite_merge_refs(where[last:])) | |
| return ''.join(parts) | |
| def _rewrite_merge_refs(text): | |
| return _MERGE_REF_PATTERN.sub( | |
| lambda m: "%s.%s" % ( | |
| "s" if m.group(1) == "source" else "t", | |
| m.group(2), | |
| ), | |
| text, | |
| ) | |
| def _normalize_index_type(index_type): | |
| if not isinstance(index_type, str): | |
| return index_type | |
| normalized = index_type.strip().lower() | |
| if normalized == "full-text": | |
| return "full-text" | |
| return index_type | |
| def _coerce_vector(value, method): | |
| if value is None or isinstance(value, str): | |
| raise ValueError("%s requires a vector query." % method) | |
| if hasattr(value, "tolist"): | |
| value = value.tolist() | |
| try: | |
| return list(value) | |
| except TypeError as e: | |
| raise ValueError("%s requires a vector query." % method) from e | |
| def _coerce_vectors(values): | |
| method = "search_vectors" | |
| if values is None or isinstance(values, str): | |
| raise ValueError("%s requires query vectors." % method) | |
| if hasattr(values, "tolist"): | |
| values = values.tolist() | |
| try: | |
| vectors = list(values) | |
| except TypeError as e: | |
| raise ValueError("%s requires query vectors." % method) from e | |
| if not vectors: | |
| raise ValueError("%s requires at least one vector query." % method) | |
| return [_coerce_vector(vector, method) for vector in vectors] | |
| def _looks_like_batch_vectors(value): | |
| if not value: | |
| return False | |
| first = value[0] | |
| return not isinstance(first, (int, float)) and hasattr(first, "__iter__") | |
| def _normalize_hybrid_routes( | |
| schema, | |
| routes, | |
| method): | |
| vector_routes, text_routes = [], [] | |
| for route in _route_specs(routes): | |
| if _is_text_route_spec(route): | |
| text_routes.append(_normalize_text_route(route, method, schema)) | |
| else: | |
| vector_routes.append( | |
| _normalize_vector_route(route, method, schema)) | |
| return vector_routes, text_routes | |
| def _route_specs(routes): | |
| if routes is None: | |
| return [] | |
| if isinstance(routes, (VectorRoute, TextRoute, Mapping)): | |
| return [routes] | |
| if _is_column_vector_pair(routes): | |
| return [routes] | |
| return list(routes) | |
| def _is_text_route_spec(route): | |
| if isinstance(route, TextRoute): | |
| return True | |
| if isinstance(route, VectorRoute): | |
| return False | |
| if isinstance(route, Mapping): | |
| route_type = route.get("type") or route.get("route_type") | |
| if route_type is not None: | |
| return str(route_type).replace("-", "_").lower() in ( | |
| "text", "full_text", "fulltext") | |
| return ( | |
| "query" in route | |
| or "text" in route | |
| or "terms" in route | |
| or "full_text_query" in route | |
| ) and not _has_vector_query(route) | |
| return False | |
| def _normalize_vector_route(route, method, schema): | |
| if isinstance(route, VectorRoute): | |
| column = route.column | |
| vector = route.vector | |
| weight = route.weight | |
| limit = route.limit | |
| options = route.options | |
| elif isinstance(route, Mapping): | |
| column = ( | |
| route.get("column") | |
| or route.get("vector_column") | |
| or route.get("field") | |
| or route.get("anns_field") | |
| ) | |
| vector = _mapping_vector(route) | |
| weight = route.get("weight", 1.0) | |
| limit = route.get("limit") | |
| options = _mapping_options(route) | |
| elif _is_column_vector_pair(route): | |
| column, vector = route | |
| weight = 1.0 | |
| limit = None | |
| options = None | |
| else: | |
| raise ValueError( | |
| "%s vector routes require a route spec with a column and vector." | |
| % method) | |
| if not column: | |
| raise ValueError("%s vector routes require a column." % method) | |
| return { | |
| "column": column, | |
| "vector": _coerce_vector(vector, method), | |
| "weight": weight, | |
| "limit": limit, | |
| "options": dict(options or {}), | |
| } | |
| def _normalize_text_route(route, method, schema): | |
| if isinstance(route, TextRoute): | |
| column = route.column | |
| query = route.query | |
| weight = route.weight | |
| limit = route.limit | |
| options = route.options | |
| elif isinstance(route, Mapping): | |
| column = route.get("column") | |
| query = ( | |
| route.get("query") | |
| or route.get("text") | |
| or route.get("terms") | |
| or route.get("full_text_query") | |
| ) | |
| weight = route.get("weight", 1.0) | |
| limit = route.get("limit") | |
| options = route.get("options") | |
| else: | |
| raise ValueError( | |
| "%s text routes require a route spec with a query." | |
| % method) | |
| text_query = _coerce_full_text_query(query, method, schema, column=column) | |
| return { | |
| "column": text_query["column"], | |
| "query": text_query["query"], | |
| "weight": weight, | |
| "limit": limit, | |
| "options": dict(options or {}), | |
| } | |
| def _has_vector_query(route): | |
| return any(key in route for key in ("vector", "query_vector", "data")) | |
| def _mapping_vector(route): | |
| if "vector" in route: | |
| return route["vector"] | |
| if "query_vector" in route: | |
| return route["query_vector"] | |
| if "data" in route: | |
| data = route["data"] | |
| if hasattr(data, "tolist"): | |
| data = data.tolist() | |
| values = list(data) | |
| if len(values) == 1 and hasattr(values[0], "__iter__"): | |
| return values[0] | |
| return values | |
| raise ValueError("Vector route requires vector, query_vector, or data.") | |
| def _mapping_options(route): | |
| if "options" in route: | |
| return route["options"] | |
| if "param" in route: | |
| param = route["param"] | |
| if isinstance(param, Mapping) and isinstance(param.get("params"), Mapping): | |
| return param["params"] | |
| return param | |
| if "params" in route: | |
| return route["params"] | |
| return None | |
| def _is_column_vector_pair(route): | |
| return ( | |
| isinstance(route, (list, tuple)) | |
| and len(route) == 2 | |
| and isinstance(route[0], str) | |
| ) | |
| def _coerce_full_text_query(query, method, schema, column=None): | |
| field_name = column or _infer_text_column(schema, "text") | |
| if isinstance(query, str): | |
| stripped = query.lstrip() | |
| query_json = ( | |
| query | |
| if stripped.startswith("{") | |
| else json.dumps({"match": {"query": query}}, separators=(",", ":")) | |
| ) | |
| return {"column": field_name, "query": query_json} | |
| if isinstance(query, Mapping): | |
| return { | |
| "column": field_name, | |
| "query": json.dumps(query, separators=(",", ":")), | |
| } | |
| raise ValueError("%s requires a text string or query mapping." % method) | |
| def _resolve_vector_column( | |
| schema: pa.Schema, | |
| column: Optional[str], | |
| dimension: int) -> str: | |
| if column is not None: | |
| try: | |
| field = schema.field(column) | |
| except KeyError as e: | |
| raise ValueError( | |
| "Vector column '%s' not found in table schema." % column) from e | |
| if not pa.types.is_fixed_size_list(field.type): | |
| raise ValueError( | |
| "Column '%s' is not a fixed-size vector column." % column) | |
| if field.type.list_size != dimension: | |
| raise ValueError( | |
| "Vector dimension %d does not match column '%s' dimension %d." | |
| % (dimension, column, field.type.list_size)) | |
| return column | |
| candidates = [ | |
| (field.name, field.type.list_size) | |
| for field in schema | |
| if pa.types.is_fixed_size_list(field.type) | |
| ] | |
| matches = [ | |
| name | |
| for name, column_dimension in candidates | |
| if column_dimension == dimension | |
| ] | |
| if len(matches) == 1: | |
| return matches[0] | |
| if not matches: | |
| available = ", ".join( | |
| "%s(%d)" % candidate for candidate in candidates) or "none" | |
| raise ValueError( | |
| "No vector column found with dimension %d; available vector " | |
| "columns: %s." | |
| % (dimension, available)) | |
| raise ValueError( | |
| "Multiple vector columns found with dimension %d: %s; pass column." | |
| % (dimension, ", ".join(matches))) | |
| def _infer_text_column(schema: pa.Schema, parameter: str = "text_column"): | |
| columns = [ | |
| field.name | |
| for field in schema | |
| if pa.types.is_string(field.type) or pa.types.is_large_string(field.type) | |
| ] | |
| return _infer_single_column(columns, "text", parameter) | |
| def _infer_single_column(columns, kind, parameter): | |
| if len(columns) == 1: | |
| return columns[0] | |
| if not columns: | |
| raise ValueError("No %s column found; pass %s." % (kind, parameter)) | |
| raise ValueError( | |
| "Multiple %s columns found %s; pass %s." | |
| % (kind, columns, parameter)) |