Skip to content
12 changes: 8 additions & 4 deletions cognite/client/_api/data_modeling/instances.py
Original file line number Diff line number Diff line change
Expand Up @@ -894,6 +894,7 @@ async def subscribe(
>>> subscription_context.cancel()

"""
query._prepare_sorts()
subscription_context = SubscriptionContext()

async def _poll_loop() -> None:
Expand Down Expand Up @@ -968,10 +969,11 @@ def _create_other_params(
f"Received in `sources` argument for views: {with_properties}."
)
if sort:
if isinstance(sort, (InstanceSort, dict)):
other_params["sort"] = [cls._dump_instance_sort(sort)]
else:
other_params["sort"] = [cls._dump_instance_sort(s) for s in sort]
sorts_seq = [sort] if isinstance(sort, (InstanceSort, dict)) else list(sort)
for s in sorts_seq:
if isinstance(s, InstanceSort):
s._apply_postgres_defaults_or_maybe_warn()
other_params["sort"] = [cls._dump_instance_sort(s) for s in sorts_seq]
if instance_type:
other_params["instanceType"] = instance_type
if debug:
Expand Down Expand Up @@ -1647,6 +1649,7 @@ async def query(
>>> res = client.data_modeling.instances.query(query, debug=debug_params)
>>> print(res.debug)
"""
query._prepare_sorts()
return await self._query_or_sync(query, "query", include_typing=include_typing, debug=debug)

async def sync(
Expand Down Expand Up @@ -1749,6 +1752,7 @@ async def sync(
>>> res = client.data_modeling.instances.sync(query, debug=debug_params)
>>> print(res.debug)
"""
query._prepare_sorts()
return await self._query_or_sync(query, "sync", include_typing=include_typing, debug=debug)

async def _query_or_sync(
Expand Down
2 changes: 1 addition & 1 deletion cognite/client/_sync_api/data_modeling/instances.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
"""
===============================================================================
c0af4f83cd8ffa0aaed6fa641bde9a98
f57e38da23620654d634c19560399f41
This file is auto-generated from the Async API modules, - do not edit manually!
===============================================================================
"""
Expand Down
104 changes: 101 additions & 3 deletions cognite/client/data_classes/data_modeling/instances.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@
overload,
)

from typing_extensions import Self
from typing_extensions import Self, override

from cognite.client._constants import OMITTED, Omitted
from cognite.client.data_classes._base import (
Expand Down Expand Up @@ -1344,13 +1344,111 @@ class InstancesApply:


class InstanceSort(DataModelingSort):
"""Sort order for an instance query.

Args:
property (list[str] | tuple[str, str] | tuple[str, str, str]): The property to sort by, given as a path, e.g.
``("mySpace", "myView/v1", "myProperty")`` or ``["node", "externalId"]``.
direction (Literal['ascending', 'descending']): Sort direction. Case-insensitive. Defaults to ``"ascending"``.
nulls_first (bool | None): Where to place ``null`` values. Defaults to ``None`` (auto). See tip below.

Tip:
For the backend database to use an index when sorting nullable properties, the ``nulls_first`` setting
must match the sort direction:

- ``ascending`` → nulls last (``nulls_first=False``)
- ``descending`` → nulls first (``nulls_first=True``)

When ``nulls_first=None`` (the default), the correct value is chosen automatically. Passing the
opposite combination is still accepted and sent to the API as-is, but may trigger a warning
for API endpoints that support index utilization if an unsupported combination is used.

Examples:

Sort by a view property ascending (default):

>>> from cognite.client.data_classes.data_modeling import InstanceSort
>>> sort = InstanceSort(("mySpace", "myView/v1", "myProperty"))

Can also use a ViewId to simplify the property path:

>>> from cognite.client.data_classes.data_modeling import ViewId
>>> view_id = ViewId("mySpace", "myView", "v1")
>>> sort = InstanceSort(view_id.as_property_ref("myProperty"))

Sort descending:

>>> sort = InstanceSort(
... view_id.as_property_ref("myProperty"),
... direction="descending",
... )

Sort by a base property:

>>> sort = InstanceSort(["node", "externalId"], direction="ascending")

Force a specific null placement (first/last). A UserWarning will fire at relevant API call
sites when this conflicts with index alignment:

>>> sort = InstanceSort(
... ("mySpace", "myView/v1", "myProperty"),
... direction="descending",
... nulls_first=True,
... )
"""

def __init__(
self,
property: list[str] | tuple[str, ...],
property: list[str] | tuple[str, str] | tuple[str, str, str],
direction: Literal["ascending", "descending"] = "ascending",
nulls_first: bool | None = None,
) -> None:
super().__init__(property, direction, nulls_first)
normalized = direction.casefold()
if normalized not in ("ascending", "descending"):
raise ValueError(f"direction must be 'ascending' or 'descending', got {direction!r}")

super().__init__(property, normalized, nulls_first) # type: ignore [arg-type]

# We override _load to get the more strict __init__ validation on 'direction' because we need it to
# be valid for the possible later automatic choice of nulls_first:
@override
@classmethod
def _load(cls, resource: dict[str, Any]) -> Self:
if not isinstance(resource, dict):
raise TypeError(f"Resource must be mapping, not {type(resource)}")

return cls(
property=resource["property"],
direction=resource.get("direction", "ascending"),
nulls_first=resource.get("nullsFirst"),
)

@property
def is_index_aligned(self) -> bool:
"""True when nulls_first matches the direction for PostgreSQL index utilization (None counts as aligned)."""
if self.nulls_first is None:
return True
return self.nulls_first is (self.direction == "descending")

def _apply_postgres_defaults_or_maybe_warn(self) -> None:
Comment thread
haakonvt marked this conversation as resolved.
Outdated
"""Resolve nulls_first for PostgreSQL index alignment, warning if the explicit value is misaligned.

When nulls_first is None, sets it to the index-compatible value. When explicitly set but
misaligned, emits a UserWarning.
"""
if self.nulls_first is None:
self.nulls_first = self.direction == "descending"

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

The _apply_postgres_defaults_or_maybe_warn method mutates the InstanceSort object by assigning a value to self.nulls_first when it is None. This side effect can cause subsequent calls to instances.search() to fail with a ValueError if the same InstanceSort object is reused, because search() (which uses Elasticsearch) explicitly prohibits non-null nulls_first values (see cognite/client/_api/data_modeling/instances.py:1311).

To avoid this, consider resolving the default value during the payload construction phase in the API client, or providing a way to get the resolved value without modifying the instance state.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good callout, will consider

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@gemini-code-assist I've adressed this now, please do a re-review

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you for addressing the concerns regarding object mutation and the potential side effects on instances.search(). The implementation of _get_query_with_defaults_applied() using deepcopy and the targeted application of defaults across query traversals effectively resolves the issue while maintaining the integrity of the original InstanceSort objects. The addition of comprehensive unit tests covering these scenarios is also appreciated.


elif not self.is_index_aligned:
import warnings

warnings.warn(
f"InstanceSort: direction={self.direction!r} with nulls_first={self.nulls_first} is not "
f"index-aligned and will likely prevent database index utilization. "
f"Use nulls_first={self.direction == 'descending'} (or omit it) for optimal performance.",
UserWarning,
stacklevel=3,
)


@dataclass
Expand Down
31 changes: 30 additions & 1 deletion cognite/client/data_classes/data_modeling/query.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@

from abc import ABC
from collections import UserDict
from collections.abc import Mapping, MutableMapping, Sequence
from collections.abc import Iterator, Mapping, MutableMapping, Sequence
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Generic, Literal, TypeVar

Expand Down Expand Up @@ -85,6 +85,9 @@ def _load_list(
class SelectBase(CogniteResource, ABC):
sources: list[SourceSelector] = field(default_factory=list)

def _iter_sorts(self) -> Iterator[InstanceSort]:
yield from ()

def dump(self, camel_case: bool = True) -> dict[str, Any]:
output: dict[str, Any] = {}
if self.sources:
Expand Down Expand Up @@ -133,6 +136,9 @@ class Select(SelectBase):
sort: list[InstanceSort] = field(default_factory=list)
limit: int | None = None

def _iter_sorts(self) -> Iterator[InstanceSort]:
yield from self.sort

def dump(self, camel_case: bool = True) -> dict[str, Any]:
output = super().dump(camel_case)
if self.sort:
Expand Down Expand Up @@ -179,6 +185,16 @@ def instance_type_by_result_expression(self) -> dict[str, type[NodeListWithCurso
for k, v in self.with_.items()
}

def _iter_sorts(self) -> Iterator[InstanceSort]:
for expr in self.with_.values():
yield from expr._iter_sorts()
for sel in self.select.values():
yield from sel._iter_sorts()

def _prepare_sorts(self) -> None:
for sort in self._iter_sorts():
sort._apply_postgres_defaults_or_maybe_warn()

def dump(self, camel_case: bool = True) -> dict[str, Any]:
output: dict[str, Any] = {
"with": {k: v.dump(camel_case) for k, v in self.with_.items()},
Expand Down Expand Up @@ -281,6 +297,9 @@ class ResultSetExpressionBase(CogniteResource, ABC):
def _load_sort(resource: dict[str, Any], name: str) -> list[InstanceSort]:
return [InstanceSort.load(sort) for sort in resource.get(name, [])]

def _iter_sorts(self) -> Iterator[InstanceSort]:
yield from ()

@staticmethod
def _init_through(through: list[str] | tuple[str, str, str] | PropertyId | None) -> PropertyId | None:
def error() -> Never:
Expand Down Expand Up @@ -336,6 +355,9 @@ def __eq__(self, other: object) -> bool:
return NotImplemented
return type(self) is type(other) and self.dump() == other.dump()

def _iter_sorts(self) -> Iterator[InstanceSort]:
yield from self.sort


@dataclass(eq=False) # Prevents @dataclass from generating its own __eq__, so the parent's is used
class NodeResultSetExpression(NodeOrEdgeResultSetExpression):
Expand Down Expand Up @@ -429,6 +451,10 @@ class EdgeResultSetExpression(NodeOrEdgeResultSetExpression):
limit_each: int | None = None
post_sort: list[InstanceSort] = field(default_factory=list)

def _iter_sorts(self) -> Iterator[InstanceSort]:
yield from self.sort
yield from self.post_sort

@classmethod
def _load(cls, resource: dict[str, Any]) -> Self:
query_edge = resource["edges"]
Expand Down Expand Up @@ -494,6 +520,9 @@ def __eq__(self, other: object) -> bool:
return NotImplemented
return type(self) is type(other) and self.dump() == other.dump()

def _iter_sorts(self) -> Iterator[InstanceSort]:
yield from self.backfill_sort

@classmethod
def _load(cls, resource: dict[str, Any]) -> ResultSetExpressionSync:
if "nodes" in resource:
Expand Down
Loading
Loading