[python] Read scan.fallback-branch partitions in batch reads

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)

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':

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()