The current implementation fetches blocks sequentially within the remote task, which can be inefficient. You can improve performance by fetching all blocks in parallel using a single ray.get() call before merging them.

The TARGET_NUM_ROWS_PER_BLOCK is hardcoded to 128. This might not be optimal for all use cases. Consider making this value configurable, for example, through the DataContext. This would provide more flexibility for users to tune performance based on their specific data characteristics (e.g., row size).

TypeError when bundle has unknown row count

RefBundle.num_rows() returns Optional[int] and can return None when block metadata doesn't have a row count. The code at line 131 performs self._pending_rows[input_index] += refs.num_rows() which raises a TypeError if num_rows() is None. This propagates to line 186 where min(self._pending_rows) fails with mixed int/None values, and line 243 where bundle_rows <= rows_remaining comparison fails. Several data sources like BigQuery, Delta Sharing, and Hudi legitimately set num_rows=None in metadata. The previous bulk implementation handled this by fetching row counts remotely, but the new streaming implementation lacks this safeguard.

Output bundle size_bytes is zero when slicing occurs

The total_size_bytes calculation only counts bundles that are fully consumed, but when bundles are sliced (split across multiple tasks), the consumed list is empty until the final slice is used. This means output bundles created from sliced data will have size_bytes=0 in their metadata, while the final task gets the entire original bundle's size. This causes inaccurate memory accounting despite implements_accurate_memory_accounting() returning True. The fix is to calculate size from merged_bundles (which contains the actual processed data) instead of from consumed bundles.

Output bundle ownership incorrectly based on input bundles

The output bundle's owns_blocks is derived from consumed input bundles rather than recognizing that the remote task creates a new block. When input bundles don't own their blocks (e.g., shared bundles) and are consumed whole without slicing, owns_blocks becomes False for the output. However, the output contains a newly created block from _zip_blocks_with_slices that has no other references, so it should always be owned. This could prevent eager memory cleanup when destroy_if_owned is called, since the output incorrectly believes it doesn't own its block.