def _maybe_slice_large_record_batch(
    record_batch: pa.RecordBatch,
) -> Iterable[pa.RecordBatch]:
  """Slices large batches into smaller chunks."""
  if record_batch.nbytes > _MAX_TRANSFORMED_BATCH_BYTES_SIZE:
    if record_batch.num_rows < 2:
      logging.warning(
          'Transformed data row may be too large: %d bytes. '
          'Consider reshaping outputs to distribute elements over a larger '
          'number of rows to allow automatic slicing.',
          record_batch.nbytes,
      )
      yield record_batch
      return
    # Note that slicing is a zero-copy operation, so the produced batches will
    # still share memory with the original one up to the materialization
    # boundary.
    mid_point = record_batch.num_rows // 2
    yield from _maybe_slice_large_record_batch(
        record_batch.slice(offset=0, length=mid_point)
    )
    yield from _maybe_slice_large_record_batch(
        record_batch.slice(offset=mid_point)
    )
  else:
    yield record_batch
Slices large batches into smaller chunks.
tensorflow_transform/beam/impl.py