def map_many_values(
root: expression.Expression, parent_path: path.Path,
source_fields: Sequence[path.Step],
operation: Callable[..., tf.Tensor], dtype: tf.DType,
new_field_name: path.Step) -> Tuple[expression.Expression, path.Path]:
"""Map multiple sibling fields into a new sibling.
All source fields must have the same shape, and the shape of the output
must be the same as well.
Args:
root: original root.
parent_path: parent path of all sources and the new field.
source_fields: source fields of the operation. Must have the same shape.
operation: operation from source_fields to new field.
dtype: type of new field.
new_field_name: name of the new field.
Returns:
The new expression and the new path as a pair.
"""
new_path = parent_path.get_child(new_field_name)
return expression_add.add_paths(
root, {
new_path:
_MapValuesExpression([
root.get_descendant_or_error(parent_path.get_child(f))
for f in source_fields
], operation, dtype)
}), new_path
def promote_and_broadcast(
root: expression.Expression, path_dictionary: Mapping[path.Step,
path.Path],
dest_path_parent: path.Path) -> expression.Expression:
"""Promote and broadcast a set of paths to a particular location.
Args:
root: the original expression.
path_dictionary: a map from destination fields to origin paths.
dest_path_parent: a map from destination strings.
Returns:
A new expression, where all the origin paths are promoted and broadcast
until they are children of dest_path_parent.
"""
result_paths = {}
# Here, we branch out and create a different tree for each field that is
# promoted and broadcast.
for field_name, origin_path in path_dictionary.items():
result_path = dest_path_parent.get_child(field_name)
new_root = _promote_and_broadcast_name(root, origin_path,
dest_path_parent, field_name)
result_paths[result_path] = new_root
# We create a new tree that has all of the generated fields from the older
# trees.
return expression_add.add_to(root, result_paths)
def _promote_and_broadcast_name(
root: expression.Expression, origin: path.Path,
dest_path_parent: path.Path,
field_name: path.Step) -> expression.Expression:
new_root, anonymous_path = promote_and_broadcast_anonymous(
root, origin, dest_path_parent)
path_result = dest_path_parent.get_child(field_name)
return expression_add.add_paths(
new_root,
{path_result: new_root.get_descendant_or_error(anonymous_path)})
def _map_prensor_impl(
root: expression.Expression, root_path: path.Path,
paths_needed: Sequence[path.Path],
operation: Callable[[prensor.Prensor, calculate_options.Options],
prensor.LeafNodeTensor], is_repeated: bool,
dtype: tf.DType,
new_field_name: path.Step) -> Tuple[expression.Expression, path.Path]:
"""Map prensor implementation."""
child_expr = root.get_descendant_or_error(root_path)
sibling_child_expr = project.project(child_expr, paths_needed)
new_field_expr = _MapPrensorExpression(sibling_child_expr, operation,
is_repeated, dtype)
new_path = root_path.get_child(new_field_name)
return expression_add.add_paths(root, {new_path: new_field_expr}), new_path
def map_prensor_to_prensor(root_expr: expression.Expression, source: path.Path,
paths_needed: Sequence[path.Path],
prensor_op: Callable[[prensor.Prensor],
prensor.Prensor],
output_schema: Schema) -> expression.Expression:
r"""Maps an expression to a prensor, and merges that prensor.
For example, suppose you have an op my_op, that takes a prensor of the form:
event
/ \
foo bar
and produces a prensor of the form my_result_schema:
event
/ \
foo2 bar2
If you give it an expression original with the schema:
session
|
event
/ \
foo bar
result = map_prensor_to_prensor(
original,
path.Path(["session","event"]),
my_op,
my_output_schema)
Result will have the schema:
session
|
event--------
/ \ \ \
foo bar foo2 bar2
Args:
root_expr: the root expression
source: the path where the prensor op is applied.
paths_needed: the paths needed for the op.
prensor_op: the prensor op
output_schema: the output schema of the op.
Returns:
A new expression where the prensor is merged.
"""
original_child = root_expr.get_descendant_or_error(source).project(
paths_needed)
prensor_child = _PrensorOpExpression(original_child, prensor_op,
output_schema)
paths_map = {
source.get_child(k): prensor_child.get_child_or_error(k)
for k in prensor_child.known_field_names()
}
result = expression_add.add_paths(root_expr, paths_map)
return result