Add kd.to_proto_any function.

Fix sparsity handling in `kd.proto.get_proto_full_name`.

PiperOrigin-RevId: 857219089
Change-Id: I0e968a5f715ba29bd599c7e4e8280c442cbc96b8

Author: nbjohnson0

docs/api_reference.md

@@ -10804,6 +10804,25 @@ Args:
 Returns:
   A converted proto message or list of converted proto messages.</code></pre>
 
+### `kd.to_proto_any(x: DataSlice, *, descriptor_pool: DescriptorPool | None = None, deterministic: bool = False) -> Any | list[_NestedAnyMessageList] | None` {#kd.to_proto_any}
+
+<pre class="no-copy"><code class="lang-text no-auto-prettify">Converts a DataSlice or DataItem to proto Any messages.
+
+The schemas of all present values in `x` must have been derived from a proto
+schema using `from_proto` or `schema_from_proto`, so that the original names
+of the message types are embedded in the schema. Otherwise, this will fail.
+
+Args:
+  x: DataSlice to convert.
+  descriptor_pool: Overrides the descriptor pool used to look up python proto
+    message classes based on proto message type full name. If None, the
+    default descriptor pool is used.
+  deterministic: Passed to Any.Pack.
+
+Returns:
+  A proto Any message or nested list of proto Any messages with the same
+  shape as the input. Missing elements in the input are None in the output.</code></pre>
+
 ### `kd.to_proto_bytes(x, proto_path, /)` {#kd.to_proto_bytes}
 
 Alias for [kd.proto.to_proto_bytes](#kd.proto.to_proto_bytes) operator.

koladata/operators/proto.cc

@@ -304,7 +304,11 @@ absl::StatusOr<DataSlice> GetProtoFullName(const DataSlice& x) {
     // appropriate all-missing slice, we can reuse it as the return value.
     return x;
   } else {
-    ASSIGN_OR_RETURN(schema, ExpandToShape(x.GetSchema(), x.GetShape(), 0));
+    // schema = x.get_schema() & kd.has(x)
+    schema = x.GetSchema();
+    ASSIGN_OR_RETURN(auto has_x, ops::Has(x));
+    ASSIGN_OR_RETURN(schema,
+                     ops::ApplyMask(std::move(schema), std::move(has_x)));
   }
 
   ASSIGN_OR_RETURN(auto schema_metadata,

py/koladata/functions/functions.py

@@ -114,6 +114,7 @@
 get_attr_names = _attrs.get_attr_names
 # TODO: b/435124266 - Add to_proto explicitly to kd.py.
 to_proto = _proto_conversions.to_proto
+to_proto_any = _proto_conversions.to_proto_any
 
 slice = _data_slice.DataSlice.from_vals  # pylint: disable=redefined-builtin
 item = _data_item.DataItem.from_vals

py/koladata/functions/proto_conversions.py

@@ -15,7 +15,7 @@
 """Koda functions for converting to and from protocol buffers."""
 
 from collections.abc import Iterator
-from typing import Type, TypeAlias, TypeVar, cast
+from typing import Any, Type, TypeAlias, TypeVar, cast
 
 from google.protobuf import any as protobuf_any
 from google.protobuf import descriptor_pool as protobuf_descriptor_pool
@@ -47,6 +47,9 @@
     | tuple['_NestedMessageContainer', ...]
     | None
 )
+_NestedAnyMessageList: TypeAlias = (
+    any_pb2.Any | list['_NestedAnyMessageList'] | None
+)
 _NestedAnyMessageContainer: TypeAlias = (
     any_pb2.Any
     | list['_NestedAnyMessageContainer']
@@ -73,7 +76,7 @@ def _flatten(x: _NestedMessageContainer) -> Iterator[message.Message | None]:
 
 # Note: could use `tree.unflatten_as`, but it's not worth adding an additional
 # third-party dependency just for this.
-def _unflatten(shape: _NestedNoneList, it: Iterator[_T]) -> list[_T]:
+def _unflatten(shape: _NestedNoneList, it: Iterator[Any]) -> Any:
   """Unflattens an iterator into a shape given by a nested list with None leaves."""
   return [_unflatten(x, it) for x in shape] if shape is not None else next(it)
 
@@ -372,3 +375,58 @@ def to_proto(
   x_shape = (x & mask_constants.missing).to_py()
   results_flat = x.flatten()._to_proto(message_class)  # pylint: disable=protected-access
   return _unflatten(x_shape, iter(results_flat))
+
+
+def to_proto_any(
+    x: data_slice.DataSlice,
+    *,
+    descriptor_pool: protobuf_descriptor_pool.DescriptorPool | None = None,
+    deterministic: bool = False,
+) -> _NestedAnyMessageList:
+  """Converts a DataSlice or DataItem to proto Any messages.
+
+  The schemas of all present values in `x` must have been derived from a proto
+  schema using `from_proto` or `schema_from_proto`, so that the original names
+  of the message types are embedded in the schema. Otherwise, this will fail.
+
+  Args:
+    x: DataSlice to convert.
+    descriptor_pool: Overrides the descriptor pool used to look up python proto
+      message classes based on proto message type full name. If None, the
+      default descriptor pool is used.
+    deterministic: Passed to Any.Pack.
+
+  Returns:
+    A proto Any message or nested list of proto Any messages with the same
+    shape as the input. Missing elements in the input are None in the output.
+  """
+  if descriptor_pool is None:
+    descriptor_pool = protobuf_descriptor_pool.Default()
+
+  x_shape = (x & mask_constants.missing).to_py()
+  x_flat = x.flatten()
+  full_names = expr_eval.eval(
+      kde_operators.kde.proto.get_proto_full_name(x_flat)
+  ).to_py()
+  results = []
+  for i, full_name in enumerate(full_names):
+    x_item = x_flat.S[i]
+    if full_name is None:
+      if not x_item.is_empty():
+        raise ValueError(
+            'to_proto_any expects entities converted from proto messages, got'
+            f' {x_item}'
+        )
+      else:
+        results.append(None)
+    else:
+      message_descriptor = descriptor_pool.FindMessageTypeByName(full_name)
+      message_type = protobuf_message_factory.GetMessageClass(
+          message_descriptor
+      )
+      m = to_proto(x_item, message_type)
+      result = any_pb2.Any()
+      result.Pack(m, deterministic=deterministic)
+      results.append(result)
+
+  return _unflatten(x_shape, iter(results))

py/koladata/functions/tests/BUILD

@@ -986,6 +986,20 @@ py_test(
     ],
 )
 
+py_test(
+    name = "to_proto_any_test",
+    srcs = ["to_proto_any_test.py"],
+    deps = [
+        ":test_py_pb2",
+        "//py:python_path",
+        "//py/koladata/functions:proto_conversions",
+        "//py/koladata/testing",
+        "//py/koladata/types:data_slice",
+        "@com_google_absl_py//absl/testing:absltest",
+        "@com_google_protobuf//:protobuf_python",
+    ],
+)
+
 proto_library(
     name = "test_proto",
     srcs = ["test.proto"],

py/koladata/functions/tests/to_proto_any_test.py

@@ -0,0 +1,125 @@
+# Copyright 2025 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.
+
+from absl.testing import absltest
+from google.protobuf import descriptor_pool
+from koladata.functions import proto_conversions
+from koladata.functions.tests import test_pb2
+from koladata.testing import testing
+from koladata.types import data_slice
+from google.protobuf import any_pb2
+
+ds = data_slice.DataSlice.from_vals
+
+
+class ToProtoAnyTest(absltest.TestCase):
+
+  def test_zero_items(self):
+    x = ds([])
+    res = proto_conversions.to_proto_any(x)
+    self.assertEqual(res, [])
+
+  def test_single_item(self):
+    m = test_pb2.MessageA(some_text='thing 1')
+    x = proto_conversions.from_proto(m)
+    res = proto_conversions.to_proto_any(x)
+    self.assertIsInstance(res, any_pb2.Any)
+    unpacked_m = test_pb2.MessageA()
+    res.Unpack(unpacked_m)
+    self.assertEqual(unpacked_m, m)
+
+  def test_single_none(self):
+    x = ds(None)
+    res = proto_conversions.to_proto_any(x)
+    self.assertIsNone(res)
+
+  def test_list_with_none(self):
+    m1 = test_pb2.MessageA(some_text='thing 1')
+    x = proto_conversions.from_proto([m1, None])
+    res = proto_conversions.to_proto_any(x)
+    self.assertIsInstance(res, list)
+    self.assertLen(res, 2)
+    self.assertIsInstance(res[0], any_pb2.Any)
+    self.assertIsNone(res[1])
+    unpacked_m1 = test_pb2.MessageA()
+    res[0].Unpack(unpacked_m1)
+    self.assertEqual(unpacked_m1, m1)
+
+  def test_multiple_messages_same_type(self):
+    m1 = test_pb2.MessageA(some_text='thing 1')
+    m2 = test_pb2.MessageA(some_text='thing 2')
+    x = proto_conversions.from_proto([m1, m2])
+    res = proto_conversions.to_proto_any(x)
+    self.assertIsInstance(res, list)
+    self.assertLen(res, 2)
+    unpacked_m1 = test_pb2.MessageA()
+    res[0].Unpack(unpacked_m1)
+    self.assertEqual(unpacked_m1, m1)
+    unpacked_m2 = test_pb2.MessageA()
+    res[1].Unpack(unpacked_m2)
+    self.assertEqual(unpacked_m2, m2)
+
+  def test_multiple_messages_different_types(self):
+    m1 = test_pb2.MessageA(some_text='thing 1')
+    m2 = test_pb2.MessageB(text='thing 2')
+    # Need to use from_proto_any to get them into the same DataSlice
+    any_m1 = any_pb2.Any()
+    any_m1.Pack(m1)
+    any_m2 = any_pb2.Any()
+    any_m2.Pack(m2)
+    x = proto_conversions.from_proto_any([any_m1, any_m2])
+
+    res = proto_conversions.to_proto_any(x)
+    self.assertIsInstance(res, list)
+    self.assertLen(res, 2)
+    unpacked_m1 = test_pb2.MessageA()
+    res[0].Unpack(unpacked_m1)
+    self.assertEqual(unpacked_m1, m1)
+    unpacked_m2 = test_pb2.MessageB()
+    res[1].Unpack(unpacked_m2)
+    self.assertEqual(unpacked_m2, m2)
+
+  def test_nested_list_input(self):
+    m1 = test_pb2.MessageA(some_text='1')
+    any_m1 = any_pb2.Any()
+    any_m1.Pack(m1)
+    m2 = test_pb2.MessageA(some_text='2')
+    any_m2 = any_pb2.Any()
+    any_m2.Pack(m2)
+    m3 = test_pb2.MessageB(text='3')
+    any_m3 = any_pb2.Any()
+    any_m3.Pack(m3)
+
+    x = proto_conversions.from_proto_any([[any_m1, None, any_m2], [], [any_m3]])
+    res = proto_conversions.to_proto_any(x)
+    testing.assert_equivalent(proto_conversions.from_proto_any(res), x)
+
+  def test_not_from_proto(self):
+    x = ds([1, 2, 3]).implode()
+    with self.assertRaisesRegex(
+        ValueError,
+        'to_proto_any expects entities converted from proto messages',
+    ):
+      proto_conversions.to_proto_any(x)
+
+  def test_empty_descriptor_pool(self):
+    m = test_pb2.MessageA(some_text='thing 1')
+    x = proto_conversions.from_proto(m)
+    pool = descriptor_pool.DescriptorPool()
+    with self.assertRaisesRegex(KeyError, 'MessageA'):
+      proto_conversions.to_proto_any(x, descriptor_pool=pool)
+
+
+if __name__ == '__main__':
+  absltest.main()

py/koladata/kd_dynamic.pyi

@@ -831,6 +831,7 @@ to_json = _operators_json.to_json
 to_none = _operators_schema.to_none
 to_object = _operators_schema.to_object
 to_proto = _functions_proto_conversions.to_proto
+to_proto_any = _functions_proto_conversions.to_proto_any
 to_proto_bytes = _operators_proto.to_proto_bytes
 to_proto_json = _operators_proto.to_proto_json
 to_schema = _operators_schema.to_schema

py/koladata/operators/kde_operators_test.py

@@ -38,6 +38,7 @@
     'from_pytree',
     'py_reference',
     'to_proto',
+    'to_proto_any',
     'to_py',
     'to_pylist',
     'to_pytree',

py/koladata/operators/tests/proto_get_proto_full_name_test.py

@@ -54,12 +54,13 @@ def test_from_proto(self):
 
   def test_from_proto_slice(self):
     m = test_pb2.MessageA(some_text='hello')
-    kd_m = proto_conversions.from_proto([m, m])
+    kd_m = proto_conversions.from_proto([m, None, m])
     result = expr_eval.eval(kde.proto.get_proto_full_name(kd_m))
     testing.assert_equal(
         result.no_bag(),
         ds([
             'koladata.functions.testing.MessageA',
+            None,
             'koladata.functions.testing.MessageA',
         ]),
     )
@@ -69,13 +70,15 @@ def test_from_proto_object_slice(self):
     m2 = test_pb2.MessageB(text='bar')
     kd_m = ds([
         fns.obj(proto_conversions.from_proto(m1)),
+        None,
         fns.obj(proto_conversions.from_proto(m2)),
     ])
     result = expr_eval.eval(kde.proto.get_proto_full_name(kd_m))
     testing.assert_equal(
         result.no_bag(),
         ds([
             'koladata.functions.testing.MessageA',
+            None,
             'koladata.functions.testing.MessageB',
         ]),
     )