def _map_prensor_impl(root, root_path, paths_needed, operation, is_repeated,
dtype, new_field_name):
"""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 testPromoteAndProjectExpression(self):
filenames = [
"struct2tensor/testdata/parquet_testdata/dremel_example.parquet"
]
batch_size = 2
exp = parquet.create_expression_from_parquet_file(filenames)
new_exp = promote.promote(exp, path.Path(["Name", "Language", "Code"]),
"new_code")
new_code_project_exp = project.project(
new_exp, [path.Path(["Name", "new_code"])])
docid_project_exp = project.project(exp, [path.Path(["DocId"])])
pqds = parquet.calculate_parquet_values(
[new_code_project_exp, docid_project_exp], exp, filenames,
batch_size)
new_code_expected = prensor.create_prensor_from_descendant_nodes({
path.Path([]):
prensor.RootNodeTensor(tf.constant(2, dtype=tf.int64)),
path.Path(["Name"]):
prensor.ChildNodeTensor(tf.constant([0, 0, 0, 1], dtype=tf.int64),
True),
path.Path(["Name", "new_code"]):
prensor.LeafNodeTensor(tf.constant([0, 0, 2], dtype=tf.int64),
tf.constant([b"en-us", b"en", b"en-gb"]),
True)
})
docid_expected = prensor.create_prensor_from_descendant_nodes({
path.Path([]):
prensor.RootNodeTensor(tf.constant(2, dtype=tf.int64)),
path.Path(["DocId"]):
prensor.LeafNodeTensor(tf.constant([0, 1], dtype=tf.int64),
tf.constant([10, 20], dtype=tf.int64),
False)
})
for ele in pqds:
new_code_pren = ele[0]
docid_pren = ele[1]
self._assertPrensorEqual(new_code_pren, new_code_expected)
self._assertPrensorEqual(docid_pren, docid_expected)
def testPlaceholderExpression(self):
pren = prensor_test_util.create_nested_prensor()
expected_pren = prensor.create_prensor_from_descendant_nodes({
path.Path([]):
prensor.RootNodeTensor(tf.constant(3, dtype=tf.int64)),
path.Path(["new_friends"]):
prensor.LeafNodeTensor(
tf.constant([0, 1, 1, 1, 2], dtype=tf.int64),
tf.constant(["a", "b", "c", "d", "e"], dtype=tf.string), True)
})
root_schema = mpp.create_schema(is_repeated=True,
children={
"doc": {
"is_repeated": True,
"children": {
"bar": {
"is_repeated": True,
"dtype": tf.string
},
"keep_me": {
"is_repeated": False,
"dtype": tf.bool
}
}
},
"user": {
"is_repeated": True,
"children": {
"friends": {
"is_repeated": True,
"dtype": tf.string
}
}
}
})
exp = placeholder.create_expression_from_schema(root_schema)
promote_exp = promote.promote(exp, path.Path(["user", "friends"]),
"new_friends")
project_exp = project.project(promote_exp,
[path.Path(["new_friends"])])
new_friends_exp = project_exp.get_descendant(path.Path(["new_friends"
]))
result = calculate.calculate_values([new_friends_exp],
feed_dict={exp: pren})
res_node = result[0]
exp_node = expected_pren.get_descendant(path.Path(["new_friends"
])).node
self.assertAllEqual(res_node.is_repeated, exp_node.is_repeated)
self.assertAllEqual(res_node.values, exp_node.values)
self.assertAllEqual(res_node.parent_index, exp_node.parent_index)
def test_project(self):
expr = create_expression.create_expression_from_prensor(
prensor_test_util.create_nested_prensor())
projected = project.project(
expr,
[path.Path(["user", "friends"]),
path.Path(["doc", "keep_me"])])
self.assertIsNotNone(
projected.get_descendant(path.Path(["user", "friends"])))
self.assertIsNotNone(
projected.get_descendant(path.Path(["doc", "keep_me"])))
self.assertIsNone(projected.get_descendant(path.Path(["doc", "bar"])))
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