class_docstr: Samples rows from record batches.
  
  The record batches in the partition can vary in the number of rows.
  SamplePartition guarantees that the sample returned is always going to be
  <= sample_size.
  
  The actual sampling occurs in `compact`. It uses np.partition to calculate
  the top-k of record batch's rows. Where the top-k is a random number assigned
  to each row. Given a uniform distribution of the random number, we can keep a
  running sample of the partition of size k. This gives each row an equal
  probability of being selected.
class_docstr_tok: ["Samples","rows","from","record","batches",".","The","record","batches","in","the","partition","can","vary","in","the","number","of","rows",".","SamplePartition","guarantees","that","the","sample","returned","is","always","going","to","be","<","=","sample_size",".","The","actual","sampling","occurs","in","`","compact","`",".","It","uses","np.partition","to","calculate","the","top-k","of","record","batch","'s","rows",".","Where","the","top-k","is","a","random","number","assigned","to","each","row",".","Given","a","uniform","distribution","of","the","random","number",",","we","can","keep","a","running","sample","of","the","partition","of","size","k.","This","gives","each","row","an","equal","probability","of","being","selected","."]
docstr: Compacts the accumulator.
  
  This compact selects samples rows from the record batch, and merges them
  into one record batch. We can then clear the cache of all record batches
  seen so far. If the accumulator holds too few record batches, then nothing
  will be compacted.
  
  The sampling is done by assigning each row in the record batch a random
  number. Then we choose the top-k of the random numbers to get a sample of
  size k.
  
  Args:
    accumulator: The _SampleRecordBatchRowsAccumulator to compact.
  
  Returns:
    A _SampleRecordBatchRowsAccumulator that contains one or a list of record
    batch.
docstr_tok: ["Compacts","the","accumulator",".","This","compact","selects","samples","rows","from","the","record","batch",",","and","merges","them","into","one","record","batch",".","We","can","then","clear","the","cache","of","all","record","batches","seen","so","far",".","If","the","accumulator","holds","too","few","record","batches",",","then","nothing","will","be","compacted",".","The","sampling","is","done","by","assigning","each","row","in","the","record","batch","a","random","number",".","Then","we","choose","the","top-k","of","the","random","numbers","to","get","a","sample","of","size","k.","Args",":","accumulator",":","The","_SampleRecordBatchRowsAccumulator","to","compact",".","Returns",":","A","_SampleRecordBatchRowsAccumulator","that","contains","one","or","a","list","of","record","batch","."]
code: def _compact_impl(
      self, accumulator: _SampleRecordBatchRowsAccumulator
  ) -> _SampleRecordBatchRowsAccumulator:
      """Compacts the accumulator.
  
      This compact selects samples rows from the record batch, and merges them
      into one record batch. We can then clear the cache of all record batches
      seen so far. If the accumulator holds too few record batches, then nothing
      will be compacted.
  
      The sampling is done by assigning each row in the record batch a random
      number. Then we choose the top-k of the random numbers to get a sample of
      size k.
  
      Args:
        accumulator: The _SampleRecordBatchRowsAccumulator to compact.
  
      Returns:
        A _SampleRecordBatchRowsAccumulator that contains one or a list of record
        batch.
      """
      self._combine_num_record_batches.update(
          len(accumulator.record_batches)
      )
  
      # There is nothing to compact.
      if accumulator.curr_num_rows <= 1:
          return accumulator
  
      # There is no need to compact yet.
      if (
          len(accumulator.record_batches) <= 1
          and accumulator.curr_num_rows <= self._sample_size
      ):
          return accumulator
      self._num_compacts.inc(1)
      k = min(self._sample_size, accumulator.curr_num_rows)
  
      rand_ints = np.concatenate(accumulator.random_ints)
  
      # Find the value that is the breakpoint for the top-k.
      kth_value = np.partition(rand_ints, k - 1)[k - 1]
  
      # This mask will always have >= 1 Trues.
      equals_to_kth = rand_ints == kth_value
  
      # This mask will always have < k Trues.
      less_than_kth = rand_ints < kth_value
  
      # Since there may be duplicate values, `equals_to_kth + less_than_kth` might
      # be greater than `k`. We need to keep track of how many to add, without
      # surpassing `k`.
      kth_to_add = k - np.sum(less_than_kth)
  
      # Preserve the random integers that we had assigned to each row.
      sample_random_ints = rand_ints[rand_ints <= kth_value][:k]
  
      beg = 0
      sample_indices = []
      for rb in accumulator.record_batches:
          size = rb.num_rows
          end = beg + size
          less_than_kth_indices = np.nonzero(less_than_kth[beg:end])[0]
          indices = less_than_kth_indices
  
          # Add indices of any duplicate values that are equal to `k`.
          if kth_to_add > 0:
              equals_to_kth_indices = np.nonzero(
                  equals_to_kth[beg:end]
              )[0]
              if equals_to_kth_indices.size > 0:
                  if equals_to_kth_indices.size >= kth_to_add:
                      indices = np.concatenate(
                          [
                              less_than_kth_indices,
                              equals_to_kth_indices[:kth_to_add],
                          ]
                      )
                      kth_to_add = 0
                  else:
                      indices = np.concatenate(
                          [less_than_kth_indices, equals_to_kth_indices]
                      )
                      kth_to_add -= equals_to_kth_indices.size
  
          sample_indices.append(indices)
          beg += size
  
      result = _SampleRecordBatchRowsAccumulator()
  
      # Take and merge the record batches, based on the sampled indices.
      rbs = []
      for rb, indices in zip(
          accumulator.record_batches, sample_indices
      ):
          rbs.append(table_util.RecordBatchTake(rb, pa.array(indices)))
      compressed_rb = table_util.MergeRecordBatches(rbs)
      result.record_batches = [compressed_rb]
      result.curr_num_rows = compressed_rb.num_rows
      result.curr_byte_size = compressed_rb.nbytes
      result.random_ints = [sample_random_ints]
  
      self._combine_byte_size.update(result.curr_byte_size)
  
      return result
  
imports: import collections
  import functools
  from typing import Dict, Iterable, Text, Tuple
  import apache_beam
  import numpy
  import pyarrow
  from tensorflow_data_validation import constants
  from tensorflow_data_validation import types
  from tensorflow_data_validation.statistics.generators import stats_generator
  from tensorflow_data_validation.utils import stats_util
  from tfx_bsl.arrow import table_util
  from tensorflow_metadata.proto.v0 import statistics_pb2
label_desc: Use image node_id 7 for calling the _SampleRecordBatchRows obj's underlying member method code with example usage: obj._compact_impl(accumulator) and returns: result, accumulator, accumulator