code: def test_generate_statistics_in_memory_valid_custom_generator(self):
      # CombinerStatsGenerator that returns a DatasetFeatureStatistic proto with
      # custom stat.
      class CustomCombinerStatsGenerator(
          stats_generator.CombinerStatsGenerator
      ):
          def create_accumulator(self):
              return 0
  
          def add_input(self, accumulator, input_batch):
              return 0
  
          def merge_accumulators(self, accumulators):
              return 0
  
          def extract_output(self, accumulator):
              stats_proto = statistics_pb2.DatasetFeatureStatistics()
              proto_feature = stats_proto.features.add()
              proto_feature.path.step[:] = ["a"]
              custom_stat = proto_feature.custom_stats.add()
              custom_stat.name = "custom_stat"
              custom_stat.str = "custom_stat_value"
              return stats_proto
  
      record_batch = pa.RecordBatch.from_arrays(
          [pa.array([[b"xyz", b"qwe"], [b"qwe"], [b"qwe"]])], ["a"]
      )
  
      expected_result = text_format.Parse(
          """
          datasets {
            num_examples: 3
            features {
              path {
                step: "a"
              }
              type: STRING
              custom_stats {
                name: 'custom_stat'
                str: 'custom_stat_value'
              }
              string_stats {
                avg_length: 3
                unique: 2
                common_stats {
                  num_non_missing: 3
                  min_num_values: 1
                  max_num_values: 2
                  avg_num_values: 1.333333
                  tot_num_values: 4
                }
                top_values {
                  value: 'qwe'
                  frequency: 3
                }
                top_values {
                  value: 'xyz'
                  frequency: 1
                }
                rank_histogram {
                  buckets {
                    low_rank: 0
                    high_rank: 0
                    label: "qwe"
                    sample_count: 3.0
                  }
                  buckets {
                    low_rank: 1
                    high_rank: 1
                    label: "xyz"
                    sample_count: 1.0
                  }
                }
              }
            }
          }""",
          statistics_pb2.DatasetFeatureStatisticsList(),
      )
  
      options = stats_options.StatsOptions(
          generators=[
              CustomCombinerStatsGenerator("CustomStatsGenerator")
          ],
          num_top_values=4,
          num_rank_histogram_buckets=3,
          num_values_histogram_buckets=3,
          enable_semantic_domain_stats=True,
      )
      result = stats_impl.generate_statistics_in_memory(
          record_batch, options
      )
      test_util.assert_dataset_feature_stats_proto_equal(
          self,
          result.datasets[0],
          expected_result.datasets[0],
          check_histograms=False,
      )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import copy
  import sys
  from typing import Iterable
  import unittest
  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.statistics import stats_impl
  from tensorflow_data_validation.statistics import stats_options
  from tensorflow_data_validation.statistics.generators import cross_feature_stats_generator
  from tensorflow_data_validation.statistics.generators import stats_generator
  from tensorflow_data_validation.utils import slicing_util
  from tensorflow_data_validation.utils import test_util
  from tfx_bsl.arrow import array_util
  from tfx_bsl.arrow import table_util
  from tfx_bsl.public.proto import slicing_spec_pb2
  from tfx_bsl.statistics import merge_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 9 for calling the StatsImplTest obj's underlying member method code with example usage: obj.test_generate_statistics_in_memory_valid_custom_generator() and returns: int, int, int, stats_proto