code: def enumerate_arrays(
      record_batch: pa.RecordBatch,
      example_weight_map: Optional[ExampleWeightMap],
      enumerate_leaves_only: bool,
      wrap_flat_struct_in_list: bool = True,
      column_select_fn: Optional[
          Callable[[types.FeatureName], bool]
      ] = None,
  ) -> Iterable[
      Tuple[types.FeaturePath, pa.Array, Optional[np.ndarray]]
  ]:
      """Enumerates arrays in a RecordBatch.
  
      Define:
        primitive: primitive arrow arrays (e.g. Int64Array).
        nested_list := list<nested_list> | list<primitive> | null
        # note: a null array can be seen as a list<primitive>, which contains only
        #   nulls and the type of the primitive is unknown.
        # example:
        #   null,
        #   list<null>,  # like list<list<unknown_type>> with only null values.
        #   list<list<int64>>,
        struct := struct<{field: nested_list | struct}> | list<struct>
        # example:
        #   struct<{"foo": list<int64>},
        #   list<struct<{"foo": list<int64>}>>,
        #   struct<{"foo": struct<{"bar": list<list<int64>>}>}>
  
      This function assumes `record_batch` contains only nested_list and struct
      columns. It enumerates each column in `record_batch`, and if that column is
      a struct, it flattens the outer lists wrapping it (if any), and recursively
      enumerates the array of each field in the struct (also see
      `enumerate_leaves_only`).
  
      The weights get "aligned" automatically in this process, therefore weights,
      the third term in the returned tuple always has enumerated_array[i]'s weight
      being weights[i].
  
      A FeaturePath is included in the result to address the enumerated array.
      Note that the FeaturePath merely addresses in the `record_batch` and struct
      arrays. It does not indicate whether / how a struct array is nested.
  
      Args:
        record_batch: The RecordBatch whose arrays to be visited.
        example_weight_map: an ExampleWeightMap that maps a FeaturePath to its
          corresponding weight column.
        enumerate_leaves_only: If True, only enumerate leaf arrays. A leaf array
          is an array whose type does not have any struct nested in.
          Otherwise, also enumerate the struct arrays where the leaf arrays are
          contained.
        wrap_flat_struct_in_list: if True, and if a struct<[Ts]> array is
          encountered, it will be wrapped in a list array, so it becomes a
          list<struct<[Ts]>>, in which each sub-list contains one element.
          A caller can make use of this option to assume all the arrays enumerated
          here are list<inner_type>.
        column_select_fn: If provided, only enumerates leaf arrays of columns with
          names for which this function evaluates to True.
      Yields:
        A tuple. The first term is the path of the feature; the second term is
        the feature array and the third term is the weight array for the feature
        array (i.e. weights[i] is the weight for array[i]).
  
      Raises:
        ValueError: When the weight column is not a list array whose elements are
          1-element lists.
      """
  
      def _recursion_helper(
          feature_path: types.FeaturePath,
          array: pa.Array,
          all_weights: Dict[types.FeatureName, np.ndarray],
      ) -> Iterable[
          Tuple[types.FeaturePath, pa.Array, Optional[np.ndarray]]
      ]:
          """Recursion helper."""
          array_type = array.type
          innermost_nested_type = array_util.get_innermost_nested_type(
              array_type
          )
          if pa.types.is_struct(innermost_nested_type):
              if not enumerate_leaves_only:
                  weights = all_weights.get(
                      example_weight_map.get(feature_path)
                  )
                  # special handing for a flat struct array -- wrap it in a ListArray
                  # whose elements are singleton lists. This way downstream can keep
                  # assuming the enumerated arrays are list<*>.
                  to_yield = array
                  if (
                      pa.types.is_struct(array_type)
                      and wrap_flat_struct_in_list
                  ):
                      to_yield = array_util.ToSingletonListArray(array)
                  yield (feature_path, to_yield, weights)
              (
                  flat_struct_array,
                  parent_indices,
              ) = array_util.flatten_nested(array, bool(all_weights))
              # Potential optimization:
              # Only flatten weights that we know will be used in the recursion.
              flat_all_weights = {
                  weight_feature_name: w[parent_indices]
                  for weight_feature_name, w in all_weights.items()
              }
              for field in flat_struct_array.type:
                  field_name = field.name
                  yield from _recursion_helper(
                      feature_path.child(field_name),
                      array_util.get_field(
                          flat_struct_array, field_name
                      ),
                      flat_all_weights,
                  )
          else:
              weights = all_weights.get(
                  example_weight_map.get(feature_path)
              )
              yield (feature_path, array, weights)
  
      if example_weight_map is None:
          example_weight_map = ExampleWeightMap(
              weight_feature=None, per_feature_override=None
          )
      all_weights = {
          weight_column: get_weight_feature(record_batch, weight_column)
          for weight_column in example_weight_map.all_weight_features()
      }
  
      for column_name, column in zip(
          record_batch.schema.names, record_batch.columns
      ):
          if column_select_fn and not column_select_fn(column_name):
              continue
          yield from _recursion_helper(
              types.FeaturePath([column_name]), column, all_weights
          )
  
file_path: data-validation/tensorflow_data_validation/arrow/arrow_util.py
imports: from typing import Callable, Dict, Iterable, Optional, Text, Tuple
  import numpy
  import pyarrow
  from tensorflow_data_validation import types
  from tensorflow_data_validation.utils.example_weight_map import ExampleWeightMap
  from tfx_bsl.arrow import array_util
label_desc: Use image node_id 3 for calling a global function with example usage: enumerate_arrays(record_batch, example_weight_map, enumerate_leaves_only, wrap_flat_struct_in_list, column_select_fn) without return types