From 29650176338983333863e4d5e6481ca8d934b13f Mon Sep 17 00:00:00 2001 From: henriquejsza Date: Fri, 30 Oct 2026 10:32:29 -0300 Subject: [PATCH] [python] Read scan.fallback-branch partitions in batch reads Batch reads in pypaimon only rejected scan.fallback-branch as a catalog option and never read it from the table, so a read left out the partitions that exist only in the fallback branch. Spark returns those rows, as the issue reports. FileStoreTable.new_read_builder now returns a FallbackReadBuilder when the table option scan.fallback-branch is set and the current branch is not the fallback branch itself. Mirroring Java's FallbackReadFileStoreTable, partitions present in the main branch are read from main and partitions that exist only in the fallback branch are read from the fallback branch. Filter, projection, read type and limit are applied to both branch scans, and the fallback branch must have the same row type as main. All branches share the same physical data files, so the merged splits are read with a single TableRead. Closes #10884 Signed-off-by: henriquejsza --- .../pypaimon/read/fallback_read_builder.py | 142 ++++++++++++++++++ .../pypaimon/table/file_store_table.py | 4 + .../tests/batch_read_fallback_branch_test.py | 111 ++++++++++++++ 3 files changed, 257 insertions(+) create mode 100644 paimon-python/pypaimon/read/fallback_read_builder.py create mode 100644 paimon-python/pypaimon/tests/batch_read_fallback_branch_test.py 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()