def METHOD_NAME(
    extracts: util.StandardExtracts,
    num_rows_key: str,
    array_to_truncate: np.ndarray,
    model_name: Optional[str] = None,
    allow_missing_key: Optional[bool] = False,
) -> np.ndarray:
  """Get the array to be truncated by the number of rows.
  Args:
   extracts: TFMA extracts that stores the keys.
   num_rows_key: Number of rows in each column except the paddings. For
     multi-dimensional input, it will truncate on the first dimension.
   array_to_truncate: the array to be truncated te
   model_name: The name of the model for outputs.
   allow_missing_key: (Optional) If true, it will do nothing instead of
     raising errors when col_names are not found.
  Returns:
   The array truncated by the number of rows.
  Raises:
   KeyError: The num_rows_key is not found in extracts.
  """
  num_of_rows = None
  dict_to_search = collections.ChainMap(
      extracts.get_predictions(model_name) or {},
      extracts.get_features() or {},
      extracts.get_transformed_features(model_name) or {})
  if num_rows_key:
    if dict_to_search and num_rows_key in dict_to_search:
      num_of_rows = dict_to_search[num_rows_key]
      if isinstance(num_of_rows, np.ndarray):
        num_of_rows = num_of_rows.item()
    else:
      if not allow_missing_key:
        raise KeyError(f"Key {num_rows_key} is not found under predictions, "
                       "features, or transformed features of the extracts."
                       "Please set allow_missing_key to True, if you want to "
                       "skip truncation instead.")
  result = array_to_truncate
  if num_of_rows and num_of_rows > 0 and len(result) > num_of_rows:
    result = result[:num_of_rows]
  return result