| # 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 os | |
| import shutil | |
| import tempfile | |
| import unittest | |
| import pyarrow as pa | |
| from pypaimon import CatalogFactory, Schema | |
| from pypaimon.manifest.schema.data_file_meta import DataFileMeta | |
| from pypaimon.table.row.offset_row import OffsetRow | |
| from pypaimon.table.row.vector import Vector | |
| class VectorClassTest(unittest.TestCase): | |
| def test_basic_operations(self): | |
| v = Vector([1.0, 2.0, 3.0]) | |
| self.assertEqual(len(v), 3) | |
| self.assertEqual(v[0], 1.0) | |
| self.assertEqual(v[1], 2.0) | |
| self.assertEqual(v[2], 3.0) | |
| def test_to_list(self): | |
| v = Vector([1.0, 2.0, 3.0]) | |
| self.assertEqual(v.to_list(), [1.0, 2.0, 3.0]) | |
| def test_from_list(self): | |
| v = Vector.from_list([4.0, 5.0, 6.0]) | |
| self.assertEqual(v.to_list(), [4.0, 5.0, 6.0]) | |
| def test_equality(self): | |
| v1 = Vector([1.0, 2.0, 3.0]) | |
| v2 = Vector([1.0, 2.0, 3.0]) | |
| v3 = Vector([4.0, 5.0, 6.0]) | |
| self.assertEqual(v1, v2) | |
| self.assertNotEqual(v1, v3) | |
| def test_hash(self): | |
| v1 = Vector([1.0, 2.0, 3.0]) | |
| v2 = Vector([1.0, 2.0, 3.0]) | |
| self.assertEqual(hash(v1), hash(v2)) | |
| def test_str_repr(self): | |
| v = Vector([1.0, 2.0]) | |
| self.assertEqual(str(v), "Vector([1.0, 2.0])") | |
| self.assertEqual(repr(v), "Vector([1.0, 2.0])") | |
| def test_integer_vector(self): | |
| v = Vector([1, 2, 3]) | |
| self.assertEqual(v.to_list(), [1.0, 2.0, 3.0]) | |
| def test_empty_vector(self): | |
| v = Vector([]) | |
| self.assertEqual(len(v), 0) | |
| self.assertEqual(v.to_list(), []) | |
| def test_offset_row_legacy_positional_vector_indices(self): | |
| row = OffsetRow(([1.0, 2.0],), 0, 1, None, None, {0}) | |
| self.assertEqual(row.get_vector(0).to_list(), [1.0, 2.0]) | |
| self.assertEqual(row._descriptor_field_indices, frozenset()) | |
| class VectorFileDetectionTest(unittest.TestCase): | |
| def test_is_vector_file(self): | |
| self.assertTrue(DataFileMeta.is_vector_file("data-uuid-0.vector.lance")) | |
| self.assertTrue(DataFileMeta.is_vector_file("data-uuid-0.vector.parquet")) | |
| self.assertFalse(DataFileMeta.is_vector_file("data-uuid-0.parquet")) | |
| self.assertFalse(DataFileMeta.is_vector_file("data-uuid-0.lance")) | |
| self.assertFalse(DataFileMeta.is_vector_file("data-uuid-0.blob")) | |
| class VectorOnlyTableTest(unittest.TestCase): | |
| """Vector-only tables (no normal columns) must be rejected at schema creation.""" | |
| def test_vector_only_table_rejected(self): | |
| pa_schema = pa.schema([ | |
| ('embed1', pa.list_(pa.float32(), 3)), | |
| ('embed2', pa.list_(pa.float32(), 2)), | |
| ]) | |
| with self.assertRaises(ValueError) as ctx: | |
| Schema.from_pyarrow_schema( | |
| pa_schema, | |
| options={ | |
| 'vector.file.format': 'vortex', | |
| 'row-tracking.enabled': 'true', | |
| 'data-evolution.enabled': 'true', | |
| 'bucket': '-1', | |
| } | |
| ) | |
| self.assertIn("must have other normal columns", str(ctx.exception)) | |
| def test_vector_dedicated_missing_row_tracking_rejected(self): | |
| pa_schema = pa.schema([ | |
| ('id', pa.int32()), | |
| ('embed', pa.list_(pa.float32(), 3)), | |
| ]) | |
| with self.assertRaises(ValueError) as ctx: | |
| Schema.from_pyarrow_schema( | |
| pa_schema, | |
| options={ | |
| 'vector.file.format': 'vortex', | |
| 'data-evolution.enabled': 'true', | |
| 'bucket': '-1', | |
| } | |
| ) | |
| self.assertIn("row-tracking.enabled", str(ctx.exception)) | |
| class VectorTableWriteReadTest(unittest.TestCase): | |
| @classmethod | |
| def setUpClass(cls): | |
| cls.temp_dir = tempfile.mkdtemp() | |
| cls.warehouse = os.path.join(cls.temp_dir, 'warehouse') | |
| cls.catalog = CatalogFactory.create({'warehouse': cls.warehouse}) | |
| cls.catalog.create_database('test_db', False) | |
| @classmethod | |
| def tearDownClass(cls): | |
| try: | |
| shutil.rmtree(cls.temp_dir) | |
| except OSError: | |
| pass | |
| def test_inline_vector_write_read(self): | |
| """Write and read vector data stored inline (no vector.file.format).""" | |
| pa_schema = pa.schema([ | |
| ('id', pa.int64()), | |
| ('embed', pa.list_(pa.float32(), 3)), | |
| ]) | |
| schema = Schema.from_pyarrow_schema(pa_schema) | |
| self.catalog.create_table('test_db.inline_vector', schema, False) | |
| table = self.catalog.get_table('test_db.inline_vector') | |
| test_data = pa.table({ | |
| 'id': pa.array([1, 2, 3], type=pa.int64()), | |
| 'embed': pa.FixedSizeListArray.from_arrays( | |
| pa.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0], type=pa.float32()), | |
| 3 | |
| ), | |
| }) | |
| write_builder = table.new_batch_write_builder() | |
| writer = write_builder.new_write() | |
| writer.write_arrow(test_data) | |
| commit_messages = writer.prepare_commit() | |
| write_builder.new_commit().commit(commit_messages) | |
| writer.close() | |
| # Read back | |
| read_builder = table.new_read_builder() | |
| splits = read_builder.new_scan().plan().splits() | |
| result = read_builder.new_read().to_arrow(splits) | |
| self.assertEqual(result.num_rows, 3) | |
| self.assertEqual(result.column('id').to_pylist(), [1, 2, 3]) | |
| embed_col = result.column('embed') | |
| self.assertTrue(pa.types.is_fixed_size_list(embed_col.type)) | |
| self.assertEqual(embed_col.type.list_size, 3) | |
| self.assertEqual(embed_col[0].as_py(), [1.0, 2.0, 3.0]) | |
| self.assertEqual(embed_col[1].as_py(), [4.0, 5.0, 6.0]) | |
| self.assertEqual(embed_col[2].as_py(), [7.0, 8.0, 9.0]) | |
| def test_get_vector_row_access(self): | |
| """Test get_vector() returns Vector objects from InternalRow.""" | |
| pa_schema = pa.schema([ | |
| ('id', pa.int64()), | |
| ('embed', pa.list_(pa.float32(), 3)), | |
| ]) | |
| schema = Schema.from_pyarrow_schema(pa_schema) | |
| self.catalog.create_table('test_db.row_vector', schema, False) | |
| table = self.catalog.get_table('test_db.row_vector') | |
| test_data = pa.table({ | |
| 'id': pa.array([1, 2], type=pa.int64()), | |
| 'embed': pa.FixedSizeListArray.from_arrays( | |
| pa.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], type=pa.float32()), | |
| 3 | |
| ), | |
| }) | |
| write_builder = table.new_batch_write_builder() | |
| writer = write_builder.new_write() | |
| writer.write_arrow(test_data) | |
| commit_messages = writer.prepare_commit() | |
| write_builder.new_commit().commit(commit_messages) | |
| writer.close() | |
| # Read rows and use get_vector() — collect vectors eagerly since OffsetRow is reused | |
| read_builder = table.new_read_builder() | |
| splits = read_builder.new_scan().plan().splits() | |
| vectors = set() | |
| count = 0 | |
| for row in read_builder.new_read().to_iterator(splits): | |
| vec = row.get_vector(1) | |
| self.assertIsInstance(vec, Vector) | |
| self.assertEqual(len(vec), 3) | |
| vectors.add(tuple(vec.to_list())) | |
| count += 1 | |
| self.assertEqual(count, 2) | |
| self.assertIn((1.0, 2.0, 3.0), vectors) | |
| self.assertIn((4.0, 5.0, 6.0), vectors) | |
| @unittest.skipUnless( | |
| __import__('importlib').util.find_spec('vortex') is not None, | |
| "vortex not installed" | |
| ) | |
| def test_vector_dedicated_format_write_read_vortex(self): | |
| """Write vector data to separate .vector.vortex files.""" | |
| pa_schema = pa.schema([ | |
| ('id', pa.int64()), | |
| ('embed', pa.list_(pa.float32(), 3)), | |
| ]) | |
| schema = Schema.from_pyarrow_schema( | |
| pa_schema, | |
| options={ | |
| 'file.format': 'vortex', | |
| 'row-tracking.enabled': 'true', | |
| 'data-evolution.enabled': 'true', | |
| 'vector.file.format': 'vortex', | |
| } | |
| ) | |
| self.catalog.create_table('test_db.dedicated_vector_vortex', schema, False) | |
| table = self.catalog.get_table('test_db.dedicated_vector_vortex') | |
| test_data = pa.table({ | |
| 'id': pa.array([1, 2, 3], type=pa.int64()), | |
| 'embed': pa.FixedSizeListArray.from_arrays( | |
| pa.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0], type=pa.float32()), | |
| 3 | |
| ), | |
| }) | |
| write_builder = table.new_batch_write_builder() | |
| writer = write_builder.new_write() | |
| writer.write_arrow(test_data) | |
| commit_messages = writer.prepare_commit() | |
| # Verify file names: should have both normal files and .vector.vortex | |
| all_files = [] | |
| for msg in commit_messages: | |
| all_files.extend(msg.new_files) | |
| normal_files = [f for f in all_files if not DataFileMeta.is_vector_file(f.file_name)] | |
| vector_files = [f for f in all_files if DataFileMeta.is_vector_file(f.file_name)] | |
| self.assertGreater(len(normal_files), 0, "Should have normal data files") | |
| self.assertGreater(len(vector_files), 0, "Should have vector files") | |
| for vf in vector_files: | |
| self.assertIn('.vector.vortex', vf.file_name) | |
| write_builder.new_commit().commit(commit_messages) | |
| writer.close() | |
| # Read back | |
| read_builder = table.new_read_builder() | |
| splits = read_builder.new_scan().plan().splits() | |
| result = read_builder.new_read().to_arrow(splits) | |
| self.assertEqual(result.num_rows, 3) | |
| ids = sorted(result.column('id').to_pylist()) | |
| self.assertEqual(ids, [1, 2, 3]) | |
| embed_by_id = {} | |
| for i in range(result.num_rows): | |
| embed_by_id[result.column('id')[i].as_py()] = result.column('embed')[i].as_py() | |
| self.assertEqual(embed_by_id[1], [1.0, 2.0, 3.0]) | |
| self.assertEqual(embed_by_id[2], [4.0, 5.0, 6.0]) | |
| self.assertEqual(embed_by_id[3], [7.0, 8.0, 9.0]) | |
| @unittest.skipUnless( | |
| __import__('importlib').util.find_spec('lance') is not None, | |
| "lance not installed" | |
| ) | |
| def test_vector_dedicated_format_write_read_lance(self): | |
| """Write vector data to separate .vector.lance files.""" | |
| pa_schema = pa.schema([ | |
| ('id', pa.int64()), | |
| ('embed', pa.list_(pa.float32(), 3)), | |
| ]) | |
| schema = Schema.from_pyarrow_schema( | |
| pa_schema, | |
| options={ | |
| 'row-tracking.enabled': 'true', | |
| 'data-evolution.enabled': 'true', | |
| 'vector.file.format': 'lance', | |
| } | |
| ) | |
| self.catalog.create_table('test_db.dedicated_vector_lance', schema, False) | |
| table = self.catalog.get_table('test_db.dedicated_vector_lance') | |
| test_data = pa.table({ | |
| 'id': pa.array([1, 2, 3], type=pa.int64()), | |
| 'embed': pa.FixedSizeListArray.from_arrays( | |
| pa.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0], type=pa.float32()), | |
| 3 | |
| ), | |
| }) | |
| write_builder = table.new_batch_write_builder() | |
| writer = write_builder.new_write() | |
| writer.write_arrow(test_data) | |
| commit_messages = writer.prepare_commit() | |
| all_files = [] | |
| for msg in commit_messages: | |
| all_files.extend(msg.new_files) | |
| vector_files = [f for f in all_files if DataFileMeta.is_vector_file(f.file_name)] | |
| self.assertGreater(len(vector_files), 0, "Should have vector files") | |
| for vf in vector_files: | |
| self.assertIn('.vector.lance', vf.file_name) | |
| write_builder.new_commit().commit(commit_messages) | |
| writer.close() | |
| # Read back | |
| read_builder = table.new_read_builder() | |
| splits = read_builder.new_scan().plan().splits() | |
| result = read_builder.new_read().to_arrow(splits) | |
| self.assertEqual(result.num_rows, 3) | |
| def test_vector_table_partial_update_non_vector_column(self): | |
| vector_schema = pa.schema([ | |
| ('id', pa.int32()), | |
| ('name', pa.string()), | |
| ('embedding', pa.list_(pa.float32(), 4)), | |
| ]) | |
| opts = { | |
| 'row-tracking.enabled': 'true', | |
| 'data-evolution.enabled': 'true', | |
| 'data-evolution.write-cols-optimization.enabled': 'true', | |
| 'vector.file.format': 'parquet', | |
| } | |
| s = Schema.from_pyarrow_schema(vector_schema, options=opts) | |
| table_name = 'test_db.vector_de_seq' | |
| self.catalog.create_table(table_name, s, False) | |
| table = self.catalog.get_table(table_name) | |
| wb = table.new_batch_write_builder() | |
| w = wb.new_write() | |
| w.write_arrow(pa.Table.from_pydict( | |
| { | |
| 'id': [1, 2], | |
| 'name': ['a', 'b'], | |
| 'embedding': [[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]], | |
| }, | |
| schema=vector_schema, | |
| )) | |
| initial_messages = w.prepare_commit() | |
| initial_files = [ | |
| file for message in initial_messages for file in message.new_files | |
| ] | |
| normal_file = next( | |
| file for file in initial_files | |
| if not DataFileMeta.is_vector_file(file.file_name) | |
| ) | |
| vector_file = next( | |
| file for file in initial_files | |
| if DataFileMeta.is_vector_file(file.file_name) | |
| ) | |
| self.assertIsNone(normal_file.write_cols) | |
| self.assertEqual(['embedding'], vector_file.write_cols) | |
| wb.new_commit().commit(initial_messages) | |
| w.close() | |
| table = self.catalog.get_table(table_name) | |
| rb = table.new_read_builder() | |
| rb = rb.with_projection(['name', '_ROW_ID']) | |
| splits = rb.new_scan().plan().splits() | |
| source = rb.new_read().to_arrow(splits) | |
| update_data = pa.table({ | |
| '_ROW_ID': source.column('_ROW_ID'), | |
| 'name': pa.array(['updated', 'updated'], type=pa.string()), | |
| }) | |
| update_builder = table.new_batch_write_builder() | |
| updater = update_builder.new_update().new_update_by_row_id() | |
| msgs = updater.update_columns(update_data, ['name']) | |
| update_builder.new_commit().commit(msgs) | |
| table = self.catalog.get_table(table_name) | |
| rb = table.new_read_builder() | |
| splits = rb.new_scan().plan().splits() | |
| result = rb.new_read().to_arrow(splits).sort_by('id').to_pydict() | |
| self.assertEqual(result['name'], ['updated', 'updated']) | |
| def test_backfill_vector_column_added_after_existing_rows(self): | |
| from pypaimon.schema.data_types import AtomicType, VectorType | |
| from pypaimon.schema.schema_change import SchemaChange | |
| from pypaimon.utils.range import Range | |
| table_name = 'test_db.vector_backfill_added_column' | |
| base_schema = pa.schema([ | |
| ('id', pa.int64()), | |
| ('image', pa.large_binary()), | |
| ]) | |
| opts = { | |
| 'row-tracking.enabled': 'true', | |
| 'data-evolution.enabled': 'true', | |
| 'vector.file.format': 'parquet', | |
| } | |
| self.catalog.create_table( | |
| table_name, | |
| Schema.from_pyarrow_schema(base_schema, options=opts), | |
| False, | |
| ) | |
| table = self.catalog.get_table(table_name) | |
| wb = table.new_batch_write_builder() | |
| writer = wb.new_write() | |
| writer.write_arrow(pa.table({ | |
| 'id': pa.array([1, 2, 3], type=pa.int64()), | |
| 'image': pa.array([b'a', b'b', b'c'], type=pa.large_binary()), | |
| })) | |
| commit = wb.new_commit() | |
| commit.commit(writer.prepare_commit()) | |
| writer.close() | |
| commit.close() | |
| self.catalog.alter_table( | |
| table_name, | |
| [SchemaChange.add_column( | |
| 'embedding', VectorType(True, AtomicType('FLOAT'), 2))], | |
| False, | |
| ) | |
| table = self.catalog.get_table(table_name) | |
| read_builder = table.new_read_builder().with_projection( | |
| ['image', 'embedding', '_ROW_ID']) | |
| read_builder.with_filter( | |
| read_builder.new_predicate_builder().is_null('embedding')) | |
| splits = ( | |
| read_builder.new_scan() | |
| .with_row_ranges([Range(0, 2)]) | |
| .plan() | |
| .splits() | |
| ) | |
| pending = read_builder.new_read().to_arrow(splits).sort_by('_ROW_ID') | |
| self.assertEqual([0, 1, 2], pending['_ROW_ID'].to_pylist()) | |
| self.assertEqual([b'a', b'b', b'c'], pending['image'].to_pylist()) | |
| updates = pa.table({ | |
| '_ROW_ID': pending['_ROW_ID'], | |
| 'embedding': pa.array( | |
| [[1.0, 0.0], [0.5, 0.5], [0.0, 1.0]], | |
| type=pa.list_(pa.float32(), 2), | |
| ), | |
| }) | |
| wb = table.new_batch_write_builder() | |
| messages = ( | |
| wb.new_update() | |
| .with_update_type(['embedding']) | |
| .update_by_arrow_with_row_id(updates) | |
| ) | |
| commit = wb.new_commit() | |
| commit.commit(messages) | |
| commit.close() | |
| table = self.catalog.get_table(table_name) | |
| read_builder = table.new_read_builder().with_projection( | |
| ['id', 'embedding']) | |
| result = read_builder.new_read().to_arrow( | |
| read_builder.new_scan().plan().splits()).sort_by('id') | |
| self.assertEqual( | |
| [[1.0, 0.0], [0.5, 0.5], [0.0, 1.0]], | |
| result['embedding'].to_pylist(), | |
| ) | |
| retry_builder = table.new_read_builder().with_projection( | |
| ['embedding', '_ROW_ID']) | |
| retry_builder.with_filter( | |
| retry_builder.new_predicate_builder().is_null('embedding')) | |
| retry_splits = ( | |
| retry_builder.new_scan() | |
| .with_row_ranges([Range(0, 2)]) | |
| .plan() | |
| .splits() | |
| ) | |
| self.assertEqual( | |
| 0, | |
| retry_builder.new_read().to_arrow(retry_splits).num_rows, | |
| ) | |
| def test_vector_table_partial_update_non_vector_column_with_rolling_files(self): | |
| vector_schema = pa.schema([ | |
| ('id', pa.int32()), | |
| ('name', pa.string()), | |
| ('embedding', pa.list_(pa.float32(), 4)), | |
| ]) | |
| opts = { | |
| 'row-tracking.enabled': 'true', | |
| 'data-evolution.enabled': 'true', | |
| 'vector.file.format': 'parquet', | |
| 'target-file-size': '1KB', | |
| } | |
| s = Schema.from_pyarrow_schema(vector_schema, options=opts) | |
| table_name = 'test_db.vector_de_seq_rolling' | |
| self.catalog.create_table(table_name, s, False) | |
| write_schema = pa.schema([ | |
| ('id', pa.int32()), | |
| ('name', pa.string()), | |
| ]) | |
| table = self.catalog.get_table(table_name) | |
| wb = table.new_batch_write_builder() | |
| w = wb.new_write().with_write_type(['id', 'name']) | |
| for start in (0, 1000): | |
| ids = list(range(start, start + 1000)) | |
| w.write_arrow(pa.Table.from_pydict( | |
| { | |
| 'id': ids, | |
| 'name': [f'name_{i}_' + 'x' * 2048 for i in ids], | |
| }, | |
| schema=write_schema, | |
| )) | |
| commit_messages = w.prepare_commit() | |
| normal_files = [ | |
| f for msg in commit_messages for f in msg.new_files | |
| if not DataFileMeta.is_vector_file(f.file_name) | |
| ] | |
| self.assertGreaterEqual(len(normal_files), 2) | |
| for file in normal_files: | |
| self.assertEqual(file.min_sequence_number, 0) | |
| self.assertEqual(file.max_sequence_number, file.row_count - 1) | |
| wb.new_commit().commit(commit_messages) | |
| w.close() | |
| table = self.catalog.get_table(table_name) | |
| rb = table.new_read_builder().with_projection(['id', 'name', '_ROW_ID']) | |
| splits = rb.new_scan().plan().splits() | |
| source = rb.new_read().to_arrow(splits).sort_by('id') | |
| update_data = pa.table({ | |
| '_ROW_ID': source.column('_ROW_ID'), | |
| 'name': pa.array(['updated'] * source.num_rows, type=pa.string()), | |
| }) | |
| update_builder = table.new_batch_write_builder() | |
| updater = update_builder.new_update().new_update_by_row_id() | |
| msgs = updater.update_columns(update_data, ['name']) | |
| update_normal_files = [ | |
| f for msg in msgs for f in msg.new_files | |
| if not DataFileMeta.is_vector_file(f.file_name) | |
| ] | |
| self.assertGreaterEqual(len(update_normal_files), 2) | |
| for file in update_normal_files: | |
| self.assertEqual(file.min_sequence_number, 0) | |
| self.assertEqual(file.max_sequence_number, file.row_count - 1) | |
| update_builder.new_commit().commit(msgs) | |
| table = self.catalog.get_table(table_name) | |
| rb = table.new_read_builder().with_projection(['id', 'name']) | |
| splits = rb.new_scan().plan().splits() | |
| result = rb.new_read().to_arrow(splits).sort_by('id').to_pydict() | |
| self.assertEqual(result['id'], list(range(2000))) | |
| self.assertEqual(result['name'], ['updated'] * 2000) | |
| if __name__ == '__main__': | |
| unittest.main() |