| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, |
| # software distributed under the License is distributed on an |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| # KIND, either express or implied. See the License for the |
| # specific language governing permissions and limitations |
| # under the License. |
| |
| import unittest |
| from datetime import date, datetime |
| from decimal import Decimal |
| import os |
| import struct |
| import sys |
| import types |
| |
| import pyarrow as pa |
| |
| from pypaimon.globalindex.build_plan import ( |
| filter_non_indexable_splits as _filter_non_indexable_splits, |
| split_by_global_index_shard as _split_by_global_index_shard, |
| split_one_by_contiguous_row_range as _split_one_by_contiguous_row_range, |
| ) |
| from pypaimon.globalindex.create_global_index import GlobalIndexBuilder |
| from pypaimon.globalindex.key_serializer import create_serializer |
| from pypaimon.globalindex.full_text.native_full_text_global_index_reader import ( |
| FULL_TEXT_IDENTIFIER, |
| NativeFullTextIndexOptions, |
| ) |
| from pypaimon.globalindex.full_text.native_full_text_index_writer import ( |
| NativeFullTextIndexWriter, |
| ) |
| from pypaimon.globalindex.vindex.vindex_vector_index_writer import ( |
| VindexVectorIndexWriter, |
| _sample_training_vectors, |
| native_options, |
| train_sample_ratio, |
| ) |
| from pypaimon.globalindex.data_evolution_global_index_scanner import DataEvolutionGlobalIndexScanner |
| from pypaimon.index.index_file_handler import IndexFileHandler |
| from pypaimon.schema.data_types import ArrayType, AtomicType, RowType |
| from pypaimon.tests.data_evolution_test_helpers import ( |
| BatchModeMixin, |
| DataEvolutionTestBase, |
| ) |
| from pypaimon.table.row.generic_row import GenericRow |
| from pypaimon.utils.range import Range |
| |
| |
| class _FakeFile: |
| |
| 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) |
| |
| |
| class _FakeSplit: |
| |
| def __init__(self, files): |
| self.files = files |
| self.partition = GenericRow([], []) |
| self.bucket = 0 |
| self.raw_convertible = False |
| |
| |
| class _FakeVectorIndexTraining: |
| |
| def __init__(self, options, data): |
| self.options = dict(options) |
| self.trained = data.tolist() |
| self.closed = False |
| |
| def close(self): |
| self.closed = True |
| |
| |
| class _FakeVectorIndexTrainer: |
| |
| @classmethod |
| def train(cls, options, data): |
| return _FakeVectorIndexTraining(options, data) |
| |
| |
| class _FakeVectorIndexWriter: |
| instances = [] |
| |
| def __init__(self, training): |
| self.options = dict(training.options) |
| self.trained = training.trained |
| self.added_ids = None |
| self.added_vectors = None |
| self.closed = False |
| _FakeVectorIndexWriter.instances.append(self) |
| |
| 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 _FakeFullTextForBuild(types.SimpleNamespace): |
| |
| def __init__(self): |
| super().__init__() |
| self.writers = [] |
| parent = self |
| |
| class FullTextIndexWriter: |
| |
| def __init__(self_inner, options=None): |
| self_inner.options = dict(options or {}) |
| self_inner.documents = [] |
| self_inner.written = False |
| self_inner.closed = False |
| parent.writers.append(self_inner) |
| |
| def add_document(self_inner, row_id, text): |
| self_inner.documents.append({ |
| "row_id": int(row_id), |
| "text": str(text), |
| }) |
| |
| def write(self_inner, output): |
| self_inner.written = True |
| output.write(b"fake-ftindex") |
| |
| def close(self_inner): |
| self_inner.closed = True |
| |
| self.FullTextIndexWriter = FullTextIndexWriter |
| |
| |
| def _archive_file_names(file_io, file_path): |
| stream = file_io.new_input_stream(file_path) |
| try: |
| file_count = struct.unpack(">i", stream.read(4))[0] |
| names = [] |
| for _ in range(file_count): |
| name_len = struct.unpack(">i", stream.read(4))[0] |
| names.append(stream.read(name_len).decode("utf-8")) |
| data_len = struct.unpack(">q", stream.read(8))[0] |
| stream.seek(stream.tell() + data_len) |
| return names |
| finally: |
| stream.close() |
| |
| |
| 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): |
| |
| table_options = { |
| 'row-tracking.enabled': 'true', |
| 'data-evolution.enabled': 'true', |
| 'global-index.enabled': 'true', |
| 'bucket': '-1', |
| 'file.format': 'parquet', |
| } |
| |
| def test_create_btree_global_index_from_python(self): |
| table = self._create_table() |
| self._write_arrow(table, pa.table( |
| { |
| 'id': [3, 1, 2, 2], |
| 'name': ['c', 'a', 'b1', 'b2'], |
| 'age': [30, 10, 20, 21], |
| 'city': ['z', 'x', 'y', 'y2'], |
| }, |
| schema=self.pa_schema, |
| )) |
| |
| added = table.create_global_index( |
| 'id', |
| options={'sorted-index.records-per-range': '2'}, |
| ) |
| |
| self.assertEqual(2, added) |
| snapshot = table.snapshot_manager().get_latest_snapshot() |
| self.assertIsNotNone(snapshot.index_manifest) |
| |
| entries = IndexFileHandler(table).scan(snapshot) |
| self.assertEqual(2, len(entries)) |
| self.assertEqual({'btree'}, {e.index_file.index_type for e in entries}) |
| self.assertEqual({0}, {e.index_file.global_index_meta.row_range_start for e in entries}) |
| self.assertEqual({3}, {e.index_file.global_index_meta.row_range_end for e in entries}) |
| |
| read_builder = table.new_read_builder() |
| predicate = read_builder.new_predicate_builder().equal('id', 2) |
| with DataEvolutionGlobalIndexScanner.create( |
| table, |
| predicate=predicate, |
| snapshot=snapshot) as scanner: |
| result = scanner.scan(predicate) |
| |
| self.assertEqual( |
| [Range(2, 3)], |
| result.results().to_range_list(), |
| ) |
| |
| table_name = table.identifier.get_full_name() |
| table_indexes = self.catalog.get_table(table_name + '$table_indexes') |
| index_read_builder = table_indexes.new_read_builder().with_projection([ |
| 'index_type', |
| 'index_field_name', |
| 'row_range_start', |
| 'row_range_end', |
| ]) |
| index_table = index_read_builder.new_read().to_arrow( |
| index_read_builder.new_scan().plan().splits()) |
| self.assertEqual(2, index_table.num_rows) |
| self.assertEqual(['btree', 'btree'], |
| index_table.column('index_type').to_pylist()) |
| self.assertEqual(['id', 'id'], |
| index_table.column('index_field_name').to_pylist()) |
| self.assertEqual([0, 0], |
| index_table.column('row_range_start').to_pylist()) |
| self.assertEqual([3, 3], |
| index_table.column('row_range_end').to_pylist()) |
| |
| key_ranges = self.catalog.get_table(table_name + '$file_key_ranges') |
| range_read_builder = key_ranges.new_read_builder().with_projection([ |
| 'file_path', |
| 'record_count', |
| 'first_row_id', |
| ]) |
| range_table = range_read_builder.new_read().to_arrow( |
| range_read_builder.new_scan().plan().splits()) |
| self.assertEqual(1, range_table.num_rows) |
| file_path = range_table.column('file_path').to_pylist()[0] |
| self.assertIn('/bucket-0/', file_path) |
| self.assertEqual([4], range_table.column('record_count').to_pylist()) |
| self.assertEqual([0], range_table.column('first_row_id').to_pylist()) |
| |
| dv_ranges_type = table_indexes.row_type().fields[6].type |
| self.assertIsInstance(dv_ranges_type, ArrayType) |
| self.assertIsInstance(dv_ranges_type.element, RowType) |
| self.assertEqual( |
| ['f0', 'f1', 'f2', '_CARDINALITY'], |
| [field.name for field in dv_ranges_type.element.fields], |
| ) |
| |
| self.assertEqual(2, table.drop_global_index('id', dry_run=True)) |
| self.assertEqual(2, len(IndexFileHandler(table).scan( |
| table.snapshot_manager().get_latest_snapshot()))) |
| |
| self.assertEqual(2, table.drop_global_index('id')) |
| latest_snapshot = table.snapshot_manager().get_latest_snapshot() |
| self.assertEqual([], IndexFileHandler(table).scan(latest_snapshot)) |
| |
| index_read_builder = table_indexes.new_read_builder() |
| index_table = index_read_builder.new_read().to_arrow( |
| index_read_builder.new_scan().plan().splits()) |
| self.assertEqual(0, index_table.num_rows) |
| |
| def test_stale_global_index_commit_conflicts_when_data_files_removed(self): |
| table = self._create_table() |
| self._write_arrow(table, pa.table( |
| { |
| 'id': [1, 2, 3, 4], |
| 'name': ['a', 'b', 'c', 'd'], |
| 'age': [10, 20, 30, 40], |
| 'city': ['x', 'y', 'z', 'w'], |
| }, |
| schema=self.pa_schema, |
| )) |
| |
| messages = GlobalIndexBuilder( |
| table, |
| 'id', |
| options={'sorted-index.records-per-range': '2'}, |
| ).build() |
| self.assertGreater(sum(len(message.index_adds) for message in messages), 0) |
| |
| overwrite_wb = table.new_batch_write_builder().overwrite({}) |
| overwrite_write = overwrite_wb.new_write() |
| overwrite_commit = overwrite_wb.new_commit() |
| overwrite_write.write_arrow(pa.table( |
| { |
| 'id': [5, 6], |
| 'name': ['e', 'f'], |
| 'age': [50, 60], |
| 'city': ['new1', 'new2'], |
| }, |
| schema=self.pa_schema, |
| )) |
| overwrite_commit.commit(overwrite_write.prepare_commit()) |
| overwrite_write.close() |
| overwrite_commit.close() |
| |
| stale_commit = table.new_batch_write_builder().new_commit() |
| with self.assertRaises(RuntimeError) as ctx: |
| stale_commit.commit(messages) |
| stale_commit.close() |
| self.assertIn( |
| 'Global index row ID existence conflict', |
| str(ctx.exception), |
| ) |
| self.assertEqual( |
| [], |
| IndexFileHandler(table).scan(table.snapshot_manager().get_latest_snapshot()), |
| ) |
| |
| def test_create_bitmap_global_index_from_python(self): |
| table = self._create_table() |
| self._write_arrow(table, pa.table( |
| { |
| 'id': [1, 2, 3, 4, 5], |
| 'name': ['a', 'b', 'c', 'd', 'e'], |
| 'age': [10, 20, 30, 40, 50], |
| 'city': ['vip', 'trial', None, 'vip', 'blocked'], |
| }, |
| schema=self.pa_schema, |
| )) |
| |
| added = table.create_global_index( |
| 'city', |
| index_type='bitmap', |
| options={ |
| 'sorted-index.records-per-range': '2', |
| 'bitmap-index.dictionary-block-size': '1 b', |
| }, |
| ) |
| |
| self.assertEqual(3, added) |
| snapshot = table.snapshot_manager().get_latest_snapshot() |
| entries = IndexFileHandler(table).scan(snapshot) |
| self.assertEqual(3, len(entries)) |
| self.assertEqual({'bitmap'}, {e.index_file.index_type for e in entries}) |
| self.assertEqual([1, 2, 2], |
| sorted(e.index_file.row_count for e in entries)) |
| self.assertEqual( |
| {0}, |
| {e.index_file.global_index_meta.row_range_start for e in entries}, |
| ) |
| self.assertEqual( |
| {4}, |
| {e.index_file.global_index_meta.row_range_end for e in entries}, |
| ) |
| |
| read_builder = table.new_read_builder() |
| predicate_builder = read_builder.new_predicate_builder() |
| cases = [ |
| (predicate_builder.is_in('city', ['vip', 'trial']), |
| [Range(0, 1), Range(3, 3)]), |
| (predicate_builder.is_null('city'), [Range(2, 2)]), |
| (predicate_builder.not_equal('city', 'blocked'), |
| [Range(0, 1), Range(3, 3)]), |
| ] |
| for predicate, expected in cases: |
| with DataEvolutionGlobalIndexScanner.create( |
| table, |
| predicate=predicate, |
| snapshot=snapshot) as scanner: |
| result = scanner.scan(predicate) |
| self.assertEqual(expected, result.results().to_range_list()) |
| |
| def test_create_bitmap_global_index_rejects_unsupported_compression(self): |
| table = self._create_table() |
| self._write_arrow(table, pa.table( |
| { |
| 'id': [1], |
| 'name': ['a'], |
| 'age': [10], |
| 'city': ['vip'], |
| }, |
| schema=self.pa_schema, |
| )) |
| |
| with self.assertRaisesRegex(ValueError, 'bitmap-index.compression=none'): |
| table.create_global_index( |
| 'city', |
| index_type='bitmap', |
| options={'bitmap-index.compression': 'lz4'}, |
| ) |
| |
| def test_sorted_index_records_per_range_matches_java_floating_factor(self): |
| table = self._create_table() |
| rows = list(range(12)) |
| self._write_arrow(table, pa.table( |
| { |
| 'id': rows, |
| 'name': ['n%s' % i for i in rows], |
| 'age': rows, |
| 'city': ['c%s' % i for i in rows], |
| }, |
| schema=self.pa_schema, |
| )) |
| |
| added = table.create_global_index( |
| 'id', |
| options={'sorted-index.records-per-range': '10'}, |
| ) |
| |
| self.assertEqual(1, added) |
| |
| def test_create_global_index_uses_external_path(self): |
| external_root = 'file://%s' % os.path.join( |
| self.tempdir, 'global-index-external') |
| options = dict(self.table_options) |
| options['global-index.external-path'] = external_root |
| table = self._create_table(options=options) |
| self._write_arrow(table, pa.table( |
| { |
| 'id': [1, 2], |
| 'name': ['a', 'b'], |
| 'age': [10, 20], |
| 'city': ['x', 'y'], |
| }, |
| schema=self.pa_schema, |
| )) |
| |
| self.assertEqual(1, table.create_global_index('id')) |
| |
| snapshot = table.snapshot_manager().get_latest_snapshot() |
| entries = IndexFileHandler(table).scan(snapshot) |
| self.assertEqual(1, len(entries)) |
| external_path = entries[0].index_file.external_path |
| self.assertIsNotNone(external_path) |
| self.assertTrue(external_path.startswith(external_root + '/')) |
| self.assertTrue(table.file_io.exists(external_path)) |
| |
| def test_create_global_index_skips_existing_ranges(self): |
| table = self._create_table() |
| self._write_arrow(table, pa.table( |
| { |
| 'id': [3, 1, 2, 2], |
| 'name': ['c', 'a', 'b1', 'b2'], |
| 'age': [30, 10, 20, 21], |
| 'city': ['z', 'x', 'y', 'y2'], |
| }, |
| schema=self.pa_schema, |
| )) |
| options = {'sorted-index.records-per-range': '2'} |
| |
| self.assertEqual(2, table.create_global_index('id', options=options)) |
| snapshot = table.snapshot_manager().get_latest_snapshot() |
| self.assertEqual(2, len(IndexFileHandler(table).scan(snapshot))) |
| |
| self.assertEqual(0, table.create_global_index('id', options=options)) |
| |
| latest_snapshot = table.snapshot_manager().get_latest_snapshot() |
| self.assertEqual(2, len(IndexFileHandler(table).scan(latest_snapshot))) |
| |
| self._write_arrow(table, pa.table( |
| { |
| 'id': [5, 4], |
| 'name': ['e', 'd'], |
| 'age': [50, 40], |
| 'city': ['v', 'u'], |
| }, |
| schema=self.pa_schema, |
| )) |
| |
| self.assertEqual(1, table.create_global_index('id', options=options)) |
| latest_snapshot = table.snapshot_manager().get_latest_snapshot() |
| entries = sorted( |
| IndexFileHandler(table).scan(latest_snapshot), |
| key=lambda entry: ( |
| entry.index_file.global_index_meta.row_range_start, |
| entry.index_file.file_name, |
| ), |
| ) |
| self.assertEqual( |
| [(0, 3), (0, 3), (4, 5)], |
| [ |
| ( |
| entry.index_file.global_index_meta.row_range_start, |
| entry.index_file.global_index_meta.row_range_end, |
| ) |
| for entry in entries |
| ], |
| ) |
| self.assertEqual(0, table.create_global_index('id', options=options)) |
| self.assertEqual( |
| 3, |
| len(IndexFileHandler(table).scan( |
| table.snapshot_manager().get_latest_snapshot())), |
| ) |
| |
| def test_create_btree_global_index_for_java_scalar_types(self): |
| schema = pa.schema([ |
| ('flag', pa.bool_()), |
| ('amount', pa.decimal128(10, 2)), |
| ('dt', pa.date32()), |
| ('ts', pa.timestamp('us')), |
| ('payload', pa.string()), |
| ]) |
| table = self._create_table(pa_schema=schema, options=self.table_options) |
| self._write_arrow(table, pa.table( |
| { |
| 'flag': [True, False, True], |
| 'amount': [ |
| Decimal('10.25'), Decimal('20.50'), Decimal('30.75')], |
| 'dt': [ |
| date(2026, 6, 18), |
| date(2026, 6, 19), |
| date(2026, 6, 20), |
| ], |
| 'ts': [ |
| datetime(2026, 6, 18, 10, 0, 0, 123456), |
| datetime(2026, 6, 19, 10, 0, 0, 123456), |
| datetime(2026, 6, 20, 10, 0, 0, 123456), |
| ], |
| 'payload': ['a', 'b', 'c'], |
| }, |
| schema=schema, |
| )) |
| |
| 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( |
| VectorIndexTrainer=_FakeVectorIndexTrainer, |
| 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( |
| VectorIndexTrainer=_FakeVectorIndexTrainer, |
| 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], |
| ) |
| self.assertEqual( |
| ['2', '2', '1'], |
| [ |
| writer.options['expected-vector-count'] |
| 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_skips_existing_ranges(self): |
| schema = pa.schema([ |
| ('id', pa.int32()), |
| ('embedding', pa.list_(pa.float32())), |
| ]) |
| table = self._create_table(pa_schema=schema, options=self.table_options) |
| self._write_arrow(table, pa.table( |
| { |
| 'id': [1, 2, 3], |
| 'embedding': pa.array( |
| [[1.0, 0.0], [0.0, 1.0], [0.5, 0.5]], |
| type=pa.list_(pa.float32()), |
| ), |
| }, |
| schema=schema, |
| )) |
| |
| old_module = sys.modules.get("paimon_vindex") |
| sys.modules["paimon_vindex"] = types.SimpleNamespace( |
| VectorIndexTrainer=_FakeVectorIndexTrainer, |
| VectorIndexWriter=_FakeVectorIndexWriter) |
| _FakeVectorIndexWriter.instances = [] |
| options = { |
| 'global-index.row-count-per-shard': '2', |
| 'ivf-flat.dimension': '2', |
| } |
| try: |
| self.assertEqual( |
| 2, |
| table.create_global_index( |
| 'embedding', |
| index_type='ivf-flat', |
| options=options, |
| ), |
| ) |
| self.assertEqual( |
| 0, |
| table.create_global_index( |
| 'embedding', |
| index_type='ivf-flat', |
| options=options, |
| ), |
| ) |
| |
| self._write_arrow(table, pa.table( |
| { |
| 'id': [4, 5], |
| 'embedding': pa.array( |
| [[0.2, 0.8], [0.9, 0.1]], |
| type=pa.list_(pa.float32()), |
| ), |
| }, |
| schema=schema, |
| )) |
| |
| self.assertEqual( |
| 2, |
| table.create_global_index( |
| 'embedding', |
| index_type='ivf-flat', |
| options=options, |
| ), |
| ) |
| self.assertEqual( |
| 0, |
| table.create_global_index( |
| 'embedding', |
| index_type='ivf-flat', |
| options=options, |
| ), |
| ) |
| finally: |
| if old_module is None: |
| sys.modules.pop("paimon_vindex", None) |
| else: |
| sys.modules["paimon_vindex"] = old_module |
| |
| 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, 3), (4, 4)], |
| [ |
| ( |
| entry.index_file.global_index_meta.row_range_start, |
| entry.index_file.global_index_meta.row_range_end, |
| ) |
| for entry in entries |
| ], |
| ) |
| |
| def test_create_native_fulltext_global_index_from_python(self): |
| schema = pa.schema([ |
| ('id', pa.int32()), |
| ('content', pa.string()), |
| ]) |
| table = self._create_table(pa_schema=schema, options=self.table_options) |
| self._write_arrow(table, pa.table( |
| { |
| 'id': [1, 2, 3], |
| 'content': [ |
| 'Apache Paimon full text', |
| None, |
| 'Native full text search', |
| ], |
| }, |
| schema=schema, |
| )) |
| |
| ftindex = _FakeFullTextForBuild() |
| old_ftindex = sys.modules.get("paimon_ftindex") |
| sys.modules["paimon_ftindex"] = ftindex |
| try: |
| added = table.create_global_index( |
| 'content', |
| index_type=FULL_TEXT_IDENTIFIER, |
| options={ |
| 'global-index.row-count-per-shard': '2', |
| 'full-text.tokenizer': 'ngram', |
| 'full-text.ngram.min-gram': '2', |
| 'full-text.ngram.max-gram': '3', |
| 'full-text.ngram.prefix-only': 'true', |
| 'full-text.with-position': 'false', |
| }, |
| ) |
| finally: |
| if old_ftindex is None: |
| sys.modules.pop("paimon_ftindex", None) |
| else: |
| sys.modules["paimon_ftindex"] = old_ftindex |
| |
| self.assertEqual(2, added) |
| self.assertEqual(2, len(ftindex.writers)) |
| self.assertEqual( |
| { |
| "tokenizer": "ngram", |
| "ngram.min-gram": "2", |
| "ngram.max-gram": "3", |
| "ngram.prefix-only": "true", |
| "with-position": "false", |
| }, |
| ftindex.writers[0].options, |
| ) |
| self.assertEqual( |
| [{'row_id': 0, 'text': 'Apache Paimon full text'}], |
| ftindex.writers[0].documents, |
| ) |
| self.assertEqual( |
| [{'row_id': 0, 'text': 'Native full text search'}], |
| ftindex.writers[1].documents, |
| ) |
| self.assertTrue(all(writer.written for writer in ftindex.writers)) |
| self.assertTrue(all(writer.closed for writer in ftindex.writers)) |
| |
| 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, 2, 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 |
| ], |
| ) |
| self.assertEqual( |
| {FULL_TEXT_IDENTIFIER}, |
| {entry.index_file.index_type for entry in entries}, |
| ) |
| for entry in entries: |
| index_options = NativeFullTextIndexOptions.deserialize( |
| entry.index_file.global_index_meta.index_meta) |
| self.assertEqual( |
| { |
| "tokenizer": "ngram", |
| "ngram.min-gram": "2", |
| "ngram.max-gram": "3", |
| "ngram.prefix-only": "true", |
| "with-position": "false", |
| }, |
| index_options.to_native_options(), |
| ) |
| file_path = table.path_factory().global_index_path_factory().to_path( |
| entry.index_file.file_name) |
| self.assertEqual( |
| b"fake-ftindex", |
| table.file_io.new_input_stream(file_path).read(), |
| ) |
| |
| def test_native_fulltext_writer_accepts_jieba_tokenizer(self): |
| table = self._create_table() |
| index_path = ( |
| table.path_factory() |
| .global_index_path_factory() |
| .global_index_root_path() |
| ) |
| ftindex = _FakeFullTextForBuild() |
| old_ftindex = sys.modules.get("paimon_ftindex") |
| sys.modules["paimon_ftindex"] = ftindex |
| try: |
| writer = NativeFullTextIndexWriter( |
| table.file_io, |
| index_path, |
| AtomicType('STRING'), |
| {'full-text.tokenizer': 'jieba'}, |
| ) |
| writer.write('北京大学支持全文检索', 0) |
| entries = writer.finish() |
| finally: |
| if old_ftindex is None: |
| sys.modules.pop("paimon_ftindex", None) |
| else: |
| sys.modules["paimon_ftindex"] = old_ftindex |
| |
| self.assertEqual(1, len(entries)) |
| self.assertEqual([{"row_id": 0, "text": "北京大学支持全文检索"}], |
| ftindex.writers[0].documents) |
| self.assertEqual({"tokenizer": "jieba"}, ftindex.writers[0].options) |
| self.assertTrue(ftindex.writers[0].closed) |
| |
| def test_native_fulltext_options_are_passed_through(self): |
| ngram = NativeFullTextIndexOptions.from_options({ |
| 'full-text.tokenizer': 'ngram', |
| 'full-text.ngram.min-gram': '2', |
| 'full-text.ngram.max-gram': '3', |
| 'full-text.ngram.prefix-only': 'true', |
| 'full-text.lower-case': 'false', |
| 'full-text.custom-future-option': 'future-value', |
| 'unrelated': 'ignored', |
| }) |
| |
| self.assertEqual( |
| { |
| "tokenizer": "ngram", |
| "ngram.min-gram": "2", |
| "ngram.max-gram": "3", |
| "ngram.prefix-only": "true", |
| "lower-case": "false", |
| "custom-future-option": "future-value", |
| }, |
| ngram.to_native_options(), |
| ) |
| |
| def test_create_native_fulltext_global_index_rejects_non_string_column(self): |
| table = self._create_table() |
| self._write_arrow(table, pa.table( |
| { |
| 'id': [1], |
| 'name': ['a'], |
| 'age': [10], |
| 'city': ['x'], |
| }, |
| schema=self.pa_schema, |
| )) |
| |
| with self.assertRaisesRegex(ValueError, 'requires string type'): |
| table.create_global_index( |
| 'id', |
| index_type=FULL_TEXT_IDENTIFIER, |
| ) |
| |
| 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_vindex_native_options_support_030_indexes(self): |
| data_type = ArrayType(True, AtomicType('FLOAT')) |
| options = { |
| 'ivf-rq.rq.bits': '5', |
| 'ivf-rq.max-bytes-per-vector': '96', |
| 'diskann.build-preset': 'balanced', |
| 'diskann.pq.code-ratio': '0.0625', |
| 'diskann.raw-vector-encoding': 'f16', |
| } |
| |
| rq_result = native_options( |
| data_type, options, 'ivf-rq', 'embedding') |
| self.assertEqual('ivf_rq', rq_result['index.type']) |
| self.assertEqual('5', rq_result['rq.bits']) |
| self.assertEqual('96', rq_result['max-bytes-per-vector']) |
| self.assertEqual('inner_product', rq_result['metric']) |
| |
| diskann_result = native_options( |
| data_type, options, 'diskann', 'embedding') |
| self.assertEqual('diskann', diskann_result['index.type']) |
| self.assertEqual( |
| 'balanced', diskann_result['diskann.build-preset']) |
| self.assertEqual('0.0625', diskann_result['pq.code-ratio']) |
| self.assertEqual( |
| 'f16', diskann_result['diskann.raw-vector-encoding']) |
| |
| def test_vindex_training_sample_ratio(self): |
| options = { |
| 'ivf-rq.train.sample-ratio': '0.5', |
| 'fields.embedding.train.sample-ratio': '0.25', |
| } |
| self.assertEqual( |
| 0.25, train_sample_ratio(options, 'ivf-rq', 'embedding')) |
| self.assertEqual( |
| 0.5, train_sample_ratio(options, 'ivf-rq', 'other')) |
| |
| import numpy as np |
| vectors = np.arange(20, dtype=np.float32).reshape(10, 2) |
| sampled = _sample_training_vectors(np, vectors, 0.4) |
| self.assertEqual( |
| [[0.0, 1.0], [4.0, 5.0], [10.0, 11.0], [14.0, 15.0]], |
| sampled.tolist(), |
| ) |
| |
| def test_split_by_contiguous_row_range_matches_java_builder(self): |
| split = _FakeSplit([ |
| _FakeFile('a', 0, 2), |
| _FakeFile('b', 4, 2), |
| _FakeFile('c', 6, 1), |
| _FakeFile('d', 10, 1), |
| ]) |
| |
| splits = _split_one_by_contiguous_row_range(split) |
| |
| self.assertEqual( |
| [['a'], ['b', 'c'], ['d']], |
| [[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_uses_unindexed_ranges(self): |
| split = _FakeSplit([ |
| _FakeFile('a', 0, 100), |
| ]) |
| |
| shards = _split_by_global_index_shard( |
| [split], 100, [Range(0, 9), Range(90, 99)]) |
| |
| self.assertEqual( |
| [ |
| (['a'], 0, 9), |
| (['a'], 90, 99), |
| ], |
| [ |
| ([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), |
| ('TINYINT', -7), |
| ('SMALLINT', 1024), |
| ('INT', 42), |
| ('BIGINT', 1234567890123), |
| ('FLOAT', 1.25), |
| ('DOUBLE', 3.14159), |
| ('DECIMAL(20, 5)', Decimal('-1234567890123.45678')), |
| ('DATE', date(2026, 6, 18)), |
| ('TIME(3)', datetime(2026, 6, 18, 1, 2, 3, 456000).time()), |
| ('TIMESTAMP(6)', datetime(2026, 6, 18, 1, 2, 3, 456789)), |
| ('VARCHAR(16)', 'abc'), |
| ] |
| |
| for type_name, value in cases: |
| with self.subTest(type_name=type_name): |
| serializer = create_serializer(AtomicType(type_name)) |
| actual = serializer.deserialize(serializer.serialize(value)) |
| if type_name == 'FLOAT': |
| self.assertAlmostEqual(value, actual, places=6) |
| elif type_name == 'DOUBLE': |
| self.assertAlmostEqual(value, actual, places=12) |
| else: |
| self.assertEqual(value, actual) |
| |
| |
| if __name__ == "__main__": |
| unittest.main() |
