def _convert_to_record_batch(
    batch_dict: Dict[str, Union[common_types.TensorValueType, pa.Array]],
    converter: tensor_to_arrow.TensorsToRecordBatchConverter,
    passthrough_keys: Set[str],
    input_metadata: Union[
        TensorAdapterConfig, dataset_metadata.DatasetMetadata
    ],
    validate_varlen_sparse_values: bool = False,
) -> Iterable[Tuple[pa.RecordBatch, Dict[str, pa.Array]]]:
  """Convert batch of ndarrays to pyarrow.RecordBatches."""

  # Making a copy of batch_dict because mutating PCollection elements is not
  # allowed.
  if passthrough_keys:
    batch_dict = copy.copy(batch_dict)
  passthrough_data = {
      key: batch_dict.pop(key) for key in passthrough_keys if key in batch_dict
  }

  if validate_varlen_sparse_values:
    for name, representation in converter.tensor_representations().items():
      if representation.WhichOneof('kind') == 'varlen_sparse_tensor':
        impl_helper.validate_varlen_sparse_value(name, batch_dict[name])

  record_batch = converter.convert(batch_dict)
  arrow_columns, arrow_schema = record_batch.columns, record_batch.schema

  batch_size = len(arrow_columns[0])
  # This dict will contain pass-through data with batch size of 1 if it doesn't
  # match batch size of the transformed data.
  unary_passthrough_features = {}
  for key, data in passthrough_data.items():
    # Only raising a ValueError in case pass-through data has more than one
    # distinct value. If it has one value and batch_size>1 then it will have to
    # be handled by the user.
    # TODO(b/38376110): Restrict to matching batch dimensions and clean this up
    # once the internal feature key is deprecated.
    if len(data) not in (batch_size, 1):
      # The passthrough column should be of list<primitive> type with each
      # sub-list being either null or of length 1.
      data_set = set(
          None if elem is None else elem[0] for elem in data.to_pylist())
      if len(data_set) == 1:
        elem = data_set.pop()
        data = pa.array([None if elem is None else [elem]], type=data.type)
      else:
        raise ValueError(
            'Cannot pass-through data when input and output batch sizes '
            'are different ({} vs. {})'.format(len(data), batch_size))
    if len(data) == batch_size:
      arrow_schema = arrow_schema.append(input_metadata.arrow_schema.field(key))
      arrow_columns.append(data)
    else:
      unary_passthrough_features[key] = data
  for reccord_batch in _maybe_slice_large_record_batch(
      pa.RecordBatch.from_arrays(arrow_columns, schema=arrow_schema)
  ):
    yield reccord_batch, unary_passthrough_features
Convert batch of ndarrays to pyarrow.RecordBatches.
tensorflow_transform/beam/impl.py