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