| # 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 sys | |
| import tempfile | |
| import time | |
| import unittest | |
| import numpy as np | |
| import pandas as pd | |
| import pyarrow as pa | |
| import pytest | |
| from pypaimon import CatalogFactory, Schema | |
| from pypaimon.common.options.core_options import CoreOptions | |
| from pypaimon.manifest.schema.manifest_entry import ManifestEntry | |
| from pypaimon.snapshot.snapshot import BATCH_COMMIT_IDENTIFIER | |
| from pypaimon.table.row.generic_row import GenericRow | |
| from pypaimon.write.file_store_commit import CommitFailRetryResult | |
| class AoReaderTest(unittest.TestCase): | |
| @classmethod | |
| def setUpClass(cls): | |
| cls.tempdir = tempfile.mkdtemp() | |
| cls.warehouse = os.path.join(cls.tempdir, 'warehouse') | |
| cls.catalog = CatalogFactory.create({ | |
| 'warehouse': cls.warehouse | |
| }) | |
| cls.catalog.create_database('default', True) | |
| cls.pa_schema = pa.schema([ | |
| ('user_id', pa.int32()), | |
| ('item_id', pa.int64()), | |
| ('behavior', pa.string()), | |
| ('dt', pa.string()) | |
| ]) | |
| cls.expected = pa.Table.from_pydict({ | |
| 'user_id': [1, 2, 3, 4, 5, 6, 7, 8], | |
| 'item_id': [1001, 1002, 1003, 1004, 1005, 1006, 1007, 1008], | |
| 'behavior': ['a', 'b', 'c', None, 'e', 'f', 'g', 'h'], | |
| 'dt': ['p1', 'p1', 'p2', 'p1', 'p2', 'p1', 'p2', 'p2'], | |
| }, schema=cls.pa_schema) | |
| @classmethod | |
| def tearDownClass(cls): | |
| shutil.rmtree(cls.tempdir, ignore_errors=True) | |
| def test_parquet_ao_reader(self): | |
| schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt']) | |
| self.catalog.create_table('default.test_append_only_parquet', schema, False) | |
| table = self.catalog.get_table('default.test_append_only_parquet') | |
| self._write_test_table(table) | |
| read_builder = table.new_read_builder() | |
| actual = self._read_test_table(read_builder).sort_by('user_id') | |
| self.assertEqual(actual, self.expected) | |
| @pytest.mark.python_write | |
| def test_target_file_row_num_fails_fast(self): | |
| schema = Schema.from_pyarrow_schema( | |
| self.pa_schema, partition_keys=['dt'], | |
| options={'target-file-row-num': '5'}) | |
| self.catalog.create_table('default.test_target_file_row_num', schema, False) | |
| table = self.catalog.get_table('default.test_target_file_row_num') | |
| with self.assertRaisesRegex( | |
| NotImplementedError, 'row-count based file rolling'): | |
| write_builder = table.new_batch_write_builder() | |
| table_write = write_builder.new_write() | |
| table_write.write_arrow(self.expected) | |
| def test_orc_ao_reader(self): | |
| schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt'], options={'file.format': 'orc'}) | |
| self.catalog.create_table('default.test_append_only_orc', schema, False) | |
| table = self.catalog.get_table('default.test_append_only_orc') | |
| self._write_test_table(table) | |
| read_builder = table.new_read_builder() | |
| actual = self._read_test_table(read_builder).sort_by('user_id') | |
| self.assertEqual(actual, self.expected) | |
| def test_avro_ao_reader(self): | |
| schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt'], options={'file.format': 'avro'}) | |
| self.catalog.create_table('default.test_append_only_avro', schema, False) | |
| table = self.catalog.get_table('default.test_append_only_avro') | |
| self._write_test_table(table) | |
| read_builder = table.new_read_builder() | |
| actual = self._read_test_table(read_builder).sort_by('user_id') | |
| self.assertEqual(actual, self.expected) | |
| def test_avro_ao_reader_filter_keeps_rows_past_first_batch(self): | |
| # A filtered Avro read must not stop when a full batch_size (1024) block | |
| # matches nothing: the reader returned None there, which the caller reads as | |
| # end-of-input, silently dropping every later matching row. | |
| schema = Schema.from_pyarrow_schema( | |
| self.pa_schema, partition_keys=['dt'], options={'file.format': 'avro'}) | |
| self.catalog.create_table('default.test_avro_filter_batches', schema, False) | |
| table = self.catalog.get_table('default.test_avro_filter_batches') | |
| n = 3000 | |
| data = pa.Table.from_pydict({ | |
| 'user_id': list(range(n)), | |
| 'item_id': [1000 + i for i in range(n)], | |
| 'behavior': ['a'] * n, | |
| 'dt': ['p1'] * n, | |
| }, schema=self.pa_schema) | |
| wb = table.new_batch_write_builder() | |
| w, c = wb.new_write(), wb.new_commit() | |
| w.write_arrow(data) | |
| c.commit(w.prepare_commit()) | |
| w.close() | |
| c.close() | |
| pb = table.new_read_builder().new_predicate_builder() | |
| # >= 2500 over 3000 rows fully filters the first two 1024-row batches, so the | |
| # test distinguishes the loop from an implementation that only retries once. | |
| read_builder = table.new_read_builder().with_filter( | |
| pb.greater_or_equal('user_id', 2500)) | |
| result = self._read_test_table(read_builder) | |
| self.assertEqual( | |
| sorted(result.column('user_id').to_pylist()), list(range(2500, n))) | |
| def test_lance_ao_reader(self): | |
| schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt'], options={'file.format': 'lance'}) | |
| self.catalog.create_table('default.test_append_only_lance', schema, False) | |
| table = self.catalog.get_table('default.test_append_only_lance') | |
| self._write_test_table(table) | |
| read_builder = table.new_read_builder() | |
| actual = self._read_test_table(read_builder).sort_by('user_id') | |
| self.assertEqual(actual, self.expected) | |
| def test_isnull_partition_filter(self): | |
| schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt']) | |
| self.catalog.create_table('default.isnull_partition_filter', schema, False) | |
| table = self.catalog.get_table('default.isnull_partition_filter') | |
| data = pa.Table.from_pydict({ | |
| 'user_id': [1, 2], | |
| 'item_id': [1001, 1002], | |
| 'behavior': ['a', 'b'], | |
| 'dt': [None, 'p1'], | |
| }, schema=self.pa_schema) | |
| wb = table.new_batch_write_builder() | |
| w, c = wb.new_write(), wb.new_commit() | |
| w.write_arrow(data) | |
| c.commit(w.prepare_commit()) | |
| w.close() | |
| c.close() | |
| pb = table.new_read_builder().new_predicate_builder() | |
| rb = table.new_read_builder().with_filter(pb.is_null('dt')) | |
| result = self._read_test_table(rb) | |
| self.assertEqual(result.column('dt').to_pylist(), [None]) | |
| self.assertEqual(result.column('user_id').to_pylist(), [1]) | |
| def test_plan_snapshot_id_for_empty_and_non_empty_scan(self): | |
| schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt']) | |
| self.catalog.create_table('default.test_plan_snapshot_id', schema, False) | |
| table = self.catalog.get_table('default.test_plan_snapshot_id') | |
| empty_plan = table.new_read_builder().new_scan().plan() | |
| self.assertIsNone(empty_plan.snapshot_id) | |
| self.assertEqual(len(empty_plan.splits()), 0) | |
| self._write_test_table(table) | |
| plan = table.new_read_builder().new_scan().plan() | |
| self.assertEqual(plan.snapshot_id, 2) | |
| self.assertGreater(len(plan.splits()), 0) | |
| def test_incremental_timestamp_empty_range_keeps_end_snapshot_id(self): | |
| schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt']) | |
| self.catalog.create_table('default.test_incremental_empty_range_snapshot', schema, False) | |
| table = self.catalog.get_table('default.test_incremental_empty_range_snapshot') | |
| write_builder = table.new_batch_write_builder() | |
| table_write = write_builder.new_write() | |
| table_commit = write_builder.new_commit() | |
| pa_table = pa.Table.from_pydict({ | |
| 'user_id': [1], | |
| 'item_id': [1001], | |
| 'behavior': ['a'], | |
| 'dt': ['p1'], | |
| }, schema=self.pa_schema) | |
| table_write.write_arrow(pa_table) | |
| table_commit.commit(table_write.prepare_commit()) | |
| table_write.close() | |
| table_commit.close() | |
| snapshot_manager = table.snapshot_manager() | |
| snapshot = snapshot_manager.get_latest_snapshot() | |
| table_inc = table.copy({ | |
| CoreOptions.INCREMENTAL_BETWEEN_TIMESTAMP.key(): | |
| "{},{}".format(snapshot.time_millis, snapshot.time_millis + 1) | |
| }) | |
| plan = table_inc.new_read_builder().new_scan().plan() | |
| self.assertEqual(plan.snapshot_id, snapshot.id) | |
| self.assertEqual(len(plan.splits()), 0) | |
| @unittest.skipIf(sys.version_info < (3, 11), "vortex-data requires Python >= 3.11") | |
| def test_vortex_ao_reader(self): | |
| schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt'], options={'file.format': 'vortex'}) | |
| self.catalog.create_table('default.test_append_only_vortex', schema, False) | |
| table = self.catalog.get_table('default.test_append_only_vortex') | |
| self._write_test_table(table) | |
| read_builder = table.new_read_builder() | |
| actual = self._read_test_table(read_builder).sort_by('user_id') | |
| self.assertEqual(actual, self.expected) | |
| @unittest.skipIf(sys.version_info < (3, 11), "vortex-data requires Python >= 3.11") | |
| def test_vortex_ao_reader_with_filter(self): | |
| schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt'], options={'file.format': 'vortex'}) | |
| self.catalog.create_table('default.test_append_only_vortex_filter', schema, False) | |
| table = self.catalog.get_table('default.test_append_only_vortex_filter') | |
| self._write_test_table(table) | |
| predicate_builder = table.new_read_builder().new_predicate_builder() | |
| p1 = predicate_builder.less_than('user_id', 7) | |
| p2 = predicate_builder.greater_or_equal('user_id', 2) | |
| p3 = predicate_builder.between('user_id', 0, 6) # [2/b, 3/c, 4/d, 5/e, 6/f] left | |
| p6 = predicate_builder.is_not_null('behavior') # exclude 4/d -> [2/b, 3/c, 5/e, 6/f] | |
| g1 = predicate_builder.and_predicates([p1, p2, p3, p6]) | |
| read_builder = table.new_read_builder().with_filter(g1) | |
| actual = self._read_test_table(read_builder) | |
| expected = pa.concat_tables([ | |
| self.expected.slice(1, 2), # 2/b, 3/c | |
| self.expected.slice(4, 2), # 5/e, 6/f | |
| ]) | |
| self.assertEqual(actual.sort_by('user_id'), expected) | |
| # OR predicates with startswith, endswith, contains, equal, is_null | |
| p7 = predicate_builder.startswith('behavior', 'a') | |
| p10 = predicate_builder.equal('item_id', 1002) | |
| p11 = predicate_builder.is_null('behavior') | |
| p9 = predicate_builder.contains('behavior', 'f') | |
| p8 = predicate_builder.endswith('dt', 'p2') | |
| g2 = predicate_builder.or_predicates([p7, p8, p9, p10, p11]) | |
| read_builder = table.new_read_builder().with_filter(g2) | |
| actual = self._read_test_table(read_builder) | |
| self.assertEqual(actual.sort_by('user_id'), self.expected) | |
| # Combined AND + OR | |
| g3 = predicate_builder.and_predicates([g1, g2]) | |
| read_builder = table.new_read_builder().with_filter(g3) | |
| actual = self._read_test_table(read_builder) | |
| expected = pa.concat_tables([ | |
| self.expected.slice(1, 2), # 2/b, 3/c | |
| self.expected.slice(4, 2), # 5/e, 6/f | |
| ]) | |
| self.assertEqual(actual.sort_by('user_id'), expected) | |
| # not_equal also filters None values | |
| p12 = predicate_builder.not_equal('behavior', 'f') | |
| read_builder = table.new_read_builder().with_filter(p12) | |
| actual = self._read_test_table(read_builder) | |
| expected = pa.concat_tables([ | |
| self.expected.slice(0, 1), # 1/a | |
| self.expected.slice(1, 1), # 2/b | |
| self.expected.slice(2, 1), # 3/c | |
| self.expected.slice(4, 1), # 5/e | |
| self.expected.slice(6, 1), # 7/g | |
| self.expected.slice(7, 1), # 8/h | |
| ]) | |
| self.assertEqual(actual.sort_by('user_id'), expected) | |
| @unittest.skipIf(sys.version_info < (3, 11), "vortex-data requires Python >= 3.11") | |
| def test_vortex_ao_reader_with_projection(self): | |
| schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt'], options={'file.format': 'vortex'}) | |
| self.catalog.create_table('default.test_vortex_append_only_projection', schema, False) | |
| table = self.catalog.get_table('default.test_vortex_append_only_projection') | |
| self._write_test_table(table) | |
| read_builder = table.new_read_builder().with_projection(['dt', 'user_id']) | |
| actual = self._read_test_table(read_builder).sort_by('user_id') | |
| expected = self.expected.select(['dt', 'user_id']) | |
| self.assertEqual(actual, expected) | |
| @unittest.skipIf(sys.version_info < (3, 11), "vortex-data requires Python >= 3.11") | |
| def test_vortex_ao_reader_with_shard(self): | |
| schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt'], options={'file.format': 'vortex'}) | |
| self.catalog.create_table('default.test_append_only_vortex_shard', schema, False) | |
| table = self.catalog.get_table('default.test_append_only_vortex_shard') | |
| self._write_test_table(table) | |
| read_builder = table.new_read_builder() | |
| table_read = read_builder.new_read() | |
| shard_tables = [] | |
| total_shards = 3 | |
| for i in range(total_shards): | |
| splits = read_builder.new_scan().with_shard(i, total_shards).plan().splits() | |
| shard_tables.append(table_read.to_arrow(splits)) | |
| actual = pa.concat_tables(shard_tables).sort_by('user_id') | |
| self.assertEqual(actual, self.expected) | |
| def test_lance_ao_reader_with_filter(self): | |
| schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt'], options={'file.format': 'lance'}) | |
| self.catalog.create_table('default.test_append_only_lance_filter', schema, False) | |
| table = self.catalog.get_table('default.test_append_only_lance_filter') | |
| self._write_test_table(table) | |
| predicate_builder = table.new_read_builder().new_predicate_builder() | |
| p1 = predicate_builder.less_than('user_id', 7) | |
| p2 = predicate_builder.greater_or_equal('user_id', 2) | |
| p3 = predicate_builder.between('user_id', 0, 6) # [2/b, 3/c, 4/d, 5/e, 6/f] left | |
| p4 = predicate_builder.is_not_in('behavior', ['b', 'e']) # [3/c, 4/d, 6/f] left | |
| p5 = predicate_builder.is_in('dt', ['p1']) # exclude 3/c | |
| p6 = predicate_builder.is_not_null('behavior') # exclude 4/d | |
| g1 = predicate_builder.and_predicates([p1, p2, p3, p4, p5, p6]) | |
| read_builder = table.new_read_builder().with_filter(g1) | |
| actual = self._read_test_table(read_builder) | |
| expected = pa.concat_tables([ | |
| self.expected.slice(5, 1) # 6/f | |
| ]) | |
| self.assertEqual(actual.sort_by('user_id'), expected) | |
| def test_lance_ao_reader_with_shard(self): | |
| schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt'], options={'file.format': 'lance'}) | |
| self.catalog.create_table('default.test_append_only_lance_shard', schema, False) | |
| table = self.catalog.get_table('default.test_append_only_lance_shard') | |
| self._write_test_table(table) | |
| read_builder = table.new_read_builder() | |
| table_read = read_builder.new_read() | |
| shard_tables = [] | |
| total_shards = 3 | |
| for i in range(total_shards): | |
| splits = read_builder.new_scan().with_shard(i, total_shards).plan().splits() | |
| shard_tables.append(table_read.to_arrow(splits)) | |
| actual = pa.concat_tables(shard_tables).sort_by('user_id') | |
| self.assertEqual(actual, self.expected) | |
| def test_lance_sliced_split_row_range_pushdown(self): | |
| """ | |
| SlicedSplit with Lance format calls read_range() instead of read_all(), | |
| reading only the requested row slice from disk rather than the full file. | |
| """ | |
| import unittest.mock as mock | |
| try: | |
| import lance as _lance | |
| import lance.file # ensure submodule is loaded before write_lance uses it | |
| except ImportError: | |
| self.skipTest("lance not installed") | |
| schema = Schema.from_pyarrow_schema( | |
| pa.schema([('id', pa.int64()), ('value', pa.string())]), | |
| options={'file.format': 'lance'}) | |
| self.catalog.create_table('default.test_lance_sliced_split', schema, False) | |
| table = self.catalog.get_table('default.test_lance_sliced_split') | |
| # Write 1000 rows in a single shot so they land in one file | |
| n = 1000 | |
| pa_table = pa.table({'id': list(range(n)), 'value': [f'v{i}' for i in range(n)]}) | |
| write_builder = table.new_batch_write_builder() | |
| table_write = write_builder.new_write() | |
| table_commit = write_builder.new_commit() | |
| table_write.write_arrow(pa_table) | |
| table_commit.commit(table_write.prepare_commit()) | |
| table_write.close() | |
| table_commit.close() | |
| scan = table.new_read_builder().new_scan() | |
| splits = scan.plan().splits() | |
| self.assertEqual(len(splits), 1, "Expected single split (no partition, single bucket)") | |
| data_split = splits[0] | |
| self.assertEqual(len(data_split.files), 1, "Expected single file in split") | |
| file_name = data_split.files[0].file_name | |
| slice_start, slice_end = 200, 700 # request 500 of 1000 rows | |
| from pypaimon.read.sliced_split import SlicedSplit | |
| sliced = SlicedSplit(data_split, {file_name: (slice_start, slice_end)}) | |
| # Spy on lance.file.LanceFileReader to verify read_range is used | |
| OrigReader = _lance.file.LanceFileReader | |
| read_all_calls = [] | |
| read_range_calls = [] | |
| class SpyLanceFileReader: | |
| def __init__(self, *args, **kwargs): | |
| self._r = OrigReader(*args, **kwargs) | |
| def read_all(self, *args, **kwargs): | |
| read_all_calls.append(()) | |
| return self._r.read_all(*args, **kwargs) | |
| def read_range(self, start, count, **kwargs): | |
| read_range_calls.append((start, count)) | |
| return self._r.read_range(start, count, **kwargs) | |
| def __getattr__(self, name): | |
| return getattr(self._r, name) | |
| table_read = table.new_read_builder().new_read() | |
| with mock.patch.object(_lance.file, 'LanceFileReader', SpyLanceFileReader): | |
| result = table_read.to_arrow([sliced]) | |
| # Verify that read_range was used, not read_all | |
| self.assertEqual(len(read_all_calls), 0, | |
| "read_all() must not be called when SlicedSplit row range is available") | |
| self.assertEqual(len(read_range_calls), 1) | |
| actual_start, actual_count = read_range_calls[0] | |
| self.assertEqual(actual_start, slice_start) | |
| self.assertEqual(actual_count, slice_end - slice_start, | |
| f"read_range should request exactly {slice_end - slice_start} rows, " | |
| f"not the full {n} rows of the file") | |
| # Verify functional correctness: correct rows returned | |
| self.assertEqual(result.num_rows, slice_end - slice_start) | |
| result_ids = sorted(result.column('id').to_pylist()) | |
| self.assertEqual(result_ids, list(range(slice_start, slice_end))) | |
| def test_lance_shard_row_id_correctness(self): | |
| """ | |
| with_shard splits a file into contiguous ranges via SlicedSplit. | |
| _ROW_ID across all shards must equal the full table's _ROW_ID. | |
| """ | |
| try: | |
| import lance # noqa: F401 | |
| except ImportError: | |
| self.skipTest("lance not installed") | |
| schema = Schema.from_pyarrow_schema( | |
| pa.schema([('id', pa.int64()), ('value', pa.string())]), | |
| options={'file.format': 'lance', | |
| 'row-tracking.enabled': 'true', | |
| 'data-evolution.enabled': 'true'}) | |
| self.catalog.create_table('default.test_lance_shard_row_id', schema, False) | |
| table = self.catalog.get_table('default.test_lance_shard_row_id') | |
| n = 1000 | |
| pa_table = pa.table({'id': list(range(n)), 'value': [f'v{i}' for i in range(n)]}) | |
| write_builder = table.new_batch_write_builder() | |
| table_write = write_builder.new_write() | |
| table_commit = write_builder.new_commit() | |
| table_write.write_arrow(pa_table) | |
| table_commit.commit(table_write.prepare_commit()) | |
| table_write.close() | |
| table_commit.close() | |
| # Read full table _ROW_ID as ground truth | |
| read_builder = table.new_read_builder().with_projection(['id', '_ROW_ID']) | |
| full_splits = read_builder.new_scan().plan().splits() | |
| table_read = read_builder.new_read() | |
| full_result = table_read.to_arrow(full_splits) | |
| expected_row_ids = sorted(full_result.column('_ROW_ID').to_pylist()) | |
| # Read via with_shard and collect _ROW_ID from all shards | |
| total_shards = 3 | |
| all_row_ids = [] | |
| for i in range(total_shards): | |
| shard_splits = read_builder.new_scan().with_shard(i, total_shards).plan().splits() | |
| shard_result = table_read.to_arrow(shard_splits) | |
| all_row_ids.extend(shard_result.column('_ROW_ID').to_pylist()) | |
| self.assertEqual(sorted(all_row_ids), expected_row_ids) | |
| def test_lance_indexed_split_take_rows_pushdown(self): | |
| """ | |
| IndexedSplit row ranges (from ANN global index results) are converted to | |
| local file indices and pushed down to lance.file.LanceFileReader.take_rows(), | |
| so only the matched rows are physically read instead of the full file. | |
| """ | |
| import unittest.mock as mock | |
| try: | |
| import lance as _lance | |
| import lance.file | |
| except ImportError: | |
| self.skipTest("lance not installed") | |
| schema = Schema.from_pyarrow_schema( | |
| pa.schema([('id', pa.int64()), ('value', pa.string())]), | |
| options={'file.format': 'lance', | |
| 'row-tracking.enabled': 'true', | |
| 'data-evolution.enabled': 'true'}) | |
| self.catalog.create_table('default.test_lance_indexed_split', schema, False) | |
| table = self.catalog.get_table('default.test_lance_indexed_split') | |
| n = 1000 | |
| pa_table = pa.table({'id': list(range(n)), 'value': [f'v{i}' for i in range(n)]}) | |
| write_builder = table.new_batch_write_builder() | |
| table_write = write_builder.new_write() | |
| table_commit = write_builder.new_commit() | |
| table_write.write_arrow(pa_table) | |
| table_commit.commit(table_write.prepare_commit()) | |
| table_write.close() | |
| table_commit.close() | |
| splits = table.new_read_builder().new_scan().plan().splits() | |
| self.assertEqual(len(splits), 1) | |
| data_split = splits[0] | |
| self.assertEqual(len(data_split.files), 1) | |
| file_meta = data_split.files[0] | |
| self.assertIsNotNone(file_meta.first_row_id, "row tracking must assign first_row_id") | |
| first_row_id = file_meta.first_row_id | |
| # Simulate ANN result: 3 scattered global row IDs | |
| from pypaimon.globalindex.indexed_split import IndexedSplit | |
| from pypaimon.utils.range import Range | |
| target_global_ids = [first_row_id + 50, first_row_id + 300, first_row_id + 700] | |
| row_ranges = [Range(g, g) for g in target_global_ids] | |
| indexed = IndexedSplit(data_split, row_ranges) | |
| OrigReader = _lance.file.LanceFileReader | |
| take_rows_calls = [] | |
| read_all_calls = [] | |
| class SpyLanceFileReader: | |
| def __init__(self, *args, **kwargs): | |
| self._r = OrigReader(*args, **kwargs) | |
| def take_rows(self, indices, **kwargs): | |
| take_rows_calls.append(list(indices)) | |
| return self._r.take_rows(indices, **kwargs) | |
| def read_all(self, *args, **kwargs): | |
| read_all_calls.append(()) | |
| return self._r.read_all(*args, **kwargs) | |
| def read_range(self, *args, **kwargs): | |
| return self._r.read_range(*args, **kwargs) | |
| def __getattr__(self, name): | |
| return getattr(self._r, name) | |
| table_read = table.new_read_builder().new_read() | |
| with mock.patch.object(_lance.file, 'LanceFileReader', SpyLanceFileReader): | |
| result = table_read.to_arrow([indexed]) | |
| self.assertEqual(len(read_all_calls), 0, | |
| "read_all() must not be called when IndexedSplit row ranges are available") | |
| self.assertEqual(len(take_rows_calls), 1) | |
| # local indices = global_ids - first_row_id | |
| expected_local = [g - first_row_id for g in target_global_ids] | |
| self.assertEqual(sorted(take_rows_calls[0]), sorted(expected_local), | |
| f"take_rows should request local indices {expected_local}, " | |
| f"not read the full {n} rows") | |
| self.assertEqual(result.num_rows, len(target_global_ids)) | |
| result_ids = sorted(result.column('id').to_pylist()) | |
| self.assertEqual(result_ids, [50, 300, 700]) | |
| def test_lance_indexed_split_row_id_correctness(self): | |
| """ | |
| IndexedSplit (ANN vector search) with native take_rows pushdown must | |
| return correct _ROW_ID values matching the original global row IDs. | |
| """ | |
| try: | |
| import lance # noqa: F401 | |
| except ImportError: | |
| self.skipTest("lance not installed") | |
| schema = Schema.from_pyarrow_schema( | |
| pa.schema([('id', pa.int64()), ('value', pa.string())]), | |
| options={'file.format': 'lance', | |
| 'row-tracking.enabled': 'true', | |
| 'data-evolution.enabled': 'true'}) | |
| self.catalog.create_table('default.test_lance_indexed_row_id', schema, False) | |
| table = self.catalog.get_table('default.test_lance_indexed_row_id') | |
| n = 1000 | |
| pa_table = pa.table({'id': list(range(n)), 'value': [f'v{i}' for i in range(n)]}) | |
| write_builder = table.new_batch_write_builder() | |
| table_write = write_builder.new_write() | |
| table_commit = write_builder.new_commit() | |
| table_write.write_arrow(pa_table) | |
| table_commit.commit(table_write.prepare_commit()) | |
| table_write.close() | |
| table_commit.close() | |
| # Read full table to get ground-truth _ROW_ID for specific rows | |
| read_builder = table.new_read_builder().with_projection(['id', '_ROW_ID']) | |
| full_splits = read_builder.new_scan().plan().splits() | |
| table_read = read_builder.new_read() | |
| full_result = table_read.to_arrow(full_splits) | |
| # Build id -> _ROW_ID mapping from full scan | |
| id_to_row_id = dict(zip( | |
| full_result.column('id').to_pylist(), | |
| full_result.column('_ROW_ID').to_pylist())) | |
| # Construct IndexedSplit targeting specific rows | |
| data_split = full_splits[0] | |
| file_meta = data_split.files[0] | |
| first_row_id = file_meta.first_row_id | |
| from pypaimon.globalindex.indexed_split import IndexedSplit | |
| from pypaimon.utils.range import Range | |
| target_local_offsets = [50, 300, 700] | |
| target_global_ids = [first_row_id + o for o in target_local_offsets] | |
| row_ranges = [Range(g, g) for g in target_global_ids] | |
| indexed = IndexedSplit(data_split, row_ranges) | |
| result = table_read.to_arrow([indexed]) | |
| self.assertEqual(result.num_rows, len(target_global_ids)) | |
| for i in range(result.num_rows): | |
| row_id = result.column('_ROW_ID')[i].as_py() | |
| data_id = result.column('id')[i].as_py() | |
| self.assertEqual(row_id, id_to_row_id[data_id], | |
| f"row id={data_id}: _ROW_ID should be {id_to_row_id[data_id]}, got {row_id}") | |
| def test_append_only_multi_write_once_commit(self): | |
| schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt']) | |
| self.catalog.create_table('default.test_append_only_multi_once_commit', schema, False) | |
| table = self.catalog.get_table('default.test_append_only_multi_once_commit') | |
| write_builder = table.new_batch_write_builder() | |
| table_write = write_builder.new_write() | |
| table_commit = write_builder.new_commit() | |
| data1 = { | |
| 'user_id': [1, 2, 3, 4], | |
| 'item_id': [1001, 1002, 1003, 1004], | |
| 'behavior': ['a', 'b', 'c', None], | |
| 'dt': ['p1', 'p1', 'p2', 'p1'], | |
| } | |
| pa_table1 = pa.Table.from_pydict(data1, schema=self.pa_schema) | |
| data2 = { | |
| 'user_id': [5, 6, 7, 8], | |
| 'item_id': [1005, 1006, 1007, 1008], | |
| 'behavior': ['e', 'f', 'g', 'h'], | |
| 'dt': ['p2', 'p1', 'p2', 'p2'], | |
| } | |
| pa_table2 = pa.Table.from_pydict(data2, schema=self.pa_schema) | |
| table_write.write_arrow(pa_table1) | |
| table_write.write_arrow(pa_table2) | |
| table_commit.commit(table_write.prepare_commit()) | |
| table_write.close() | |
| table_commit.close() | |
| read_builder = table.new_read_builder() | |
| actual = self._read_test_table(read_builder).sort_by('user_id') | |
| self.assertEqual(actual, self.expected) | |
| def test_commit_retry_filter(self): | |
| schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt']) | |
| self.catalog.create_table('default.test_commit_retry_filter', schema, False) | |
| table = self.catalog.get_table('default.test_commit_retry_filter') | |
| write_builder = table.new_batch_write_builder() | |
| table_write = write_builder.new_write() | |
| table_commit = write_builder.new_commit() | |
| data1 = { | |
| 'user_id': [1, 2, 3, 4], | |
| 'item_id': [1001, 1002, 1003, 1004], | |
| 'behavior': ['a', 'b', 'c', None], | |
| 'dt': ['p1', 'p1', 'p2', 'p1'], | |
| } | |
| pa_table1 = pa.Table.from_pydict(data1, schema=self.pa_schema) | |
| data2 = { | |
| 'user_id': [5, 6, 7, 8], | |
| 'item_id': [1005, 1006, 1007, 1008], | |
| 'behavior': ['e', 'f', 'g', 'h'], | |
| 'dt': ['p2', 'p1', 'p2', 'p2'], | |
| } | |
| pa_table2 = pa.Table.from_pydict(data2, schema=self.pa_schema) | |
| table_write.write_arrow(pa_table1) | |
| table_write.write_arrow(pa_table2) | |
| messages = table_write.prepare_commit() | |
| table_commit.commit(messages) | |
| table_write.close() | |
| snapshot_manager = table.snapshot_manager() | |
| latest_snapshot = snapshot_manager.get_latest_snapshot() | |
| commit_entries = [] | |
| for msg in messages: | |
| partition = GenericRow(list(msg.partition), table.partition_keys_fields) | |
| for file in msg.new_files: | |
| commit_entries.append(ManifestEntry( | |
| kind=0, | |
| partition=partition, | |
| bucket=msg.bucket, | |
| total_buckets=table.total_buckets, | |
| file=file | |
| )) | |
| # mock retry | |
| success = table_commit.file_store_commit._try_commit_once( | |
| CommitFailRetryResult(None), | |
| "APPEND", | |
| commit_entries, | |
| [], | |
| BATCH_COMMIT_IDENTIFIER, | |
| latest_snapshot) | |
| self.assertTrue(success.is_success()) | |
| table_commit.close() | |
| read_builder = table.new_read_builder() | |
| actual = self._read_test_table(read_builder).sort_by('user_id') | |
| self.assertEqual(actual, self.expected) | |
| def test_over_1000_cols_read(self): | |
| num_rows = 1 | |
| num_cols = 10 | |
| table_name = "default.testBug" | |
| # Generate dynamic schema based on column count | |
| schema_fields = [] | |
| for i in range(1, num_cols + 1): | |
| col_name = f'c{i:03d}' | |
| if i == 1: | |
| schema_fields.append((col_name, pa.string())) # ID column | |
| elif i == 2: | |
| schema_fields.append((col_name, pa.string())) # Name column | |
| elif i == 3: | |
| schema_fields.append((col_name, pa.string())) # Category column (partition key) | |
| elif i % 4 == 0: | |
| schema_fields.append((col_name, pa.float64())) # Float columns | |
| elif i % 4 == 1: | |
| schema_fields.append((col_name, pa.int32())) # Int columns | |
| elif i % 4 == 2: | |
| schema_fields.append((col_name, pa.string())) # String columns | |
| else: | |
| schema_fields.append((col_name, pa.int64())) # Long columns | |
| pa_schema = pa.schema(schema_fields) | |
| schema = Schema.from_pyarrow_schema( | |
| pa_schema, | |
| partition_keys=['c003'], # Use c003 as partition key | |
| ) | |
| # Create table | |
| self.catalog.create_table(table_name, schema, False) | |
| table = self.catalog.get_table(table_name) | |
| # Generate test data | |
| np.random.seed(42) # For reproducible results | |
| categories = ['Electronics', 'Clothing', 'Books', 'Home', 'Sports', 'Food', 'Toys', 'Beauty', 'Health', 'Auto'] | |
| statuses = ['Active', 'Inactive', 'Pending', 'Completed'] | |
| # Generate data dictionary | |
| test_data = {} | |
| for i in range(1, num_cols + 1): | |
| col_name = f'c{i:03d}' | |
| if i == 1: | |
| test_data[col_name] = [f'Product_{j}' for j in range(1, num_rows + 1)] | |
| elif i == 2: | |
| test_data[col_name] = [f'Product_{j}' for j in range(1, num_rows + 1)] | |
| elif i == 3: | |
| test_data[col_name] = np.random.choice(categories, num_rows) | |
| elif i % 4 == 0: | |
| test_data[col_name] = np.random.uniform(1.0, 1000.0, num_rows).round(2) | |
| elif i % 4 == 1: | |
| test_data[col_name] = np.random.randint(1, 100, num_rows) | |
| elif i % 4 == 2: | |
| test_data[col_name] = np.random.choice(statuses, num_rows) | |
| else: | |
| test_data[col_name] = np.random.randint(1640995200, 1672531200, num_rows) | |
| test_df = pd.DataFrame(test_data) | |
| write_builder = table.new_batch_write_builder() | |
| table_write = write_builder.new_write() | |
| table_commit = write_builder.new_commit() | |
| table_write.write_pandas(test_df) | |
| table_commit.commit(table_write.prepare_commit()) | |
| table_write.close() | |
| table_commit.close() | |
| read_builder = table.new_read_builder() | |
| table_scan = read_builder.new_scan() | |
| table_read = read_builder.new_read() | |
| result = table_read.to_pandas(table_scan.plan().splits()) | |
| self.assertEqual(result.to_dict(), test_df.to_dict()) | |
| def test_ao_reader_with_filter(self): | |
| schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt']) | |
| self.catalog.create_table('default.test_append_only_filter', schema, False) | |
| table = self.catalog.get_table('default.test_append_only_filter') | |
| self._write_test_table(table) | |
| predicate_builder = table.new_read_builder().new_predicate_builder() | |
| p1 = predicate_builder.less_than('user_id', 7) | |
| p2 = predicate_builder.greater_or_equal('user_id', 2) | |
| p3 = predicate_builder.between('user_id', 0, 6) # [2/b, 3/c, 4/d, 5/e, 6/f] left | |
| p4 = predicate_builder.is_not_in('behavior', ['b', 'e']) # [3/c, 4/d, 6/f] left | |
| p5 = predicate_builder.is_in('dt', ['p1']) # exclude 3/c | |
| p6 = predicate_builder.is_not_null('behavior') # exclude 4/d | |
| g1 = predicate_builder.and_predicates([p1, p2, p3, p4, p5, p6]) | |
| read_builder = table.new_read_builder().with_filter(g1) | |
| actual = self._read_test_table(read_builder) | |
| expected = pa.concat_tables([ | |
| self.expected.slice(5, 1) # 6/f | |
| ]) | |
| self.assertEqual(actual.sort_by('user_id'), expected) | |
| p7 = predicate_builder.startswith('behavior', 'a') | |
| p10 = predicate_builder.equal('item_id', 1002) | |
| p11 = predicate_builder.is_null('behavior') | |
| p9 = predicate_builder.contains('behavior', 'f') | |
| p8 = predicate_builder.endswith('dt', 'p2') | |
| g2 = predicate_builder.or_predicates([p7, p8, p9, p10, p11]) | |
| read_builder = table.new_read_builder().with_filter(g2) | |
| actual = self._read_test_table(read_builder) | |
| self.assertEqual(actual.sort_by('user_id'), self.expected) | |
| g3 = predicate_builder.and_predicates([g1, g2]) | |
| read_builder = table.new_read_builder().with_filter(g3) | |
| actual = self._read_test_table(read_builder) | |
| expected = pa.concat_tables([ | |
| self.expected.slice(5, 1) # 6/f | |
| ]) | |
| self.assertEqual(actual.sort_by('user_id'), expected) | |
| # Same as java, 'not_equal' will also filter records of 'None' value | |
| p12 = predicate_builder.not_equal('behavior', 'f') | |
| read_builder = table.new_read_builder().with_filter(p12) | |
| actual = self._read_test_table(read_builder) | |
| expected = pa.concat_tables([ | |
| # not only 6/f, but also 4/d will be filtered | |
| self.expected.slice(0, 1), # 1/a | |
| self.expected.slice(1, 1), # 2/b | |
| self.expected.slice(2, 1), # 3/c | |
| self.expected.slice(4, 1), # 5/e | |
| self.expected.slice(6, 1), # 7/g | |
| self.expected.slice(7, 1), # 8/h | |
| ]) | |
| self.assertEqual(actual.sort_by('user_id'), expected) | |
| def test_ao_reader_with_large_filter(self): | |
| schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt']) | |
| self.catalog.create_table('default.test_append_only_large_filter', schema, False) | |
| table = self.catalog.get_table('default.test_append_only_large_filter') | |
| self._write_test_table(table) | |
| predicate_builder = table.new_read_builder().new_predicate_builder() | |
| predicate = predicate_builder.or_predicates([ | |
| predicate_builder.equal('user_id', value) | |
| for value in list(range(100, 355)) + [2, 6] | |
| ]) | |
| read_builder = table.new_read_builder().with_filter(predicate) | |
| actual = self._read_test_table(read_builder).sort_by('user_id') | |
| expected = pa.concat_tables([ | |
| self.expected.slice(1, 1), | |
| self.expected.slice(5, 1), | |
| ]) | |
| self.assertEqual(actual, expected) | |
| def test_ao_reader_with_projection(self): | |
| schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt']) | |
| self.catalog.create_table('default.test_append_only_projection', schema, False) | |
| table = self.catalog.get_table('default.test_append_only_projection') | |
| self._write_test_table(table) | |
| read_builder = table.new_read_builder().with_projection(['dt', 'user_id']) | |
| actual = self._read_test_table(read_builder).sort_by('user_id') | |
| expected = self.expected.select(['dt', 'user_id']) | |
| self.assertEqual(actual, expected) | |
| def test_avro_ao_reader_with_projection(self): | |
| schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt'], options={'file.format': 'avro'}) | |
| self.catalog.create_table('default.test_avro_append_only_projection', schema, False) | |
| table = self.catalog.get_table('default.test_avro_append_only_projection') | |
| self._write_test_table(table) | |
| read_builder = table.new_read_builder().with_projection(['dt', 'user_id']) | |
| actual = self._read_test_table(read_builder).sort_by('user_id') | |
| expected = self.expected.select(['dt', 'user_id']) | |
| self.assertEqual(actual, expected) | |
| def test_ao_reader_with_limit(self): | |
| schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt']) | |
| self.catalog.create_table('default.test_append_only_limit', schema, False) | |
| table = self.catalog.get_table('default.test_append_only_limit') | |
| self._write_test_table(table) | |
| # Row-level limit: the reader stops at exactly N rows (not "first | |
| # split's full row count"). Scan still keeps the first split that | |
| # covers the limit; the reader short-circuits inside it. | |
| read_builder = table.new_read_builder().with_limit(1) | |
| actual = self._read_test_table(read_builder) | |
| self.assertEqual(actual.num_rows, 1) | |
| def test_incremental_timestamp(self): | |
| schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt']) | |
| self.catalog.create_table('default.test_incremental_parquet', schema, False) | |
| table = self.catalog.get_table('default.test_incremental_parquet') | |
| timestamp = int(time.time() * 1000) | |
| self._write_test_table(table) | |
| snapshot_manager = table.snapshot_manager() | |
| t1 = snapshot_manager.get_snapshot_by_id(1).time_millis | |
| t2 = snapshot_manager.get_snapshot_by_id(2).time_millis | |
| # test 1 | |
| table = table.copy({CoreOptions.INCREMENTAL_BETWEEN_TIMESTAMP.key(): str(timestamp - 1) + ',' + str(timestamp)}) | |
| read_builder = table.new_read_builder() | |
| actual = self._read_test_table(read_builder) | |
| self.assertEqual(len(actual), 0) | |
| # test 2 | |
| table = table.copy({CoreOptions.INCREMENTAL_BETWEEN_TIMESTAMP.key(): str(timestamp) + ',' + str(t2)}) | |
| read_builder = table.new_read_builder() | |
| actual = self._read_test_table(read_builder).sort_by('user_id') | |
| self.assertEqual(self.expected, actual) | |
| # test 3 | |
| table = table.copy({CoreOptions.INCREMENTAL_BETWEEN_TIMESTAMP.key(): str(t1) + ',' + str(t2)}) | |
| read_builder = table.new_read_builder() | |
| actual = self._read_test_table(read_builder).sort_by('user_id') | |
| expected = self.expected.slice(4, 4) | |
| self.assertEqual(expected, actual) | |
| def test_incremental_read_multi_snapshots(self): | |
| schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt']) | |
| self.catalog.create_table('default.test_incremental_100', schema, False) | |
| table = self.catalog.get_table('default.test_incremental_100') | |
| write_builder = table.new_batch_write_builder() | |
| for i in range(1, 101): | |
| table_write = write_builder.new_write() | |
| table_commit = write_builder.new_commit() | |
| pa_table = pa.Table.from_pydict({ | |
| 'user_id': [i], | |
| 'item_id': [1000 + i], | |
| 'behavior': [f'snap{i}'], | |
| 'dt': ['p1' if i % 2 == 1 else 'p2'], | |
| }, schema=self.pa_schema) | |
| table_write.write_arrow(pa_table) | |
| table_commit.commit(table_write.prepare_commit()) | |
| table_write.close() | |
| table_commit.close() | |
| snapshot_manager = table.snapshot_manager() | |
| t10 = snapshot_manager.get_snapshot_by_id(10).time_millis | |
| t20 = snapshot_manager.get_snapshot_by_id(20).time_millis | |
| table_inc = table.copy({CoreOptions.INCREMENTAL_BETWEEN_TIMESTAMP.key(): f"{t10},{t20}"}) | |
| read_builder = table_inc.new_read_builder() | |
| actual = self._read_test_table(read_builder).sort_by('user_id') | |
| expected = pa.Table.from_pydict({ | |
| 'user_id': list(range(11, 21)), | |
| 'item_id': [1000 + i for i in range(11, 21)], | |
| 'behavior': [f'snap{i}' for i in range(11, 21)], | |
| 'dt': ['p1' if i % 2 == 1 else 'p2' for i in range(11, 21)], | |
| }, schema=self.pa_schema).sort_by('user_id') | |
| self.assertEqual(expected, actual) | |
| def test_concurrent_writes_with_retry(self): | |
| """Test concurrent writes to verify retry mechanism works correctly.""" | |
| import threading | |
| # Run the test 10 times to verify stability | |
| iter_num = 5 | |
| for test_iteration in range(iter_num): | |
| # Create a unique table for each iteration | |
| table_name = f'default.test_concurrent_writes_{test_iteration}' | |
| # Concurrent commits are expected here; enlarge the retry budget so the | |
| # default (commit.max-retries=10, commit.max-retry-wait=1s) does not | |
| # exhaust under heavy CI load and produce a flaky failure. | |
| schema = Schema.from_pyarrow_schema( | |
| self.pa_schema, | |
| options={ | |
| 'commit.max-retries': '50', | |
| 'commit.max-retry-wait': '30s', | |
| }, | |
| ) | |
| self.catalog.create_table(table_name, schema, False) | |
| table = self.catalog.get_table(table_name) | |
| write_results = [] | |
| write_errors = [] | |
| def write_data(thread_id, start_user_id): | |
| """Write data in a separate thread.""" | |
| try: | |
| threading.current_thread().name = f"Iter{test_iteration}-Thread-{thread_id}" | |
| write_builder = table.new_batch_write_builder() | |
| table_write = write_builder.new_write() | |
| table_commit = write_builder.new_commit() | |
| # Create unique data for this thread | |
| data = { | |
| 'user_id': list(range(start_user_id, start_user_id + 5)), | |
| 'item_id': [1000 + i for i in range(start_user_id, start_user_id + 5)], | |
| 'behavior': [f'thread{thread_id}_{i}' for i in range(5)], | |
| 'dt': ['p1' if i % 2 == 0 else 'p2' for i in range(5)], | |
| } | |
| pa_table = pa.Table.from_pydict(data, schema=self.pa_schema) | |
| table_write.write_arrow(pa_table) | |
| commit_messages = table_write.prepare_commit() | |
| table_commit.commit(commit_messages) | |
| table_write.close() | |
| table_commit.close() | |
| write_results.append({ | |
| 'thread_id': thread_id, | |
| 'start_user_id': start_user_id, | |
| 'success': True | |
| }) | |
| except Exception as e: | |
| write_errors.append({ | |
| 'thread_id': thread_id, | |
| 'error': str(e) | |
| }) | |
| # Create and start multiple threads. Keep this modest (3 vs. the | |
| # original 10) because GHA runners under load can't drain 10 | |
| # simultaneously-conflicting commits even with | |
| # ``commit.max-retries=50`` (50 attempts * 30s back-off ~25 min, | |
| # still timing out in CI). Three threads exercises the retry path | |
| # without pushing each iteration past the per-test wall-time | |
| # budget. | |
| threads = [] | |
| num_threads = 3 | |
| for i in range(num_threads): | |
| thread = threading.Thread( | |
| target=write_data, | |
| args=(i, i * 10) | |
| ) | |
| threads.append(thread) | |
| thread.start() | |
| # Wait for all threads to complete | |
| for thread in threads: | |
| thread.join() | |
| # Verify all writes succeeded (retry mechanism should handle conflicts) | |
| self.assertEqual(num_threads, len(write_results), | |
| f"Iteration {test_iteration}: Expected {num_threads} successful writes, " | |
| f"got {len(write_results)}. Errors: {write_errors}") | |
| self.assertEqual(0, len(write_errors), | |
| f"Iteration {test_iteration}: Expected no errors, but got: {write_errors}") | |
| read_builder = table.new_read_builder() | |
| actual = self._read_test_table(read_builder).sort_by('user_id') | |
| # Verify data rows | |
| self.assertEqual(num_threads * 5, actual.num_rows, | |
| f"Iteration {test_iteration}: Expected {num_threads * 5} rows") | |
| # Verify user_id | |
| user_ids = actual.column('user_id').to_pylist() | |
| expected_user_ids = [] | |
| for i in range(num_threads): | |
| expected_user_ids.extend(range(i * 10, i * 10 + 5)) | |
| expected_user_ids.sort() | |
| self.assertEqual(user_ids, expected_user_ids, | |
| f"Iteration {test_iteration}: User IDs mismatch") | |
| # Verify snapshot count (should have num_threads snapshots) | |
| snapshot_manager = table.snapshot_manager() | |
| latest_snapshot = snapshot_manager.get_latest_snapshot() | |
| self.assertIsNotNone(latest_snapshot, | |
| f"Iteration {test_iteration}: Latest snapshot should not be None") | |
| self.assertEqual(latest_snapshot.id, num_threads, | |
| f"Iteration {test_iteration}: Expected snapshot ID {num_threads}, " | |
| f"got {latest_snapshot.id}") | |
| print(f"✓ Iteration {test_iteration + 1}/{iter_num} completed successfully") | |
| @pytest.mark.python_plan | |
| def test_is_in_with_partitions(self): | |
| from pypaimon.manifest.manifest_file_manager import ManifestFileManager | |
| from io import BytesIO | |
| import fastavro | |
| num_partitions = 5000 | |
| in_size = 2000 | |
| schema = Schema.from_pyarrow_schema( | |
| pa.schema([('pk', pa.int32()), ('val', pa.string()), | |
| ('pt', pa.string())]), | |
| partition_keys=['pt']) | |
| self.catalog.create_table( | |
| 'default.test_is_in_partition_bench', schema, False) | |
| table = self.catalog.get_table( | |
| 'default.test_is_in_partition_bench') | |
| wb = table.new_batch_write_builder() | |
| w = wb.new_write() | |
| bench_schema = pa.schema( | |
| [('pk', pa.int32()), ('val', pa.string()), | |
| ('pt', pa.string())]) | |
| for i in range(num_partitions): | |
| w.write_arrow(pa.Table.from_pydict( | |
| {'pk': [i], 'val': [f'v_{i}'], 'pt': [f'pt_{i}']}, | |
| schema=bench_schema)) | |
| wb.new_commit().commit(w.prepare_commit()) | |
| w.close() | |
| builder = table.new_read_builder().new_predicate_builder() | |
| in_values = [f'pt_{i}' for i in range(in_size)] | |
| pred = builder.is_in('pt', in_values) | |
| splits = table.new_read_builder().with_filter( | |
| pred).new_scan().plan().splits() | |
| self.assertEqual(len(splits), in_size) | |
| from pypaimon.manifest.schema.data_file_meta import DataFileMeta | |
| entry_counts = {'avro_total': 0, 'constructed': 0} | |
| original_read = ManifestFileManager.read | |
| original_dfm_init = DataFileMeta.__init__ | |
| def counting_read(self_mgr, manifest_file_name, | |
| manifest_entry_filter=None, | |
| drop_stats=True, early_entry_filter=None, | |
| early_record_filter=None, partition_filter=None, | |
| selected_blocks=None, file_size=None): | |
| # avro_total = every entry in the manifest (no manifest-file pruning | |
| # here: single file, is_in spans its partition stats). | |
| path = f"{self_mgr.manifest_path}/{manifest_file_name}" | |
| with self_mgr.file_io.new_input_stream(path) as s: | |
| for _ in fastavro.reader(BytesIO(s.read())): | |
| entry_counts['avro_total'] += 1 | |
| return original_read( | |
| self_mgr, manifest_file_name, | |
| manifest_entry_filter, drop_stats, | |
| early_entry_filter, early_record_filter, partition_filter, | |
| selected_blocks=selected_blocks, file_size=file_size) | |
| def counting_dfm_init(self_dfm, *args, **kwargs): | |
| entry_counts['constructed'] += 1 | |
| original_dfm_init(self_dfm, *args, **kwargs) | |
| ManifestFileManager.read = counting_read | |
| DataFileMeta.__init__ = counting_dfm_init | |
| try: | |
| table.new_read_builder().with_filter( | |
| pred).new_scan().plan() | |
| finally: | |
| ManifestFileManager.read = original_read | |
| DataFileMeta.__init__ = original_dfm_init | |
| # all entries scanned, but only matching ones build a DataFileMeta | |
| # (_FILE + drop_stats copy = 2 each); non-matching skipped pre-construct. | |
| self.assertEqual(entry_counts['avro_total'], num_partitions) | |
| self.assertEqual(entry_counts['constructed'], in_size * 2) | |
| def test_partition_filter_test_handles_isnull(self): | |
| # The early filter keeps/drops an entry via Predicate.test on its exact | |
| # partition value: isNull keeps null, drops non-null; equal is the reverse. | |
| from pypaimon.table.row.generic_row import GenericRow | |
| from pypaimon.table.row.row_kind import RowKind | |
| schema = Schema.from_pyarrow_schema( | |
| pa.schema([('pk', pa.int32()), ('pt', pa.string())]), | |
| partition_keys=['pt']) | |
| self.catalog.create_table('default.test_isnull_pt_pred', schema, False) | |
| table = self.catalog.get_table('default.test_isnull_pt_pred') | |
| fields = table.table_schema.fields # predicate index is table-schema space | |
| row_null = GenericRow([1, None], fields, RowKind.INSERT) | |
| row_x = GenericRow([2, 'x'], fields, RowKind.INSERT) | |
| pb = table.new_read_builder().new_predicate_builder() | |
| self.assertTrue(pb.is_null('pt').test(row_null)) | |
| self.assertFalse(pb.is_null('pt').test(row_x)) | |
| self.assertFalse(pb.equal('pt', 'x').test(row_null)) | |
| self.assertTrue(pb.equal('pt', 'x').test(row_x)) | |
| @pytest.mark.python_plan | |
| def test_early_partition_filter_isnull_scan(self): | |
| # Black-box: drive the early manifest partition filter through an actual | |
| # scan (not just Predicate.test), over a null + non-null partition mix. | |
| # isNull must keep only the null partition, and the non-matching entries | |
| # must be skipped before their _FILE block is deserialized. | |
| from pypaimon.manifest.schema.data_file_meta import DataFileMeta | |
| schema = Schema.from_pyarrow_schema( | |
| pa.schema([('pk', pa.int32()), ('pt', pa.string())]), | |
| partition_keys=['pt']) | |
| self.catalog.create_table('default.early_isnull_pt', schema, False) | |
| table = self.catalog.get_table('default.early_isnull_pt') | |
| ws = pa.schema([('pk', pa.int32()), ('pt', pa.string())]) | |
| wb = table.new_batch_write_builder() | |
| w = wb.new_write() | |
| w.write_arrow(pa.Table.from_pydict({'pk': [0], 'pt': [None]}, schema=ws)) | |
| for i in range(1, 6): | |
| w.write_arrow(pa.Table.from_pydict( | |
| {'pk': [i], 'pt': [f'p{i}']}, schema=ws)) | |
| wb.new_commit().commit(w.prepare_commit()) | |
| w.close() | |
| pb = table.new_read_builder().new_predicate_builder() | |
| # isNull keeps only the null partition; count DataFileMeta constructions to | |
| # prove the 5 non-null entries were skipped by the early filter pre-_FILE | |
| # (each kept entry builds 2: _FILE + drop_stats copy). | |
| constructed = {'n': 0} | |
| orig_init = DataFileMeta.__init__ | |
| def counting_init(self_dfm, *args, **kwargs): | |
| constructed['n'] += 1 | |
| orig_init(self_dfm, *args, **kwargs) | |
| DataFileMeta.__init__ = counting_init | |
| try: | |
| null_splits = table.new_read_builder().with_filter( | |
| pb.is_null('pt')).new_scan().plan().splits() | |
| finally: | |
| DataFileMeta.__init__ = orig_init | |
| self.assertEqual([s.partition.values for s in null_splits], [[None]]) | |
| self.assertEqual(constructed['n'], 2) | |
| # complement: equal keeps only the matching non-null partition | |
| p3_splits = table.new_read_builder().with_filter( | |
| pb.equal('pt', 'p3')).new_scan().plan().splits() | |
| self.assertEqual([s.partition.values for s in p3_splits], [['p3']]) | |
| def _write_test_table(self, table): | |
| write_builder = table.new_batch_write_builder() | |
| # first write | |
| table_write = write_builder.new_write() | |
| table_commit = write_builder.new_commit() | |
| data1 = { | |
| 'user_id': [1, 2, 3, 4], | |
| 'item_id': [1001, 1002, 1003, 1004], | |
| 'behavior': ['a', 'b', 'c', None], | |
| 'dt': ['p1', 'p1', 'p2', 'p1'], | |
| } | |
| pa_table = pa.Table.from_pydict(data1, schema=self.pa_schema) | |
| table_write.write_arrow(pa_table) | |
| table_commit.commit(table_write.prepare_commit()) | |
| table_write.close() | |
| table_commit.close() | |
| # second write | |
| table_write = write_builder.new_write() | |
| table_commit = write_builder.new_commit() | |
| data2 = { | |
| 'user_id': [5, 6, 7, 8], | |
| 'item_id': [1005, 1006, 1007, 1008], | |
| 'behavior': ['e', 'f', 'g', 'h'], | |
| 'dt': ['p2', 'p1', 'p2', 'p2'], | |
| } | |
| pa_table = pa.Table.from_pydict(data2, schema=self.pa_schema) | |
| table_write.write_arrow(pa_table) | |
| table_commit.commit(table_write.prepare_commit()) | |
| table_write.close() | |
| table_commit.close() | |
| def _read_test_table(self, read_builder): | |
| table_read = read_builder.new_read() | |
| splits = read_builder.new_scan().plan().splits() | |
| return table_read.to_arrow(splits) |