def test_reroot_and_create_proto_index(self):
expr = create_expression.create_expression_from_prensor(
prensor_test_util.create_big_prensor()).reroot(
"doc").create_proto_index("proto_index")
proto_index = expr.get_child("proto_index")
new_field = expr.get_child("bar")
leaf_node = expression_test_util.calculate_value_slowly(new_field)
proto_index_node = expression_test_util.calculate_value_slowly(
proto_index)
self.assertIsNotNone(new_field)
self.assertTrue(new_field.is_repeated)
self.assertEqual(new_field.type, tf.string)
self.assertTrue(new_field.is_leaf)
self.assertEqual(new_field.known_field_names(), frozenset())
self.assertEqual(leaf_node.values.dtype, tf.string)
self.assertIsNotNone(proto_index)
self.assertFalse(proto_index.is_repeated)
self.assertEqual(proto_index.type, tf.int64)
self.assertTrue(proto_index.is_leaf)
self.assertEqual(proto_index.known_field_names(), frozenset())
self.assertEqual(proto_index_node.values.dtype, tf.int64)
self.assertAllEqual([b"a", b"b", b"c", b"d"], leaf_node.values)
self.assertAllEqual([0, 1, 1, 2], leaf_node.parent_index)
self.assertAllEqual([0, 1, 1], proto_index_node.values)
self.assertAllEqual([0, 1, 2], proto_index_node.parent_index)
def test_promote_with_schema_dense_fraction(self):
"""Test when min_fraction is not 1."""
s = prensor_test_util.create_big_prensor_schema()
feature_dict = {feature.name: feature for feature in s.feature}
user_feature = feature_dict["user"]
user_feature.value_count.min = 3
user_feature.value_count.max = 3
user_feature.presence.min_fraction = 1
user_dict = {
feature.name: feature
for feature in user_feature.struct_domain.feature
}
friends_feature = user_dict["friends"]
friends_feature.presence.min_fraction = 0.9
expr = create_expression.create_expression_from_prensor(
prensor_test_util.create_big_prensor()).apply_schema(s)
new_root, new_field = promote.promote_anonymous(
expr, path.Path(["user", "friends"]))
new_field = new_root.get_descendant_or_error(new_field)
new_schema_feature = new_field.schema_feature
self.assertIsNotNone(new_schema_feature)
self.assertEqual(new_schema_feature.presence.min_fraction, 0.3)
def test_filter_by_child(self):
"""Tests filter_by_child."""
root = create_expression.create_expression_from_prensor(
prensor_test_util.create_big_prensor())
root_2 = filter_expression.filter_by_child(root,
path.create_path("doc"),
"keep_me", "new_doc")
[result] = calculate.calculate_prensors([root_2])
self.assertAllEqual(
result.get_descendant_or_error(path.Path(["new_doc"
])).node.parent_index,
[1])
self.assertAllEqual(
result.get_descendant_or_error(path.Path(["new_doc", "keep_me"
])).node.parent_index,
[0])
self.assertAllEqual(
result.get_descendant_or_error(path.Path(["new_doc",
"keep_me"])).node.values,
[True])
self.assertAllEqual(
result.get_descendant_or_error(path.Path(["new_doc", "bar"
])).node.parent_index,
[0, 0])
self.assertAllEqual(
result.get_descendant_or_error(path.Path(["new_doc",
"bar"])).node.values,
[b"b", b"c"])
def test_apply_schema(self):
expr = create_expression.create_expression_from_prensor(
prensor_test_util.create_big_prensor())
expr2 = expr.apply_schema(
prensor_test_util.create_big_prensor_schema())
foo_expr = expr2.get_descendant(path.Path(["foo"]))
self.assertIsNotNone(foo_expr)
foorepeated_expr = expr2.get_descendant(path.Path(["foorepeated"]))
self.assertIsNotNone(foorepeated_expr)
doc_bar_expr = expr2.get_descendant(path.Path(["doc", "bar"]))
self.assertIsNotNone(doc_bar_expr)
# Test that a domain already in the feature is maintained.
self.assertEqual(foo_expr.schema_feature.int_domain.max, 10)
# Test that an int_domain specified at the schema level is inserted
# correctly.
self.assertEqual(foorepeated_expr.schema_feature.int_domain.max, 10)
# Test that a string_domain specified at the schema level is inserted
# correctly.
self.assertEqual(doc_bar_expr.schema_feature.string_domain.value[0],
"a")
self.assertIsNotNone(
expr2.get_descendant(path.Path(["user", "friends"])))
self.assertIsNotNone(
expr2.get_descendant(path.Path(["doc", "keep_me"])))
def test_prensor_children_ordered(self):
def _recursively_check_sorted(p):
self.assertEqual(list(p.get_children().keys()),
sorted(p.get_children().keys()))
for c in p.get_children().values():
_recursively_check_sorted(c)
for pren in [
prensor_test_util.create_nested_prensor(),
prensor_test_util.create_big_prensor(),
prensor_test_util.create_deep_prensor()
]:
_recursively_check_sorted(pren)
p = prensor.create_prensor_from_descendant_nodes({
path.Path([]):
prensor_test_util.create_root_node(1),
path.Path(["d"]):
prensor_test_util.create_optional_leaf_node([0], [True]),
path.Path(["c"]):
prensor_test_util.create_optional_leaf_node([0], [True]),
path.Path(["b"]):
prensor_test_util.create_optional_leaf_node([0], [True]),
path.Path(["a"]):
prensor_test_util.create_optional_leaf_node([0], [True]),
})
self.assertEqual(["a", "b", "c", "d"], list(p.get_children().keys()))
def test_get_schema_missing_features(self):
# The expr has a number of features: foo, foorepeated, doc, user.
expr = create_expression.create_expression_from_prensor(
prensor_test_util.create_big_prensor())
# The schema has only a subset of the features on the expr.
schema = schema_pb2.Schema()
feature = schema.feature.add()
feature.name = "foo"
feature.type = schema_pb2.FeatureType.INT
feature.value_count.min = 1
feature.value_count.max = 1
feature = schema.feature.add()
feature.name = "foorepeated"
feature.type = schema_pb2.FeatureType.INT
feature.value_count.min = 0
feature.value_count.max = 5
feature = schema.feature.add()
feature.name = "doc"
feature.type = schema_pb2.FeatureType.STRUCT
feature.struct_domain.feature.append(
schema_pb2.Feature(name="keep_me",
type=schema_pb2.FeatureType.INT))
# By default, the output schema has all features present in the expr.
expr = expr.apply_schema(schema)
output_schema = expr.get_schema()
self.assertNotEqual(schema, output_schema)
self.assertLen(schema.feature, 3)
self.assertLen(output_schema.feature, 4)
# With create_schema_features = False, only features on the original schema
# propogate to the new schema.
output_schema = expr.get_schema(create_schema_features=False)
self.assertLen(output_schema.feature, 3)
def test_slice_end(self):
with self.session(use_gpu=False) as sess:
root = create_expression.create_expression_from_prensor(
prensor_test_util.create_big_prensor())
root_2 = slice_expression.slice_expression(root,
path.Path(["doc"]),
"new_doc", None, 1)
result = prensor_value.materialize(
calculate.calculate_prensors([root_2])[0], sess)
self.assertAllEqual(
result.get_descendant_or_error(path.Path(
["new_doc"])).node.parent_index, [0, 1])
self.assertAllEqual(
result.get_descendant_or_error(
path.Path(["new_doc", "keep_me"])).node.parent_index,
[0, 1])
self.assertAllEqual(
result.get_descendant_or_error(
path.Path(["new_doc", "keep_me"])).node.values,
[False, True])
self.assertAllEqual(
result.get_descendant_or_error(path.Path(
["new_doc", "bar"])).node.parent_index, [0, 1, 1])
self.assertAllEqual(
result.get_descendant_or_error(path.Path(["new_doc",
"bar"])).node.values,
[b"a", b"b", b"c"])
def test_create_has_field(self):
expr = create_expression.create_expression_from_prensor(
prensor_test_util.create_big_prensor())
new_root = expr.create_has_field("doc.keep_me", "result")
new_field = new_root.get_descendant_or_error(
path.Path(["doc", "result"]))
leaf_node = expression_test_util.calculate_value_slowly(new_field)
self.assertAllEqual(leaf_node.parent_index, [0, 1, 2])
self.assertAllEqual(leaf_node.values, [True, True, False])
def test_size_missing_value(self):
expr = create_expression.create_expression_from_prensor(
prensor_test_util.create_big_prensor())
new_root = size.size(expr, path.Path(["doc", "keep_me"]), "result")
new_field = new_root.get_descendant_or_error(
path.Path(["doc", "result"]))
leaf_node = expression_test_util.calculate_value_slowly(new_field)
self.assertAllEqual(leaf_node.parent_index, [0, 1, 2])
self.assertAllEqual(leaf_node.values, [1, 1, 0])
def test_size_anonymous(self):
expr = create_expression.create_expression_from_prensor(
prensor_test_util.create_big_prensor())
new_root, new_path = size.size_anonymous(expr,
path.Path(["doc", "bar"]))
new_field = new_root.get_descendant_or_error(new_path)
leaf_node = expression_test_util.calculate_value_slowly(new_field)
self.assertAllEqual(leaf_node.parent_index, [0, 1, 2])
self.assertAllEqual(leaf_node.values, [1, 2, 1])
def test_broadcast_substructure(self):
"""Tests broadcast of a submessage.
The result of broadcasting `user` into `doc` looks like:
{
foo: 9,
foorepeated: [9],
doc: [{bar:["a"], keep_me:False, new_user: [{friends:["a"]}]}],
user: [{friends:["a"]}]
},
{
foo: 8,
foorepeated: [8, 7],
doc: [
{
bar: ["b","c"],
keep_me: True,
new_user: [{friends:["b", "c"]},{friends:["d"]}]
},
{
bar: ["d"],
new_user: [{friends:["b", "c"]},{friends:["d"]}]
}
],
user: [{friends:["b", "c"]},{friends:["d"]}],
},
{
foo: 7,
foorepeated: [6],
user: [{friends:["e"]}]
}
"""
expr = create_expression.create_expression_from_prensor(
prensor_test_util.create_big_prensor())
new_root = broadcast.broadcast(expr, path.Path(["user"]), "doc", "new_user")
new_user = new_root.get_child("doc").get_child("new_user")
self.assertIsNotNone(new_user)
self.assertTrue(new_user.is_repeated)
self.assertIsNone(new_user.type)
self.assertFalse(new_user.is_leaf)
new_user_node = expression_test_util.calculate_value_slowly(new_user)
self.assertAllEqual(new_user_node.parent_index, [0, 1, 1, 2, 2])
self.assertAllEqual(new_user_node.index_to_value, [0, 1, 2, 1, 2])
new_friends = new_user.get_child("friends")
self.assertIsNotNone(new_friends)
self.assertTrue(new_friends.is_repeated)
self.assertEqual(new_friends.type, tf.string)
self.assertTrue(new_friends.is_leaf)
new_friends_node = expression_test_util.calculate_value_slowly(new_friends)
self.assertEqual(new_friends_node.values.dtype, tf.string)
self.assertAllEqual(new_friends_node.values,
["a", "b", "c", "d", "b", "c", "d"])
self.assertAllEqual(new_friends_node.parent_index, [0, 1, 1, 2, 3, 3, 4])
def test_broadcast_and_calculate(self):
"""Tests get_sparse_tensors on a deep tree."""
expr = create_expression.create_expression_from_prensor(
prensor_test_util.create_big_prensor())
new_root, new_path = broadcast.broadcast_anonymous(
expr, path.Path(["foo"]), "user")
new_field = new_root.get_descendant_or_error(new_path)
leaf_node = expression_test_util.calculate_value_slowly(new_field)
self.assertAllEqual(leaf_node.parent_index, [0, 1, 2, 3])
self.assertAllEqual(leaf_node.values, [9, 8, 8, 7])
def test_get_schema_no_schema(self):
expr = create_expression.create_expression_from_prensor(
prensor_test_util.create_big_prensor())
output_schema = expr.get_schema()
self.assertLen(output_schema.feature, 4)
# With create_schema_features = False, the schema will be empty.
output_schema = expr.get_schema(create_schema_features=False)
self.assertEmpty(output_schema.feature)
self.assertEqual(schema_pb2.Schema(), output_schema)
def test_promote_with_schema(self):
expr = create_expression.create_expression_from_prensor(
prensor_test_util.create_big_prensor()).apply_schema(
prensor_test_util.create_big_prensor_schema())
new_root, new_field = promote.promote_anonymous(
expr, path.Path(["user", "friends"]))
new_field = new_root.get_descendant_or_error(new_field)
new_schema_feature = new_field.schema_feature
self.assertIsNotNone(new_schema_feature)
self.assertEqual(new_schema_feature.string_domain.value[0], "a")
def test_create_size_field(self):
with self.session(use_gpu=False) as sess:
expr = create_expression.create_expression_from_prensor(
prensor_test_util.create_big_prensor())
new_root = expr.create_size_field("doc.bar", "result")
new_field = new_root.get_descendant_or_error(path.Path(["doc", "result"]))
leaf_node = expression_test_util.calculate_value_slowly(new_field)
[parent_index,
values] = sess.run([leaf_node.parent_index, leaf_node.values])
self.assertAllEqual(parent_index, [0, 1, 2])
self.assertAllEqual(values, [1, 2, 1])
def test_slice_mask_begin_negative(self):
root = create_expression.create_expression_from_prensor(
prensor_test_util.create_big_prensor())
root_2, new_path = slice_expression._get_slice_mask(
root, path.Path(["doc"]), -1, None)
result = calculate.calculate_prensors([root_2])[0]
self.assertAllEqual(
result.get_descendant_or_error(new_path).node.parent_index,
[0, 1, 1])
self.assertAllEqual(
result.get_descendant_or_error(new_path).node.values,
[True, False, True])
def test_get_schema(self):
"""Integration test between get_schema and apply_schema."""
expr = create_expression.create_expression_from_prensor(
prensor_test_util.create_big_prensor())
expr2 = expr.apply_schema(prensor_test_util.create_big_prensor_schema())
schema_result = expr2.get_schema()
feature_map = _features_as_map(schema_result.feature)
self.assertIn("foo", feature_map)
# Test that a domain already in the feature is maintained.
self.assertEqual(feature_map["foo"].int_domain.max, 10)
# Test that an int_domain specified at the schema level is inserted
# correctly.
self.assertEqual(feature_map["foorepeated"].int_domain.max, 10)
def test_broadcast(self):
"""Tests broadcast.broadcast(...), and indirectly tests set_path."""
expr = create_expression.create_expression_from_prensor(
prensor_test_util.create_big_prensor())
new_root = expr.broadcast("foo", "user", "new_field")
new_field = new_root.get_child("user").get_child("new_field")
self.assertIsNotNone(new_field)
self.assertFalse(new_field.is_repeated)
self.assertEqual(new_field.type, tf.int32)
self.assertTrue(new_field.is_leaf)
leaf_node = expression_test_util.calculate_value_slowly(new_field)
self.assertEqual(leaf_node.values.dtype, tf.int32)
self.assertEqual(new_field.known_field_names(), frozenset())
def test_get_schema_empty_schema(self):
expr = create_expression.create_expression_from_prensor(
prensor_test_util.create_big_prensor())
schema = schema_pb2.Schema()
# By default, the output schema has all features present in the proto.
expr = expr.apply_schema(schema)
output_schema = expr.get_schema()
self.assertEmpty(schema.feature)
self.assertLen(output_schema.feature, 4)
# With create_schema_features = False, the schema will be empty.
output_schema = expr.get_schema(create_schema_features=False)
self.assertEmpty(output_schema.feature)
self.assertEqual(schema, output_schema)
def test_slice_mask_begin_positive(self):
with self.session(use_gpu=False) as sess:
root = create_expression.create_expression_from_prensor(
prensor_test_util.create_big_prensor())
root_2, new_path = slice_expression._get_slice_mask(
root, path.Path(["doc"]), 1, None)
result = prensor_value.materialize(
calculate.calculate_prensors([root_2])[0], sess)
self.assertAllEqual(
result.get_descendant_or_error(new_path).node.parent_index,
[0, 1, 1])
self.assertAllEqual(
result.get_descendant_or_error(new_path).node.values,
[False, False, True])
def test_prensor_is_composite_tensor(self):
for pren in [
prensor_test_util.create_nested_prensor(),
prensor_test_util.create_big_prensor(),
prensor_test_util.create_deep_prensor()
]:
flattened_tensors = tf.nest.flatten(pren, expand_composites=True)
self.assertIsInstance(flattened_tensors, list)
for t in flattened_tensors:
self.assertIsInstance(t, tf.Tensor)
packed_pren = tf.nest.pack_sequence_as(pren,
flattened_tensors,
expand_composites=True)
self._assert_prensor_equals(pren, packed_pren)
def test_promote_and_broadcast_test(self):
expr = create_expression.create_expression_from_prensor(
prensor_test_util.create_big_prensor())
new_root = expr.promote_and_broadcast({"new_field": "user.friends"},
path.Path(["doc"]))
new_field = new_root.get_descendant_or_error(
path.Path(["doc", "new_field"]))
self.assertTrue(new_field.is_repeated)
self.assertEqual(new_field.type, tf.string)
self.assertTrue(new_field.is_leaf)
self.assertTrue(new_field.calculation_equal(new_field))
self.assertFalse(new_field.calculation_equal(expr))
leaf_node = expression_test_util.calculate_value_slowly(new_field)
self.assertEqual(leaf_node.values.dtype, tf.string)
self.assertEqual(new_field.known_field_names(), frozenset())
def test_promote_and_broadcast_anonymous(self):
"""A basic promote and broadcast."""
expr = create_expression.create_expression_from_prensor(
prensor_test_util.create_big_prensor())
new_root, p = promote_and_broadcast.promote_and_broadcast_anonymous(
expr, path.Path(["user", "friends"]), path.Path(["doc"]))
new_field = new_root.get_descendant_or_error(p)
self.assertTrue(new_field.is_repeated)
self.assertEqual(new_field.type, tf.string)
self.assertTrue(new_field.is_leaf)
self.assertTrue(new_field.calculation_equal(new_field))
self.assertFalse(new_field.calculation_equal(expr))
leaf_node = expression_test_util.calculate_value_slowly(new_field)
self.assertEqual(leaf_node.values.dtype, tf.string)
self.assertEqual(new_field.known_field_names(), frozenset())
def test_apply_empty_schema(self):
"""Test that applying an empty schema does not filter out paths."""
expr = create_expression.create_expression_from_prensor(
prensor_test_util.create_big_prensor())
expr2 = expr.apply_schema(schema_pb2.Schema())
foo_expr = expr2.get_descendant(path.Path(["foo"]))
self.assertIsNotNone(foo_expr)
foorepeated_expr = expr2.get_descendant(path.Path(["foorepeated"]))
self.assertIsNotNone(foorepeated_expr)
doc_bar_expr = expr2.get_descendant(path.Path(["doc", "bar"]))
self.assertIsNotNone(doc_bar_expr)
known_field_names = expr2.known_field_names()
self.assertIn("doc", known_field_names)
self.assertIn("foo", known_field_names)
self.assertIn("foorepeated", known_field_names)
self.assertIn("user", known_field_names)
def test_broadcast_anonymous(self):
expr = create_expression.create_expression_from_prensor(
prensor_test_util.create_big_prensor())
new_root, p = broadcast.broadcast_anonymous(expr, path.Path(["foo"]),
"user")
new_field = new_root.get_descendant_or_error(p)
self.assertFalse(new_field.is_repeated)
self.assertEqual(new_field.type, tf.int32)
self.assertTrue(new_field.is_leaf)
self.assertTrue(new_field.calculation_equal(new_field))
self.assertFalse(new_field.calculation_equal(expr))
leaf_node = expression_test_util.calculate_value_slowly(new_field)
self.assertEqual(leaf_node.values.dtype, tf.int32)
self.assertEqual(new_field.known_field_names(), frozenset())
sources = new_field.get_source_expressions()
self.assertLen(sources, 2)
self.assertIs(expr.get_child("foo"), sources[0])
self.assertIs(expr.get_child("user"), sources[1])
def test_big_prensor(self):
"""Test the big prensor.
a prensor expression representing:
{foo:9, foorepeated:[9], doc:[{bar:["a"], keep_me:False}],
user:[{friends:["a"]}]}
{foo:8, foorepeated:[8,7],
doc:[{bar:["b","c"],keep_me:True},{bar:["d"]}],
user:[{friends:["b", "c"]},{friends:["d"]}],}
{foo:7, foorepeated:[6], user:[friends:["e"]]}
"""
pren = prensor_test_util.create_big_prensor()
st = prensor_to_structured_tensor.prensor_to_structured_tensor(pren)
self.assertAllEqual(st.field_value("foo"), [[9], [8], [7]])
self.assertAllEqual(st.field_value("foorepeated"), [[9], [8, 7], [6]])
self.assertAllEqual(st.field_value(["doc", "keep_me"]),
[[[False]], [[True], []], []])
self.assertAllEqual(st.field_value(["user", "friends"]),
[[[b"a"]], [[b"b", b"c"], [b"d"]], [[b"e"]]])
self.assertAllEqual(st.field_value(["doc", "bar"]),
[[[b"a"]], [[b"b", b"c"], [b"d"]], []])
def test_promote_with_schema_dense_parent(self):
s = prensor_test_util.create_big_prensor_schema()
feature_dict = {feature.name: feature for feature in s.feature}
user_feature = feature_dict["user"]
user_feature.value_count.min = 3
user_feature.value_count.max = 3
user_feature.presence.min_fraction = 1
user_feature.lifecycle_stage = schema_pb2.LifecycleStage.ALPHA
user_dict = {
feature.name: feature
for feature in user_feature.struct_domain.feature
}
friends_feature = user_dict["friends"]
friends_feature.value_count.min = 2
friends_feature.value_count.max = 2
friends_feature.presence.min_fraction = 1
friends_feature.presence.min_count = 10
friends_feature.lifecycle_stage = schema_pb2.LifecycleStage.BETA
friends_feature.distribution_constraints.min_domain_mass = 0.5
expr = create_expression.create_expression_from_prensor(
prensor_test_util.create_big_prensor()).apply_schema(s)
new_root, new_field = promote.promote_anonymous(
expr, path.Path(["user", "friends"]))
[new_field] = new_root.get_descendant_or_error(
new_field).get_source_expressions()
new_schema_feature = new_field.schema_feature
self.assertIsNotNone(new_schema_feature)
self.assertEqual(new_schema_feature.string_domain.value[0], "a")
self.assertEqual(new_schema_feature.value_count.max, 6)
self.assertEqual(new_schema_feature.value_count.min, 6)
self.assertEqual(new_schema_feature.presence.min_fraction, 1)
self.assertEqual(new_schema_feature.presence.min_count, 3)
self.assertEqual(new_schema_feature.lifecycle_stage,
schema_pb2.LifecycleStage.ALPHA)
self.assertEqual(
new_schema_feature.distribution_constraints.min_domain_mass, 0.5)
def _check_lifecycle_stage(a, b):
s = prensor_test_util.create_big_prensor_schema()
feature_dict = {feature.name: feature for feature in s.feature}
user_feature = feature_dict["user"]
if a is not None:
user_feature.lifecycle_stage = a
user_dict = {
feature.name: feature
for feature in user_feature.struct_domain.feature
}
friends_feature = user_dict["friends"]
if b is not None:
friends_feature.lifecycle_stage = b
expr = create_expression.create_expression_from_prensor(
prensor_test_util.create_big_prensor()).apply_schema(s)
new_root, new_field = promote.promote_anonymous(
expr, path.Path(["user", "friends"]))
new_field = new_root.get_descendant_or_error(new_field)
return new_field.schema_feature.lifecycle_stage
def test_reroot_and_create_proto_index(self):
expr = create_expression.create_expression_from_prensor(
prensor_test_util.create_big_prensor())
new_root = reroot.reroot(expr, path.Path(["doc"]))
proto_index = reroot.create_proto_index_field(
new_root, "proto_index").get_child("proto_index")
new_field = new_root.get_child("bar")
leaf_node = expression_test_util.calculate_value_slowly(new_field)
proto_index_node = expression_test_util.calculate_value_slowly(proto_index)
self.assertIsNotNone(new_field)
self.assertTrue(new_field.is_repeated)
self.assertEqual(new_field.type, tf.string)
self.assertTrue(new_field.is_leaf)
self.assertEqual(new_field.known_field_names(), frozenset())
self.assertEqual(leaf_node.values.dtype, tf.string)
self.assertIsNotNone(proto_index)
self.assertFalse(proto_index.is_repeated)
self.assertEqual(proto_index.type, tf.int64)
self.assertTrue(proto_index.is_leaf)
self.assertEqual(proto_index.known_field_names(), frozenset())
self.assertEqual(proto_index_node.values.dtype, tf.int64)
with self.session(use_gpu=False) as sess:
[
leaf_node_values, leaf_node_parent_index, proto_index_values,
proto_index_parent_index
] = sess.run([
leaf_node.values, leaf_node.parent_index, proto_index_node.values,
proto_index_node.parent_index
])
self.assertAllEqual([b"a", b"b", b"c", b"d"], leaf_node_values)
self.assertAllEqual([0, 1, 1, 2], leaf_node_parent_index)
self.assertAllEqual([0, 1, 1], proto_index_values)
self.assertAllEqual([0, 1, 2], proto_index_parent_index)
def test_slice_begin(self):
"""Test slice with only begin specified.
Starts with:
{
foo:9,
foorepeated:[9],
doc:[{
bar:["a"],
keep_me:False
}],
user:[
{
friends:["a"]
}]
}
{foo:8,
foorepeated:[8,7],
doc:[{
bar:["b","c"],
keep_me:True
},{
bar:["d"]
}],
user:[{
friends:["b", "c"]
},{
friends:["d"]
}],
}
{foo:7,
foorepeated:[6],
user:[{friends:["e"]}]}
Creates new_doc by slicing doc[1:]:
{foo:9,
foorepeated:[9],
doc:[{
bar:["a"],
keep_me:False
}],
user:[{
friends:["a"]
}]}
{foo:8,
foorepeated:[8,7],
doc:[{
bar:["b","c"],
keep_me:True
},{
bar:["d"]
}],
new_doc[{
bar:["d"]
}],
user:[{
friends:["b", "c"]
},{
friends:["d"]}],}
{foo:7,
foorepeated:[6],
user:[{
friends:["e"]
}]}
"""
with self.session(use_gpu=False) as sess:
root = create_expression.create_expression_from_prensor(
prensor_test_util.create_big_prensor())
root_2 = slice_expression.slice_expression(root,
path.Path(["doc"]),
"new_doc", 1, None)
result = prensor_value.materialize(
calculate.calculate_prensors([root_2])[0], sess)
self.assertAllEqual(
result.get_descendant_or_error(path.Path(
["new_doc"])).node.parent_index, [1])
self.assertAllEqual(
result.get_descendant_or_error(
path.Path(["new_doc", "keep_me"])).node.parent_index, [])
self.assertAllEqual(
result.get_descendant_or_error(
path.Path(["new_doc", "keep_me"])).node.values, [])
self.assertAllEqual(
result.get_descendant_or_error(path.Path(
["new_doc", "bar"])).node.parent_index, [0])
self.assertAllEqual(
result.get_descendant_or_error(path.Path(["new_doc",
"bar"])).node.values,
[b"d"])
