def testConvertToLargeRecordBatch(
      self, passthrough_data, expected_unary_features
  ):
    """Tests slicing of large transformed batches during conversion."""
    # Any Beam test pipeline handling elements this large crashes the program
    # with OOM (even with 28GB memory available), so we test the conversion
    # pretty narrowly.

    # 2^31 elements in total.
    num_values = 1 << 21
    batch_dict = {
        'a': np.zeros([_LARGE_BATCH_SIZE, num_values], np.float32),
        **passthrough_data,
    }
    schema = schema_utils.schema_from_feature_spec(
        {'a': tf.io.FixedLenFeature([num_values], tf.float32)}
    )
    converter = impl_helper.make_tensor_to_arrow_converter(schema)
    arrow_schema = pa.schema(
        [
            ('a', pa.large_list(pa.float32())),
        ]
        + [(key, value.type) for key, value in passthrough_data.items()]
    )
    input_metadata = tensor_adapter.TensorAdapterConfig(arrow_schema, {})
    actual_num_rows = 0
    actual_num_batches = 0
    # Features are either going to be in the `record_batch` or in
    # `unary_features`.
    record_batch_features = set(batch_dict.keys()) - set(
        expected_unary_features.keys()
    )
    for record_batch, unary_features in impl._convert_to_record_batch(
        batch_dict, converter, set(passthrough_data.keys()), input_metadata
    ):
      self.assertEqual(set(record_batch.schema.names), record_batch_features)
      self.assertEqual(unary_features, expected_unary_features)
      self.assertLessEqual(
          record_batch.nbytes, impl._MAX_TRANSFORMED_BATCH_BYTES_SIZE
      )
      actual_num_rows += record_batch.num_rows
      actual_num_batches += 1
    self.assertEqual(actual_num_rows, _LARGE_BATCH_SIZE)
    self.assertGreater(actual_num_batches, 1)
Tests slicing of large transformed batches during conversion.
tensorflow_transform/beam/impl_output_record_batches_test.py