def _recursion_helper(
feature_path: types.FeaturePath, array: pa.Array,
weights: Optional[np.ndarray]
) -> Iterable[Tuple[types.FeaturePath, pa.Array, Optional[np.ndarray]]]:
"""Recursion helper."""
array_type = array.type
if is_list_like(array_type) and pa.types.is_struct(
array_type.value_type):
if not enumerate_leaves_only:
yield (feature_path, array, weights)
flat_struct_array = array.flatten()
flat_weights = None
if weights is not None:
flat_weights = weights[
array_util.GetFlattenedArrayParentIndices(
array).to_numpy()]
for field in flat_struct_array.type:
field_name = field.name
# use "yield from" after PY 3.3.
for e in _recursion_helper(feature_path.child(field_name),
flat_struct_array.field(field_name),
flat_weights):
yield e
else:
yield (feature_path, array, weights)
def _recursion_helper(
feature_path: types.FeaturePath, array: pa.Array,
weights: Optional[np.ndarray]
) -> Iterable[Tuple[types.FeaturePath, pa.Array, Optional[np.ndarray]]]:
"""Recursion helper."""
array_type = array.type
innermost_nested_type = get_innermost_nested_type(array_type)
if pa.types.is_struct(innermost_nested_type):
if not enumerate_leaves_only:
# 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 = flatten_nested(
array, weights is not None)
flat_weights = None if weights is None else weights[parent_indices]
for field in flat_struct_array.type:
field_name = field.name
# use "yield from" after PY 3.3.
for e in _recursion_helper(feature_path.child(field_name),
flat_struct_array.field(field_name),
flat_weights):
yield e
else:
yield (feature_path, array, weights)
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 = 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 = 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), flat_struct_array.field(field_name),
flat_all_weights)
else:
weights = all_weights.get(example_weight_map.get(feature_path))
yield (feature_path, array, weights)
def _recursion_helper(
parent_path: types.FeaturePath,
container: Union[schema_pb2.Schema, schema_pb2.StructDomain]
) -> List[Tuple[types.FeaturePath, schema_pb2.SparseFeature]]:
"""Helper function that is used in finding sparse features in a tree."""
result = []
for sf in container.sparse_feature:
# Sparse features do not have a struct_domain, so they cannot be parent
# features. Thus, once this reaches a sparse feature, add it to the
# result.
result.append((parent_path.child(sf.name), sf))
for f in container.feature:
if f.type == schema_pb2.STRUCT:
result.extend(
_recursion_helper(parent_path.child(f.name),
f.struct_domain))
return result
def _recursion_helper(
parent_path: types.FeaturePath,
feature_container: Iterable[schema_pb2.Feature],
result: List[Tuple[types.FeaturePath, schema_pb2.Feature]]):
for f in feature_container:
feature_path = parent_path.child(f.name)
if f.type != schema_pb2.STRUCT:
result.append((feature_path, f))
else:
_recursion_helper(feature_path, f.struct_domain.feature, result)
