[Python] Enable field merge read in row-tracking table (#6399)

Author: discivigour

paimon-python/pypaimon/manifest/manifest_file_manager.py

@@ -60,9 +60,13 @@ def read(self, manifest_file_name: str, bucket_filter=None) -> List[ManifestEntr
                 null_counts=key_dict['_NULL_COUNTS'],
             )
             value_dict = dict(file_dict['_VALUE_STATS'])
-            if file_dict.get('_VALUE_STATS_COLS') is None:
-                fields = self.table.table_schema.fields
-            elif not file_dict.get('_VALUE_STATS_COLS'):
+            if file_dict['_VALUE_STATS_COLS'] is None:
+                if file_dict['_WRITE_COLS'] is None:
+                    fields = self.table.table_schema.fields
+                else:
+                    read_fields = file_dict['_WRITE_COLS']
+                    fields = [self.table.field_dict[col] for col in read_fields]
+            elif not file_dict['_VALUE_STATS_COLS']:
                 fields = []
             else:
                 fields = [self.table.field_dict[col] for col in file_dict['_VALUE_STATS_COLS']]
@@ -89,6 +93,9 @@ def read(self, manifest_file_name: str, bucket_filter=None) -> List[ManifestEntr
                 embedded_index=file_dict['_EMBEDDED_FILE_INDEX'],
                 file_source=file_dict['_FILE_SOURCE'],
                 value_stats_cols=file_dict.get('_VALUE_STATS_COLS'),
+                external_path=file_dict.get('_EXTERNAL_PATH'),
+                first_row_id=file_dict['_FIRST_ROW_ID'],
+                write_cols=file_dict['_WRITE_COLS'],
             )
             entry = ManifestEntry(
                 kind=record['_KIND'],
@@ -137,6 +144,9 @@ def write(self, file_name, entries: List[ManifestEntry]):
                     "_EMBEDDED_FILE_INDEX": entry.file.embedded_index,
                     "_FILE_SOURCE": entry.file.file_source,
                     "_VALUE_STATS_COLS": entry.file.value_stats_cols,
+                    "_EXTERNAL_PATH": entry.file.external_path,
+                    "_FIRST_ROW_ID": entry.file.first_row_id,
+                    "_WRITE_COLS": entry.file.write_cols,
                 }
             }
             avro_records.append(avro_record)

paimon-python/pypaimon/manifest/schema/data_file_meta.py

@@ -47,6 +47,8 @@ class DataFileMeta:
     file_source: Optional[str] = None
     value_stats_cols: Optional[List[str]] = None
     external_path: Optional[str] = None
+    first_row_id: Optional[int] = None
+    write_cols: Optional[List[str]] = None
 
     # not a schema field, just for internal usage
     file_path: str = None
@@ -59,6 +61,58 @@ def set_file_path(self, table_path: Path, partition: GenericRow, bucket: int):
         path_builder = path_builder / ("bucket-" + str(bucket)) / self.file_name
         self.file_path = str(path_builder)
 
+    def assign_first_row_id(self, first_row_id: int) -> 'DataFileMeta':
+        """Create a new DataFileMeta with the assigned first_row_id."""
+        return DataFileMeta(
+            file_name=self.file_name,
+            file_size=self.file_size,
+            row_count=self.row_count,
+            min_key=self.min_key,
+            max_key=self.max_key,
+            key_stats=self.key_stats,
+            value_stats=self.value_stats,
+            min_sequence_number=self.min_sequence_number,
+            max_sequence_number=self.max_sequence_number,
+            schema_id=self.schema_id,
+            level=self.level,
+            extra_files=self.extra_files,
+            creation_time=self.creation_time,
+            delete_row_count=self.delete_row_count,
+            embedded_index=self.embedded_index,
+            file_source=self.file_source,
+            value_stats_cols=self.value_stats_cols,
+            external_path=self.external_path,
+            first_row_id=first_row_id,
+            write_cols=self.write_cols,
+            file_path=self.file_path
+        )
+
+    def assign_sequence_number(self, min_sequence_number: int, max_sequence_number: int) -> 'DataFileMeta':
+        """Create a new DataFileMeta with the assigned sequence numbers."""
+        return DataFileMeta(
+            file_name=self.file_name,
+            file_size=self.file_size,
+            row_count=self.row_count,
+            min_key=self.min_key,
+            max_key=self.max_key,
+            key_stats=self.key_stats,
+            value_stats=self.value_stats,
+            min_sequence_number=min_sequence_number,
+            max_sequence_number=max_sequence_number,
+            schema_id=self.schema_id,
+            level=self.level,
+            extra_files=self.extra_files,
+            creation_time=self.creation_time,
+            delete_row_count=self.delete_row_count,
+            embedded_index=self.embedded_index,
+            file_source=self.file_source,
+            value_stats_cols=self.value_stats_cols,
+            external_path=self.external_path,
+            first_row_id=self.first_row_id,
+            write_cols=self.write_cols,
+            file_path=self.file_path
+        )
+
 
 DATA_FILE_META_SCHEMA = {
     "type": "record",
@@ -83,9 +137,14 @@ def set_file_path(self, table_path: Path, partition: GenericRow, bucket: int):
          "default": None},
         {"name": "_DELETE_ROW_COUNT", "type": ["null", "long"], "default": None},
         {"name": "_EMBEDDED_FILE_INDEX", "type": ["null", "bytes"], "default": None},
-        {"name": "_FILE_SOURCE", "type": ["null", "int"], "default": None},
+        {"name": "_FILE_SOURCE", "type": ["null", "string"], "default": None},
         {"name": "_VALUE_STATS_COLS",
          "type": ["null", {"type": "array", "items": "string"}],
          "default": None},
+        {"name": "_EXTERNAL_PATH", "type": ["null", "string"], "default": None},
+        {"name": "_FIRST_ROW_ID", "type": ["null", "long"], "default": None},
+        {"name": "_WRITE_COLS",
+         "type": ["null", {"type": "array", "items": "string"}],
+         "default": None},
     ]
 }

paimon-python/pypaimon/manifest/schema/manifest_entry.py

@@ -31,6 +31,26 @@ class ManifestEntry:
     total_buckets: int
     file: DataFileMeta
 
+    def assign_first_row_id(self, first_row_id: int) -> 'ManifestEntry':
+        """Create a new ManifestEntry with the assigned first_row_id."""
+        return ManifestEntry(
+            kind=self.kind,
+            partition=self.partition,
+            bucket=self.bucket,
+            total_buckets=self.total_buckets,
+            file=self.file.assign_first_row_id(first_row_id)
+        )
+
+    def assign_sequence_number(self, min_sequence_number: int, max_sequence_number: int) -> 'ManifestEntry':
+        """Create a new ManifestEntry with the assigned sequence numbers."""
+        return ManifestEntry(
+            kind=self.kind,
+            partition=self.partition,
+            bucket=self.bucket,
+            total_buckets=self.total_buckets,
+            file=self.file.assign_sequence_number(min_sequence_number, max_sequence_number)
+        )
+
 
 MANIFEST_ENTRY_SCHEMA = {
     "type": "record",

paimon-python/pypaimon/read/reader/data_evolution_merge_reader.py

@@ -0,0 +1,85 @@
+################################################################################
+#  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.
+################################################################################
+from typing import List, Optional
+
+import pyarrow as pa
+from pyarrow import RecordBatch
+
+from pypaimon.read.reader.iface.record_batch_reader import RecordBatchReader
+
+
+class DataEvolutionMergeReader(RecordBatchReader):
+    """
+    This is a union reader which contains multiple inner readers, Each reader is responsible for reading one file.
+
+    This reader, assembling multiple reader into one big and great reader, will merge the batches from all readers.
+
+    For example, if rowOffsets is {0, 2, 0, 1, 2, 1} and fieldOffsets is {0, 0, 1, 1, 1, 0}, it means:
+     - The first field comes from batch0, and it is at offset 0 in batch0.
+     - The second field comes from batch2, and it is at offset 0 in batch2.
+     - The third field comes from batch0, and it is at offset 1 in batch0.
+     - The fourth field comes from batch1, and it is at offset 1 in batch1.
+     - The fifth field comes from batch2, and it is at offset 1 in batch2.
+     - The sixth field comes from batch1, and it is at offset 0 in batch1.
+    """
+
+    def __init__(self, row_offsets: List[int], field_offsets: List[int], readers: List[Optional[RecordBatchReader]]):
+        if row_offsets is None:
+            raise ValueError("Row offsets must not be null")
+        if field_offsets is None:
+            raise ValueError("Field offsets must not be null")
+        if len(row_offsets) != len(field_offsets):
+            raise ValueError("Row offsets and field offsets must have the same length")
+        if not row_offsets:
+            raise ValueError("Row offsets must not be empty")
+        if not readers or len(readers) < 1:
+            raise ValueError("Readers should be more than 0")
+        self.row_offsets = row_offsets
+        self.field_offsets = field_offsets
+        self.readers = readers
+
+    def read_arrow_batch(self) -> Optional[RecordBatch]:
+        batches: List[Optional[RecordBatch]] = [None] * len(self.readers)
+        for i, reader in enumerate(self.readers):
+            if reader is not None:
+                batch = reader.read_arrow_batch()
+                if batch is None:
+                    # all readers are aligned, as long as one returns null, the others will also have no data
+                    return None
+                batches[i] = batch
+        # Assemble record batches from batches based on row_offsets and field_offsets
+        columns = []
+        names = []
+        for i in range(len(self.row_offsets)):
+            batch_index = self.row_offsets[i]
+            field_index = self.field_offsets[i]
+            if batches[batch_index] is not None:
+                column = batches[batch_index].column(field_index)
+                columns.append(column)
+                names.append(batches[batch_index].schema.names[field_index])
+        if columns:
+            return pa.RecordBatch.from_arrays(columns, names)
+        return None
+
+    def close(self) -> None:
+        try:
+            for reader in self.readers:
+                if reader is not None:
+                    reader.close()
+        except Exception as e:
+            raise IOError("Failed to close inner readers") from e

…n/read/reader/data_file_record_reader.py …on/read/reader/data_file_batch_reader.pypaimon-python/pypaimon/read/reader/data_file_record_reader.py renamed to paimon-python/pypaimon/read/reader/data_file_batch_reader.py paimon-python/pypaimon/read/reader/data_file_record_reader.py renamed to paimon-python/pypaimon/read/reader/data_file_batch_reader.py

@@ -0,0 +1,120 @@
+"""
+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.
+"""
+"""
+FieldBunch classes for organizing files by field in data evolution.
+
+These classes help organize DataFileMeta objects into groups based on their field content,
+supporting both regular data files and blob files.
+"""
+from abc import ABC
+from typing import List
+from pypaimon.manifest.schema.data_file_meta import DataFileMeta
+
+
+class FieldBunch(ABC):
+    """Interface for files organized by field."""
+
+    def row_count(self) -> int:
+        """Return the total row count for this bunch."""
+        ...
+
+    def files(self) -> List[DataFileMeta]:
+        """Return the list of files in this bunch."""
+        ...
+
+
+class DataBunch(FieldBunch):
+    """Files for a single data file."""
+
+    def __init__(self, data_file: DataFileMeta):
+        self.data_file = data_file
+
+    def row_count(self) -> int:
+        return self.data_file.row_count
+
+    def files(self) -> List[DataFileMeta]:
+        return [self.data_file]
+
+
+class BlobBunch(FieldBunch):
+    """Files for partial field (blob files)."""
+
+    def __init__(self, expected_row_count: int):
+        self._files: List[DataFileMeta] = []
+        self.expected_row_count = expected_row_count
+        self.latest_first_row_id = -1
+        self.expected_next_first_row_id = -1
+        self.latest_max_sequence_number = -1
+        self._row_count = 0
+
+    def add(self, file: DataFileMeta) -> None:
+        """Add a blob file to this bunch."""
+        if not self._is_blob_file(file.file_name):
+            raise ValueError("Only blob file can be added to a blob bunch.")
+
+        if file.first_row_id == self.latest_first_row_id:
+            if file.max_sequence_number >= self.latest_max_sequence_number:
+                raise ValueError(
+                    "Blob file with same first row id should have decreasing sequence number."
+                )
+            return
+
+        if self._files:
+            first_row_id = file.first_row_id
+            if first_row_id < self.expected_next_first_row_id:
+                if file.max_sequence_number >= self.latest_max_sequence_number:
+                    raise ValueError(
+                        "Blob file with overlapping row id should have decreasing sequence number."
+                    )
+                return
+            elif first_row_id > self.expected_next_first_row_id:
+                raise ValueError(
+                    f"Blob file first row id should be continuous, expect "
+                    f"{self.expected_next_first_row_id} but got {first_row_id}"
+                )
+
+            if file.schema_id != self._files[0].schema_id:
+                raise ValueError(
+                    "All files in a blob bunch should have the same schema id."
+                )
+            if file.write_cols != self._files[0].write_cols:
+                raise ValueError(
+                    "All files in a blob bunch should have the same write columns."
+                )
+
+        self._files.append(file)
+        self._row_count += file.row_count
+        if self._row_count > self.expected_row_count:
+            raise ValueError(
+                f"Blob files row count exceed the expect {self.expected_row_count}"
+            )
+
+        self.latest_max_sequence_number = file.max_sequence_number
+        self.latest_first_row_id = file.first_row_id
+        self.expected_next_first_row_id = self.latest_first_row_id + file.row_count
+
+    def row_count(self) -> int:
+        return self._row_count
+
+    def files(self) -> List[DataFileMeta]:
+        return self._files
+
+    @staticmethod
+    def _is_blob_file(file_name: str) -> bool:
+        """Check if a file is a blob file based on its extension."""
+        return file_name.endswith('.blob')