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 a34c8a1ff8 [python] Support vindex global index build (#8398)
a34c8a1ff8 is described below
commit a34c8a1ff81939e7d5db301695f2e9a51968729e
Author: Jingsong Lee <[email protected]>
AuthorDate: Tue Jun 30 23:09:57 2026 +0800
[python] Support vindex global index build (#8398)
Builds on #8393 to add PyPaimon support for building paimon-vindex
global vector indexes from Python, with behavior aligned to the Java
default global index builder. Also updates the global index docs with
SQL/Python SDK TabItem examples for BTree and vector index lifecycle
APIs.
---
docs/docs/multimodal-table/global-index.mdx | 5 +-
docs/docs/multimodal-table/global-index/btree.mdx | 75 +++++
docs/docs/multimodal-table/global-index/vector.mdx | 130 +++++++++
.../pypaimon/common/options/core_options.py | 10 +
.../pypaimon/globalindex/create_global_index.py | 275 +++++++++++++++++-
.../pypaimon/globalindex/vindex/__init__.py | 16 +-
.../vindex/vindex_vector_index_writer.py | 300 ++++++++++++++++++++
.../pypaimon/tests/global_index_build_test.py | 306 ++++++++++++++++++++-
8 files changed, 1102 insertions(+), 15 deletions(-)
diff --git a/docs/docs/multimodal-table/global-index.mdx
b/docs/docs/multimodal-table/global-index.mdx
index a163390e0b..6ee0882707 100644
--- a/docs/docs/multimodal-table/global-index.mdx
+++ b/docs/docs/multimodal-table/global-index.mdx
@@ -157,8 +157,9 @@ added_files = table.create_global_index(
</Tabs>
-PyPaimon global index build currently supports single-column BTree indexes on
-tables with row tracking enabled.
+PyPaimon global index build currently supports single-column BTree indexes and
+single-column paimon-vindex IVF vector indexes on tables with row tracking
+enabled.
Drop index files:
diff --git a/docs/docs/multimodal-table/global-index/btree.mdx
b/docs/docs/multimodal-table/global-index/btree.mdx
index ee50db0aed..afc953ebe5 100644
--- a/docs/docs/multimodal-table/global-index/btree.mdx
+++ b/docs/docs/multimodal-table/global-index/btree.mdx
@@ -48,6 +48,10 @@ For keyword-style text retrieval, use [Full-Text
Index](./full-text) instead.
## Build BTree Index
+<Tabs groupId="btree-build">
+
+<TabItem value="sql" label="SQL">
+
```sql
-- Create BTree index on 'name' column
CALL sys.create_global_index(
@@ -68,6 +72,77 @@ CALL sys.create_global_index(
);
```
+</TabItem>
+
+<TabItem value="python-sdk" label="Python SDK">
+
+```python
+table = catalog.get_table("db.my_table")
+
+# Create BTree index on 'name' column.
+added_files = table.create_global_index("name", index_type="btree")
+print(added_files)
+```
+
+The API returns the number of committed index files. You can pass BTree build
+options and restrict the build to selected partitions:
+
+```python
+added_files = table.create_global_index(
+ "name",
+ index_type="btree",
+ partitions=[{"dt": "2026-06-18"}, {"dt": "2026-06-19"}],
+ options={"sorted-index.records-per-range": "10000000"},
+)
+print(added_files)
+```
+
+</TabItem>
+
+</Tabs>
+
+## Drop BTree Index
+
+<Tabs groupId="btree-drop">
+
+<TabItem value="sql" label="SQL">
+
+```sql
+CALL sys.drop_global_index(
+ table => 'db.my_table',
+ index_column => 'name',
+ index_type => 'btree'
+);
+```
+
+</TabItem>
+
+<TabItem value="python-sdk" label="Python SDK">
+
+```python
+table = catalog.get_table("db.my_table")
+
+dropped_files = table.drop_global_index("name", index_type="btree")
+print(dropped_files)
+```
+
+You can also restrict the drop to selected partitions, or count matched files
+without committing:
+
+```python
+matched_files = table.drop_global_index(
+ "name",
+ index_type="btree",
+ partitions=[{"dt": "2026-06-18"}, {"dt": "2026-06-19"}],
+ dry_run=True,
+)
+print(matched_files)
+```
+
+</TabItem>
+
+</Tabs>
+
## BTree Options
| Option | Default | Description |
diff --git a/docs/docs/multimodal-table/global-index/vector.mdx
b/docs/docs/multimodal-table/global-index/vector.mdx
index 1fa560645e..9e25b8f983 100644
--- a/docs/docs/multimodal-table/global-index/vector.mdx
+++ b/docs/docs/multimodal-table/global-index/vector.mdx
@@ -52,6 +52,10 @@ Choose the index type based on the trade-off you want:
## Build Vector Index
+<Tabs groupId="vector-build">
+
+<TabItem value="sql" label="SQL">
+
```sql
-- Create IVF-PQ vector index on 'embedding' column
CALL sys.create_global_index(
@@ -70,6 +74,52 @@ CALL sys.create_global_index(
);
```
+</TabItem>
+
+<TabItem value="python-sdk" label="Python SDK">
+
+```python
+table = catalog.get_table("db.my_table")
+
+# Create IVF-PQ vector index on 'embedding' column.
+added_files = table.create_global_index(
+ "embedding",
+ index_type="ivf-pq",
+ options={
+ "ivf-pq.dimension": "768",
+ "ivf-pq.distance.metric": "cosine",
+ "ivf-pq.nlist": "256",
+ "ivf-pq.pq.m": "16",
+ },
+)
+print(added_files)
+```
+
+You can also build only selected partitions or tune the row-id shard size used
+by non-sorted global indexes:
+
+```python
+added_files = table.create_global_index(
+ "embedding",
+ index_type="ivf-flat",
+ partitions=[{"dt": "2026-06-18"}, {"dt": "2026-06-19"}],
+ options={
+ "ivf-flat.dimension": "768",
+ "global-index.row-count-per-shard": "100000",
+ },
+)
+print(added_files)
+```
+
+PyPaimon vector index build currently supports the paimon-vindex IVF
identifiers:
+`ivf-flat`, `ivf-pq`, `ivf-hnsw-flat`, and `ivf-hnsw-sq`. Use SQL to build
+Lumina indexes. Install `paimon-vindex` or `pypaimon[vindex]` before building
+or querying IVF vector indexes from Python.
+
+</TabItem>
+
+</Tabs>
+
Use `index_type => 'lumina'` for new Lumina indexes. The legacy
`lumina-vector-ann` identifier is
kept only so existing tables can still load old indexes.
@@ -78,6 +128,48 @@ IVF indexes or `lumina.index.dimension` for Lumina indexes.
For `VECTOR<FLOAT>`
the dimension from the column type. When `lumina.index.dimension` is
explicitly set for a
`VECTOR<FLOAT>` column, it must match the vector type length.
+## Drop Vector Index
+
+<Tabs groupId="vector-drop">
+
+<TabItem value="sql" label="SQL">
+
+```sql
+CALL sys.drop_global_index(
+ table => 'db.my_table',
+ index_column => 'embedding',
+ index_type => 'ivf-pq'
+);
+```
+
+</TabItem>
+
+<TabItem value="python-sdk" label="Python SDK">
+
+```python
+table = catalog.get_table("db.my_table")
+
+dropped_files = table.drop_global_index("embedding", index_type="ivf-pq")
+print(dropped_files)
+```
+
+You can also restrict the drop to selected partitions, or count matched files
+without committing:
+
+```python
+matched_files = table.drop_global_index(
+ "embedding",
+ index_type="ivf-pq",
+ partitions=[{"dt": "2026-06-18"}, {"dt": "2026-06-19"}],
+ dry_run=True,
+)
+print(matched_files)
+```
+
+</TabItem>
+
+</Tabs>
+
Supported IVF vector index options:
| Option | Default | Description |
@@ -123,6 +215,10 @@ as `<field-name>`. Field-level vector options do not
include the index-type pref
use `fields.image_embedding.nlist` to override the shared `ivf-pq.nlist`
option for
`image_embedding`:
+<Tabs groupId="vector-per-field-options">
+
+<TabItem value="sql" label="SQL">
+
```sql
CREATE TABLE my_table (
id INT,
@@ -142,6 +238,40 @@ CREATE TABLE my_table (
);
```
+</TabItem>
+
+<TabItem value="python-sdk" label="Python SDK">
+
+```python
+import pyarrow as pa
+
+from pypaimon import Schema
+
+schema = Schema.from_pyarrow_schema(
+ pa.schema([
+ pa.field("id", pa.int32()),
+ pa.field("title_embedding", pa.list_(pa.float32())),
+ pa.field("image_embedding", pa.list_(pa.float32())),
+ ]),
+ options={
+ "bucket": "-1",
+ "row-tracking.enabled": "true",
+ "data-evolution.enabled": "true",
+ "global-index.enabled": "true",
+ "fields.title_embedding.dimension": "768",
+ "fields.image_embedding.dimension": "512",
+ "ivf-pq.nlist": "256",
+ "fields.image_embedding.nlist": "512",
+ },
+)
+
+catalog.create_table("db.my_table", schema, ignore_if_exists=False)
+```
+
+</TabItem>
+
+</Tabs>
+
With the properties above, `title_embedding` is indexed with `nlist=256` while
`image_embedding`
uses `nlist=512`.
diff --git a/paimon-python/pypaimon/common/options/core_options.py
b/paimon-python/pypaimon/common/options/core_options.py
index 9ae32c187b..d76785198b 100644
--- a/paimon-python/pypaimon/common/options/core_options.py
+++ b/paimon-python/pypaimon/common/options/core_options.py
@@ -667,6 +667,13 @@ class CoreOptions:
)
)
+ GLOBAL_INDEX_ROW_COUNT_PER_SHARD: ConfigOption[int] = (
+ ConfigOptions.key("global-index.row-count-per-shard")
+ .long_type()
+ .default_value(100000)
+ .with_description("Row count per shard for global index.")
+ )
+
GLOBAL_INDEX_COLUMN_UPDATE_ACTION:
ConfigOption[GlobalIndexColumnUpdateAction] = (
ConfigOptions.key("global-index.column-update-action")
.enum_type(GlobalIndexColumnUpdateAction)
@@ -1197,6 +1204,9 @@ class CoreOptions:
def global_index_thread_num(self) -> Optional[int]:
return self.options.get(CoreOptions.GLOBAL_INDEX_THREAD_NUM)
+ def global_index_row_count_per_shard(self) -> int:
+ return self.options.get(CoreOptions.GLOBAL_INDEX_ROW_COUNT_PER_SHARD)
+
def btree_index_fallback_scan_max_size(self) -> int:
return self.options.get(
CoreOptions.BTREE_INDEX_FALLBACK_SCAN_MAX_SIZE
diff --git a/paimon-python/pypaimon/globalindex/create_global_index.py
b/paimon-python/pypaimon/globalindex/create_global_index.py
index 3aa3bb919f..b1c577005b 100644
--- a/paimon-python/pypaimon/globalindex/create_global_index.py
+++ b/paimon-python/pypaimon/globalindex/create_global_index.py
@@ -31,6 +31,12 @@ from pypaimon.globalindex.btree.btree_index_writer import (
)
from pypaimon.globalindex.global_index_meta import GlobalIndexMeta
from pypaimon.globalindex.key_serializer import create_serializer
+from pypaimon.globalindex.vindex.vindex_vector_global_index_reader import (
+ VINDEX_IDENTIFIERS,
+)
+from pypaimon.globalindex.vindex.vindex_vector_index_writer import (
+ VindexVectorIndexWriter,
+)
from pypaimon.index.index_file_meta import IndexFileMeta
from pypaimon.manifest.index_manifest_entry import IndexManifestEntry
from pypaimon.read.split import DataSplit
@@ -76,7 +82,7 @@ def create_global_index(
class GlobalIndexBuilder:
- """Small Python builder for sorted global indexes."""
+ """Small Python builder for global indexes."""
def __init__(
self,
@@ -95,15 +101,15 @@ class GlobalIndexBuilder:
self._options = _merged_options(table, options)
self._core_options = CoreOptions(self._options)
- if self._index_type != BTREE_IDENTIFIER:
+ if self._index_type != BTREE_IDENTIFIER and self._index_type not in
VINDEX_IDENTIFIERS:
raise ValueError(
- "Python global index build currently supports only '%s', got
'%s'."
- % (BTREE_IDENTIFIER, index_type)
+ "Python global index build currently supports '%s' and %s, got
'%s'."
+ % (BTREE_IDENTIFIER, VINDEX_IDENTIFIERS, index_type)
)
if len(self._index_columns) != 1:
raise ValueError(
- "Python '%s' global index build currently supports one column,
got %s."
- % (BTREE_IDENTIFIER, self._index_columns)
+ "Python global index build currently supports one column, got
%s."
+ % self._index_columns
)
if not self._table.options.row_tracking_enabled():
raise ValueError(
@@ -116,6 +122,8 @@ class GlobalIndexBuilder:
"Column '%s' does not exist in table '%s'."
% (column, self._table.identifier)
)
+ if self._index_type in VINDEX_IDENTIFIERS:
+ self._validate_vindex_table()
def build(self) -> List[CommitMessage]:
read_builder = self._table.new_read_builder()
@@ -129,11 +137,11 @@ class GlobalIndexBuilder:
return []
index_field = self._table.field_dict[self._index_columns[0]]
- key_serializer = create_serializer(index_field.type)
- block_size = self._core_options.btree_index_block_size()
- records_per_range = self._core_options.sorted_index_records_per_range()
- if records_per_range <= 0:
- raise ValueError("sorted-index.records-per-range must be
positive.")
+ if self._index_type in VINDEX_IDENTIFIERS:
+ splits = _filter_non_indexable_splits(
+ self._table, splits, self._index_columns)
+ if not splits:
+ return []
read_type = [index_field, SpecialFields.ROW_ID]
from pypaimon.read.table_read import TableRead
@@ -146,6 +154,19 @@ class GlobalIndexBuilder:
index_path_factory =
self._table.path_factory().global_index_path_factory()
index_path = index_path_factory.global_index_root_path()
+ if self._index_type == BTREE_IDENTIFIER:
+ return self._build_btree(splits, index_field, table_read,
index_path)
+ return self._build_vindex(splits, index_field, table_read, index_path)
+
+ def _build_btree(
+ self, splits, index_field, table_read, index_path: str
+ ) -> List[CommitMessage]:
+ key_serializer = create_serializer(index_field.type)
+ block_size = self._core_options.btree_index_block_size()
+ records_per_range = self._core_options.sorted_index_records_per_range()
+ if records_per_range <= 0:
+ raise ValueError("sorted-index.records-per-range must be
positive.")
+
messages = []
for split in _split_by_contiguous_row_range(splits):
row_range = _calc_row_range(split)
@@ -192,6 +213,77 @@ class GlobalIndexBuilder:
)
return messages
+ def _build_vindex(
+ self, splits, index_field, table_read, index_path: str
+ ) -> List[CommitMessage]:
+ rows_per_shard = self._core_options.global_index_row_count_per_shard()
+ if rows_per_shard <= 0:
+ raise ValueError(
+ "Option 'global-index.row-count-per-shard' must be greater
than 0."
+ )
+
+ messages = []
+ for split, row_range in _split_by_global_index_shard(
+ splits, rows_per_shard
+ ):
+ table = table_read.to_arrow([split])
+ if table is None or table.num_rows == 0:
+ continue
+
+ writer = VindexVectorIndexWriter(
+ self._table.file_io,
+ index_path,
+ index_field.type,
+ self._index_type,
+ self._options.to_map(),
+ index_field.name,
+ )
+ try:
+ for vector, row_id in _extract_vector_rows(
+ table,
+ self._index_columns[0],
+ SpecialFields.ROW_ID.name,
+ row_range,
+ ):
+ writer.write(vector, row_id - row_range.from_)
+
+ index_adds = _to_index_manifest_entries(
+ self._table,
+ split.partition,
+ row_range,
+ index_field.id,
+ self._index_type,
+ writer.finish(),
+ )
+ finally:
+ writer.close()
+ if index_adds:
+ messages.append(
+ CommitMessage(
+ partition=tuple(split.partition.values),
+ bucket=0,
+ new_files=[],
+ index_adds=index_adds,
+ )
+ )
+ return messages
+
+ def _validate_vindex_table(self) -> None:
+ bucket = self._core_options.bucket()
+ if bucket != -1:
+ raise ValueError(
+ "Generic global index only supports unaware-bucket tables "
+ "(bucket = -1), but table '%s' has bucket = %s."
+ % (self._table.identifier, bucket)
+ )
+ if self._core_options.deletion_vectors_enabled():
+ raise ValueError(
+ "Generic global index does not support tables with deletion "
+ "vectors enabled. Table '%s' has "
+ "'deletion-vectors.enabled' = true."
+ % self._table.identifier
+ )
+
def _resolve_partition_filter(self, read_builder) -> Optional[Predicate]:
if self._partition_filter is not None:
return self._partition_filter
@@ -259,6 +351,148 @@ def _split_by_contiguous_row_range(splits):
return result
+def _filter_non_indexable_splits(table, splits, index_columns):
+ boundary = _find_min_non_indexable_row_id(
+ table.schema_manager,
+ [file for split in splits for file in split.files],
+ index_columns,
+ )
+ if boundary is None:
+ return splits
+
+ result = []
+ for split in splits:
+ files = [
+ file for file in split.files
+ if file.row_id_range() is not None
+ and file.row_id_range().from_ < boundary
+ ]
+ if files:
+ result.append(_copy_split_with_files(split, files))
+ return result
+
+
+def _find_min_non_indexable_row_id(schema_manager, files, index_columns):
+ schema_contains_columns = {}
+ index_column_set = set(index_columns)
+ boundary = None
+ for file in files:
+ row_range = file.row_id_range()
+ if row_range is None:
+ continue
+
+ schema_id = file.schema_id
+ if schema_id not in schema_contains_columns:
+ schema = schema_manager.get_schema(schema_id)
+ schema_field_names = {field.name for field in schema.fields}
+ schema_contains_columns[schema_id] = (
+ index_column_set.issubset(schema_field_names)
+ )
+ if not schema_contains_columns[schema_id]:
+ if boundary is None or row_range.from_ < boundary:
+ boundary = row_range.from_
+ return boundary
+
+
+def _split_by_global_index_shard(splits, rows_per_shard):
+ if rows_per_shard <= 0:
+ raise ValueError(
+ "Option 'global-index.row-count-per-shard' must be greater than 0."
+ )
+
+ groups = {}
+ for split in splits:
+ key = (_partition_key(split.partition), split.bucket)
+ if key not in groups:
+ groups[key] = {
+ "partition": split.partition,
+ "bucket": split.bucket,
+ "files": [],
+ }
+ for file in split.files:
+ if file.row_id_range() is None:
+ continue
+ groups[key]["files"].append(file)
+
+ result = []
+ for group in groups.values():
+ files_by_shard = {}
+ for file in group["files"]:
+ file_range = file.row_id_range()
+ start_shard = file_range.from_ // rows_per_shard
+ end_shard = file_range.to // rows_per_shard
+ for shard_id in range(start_shard, end_shard + 1):
+ shard_start = shard_id * rows_per_shard
+ files_by_shard.setdefault(shard_start, []).append(file)
+
+ for shard_start in sorted(files_by_shard):
+ shard_end = shard_start + rows_per_shard - 1
+ shard_files = sorted(
+ files_by_shard[shard_start],
+ key=lambda file: file.row_id_range().from_,
+ )
+ current_group = []
+ current_group_end = None
+ for file in shard_files:
+ file_range = file.row_id_range()
+ if not current_group:
+ current_group.append(file)
+ current_group_end = file_range.to
+ elif file_range.from_ <= current_group_end + 1:
+ current_group.append(file)
+ current_group_end = max(current_group_end, file_range.to)
+ else:
+ _append_shard_split(
+ result,
+ current_group,
+ shard_start,
+ shard_end,
+ group["partition"],
+ group["bucket"],
+ )
+ current_group = [file]
+ current_group_end = file_range.to
+
+ if current_group:
+ _append_shard_split(
+ result,
+ current_group,
+ shard_start,
+ shard_end,
+ group["partition"],
+ group["bucket"],
+ )
+
+ return result
+
+
+def _partition_key(partition):
+ values = getattr(partition, "values", None)
+ return tuple(values) if values is not None else partition
+
+
+def _append_shard_split(
+ result,
+ files,
+ shard_start,
+ shard_end,
+ partition,
+ bucket,
+):
+ group_start = min(file.row_id_range().from_ for file in files)
+ group_end = max(file.row_id_range().to for file in files)
+ row_range = Range(max(group_start, shard_start), min(group_end, shard_end))
+ result.append((
+ DataSplit(
+ files=list(files),
+ partition=partition,
+ bucket=bucket,
+ raw_convertible=False,
+ ),
+ row_range,
+ ))
+
+
def _split_one_by_contiguous_row_range(split):
for file in split.files:
if file.row_id_range() is None:
@@ -342,6 +576,25 @@ def _extract_sorted_rows(
return sorted(rows, key=cmp_to_key(compare))
+def _extract_vector_rows(
+ table: pa.Table,
+ index_column: str,
+ row_id_column: str,
+ row_range: Optional[Range] = None,
+):
+ vectors = table.column(index_column).to_pylist()
+ row_ids = table.column(row_id_column).to_pylist()
+ rows = []
+ for vector, row_id in zip(vectors, row_ids):
+ if row_id is None:
+ raise ValueError("Cannot build global index because _ROW_ID is
null.")
+ row_id = int(row_id)
+ if row_range is not None and not row_range.contains(row_id):
+ continue
+ rows.append((vector, row_id))
+ return rows
+
+
def _chunks(rows, size):
for start in range(0, len(rows), size):
yield rows[start:start + size]
diff --git a/paimon-python/pypaimon/globalindex/vindex/__init__.py
b/paimon-python/pypaimon/globalindex/vindex/__init__.py
index 76373df86c..2ef60be3b8 100644
--- a/paimon-python/pypaimon/globalindex/vindex/__init__.py
+++ b/paimon-python/pypaimon/globalindex/vindex/__init__.py
@@ -15,4 +15,18 @@
# specific language governing permissions and limitations
# under the License.
-"""paimon-vindex based global index readers."""
+"""paimon-vindex based global index readers and writers."""
+
+from pypaimon.globalindex.vindex.vindex_vector_global_index_reader import (
+ VINDEX_IDENTIFIERS,
+ VindexVectorGlobalIndexReader,
+)
+from pypaimon.globalindex.vindex.vindex_vector_index_writer import (
+ VindexVectorIndexWriter,
+)
+
+__all__ = [
+ 'VINDEX_IDENTIFIERS',
+ 'VindexVectorGlobalIndexReader',
+ 'VindexVectorIndexWriter',
+]
diff --git
a/paimon-python/pypaimon/globalindex/vindex/vindex_vector_index_writer.py
b/paimon-python/pypaimon/globalindex/vindex/vindex_vector_index_writer.py
new file mode 100644
index 0000000000..2a6ed8670e
--- /dev/null
+++ b/paimon-python/pypaimon/globalindex/vindex/vindex_vector_index_writer.py
@@ -0,0 +1,300 @@
+# 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.
+
+"""paimon-vindex global index writer."""
+
+import math
+import os
+import tempfile
+import uuid
+from array import array
+from typing import BinaryIO, Dict, List, Mapping, Optional
+
+from pypaimon.globalindex.result_entry import ResultEntry
+from pypaimon.schema.data_types import ArrayType, AtomicType, DataType,
VectorType
+
+
+FILE_NAME_PREFIX = "vector"
+DEFAULT_DIMENSION = 128
+ADD_BATCH_SIZE = 10000
+
+
+class VindexVectorIndexWriter:
+ """Writer for one paimon-vindex global index file."""
+
+ def __init__(
+ self,
+ file_io,
+ index_path: str,
+ data_type: DataType,
+ index_type: str,
+ options: Mapping[str, object],
+ field_name: str,
+ ):
+ self.file_name = (
+ "%s-%s-global-index-%s.index"
+ % (FILE_NAME_PREFIX, index_type, uuid.uuid4())
+ )
+ self._file_io = file_io
+ self._index_path = index_path.rstrip('/')
+ self._index_type = index_type
+ self._native_options = native_options(
+ data_type, options, index_type, field_name)
+ self._dimension = int(self._native_options["dimension"])
+ self._row_count = 0
+ self._vector_count = 0
+ self._row_id_temp: Optional[BinaryIO] = None
+ self._vector_temp: Optional[BinaryIO] = None
+ self._row_id_temp_path: Optional[str] = None
+ self._vector_temp_path: Optional[str] = None
+ self._closed = False
+
+ validate_vector_type(data_type)
+
+ @property
+ def native_options(self) -> Dict[str, str]:
+ return dict(self._native_options)
+
+ def write(self, vector, relative_row_id: int) -> None:
+ if self._closed:
+ raise RuntimeError("VindexVectorIndexWriter is already closed.")
+
+ self._row_count += 1
+ if vector is None:
+ return
+
+ materialized = _materialize_vector(
+ vector, self._dimension, relative_row_id)
+ self._ensure_temp_files()
+ self._row_id_temp.write(array("q", [int(relative_row_id)]).tobytes())
+ self._vector_temp.write(array("f", materialized).tobytes())
+ self._vector_count += 1
+
+ def finish(self) -> List[ResultEntry]:
+ if self._closed:
+ raise RuntimeError("VindexVectorIndexWriter is already closed.")
+ self._closed = True
+
+ file_path = self._file_path()
+ try:
+ if self._vector_count == 0:
+ return []
+
+ try:
+ import numpy as np
+ from paimon_vindex import VectorIndexWriter
+ except ImportError as e:
+ raise ImportError(
+ "paimon-vindex is required to build vindex vector indexes.
"
+ "Install paimon-vindex==0.1.0 or pypaimon[vindex].") from e
+
+ self._close_temp_files()
+ self._file_io.check_or_mkdirs(self._index_path)
+ vectors = np.fromfile(
+ self._vector_temp_path,
+ dtype=np.float32,
+ count=self._vector_count * self._dimension,
+ ).reshape(self._vector_count, self._dimension)
+ with VectorIndexWriter(self._native_options) as writer:
+ writer.train(vectors)
+ del vectors
+ self._add_vectors_in_batches(np, writer)
+ with self._file_io.new_output_stream(file_path) as
output_stream:
+ writer.write(output_stream)
+ except Exception:
+ self._file_io.delete_quietly(file_path)
+ raise
+ finally:
+ self._delete_temp_files()
+
+ return [ResultEntry(self.file_name, self._row_count, b"{}")]
+
+ def _file_path(self) -> str:
+ return "%s/%s" % (self._index_path, self.file_name)
+
+ def close(self) -> None:
+ if not self._closed:
+ self._closed = True
+ self._delete_temp_files()
+
+ def _ensure_temp_files(self) -> None:
+ if self._row_id_temp is not None:
+ return
+
+ row_id_temp = tempfile.NamedTemporaryFile(
+ prefix="paimon-vindex-row-ids-", suffix=".bin", delete=False)
+ vector_temp = tempfile.NamedTemporaryFile(
+ prefix="paimon-vindex-vectors-", suffix=".bin", delete=False)
+ self._row_id_temp = row_id_temp
+ self._vector_temp = vector_temp
+ self._row_id_temp_path = row_id_temp.name
+ self._vector_temp_path = vector_temp.name
+
+ def _close_temp_files(self) -> None:
+ for temp_file in (self._row_id_temp, self._vector_temp):
+ if temp_file is not None and not temp_file.closed:
+ temp_file.flush()
+ temp_file.close()
+ self._row_id_temp = None
+ self._vector_temp = None
+
+ def _delete_temp_files(self) -> None:
+ self._close_temp_files()
+ for path in (self._row_id_temp_path, self._vector_temp_path):
+ if path is not None and os.path.exists(path):
+ try:
+ os.remove(path)
+ except OSError:
+ pass
+ self._row_id_temp_path = None
+ self._vector_temp_path = None
+
+ def _add_vectors_in_batches(self, np, writer) -> None:
+ with open(self._row_id_temp_path, "rb") as row_id_file, open(
+ self._vector_temp_path, "rb"
+ ) as vector_file:
+ remaining = self._vector_count
+ while remaining > 0:
+ batch_size = min(ADD_BATCH_SIZE, remaining)
+ row_ids = np.fromfile(
+ row_id_file, dtype=np.int64, count=batch_size)
+ vectors = np.fromfile(
+ vector_file,
+ dtype=np.float32,
+ count=batch_size * self._dimension,
+ ).reshape(batch_size, self._dimension)
+ writer.add_vectors(row_ids, vectors)
+ remaining -= batch_size
+
+
+def native_options(
+ data_type: DataType,
+ options: Mapping[str, object],
+ index_type: str,
+ field_name: str,
+) -> Dict[str, str]:
+ result: Dict[str, str] = {}
+ option_prefix = "%s." % index_type
+ field_prefix = "fields.%s." % field_name
+
+ for key, value in options.items():
+ key = str(key)
+ if key.startswith(option_prefix):
+ native_key = _native_option_key(key[len(option_prefix):])
+ if native_key is not None:
+ result[native_key] = str(value)
+
+ for key, value in options.items():
+ key = str(key)
+ if key.startswith(field_prefix):
+ native_key = _native_option_key(key[len(field_prefix):])
+ if native_key is not None:
+ result[native_key] = str(value)
+
+ result["index.type"] = index_type.replace('-', '_')
+ result["dimension"] = str(_dimension(data_type, result, index_type))
+ return result
+
+
+def validate_vector_type(data_type: DataType) -> None:
+ if isinstance(data_type, VectorType):
+ element_type = data_type.element
+ if _is_float_type(element_type):
+ return
+ raise ValueError(
+ "Vector index requires float vector, but got: %s" % element_type)
+
+ if isinstance(data_type, ArrayType):
+ element_type = data_type.element
+ if _is_float_type(element_type):
+ return
+ raise ValueError(
+ "Vector index requires float array, but got: %s" % element_type)
+
+ raise ValueError(
+ "Vector index requires VectorType or ArrayType<FLOAT>, but got: %s"
+ % data_type)
+
+
+def _native_option_key(option_key: str) -> Optional[str]:
+ if option_key in ("index.dimension", "dimension"):
+ return "dimension"
+ if option_key in ("distance.metric", "metric"):
+ return "metric"
+ if option_key in (
+ "nlist",
+ "pq.m",
+ "hnsw.m",
+ "hnsw.ef-construction",
+ "hnsw.max-level",
+ ):
+ return option_key
+ if option_key in ("pq.use-opq", "use-opq"):
+ return "use-opq"
+ return None
+
+
+def _dimension(
+ data_type: DataType, native_options_map: Mapping[str, str], index_type: str
+) -> int:
+ if isinstance(data_type, VectorType):
+ return data_type.length
+
+ dimension = native_options_map.get("dimension")
+ value = DEFAULT_DIMENSION if dimension is None else int(dimension)
+ if value <= 0:
+ raise ValueError(
+ "Invalid value for '%s.dimension': %s. Must be a positive integer."
+ % (index_type, value)
+ )
+ return value
+
+
+def _is_float_type(data_type: DataType) -> bool:
+ return (
+ isinstance(data_type, AtomicType)
+ and data_type.type.upper() == "FLOAT"
+ )
+
+
+def _materialize_vector(
+ value, dimension: int, relative_row_id: int
+) -> List[float]:
+ if hasattr(value, "as_py"):
+ value = value.as_py()
+ if hasattr(value, "tolist"):
+ value = value.tolist()
+ if not isinstance(value, (list, tuple)):
+ raise ValueError("Unsupported vector value: %s" % type(value).__name__)
+ if len(value) != dimension:
+ raise ValueError(
+ "Vector dimension mismatch: expected %d, but got %d"
+ % (dimension, len(value))
+ )
+
+ vector = []
+ for index, element in enumerate(value):
+ if element is None:
+ raise ValueError("Vector element at index %d is null" % index)
+ element = float(element)
+ if not math.isfinite(element):
+ raise ValueError(
+ "Vector element at rowId=%d, index=%d is %s"
+ % (relative_row_id, index, element)
+ )
+ vector.append(element)
+ return vector
diff --git a/paimon-python/pypaimon/tests/global_index_build_test.py
b/paimon-python/pypaimon/tests/global_index_build_test.py
index 15fd3bfd39..00279d01d7 100644
--- a/paimon-python/pypaimon/tests/global_index_build_test.py
+++ b/paimon-python/pypaimon/tests/global_index_build_test.py
@@ -19,13 +19,21 @@ import unittest
from datetime import date, datetime
from decimal import Decimal
import os
+import sys
+import types
import pyarrow as pa
from pypaimon.globalindex.create_global_index import (
+ _filter_non_indexable_splits,
+ _split_by_global_index_shard,
_split_one_by_contiguous_row_range,
)
from pypaimon.globalindex.key_serializer import create_serializer
+from pypaimon.globalindex.vindex.vindex_vector_index_writer import (
+ VindexVectorIndexWriter,
+ native_options,
+)
from pypaimon.globalindex.global_index_scanner import GlobalIndexScanner
from pypaimon.index.index_file_handler import IndexFileHandler
from pypaimon.schema.data_types import ArrayType, AtomicType, RowType
@@ -39,12 +47,15 @@ from pypaimon.utils.range import Range
class _FakeFile:
- def __init__(self, file_name, first_row_id, row_count):
+ def __init__(self, file_name, first_row_id, row_count, schema_id=0):
self.file_name = file_name
self.first_row_id = first_row_id
self.row_count = row_count
+ self.schema_id = schema_id
def row_id_range(self):
+ if self.first_row_id is None:
+ return None
return Range(self.first_row_id,
self.first_row_id + self.row_count - 1)
@@ -58,6 +69,52 @@ class _FakeSplit:
self.raw_convertible = False
+class _FakeVectorIndexWriter:
+ instances = []
+
+ def __init__(self, options):
+ self.options = dict(options)
+ self.trained = None
+ self.added_ids = None
+ self.added_vectors = None
+ self.closed = False
+ _FakeVectorIndexWriter.instances.append(self)
+
+ def train(self, data):
+ self.trained = data.tolist()
+
+ def add_vectors(self, ids, data):
+ self.added_ids = ids.tolist()
+ self.added_vectors = data.tolist()
+
+ def write(self, file):
+ file.write(b"fake-vindex")
+
+ def close(self):
+ self.closed = True
+
+ def __enter__(self):
+ return self
+
+ def __exit__(self, exc_type, exc_val, exc_tb):
+ self.close()
+ return False
+
+
+class _FakeSchemaManager:
+
+ def __init__(self, fields_by_schema_id):
+ self.fields_by_schema_id = fields_by_schema_id
+
+ def get_schema(self, schema_id):
+ return types.SimpleNamespace(
+ fields=[
+ types.SimpleNamespace(name=name)
+ for name in self.fields_by_schema_id[schema_id]
+ ]
+ )
+
+
class GlobalIndexBuildTest(
BatchModeMixin, DataEvolutionTestBase, unittest.TestCase):
@@ -245,6 +302,168 @@ class GlobalIndexBuildTest(
for column in ['flag', 'amount', 'dt', 'ts']:
self.assertEqual(1, table.create_global_index(column))
+ def test_create_vindex_global_index_from_python(self):
+ schema = pa.schema([
+ ('id', pa.int32()),
+ ('embedding', pa.list_(pa.float32())),
+ ])
+ table = self._create_table(pa_schema=schema,
options=self.table_options)
+ vectors = pa.array(
+ [[1.0, 0.0], [0.0, 1.0], None],
+ type=pa.list_(pa.float32()),
+ )
+ self._write_arrow(table, pa.table(
+ {'id': [1, 2, 3], 'embedding': vectors},
+ schema=schema,
+ ))
+
+ old_module = sys.modules.get("paimon_vindex")
+ sys.modules["paimon_vindex"] = types.SimpleNamespace(
+ VectorIndexWriter=_FakeVectorIndexWriter)
+ _FakeVectorIndexWriter.instances = []
+ try:
+ added = table.create_global_index(
+ 'embedding',
+ index_type='ivf-flat',
+ options={
+ 'ivf-flat.dimension': '2',
+ 'ivf-flat.distance.metric': 'l2',
+ 'ivf-flat.nlist': '1',
+ },
+ )
+ finally:
+ if old_module is None:
+ sys.modules.pop("paimon_vindex", None)
+ else:
+ sys.modules["paimon_vindex"] = old_module
+
+ self.assertEqual(1, added)
+ self.assertEqual(1, len(_FakeVectorIndexWriter.instances))
+ fake_writer = _FakeVectorIndexWriter.instances[0]
+ self.assertEqual('ivf_flat', fake_writer.options['index.type'])
+ self.assertEqual('2', fake_writer.options['dimension'])
+ self.assertEqual('l2', fake_writer.options['metric'])
+ self.assertEqual('1', fake_writer.options['nlist'])
+ self.assertEqual([[1.0, 0.0], [0.0, 1.0]], fake_writer.trained)
+ self.assertEqual([0, 1], fake_writer.added_ids)
+ self.assertEqual([[1.0, 0.0], [0.0, 1.0]], fake_writer.added_vectors)
+ self.assertTrue(fake_writer.closed)
+
+ snapshot = table.snapshot_manager().get_latest_snapshot()
+ entries = IndexFileHandler(table).scan(snapshot)
+ self.assertEqual(1, len(entries))
+ entry = entries[0]
+ self.assertEqual('ivf-flat', entry.index_file.index_type)
+ self.assertEqual(3, entry.index_file.row_count)
+ self.assertEqual(b'{}',
bytes(entry.index_file.global_index_meta.index_meta))
+ self.assertTrue(table.file_io.exists(
+ table.path_factory().global_index_path_factory().to_path(
+ entry.index_file.file_name)))
+
+ def test_create_vindex_global_index_respects_row_count_per_shard(self):
+ schema = pa.schema([
+ ('id', pa.int32()),
+ ('embedding', pa.list_(pa.float32())),
+ ])
+ table = self._create_table(pa_schema=schema,
options=self.table_options)
+ vectors = pa.array(
+ [[1.0, 0.0], [0.0, 1.0], [0.5, 0.5], [0.2, 0.8], [0.9, 0.1]],
+ type=pa.list_(pa.float32()),
+ )
+ self._write_arrow(table, pa.table(
+ {'id': [1, 2, 3, 4, 5], 'embedding': vectors},
+ schema=schema,
+ ))
+
+ old_module = sys.modules.get("paimon_vindex")
+ sys.modules["paimon_vindex"] = types.SimpleNamespace(
+ VectorIndexWriter=_FakeVectorIndexWriter)
+ _FakeVectorIndexWriter.instances = []
+ try:
+ added = table.create_global_index(
+ 'embedding',
+ index_type='ivf-flat',
+ options={
+ 'global-index.row-count-per-shard': '2',
+ 'ivf-flat.dimension': '2',
+ },
+ )
+ finally:
+ if old_module is None:
+ sys.modules.pop("paimon_vindex", None)
+ else:
+ sys.modules["paimon_vindex"] = old_module
+
+ self.assertEqual(3, added)
+ self.assertEqual(3, len(_FakeVectorIndexWriter.instances))
+ self.assertEqual(
+ [[0, 1], [0, 1], [0]],
+ [writer.added_ids for writer in _FakeVectorIndexWriter.instances],
+ )
+
+ snapshot = table.snapshot_manager().get_latest_snapshot()
+ entries = sorted(
+ IndexFileHandler(table).scan(snapshot),
+ key=lambda entry:
entry.index_file.global_index_meta.row_range_start,
+ )
+ self.assertEqual(
+ [(0, 1, 2), (2, 3, 2), (4, 4, 1)],
+ [
+ (
+ entry.index_file.global_index_meta.row_range_start,
+ entry.index_file.global_index_meta.row_range_end,
+ entry.index_file.row_count,
+ )
+ for entry in entries
+ ],
+ )
+
+ def
test_create_vindex_global_index_rejects_generic_unsupported_tables(self):
+ schema = pa.schema([
+ ('id', pa.int32()),
+ ('embedding', pa.list_(pa.float32())),
+ ])
+ bucket_options = dict(self.table_options)
+ bucket_options['bucket'] = '1'
+ bucket_table = self._create_table(pa_schema=schema,
options=bucket_options)
+ with self.assertRaisesRegex(ValueError, 'unaware-bucket'):
+ bucket_table.create_global_index(
+ 'embedding',
+ index_type='ivf-flat',
+ options={'ivf-flat.dimension': '2'},
+ )
+
+ dv_options = dict(self.table_options)
+ dv_options['deletion-vectors.enabled'] = 'true'
+ dv_table = self._create_table(pa_schema=schema, options=dv_options)
+ with self.assertRaisesRegex(ValueError, 'deletion vectors'):
+ dv_table.create_global_index(
+ 'embedding',
+ index_type='ivf-flat',
+ options={'ivf-flat.dimension': '2'},
+ )
+
+ def test_vindex_native_options_follow_java_mapping(self):
+ data_type = ArrayType(True, AtomicType('FLOAT'))
+ options = {
+ 'ivf-pq.dimension': '128',
+ 'ivf-pq.distance.metric': 'cosine',
+ 'ivf-pq.nlist': '256',
+ 'ivf-pq.pq.m': '16',
+ 'fields.embedding.dimension': '64',
+ 'fields.embedding.nlist': '512',
+ 'fields.embedding.pq.use-opq': 'true',
+ }
+
+ result = native_options(data_type, options, 'ivf-pq', 'embedding')
+
+ self.assertEqual('ivf_pq', result['index.type'])
+ self.assertEqual('64', result['dimension'])
+ self.assertEqual('cosine', result['metric'])
+ self.assertEqual('512', result['nlist'])
+ self.assertEqual('16', result['pq.m'])
+ self.assertEqual('true', result['use-opq'])
+
def test_split_by_contiguous_row_range_matches_java_builder(self):
split = _FakeSplit([
_FakeFile('a', 0, 2),
@@ -260,6 +479,91 @@ class GlobalIndexBuildTest(
[[file.file_name for file in s.files] for s in splits],
)
+ def test_split_by_global_index_shard_matches_java_default_builder(self):
+ split = _FakeSplit([
+ _FakeFile('a', 0, 3),
+ _FakeFile('b', 3, 2),
+ _FakeFile('c', 6, 3),
+ ])
+
+ shards = _split_by_global_index_shard([split], 4)
+
+ self.assertEqual(
+ [
+ (['a', 'b'], 0, 3),
+ (['b'], 4, 4),
+ (['c'], 6, 7),
+ (['c'], 8, 8),
+ ],
+ [
+ ([file.file_name for file in shard.files], row_range.from_,
row_range.to)
+ for shard, row_range in shards
+ ],
+ )
+
+ def test_split_by_global_index_shard_skips_files_without_row_ids(self):
+ split = _FakeSplit([
+ _FakeFile('no-row-id', None, 2),
+ _FakeFile('indexed', 4, 2),
+ ])
+
+ shards = _split_by_global_index_shard([split], 4)
+
+ self.assertEqual(
+ [(['indexed'], 4, 5)],
+ [
+ ([file.file_name for file in shard.files], row_range.from_,
row_range.to)
+ for shard, row_range in shards
+ ],
+ )
+
+ def test_filter_non_indexable_splits_matches_java_generic_builder(self):
+ split = _FakeSplit([
+ _FakeFile('indexable-before', 0, 2, schema_id=0),
+ _FakeFile('non-indexable', 2, 2, schema_id=1),
+ _FakeFile('indexable-after-boundary', 4, 2, schema_id=0),
+ ])
+ table = types.SimpleNamespace(
+ schema_manager=_FakeSchemaManager({
+ 0: ['id', 'embedding'],
+ 1: ['id'],
+ })
+ )
+
+ splits = _filter_non_indexable_splits(
+ table, [split], ['embedding'])
+
+ self.assertEqual(1, len(splits))
+ self.assertEqual(
+ ['indexable-before'],
+ [file.file_name for file in splits[0].files],
+ )
+
+ def test_vindex_writer_close_cleans_temp_files(self):
+ schema = pa.schema([
+ ('id', pa.int32()),
+ ('embedding', pa.list_(pa.float32())),
+ ])
+ table = self._create_table(pa_schema=schema,
options=self.table_options)
+ writer = VindexVectorIndexWriter(
+ table.file_io,
+
table.path_factory().global_index_path_factory().global_index_root_path(),
+ ArrayType(True, AtomicType('FLOAT')),
+ 'ivf-flat',
+ {'ivf-flat.dimension': '2'},
+ 'embedding',
+ )
+ writer.write([1.0, 0.0], 0)
+ row_id_temp_path = writer._row_id_temp_path
+ vector_temp_path = writer._vector_temp_path
+ self.assertTrue(os.path.exists(row_id_temp_path))
+ self.assertTrue(os.path.exists(vector_temp_path))
+
+ writer.close()
+
+ self.assertFalse(os.path.exists(row_id_temp_path))
+ self.assertFalse(os.path.exists(vector_temp_path))
+
def test_java_scalar_key_serializers_round_trip(self):
cases = [
('BOOLEAN', True),
