docstr: Tests slicing of large transformed batches during conversion.
docstr_tok: ["Tests","slicing","of","large","transformed","batches","during","conversion","."]
code: 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)
  
file_path: transform/tensorflow_transform/beam/impl_output_record_batches_test.py
imports: import collections
  import numpy
  import pyarrow
  import tensorflow
  from tensorflow_transform import impl_helper
  from tensorflow_transform.beam import impl
  from tensorflow_transform.beam import impl_test
  from tensorflow_transform.beam import tft_unit
  from tensorflow_transform.tf_metadata import schema_utils
  from tfx_bsl.tfxio import tensor_adapter
label_desc: Use image node_id 4 for calling the BeamImplOutputRecordBatchesTest obj's underlying member method code with example usage: obj.testConvertToLargeRecordBatch(passthrough_data, expected_unary_features) without return types