code: def _flatten_and_impute(
      examples: pa.RecordBatch,
      categorical_features: Set[types.FeaturePath],
  ) -> Dict[types.FeaturePath, np.ndarray]:
      """Flattens and imputes the values in the input Arrow RecordBatch.
  
      Replaces missing values with _CATEGORICAL_FEATURE_IMPUTATION_FILL_VALUE
      for categorical features and 10*max(feature_values) for numeric features.
      We impute missing values with an extreme value that is far from observed
      values so it does not incorrectly impact KNN results. 10*max(feature_values)
      is used instead of sys.max_float because max_float is large enough to cause
      unexpected float arithmetic errors.
  
      Args:
        examples: Arrow RecordBatch containing a batch of examples where all
          features are univalent.
        categorical_features: Set of categorical feature names.
  
      Returns:
        A Dict[FeaturePath, np.ndarray] where the key is the feature path and the
        value is a 1D numpy array corresponding to the feature values.
      """
      num_rows = examples.num_rows
      result = {}
      for column_name, feature_array in zip(
          examples.schema.names, examples.columns
      ):
          feature_path = types.FeaturePath([column_name])
          imputation_fill_value = (
              _CATEGORICAL_FEATURE_IMPUTATION_FILL_VALUE
              if feature_path in categorical_features
              else sys.maxsize
          )
          if pa.types.is_null(feature_array.type):
              # If null array, impute all values.
              imputed_values_array = np.full(
                  shape=num_rows, fill_value=imputation_fill_value
              )
              result[feature_path] = imputed_values_array
          else:
              # to_pandas returns a readonly array. Create a copy as we will be imputing
              # the NaN values.
              (
                  flattened_array,
                  non_missing_parent_indices,
              ) = array_util.flatten_nested(
                  feature_array, return_parent_indices=True
              )
              assert non_missing_parent_indices is not None
              non_missing_values = np.copy(np.asarray(flattened_array))
              is_categorical_feature = (
                  feature_path in categorical_features
              )
              result_dtype = non_missing_values.dtype
              if (
                  non_missing_parent_indices.size < num_rows
                  and is_categorical_feature
              ):
                  result_dtype = object
              flattened_array = np.ndarray(
                  shape=num_rows, dtype=result_dtype
              )
              num_values = np.asarray(
                  array_util.ListLengthsFromListArray(feature_array)
              )
              missing_parent_indices = np.where(num_values == 0)[0]
              if feature_path not in categorical_features:
                  # Also impute any NaN values.
                  nan_mask = np.isnan(non_missing_values)
                  if not np.all(nan_mask):
                      imputation_fill_value = (
                          non_missing_values[~nan_mask].max() * 10
                      )
                  non_missing_values[
                      nan_mask.nonzero()[0]
                  ] = imputation_fill_value
              flattened_array[
                  non_missing_parent_indices
              ] = non_missing_values
              if missing_parent_indices.any():
                  flattened_array[
                      missing_parent_indices
                  ] = imputation_fill_value
              result[feature_path] = flattened_array
      return result
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import sys
  from typing import Dict, List, Optional, Sequence, Set, Text, Union
  import numpy
  import pandas
  import pyarrow
  from tensorflow_data_validation import types
  from tensorflow_data_validation.arrow import arrow_util
  from tensorflow_data_validation.statistics.generators import partitioned_stats_generator
  from tensorflow_data_validation.utils import schema_util
  from tensorflow_data_validation.utils import stats_util
  from tfx_bsl.arrow import array_util
  from tensorflow_metadata.proto.v0 import schema_pb2
  from tensorflow_metadata.proto.v0 import statistics_pb2
label_desc: Use image node_id 1 for calling a global function with example usage: _flatten_and_impute(examples, categorical_features) and returns: result