code: def _ExtractFeatures(  # pylint: disable=invalid-name
      extracts: beam.pvalue.PCollection,
      tensor_representations: Mapping[
          str, schema_pb2.TensorRepresentation
      ],
  ) -> beam.pvalue.PCollection:
      """Extracts features from extracts.
  
      Args:
        extracts: PCollection containing features under tfma.ARROW_RECORD_BATCH_KEY
          or tfma.FEATURES_KEY.
        tensor_representations: Tensor representations.
  
      Returns:
        PCollection of extracts with additional features added under the key
        tfma.FEATURES_KEY and optionally inputs added under the tfma.INPUTS_KEY.
  
      Raises:
        ValueError: If incoming extracts contains neither tfma.FEATURES_KEY nor
          tfma.ARROW_RECORD_BATCH_KEY.
      """
  
      def extract_features(extracts: types.Extracts) -> types.Extracts:
          """Extract features from extracts containing arrow table."""
          result = copy.copy(extracts)
          if constants.ARROW_RECORD_BATCH_KEY in extracts:
              (
                  record_batch,
                  serialized_examples,
              ) = _drop_unsupported_columns_and_fetch_raw_data_column(
                  extracts[constants.ARROW_RECORD_BATCH_KEY]
              )
              del result[constants.ARROW_RECORD_BATCH_KEY]
              features = (
                  result[constants.FEATURES_KEY]
                  if constants.FEATURES_KEY in result
                  else {}
              )
              features.update(
                  util.record_batch_to_tensor_values(
                      record_batch, tensor_representations
                  )
              )
              result[constants.FEATURES_KEY] = features
              result[constants.INPUT_KEY] = serialized_examples
          elif constants.FEATURES_KEY not in extracts:
              raise ValueError(
                  "Incoming extracts must contain either tfma.ARROW_RECORD_BATCH_KEY "
                  f"or tfma.FEATURES_KEY, but extracts={extracts}"
              )
          return result
  
      return extracts | "ExtractFeatures" >> beam.Map(extract_features)
  
file_path: model-analysis/tensorflow_model_analysis/extractors/features_extractor.py
imports: import copy
  from typing import Mapping, Optional, Text, Tuple
  import apache_beam
  import numpy
  import pyarrow
  from tensorflow_model_analysis import constants
  from tensorflow_model_analysis.api import types
  from tensorflow_model_analysis.extractors import extractor
  from tensorflow_model_analysis.proto import config_pb2
  from tensorflow_model_analysis.utils import util
  from tensorflow_metadata.proto.v0 import schema_pb2
label_desc: Use image node_id 6 for calling a global function with example usage: _ExtractFeatures(extracts, tensor_representations) and returns: unknown, result