Maintain datatype retrieve rows (#2163)

Author: doctrino

CHANGELOG.md

@@ -21,6 +21,11 @@ Changes are grouped as follows
 ### Fixed
 - Fixes type annotations for Functions API. Adds new `FunctionHandle` type for annotating function handles.
 
+## [7.75.2] - 2025-06-05
+### Fixed
+- The `client.raw.rows.retrieve_dataframe` method now has a new parameter `infer_dtypes` that allows
+  you to not infer the data types of column types in the returning dataframe.
+
 ## [7.75.1] - 2025-05-15
 ### Fixed
 - Fixes missing `instance_id` field in `Document` class returned from `client.documents.list()` and `client.documents.search()`.

cognite/client/_api/raw.py

@@ -717,6 +717,7 @@ def retrieve_dataframe(
         limit: int | None = DEFAULT_LIMIT_READ,
         partitions: int | None = None,
         last_updated_time_in_index: bool = False,
+        infer_dtypes: bool = True,
     ) -> pd.DataFrame:
         """`Retrieve rows in a table as a pandas dataframe. <https://developer.cognite.com/api#tag/Raw/operation/getRows>`_
 
@@ -734,6 +735,7 @@ def retrieve_dataframe(
                 (will be capped at this value). To prevent unexpected problems and maximize read throughput, check out
                 `concurrency limits in the API documentation. <https://developer.cognite.com/api#tag/Raw/#section/Request-and-concurrency-limits>`_
             last_updated_time_in_index (bool): Use a MultiIndex with row keys and last_updated_time as index.
+            infer_dtypes (bool): If True, pandas will try to infer dtypes of the columns. Defaults to True.
 
         Returns:
             pd.DataFrame: The requested rows in a pandas dataframe.
@@ -756,7 +758,7 @@ def retrieve_dataframe(
         else:
             idx = [r.key for r in rows]
         cols = [r.columns for r in rows]
-        return pd.DataFrame(cols, index=idx)
+        return pd.DataFrame(cols, index=idx, dtype=object if not infer_dtypes else None)
 
     def _get_parallel_cursors(
         self,

cognite/client/_version.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
 
-__version__ = "7.75.1"
+__version__ = "7.75.2"
 
 __api_subversion__ = "20230101"

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "cognite-sdk"
-version = "7.75.1"
+version = "7.75.2"
 
 description = "Cognite Python SDK"
 readme = "README.md"

tests/tests_unit/test_api/test_raw.py

@@ -3,10 +3,11 @@
 
 import pytest
 
+from cognite.client import CogniteClient
 from cognite.client._api.raw import RawRowsAPI
 from cognite.client.data_classes import Database, DatabaseList, Row, RowList, RowWrite, RowWriteList, Table, TableList
 from cognite.client.exceptions import CogniteAPIError
-from tests.utils import jsgz_load
+from tests.utils import assert_all_value_types_equal, jsgz_load
 
 
 @pytest.fixture
@@ -81,6 +82,29 @@ def mock_retrieve_raw_rows_response_two_rows(rsps, cognite_client):
     yield rsps
 
 
+@pytest.fixture
+def integer_rows_response() -> dict:
+    return {
+        "items": [
+            {"key": "row1", "columns": {"c1": 1, "c2": 4.0}, "lastUpdatedTime": 0},
+            {"key": "row2", "columns": {"c1": 2, "c2": 3}, "lastUpdatedTime": 1},
+            {"key": "row3", "columns": {"c1": 3, "c2": 1}, "lastUpdatedTime": 2},
+            {"key": "row4", "columns": {"c1": None, "c2": 0.1}, "lastUpdatedTime": 3},
+        ]
+    }
+
+
+@pytest.fixture
+def mock_retrieve_integer_rows(rsps, integer_rows_response: dict, cognite_client: CogniteClient):
+    rsps.add(
+        rsps.GET,
+        cognite_client.raw._get_base_url_with_base_path() + "/raw/dbs/db1/tables/table1/rows",
+        status=200,
+        json=integer_rows_response,
+    )
+    yield rsps
+
+
 @pytest.fixture
 def mock_retrieve_raw_rows_response_one_row(rsps, cognite_client):
     response_body = {"items": [{"key": "row1", "columns": {"c1": 1, "c2": "2"}, "lastUpdatedTime": 0}]}
@@ -378,6 +402,19 @@ def test_retrieve_dataframe_two_rows(self, cognite_client, mock_retrieve_raw_row
             pd.Timestamp(1, unit="ms"),
         ]
 
+    def test_retrieve_dataframe_integers(
+        self,
+        cognite_client: CogniteClient,
+        mock_retrieve_integer_rows,
+        integer_rows_response: dict,
+    ) -> None:
+        result = cognite_client.raw.rows.retrieve_dataframe(db_name="db1", table_name="table1", infer_dtypes=False)
+
+        actual = result.to_dict(orient="index")
+        expected = {row["key"]: row["columns"] for row in integer_rows_response["items"]}
+        assert_all_value_types_equal(actual, expected)
+        assert actual == expected
+
 
 @pytest.mark.parametrize("raw_cls", (Row, RowWrite))
 def test_raw_row__direct_column_access(raw_cls):

tests/utils.py

@@ -630,3 +630,12 @@ def get_or_raise(obj: T | None) -> T:
     if obj is None:
         raise ValueError("Object is None")
     return obj
+
+
+def assert_all_value_types_equal(d1: dict, d2: dict) -> None:
+    assert d1.keys() == d2.keys()
+    for k in d1.keys():
+        v1, v2 = d1[k], d2[k]
+        assert type(v1) is type(v2), f"Type mismatch for key '{k}': {type(v1)} != {type(v2)}"
+        if isinstance(v1, dict) and isinstance(v2, dict):
+            assert_all_value_types_equal(v1, v2)