class_docstr: Tests for _asssign_to_partition.
docstr: Tests that batches are randomly partitioned.
  
  Tests an input batch with one univalent feature taking on values in {0,1,2}.
  The input batch has 4500 examples. The partitioner is configured to have 3
  partitions. So, we expect there to be around 4500/3/3 = 500 of each.
  {0,1,2} in each partition.
docstr_tok: ["Tests","that","batches","are","randomly","partitioned",".","Tests","an","input","batch","with","one","univalent","feature","taking","on","values","in","{","0,1,2","}",".","The","input","batch","has","4500","examples",".","The","partitioner","is","configured","to","have","3","partitions",".","So",",","we","expect","there","to","be","around","4500\/3\/3","=","500","of","each",".","{","0,1,2","}","in","each","partition","."]
code: def test_partitioner(self):
      """Tests that batches are randomly partitioned.
  
      Tests an input batch with one univalent feature taking on values in {0,1,2}.
      The input batch has 4500 examples. The partitioner is configured to have 3
      partitions. So, we expect there to be around 4500/3/3 = 500 of each.
      {0,1,2} in each partition.
      """
  
      np.random.seed(TEST_SEED)
      record_batches = [
          pa.RecordBatch.from_arrays([pa.array([x])], ["a"])
          for x in np.random.randint(0, 3, (4500, 1))
      ]
      record_batches = [
          (constants.DEFAULT_SLICE_KEY, record_batch)
          for record_batch in record_batches
      ]
      num_partitions = 3
  
      # The i,jth value of result represents the number of examples with value j
      # assigned to partition i.
      result = [[0, 0, 0], [0, 0, 0], [0, 0, 0]]
  
      partitioned_record_batches = [
          partitioned_stats_generator._default_assign_to_partition(
              record_batch, num_partitions
          )
          for record_batch in record_batches
      ]
      for (
          unused_slice_key,
          partition_key,
      ), record_batch in partitioned_record_batches:
          result[partition_key][
              record_batch.column(0).to_pylist()[0][0]
          ] += 1
  
      for partition in result:
          for count in partition:
              self.assertBetween(count, 400, 600)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl.testing import absltest
  from absl.testing import parameterized
  import apache_beam
  from apache_beam.testing import util
  import numpy
  import pyarrow
  from tensorflow_data_validation import constants
  from tensorflow_data_validation import types
  from tensorflow_data_validation.statistics.generators import partitioned_stats_generator
  from tensorflow_data_validation.statistics.generators import sklearn_mutual_information
  from tensorflow_data_validation.utils import test_util
  from google.protobuf import text_format
  from tensorflow.python.util.protobuf import compare
  from tensorflow_metadata.proto.v0 import schema_pb2
  from tensorflow_metadata.proto.v0 import statistics_pb2
label_desc: Use image node_id 1 for calling the AssignToPartitionTest obj's underlying member method code with example usage: obj.test_partitioner() without return types