Optimize table fetches for explicit attribute lists via entries-search projection (#210)

PatrykGala · patrykgala-oai · web-flow · commit 57012693d5aa · 2026-02-12T16:43:43.000Z
Co-authored-by: Patryk Gała &lt;patryk.gala@openai.com&gt;
diff --git a/src/neptune_query/__init__.py b/src/neptune_query/__init__.py
@@ -304,6 +304,7 @@ def fetch_experiments_table(
         sort_direction=sort_direction,
         limit=limit,
         type_suffix_in_column_names=type_suffix_in_column_names,
+        exact_attribute_names=attributes if isinstance(attributes, list) else None,
         container_type=_search.ContainerType.EXPERIMENT,
     )
 
diff --git a/src/neptune_query/internal/composition/fetch_table.py b/src/neptune_query/internal/composition/fetch_table.py
@@ -24,7 +24,10 @@
 from ...exceptions import NeptuneUserError
 from .. import client as _client
 from .. import context as _context
-from .. import identifiers
+from .. import (
+    env,
+    identifiers,
+)
 from ..composition import attribute_components as _components
 from ..composition import (
     concurrency,
@@ -61,6 +64,7 @@ def fetch_table(
     sort_direction: Literal["asc", "desc"],
     limit: Optional[int],
     type_suffix_in_column_names: bool,
+    exact_attribute_names: Optional[list[str]] = None,
     context: Optional[_context.Context] = None,
     container_type: search.ContainerType,
 ) -> pd.DataFrame:
@@ -93,6 +97,30 @@ def fetch_table(
         sort_by = sort_by_inference_result.result
         sort_by_inference_result.emit_warnings()
 
+        if exact_attribute_names is not None:
+            exact_attribute_names_set = set(exact_attribute_names)
+            if len(exact_attribute_names) <= env.NEPTUNE_QUERY_ENTRIES_SEARCH_MAX_PROJECTION_ATTRIBUTES.get():
+                table_rows: list[TableRow] = []
+                for page in search.fetch_table_rows_exact_attributes(
+                    client=client,
+                    project_identifier=project_identifier,
+                    filter_=filter_,
+                    requested_attribute_names=exact_attribute_names_set,
+                    sort_by=sort_by,
+                    sort_direction=_sort_direction,
+                    limit=limit,
+                    container_type=container_type,
+                ):
+                    for item in page.items:
+                        table_rows.append(
+                            TableRow(values=item.values, label=item.label, project_identifier=project_identifier)
+                        )
+                return create_runs_table(
+                    table_rows=table_rows,
+                    type_suffix_in_column_names=type_suffix_in_column_names,
+                    container_type=container_type,
+                )
+
         sys_id_label_mapping: dict[identifiers.SysId, str] = {}
         result_by_id: dict[identifiers.SysId, list[att_vals.AttributeValue]] = {}
 
diff --git a/src/neptune_query/internal/env.py b/src/neptune_query/internal/env.py
@@ -27,6 +27,7 @@
     "NEPTUNE_PROJECT",
     "NEPTUNE_QUERY_ATTRIBUTE_DEFINITIONS_BATCH_SIZE",
     "NEPTUNE_QUERY_ATTRIBUTE_VALUES_BATCH_SIZE",
+    "NEPTUNE_QUERY_ENTRIES_SEARCH_MAX_PROJECTION_ATTRIBUTES",
     "NEPTUNE_QUERY_FILES_MAX_CONCURRENCY",
     "NEPTUNE_QUERY_FILES_TIMEOUT",
     "NEPTUNE_QUERY_MAX_REQUEST_SIZE",
@@ -95,6 +96,9 @@ def _map_logging_level(value: str) -> str:
     "NEPTUNE_QUERY_ATTRIBUTE_DEFINITIONS_BATCH_SIZE", int, 10_000
 )
 NEPTUNE_QUERY_ATTRIBUTE_VALUES_BATCH_SIZE = EnvVariable[int]("NEPTUNE_QUERY_ATTRIBUTE_VALUES_BATCH_SIZE", int, 10_000)
+NEPTUNE_QUERY_ENTRIES_SEARCH_MAX_PROJECTION_ATTRIBUTES = EnvVariable[int](
+    "NEPTUNE_QUERY_ENTRIES_SEARCH_MAX_PROJECTION_ATTRIBUTES", int, 1000
+)
 NEPTUNE_QUERY_FILES_BATCH_SIZE = EnvVariable[int]("NEPTUNE_QUERY_FILES_BATCH_SIZE", int, 200)
 NEPTUNE_QUERY_FILES_MAX_CONCURRENCY = EnvVariable[int]("NEPTUNE_QUERY_FILES_MAX_CONCURRENCY", int, 1)
 NEPTUNE_QUERY_FILES_TIMEOUT = EnvVariable[Optional[int]]("NEPTUNE_QUERY_FILES_TIMEOUT", _lift_optional(int), None)
diff --git a/src/neptune_query/internal/retrieval/search.py b/src/neptune_query/internal/retrieval/search.py
@@ -43,11 +43,16 @@
     _Filter,
 )
 from ..logger import get_logger
+from ..retrieval import attribute_values as att_vals
 from ..retrieval import (
     retry,
     util,
 )
-from ..retrieval.attribute_types import map_attribute_type_python_to_backend
+from ..retrieval.attribute_types import (
+    extract_value,
+    map_attribute_type_backend_to_python,
+    map_attribute_type_python_to_backend,
+)
 
 logger = get_logger()
 
@@ -77,6 +82,13 @@ def sys_id(self) -> identifiers.SysId: ...
     def label(self) -> str: ...
 
 
+@dataclass(frozen=True)
+class TableSearchEntry:
+    sys_id: identifiers.SysId
+    label: str
+    values: list[att_vals.AttributeValue]
+
+
 @dataclass(frozen=True)
 class ExperimentSysAttrs:
     sys_id: identifiers.SysId
@@ -137,6 +149,65 @@ def __call__(
     ) -> Generator[util.Page[T], None, None]: ...
 
 
+def _build_entries_search_params(
+    *,
+    attribute_projection: list[str],
+    batch_size: int,
+    container_type: ContainerType,
+    filter_: Optional[_Filter],
+    sort_by: _Attribute,
+    sort_direction: Literal["asc", "desc"],
+) -> dict[str, Any]:
+    params: dict[str, Any] = {
+        "attributeFilters": [{"path": attribute_name} for attribute_name in attribute_projection],
+        "pagination": {"limit": batch_size},
+        "experimentLeader": container_type == ContainerType.EXPERIMENT,
+        "sorting": {
+            "dir": _map_direction(sort_direction),
+            "sortBy": {"name": sort_by.name},
+        },
+    }
+    if filter_ is not None:
+        params["query"] = {"query": str(filter_)}
+    if sort_by.aggregation is not None:
+        params["sorting"]["aggregationMode"] = sort_by.aggregation
+    if sort_by.type is not None:
+        params["sorting"]["sortBy"]["type"] = map_attribute_type_python_to_backend(sort_by.type)
+
+    return params
+
+
+def _fetch_entries_with_projection(
+    *,
+    client: AuthenticatedClient,
+    project_identifier: identifiers.ProjectIdentifier,
+    attribute_projection: list[str],
+    process_page: Callable[[ProtoLeaderboardEntriesSearchResultDTO], util.Page[T]],
+    filter_: Optional[_Filter],
+    sort_by: _Attribute,
+    sort_direction: Literal["asc", "desc"],
+    limit: Optional[int],
+    batch_size: int,
+    container_type: ContainerType,
+) -> Generator[util.Page[T], None, None]:
+    params = _build_entries_search_params(
+        attribute_projection=attribute_projection,
+        batch_size=batch_size,
+        container_type=container_type,
+        filter_=filter_,
+        sort_by=sort_by,
+        sort_direction=sort_direction,
+    )
+
+    return util.fetch_pages(
+        client=client,
+        fetch_page=ft.partial(_fetch_sys_attrs_page, project_identifier=project_identifier),
+        process_page=process_page,
+        make_new_page_params=ft.partial(_make_new_sys_attrs_page_params, batch_size=batch_size, limit=limit),
+        initial_params=params,
+    )
+
+
 def _create_fetch_sys_attrs(
     attribute_names: List[str],
     make_record: Callable[[dict[str, Any]], T],
@@ -152,28 +223,17 @@ def fetch_sys_attrs(
         batch_size: int = env.NEPTUNE_QUERY_SYS_ATTRS_BATCH_SIZE.get(),
         container_type: ContainerType = default_container_type,
     ) -> Generator[util.Page[T], None, None]:
-        params: dict[str, Any] = {
-            "attributeFilters": [{"path": attribute_name} for attribute_name in attribute_names],
-            "pagination": {"limit": batch_size},
-            "experimentLeader": container_type == ContainerType.EXPERIMENT,
-            "sorting": {
-                "dir": _map_direction(sort_direction),
-                "sortBy": {"name": sort_by.name},
-            },
-        }
-        if filter_ is not None:
-            params["query"] = {"query": str(filter_)}
-        if sort_by.aggregation is not None:
-            params["sorting"]["aggregationMode"] = sort_by.aggregation
-        if sort_by.type is not None:
-            params["sorting"]["sortBy"]["type"] = map_attribute_type_python_to_backend(sort_by.type)
-
-        return util.fetch_pages(
+        return _fetch_entries_with_projection(
             client=client,
-            fetch_page=ft.partial(_fetch_sys_attrs_page, project_identifier=project_identifier),
+            project_identifier=project_identifier,
+            attribute_projection=attribute_names,
             process_page=ft.partial(_process_sys_attrs_page, make_record=make_record),
-            make_new_page_params=ft.partial(_make_new_sys_attrs_page_params, batch_size=batch_size, limit=limit),
-            initial_params=params,
+            filter_=filter_,
+            sort_by=sort_by,
+            sort_direction=sort_direction,
+            limit=limit,
+            batch_size=batch_size,
+            container_type=container_type,
         )
 
     return fetch_sys_attrs
@@ -215,6 +275,42 @@ def fetch_sys_id_labels(container_type: ContainerType) -> FetchSysAttrs[SysIdLab
 fetch_sys_ids = fetch_experiment_sys_ids
 
 
+def fetch_table_rows_exact_attributes(
+    *,
+    client: AuthenticatedClient,
+    project_identifier: identifiers.ProjectIdentifier,
+    filter_: Optional[_Filter],
+    requested_attribute_names: set[str],
+    sort_by: _Attribute,
+    sort_direction: Literal["asc", "desc"],
+    limit: Optional[int],
+    container_type: ContainerType,
+) -> Generator[util.Page[TableSearchEntry], None, None]:
+    batch_size = env.NEPTUNE_QUERY_SYS_ATTRS_BATCH_SIZE.get()
+
+    label_attribute_name = "sys/name" if container_type == ContainerType.EXPERIMENT else "sys/custom_run_id"
+    projection_attribute_names = set(requested_attribute_names)
+    projection_attribute_names.update({"sys/id", label_attribute_name})
+
+    yield from _fetch_entries_with_projection(
+        client=client,
+        project_identifier=project_identifier,
+        attribute_projection=list(projection_attribute_names),
+        process_page=ft.partial(
+            _process_table_rows_exact_attributes_page,
+            project_identifier=project_identifier,
+            label_attribute_name=label_attribute_name,
+            requested_attribute_names=requested_attribute_names,
+        ),
+        filter_=filter_,
+        sort_by=sort_by,
+        sort_direction=sort_direction,
+        limit=limit,
+        batch_size=batch_size,
+        container_type=container_type,
+    )
+
+
 def _fetch_sys_attrs_page(
     client: AuthenticatedClient,
     params: dict[str, Any],
@@ -249,6 +345,49 @@ def _process_sys_attrs_page(
     return util.Page(items=items)
 
 
+def _process_table_rows_exact_attributes_page(
+    data: ProtoLeaderboardEntriesSearchResultDTO,
+    project_identifier: identifiers.ProjectIdentifier,
+    label_attribute_name: str,
+    requested_attribute_names: set[str],
+) -> util.Page[TableSearchEntry]:
+    items: list[TableSearchEntry] = []
+
+    for entry in data.entries:
+        attributes_by_name = {
+            attr.name: attr
+            for attr in entry.attributes
+            if attr.name in ("sys/id", label_attribute_name) and attr.HasField("string_properties")
+        }
+        label = attributes_by_name[label_attribute_name].string_properties.value
+        sys_id = identifiers.SysId(attributes_by_name["sys/id"].string_properties.value)
+        run_identifier = identifiers.RunIdentifier(project_identifier=project_identifier, sys_id=sys_id)
+
+        values: list[att_vals.AttributeValue] = []
+        for attr in entry.attributes:
+            if attr.name not in requested_attribute_names:
+                continue
+
+            item_value = extract_value(attr)
+            if item_value is None:
+                continue
+
+            values.append(
+                att_vals.AttributeValue(
+                    attribute_definition=identifiers.AttributeDefinition(
+                        name=attr.name,
+                        type=map_attribute_type_backend_to_python(attr.type),
+                    ),
+                    value=item_value,
+                    run_identifier=run_identifier,
+                )
+            )
+
+        items.append(TableSearchEntry(sys_id=sys_id, label=label, values=values))
+
+    return util.Page(items=items)
+
+
 def _make_new_sys_attrs_page_params(
     params: dict[str, Any],
     data: Optional[ProtoLeaderboardEntriesSearchResultDTO],
diff --git a/src/neptune_query/runs.py b/src/neptune_query/runs.py
@@ -288,6 +288,7 @@ def fetch_runs_table(
         sort_direction=sort_direction,
         limit=limit,
         type_suffix_in_column_names=type_suffix_in_column_names,
+        exact_attribute_names=attributes if isinstance(attributes, list) else None,
         container_type=_search.ContainerType.RUN,
     )
 
diff --git a/tests/unit/internal/test_output_format.py b/tests/unit/internal/test_output_format.py
diff --git a/tests/unit/v1/test_split.py b/tests/unit/v1/test_split.py

Original file line number	Diff line number	Diff line change
`@@ -304,6 +304,7 @@ def fetch_experiments_table(`
`304`	`304`	`sort_direction=sort_direction,`
`305`	`305`	`limit=limit,`
`306`	`306`	`type_suffix_in_column_names=type_suffix_in_column_names,`
	`307`	`+ exact_attribute_names=attributes if isinstance(attributes, list) else None,`
`307`	`308`	`container_type=_search.ContainerType.EXPERIMENT,`
`308`	`309`	`)`
`309`	`310`
Original file line number	Diff line number	Diff line change
`@@ -288,6 +288,7 @@ def fetch_runs_table(`
`288`	`288`	`sort_direction=sort_direction,`
`289`	`289`	`limit=limit,`
`290`	`290`	`type_suffix_in_column_names=type_suffix_in_column_names,`
	`291`	`+ exact_attribute_names=attributes if isinstance(attributes, list) else None,`
`291`	`292`	`container_type=_search.ContainerType.RUN,`
`292`	`293`	`)`
`293`	`294`