def get_index_from_end(
t: expression.Expression, source_path: path.Path,
new_field_name: path.Step) -> Tuple[expression.Expression, path.Path]:
"""Gets the number of steps from the end of the array.
Given an array ["a", "b", "c"], with indices [0, 1, 2], the result of this
is [-3,-2,-1].
Args:
t: original expression
source_path: path in expression to get index of.
new_field_name: the name of the new field.
Returns:
The new expression and the new path as a pair.
"""
new_path = source_path.get_parent().get_child(new_field_name)
work_expr, positional_index_path = get_positional_index(
t, source_path, path.get_anonymous_field())
work_expr, size_path = size.size_anonymous(work_expr, source_path)
work_expr = expression_add.add_paths(
work_expr, {
new_path:
_PositionalIndexFromEndExpression(
work_expr.get_descendant_or_error(positional_index_path),
work_expr.get_descendant_or_error(size_path))
})
# Removing the intermediate anonymous nodes.
result = expression_add.add_to(t, {new_path: work_expr})
return result, 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 slice_expression(expr: expression.Expression, p: path.Path,
new_field_name: path.Step, begin: Optional[IndexValue],
end: Optional[IndexValue]) -> expression.Expression:
"""Creates a new subtree with a sliced expression.
This follows the pattern of python slice() method.
See module-level comments for examples.
Args:
expr: the original root expression
p: the path to the source to be sliced.
new_field_name: the name of the new subtree.
begin: beginning index
end: end index.
Returns:
A new root expression.
"""
work_expr, mask_anonymous_path = _get_slice_mask(expr, p, begin, end)
work_expr = filter_expression.filter_by_sibling(
work_expr, p, mask_anonymous_path.field_list[-1], new_field_name)
new_path = p.get_parent().get_child(new_field_name)
# We created a lot of anonymous fields and intermediate expressions. Just grab
# the final result (and its children).
return expression_add.add_to(expr, {new_path: work_expr})
def test_add_to_already_existing_path(self):
with self.assertRaises(ValueError):
root = create_expression.create_expression_from_prensor(
prensor_test_util.create_nested_prensor())
root_1 = expression_add.add_paths(
root, {
path.Path(["user", "friends_2"]):
root.get_descendant_or_error(path.Path(["user", "friends"
]))
})
root_2 = expression_add.add_paths(
root_1, {
path.Path(["user", "friends_3"]):
root_1.get_descendant_or_error(
path.Path(["user", "friends_2"]))
})
expression_add.add_to(root,
{path.Path(["user", "friends"]): root_2})
def test_add_to(self):
root = create_expression.create_expression_from_prensor(
prensor_test_util.create_nested_prensor())
root_1 = expression_add.add_paths(
root, {
path.Path(["user", "friends_2"]):
root.get_descendant_or_error(path.Path(["user", "friends"]))
})
root_2 = expression_add.add_paths(
root_1, {
path.Path(["user", "friends_3"]):
root_1.get_descendant_or_error(path.Path(["user", "friends_2"
]))
})
root_3 = expression_add.add_to(
root, {path.Path(["user", "friends_3"]): root_2})
new_field = root_3.get_descendant_or_error(
path.Path(["user", "friends_3"]))
self.assertIsNotNone(new_field)
self.assertTrue(new_field.is_repeated)
self.assertEqual(new_field.type, tf.string)
leaf_node = expression_test_util.calculate_value_slowly(new_field)
self.assertEqual(leaf_node.values.dtype, tf.string)
