diff --git a/dpsynth/adapters/README.md b/dpsynth/adapters/README.md index 4f35e8a..f81b26d 100644 --- a/dpsynth/adapters/README.md +++ b/dpsynth/adapters/README.md @@ -22,3 +22,12 @@ require remain optional for the core library to keep it lightweight). > between the pipeline DP implementations and the local-mode NumPy-based > implementations. How it fits within the broader ecosystem long-term is an > open question. + +## Optional Dependencies + +To keep the core DPSynth package lightweight, adapter dependencies are optional: + +- **Pydantic**: Included in core dependencies (`pip install dpsynth`). +- **Protobuf**: Install via `pip install protobuf`. +- **Beam**: Install via `pip install dpsynth[pipeline]` or `pip install apache-beam`. + diff --git a/dpsynth/adapters/protobuf.py b/dpsynth/adapters/protobuf.py new file mode 100644 index 0000000..144b918 --- /dev/null +++ b/dpsynth/adapters/protobuf.py @@ -0,0 +1,182 @@ +# Copyright 2026 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. +"""Protobuf descriptor to dpsynth domain adapter.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from typing import Any, Literal + +from dpsynth import domain +from google.protobuf import descriptor +from google.protobuf import message +from google.protobuf import message_factory + +OpenSetAttribute = domain.OpenSetCategoricalAttribute + +_INTEGER_TYPES = { + descriptor.FieldDescriptor.TYPE_INT32, + descriptor.FieldDescriptor.TYPE_INT64, + descriptor.FieldDescriptor.TYPE_UINT32, + descriptor.FieldDescriptor.TYPE_UINT64, + descriptor.FieldDescriptor.TYPE_SINT32, + descriptor.FieldDescriptor.TYPE_SINT64, + descriptor.FieldDescriptor.TYPE_FIXED32, + descriptor.FieldDescriptor.TYPE_FIXED64, + descriptor.FieldDescriptor.TYPE_SFIXED32, + descriptor.FieldDescriptor.TYPE_SFIXED64, +} + +_FLOAT_TYPES = { + descriptor.FieldDescriptor.TYPE_FLOAT, + descriptor.FieldDescriptor.TYPE_DOUBLE, +} + + +def _is_repeated(field) -> bool: + return ( + getattr(field, "is_repeated", False) + or getattr(field, "label", None) + == descriptor.FieldDescriptor.LABEL_REPEATED + ) + + +def _resolve_descriptor(proto): + """Resolves input to a protobuf Message Descriptor.""" + if isinstance(proto, descriptor.Descriptor): + return proto + desc = getattr(proto, "DESCRIPTOR", None) + if isinstance(desc, descriptor.Descriptor): + return desc + raise TypeError( + "Expected a Protobuf Message class, instance, or Descriptor, got" + f" {type(proto)}." + ) + + +def infer_domain( + proto: descriptor.Descriptor | type[message.Message] | message.Message, + *, + numerical_bounds: Mapping[str, tuple[float, float]] | None = None, + enum_format: Literal["name", "number"] = "name", + ignore_unsupported_fields: bool = False, +) -> dict[str, domain.AttributeType]: + """Infers attribute domains from a Protobuf definition for a flat schema.""" + if enum_format not in ("name", "number"): + raise ValueError(f"Unknown enum_format '{enum_format}'.") + + msg_desc = _resolve_descriptor(proto) + attributes: dict[str, domain.AttributeType] = {} + numerical_bounds = numerical_bounds or {} + + for field in msg_desc.fields: + if _is_repeated(field): + if ignore_unsupported_fields: + continue + raise ValueError(f"Repeated field '{field.name}' is not supported.") + + if field.type == descriptor.FieldDescriptor.TYPE_MESSAGE: + # To extend to nested schemas, message fields could be recursively + # flattened with delimited column names or modeled as linked sub-tables. + if ignore_unsupported_fields: + continue + raise ValueError( + f"Nested message field '{field.name}' ({field.message_type.name}) is" + " not supported in flat schema." + ) + + if field.type == descriptor.FieldDescriptor.TYPE_ENUM: + values = [ + v.name if enum_format == "name" else v.number + for v in field.enum_type.values + ] + attributes[field.name] = domain.CategoricalAttribute( + possible_values=values, out_of_domain_index=0 + ) + elif field.type in _INTEGER_TYPES or field.type in _FLOAT_TYPES: + dtype = "int" if field.type in _INTEGER_TYPES else "float" + if field.name not in numerical_bounds: + raise ValueError( + f"Numerical bounds must be specified for field '{field.name}'." + ) + bounds = numerical_bounds[field.name] + attributes[field.name] = domain.NumericalAttribute( + min_value=bounds[0], max_value=bounds[1], dtype=dtype + ) + elif field.type == descriptor.FieldDescriptor.TYPE_STRING: + attributes[field.name] = domain.OpenSetCategoricalAttribute() + elif field.type == descriptor.FieldDescriptor.TYPE_BOOL: + attributes[field.name] = domain.CategoricalAttribute( + possible_values=[False, True], out_of_domain_index=0 + ) + else: + if ignore_unsupported_fields: + continue + raise ValueError( + f"Field '{field.name}' of type {field.type} is not supported." + ) + + return attributes + + +def infer_schema( + proto: descriptor.Descriptor | type[message.Message] | message.Message, + *, + numerical_bounds: Mapping[str, tuple[float, float]] | None = None, + enum_format: Literal["name", "number"] = "name", + ignore_unsupported_fields: bool = False, +) -> domain.Schema: + """Derives a dpsynth.Schema from a Protobuf definition for a flat schema.""" + return domain.Schema( + infer_domain( + proto, + numerical_bounds=numerical_bounds, + enum_format=enum_format, + ignore_unsupported_fields=ignore_unsupported_fields, + ) + ) + + +def to_tuple( + msg: message.Message, + *, + enum_format: Literal["name", "number"] = "name", +) -> tuple[Any, ...]: + """Converts a flat Protobuf message instance to a tuple of field values.""" + result = [] + for field in msg.DESCRIPTOR.fields: + val = getattr(msg, field.name) + if ( + field.type == descriptor.FieldDescriptor.TYPE_ENUM + and enum_format == "name" + ): + val = field.enum_type.values_by_number[val].name + result.append(val) + return tuple(result) + + +def from_tuple( + values: Sequence[Any], + proto: descriptor.Descriptor | type[message.Message] | message.Message, +) -> message.Message: + """Populates a flat Protobuf message from a sequence of field values.""" + desc = _resolve_descriptor(proto) + msg = message_factory.GetMessageClass(desc)() + for field, val in zip(desc.fields, values): + setattr(msg, field.name, val) + return msg + + +infer_domain_from_proto = infer_domain +schema_from_proto = infer_schema diff --git a/tests/adapters/protobuf_test.py b/tests/adapters/protobuf_test.py new file mode 100644 index 0000000..35cf297 --- /dev/null +++ b/tests/adapters/protobuf_test.py @@ -0,0 +1,297 @@ +# Copyright 2026 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 protobuf adapter.""" + +from __future__ import annotations + +import importlib + +from absl.testing import absltest +from absl.testing import parameterized +from dpsynth import domain +from dpsynth.adapters import protobuf +from google.protobuf import descriptor_pool +from google.protobuf import message_factory + +_desc_module_name = "google.protobuf.descriptor" + "_pb2" +_descriptor_pb2 = importlib.import_module(_desc_module_name) + + +def _build_test_messages(): + """Constructs test protobuf message classes dynamically.""" + fdp = _descriptor_pb2.FieldDescriptorProto + file_proto = _descriptor_pb2.FileDescriptorProto( + name="test_sample.proto", package="test_package" + ) + + enum_proto = file_proto.enum_type.add(name="Status") + for name, num in [ + ("UNKNOWN", 0), + ("PENDING", 1), + ("ACTIVE", 2), + ("COMPLETED", 3), + ]: + enum_proto.value.add(name=name, number=num) + + msg_proto = file_proto.message_type.add(name="FlatUser") + msg_proto.field.add(name="user_id", number=1, type=fdp.TYPE_INT64) + msg_proto.field.add(name="username", number=2, type=fdp.TYPE_STRING) + msg_proto.field.add( + name="status", + number=3, + type=fdp.TYPE_ENUM, + type_name=".test_package.Status", + ) + msg_proto.field.add(name="score", number=4, type=fdp.TYPE_FLOAT) + msg_proto.field.add(name="is_verified", number=5, type=fdp.TYPE_BOOL) + + nested_proto = file_proto.message_type.add(name="NestedMessage") + nested_proto.field.add( + name="user", + number=1, + type=fdp.TYPE_MESSAGE, + type_name=".test_package.FlatUser", + ) + + rep_proto = file_proto.message_type.add(name="RepeatedMessage") + rep_proto.field.add( + name="values", number=1, type=fdp.TYPE_INT32, label=fdp.LABEL_REPEATED + ) + + unsupp_proto = file_proto.message_type.add(name="UnsupportedFieldMessage") + unsupp_proto.field.add(name="raw_data", number=1, type=fdp.TYPE_BYTES) + + non_num_proto = file_proto.message_type.add(name="NonNumericalMessage") + non_num_proto.field.add(name="username", number=1, type=fdp.TYPE_STRING) + + pool = descriptor_pool.DescriptorPool() + file_desc = pool.Add(file_proto) + return tuple( + message_factory.GetMessageClass(file_desc.message_types_by_name[name]) + for name in [ + "FlatUser", + "NestedMessage", + "RepeatedMessage", + "UnsupportedFieldMessage", + "NonNumericalMessage", + ] + ) + + +_DEFAULT_BOUNDS = {"user_id": (0.0, 100.0), "score": (0.0, 100.0)} + + +class ProtobufAdapterTest(parameterized.TestCase): + + @classmethod + def setUpClass(cls): + super().setUpClass() + ( + cls.flat_user_cls, + cls.nested_msg_cls, + cls.repeated_msg_cls, + cls.unsupported_msg_cls, + cls.non_numerical_msg_cls, + ) = _build_test_messages() + + def test_infer_domain_flat_user(self): + attributes = protobuf.infer_domain_from_proto( + self.flat_user_cls, numerical_bounds=_DEFAULT_BOUNDS + ) + + self.assertEqual( + attributes["user_id"], + domain.NumericalAttribute(min_value=0.0, max_value=100.0, dtype="int"), + ) + self.assertEqual( + attributes["username"], domain.OpenSetCategoricalAttribute() + ) + self.assertEqual( + attributes["status"], + domain.CategoricalAttribute( + possible_values=["UNKNOWN", "PENDING", "ACTIVE", "COMPLETED"], + out_of_domain_index=0, + ), + ) + self.assertEqual( + attributes["score"], + domain.NumericalAttribute( + min_value=0.0, max_value=100.0, dtype="float" + ), + ) + self.assertEqual( + attributes["is_verified"], + domain.CategoricalAttribute( + possible_values=[False, True], out_of_domain_index=0 + ), + ) + + def test_infer_schema(self): + schema = protobuf.infer_schema( + self.flat_user_cls, numerical_bounds=_DEFAULT_BOUNDS + ) + self.assertIsInstance(schema, domain.Schema) + self.assertLen(schema, 5) + self.assertEqual(schema["username"], domain.OpenSetCategoricalAttribute()) + + def test_input_type_flexibility(self): + for proto_input in [ + self.flat_user_cls, + self.flat_user_cls(), + self.flat_user_cls.DESCRIPTOR, + ]: + schema = protobuf.infer_schema( + proto_input, numerical_bounds=_DEFAULT_BOUNDS + ) + self.assertIsInstance(schema, domain.Schema) + self.assertLen(schema, 5) + + def test_invalid_input_type_raises(self): + with self.assertRaisesRegex(TypeError, "Expected a Protobuf Message class"): + protobuf.infer_schema("not a proto") # pyrefly: ignore[bad-argument-type] + + def test_enum_format_number(self): + attributes = protobuf.infer_domain( + self.flat_user_cls, + numerical_bounds=_DEFAULT_BOUNDS, + enum_format="number", + ) + self.assertEqual( + attributes["status"], + domain.CategoricalAttribute( + possible_values=[0, 1, 2, 3], out_of_domain_index=0 + ), + ) + + def test_invalid_enum_format_raises(self): + with self.assertRaisesRegex(ValueError, "Unknown enum_format"): + protobuf.infer_domain( + self.flat_user_cls, + numerical_bounds=_DEFAULT_BOUNDS, + enum_format="other", # pyrefly: ignore[bad-argument-type] + ) + + def test_numerical_bounds(self): + bounds = {"user_id": (1.0, 1000.0), "score": (-5.0, 5.0)} + attributes = protobuf.infer_domain( + self.flat_user_cls, + numerical_bounds=bounds, + ) + self.assertEqual( + attributes["user_id"], + domain.NumericalAttribute(min_value=1.0, max_value=1000.0, dtype="int"), + ) + self.assertEqual( + attributes["score"], + domain.NumericalAttribute(min_value=-5.0, max_value=5.0, dtype="float"), + ) + + def test_missing_bounds_raises(self): + with self.assertRaisesRegex( + ValueError, "Numerical bounds must be specified for field 'user_id'." + ): + protobuf.infer_domain(self.flat_user_cls) + + with self.assertRaisesRegex( + ValueError, "Numerical bounds must be specified for field 'score'." + ): + protobuf.infer_domain( + self.flat_user_cls, + numerical_bounds={"user_id": (0.0, 100.0)}, + ) + + def test_non_numerical_proto_without_bounds(self): + attributes = protobuf.infer_domain(self.non_numerical_msg_cls) + self.assertIn("username", attributes) + self.assertIsInstance( + attributes["username"], domain.OpenSetCategoricalAttribute + ) + + def test_nested_message_rejected(self): + with self.assertRaisesRegex(ValueError, "Nested message field 'user'"): + protobuf.infer_domain(self.nested_msg_cls) + + attrs = protobuf.infer_domain( + self.nested_msg_cls, ignore_unsupported_fields=True + ) + self.assertEmpty(attrs) + + def test_repeated_field_rejected(self): + with self.assertRaisesRegex(ValueError, "Repeated field 'values'"): + protobuf.infer_domain(self.repeated_msg_cls) + + attrs = protobuf.infer_domain( + self.repeated_msg_cls, ignore_unsupported_fields=True + ) + self.assertEmpty(attrs) + + def test_unsupported_field_rejected(self): + with self.assertRaisesRegex(ValueError, "Field 'raw_data' of type"): + protobuf.infer_domain(self.unsupported_msg_cls) + + attrs = protobuf.infer_domain( + self.unsupported_msg_cls, + ignore_unsupported_fields=True, + ) + self.assertEmpty(attrs) + + def test_to_tuple(self): + user = self.flat_user_cls( + user_id=42, + username="alice", + status=2, # ACTIVE + score=95.5, + is_verified=True, + ) + t = protobuf.to_tuple(user) + self.assertEqual(t, (42, "alice", "ACTIVE", 95.5, True)) + + def test_to_tuple_number_enum(self): + user = self.flat_user_cls( + user_id=42, + username="alice", + status=2, + score=95.5, + is_verified=True, + ) + t = protobuf.to_tuple(user, enum_format="number") + self.assertEqual(t, (42, "alice", 2, 95.5, True)) + + def test_from_tuple(self): + t = (42, "alice", "ACTIVE", 95.5, True) + user = protobuf.from_tuple(t, self.flat_user_cls) + self.assertIsInstance(user, self.flat_user_cls) + self.assertEqual(user.user_id, 42) + self.assertEqual(user.username, "alice") + self.assertEqual(user.status, 2) + self.assertEqual(user.score, 95.5) + self.assertTrue(user.is_verified) + + def test_from_tuple_from_instance_and_desc(self): + t = (42, "alice", 2, 95.5, True) + user1 = protobuf.from_tuple(t, self.flat_user_cls()) + user2 = protobuf.from_tuple(t, self.flat_user_cls.DESCRIPTOR) + self.assertEqual(user1.user_id, 42) + self.assertEqual(user2.user_id, 42) + + def test_tuple_proto_roundtrip(self): + original_tuple = (100, "bob", "PENDING", 12.5, False) + user = protobuf.from_tuple(original_tuple, self.flat_user_cls) + roundtripped = protobuf.to_tuple(user) + self.assertEqual(roundtripped, original_tuple) + + +if __name__ == "__main__": + absltest.main()