feat(state time series): add retrieve_arrays support for raw state datapoints

cognite/client/_api/datapoint_tasks.py

@@ -730,14 +730,23 @@ def _create_empty_result(self) -> Datapoints | DatapointsArray:
             else:
                 return Datapoints(**self.ts_info, timestamp=[], value=[], **status_cols)
 
-        if self.is_state_dps:
-            raise NotImplementedError(
-                "State datapoints are not yet supported when using `retrieve_arrays(...)`. "
-                "Please use `retrieve(...)` instead"
-            )
-
         if self.query.include_status:
             status_cols.update(status_code=np.array([], dtype=np.int32), status_symbol=np.array([], dtype=np.object_))
+
+        if self.is_state_dps:
+            # Numpy has no notion of a nullable int32 array. Since bad status datapoints may be missing their numeric
+            # value, we always use float64 (NaN for missing) whenever the caller includes bad datapoints, regardless
+            # of whether any actually are missing so that the dtype stays consistent:
+            numeric_dtype = np.int32 if self.query.ignore_bad_datapoints else np.float64
+            return DatapointsArray._load_from_arrays(
+                {
+                    **self.ts_info,
+                    "timestamp": np.array([], dtype=np.int64),
+                    "numeric_states": np.array([], dtype=numeric_dtype),
+                    "string_states": np.array([], dtype=np.object_),
+                    **status_cols,
+                }
+            )
         return DatapointsArray._load_from_arrays(
             {
                 **self.ts_info,
@@ -765,11 +774,33 @@ def _get_result(self) -> Datapoints | DatapointsArray:
                 )
             if not self.query.ignore_bad_datapoints:
                 status_columns["null_timestamps"] = self.null_timestamps
+
+            data_columns: dict[str, Any]
+            if not self.is_state_dps:
+                data_columns = {"value": create_array_from_dps_container(self.dps_data)}
+            else:
+                # Numeric dtype for state dps depends on `ignore_bad_datapoints` setting (nullable or not), so we always
+                # warn the user about this. TODO: Maybe revisit this decision? Most users just call to_pandas() and then
+                # they get pandas extension dtype Int32, which is nullable...
+                numeric_dtype = np.int32 if self.query.ignore_bad_datapoints else np.float64
+                if not self.query.ignore_bad_datapoints:
+                    warnings.warn(
+                        "The setting `ignore_bad_datapoints=False` means a state time series' numeric state "
+                        "values can be missing. Since numpy has no notion of a nullable int32 array, the "
+                        "'numeric_states' array is upcast to float64, which can perfectly represent any int32 "
+                        "value and uses NaN for the missing ones.",
+                        UserWarning,
+                    )
+                num_list, str_list = create_state_lists_from_dps_container(self.dps_data)
+                data_columns = {
+                    "numeric_states": np.array(num_list, dtype=numeric_dtype),
+                    "string_states": np.array(str_list, dtype=np.object_),
+                }
             return DatapointsArray._load_from_arrays(
                 {
                     **self.ts_info,
                     "timestamp": create_array_from_dps_container(self.ts_data),
-                    "value": create_array_from_dps_container(self.dps_data),
+                    **data_columns,
                     **status_columns,
                 }
             )
@@ -828,25 +859,30 @@ def _unpack_and_store(self, idx: tuple[float, ...], dps: DatapointsRaw) -> None:
             self._unpack_and_store_basic(idx, dps)
 
     def _unpack_and_store_numpy(self, idx: tuple[float, ...], dps: DatapointsRaw) -> None:
-        if self.is_state_dps:
-            raise NotImplementedError(
-                "Retrieving raw state datapoints using `retrieve_arrays(...)` is not yet supported. "
-                "Please use `retrieve(...)` instead."
-            )
         self.ts_data[idx].append(DpsUnpackFns.extract_timestamps_numpy(dps))
 
-        assert self.raw_dtype_numpy is not None
-        if self.query.ignore_bad_datapoints:
-            self.dps_data[idx].append(DpsUnpackFns.extract_raw_dps_numpy(dps, self.raw_dtype_numpy))
+        if self.is_state_dps:
+            # Performance note: We don't materialize numpy arrays per-batch here like we do for "normal raw" datapoints
+            # to keep things simple (allows easy reuse of 'self.dps_data'). This gives a slightly higher-than-necessary
+            # memory footprint. Thus we do one final array conversion in `_get_result` instead.
+            dps = cast(StateDatapoints, dps)
+            if self.query.ignore_bad_datapoints:
+                self.dps_data[idx].append(DpsUnpackFns.extract_raw_num_and_str_state_dps(dps))
+            else:
+                self.dps_data[idx].append(DpsUnpackFns.extract_nullable_raw_num_and_str_state_dps(dps))
         else:
-            # After this step, missing values (represented with None) will become NaNs and thus become
-            # indistinguishable from any NaNs that was returned! We need to store these timestamps in a property
-            # to allow our users to inspect them - but maybe even more important, allow the SDK to accurately
-            # use the DatapointsArray to replicate datapoints (exactly).
-            arr, missing_idxs = DpsUnpackFns.extract_nullable_raw_dps_numpy(dps, self.raw_dtype_numpy)
-            self.dps_data[idx].append(arr)
-            if missing_idxs:
-                self.null_timestamps.update(self.ts_data[idx][-1][missing_idxs].tolist())
+            assert self.raw_dtype_numpy is not None
+            if self.query.ignore_bad_datapoints:
+                self.dps_data[idx].append(DpsUnpackFns.extract_raw_dps_numpy(dps, self.raw_dtype_numpy))
+            else:
+                # After this step, missing values (represented with None) will become NaNs and thus become
+                # indistinguishable from any NaNs that was returned! We need to store these timestamps in a property
+                # to allow our users to inspect them - but maybe even more important, allow the SDK to accurately
+                # use the DatapointsArray to replicate datapoints (exactly).
+                arr, missing_idxs = DpsUnpackFns.extract_nullable_raw_dps_numpy(dps, self.raw_dtype_numpy)
+                self.dps_data[idx].append(arr)
+                if missing_idxs:
+                    self.null_timestamps.update(self.ts_data[idx][-1][missing_idxs].tolist())
 
         if self.query.include_status:
             self.status_code[idx].append(DpsUnpackFns.extract_status_code_numpy(dps))

cognite/client/data_classes/datapoints.py

@@ -68,6 +68,7 @@
 
     NumpyDatetime64NSArray: TypeAlias = npt.NDArray[np.datetime64]
     NumpyUInt32Array: TypeAlias = npt.NDArray[np.uint32]
+    NumpyInt32Array: TypeAlias = npt.NDArray[np.int32]
     NumpyInt64Array: TypeAlias = npt.NDArray[np.int64]
     NumpyFloat64Array: TypeAlias = npt.NDArray[np.float64]
     NumpyObjArray: TypeAlias = npt.NDArray[np.object_]
@@ -758,6 +759,8 @@ def __init__(
         granularity: str | None = None,
         timestamp: NumpyDatetime64NSArray | None = None,
         value: NumpyFloat64Array | NumpyObjArray | None = None,
+        numeric_states: NumpyInt32Array | NumpyFloat64Array | None = None,
+        string_states: NumpyObjArray | None = None,
         average: NumpyFloat64Array | None = None,
         max: NumpyFloat64Array | None = None,
         max_datapoint: NumpyObjArray | None = None,
@@ -794,6 +797,8 @@ def __init__(
             timestamp if timestamp is not None else np.array([], dtype="datetime64[ns]")
         )
         self.value = value
+        self.numeric_states = numeric_states
+        self.string_states = string_states
         self.average = average
         self.max = max
         self.max_datapoint = max_datapoint

tests/tests_integration/test_api/test_datapoints.py

@@ -783,6 +783,32 @@ def test_insert_state_dps_then_modify_state_set(
         assert dps_with_bad.string_states == [None, *exp_string_states]
         assert dps_with_bad.numeric_states == [None, *exp_numeric_states]
 
+        # ------------------- #
+        # Now let's repeat the exact same using retrieve_arrays():
+        arr_no_bad = cognite_client.time_series.data.retrieve_arrays(
+            instance_id=node_id, ignore_bad_datapoints=True, limit=50
+        )
+        assert arr_no_bad is not None
+        assert arr_no_bad.numeric_states is not None
+        assert arr_no_bad.string_states is not None
+        assert arr_no_bad.numeric_states.dtype == np.int32
+        np.testing.assert_array_equal(arr_no_bad.numeric_states, np.array(exp_numeric_states, dtype=np.int32))
+        np.testing.assert_array_equal(arr_no_bad.string_states, np.array(exp_string_states, dtype=object))
+
+        # Here we also ensure that we get the upcast warning for numeric_states (possibly containing NaNs):
+        with pytest.warns(UserWarning, match="upcast to float64"):
+            arr_with_bad = cognite_client.time_series.data.retrieve_arrays(
+                instance_id=node_id, ignore_bad_datapoints=False, limit=50
+            )
+
+        assert arr_with_bad is not None
+        assert arr_with_bad.numeric_states is not None
+        assert arr_with_bad.string_states is not None
+        assert arr_with_bad.numeric_states.dtype == np.float64
+        np.testing.assert_array_equal(arr_with_bad.numeric_states, np.array([np.nan, *exp_numeric_states]))
+        np.testing.assert_array_equal(arr_with_bad.string_states, np.array([None, *exp_string_states], dtype=object))
+
+        # ------------------- #
         # Ensure writing state dp with 20 or "twenty" now fails:
         with pytest.raises(CogniteAPIError) as e:
             cognite_client.time_series.data.insert_states(
@@ -897,6 +923,40 @@ def test_retrieve_state_datapoints_empty(
             assert df.empty
             assert list(df.columns) == [(ts_id, "numeric"), (ts_id, "string")]
 
+    @pytest.mark.parametrize(
+        "ignore_bad_datapoints, exp_numeric_dtype",
+        [(True, np.int32), (False, np.float64)],
+    )
+    def test_retrieve_arrays_state_datapoints_empty(
+        self,
+        cognite_client: CogniteClient,
+        empty_state_ts: NodeApplyResult,
+        ignore_bad_datapoints: bool,
+        exp_numeric_dtype: type,
+    ) -> None:
+        ts_id = empty_state_ts.as_id()
+        arr_lst = cognite_client.time_series.data.retrieve_arrays(
+            instance_id=[
+                DatapointsQuery(instance_id=ts_id, limit=0),
+                DatapointsQuery(instance_id=ts_id, limit=1),
+            ],
+            start="1h-ahead",
+            end="2h-ahead",
+            ignore_bad_datapoints=ignore_bad_datapoints,
+        )
+        for arr in arr_lst:
+            assert isinstance(arr, DatapointsArray)
+            assert arr.type == "state"
+            assert len(arr) == 0
+            assert arr.value is None
+            assert arr.numeric_states is not None
+            assert len(arr.numeric_states) == 0
+            assert arr.numeric_states.dtype == exp_numeric_dtype
+            assert arr.string_states is not None
+            assert len(arr.string_states) == 0
+
+            # TODO: awaiting implementation: `df = dps.to_pandas() & assert df.empty`
+
     @pytest.mark.parametrize(
         "retrieve_call",
         [
@@ -912,10 +972,6 @@ def test_retrieve_state_datapoints_empty(
                 ),
                 id="aggregate states retrieve_arrays",
             ),
-            pytest.param(
-                lambda client, ts_id: client.time_series.data.retrieve_arrays(instance_id=ts_id, limit=1),
-                id="raw states retrieve_arrays",
-            ),
             pytest.param(
                 lambda client, ts_id: client.time_series.data.retrieve_dataframe(instance_id=ts_id, limit=1),
                 id="raw states retrieve_dataframe",

tests/utils.py

@@ -530,6 +530,8 @@ def create_value(self, type_: Any, var_name: str | None = None) -> Any:
                 return np.array([self._random.random() for _ in range(3)], dtype=np.float64)
             elif type_ == NDArray[np.uint32]:
                 return np.array([self._random.randint(1, 100) for _ in range(3)], dtype=np.uint32)
+            elif type_ == NDArray[np.int32]:
+                return np.array([self._random.randint(1, 100) for _ in range(3)], dtype=np.int32)
             elif type_ == NDArray[np.int64]:
                 return np.array([self._random.randint(1, 100) for _ in range(3)], dtype=np.int64)
             elif type_ == NDArray[np.datetime64]:
@@ -658,6 +660,7 @@ def _type_checking(cls) -> dict[str, type]:
             "AsyncCogniteClient": AsyncCogniteClient,
             "NumpyDatetime64NSArray": npt.NDArray[np.datetime64],
             "NumpyUInt32Array": npt.NDArray[np.uint32],
+            "NumpyInt32Array": npt.NDArray[np.int32],
             "NumpyInt64Array": npt.NDArray[np.int64],
             "NumpyFloat64Array": npt.NDArray[np.float64],
             "NumpyObjArray": npt.NDArray[np.object_],