Function _split_all_blocks

python/ray/data/_internal/split.py:155–217

Split all the input blocks based on the split indices

(
    blocks_with_metadata: List[Tuple[ObjectRef[Block], BlockMetadata]],
    per_block_split_indices: List[List[int]],
    owned_by_consumer: bool,
    label_selector: Optional[Dict[str, str]] = None,
)

def _split_all_blocks(
    blocks_with_metadata: List[Tuple[ObjectRef[Block], BlockMetadata]],
    per_block_split_indices: List[List[int]],
    owned_by_consumer: bool,
    label_selector: Optional[Dict[str, str]] = None,
) -> Iterable[Tuple[ObjectRef[Block], BlockMetadata]]:
    """Split all the input blocks based on the split indices"""
    split_single_block = cached_remote_fn(_split_single_block)

    all_blocks_split_results: List[BlockPartition] = [None] * len(blocks_with_metadata)

    per_block_split_metadata_futures = []
    per_block_split_block_refs = []

    # tracking splitted blocks for gc.
    blocks_splitted = []
    for block_id, block_split_indices in enumerate(per_block_split_indices):
        (block_ref, meta) = blocks_with_metadata[block_id]
        block_row = meta.num_rows
        block_split_indices = _drop_empty_block_split(block_split_indices, block_row)
        if len(block_split_indices) == 0:
            # optimization: if no split is needed, we just need to add it to the
            # result
            all_blocks_split_results[block_id] = [(block_ref, meta)]
        else:
            # otherwise call split remote function.
            options = {
                "scheduling_strategy": "SPREAD",
                "num_returns": 2 + len(block_split_indices),
            }
            if label_selector:
                options["label_selector"] = label_selector
            object_refs = split_single_block.options(**options).remote(
                block_id,
                block_ref,
                meta,
                block_split_indices,
            )
            per_block_split_metadata_futures.append(object_refs[0])
            per_block_split_block_refs.append(object_refs[1:])

            blocks_splitted.append(block_ref)

    if per_block_split_metadata_futures:
        # only get metadata.
        per_block_split_metadata = ray.get(per_block_split_metadata_futures)
        for (block_id, meta), block_refs in zip(
            per_block_split_metadata, per_block_split_block_refs
        ):
            assert len(meta) == len(block_refs)
            all_blocks_split_results[block_id] = zip(block_refs, meta)

    # We make a copy for the blocks that have been splitted, so the input blocks
    # can be cleared if they are owned by consumer (consumer-owned blocks will
    # only be consumed by the owner).
    if owned_by_consumer:
        for b in blocks_splitted:
            trace_deallocation(b, "split._split_all_blocks")