| # 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 tempfile |
| import unittest |
| |
| import pyarrow as pa |
| |
| from pypaimon import CatalogFactory, Schema |
| |
| |
| def _has_native_planner(): |
| try: |
| from pypaimon_rust.datafusion import PaimonCatalog, Split |
| except Exception: |
| return False |
| return hasattr(PaimonCatalog, 'get_table') and hasattr(Split, 'serialize') |
| |
| |
| @unittest.skipUnless(_has_native_planner(), |
| "pypaimon_rust with split-planning API not installed") |
| class NativePlanIntegrationTest(unittest.TestCase): |
| """Live round-trip guarding the cross-language SplitSerializer against drift: |
| plan via pypaimon_rust, decode, and require the same rows as the normal plan. |
| The golden unit tests only prove self-consistency; this proves byte-compat |
| with the real producer.""" |
| |
| def setUp(self): |
| self.cat = CatalogFactory.create({'warehouse': tempfile.mkdtemp(prefix='np_it_')}) |
| self.cat.create_database('default', True) |
| self.schema = pa.schema([('k', pa.int64()), ('v', pa.string())]) |
| |
| def _write(self, name, rows): |
| t = self.cat.get_table('default.%s' % name) |
| wb = t.new_batch_write_builder() |
| w, c = wb.new_write(), wb.new_commit() |
| w.write_arrow(pa.Table.from_pylist(rows, schema=self.schema)) |
| c.commit(w.prepare_commit()) |
| w.close() |
| c.close() |
| |
| def _plan_and_read(self, name, native): |
| t = self.cat.get_table('default.%s' % name) |
| if native: |
| t = t.copy({'scan.native-plan.enabled': 'true'}) |
| rb = t.new_read_builder() |
| plan = rb.new_scan().plan() |
| rows = rb.new_read().to_arrow(plan.splits()).to_pylist() |
| return plan.snapshot_id, sorted(rows, key=lambda r: r['k']) |
| |
| def _assert_matches(self, name, expect_native=True): |
| sid_n, rows_n = self._plan_and_read(name, native=False) |
| sid_r, rows_r = self._plan_and_read(name, native=True) |
| self.assertEqual(rows_r, rows_n) |
| self.assertEqual(sid_r, sid_n) # snapshot id preserved through native plan |
| self.assertIsNotNone(sid_r) |
| # Guard against a false green where native silently fell back to Python: assert the |
| # native planner was (or was not) actually used, as expected for this table. |
| native_table = self.cat.get_table('default.%s' % name).copy( |
| {'scan.native-plan.enabled': 'true'}) |
| self.assertEqual( |
| native_table.new_read_builder().explain().native_planned, expect_native) |
| |
| def test_primary_key_matches_normal_plan(self): |
| self.cat.create_table('default.pk_t', Schema.from_pyarrow_schema( |
| self.schema, primary_keys=['k'], options={'bucket': '1'}), False) |
| self._write('pk_t', [{'k': 1, 'v': 'a1'}, {'k': 2, 'v': 'b1'}]) |
| self._write('pk_t', [{'k': 2, 'v': 'b2'}, {'k': 3, 'v': 'c1'}]) # k=2 updated |
| self._assert_matches('pk_t') |
| |
| def test_pk_equal_to_partition_key_falls_back(self): |
| # Empty trimmed PK: native would skip merge and return duplicates -> must fall back. |
| self.cat.create_table('default.pkpart_t', Schema.from_pyarrow_schema( |
| self.schema, partition_keys=['k'], primary_keys=['k'], options={'bucket': '1'}), False) |
| self._write('pkpart_t', [{'k': 1, 'v': 'a1'}, {'k': 2, 'v': 'b1'}]) |
| self._write('pkpart_t', [{'k': 2, 'v': 'b2'}]) |
| native_table = self.cat.get_table('default.pkpart_t').copy( |
| {'scan.native-plan.enabled': 'true'}) |
| self.assertFalse(native_table.new_read_builder().explain().native_planned) |
| with self.assertRaises(ValueError): |
| rb = native_table.new_read_builder() |
| rb.new_read().to_arrow(rb.new_scan().plan().splits()) |
| |
| def test_copy_removed_persisted_scan_option_falls_back(self): |
| # copy() removes a persisted scan.snapshot-id that Rust would still reload -> fall back. |
| self.cat.create_table('default.snapopt_t', Schema.from_pyarrow_schema( |
| self.schema, options={'scan.snapshot-id': '1'}), False) |
| self._write('snapopt_t', [{'k': 1, 'v': 'a'}]) # snapshot 1 |
| self._write('snapopt_t', [{'k': 2, 'v': 'b'}]) # snapshot 2 |
| native = self.cat.get_table('default.snapopt_t').copy( |
| {'scan.snapshot-id': None, 'scan.native-plan.enabled': 'true'}) |
| self.assertFalse(native.new_read_builder().explain().native_planned) |
| |
| def test_first_row_merge_engine_falls_back(self): |
| self.cat.create_table('default.fr_t', Schema.from_pyarrow_schema( |
| self.schema, primary_keys=['k'], |
| options={'bucket': '1', 'merge-engine': 'first-row'}), False) |
| self._write('fr_t', [{'k': 1, 'v': 'a'}, {'k': 2, 'v': 'b'}]) |
| self._write('fr_t', [{'k': 1, 'v': 'X'}, {'k': 3, 'v': 'c'}]) # k=1 stays 'a' |
| self._assert_matches('fr_t', expect_native=False) |
| |
| def test_append_matches_normal_plan(self): |
| self.cat.create_table( |
| 'default.ap_t', Schema.from_pyarrow_schema(self.schema), False) |
| self._write('ap_t', [{'k': 1, 'v': 'a'}, {'k': 2, 'v': 'b'}]) |
| self._write('ap_t', [{'k': 3, 'v': 'c'}]) |
| self._assert_matches('ap_t') |
| |
| def test_dynamic_split_target_size_matches_normal_plan(self): |
| self.cat.create_table( |
| 'default.split_t', Schema.from_pyarrow_schema(self.schema), False) |
| self._write('split_t', [{'k': 1, 'v': 'a'}]) |
| self._write('split_t', [{'k': 2, 'v': 'b'}]) |
| options = {'source.split.target-size': '1b'} |
| normal_table = self.cat.get_table('default.split_t').copy(options) |
| native_table = normal_table.copy({'scan.native-plan.enabled': 'true'}) |
| |
| normal = normal_table.new_read_builder().new_scan().plan() |
| native = native_table.new_read_builder().explain() |
| |
| self.assertTrue(native.native_planned) |
| self.assertEqual(native.split_count, len(normal.splits())) |
| self.assertGreater(native.split_count, 1) |
| |
| def test_dynamic_split_open_file_cost_matches_normal_plan(self): |
| stored_options = { |
| 'source.split.target-size': '128mb', |
| 'source.split.open-file-cost': '1b', |
| } |
| self.cat.create_table('default.open_cost_t', Schema.from_pyarrow_schema( |
| self.schema, options=stored_options), False) |
| self._write('open_cost_t', [{'k': 1, 'v': 'a'}]) |
| self._write('open_cost_t', [{'k': 2, 'v': 'b'}]) |
| self._write('open_cost_t', [{'k': 3, 'v': 'c'}]) |
| base_table = self.cat.get_table('default.open_cost_t') |
| normal_table = base_table.copy({'source.split.open-file-cost': '64mb'}) |
| native_table = normal_table.copy({'scan.native-plan.enabled': 'true'}) |
| |
| baseline = base_table.new_read_builder().new_scan().plan() |
| normal = normal_table.new_read_builder().new_scan().plan() |
| native = native_table.new_read_builder().explain() |
| |
| self.assertEqual(len(baseline.splits()), 1) |
| self.assertGreater(len(normal.splits()), len(baseline.splits())) |
| self.assertTrue(native.native_planned) |
| self.assertEqual(native.split_count, len(normal.splits())) |
| |
| def test_dynamic_split_option_reset_matches_normal_plan(self): |
| stored_options = { |
| 'source.split.target-size': '1b', |
| 'source.split.open-file-cost': '1b', |
| } |
| self.cat.create_table('default.split_reset_t', Schema.from_pyarrow_schema( |
| self.schema, options=stored_options), False) |
| self._write('split_reset_t', [{'k': 1, 'v': 'a'}]) |
| self._write('split_reset_t', [{'k': 2, 'v': 'b'}]) |
| reset_options = { |
| 'source.split.target-size': None, |
| 'source.split.open-file-cost': None, |
| } |
| normal_table = self.cat.get_table('default.split_reset_t').copy(reset_options) |
| native_table = normal_table.copy({'scan.native-plan.enabled': 'true'}) |
| |
| normal = normal_table.new_read_builder().new_scan().plan() |
| native = native_table.new_read_builder().explain() |
| |
| self.assertTrue(native.native_planned) |
| self.assertEqual(native.split_count, len(normal.splits())) |
| self.assertEqual(native.split_count, 1) |
| |
| def test_partitioned_table_falls_back(self): |
| # Rust bucket_path can diverge from the writer's str(value) partition dir -> fall back. |
| schema = pa.schema([('k', pa.int64()), ('p', pa.string())]) |
| self.cat.create_table('default.pt_t', Schema.from_pyarrow_schema( |
| schema, partition_keys=['p']), False) |
| t = self.cat.get_table('default.pt_t') |
| wb = t.new_batch_write_builder() |
| w, c = wb.new_write(), wb.new_commit() |
| w.write_arrow(pa.Table.from_pylist( |
| [{'k': 1, 'p': 'a'}, {'k': 2, 'p': 'a'}, {'k': 3, 'p': 'b'}], schema=schema)) |
| c.commit(w.prepare_commit()) |
| w.close() |
| c.close() |
| |
| native_table = self.cat.get_table('default.pt_t').copy( |
| {'scan.native-plan.enabled': 'true'}) |
| self.assertFalse(native_table.new_read_builder().explain().native_planned) |
| |
| def test_explain_reflects_native_plan(self): |
| self.cat.create_table( |
| 'default.ex_t', Schema.from_pyarrow_schema(self.schema), False) |
| self._write('ex_t', [{'k': 1, 'v': 'a'}, {'k': 2, 'v': 'b'}]) |
| normal = self.cat.get_table('default.ex_t').new_read_builder().explain() |
| native = self.cat.get_table('default.ex_t').copy( |
| {'scan.native-plan.enabled': 'true'}).new_read_builder().explain() |
| self.assertFalse(normal.native_planned) |
| self.assertTrue(native.native_planned) |
| self.assertEqual(native.split_count, normal.split_count) |
| self.assertEqual(native.snapshot_id, normal.snapshot_id) |
| self.assertIn('native', str(native)) # render shows the Planner line |
| |
| def test_empty_table_explain_reflects_python_fallback(self): |
| self.cat.create_table( |
| 'default.empty_t', Schema.from_pyarrow_schema(self.schema), False) |
| normal = self.cat.get_table('default.empty_t').new_read_builder().explain() |
| fallback = self.cat.get_table('default.empty_t').copy( |
| {'scan.native-plan.enabled': 'true'}).new_read_builder().explain() |
| |
| self.assertFalse(fallback.native_planned) |
| self.assertEqual(fallback.snapshot_id, normal.snapshot_id) |
| self.assertEqual(fallback.split_count, 0) |
| self.assertNotIn('Planner:', str(fallback)) |
| |
| |
| if __name__ == '__main__': |
| unittest.main() |
