This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new 6d288e7fec [python] Execute local vector search through Rust core
(#10476)
6d288e7fec is described below
commit 6d288e7fec1b2be68ab122177df2706b33857e56
Author: Jingsong Lee <[email protected]>
AuthorDate: Sat Oct 10 11:04:31 2026 +0800
[python] Execute local vector search through Rust core (#10476)
---
paimon-python/README.md | 23 +
paimon-python/pypaimon/ray/batch_vector_search.py | 6 +-
paimon-python/pypaimon/ray/vector_search.py | 5 +
paimon-python/pypaimon/read/native_plan.py | 10 +-
paimon-python/pypaimon/schema/data_types.py | 1 -
.../table/source/batch_vector_search_builder.py | 7 +
.../pypaimon/table/source/native_vector_search.py | 93 +++
.../pypaimon/table/source/vector_search_builder.py | 7 +
.../pypaimon/table/source/vector_search_read.py | 49 +-
.../pypaimon/table/source/vector_search_scan.py | 12 +-
paimon-python/pypaimon/tests/data_types_test.py | 1 -
paimon-python/pypaimon/tests/native_plan_test.py | 14 +-
.../pypaimon/tests/native_vector_search_test.py | 629 +++++++++++++++++++++
.../pypaimon/tests/vector_search_filter_test.py | 72 ++-
14 files changed, 901 insertions(+), 28 deletions(-)
diff --git a/paimon-python/README.md b/paimon-python/README.md
index cc28e3b66d..15e9e6daf4 100644
--- a/paimon-python/README.md
+++ b/paimon-python/README.md
@@ -766,6 +766,29 @@ the fallback with `hdfs.client.fallback-to-pyarrow=false`
if you want
hard failures instead.
+# Native local vector search
+
+For REST tables, `read.native.enabled=true` also delegates local vector search
+execution to Rust. Single-vector searches support Data Evolution and configured
+PK vector indexes; batch searches use the existing Data Evolution result API.
+
+```python
+native_table = table.copy({'read.native.enabled': 'true'})
+result = (native_table.new_vector_search_builder()
+ .with_vector_column('embedding')
+ .with_query_vector([1.0, 0.0])
+ .with_limit(10)
+ .execute_local())
+```
+
+Rust plans one snapshot, applies scalar and partition filters before Top-K,
+searches indexes, scans uncovered rows and performs refinement. The returned
+Python result can be passed to `with_global_index_result` as usual, including
+scored physical-position splits for PK tables. An unavailable native backend
+falls back to the Python search reader. Distributed Ray searches keep their
+existing scan/read path.
+
+
# Vector index range reads
Native vector indexes (`ivf-flat`, `ivf-pq`, `ivf-sq`, `ivf-rq`, and `diskann`)
diff --git a/paimon-python/pypaimon/ray/batch_vector_search.py
b/paimon-python/pypaimon/ray/batch_vector_search.py
index 1a226020cb..164bd5793a 100644
--- a/paimon-python/pypaimon/ray/batch_vector_search.py
+++ b/paimon-python/pypaimon/ray/batch_vector_search.py
@@ -25,7 +25,7 @@ from pypaimon.ray.vector_search import (
_execution_options, _map_tasks, _require_ray, _shared_scalar_filters,
_worker_pre_filter,
)
from pypaimon.table.source.vector_search_read import (
- BatchVectorSearchReadImpl, _filtered_raw_row_ranges, _offer_score,
_scored_result,
+ BatchVectorSearchReadImpl, _filtered_raw_row_ranges, _offer_score,
_scored_result, _split_search_splits,
)
@@ -48,6 +48,10 @@ class _RayBatchVectorSearchRead(BatchVectorSearchReadImpl):
self._concurrency = concurrency
self._remote_args = remote_args
+ def _prepare_search_splits(self, splits, snapshot):
+ # Scalar evaluation and candidate verification belong to the workers.
+ return _split_search_splits(splits)
+
def _search_index_splits(self, splits, query, search_limit, snapshot,
batch=False):
# Send predicates and metadata, leaving each worker to build and
consume
# its own filter. Batch queries reuse the same filter for every vector.
diff --git a/paimon-python/pypaimon/ray/vector_search.py
b/paimon-python/pypaimon/ray/vector_search.py
index 0b592ec887..1a6bfad857 100644
--- a/paimon-python/pypaimon/ray/vector_search.py
+++ b/paimon-python/pypaimon/ray/vector_search.py
@@ -28,6 +28,7 @@ from pypaimon.table.source.vector_search_read import (
_filtered_raw_row_ranges,
_offer_score,
_scored_result,
+ _split_search_splits,
)
from pypaimon.table.source import global_index_live_row_filter
from pypaimon.utils.range import Range
@@ -78,6 +79,10 @@ class _RayVectorSearchRead(DataEvolutionVectorRead):
self._concurrency = concurrency
self._remote_args = remote_args
+ def _prepare_search_splits(self, splits, snapshot):
+ # Scalar evaluation and candidate verification belong to the workers.
+ return _split_search_splits(splits)
+
def _search_index_splits(self, splits, query, search_limit, snapshot,
batch=False):
# Send predicates and metadata, leaving each worker to build and
consume
# its own filter. Batch queries reuse the same filter for every vector.
diff --git a/paimon-python/pypaimon/read/native_plan.py
b/paimon-python/pypaimon/read/native_plan.py
index 58a6149b22..0a923cd701 100644
--- a/paimon-python/pypaimon/read/native_plan.py
+++ b/paimon-python/pypaimon/read/native_plan.py
@@ -330,8 +330,8 @@ def _rest_data_token(table):
return file_io._existing_valid_token()
-def _native_read_builder(table):
- """Return a fresh builder with the current schema and shared REST
FileIO."""
+def _native_table(table):
+ """Return the current schema and shared REST FileIO for Native
operations."""
rest_response = _resolved_rest_table_response(table)
file_io_options = _resolved_schema_file_io_options(table)
if rest_response is not None:
@@ -365,7 +365,11 @@ def _native_read_builder(table):
branch = getattr(rt, 'branch', None)
if not callable(branch) or branch() != table.current_branch():
raise RuntimeError("Native table did not resolve the requested
branch")
- return rt.new_read_builder()
+ return rt
+
+
+def _native_read_builder(table):
+ return _native_table(table).new_read_builder()
def _configure_native_read_builder(builder, predicate, limit, projection,
diff --git a/paimon-python/pypaimon/schema/data_types.py
b/paimon-python/pypaimon/schema/data_types.py
index 6d420f0759..893a5f2341 100755
--- a/paimon-python/pypaimon/schema/data_types.py
+++ b/paimon-python/pypaimon/schema/data_types.py
@@ -188,7 +188,6 @@ class VectorType(DataType):
"type": "VECTOR" + (" NOT NULL" if not self.nullable else ""),
"element": self.element.to_dict() if self.element else None,
"length": self.length,
- "nullable": self.nullable
}
@classmethod
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 b07cd1e835..f6fd739a5c 100644
--- a/paimon-python/pypaimon/table/source/batch_vector_search_builder.py
+++ b/paimon-python/pypaimon/table/source/batch_vector_search_builder.py
@@ -109,6 +109,13 @@ class
BatchVectorSearchBuilderImpl(AbstractVectorSearchBuilderImpl,
self._query_vectors = vectors
return self
+ def execute_batch_local(self):
+ from pypaimon.table.source.native_vector_search import
try_native_vector_search
+ results = try_native_vector_search(self, self._query_vectors,
batch=True)
+ if results is not None:
+ return results
+ return super().execute_batch_local()
+
def new_batch_vector_search_read(self):
# type: () -> BatchVectorSearchRead
reject_search_under_query_auth(self._table)
diff --git a/paimon-python/pypaimon/table/source/native_vector_search.py
b/paimon-python/pypaimon/table/source/native_vector_search.py
new file mode 100644
index 0000000000..ce61d41b53
--- /dev/null
+++ b/paimon-python/pypaimon/table/source/native_vector_search.py
@@ -0,0 +1,93 @@
+# 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.
+
+"""Local vector search through the Rust core Scan -> Plan -> Read path."""
+
+import logging
+
+from pypaimon.catalog.table_query_auth import reject_search_under_query_auth
+from pypaimon.globalindex.vector_search_result import
DictBasedScoredIndexResult
+from pypaimon.read.native_plan import (
+ _catalog_metastore, _native_table, _partition_fields, _predicate_to_native,
+ _raise_if_native_fork_safety_error, native_method_available)
+from pypaimon.read.split_serializer import deserialize_split_v1
+from pypaimon.table.source.primary_key_scored_result import (
+ PrimaryKeyScoredResult, PrimaryKeySearchPosition)
+
+logger = logging.getLogger(__name__)
+
+
+def try_native_vector_search(builder, queries, batch=False):
+ """Return existing Python result types, or None to use the Python reader.
+
+ Search, scalar filtering, refinement and physical-position selection all
+ belong to Rust core. Python only converts the scored result metadata.
+ """
+ table = builder._table
+ reject_search_under_query_auth(table)
+ try:
+ native_enabled = getattr(table.options, 'native_read_enabled', None)
+ if not callable(native_enabled) or not native_enabled():
+ return None
+ loader = getattr(getattr(table, 'catalog_environment', None),
'catalog_loader', None)
+ if loader is None or _catalog_metastore(loader) != 'rest':
+ return None
+ context_fn = getattr(loader, 'context', None)
+ if not callable(context_fn):
+ return None
+ context = context_fn()
+ if any(getattr(context, attr, None) is not None for attr in (
+ 'hadoop_conf', 'prefer_io_loader', 'fallback_io_loader')):
+ return None
+ # Python's batch reader currently supports only global row-ID results.
+ if batch and builder._is_primary_key_vector_search():
+ return None
+ method = 'new_batch_vector_search_builder' if batch else
'new_vector_search_builder'
+ if not native_method_available('Table', method):
+ return None
+ native = getattr(_native_table(table), method)()
+ if builder._vector_column is not None:
+ native.with_vector_column(builder._vector_column.name)
+ native.with_limit(builder._limit).with_options(builder._options)
+ if builder._filter is not None:
+ native.with_filter(_predicate_to_native(builder._filter))
+ if builder._partition_filter is not None:
+
native.with_partition_filter(_predicate_to_native(builder._partition_filter))
+ if batch:
+ native.with_query_vectors(queries)
+ results = native.execute_batch_local()
+ return [DictBasedScoredIndexResult(result.row_ids()) for result in
results]
+ native.with_query_vector(queries)
+ result = native.execute_local()
+ if builder._is_primary_key_vector_search():
+ return _NativePrimaryKeyScoredResult(table, result)
+ return DictBasedScoredIndexResult(result.row_ids())
+ except Exception as error:
+ _raise_if_native_fork_safety_error(error)
+ logger.warning('Native vector search failed, using the Python reader:
%s', error)
+ return None
+
+
+class _NativePrimaryKeyScoredResult(PrimaryKeyScoredResult):
+ """Metadata view of the physical selections already built by Rust core."""
+
+ def __init__(self, table, result):
+ self._snapshot_id = result.snapshot_id() or 0
+ self._positions = tuple(PrimaryKeySearchPosition(*position) for
position in result.positions())
+ self._splits = tuple(deserialize_split_v1(
+ split, _partition_fields(table), table.trimmed_primary_keys_fields)
+ for split in result.splits())
diff --git a/paimon-python/pypaimon/table/source/vector_search_builder.py
b/paimon-python/pypaimon/table/source/vector_search_builder.py
index 2c7e309a3f..76ce6575e9 100644
--- a/paimon-python/pypaimon/table/source/vector_search_builder.py
+++ b/paimon-python/pypaimon/table/source/vector_search_builder.py
@@ -250,6 +250,13 @@ class
VectorSearchBuilderImpl(AbstractVectorSearchBuilderImpl, VectorSearchBuild
self._query_vector = vector
return self
+ def execute_local(self):
+ from pypaimon.table.source.native_vector_search import
try_native_vector_search
+ result = try_native_vector_search(self, self._query_vector)
+ if result is not None:
+ return result
+ return super().execute_local()
+
def new_vector_search_read(self):
# type: () -> VectorSearchRead
reject_search_under_query_auth(self._table)
diff --git a/paimon-python/pypaimon/table/source/vector_search_read.py
b/paimon-python/pypaimon/table/source/vector_search_read.py
index f90d0a4d93..839993fbed 100644
--- a/paimon-python/pypaimon/table/source/vector_search_read.py
+++ b/paimon-python/pypaimon/table/source/vector_search_read.py
@@ -25,6 +25,7 @@ from contextlib import closing
from threading import Lock
from pypaimon.common.options.core_options import CoreOptions
+from pypaimon.common.options.options import Options
from pypaimon.globalindex.batch_vector_search import BatchVectorSearch
from pypaimon.globalindex.global_index_meta import GlobalIndexIOMeta
from pypaimon.globalindex.global_index_result import GlobalIndexResult
@@ -92,6 +93,7 @@ class AbstractVectorSearchReadImpl:
self._partition_filter = partition_filter
self._options = dict(options or {})
self._index_metric = None
+ self._scalar_pre_filter = None
def _search_metric(self, index_type=None):
if self._index_metric is not None:
@@ -124,6 +126,32 @@ class AbstractVectorSearchReadImpl:
% (self._index_metric, metric, self._vector_column.name))
self._index_metric = metric
+ def _prepare_search_splits(self, splits, snapshot):
+ """Route unevaluable scalar filters through raw scoring, as in Java."""
+ index_splits, raw_splits = _split_search_splits(splits)
+ self._scalar_pre_filter = None
+ if self._filter is None or not index_splits:
+ return index_splits, raw_splits
+ result = self._scalar_index_result(index_splits, snapshot)
+ if result is not None:
+ self._scalar_pre_filter = self._refine_scalar_result(index_splits,
result, snapshot)
+ return index_splits, raw_splits
+
+ from pypaimon.table.source.vector_search_split import
RawVectorSearchSplit
+ index_type = _vector_index_type(index_splits)
+ # Preserve the persisted metric even when these ANN segments are
+ # replaced by exact raw reads. Closing the wrapper closes its reader.
+ for split in index_splits:
+ if split.vector_index_files:
+ _, offset_reader = self._open_offset_reader(
+ split.vector_index_files, split.row_range_start,
split.row_range_end)
+ offset_reader.close()
+ raw_splits.append(RawVectorSearchSplit(
+ [Range(split.row_range_start, split.row_range_end) for split in
index_splits],
+ [index_file for split in index_splits for index_file in
split.scalar_index_files],
+ index_type))
+ return [], raw_splits
+
def _pre_filters(self, splits, snapshot=None):
# type: (list) -> List[RoaringBitmap64]
"""Evaluate live-row/scalar filters and return one bitmap per index
split."""
@@ -132,7 +160,9 @@ class AbstractVectorSearchReadImpl:
live_rows = global_index_live_row_filter.live_rows(
self._table, self._partition_filter, snapshot)
- matched_rows = self._scalar_matched_rows(splits, snapshot)
+ matched_rows = self._scalar_pre_filter
+ if matched_rows is None:
+ matched_rows = self._scalar_matched_rows(splits, snapshot)
return self._combine_pre_filters(splits, live_rows, matched_rows)
@staticmethod
@@ -165,7 +195,7 @@ class AbstractVectorSearchReadImpl:
return self._refine_scalar_result(splits,
self._scalar_index_result(splits, snapshot), snapshot)
def _scalar_index_result(self, splits, snapshot=None):
- """Evaluate scalar indexes without reading candidate rows."""
+ """Evaluate scalar indexes without data reads; None means no usable
result."""
# Collect scalar index files across splits, deduplicated by file name.
seen = set()
scalar_files = []
@@ -177,7 +207,7 @@ class AbstractVectorSearchReadImpl:
scalar_files.append(index_file)
if not scalar_files:
- return GlobalIndexResult.create_empty()
+ return None
from pypaimon.globalindex.data_evolution_global_index_scanner import
DataEvolutionGlobalIndexScanner
scanner = DataEvolutionGlobalIndexScanner.create(
@@ -187,7 +217,7 @@ class AbstractVectorSearchReadImpl:
snapshot=snapshot,
)
if scanner is None:
- return GlobalIndexResult.create_empty()
+ return None
try:
return scanner.scan(self._filter)
finally:
@@ -261,10 +291,15 @@ class AbstractVectorSearchReadImpl:
if evaluation is None:
return None
include_ranges = evaluation.result.results().to_range_list()
+ # Use the same effective mode as the scan so this optimization
+ # cannot remove raw ranges admitted by a query option override.
+ effective_options = dict(_table_options_map(self._table))
+ effective_options.update(self._options)
+ effective_core = CoreOptions(Options(effective_options))
include_ranges.extend(
scanner.unindexed_ranges(
self._filter,
- search_mode=self._table.options.scalar_index_search_mode(),
+ search_mode=effective_core.scalar_index_search_mode(),
contributing_field_ids=(
evaluation.contributing_field_ids),
))
@@ -671,7 +706,7 @@ class DataEvolutionVectorRead(AbstractVectorSearchReadImpl,
VectorSearchRead):
def _read(self, splits, snapshot):
self._index_metric = None
- index_splits, raw_splits = _split_search_splits(splits)
+ index_splits, raw_splits = self._prepare_search_splits(splits,
snapshot)
if not index_splits and not raw_splits:
return GlobalIndexResult.create_empty()
@@ -719,7 +754,7 @@ class
BatchVectorSearchReadImpl(AbstractVectorSearchReadImpl,
def _read_batch(self, splits, snapshot):
self._index_metric = None
n = len(self._query_vectors)
- index_splits, raw_splits = _split_search_splits(splits)
+ index_splits, raw_splits = self._prepare_search_splits(splits,
snapshot)
if not index_splits and not raw_splits:
return [GlobalIndexResult.create_empty() for _ in range(n)]
diff --git a/paimon-python/pypaimon/table/source/vector_search_scan.py
b/paimon-python/pypaimon/table/source/vector_search_scan.py
index 8484588f35..4203fe6235 100644
--- a/paimon-python/pypaimon/table/source/vector_search_scan.py
+++ b/paimon-python/pypaimon/table/source/vector_search_scan.py
@@ -20,7 +20,8 @@
from abc import ABC, abstractmethod
from collections import defaultdict
-from pypaimon.common.options.core_options import GlobalIndexSearchMode
+from pypaimon.common.options.core_options import CoreOptions,
GlobalIndexSearchMode
+from pypaimon.common.options.options import Options
from pypaimon.globalindex.data_evolution_global_index_coverage import
DataEvolutionGlobalIndexCoverage
from pypaimon.globalindex.data_evolution_global_index_scanner import (
is_supported_scalar_index,
@@ -169,7 +170,12 @@ class DataEvolutionVectorScan(VectorSearchScan):
)
)
- vector_search_mode = self._table.options.vector_index_search_mode()
+ # Java's DataEvolutionVectorScan overlays query options before
resolving
+ # the family search modes, without changing the table's configuration.
+ effective_options = dict(self._table.options.options.to_map())
+ effective_options.update(self._options)
+ effective_core = CoreOptions(Options(effective_options))
+ vector_search_mode = effective_core.vector_index_search_mode()
raw_row_ranges = DataEvolutionGlobalIndexCoverage(
self._table,
snapshot,
@@ -194,7 +200,7 @@ class DataEvolutionVectorScan(VectorSearchScan):
).unindexed_ranges(
self._table.fields,
self._filter,
- search_mode=self._table.options.scalar_index_search_mode(),
+ search_mode=effective_core.scalar_index_search_mode(),
)
if vector_search_mode == GlobalIndexSearchMode.FAST:
scalar_unindexed_ranges = Range.and_(
diff --git a/paimon-python/pypaimon/tests/data_types_test.py
b/paimon-python/pypaimon/tests/data_types_test.py
index e106f1fdf3..ae41cfe918 100755
--- a/paimon-python/pypaimon/tests/data_types_test.py
+++ b/paimon-python/pypaimon/tests/data_types_test.py
@@ -242,7 +242,6 @@ class DataTypesTest(unittest.TestCase):
"type": "VECTOR",
"element": "FLOAT",
"length": 3,
- "nullable": True
}
)
self.assertEqual(vector_type,
VectorType.from_dict(vector_type.to_dict()))
diff --git a/paimon-python/pypaimon/tests/native_plan_test.py
b/paimon-python/pypaimon/tests/native_plan_test.py
index 19bc268b70..7a56edac8d 100644
--- a/paimon-python/pypaimon/tests/native_plan_test.py
+++ b/paimon-python/pypaimon/tests/native_plan_test.py
@@ -1367,7 +1367,7 @@ class NativePlanTest(unittest.TestCase):
fs._apply_push_down_limit.assert_not_called()
fs.scan.assert_not_called()
- def
test_primary_key_shard_defers_limit_until_python_index_refinement(self):
+ def test_primary_key_shard_uses_native_index_refinement(self):
fs = Mock(partition_key_predicate=None)
scan = _scan(True, fs)
scan.table.is_primary_key_table = True
@@ -1376,14 +1376,14 @@ class NativePlanTest(unittest.TestCase):
fs.idx_of_this_subtask, fs.number_of_para_subtasks = 1, 2
scan.predicate = Mock()
splits = [Mock(bucket=0), Mock(bucket=0)]
- fs._apply_primary_key_sorted_indexes.return_value = splits[1:]
- fs._apply_push_down_limit.side_effect = lambda selected: selected
- with patch('pypaimon.read.native_plan.native_plan',
return_value=Plan(splits, 3)) as native:
+ # Rust receives the predicate and owns refinement before LIMIT.
+ with patch('pypaimon.read.native_plan.native_plan',
return_value=Plan(splits[1:], 3)) as native:
self.assertEqual(scan.plan().splits(), splits[1:])
self.assertEqual(native.call_args[1]['shard'], (1, 2))
- self.assertIsNone(native.call_args[1]['limit'])
- fs._apply_primary_key_sorted_indexes.assert_called_once()
- fs._apply_push_down_limit.assert_called_once_with(splits[1:])
+ self.assertEqual(native.call_args[1]['limit'], 1)
+ self.assertIs(native.call_args[1]['predicate'], scan.predicate)
+ fs._apply_primary_key_sorted_indexes.assert_not_called()
+ fs._apply_push_down_limit.assert_not_called()
fs.scan.assert_not_called()
def test_native_plan_requires_split_api(self):
diff --git a/paimon-python/pypaimon/tests/native_vector_search_test.py
b/paimon-python/pypaimon/tests/native_vector_search_test.py
new file mode 100644
index 0000000000..cecd4ca927
--- /dev/null
+++ b/paimon-python/pypaimon/tests/native_vector_search_test.py
@@ -0,0 +1,629 @@
+# 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.
+
+"""Existing local vector searches delegated to Rust with a REST table
context."""
+
+import shutil
+from pathlib import Path
+from unittest.mock import patch
+
+import pyarrow as pa
+import pytest
+
+from pypaimon import Schema
+from pypaimon.read.native_plan import _native_table, _predicate_to_native,
native_method_available
+from pypaimon.table.source.primary_key_scored_result import
PrimaryKeyScoredResult
+from pypaimon.tests import native_plan_rest_test,
primary_key_global_index_golden_test
+
+rest_catalog = native_plan_rest_test.rest_catalog
+golden_catalog = primary_key_global_index_golden_test.catalog
+
+pytestmark = [pytest.mark.native_plan, pytest.mark.skipif(
+ not native_method_available('Table', 'new_vector_search_builder'),
reason='Rust vector API required')]
+
+
+def _write(table, rows):
+ builder = table.new_batch_write_builder()
+ writer, commit = builder.new_write(), builder.new_commit()
+ try:
+ writer.write_arrow(pa.Table.from_pylist(rows, schema=pa.schema([
+ ('id', pa.int32()), ('embedding', pa.list_(pa.float32())), ('pt',
pa.int32())])))
+ commit.commit(writer.prepare_commit())
+ finally:
+ writer.close()
+ commit.close()
+
+
[email protected]
+def vector_table(rest_catalog):
+ catalog, _ = rest_catalog
+ catalog.create_table('default.vectors', Schema.from_pyarrow_schema(
+ pa.schema([('id', pa.int32()), ('embedding', pa.list_(pa.float32())),
('pt', pa.int32())]),
+ partition_keys=['pt'], options={
+ 'file.format': 'parquet', 'bucket': '-1', 'row-tracking.enabled':
'true',
+ 'data-evolution.enabled': 'true', 'global-index.enabled': 'true',
+ 'vector-index.search-mode': 'full', 'read.batch-size': '1',
+ }), False)
+ table = catalog.get_table('default.vectors')
+ _write(table, [
+ {'id': 1, 'embedding': [1, 0], 'pt': 0},
+ {'id': 2, 'embedding': [0, 1], 'pt': 0},
+ {'id': 3, 'embedding': None, 'pt': 0},
+ {'id': 4, 'embedding': [1, 0], 'pt': 1},
+ {'id': 5, 'embedding': [0, 1], 'pt': 1},
+ {'id': 6, 'embedding': [1, 1], 'pt': 1},
+ ])
+ return table
+
+
+def _single(table, query=(1, 0), limit=3):
+ return
table.new_vector_search_builder().with_vector_column('embedding').with_query_vector(query).with_limit(limit)
+
+
+def _scores(result):
+ getter = result.score_getter()
+ return {row_id: getter(row_id) for row_id in result.results()}
+
+
+def _assert_scores(actual, expected):
+ assert set(actual) == set(expected)
+ assert actual == pytest.approx(expected, abs=1e-6)
+
+
+def _rows(table, result):
+ builder = table.new_read_builder()
+ splits =
builder.new_scan().with_global_index_result(result).plan().splits()
+ return builder.new_read().to_arrow(splits)
+
+
[email protected]('metric', ['l2', 'cosine', 'inner_product'])
[email protected]('filtered', [False, True])
+def test_single_and_batch_search_match_python(vector_table, metric, filtered):
+ classic = vector_table.copy({'read.native.enabled': 'false',
'scan.native-plan.enabled': 'false'})
+ native = vector_table.copy({'read.native.enabled': 'true',
'scan.native-plan.enabled': 'true'})
+ predicate = classic.new_read_builder().new_predicate_builder()
+ queries = [[1, 0], [0, 1], [1, 0]]
+
+ def configure(builder):
+ builder.with_options({'ivf-flat.metric': metric})
+ if filtered:
+ builder.with_filter(predicate.greater_than('id',
1)).with_filter(predicate.less_than('id', 6))
+ # Partition is field 2 in the table, and field 0 in the partition
row.
+ builder.with_partition_filter(predicate.equal('pt', 1))
+ return builder
+
+ expected = [_scores(configure(_single(classic, query)).execute_local())
for query in queries]
+ from pypaimon.table.source.vector_search_read import
DataEvolutionVectorRead, BatchVectorSearchReadImpl
+ with patch.object(DataEvolutionVectorRead, 'read_plan',
side_effect=AssertionError('Python search ran')), \
+ patch.object(BatchVectorSearchReadImpl, 'read_batch_plan',
+ side_effect=AssertionError('Python batch search
ran')):
+ actual = [configure(_single(native, query)).execute_local() for query
in queries]
+ batch =
configure(native.new_batch_vector_search_builder().with_vector_column('embedding')
+
.with_query_vectors(queries).with_limit(3)).execute_batch_local()
+ assert len(batch) == len(queries)
+ for single, batched, scores in zip(actual, batch, expected):
+ _assert_scores(_scores(single), scores)
+ _assert_scores(_scores(batched), scores)
+ rows = _rows(native, actual[0]).to_pylist()
+ if filtered:
+ assert rows and all(row['pt'] == 1 for row in rows)
+ else:
+ assert len(rows) == 3
+
+
[email protected]('empty', ['table', 'filter', 'partition'])
+def test_empty_search_keeps_batch_arity(vector_table, rest_catalog, empty):
+ native = vector_table.copy({'read.native.enabled': 'true'})
+ if empty == 'table':
+ catalog, _ = rest_catalog
+ schema = vector_table.table_schema
+ catalog.create_table('default.empty_vectors', Schema(
+ schema.fields, schema.partition_keys, schema.primary_keys,
schema.options), False)
+ native =
catalog.get_table('default.empty_vectors').copy({'read.native.enabled': 'true'})
+ predicate = native.new_read_builder().new_predicate_builder()
+ builder =
(native.new_batch_vector_search_builder().with_vector_column('embedding')
+ .with_query_vectors([[1, 0], [0, 1]]).with_limit(2))
+ if empty == 'filter':
+ builder.with_filter(predicate.equal('id', 999))
+ elif empty == 'partition':
+ builder.with_partition_filter(predicate.equal('pt', 999))
+ from pypaimon.table.source.vector_search_read import
BatchVectorSearchReadImpl
+ with patch.object(BatchVectorSearchReadImpl, 'read_batch_plan',
side_effect=AssertionError('Python search ran')):
+ results = builder.execute_batch_local()
+ assert [_scores(result) for result in results] == [{}, {}]
+
+
+def test_scan_plan_is_reusable_and_snapshot_pinned(vector_table):
+ table = _native_table(vector_table.copy({'read.native.enabled': 'true'}))
+ builder = table.new_vector_search_builder().with_vector_column('embedding')
+ scan = builder.new_vector_search_scan()
+ plan = scan.scan()
+ snapshot = plan.snapshot_id()
+ _write(vector_table, [{'id': 7, 'embedding': [1, 0], 'pt': 0}])
+ builder.with_limit(10).with_query_vector([1, 0])
+ reader = builder.new_vector_search_read()
+ result = reader.read_plan(plan)
+ assert result.snapshot_id() == snapshot
+ assert len(result) == 5
+ assert reader.read_plan(plan).row_ids() == result.row_ids()
+ assert len(builder.execute_local()) == 6
+ filtered =
table.new_vector_search_builder().with_vector_column('embedding').with_query_vector([1,
0]).with_limit(2)
+ predicate =
vector_table.new_read_builder().new_predicate_builder().equal('id', 1)
+ filtered.with_filter(_predicate_to_native(predicate))
+ with pytest.raises(Exception, match='plan.*reader|context|different'):
+ filtered.new_vector_search_read().read_plan(plan)
+
+
[email protected]
+def pk_vector_table(rest_catalog, golden_catalog):
+ catalog, server = rest_catalog
+ golden = golden_catalog.get_table('default.test_pk_vector_golden')
+ schema = golden.table_schema
+ catalog.create_table('default.pk_vectors', Schema(
+ schema.fields, schema.partition_keys, schema.primary_keys,
schema.options), False)
+ source = Path(golden.table_path)
+ destination = Path(server.data_path) / server.warehouse / 'default' /
'pk_vectors'
+ for directory in ('bucket-0', 'index', 'manifest', 'snapshot'):
+ shutil.copytree(str(source / directory), str(destination / directory))
+ return catalog.get_table('default.pk_vectors')
+
+
[email protected]('filtered', [False, True])
[email protected]('refine', [False, True])
+def test_java_pk_index_results_are_physical_selections(pk_vector_table,
filtered, refine):
+ classic = pk_vector_table.copy({'read.native.enabled': 'false',
'scan.native-plan.enabled': 'false'})
+ native = pk_vector_table.copy({'read.native.enabled': 'true',
'scan.native-plan.enabled': 'true'})
+ predicate = classic.new_read_builder().new_predicate_builder().equal('id',
3)
+
+ def configure(table):
+ builder = _single(table, [1, 0, 0, 0], limit=2)
+ if filtered:
+ builder.with_filter(predicate)
+ if refine:
+ builder.with_option('ivf.refine_factor', '2')
+ return builder
+
+ expected = configure(classic).execute_local()
+ from pypaimon.table.source.primary_key_vector_read import
PrimaryKeyVectorRead
+ with patch.object(PrimaryKeyVectorRead, 'read_plan',
side_effect=AssertionError('Python PK search ran')):
+ result = configure(native).execute_local()
+ assert isinstance(result, PrimaryKeyScoredResult)
+ assert result.snapshot_id == expected.snapshot_id
+ assert [(p.bucket, p.data_file_name, p.row_position) for p in
result.positions] == [
+ (p.bucket, p.data_file_name, p.row_position) for p in
expected.positions]
+ assert [p.score for p in result.positions] == pytest.approx([p.score for p
in expected.positions])
+ rows = _rows(native, result)
+ assert sorted(rows.column('id').to_pylist()) == sorted(_rows(classic,
expected).column('id').to_pylist())
+ assert len(result.splits) == len(expected.splits)
+ for split in result.splits:
+ assert len(split.files) == 1
+ assert not split.raw_convertible
+ assert split.scores is not None
+
+
+def test_native_search_failure_falls_back(vector_table):
+ native = vector_table.copy({'read.native.enabled': 'true'})
+ expected = _scores(_single(native.copy({'read.native.enabled':
'false'})).execute_local())
+ from pypaimon.table.source import native_vector_search
+ with patch.object(native_vector_search, '_native_table',
side_effect=RuntimeError('unavailable')):
+ _assert_scores(_scores(_single(native).execute_local()), expected)
+
+
+def test_explicit_disable_uses_python(vector_table):
+ table = vector_table.copy({'read.native.enabled': 'false'})
+ from pypaimon.table.source import native_vector_search
+ with patch.object(native_vector_search, '_native_table',
side_effect=AssertionError('Native search ran')):
+ assert len(_single(table).execute_local().results()) == 3
+
+
[email protected]('mode', ['fast', 'full', 'detail'])
[email protected]('filtered', [False, True])
+def test_indexed_search_and_raw_tail_match_python(vector_table, mode,
filtered):
+ from pypaimon.globalindex.create_global_index import create_global_index
+ from pypaimon.table.source.vector_search_read import
DataEvolutionVectorRead, BatchVectorSearchReadImpl
+ create_global_index(vector_table, 'embedding', 'ivf-flat', options={
+ 'ivf-flat.dimension': '2', 'ivf-flat.nlist': '1',
'ivf-flat.distance.metric': 'cosine'})
+ _write(vector_table, [{'id': 7, 'embedding': [1, 0], 'pt': 1}])
+ classic = vector_table.copy({'read.native.enabled': 'false',
'scan.native-plan.enabled': 'false',
+ 'vector-index.search-mode': mode})
+ native = vector_table.copy({'read.native.enabled': 'true',
'scan.native-plan.enabled': 'true',
+ 'vector-index.search-mode': mode})
+ predicate = classic.new_read_builder().new_predicate_builder()
+
+ def configure(builder):
+ builder.with_option('refine_factor', '2').with_option('ivf.nprobe',
'1')
+ if filtered:
+ builder.with_filter(predicate.greater_or_equal('id', 4))
+ builder.with_partition_filter(predicate.equal('pt', 1))
+ return builder
+
+ queries = [[1, 0], [0, 1]]
+ expected = [_scores(configure(_single(classic, query, 3)).execute_local())
for query in queries]
+ with patch.object(DataEvolutionVectorRead, 'read_plan',
side_effect=AssertionError('Python search ran')), \
+ patch.object(BatchVectorSearchReadImpl, 'read_batch_plan',
side_effect=AssertionError('Python search ran')):
+ actual = [configure(_single(native, query, 3)).execute_local() for
query in queries]
+ batch =
configure(native.new_batch_vector_search_builder().with_vector_column('embedding')
+
.with_query_vectors(queries).with_limit(3)).execute_batch_local()
+ for single, batched, scores in zip(actual, batch, expected):
+ _assert_scores(_scores(single), scores)
+ _assert_scores(_scores(batched), scores)
+
+
+def test_fork_safety_error_does_not_fall_back(vector_table):
+ from pypaimon_rust import ForkSafetyError
+ from pypaimon.table.source import native_vector_search
+ table = vector_table.copy({'read.native.enabled': 'true'})
+ with patch.object(native_vector_search, '_native_table',
side_effect=ForkSafetyError('unsafe fork')):
+ with pytest.raises(ForkSafetyError, match='unsafe fork'):
+ _single(table).execute_local()
+
+
[email protected]('filtered', [False, True])
[email protected]('mode', ['fast', 'full'])
+def test_mixed_index_metrics_are_rejected(vector_table, filtered, mode):
+ from pypaimon.globalindex.create_global_index import create_global_index
+ for metric, identifier in [('l2', 7), ('cosine', 8)]:
+ _write(vector_table, [{'id': identifier, 'embedding':
[float(identifier), 0], 'pt': 1}])
+ create_global_index(vector_table, 'embedding', 'ivf-flat', options={
+ 'ivf-flat.dimension': '2', 'ivf-flat.nlist': '1',
'ivf-flat.distance.metric': metric})
+ predicate =
vector_table.new_read_builder().new_predicate_builder().greater_than('id', 0)
+ for native in (False, True):
+ table = vector_table.copy({'read.native.enabled': str(native).lower(),
'vector-index.search-mode': mode})
+ if native:
+ direct =
(_native_table(table).new_vector_search_builder().with_vector_column('embedding')
+ .with_query_vector([1, 0]).with_limit(3))
+ if filtered:
+ direct.with_filter(_predicate_to_native(predicate))
+ with pytest.raises(Exception, match='Cannot merge vector indexes
with different metrics'):
+ direct.execute_local()
+ builders = [_single(table),
table.new_batch_vector_search_builder().with_vector_column('embedding')
+ .with_query_vectors([[1, 0], [0, 1]]).with_limit(3)]
+ for builder in builders:
+ if filtered:
+ builder.with_filter(predicate)
+ with pytest.raises(Exception, match='Cannot merge vector indexes
with different metrics'):
+ builder.execute_batch_local() if hasattr(builder,
'execute_batch_local') else builder.execute_local()
+
+
[email protected]('option', ['hadoop_conf', 'prefer_io_loader',
'fallback_io_loader'])
+def test_custom_io_context_keeps_python_search(vector_table, option):
+ from copy import copy
+ from pypaimon.table.source import native_vector_search
+ table = vector_table.copy({'read.native.enabled': 'true'})
+ loader = table.catalog_environment.catalog_loader
+ context = copy(loader.context())
+ setattr(context, option, object())
+ expected = _scores(_single(table.copy({'read.native.enabled':
'false'})).execute_local())
+ with patch.object(loader, 'context', return_value=context), \
+ patch.object(native_vector_search, '_native_table',
+ side_effect=AssertionError('Unsupported Native
context')):
+ _assert_scores(_scores(_single(table).execute_local()), expected)
+
+
[email protected]('mode', ['fast', 'full', 'detail'])
+def test_unevaluable_scalar_filter_scores_indexed_ranges_raw(vector_table,
mode):
+ from pypaimon.globalindex.create_global_index import create_global_index
+ from pypaimon.table.source.vector_search_read import
DataEvolutionVectorRead, BatchVectorSearchReadImpl
+ # The matching vectors live outside the nearest IVF cluster. Sending a
+ # data-derived allow-list to ANN with nprobe=1 silently loses all matches.
+ _write(vector_table, [
+ {'id': i, 'embedding': [i / 1000, 0], 'pt': 2} for i in range(50)
+ ] + [
+ {'id': i, 'embedding': [float(i), 0], 'pt': 2} for i in range(100, 150)
+ ])
+ create_global_index(vector_table, 'embedding', 'ivf-flat', options={
+ 'ivf-flat.dimension': '2', 'ivf-flat.nlist': '4',
'ivf-flat.distance.metric': 'l2'})
+ classic = vector_table.copy({'read.native.enabled': 'false',
'scan.native-plan.enabled': 'false',
+ 'vector-index.search-mode': mode})
+ native = vector_table.copy({'read.native.enabled': 'true',
'scan.native-plan.enabled': 'true',
+ 'vector-index.search-mode': mode})
+ predicates = classic.new_read_builder().new_predicate_builder()
+ predicate = predicates.greater_or_equal('id', 100)
+ partition = predicates.equal('pt', 2)
+
+ def configure(builder, direct=False):
+ return (builder.with_option('ivf.nprobe', '1')
+ .with_filter(_predicate_to_native(predicate) if direct else
predicate)
+ .with_partition_filter(_predicate_to_native(partition) if
direct else partition))
+
+ queries = [[0, 0], [1, 0]]
+ expected = [_scores(configure(_single(classic, query)).execute_local())
for query in queries]
+ assert all(len(scores) == 3 for scores in expected)
+ direct_table = _native_table(native)
+ direct =
configure(direct_table.new_batch_vector_search_builder().with_vector_column('embedding')
+ .with_query_vectors(queries).with_limit(3),
direct=True).execute_batch_local()
+ for result, scores in zip(direct, expected):
+ _assert_scores(result.row_ids(), scores)
+ with patch.object(DataEvolutionVectorRead, 'read_plan',
side_effect=AssertionError('Python search ran')), \
+ patch.object(BatchVectorSearchReadImpl, 'read_batch_plan',
side_effect=AssertionError('Python search ran')):
+ singles = [configure(_single(native, query)).execute_local() for query
in queries]
+ batch =
configure(native.new_batch_vector_search_builder().with_vector_column('embedding')
+
.with_query_vectors(queries).with_limit(3)).execute_batch_local()
+ for single, batched, scores in zip(singles, batch, expected):
+ _assert_scores(_scores(single), scores)
+ _assert_scores(_scores(batched), scores)
+ assert sorted(_rows(native, singles[0]).column('id').to_pylist()) == [100,
101, 102]
+
+
[email protected]('mode', ['fast', 'full', 'detail'])
[email protected]('refine', [False, True])
[email protected]('unindexed_tail', [False, True])
+def test_scalar_candidates_obey_refinement_switch(rest_catalog, mode, refine,
unindexed_tail):
+ from pypaimon.globalindex.create_global_index import create_global_index
+ from pypaimon.table.source.vector_search_read import
DataEvolutionVectorRead, BatchVectorSearchReadImpl
+ catalog, _ = rest_catalog
+ schema = pa.schema([('id', pa.int32()), ('name', pa.string()),
('embedding', pa.list_(pa.float32()))])
+ catalog.create_table('default.scalar_vectors',
Schema.from_pyarrow_schema(schema, options={
+ 'bucket': '-1', 'row-tracking.enabled': 'true',
'data-evolution.enabled': 'true',
+ 'global-index.enabled': 'true', 'file.format': 'parquet'}), False)
+ table = catalog.get_table('default.scalar_vectors')
+
+ def write(rows):
+ builder = table.new_batch_write_builder()
+ writer, commit = builder.new_write(), builder.new_commit()
+ try:
+ writer.write_arrow(pa.Table.from_pylist(rows, schema=schema))
+ commit.commit(writer.prepare_commit())
+ finally:
+ writer.close()
+ commit.close()
+
+ write([{'id': i, 'name': name, 'embedding': [float(i), 0]}
+ for i, name in enumerate(['alpha', 'beta zeta', 'gamma'])])
+ create_global_index(table, 'embedding', 'ivf-flat', options={
+ 'ivf-flat.dimension': '2', 'ivf-flat.nlist': '1',
'ivf-flat.distance.metric': 'l2'})
+ create_global_index(table, 'name', 'btree')
+ if unindexed_tail:
+ # Refinement=false excludes inexact indexed candidates, while an
+ # unindexed tail in FULL/DETAIL must still execute the row predicate.
+ write([{'id': 3, 'name': 'tail zeta', 'embedding': [3, 0]}])
+ options = {'vector-index.search-mode': mode,
'global-index.filter.refine-from-data': str(refine).lower()}
+ classic = table.copy(dict(options, **{'read.native.enabled': 'false',
'scan.native-plan.enabled': 'false'}))
+ native = table.copy(dict(options, **{'read.native.enabled': 'true',
'scan.native-plan.enabled': 'true'}))
+ predicate =
classic.new_read_builder().new_predicate_builder().like('name', '%zeta%')
+
+ def configure(builder, direct=False):
+ return builder.with_filter(_predicate_to_native(predicate) if direct
else predicate)
+
+ queries = [[0, 0], [1, 0]]
+ expected = [_scores(configure(_single(classic, query, 1)).execute_local())
for query in queries]
+ count = 1 if refine or (unindexed_tail and mode != 'fast') else 0
+ assert [len(scores) for scores in expected] == [count, count]
+ direct_table = _native_table(native)
+ for query, scores in zip(queries, expected):
+ result =
configure(direct_table.new_vector_search_builder().with_vector_column('embedding')
+ .with_query_vector(vector=query).with_limit(1),
direct=True).execute_local()
+ _assert_scores(result.row_ids(), scores)
+ with patch.object(DataEvolutionVectorRead, 'read_plan',
side_effect=AssertionError('Python search ran')), \
+ patch.object(BatchVectorSearchReadImpl, 'read_batch_plan',
side_effect=AssertionError('Python search ran')):
+ singles = [configure(_single(native, query, 1)).execute_local() for
query in queries]
+ batch =
configure(native.new_batch_vector_search_builder().with_vector_column('embedding')
+
.with_query_vectors(queries).with_limit(1)).execute_batch_local()
+ for single, batched, scores in zip(singles, batch, expected):
+ _assert_scores(_scores(single), scores)
+ _assert_scores(_scores(batched), scores)
+
+
[email protected]('mode', ['fast', 'full', 'detail'])
+def
test_partition_filter_does_not_make_exact_scalar_index_inexact(vector_table,
mode):
+ from pypaimon.globalindex.create_global_index import create_global_index
+ from pypaimon.table.source.vector_search_read import
DataEvolutionVectorRead, BatchVectorSearchReadImpl
+ create_global_index(vector_table, 'embedding', 'ivf-flat', options={
+ 'ivf-flat.dimension': '2', 'ivf-flat.nlist': '1',
'ivf-flat.distance.metric': 'l2'})
+ create_global_index(vector_table, 'id', 'btree')
+ options = {'vector-index.search-mode': mode,
'global-index.filter.refine-from-data': 'false'}
+ classic = vector_table.copy(dict(options, **{'read.native.enabled':
'false', 'scan.native-plan.enabled': 'false'}))
+ native = vector_table.copy(dict(options, **{'read.native.enabled': 'true',
'scan.native-plan.enabled': 'true'}))
+ predicates = classic.new_read_builder().new_predicate_builder()
+ predicate, partition = predicates.equal('id', 4), predicates.equal('pt', 1)
+
+ def configure(builder, direct=False):
+ return (builder.with_filter(_predicate_to_native(predicate) if direct
else predicate)
+ .with_partition_filter(_predicate_to_native(partition) if
direct else partition))
+
+ expected = _scores(configure(_single(classic)).execute_local())
+ assert len(expected) == 1
+ direct =
configure(_native_table(native).new_vector_search_builder().with_vector_column('embedding')
+ .with_query_vector(vector=[1, 0]).with_limit(3),
direct=True).execute_local()
+ _assert_scores(direct.row_ids(), expected)
+ with patch.object(DataEvolutionVectorRead, 'read_plan',
side_effect=AssertionError('Python search ran')), \
+ patch.object(BatchVectorSearchReadImpl, 'read_batch_plan',
side_effect=AssertionError('Python search ran')):
+ single = configure(_single(native)).execute_local()
+ batch =
configure(native.new_batch_vector_search_builder().with_vector_column('embedding')
+ .with_query_vectors([[1, 0], [1,
0]]).with_limit(3)).execute_batch_local()
+ _assert_scores(_scores(single), expected)
+ for result in batch:
+ _assert_scores(_scores(result), expected)
+
+
[email protected]('mode', ['fast', 'full', 'detail'])
[email protected]('refine', [False, True])
+def test_unsupported_conjunct_on_same_indexed_field_is_not_exact(rest_catalog,
mode, refine):
+ from pypaimon.globalindex.create_global_index import create_global_index
+ from pypaimon.table.source.vector_search_read import
DataEvolutionVectorRead, BatchVectorSearchReadImpl
+ catalog, server = rest_catalog
+ schema = pa.schema([('id', pa.int32()), ('name', pa.string()),
('embedding', pa.list_(pa.float32()))])
+ catalog.create_table('default.fm_vectors',
Schema.from_pyarrow_schema(schema, options={
+ 'bucket': '-1', 'row-tracking.enabled': 'true',
'data-evolution.enabled': 'true',
+ 'global-index.enabled': 'true', 'file.format': 'parquet'}), False)
+ table = catalog.get_table('default.fm_vectors')
+ builder = table.new_batch_write_builder()
+ writer, commit = builder.new_write(), builder.new_commit()
+ try:
+ writer.write_arrow(pa.Table.from_pylist([
+ {'id': 1, 'name': None, 'embedding': [1, 0]},
+ {'id': 2, 'name': 'alpha', 'embedding': [0, 1]},
+ ], schema=schema))
+ commit.commit(writer.prepare_commit())
+ finally:
+ writer.close()
+ commit.close()
+ create_global_index(table, 'embedding', 'ivf-flat', options={
+ 'ivf-flat.dimension': '2', 'ivf-flat.nlist': '1',
'ivf-flat.distance.metric': 'l2'})
+ # Build FM through Rust's existing procedure; PyPaimon can read this index
+ # but its Python index-building API does not create FM definitions.
+ from pypaimon_rust.datafusion import SQLContext
+ context = SQLContext()
+ context.register_catalog('paimon', {
+ 'metastore': 'rest', 'uri': server.get_url(), 'warehouse': 'warehouse',
+ 'token.provider': 'bear', 'token': 'test-token', 'data-token.enabled':
'false'})
+ context.sql("CALL sys.create_global_index(table => 'default.fm_vectors', "
+ "index_column => 'name', index_type => 'fm')")
+ native = table.copy({'read.native.enabled': 'true',
'scan.native-plan.enabled': 'true',
+ 'vector-index.search-mode': mode,
+ 'global-index.filter.refine-from-data':
str(refine).lower()})
+ predicates = native.new_read_builder().new_predicate_builder()
+ is_null = predicates.is_null('name')
+ # FM can evaluate isNull but declines startsWith on the SAME field.
+ # Its contributing field ID must not mark this conjunction exact.
+ impossible = predicates.and_predicates([is_null,
predicates.startswith('name', 'a')])
+ direct_table = _native_table(native)
+ for predicate, count in [(is_null, 1), (impossible, 0)]:
+ direct =
(direct_table.new_vector_search_builder().with_vector_column('embedding')
+ .with_query_vector(vector=[1, 0]).with_limit(2)
+
.with_filter(_predicate_to_native(predicate)).execute_local())
+ assert len(direct) == count
+ with patch.object(DataEvolutionVectorRead, 'read_plan',
side_effect=AssertionError('Python search ran')), \
+ patch.object(BatchVectorSearchReadImpl, 'read_batch_plan',
+ side_effect=AssertionError('Python search ran')):
+ single = _single(native,
limit=2).with_filter(predicate).execute_local()
+ batch =
(native.new_batch_vector_search_builder().with_vector_column('embedding')
+ .with_query_vectors([[1, 0], [0,
1]]).with_limit(2).with_filter(predicate).execute_batch_local())
+ assert len(single.results()) == count
+ assert [len(result.results()) for result in batch] == [count, count]
+ assert len(_rows(native, _single(native,
limit=2).with_filter(is_null).execute_local())) == 1
+
+
[email protected]('scalar_mode', ['fast', 'full', 'detail'])
[email protected]('vector_mode', ['fast', 'full', 'detail'])
[email protected]('vector_covers_tail', [False, True])
[email protected]('builder_options', [False, True])
+def test_scalar_unindexed_ranges_follow_both_search_modes(
+ vector_table, scalar_mode, vector_mode, vector_covers_tail,
builder_options):
+ from pypaimon.globalindex.create_global_index import create_global_index
+ from pypaimon.table.source.vector_search_read import
DataEvolutionVectorRead, BatchVectorSearchReadImpl
+ create_global_index(vector_table, 'id', 'btree')
+ index_options = {'ivf-flat.dimension': '2', 'ivf-flat.nlist': '1',
'ivf-flat.distance.metric': 'l2'}
+ if not vector_covers_tail:
+ create_global_index(vector_table, 'embedding', 'ivf-flat',
options=index_options)
+ _write(vector_table, [{'id': 7, 'embedding': [3, 3], 'pt': 1}])
+ if vector_covers_tail:
+ create_global_index(vector_table, 'embedding', 'ivf-flat',
options=index_options)
+ modes = {'scalar-index.search-mode': scalar_mode,
'vector-index.search-mode': vector_mode}
+ table_modes = {'scalar-index.search-mode': 'fast',
'vector-index.search-mode': 'fast'} if builder_options else modes
+ # Exact scalar matches do not require refinement; uncovered scalar ranges
+ # still execute the residual even when candidate refinement is disabled.
+ options = dict(table_modes, **{'global-index.filter.refine-from-data':
'false'})
+ classic = vector_table.copy(dict(options, **{'read.native.enabled':
'false', 'scan.native-plan.enabled': 'false'}))
+ native = vector_table.copy(dict(options, **{'read.native.enabled': 'true',
'scan.native-plan.enabled': 'true'}))
+ predicate = classic.new_read_builder().new_predicate_builder().equal('id',
7)
+
+ def configure(builder, direct=False):
+ if builder_options:
+ builder.with_options(modes)
+ return builder.with_filter(_predicate_to_native(predicate) if direct
else predicate)
+
+ queries = [[0, 0], [1, 0]]
+ expected = [_scores(configure(_single(classic, query,
10)).execute_local()) for query in queries]
+ # Vector FAST constrains scalar fallback to vector coverage. Vector
+ # FULL/DETAIL retain their own uncovered ranges independently of scalar
FAST.
+ count = int(scalar_mode != 'fast' if vector_covers_tail else vector_mode
!= 'fast')
+ assert [len(scores) for scores in expected] == [count, count]
+ direct_table = _native_table(native)
+ for query, scores in zip(queries, expected):
+ result =
configure(direct_table.new_vector_search_builder().with_vector_column('embedding')
+ .with_query_vector(vector=query).with_limit(10),
direct=True).execute_local()
+ _assert_scores(result.row_ids(), scores)
+ direct_batch =
configure(direct_table.new_batch_vector_search_builder().with_vector_column('embedding')
+ .with_query_vectors(queries).with_limit(10),
direct=True).execute_batch_local()
+ for result, scores in zip(direct_batch, expected):
+ _assert_scores(result.row_ids(), scores)
+ with patch.object(DataEvolutionVectorRead, 'read_plan',
side_effect=AssertionError('Python search ran')), \
+ patch.object(BatchVectorSearchReadImpl, 'read_batch_plan',
side_effect=AssertionError('Python search ran')):
+ singles = [configure(_single(native, query, 10)).execute_local() for
query in queries]
+ batch =
configure(native.new_batch_vector_search_builder().with_vector_column('embedding')
+
.with_query_vectors(queries).with_limit(10)).execute_batch_local()
+ for single, batched, scores in zip(singles, batch, expected):
+ _assert_scores(_scores(single), scores)
+ _assert_scores(_scores(batched), scores)
+ if count:
+ assert _rows(native, singles[0]).column('id').to_pylist() == [7]
+
+
[email protected]('scalar_mode', ['full', 'detail'])
[email protected]('builder_options', [False, True])
+def
test_raw_scalar_prefilter_keeps_uncovered_rows_with_overlapping_scalar_file(
+ rest_catalog, scalar_mode, builder_options):
+ from pypaimon.globalindex.create_global_index import create_global_index
+ from pypaimon.table.source.vector_search_read import
DataEvolutionVectorRead, BatchVectorSearchReadImpl
+ catalog, _ = rest_catalog
+ schema = pa.schema([('id', pa.int32()), ('name', pa.string()),
('embedding', pa.list_(pa.float32()))])
+ catalog.create_table('default.partial_scalar_vectors',
Schema.from_pyarrow_schema(schema, options={
+ 'bucket': '-1', 'row-tracking.enabled': 'true',
'data-evolution.enabled': 'true',
+ 'global-index.enabled': 'true', 'file.format': 'parquet'}), False)
+ table = catalog.get_table('default.partial_scalar_vectors')
+
+ def write(rows):
+ builder = table.new_batch_write_builder()
+ writer, commit = builder.new_write(), builder.new_commit()
+ try:
+ writer.write_arrow(pa.Table.from_pylist(rows, schema=schema))
+ commit.commit(writer.prepare_commit())
+ finally:
+ writer.close()
+ commit.close()
+
+ write([{'id': i, 'name': 'old', 'embedding': [float(i), 0]} for i in
range(6)])
+ create_global_index(table, 'name', 'btree')
+ write([{'id': 6, 'name': 'target', 'embedding': [6, 0]}])
+ create_global_index(table, 'id', 'btree')
+ write([{'id': 7, 'name': 'other', 'embedding': [7, 0]},
+ {'id': 8, 'name': 'target', 'embedding': [8, 0]}])
+ create_global_index(table, 'embedding', 'ivf-flat', options={
+ 'ivf-flat.dimension': '2', 'ivf-flat.nlist': '1',
'ivf-flat.distance.metric': 'l2'})
+ options = {'vector-index.search-mode': 'fast',
+ 'scalar-index.search-mode': 'fast' if builder_options else
scalar_mode}
+ classic = table.copy(dict(options, **{'read.native.enabled': 'false',
'scan.native-plan.enabled': 'false'}))
+ native = table.copy(dict(options, **{'read.native.enabled': 'true',
'scan.native-plan.enabled': 'true'}))
+ predicates = classic.new_read_builder().new_predicate_builder()
+ predicate = predicates.and_predicates([
+ predicates.greater_or_equal('id', 6), predicates.equal('name',
'target')])
+
+ def configure(builder, direct=False):
+ if builder_options:
+ builder.with_option('scalar-index.search-mode', scalar_mode)
+ return builder.with_filter(_predicate_to_native(predicate) if direct
else predicate)
+
+ # The raw gap [6,8] still overlaps id's index [0,6]. Both the indexed
+ # candidate 6 and the unindexed match 8 must survive its raw pre-filter.
+ queries = [[0, 0], [1, 0]]
+ expected = [_scores(configure(_single(classic, query,
10)).execute_local()) for query in queries]
+ assert [len(scores) for scores in expected] == [2, 2]
+ direct_table = _native_table(native)
+ direct =
configure(direct_table.new_batch_vector_search_builder().with_vector_column('embedding')
+ .with_query_vectors(queries).with_limit(10),
direct=True).execute_batch_local()
+ for result, scores in zip(direct, expected):
+ _assert_scores(result.row_ids(), scores)
+ with patch.object(DataEvolutionVectorRead, 'read_plan',
side_effect=AssertionError('Python search ran')), \
+ patch.object(BatchVectorSearchReadImpl, 'read_batch_plan',
side_effect=AssertionError('Python search ran')):
+ singles = [configure(_single(native, query, 10)).execute_local() for
query in queries]
+ batch =
configure(native.new_batch_vector_search_builder().with_vector_column('embedding')
+
.with_query_vectors(queries).with_limit(10)).execute_batch_local()
+ for single, batched, scores in zip(singles, batch, expected):
+ _assert_scores(_scores(single), scores)
+ _assert_scores(_scores(batched), scores)
+ assert sorted(_rows(native, singles[0]).column('id').to_pylist()) == [6, 8]
diff --git a/paimon-python/pypaimon/tests/vector_search_filter_test.py
b/paimon-python/pypaimon/tests/vector_search_filter_test.py
index f707189fae..655080f8bc 100644
--- a/paimon-python/pypaimon/tests/vector_search_filter_test.py
+++ b/paimon-python/pypaimon/tests/vector_search_filter_test.py
@@ -1588,6 +1588,38 @@ class VectorSearchFilterTest(unittest.TestCase):
self.assertFalse(any(isinstance(split, RawVectorSearchSplit)
for split in splits))
+ def
test_query_search_modes_override_table_modes_without_mutating_options(self):
+ from pypaimon.table.source.vector_search_split import
RawVectorSearchSplit
+
+ for mode in ('full', 'detail'):
+ for table_mode, scalar_mode, vector_mode, vector_end, expected in [
+ ('fast', mode, 'fast', 9, [Range(5, 9)]),
+ ('fast', 'fast', mode, 4, [Range(5, 9)]),
+ (mode, 'fast', 'fast', 9, []),
+ ]:
+ with self.subTest(mode=mode, scalar=scalar_mode,
vector=vector_mode, table=table_mode):
+ entries = [
+ _entry(None, field_id=1,
index_type='lumina-vector-ann', file_name='vec.index',
+ row_range_start=0, row_range_end=vector_end),
+ _entry(None, field_id=0, index_type='btree',
file_name='id.index',
+ row_range_start=0, row_range_end=4),
+ ]
+ table = _StubTable(fields=[self.id_field,
self.embedding_field], entries=entries)
+ original_options = {'vector-index.search-mode':
table_mode, 'scalar-index.search-mode': table_mode}
+ table.options =
CoreOptions(Options(dict(original_options)))
+ table.data_ranges_for_data_evolution_global_index_coverage
= (
+ lambda snapshot, partition: [Range(0, 9)])
+ _patch_snapshot(self, entries,
types.SimpleNamespace(next_row_id=10))
+ predicate = Predicate(method='equal', index=0, field='id',
literals=[7])
+ splits =
(VectorSearchBuilderImpl(table).with_vector_column('embedding')
+ .with_filter(predicate).with_options({
+ 'scalar-index.search-mode': scalar_mode,
'vector-index.search-mode': vector_mode})
+ .new_vector_search_scan().scan().splits())
+ raw_ranges = [row_range for split in splits if
isinstance(split, RawVectorSearchSplit)
+ for row_range in split.row_ranges]
+ self.assertEqual(expected, raw_ranges)
+ self.assertEqual(original_options,
table.options.options.to_map())
+
def test_fast_vector_mode_limits_scalar_fallback_to_vector_coverage(self):
from pypaimon.table.source.vector_search_split import
RawVectorSearchSplit
@@ -1708,6 +1740,32 @@ class VectorSearchFilterTest(unittest.TestCase):
contributing_field_ids=frozenset([0]),
)
+ def test_raw_pre_filter_query_mode_overrides_table_mode(self):
+ from pypaimon.table.source.vector_search_read import
DataEvolutionVectorRead
+ from pypaimon.table.source.vector_search_split import
RawVectorSearchSplit
+
+ table = _StubTable(fields=[self.id_field, self.embedding_field],
entries=[])
+ table.options = CoreOptions(Options({'scalar-index.search-mode':
'fast'}))
+ predicate = Predicate(method='equal', index=0, field='id',
literals=[5])
+ scanner = mock.MagicMock()
+ scanner.scan_with_coverage.return_value = GlobalIndexEvaluation(
+ GlobalIndexResult.create_empty(), frozenset([0]))
+ scanner.unindexed_ranges.return_value = [Range(5, 9)]
+ for mode in ('full', 'detail'):
+ with self.subTest(mode=mode):
+ scanner.reset_mock()
+ reader = DataEvolutionVectorRead(table, limit=3,
vector_column=self.embedding_field,
+ query_vector=[1, 0, 0, 0],
filter_=predicate,
+
options={'scalar-index.search-mode': mode})
+ with
mock.patch('pypaimon.globalindex.data_evolution_global_index_scanner.'
+ 'DataEvolutionGlobalIndexScanner.create',
return_value=scanner):
+ ranges = reader._raw_pre_filter([
+ RawVectorSearchSplit([Range(5, 9)],
[self.entries[2].index_file])])
+ self.assertEqual([Range(5, 9)], ranges)
+ scanner.unindexed_ranges.assert_called_once_with(
+ predicate, search_mode=GlobalIndexSearchMode(mode),
contributing_field_ids=frozenset([0]))
+ self.assertEqual(GlobalIndexSearchMode.FAST,
table.options.scalar_index_search_mode())
+
def test_raw_vector_pre_filter_keeps_full_fallback_as_ranges(self):
from pypaimon.table.source.vector_search_read import
DataEvolutionVectorRead
from pypaimon.table.source.vector_search_split import
RawVectorSearchSplit
@@ -2950,7 +3008,7 @@ class VectorSearchManySplitsTest(unittest.TestCase):
self.assertEqual([1, 8], sorted(list(results[0].results())))
self.assertEqual([2, 8], sorted(list(results[1].results())))
- def test_read_uses_empty_index_prefilter_when_scalar_index_missing(self):
+ def test_read_routes_index_ranges_to_raw_when_scalar_index_missing(self):
from pypaimon.table.source.vector_search_read import
DataEvolutionVectorRead
from pypaimon.table.source.vector_search_split import
IndexVectorSearchSplit
@@ -2973,10 +3031,14 @@ class VectorSearchManySplitsTest(unittest.TestCase):
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())
+ offset_reader = mock.Mock()
+ with mock.patch.object(reader, '_open_offset_reader',
return_value=(None, offset_reader)):
+ indexed, raw = reader._prepare_search_splits([split], None)
+ self.assertEqual([], indexed)
+ self.assertEqual(1, len(raw))
+ self.assertEqual([Range(0, 4)], raw[0].row_ranges)
+ self.assertEqual('lumina-vector-ann', raw[0].index_type)
+ offset_reader.close.assert_called_once()
def test_raw_search_uses_partition_filter_and_index_type_metric(self):
import pyarrow as pa
