Add support for protos in from_py.

PiperOrigin-RevId: 858223810
Change-Id: Ide1ea43f958851e748431dde471801d611dd99d6

Author: atorero

py/koladata/base/BUILD

@@ -227,6 +227,31 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "py_proto_utils",
+    srcs = ["py_proto_utils.cc"],
+    hdrs = ["py_proto_utils.h"],
+    deps = [
+        "//koladata:data_bag",
+        "//koladata:data_slice",
+        "//koladata/internal:data_item",
+        "//koladata/internal:data_slice",
+        "//koladata/internal:dtype",
+        "//koladata/operators:lib",
+        "//koladata/proto:from_proto",
+        "//py/koladata/types:pybind11_protobuf_wrapper",
+        "@com_google_absl//absl/base:nullability",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings:string_view",
+        "@com_google_absl//absl/types:span",
+        "@com_google_arolla//arolla/util",
+        "@com_google_arolla//arolla/util:status_backport",
+        "@com_google_arolla//py/arolla/py_utils",
+        "@com_google_protobuf//:protobuf",
+        "@rules_python//python/cc:current_py_cc_headers",  # buildcleaner: keep
+    ],
+)
+
 koladata_py_extension(
     name = "py_functors_base_py_ext",
     srcs = ["py_functors_base_module.cc"],

py/koladata/base/py_conversions/BUILD

@@ -76,6 +76,8 @@ cc_library(
         "//koladata/internal:schema_attrs",
         "//koladata/internal/op_utils:trampoline_executor",
         "//py/koladata/base:boxing",
+        "//py/koladata/base:py_proto_utils",
+        "//py/koladata/types:pybind11_protobuf_wrapper",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/base:no_destructor",
         "@com_google_absl//absl/base:nullability",

py/koladata/base/py_conversions/from_py.cc

@@ -53,6 +53,8 @@
 #include "koladata/uuid_utils.h"
 #include "py/koladata/base/boxing.h"
 #include "py/koladata/base/py_conversions/dataclasses_util.h"
+#include "py/koladata/base/py_proto_utils.h"
+#include "py/koladata/types/pybind11_protobuf_wrapper.h"
 #include "arolla/util/status_macros_backport.h"
 
 namespace koladata::python {
@@ -228,6 +230,21 @@ class FromPyConverter {
     return false;
   }
 
+  // Returns true if the given Python objects should be treated as a proto.
+  // Checks the first non-None object.
+  absl::StatusOr<bool> IsProto(const std::vector<PyObject*>& py_objects) {
+    if (py_objects.empty()) {
+      return false;
+    }
+    for (PyObject* py_obj : py_objects) {
+      if (Py_IsNone(py_obj)) {
+        continue;
+      }
+      return IsPyProtoMessage(py_obj);
+    }
+    return false;
+  }
+
   // Verifies that dict_as_obj is not set for dict schema.
   absl::Status VerifyDictAsObj(const std::optional<DataSlice>& schema) {
     if (dict_as_obj_ && schema && schema->IsDictSchema()) {
@@ -512,6 +529,12 @@ class FromPyConverter {
                           cur_depth, executor, result);
     }
 
+    ASSIGN_OR_RETURN(bool is_proto, IsProto(py_objects));
+    if (is_proto) {
+      return ConvertProto(py_objects, std::move(cur_shape), schema,
+                          std::move(itemid), result);
+    }
+
     if (IsEntity(py_objects, schema)) {
       DCHECK(schema.has_value());
       return ConvertEntities(py_objects, std::move(cur_shape), *schema,
@@ -1002,6 +1025,23 @@ class FromPyConverter {
     return absl::OkStatus();
   }
 
+  absl::Status ConvertProto(const std::vector<PyObject*>& py_objects,
+                            DataSlice::JaggedShape cur_shape,
+                            std::optional<DataSlice> schema,
+                            std::optional<DataSlice> itemid,
+                            std::optional<DataSlice>& result) {
+    if (itemid.has_value()) {
+      itemid = itemid->Flatten();
+    }
+
+    ASSIGN_OR_RETURN(DataSlice proto_slice,
+                     FromProtoObjects(GetBag(), py_objects, /*extensions=*/{},
+                                      /*itemid=*/itemid, /*schema=*/schema));
+
+    ASSIGN_OR_RETURN(result, proto_slice.Reshape(std::move(cur_shape)));
+    return absl::OkStatus();
+  }
+
   // Converts Python objects (dataclasses) to Objects/Entities.
   // If the schema is kObject, creates an Object DataSlice. Otherwise (which
   // can only happen if the schema is nullopt), creates an Entity DataSlice.

py/koladata/base/py_proto_utils.cc

@@ -0,0 +1,77 @@
+// 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.
+//
+#include "py/koladata/base/py_proto_utils.h"
+
+#include <any>
+#include <optional>
+#include <utility>
+#include <vector>
+
+#include "absl/base/nullability.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
+#include "absl/types/span.h"
+#include "arolla/util/unit.h"
+#include "koladata/data_bag.h"
+#include "koladata/data_slice.h"
+#include "koladata/internal/data_item.h"
+#include "koladata/internal/dtype.h"
+#include "koladata/internal/slice_builder.h"
+#include "koladata/operators/slices.h"
+#include "koladata/proto/from_proto.h"
+#include "google/protobuf/message.h"
+#include "py/arolla/py_utils/py_utils.h"
+#include "py/koladata/types/pybind11_protobuf_wrapper.h"
+#include "arolla/util/status_macros_backport.h"
+
+namespace koladata::python {
+absl::StatusOr<DataSlice> FromProtoObjects(
+    const absl_nonnull DataBagPtr& db, const std::vector<PyObject*>& py_objects,
+    absl::Span<const absl::string_view> extensions,
+    const std::optional<DataSlice>& itemid,
+    const std::optional<DataSlice>& schema) {
+  arolla::python::DCheckPyGIL();
+  const Py_ssize_t messages_list_len = py_objects.size();
+
+  internal::SliceBuilder message_mask_builder(messages_list_len);
+  auto typed_message_mask_builder = message_mask_builder.typed<arolla::Unit>();
+  std::vector<std::any> message_owners;
+  message_owners.reserve(messages_list_len);
+  std::vector<const ::google::protobuf::Message* absl_nonnull> message_ptrs;
+  message_ptrs.reserve(messages_list_len);
+  for (Py_ssize_t i = 0; i < messages_list_len; ++i) {
+    PyObject* py_message = py_objects[i];  // Borrowed.
+    if (py_message != Py_None) {
+      typed_message_mask_builder.InsertIfNotSet(i, arolla::kUnit);
+
+      ASSIGN_OR_RETURN((auto [message_ptr, message_owner]),
+                       python::UnwrapPyProtoMessage(py_message));
+      message_owners.push_back(std::move(message_owner));
+      message_ptrs.push_back(message_ptr);
+    }
+  }
+  ASSIGN_OR_RETURN(
+      auto message_mask,
+      DataSlice::Create(std::move(message_mask_builder).Build(),
+                        DataSlice::JaggedShape::FlatFromSize(messages_list_len),
+                        internal::DataItem(schema::kMask)));
+
+  ASSIGN_OR_RETURN(DataSlice dense_result,
+                   FromProto(db, message_ptrs, extensions, itemid, schema));
+  ASSIGN_OR_RETURN(DataSlice result,
+                   ops::InverseSelect(dense_result, message_mask));
+  return result;
+}
+}  // namespace koladata::python

py/koladata/base/py_proto_utils.h

@@ -0,0 +1,44 @@
+// 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.
+//
+// Utilities for converting Python proto messages to Koda objects.
+
+#ifndef KOLADATA_BASE_PY_PROTO_UTILS_H_
+#define KOLADATA_BASE_PY_PROTO_UTILS_H_
+
+#include <optional>
+#include <vector>
+
+#include "Python.h"
+#include "absl/base/nullability.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
+#include "absl/types/span.h"
+#include "koladata/data_bag.h"
+#include "koladata/data_slice.h"
+
+namespace koladata::python {
+// Treats Python objects (potentially sparse) as proto Messages and converts
+// them to Koda objects (to missing elements if None).
+// Basically, calls `FromProto` on the non-None objects, and sets the missing
+// elements for None elements. If an element is neither a proto Message nor
+// None, raises an error.
+absl::StatusOr<DataSlice> FromProtoObjects(
+    const absl_nonnull DataBagPtr& db, const std::vector<PyObject*>& py_objects,
+    absl::Span<const absl::string_view> extensions,
+    const std::optional<DataSlice>& itemid,
+    const std::optional<DataSlice>& schema);
+}  // namespace koladata::python
+
+#endif  // KOLADATA_BASE_PY_PROTO_UTILS_H_

py/koladata/functions/tests/BUILD

@@ -565,10 +565,12 @@ py_test(
     name = "from_py_test",
     srcs = ["from_py_test.py"],
     deps = [
+        ":test_py_pb2",
         "//py:python_path",
         "//py/koladata/functions",
         "//py/koladata/functions:attrs",
         "//py/koladata/functions:object_factories",
+        "//py/koladata/functions:proto_conversions",
         "//py/koladata/functions:py_conversions",
         "//py/koladata/operators:kde_operators",
         "//py/koladata/testing",