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 9c9a78670c [Python] Delegate local vindex builds to native Rust
(#10488)
9c9a78670c is described below
commit 9c9a78670c03c188dd878e5569731b2f52ec14e0
Author: Jingsong Lee <[email protected]>
AuthorDate: Sat Oct 10 17:41:41 2026 +0800
[Python] Delegate local vindex builds to native Rust (#10488)
---
.../pypaimon/globalindex/native_index_build.py | 20 +-
.../tests/native_generic_index_build_test.py | 18 --
.../tests/native_vector_index_build_test.py | 254 +++++++++++++++++++++
3 files changed, 271 insertions(+), 21 deletions(-)
diff --git a/paimon-python/pypaimon/globalindex/native_index_build.py
b/paimon-python/pypaimon/globalindex/native_index_build.py
index 48ecf421e3..ad7e6f404a 100644
--- a/paimon-python/pypaimon/globalindex/native_index_build.py
+++ b/paimon-python/pypaimon/globalindex/native_index_build.py
@@ -29,9 +29,8 @@ def build_native_global_index(builder, partition_filter):
Build errors propagate: a failed native build must not start a second
build.
"""
table = builder._table
- # The Rust vector pipeline requires dense, non-null vectors. Keep Python
- # vector construction until it supports Java's sparse relative row IDs.
- if (builder._index_type not in ('btree', 'bitmap', 'full-text')
+ from pypaimon.globalindex.vindex.vindex_vector_global_index_reader import
VINDEX_IDENTIFIERS
+ if (builder._index_type not in ('btree', 'bitmap', 'full-text') +
tuple(VINDEX_IDENTIFIERS)
or not table.options.native_write_enabled()
or not table.options.data_evolution_enabled()
or not table.options.global_index_enabled()
@@ -53,6 +52,21 @@ def build_native_global_index(builder, partition_filter):
from
pypaimon.globalindex.full_text.native_full_text_global_index_reader import
NativeFullTextIndexOptions
options.update({'full-text.' + key: value for key, value in
NativeFullTextIndexOptions.from_options(builder._options.to_map()).to_native_options().items()})
+ elif builder._index_type in VINDEX_IDENTIFIERS:
+ from pypaimon.globalindex.vindex.vindex_vector_index_writer import
native_options, train_sample_ratio
+ column = builder._index_columns[0]
+ # Reuse the Java-compatible option mapping; Rust's low-level builder
also
+ # accepts native aliases which Python's public API deliberately
ignores.
+ mapped = native_options(table.field_dict[column].type,
table.options.options.to_map(),
+ builder._index_type, column,
builder._user_options)
+ options = {key: value for key, value in options.items()
+ if key in ('global-index.row-count-per-shard',
'global-index.build.parallelism',
+ 'vindex.build.granule.enabled')
+ and (key != 'vindex.build.granule.enabled' or
builder._index_type != 'diskann')}
+ options.update({'fields.' + column + '.' + key: value for key, value
in mapped.items()
+ if key != 'index.type'})
+ options['fields.' + column + '.train.sample-ratio'] =
str(train_sample_ratio(
+ table.options.options.to_map(), builder._index_type, column,
builder._user_options))
native.with_options(options)
if partition_filter is not None:
native.with_partition_filter(_predicate_to_native(partition_filter))
diff --git a/paimon-python/pypaimon/tests/native_generic_index_build_test.py
b/paimon-python/pypaimon/tests/native_generic_index_build_test.py
index b700ae53c5..425f92ee3c 100644
--- a/paimon-python/pypaimon/tests/native_generic_index_build_test.py
+++ b/paimon-python/pypaimon/tests/native_generic_index_build_test.py
@@ -155,24 +155,6 @@ def
test_python_full_text_build_delegates_to_common_native_entry(rest_catalog):
index_tests._commit(table, messages)
-def test_python_nullable_vector_build_keeps_sparse_row_ids(rest_catalog):
- catalog, _ = rest_catalog
- schema = pa.schema([('id', pa.int32()), ('name', pa.string()),
- ('embedding', pa.list_(pa.float32())), ('pt',
pa.int32())])
- table = index_tests._create(catalog, schema=schema)
- rows = [dict(ROWS[0]), dict(ROWS[1]), dict(ROWS[2])]
- rows[1]['embedding'] = None
- index_tests._append(table, rows, schema=schema)
- with patch.object(GlobalIndexBuilder, '_build_generic_index',
- wraps=GlobalIndexBuilder(table, 'embedding',
index_type='ivf-flat',
-
options=_options('ivf-flat'))._build_generic_index) as build:
- messages = GlobalIndexBuilder(table, 'embedding',
index_type='ivf-flat',
- options=_options('ivf-flat')).build()
- assert build.call_count == 1
- assert sum(file.row_count for file in
index_tests._message_files(messages)) == 3
- index_tests._commit(table, messages)
-
-
@pytest.mark.parametrize('option_source', ['build', 'table'])
@pytest.mark.parametrize('stop_words', [['paimon', 'apache'], 'paimon;apache'])
def test_native_full_text_options_match_python_list_encoding(rest_catalog,
option_source, stop_words):
diff --git a/paimon-python/pypaimon/tests/native_vector_index_build_test.py
b/paimon-python/pypaimon/tests/native_vector_index_build_test.py
new file mode 100644
index 0000000000..4e6531150d
--- /dev/null
+++ b/paimon-python/pypaimon/tests/native_vector_index_build_test.py
@@ -0,0 +1,254 @@
+# 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.
+
+"""REST vector builds delegated to Rust preserve Java writer semantics."""
+
+from unittest.mock import patch
+
+import pyarrow as pa
+import pytest
+
+from pypaimon.globalindex.create_global_index import GlobalIndexBuilder
+from pypaimon.globalindex.vindex.vindex_vector_global_index_reader import (
+ PaimonVindexInput, VINDEX_IDENTIFIERS,
+)
+from pypaimon.tests import native_generic_index_build_test as generic
+from pypaimon.tests import native_sorted_index_build_test as indexes
+from pypaimon.read.native_plan import _native_table
+
+rest_catalog = indexes.rest_catalog
+pytestmark = generic.pytestmark
+SCHEMA = pa.schema([('id', pa.int32()), ('embedding', pa.list_(pa.float32())),
('pt', pa.int32())])
+
+
+def _rows(vectors, start=0, partition=0):
+ return [{'id': start + i, 'embedding': vector, 'pt': partition} for i,
vector in enumerate(vectors)]
+
+
+def _table(catalog, vectors, options=None, partitioned=False):
+ settings = {'vector-index.search-mode': 'fast'}
+ settings.update(options or {})
+ table = indexes._create(catalog, schema=SCHEMA, options=settings,
partitioned=partitioned)
+ indexes._append(table, _rows(vectors), schema=SCHEMA)
+ return table
+
+
+def _build(table, kind, options=None, **kwargs):
+ with patch.object(GlobalIndexBuilder, '_build_generic_index',
+ side_effect=AssertionError('Python vector rows were
materialized')):
+ return GlobalIndexBuilder(table, 'embedding', index_type=kind,
options=options, **kwargs).build()
+
+
+def _metadata(table, file):
+ from paimon_vindex import VectorIndexReader
+ with table.file_io.new_input_stream(indexes._path(table, file)) as stream:
+ adapter = PaimonVindexInput(stream)
+ try:
+ with VectorIndexReader(adapter) as reader:
+ return reader.metadata()
+ finally:
+ adapter.close()
+
+
+def _bytes(table, file):
+ with table.file_io.new_input_stream(indexes._path(table, file)) as stream:
+ return stream.read()
+
+
+def _assert_ids(table, expected, kind):
+ # Both readers must observe original IDs rather than a compacted 0..N
sequence.
+ search_options = {'diskann.l_search': '64'} if kind == 'diskann' else
{'ivf.nprobe': '64'}
+ classic = table.copy({'read.native.enabled': 'false',
'scan.native-plan.enabled': 'false'})
+ result =
(classic.new_vector_search_builder().with_vector_column('embedding')
+ .with_query_vector([1.,
1.]).with_limit(64).with_options(search_options).execute_local())
+ assert sorted(result.results()) == expected
+ result =
(_native_table(table).new_vector_search_builder().with_vector_column('embedding')
+ .with_query_vector([1.,
1.]).with_limit(64).with_options(search_options).execute_local())
+ assert sorted(result.row_ids()) == expected
+
+
[email protected]('kind', VINDEX_IDENTIFIERS)
[email protected]('granule', [False, True])
+def
test_nullable_vectors_preserve_row_ids_and_logical_cardinality(rest_catalog,
kind, granule):
+ catalog, _ = rest_catalog
+ vectors = [None, [1., 0.], None, [0., 1.], None, None, [1., 1.], None]
+ table = _table(catalog, vectors)
+ options = generic._options(kind)
+ options[kind + '.train.sample-ratio'] = '0.5'
+ if kind != 'diskann':
+ options['vindex.build.granule.enabled'] = str(granule).lower()
+ before = table.snapshot_manager().get_latest_snapshot().id
+ messages = _build(table, kind, options)
+ files = indexes._message_files(messages)
+ assert len(files) == 1
+ assert files[0].row_count == 8
+ assert files[0].global_index_meta.row_range_start == 0
+ assert files[0].global_index_meta.row_range_end == 7
+ assert _metadata(table, files[0]).total_vectors == 3
+ assert table.snapshot_manager().get_latest_snapshot().id == before
+ indexes._commit(table, messages)
+ _assert_ids(table, [1, 3, 6], kind)
+ assert _build(table, kind, options) == []
+
+
[email protected]('kind', VINDEX_IDENTIFIERS)
+def test_all_null_shards_create_no_file_or_snapshot(rest_catalog, kind):
+ catalog, _ = rest_catalog
+ table = _table(catalog, [None] * 9)
+ before = table.snapshot_manager().get_latest_snapshot().id
+ assert _build(table, kind, generic._options(kind, 4)) == []
+ assert table.create_global_index('embedding', index_type=kind,
options=generic._options(kind, 4)) == 0
+ assert table.snapshot_manager().get_latest_snapshot().id == before
+ assert indexes._files(table) == []
+ root =
table.path_factory().global_index_path_factory().global_index_root_path()
+ assert not table.file_io.exists(root) or table.file_io.list_status(root)
== []
+
+
[email protected]('kind', VINDEX_IDENTIFIERS)
[email protected]('native_commit', [False, True])
+def
test_empty_shard_does_not_shift_next_shard_or_incremental_ids(rest_catalog,
kind, native_commit):
+ catalog, _ = rest_catalog
+ table = _table(catalog, [None] * 4 + [None, [1., 0.], None, [0., 1.]],
+ {'commit.native.enabled': str(native_commit).lower()})
+ options = generic._options(kind, 4)
+ messages = _build(table, kind, options)
+ files = indexes._message_files(messages)
+ assert len(files) == 1
+ assert files[0].row_count == 4
+ assert files[0].global_index_meta.row_range_start == 4
+ assert files[0].global_index_meta.row_range_end == 7
+ indexes._commit(table, messages)
+ _assert_ids(table, [5, 7], kind)
+ indexes._append(table, _rows([None, [1., 1.], None], 8), schema=SCHEMA)
+ messages = _build(table, kind, options)
+ files = indexes._message_files(messages)
+ assert len(files) == 1
+ assert files[0].global_index_meta.row_range_start == 8
+ assert files[0].row_count == 3
+ indexes._commit(table, messages)
+ _assert_ids(table, [5, 7, 9], kind)
+
+
[email protected]('kind', VINDEX_IDENTIFIERS)
+def test_sparse_full_spill_bytes_match_python_writer(rest_catalog, kind):
+ catalog, _ = rest_catalog
+ vectors = [None if i % 5 == 0 else [float(i % 7), float(i // 7)] for i in
range(97)]
+ table = _table(catalog, vectors, {'read.batch-size': '7'})
+ options = generic._options(kind)
+ options[kind + '.train.sample-ratio'] = ' 0.37 '
+ if kind != 'diskann':
+ options['vindex.build.granule.enabled'] = 'false'
+ native = _build(table, kind, options)
+ classic = GlobalIndexBuilder(table.copy({'write.native.enabled':
'false'}), 'embedding',
+ index_type=kind, options=options).build()
+ native_files, classic_files = indexes._message_files(native),
indexes._message_files(classic)
+ assert len(native_files) == len(classic_files) == 1
+ assert native_files[0].row_count == classic_files[0].row_count ==
len(vectors)
+ assert _bytes(table, native_files[0]) == _bytes(table, classic_files[0])
+
+
[email protected]('kind', VINDEX_IDENTIFIERS)
[email protected]('budget', [None, 4096])
+def
test_default_training_config_matches_python_without_fixed_nlist_or_pq_m(rest_catalog,
kind, budget):
+ catalog, _ = rest_catalog
+ table = _table(catalog, [None, [1., 0.], None, [0., 1.], [1., 1.]])
+ options = {kind + '.dimension': '2'}
+ if budget is not None:
+ options[kind + '.max-bytes-per-vector'] = str(budget)
+ native = indexes._message_files(_build(table, kind, options))
+ classic_builder = GlobalIndexBuilder(table.copy({'write.native.enabled':
'false'}),
+ 'embedding', index_type=kind,
options=options)
+ classic = indexes._message_files(classic_builder.build())
+ assert _bytes(table, native[0]) == _bytes(table, classic[0])
+ assert _metadata(table, native[0]).total_vectors == 3
+
+
[email protected]('kind', ['ivf-flat', 'ivf-pq'])
+def test_field_aliases_and_build_options_override_table_options(rest_catalog,
kind):
+ catalog, _ = rest_catalog
+ table = _table(catalog, [None, [1., 0.], [0., 1.]],
+ {kind + '.dimension': '4',
'fields.embedding.distance.metric': 'inner_product'})
+ options = {kind + '.index.dimension': '2', 'fields.embedding.metric': 'l2',
+ kind + '.expected-vector-count': '32', 'ivf.coarse-assignment':
'exact',
+ 'global-index.build.parallelism': ' 2 ', 'metric': 'invalid',
'dimension': '999',
+ kind + '.ignored-option': 'ignored',
+ 'fields.embedding.ivf.train.max-points-per-centroid': '16'}
+ if kind == 'ivf-pq':
+ options.update({'fields.embedding.pq.code-ratio': '0.25', kind +
'.ivf.pq-encoding': 'canonical',
+ kind + '.use-opq': 'true'})
+ native = indexes._message_files(_build(table, kind, options))
+ classic_builder = GlobalIndexBuilder(table.copy({'write.native.enabled':
'false'}),
+ 'embedding', index_type=kind,
options=options)
+ classic = indexes._message_files(classic_builder.build())
+ assert _bytes(table, native[0]) == _bytes(table, classic[0])
+ meta = _metadata(table, native[0])
+ assert meta.dimension == 2
+ assert meta.metric == 'l2'
+
+
[email protected]('vector', [[1.], [1., None]])
+def test_invalid_non_null_vectors_fail_without_python_retry(rest_catalog,
vector):
+ catalog, _ = rest_catalog
+ table = _table(catalog, [[1., 0.], None, vector])
+ before = table.snapshot_manager().get_latest_snapshot().id
+ with pytest.raises(Exception, match='dimension|null'):
+ _build(table, 'ivf-flat', generic._options('ivf-flat', 1))
+ assert table.snapshot_manager().get_latest_snapshot().id == before
+ assert indexes._files(table) == []
+ root =
table.path_factory().global_index_path_factory().global_index_root_path()
+ assert not table.file_io.exists(root) or table.file_io.list_status(root)
== []
+
+
[email protected]('value', ['0', '-1', 'invalid'])
+def
test_invalid_build_parallelism_is_rejected_before_writing_files(rest_catalog,
value):
+ catalog, _ = rest_catalog
+ table = _table(catalog, [[1., 0.]])
+ options = generic._options('ivf-flat')
+ options['global-index.build.parallelism'] = value
+ before = table.snapshot_manager().get_latest_snapshot().id
+ with pytest.raises(Exception, match='parallelism'):
+ _build(table, 'ivf-flat', options)
+ assert table.snapshot_manager().get_latest_snapshot().id == before
+ assert indexes._files(table) == []
+
+
[email protected]('kind', VINDEX_IDENTIFIERS)
+def test_all_null_build_does_not_instantiate_native_training(rest_catalog,
kind):
+ catalog, _ = rest_catalog
+ table = _table(catalog, [None] * 4)
+ options = {kind + '.dimension': '2', kind + '.distance.metric': 'invalid'}
+ if kind != 'diskann':
+ options[kind + '.nlist'] = '0'
+ assert _build(table, kind, options) == []
+ classic = table.copy({'write.native.enabled': 'false'})
+ assert GlobalIndexBuilder(classic, 'embedding', index_type=kind,
options=options).build() == []
+ assert indexes._files(table) == []
+
+
+def test_automatic_pq_budget_uses_real_non_null_cardinality(rest_catalog):
+ catalog, _ = rest_catalog
+ vectors = [None if i % 3 == 0 else [float(i % 7), float(i // 7)] for i in
range(1024)]
+ table = _table(catalog, vectors)
+ options = {'ivf-pq.dimension': '2', 'ivf-pq.max-bytes-per-vector': '64',
+ 'global-index.row-count-per-shard': '2048'}
+ native = indexes._message_files(_build(table, 'ivf-pq', options))
+ classic_table = table.copy({'write.native.enabled': 'false'})
+ classic_builder = GlobalIndexBuilder(classic_table, 'embedding',
index_type='ivf-pq', options=options)
+ classic = indexes._message_files(classic_builder.build())
+ assert _bytes(table, native[0]) == _bytes(table, classic[0])
+ assert native[0].row_count == 1024
+ assert _metadata(table, native[0]).total_vectors == 682
