diff --git a/cognite/client/data_classes/datapoints.py b/cognite/client/data_classes/datapoints.py index a6cfeb1674..d778fa874d 100644 --- a/cognite/client/data_classes/datapoints.py +++ b/cognite/client/data_classes/datapoints.py @@ -1056,6 +1056,8 @@ def to_pandas( # type: ignore [override] include_granularity_name=include_granularity_name, include_status=include_status, include_unit=include_unit, + include_numeric_states=False, # not implemented + include_string_states=False, # not implemented ) @@ -1255,6 +1257,8 @@ def to_pandas( # type: ignore [override] include_granularity_name: bool = False, include_unit: bool = True, include_status: bool = True, + include_numeric_states: bool = True, + include_string_states: bool = True, ) -> pandas.DataFrame: """Convert the datapoints into a pandas DataFrame. @@ -1264,6 +1268,8 @@ def to_pandas( # type: ignore [override] include_unit (bool): Include the unit_external_id in the dataframe columns, if present (separate MultiIndex level) include_status (bool): Include status code and status symbol as separate columns, if available. Also adds the status info as a separate level in the columns (MultiIndex). + include_numeric_states (bool): For state time series, include the numeric states in the dataframe columns. Defaults to True. + include_string_states (bool): For state time series, include the string states in the dataframe columns. Defaults to True. Returns: pandas.DataFrame: The dataframe. @@ -1276,6 +1282,8 @@ def to_pandas( # type: ignore [override] include_granularity_name=include_granularity_name, include_status=include_status, include_unit=include_unit, + include_numeric_states=include_numeric_states, + include_string_states=include_string_states, ) @classmethod @@ -1589,6 +1597,8 @@ def to_pandas( # type: ignore [override] include_granularity_name=include_granularity_name, include_status=include_status, include_unit=include_unit, + include_numeric_states=False, # not implemented yet + include_string_states=False, # not implemented yet ) def dump(self, camel_case: bool = True, convert_timestamps: bool = False) -> list[dict[str, Any]]: @@ -1654,6 +1664,8 @@ def to_pandas( # type: ignore [override] include_granularity_name: bool = False, include_unit: bool = True, include_status: bool = True, + include_numeric_states: bool = True, + include_string_states: bool = True, ) -> pandas.DataFrame: """Convert the datapoints list into a pandas DataFrame. @@ -1663,6 +1675,8 @@ def to_pandas( # type: ignore [override] include_unit (bool): Include the unit_external_id in the dataframe columns, if present (separate MultiIndex level) include_status (bool): Include status code and status symbol as separate columns, if available. Also adds the status info as a separate level in the columns (MultiIndex). + include_numeric_states (bool): For state time series, include the numeric states in the dataframe columns. Defaults to True. + include_string_states (bool): For state time series, include the string states in the dataframe columns. Defaults to True. Returns: pandas.DataFrame: The datapoints list as a pandas DataFrame. @@ -1673,6 +1687,8 @@ def to_pandas( # type: ignore [override] include_granularity_name=include_granularity_name, include_status=include_status, include_unit=include_unit, + include_numeric_states=include_numeric_states, + include_string_states=include_string_states, ) diff --git a/cognite/client/utils/_pandas_helpers.py b/cognite/client/utils/_pandas_helpers.py index 26fb7e24d4..c103265f50 100644 --- a/cognite/client/utils/_pandas_helpers.py +++ b/cognite/client/utils/_pandas_helpers.py @@ -75,6 +75,10 @@ def is_pandas_v2_or_lower() -> bool: return pandas_major_version() < 3 +def pandas_string_array_dtype() -> Literal["object", "str"]: + return "object" if is_pandas_v2_or_lower() else "str" + + @cache def timestamp_dtype_unit() -> Literal["ns", "ms"]: """The datetime64 resolution to use for all timestamp columns/indices produced by the SDK, see @@ -136,6 +140,8 @@ def concat_dps_dataframe_list( include_granularity_name: bool, include_status: bool, include_unit: bool, + include_numeric_states: bool, + include_string_states: bool, ) -> pd.DataFrame: import pandas as pd @@ -153,7 +159,15 @@ def concat_dps_dataframe_list( ) # Since we use a MultiIndex for the dataframe columns, these do not join nicely in pd.concat, so we need # to do that manually ourselves after combining. - columns_lst = [_extract_column_info_from_dps_for_dataframe(dps, include_status=include_status) for dps in dps_lst] + columns_lst = [ + _extract_column_info_from_dps_for_dataframe( + dps, + include_status=include_status, + include_numeric_states=include_numeric_states, + include_string_states=include_string_states, + ) + for dps in dps_lst + ] counter = itertools.count() # Ensure unique column names initially dfs = [ pd.DataFrame( @@ -222,10 +236,16 @@ def convert_timestamp_columns_to_datetime(df: pd.DataFrame) -> pd.DataFrame: def concat_dataframes_with_nullable_int_cols(dfs: Sequence[pd.DataFrame]) -> pd.DataFrame: import pandas as pd + # Columns already using a pandas nullable integer extension dtype (e.g. the Int32 dtype used + # for numeric state datapoints) survive pd.concat's outer-join just fine (missing rows are + # filled with pd.NA, dtype is preserved). Only plain numpy int/uint columns need help here, + # since those silently upcast to float64 if the join introduces missing rows for that column: + # TODO: status_code is still a plain numpy uint32 column, so it always lands here and gets + # blanket-cast to Int64 below. We should switch it to the nullable UInt32. int_cols = [ i for i, dtype in enumerate(itertools.chain.from_iterable(df.dtypes for df in dfs)) - if issubclass(dtype.type, Integral) + if not pd.api.types.is_extension_array_dtype(dtype) and issubclass(dtype.type, Integral) ] # TODO: Performance optimization possible: The more unique each df.index is to the rest of the dfs, the # slower `pd.concat` scales. A manual "union(df.index for df in dfs)" + column insertion is faster for large @@ -237,6 +257,7 @@ def concat_dataframes_with_nullable_int_cols(dfs: Sequence[pd.DataFrame]) -> pd. if pandas_major_version() >= 2: df.isetitem(int_cols, df.iloc[:, int_cols].astype("Int64")) else: + # TODO: We specify pandas >= 2.1, so we can remove this branch. # As of pandas >=1.5.0, <2, converting float cols (that used to be int) to nullable int using iloc raises FutureWarning, # but the suggested code change (to use `frame.isetitem(...)`) results in the wrong dtype (object). # See Github Issue: https://github.com/pandas-dev/pandas/issues/49922 @@ -269,9 +290,16 @@ def convert_dps_to_dataframe( include_granularity_name: bool, include_status: bool, include_unit: bool, + include_numeric_states: bool, + include_string_states: bool, ) -> pd.DataFrame: pd = local_import("pandas") - columns = _extract_column_info_from_dps_for_dataframe(dps, include_status=include_status) + columns = _extract_column_info_from_dps_for_dataframe( + dps, + include_status=include_status, + include_numeric_states=include_numeric_states, + include_string_states=include_string_states, + ) df = pd.DataFrame( # We initially use integer indexing to allow duplicate column names: {i: col.as_array() for i, col in enumerate(columns)}, @@ -291,33 +319,52 @@ def convert_dps_to_dataframe( class _DpsColumnInfo: """ Used when converting Datapoints/DatapointsArray/DatapointsList/DatapointsArrayList to pandas DataFrame to help - avoid the madness of how many columns we should end up with based on status codes/symbols, number of aggregates etc. + avoid the absolute madness of how many columns we should end up: + - the raw datapoints (easy!) + - the number of classic aggregates (10+, but just 1 value per granularity interval) + - status codes & symbols (2 extra columns, if requested) + - state datapoints (2 columns, numeric and string states) + + ...and not yet implemented, but I'm scared: + - state aggregate datapoints ("arbitrary" number of states, leading to possibly 200+ columns (1 value per state per gran. interval)) - A single Datapoints/DatapointsArray can result in 10+ columns from aggregates, and 1 or 3 columns from raw datapoints, - with or without the 2 extra status info columns. + Thus, a single Datapoints/DatapointsArray can result in anything between 1 to 300 columns. Yey. """ column_id: NodeId | str | int - data: list[float] | list[str] | list[int] | NumpyUInt32Array | NumpyInt64Array | NumpyFloat64Array | NumpyObjArray + data: ( + list[float] + | list[str] + | list[str | None] + | list[int] + | NumpyUInt32Array + | NumpyInt64Array + | NumpyFloat64Array + | NumpyObjArray + ) is_string: bool | None = None is_array: bool = False aggregate: str | None = None granularity: str | None = None unit_xid: str | None = None status_info: Literal["code", "symbol"] | None = None + state_type: Literal["numeric", "string"] | None = None def as_multi_index_tuple(self, include_aggregate: bool, include_granularity: bool, include_unit: bool) -> tuple: return ( self.column_id, + self.state_type, self.status_info, # since these split to separate cols, they are already filtered out if not wanted self.aggregate if include_aggregate else None, self.granularity if include_granularity else None, self.unit_xid if include_unit else None, ) - def as_array(self) -> NumpyObjArray | NumpyFloat64Array | NumpyInt64Array | NumpyUInt32Array: + def as_array( + self, + ) -> NumpyObjArray | NumpyFloat64Array | NumpyInt64Array | NumpyUInt32Array | pd.arrays.IntegerArray: if self.is_array: - return self.data # type: ignore [return-value] + return self.data elif self.aggregate is None: return self._convert_to_array_for_raw_dps() @@ -326,9 +373,15 @@ def as_array(self) -> NumpyObjArray | NumpyFloat64Array | NumpyInt64Array | Nump def _convert_to_array_for_raw_dps( self, - ) -> npt.NDArray[np.object_] | npt.NDArray[np.float64] | npt.NDArray[np.uint32]: + ) -> npt.NDArray[np.object_] | npt.NDArray[np.float64] | npt.NDArray[np.uint32] | pd.arrays.IntegerArray: import numpy as np + if self.state_type == "numeric": + # Numeric states are guaranteed to be valid 32-bit ints, but may contain missing values due to "bad status", + # so we use the pandas extension dtype which is nullable: + pd = local_import("pandas") + return pd.array(self.data, dtype="Int32") + match self.is_string, self.status_info: case True, None: return np.array(self.data, dtype=np.object_) @@ -361,6 +414,45 @@ def _convert_to_array_for_agg_dps( return np.array(self.data, dtype=np.float64) +def _extract_raw_states_column_info( + dps: Datapoints, + identifier: NodeId | str | int, + include_status: bool, + include_numeric_states: bool, + include_string_states: bool, +) -> list[_DpsColumnInfo]: + columns = [] + if include_numeric_states: + assert dps.numeric_states is not None + columns.append( + _DpsColumnInfo( + identifier, + data=dps.numeric_states, + is_string=False, + is_array=False, + state_type="numeric", + ) + ) + if include_string_states: + assert dps.string_states is not None + columns.append( + _DpsColumnInfo( + identifier, + data=dps.string_states, + is_string=True, + is_array=False, + state_type="string", + ) + ) + if include_status: + if dps.status_code is not None: + columns.append(_DpsColumnInfo(identifier, data=dps.status_code, is_array=False, status_info="code")) + if dps.status_symbol is not None: + columns.append(_DpsColumnInfo(identifier, data=dps.status_symbol, is_array=False, status_info="symbol")) + + return columns + + def _extract_raw_column_info( dps: Datapoints | DatapointsArray, identifier: NodeId | str | int, @@ -406,15 +498,33 @@ def _extract_aggregate_column_info_from_dps( def _extract_column_info_from_dps_for_dataframe( - dps: Datapoints | DatapointsArray, include_status: bool + dps: Datapoints | DatapointsArray, include_status: bool, include_numeric_states: bool, include_string_states: bool ) -> list[_DpsColumnInfo]: - from cognite.client.data_classes import DatapointsArray + from cognite.client.data_classes import Datapoints, DatapointsArray identifier = _resolve_ts_identifier_as_df_column_name(dps) is_array = isinstance(dps, DatapointsArray) - if dps.value is not None: + # TODO: State raw vs aggregate dps must be routed differently (when we have support for the latter...) + if dps.type == "state": + if is_array: + # Unreachable state in the SDK, but users may instantiate manually, so we need to handle it: + raise NotImplementedError( + "State datapoints stored as DatapointsArray are not supported yet for conversion to pandas DataFrame" + ) + assert isinstance(dps, Datapoints) # mypy doesn't understand the is-array-raise-check above... + if dps.numeric_states is None or dps.string_states is None: + # ...also unreachable, but same gotcha as above: + raise NotImplementedError( + "State aggregate datapoints are not yet supported for conversion to pandas DataFrame" + ) + else: + return _extract_raw_states_column_info( + dps, identifier, include_status, include_numeric_states, include_string_states + ) + elif dps.value is not None: return _extract_raw_column_info(dps, identifier, is_array, include_status) - return _extract_aggregate_column_info_from_dps(dps, identifier, is_array) + else: + return _extract_aggregate_column_info_from_dps(dps, identifier, is_array) def _create_multi_index_from_columns( @@ -434,7 +544,7 @@ def _create_multi_index_from_columns( ) for col in columns ], - columns=["identifier", "status", "aggregate", "granularity", "unit"], + columns=["identifier", "state", "status", "aggregate", "granularity", "unit"], ) # Key operation is to drop all-nan columns, which in the multi-index translates to dropping # the corresponding levels: diff --git a/tests/tests_integration/test_api/test_datapoints.py b/tests/tests_integration/test_api/test_datapoints.py index bf38f98aeb..c2a6bc1897 100644 --- a/tests/tests_integration/test_api/test_datapoints.py +++ b/tests/tests_integration/test_api/test_datapoints.py @@ -893,6 +893,10 @@ def test_retrieve_state_datapoints_empty( assert dps.numeric_states == [] assert dps.string_states == [] + df = dps.to_pandas() + assert df.empty + assert list(df.columns) == [(ts_id, "numeric"), (ts_id, "string")] + @pytest.mark.parametrize( "retrieve_call", [ diff --git a/tests/tests_unit/test_data_classes/test_datapoints.py b/tests/tests_unit/test_data_classes/test_datapoints.py index 96b1f739a8..0dd220badb 100644 --- a/tests/tests_unit/test_data_classes/test_datapoints.py +++ b/tests/tests_unit/test_data_classes/test_datapoints.py @@ -10,8 +10,8 @@ from cognite.client.data_classes import Datapoint, DatapointsArray, StateDatapointsInsert, StateDatapointWrite from cognite.client.data_classes._base import CogniteResourceList from cognite.client.data_classes.data_modeling.ids import NodeId -from cognite.client.data_classes.datapoints import DatapointsArrayList, DatapointsList -from tests.utils import PANDAS_TS_UNIT +from cognite.client.data_classes.datapoints import Datapoints, DatapointsArrayList, DatapointsList +from tests.utils import PANDAS_STR_DTYPE, PANDAS_TS_UNIT class TestDatapoint: @@ -120,6 +120,152 @@ def test_identifier_priority(self, dps_lst_cls: type[CogniteResourceList]) -> No pd.testing.assert_frame_equal(df, exp_df) +@pytest.mark.dsl +class TestStateDatapointsToPandas: + @pytest.fixture + def node_id(self) -> NodeId: + return NodeId("ss", "xx") + + @pytest.fixture + def state_dps(self, node_id: NodeId) -> Datapoints: + return Datapoints( + id=123, + instance_id=node_id, + is_string=False, + is_step=True, + type="state", + timestamp=[1000, 2000, 3000, 4000], + # For bad datapoints, even numeric can be missing (None): + numeric_states=[0, 1, 0, None], # type: ignore [list-item] + string_states=["off", "on", None, None], + ) + + def test_default_includes_both_state_columns(self, state_dps: Datapoints, node_id: NodeId) -> None: + import pandas as pd + + df = state_dps.to_pandas() + + assert list(df.columns) == [(node_id, "numeric"), (node_id, "string")] + assert df.columns.names == ["identifier", "state"] + + numeric_values = df[node_id, "numeric"].tolist() + assert numeric_values[:-1] == [0, 1, 0] + assert numeric_values[3] is pd.NA + assert df[node_id, "numeric"].dtype == "Int32" + + # Missing string states are represented as None on pandas v2 (object dtype) and as + # NaN on pandas v3 (its new native 'str' dtype), see PANDAS_STR_DTYPE: + string_values = df[node_id, "string"].tolist() + assert string_values[:2] == ["off", "on"] + assert all(pd.isna(v) for v in string_values[2:]) + assert df[node_id, "string"].dtype == PANDAS_STR_DTYPE + + def test_exclude_numeric_states(self, state_dps: Datapoints, node_id: NodeId) -> None: + import pandas as pd + + df = state_dps.to_pandas(include_numeric_states=False) + + assert list(df.columns) == [(node_id, "string")] + string_values = df[node_id, "string"].tolist() + assert string_values[:2] == ["off", "on"] + assert all(pd.isna(v) for v in string_values[2:]) + + def test_exclude_string_states(self, state_dps: Datapoints, node_id: NodeId) -> None: + import pandas as pd + + df = state_dps.to_pandas(include_string_states=False) + + assert list(df.columns) == [(node_id, "numeric")] + numeric_values = df[node_id, "numeric"].tolist() + assert numeric_values[:3] == [0, 1, 0] + assert numeric_values[3] is pd.NA + + def test_exclude_both_states_without_status_gives_empty_dataframe( + self, state_dps: Datapoints, node_id: NodeId + ) -> None: + df = state_dps.to_pandas(include_numeric_states=False, include_string_states=False, include_status=False) + + assert df.shape == (4, 0) + + def test_status_columns_included_alongside_state_columns(self) -> None: + dps = Datapoints( + id=123, + is_string=False, + is_step=False, + type="state", + timestamp=[1000, 2000], + numeric_states=[0, 1], + string_states=["off", "on"], + status_code=[0, 2147483648], + status_symbol=["Good", "Bad"], + ) + df = dps.to_pandas() + + assert set(df.columns) == { + (123, "numeric", ""), + (123, "string", ""), + (123, "", "code"), + (123, "", "symbol"), + } + assert df[123, "", "code"].tolist() == [0, 2147483648] + assert df[123, "", "symbol"].tolist() == ["Good", "Bad"] + + def test_datapoints_array_with_state_type_raises(self) -> None: + import numpy as np + + arr = DatapointsArray( + id=123, + is_string=False, + is_step=False, + type="state", + timestamp=np.array([1000], dtype="datetime64[ns]"), + ) + arr_lst = DatapointsArrayList([arr]) + + for dps in [arr, arr_lst]: + with pytest.raises(NotImplementedError, match="DatapointsArray are not supported"): + dps.to_pandas() # type: ignore [attr-defined] + + def test_mixed_state_and_numeric_dps_list_to_pandas(self, state_dps: Datapoints, node_id: NodeId) -> None: + numeric_dps = Datapoints( + id=456, + is_string=False, + is_step=False, + type="numeric", + timestamp=[1000, 2000, 3000, 4000], + value=[1.5, 2.5, 3.5, 4.5], + ) + df = DatapointsList([state_dps, numeric_dps]).to_pandas() + + assert set(df.columns) == {(node_id, "numeric"), (node_id, "string"), (456, "")} + assert df[(456, "")].tolist() == [1.5, 2.5, 3.5, 4.5] + + def test_datapoints_list_to_pandas(self, state_dps: Datapoints, node_id: NodeId) -> None: + import numpy as np + import pandas as pd + + other_state_dps = Datapoints( + id=456, + is_string=False, + is_step=True, + type="state", + # Ensure some timestamps align and others don't, to test the outer join behavior: + timestamp=[1000, 2000, 3500, 5500], + numeric_states=[10, 11, 10, 11], + string_states=["idle", "running", "idle", "running"], + ) + df = DatapointsList([state_dps, other_state_dps]).to_pandas() + + assert set(df.columns) == {(node_id, "numeric"), (node_id, "string"), (456, "numeric"), (456, "string")} + assert df[456, "numeric"].tolist() == [10, 11, pd.NA, 10, pd.NA, 11] + np.testing.assert_array_equal( # easy way to make nans compare equal + df[456, "string"].tolist(), + ["idle", "running", math.nan, "idle", math.nan, "running"], + ) + assert df[node_id, "numeric"].dtype == "Int32" + assert df[456, "numeric"].dtype == "Int32" + + class TestStateDatapointWrite: @pytest.mark.parametrize( "kwargs, expected", diff --git a/tests/utils.py b/tests/utils.py index c42112f4e3..90884755a2 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -76,20 +76,28 @@ from cognite.client.testing import AsyncCogniteClientMock from cognite.client.utils import _json_extended as _json from cognite.client.utils._concurrency import CRUDConcurrency -from cognite.client.utils._pandas_helpers import timestamp_dtype_unit +from cognite.client.utils._pandas_helpers import pandas_string_array_dtype, timestamp_dtype_unit from cognite.client.utils._text import random_string REPO_ROOT = Path(__file__).resolve().parent.parent -# The resolution the SDK uses for pandas timestamp columns/indices depends on the installed pandas -# major version, see `timestamp_dtype_unit` for the reasoning. Tests asserting exact dtypes should -# build their expectations relative to this instead of hardcoding "ms" (or "ns"): try: + # The resolution the SDK uses for pandas timestamp columns/indices depends on the installed pandas + # major version, see `timestamp_dtype_unit` for the reasoning. Tests asserting exact dtypes should + # build their expectations relative to this instead of hardcoding "ms" (or "ns"): PANDAS_TS_UNIT: Literal["ns", "ms"] = timestamp_dtype_unit() + + # We respect the new pandas v3 default of inferring a native 'str' dtype for object columns of + # strings (we could force the old 'object' dtype via 'future.infer_string'). Thus the SDK's own + # df string columns (e.g. state datapoints' string representations) end up with a dtype that depends + # on the installed pandas major version. + PANDAS_STR_DTYPE: Literal["object", "str"] = pandas_string_array_dtype() except ModuleNotFoundError: # When we test with '--test-deps-only-core', pandas is not installed. For simplicity of imports, we still - # define PANDAS_TS_UNIT here: + # define these constants here: PANDAS_TS_UNIT = "ms" + PANDAS_STR_DTYPE = "str" + T_Type = TypeVar("T_Type", bound=type)