diff --git a/paimon-python/pypaimon/read/fallback_read_builder.py b/paimon-python/pypaimon/read/fallback_read_builder.py new file mode 100644 index 000000000000..fd998ba9a3f5 --- /dev/null +++ b/paimon-python/pypaimon/read/fallback_read_builder.py @@ -0,0 +1,142 @@ +# 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. + +"""Batch read builder honouring the ``scan.fallback-branch`` table option. + +Mirrors Java's ``org.apache.paimon.table.FallbackReadFileStoreTable``: +partitions that exist in the main branch are read from main; partitions +that exist only in the fallback branch are read from the fallback branch. +All branches of a table share the same physical data files (only snapshot, +tag and schema metadata are branch scoped), and the fallback branch must +have the same row type as main, so a split planned on either branch can be +read with a single :class:`~pypaimon.read.table_read.TableRead`. +""" + +from typing import Dict, List, TYPE_CHECKING, Union + +from pypaimon.read.plan import Plan +from pypaimon.read.read_builder import ReadBuilder + + +if TYPE_CHECKING: + from pypaimon.read.split import Split + from pypaimon.read.table_read import TableRead + from pypaimon.read.table_scan import TableScan + from pypaimon.table.file_store_table import FileStoreTable + + +def _partition_key(split: "Split") -> tuple: + """Return a hashable, branch-independent key for a split's partition.""" + partition = split.partition + values = getattr(partition, "values", None) + if values is None: + return (partition,) + return tuple(values) + + +class FallbackReadBuilder(ReadBuilder): + """A :class:`ReadBuilder` that merges main and fallback-branch partitions. + + Configuration (filter, projection, read type, limit) is applied to both + the main and the fallback builders so the fallback scan reads the same + columns and rows as main. + """ + + def __init__(self, table: "FileStoreTable", fallback_branch: str): + super().__init__(table) + self._fallback_table = table.copy({"branch": fallback_branch}) + self._validate_same_row_type(table, self._fallback_table, fallback_branch) + self._fallback_builder = ReadBuilder(self._fallback_table) + + @staticmethod + def _validate_same_row_type(main_table, fallback_table, fallback_branch): + main_fields = main_table.table_schema.fields + fallback_fields = fallback_table.table_schema.fields + same = len(main_fields) == len(fallback_fields) and all( + m.name == f.name and str(m.type) == str(f.type) + for m, f in zip(main_fields, fallback_fields) + ) + if not same: + raise ValueError( + "The fallback branch must have the same schema as the main " + "branch. This validation is triggered because " + "'scan.fallback-branch' is configured to '{}'.".format( + fallback_branch)) + + # -- configuration is mirrored onto the fallback builder --------------- + + def with_filter(self, predicate) -> "FallbackReadBuilder": + super().with_filter(predicate) + self._fallback_builder.with_filter(predicate) + return self + + def with_partition_filter(self, partition_filter) -> "FallbackReadBuilder": + super().with_partition_filter(partition_filter) + self._fallback_builder.with_partition_filter(partition_filter) + return self + + def with_projection( + self, projection: Union[List[str], Dict[str, str]] + ) -> "FallbackReadBuilder": + super().with_projection(projection) + self._fallback_builder.with_projection(projection) + return self + + def with_read_type(self, read_type) -> "FallbackReadBuilder": + super().with_read_type(read_type) + self._fallback_builder.with_read_type(read_type) + return self + + def with_limit(self, limit: int) -> "FallbackReadBuilder": + super().with_limit(limit) + self._fallback_builder.with_limit(limit) + return self + + # -- planning merges the two branches ---------------------------------- + + def new_scan(self) -> "TableScan": + return _FallbackTableScan( + super().new_scan(), self._fallback_builder.new_scan()) + + def new_read(self) -> "TableRead": + return super().new_read() + + +class _FallbackTableScan: + """Wraps the main and fallback scans and merges their plans by partition.""" + + def __init__(self, main_scan: "TableScan", fallback_scan: "TableScan"): + self._main_scan = main_scan + self._fallback_scan = fallback_scan + + def plan(self) -> Plan: + main_plan = self._main_scan.plan() + main_splits = list(main_plan.splits()) + main_partitions = {_partition_key(split) for split in main_splits} + + fallback_splits = [ + split + for split in self._fallback_scan.plan().splits() + if _partition_key(split) not in main_partitions + ] + + return Plan(_splits=main_splits + fallback_splits, + snapshot_id=main_plan.snapshot_id) + + def __getattr__(self, name): + # Forward any other scan API (shard/slice/stats/...) to the main scan. + return getattr(self._main_scan, name) diff --git a/paimon-python/pypaimon/table/file_store_table.py b/paimon-python/pypaimon/table/file_store_table.py index 30111c91e95f..538991e55224 100644 --- a/paimon-python/pypaimon/table/file_store_table.py +++ b/paimon-python/pypaimon/table/file_store_table.py @@ -417,6 +417,10 @@ def bucket_mode(self) -> BucketMode: return BucketMode.HASH_FIXED def new_read_builder(self) -> 'ReadBuilder': + fallback_branch = self.options.scan_fallback_branch() + if fallback_branch and fallback_branch != self.current_branch(): + from pypaimon.read.fallback_read_builder import FallbackReadBuilder + return FallbackReadBuilder(self, fallback_branch) return ReadBuilder(self) def new_stream_read_builder(self) -> 'StreamReadBuilder': diff --git a/paimon-python/pypaimon/tests/batch_read_fallback_branch_test.py b/paimon-python/pypaimon/tests/batch_read_fallback_branch_test.py new file mode 100644 index 000000000000..008ba0eb36e1 --- /dev/null +++ b/paimon-python/pypaimon/tests/batch_read_fallback_branch_test.py @@ -0,0 +1,111 @@ +# 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. + +"""Batch-read coverage for the ``scan.fallback-branch`` table option. + +Mirrors Java's ``FallbackReadFileStoreTable``: partitions that exist in +the main branch are read from main; partitions that exist only in the +fallback branch are read from the fallback branch. +""" + +import os +import shutil +import tempfile +import unittest + +import pyarrow as pa + +from pypaimon import CatalogFactory, Schema +from pypaimon.schema.schema_change import SchemaChange + + +class BatchReadFallbackBranchTest(unittest.TestCase): + + def setUp(self): + self.temp_dir = tempfile.mkdtemp(prefix="unittest_fallback_branch_") + warehouse = os.path.join(self.temp_dir, "warehouse") + os.makedirs(warehouse, exist_ok=True) + self.catalog = CatalogFactory.create({"warehouse": warehouse}) + self.catalog.create_database("db", True) + self.pa_schema = pa.schema([ + ("k", pa.int64()), + ("p", pa.string()), + ]) + self.catalog.create_table( + "db.t", + Schema.from_pyarrow_schema(self.pa_schema, partition_keys=["p"]), + False, + ) + + def tearDown(self): + shutil.rmtree(self.temp_dir, ignore_errors=True) + + def _write(self, table_name, rows): + wb = self.catalog.get_table(table_name).new_batch_write_builder() + writer = wb.new_write() + try: + writer.write_arrow(pa.table(rows, schema=self.pa_schema)) + wb.new_commit().commit(writer.prepare_commit()) + finally: + writer.close() + + def _read(self, table_name): + rb = self.catalog.get_table(table_name).new_read_builder() + rows = rb.new_read().to_arrow(rb.new_scan().plan().splits()).to_pylist() + return sorted(rows, key=lambda row: (row["p"], row["k"])) + + def _set_fallback(self, branch): + self.catalog.alter_table( + "db.t", [SchemaChange.set_option("scan.fallback-branch", branch)]) + + def test_batch_read_includes_fallback_only_partitions(self): + """Partition that is missing from main is read from the fallback branch.""" + self.catalog.create_branch("db.t", "b") + self._write("db.t", {"k": [1], "p": ["a"]}) + self._write("db.t$branch_b", {"k": [2], "p": ["b"]}) + self._set_fallback("b") + + self.assertEqual( + [{"k": 1, "p": "a"}, {"k": 2, "p": "b"}], + self._read("db.t"), + ) + + def test_main_partition_shadows_fallback(self): + """A partition present in main is read from main, not the fallback.""" + self.catalog.create_branch("db.t", "b") + self._write("db.t", {"k": [1], "p": ["a"]}) + # Same partition "a" exists on the fallback branch with a different row. + self._write("db.t$branch_b", {"k": [99], "p": ["a"]}) + self._write("db.t$branch_b", {"k": [2], "p": ["b"]}) + self._set_fallback("b") + + self.assertEqual( + [{"k": 1, "p": "a"}, {"k": 2, "p": "b"}], + self._read("db.t"), + ) + + def test_no_fallback_option_reads_main_only(self): + """Without the option the read stays on the main branch.""" + self.catalog.create_branch("db.t", "b") + self._write("db.t", {"k": [1], "p": ["a"]}) + self._write("db.t$branch_b", {"k": [2], "p": ["b"]}) + + self.assertEqual([{"k": 1, "p": "a"}], self._read("db.t")) + + +if __name__ == "__main__": + unittest.main()