Transform prensors to structured prensors.

There are three limitations:
1. The optional/repeated distinction is lost.
2. If the prensor has fields that do not conform to the structured tensor requirements,
   the behavior is undefined.
3. There is no way yet to represent multidimensional fields. This can be handled in a
   future extension.

PiperOrigin-RevId: 309090341

Author: mzinkevi

google3/third_party/py/struct2tensor/prensor_to_structured_tensor.py

@@ -0,0 +1,113 @@
+# Copyright 2020 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Converting a prensor to a structured tensor.
+
+Structured tensors are almost completely more general than a prensor.
+There are only four issues:
+
+1. There is no disambiguation between an optional and a repeated field in a
+   structured test. This is lost.
+2. The set of field names is more restricted in a structured tensor. This is
+   ignored for now: behavior of disallowed field names is undefined.
+3. We could try to keep more shape information, especially if the size of the
+   root prensor is known statically.
+4. TODO:interpret a given field name to represent "anonymous" dimensions, and
+   create multidimensional fields when that field name is used.
+
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+
+from __future__ import print_function
+import six
+from struct2tensor import path
+from struct2tensor import prensor
+import tensorflow as tf
+from typing import Mapping, Union
+
+from tensorflow.python.ops.ragged.row_partition import RowPartition  # pylint: disable=g-direct-tensorflow-import
+from tensorflow.python.ops.structured import structured_tensor  # pylint: disable=g-direct-tensorflow-import
+
+
+def prensor_to_structured_tensor(
+    p: prensor.Prensor) -> structured_tensor.StructuredTensor:
+  """Creates a structured tensor from a prensor.
+
+  All information about optional and repeated fields is dropped.
+  If the field names in the proto do not meet the specifications for
+  StructuredTensor, the behavior is undefined.
+
+  Args:
+    p: the prensor to convert.
+
+  Returns:
+    An equivalent StructuredTensor.
+
+  Raises:
+    ValueError: if the root of the prensor is not a RootNodeTensor.
+  """
+  node = p.node
+  if isinstance(node, prensor.RootNodeTensor):
+    return _root_node_to_structured_tensor(
+        _prensor_to_field_map(p.get_children(), node.size))
+  raise ValueError("Must be a root prensor")
+
+
+def _root_node_to_structured_tensor(
+    fields: Mapping[path.Step, prensor.Prensor]
+) -> structured_tensor.StructuredTensor:
+  """Convert a map of prensors to a structured tensor."""
+  return structured_tensor.StructuredTensor.from_fields(
+      fields=fields, shape=tf.TensorShape([None]))
+
+
+def _prensor_to_structured_tensor_helper(
+    p: prensor.Prensor, nrows: tf.Tensor
+) -> Union[tf.RaggedTensor, structured_tensor.StructuredTensor]:
+  """Convert a prensor to a structured tensor with a certain number of rows."""
+  node = p.node
+  if isinstance(node, prensor.LeafNodeTensor):
+    return _leaf_node_to_ragged_tensor(node, nrows)
+  assert isinstance(node, prensor.ChildNodeTensor)
+  return _child_node_to_structured_tensor(
+      node, _prensor_to_field_map(p.get_children(), node.size), nrows)
+
+
+def _prensor_to_field_map(
+    p_fields: Mapping[path.Step, prensor.Prensor],
+    nrows: tf.Tensor) -> Mapping[path.Step, structured_tensor.StructuredTensor]:
+  """Convert a map of prensors to map of structured tensors."""
+  return {
+      step: _prensor_to_structured_tensor_helper(child, nrows)
+      for step, child in six.iteritems(p_fields)
+  }
+
+
+def _child_node_to_structured_tensor(
+    node: prensor.ChildNodeTensor, fields: Mapping[path.Step, prensor.Prensor],
+    nrows: tf.Tensor) -> structured_tensor.StructuredTensor:
+  """Convert a map of prensors to map of structured tensors."""
+  st = structured_tensor.StructuredTensor.from_fields(
+      fields=fields, shape=tf.TensorShape([None]))
+  row_partition = RowPartition.from_value_rowids(
+      value_rowids=node.parent_index, nrows=nrows)
+  return st.partition_outer_dimension(row_partition)
+
+
+def _leaf_node_to_ragged_tensor(node: prensor.LeafNodeTensor,
+                                nrows: tf.Tensor) -> tf.RaggedTensor:
+  """Converts a LeafNodeTensor to a 2D ragged tensor."""
+  return tf.RaggedTensor.from_value_rowids(
+      values=node.values, value_rowids=node.parent_index, nrows=nrows)

google3/third_party/py/struct2tensor/prensor_to_structured_tensor_test.py

@@ -0,0 +1,108 @@
+# Lint as: python3
+# Copyright 2020 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for StructuredTensor."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from struct2tensor import prensor
+from struct2tensor import prensor_to_structured_tensor
+from struct2tensor.test import prensor_test_util
+
+from tensorflow.python.framework import test_util  # pylint: disable=g-direct-tensorflow-import
+from tensorflow.python.ops.structured import structured_tensor  # pylint: disable=g-direct-tensorflow-import
+from tensorflow.python.platform import unittest  # pylint: disable=g-direct-tensorflow-import
+
+
+def _make_structured_tensor(shape, fields):
+  return structured_tensor.StructuredTensor.from_fields(
+      fields=fields, shape=shape)
+
+
+# pylint: disable=g-long-lambda
+# @test_util.run_all_in_graph_and_eager_modes
+class PrensorToStructuredTensorTest(test_util.TensorFlowTestCase):
+
+  def test_simple_prensor(self):
+    pren = prensor_test_util.create_simple_prensor()
+    st = prensor_to_structured_tensor.prensor_to_structured_tensor(pren)
+    self.assertAllEqual(st._fields["foo"], [[9], [8], [7]])
+    self.assertAllEqual(st._fields["foorepeated"], [[9], [8, 7], [6]])
+
+  def test_nested_prensor(self):
+    """Tests on a deep expression."""
+    pren = prensor_test_util.create_nested_prensor()
+    st = prensor_to_structured_tensor.prensor_to_structured_tensor(pren)
+    self.assertAllEqual(st._fields["doc"]._fields["bar"],
+                        [[[b"a"]], [[b"b", b"c"], [b"d"]], []])
+    self.assertAllEqual(st._fields["doc"]._fields["keep_me"],
+                        [[[False]], [[True], []], []])
+    self.assertAllEqual(st._fields["user"]._fields["friends"],
+                        [[[b"a"]], [[b"b", b"c"], [b"d"]], [[b"e"]]])
+
+  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._fields["foo"], [[9], [8], [7]])
+    self.assertAllEqual(st._fields["foorepeated"], [[9], [8, 7], [6]])
+    self.assertAllEqual(st._fields["doc"]._fields["keep_me"],
+                        [[[False]], [[True], []], []])
+    self.assertAllEqual(st._fields["user"]._fields["friends"],
+                        [[[b"a"]], [[b"b", b"c"], [b"d"]], [[b"e"]]])
+    self.assertAllEqual(st._fields["doc"]._fields["bar"],
+                        [[[b"a"]], [[b"b", b"c"], [b"d"]], []])
+
+  def test_deep_prensor(self):
+    """Test a prensor with three layers: root, event, and doc.
+
+      a prensor expression representing:
+      {foo:9, foorepeated:[9], user:[{friends:["a"]}],
+       event:{doc:[{bar:["a"], keep_me:False}]}}
+      {foo:8, foorepeated:[8,7],
+       event:{doc:[{bar:["b","c"], keep_me:True},{bar:["d"]}]},
+       user:[{friends:["b", "c"]}, {friends:["d"]}]}
+      {foo:7, foorepeated:[6], user:[friends:["e"]], event:{}}
+    """
+    pren = prensor_test_util.create_deep_prensor()
+    st = prensor_to_structured_tensor.prensor_to_structured_tensor(pren)
+    self.assertAllEqual(st._fields["foo"], [[9], [8], [7]])
+    self.assertAllEqual(st._fields["foorepeated"], [[9], [8, 7], [6]])
+    self.assertAllEqual(st._fields["user"]._fields["friends"],
+                        [[[b"a"]], [[b"b", b"c"], [b"d"]], [[b"e"]]])
+    self.assertAllEqual(st._fields["event"]._fields["doc"]._fields["bar"],
+                        [[[[b"a"]]], [[[b"b", b"c"], [b"d"]]], [[]]])
+    self.assertAllEqual(st._fields["event"]._fields["doc"]._fields["keep_me"],
+                        [[[[False]]], [[[True], []]], [[]]])
+
+  def test_non_root_prensor(self):
+    child_prensor = prensor.create_prensor_from_root_and_children(
+        prensor_test_util.create_child_node([0, 0, 1, 3, 7], True), {})
+    with self.assertRaisesRegexp(ValueError, "Must be a root prensor"):
+      prensor_to_structured_tensor.prensor_to_structured_tensor(child_prensor)
+
+
+if __name__ == "__main__":
+  unittest.main()