diff --git a/paimon-python/pypaimon/globalindex/create_global_index.py b/paimon-python/pypaimon/globalindex/create_global_index.py index b2e125e42159..d68221a0ae6c 100644 --- a/paimon-python/pypaimon/globalindex/create_global_index.py +++ b/paimon-python/pypaimon/globalindex/create_global_index.py @@ -288,6 +288,8 @@ def _create_sorted_index_writer(self, index_path: str, key_serializer): def _build_generic_index( self, splits, unindexed_ranges, index_field, table_read, index_path: str ) -> List[CommitMessage]: + from pypaimon.read.table_read import _ClosableArrowBatchReader + rows_per_shard = self._core_options.global_index_row_count_per_shard() if rows_per_shard <= 0: raise ValueError( @@ -298,26 +300,38 @@ def _build_generic_index( for index_split, index_range in _split_by_global_index_shard( splits, rows_per_shard, unindexed_ranges ): - table = table_read.to_arrow([index_split]) - if table is None or table.num_rows == 0: - continue - - writer = self._create_generic_index_writer(index_path, index_field) + writer = None try: - if self._index_type in VINDEX_IDENTIFIERS: - if table.column(SpecialFields.ROW_ID.name).null_count: - raise ValueError("Cannot build global index because _ROW_ID is null.") - for batch in table.to_batches(max_chunksize=ADD_BATCH_SIZE): - _write_vector_batch( - writer, batch, self._index_columns[0], index_range) - else: - for value, row_id in _extract_index_rows( - table, - self._index_columns[0], - SpecialFields.ROW_ID.name, - index_range, - ): - writer.write(value, row_id - index_range.from_) + reader, batches = table_read._new_arrow_batch_reader([index_split]) + # Close the Python iterator explicitly on failure as well as + # the Arrow reader, which may retain a suspended generator. + with _ClosableArrowBatchReader(reader, batches) as batch_reader: + for batch in batch_reader: + if batch.num_rows == 0: + continue + if writer is None: + writer = self._create_generic_index_writer( + index_path, index_field) + if self._index_type in VINDEX_IDENTIFIERS: + if batch.column(SpecialFields.ROW_ID.name).null_count: + raise ValueError( + "Cannot build global index because _ROW_ID is null.") + for offset in range(0, batch.num_rows, ADD_BATCH_SIZE): + _write_vector_batch( + writer, batch.slice(offset, ADD_BATCH_SIZE), + self._index_columns[0], index_range) + else: + for value, row_id in _extract_index_rows( + batch, + self._index_columns[0], + SpecialFields.ROW_ID.name, + index_range, + ): + writer.write(value, row_id - index_range.from_) + del batch + + if writer is None: + continue index_adds = _to_index_manifest_entries( self._table, @@ -328,7 +342,8 @@ def _build_generic_index( writer.finish(), ) finally: - writer.close() + if writer is not None: + writer.close() if index_adds: messages.append( CommitMessage( @@ -469,7 +484,7 @@ def _write_vector_batch(writer, batch, index_column, row_range): def _extract_index_rows( - table: pa.Table, + table: Union[pa.Table, pa.RecordBatch], index_column: str, row_id_column: str, row_range: Optional[Range] = None, diff --git a/paimon-python/pypaimon/tests/global_index_build_test.py b/paimon-python/pypaimon/tests/global_index_build_test.py index 186270fcb79a..1633b5270f58 100644 --- a/paimon-python/pypaimon/tests/global_index_build_test.py +++ b/paimon-python/pypaimon/tests/global_index_build_test.py @@ -22,6 +22,7 @@ import struct import sys import types +from unittest.mock import Mock, patch import pyarrow as pa @@ -556,7 +557,8 @@ def test_create_vindex_global_index_from_python(self): ('id', pa.int32()), ('embedding', pa.list_(pa.float32())), ]) - table = self._create_table(pa_schema=schema, options=self.table_options) + table = self._create_table(pa_schema=schema, options=dict( + self.table_options, **{'read.batch-size': '1'})) vectors = pa.array( [[1.0, 0.0], [0.0, 1.0], None], type=pa.list_(pa.float32()), @@ -610,12 +612,56 @@ def test_create_vindex_global_index_from_python(self): table.path_factory().global_index_path_factory().to_path( entry.index_file.file_name))) + def test_create_vindex_streaming_failure_cleans_resources(self): + from pypaimon.read.table_read import TableRead + + schema = pa.schema([('embedding', pa.list_(pa.float32()))]) + table = self._create_table(pa_schema=schema, options=dict( + self.table_options, **{'read.batch-size': '1'})) + self._write_arrow(table, pa.table( + {'embedding': [[1.0, 0.0], [0.0, 1.0], [0.5, 0.5]]}, schema=schema)) + snapshot_id = table.snapshot_manager().get_latest_snapshot().id + original_write = VindexVectorIndexWriter.write_batch + original_batches = TableRead._arrow_batch_generator + temp_paths = [] + closed = [] + generators = [] + + def failing_write(writer, vector, row_id): + original_write(writer, vector, row_id) + if 1 in row_id.to_pylist(): + temp_paths.extend([writer._row_id_temp_path, writer._vector_temp_path]) + raise RuntimeError('injected write failure') + + def tracked_batches(reader, *args): + def generate(): + try: + yield from original_batches(reader, *args) + finally: + closed.append(True) + generator = generate() + generators.append(generator) + return generator + + with patch.object(VindexVectorIndexWriter, 'write_batch', failing_write), \ + patch.object(TableRead, '_arrow_batch_generator', tracked_batches): + with self.assertRaisesRegex(RuntimeError, 'injected write failure'): + table.create_global_index('embedding', index_type='ivf-flat', options={ + 'ivf-flat.dimension': '2', + }) + + self.assertEqual([True], closed) + self.assertEqual(2, len(temp_paths)) + self.assertTrue(all(not os.path.exists(path) for path in temp_paths)) + self.assertEqual(snapshot_id, table.snapshot_manager().get_latest_snapshot().id) + 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) + table = self._create_table(pa_schema=schema, options=dict( + self.table_options, **{'read.batch-size': '1'})) 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()), @@ -774,7 +820,8 @@ def test_create_native_fulltext_global_index_from_python(self): ('id', pa.int32()), ('content', pa.string()), ]) - table = self._create_table(pa_schema=schema, options=self.table_options) + table = self._create_table(pa_schema=schema, options=dict( + self.table_options, **{'read.batch-size': '1'})) self._write_arrow(table, pa.table( { 'id': [1, 2, 3], @@ -1181,5 +1228,141 @@ def test_java_scalar_key_serializers_round_trip(self): self.assertEqual(value, actual) +class GenericIndexStreamingTest(unittest.TestCase): + + schema = pa.schema([ + ('embedding', pa.list_(pa.float32())), + ('_ROW_ID', pa.int64()), + ]) + + def setUp(self): + self.builder = object.__new__(GlobalIndexBuilder) + self.builder._table = Mock() + self.builder._core_options = Mock() + self.builder._core_options.global_index_row_count_per_shard.return_value = 10 + self.builder._index_columns = ['embedding'] + self.builder._index_type = 'ivf-flat' + self.writer = Mock() + self.writer.finish.return_value = [] + self.builder._create_generic_index_writer = Mock(return_value=self.writer) + self.read = Mock() + self.read.to_arrow.side_effect = AssertionError('Must not materialize a shard') + self.events = [] + + def _batch(self, values, row_ids): + return pa.RecordBatch.from_arrays([ + pa.array(values, type=self.schema.field(0).type), + pa.array(row_ids, type=pa.int64()), + ], schema=self.schema) + + def _build(self, batches): + def generate(): + try: + for batch in batches: + self.events.append('read') + yield batch + finally: + self.events.append('reader closed') + + # Retain the generator: cleanup must be explicit, not depend on GC. + self.generator = generate() + reader = pa.RecordBatchReader.from_batches(self.schema, self.generator) + self.read._new_arrow_batch_reader.return_value = reader, self.generator + module = 'pypaimon.globalindex.create_global_index' + with patch(module + '._split_by_global_index_shard', return_value=[ + (_FakeSplit([]), Range(10, 19)), + ]), patch(module + '._to_index_manifest_entries', return_value=[]): + return self.builder._build_generic_index( + [], [], Mock(), self.read, '/unused') + + def test_batches_are_written_before_reading_the_next_batch(self): + self.builder._index_type = 'lucene' + written = [] + + def write(value, row_id): + self.events.append('write') + written.append((value, row_id)) + + def finish(): + self.assertEqual('reader closed', self.events[-1]) + return [] + + self.writer.write.side_effect = write + self.writer.finish.side_effect = finish + self._build([ + self._batch([], []), + self._batch([[9.0], [10.0], None], [9, 10, 11]), + self._batch([[19.0], [20.0]], [19, 20]), + ]) + self.assertEqual([([10.0], 0), (None, 1), ([19.0], 9)], written) + self.assertEqual([ + 'read', 'read', 'write', 'write', 'read', 'write', 'reader closed', + ], self.events) + self.writer.finish.assert_called_once() + self.writer.close.assert_called_once() + self.read.to_arrow.assert_not_called() + + def test_vector_batches_are_bounded_and_written_before_next_read(self): + written = [] + + def write_batch(vectors, row_ids): + self.assertLessEqual(len(row_ids), 2) + self.events.append('write batch') + written.extend(zip(vectors.to_pylist(), row_ids.to_pylist())) + + self.writer.write_batch.side_effect = write_batch + with patch('pypaimon.globalindex.create_global_index.ADD_BATCH_SIZE', 2): + self._build([ + self._batch([], []), + self._batch([[9.0], [10.0], None, [12.0]], [9, 10, 11, 12]), + self._batch([[19.0], [20.0]], [19, 20]), + ]) + self.assertEqual([([10.0], 0), (None, 1), ([12.0], 2), ([19.0], 9)], written) + self.assertEqual([ + 'read', 'read', 'write batch', 'write batch', 'read', + 'write batch', 'reader closed', + ], self.events) + self.writer.write.assert_not_called() + self.writer.finish.assert_called_once() + self.writer.close.assert_called_once() + self.read.to_arrow.assert_not_called() + + def test_empty_input_does_not_create_a_writer(self): + for batches in ([], [self._batch([], [])]): + with self.subTest(batches=len(batches)): + self.assertEqual([], self._build(batches)) + self.builder._create_generic_index_writer.assert_not_called() + self.assertEqual('reader closed', self.events[-1]) + + def test_failures_close_reader_and_writer(self): + for failure in ('create', 'read', 'write', 'finish', 'null_row_id'): + with self.subTest(failure=failure): + self.setUp() + error = RuntimeError('injected failure') + if failure == 'create': + self.builder._create_generic_index_writer.side_effect = error + elif failure in ('write', 'finish'): + getattr(self.writer, 'write_batch' if failure == 'write' else failure).side_effect = error + + def batches(): + yield self._batch([[10.0]], [10]) + if failure == 'read': + raise error + yield self._batch([[11.0]], [ + None if failure == 'null_row_id' else 11]) + + exception = ValueError if failure == 'null_row_id' else RuntimeError + message = '_ROW_ID is null' if failure == 'null_row_id' else 'injected failure' + with self.assertRaisesRegex(exception, message): + self._build(batches()) + self.assertEqual('reader closed', self.events[-1]) + if failure == 'create': + self.writer.close.assert_not_called() + else: + self.writer.close.assert_called_once() + if failure != 'finish': + self.writer.finish.assert_not_called() + + if __name__ == "__main__": unittest.main() diff --git a/paimon-python/pypaimon/tests/vindex_batch_write_test.py b/paimon-python/pypaimon/tests/vindex_batch_write_test.py index 72da3bb74531..def8858c0475 100644 --- a/paimon-python/pypaimon/tests/vindex_batch_write_test.py +++ b/paimon-python/pypaimon/tests/vindex_batch_write_test.py @@ -167,7 +167,7 @@ def test_builder_filters_ranges_before_vector_validation(self): with self.assertRaisesRegex(ValueError, '_ROW_ID is null'): _write_vector_batch(self._writer(), batch, 'embedding', Range(10, 19)) - def test_null_row_ids_are_rejected_before_writing_any_batch(self): + def test_null_row_ids_are_rejected_before_writing_source_batch(self): builder = object.__new__(GlobalIndexBuilder) builder._core_options = Mock() builder._core_options.global_index_row_count_per_shard.return_value = 10 @@ -176,10 +176,13 @@ def test_null_row_ids_are_rejected_before_writing_any_batch(self): writer = Mock() builder._create_generic_index_writer = Mock(return_value=writer) read = Mock() - read.to_arrow.return_value = pa.table({ + table = pa.table({ 'embedding': pa.array([[1], [3, 4]], type=pa.list_(pa.float32())), '_ROW_ID': pa.array([0, None], type=pa.int64()), }) + batches = iter(table.to_batches()) + read._new_arrow_batch_reader.return_value = ( + pa.RecordBatchReader.from_batches(table.schema, batches), batches) module = 'pypaimon.globalindex.create_global_index' with patch(module + '._split_by_global_index_shard', return_value=[ (Mock(), Range(0, 9)),