class_docstr: Tests for NonStreamingCustomStatsGenerator.
class_docstr_tok: ["Tests","for","NonStreamingCustomStatsGenerator","."]
code: def test_row_and_column_partitions_reassemble(self):
      # We'd like to test the row/column partitioning behavior in a non-trivial
      # condition for column partitioning. This test skips the actual MI
      # calculation, and just verifies that RecordBatches passed to it are as we
      # expect.
  
      # Column names chosen so that
      # yams:     partition 0
      # arugula:  partition 0
      # apple:    partition 1
      #
      # Note that partition indices should be deterministic.
      batch1 = pa.RecordBatch.from_arrays(
          [
              pa.array([1]),
              pa.array([2]),
              pa.array(["a"]),
          ],
          ["yams", "arugula", "label_key"],
      )
      batch2 = pa.RecordBatch.from_arrays(
          [
              pa.array([3]),
              pa.array(["b"]),
          ],
          ["yams", "label_key"],
      )
      batch3 = pa.RecordBatch.from_arrays(
          [
              pa.array([4]),
              pa.array(["c"]),
          ],
          ["apple", "label_key"],
      )
  
      merged = table_util.MergeRecordBatches(
          [batch1, batch2, batch3]
      ).to_pandas()
  
      mi = mutual_information.MutualInformation(
          label_feature=types.FeaturePath(["label_key"]),
          schema=self.schema,
          max_encoding_length=TEST_MAX_ENCODING_LENGTH,
          column_partitions=3,
          seed=TEST_SEED,
      )
  
      def _make_equal_dataframe_items(expected):
          """Compare lists of dataframes without considering order or count."""
  
          def _assert_fn(dataframes):
              got_expected = [False] * len(expected)
              got_actual = [False] * len(dataframes)
              for i, dfi in enumerate(expected):
                  for j, dfj in enumerate(dataframes):
                      # Sort by the label to account for non-deterministic PCollection
                      # order, and reorder columns for consistency.
                      dfi = dfi.sort_values("label_key")
                      dfi = dfi[list(sorted(dfi.columns))].reset_index(
                          drop=True
                      )
  
                      dfj = dfj.sort_values("label_key")
                      dfj = dfj[list(sorted(dfj.columns))].reset_index(
                          drop=True
                      )
                      if dfi.equals(dfj):
                          got_expected[i] = True
                          got_actual[j] = True
              self.assertTrue(
                  min(got_expected),
                  msg="some expected outputs missing\ngot: %s\nexpected: %s"
                  % (dataframes, expected),
              )
              self.assertTrue(
                  min(got_actual),
                  msg="some actual outputs not expected\ngot: %s\nexpected: %s"
                  % (dataframes, expected),
              )
  
          return _assert_fn
  
      with beam.Pipeline() as p:
          result = (
              p
              | beam.Create([("", batch1), ("", batch2), ("", batch3)])
              | mi.partitioner(1)
              | beam.CombinePerKey(
                  partitioned_stats_generator._SampleRecordBatchRows(
                      999
                  )
              )
              | beam.Map(lambda x: x[1].to_pandas())
          )
          # Note that the batches passed to MI compute are column-wise slices of
          # the merged RecordBatch.
          beam_test_util.assert_that(
              result,
              _make_equal_dataframe_items(
                  [
                      merged[["yams", "arugula", "label_key"]],
                      merged[["apple", "label_key"]],
                      merged[["label_key"]],
                  ]
              ),
          )
  
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 types
  from tensorflow_data_validation.statistics.generators import mutual_information
  from tensorflow_data_validation.statistics.generators import partitioned_stats_generator
  from tensorflow_data_validation.utils import test_util
  from tfx_bsl.arrow import table_util
  from google.protobuf import text_format
  from tensorflow_metadata.proto.v0 import schema_pb2
  from tensorflow_metadata.proto.v0 import statistics_pb2
label_desc: Use image node_id 5 for calling the NonStreamingCustomStatsGeneratorTest obj's underlying member method code with example usage: obj.test_row_and_column_partitions_reassemble() and returns: _assert_fn