import itertools from collections import defaultdict from dataclasses import dataclass from typing import Dict, Iterable, Iterator, List, Optional, Tuple import ray from .common import NodeIdStr from ray.data._internal.memory_tracing import trace_deallocation from ray.data.block import ( Block, BlockAccessor, BlockMetadata, Schema, _take_first_non_empty_schema, ) from ray.data.context import DataContext from ray.types import ObjectRef @dataclass(frozen=True) class BlockSlice: """A slice of a block.""" # Starting row offset (inclusive) within the block. start_offset: int # Ending row offset (exclusive) within the block. end_offset: int @property def num_rows(self) -> int: return self.end_offset - self.start_offset @dataclass(frozen=True, slots=True) class BlockEntry: """One block delivery: the ref + the block's measured metadata. Used as the element type of ``RefBundle.blocks`` (replaces the legacy ``(ObjectRef, BlockMetadata)`` 2-tuple shape). Naming the fields makes every call site self-describing and reserves room for the bundle entry to grow without disturbing the surrounding shape. """ ref: ObjectRef[Block] metadata: BlockMetadata @dataclass(frozen=True) class RefBundle: """A group of data block references and their metadata. Operators take in and produce streams of RefBundles. Most commonly a RefBundle consists of a single block object reference. In some cases, e.g., due to block splitting, or for a reduce task, there may be more than one block. Block bundles have ownership semantics, i.e., shared ownership (similar to C++ shared_ptr, multiple operators share the same block bundle), or unique ownership (similar to C++ unique_ptr, only one operator owns the block bundle). This allows operators to know whether they can destroy blocks when they don't need them. Destroying blocks eagerly is more efficient than waiting for Python GC / Ray reference counting to kick in. """ # Per-block entries. The size_bytes must be known in the metadata, # num_rows is optional. Legacy ``(ref, metadata)`` 2-tuples are no longer # accepted at construction and must be explicitly wrapped in ``BlockEntry`` # (``__post_init__`` rejects anything else with an actionable assertion). blocks: Tuple[BlockEntry, ...] # The schema of the blocks in this bundle. This is optional, and may be None # if blocks are empty. schema: Optional["Schema"] # Whether we own the blocks (can safely destroy them). owns_blocks: bool # The slices of the blocks in this bundle. After __post_init__, this is always # a list with length equal to len(blocks). Individual entries can be None to # represent a full block (equivalent to BlockSlice(0, num_rows)). # Pass None during construction to initialize all slices as None (full blocks). slices: Optional[Tuple[Optional[BlockSlice], ...]] = None # This attribute is used by the split() operator to assign bundles to logical # output splits. It is otherwise None. output_split_idx: Optional[int] = None # Object metadata (size, locations, spilling status) _cached_object_meta: Optional[Dict[ObjectRef, "_ObjectMetadata"]] = None # Preferred locations for this bundle determined based on the locations # of individual objects and their corresponding size, ie location with the # largest total number of bytes present there has the highest preference. _cached_preferred_locations: Optional[Dict[NodeIdStr, int]] = None def __post_init__(self): if self.schema is not None: import pyarrow as pa from ray.data._internal.pandas_block import PandasBlockSchema assert isinstance( self.schema, (pa.lib.Schema, PandasBlockSchema) ), f"Schema must be a pyarrow or PandasBlockSchema, got {type(self.schema)}" if not isinstance(self.blocks, tuple): object.__setattr__(self, "blocks", tuple(self.blocks)) for entry in self.blocks: assert isinstance(entry, BlockEntry), ( f"RefBundle.blocks must contain BlockEntry instances; got {type(entry).__name__}. " "Construct entries with `BlockEntry(ref=..., metadata=...)`." ) if self.slices is None: object.__setattr__(self, "slices", (None,) * len(self.blocks)) else: if not isinstance(self.slices, tuple): object.__setattr__(self, "slices", tuple(self.slices)) assert len(self.blocks) == len( self.slices ), "Number of blocks and slices must match" # Validate slice ranges for entry, block_slice in zip(self.blocks, self.slices): if block_slice is not None: assert ( block_slice.start_offset >= 0 ), f"Slice start_offset must be non-negative: {block_slice.start_offset}" assert ( block_slice.end_offset >= block_slice.start_offset ), f"Slice end_offset must be >= start_offset: [{block_slice.start_offset}, {block_slice.end_offset})" if entry.metadata.num_rows is not None: assert ( block_slice.end_offset <= entry.metadata.num_rows ), f"Slice range [{block_slice.start_offset}, {block_slice.end_offset}) exceeds block num_rows: {entry.metadata.num_rows}" for entry in self.blocks: if entry.metadata.size_bytes is None: raise ValueError( "The size in bytes of the block must be known: {}".format(entry) ) @property def block_refs(self) -> List[ObjectRef[Block]]: """List of block references in this bundle.""" return [entry.ref for entry in self.blocks] @property def metadata(self) -> List[BlockMetadata]: """List of block metadata in this bundle.""" return [entry.metadata for entry in self.blocks] def num_rows(self) -> Optional[int]: """Number of rows present in this bundle, if known. Iterates through blocks and their corresponding slices to calculate the total. Note: Block metadata always refers to the full block, not the slice. - If block_slice is None, uses the full block's metadata.num_rows - If block_slice is present, uses the slice's num_rows (partial block portion) - Returns None if any full block has unknown row count (metadata.num_rows is None) """ total = 0 for metadata, block_slice in zip(self.metadata, self.slices): if block_slice is None: if metadata.num_rows is None: return None total += metadata.num_rows else: total += block_slice.num_rows return total def size_bytes(self) -> int: """Size of the blocks of this bundle in bytes. Iterates through blocks and their corresponding slices to calculate the total size. Note: Block metadata always refers to the full block, not the slice. - If block_slice is None, uses the full block's metadata.size_bytes - If block_slice is present but num_rows is unknown or zero, uses full metadata.size_bytes - If block_slice represents a partial block, estimates size proportionally based on (metadata.size_bytes / metadata.num_rows) * block_slice.num_rows - Otherwise, uses the full metadata.size_bytes """ total = 0 for entry, block_slice in zip(self.blocks, self.slices): metadata = entry.metadata if block_slice is None: # Full block total += metadata.size_bytes elif metadata.num_rows is None or metadata.num_rows == 0: # Unknown num_rows or empty block - use full metadata size total += metadata.size_bytes elif metadata.num_rows != block_slice.num_rows: # Partial block - estimate size based on rows per_row = metadata.size_bytes / metadata.num_rows total += max(1, round(per_row * block_slice.num_rows)) else: total += metadata.size_bytes return total def destroy_if_owned(self) -> int: """Clears the object store memory for these blocks if owned. Returns: The number of bytes freed. """ should_free = self.owns_blocks and DataContext.get_current().eager_free for block_ref in self.block_refs: trace_deallocation( block_ref, "RefBundle.destroy_if_owned", free=should_free ) return self.size_bytes() if should_free else 0 def get_preferred_object_locations(self) -> Dict[NodeIdStr, int]: """Returns a mapping of node IDs to total bytes stored on each node. Returns: Dict mapping node ID to total bytes stored on that node """ meta = self._get_cached_metadata() if self._cached_preferred_locations is None: preferred_locs: Dict[NodeIdStr, int] = defaultdict(int) for ref, obj_meta in meta.items(): for loc in obj_meta.locs: preferred_locs[loc] += obj_meta.size # NOTE: We're working around object being immutable to update cached # values (safe) object.__setattr__(self, "_cached_preferred_locations", preferred_locs) return self._cached_preferred_locations def _get_cached_metadata(self) -> Dict[ObjectRef, "_ObjectMetadata"]: if self._cached_object_meta is None: # This call is pretty fast for owned objects (~5k/s), so we don't need to # batch it for now. meta = ray.experimental.get_local_object_locations(self.block_refs) # Extract locations object_metas: Dict[ObjectRef, _ObjectMetadata] = { ref: _ObjectMetadata( size=meta[ref]["object_size"], spilled=meta[ref]["did_spill"], locs=meta[ref]["node_ids"], ) for ref in self.block_refs } # NOTE: We're working around object being immutable to update cached # values (safe) object.__setattr__(self, "_cached_object_meta", object_metas) return self._cached_object_meta def slice(self, needed_rows: int) -> Tuple["RefBundle", "RefBundle"]: """Slice a Ref Bundle into the first bundle containing the first `needed_rows` rows and the remaining bundle containing the remaining rows. Args: needed_rows: Number of rows to take from the head of the bundle. Returns: A tuple of (sliced_bundle, remaining_bundle). The needed rows must be less than the number of rows in the bundle. """ assert needed_rows > 0, "needed_rows must be positive." assert ( self.num_rows() is not None ), "Cannot slice a RefBundle with unknown number of rows." assert ( needed_rows < self.num_rows() ), f"To slice a RefBundle, the number of requested rows must be less than the number of rows in the bundle. Requested {needed_rows} rows but bundle only has {self.num_rows()} rows." block_slices = [] for metadata, block_slice in zip(self.metadata, self.slices): if block_slice is None: # None represents a full block, convert to explicit BlockSlice assert ( metadata.num_rows is not None ), "Cannot derive block slice for a RefBundle with unknown block row counts." block_slices.append( BlockSlice(start_offset=0, end_offset=metadata.num_rows) ) else: block_slices.append(block_slice) consumed_blocks: List[BlockEntry] = [] consumed_slices: List[BlockSlice] = [] remaining_blocks: List[BlockEntry] = [] remaining_slices: List[BlockSlice] = [] rows_to_take = needed_rows for entry, block_slice in zip(self.blocks, block_slices): block_rows = block_slice.num_rows if rows_to_take >= block_rows: consumed_blocks.append(entry) consumed_slices.append(block_slice) rows_to_take -= block_rows else: if rows_to_take == 0: remaining_blocks.append(entry) remaining_slices.append(block_slice) continue consume_slice = BlockSlice( start_offset=block_slice.start_offset, end_offset=block_slice.start_offset + rows_to_take, ) consumed_blocks.append(entry) consumed_slices.append(consume_slice) leftover_rows = block_rows - rows_to_take if leftover_rows > 0: remainder_slice = BlockSlice( start_offset=consume_slice.end_offset, end_offset=block_slice.end_offset, ) remaining_blocks.append(entry) remaining_slices.append(remainder_slice) rows_to_take = 0 sliced_bundle = RefBundle( blocks=tuple(consumed_blocks), schema=self.schema, owns_blocks=False, slices=tuple(consumed_slices) if consumed_slices else None, ) remaining_bundle = RefBundle( blocks=tuple(remaining_blocks), schema=self.schema, owns_blocks=False, slices=tuple(remaining_slices) if remaining_slices else None, ) return sliced_bundle, remaining_bundle @classmethod def merge_ref_bundles(cls, bundles: Iterable["RefBundle"]) -> "RefBundle": """Merge multiple RefBundles into a single RefBundle. Args: bundles: An iterable of RefBundles to merge. Returns: A single RefBundle containing all blocks from the input bundles. owns_blocks is True only if all input bundles own their blocks. schema is the first non-empty schema found. """ bundles = list(bundles) if not bundles: return cls(blocks=(), owns_blocks=True, schema=None) merged_blocks = list( itertools.chain.from_iterable(bundle.blocks for bundle in bundles) ) merged_slices = list( itertools.chain.from_iterable(bundle.slices for bundle in bundles) ) # Ray Data uses the `owns_blocks` flag to determine if the system can eagerly # destroy blocks when they're no longer needed. To be safe, we only set this # to True if all input bundles own their blocks. owns_blocks = all(bundle.owns_blocks for bundle in bundles) # TODO: Reconcile the schemas rather than taking the first non-empty schema. schema = _take_first_non_empty_schema(bundle.schema for bundle in bundles) return cls( blocks=tuple(merged_blocks), schema=schema, owns_blocks=owns_blocks, slices=merged_slices, ) def __eq__(self, other: "RefBundle"): if self is other: return True elif not isinstance(other, RefBundle): return False return ( self.blocks == other.blocks and self.slices == other.slices # NOTE: We're establishing a requirement of schemas for `RefBundle` # to be exactly the same object for it to be considered equal. # # This is necessary to avoid a full schema equality check that # is computationally intensive. and self.schema is other.schema and self.owns_blocks == other.owns_blocks and self.output_split_idx == other.output_split_idx ) def __hash__(self) -> int: return hash( ( # Only hash block refs *[entry.ref for entry in self.blocks], *self.slices, # Check out comment in ``__eq__`` id(self.schema), self.owns_blocks, self.output_split_idx, ) ) def __len__(self) -> int: return len(self.blocks) def __str__(self) -> str: lines = [ f"RefBundle({len(self.blocks)} blocks,", f" {self.num_rows()} rows,", f" schema={self.schema},", f" owns_blocks={self.owns_blocks},", " blocks=(", ] # Loop through each block and show details for i, (entry, block_slice) in enumerate(zip(self.blocks, self.slices)): metadata = entry.metadata row_str = ( f"{metadata.num_rows} rows" if metadata.num_rows is not None else "unknown rows" ) bytes_str = f"{metadata.size_bytes} bytes" slice_str = ( f"slice={block_slice}" if block_slice is not None else "slice=None (full block)" ) lines.append(f" {i}: {row_str}, {bytes_str}, {slice_str}") lines.append(" )") lines.append(")") return "\n".join(lines) @dataclass class _ObjectMetadata: # Object size in bytes size: int # Flag whether object has been spilled spilled: bool # List of nodes object exists on locs: List[NodeIdStr] = None def _ref_bundles_iterator_to_block_refs_list( ref_bundles: Iterator[RefBundle], ) -> List[ObjectRef[Block]]: """Convert an iterator of RefBundles to a list of Block object references.""" return [ block_ref for ref_bundle in ref_bundles for block_ref in ref_bundle.block_refs ] def _iter_sliced_blocks( blocks: Iterable[Block], slices: List[Optional[BlockSlice]], ) -> Iterator[Block]: blocks_list = list(blocks) for block, block_slice in zip(blocks_list, slices): if block_slice is None: # None represents a full block - yield it as is yield block else: accessor = BlockAccessor.for_block(block) start = block_slice.start_offset end = block_slice.end_offset assert start <= end, "start must be less than end" assert start >= 0, "start must be non-negative" assert ( end <= accessor.num_rows() ), "end must be less than or equal to the number of rows in the block" yield accessor.slice(start, end, copy=False)