| # 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 unittest |
| from types import SimpleNamespace |
| |
| import pandas as pd |
| import pyarrow as pa |
| import pyarrow.dataset as ds |
| |
| from pypaimon import CatalogFactory, Schema |
| from pypaimon.common.predicate import Predicate |
| from pypaimon.manifest.manifest_list_manager import ManifestListManager |
| from pypaimon.table.row.offset_row import OffsetRow |
| |
| |
| def _filter_batch_arrow(batch, predicate): |
| expr = predicate.to_arrow() |
| table = ds.InMemoryDataset(pa.Table.from_batches([batch])).scanner(filter=expr).to_table() |
| if table.num_rows == 0: |
| return batch.slice(0, 0) |
| batches = table.to_batches() |
| if len(batches) == 1: |
| return batches[0] |
| return pa.RecordBatch.from_arrays( |
| [table.column(i) for i in range(table.num_columns)], schema=table.schema |
| ) |
| |
| |
| def _filter_batch_row_by_row(batch, predicate, ncols): |
| nrows = batch.num_rows |
| mask = [] |
| row_tuple = [None] * ncols |
| offset_row = OffsetRow(row_tuple, 0, ncols) |
| for i in range(nrows): |
| for j in range(ncols): |
| row_tuple[j] = batch.column(j)[i].as_py() |
| offset_row.replace(tuple(row_tuple)) |
| try: |
| mask.append(predicate.test(offset_row)) |
| except (TypeError, ValueError): |
| mask.append(False) |
| if not any(mask): |
| return batch.slice(0, 0) |
| return batch.filter(pa.array(mask)) |
| |
| |
| def _batches_equal(a, b): |
| if a.num_rows != b.num_rows or a.num_columns != b.num_columns: |
| return False |
| for i in range(a.num_columns): |
| col_a, col_b = a.column(i), b.column(i) |
| for j in range(a.num_rows): |
| va_py = col_a[j].as_py() if hasattr(col_a[j], "as_py") else col_a[j] |
| vb_py = col_b[j].as_py() if hasattr(col_b[j], "as_py") else col_b[j] |
| if va_py != vb_py: |
| return False |
| return True |
| |
| |
| class DataEvolutionTest(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', False) |
| |
| @classmethod |
| def tearDownClass(cls): |
| shutil.rmtree(cls.tempdir, ignore_errors=True) |
| |
| def test_basic(self): |
| simple_pa_schema = pa.schema([ |
| ('f0', pa.int8()), |
| ('f1', pa.int16()), |
| ]) |
| schema = Schema.from_pyarrow_schema(simple_pa_schema, |
| options={'row-tracking.enabled': 'true', 'data-evolution.enabled': 'true'}) |
| self.catalog.create_table('default.test_row_tracking', schema, False) |
| table = self.catalog.get_table('default.test_row_tracking') |
| |
| # write 1 |
| write_builder = table.new_batch_write_builder() |
| table_write = write_builder.new_write() |
| table_commit = write_builder.new_commit() |
| expect_data = pa.Table.from_pydict({ |
| 'f0': [-1, 2], |
| 'f1': [-1001, 1002] |
| }, schema=simple_pa_schema) |
| table_write.write_arrow(expect_data) |
| table_commit.commit(table_write.prepare_commit()) |
| table_write.close() |
| table_commit.close() |
| |
| # write 2 |
| table_write = write_builder.new_write().with_write_type(['f0']) |
| table_commit = write_builder.new_commit() |
| data2 = pa.Table.from_pydict({ |
| 'f0': [3, 4], |
| }, schema=pa.schema([ |
| ('f0', pa.int8()), |
| ])) |
| table_write.write_arrow(data2) |
| cmts = table_write.prepare_commit() |
| cmts[0].new_files[0].first_row_id = 0 |
| table_commit.commit(cmts) |
| table_write.close() |
| table_commit.close() |
| |
| read_builder = table.new_read_builder() |
| table_scan = read_builder.new_scan() |
| table_read = read_builder.new_read() |
| actual_data = table_read.to_arrow(table_scan.plan().splits()) |
| expect_data = pa.Table.from_pydict({ |
| 'f0': [3, 4], |
| 'f1': [-1001, 1002] |
| }, schema=pa.schema([ |
| ('f0', pa.int8()), |
| ('f1', pa.int16()), |
| ])) |
| self.assertEqual(actual_data, expect_data) |
| self.assertEqual( |
| len(actual_data.schema), len(expect_data.schema), |
| 'Read output column count must match schema') |
| self.assertEqual( |
| actual_data.schema.names, expect_data.schema.names, |
| 'Read output column names must match schema') |
| |
| def test_partitioned_read_requested_column_missing_in_file(self): |
| pa_schema = pa.schema([('f0', pa.int32()), ('f1', pa.string()), ('dt', pa.string())]) |
| schema = Schema.from_pyarrow_schema( |
| pa_schema, |
| partition_keys=['dt'], |
| options={'row-tracking.enabled': 'true', 'data-evolution.enabled': 'true'} |
| ) |
| self.catalog.create_table('default.test_partition_missing_col', schema, False) |
| table = self.catalog.get_table('default.test_partition_missing_col') |
| wb = table.new_batch_write_builder() |
| |
| tw1 = wb.new_write() |
| tc1 = wb.new_commit() |
| tw1.write_arrow(pa.Table.from_pydict( |
| {'f0': [1, 2], 'f1': ['a', 'b'], 'dt': ['p1', 'p1']}, |
| schema=pa_schema |
| )) |
| tc1.commit(tw1.prepare_commit()) |
| tw1.close() |
| tc1.close() |
| |
| tw2 = wb.new_write().with_write_type(['f0', 'dt']) |
| tc2 = wb.new_commit() |
| # Row key extractor uses table column indices; pass table-ordered data with null for f1 |
| tw2.write_arrow(pa.Table.from_pydict( |
| {'f0': [3, 4], 'f1': [None, None], 'dt': ['p1', 'p1']}, |
| schema=pa_schema |
| )) |
| tc2.commit(tw2.prepare_commit()) |
| tw2.close() |
| tc2.close() |
| |
| actual = table.new_read_builder().new_read().to_arrow(table.new_read_builder().new_scan().plan().splits()) |
| self.assertEqual(len(actual.schema), 3, 'Must have f0, f1, dt (no silent drop when f1 missing in file)') |
| self.assertEqual(actual.schema.names, ['f0', 'f1', 'dt']) |
| self.assertEqual(actual.num_rows, 4) |
| f1_col = actual.column('f1') |
| self.assertEqual(f1_col[0].as_py(), 'a') |
| self.assertEqual(f1_col[1].as_py(), 'b') |
| self.assertIsNone(f1_col[2].as_py()) |
| self.assertIsNone(f1_col[3].as_py()) |
| |
| # Assert manifest file meta contains min and max row id |
| manifest_list_manager = ManifestListManager(table) |
| snapshot_manager = table.snapshot_manager() |
| all_manifests = manifest_list_manager.read_all(snapshot_manager.get_latest_snapshot()) |
| first_commit = next((m for m in all_manifests if m.min_row_id == 0 and m.max_row_id == 1), None) |
| self.assertIsNotNone(first_commit, "Should have a manifest with min_row_id=0, max_row_id=1") |
| second_commit = next((m for m in all_manifests if m.min_row_id == 2 and m.max_row_id == 3), None) |
| self.assertIsNotNone(second_commit, "Should have a manifest with min_row_id=2, max_row_id=3") |
| |
| def test_merge_reader(self): |
| from pypaimon.read.reader.concat_batch_reader import MergeAllBatchReader |
| |
| simple_pa_schema = pa.schema([ |
| ('f0', pa.int32()), |
| ('f1', pa.string()), |
| ('f2', pa.string()), |
| ]) |
| schema = Schema.from_pyarrow_schema( |
| simple_pa_schema, |
| options={ |
| 'row-tracking.enabled': 'true', |
| 'data-evolution.enabled': 'true', |
| 'read.batch-size': '4096', |
| }, |
| ) |
| self.catalog.create_table('default.test_merge_reader_batch_sizes', schema, False) |
| table = self.catalog.get_table('default.test_merge_reader_batch_sizes') |
| |
| write_builder = table.new_batch_write_builder() |
| size = 5000 |
| w0 = write_builder.new_write().with_write_type(['f0', 'f1']) |
| w1 = write_builder.new_write().with_write_type(['f2']) |
| c = write_builder.new_commit() |
| d0 = pa.Table.from_pydict( |
| {'f0': list(range(size)), 'f1': [f'a{i}' for i in range(size)]}, |
| schema=pa.schema([('f0', pa.int32()), ('f1', pa.string())]), |
| ) |
| d1 = pa.Table.from_pydict( |
| {'f2': [f'b{i}' for i in range(size)]}, |
| schema=pa.schema([('f2', pa.string())]), |
| ) |
| w0.write_arrow(d0) |
| w1.write_arrow(d1) |
| cmts = w0.prepare_commit() + w1.prepare_commit() |
| for msg in cmts: |
| for nf in msg.new_files: |
| nf.first_row_id = 0 |
| c.commit(cmts) |
| w0.close() |
| w1.close() |
| c.close() |
| |
| original_merge_all = MergeAllBatchReader |
| call_count = [0] |
| |
| def patched_merge_all(reader_suppliers, batch_size=1024): |
| call_count[0] += 1 |
| if call_count[0] == 2: |
| batch_size = 999 |
| return original_merge_all(reader_suppliers, batch_size=batch_size) |
| |
| import pypaimon.read.split_read as split_read_module |
| split_read_module.MergeAllBatchReader = patched_merge_all |
| try: |
| read_builder = table.new_read_builder() |
| table_scan = read_builder.new_scan() |
| table_read = read_builder.new_read() |
| splits = table_scan.plan().splits() |
| actual_data = table_read.to_arrow(splits) |
| expect_data = pa.Table.from_pydict({ |
| 'f0': list(range(size)), |
| 'f1': [f'a{i}' for i in range(size)], |
| 'f2': [f'b{i}' for i in range(size)], |
| }, schema=simple_pa_schema) |
| self.assertEqual(actual_data.num_rows, size) |
| self.assertEqual(actual_data, expect_data) |
| finally: |
| split_read_module.MergeAllBatchReader = original_merge_all |
| |
| def test_with_slice(self): |
| pa_schema = pa.schema([ |
| ("id", pa.int64()), |
| ("b", pa.int32()), |
| ("c", pa.int32()), |
| ]) |
| schema = Schema.from_pyarrow_schema( |
| pa_schema, |
| options={ |
| "row-tracking.enabled": "true", |
| "data-evolution.enabled": "true", |
| "source.split.target-size": "512m", |
| }, |
| ) |
| table_name = "default.test_with_slice_data_evolution" |
| self.catalog.create_table(table_name, schema, ignore_if_exists=True) |
| table = self.catalog.get_table(table_name) |
| |
| for batch in [ |
| {"id": [1, 2], "b": [10, 20], "c": [100, 200]}, |
| {"id": [1001, 2001], "b": [1011, 2011], "c": [1001, 2001]}, |
| {"id": [-1, -2], "b": [-10, -20], "c": [-100, -200]}, |
| ]: |
| wb = table.new_batch_write_builder() |
| tw = wb.new_write() |
| tc = wb.new_commit() |
| tw.write_arrow(pa.Table.from_pydict(batch, schema=pa_schema)) |
| tc.commit(tw.prepare_commit()) |
| tw.close() |
| tc.close() |
| |
| rb = table.new_read_builder() |
| full_splits = rb.new_scan().plan().splits() |
| full_result = rb.new_read().to_pandas(full_splits) |
| self.assertEqual( |
| len(full_result), |
| 6, |
| "Full scan should return 6 rows", |
| ) |
| self.assertEqual( |
| sorted(full_result["id"].tolist()), |
| [-2, -1, 1, 2, 1001, 2001], |
| "Full set ids mismatch", |
| ) |
| |
| # with_slice(1, 4) -> row indices [1, 2, 3] -> 3 rows with id in (2, 1001, 2001) |
| scan = rb.new_scan().with_slice(1, 4) |
| splits = scan.plan().splits() |
| result = rb.new_read().to_pandas(splits) |
| self.assertEqual( |
| len(result), |
| 3, |
| "with_slice(1, 4) should return 3 rows (indices 1,2,3). " |
| "Bug: DataEvolutionSplitGenerator returns 2 when split has multiple data files.", |
| ) |
| ids = result["id"].tolist() |
| self.assertEqual( |
| sorted(ids), |
| [2, 1001, 2001], |
| "with_slice(1, 4) should return id in (2, 1001, 2001). Got ids=%s" % ids, |
| ) |
| scan_oob = rb.new_scan().with_slice(10, 12) |
| splits_oob = scan_oob.plan().splits() |
| result_oob = rb.new_read().to_pandas(splits_oob) |
| self.assertEqual( |
| len(result_oob), |
| 0, |
| "with_slice(10, 12) on 6 rows should return 0 rows (out of bounds), got %d" % len(result_oob), |
| ) |
| |
| # Out-of-bounds slice: 6 rows total, slice(10, 12) should return 0 rows |
| scan_oob = rb.new_scan().with_slice(10, 12) |
| splits_oob = scan_oob.plan().splits() |
| result_oob = rb.new_read().to_pandas(splits_oob) |
| self.assertEqual( |
| len(result_oob), |
| 0, |
| "with_slice(10, 12) on 6 rows should return 0 rows (out of bounds), got %d" |
| % len(result_oob), |
| ) |
| |
| def test_with_slice_partitioned_table(self): |
| pa_schema = pa.schema([ |
| ("pt", pa.int64()), |
| ("b", pa.int32()), |
| ("c", pa.int32()), |
| ]) |
| schema = Schema.from_pyarrow_schema( |
| pa_schema, |
| partition_keys=["pt"], |
| options={ |
| "row-tracking.enabled": "true", |
| "data-evolution.enabled": "true", |
| "source.split.target-size": "512m", |
| }, |
| ) |
| table_name = "default.test_with_slice_partitioned_table" |
| self.catalog.create_table(table_name, schema, ignore_if_exists=True) |
| table = self.catalog.get_table(table_name) |
| |
| for batch in [ |
| {"pt": [1, 1], "b": [10, 20], "c": [100, 200]}, |
| {"pt": [2, 2], "b": [1011, 2011], "c": [1001, 2001]}, |
| {"pt": [2, 2], "b": [-10, -20], "c": [-100, -200]}, |
| ]: |
| wb = table.new_batch_write_builder() |
| tw = wb.new_write() |
| tc = wb.new_commit() |
| tw.write_arrow(pa.Table.from_pydict(batch, schema=pa_schema)) |
| tc.commit(tw.prepare_commit()) |
| tw.close() |
| tc.close() |
| |
| rb: ReadBuilder = table.new_read_builder() |
| full_splits = rb.new_scan().plan().splits() |
| full_result = rb.new_read().to_pandas(full_splits) |
| self.assertEqual( |
| len(full_result), |
| 6, |
| "Full scan should return 6 rows", |
| ) |
| |
| predicate_builder = rb.new_predicate_builder() |
| rb.with_filter(predicate_builder.equal("pt", 2)) |
| |
| # 0 to 2 |
| scan_oob = rb.new_scan().with_slice(0, 2) |
| splits_oob = scan_oob.plan().splits() |
| result_oob = rb.new_read().to_pandas(splits_oob) |
| self.assertEqual( |
| sorted(result_oob["b"].tolist()), |
| [1011, 2011], |
| "Full set b mismatch", |
| ) |
| |
| # 2 to 4 |
| scan_oob = rb.new_scan().with_slice(2, 4) |
| splits_oob = scan_oob.plan().splits() |
| result_oob = rb.new_read().to_pandas(splits_oob) |
| self.assertEqual( |
| sorted(result_oob["b"].tolist()), |
| [-20, -10], |
| "Full set b mismatch", |
| ) |
| |
| def test_multiple_appends(self): |
| simple_pa_schema = pa.schema([ |
| ('f0', pa.int32()), |
| ('f1', pa.string()), |
| ('f2', pa.string()), |
| ]) |
| schema = Schema.from_pyarrow_schema( |
| simple_pa_schema, |
| options={'row-tracking.enabled': 'true', 'data-evolution.enabled': 'true'} |
| ) |
| self.catalog.create_table('default.test_multiple_appends', schema, False) |
| table = self.catalog.get_table('default.test_multiple_appends') |
| |
| write_builder = table.new_batch_write_builder() |
| |
| # write 100 rows: (1, "a", "b") |
| table_write = write_builder.new_write() |
| table_commit = write_builder.new_commit() |
| init_data = pa.Table.from_pydict({ |
| 'f0': [1] * 100, |
| 'f1': ['a'] * 100, |
| 'f2': ['b'] * 100, |
| }, schema=simple_pa_schema) |
| table_write.write_arrow(init_data) |
| table_commit.commit(table_write.prepare_commit()) |
| table_write.close() |
| table_commit.close() |
| # append：write (2, "x") and ("y"), set first_row_id = 100 |
| write0 = write_builder.new_write().with_write_type(['f0', 'f1']) |
| write1 = write_builder.new_write().with_write_type(['f2']) |
| commit = write_builder.new_commit() |
| data0 = pa.Table.from_pydict({'f0': [2], 'f1': ['x']}, |
| schema=pa.schema([('f0', pa.int32()), ('f1', pa.string())])) |
| data1 = pa.Table.from_pydict({'f2': ['y']}, schema=pa.schema([('f2', pa.string())])) |
| write0.write_arrow(data0) |
| write1.write_arrow(data1) |
| cmts = write0.prepare_commit() + write1.prepare_commit() |
| for c in cmts: |
| for nf in c.new_files: |
| nf.first_row_id = 100 |
| commit.commit(cmts) |
| write0.close() |
| write1.close() |
| commit.close() |
| |
| # append：write (3, "c") and ("d"), set first_row_id = 101 |
| write0 = write_builder.new_write().with_write_type(['f0', 'f1']) |
| commit0 = write_builder.new_commit() |
| data0 = pa.Table.from_pydict({'f0': [3], 'f1': ['c']}, |
| schema=pa.schema([('f0', pa.int32()), ('f1', pa.string())])) |
| write0.write_arrow(data0) |
| cmts0 = write0.prepare_commit() |
| for c in cmts0: |
| for nf in c.new_files: |
| nf.first_row_id = 101 |
| commit0.commit(cmts0) |
| write0.close() |
| commit0.close() |
| |
| write1 = write_builder.new_write().with_write_type(['f2']) |
| commit1 = write_builder.new_commit() |
| data1 = pa.Table.from_pydict({'f2': ['d']}, schema=pa.schema([('f2', pa.string())])) |
| write1.write_arrow(data1) |
| cmts1 = write1.prepare_commit() |
| for c in cmts1: |
| for nf in c.new_files: |
| nf.first_row_id = 101 |
| commit1.commit(cmts1) |
| write1.close() |
| commit1.close() |
| |
| read_builder = table.new_read_builder() |
| table_scan = read_builder.new_scan() |
| table_read = read_builder.new_read() |
| actual = table_read.to_arrow(table_scan.plan().splits()) |
| |
| self.assertEqual(actual.num_rows, 102) |
| expect = pa.Table.from_pydict({ |
| 'f0': [1] * 100 + [2] + [3], |
| 'f1': ['a'] * 100 + ['x'] + ['c'], |
| 'f2': ['b'] * 100 + ['y'] + ['d'], |
| }, schema=simple_pa_schema) |
| self.assertEqual(actual, expect) |
| self.assertEqual(len(actual.schema), len(expect.schema), 'Merge read output column count must match schema') |
| self.assertEqual(actual.schema.names, expect.schema.names, 'Merge read output column names must match schema') |
| |
| def test_disorder_cols_append(self): |
| simple_pa_schema = pa.schema([ |
| ('f0', pa.int32()), |
| ('f1', pa.string()), |
| ('f2', pa.string()), |
| ]) |
| schema = Schema.from_pyarrow_schema( |
| simple_pa_schema, |
| options={'row-tracking.enabled': 'true', 'data-evolution.enabled': 'true'} |
| ) |
| self.catalog.create_table('default.test_disorder_cols_append', schema, False) |
| table = self.catalog.get_table('default.test_disorder_cols_append') |
| |
| write_builder = table.new_batch_write_builder() |
| num_rows = 100 |
| # write 1 rows: (1, "a", "b") |
| table_write = write_builder.new_write() |
| table_commit = write_builder.new_commit() |
| init_data = pa.Table.from_pydict({ |
| 'f0': [1] * num_rows, |
| 'f1': ['a'] * num_rows, |
| 'f2': ['b'] * num_rows, |
| }, schema=simple_pa_schema) |
| table_write.write_arrow(init_data) |
| table_commit.commit(table_write.prepare_commit()) |
| table_write.close() |
| table_commit.close() |
| |
| # append：set first_row_id = 0 to modify the row with columns write |
| write0 = write_builder.new_write().with_write_type(['f0', 'f2']) |
| write1 = write_builder.new_write().with_write_type(['f1']) |
| commit = write_builder.new_commit() |
| data0 = pa.Table.from_pydict({'f0': [2] * num_rows, 'f2': ['y'] * num_rows}, |
| schema=pa.schema([('f0', pa.int32()), ('f2', pa.string())])) |
| data1 = pa.Table.from_pydict({'f1': ['x'] * num_rows}, schema=pa.schema([('f1', pa.string())])) |
| write0.write_arrow(data0) |
| write1.write_arrow(data1) |
| cmts = write0.prepare_commit() + write1.prepare_commit() |
| for c in cmts: |
| for nf in c.new_files: |
| nf.first_row_id = 0 |
| commit.commit(cmts) |
| write0.close() |
| write1.close() |
| commit.close() |
| |
| read_builder = table.new_read_builder() |
| table_scan = read_builder.new_scan() |
| table_read = read_builder.new_read() |
| actual = table_read.to_arrow(table_scan.plan().splits()) |
| |
| self.assertEqual(actual.num_rows, 100) |
| expect = pa.Table.from_pydict({ |
| 'f0': [2] * num_rows, |
| 'f1': ['x'] * num_rows, |
| 'f2': ['y'] * num_rows, |
| }, schema=simple_pa_schema) |
| self.assertEqual(actual, expect) |
| |
| read_builder2 = table.new_read_builder() |
| read_builder2.with_projection(['f2', 'f0', 'f1']) |
| actual2 = read_builder2.new_read().to_arrow( |
| read_builder2.new_scan().plan().splits()) |
| self.assertEqual(actual2.column('f0').to_pylist(), [2] * num_rows) |
| self.assertEqual(actual2.column('f1').to_pylist(), ['x'] * num_rows) |
| self.assertEqual(actual2.column('f2').to_pylist(), ['y'] * num_rows) |
| |
| def test_only_some_columns(self): |
| simple_pa_schema = pa.schema([ |
| ('f0', pa.int32()), |
| ('f1', pa.string()), |
| ('f2', pa.string()), |
| ]) |
| schema = Schema.from_pyarrow_schema( |
| simple_pa_schema, |
| options={'row-tracking.enabled': 'true', 'data-evolution.enabled': 'true'} |
| ) |
| self.catalog.create_table('default.test_only_some_columns', schema, False) |
| table = self.catalog.get_table('default.test_only_some_columns') |
| |
| write_builder = table.new_batch_write_builder() |
| |
| # Commit 1: f0 |
| w0 = write_builder.new_write().with_write_type(['f0']) |
| c0 = write_builder.new_commit() |
| d0 = pa.Table.from_pydict({'f0': [1]}, schema=pa.schema([('f0', pa.int32())])) |
| w0.write_arrow(d0) |
| c0.commit(w0.prepare_commit()) |
| w0.close() |
| c0.close() |
| |
| # Commit 2: f1, first_row_id = 0 |
| w1 = write_builder.new_write().with_write_type(['f1']) |
| c1 = write_builder.new_commit() |
| d1 = pa.Table.from_pydict({'f1': ['a']}, schema=pa.schema([('f1', pa.string())])) |
| w1.write_arrow(d1) |
| cmts1 = w1.prepare_commit() |
| for c in cmts1: |
| for nf in c.new_files: |
| nf.first_row_id = 0 |
| c1.commit(cmts1) |
| w1.close() |
| c1.close() |
| |
| # Commit 3: f2, first_row_id = 0 |
| w2 = write_builder.new_write().with_write_type(['f2']) |
| c2 = write_builder.new_commit() |
| d2 = pa.Table.from_pydict({'f2': ['b']}, schema=pa.schema([('f2', pa.string())])) |
| w2.write_arrow(d2) |
| cmts2 = w2.prepare_commit() |
| for c in cmts2: |
| for nf in c.new_files: |
| nf.first_row_id = 0 |
| c2.commit(cmts2) |
| w2.close() |
| c2.close() |
| |
| read_builder = table.new_read_builder() |
| table_scan = read_builder.new_scan() |
| table_read = read_builder.new_read() |
| actual = table_read.to_arrow(table_scan.plan().splits()) |
| |
| expect = pa.Table.from_pydict({ |
| 'f0': [1], |
| 'f1': ['a'], |
| 'f2': ['b'], |
| }, schema=simple_pa_schema) |
| self.assertEqual(actual, expect) |
| |
| def _create_filter_test_table(self, table_name: str): |
| pa_schema = pa.schema([ |
| ("id", pa.int64()), |
| ("b", pa.int32()), |
| pa.field("c", pa.int32(), nullable=True), |
| ]) |
| schema = Schema.from_pyarrow_schema( |
| pa_schema, options={"row-tracking.enabled": "true", "data-evolution.enabled": "true"}, |
| ) |
| self.catalog.create_table(table_name, schema, ignore_if_exists=True) |
| table = self.catalog.get_table(table_name) |
| wb = table.new_batch_write_builder() |
| w0, c0 = wb.new_write().with_write_type(["id", "b"]), wb.new_commit() |
| w0.write_arrow(pa.Table.from_pydict( |
| {"id": [1, 2, 3], "b": [10, 20, 30]}, |
| schema=pa.schema([("id", pa.int64()), ("b", pa.int32())]), |
| )) |
| c0.commit(w0.prepare_commit()) |
| w0.close() |
| c0.close() |
| w1, c1 = wb.new_write().with_write_type(["c"]), wb.new_commit() |
| w1.write_arrow(pa.Table.from_pydict( |
| {"c": [100, None, 200]}, |
| schema=pa.schema([pa.field("c", pa.int32(), nullable=True)]), |
| )) |
| cmts1 = w1.prepare_commit() |
| for cmt in cmts1: |
| for nf in cmt.new_files: |
| nf.first_row_id = 0 |
| c1.commit(cmts1) |
| w1.close() |
| c1.close() |
| return table |
| |
| def test_with_filter(self): |
| table = self._create_filter_test_table("default.test_filter_on_evolved_column") |
| rb = table.new_read_builder() |
| splits = rb.new_scan().plan().splits() |
| |
| full_df = rb.new_read().to_pandas(splits) |
| self.assertEqual(len(full_df), 3, "Full scan must return 3 rows") |
| full_sorted = full_df.sort_values("id").reset_index(drop=True) |
| self.assertEqual(full_sorted["id"].tolist(), [1, 2, 3]) |
| self.assertEqual(full_sorted["b"].tolist(), [10, 20, 30]) |
| self.assertEqual(full_sorted["c"].iloc[0], 100) |
| self.assertTrue(pd.isna(full_sorted["c"].iloc[1]), "Row id=2 must have NULL c") |
| self.assertEqual(full_sorted["c"].iloc[2], 200) |
| |
| predicate_gt = rb.new_predicate_builder().greater_than("c", 150) |
| rb_gt = table.new_read_builder().with_filter(predicate_gt) |
| result_gt = rb_gt.new_read().to_pandas(rb_gt.new_scan().plan().splits()) |
| self.assertEqual(len(result_gt), 1, "Filter c > 150 should return 1 row (c=200)") |
| self.assertEqual(result_gt["id"].iloc[0], 3, "Row with c=200 must have id=3") |
| self.assertEqual(result_gt["b"].iloc[0], 30, "Row with c=200 must have b=30") |
| self.assertEqual(result_gt["c"].iloc[0], 200, "Filtered row must have c=200") |
| |
| predicate_lt = rb.new_predicate_builder().less_than("c", 150) |
| rb_lt = table.new_read_builder().with_filter(predicate_lt) |
| result_lt = rb_lt.new_read().to_pandas(rb_lt.new_scan().plan().splits()) |
| self.assertEqual(len(result_lt), 1, "Filter c < 150 should return 1 row (c=100)") |
| self.assertEqual(result_lt["id"].iloc[0], 1, "Row with c=100 must have id=1") |
| self.assertEqual(result_lt["c"].iloc[0], 100, "Filtered row must have c=100") |
| |
| predicate_id = rb.new_predicate_builder().equal("id", 2) |
| rb_id = table.new_read_builder().with_filter(predicate_id) |
| result_id = rb_id.new_read().to_pandas(rb_id.new_scan().plan().splits()) |
| self.assertEqual(len(result_id), 1, "Filter id == 2 should return 1 row") |
| self.assertEqual(result_id["id"].iloc[0], 2, "Filtered row must have id=2") |
| self.assertTrue(pd.isna(result_id["c"].iloc[0]), "Row id=2 must have c=NULL") |
| |
| pb = rb.new_predicate_builder() |
| predicate_and = pb.and_predicates([ |
| pb.greater_than("c", 50), |
| pb.less_than("c", 150), |
| ]) |
| rb_and = table.new_read_builder().with_filter(predicate_and) |
| result_and = rb_and.new_read().to_pandas(rb_and.new_scan().plan().splits()) |
| self.assertEqual( |
| len(result_and), 1, |
| "Filter c>50 AND c<150 should return 1 row (c=100)", |
| ) |
| self.assertEqual(result_and["id"].iloc[0], 1, "Row with c=100 must have id=1") |
| self.assertEqual(result_and["c"].iloc[0], 100, "Filtered row must have c=100") |
| |
| predicate_is_null = rb.new_predicate_builder().is_null("c") |
| rb_null = table.new_read_builder().with_filter(predicate_is_null) |
| result_null = rb_null.new_read().to_pandas(rb_null.new_scan().plan().splits()) |
| self.assertEqual(len(result_null), 1, "Filter c IS NULL should return 1 row (id=2)") |
| self.assertEqual(result_null["id"].iloc[0], 2, "NULL row must have id=2") |
| self.assertTrue(pd.isna(result_null["c"].iloc[0]), "Filtered row c must be NULL") |
| |
| predicate_not_null = rb.new_predicate_builder().is_not_null("c") |
| rb_not_null = table.new_read_builder().with_filter(predicate_not_null) |
| result_not_null = rb_not_null.new_read().to_pandas( |
| rb_not_null.new_scan().plan().splits()) |
| self.assertEqual( |
| len(result_not_null), 2, |
| "Filter c IS NOT NULL should return 2 rows (id=1, id=3)", |
| ) |
| result_not_null_sorted = result_not_null.sort_values("id").reset_index(drop=True) |
| self.assertEqual(result_not_null_sorted["id"].tolist(), [1, 3]) |
| self.assertEqual(result_not_null_sorted["c"].tolist(), [100, 200]) |
| |
| predicate_or = pb.or_predicates([ |
| pb.greater_than("c", 150), |
| pb.less_than("c", 100), |
| ]) |
| rb_or = table.new_read_builder().with_filter(predicate_or) |
| result_or = rb_or.new_read().to_pandas(rb_or.new_scan().plan().splits()) |
| self.assertEqual( |
| len(result_or), 1, |
| "Filter c>150 OR c<100 should return 1 row (id=3, c=200)", |
| ) |
| self.assertEqual(result_or["id"].iloc[0], 3, "Row with c=200 must have id=3") |
| self.assertEqual(result_or["c"].iloc[0], 200, "Filtered row must have c=200") |
| |
| def test_with_filter_and_projection(self): |
| table = self._create_filter_test_table("default.test_filter_and_projection_evolved") |
| rb_full = table.new_read_builder() |
| predicate = rb_full.new_predicate_builder().greater_than("c", 150) |
| rb_filtered = table.new_read_builder().with_projection(["c", "id"]).with_filter(predicate) |
| result = rb_filtered.new_read().to_pandas(rb_filtered.new_scan().plan().splits()) |
| self.assertEqual(len(result), 1, "Filter c > 150 with projection [c, id] should return 1 row") |
| self.assertEqual(result["id"].iloc[0], 3) |
| self.assertEqual(result["c"].iloc[0], 200) |
| for _, row in result.iterrows(): |
| self.assertGreater( |
| row["c"], |
| 150, |
| "Each row must satisfy predicate c > 150 (row-by-row path uses predicate.index; " |
| "if schema_fields != read_type, wrong column is compared).", |
| ) |
| |
| predicate2 = rb_full.new_predicate_builder().is_null("c") |
| rb2_filtered = table.new_read_builder().with_projection(["id", "c"]).with_filter(predicate2) |
| result2 = rb2_filtered.new_read().to_pandas(rb2_filtered.new_scan().plan().splits()) |
| self.assertEqual(len(result2), 1, "Filter c IS NULL with projection [id, c] should return 1 row") |
| self.assertEqual(result2["id"].iloc[0], 2) |
| self.assertTrue(pd.isna(result2["c"].iloc[0])) |
| |
| predicate3 = rb_full.new_predicate_builder().greater_than("c", 50) |
| rb3_filtered = table.new_read_builder().with_projection(["c"]).with_filter(predicate3) |
| result3 = rb3_filtered.new_read().to_pandas(rb3_filtered.new_scan().plan().splits()) |
| self.assertEqual(len(result3), 2, "Filter c > 50 with projection [c] should return 2 rows (c=100, 200)") |
| self.assertEqual(sorted(result3["c"].tolist()), [100, 200]) |
| |
| # Build predicate from same read_type as projection [id, c] so indices match (c at index 1). |
| rb4 = table.new_read_builder().with_projection(["id", "c"]) |
| pb4 = rb4.new_predicate_builder() |
| predicate_compound = pb4.and_predicates([ |
| pb4.greater_than("c", 150), |
| pb4.is_not_null("c"), |
| ]) |
| rb4_filtered = rb4.with_filter(predicate_compound) |
| result4 = rb4_filtered.new_read().to_pandas(rb4_filtered.new_scan().plan().splits()) |
| self.assertEqual(len(result4), 1, "Filter c>150 AND c IS NOT NULL with projection [id,c] should return 1 row") |
| self.assertEqual(result4["id"].iloc[0], 3) |
| self.assertEqual(result4["c"].iloc[0], 200) |
| |
| predicate_filter_on_non_projected = rb_full.new_predicate_builder().greater_than("c", 150) |
| rb_non_projected = table.new_read_builder().with_projection(["id"]).with_filter( |
| predicate_filter_on_non_projected |
| ) |
| result_non_projected = rb_non_projected.new_read().to_pandas( |
| rb_non_projected.new_scan().plan().splits() |
| ) |
| self.assertEqual( |
| len(result_non_projected), |
| 3, |
| "Filter c > 150 with projection [id]: c not in read_type so filter is dropped, all 3 rows returned.", |
| ) |
| self.assertEqual( |
| list(result_non_projected.columns), |
| ["id"], |
| "Projection [id] should return only id column.", |
| ) |
| table_read = rb_non_projected.new_read() |
| splits = rb_non_projected.new_scan().plan().splits() |
| expected_output_arity = len(table_read.read_type) |
| try: |
| rows_from_iterator = list(table_read.to_iterator(splits)) |
| except ValueError as e: |
| if "Expected Arrow table or array" in str(e): |
| self.skipTest( |
| "RecordBatchReader path uses polars.from_arrow(RecordBatch) which fails; " |
| "skip to_iterator projection assertion on this path" |
| ) |
| raise |
| self.assertEqual(len(rows_from_iterator), 3, "to_iterator should return same row count as to_pandas") |
| for row in rows_from_iterator: |
| self.assertIsInstance(row, OffsetRow) |
| self.assertEqual( |
| row.arity, |
| expected_output_arity, |
| "to_iterator must yield rows with only read_type columns (arity=%d)." |
| % expected_output_arity, |
| ) |
| |
| def test_null_predicate_arrow_vs_row_by_row(self): |
| schema = pa.schema([("id", pa.int64()), ("c", pa.int64())]) |
| batch = pa.RecordBatch.from_pydict( |
| {"id": [1, 2, 3], "c": [10, None, 20]}, |
| schema=schema, |
| ) |
| ncols = 2 |
| |
| # is_null('c'): Arrow and row-by-row must return same rows |
| pred_is_null = Predicate(method="isNull", index=1, field="c", literals=None) |
| arrow_res = _filter_batch_arrow(batch, pred_is_null) |
| row_res = _filter_batch_row_by_row(batch, pred_is_null, ncols) |
| self.assertEqual(arrow_res.num_rows, row_res.num_rows) |
| self.assertTrue(_batches_equal(arrow_res, row_res)) |
| self.assertEqual(arrow_res.num_rows, 1) |
| self.assertEqual(arrow_res.column("id")[0].as_py(), 2) |
| self.assertIsNone(arrow_res.column("c")[0].as_py()) |
| |
| # is_not_null('c'): Arrow and row-by-row must return same rows |
| pred_not_null = Predicate(method="isNotNull", index=1, field="c", literals=None) |
| arrow_res2 = _filter_batch_arrow(batch, pred_not_null) |
| row_res2 = _filter_batch_row_by_row(batch, pred_not_null, ncols) |
| self.assertEqual(arrow_res2.num_rows, row_res2.num_rows) |
| self.assertTrue(_batches_equal(arrow_res2, row_res2)) |
| self.assertEqual(arrow_res2.num_rows, 2) |
| |
| pred_eq_null = Predicate(method="equal", index=1, field="c", literals=[None]) |
| row_res3 = _filter_batch_row_by_row(batch, pred_eq_null, ncols) |
| self.assertEqual(row_res3.num_rows, 0) # Paimon: val is None -> False, no row matches |
| arrow_res3 = _filter_batch_arrow(batch, pred_eq_null) |
| self.assertEqual(arrow_res3.num_rows, 0) # Arrow: NULL==NULL is null, filtered out |
| self.assertEqual(arrow_res3.num_rows, row_res3.num_rows) |
| |
| def test_filter_row_by_row_mismatched_schema(self): |
| batch = pa.RecordBatch.from_pydict( |
| {"c": [1, 200, 50], "id": [100, 2, 3]}, |
| schema=pa.schema([("c", pa.int64()), ("id", pa.int64())]), |
| ) |
| pred = Predicate(method="greaterThan", index=0, field="c", literals=[150]) |
| |
| ncols = 3 |
| nrows = batch.num_rows |
| id_col = batch.column("id") |
| c_col = batch.column("c") |
| row_tuple = [None] * ncols |
| offset_row = OffsetRow(row_tuple, 0, ncols) |
| mask = [] |
| for i in range(nrows): |
| row_tuple[0] = id_col[i].as_py() |
| row_tuple[1] = None |
| row_tuple[2] = c_col[i].as_py() |
| offset_row.replace(tuple(row_tuple)) |
| try: |
| mask.append(pred.test(offset_row)) |
| except (TypeError, ValueError): |
| mask.append(False) |
| rows_passing_wrong_layout = sum(mask) |
| self.assertEqual( |
| rows_passing_wrong_layout, |
| 0, |
| "With wrong layout (position 0 = id), predicate c > 150 becomes id > 150 -> 0 rows. " |
| "This reproduces FilterRecordBatchReader bug when schema_fields=table.fields.", |
| ) |
| ncols_right = 2 |
| row_tuple_right = [None] * ncols_right |
| offset_row_right = OffsetRow(row_tuple_right, 0, ncols_right) |
| mask_right = [] |
| for i in range(nrows): |
| row_tuple_right[0] = c_col[i].as_py() |
| row_tuple_right[1] = id_col[i].as_py() |
| offset_row_right.replace(tuple(row_tuple_right)) |
| try: |
| mask_right.append(pred.test(offset_row_right)) |
| except (TypeError, ValueError): |
| mask_right.append(False) |
| rows_passing_right_layout = sum(mask_right) |
| self.assertEqual( |
| rows_passing_right_layout, |
| 1, |
| "With correct layout (position 0 = c), predicate c > 150 -> 1 row (c=200).", |
| ) |
| |
| def test_null_values(self): |
| simple_pa_schema = pa.schema([ |
| ('f0', pa.int32()), |
| ('f1', pa.string()), |
| ('f2', pa.string()), |
| ]) |
| schema = Schema.from_pyarrow_schema( |
| simple_pa_schema, |
| options={'row-tracking.enabled': 'true', 'data-evolution.enabled': 'true'} |
| ) |
| self.catalog.create_table('default.test_null_values', schema, False) |
| table = self.catalog.get_table('default.test_null_values') |
| |
| write_builder = table.new_batch_write_builder() |
| |
| # Commit 1: some cols are null |
| w0 = write_builder.new_write().with_write_type(['f0', 'f1']) |
| w1 = write_builder.new_write().with_write_type(['f2']) |
| c = write_builder.new_commit() |
| |
| d0 = pa.Table.from_pydict({'f0': [1], 'f1': [None]}, |
| schema=pa.schema([('f0', pa.int32()), ('f1', pa.string())])) |
| d1 = pa.Table.from_pydict({'f2': [None]}, schema=pa.schema([('f2', pa.string())])) |
| w0.write_arrow(d0) |
| w1.write_arrow(d1) |
| cmts = w0.prepare_commit() + w1.prepare_commit() |
| for msg in cmts: |
| for nf in msg.new_files: |
| nf.first_row_id = 0 |
| c.commit(cmts) |
| w0.close() |
| w1.close() |
| c.close() |
| |
| # Commit 2 |
| w1 = write_builder.new_write().with_write_type(['f2']) |
| c1 = write_builder.new_commit() |
| d1 = pa.Table.from_pydict({'f2': ['c']}, schema=pa.schema([('f2', pa.string())])) |
| w1.write_arrow(d1) |
| cmts1 = w1.prepare_commit() |
| for msg in cmts1: |
| for nf in msg.new_files: |
| nf.first_row_id = 0 |
| c1.commit(cmts1) |
| w1.close() |
| c1.close() |
| |
| read_builder = table.new_read_builder() |
| table_scan = read_builder.new_scan() |
| table_read = read_builder.new_read() |
| actual = table_read.to_arrow(table_scan.plan().splits()) |
| expect = pa.Table.from_pydict({ |
| 'f0': [1], |
| 'f1': [None], |
| 'f2': ['c'], |
| }, schema=simple_pa_schema) |
| self.assertEqual(actual, expect) |
| |
| # different first_row_id append multiple times |
| def test_multiple_appends_different_first_row_ids(self): |
| simple_pa_schema = pa.schema([ |
| ('f0', pa.int32()), |
| ('f1', pa.string()), |
| ('f2', pa.string()), |
| ]) |
| schema = Schema.from_pyarrow_schema( |
| simple_pa_schema, |
| options={'row-tracking.enabled': 'true', 'data-evolution.enabled': 'true'} |
| ) |
| self.catalog.create_table('default.test_multiple_appends_diff_rowid', schema, False) |
| table = self.catalog.get_table('default.test_multiple_appends_diff_rowid') |
| |
| write_builder = table.new_batch_write_builder() |
| |
| # commit 1 |
| w0 = write_builder.new_write().with_write_type(['f0', 'f1']) |
| w1 = write_builder.new_write().with_write_type(['f2']) |
| c = write_builder.new_commit() |
| d0 = pa.Table.from_pydict({'f0': [1], 'f1': ['a']}, |
| schema=pa.schema([('f0', pa.int32()), ('f1', pa.string())])) |
| d1 = pa.Table.from_pydict({'f2': ['b']}, schema=pa.schema([('f2', pa.string())])) |
| w0.write_arrow(d0) |
| w1.write_arrow(d1) |
| cmts = w0.prepare_commit() + w1.prepare_commit() |
| for msg in cmts: |
| for nf in msg.new_files: |
| nf.first_row_id = 0 |
| c.commit(cmts) |
| w0.close() |
| w1.close() |
| c.close() |
| |
| # commit 2 |
| w0 = write_builder.new_write().with_write_type(['f0', 'f1']) |
| c0 = write_builder.new_commit() |
| d0 = pa.Table.from_pydict({'f0': [2], 'f1': ['c']}, |
| schema=pa.schema([('f0', pa.int32()), ('f1', pa.string())])) |
| w0.write_arrow(d0) |
| cmts0 = w0.prepare_commit() |
| for msg in cmts0: |
| for nf in msg.new_files: |
| nf.first_row_id = 1 |
| c0.commit(cmts0) |
| w0.close() |
| c0.close() |
| |
| # commit 3 |
| w1 = write_builder.new_write().with_write_type(['f2']) |
| c1 = write_builder.new_commit() |
| d1 = pa.Table.from_pydict({'f2': ['d']}, schema=pa.schema([('f2', pa.string())])) |
| w1.write_arrow(d1) |
| cmts1 = w1.prepare_commit() |
| for msg in cmts1: |
| for nf in msg.new_files: |
| nf.first_row_id = 1 |
| c1.commit(cmts1) |
| w1.close() |
| c1.close() |
| |
| read_builder = table.new_read_builder() |
| table_scan = read_builder.new_scan() |
| table_read = read_builder.new_read() |
| actual = table_read.to_arrow(table_scan.plan().splits()) |
| |
| expect = pa.Table.from_pydict({ |
| 'f0': [1, 2], |
| 'f1': ['a', 'c'], |
| 'f2': ['b', 'd'], |
| }, schema=simple_pa_schema) |
| self.assertEqual(actual, expect) |
| |
| read_builder = table.new_read_builder() |
| read_builder.with_projection(['f0', 'f1', 'f2', '_ROW_ID', '_SEQUENCE_NUMBER']) |
| table_scan = read_builder.new_scan() |
| table_read = read_builder.new_read() |
| actual_with_meta = table_read.to_arrow(table_scan.plan().splits()) |
| self.assertFalse( |
| actual_with_meta.schema.field('_ROW_ID').nullable, |
| '_ROW_ID must be non-nullable per SpecialFields', |
| ) |
| self.assertFalse( |
| actual_with_meta.schema.field('_SEQUENCE_NUMBER').nullable, |
| '_SEQUENCE_NUMBER must be non-nullable per SpecialFields', |
| ) |
| |
| rb_with_row_id = table.new_read_builder().with_projection(['f0', 'f1', 'f2', '_ROW_ID']) |
| pb = rb_with_row_id.new_predicate_builder() |
| rb_eq0 = table.new_read_builder().with_filter(pb.equal('_ROW_ID', 0)) |
| result_eq0 = rb_eq0.new_read().to_arrow(rb_eq0.new_scan().plan().splits()) |
| self.assertEqual(result_eq0, pa.Table.from_pydict( |
| {'f0': [1], 'f1': ['a'], 'f2': ['b']}, schema=simple_pa_schema)) |
| rb_eq1 = table.new_read_builder().with_filter(pb.equal('_ROW_ID', 1)) |
| result_eq1 = rb_eq1.new_read().to_arrow(rb_eq1.new_scan().plan().splits()) |
| self.assertEqual(result_eq1, pa.Table.from_pydict( |
| {'f0': [2], 'f1': ['c'], 'f2': ['d']}, schema=simple_pa_schema)) |
| rb_in = table.new_read_builder().with_filter(pb.is_in('_ROW_ID', [0, 1])) |
| result_in = rb_in.new_read().to_arrow(rb_in.new_scan().plan().splits()) |
| self.assertEqual(result_in, expect) |
| |
| def test_filter_by_row_id(self): |
| simple_pa_schema = pa.schema([('f0', pa.int32())]) |
| schema = Schema.from_pyarrow_schema( |
| simple_pa_schema, |
| options={'row-tracking.enabled': 'true', 'data-evolution.enabled': 'true'}, |
| ) |
| self.catalog.create_table('default.test_row_id_filter_empty_and_or', schema, False) |
| table = self.catalog.get_table('default.test_row_id_filter_empty_and_or') |
| write_builder = table.new_batch_write_builder() |
| |
| # Commit 1: _ROW_ID 0, 1 with f0=1, 2 |
| w = write_builder.new_write().with_write_type(['f0']) |
| c = write_builder.new_commit() |
| w.write_arrow(pa.Table.from_pydict( |
| {'f0': [1, 2]}, schema=pa.schema([('f0', pa.int32())]))) |
| cmts = w.prepare_commit() |
| for msg in cmts: |
| for nf in msg.new_files: |
| nf.first_row_id = 0 |
| c.commit(cmts) |
| w.close() |
| c.close() |
| |
| # Commit 2: _ROW_ID 2, 3 with f0=101, 102 |
| w = write_builder.new_write().with_write_type(['f0']) |
| c = write_builder.new_commit() |
| w.write_arrow(pa.Table.from_pydict( |
| {'f0': [101, 102]}, schema=pa.schema([('f0', pa.int32())]))) |
| cmts = w.prepare_commit() |
| for msg in cmts: |
| for nf in msg.new_files: |
| nf.first_row_id = 2 |
| c.commit(cmts) |
| w.close() |
| c.close() |
| |
| rb_with_row_id = table.new_read_builder().with_projection(['f0', '_ROW_ID']) |
| pb = rb_with_row_id.new_predicate_builder() |
| |
| # 1. Non-existent _ROW_ID -> empty |
| rb_eq999 = table.new_read_builder().with_filter(pb.equal('_ROW_ID', 999)) |
| result_eq999 = rb_eq999.new_read().to_arrow(rb_eq999.new_scan().plan().splits()) |
| self.assertEqual(len(result_eq999), 0, "Non-existent _ROW_ID should return empty") |
| |
| # 2. AND: _ROW_ID=0 AND f0=1 -> 1 row |
| rb_and = table.new_read_builder().with_filter( |
| pb.and_predicates([pb.equal('_ROW_ID', 0), pb.equal('f0', 1)]) |
| ) |
| result_and = rb_and.new_read().to_arrow(rb_and.new_scan().plan().splits()) |
| self.assertEqual(len(result_and), 1) |
| self.assertEqual(result_and['f0'][0].as_py(), 1) |
| |
| # 3. OR: _ROW_ID=0 OR f0>100 -> at least row with _ROW_ID=0 and all f0>100 |
| rb_or = table.new_read_builder().with_filter( |
| pb.or_predicates([pb.equal('_ROW_ID', 0), pb.greater_than('f0', 100)]) |
| ) |
| result_or = rb_or.new_read().to_arrow(rb_or.new_scan().plan().splits()) |
| f0_vals = set(result_or['f0'][i].as_py() for i in range(len(result_or))) |
| self.assertGreaterEqual(len(result_or), 3, "OR should return _ROW_ID=0 row and f0>100 rows") |
| self.assertIn(1, f0_vals, "_ROW_ID=0 row has f0=1") |
| self.assertIn(101, f0_vals) |
| self.assertIn(102, f0_vals) |
| |
| def test_filter_manifest_entries_by_row_ranges(self): |
| from pypaimon.read.scanner.file_scanner import _filter_manifest_entries_by_row_ranges |
| |
| entry_0 = SimpleNamespace(file=SimpleNamespace(first_row_id=0, row_count=1)) |
| entries = [entry_0] |
| row_ranges = [] |
| |
| filtered = _filter_manifest_entries_by_row_ranges(entries, row_ranges) |
| self.assertEqual(filtered, [], "empty row_ranges must return no entries, not all entries") |
| |
| def test_more_data(self): |
| simple_pa_schema = pa.schema([ |
| ('f0', pa.int32()), |
| ('f1', pa.string()), |
| ('f2', pa.string()), |
| ]) |
| schema = Schema.from_pyarrow_schema( |
| simple_pa_schema, |
| options={'row-tracking.enabled': 'true', 'data-evolution.enabled': 'true'} |
| ) |
| self.catalog.create_table('default.test_more_data', schema, False) |
| table = self.catalog.get_table('default.test_more_data') |
| |
| write_builder = table.new_batch_write_builder() |
| |
| # first commit：100k rows |
| w0 = write_builder.new_write().with_write_type(['f0', 'f1']) |
| w1 = write_builder.new_write().with_write_type(['f2']) |
| c = write_builder.new_commit() |
| size = 100000 |
| d0 = pa.Table.from_pydict({ |
| 'f0': list(range(size)), |
| 'f1': [f'a{i}' for i in range(size)], |
| }, schema=pa.schema([('f0', pa.int32()), ('f1', pa.string())])) |
| d1 = pa.Table.from_pydict({ |
| 'f2': [f'b{i}' for i in range(size)], |
| }, schema=pa.schema([('f2', pa.string())])) |
| w0.write_arrow(d0) |
| w1.write_arrow(d1) |
| cmts = w0.prepare_commit() + w1.prepare_commit() |
| for msg in cmts: |
| for nf in msg.new_files: |
| nf.first_row_id = 0 |
| c.commit(cmts) |
| w0.close() |
| w1.close() |
| c.close() |
| |
| # second commit：overwrite f2 to 'c{i}' |
| w1 = write_builder.new_write().with_write_type(['f2']) |
| c1 = write_builder.new_commit() |
| d1 = pa.Table.from_pydict({ |
| 'f2': [f'c{i}' for i in range(size)], |
| }, schema=pa.schema([('f2', pa.string())])) |
| w1.write_arrow(d1) |
| cmts1 = w1.prepare_commit() |
| c1.commit(cmts1) |
| w1.close() |
| c1.close() |
| |
| read_builder = table.new_read_builder() |
| table_scan = read_builder.new_scan() |
| table_read = read_builder.new_read() |
| actual = table_read.to_arrow(table_scan.plan().splits()) |
| |
| expect = pa.Table.from_pydict({ |
| 'f0': list(range(size)), |
| 'f1': [f'a{i}' for i in range(size)], |
| 'f2': [f'c{i}' for i in range(size)], |
| }, schema=simple_pa_schema) |
| self.assertEqual(actual, expect) |
| |
| def test_read_row_tracking_metadata(self): |
| simple_pa_schema = pa.schema([ |
| ('f0', pa.int8()), |
| ('f1', pa.int16()), |
| ]) |
| schema = Schema.from_pyarrow_schema(simple_pa_schema, |
| options={'row-tracking.enabled': 'true', 'data-evolution.enabled': 'true'}) |
| self.catalog.create_table('default.test_row_tracking_meta', schema, False) |
| table = self.catalog.get_table('default.test_row_tracking_meta') |
| |
| # write 1 |
| write_builder = table.new_batch_write_builder() |
| table_write = write_builder.new_write() |
| table_commit = write_builder.new_commit() |
| expect_data = pa.Table.from_pydict({ |
| 'f0': [-1, 2], |
| 'f1': [-1001, 1002] |
| }, schema=simple_pa_schema) |
| table_write.write_arrow(expect_data) |
| table_commit.commit(table_write.prepare_commit()) |
| table_write.close() |
| table_commit.close() |
| |
| read_builder = table.new_read_builder() |
| read_builder.with_projection(['f0', '_ROW_ID', 'f1', '_SEQUENCE_NUMBER']) |
| table_scan = read_builder.new_scan() |
| table_read = read_builder.new_read() |
| actual_data = table_read.to_arrow(table_scan.plan().splits()) |
| expect_data = pa.Table.from_pydict({ |
| 'f0': [-1, 2], |
| '_ROW_ID': [0, 1], |
| 'f1': [-1001, 1002], |
| '_SEQUENCE_NUMBER': [1, 1], |
| }, schema=pa.schema([ |
| ('f0', pa.int8()), |
| pa.field('_ROW_ID', pa.int64(), nullable=False), |
| ('f1', pa.int16()), |
| pa.field('_SEQUENCE_NUMBER', pa.int64(), nullable=False), |
| ])) |
| self.assertEqual(actual_data, expect_data) |
| self.assertEqual(len(actual_data.schema), len(expect_data.schema), 'Read output column count must match schema') |
| |
| # write 2 |
| table_write = write_builder.new_write().with_write_type(['f0']) |
| table_commit = write_builder.new_commit() |
| data2 = pa.Table.from_pydict({ |
| 'f0': [3, 4], |
| }, schema=pa.schema([ |
| ('f0', pa.int8()), |
| ])) |
| table_write.write_arrow(data2) |
| cmts = table_write.prepare_commit() |
| cmts[0].new_files[0].first_row_id = 0 |
| table_commit.commit(cmts) |
| table_write.close() |
| table_commit.close() |
| |
| read_builder = table.new_read_builder() |
| read_builder.with_projection(['f0', 'f1', '_ROW_ID', '_SEQUENCE_NUMBER']) |
| table_scan = read_builder.new_scan() |
| table_read = read_builder.new_read() |
| actual_data = table_read.to_arrow(table_scan.plan().splits()) |
| self.assertFalse(actual_data.schema.field('_ROW_ID').nullable) |
| self.assertFalse(actual_data.schema.field('_SEQUENCE_NUMBER').nullable) |
| expect_data = pa.Table.from_pydict({ |
| 'f0': [3, 4], |
| 'f1': [-1001, 1002], |
| '_ROW_ID': [0, 1], |
| '_SEQUENCE_NUMBER': [2, 2], |
| }, schema=pa.schema([ |
| ('f0', pa.int8()), |
| ('f1', pa.int16()), |
| pa.field('_ROW_ID', pa.int64(), nullable=False), |
| pa.field('_SEQUENCE_NUMBER', pa.int64(), nullable=False), |
| ])) |
| self.assertEqual(actual_data, expect_data) |
| self.assertEqual(len(actual_data.schema), len(expect_data.schema), 'Read output column count must match schema') |
| |
| def test_with_blob(self): |
| from pypaimon.table.row.blob import BlobDescriptor |
| |
| pa_schema = pa.schema([ |
| ('id', pa.int32()), |
| ('picture', pa.large_binary()), |
| ]) |
| schema = Schema.from_pyarrow_schema( |
| pa_schema, |
| options={ |
| 'row-tracking.enabled': 'true', |
| 'data-evolution.enabled': 'true', |
| 'blob-as-descriptor': 'true', |
| }, |
| ) |
| self.catalog.create_table('default.test_with_blob', schema, False) |
| table = self.catalog.get_table('default.test_with_blob') |
| |
| blob_path = os.path.join(self.tempdir, 'blob_ev') |
| with open(blob_path, 'wb') as f: |
| f.write(b'x') |
| descriptor = BlobDescriptor(blob_path, 0, 1) |
| |
| wb = table.new_batch_write_builder() |
| tw = wb.new_write() |
| tc = wb.new_commit() |
| tw.write_arrow(pa.Table.from_pydict( |
| {'id': [1], 'picture': [descriptor.serialize()]}, |
| schema=pa_schema, |
| )) |
| cmts = tw.prepare_commit() |
| if cmts and cmts[0].new_files: |
| for nf in cmts[0].new_files: |
| nf.first_row_id = 0 |
| tc.commit(cmts) |
| tw.close() |
| tc.close() |
| |
| tw = wb.new_write() |
| tc = wb.new_commit() |
| tw.write_arrow(pa.Table.from_pydict( |
| {'id': [2], 'picture': [descriptor.serialize()]}, |
| schema=pa_schema, |
| )) |
| cmts = tw.prepare_commit() |
| if cmts and cmts[0].new_files: |
| for nf in cmts[0].new_files: |
| nf.first_row_id = 1 |
| tc.commit(cmts) |
| tw.close() |
| tc.close() |
| |
| rb = table.new_read_builder() |
| rb.with_projection(['id', '_ROW_ID', 'picture', '_SEQUENCE_NUMBER']) |
| actual = rb.new_read().to_arrow(rb.new_scan().plan().splits()) |
| self.assertEqual(actual.num_rows, 2) |
| self.assertEqual(actual.column('id').to_pylist(), [1, 2]) |
| self.assertEqual(actual.column('_ROW_ID').to_pylist(), [0, 1]) |
| |
| def test_from_arrays_without_schema(self): |
| schema = pa.schema([ |
| ('f0', pa.int8()), |
| pa.field('_ROW_ID', pa.int64(), nullable=False), |
| pa.field('_SEQUENCE_NUMBER', pa.int64(), nullable=False), |
| ]) |
| batch = pa.RecordBatch.from_pydict( |
| {'f0': [1], '_ROW_ID': [0], '_SEQUENCE_NUMBER': [1]}, |
| schema=schema |
| ) |
| self.assertFalse(batch.schema.field('_ROW_ID').nullable) |
| self.assertFalse(batch.schema.field('_SEQUENCE_NUMBER').nullable) |
| |
| arrays = list(batch.columns) |
| rebuilt = pa.RecordBatch.from_arrays(arrays, names=batch.schema.names) |
| self.assertTrue(rebuilt.schema.field('_ROW_ID').nullable) |
| self.assertTrue(rebuilt.schema.field('_SEQUENCE_NUMBER').nullable) |
| |
| def test_read_full_schema_on_write_before_evolution(self): |
| from pypaimon.schema.schema_change import SchemaChange |
| from pypaimon.schema.data_types import AtomicType |
| |
| # Step 1: Create table with [f0, f1] |
| initial_schema = pa.schema([ |
| ('f0', pa.int32()), |
| ('f1', pa.string()), |
| ]) |
| schema = Schema.from_pyarrow_schema( |
| initial_schema, |
| options={'row-tracking.enabled': 'true', 'data-evolution.enabled': 'true'}, |
| ) |
| table_name = 'default.test_no_write_cols_schema_evo' |
| self.catalog.create_table(table_name, schema, False) |
| table = self.catalog.get_table(table_name) |
| |
| # Step 2: Write data with ALL columns of old schema → write_cols = None |
| 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.from_pydict( |
| {'f0': [1, 2], 'f1': ['a', 'b']}, |
| schema=initial_schema, |
| )) |
| cmts = table_write.prepare_commit() |
| for c in cmts: |
| for nf in c.new_files: |
| self.assertIsNone(nf.write_cols) |
| table_commit.commit(cmts) |
| table_write.close() |
| table_commit.close() |
| |
| # Step 3: Schema evolution - add column f2 |
| self.catalog.alter_table( |
| table_name, |
| [SchemaChange.add_column('f2', AtomicType('STRING'))], |
| ) |
| table = self.catalog.get_table(table_name) |
| |
| # Step 4: Write f2 only for same rows (first_row_id = 0) |
| write_builder = table.new_batch_write_builder() |
| table_write = write_builder.new_write().with_write_type(['f2']) |
| table_commit = write_builder.new_commit() |
| table_write.write_arrow(pa.Table.from_pydict( |
| {'f2': ['x', 'y']}, |
| schema=pa.schema([('f2', pa.string())]), |
| )) |
| cmts_f2 = table_write.prepare_commit() |
| for c in cmts_f2: |
| for nf in c.new_files: |
| nf.first_row_id = 0 |
| table_commit.commit(cmts_f2) |
| table_write.close() |
| table_commit.close() |
| |
| # Step 5: Read all columns |
| read_builder = table.new_read_builder() |
| table_scan = read_builder.new_scan() |
| table_read = read_builder.new_read() |
| splits = table_scan.plan().splits() |
| |
| for split in splits: |
| for f in split.files: |
| if f.write_cols is None: |
| f.max_sequence_number = 999999 |
| |
| actual = table_read.to_arrow(splits) |
| |
| expect = pa.Table.from_pydict({ |
| 'f0': [1, 2], |
| 'f1': ['a', 'b'], |
| 'f2': ['x', 'y'], |
| }, schema=pa.schema([ |
| ('f0', pa.int32()), |
| ('f1', pa.string()), |
| ('f2', pa.string()), |
| ])) |
| self.assertEqual(actual.num_rows, 2) |
| self.assertEqual(actual, expect) |
| |
| @unittest.skipIf(sys.version_info < (3, 11), "vortex-data requires Python >= 3.11") |
| def test_vortex_basic(self): |
| """Test basic data evolution read/write with Vortex format.""" |
| pa_schema = pa.schema([ |
| ('f0', pa.int32()), |
| ('f1', pa.string()), |
| ('f2', pa.string()), |
| ]) |
| schema = Schema.from_pyarrow_schema( |
| pa_schema, |
| options={ |
| 'row-tracking.enabled': 'true', |
| 'data-evolution.enabled': 'true', |
| 'file.format': 'vortex', |
| } |
| ) |
| self.catalog.create_table('default.test_vortex_basic', schema, False) |
| table = self.catalog.get_table('default.test_vortex_basic') |
| |
| write_builder = table.new_batch_write_builder() |
| |
| # Commit 1: write f0, f1 |
| w0 = write_builder.new_write().with_write_type(['f0', 'f1']) |
| w1 = write_builder.new_write().with_write_type(['f2']) |
| c = write_builder.new_commit() |
| d0 = pa.Table.from_pydict( |
| {'f0': [1, 2, 3], 'f1': ['a', 'b', 'c']}, |
| schema=pa.schema([('f0', pa.int32()), ('f1', pa.string())])) |
| d1 = pa.Table.from_pydict( |
| {'f2': ['x', 'y', 'z']}, |
| schema=pa.schema([('f2', pa.string())])) |
| w0.write_arrow(d0) |
| w1.write_arrow(d1) |
| cmts = w0.prepare_commit() + w1.prepare_commit() |
| for msg in cmts: |
| for nf in msg.new_files: |
| nf.first_row_id = 0 |
| c.commit(cmts) |
| w0.close() |
| w1.close() |
| c.close() |
| |
| read_builder = table.new_read_builder() |
| actual = read_builder.new_read().to_arrow(read_builder.new_scan().plan().splits()) |
| expected = pa.Table.from_pydict({ |
| 'f0': [1, 2, 3], |
| 'f1': ['a', 'b', 'c'], |
| 'f2': ['x', 'y', 'z'], |
| }, schema=pa_schema) |
| self.assertEqual(actual, expected) |
| |
| @unittest.skipIf(sys.version_info < (3, 11), "vortex-data requires Python >= 3.11") |
| def test_vortex_row_id_filter(self): |
| """Test that Vortex row_indices pushdown works via file_reader_supplier row_ranges.""" |
| pa_schema = pa.schema([ |
| ('f0', pa.int32()), |
| ('f1', pa.string()), |
| ]) |
| schema = Schema.from_pyarrow_schema( |
| pa_schema, |
| options={ |
| 'row-tracking.enabled': 'true', |
| 'data-evolution.enabled': 'true', |
| 'file.format': 'vortex', |
| } |
| ) |
| self.catalog.create_table('default.test_vortex_row_id_filter', schema, False) |
| table = self.catalog.get_table('default.test_vortex_row_id_filter') |
| |
| write_builder = table.new_batch_write_builder() |
| |
| # Commit 1: rows 0-4 |
| w = write_builder.new_write() |
| c = write_builder.new_commit() |
| w.write_arrow(pa.Table.from_pydict( |
| {'f0': list(range(5)), 'f1': [f'v{i}' for i in range(5)]}, |
| schema=pa_schema)) |
| c.commit(w.prepare_commit()) |
| w.close() |
| c.close() |
| |
| # Commit 2: rows 5-9 |
| w = write_builder.new_write() |
| c = write_builder.new_commit() |
| w.write_arrow(pa.Table.from_pydict( |
| {'f0': list(range(5, 10)), 'f1': [f'v{i}' for i in range(5, 10)]}, |
| schema=pa_schema)) |
| c.commit(w.prepare_commit()) |
| w.close() |
| c.close() |
| |
| # Full read |
| rb = table.new_read_builder() |
| full = rb.new_read().to_arrow(rb.new_scan().plan().splits()) |
| self.assertEqual(full.num_rows, 10) |
| |
| # Filter by _ROW_ID using predicate — triggers row_ranges pushdown in Vortex |
| rb_with_rowid = table.new_read_builder().with_projection(['f0', 'f1', '_ROW_ID']) |
| pb = rb_with_rowid.new_predicate_builder() |
| rb_filtered = table.new_read_builder().with_filter(pb.equal('_ROW_ID', 3)) |
| filtered = rb_filtered.new_read().to_arrow(rb_filtered.new_scan().plan().splits()) |
| self.assertEqual(filtered.num_rows, 1) |
| self.assertEqual(filtered.column('f0')[0].as_py(), 3) |
| self.assertEqual(filtered.column('f1')[0].as_py(), 'v3') |
| |
| # Filter by _ROW_ID range spanning two files |
| rb_range = table.new_read_builder().with_filter(pb.between('_ROW_ID', 3, 6)) |
| range_result = rb_range.new_read().to_arrow(rb_range.new_scan().plan().splits()) |
| self.assertEqual(range_result.num_rows, 4) |
| self.assertEqual(sorted(range_result.column('f0').to_pylist()), [3, 4, 5, 6]) |
| |
| @unittest.skipIf(sys.version_info < (3, 11), "vortex-data requires Python >= 3.11") |
| def test_vortex_with_slice(self): |
| """Test with_slice on Vortex data evolution table.""" |
| pa_schema = pa.schema([ |
| ('id', pa.int64()), |
| ('val', pa.int32()), |
| ]) |
| schema = Schema.from_pyarrow_schema( |
| pa_schema, |
| options={ |
| 'row-tracking.enabled': 'true', |
| 'data-evolution.enabled': 'true', |
| 'file.format': 'vortex', |
| 'source.split.target-size': '512m', |
| } |
| ) |
| self.catalog.create_table('default.test_vortex_with_slice', schema, False) |
| table = self.catalog.get_table('default.test_vortex_with_slice') |
| |
| for batch in [ |
| {'id': [1, 2], 'val': [10, 20]}, |
| {'id': [3, 4], 'val': [30, 40]}, |
| {'id': [5, 6], 'val': [50, 60]}, |
| ]: |
| wb = table.new_batch_write_builder() |
| tw = wb.new_write() |
| tc = wb.new_commit() |
| tw.write_arrow(pa.Table.from_pydict(batch, schema=pa_schema)) |
| tc.commit(tw.prepare_commit()) |
| tw.close() |
| tc.close() |
| |
| rb = table.new_read_builder() |
| |
| # Full read |
| full = rb.new_read().to_arrow(rb.new_scan().plan().splits()) |
| self.assertEqual(full.num_rows, 6) |
| |
| # with_slice(1, 4) -> rows at index 1,2,3 -> id in (2,3,4) |
| scan = rb.new_scan().with_slice(1, 4) |
| sliced = rb.new_read().to_arrow(scan.plan().splits()) |
| self.assertEqual(sliced.num_rows, 3) |
| self.assertEqual(sorted(sliced.column('id').to_pylist()), [2, 3, 4]) |
| |
| def test_large_file_read(self): |
| pa_schema = pa.schema([ |
| ('id', pa.int32()), |
| ('name', pa.string()), |
| ]) |
| schema = Schema.from_pyarrow_schema(pa_schema, options={ |
| 'row-tracking.enabled': 'true', |
| 'data-evolution.enabled': 'true', |
| }) |
| self.catalog.create_table('default.test_large_file_row_id', schema, False) |
| table = self.catalog.get_table('default.test_large_file_row_id') |
| |
| # Write >1024 rows in a single file |
| num_rows = 2000 |
| data = pa.Table.from_pydict({ |
| 'id': list(range(num_rows)), |
| 'name': [f'name_{i}' for i in range(num_rows)], |
| }, schema=pa_schema) |
| |
| wb = table.new_batch_write_builder() |
| tw = wb.new_write() |
| tc = wb.new_commit() |
| tw.write_arrow(data) |
| tc.commit(tw.prepare_commit()) |
| tw.close() |
| tc.close() |
| |
| update_ids = list(range(0, 1500)) |
| upsert_data = pa.Table.from_pydict({ |
| 'id': update_ids, |
| 'name': [f'upsert_{i}' for i in update_ids], |
| }, schema=pa_schema) |
| |
| wb = table.new_batch_write_builder() |
| tu = wb.new_update() |
| msgs = tu.upsert_by_arrow_with_key(upsert_data, ['id']) |
| tc = wb.new_commit() |
| tc.commit(msgs) |
| tc.close() |
