class_docstr: Implementation of `Dataset.zip()`.
code: def __init__(self, other: "Dataset"):
      def do_zip_all(
          block_list: BlockList, clear_input_blocks: bool, *_
      ):
          # Repartition other to align with the base dataset, and then zip together
          # the blocks in parallel.
          # TODO(Clark): Port this to a streaming zip, e.g. push block pairs through
          # an actor that buffers and zips.
          base_block_list = block_list
          base_blocks_with_metadata = (
              block_list.get_blocks_with_metadata()
          )
          (
              base_block_rows,
              base_block_bytes,
          ) = _calculate_blocks_rows_and_bytes(
              base_blocks_with_metadata
          )
          # Execute other to a block list.
          # NOTE: Require to preserve order when executing the other side,
          # because streaming execution does not preserve order by default.
          other_block_list = other._plan.execute(preserve_order=True)
          other_blocks_with_metadata = (
              other_block_list.get_blocks_with_metadata()
          )
          (
              other_block_rows,
              other_block_bytes,
          ) = _calculate_blocks_rows_and_bytes(
              other_blocks_with_metadata
          )
          inverted = False
          if sum(other_block_bytes) > sum(base_block_bytes):
              # Make sure that other is the smaller dataset, so we minimize
              # splitting work when aligning other with base.
              # TODO(Clark): Improve this heuristic for minimizing splitting work,
              # e.g. by generating the splitting plans for each route (via
              # _generate_per_block_split_indices) and choosing the plan that splits
              # the least cumulative bytes.
              base_block_list, other_block_list = (
                  other_block_list,
                  base_block_list,
              )
              base_blocks_with_metadata, other_blocks_with_metadata = (
                  other_blocks_with_metadata,
                  base_blocks_with_metadata,
              )
              base_block_rows, other_block_rows = (
                  other_block_rows,
                  base_block_rows,
              )
              inverted = True
          # Get the split indices that will align other with base.
          indices = list(itertools.accumulate(base_block_rows))
          indices.pop(-1)
  
          # Check that each dataset has the same number of rows.
          # TODO(Clark): Support different number of rows via user-directed
          # dropping/padding.
          total_base_rows = sum(base_block_rows)
          total_other_rows = sum(other_block_rows)
          if total_base_rows != total_other_rows:
              raise ValueError(
                  "Cannot zip datasets of different number of rows: "
                  f"{total_base_rows}, {total_other_rows}"
              )
  
          # Split other at the alignment indices, such that for every block in
          # block_list, we have a list of blocks from other that has the same
          # cumulative number of rows as that block.
          # NOTE: _split_at_indices has a no-op fastpath if the blocks are already
          # aligned.
          aligned_other_blocks_with_metadata = _split_at_indices(
              other_blocks_with_metadata,
              indices,
              other_block_list._owned_by_consumer,
              other_block_rows,
          )
          del other_blocks_with_metadata
  
          base_blocks = [b for b, _ in base_blocks_with_metadata]
          other_blocks = aligned_other_blocks_with_metadata[0]
          del (
              base_blocks_with_metadata,
              aligned_other_blocks_with_metadata,
          )
          if clear_input_blocks:
              base_block_list.clear()
              other_block_list.clear()
  
          do_zip = cached_remote_fn(_do_zip, num_returns=2)
  
          out_blocks = []
          out_metadata = []
          for base_block, other_blocks in zip(
              base_blocks, other_blocks
          ):
              # For each block in base, zip it together with 1 or more blocks from
              # other. We're guaranteed to have that base_block has the same number of
              # rows as other_blocks has cumulatively.
              res, meta = do_zip.remote(
                  base_block, *other_blocks, inverted=inverted
              )
              out_blocks.append(res)
              out_metadata.append(meta)
  
          # Early release memory.
          del base_blocks, other_blocks
  
          # TODO(ekl) it might be nice to have a progress bar here.
          out_metadata = ray.get(out_metadata)
          blocks = BlockList(
              out_blocks,
              out_metadata,
              owned_by_consumer=base_block_list._owned_by_consumer,
          )
          return blocks, {}
  
      super().__init__("Zip", None, do_zip_all)
  
imports: import itertools
  from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
  import ray
  from ray.data._internal.block_list import BlockList
  from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
  from ray.data._internal.execution.interfaces import TaskContext
  from ray.data._internal.fast_repartition import fast_repartition
  from ray.data._internal.plan import AllToAllStage
  from ray.data._internal.remote_fn import cached_remote_fn
  from ray.data._internal.shuffle_and_partition import PushBasedShufflePartitionOp, SimpleShufflePartitionOp
  from ray.data._internal.sort import SortKey, sort_impl
  from ray.data._internal.split import _split_at_index, _split_at_indices
  from ray.data.block import Block, BlockAccessor, BlockExecStats, BlockMetadata, BlockPartition
  from ray.data.context import DataContext
label_desc: Use image node_id 1 to create a new ZipStage object from inherited base classes: AllToAllStage with example: obj = ZipStage(other)