Allow kd.from_proto to use extensions from both an explicit schema and the extensions argument.

Previously, the `extensions` argument was ignored if there was an explicit schema. This behavior also wasn't clearly documented at the python layer.

Also correct some outdated statements in the kd.from_proto docstring.

PiperOrigin-RevId: 857273382
Change-Id: I97c35fa4d08b241721a9b17e81767f47e597be76

Author: nbjohnson0

docs/api_reference.md

@@ -9820,14 +9820,14 @@ respectively. Enums are converted to INT32. The attribute names on the Koda
 objects match the field names in the proto definition. See below for methods
 to convert proto extensions to attributes alongside regular fields.
 
-Messages, primitive fields, repeated fields, and maps are converted to
-equivalent Koda structures. Enums are converted to ints.
+If schema is not specified or schema is kd.OBJECT, only present fields in
+`messages` are loaded and included in the converted schema. To get a schema
+that contains all fields independent of the data, use `kd.schema_from_proto`.
 
-Only present values in `messages` are added. Default and missing values are
-not used.
+Proto extensions are ignored by default unless `extensions` is specified, or
+if an explicit entity schema with parenthesized attr names is specified. If
+both are specified, we use the union of the two extension sets.
 
-Proto extensions are ignored by default unless `extensions` is specified (or
-if an explicit entity schema with parenthesized attrs is used).
 The format of each extension specified in `extensions` is a dot-separated
 sequence of field names and/or extension names, where extension names are
 fully-qualified extension paths surrounded by parentheses. This sequence of
@@ -9839,6 +9839,9 @@ default behavior of traversing all fields. For example:
   "path.to.map_field.values.(package_name.some_extension)"
   "path.(package_name.some_extension).(package_name2.nested_extension)"
 
+If an explicit entity schema attr name starts with "(" and ends with ")" it is
+also interpreted as an extension name.
+
 Extensions are looked up using the C++ generated descriptor pool, using
 `DescriptorPool::FindExtensionByName`, which requires that all extensions are
 compiled in as C++ protos. The Koda attribute names for the extension fields
@@ -9852,7 +9855,8 @@ be accessed using `.get_attr(name)'. For example,
   ds.optional_field.get_attr('(package_name.DefExtension.def_extension)')
 
 If `messages` is a single proto Message, the result is a DataItem. If it is a
-list of proto Messages, the result is an 1D DataSlice.
+nested list of proto Messages, the result is a DataSlice with the same number
+of dimensions as the nesting level.
 
 Args:
   messages: Message or nested list/tuple of Message of the same type. Any of

koladata/proto/from_proto.cc

@@ -1007,6 +1007,8 @@ absl::Status FromProtoMessage(
   struct FieldVars {
     const FieldDescriptor* absl_nonnull field_descriptor;
     absl::string_view attr_name;
+    // Whether this is an extension field that is not on the requested schema.
+    bool is_non_schema_extension = false;
     std::optional<DataSlice> value;
   };
 
@@ -1033,7 +1035,11 @@ absl::Status FromProtoMessage(
     // For explicit entity schemas, use the schema attr names as the list of
     // fields and extensions to convert.
     ASSIGN_OR_RETURN(vars->schema_attr_names, schema->GetAttrNames());
-    vars->fields.reserve(vars->schema_attr_names.size());
+    const int64_t max_num_fields =
+        vars->schema_attr_names.size() +
+        ((extension_map != nullptr) ? extension_map->extension_fields.size()
+                                    : 0);
+    vars->fields.reserve(max_num_fields);
     for (const auto& attr_name : vars->schema_attr_names) {
       if (attr_name.starts_with('(') && attr_name.ends_with(')')) {
         // Interpret attrs with parentheses as fully-qualified extension paths.
@@ -1053,11 +1059,11 @@ absl::Status FromProtoMessage(
               field->full_name(), message_descriptor.full_name(),
               field->containing_type()->full_name()));
         }
-        vars->fields.emplace_back(field, attr_name);
+        vars->fields.emplace_back(field, attr_name, false);
       } else {
         const auto* field = message_descriptor.FindFieldByName(attr_name);
         if (field != nullptr) {
-          vars->fields.emplace_back(field, attr_name);
+          vars->fields.emplace_back(field, attr_name, false);
         }
       }
     }
@@ -1071,19 +1077,23 @@ absl::Status FromProtoMessage(
     for (int i_field = 0; i_field < message_descriptor.field_count();
          ++i_field) {
       const auto* field = message_descriptor.field(i_field);
-      vars->fields.emplace_back(field, field->name());
+      vars->fields.emplace_back(field, field->name(), false);
     }
-    if (extension_map != nullptr) {
-      for (const auto& [attr_name, field] : extension_map->extension_fields) {
-        if (field->containing_type() != &message_descriptor) {
-          return absl::InvalidArgumentError(absl::StrFormat(
-              "extension \"%s\" exists, but isn't an extension "
-              "on target message type \"%s\", expected \"%s\"",
-              field->full_name(), message_descriptor.full_name(),
-              field->containing_type()->full_name()));
-        }
-        vars->fields.emplace_back(field, attr_name);
+  }
+  if (extension_map != nullptr) {
+    for (const auto& [attr_name, field] : extension_map->extension_fields) {
+      if (vars->schema_attr_names.contains(attr_name)) {
+        // Already handled by the schema.
+        continue;
+      }
+      if (field->containing_type() != &message_descriptor) {
+        return absl::InvalidArgumentError(absl::StrFormat(
+            "extension \"%s\" exists, but isn't an extension "
+            "on target message type \"%s\", expected \"%s\"",
+            field->full_name(), message_descriptor.full_name(),
+            field->containing_type()->full_name()));
       }
+      vars->fields.emplace_back(field, attr_name, true);
     }
   }
 
@@ -1098,7 +1108,8 @@ absl::Status FromProtoMessage(
   for (auto& field_vars : vars->fields) {
     RETURN_IF_ERROR(FromProtoField(
         db, field_vars.attr_name, field_vars.attr_name,
-        *field_vars.field_descriptor, messages, itemid, schema,
+        *field_vars.field_descriptor, messages, itemid,
+        field_vars.is_non_schema_extension ? std::nullopt : schema,
         allocated_schema_metadata, extension_map,
         /*ignore_field_presence=*/false, executor, field_vars.value));
   }
@@ -1107,10 +1118,12 @@ absl::Status FromProtoMessage(
                     &result]() -> absl::Status {
     std::vector<absl::string_view> value_attr_names;
     std::vector<DataSlice> values;
+    bool has_non_schema_extensions = false;
     for (auto& field_vars : vars->fields) {
       if (field_vars.value.has_value()) {
         values.push_back(*std::move(field_vars.value));
         value_attr_names.push_back(field_vars.attr_name);
+        has_non_schema_extensions |= field_vars.is_non_schema_extension;
       }
     }
 
@@ -1126,13 +1139,13 @@ absl::Status FromProtoMessage(
                                      /*itemid=*/vars->itemid));
       } else {  // schema != OBJECT
         ASSIGN_OR_RETURN(
-            result,
-            EntityCreator::Shaped(db, std::move(result_shape),
-                                  /*attr_names=*/std::move(value_attr_names),
-                                  /*values=*/std::move(values),
-                                  /*schema=*/std::move(vars->requested_schema),
-                                  /*overwrite_schema=*/false,
-                                  /*itemid=*/vars->itemid));
+            result, EntityCreator::Shaped(
+                        db, std::move(result_shape),
+                        /*attr_names=*/std::move(value_attr_names),
+                        /*values=*/std::move(values),
+                        /*schema=*/std::move(vars->requested_schema),
+                        /*overwrite_schema=*/has_non_schema_extensions,
+                        /*itemid=*/vars->itemid));
       }
     } else {  // schema == nullopt
       ASSIGN_OR_RETURN(result, EntityCreator::Shaped(

koladata/proto/from_proto.h

@@ -59,9 +59,9 @@ namespace koladata {
 // `schema` were nullopt, but converted to OBJECT recursively.
 //
 // If `schema` is an entity schema, the set of attr names on that schema is the
-// set of fields and extensions that are converted, overriding the default
-// behavior and ignoring `extensions`. Attr names that start with '(' and end
-// with ')' are interpreted as fully-qualified extension names and cause the
+// set of fields and extensions that are converted, in addition to the
+// extensions in `extensions`. Attr names that start with '(' and end with ')'
+// are interpreted as fully-qualified extension names and cause the
 // corresponding extension to be converted if present on the message. If a
 // sub-schema of this schema is OBJECT, the corresponding sub-messages are
 // converted using the OBJECT rules above.

koladata/proto/from_proto_test.cc

@@ -1148,6 +1148,56 @@ TEST(FromProtoTest, ExtensionViaSchema) {
       IsOkAndHolds(IsEquivalentTo(test::DataSlice<bool>({true}, db))));
 }
 
+TEST(FromProtoTest, ExtensionViaSchemaExtensionsAndExtensionsListUnion) {
+  testing::ExampleMessage2 message;
+  message.SetExtension(koladata::testing::m2_bool_extension_field, false);
+  message.MutableExtension(koladata::testing::m2_message2_extension_field)
+      ->SetExtension(koladata::testing::m2_bool_extension_field, true);
+
+  auto db = DataBag::EmptyMutable();
+  auto schema = test::EntitySchema(
+      {
+          "(koladata.testing.m2_message2_extension_field)",
+      },
+      {
+          test::EntitySchema(
+              {
+                  "(koladata.testing.m2_bool_extension_field)",
+              },
+              {
+                  test::Schema(schema::kBool),
+              },
+              db),
+      },
+      db);
+
+  ASSERT_OK_AND_ASSIGN(
+      auto result, FromProto(db, {&message},
+                             {
+                                 "(koladata.testing.m2_bool_extension_field)",
+                             },
+                             std::nullopt, schema));
+
+  EXPECT_THAT(result.GetAttrNames(),
+              IsOkAndHolds(UnorderedElementsAreArray({
+                  "(koladata.testing.m2_bool_extension_field)",
+                  "(koladata.testing.m2_message2_extension_field)",
+              })));
+  EXPECT_THAT(result.GetAttr("(koladata.testing.m2_bool_extension_field)"),
+              IsOkAndHolds(IsEquivalentTo(test::DataSlice<bool>({false}, db))));
+
+  ASSERT_OK_AND_ASSIGN(
+      auto message_ext_field,
+      result.GetAttr("(koladata.testing.m2_message2_extension_field)"));
+  EXPECT_THAT(message_ext_field.GetAttrNames(),
+              IsOkAndHolds(UnorderedElementsAreArray({
+                  "(koladata.testing.m2_bool_extension_field)",
+              })));
+  EXPECT_THAT(
+      message_ext_field.GetAttr("(koladata.testing.m2_bool_extension_field)"),
+      IsOkAndHolds(IsEquivalentTo(test::DataSlice<bool>({true}, db))));
+}
+
 TEST(FromProtoTest, InvalidExtensionPath) {
   testing::ExampleMessage message;
 

py/koladata/functions/proto_conversions.py

@@ -97,14 +97,14 @@ def from_proto(
   objects match the field names in the proto definition. See below for methods
   to convert proto extensions to attributes alongside regular fields.
 
-  Messages, primitive fields, repeated fields, and maps are converted to
-  equivalent Koda structures. Enums are converted to ints.
+  If schema is not specified or schema is kd.OBJECT, only present fields in
+  `messages` are loaded and included in the converted schema. To get a schema
+  that contains all fields independent of the data, use `kd.schema_from_proto`.
 
-  Only present values in `messages` are added. Default and missing values are
-  not used.
+  Proto extensions are ignored by default unless `extensions` is specified, or
+  if an explicit entity schema with parenthesized attr names is specified. If
+  both are specified, we use the union of the two extension sets.
 
-  Proto extensions are ignored by default unless `extensions` is specified (or
-  if an explicit entity schema with parenthesized attrs is used).
   The format of each extension specified in `extensions` is a dot-separated
   sequence of field names and/or extension names, where extension names are
   fully-qualified extension paths surrounded by parentheses. This sequence of
@@ -116,6 +116,9 @@ def from_proto(
     "path.to.map_field.values.(package_name.some_extension)"
     "path.(package_name.some_extension).(package_name2.nested_extension)"
 
+  If an explicit entity schema attr name starts with "(" and ends with ")" it is
+  also interpreted as an extension name.
+
   Extensions are looked up using the C++ generated descriptor pool, using
   `DescriptorPool::FindExtensionByName`, which requires that all extensions are
   compiled in as C++ protos. The Koda attribute names for the extension fields
@@ -129,7 +132,8 @@ def from_proto(
     ds.optional_field.get_attr('(package_name.DefExtension.def_extension)')
 
   If `messages` is a single proto Message, the result is a DataItem. If it is a
-  list of proto Messages, the result is an 1D DataSlice.
+  nested list of proto Messages, the result is a DataSlice with the same number
+  of dimensions as the nesting level.
 
   Args:
     messages: Message or nested list/tuple of Message of the same type. Any of

py/koladata/functions/tests/from_proto_test.py

@@ -398,6 +398,59 @@ def test_extensions(self):
         'koladata.functions.testing.MessageAExtension',
     )
 
+  def test_extensions_from_arg_and_schema(self):
+    m = test_pb2.MessageA(
+        some_text='thing 1',
+    )
+
+    m.message_set_extensions.Extensions[
+        test_pb2.MessageAExtension.message_set_extension
+    ].extra = 1
+    m.Extensions[test_pb2.MessageAExtension.message_a_extension].extra = 2
+
+    s = proto_conversions.schema_from_proto(
+        test_pb2.MessageA,
+        extensions=[
+            '(koladata.functions.testing.MessageAExtension.message_a_extension)'
+        ],
+    )
+    x = proto_conversions.from_proto(
+        m,
+        schema=s,
+        extensions=[
+            'message_set_extensions.(koladata.functions.testing.MessageAExtension.message_set_extension)'
+        ],
+    )
+
+    self.assertCountEqual(
+        x.get_attr_names(intersection=True),
+        [
+            'some_text',
+            'some_float',
+            'message_b_list',
+            'message_set_extensions',
+            '(koladata.functions.testing.MessageAExtension.message_a_extension)',
+        ],
+    )
+    self.assertEqual(
+        x.get_attr(
+            '(koladata.functions.testing.MessageAExtension.message_a_extension)'
+        ).extra,
+        2,
+    )
+    self.assertCountEqual(
+        x.message_set_extensions.get_attr_names(intersection=True),
+        [
+            '(koladata.functions.testing.MessageAExtension.message_set_extension)'
+        ],
+    )
+    self.assertEqual(
+        x.message_set_extensions.get_attr(
+            '(koladata.functions.testing.MessageAExtension.message_set_extension)'
+        ).extra,
+        1,
+    )
+
   def test_extension_on_wrong_message_error(self):
     with self.assertRaisesRegex(
         ValueError,