code: def _GroupByQueryKey(  # pylint: disable=invalid-name
      extracts: beam.pvalue.PCollection,
      query_key: str,
  ) -> beam.pvalue.PCollection:
      """PTransform for grouping extracts by a query key.
  
      Args:
        extracts: Incoming PCollection consisting of extracts.
        query_key: Query key to group extracts by. Must be a member of the dict of
          features stored under tfma.FEATURES_KEY.
  
      Returns:
        PCollection of lists of extracts where each list is associated with same
        query key.
      """
      missing_query_key_counter = beam.metrics.Metrics.counter(
          constants.METRICS_NAMESPACE, "missing_query_key"
      )
  
      def key_by_query_key(
          extracts: types.Extracts, query_key: str
      ) -> Tuple[str, types.Extracts]:
          """Extract the query key from the extract and key by that."""
          value = metric_util.to_scalar(
              util.get_by_keys(
                  extracts,
                  [constants.FEATURES_KEY, query_key],
                  optional=True,
              ),
              tensor_name=query_key,
          )
          if value is None:
              missing_query_key_counter.inc()
              return ("", extracts)
          return ("{}".format(value), extracts)
  
      # pylint: disable=no-value-for-parameter
      return (
          extracts
          | "KeyByQueryId" >> beam.Map(key_by_query_key, query_key)
          | "GroupByKey"
          >> beam.CombinePerKey(beam.combiners.ToListCombineFn())
          | "DropQueryId" >> beam.Map(lambda kv: kv[1])
          | "MergeExtracts" >> beam.Map(util.merge_extracts)
      )
  
imports: import copy
  import datetime
  import itertools
  import numbers
  from typing import Any, Dict, Iterable, Iterator, List, NamedTuple, Optional, Tuple, Type, TypeVar, Union
  import apache_beam
  import numpy
  from tensorflow_model_analysis import constants
  from tensorflow_model_analysis.api import types
  from tensorflow_model_analysis.evaluators import counter_util
  from tensorflow_model_analysis.evaluators import eval_saved_model_util
  from tensorflow_model_analysis.evaluators import evaluator
  from tensorflow_model_analysis.evaluators import jackknife
  from tensorflow_model_analysis.evaluators import keras_util
  from tensorflow_model_analysis.evaluators import metrics_validator
  from tensorflow_model_analysis.evaluators import poisson_bootstrap
  from tensorflow_model_analysis.extractors import slice_key_extractor
  from tensorflow_model_analysis.metrics import metric_specs
  from tensorflow_model_analysis.metrics import metric_types
  from tensorflow_model_analysis.metrics import metric_util
  from tensorflow_model_analysis.proto import config_pb2
  from tensorflow_model_analysis.slicer import slicer_lib
  from tensorflow_model_analysis.utils import model_util
  from tensorflow_model_analysis.utils import util
  from tensorflow_metadata.proto.v0 import schema_pb2
label_desc: Use image node_id 3 for calling a global function with example usage: _GroupByQueryKey(extracts, query_key) and returns: unknown, str,  extracts, str,  extracts