feat: directly fetch run_attribute_definitions

Author: michalsosn

src/neptune_query/internal/composition/attribute_components.py

@@ -37,6 +37,7 @@
     split,
     util,
 )
+from .run_attributes import fetch_run_attribute_definitions
 
 
 def fetch_attribute_definitions_split(
@@ -163,3 +164,29 @@ def fetch_attribute_values_split(
             downstream=downstream,
         ),
     )
+
+
+def fetch_run_attribute_definitions_split(
+    client: AuthenticatedClient,
+    project_identifier: identifiers.ProjectIdentifier,
+    attribute_filter: filters._BaseAttributeFilter,
+    executor: Executor,
+    fetch_attribute_definitions_executor: Executor,
+    sys_ids: list[identifiers.SysId],
+    downstream: Callable[[util.Page[identifiers.RunAttributeDefinition]], concurrency.OUT],
+) -> concurrency.OUT:
+    return concurrency.generate_concurrently(
+        items=split.split_sys_ids(sys_ids),
+        executor=executor,
+        downstream=lambda sys_ids_split: concurrency.generate_concurrently(
+            fetch_run_attribute_definitions(
+                client=client,
+                project_identifier=project_identifier,
+                run_identifiers=[identifiers.RunIdentifier(project_identifier, sys_id) for sys_id in sys_ids_split],
+                attribute_filter=attribute_filter,
+                executor=fetch_attribute_definitions_executor,
+            ),
+            executor=executor,
+            downstream=lambda run_definitions: downstream(run_definitions),
+        ),
+    )

src/neptune_query/internal/composition/attributes.py

@@ -32,6 +32,7 @@
 from ..composition import concurrency
 from ..retrieval import attribute_definitions as att_defs
 from ..retrieval import util
+from ..retrieval.attribute_filter import split_attribute_filters
 from ..retrieval.attribute_types import TYPE_AGGREGATIONS
 
 
@@ -126,7 +127,7 @@ def go_fetch_single(
             batch_size=batch_size,
         )
 
-    filters_ = att_defs.split_attribute_filters(attribute_filter)
+    filters_ = split_attribute_filters(attribute_filter)
 
     output = concurrency.generate_concurrently(
         items=(filter_ for filter_ in filters_),

src/neptune_query/internal/composition/fetch_metrics.py

@@ -30,7 +30,7 @@
     type_inference,
     validation,
 )
-from ..composition.attribute_components import fetch_attribute_definitions_split
+from ..composition.attribute_components import fetch_run_attribute_definitions_split
 from ..context import (
     Context,
     get_context,
@@ -145,25 +145,15 @@ def go_fetch_sys_attrs() -> Generator[list[identifiers.SysId], None, None]:
     output = concurrency.generate_concurrently(
         items=go_fetch_sys_attrs(),
         executor=executor,
-        downstream=lambda sys_ids: fetch_attribute_definitions_split(
+        downstream=lambda sys_ids: fetch_run_attribute_definitions_split(
             client=client,
             project_identifier=project_identifier,
             attribute_filter=attributes,
             executor=executor,
             fetch_attribute_definitions_executor=fetch_attribute_definitions_executor,
             sys_ids=sys_ids,
-            downstream=lambda sys_ids_split, definitions_page: concurrency.generate_concurrently(
-                items=split.split_series_attributes(
-                    items=(
-                        identifiers.RunAttributeDefinition(
-                            run_identifier=identifiers.RunIdentifier(project_identifier, sys_id),
-                            attribute_definition=definition,
-                        )
-                        for sys_id in sys_ids_split
-                        for definition in definitions_page.items
-                        if definition.type == "float_series"
-                    )
-                ),
+            downstream=lambda run_definitions_page: concurrency.generate_concurrently(
+                items=split.split_series_attributes(items=run_definitions_page.items),
                 executor=executor,
                 downstream=lambda run_attribute_definitions_split: concurrency.return_value(
                     fetch_multiple_series_values(

src/neptune_query/internal/composition/run_attributes.py

@@ -14,11 +14,9 @@
 # limitations under the License.
 
 from concurrent.futures import Executor
-from dataclasses import dataclass
 from typing import (
     Generator,
     Iterable,
-    Literal,
     Optional,
 )
 
@@ -30,50 +28,50 @@
     identifiers,
 )
 from ..composition import concurrency
-from ..retrieval import attribute_definitions as att_defs
+from ..retrieval import attribute_values as att_vals
 from ..retrieval import util
-from ..retrieval.attribute_types import TYPE_AGGREGATIONS
+from ..retrieval.attribute_filter import split_attribute_filters
 
 
 def fetch_run_attribute_definitions(
     client: AuthenticatedClient,
-    project_identifiers: Iterable[identifiers.ProjectIdentifier],
+    project_identifier: identifiers.ProjectIdentifier,
     run_identifiers: Optional[Iterable[identifiers.RunIdentifier]],
     attribute_filter: filters._BaseAttributeFilter,
     executor: Executor,
     batch_size: int = env.NEPTUNE_QUERY_ATTRIBUTE_DEFINITIONS_BATCH_SIZE.get(),
 ) -> Generator[util.Page[identifiers.RunAttributeDefinition], None, None]:
     pages_filters = _fetch_run_attribute_definitions(
-        client, project_identifiers, run_identifiers, attribute_filter, batch_size, executor
+        client, project_identifier, run_identifiers, attribute_filter, batch_size, executor
     )
 
     seen_items: set[identifiers.RunAttributeDefinition] = set()
-    for page, filter_ in pages_filters:
+    for page in pages_filters:
         new_items = [item for item in page.items if item not in seen_items]
         seen_items.update(new_items)
         yield util.Page(items=new_items)
 
 
 def _fetch_run_attribute_definitions(
     client: AuthenticatedClient,
-    project_identifiers: Iterable[identifiers.ProjectIdentifier],
+    project_identifier: identifiers.ProjectIdentifier,
     run_identifiers: Optional[Iterable[identifiers.RunIdentifier]],
     attribute_filter: filters._BaseAttributeFilter,
     batch_size: int,
     executor: Executor,
 ) -> Generator[util.Page[identifiers.RunAttributeDefinition], None, None]:
     def go_fetch_single(
         filter_: filters._AttributeFilter,
-    ) -> Generator[util.Page[identifiers.AttributeDefinition], None, None]:
-        return att_defs.fetch_attribute_definitions_single_filter(
+    ) -> Generator[util.Page[identifiers.RunAttributeDefinition], None, None]:
+        return att_vals.fetch_run_attribute_definitions_single_filter(
             client=client,
-            project_identifiers=project_identifiers,
+            project_identifier=project_identifier,
             run_identifiers=run_identifiers,
             attribute_filter=filter_,
             batch_size=batch_size,
         )
 
-    filters_ = att_defs.split_attribute_filters(attribute_filter)
+    filters_ = split_attribute_filters(attribute_filter)
 
     output = concurrency.generate_concurrently(
         items=(filter_ for filter_ in filters_),

src/neptune_query/internal/retrieval/attribute_definitions.py

@@ -13,14 +13,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import functools as ft
-import itertools as it
-import re
 from typing import (
     Any,
     Generator,
     Iterable,
     Optional,
-    Union,
 )
 
 from neptune_api.api.retrieval import query_attribute_definitions_within_project
@@ -40,17 +37,7 @@
 from ..retrieval import attribute_types as types  # noqa: E402
 from ..retrieval import util  # noqa: E402
 from ..retrieval import retry
-
-
-def split_attribute_filters(
-    _attribute_filter: filters._BaseAttributeFilter,
-) -> list[filters._AttributeFilter]:
-    if isinstance(_attribute_filter, filters._AttributeFilter):
-        return [_attribute_filter]
-    elif isinstance(_attribute_filter, filters._AttributeFilterAlternative):
-        return list(it.chain.from_iterable(split_attribute_filters(child) for child in _attribute_filter.filters))
-    else:
-        raise RuntimeError(f"Unexpected filter type: {type(_attribute_filter)}")
+from .attribute_filter import transform_attribute_filter_into_params
 
 
 def fetch_attribute_definitions_single_filter(
@@ -69,34 +56,8 @@ def fetch_attribute_definitions_single_filter(
     if run_identifiers is not None:
         params["experimentIdsFilter"] = [str(e) for e in run_identifiers]
 
-    # Convert name_eq to an additional condition added to each of must_match_any alternatives.
-    name_regexes = None
-    if attribute_filter.name_eq is not None:
-        name_regexes = _escape_name_eq(_variants_to_list(attribute_filter.name_eq))
-
-    if attribute_filter.must_match_any is not None:
-        attribute_name_filter_dtos = []
-        for alternative in attribute_filter.must_match_any:
-            attribute_name_filter_dto = {}
-            must_match_regexes = _union_options([name_regexes, alternative.must_match_regexes])
-            if must_match_regexes is not None:
-                attribute_name_filter_dto["mustMatchRegexes"] = must_match_regexes
-            if alternative.must_not_match_regexes is not None:
-                attribute_name_filter_dto["mustNotMatchRegexes"] = alternative.must_not_match_regexes
-            if attribute_name_filter_dto:
-                attribute_name_filter_dtos.append(attribute_name_filter_dto)
-        params["attributeNameFilter"]["mustMatchAny"] = attribute_name_filter_dtos
-
-    elif name_regexes is not None:
-        params["attributeNameFilter"]["mustMatchAny"] = [{"mustMatchRegexes": name_regexes}]
-
-    attribute_types = _variants_to_list(attribute_filter.type_in)
-    if attribute_types is not None:
-        params["attributeFilter"] = [
-            {"attributeType": types.map_attribute_type_python_to_backend(_type)} for _type in attribute_types
-        ]
-
-    # note: attribute_filter.aggregations is intentionally ignored
+    attribute_filter_params = transform_attribute_filter_into_params(attribute_filter)
+    params.update(attribute_filter_params)
 
     return util.fetch_pages(
         client=client,
@@ -149,36 +110,3 @@ def _make_new_attribute_definitions_page_params(
 
     params["nextPage"]["nextPageToken"] = next_page_token
     return params
-
-
-def _escape_name_eq(names: Optional[list[str]]) -> Optional[list[str]]:
-    if names is None:
-        return None
-
-    escaped = [f"{re.escape(name)}" for name in names]
-
-    if len(escaped) == 1:
-        return [f"^{escaped[0]}$"]
-    else:
-        joined = "|".join(escaped)
-        return [f"^({joined})$"]
-
-
-def _variants_to_list(param: Union[str, Iterable[str], None]) -> Optional[list[str]]:
-    if param is None:
-        return None
-    if isinstance(param, str):
-        return [param]
-    return list(param)
-
-
-def _union_options(options: list[Optional[list[str]]]) -> Optional[list[str]]:
-    result = None
-
-    for option in options:
-        if option is not None:
-            if result is None:
-                result = []
-            result.extend(option)
-
-    return result

src/neptune_query/internal/retrieval/attribute_filter.py

@@ -0,0 +1,107 @@
+#
+# Copyright (c) 2025, Neptune Labs Sp. z o.o.
+#
+# 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.
+import itertools as it
+import re
+from typing import (
+    Any,
+    Iterable,
+    Optional,
+    Union,
+)
+
+from .. import filters  # noqa: E402
+from ..retrieval import attribute_types as types  # noqa: E402
+
+
+def split_attribute_filters(
+    _attribute_filter: filters._BaseAttributeFilter,
+) -> list[filters._AttributeFilter]:
+    if isinstance(_attribute_filter, filters._AttributeFilter):
+        return [_attribute_filter]
+    elif isinstance(_attribute_filter, filters._AttributeFilterAlternative):
+        return list(it.chain.from_iterable(split_attribute_filters(child) for child in _attribute_filter.filters))
+    else:
+        raise RuntimeError(f"Unexpected filter type: {type(_attribute_filter)}")
+
+
+def transform_attribute_filter_into_params(
+    attribute_filter: filters._AttributeFilter,
+) -> dict[str, Any]:
+    params: dict[str, Any] = {
+        "attributeNameFilter": {},
+    }
+
+    name_regexes = None
+    if attribute_filter.name_eq is not None:
+        name_regexes = _escape_name_eq(_variants_to_list(attribute_filter.name_eq))
+
+    if attribute_filter.must_match_any is not None:
+        attribute_name_filter_dtos = []
+        for alternative in attribute_filter.must_match_any:
+            attribute_name_filter_dto = {}
+            must_match_regexes = _union_options([name_regexes, alternative.must_match_regexes])
+            if must_match_regexes is not None:
+                attribute_name_filter_dto["mustMatchRegexes"] = must_match_regexes
+            if alternative.must_not_match_regexes is not None:
+                attribute_name_filter_dto["mustNotMatchRegexes"] = alternative.must_not_match_regexes
+            if attribute_name_filter_dto:
+                attribute_name_filter_dtos.append(attribute_name_filter_dto)
+        params["attributeNameFilter"]["mustMatchAny"] = attribute_name_filter_dtos
+
+    elif name_regexes is not None:
+        params["attributeNameFilter"]["mustMatchAny"] = [{"mustMatchRegexes": name_regexes}]
+
+    attribute_types = _variants_to_list(attribute_filter.type_in)
+    if attribute_types is not None:
+        params["attributeFilter"] = [
+            {"attributeType": types.map_attribute_type_python_to_backend(_type)} for _type in attribute_types
+        ]
+
+    # note: attribute_filter.aggregations is intentionally ignored
+
+    return params
+
+
+def _escape_name_eq(names: Optional[list[str]]) -> Optional[list[str]]:
+    if names is None:
+        return None
+
+    escaped = [f"{re.escape(name)}" for name in names]
+
+    if len(escaped) == 1:
+        return [f"^{escaped[0]}$"]
+    else:
+        joined = "|".join(escaped)
+        return [f"^({joined})$"]
+
+
+def _variants_to_list(param: Union[str, Iterable[str], None]) -> Optional[list[str]]:
+    if param is None:
+        return None
+    if isinstance(param, str):
+        return [param]
+    return list(param)
+
+
+def _union_options(options: list[Optional[list[str]]]) -> Optional[list[str]]:
+    result = None
+
+    for option in options:
+        if option is not None:
+            if result is None:
+                result = []
+            result.extend(option)
+
+    return result