[python] Fix with_shard feature for blob data (#6691)

Author: discivigour

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

@@ -75,6 +75,9 @@ def read_arrow_batch(self) -> Optional[RecordBatch]:
                                    min(self.split_end_row, self.cur_end) - self.split_start_row)
             elif cur_begin < self.split_end_row <= self.cur_end:
                 return batch.slice(0, self.split_end_row - cur_begin)
+            else:
+                # return empty RecordBatch if the batch size has not reached split_start_row
+                return pa.RecordBatch.from_arrays([], [])
         else:
             return batch
 

paimon-python/pypaimon/read/scanner/full_starting_scanner.py

@@ -117,14 +117,20 @@ def _append_only_filter_by_shard(self, partitioned_files: defaultdict) -> (defau
             for entry in file_entries:
                 total_row += entry.file.row_count
 
-        # Calculate number of rows this shard should process
-        # Last shard handles all remaining rows (handles non-divisible cases)
-        if self.idx_of_this_subtask == self.number_of_para_subtasks - 1:
-            num_row = total_row - total_row // self.number_of_para_subtasks * self.idx_of_this_subtask
+        # Calculate number of rows this shard should process using balanced distribution
+        # Distribute remainder evenly among first few shards to avoid last shard overload
+        base_rows_per_shard = total_row // self.number_of_para_subtasks
+        remainder = total_row % self.number_of_para_subtasks
+
+        # Each of the first 'remainder' shards gets one extra row
+        if self.idx_of_this_subtask < remainder:
+            num_row = base_rows_per_shard + 1
+            start_row = self.idx_of_this_subtask * (base_rows_per_shard + 1)
         else:
-            num_row = total_row // self.number_of_para_subtasks
-        # Calculate start row and end row position for current shard in all data
-        start_row = self.idx_of_this_subtask * (total_row // self.number_of_para_subtasks)
+            num_row = base_rows_per_shard
+            start_row = (remainder * (base_rows_per_shard + 1) +
+                         (self.idx_of_this_subtask - remainder) * base_rows_per_shard)
+
         end_row = start_row + num_row
 
         plan_start_row = 0
@@ -160,24 +166,25 @@ def _append_only_filter_by_shard(self, partitioned_files: defaultdict) -> (defau
 
     def _data_evolution_filter_by_shard(self, partitioned_files: defaultdict) -> (defaultdict, int, int):
         total_row = 0
-        first_row_id_set = set()
-        # Sort by file creation time to ensure consistent sharding
         for key, file_entries in partitioned_files.items():
             for entry in file_entries:
-                if entry.file.first_row_id is None:
-                    total_row += entry.file.row_count
-                elif entry.file.first_row_id not in first_row_id_set:
-                    first_row_id_set.add(entry.file.first_row_id)
+                if not self._is_blob_file(entry.file.file_name):
                     total_row += entry.file.row_count
 
-        # Calculate number of rows this shard should process
-        # Last shard handles all remaining rows (handles non-divisible cases)
-        if self.idx_of_this_subtask == self.number_of_para_subtasks - 1:
-            num_row = total_row - total_row // self.number_of_para_subtasks * self.idx_of_this_subtask
+        # Calculate number of rows this shard should process using balanced distribution
+        # Distribute remainder evenly among first few shards to avoid last shard overload
+        base_rows_per_shard = total_row // self.number_of_para_subtasks
+        remainder = total_row % self.number_of_para_subtasks
+
+        # Each of the first 'remainder' shards gets one extra row
+        if self.idx_of_this_subtask < remainder:
+            num_row = base_rows_per_shard + 1
+            start_row = self.idx_of_this_subtask * (base_rows_per_shard + 1)
         else:
-            num_row = total_row // self.number_of_para_subtasks
-        # Calculate start row and end row position for current shard in all data
-        start_row = self.idx_of_this_subtask * (total_row // self.number_of_para_subtasks)
+            num_row = base_rows_per_shard
+            start_row = (remainder * (base_rows_per_shard + 1) +
+                         (self.idx_of_this_subtask - remainder) * base_rows_per_shard)
+
         end_row = start_row + num_row
 
         plan_start_row = 0
@@ -188,14 +195,13 @@ def _data_evolution_filter_by_shard(self, partitioned_files: defaultdict) -> (de
         # Iterate through all file entries to find files that overlap with current shard range
         for key, file_entries in partitioned_files.items():
             filtered_entries = []
-            first_row_id_set = set()
+            blob_added = False  # If it is true, all blobs corresponding to this data file are added
             for entry in file_entries:
-                if entry.file.first_row_id is not None:
-                    if entry.file.first_row_id in first_row_id_set:
+                if self._is_blob_file(entry.file.file_name):
+                    if blob_added:
                         filtered_entries.append(entry)
-                        continue
-                    else:
-                        first_row_id_set.add(entry.file.first_row_id)
+                    continue
+                blob_added = False
                 entry_begin_row = entry_end_row  # Starting row position of current file in all data
                 entry_end_row += entry.file.row_count  # Update to row position after current file
 
@@ -213,18 +219,24 @@ def _data_evolution_filter_by_shard(self, partitioned_files: defaultdict) -> (de
                     plan_end_row = end_row - splits_start_row
                 # Add files that overlap with shard range to result
                 filtered_entries.append(entry)
+                blob_added = True
             if filtered_entries:
                 filtered_partitioned_files[key] = filtered_entries
 
         return filtered_partitioned_files, plan_start_row, plan_end_row
 
     def _compute_split_start_end_row(self, splits: List[Split], plan_start_row, plan_end_row):
         file_end_row = 0  # end row position of current file in all data
+
         for split in splits:
+            row_cnt = 0
             files = split.files
             split_start_row = file_end_row
             # Iterate through all file entries to find files that overlap with current shard range
             for file in files:
+                if self._is_blob_file(file.file_name):
+                    continue
+                row_cnt += file.row_count
                 file_begin_row = file_end_row  # Starting row position of current file in all data
                 file_end_row += file.row_count  # Update to row position after current file
 
@@ -238,7 +250,7 @@ def _compute_split_start_end_row(self, splits: List[Split], plan_start_row, plan
             if split.split_start_row is None:
                 split.split_start_row = 0
             if split.split_end_row is None:
-                split.split_end_row = split.row_count
+                split.split_end_row = row_cnt
 
     def _primary_key_filter_by_shard(self, file_entries: List[ManifestEntry]) -> List[ManifestEntry]:
         filtered_entries = []
@@ -359,61 +371,20 @@ def weight_func(fl: List[DataFileMeta]) -> int:
             splits += self._build_split_from_pack(flatten_packed_files, file_entries, True)
         return splits
 
-    def _build_split_from_pack(self, packed_files, file_entries, for_primary_key_split: bool) -> List['Split']:
-        splits = []
-        for file_group in packed_files:
-            raw_convertible = True
-            if for_primary_key_split:
-                raw_convertible = len(file_group) == 1
-
-            file_paths = []
-            total_file_size = 0
-            total_record_count = 0
-
-            for data_file in file_group:
-                data_file.set_file_path(self.table.table_path, file_entries[0].partition,
-                                        file_entries[0].bucket)
-                file_paths.append(data_file.file_path)
-                total_file_size += data_file.file_size
-                total_record_count += data_file.row_count
-
-            if file_paths:
-                split = Split(
-                    files=file_group,
-                    partition=file_entries[0].partition,
-                    bucket=file_entries[0].bucket,
-                    _file_paths=file_paths,
-                    _row_count=total_record_count,
-                    _file_size=total_file_size,
-                    raw_convertible=raw_convertible
-                )
-                splits.append(split)
-        return splits
-
-    @staticmethod
-    def _pack_for_ordered(items: List, weight_func: Callable, target_weight: int) -> List[List]:
-        packed = []
-        bin_items = []
-        bin_weight = 0
-
-        for item in items:
-            weight = weight_func(item)
-            if bin_weight + weight > target_weight and len(bin_items) > 0:
-                packed.append(list(bin_items))
-                bin_items.clear()
-                bin_weight = 0
-
-            bin_weight += weight
-            bin_items.append(item)
-
-        if len(bin_items) > 0:
-            packed.append(bin_items)
+    def _create_data_evolution_splits(self, file_entries: List[ManifestEntry]) -> List['Split']:
+        def sort_key(manifest_entry: ManifestEntry) -> tuple:
+            first_row_id = manifest_entry.file.first_row_id if manifest_entry.file.first_row_id is not None else float(
+                '-inf')
+            is_blob = 1 if self._is_blob_file(manifest_entry.file.file_name) else 0
+            # For files with same firstRowId, sort by maxSequenceNumber in descending order
+            # (larger sequence number means more recent data)
+            max_seq = manifest_entry.file.max_sequence_number
+            return first_row_id, is_blob, -max_seq
 
-        return packed
+        sorted_entries = sorted(file_entries, key=sort_key)
 
-    def _create_data_evolution_splits(self, file_entries: List[ManifestEntry]) -> List['Split']:
         partitioned_files = defaultdict(list)
-        for entry in file_entries:
+        for entry in sorted_entries:
             partitioned_files[(tuple(entry.partition.values), entry.bucket)].append(entry)
 
         if self.idx_of_this_subtask is not None:
@@ -423,11 +394,11 @@ def weight_func(file_list: List[DataFileMeta]) -> int:
             return max(sum(f.file_size for f in file_list), self.open_file_cost)
 
         splits = []
-        for key, file_entries in partitioned_files.items():
-            if not file_entries:
+        for key, sorted_entries in partitioned_files.items():
+            if not sorted_entries:
                 continue
 
-            data_files: List[DataFileMeta] = [e.file for e in file_entries]
+            data_files: List[DataFileMeta] = [e.file for e in sorted_entries]
 
             # Split files by firstRowId for data evolution
             split_by_row_id = self._split_by_row_id(data_files)
@@ -442,7 +413,7 @@ def weight_func(file_list: List[DataFileMeta]) -> int:
                 for pack in packed_files
             ]
 
-            splits += self._build_split_from_pack(flatten_packed_files, file_entries, False)
+            splits += self._build_split_from_pack(flatten_packed_files, sorted_entries, False)
 
         if self.idx_of_this_subtask is not None:
             self._compute_split_start_end_row(splits, plan_start_row, plan_end_row)
@@ -451,18 +422,8 @@ def weight_func(file_list: List[DataFileMeta]) -> int:
     def _split_by_row_id(self, files: List[DataFileMeta]) -> List[List[DataFileMeta]]:
         split_by_row_id = []
 
-        def sort_key(file: DataFileMeta) -> tuple:
-            first_row_id = file.first_row_id if file.first_row_id is not None else float('-inf')
-            is_blob = 1 if self._is_blob_file(file.file_name) else 0
-            # For files with same firstRowId, sort by maxSequenceNumber in descending order
-            # (larger sequence number means more recent data)
-            max_seq = file.max_sequence_number
-            return (first_row_id, is_blob, -max_seq)
-
-        sorted_files = sorted(files, key=sort_key)
-
         # Filter blob files to only include those within the row ID range of non-blob files
-        sorted_files = self._filter_blob(sorted_files)
+        sorted_files = self._filter_blob(files)
 
         # Split files by firstRowId
         last_row_id = -1
@@ -499,6 +460,58 @@ def sort_key(file: DataFileMeta) -> tuple:
 
         return split_by_row_id
 
+    def _build_split_from_pack(self, packed_files, file_entries, for_primary_key_split: bool) -> List['Split']:
+        splits = []
+        for file_group in packed_files:
+            raw_convertible = True
+            if for_primary_key_split:
+                raw_convertible = len(file_group) == 1
+
+            file_paths = []
+            total_file_size = 0
+            total_record_count = 0
+
+            for data_file in file_group:
+                data_file.set_file_path(self.table.table_path, file_entries[0].partition,
+                                        file_entries[0].bucket)
+                file_paths.append(data_file.file_path)
+                total_file_size += data_file.file_size
+                total_record_count += data_file.row_count
+
+            if file_paths:
+                split = Split(
+                    files=file_group,
+                    partition=file_entries[0].partition,
+                    bucket=file_entries[0].bucket,
+                    _file_paths=file_paths,
+                    _row_count=total_record_count,
+                    _file_size=total_file_size,
+                    raw_convertible=raw_convertible
+                )
+                splits.append(split)
+        return splits
+
+    @staticmethod
+    def _pack_for_ordered(items: List, weight_func: Callable, target_weight: int) -> List[List]:
+        packed = []
+        bin_items = []
+        bin_weight = 0
+
+        for item in items:
+            weight = weight_func(item)
+            if bin_weight + weight > target_weight and len(bin_items) > 0:
+                packed.append(list(bin_items))
+                bin_items.clear()
+                bin_weight = 0
+
+            bin_weight += weight
+            bin_items.append(item)
+
+        if len(bin_items) > 0:
+            packed.append(bin_items)
+
+        return packed
+
     @staticmethod
     def _is_blob_file(file_name: str) -> bool:
         return file_name.endswith('.blob')

paimon-python/pypaimon/read/table_read.py

@@ -80,6 +80,8 @@ def to_arrow(self, splits: List[Split]) -> Optional[pyarrow.Table]:
         schema = PyarrowFieldParser.from_paimon_schema(self.read_type)
         table_list = []
         for batch in iter(batch_reader.read_next_batch, None):
+            if batch.num_rows == 0:
+                continue
             table_list.append(self._try_to_pad_batch_by_schema(batch, schema))
 
         if not table_list: