From c701653559d85461eb5ce165f61d23bb1191e39f Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Sun, 21 Jun 2026 20:31:04 +0800 Subject: [PATCH 1/4] [python] Support raw fallback for global indexes --- .../pypaimon/common/options/core_options.py | 19 + .../globalindex/global_index_coverage.py | 145 ++++++ .../globalindex/global_index_scanner.py | 102 ++++- paimon-python/pypaimon/read/read_builder.py | 8 +- .../pypaimon/read/scanner/file_scanner.py | 5 +- paimon-python/pypaimon/read/table_scan.py | 30 +- .../source/batch_vector_search_builder.py | 1 + .../table/source/vector_search_builder.py | 2 + .../table/source/vector_search_read.py | 314 ++++++++++++- .../table/source/vector_search_scan.py | 123 ++++- .../table/source/vector_search_split.py | 44 +- .../pypaimon/tests/global_index_test.py | 194 ++++++++ .../tests/vector_search_filter_test.py | 419 ++++++++++++++++++ 13 files changed, 1353 insertions(+), 53 deletions(-) create mode 100644 paimon-python/pypaimon/globalindex/global_index_coverage.py diff --git a/paimon-python/pypaimon/common/options/core_options.py b/paimon-python/pypaimon/common/options/core_options.py index 93b7b92b49f9..2b24bc1be469 100644 --- a/paimon-python/pypaimon/common/options/core_options.py +++ b/paimon-python/pypaimon/common/options/core_options.py @@ -90,6 +90,12 @@ class GlobalIndexColumnUpdateAction(str, Enum): DROP_PARTITION_INDEX = "DROP_PARTITION_INDEX" +class GlobalIndexSearchMode(str, Enum): + FAST = "fast" + FULL = "full" + DETAIL = "detail" + + class CoreOptions: """Core options for Paimon tables.""" # File format constants @@ -618,6 +624,16 @@ class CoreOptions: .with_description("Whether to enable global index for scan.") ) + GLOBAL_INDEX_SEARCH_MODE: ConfigOption[GlobalIndexSearchMode] = ( + ConfigOptions.key("global-index.search-mode") + .enum_type(GlobalIndexSearchMode) + .default_value(GlobalIndexSearchMode.FAST) + .with_description( + "Search mode for global index queries. " + "Supported values are 'fast', 'full', and 'detail'." + ) + ) + GLOBAL_INDEX_THREAD_NUM: ConfigOption[int] = ( ConfigOptions.key("global-index.thread-num") .int_type() @@ -1109,6 +1125,9 @@ def commit_max_retry_wait(self) -> int: def global_index_enabled(self, default=None): return self.options.get(CoreOptions.GLOBAL_INDEX_ENABLED, default) + def global_index_search_mode(self): + return self.options.get(CoreOptions.GLOBAL_INDEX_SEARCH_MODE) + def global_index_thread_num(self) -> Optional[int]: return self.options.get(CoreOptions.GLOBAL_INDEX_THREAD_NUM) diff --git a/paimon-python/pypaimon/globalindex/global_index_coverage.py b/paimon-python/pypaimon/globalindex/global_index_coverage.py new file mode 100644 index 000000000000..9f94f8949510 --- /dev/null +++ b/paimon-python/pypaimon/globalindex/global_index_coverage.py @@ -0,0 +1,145 @@ +# 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. + +"""Row ranges covered and not covered by global index files.""" + +from typing import Collection, Dict, List, Optional, Union + +from pypaimon.common.options.core_options import CoreOptions, GlobalIndexSearchMode +from pypaimon.common.options.options import Options +from pypaimon.common.predicate import Predicate +from pypaimon.read.push_down_utils import _get_all_fields +from pypaimon.schema.data_types import DataField +from pypaimon.utils.range import Range + + +class GlobalIndexCoverage: + """Computes global-index coverage by field id.""" + + def __init__( + self, + table, + snapshot, + partition_filter, + index_files: Collection['IndexFileMeta'], + ): + self._table = table + self._snapshot = snapshot + self._partition_filter = partition_filter + self._coverage_by_field: Dict[int, List[Range]] = {} + for index_file in index_files: + meta = index_file.global_index_meta + if meta is None: + continue + row_range = Range(meta.row_range_start, meta.row_range_end) + self._add_coverage(meta.index_field_id, row_range) + if meta.extra_field_ids is not None: + for extra_field_id in meta.extra_field_ids: + self._add_coverage(extra_field_id, row_range) + + def unindexed_ranges( + self, + fields_or_field_id: Union[List[DataField], int], + predicate: Optional[Predicate] = None, + ) -> List[Range]: + if isinstance(fields_or_field_id, int): + field_ids = {fields_or_field_id} + else: + field_by_name = {f.name: f for f in fields_or_field_id} + field_ids = set() + for name in _get_all_fields(predicate): + field = field_by_name.get(name) + if field is not None: + field_ids.add(field.id) + return self._unindexed_ranges(field_ids) + + def _add_coverage(self, field_id: int, row_range: Range) -> None: + self._coverage_by_field.setdefault(field_id, []).append(row_range) + + def _indexed_ranges(self, field_ids: Collection[int]) -> List[Range]: + ranges = None + for field_id in field_ids: + field_ranges = self._coverage_by_field.get(field_id) + if not field_ranges: + return [] + field_ranges = Range.sort_and_merge_overlap(field_ranges, True) + ranges = field_ranges if ranges is None else Range.and_(ranges, field_ranges) + if ranges is None: + return [] + return Range.sort_and_merge_overlap(ranges, True) + + def _unindexed_ranges(self, field_ids: Collection[int]) -> List[Range]: + search_mode = _global_index_search_mode(self._table) + if search_mode == GlobalIndexSearchMode.FAST: + return [] + next_row_id = getattr(self._snapshot, "next_row_id", None) + if self._snapshot is None or next_row_id is None: + return [] + if next_row_id <= 0: + return [] + + if search_mode == GlobalIndexSearchMode.DETAIL: + data_ranges = self._data_ranges_by_data_files() + else: + data_ranges = [Range(0, next_row_id - 1)] + + indexed_ranges = Range.sort_and_merge_overlap( + self._indexed_ranges(field_ids), True) + unindexed = [] + for data_range in Range.sort_and_merge_overlap(data_ranges, True): + unindexed.extend(data_range.exclude(indexed_ranges)) + return Range.sort_and_merge_overlap(unindexed, True) + + def _data_ranges_by_data_files(self) -> List[Range]: + if hasattr(self._table, "data_ranges_for_global_index_coverage"): + return self._table.data_ranges_for_global_index_coverage( + self._snapshot, + self._partition_filter, + ) + + manifest_list_manager = getattr(self._table, "manifest_list_manager", None) + if manifest_list_manager is None: + from pypaimon.manifest.manifest_list_manager import ManifestListManager + manifest_list_manager = ManifestListManager(self._table) + manifest_files = manifest_list_manager.read_all(self._snapshot) + from pypaimon.manifest.manifest_file_manager import ManifestFileManager + manager = ManifestFileManager(self._table) + entries = manager.read_entries_parallel( + manifest_files, + self._entry_matches_partition, + max_workers=self._table.options.scan_manifest_parallelism(), + ) + data_ranges = [] + for entry in entries: + row_range = entry.file.row_id_range() + if row_range is not None: + data_ranges.append(row_range) + return data_ranges + + def _entry_matches_partition(self, entry) -> bool: + if self._partition_filter is None: + return True + return self._partition_filter.test(entry.partition) + + +def _global_index_search_mode(table): + options = getattr(table, "options", None) + if options is None: + return GlobalIndexSearchMode.FAST + if hasattr(options, "global_index_search_mode"): + return options.global_index_search_mode() + return CoreOptions(Options.from_none()).global_index_search_mode() diff --git a/paimon-python/pypaimon/globalindex/global_index_scanner.py b/paimon-python/pypaimon/globalindex/global_index_scanner.py index 23af17e405a8..7f2470d382e8 100644 --- a/paimon-python/pypaimon/globalindex/global_index_scanner.py +++ b/paimon-python/pypaimon/globalindex/global_index_scanner.py @@ -27,6 +27,7 @@ from pypaimon.common.options.core_options import CoreOptions from pypaimon.common.options.options import Options from pypaimon.common.predicate import Predicate +from pypaimon.globalindex.global_index_coverage import GlobalIndexCoverage from pypaimon.read.push_down_utils import _get_all_fields from pypaimon.schema.data_types import DataField from pypaimon.utils.range import Range @@ -43,11 +44,20 @@ def __init__( index_files: Collection['IndexFileMeta'], thread_num: Optional[int] = None, options: Optional[CoreOptions] = None, + table=None, + snapshot=None, + partition_filter=None, ): self._options = options or CoreOptions(Options.from_none()) self._executor = ThreadPoolExecutor( max_workers=thread_num or 32 ) + self._fields = fields + self._coverage = ( + GlobalIndexCoverage(table, snapshot, partition_filter, index_files) + if table is not None + else None + ) self._evaluator = self._create_evaluator( fields, file_io, index_path, index_files ) @@ -59,17 +69,10 @@ def _create_evaluator(self, fields, file_io, index_path, index_files): if global_index_meta is None: continue - field_id = global_index_meta.index_field_id index_type = index_file.index_type - - if field_id not in index_metas: - index_metas[field_id] = {} - if index_type not in index_metas[field_id]: - index_metas[field_id][index_type] = {} - - range_key = Range(global_index_meta.row_range_start, global_index_meta.row_range_end) - if range_key not in index_metas[field_id][index_type]: - index_metas[field_id][index_type][range_key] = [] + field_ids = [global_index_meta.index_field_id] + if global_index_meta.extra_field_ids is not None: + field_ids.extend(global_index_meta.extra_field_ids) io_meta = GlobalIndexIOMeta( file_name=index_file.file_name, @@ -77,7 +80,17 @@ def _create_evaluator(self, fields, file_io, index_path, index_files): metadata=global_index_meta.index_meta, external_path=index_file.external_path, ) - index_metas[field_id][index_type][range_key].append(io_meta) + range_key = Range( + global_index_meta.row_range_start, + global_index_meta.row_range_end) + for field_id in field_ids: + if field_id not in index_metas: + index_metas[field_id] = {} + if index_type not in index_metas[field_id]: + index_metas[field_id][index_type] = {} + if range_key not in index_metas[field_id][index_type]: + index_metas[field_id][index_type][range_key] = [] + index_metas[field_id][index_type][range_key].append(io_meta) executor = self._executor options = self._options @@ -105,13 +118,17 @@ def create(table, index_files=None, partition_filter=None, predicate=None, if index_files is not None: if len(index_files) == 0: return None + core_options = _core_options(table) return GlobalIndexScanner( fields=table.fields, file_io=table.file_io, index_path=table.path_factory().global_index_path_factory().index_path(), index_files=index_files, - thread_num=table.options.global_index_thread_num(), - options=table.options, + thread_num=core_options.global_index_thread_num(), + options=core_options, + table=table, + snapshot=_resolve_snapshot(table, snapshot), + partition_filter=partition_filter, ) # Scan index files from snapshot using partition_filter and predicate @@ -129,29 +146,47 @@ def index_file_filter(entry): global_index_meta = entry.index_file.global_index_meta if global_index_meta is None: return False - return global_index_meta.index_field_id in filter_field_ids + if global_index_meta.index_field_id in filter_field_ids: + return True + if global_index_meta.extra_field_ids is not None: + return any( + field_id in filter_field_ids + for field_id in global_index_meta.extra_field_ids + ) + return False if snapshot is None: - snapshot = table.snapshot_manager().get_latest_snapshot() + snapshot = _resolve_snapshot(table, None) index_file_handler = IndexFileHandler(table=table) entries = index_file_handler.scan(snapshot, index_file_filter) scanned_index_files = [entry.index_file for entry in entries] if len(scanned_index_files) == 0: return None + core_options = _core_options(table) return GlobalIndexScanner( fields=table.fields, file_io=table.file_io, index_path=table.path_factory().global_index_path_factory().index_path(), index_files=scanned_index_files, - thread_num=table.options.global_index_thread_num(), - options=table.options, + thread_num=core_options.global_index_thread_num(), + options=core_options, + table=table, + snapshot=snapshot, + partition_filter=partition_filter, ) def scan(self, predicate: Optional[Predicate]) -> Optional[GlobalIndexResult]: """Scan the global index with the given predicate.""" return self._evaluator.evaluate(predicate) + def unindexed_rows(self, predicate: Optional[Predicate]) -> GlobalIndexResult: + """Return coarse row ids not covered by global indexes.""" + if self._coverage is None: + return GlobalIndexResult.create_empty() + return GlobalIndexResult.from_ranges( + self._coverage.unindexed_ranges(self._fields, predicate)) + def close(self): """Close the scanner and release resources.""" self._evaluator.close() @@ -164,6 +199,39 @@ def __exit__(self, exc_type, exc_val, exc_tb) -> None: self.close() +def _resolve_snapshot(table, snapshot): + if snapshot is not None: + return snapshot + snapshot_manager = table.snapshot_manager() + if snapshot_manager is None: + return None + try: + from pypaimon.snapshot.time_travel_util import TimeTravelUtil + table_options = getattr(table.table_schema, "options", {}) + scan_keys = getattr(TimeTravelUtil, "SCAN_KEYS", None) + if scan_keys is None: + from pypaimon.snapshot.time_travel_util import SCAN_KEYS as scan_keys + has_time_travel = any(key in table_options for key in scan_keys) + resolved = TimeTravelUtil.try_travel_to_snapshot( + Options(table.table_schema.options), + table.tag_manager(), + snapshot_manager, + ) + if resolved is not None and ( + has_time_travel or getattr(resolved, "next_row_id", None) is not None): + return resolved + except Exception: + pass + return snapshot_manager.get_latest_snapshot() + + +def _core_options(table): + options = getattr(table, "options", None) + if options is None: + return CoreOptions(Options.from_none()) + return options + + def _create_readers(file_io, index_path, index_type_metas, field, executor=None, options=None): """Create readers for a specific field, dispatched by index_type. diff --git a/paimon-python/pypaimon/read/read_builder.py b/paimon-python/pypaimon/read/read_builder.py index 5537c97dd647..eea670db533a 100644 --- a/paimon-python/pypaimon/read/read_builder.py +++ b/paimon-python/pypaimon/read/read_builder.py @@ -44,12 +44,17 @@ def __init__(self, table): # in ``read_type()`` and downstream consumers. self._projection: Optional[List[str]] = None self._nested_paths: Optional[List[List[int]]] = None + self._partition_filter: Optional[Predicate] = None self._limit: Optional[int] = None def with_filter(self, predicate: Predicate) -> 'ReadBuilder': self._predicate = predicate return self + def with_partition_filter(self, partition_filter: Predicate) -> 'ReadBuilder': + self._partition_filter = partition_filter + return self + def with_projection(self, projection: List[str]) -> 'ReadBuilder': """Project to the given column names. @@ -77,7 +82,8 @@ def new_scan(self) -> TableScan: return TableScan( table=self.table, predicate=self._predicate, - limit=self._limit + limit=self._limit, + partition_predicate=self._partition_filter, ) def new_read(self) -> TableRead: diff --git a/paimon-python/pypaimon/read/scanner/file_scanner.py b/paimon-python/pypaimon/read/scanner/file_scanner.py index e5f4e332a144..3facb2f7d524 100755 --- a/paimon-python/pypaimon/read/scanner/file_scanner.py +++ b/paimon-python/pypaimon/read/scanner/file_scanner.py @@ -384,7 +384,10 @@ def _eval_global_index(self, snapshot=None): if scanner is None: return None with scanner: - return scanner.scan(self.predicate) + result = scanner.scan(self.predicate) + if result is None: + return None + return result.or_(scanner.unindexed_rows(self.predicate)) except Exception: return None diff --git a/paimon-python/pypaimon/read/table_scan.py b/paimon-python/pypaimon/read/table_scan.py index 36568c618fa5..b2fcaa860452 100755 --- a/paimon-python/pypaimon/read/table_scan.py +++ b/paimon-python/pypaimon/read/table_scan.py @@ -33,13 +33,15 @@ def __init__( self, table, predicate: Optional[Predicate], - limit: Optional[int] + limit: Optional[int], + partition_predicate: Optional[Predicate] = None, ): from pypaimon.table.file_store_table import FileStoreTable self.table: FileStoreTable = table self.predicate = predicate self.limit = limit + self.partition_predicate = partition_predicate self.file_scanner = self._create_file_scanner() def plan(self) -> Plan: @@ -79,7 +81,11 @@ def _create_file_scanner(self) -> FileScanner: earliest_snapshot = snapshot_manager.try_get_earliest_snapshot() latest_snapshot = snapshot_manager.get_latest_snapshot() if earliest_snapshot is None or latest_snapshot is None: - return FileScanner(self.table, lambda: ([], None)) + return FileScanner( + self.table, + lambda: ([], None), + partition_predicate=self.partition_predicate, + ) start_timestamp = int(ts[0]) end_timestamp = int(ts[1]) if start_timestamp >= end_timestamp: @@ -87,7 +93,11 @@ def _create_file_scanner(self) -> FileScanner: "Ending timestamp %s should be >= starting timestamp %s." % (end_timestamp, start_timestamp)) if (start_timestamp == end_timestamp or start_timestamp > latest_snapshot.time_millis or end_timestamp < earliest_snapshot.time_millis): - return FileScanner(self.table, lambda: ([], None)) + return FileScanner( + self.table, + lambda: ([], None), + partition_predicate=self.partition_predicate, + ) starting_snapshot = snapshot_manager.earlier_or_equal_time_mills(start_timestamp) earliest_snapshot = snapshot_manager.try_get_earliest_snapshot() @@ -118,7 +128,13 @@ def incremental_manifest(): manifests.extend(manifest_files) return manifests, end_snapshot - return FileScanner(self.table, incremental_manifest, self.predicate, self.limit) + return FileScanner( + self.table, + incremental_manifest, + self.predicate, + self.limit, + partition_predicate=self.partition_predicate, + ) if has_time_travel: def time_travel_manifest_scanner(): @@ -135,7 +151,8 @@ def time_travel_manifest_scanner(): self.table, time_travel_manifest_scanner, self.predicate, - self.limit + self.limit, + partition_predicate=self.partition_predicate, ) def all_manifests(): @@ -146,7 +163,8 @@ def all_manifests(): self.table, all_manifests, self.predicate, - self.limit + self.limit, + partition_predicate=self.partition_predicate, ) def with_shard(self, idx_of_this_subtask, number_of_para_subtasks) -> 'TableScan': diff --git a/paimon-python/pypaimon/table/source/batch_vector_search_builder.py b/paimon-python/pypaimon/table/source/batch_vector_search_builder.py index f987680a3661..199759487b14 100644 --- a/paimon-python/pypaimon/table/source/batch_vector_search_builder.py +++ b/paimon-python/pypaimon/table/source/batch_vector_search_builder.py @@ -123,5 +123,6 @@ def new_batch_vector_search_read(self): self._vector_column, self._query_vectors, filter_=self._filter, + partition_filter=self._partition_filter, options=self._options, ) diff --git a/paimon-python/pypaimon/table/source/vector_search_builder.py b/paimon-python/pypaimon/table/source/vector_search_builder.py index f98547192903..25b2d2e4c529 100644 --- a/paimon-python/pypaimon/table/source/vector_search_builder.py +++ b/paimon-python/pypaimon/table/source/vector_search_builder.py @@ -216,6 +216,7 @@ def new_vector_search_scan(self): self._vector_column, filter_=self._filter, partition_filter=self._partition_filter, + options=self._options, ) @@ -245,5 +246,6 @@ def new_vector_search_read(self): self._vector_column, self._query_vector, filter_=self._filter, + partition_filter=self._partition_filter, options=self._options, ) diff --git a/paimon-python/pypaimon/table/source/vector_search_read.py b/paimon-python/pypaimon/table/source/vector_search_read.py index 566b093cd6b2..965630807172 100644 --- a/paimon-python/pypaimon/table/source/vector_search_read.py +++ b/paimon-python/pypaimon/table/source/vector_search_read.py @@ -25,6 +25,12 @@ from pypaimon.globalindex.offset_global_index_reader import OffsetGlobalIndexReader from pypaimon.globalindex.vector_search import VectorSearch from pypaimon.globalindex.vector_search_result import DictBasedScoredIndexResult +from pypaimon.table.source.vector_search_split import ( + IndexVectorSearchSplit, + RawVectorSearchSplit, +) +from pypaimon.utils.range import Range +from pypaimon.utils.roaring_bitmap import RoaringBitmap64 class VectorSearchRead(ABC): @@ -56,18 +62,27 @@ def read_batch(self, splits): class AbstractVectorSearchReadImpl: """Base implementation for vector search reads.""" - def __init__(self, table, limit, vector_column, filter_=None, options=None): + def __init__( + self, + table, + limit, + vector_column, + filter_=None, + partition_filter=None, + options=None, + ): self._table = table self._limit = limit self._vector_column = vector_column self._filter = filter_ + self._partition_filter = partition_filter self._options = dict(options or {}) - def _pre_filter(self, splits): - # type: (list) -> Optional[RoaringBitmap64] - """Evaluate the scalar filter against scalar global indexes to produce a row-id bitmap.""" + def _pre_filters(self, splits): + # type: (list) -> List[RoaringBitmap64] + """Evaluate scalar indexes and return one include bitmap per index split.""" if self._filter is None: - return None + return [] # Collect scalar index files across splits, deduplicated by file name. seen = set() @@ -79,18 +94,77 @@ def _pre_filter(self, splits): seen.add(index_file.file_name) scalar_files.append(index_file) + if not scalar_files: + return _empty_bitmaps(len(splits)) + + from pypaimon.globalindex.global_index_scanner import GlobalIndexScanner + scanner = GlobalIndexScanner.create( + self._table, + index_files=scalar_files, + partition_filter=self._partition_filter, + ) + if scanner is None: + return _empty_bitmaps(len(splits)) + try: + result = scanner.scan(self._filter) + if result is None: + return _empty_bitmaps(len(splits)) + matched_rows = result.results() + finally: + scanner.close() + + include_row_ids = [] + for split in splits: + split_rows = _bitmap_of_range( + Range(split.row_range_start, split.row_range_end)) + include_row_ids.append(RoaringBitmap64.and_(matched_rows, split_rows)) + return include_row_ids + + def _pre_filter(self, splits): + # Backwards-compatible helper used by older tests/callers. + pre_filters = self._pre_filters(splits) + if not pre_filters: + return None + merged = RoaringBitmap64() + for bitmap in pre_filters: + merged = RoaringBitmap64.or_(merged, bitmap) + return merged + + def _raw_pre_filter(self, splits): + if self._filter is None: + return None + raw_rows = _bitmap_of_ranges(_raw_row_ranges(splits)) + if raw_rows.is_empty(): + return None + + seen = set() + scalar_files = [] + for split in splits: + for index_file in split.scalar_index_files: + if index_file.file_name in seen: + continue + seen.add(index_file.file_name) + scalar_files.append(index_file) if not scalar_files: return None from pypaimon.globalindex.global_index_scanner import GlobalIndexScanner - scanner = GlobalIndexScanner.create(self._table, index_files=scalar_files) + scanner = GlobalIndexScanner.create( + self._table, + index_files=scalar_files, + partition_filter=self._partition_filter, + ) if scanner is None: return None try: result = scanner.scan(self._filter) if result is None: return None - return result.results() + include = result.results() + include = RoaringBitmap64.or_( + include, + scanner.unindexed_rows(self._filter).results()) + return RoaringBitmap64.and_(include, raw_rows) finally: scanner.close() @@ -136,28 +210,90 @@ def _eval(self, row_range_start, row_range_end, vector_index_files, future.add_done_callback(lambda _: reader.close()) return future + def _read_raw_search(self, raw_row_ranges, pre_filter, query_vector, index_type=None): + raw_row_ranges = Range.sort_and_merge_overlap(raw_row_ranges, True) + if pre_filter is not None: + raw_row_ranges = Range.and_( + raw_row_ranges, + Range.sort_and_merge_overlap(pre_filter.to_range_list(), True), + ) + if not raw_row_ranges: + return DictBasedScoredIndexResult({}) + + read_builder = self._table.new_read_builder() + if self._partition_filter is not None: + read_builder = read_builder.with_partition_filter( + self._partition_filter) + if self._filter is not None: + read_builder = read_builder.with_filter(self._filter) + from pypaimon.table.special_fields import SpecialFields + projection = [f.name for f in self._table.fields] + if SpecialFields.ROW_ID.name not in projection: + projection.append(SpecialFields.ROW_ID.name) + read_builder = read_builder.with_projection(projection) + plan = read_builder.new_scan().with_global_index_result( + GlobalIndexResult.from_ranges(raw_row_ranges)).plan() + table = read_builder.new_read().to_arrow(plan.splits()) + if table is None or table.num_rows == 0: + return DictBasedScoredIndexResult({}) + + row_ids = table.column(SpecialFields.ROW_ID.name).to_pylist() + vectors = table.column(self._vector_column.name).to_pylist() + metric = _raw_search_metric( + self._table, self._vector_column, self._options, index_type) + scores = {} + for row_id, stored in zip(row_ids, vectors): + if stored is None: + continue + stored_vector = _to_vector_list(stored) + if len(stored_vector) != len(query_vector): + raise ValueError( + "Query vector dimension mismatch: expected %d, got %d" + % (len(stored_vector), len(query_vector))) + scores[row_id] = _compute_score(query_vector, stored_vector, metric) + return DictBasedScoredIndexResult(scores).top_k(self._limit) + class VectorSearchReadImpl(AbstractVectorSearchReadImpl, VectorSearchRead): """Implementation for VectorSearchRead.""" def __init__(self, table, limit, vector_column, query_vector, filter_=None, - options=None): + partition_filter=None, options=None): super().__init__(table, limit, vector_column, - filter_=filter_, options=options) + filter_=filter_, + partition_filter=partition_filter, + options=options) self._query_vector = query_vector def read(self, splits): # type: (List[VectorSearchSplit]) -> GlobalIndexResult - if not splits: + index_splits, raw_splits = _split_search_splits(splits) + if not index_splits and not raw_splits: return GlobalIndexResult.create_empty() - pre_filter = self._pre_filter(splits) + indexed = ( + DictBasedScoredIndexResult({}) + if not index_splits + else self._read_indexed(index_splits, self._query_vector) + ) + raw_result = self._read_raw_search( + _raw_row_ranges(raw_splits), + self._raw_pre_filter(raw_splits), + self._query_vector, + _raw_search_index_type(raw_splits), + ) + return indexed.or_(raw_result).top_k(self._limit) + + def _read_indexed(self, splits, query_vector): + pre_filters = self._pre_filters(splits) futures = [ self._eval( split.row_range_start, split.row_range_end, - split.vector_index_files, self._query_vector, pre_filter + split.vector_index_files, + query_vector, + None if not pre_filters else pre_filters[i] ) - for split in splits + for i, split in enumerate(splits) ] wait(futures) @@ -179,25 +315,30 @@ class BatchVectorSearchReadImpl(AbstractVectorSearchReadImpl, """Batch vector search read; result ``i`` corresponds to query vector ``i``.""" def __init__(self, table, limit, vector_column, query_vectors, - filter_=None, options=None): + filter_=None, partition_filter=None, options=None): super().__init__(table, limit, vector_column, - filter_=filter_, options=options) + filter_=filter_, + partition_filter=partition_filter, + options=options) self._query_vectors = list(query_vectors) def read_batch(self, splits): # type: (List[VectorSearchSplit]) -> List[GlobalIndexResult] n = len(self._query_vectors) - if not splits: + index_splits, raw_splits = _split_search_splits(splits) + if not index_splits and not raw_splits: return [GlobalIndexResult.create_empty() for _ in range(n)] - pre_filter = self._pre_filter(splits) + pre_filters = self._pre_filters(index_splits) futures_by_vector = [ [ self._eval( split.row_range_start, split.row_range_end, - split.vector_index_files, vector, pre_filter + split.vector_index_files, + vector, + None if not pre_filters else pre_filters[i] ) - for split in splits + for i, split in enumerate(index_splits) ] for vector in self._query_vectors ] @@ -206,6 +347,9 @@ def read_batch(self, splits): wait(futures) results = [] + raw_pre_filter = self._raw_pre_filter(raw_splits) + raw_ranges = _raw_row_ranges(raw_splits) + raw_index_type = _raw_search_index_type(raw_splits) for futures in futures_by_vector: merged_scores = {} for future in futures: @@ -215,7 +359,11 @@ def read_batch(self, splits): for row_id in split_result.results(): if row_id not in merged_scores: merged_scores[row_id] = score_getter(row_id) - results.append(DictBasedScoredIndexResult(merged_scores).top_k(self._limit)) + indexed = DictBasedScoredIndexResult(merged_scores) + vector = self._query_vectors[len(results)] + raw = self._read_raw_search( + raw_ranges, raw_pre_filter, vector, raw_index_type) + results.append(indexed.or_(raw).top_k(self._limit)) return results @@ -238,3 +386,129 @@ def _create_vector_reader(index_type, file_io, index_path, index_io_meta_list, o file_io, index_path, index_io_meta_list, options ) raise ValueError("Unsupported vector index type: '%s'" % index_type) + + +def _split_search_splits(splits): + index_splits = [] + raw_splits = [] + for split in splits: + if isinstance(split, IndexVectorSearchSplit): + index_splits.append(split) + elif isinstance(split, RawVectorSearchSplit): + raw_splits.append(split) + return index_splits, raw_splits + + +def _raw_row_ranges(raw_splits): + ranges = [] + for split in raw_splits: + ranges.extend(split.row_ranges) + return Range.sort_and_merge_overlap(ranges, True) + + +def _raw_search_index_type(raw_splits): + for split in raw_splits: + if split.index_type is not None: + return split.index_type + return None + + +def _empty_bitmaps(size): + return [RoaringBitmap64() for _ in range(size)] + + +def _bitmap_of_range(row_range): + bitmap = RoaringBitmap64() + bitmap.add_range(row_range.from_, row_range.to) + return bitmap + + +def _bitmap_of_ranges(ranges): + bitmap = RoaringBitmap64() + for row_range in ranges: + bitmap.add_range(row_range.from_, row_range.to) + return bitmap + + +def _to_vector_list(value): + if hasattr(value, "to_list"): + return value.to_list() + if hasattr(value, "as_py"): + value = value.as_py() + return list(value) + + +def _raw_search_metric(table, vector_column, options, index_type=None): + candidates = [] + field_prefix = "fields.%s." % vector_column.name + index_prefix = "%s." % index_type if index_type else None + for key in [ + field_prefix + "distance.metric", + field_prefix + "metric", + *(([ + index_prefix + "distance.metric", + index_prefix + "metric", + ]) if index_prefix is not None else []), + "test.vector.metric", + "lumina.distance.metric", + "distance.metric", + "metric", + ]: + if key in options: + candidates.append(options[key]) + table_options = getattr(getattr(table, "options", None), "options", None) + table_map = table_options.to_map() if table_options is not None else {} + for key in [ + field_prefix + "distance.metric", + field_prefix + "metric", + *(([ + index_prefix + "distance.metric", + index_prefix + "metric", + ]) if index_prefix is not None else []), + "test.vector.metric", + "lumina.distance.metric", + "distance.metric", + "metric", + ]: + if key in table_map: + candidates.append(table_map[key]) + if candidates: + return _normalize_metric(candidates[0]) + + inferred = None + for key, value in list(options.items()) + list(table_map.items()): + if key.endswith(".distance.metric") or key.endswith(".metric"): + metric = _normalize_metric(value) + if metric in ("l2", "cosine", "inner_product"): + if inferred is not None and inferred != metric: + return "l2" + inferred = metric + return inferred or "l2" + + +def _normalize_metric(metric): + return str(metric).lower().replace("-", "_") + + +def _compute_score(query, stored, metric): + if metric == "l2": + sum_sq = 0.0 + for q, s in zip(query, stored): + diff = float(q) - float(s) + sum_sq += diff * diff + return 1.0 / (1.0 + sum_sq) + if metric == "cosine": + dot = 0.0 + norm_a = 0.0 + norm_b = 0.0 + for q, s in zip(query, stored): + q = float(q) + s = float(s) + dot += q * s + norm_a += q * q + norm_b += s * s + denominator = (norm_a ** 0.5) * (norm_b ** 0.5) + return 0.0 if denominator == 0 else dot / denominator + if metric == "inner_product": + return sum(float(q) * float(s) for q, s in zip(query, stored)) + raise ValueError("Unknown vector search metric: %s" % metric) diff --git a/paimon-python/pypaimon/table/source/vector_search_scan.py b/paimon-python/pypaimon/table/source/vector_search_scan.py index 5b8b300d5c19..1b6f9bc93051 100644 --- a/paimon-python/pypaimon/table/source/vector_search_scan.py +++ b/paimon-python/pypaimon/table/source/vector_search_scan.py @@ -20,7 +20,12 @@ from abc import ABC, abstractmethod from collections import defaultdict -from pypaimon.table.source.vector_search_split import VectorSearchSplit +from pypaimon.globalindex.global_index_coverage import GlobalIndexCoverage +from pypaimon.table.source.vector_search_split import ( + IndexVectorSearchSplit, + RawVectorSearchSplit, + VectorSearchSplit, +) from pypaimon.utils.range import Range @@ -48,11 +53,19 @@ def scan(self): class VectorSearchScanImpl(VectorSearchScan): """Implementation for VectorSearchScan.""" - def __init__(self, table, vector_column, filter_=None, partition_filter=None): + def __init__( + self, + table, + vector_column, + filter_=None, + partition_filter=None, + options=None, + ): self._table = table self._vector_column = vector_column self._filter = filter_ self._partition_filter = partition_filter + self._options = dict(options or {}) def scan(self): # type: () -> VectorSearchScanPlan @@ -86,6 +99,13 @@ def scan(self): partition_filter = self._partition_filter + def contains_field(global_index_meta, field_id): + if global_index_meta.index_field_id == field_id: + return True + if global_index_meta.extra_field_ids is not None: + return field_id in global_index_meta.extra_field_ids + return False + def index_file_filter(entry): if partition_filter is not None: if not partition_filter.test(entry.partition): @@ -96,13 +116,17 @@ def index_file_filter(entry): field_id = global_index_meta.index_field_id if vector_column.id == field_id: return True - return field_id in filter_field_ids + for filter_field_id in filter_field_ids: + if contains_field(global_index_meta, filter_field_id): + return True + return False entries = index_file_handler.scan(snapshot, index_file_filter) all_index_files = [entry.index_file for entry in entries] # Group vector index files by (rowRangeStart, rowRangeEnd). vector_by_range = defaultdict(list) + vector_index_files = [] for index_file in all_index_files: meta = index_file.global_index_meta assert meta is not None @@ -110,6 +134,12 @@ def index_file_filter(entry): continue range_key = Range(meta.row_range_start, meta.row_range_end) vector_by_range[range_key].append(index_file) + vector_index_files.append(index_file) + + vector_index_type = _vector_index_type(vector_column, all_index_files) + if vector_index_type is None: + vector_index_type = _configured_vector_index_type( + self._table, vector_column, self._options) # For each vector range, attach matching scalar index files whose # row range intersects the vector range. @@ -125,7 +155,7 @@ def index_file_filter(entry): if range_key.overlaps(scalar_range): scalar_files.append(index_file) splits.append( - VectorSearchSplit( + IndexVectorSearchSplit( range_key.from_, range_key.to, vector_files, @@ -133,4 +163,89 @@ def index_file_filter(entry): ) ) + raw_row_ranges = GlobalIndexCoverage( + self._table, + snapshot, + partition_filter, + vector_index_files, + ).unindexed_ranges(vector_column.id) + scalar_index_files = [ + f for f in all_index_files + if f.global_index_meta is not None + and f.global_index_meta.index_field_id != vector_column.id + ] + if self._filter is not None: + raw_row_ranges = Range.sort_and_merge_overlap( + raw_row_ranges + + GlobalIndexCoverage( + self._table, + snapshot, + partition_filter, + scalar_index_files, + ).unindexed_ranges(self._table.fields, self._filter), + True, + ) + if raw_row_ranges: + splits.append( + RawVectorSearchSplit( + raw_row_ranges, + _scalar_index_files_for_ranges( + all_index_files, + raw_row_ranges, + vector_column.id, + ), + vector_index_type, + ) + ) + return VectorSearchScanPlan(splits) + + +def _has_intersection(ranges, row_range): + for r in ranges: + if r.overlaps(row_range): + return True + return False + + +def _scalar_index_files_for_ranges(all_index_files, row_ranges, vector_field_id): + scalar_files = [] + for index_file in all_index_files: + meta = index_file.global_index_meta + if meta is None or meta.index_field_id == vector_field_id: + continue + if _has_intersection(row_ranges, Range(meta.row_range_start, meta.row_range_end)): + scalar_files.append(index_file) + return scalar_files + + +def _vector_index_type(vector_column, index_files): + index_type = None + for index_file in index_files: + meta = index_file.global_index_meta + if meta is None or meta.index_field_id != vector_column.id: + continue + if index_type is None: + index_type = index_file.index_type + elif index_type != index_file.index_type: + raise ValueError( + "Vector column '%s' has multiple index types: %s and %s." + % (vector_column.name, index_type, index_file.index_type) + ) + return index_type + + +def _configured_vector_index_type(table, vector_column, options): + keys = [ + "index_type", + "index-type", + "vector.index-type", + "fields.%s.index-type" % vector_column.name, + ] + table_options = getattr(getattr(table, "options", None), "options", None) + table_map = table_options.to_map() if table_options is not None else {} + for key in keys: + value = options.get(key) or table_map.get(key) + if value is not None: + return str(value).lower().strip() + return None diff --git a/paimon-python/pypaimon/table/source/vector_search_split.py b/paimon-python/pypaimon/table/source/vector_search_split.py index f8279d5c2032..a8f2428688c3 100644 --- a/paimon-python/pypaimon/table/source/vector_search_split.py +++ b/paimon-python/pypaimon/table/source/vector_search_split.py @@ -15,19 +15,55 @@ # specific language governing permissions and limitations # under the License. -"""Split of vector search.""" +"""Splits of vector search.""" from dataclasses import dataclass, field -from typing import List +from typing import List, Optional from pypaimon.index.index_file_meta import IndexFileMeta +from pypaimon.utils.range import Range -@dataclass class VectorSearchSplit: - """Split of vector search.""" + """Base split of vector search. + + The constructor keeps the historical ``VectorSearchSplit(...)`` shape as + an alias for an indexed split so existing Python callers/tests continue to + work while new code can distinguish split implementations with + ``isinstance``. + """ + + def __new__( + cls, + row_range_start=None, + row_range_end=None, + vector_index_files=None, + scalar_index_files=None, + ): + if cls is VectorSearchSplit: + return IndexVectorSearchSplit( + row_range_start, + row_range_end, + vector_index_files, + scalar_index_files, + ) + return object.__new__(cls) + + +@dataclass +class IndexVectorSearchSplit(VectorSearchSplit): + """Split to read vector index files.""" row_range_start: int row_range_end: int vector_index_files: List[IndexFileMeta] scalar_index_files: List[IndexFileMeta] = field(default_factory=list) + + +@dataclass +class RawVectorSearchSplit(VectorSearchSplit): + """Split to scan raw vectors.""" + + row_ranges: List[Range] + scalar_index_files: List[IndexFileMeta] = field(default_factory=list) + index_type: Optional[str] = None diff --git a/paimon-python/pypaimon/tests/global_index_test.py b/paimon-python/pypaimon/tests/global_index_test.py index 873e93358455..60ca559cda0c 100644 --- a/paimon-python/pypaimon/tests/global_index_test.py +++ b/paimon-python/pypaimon/tests/global_index_test.py @@ -20,8 +20,15 @@ import pyarrow as pa +from pypaimon.common.options.core_options import CoreOptions, GlobalIndexSearchMode +from pypaimon.common.options.options import Options +from pypaimon.common.predicate import Predicate +from pypaimon.common.predicate_builder import PredicateBuilder +from pypaimon.globalindex.global_index_meta import GlobalIndexMeta from pypaimon.globalindex.global_index_result import GlobalIndexResult +from pypaimon.index.index_file_meta import IndexFileMeta from pypaimon.index.index_file_handler import IndexFileHandler +from pypaimon.schema.data_types import AtomicType, DataField from pypaimon.snapshot.snapshot_manager import SnapshotManager from pypaimon.tests.data_evolution_test_helpers import ( BatchModeMixin, @@ -49,6 +56,193 @@ def test_chained_and(self): self.assertEqual(result.results().cardinality(), 10001) +class _CoverageOptions: + def __init__(self, mode): + self.options = Options({"global-index.search-mode": mode}) + + def global_index_search_mode(self): + return CoreOptions(self.options).global_index_search_mode() + + +class _CoverageTable: + def __init__(self, mode, data_ranges=None): + self.options = _CoverageOptions(mode) + self.fields = [ + DataField(0, "id", AtomicType("INT")), + DataField(1, "name", AtomicType("STRING")), + ] + self._data_ranges = data_ranges or [] + + def data_ranges_for_global_index_coverage(self, snapshot, partition_filter): + return self._data_ranges + + +class _CoverageSnapshot: + def __init__(self, next_row_id): + self.next_row_id = next_row_id + + +def _coverage_index_file(field_id, start, end, extra_field_ids=None): + return IndexFileMeta( + index_type="btree", + file_name="idx-%s-%s-%s" % (field_id, start, end), + file_size=1, + row_count=end - start + 1, + global_index_meta=GlobalIndexMeta( + row_range_start=start, + row_range_end=end, + index_field_id=field_id, + extra_field_ids=extra_field_ids, + index_meta=b"", + ), + ) + + +class GlobalIndexCoverageTest(unittest.TestCase): + + def test_fast_mode_does_not_return_unindexed_ranges(self): + from pypaimon.globalindex.global_index_coverage import GlobalIndexCoverage + + table = _CoverageTable(GlobalIndexSearchMode.FAST) + coverage = GlobalIndexCoverage( + table, + _CoverageSnapshot(10), + None, + [_coverage_index_file(0, 0, 4)], + ) + + self.assertEqual([], coverage.unindexed_ranges(0)) + + def test_full_mode_uses_snapshot_next_row_id(self): + from pypaimon.globalindex.global_index_coverage import GlobalIndexCoverage + + table = _CoverageTable("full") + coverage = GlobalIndexCoverage( + table, + _CoverageSnapshot(10), + None, + [_coverage_index_file(0, 0, 4)], + ) + + self.assertEqual([Range(5, 9)], coverage.unindexed_ranges(0)) + + def test_full_mode_intersects_coverage_for_all_predicate_fields(self): + from pypaimon.globalindex.global_index_coverage import GlobalIndexCoverage + + table = _CoverageTable("full") + coverage = GlobalIndexCoverage( + table, + _CoverageSnapshot(10), + None, + [ + _coverage_index_file(0, 0, 9), + _coverage_index_file(1, 0, 4), + ], + ) + predicate = PredicateBuilder.and_predicates( + [ + Predicate(method="equal", index=0, field="id", literals=[1]), + Predicate(method="equal", index=1, field="name", literals=["a"]), + ] + ) + + self.assertEqual( + [Range(5, 9)], + coverage.unindexed_ranges(table.fields, predicate), + ) + + def test_extra_fields_count_as_index_coverage(self): + from pypaimon.globalindex.global_index_coverage import GlobalIndexCoverage + + table = _CoverageTable("full") + coverage = GlobalIndexCoverage( + table, + _CoverageSnapshot(10), + None, + [_coverage_index_file(0, 0, 9, extra_field_ids=[1])], + ) + + self.assertEqual([], coverage.unindexed_ranges(1)) + + def test_detail_mode_uses_table_data_ranges(self): + from pypaimon.globalindex.global_index_coverage import GlobalIndexCoverage + + table = _CoverageTable("detail", data_ranges=[Range(0, 2), Range(7, 9)]) + coverage = GlobalIndexCoverage( + table, + _CoverageSnapshot(10), + None, + [_coverage_index_file(0, 0, 4)], + ) + + self.assertEqual([Range(7, 9)], coverage.unindexed_ranges(0)) + + +class GlobalIndexScalarFallbackTest(unittest.TestCase): + + def test_eval_global_index_merges_unindexed_rows_when_index_scan_succeeds(self): + from pypaimon.read.scanner.file_scanner import FileScanner + + class _Options: + def global_index_enabled(self): + return True + + class _Table: + options = _Options() + + predicate = Predicate(method="equal", index=0, field="id", literals=[1]) + scanner = FileScanner.__new__(FileScanner) + scanner.predicate = predicate + scanner.partition_key_predicate = None + scanner.table = _Table() + + index_result = GlobalIndexResult.from_range(Range(1, 1)) + unindexed = GlobalIndexResult.from_range(Range(5, 6)) + fake_scanner = unittest.mock.MagicMock() + fake_scanner.scan.return_value = index_result + fake_scanner.unindexed_rows.return_value = unindexed + fake_scanner.__enter__.return_value = fake_scanner + fake_scanner.__exit__.return_value = None + + with unittest.mock.patch( + "pypaimon.globalindex.global_index_scanner.GlobalIndexScanner.create", + return_value=fake_scanner): + result = scanner._eval_global_index(snapshot=object()) + + self.assertEqual( + [Range(1, 1), Range(5, 6)], + result.results().to_range_list(), + ) + + def test_eval_global_index_keeps_none_as_full_scan(self): + from pypaimon.read.scanner.file_scanner import FileScanner + + class _Options: + def global_index_enabled(self): + return True + + class _Table: + options = _Options() + + scanner = FileScanner.__new__(FileScanner) + scanner.predicate = Predicate( + method="equal", index=0, field="id", literals=[1]) + scanner.partition_key_predicate = None + scanner.table = _Table() + + fake_scanner = unittest.mock.MagicMock() + fake_scanner.scan.return_value = None + fake_scanner.__enter__.return_value = fake_scanner + fake_scanner.__exit__.return_value = None + + with unittest.mock.patch( + "pypaimon.globalindex.global_index_scanner.GlobalIndexScanner.create", + return_value=fake_scanner): + result = scanner._eval_global_index(snapshot=object()) + + self.assertIsNone(result) + + class PlanSnapshotFetchRegressionTest( BatchModeMixin, DataEvolutionTestBase, unittest.TestCase): diff --git a/paimon-python/pypaimon/tests/vector_search_filter_test.py b/paimon-python/pypaimon/tests/vector_search_filter_test.py index f2b453346ecb..8c3a5261a4c8 100644 --- a/paimon-python/pypaimon/tests/vector_search_filter_test.py +++ b/paimon-python/pypaimon/tests/vector_search_filter_test.py @@ -44,6 +44,7 @@ from pypaimon.table.row.generic_row import GenericRow from pypaimon.table.source.vector_search_builder import VectorSearchBuilderImpl from pypaimon.utils.roaring_bitmap import RoaringBitmap64 +from pypaimon.utils.range import Range # ----------------------------- table stubs --------------------------------- @@ -122,6 +123,15 @@ def _entry(partition_row, field_id, index_type, file_name, def _patch_snapshot(testcase, entries): """Stub IndexFileHandler.scan + snapshot resolution.""" + mock.patch.stopall() + for attr in ("_scan_patch", "_travel_patch"): + patcher = getattr(testcase, attr, None) + if patcher is not None: + try: + patcher.stop() + except RuntimeError: + pass + def _scan(snapshot, entry_filter=None): if entry_filter is None: return list(entries) @@ -1005,6 +1015,172 @@ def close(self_inner): self.assertEqual("oss://bucket/id-btree-0.index", captured_io_metas[0][0].external_path) + def test_full_mode_scan_adds_raw_split_for_unindexed_vector_rows(self): + from pypaimon.common.options.core_options import CoreOptions + from pypaimon.common.options.options import Options + from pypaimon.table.source.vector_search_split import ( + IndexVectorSearchSplit, + RawVectorSearchSplit, + ) + + class _Options: + options = Options({"global-index.search-mode": "full"}) + + def global_index_search_mode(self_inner): + return CoreOptions(self_inner.options).global_index_search_mode() + + class _Snapshots: + def get_latest_snapshot(self_inner): + return types.SimpleNamespace(next_row_id=10) + + table = _StubTable(fields=[self.id_field, self.embedding_field], + entries=[self.entries[0]]) + table.options = _Options() + table.snapshot_manager = lambda: _Snapshots() + self._scan_patch.stop() + self._travel_patch.stop() + _patch_snapshot(self, [self.entries[0]]) + self._travel_patch.stop() + + splits = ( + VectorSearchBuilderImpl(table) + .with_vector_column("embedding") + .with_query_vector([1.0, 0.0, 0.0, 0.0]) + .with_limit(3) + .new_vector_search_scan() + .scan() + .splits() + ) + + self.assertEqual(2, len(splits)) + self.assertTrue(any(isinstance(s, IndexVectorSearchSplit) + for s in splits)) + raw = [s for s in splits if isinstance(s, RawVectorSearchSplit)] + self.assertEqual(1, len(raw)) + self.assertEqual([Range(5, 9)], raw[0].row_ranges) + + def test_full_mode_scan_adds_raw_split_for_uncovered_scalar_filter(self): + from pypaimon.common.options.core_options import CoreOptions + from pypaimon.common.options.options import Options + from pypaimon.table.source.vector_search_split import RawVectorSearchSplit + + class _Options: + options = Options({"global-index.search-mode": "full"}) + + def global_index_search_mode(self_inner): + return CoreOptions(self_inner.options).global_index_search_mode() + + class _Snapshots: + def get_latest_snapshot(self_inner): + return types.SimpleNamespace(next_row_id=10) + + # Vector index covers all rows, but there is no scalar id index, so + # full mode must produce a raw split for the scalar-filtered path. + table = _StubTable(fields=[self.id_field, self.embedding_field], + entries=[ + _entry(None, field_id=1, + index_type="lumina-vector-ann", + file_name="vec-all.index", + row_range_start=0, + row_range_end=9) + ]) + table.options = _Options() + table.snapshot_manager = lambda: _Snapshots() + self._scan_patch.stop() + self._travel_patch.stop() + _patch_snapshot(self, table._entries) + self._travel_patch.stop() + filter_pred = Predicate(method="greaterOrEqual", index=0, field="id", + literals=[5]) + + splits = ( + VectorSearchBuilderImpl(table) + .with_vector_column("embedding") + .with_query_vector([1.0, 0.0, 0.0, 0.0]) + .with_limit(3) + .with_filter(filter_pred) + .new_vector_search_scan() + .scan() + .splits() + ) + + raw = [s for s in splits if isinstance(s, RawVectorSearchSplit)] + self.assertEqual(1, len(raw)) + self.assertEqual([Range(0, 9)], raw[0].row_ranges) + + def test_scan_threads_builder_options_to_raw_split_index_type(self): + from pypaimon.common.options.core_options import CoreOptions + from pypaimon.common.options.options import Options + from pypaimon.table.source.vector_search_split import RawVectorSearchSplit + + class _Options: + options = Options({"global-index.search-mode": "full"}) + + def global_index_search_mode(self_inner): + return CoreOptions(self_inner.options).global_index_search_mode() + + class _Snapshots: + def get_latest_snapshot(self_inner): + return types.SimpleNamespace(next_row_id=10) + + table = _StubTable(fields=[self.id_field, self.embedding_field], + entries=[]) + table.options = _Options() + table.snapshot_manager = lambda: _Snapshots() + self._scan_patch.stop() + self._travel_patch.stop() + _patch_snapshot(self, []) + self._travel_patch.stop() + + splits = ( + VectorSearchBuilderImpl(table) + .with_vector_column("embedding") + .with_query_vector([1.0, 0.0, 0.0, 0.0]) + .with_limit(3) + .with_option("index-type", "ivf-flat") + .new_vector_search_scan() + .scan() + .splits() + ) + + raw = [s for s in splits if isinstance(s, RawVectorSearchSplit)] + self.assertEqual(1, len(raw)) + self.assertEqual("ivf-flat", raw[0].index_type) + + def test_scan_attaches_scalar_index_when_filter_hits_extra_field(self): + id_name_index = _entry(None, field_id=2, index_type="btree", + file_name="name-id.index", + row_range_start=0, + row_range_end=9) + id_name_index.index_file.global_index_meta.extra_field_ids = [0] + table = _StubTable(fields=[ + self.id_field, + self.embedding_field, + _field(2, "name", "STRING"), + ], entries=[ + self.entries[0], + id_name_index, + ]) + self._scan_patch.stop() + self._travel_patch.stop() + _patch_snapshot(self, table._entries) + filter_pred = Predicate(method="equal", index=0, field="id", + literals=[5]) + + splits = ( + VectorSearchBuilderImpl(table) + .with_vector_column("embedding") + .with_query_vector([1.0, 0.0, 0.0, 0.0]) + .with_limit(3) + .with_filter(filter_pred) + .new_vector_search_scan() + .scan() + .splits() + ) + + self.assertEqual(["name-id.index"], + [f.file_name for f in splits[0].scalar_index_files]) + class VectorSearchMultiShardScalarTest(unittest.TestCase): """Scalar pre-filter across multiple btree shards of the same field. @@ -1016,6 +1192,9 @@ class VectorSearchMultiShardScalarTest(unittest.TestCase): - An empty first shard does NOT short-circuit subsequent shards. """ + def tearDown(self): + mock.patch.stopall() + def test_hit_only_in_later_shard_returns_global_row_id(self): from pypaimon.globalindex.global_index_result import GlobalIndexResult from pypaimon.globalindex.global_index_scanner import ( @@ -1202,6 +1381,93 @@ def close(self_inner): self.assertIsNotNone(result) self.assertEqual([3], sorted(list(result.results()))) + def test_scanner_reports_unindexed_rows_for_full_mode(self): + from pypaimon.common.options.core_options import CoreOptions + from pypaimon.common.options.options import Options + from pypaimon.globalindex.global_index_scanner import ( + GlobalIndexScanner, + ) + + class _Options: + options = Options({"global-index.search-mode": "full"}) + + def global_index_search_mode(self_inner): + return CoreOptions(self_inner.options).global_index_search_mode() + + def global_index_thread_num(self_inner): + return 32 + + class _Snapshots: + def get_latest_snapshot(self_inner): + return types.SimpleNamespace(next_row_id=10) + + id_field = _field(0, "id") + emb_field = _field(1, "embedding", "FLOAT") + indexed = _entry(None, field_id=0, index_type="btree", + file_name="id-0.index", + row_range_start=0, row_range_end=4).index_file + table = _StubTable(fields=[id_field, emb_field], entries=[]) + table.options = _Options() + table.snapshot_manager = lambda: _Snapshots() + + scanner = GlobalIndexScanner.create(table, index_files=[indexed]) + try: + result = scanner.unindexed_rows( + Predicate(method="equal", index=0, field="id", literals=[7])) + finally: + scanner.close() + + self.assertEqual([Range(5, 9)], result.results().to_range_list()) + + def test_scanner_create_selects_extra_field_indexes(self): + from pypaimon.globalindex.global_index_scanner import ( + GlobalIndexScanner, + ) + + name_field = _field(0, "name", "STRING") + id_field = _field(1, "id") + emb_field = _field(2, "embedding", "FLOAT") + multi = _entry(None, field_id=0, index_type="btree", + file_name="name-id.index", + row_range_start=0, row_range_end=9).index_file + multi.global_index_meta.extra_field_ids = [1] + table = _StubTable(fields=[name_field, id_field, emb_field], + entries=[ + IndexManifestEntry(kind=0, partition=None, + bucket=0, index_file=multi) + ]) + _patch_snapshot(self, table._entries) + + class _StubBTreeReader: + def __init__(self_inner, key_serializer, file_io, index_path, + io_meta): + pass + + def visit_equal(self_inner, literal): + return GlobalIndexResult.create_empty() + + def close(self_inner): + pass + + with mock.patch( + "pypaimon.globalindex.btree.lazy_filtered_btree_reader.BTreeIndexReader", + _StubBTreeReader): + with mock.patch( + "pypaimon.globalindex.sorted_file_global_index_reader.SortedIndexFileMeta.deserialize", + return_value=BTreeIndexMeta(first_key=b'', last_key=b'zzzz', has_nulls=False)): + scanner = GlobalIndexScanner.create( + table, + predicate=Predicate(method="equal", index=1, field="id", + literals=[3]), + ) + try: + self.assertIsNotNone(scanner) + readers = scanner._evaluator._readers_function(id_field) + self.assertTrue(readers) + finally: + if scanner is not None: + scanner.close() + class VectorSearchPartitionedFilterTest(unittest.TestCase): """Partitioned-table paths: with_filter auto-split + partition-filter @@ -1489,6 +1755,159 @@ def __exit__(self_inner, *a): scores = sorted(result.score_getter()(rid) for rid in result.results()) self.assertEqual(scores, [float(i) for i in range(1190, 1200)]) + def test_read_merges_raw_search_results(self): + from pypaimon.globalindex.vector_search_result import ( + DictBasedScoredIndexResult, + ) + from pypaimon.table.source.vector_search_read import VectorSearchReadImpl + from pypaimon.table.source.vector_search_split import ( + IndexVectorSearchSplit, + RawVectorSearchSplit, + ) + + embedding_field = _field(1, "embedding", "FLOAT") + entry = _entry(None, field_id=1, index_type="lumina-vector-ann", + file_name="vec.index", + row_range_start=0, row_range_end=4) + table = _StubTable(fields=[embedding_field], entries=[entry]) + + def _fake_create(index_type, file_io, index_path, + index_io_meta_list, options=None): + class _FakeReader: + def visit_vector_search(self_inner, vs): + return _completed_future( + DictBasedScoredIndexResult({1: 0.1})) + + def close(self_inner): + pass + + return _FakeReader() + + split = IndexVectorSearchSplit( + row_range_start=0, + row_range_end=4, + vector_index_files=[entry.index_file], + ) + raw = RawVectorSearchSplit([Range(5, 9)], [], "lumina-vector-ann") + + with mock.patch( + "pypaimon.table.source.vector_search_read._create_vector_reader", + side_effect=_fake_create): + reader = VectorSearchReadImpl( + table, limit=2, vector_column=embedding_field, + query_vector=[1.0], filter_=None) + with mock.patch.object( + reader, + "_read_raw_search", + return_value=DictBasedScoredIndexResult({8: 0.9})) as raw_read: + result = reader.read([split, raw]) + + raw_read.assert_called_once() + self.assertEqual([1, 8], sorted(list(result.results()))) + + def test_read_uses_empty_index_prefilter_when_scalar_index_missing(self): + from pypaimon.table.source.vector_search_read import VectorSearchReadImpl + from pypaimon.table.source.vector_search_split import IndexVectorSearchSplit + + id_field = _field(0, "id") + embedding_field = _field(1, "embedding", "FLOAT") + entry = _entry(None, field_id=1, index_type="lumina-vector-ann", + file_name="vec.index", + row_range_start=0, row_range_end=4) + table = _StubTable(fields=[id_field, embedding_field], entries=[entry]) + filter_pred = Predicate(method="equal", index=0, field="id", + literals=[3]) + split = IndexVectorSearchSplit( + row_range_start=0, + row_range_end=4, + vector_index_files=[entry.index_file], + scalar_index_files=[], + ) + + reader = VectorSearchReadImpl( + table, limit=2, vector_column=embedding_field, + query_vector=[1.0], filter_=filter_pred) + + pre_filters = reader._pre_filters([split]) + + self.assertEqual(1, len(pre_filters)) + self.assertEqual(0, pre_filters[0].cardinality()) + + def test_raw_search_uses_partition_filter_and_index_type_metric(self): + import pyarrow as pa + + from pypaimon.table.source.vector_search_read import VectorSearchReadImpl + + id_field = _field(0, "id") + embedding_field = _field(1, "embedding", "FLOAT") + table = _StubTable(fields=[id_field, embedding_field], entries=[]) + partition_filter = Predicate(method="equal", index=0, field="pt", + literals=[1]) + filter_pred = Predicate(method="greaterOrEqual", index=0, field="id", + literals=[0]) + calls = {} + + class _Plan: + def splits(self_inner): + return ["split"] + + class _Scan: + def with_global_index_result(self_inner, result): + calls["global_index_ranges"] = result.results().to_range_list() + return self_inner + + def plan(self_inner): + return _Plan() + + class _Read: + def to_arrow(self_inner, splits): + calls["splits"] = list(splits) + return pa.table({ + "id": pa.array([5, 6], type=pa.int32()), + "embedding": pa.array([[1.0], [0.0]]), + "_ROW_ID": pa.array([5, 6], type=pa.int64()), + }) + + class _Builder: + def with_partition_filter(self_inner, predicate): + calls["partition_filter"] = predicate + return self_inner + + def with_filter(self_inner, predicate): + calls["filter"] = predicate + return self_inner + + def with_projection(self_inner, projection): + calls["projection"] = list(projection) + return self_inner + + def new_scan(self_inner): + return _Scan() + + def new_read(self_inner): + return _Read() + + table.new_read_builder = lambda: _Builder() + reader = VectorSearchReadImpl( + table, + limit=1, + vector_column=embedding_field, + query_vector=[1.0], + filter_=filter_pred, + partition_filter=partition_filter, + options={"ivf-flat.metric": "inner_product"}, + ) + + result = reader._read_raw_search( + [Range(5, 6)], None, [1.0], "ivf-flat") + + self.assertIs(partition_filter, calls["partition_filter"]) + self.assertIs(filter_pred, calls["filter"]) + self.assertEqual([Range(5, 6)], calls["global_index_ranges"]) + self.assertIn("_ROW_ID", calls["projection"]) + self.assertEqual(["split"], calls["splits"]) + self.assertEqual([5], sorted(list(result.results()))) + def tearDown(self): mock.patch.stopall() From bb2460b8a5a44accc3c0719e106414d2a3fc469c Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Sun, 21 Jun 2026 21:04:52 +0800 Subject: [PATCH 2/4] [python] Fix vector search test lint --- .../tests/vector_search_filter_test.py | 30 +++++++++++-------- 1 file changed, 17 insertions(+), 13 deletions(-) diff --git a/paimon-python/pypaimon/tests/vector_search_filter_test.py b/paimon-python/pypaimon/tests/vector_search_filter_test.py index 8c3a5261a4c8..fe63890a4287 100644 --- a/paimon-python/pypaimon/tests/vector_search_filter_test.py +++ b/paimon-python/pypaimon/tests/vector_search_filter_test.py @@ -1076,14 +1076,16 @@ def get_latest_snapshot(self_inner): # Vector index covers all rows, but there is no scalar id index, so # full mode must produce a raw split for the scalar-filtered path. - table = _StubTable(fields=[self.id_field, self.embedding_field], - entries=[ - _entry(None, field_id=1, - index_type="lumina-vector-ann", - file_name="vec-all.index", - row_range_start=0, - row_range_end=9) - ]) + table = _StubTable( + fields=[self.id_field, self.embedding_field], + entries=[ + _entry(None, field_id=1, + index_type="lumina-vector-ann", + file_name="vec-all.index", + row_range_start=0, + row_range_end=9) + ], + ) table.options = _Options() table.snapshot_manager = lambda: _Snapshots() self._scan_patch.stop() @@ -1431,11 +1433,13 @@ def test_scanner_create_selects_extra_field_indexes(self): file_name="name-id.index", row_range_start=0, row_range_end=9).index_file multi.global_index_meta.extra_field_ids = [1] - table = _StubTable(fields=[name_field, id_field, emb_field], - entries=[ - IndexManifestEntry(kind=0, partition=None, - bucket=0, index_file=multi) - ]) + table = _StubTable( + fields=[name_field, id_field, emb_field], + entries=[ + IndexManifestEntry(kind=0, partition=None, + bucket=0, index_file=multi) + ], + ) _patch_snapshot(self, table._entries) class _StubBTreeReader: From 2c5504771df4d1d160cd667ffbf4cb410c8fac49 Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Sun, 21 Jun 2026 22:03:04 +0800 Subject: [PATCH 3/4] [python] Add mixed raw fallback tests --- .../java/org/apache/paimon/JavaPyE2ETest.java | 71 +++++++++++ paimon-python/dev/run_mixed_tests.sh | 79 +++++++++++- .../tests/e2e/java_py_read_write_test.py | 62 +++++++++ .../java/org/apache/paimon/JavaPyE2ETest.java | 119 ++++++++++++++++++ 4 files changed, 330 insertions(+), 1 deletion(-) diff --git a/paimon-core/src/test/java/org/apache/paimon/JavaPyE2ETest.java b/paimon-core/src/test/java/org/apache/paimon/JavaPyE2ETest.java index 2a7584a0bcd8..695cc29d76c7 100644 --- a/paimon-core/src/test/java/org/apache/paimon/JavaPyE2ETest.java +++ b/paimon-core/src/test/java/org/apache/paimon/JavaPyE2ETest.java @@ -516,6 +516,77 @@ public void testBtreeIndexWrite() throws Exception { testBtreeIndexWriteNull(); } + @Test + @EnabledIfSystemProperty(named = "run.e2e.tests", matches = "true") + public void testBtreeRawFallbackWrite() throws Exception { + RowType rowType = + RowType.of( + new DataType[] {DataTypes.STRING(), DataTypes.STRING()}, + new String[] {"k", "v"}); + Options options = new Options(); + Path tablePath = new Path(warehouse.toString() + "/default.db/test_btree_raw_fallback"); + LocalFileIO.create().delete(tablePath, true); + options.set(PATH, tablePath.toString()); + options.set(ROW_TRACKING_ENABLED, true); + options.set(DATA_EVOLUTION_ENABLED, true); + options.set(GLOBAL_INDEX_ENABLED, true); + TableSchema tableSchema = + SchemaUtils.forceCommit( + new SchemaManager(LocalFileIO.create(), tablePath), + new Schema( + rowType.getFields(), + Collections.emptyList(), + Collections.emptyList(), + options.toMap(), + "")); + AppendOnlyFileStoreTable table = + new AppendOnlyFileStoreTable( + FileIOFinder.find(tablePath), + tablePath, + tableSchema, + CatalogEnvironment.empty()); + + BatchWriteBuilder writeBuilder = table.newBatchWriteBuilder(); + try (BatchTableWrite write = writeBuilder.newWrite(); + BatchTableCommit commit = writeBuilder.newCommit()) { + write.write( + GenericRow.of(BinaryString.fromString("k1"), BinaryString.fromString("v1"))); + write.write( + GenericRow.of(BinaryString.fromString("k2"), BinaryString.fromString("v2"))); + write.write( + GenericRow.of(BinaryString.fromString("k3"), BinaryString.fromString("v3"))); + commit.commit(write.prepareCommit()); + } + + SortedGlobalIndexBuilder builder = + new SortedGlobalIndexBuilder(table, "btree").withIndexField("k"); + try (BatchTableCommit commit = writeBuilder.newCommit()) { + commit.commit( + builder.build( + builder.scan() + .map(org.apache.paimon.utils.Pair::getValue) + .orElseThrow( + () -> + new IllegalStateException( + "Expected scan result when building index.")) + .get(0), + IOManager.create(warehouse.toString()))); + } + + try (BatchTableWrite write = writeBuilder.newWrite(); + BatchTableCommit commit = writeBuilder.newCommit()) { + write.write( + GenericRow.of(BinaryString.fromString("k4"), BinaryString.fromString("v4"))); + commit.commit(write.prepareCommit()); + } + + List indexEntries = + table.indexManifestFileReader().read(table.latestSnapshot().get().indexManifest); + assertThat(indexEntries) + .singleElement() + .matches(entry -> entry.indexFile().rowCount() == 3); + } + @Test @EnabledIfSystemProperty(named = "run.e2e.tests", matches = "true") public void testBitmapIndexWrite() throws Exception { diff --git a/paimon-python/dev/run_mixed_tests.sh b/paimon-python/dev/run_mixed_tests.sh index c3d91893df08..8da1cb345b91 100755 --- a/paimon-python/dev/run_mixed_tests.sh +++ b/paimon-python/dev/run_mixed_tests.sh @@ -219,6 +219,29 @@ run_btree_index_test() { fi } +run_btree_raw_fallback_test() { + echo -e "${YELLOW}=== Running BTree Raw Fallback Test (Java Write, Python Read) ===${NC}" + + cd "$PROJECT_ROOT" + + echo "Running Maven test for JavaPyE2ETest.testBtreeRawFallbackWrite..." + if mvn test -Dtest=org.apache.paimon.JavaPyE2ETest#testBtreeRawFallbackWrite -pl paimon-core -am -q -DfailIfNoTests=false -Drun.e2e.tests=true; then + echo -e "${GREEN}✓ Java test completed successfully${NC}" + else + echo -e "${RED}✗ Java test failed${NC}" + return 1 + fi + cd "$PAIMON_PYTHON_DIR" + echo "Running Python test for JavaPyReadWriteTest.test_read_btree_raw_fallback..." + if python -m pytest java_py_read_write_test.py::JavaPyReadWriteTest::test_read_btree_raw_fallback -v; then + echo -e "${GREEN}✓ Python test completed successfully${NC}" + return 0 + else + echo -e "${RED}✗ Python test failed${NC}" + return 1 + fi +} + run_bitmap_index_test() { echo -e "${YELLOW}=== Step 6b: Running Bitmap Index Test (Java Write, Python Read) ===${NC}" @@ -541,6 +564,32 @@ run_vindex_vector_test() { fi } +run_vindex_vector_raw_fallback_test() { + echo -e "${YELLOW}=== Running paimon-vindex Vector Raw Fallback Test (Java Write, Python Read) ===${NC}" + + cd "$PROJECT_ROOT" + + echo "Running Maven test for JavaPyE2ETest.testVindexVectorRawFallbackWrite..." + if mvn test -Dtest=org.apache.paimon.JavaPyE2ETest#testVindexVectorRawFallbackWrite -pl paimon-vector -am -q -DfailIfNoTests=false -Drun.e2e.tests=true; then + echo -e "${GREEN}✓ Java test completed successfully${NC}" + else + echo -e "${RED}✗ Java test failed${NC}" + return 1 + fi + cd "$PAIMON_PYTHON_DIR" + if ! ensure_paimon_vindex; then + return 1 + fi + echo "Running Python test for JavaPyReadWriteTest.test_read_vindex_vector_raw_fallback..." + if python -m pytest java_py_read_write_test.py::JavaPyReadWriteTest::test_read_vindex_vector_raw_fallback -v; then + echo -e "${GREEN}✓ Python test completed successfully${NC}" + return 0 + else + echo -e "${RED}✗ Python test failed${NC}" + return 1 + fi +} + run_compact_conflict_test() { echo -e "${YELLOW}=== Running Compact Conflict Test (Java Write Base, Python Shard Update + Java Compact) ===${NC}" @@ -796,6 +845,7 @@ main() { local java_read_result=0 local pk_dv_result=0 local btree_index_result=0 + local btree_raw_fallback_result=0 local bitmap_index_result=0 local compressed_global_index_result=0 local compressed_text_result=0 @@ -804,6 +854,7 @@ main() { local lumina_vector_result=0 local lumina_vector_btree_result=0 local vindex_vector_result=0 + local vindex_vector_raw_fallback_result=0 local compact_conflict_result=0 local blob_compact_conflict_result=0 local blob_alter_compact_result=0 @@ -873,6 +924,13 @@ main() { echo "" + # Run BTree raw fallback test (Java write indexed + unindexed rows, Python read) + if ! run_btree_raw_fallback_test; then + btree_raw_fallback_result=1 + fi + + echo "" + # Run Bitmap index test (Java write, Python read) if ! run_bitmap_index_test; then bitmap_index_result=1 @@ -967,9 +1025,16 @@ main() { if ! run_vindex_vector_test; then vindex_vector_result=1 fi + + echo "" + + if ! run_vindex_vector_raw_fallback_test; then + vindex_vector_raw_fallback_result=1 + fi else echo -e "${YELLOW}⏭ Skipping paimon-vindex Vector Index Test (requires Python >= 3.9, current: $PYTHON_VERSION)${NC}" vindex_vector_result=0 + vindex_vector_raw_fallback_result=0 fi echo "" @@ -1073,6 +1138,12 @@ main() { echo -e "${RED}✗ BTree Index Test (Java Write, Python Read): FAILED${NC}" fi + if [[ $btree_raw_fallback_result -eq 0 ]]; then + echo -e "${GREEN}✓ BTree Raw Fallback Test (Java Write, Python Read): PASSED${NC}" + else + echo -e "${RED}✗ BTree Raw Fallback Test (Java Write, Python Read): FAILED${NC}" + fi + if [[ $bitmap_index_result -eq 0 ]]; then echo -e "${GREEN}✓ Bitmap Index Test (Java Write, Python Read): PASSED${NC}" else @@ -1145,6 +1216,12 @@ main() { echo -e "${RED}✗ paimon-vindex Vector Index Test (Java Write, Python Read): FAILED${NC}" fi + if [[ $vindex_vector_raw_fallback_result -eq 0 ]]; then + echo -e "${GREEN}✓ paimon-vindex Vector Raw Fallback Test (Java Write, Python Read): PASSED${NC}" + else + echo -e "${RED}✗ paimon-vindex Vector Raw Fallback Test (Java Write, Python Read): FAILED${NC}" + fi + if [[ $compact_conflict_result -eq 0 ]]; then echo -e "${GREEN}✓ Compact Conflict Test (Java Write+Compact, Python Read): PASSED${NC}" else @@ -1198,7 +1275,7 @@ main() { # Clean up warehouse directory after all tests cleanup_warehouse - if [[ $java_write_result -eq 0 && $python_read_result -eq 0 && $python_write_result -eq 0 && $java_read_result -eq 0 && $pk_dv_result -eq 0 && $btree_index_result -eq 0 && $bitmap_index_result -eq 0 && $compressed_global_index_result -eq 0 && $compressed_text_result -eq 0 && $tantivy_fulltext_result -eq 0 && $lumina_vector_result -eq 0 && $lumina_vector_btree_result -eq 0 && $vindex_vector_result -eq 0 && $compact_conflict_result -eq 0 && $blob_compact_conflict_result -eq 0 && $blob_alter_compact_result -eq 0 && $data_evolution_result -eq 0 && $data_evolution_py_write_result -eq 0 && $java_variant_write_py_read_result -eq 0 && $py_variant_write_java_read_result -eq 0 && $vector_append_table_result -eq 0 && $vector_dedicated_java_write_result -eq 0 && $vector_dedicated_py_write_result -eq 0 && $multi_vector_dedicated_java_write_result -eq 0 && $multi_vector_dedicated_py_write_result -eq 0 && $row_format_result -eq 0 ]]; then + if [[ $java_write_result -eq 0 && $python_read_result -eq 0 && $python_write_result -eq 0 && $java_read_result -eq 0 && $pk_dv_result -eq 0 && $btree_index_result -eq 0 && $btree_raw_fallback_result -eq 0 && $bitmap_index_result -eq 0 && $compressed_global_index_result -eq 0 && $compressed_text_result -eq 0 && $tantivy_fulltext_result -eq 0 && $lumina_vector_result -eq 0 && $lumina_vector_btree_result -eq 0 && $vindex_vector_result -eq 0 && $vindex_vector_raw_fallback_result -eq 0 && $compact_conflict_result -eq 0 && $blob_compact_conflict_result -eq 0 && $blob_alter_compact_result -eq 0 && $data_evolution_result -eq 0 && $data_evolution_py_write_result -eq 0 && $java_variant_write_py_read_result -eq 0 && $py_variant_write_java_read_result -eq 0 && $vector_append_table_result -eq 0 && $vector_dedicated_java_write_result -eq 0 && $vector_dedicated_py_write_result -eq 0 && $multi_vector_dedicated_java_write_result -eq 0 && $multi_vector_dedicated_py_write_result -eq 0 && $row_format_result -eq 0 ]]; then echo -e "${GREEN}🎉 All tests passed! Java-Python interoperability verified.${NC}" return 0 else diff --git a/paimon-python/pypaimon/tests/e2e/java_py_read_write_test.py b/paimon-python/pypaimon/tests/e2e/java_py_read_write_test.py index f940f38a698a..691bd16dd9f9 100644 --- a/paimon-python/pypaimon/tests/e2e/java_py_read_write_test.py +++ b/paimon-python/pypaimon/tests/e2e/java_py_read_write_test.py @@ -448,6 +448,27 @@ def test_read_btree_index_table(self): if sys.version_info[:2] >= (3, 7): self._test_index_manifest_inherited_after_write() + def test_read_btree_raw_fallback(self): + table = self.catalog.get_table('default.test_btree_raw_fallback') + fast_builder = table.new_read_builder() + fast_predicate = fast_builder.new_predicate_builder().equal('k', 'k4') + fast_builder.with_filter(fast_predicate) + fast_result = fast_builder.new_read().to_arrow( + fast_builder.new_scan().plan().splits()) + self.assertEqual(0, fast_result.num_rows) + + full_table = table.copy({'global-index.search-mode': 'full'}) + read_builder = full_table.new_read_builder() + read_builder.with_filter( + read_builder.new_predicate_builder().equal('k', 'k4')) + actual = read_builder.new_read().to_arrow( + read_builder.new_scan().plan().splits()) + expected = pa.Table.from_pydict({ + 'k': ['k4'], + 'v': ['v4'], + }) + self.assertEqual(expected, actual) + def _test_read_btree_index_generic(self, table_name: str, k, k_type): table = self.catalog.get_table('default.' + table_name) read_builder: ReadBuilder = table.new_read_builder() @@ -1253,6 +1274,47 @@ def test_read_vindex_vector_index(self): print(f"paimon-vindex vector search matched rows: ids={ids}") self.assertIn(0, ids) + def test_read_vindex_vector_raw_fallback(self): + """Test raw fallback for a paimon-vindex vector index built by Java.""" + if sys.version_info < (3, 9): + self.skipTest("paimon-vindex requires Python >= 3.9") + try: + import paimon_vindex # noqa: F401 + except ImportError: + self.skipTest("paimon-vindex is not installed") + + table = self.catalog.get_table( + 'default.test_vindex_vector_raw_fallback') + fast_result = (table.new_vector_search_builder() + .with_vector_column('embedding') + .with_query_vector([1.0, 0.0, 0.0, 0.0]) + .with_limit(1) + .execute_local()) + fast_ids = sorted(list(fast_result.results())) + print( + "paimon-vindex fast-mode vector search matched rows: " + f"ids={fast_ids}") + self.assertNotIn(3, fast_ids) + + full_table = table.copy({'global-index.search-mode': 'full'}) + full_result = (full_table.new_vector_search_builder() + .with_vector_column('embedding') + .with_query_vector([1.0, 0.0, 0.0, 0.0]) + .with_limit(1) + .execute_local()) + row_ids = sorted(list(full_result.results())) + print( + "paimon-vindex full-mode vector search matched rows: " + f"ids={row_ids}") + self.assertEqual([3], row_ids) + + read_builder = full_table.new_read_builder() + scan = read_builder.new_scan().with_global_index_result(full_result) + table_read = read_builder.new_read() + pa_table = table_read.to_arrow(scan.plan().splits()) + self.assertEqual(pa_table.num_rows, 1) + self.assertEqual([3], pa_table.column('id').to_pylist()) + def test_read_lumina_vector_with_btree_filter(self): """Vector search + btree scalar pre-filter, using a table that Java populated with both a Lumina vector index on `embedding` and a BTree diff --git a/paimon-vector/src/test/java/org/apache/paimon/JavaPyE2ETest.java b/paimon-vector/src/test/java/org/apache/paimon/JavaPyE2ETest.java index 033245f9eba4..60ffac81d127 100644 --- a/paimon-vector/src/test/java/org/apache/paimon/JavaPyE2ETest.java +++ b/paimon-vector/src/test/java/org/apache/paimon/JavaPyE2ETest.java @@ -197,4 +197,123 @@ public void testVindexVectorIndexWrite() throws Exception { assertThat(indexEntries.get(0).indexFile().indexType()) .isEqualTo(IvfFlatVectorGlobalIndexerFactory.IDENTIFIER); } + + @Test + @EnabledIfSystemProperty(named = "run.e2e.tests", matches = "true") + public void testVindexVectorRawFallbackWrite() throws Exception { + String tableName = "test_vindex_vector_raw_fallback"; + Path tablePath = new Path(warehouse.toString() + "/default.db/" + tableName); + LocalFileIO fileIO = LocalFileIO.create(); + if (fileIO.exists(tablePath)) { + fileIO.delete(tablePath, true); + } + + int dimension = 4; + + RowType rowType = + RowType.of( + new DataType[] {DataTypes.INT(), new ArrayType(new FloatType())}, + new String[] {"id", "embedding"}); + + Options options = new Options(); + options.set(PATH, tablePath.toString()); + options.set(ROW_TRACKING_ENABLED, true); + options.set(DATA_EVOLUTION_ENABLED, true); + options.set(GLOBAL_INDEX_ENABLED, true); + options.setString( + IvfFlatVectorGlobalIndexerFactory.IDENTIFIER + ".dimension", + String.valueOf(dimension)); + options.setString(IvfFlatVectorGlobalIndexerFactory.IDENTIFIER + ".metric", "l2"); + options.setString(IvfFlatVectorGlobalIndexerFactory.IDENTIFIER + ".nlist", "2"); + + TableSchema tableSchema = + SchemaUtils.forceCommit( + new SchemaManager(fileIO, tablePath), + new Schema( + rowType.getFields(), + Collections.emptyList(), + Collections.emptyList(), + options.toMap(), + "")); + + AppendOnlyFileStoreTable table = + new AppendOnlyFileStoreTable( + FileIOFinder.find(tablePath), + tablePath, + tableSchema, + CatalogEnvironment.empty()); + + float[][] indexedVectors = + new float[][] { + new float[] {0.0f, 1.0f, 0.0f, 0.0f}, + new float[] {0.0f, 0.0f, 1.0f, 0.0f}, + new float[] {0.0f, 0.0f, 0.0f, 1.0f} + }; + + BatchWriteBuilder writeBuilder = table.newBatchWriteBuilder(); + try (BatchTableWrite write = writeBuilder.newWrite(); + BatchTableCommit commit = writeBuilder.newCommit()) { + for (int i = 0; i < indexedVectors.length; i++) { + write.write(GenericRow.of(i, new GenericArray(indexedVectors[i]))); + } + commit.commit(write.prepareCommit()); + } + + DataField embeddingField = table.rowType().getField("embedding"); + Options indexOptions = table.coreOptions().toConfiguration(); + + GlobalIndexSingleColumnWriter writer = + (GlobalIndexSingleColumnWriter) + GlobalIndexBuilderUtils.createIndexWriter( + table, + IvfFlatVectorGlobalIndexerFactory.IDENTIFIER, + embeddingField, + indexOptions); + + for (int i = 0; i < indexedVectors.length; i++) { + writer.write(indexedVectors[i], i); + } + + List entries = writer.finish(); + assertThat(entries).hasSize(1); + assertThat(entries.get(0).rowCount()).isEqualTo(indexedVectors.length); + + Range rowRange = new Range(0, indexedVectors.length - 1); + List indexFiles = + GlobalIndexBuilderUtils.toIndexFileMetas( + table.fileIO(), + table.store().pathFactory().globalIndexFileFactory(), + table.coreOptions(), + rowRange, + embeddingField.id(), + IvfFlatVectorGlobalIndexerFactory.IDENTIFIER, + entries); + + DataIncrement dataIncrement = DataIncrement.indexIncrement(indexFiles); + CommitMessage message = + new CommitMessageImpl( + BinaryRow.EMPTY_ROW, + 0, + null, + dataIncrement, + CompactIncrement.emptyIncrement()); + try (BatchTableCommit commit = writeBuilder.newCommit()) { + commit.commit(Collections.singletonList(message)); + } + + try (BatchTableWrite write = writeBuilder.newWrite(); + BatchTableCommit commit = writeBuilder.newCommit()) { + write.write( + GenericRow.of( + 3, new GenericArray(new float[] {1.0f, 0.0f, 0.0f, 0.0f}))); + commit.commit(write.prepareCommit()); + } + + List indexEntries = + table.indexManifestFileReader().read(table.latestSnapshot().get().indexManifest()); + assertThat(indexEntries).hasSize(1); + assertThat(indexEntries.get(0).indexFile().rowCount()).isEqualTo(indexedVectors.length); + assertThat(indexEntries.get(0).indexFile().indexType()) + .isEqualTo(IvfFlatVectorGlobalIndexerFactory.IDENTIFIER); + } } From 1e49ae276f45c28dc0678a8362b2e3bfa7494552 Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Sun, 21 Jun 2026 22:16:14 +0800 Subject: [PATCH 4/4] [python] Fix mixed vector test formatting --- .../src/test/java/org/apache/paimon/JavaPyE2ETest.java | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/paimon-vector/src/test/java/org/apache/paimon/JavaPyE2ETest.java b/paimon-vector/src/test/java/org/apache/paimon/JavaPyE2ETest.java index 60ffac81d127..4c8cbfd59cc1 100644 --- a/paimon-vector/src/test/java/org/apache/paimon/JavaPyE2ETest.java +++ b/paimon-vector/src/test/java/org/apache/paimon/JavaPyE2ETest.java @@ -303,9 +303,7 @@ public void testVindexVectorRawFallbackWrite() throws Exception { try (BatchTableWrite write = writeBuilder.newWrite(); BatchTableCommit commit = writeBuilder.newCommit()) { - write.write( - GenericRow.of( - 3, new GenericArray(new float[] {1.0f, 0.0f, 0.0f, 0.0f}))); + write.write(GenericRow.of(3, new GenericArray(new float[] {1.0f, 0.0f, 0.0f, 0.0f}))); commit.commit(write.prepareCommit()); }