feat: Raw api polars integration

cognite/client/_api/raw/rows.py

@@ -27,6 +27,7 @@
 
 if TYPE_CHECKING:
     import pandas as pd
+    import polars as pl
 
     from cognite.client import AsyncCogniteClient
     from cognite.client.config import ClientConfig
@@ -517,6 +518,51 @@ async def retrieve_dataframe(
         cols = [r.columns for r in rows]
         return pd.DataFrame(cols, index=idx, dtype=object if not infer_dtypes else None)
 
+    async def retrieve_polars_dataframe(
+        self,
+        db_name: str,
+        table_name: str,
+        min_last_updated_time: int | None = None,
+        max_last_updated_time: int | None = None,
+        columns: list[str] | None = None,
+        limit: int | None = DEFAULT_LIMIT_READ,
+        partitions: int | None = None,
+        include_last_updated_time: bool = False,
+    ) -> pl.DataFrame:
+        """Retrieve rows in a table as a polars dataframe.
+
+        Rowkeys are used as the index.
+
+        Args:
+            db_name (str): Name of the database.
+            table_name (str): Name of the table.
+            min_last_updated_time (int | None): Rows must have been last updated after this time (exclusive). Milliseconds since epoch.
+            max_last_updated_time (int | None): Rows must have been last updated before this time (inclusive). Milliseconds since epoch.
+            columns (list[str] | None): List of column keys. Set to `None` to retrieving all, use empty list, [], to retrieve only row keys.
+            limit (int | None): The number of rows to retrieve. Defaults to 25. Set to -1, float("inf") or None to return all items.
+            partitions (int | None): Retrieve rows in parallel using this number of workers. Can be used together with a (large) finite limit.
+            include_last_updated_time (bool): Whether to include the last_updated_time in the dataframe. If True, 'last_updated_time' will be added to the dataframe as a column. Defaults to False.
+        Returns:
+            pl.DataFrame: The requested rows in a polars dataframe.
+
+        Examples:
+            Get dataframe:
+
+                >>> from cognite.client import CogniteClient, AsyncCogniteClient
+                >>> client = CogniteClient()
+                >>> # async_client = AsyncCogniteClient()  # another option
+                >>> df = client.raw.rows.retrieve_polars_dataframe("db1", "t1", limit=5)
+        """
+        pl = local_import("polars")
+        rows = await self.list(
+            db_name, table_name, min_last_updated_time, max_last_updated_time, columns, limit, partitions
+        )
+        if include_last_updated_time:
+            data = [{"key": r.key, **r.columns, "last_updated_time": r.last_updated_time} for r in rows]
+        else:
+            data = [{"key": r.key, **r.columns} for r in rows]
+        return pl.DataFrame(data)
+
     async def _get_parallel_cursors(
         self,
         db_name: str,

cognite/client/data_classes/_base.py

@@ -36,6 +36,7 @@
 
 if TYPE_CHECKING:
     import pandas
+    import polars
 
     from cognite.client import AsyncCogniteClient
 
@@ -188,6 +189,14 @@ def _repr_html_(self) -> str:
     def _maybe_set_client_ref(self, client: AsyncCogniteClient) -> Self:
         return self  # Base resource has no client ref set
 
+    def to_polars(self) -> polars.DataFrame:
+        """Convert the instance into a polars DataFrame.
+
+        Returns:
+            polars.DataFrame: The polars DataFrame representing the instance.
+        """
+        raise NotImplementedError(f"to_polars is not implemented for {type(self).__name__}")
+
 
 class UnknownCogniteResource(CogniteResource):
     def __init__(self, data: dict[str, Any]) -> None:

cognite/client/data_classes/raw.py

@@ -18,6 +18,7 @@
 
 if TYPE_CHECKING:
     import pandas
+    import polars
 
 
 class RowCore(WriteableCogniteResource["RowWrite"], ABC):
@@ -66,6 +67,15 @@ def to_pandas(self) -> pandas.DataFrame:  # type: ignore[override]
         pd = local_import("pandas")
         return pd.DataFrame([self.columns], [self.key])
 
+    def to_polars(self) -> polars.DataFrame:
+        """Convert the instance into a polars DataFrame.
+
+        Returns:
+            polars.DataFrame: The polars DataFrame representing this instance.
+        """
+        pl = local_import("polars")
+        return pl.DataFrame(data={"key": self.key, **self.columns})
+
 
 T_Row = TypeVar("T_Row", bound=RowCore)
 
@@ -133,6 +143,18 @@ def to_pandas(self) -> pandas.DataFrame:  # type: ignore[override]
         index, data = zip(*((row.key, row.columns) for row in self))
         return pd.DataFrame.from_records(data, index=index)
 
+    def to_polars(self) -> polars.DataFrame:
+        """Convert the instance into a polars DataFrame.
+
+        Returns:
+            polars.DataFrame: The polars DataFrame representing this instance.
+        """
+        pl = local_import("polars")
+        if not self:
+            return pl.DataFrame(data=[])
+        data = ({"key": row.key, **row.columns} for row in self)
+        return pl.DataFrame(data=data)
+
 
 class RowWriteList(RowListCore[RowWrite]):
     _RESOURCE = RowWrite

pyproject.toml

@@ -75,7 +75,8 @@ dev = [
     "types-authlib (>=1.6.2.20250825, <2)",
     "anyio (>=4.11.0, <5)",
     "testbook (>=0.4.2, <0.5)",
-    "ipykernel (>=7.2.0, <8)"
+    "ipykernel (>=7.2.0, <8)",
+    "polars[rtcompat] (>=1.41.2,<2.0.0)",
 ]
 docs = [
     "sphinx ==8.1.*",  # Increase when Python 3.10 support is dropped