def _BatchAndPreSort(counts, sort_kwargs):  # pylint: disable=invalid-name
  """Batches vocabulary and pre-sorts the batches."""
  # This PTransform aims to partially parallelize vocabulary sorting in
  # `VocabularyOrderAndWrite`. Pre-sorting of batches in a parallel mode with
  # `beam.Map` allows to significantly speed up final single-threaded sorting of
  # a union in `_OrderElementsFn`. This is because `list.sort()` uses an
  # adaptive merge sort that identifies pre-existing order in the union. See
  # https://en.wikipedia.org/wiki/Timsort for more details.
  return (counts
          | 'BatchVocabulary' >> beam.BatchElements(
              min_batch_size=_PRESORT_BATCH_SIZE,
              max_batch_size=_PRESORT_BATCH_SIZE)
          | 'SortBatches' >> beam.Map(lambda b: sorted(b, **sort_kwargs)))  # pylint: disable=unnecessary-lambda
Batches vocabulary and pre-sorts the batches.
tensorflow_transform/beam/analyzer_impls.py