From b767ccb634ef45563b5aed5611a428b2a63aad5e Mon Sep 17 00:00:00 2001 From: Rafal Jankowski Date: Wed, 18 Sep 2024 15:42:50 +0200 Subject: [PATCH] Added support for Protocol Buffers in queryAttributesWithinProject --- src/neptune/api/models.py | 13 +- .../neptune_pb/api/model/attributes_pb2.py | 21 ++- .../neptune_pb/api/model/attributes_pb2.pyi | 68 +++++++++ .../api/model/leaderboard_entries_pb2.py | 50 +++--- .../neptune_pb/api/model/series_values_pb2.py | 14 +- .../backends/hosted_neptune_backend.py | 52 ++++--- .../internal/backends/neptune_backend.py | 1 + .../internal/backends/neptune_backend_mock.py | 1 + .../backends/offline_neptune_backend.py | 1 + tests/unit/neptune/new/api/test_models.py | 142 ++++++++++++++++-- 10 files changed, 294 insertions(+), 69 deletions(-) diff --git a/src/neptune/api/models.py b/src/neptune/api/models.py index 721d6e04c..77479346e 100644 --- a/src/neptune/api/models.py +++ b/src/neptune/api/models.py @@ -619,7 +619,7 @@ def from_model(model: Any) -> NextPage: @staticmethod def from_proto(data: Any) -> NextPage: - raise NotImplementedError() + return NextPage(limit=data.limit, next_page_token=data.nextPageToken) def to_dto(self) -> Dict[str, Any]: return { @@ -652,7 +652,11 @@ def from_model(model: Any) -> QueryFieldsExperimentResult: @staticmethod def from_proto(data: Any) -> QueryFieldsExperimentResult: - raise NotImplementedError() + return QueryFieldsExperimentResult( + object_id=data.experimentId, + object_key=data.experimentShortId, + fields=[Field.from_proto(field) for field in data.attributes], + ) @dataclass @@ -676,7 +680,10 @@ def from_model(model: Any) -> QueryFieldsResult: @staticmethod def from_proto(data: Any) -> QueryFieldsResult: - raise NotImplementedError() + return QueryFieldsResult( + entries=[QueryFieldsExperimentResult.from_proto(entry) for entry in data.entries], + next_page=NextPage.from_proto(data.nextPage), + ) @dataclass diff --git a/src/neptune/api/proto/neptune_pb/api/model/attributes_pb2.py b/src/neptune/api/proto/neptune_pb/api/model/attributes_pb2.py index 98cc5432c..846de5969 100644 --- a/src/neptune/api/proto/neptune_pb/api/model/attributes_pb2.py +++ b/src/neptune/api/proto/neptune_pb/api/model/attributes_pb2.py @@ -1,6 +1,6 @@ # -*- coding: utf-8 -*- # Generated by the protocol buffer compiler. DO NOT EDIT! -# source: model/attributes.proto +# source: neptune_pb/api/model/attributes.proto # Protobuf Python Version: 4.25.1 """Generated protocol buffer code.""" from google.protobuf import descriptor as _descriptor @@ -12,18 +12,25 @@ _sym_db = _symbol_database.Default() +from neptune.api.proto.neptune_pb.api.model import leaderboard_entries_pb2 as neptune__pb_dot_api_dot_model_dot_leaderboard__entries__pb2 -DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x16model/attributes.proto\x12\x11neptune.api.model\"a\n\x1eProtoAttributesSearchResultDTO\x12?\n\x07\x65ntries\x18\x01 \x03(\x0b\x32..neptune.api.model.ProtoAttributeDefinitionDTO\"9\n\x1bProtoAttributeDefinitionDTO\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x0c\n\x04type\x18\x02 \x01(\tB4\n0ml.neptune.leaderboard.api.model.proto.generatedP\x01\x62\x06proto3') +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n%neptune_pb/api/model/attributes.proto\x12\x11neptune.api.model\x1a.neptune_pb/api/model/leaderboard_entries.proto\"a\n\x1eProtoAttributesSearchResultDTO\x12?\n\x07\x65ntries\x18\x01 \x03(\x0b\x32..neptune.api.model.ProtoAttributeDefinitionDTO\"9\n\x1bProtoAttributeDefinitionDTO\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x0c\n\x04type\x18\x02 \x01(\t\"\xa3\x01\n\x1dProtoQueryAttributesResultDTO\x12K\n\x07\x65ntries\x18\x01 \x03(\x0b\x32:.neptune.api.model.ProtoQueryAttributesExperimentResultDTO\x12\x35\n\x08nextPage\x18\x02 \x01(\x0b\x32#.neptune.api.model.ProtoNextPageDTO\"^\n\x10ProtoNextPageDTO\x12\x1a\n\rnextPageToken\x18\x01 \x01(\tH\x00\x88\x01\x01\x12\x12\n\x05limit\x18\x02 \x01(\rH\x01\x88\x01\x01\x42\x10\n\x0e_nextPageTokenB\x08\n\x06_limit\"\x94\x01\n\'ProtoQueryAttributesExperimentResultDTO\x12\x14\n\x0c\x65xperimentId\x18\x01 \x01(\t\x12\x19\n\x11\x65xperimentShortId\x18\x02 \x01(\t\x12\x38\n\nattributes\x18\x03 \x03(\x0b\x32$.neptune.api.model.ProtoAttributeDTOB4\n0ml.neptune.leaderboard.api.model.proto.generatedP\x01\x62\x06proto3') _globals = globals() _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) -_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'model.attributes_pb2', _globals) +_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'neptune.api.proto.neptune_pb.api.model.attributes_pb2', _globals) if _descriptor._USE_C_DESCRIPTORS == False: _globals['DESCRIPTOR']._options = None _globals['DESCRIPTOR']._serialized_options = b'\n0ml.neptune.leaderboard.api.model.proto.generatedP\001' - _globals['_PROTOATTRIBUTESSEARCHRESULTDTO']._serialized_start=45 - _globals['_PROTOATTRIBUTESSEARCHRESULTDTO']._serialized_end=142 - _globals['_PROTOATTRIBUTEDEFINITIONDTO']._serialized_start=144 - _globals['_PROTOATTRIBUTEDEFINITIONDTO']._serialized_end=201 + _globals['_PROTOATTRIBUTESSEARCHRESULTDTO']._serialized_start=108 + _globals['_PROTOATTRIBUTESSEARCHRESULTDTO']._serialized_end=205 + _globals['_PROTOATTRIBUTEDEFINITIONDTO']._serialized_start=207 + _globals['_PROTOATTRIBUTEDEFINITIONDTO']._serialized_end=264 + _globals['_PROTOQUERYATTRIBUTESRESULTDTO']._serialized_start=267 + _globals['_PROTOQUERYATTRIBUTESRESULTDTO']._serialized_end=430 + _globals['_PROTONEXTPAGEDTO']._serialized_start=432 + _globals['_PROTONEXTPAGEDTO']._serialized_end=526 + _globals['_PROTOQUERYATTRIBUTESEXPERIMENTRESULTDTO']._serialized_start=529 + _globals['_PROTOQUERYATTRIBUTESEXPERIMENTRESULTDTO']._serialized_end=677 # @@protoc_insertion_point(module_scope) diff --git a/src/neptune/api/proto/neptune_pb/api/model/attributes_pb2.pyi b/src/neptune/api/proto/neptune_pb/api/model/attributes_pb2.pyi index f2acfeeb3..22c51bcdf 100644 --- a/src/neptune/api/proto/neptune_pb/api/model/attributes_pb2.pyi +++ b/src/neptune/api/proto/neptune_pb/api/model/attributes_pb2.pyi @@ -7,7 +7,9 @@ import collections.abc import google.protobuf.descriptor import google.protobuf.internal.containers import google.protobuf.message +import neptune.api.proto.neptune_pb.api.model.leaderboard_entries_pb2 import sys +import typing if sys.version_info >= (3, 8): import typing as typing_extensions @@ -49,3 +51,69 @@ class ProtoAttributeDefinitionDTO(google.protobuf.message.Message): def ClearField(self, field_name: typing_extensions.Literal["name", b"name", "type", b"type"]) -> None: ... global___ProtoAttributeDefinitionDTO = ProtoAttributeDefinitionDTO + +@typing_extensions.final +class ProtoQueryAttributesResultDTO(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + ENTRIES_FIELD_NUMBER: builtins.int + NEXTPAGE_FIELD_NUMBER: builtins.int + @property + def entries(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___ProtoQueryAttributesExperimentResultDTO]: ... + @property + def nextPage(self) -> global___ProtoNextPageDTO: ... + def __init__( + self, + *, + entries: collections.abc.Iterable[global___ProtoQueryAttributesExperimentResultDTO] | None = ..., + nextPage: global___ProtoNextPageDTO | None = ..., + ) -> None: ... + def HasField(self, field_name: typing_extensions.Literal["nextPage", b"nextPage"]) -> builtins.bool: ... + def ClearField(self, field_name: typing_extensions.Literal["entries", b"entries", "nextPage", b"nextPage"]) -> None: ... + +global___ProtoQueryAttributesResultDTO = ProtoQueryAttributesResultDTO + +@typing_extensions.final +class ProtoNextPageDTO(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + NEXTPAGETOKEN_FIELD_NUMBER: builtins.int + LIMIT_FIELD_NUMBER: builtins.int + nextPageToken: builtins.str + limit: builtins.int + def __init__( + self, + *, + nextPageToken: builtins.str | None = ..., + limit: builtins.int | None = ..., + ) -> None: ... + def HasField(self, field_name: typing_extensions.Literal["_limit", b"_limit", "_nextPageToken", b"_nextPageToken", "limit", b"limit", "nextPageToken", b"nextPageToken"]) -> builtins.bool: ... + def ClearField(self, field_name: typing_extensions.Literal["_limit", b"_limit", "_nextPageToken", b"_nextPageToken", "limit", b"limit", "nextPageToken", b"nextPageToken"]) -> None: ... + @typing.overload + def WhichOneof(self, oneof_group: typing_extensions.Literal["_limit", b"_limit"]) -> typing_extensions.Literal["limit"] | None: ... + @typing.overload + def WhichOneof(self, oneof_group: typing_extensions.Literal["_nextPageToken", b"_nextPageToken"]) -> typing_extensions.Literal["nextPageToken"] | None: ... + +global___ProtoNextPageDTO = ProtoNextPageDTO + +@typing_extensions.final +class ProtoQueryAttributesExperimentResultDTO(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + EXPERIMENTID_FIELD_NUMBER: builtins.int + EXPERIMENTSHORTID_FIELD_NUMBER: builtins.int + ATTRIBUTES_FIELD_NUMBER: builtins.int + experimentId: builtins.str + experimentShortId: builtins.str + @property + def attributes(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[neptune.api.proto.neptune_pb.api.model.leaderboard_entries_pb2.ProtoAttributeDTO]: ... + def __init__( + self, + *, + experimentId: builtins.str = ..., + experimentShortId: builtins.str = ..., + attributes: collections.abc.Iterable[neptune.api.proto.neptune_pb.api.model.leaderboard_entries_pb2.ProtoAttributeDTO] | None = ..., + ) -> None: ... + def ClearField(self, field_name: typing_extensions.Literal["attributes", b"attributes", "experimentId", b"experimentId", "experimentShortId", b"experimentShortId"]) -> None: ... + +global___ProtoQueryAttributesExperimentResultDTO = ProtoQueryAttributesExperimentResultDTO diff --git a/src/neptune/api/proto/neptune_pb/api/model/leaderboard_entries_pb2.py b/src/neptune/api/proto/neptune_pb/api/model/leaderboard_entries_pb2.py index f6a735ee7..9b1608bcc 100644 --- a/src/neptune/api/proto/neptune_pb/api/model/leaderboard_entries_pb2.py +++ b/src/neptune/api/proto/neptune_pb/api/model/leaderboard_entries_pb2.py @@ -1,6 +1,6 @@ # -*- coding: utf-8 -*- # Generated by the protocol buffer compiler. DO NOT EDIT! -# source: model/leaderboard_entries.proto +# source: neptune_pb/api/model/leaderboard_entries.proto # Protobuf Python Version: 4.25.1 """Generated protocol buffer code.""" from google.protobuf import descriptor as _descriptor @@ -14,34 +14,34 @@ -DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1fmodel/leaderboard_entries.proto\x12\x11neptune.api.model\"\xb3\x01\n&ProtoLeaderboardEntriesSearchResultDTO\x12\x1b\n\x13matching_item_count\x18\x01 \x01(\x03\x12\x1e\n\x11total_group_count\x18\x02 \x01(\x03H\x00\x88\x01\x01\x12\x36\n\x07\x65ntries\x18\x03 \x03(\x0b\x32%.neptune.api.model.ProtoAttributesDTOB\x14\n\x12_total_group_count\"\xd1\x01\n\x12ProtoAttributesDTO\x12\x15\n\rexperiment_id\x18\x01 \x01(\t\x12\x0c\n\x04type\x18\x02 \x01(\t\x12\x12\n\nproject_id\x18\x03 \x01(\t\x12\x17\n\x0forganization_id\x18\x04 \x01(\t\x12\x14\n\x0cproject_name\x18\x05 \x01(\t\x12\x19\n\x11organization_name\x18\x06 \x01(\t\x12\x38\n\nattributes\x18\x07 \x03(\x0b\x32$.neptune.api.model.ProtoAttributeDTO\"\xed\x05\n\x11ProtoAttributeDTO\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x0c\n\x04type\x18\x02 \x01(\t\x12\x44\n\x0eint_properties\x18\x03 \x01(\x0b\x32\'.neptune.api.model.ProtoIntAttributeDTOH\x00\x88\x01\x01\x12H\n\x10\x66loat_properties\x18\x04 \x01(\x0b\x32).neptune.api.model.ProtoFloatAttributeDTOH\x01\x88\x01\x01\x12J\n\x11string_properties\x18\x05 \x01(\x0b\x32*.neptune.api.model.ProtoStringAttributeDTOH\x02\x88\x01\x01\x12\x46\n\x0f\x62ool_properties\x18\x06 \x01(\x0b\x32(.neptune.api.model.ProtoBoolAttributeDTOH\x03\x88\x01\x01\x12N\n\x13\x64\x61tetime_properties\x18\x07 \x01(\x0b\x32,.neptune.api.model.ProtoDatetimeAttributeDTOH\x04\x88\x01\x01\x12Q\n\x15string_set_properties\x18\x08 \x01(\x0b\x32-.neptune.api.model.ProtoStringSetAttributeDTOH\x05\x88\x01\x01\x12U\n\x17\x66loat_series_properties\x18\t \x01(\x0b\x32/.neptune.api.model.ProtoFloatSeriesAttributeDTOH\x06\x88\x01\x01\x42\x11\n\x0f_int_propertiesB\x13\n\x11_float_propertiesB\x14\n\x12_string_propertiesB\x12\n\x10_bool_propertiesB\x16\n\x14_datetime_propertiesB\x18\n\x16_string_set_propertiesB\x1a\n\x18_float_series_properties\"U\n\x14ProtoIntAttributeDTO\x12\x16\n\x0e\x61ttribute_name\x18\x01 \x01(\t\x12\x16\n\x0e\x61ttribute_type\x18\x02 \x01(\t\x12\r\n\x05value\x18\x03 \x01(\x03\"W\n\x16ProtoFloatAttributeDTO\x12\x16\n\x0e\x61ttribute_name\x18\x01 \x01(\t\x12\x16\n\x0e\x61ttribute_type\x18\x02 \x01(\t\x12\r\n\x05value\x18\x03 \x01(\x01\"X\n\x17ProtoStringAttributeDTO\x12\x16\n\x0e\x61ttribute_name\x18\x01 \x01(\t\x12\x16\n\x0e\x61ttribute_type\x18\x02 \x01(\t\x12\r\n\x05value\x18\x03 \x01(\t\"V\n\x15ProtoBoolAttributeDTO\x12\x16\n\x0e\x61ttribute_name\x18\x01 \x01(\t\x12\x16\n\x0e\x61ttribute_type\x18\x02 \x01(\t\x12\r\n\x05value\x18\x03 \x01(\x08\"Z\n\x19ProtoDatetimeAttributeDTO\x12\x16\n\x0e\x61ttribute_name\x18\x01 \x01(\t\x12\x16\n\x0e\x61ttribute_type\x18\x02 \x01(\t\x12\r\n\x05value\x18\x03 \x01(\x03\"[\n\x1aProtoStringSetAttributeDTO\x12\x16\n\x0e\x61ttribute_name\x18\x01 \x01(\t\x12\x16\n\x0e\x61ttribute_type\x18\x02 \x01(\t\x12\r\n\x05value\x18\x03 \x03(\t\"\xd1\x02\n\x1cProtoFloatSeriesAttributeDTO\x12\x16\n\x0e\x61ttribute_name\x18\x01 \x01(\t\x12\x16\n\x0e\x61ttribute_type\x18\x02 \x01(\t\x12\x16\n\tlast_step\x18\x03 \x01(\x01H\x00\x88\x01\x01\x12\x11\n\x04last\x18\x04 \x01(\x01H\x01\x88\x01\x01\x12\x10\n\x03min\x18\x05 \x01(\x01H\x02\x88\x01\x01\x12\x10\n\x03max\x18\x06 \x01(\x01H\x03\x88\x01\x01\x12\x14\n\x07\x61verage\x18\x07 \x01(\x01H\x04\x88\x01\x01\x12\x15\n\x08variance\x18\x08 \x01(\x01H\x05\x88\x01\x01\x12\x45\n\x06\x63onfig\x18\t \x01(\x0b\x32\x35.neptune.api.model.ProtoFloatSeriesAttributeConfigDTOB\x0c\n\n_last_stepB\x07\n\x05_lastB\x06\n\x04_minB\x06\n\x04_maxB\n\n\x08_averageB\x0b\n\t_variance\"t\n\"ProtoFloatSeriesAttributeConfigDTO\x12\x10\n\x03min\x18\x01 \x01(\x01H\x00\x88\x01\x01\x12\x10\n\x03max\x18\x02 \x01(\x01H\x01\x88\x01\x01\x12\x11\n\x04unit\x18\x03 \x01(\tH\x02\x88\x01\x01\x42\x06\n\x04_minB\x06\n\x04_maxB\x07\n\x05_unitB4\n0ml.neptune.leaderboard.api.model.proto.generatedP\x01\x62\x06proto3') +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n.neptune_pb/api/model/leaderboard_entries.proto\x12\x11neptune.api.model\"\xb3\x01\n&ProtoLeaderboardEntriesSearchResultDTO\x12\x1b\n\x13matching_item_count\x18\x01 \x01(\x03\x12\x1e\n\x11total_group_count\x18\x02 \x01(\x03H\x00\x88\x01\x01\x12\x36\n\x07\x65ntries\x18\x03 \x03(\x0b\x32%.neptune.api.model.ProtoAttributesDTOB\x14\n\x12_total_group_count\"\xd1\x01\n\x12ProtoAttributesDTO\x12\x15\n\rexperiment_id\x18\x01 \x01(\t\x12\x0c\n\x04type\x18\x02 \x01(\t\x12\x12\n\nproject_id\x18\x03 \x01(\t\x12\x17\n\x0forganization_id\x18\x04 \x01(\t\x12\x14\n\x0cproject_name\x18\x05 \x01(\t\x12\x19\n\x11organization_name\x18\x06 \x01(\t\x12\x38\n\nattributes\x18\x07 \x03(\x0b\x32$.neptune.api.model.ProtoAttributeDTO\"\xed\x05\n\x11ProtoAttributeDTO\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x0c\n\x04type\x18\x02 \x01(\t\x12\x44\n\x0eint_properties\x18\x03 \x01(\x0b\x32\'.neptune.api.model.ProtoIntAttributeDTOH\x00\x88\x01\x01\x12H\n\x10\x66loat_properties\x18\x04 \x01(\x0b\x32).neptune.api.model.ProtoFloatAttributeDTOH\x01\x88\x01\x01\x12J\n\x11string_properties\x18\x05 \x01(\x0b\x32*.neptune.api.model.ProtoStringAttributeDTOH\x02\x88\x01\x01\x12\x46\n\x0f\x62ool_properties\x18\x06 \x01(\x0b\x32(.neptune.api.model.ProtoBoolAttributeDTOH\x03\x88\x01\x01\x12N\n\x13\x64\x61tetime_properties\x18\x07 \x01(\x0b\x32,.neptune.api.model.ProtoDatetimeAttributeDTOH\x04\x88\x01\x01\x12Q\n\x15string_set_properties\x18\x08 \x01(\x0b\x32-.neptune.api.model.ProtoStringSetAttributeDTOH\x05\x88\x01\x01\x12U\n\x17\x66loat_series_properties\x18\t \x01(\x0b\x32/.neptune.api.model.ProtoFloatSeriesAttributeDTOH\x06\x88\x01\x01\x42\x11\n\x0f_int_propertiesB\x13\n\x11_float_propertiesB\x14\n\x12_string_propertiesB\x12\n\x10_bool_propertiesB\x16\n\x14_datetime_propertiesB\x18\n\x16_string_set_propertiesB\x1a\n\x18_float_series_properties\"U\n\x14ProtoIntAttributeDTO\x12\x16\n\x0e\x61ttribute_name\x18\x01 \x01(\t\x12\x16\n\x0e\x61ttribute_type\x18\x02 \x01(\t\x12\r\n\x05value\x18\x03 \x01(\x03\"W\n\x16ProtoFloatAttributeDTO\x12\x16\n\x0e\x61ttribute_name\x18\x01 \x01(\t\x12\x16\n\x0e\x61ttribute_type\x18\x02 \x01(\t\x12\r\n\x05value\x18\x03 \x01(\x01\"X\n\x17ProtoStringAttributeDTO\x12\x16\n\x0e\x61ttribute_name\x18\x01 \x01(\t\x12\x16\n\x0e\x61ttribute_type\x18\x02 \x01(\t\x12\r\n\x05value\x18\x03 \x01(\t\"V\n\x15ProtoBoolAttributeDTO\x12\x16\n\x0e\x61ttribute_name\x18\x01 \x01(\t\x12\x16\n\x0e\x61ttribute_type\x18\x02 \x01(\t\x12\r\n\x05value\x18\x03 \x01(\x08\"Z\n\x19ProtoDatetimeAttributeDTO\x12\x16\n\x0e\x61ttribute_name\x18\x01 \x01(\t\x12\x16\n\x0e\x61ttribute_type\x18\x02 \x01(\t\x12\r\n\x05value\x18\x03 \x01(\x03\"[\n\x1aProtoStringSetAttributeDTO\x12\x16\n\x0e\x61ttribute_name\x18\x01 \x01(\t\x12\x16\n\x0e\x61ttribute_type\x18\x02 \x01(\t\x12\r\n\x05value\x18\x03 \x03(\t\"\xd1\x02\n\x1cProtoFloatSeriesAttributeDTO\x12\x16\n\x0e\x61ttribute_name\x18\x01 \x01(\t\x12\x16\n\x0e\x61ttribute_type\x18\x02 \x01(\t\x12\x16\n\tlast_step\x18\x03 \x01(\x01H\x00\x88\x01\x01\x12\x11\n\x04last\x18\x04 \x01(\x01H\x01\x88\x01\x01\x12\x10\n\x03min\x18\x05 \x01(\x01H\x02\x88\x01\x01\x12\x10\n\x03max\x18\x06 \x01(\x01H\x03\x88\x01\x01\x12\x14\n\x07\x61verage\x18\x07 \x01(\x01H\x04\x88\x01\x01\x12\x15\n\x08variance\x18\x08 \x01(\x01H\x05\x88\x01\x01\x12\x45\n\x06\x63onfig\x18\t \x01(\x0b\x32\x35.neptune.api.model.ProtoFloatSeriesAttributeConfigDTOB\x0c\n\n_last_stepB\x07\n\x05_lastB\x06\n\x04_minB\x06\n\x04_maxB\n\n\x08_averageB\x0b\n\t_variance\"t\n\"ProtoFloatSeriesAttributeConfigDTO\x12\x10\n\x03min\x18\x01 \x01(\x01H\x00\x88\x01\x01\x12\x10\n\x03max\x18\x02 \x01(\x01H\x01\x88\x01\x01\x12\x11\n\x04unit\x18\x03 \x01(\tH\x02\x88\x01\x01\x42\x06\n\x04_minB\x06\n\x04_maxB\x07\n\x05_unitB4\n0ml.neptune.leaderboard.api.model.proto.generatedP\x01\x62\x06proto3') _globals = globals() _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) -_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'model.leaderboard_entries_pb2', _globals) +_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'neptune.api.proto.neptune_pb.api.model.leaderboard_entries_pb2', _globals) if _descriptor._USE_C_DESCRIPTORS == False: _globals['DESCRIPTOR']._options = None _globals['DESCRIPTOR']._serialized_options = b'\n0ml.neptune.leaderboard.api.model.proto.generatedP\001' - _globals['_PROTOLEADERBOARDENTRIESSEARCHRESULTDTO']._serialized_start=55 - _globals['_PROTOLEADERBOARDENTRIESSEARCHRESULTDTO']._serialized_end=234 - _globals['_PROTOATTRIBUTESDTO']._serialized_start=237 - _globals['_PROTOATTRIBUTESDTO']._serialized_end=446 - _globals['_PROTOATTRIBUTEDTO']._serialized_start=449 - _globals['_PROTOATTRIBUTEDTO']._serialized_end=1198 - _globals['_PROTOINTATTRIBUTEDTO']._serialized_start=1200 - _globals['_PROTOINTATTRIBUTEDTO']._serialized_end=1285 - _globals['_PROTOFLOATATTRIBUTEDTO']._serialized_start=1287 - _globals['_PROTOFLOATATTRIBUTEDTO']._serialized_end=1374 - _globals['_PROTOSTRINGATTRIBUTEDTO']._serialized_start=1376 - _globals['_PROTOSTRINGATTRIBUTEDTO']._serialized_end=1464 - _globals['_PROTOBOOLATTRIBUTEDTO']._serialized_start=1466 - _globals['_PROTOBOOLATTRIBUTEDTO']._serialized_end=1552 - _globals['_PROTODATETIMEATTRIBUTEDTO']._serialized_start=1554 - _globals['_PROTODATETIMEATTRIBUTEDTO']._serialized_end=1644 - _globals['_PROTOSTRINGSETATTRIBUTEDTO']._serialized_start=1646 - _globals['_PROTOSTRINGSETATTRIBUTEDTO']._serialized_end=1737 - _globals['_PROTOFLOATSERIESATTRIBUTEDTO']._serialized_start=1740 - _globals['_PROTOFLOATSERIESATTRIBUTEDTO']._serialized_end=2077 - _globals['_PROTOFLOATSERIESATTRIBUTECONFIGDTO']._serialized_start=2079 - _globals['_PROTOFLOATSERIESATTRIBUTECONFIGDTO']._serialized_end=2195 + _globals['_PROTOLEADERBOARDENTRIESSEARCHRESULTDTO']._serialized_start=70 + _globals['_PROTOLEADERBOARDENTRIESSEARCHRESULTDTO']._serialized_end=249 + _globals['_PROTOATTRIBUTESDTO']._serialized_start=252 + _globals['_PROTOATTRIBUTESDTO']._serialized_end=461 + _globals['_PROTOATTRIBUTEDTO']._serialized_start=464 + _globals['_PROTOATTRIBUTEDTO']._serialized_end=1213 + _globals['_PROTOINTATTRIBUTEDTO']._serialized_start=1215 + _globals['_PROTOINTATTRIBUTEDTO']._serialized_end=1300 + _globals['_PROTOFLOATATTRIBUTEDTO']._serialized_start=1302 + _globals['_PROTOFLOATATTRIBUTEDTO']._serialized_end=1389 + _globals['_PROTOSTRINGATTRIBUTEDTO']._serialized_start=1391 + _globals['_PROTOSTRINGATTRIBUTEDTO']._serialized_end=1479 + _globals['_PROTOBOOLATTRIBUTEDTO']._serialized_start=1481 + _globals['_PROTOBOOLATTRIBUTEDTO']._serialized_end=1567 + _globals['_PROTODATETIMEATTRIBUTEDTO']._serialized_start=1569 + _globals['_PROTODATETIMEATTRIBUTEDTO']._serialized_end=1659 + _globals['_PROTOSTRINGSETATTRIBUTEDTO']._serialized_start=1661 + _globals['_PROTOSTRINGSETATTRIBUTEDTO']._serialized_end=1752 + _globals['_PROTOFLOATSERIESATTRIBUTEDTO']._serialized_start=1755 + _globals['_PROTOFLOATSERIESATTRIBUTEDTO']._serialized_end=2092 + _globals['_PROTOFLOATSERIESATTRIBUTECONFIGDTO']._serialized_start=2094 + _globals['_PROTOFLOATSERIESATTRIBUTECONFIGDTO']._serialized_end=2210 # @@protoc_insertion_point(module_scope) diff --git a/src/neptune/api/proto/neptune_pb/api/model/series_values_pb2.py b/src/neptune/api/proto/neptune_pb/api/model/series_values_pb2.py index 2bb20605d..78867b352 100644 --- a/src/neptune/api/proto/neptune_pb/api/model/series_values_pb2.py +++ b/src/neptune/api/proto/neptune_pb/api/model/series_values_pb2.py @@ -1,6 +1,6 @@ # -*- coding: utf-8 -*- # Generated by the protocol buffer compiler. DO NOT EDIT! -# source: model/series_values.proto +# source: neptune_pb/api/model/series_values.proto # Protobuf Python Version: 4.25.1 """Generated protocol buffer code.""" from google.protobuf import descriptor as _descriptor @@ -14,16 +14,16 @@ -DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x19model/series_values.proto\x12\x11neptune.api.model\"q\n\x19ProtoFloatSeriesValuesDTO\x12\x18\n\x10total_item_count\x18\x01 \x01(\x03\x12:\n\x06values\x18\x02 \x03(\x0b\x32*.neptune.api.model.ProtoFloatPointValueDTO\"P\n\x17ProtoFloatPointValueDTO\x12\x18\n\x10timestamp_millis\x18\x01 \x01(\x03\x12\x0c\n\x04step\x18\x02 \x01(\x01\x12\r\n\x05value\x18\x03 \x01(\x01\x42\x34\n0ml.neptune.leaderboard.api.model.proto.generatedP\x01\x62\x06proto3') +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n(neptune_pb/api/model/series_values.proto\x12\x11neptune.api.model\"q\n\x19ProtoFloatSeriesValuesDTO\x12\x18\n\x10total_item_count\x18\x01 \x01(\x03\x12:\n\x06values\x18\x02 \x03(\x0b\x32*.neptune.api.model.ProtoFloatPointValueDTO\"P\n\x17ProtoFloatPointValueDTO\x12\x18\n\x10timestamp_millis\x18\x01 \x01(\x03\x12\x0c\n\x04step\x18\x02 \x01(\x01\x12\r\n\x05value\x18\x03 \x01(\x01\x42\x34\n0ml.neptune.leaderboard.api.model.proto.generatedP\x01\x62\x06proto3') _globals = globals() _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) -_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'model.series_values_pb2', _globals) +_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'neptune.api.proto.neptune_pb.api.model.series_values_pb2', _globals) if _descriptor._USE_C_DESCRIPTORS == False: _globals['DESCRIPTOR']._options = None _globals['DESCRIPTOR']._serialized_options = b'\n0ml.neptune.leaderboard.api.model.proto.generatedP\001' - _globals['_PROTOFLOATSERIESVALUESDTO']._serialized_start=48 - _globals['_PROTOFLOATSERIESVALUESDTO']._serialized_end=161 - _globals['_PROTOFLOATPOINTVALUEDTO']._serialized_start=163 - _globals['_PROTOFLOATPOINTVALUEDTO']._serialized_end=243 + _globals['_PROTOFLOATSERIESVALUESDTO']._serialized_start=63 + _globals['_PROTOFLOATSERIESVALUESDTO']._serialized_end=176 + _globals['_PROTOFLOATPOINTVALUEDTO']._serialized_start=178 + _globals['_PROTOFLOATPOINTVALUEDTO']._serialized_end=258 # @@protoc_insertion_point(module_scope) diff --git a/src/neptune/internal/backends/hosted_neptune_backend.py b/src/neptune/internal/backends/hosted_neptune_backend.py index 7a0bd15d7..3135aebb4 100644 --- a/src/neptune/internal/backends/hosted_neptune_backend.py +++ b/src/neptune/internal/backends/hosted_neptune_backend.py @@ -61,7 +61,10 @@ StringSeriesValues, StringSetField, ) -from neptune.api.proto.neptune_pb.api.model.attributes_pb2 import ProtoAttributesSearchResultDTO +from neptune.api.proto.neptune_pb.api.model.attributes_pb2 import ( + ProtoAttributesSearchResultDTO, + ProtoQueryAttributesResultDTO, +) from neptune.api.proto.neptune_pb.api.model.leaderboard_entries_pb2 import ProtoAttributesDTO from neptune.api.proto.neptune_pb.api.model.series_values_pb2 import ProtoFloatSeriesValuesDTO from neptune.api.searching_entries import iter_over_pages @@ -1193,12 +1196,14 @@ def query_fields_within_project( experiment_ids_filter: Optional[List[str]] = None, experiment_names_filter: Optional[List[str]] = None, next_page: Optional[NextPage] = None, + use_proto: Optional[bool] = None, ) -> QueryFieldsResult: + use_proto = use_proto if use_proto is not None else self.use_proto query = { - "experimentIdsFilter": experiment_ids_filter or None, - "experimentNamesFilter": experiment_names_filter or None, - "nextPage": next_page.to_dto() if next_page else None, + "experimentIdsFilter": experiment_ids_filter or None, + "experimentNamesFilter": experiment_names_filter or None, + "nextPage": next_page.to_dto() if next_page else None, } # If we are provided with both explicit column names, and a regex, @@ -1208,29 +1213,40 @@ def query_fields_within_project( if field_names_filter: # Make sure we don't pass too broad regex for explicit column names - terms += [f'^{name}$' for name in field_names_filter] + terms += [f"^{name}$" for name in field_names_filter] - regex = '|'.join(terms) - query['attributeNameFilter'] = {'mustMatchRegexes': [regex]} + regex = "|".join(terms) + query["attributeNameFilter"] = {"mustMatchRegexes": [regex]} elif field_names_filter: - query['attributeNamesFilter'] = field_names_filter + query["attributeNamesFilter"] = field_names_filter - params = {"projectIdentifier": project_id, 'query': query} + params = {"projectIdentifier": project_id, "query": query} try: - data = ( - self.leaderboard_client.api.queryAttributesWithinProject( - **params, - **DEFAULT_REQUEST_KWARGS, + if use_proto: + result = ( + self.leaderboard_client.api.queryAttributesWithinProjectProto( + **params, + **DEFAULT_PROTO_REQUEST_KWARGS, + ) + .response() + .result ) - .response() - .result - ) - return QueryFieldsResult.from_model(data) + data = ProtoQueryAttributesResultDTO.FromString(result) + return QueryFieldsResult.from_proto(data) + else: + data = ( + self.leaderboard_client.api.queryAttributesWithinProject( + **params, + **DEFAULT_REQUEST_KWARGS, + ) + .response() + .result + ) + return QueryFieldsResult.from_model(data) except HTTPNotFound: raise ProjectNotFound(project_id=project_id) - def get_fields_definitions( self, container_id: str, diff --git a/src/neptune/internal/backends/neptune_backend.py b/src/neptune/internal/backends/neptune_backend.py index 086b321dd..578a003ca 100644 --- a/src/neptune/internal/backends/neptune_backend.py +++ b/src/neptune/internal/backends/neptune_backend.py @@ -357,4 +357,5 @@ def query_fields_within_project( experiment_ids_filter: Optional[List[str]] = None, experiment_names_filter: Optional[List[str]] = None, next_page: Optional[NextPage] = None, + use_proto: Optional[bool] = None, ) -> QueryFieldsResult: ... diff --git a/src/neptune/internal/backends/neptune_backend_mock.py b/src/neptune/internal/backends/neptune_backend_mock.py index 74db4734e..9943532ff 100644 --- a/src/neptune/internal/backends/neptune_backend_mock.py +++ b/src/neptune/internal/backends/neptune_backend_mock.py @@ -828,6 +828,7 @@ def query_fields_within_project( experiment_ids_filter: Optional[List[str]] = None, experiment_names_filter: Optional[List[str]] = None, next_page: Optional[NextPage] = None, + use_proto: Optional[bool] = None, ) -> QueryFieldsResult: return QueryFieldsResult( entries=[], diff --git a/src/neptune/internal/backends/offline_neptune_backend.py b/src/neptune/internal/backends/offline_neptune_backend.py index 69b86327e..7cd3a980d 100644 --- a/src/neptune/internal/backends/offline_neptune_backend.py +++ b/src/neptune/internal/backends/offline_neptune_backend.py @@ -195,5 +195,6 @@ def query_fields_within_project( experiment_ids_filter: Optional[List[str]] = None, experiment_names_filter: Optional[List[str]] = None, next_page: Optional[NextPage] = None, + use_proto: Optional[bool] = None, ) -> QueryFieldsResult: raise NeptuneOfflineModeFetchException diff --git a/tests/unit/neptune/new/api/test_models.py b/tests/unit/neptune/new/api/test_models.py index a3a8e436d..a6c9539ce 100644 --- a/tests/unit/neptune/new/api/test_models.py +++ b/tests/unit/neptune/new/api/test_models.py @@ -48,6 +48,11 @@ StringSeriesValues, StringSetField, ) +from neptune.api.proto.neptune_pb.api.model.attributes_pb2 import ( + ProtoNextPageDTO, + ProtoQueryAttributesExperimentResultDTO, + ProtoQueryAttributesResultDTO, +) from neptune.api.proto.neptune_pb.api.model.leaderboard_entries_pb2 import ( ProtoAttributeDTO, ProtoAttributesDTO, @@ -2340,11 +2345,14 @@ def test__next_page__from_model(): def test__next_page__from_proto(): # given - proto = Mock() + proto = ProtoNextPageDTO(nextPageToken="some-token", limit=10) + + # when + result = NextPage.from_model(proto) # then - with pytest.raises(NotImplementedError): - NextPage.from_proto(proto) + assert result.next_page_token == "some-token" + assert result.limit == 10 def test__query_field_definitions_result__from_dict(): @@ -2515,11 +2523,51 @@ def test__query_fields_experiment_result__from_model(): def test__query_fields_experiment_result__from_proto(): # given - proto = Mock() + proto = ProtoQueryAttributesExperimentResultDTO( + experimentId="some-id-1", + experimentShortId="some-key-1", + attributes=[ + ProtoAttributeDTO( + name="some/float", + type="float", + float_properties=ProtoFloatAttributeDTO( + attribute_name="some/float", + attribute_type="float", + value=18.5, + ), + ), + ProtoAttributeDTO( + name="some/int", + type="int", + int_properties=ProtoIntAttributeDTO( + attribute_name="some/int", + attribute_type="int", + value=18, + ), + ), + ], + ) + + # when + result = QueryFieldsExperimentResult.from_proto(proto) # then - with pytest.raises(NotImplementedError): - QueryFieldsExperimentResult.from_proto(proto) + assert result.object_id == "some-id-1" + assert result.object_key == "some-key-1" + + assert len(result.fields) == 2 + + field_1 = result.fields[0] + assert field_1.path == "some/float" + assert field_1.type == FieldType.FLOAT + assert isinstance(field_1, FloatField) + assert field_1.value == 18.5 + + field_2 = result.fields[1] + assert field_2.path == "some/int" + assert field_2.type == FieldType.INT + assert isinstance(field_2, IntField) + assert field_2.value == 18 def test__query_fields_result__from_dict(): @@ -2677,8 +2725,84 @@ def test__query_fields_result__from_model(): def test__query_fields_result__from_proto(): # given - proto = Mock() + proto = ProtoQueryAttributesResultDTO( + entries=[ + ProtoQueryAttributesExperimentResultDTO( + experimentId="some-id-1", + experimentShortId="some-key-1", + attributes=[ + ProtoAttributeDTO( + name="some/float", + type="float", + float_properties=ProtoFloatAttributeDTO( + attribute_name="some/float", + attribute_type="float", + value=18.5, + ), + ), + ProtoAttributeDTO( + name="some/int", + type="int", + int_properties=ProtoIntAttributeDTO( + attribute_name="some/int", + attribute_type="int", + value=18, + ), + ), + ], + ), + ProtoQueryAttributesExperimentResultDTO( + experimentId="some-id-2", + experimentShortId="some-key-2", + attributes=[ + ProtoAttributeDTO( + name="some/string", + type="string", + string_properties=ProtoStringAttributeDTO( + attribute_name="some/string", + attribute_type="string", + value="hello", + ), + ), + ], + ), + ], + nextPage=ProtoNextPageDTO(nextPageToken="some-token", limit=2), + ) + + # when + result = QueryFieldsResult.from_proto(proto) # then - with pytest.raises(NotImplementedError): - QueryFieldsResult.from_proto(proto) + assert len(result.entries) == 2 + + entry_1 = result.entries[0] + assert entry_1.object_id == "some-id-1" + assert entry_1.object_key == "some-key-1" + assert len(entry_1.fields) == 2 + + field_1_1 = entry_1.fields[0] + assert field_1_1.path == "some/float" + assert field_1_1.type == FieldType.FLOAT + assert isinstance(field_1_1, FloatField) + assert field_1_1.value == 18.5 + + field_1_2 = entry_1.fields[1] + assert field_1_2.path == "some/int" + assert field_1_2.type == FieldType.INT + assert isinstance(field_1_2, IntField) + assert field_1_2.value == 18 + + entry_2 = result.entries[1] + assert entry_2.object_id == "some-id-2" + assert entry_2.object_key == "some-key-2" + assert len(entry_2.fields) == 1 + + field_2_1 = entry_2.fields[0] + assert field_2_1.path == "some/string" + assert field_2_1.type == FieldType.STRING + assert isinstance(field_2_1, StringField) + assert field_2_1.value == "hello" + + assert result.next_page.next_page_token == "some-token" + assert result.next_page.limit == 2