Update timestamp dtypes in MetricSnapshotDataFrame and MetricRangeDataFrame (#234)

* Update timestamps to be dtype DateTime instead of float by default.

- Updates MetricRangeDataFrame to have pd.DatetimeIndex instead of
  pd.Float64Index.
- Updates MetricSnapshotDataFrame to have the column timestamp as
  dtype datetime64[ns] instead of float64.
- Adds test to ensure correct dtype conversion.

* Rename timestamp (unix time) to datetime parameter in dataframes.

Author: Karanraj Chauhan

prometheus_api_client/metric_range_df.py

@@ -1,5 +1,5 @@
 """A pandas.DataFrame subclass for Prometheus range vector responses."""
-from pandas import DataFrame
+from pandas import DataFrame, to_datetime
 from pandas._typing import Axes, Dtype
 from typing import Optional, Sequence
 
@@ -24,6 +24,9 @@ class MetricRangeDataFrame(DataFrame):
                  default to list of labels + "timestamp" + "value" if not provided.
     :param dtype: (dtype) default None. Data type to force. Only a single dtype is allowed. If None, infer.
     :param copy: (bool) default False. Copy data from inputs. Only affects DataFrame / 2d ndarray input.
+    :param ts_as_datetime: (bool) default True. Convert the timestamps returned by prometheus
+                 from float64 (unix time) to pandas datetime objects. This results in a pd.DatetimeIndex
+                 as the dtype of the index of the returned dataframe, instead of pd.Float64Index
 
     Example Usage:
       .. code-block:: python
@@ -52,6 +55,7 @@ def __init__(
         columns: Optional[Axes] = None,
         dtype: Optional[Dtype] = None,
         copy: bool = False,
+        ts_as_datetime: bool = True,
     ):
         """Functions as a constructor for MetricRangeDataFrame class."""
         if data is not None:
@@ -74,4 +78,8 @@ def __init__(
             data=row_data, index=index, columns=columns, dtype=dtype, copy=copy
         )
 
+        # convert to DateTime type instead of Float64
+        if ts_as_datetime:
+            self["timestamp"] = to_datetime(self["timestamp"], unit="s")
+
         self.set_index(["timestamp"], inplace=True)

prometheus_api_client/metric_snapshot_df.py

@@ -1,5 +1,5 @@
 """A pandas.DataFrame subclass for Prometheus query response."""
-from pandas import DataFrame
+from pandas import DataFrame, to_datetime
 from pandas._typing import Axes, Dtype
 from typing import Optional, Sequence
 
@@ -18,14 +18,18 @@ class MetricSnapshotDataFrame(DataFrame):
 
     :param data: (list|json) A single metric (json with keys "metric" and "values"/"value")
         or list of such metrics received from Prometheus as a response to query
-    :param ts_values_keep: (str) If several timestamp-value tuples are returned for a given
-        metric + label config, determine which one to keep. Currently only supports 'first', 'last'.
     :param index: (pandas.Index|array-like) Index to use for resulting dataframe. Will default to
                  pandas.RangeIndex if no indexing information part of input data and no index provided.
     :param columns: (pandas.Index|array-like) Column labels to use for resulting dataframe. Will
                  default to list of labels + "timestamp" + "value" if not provided.
     :param dtype: (dtype) default None. Data type to force. Only a single dtype is allowed. If None, infer.
     :param copy: (bool) default False. Copy data from inputs. Only affects DataFrame / 2d ndarray input.
+    :param ts_values_keep: (str) If several timestamp-value tuples are returned for a given
+                 metric + label config, determine which one to keep. Currently only supports 'first', 'last'.
+    :param ts_as_datetime: (bool) default True. Convert the timestamps returned by prometheus
+                 from float64 (unix time) to pandas datetime objects. This results in the timestamp column
+                 of the returned dataframe to be of dtype datetime64[ns] instead float64
+
 
     Example Usage:
       .. code-block:: python
@@ -49,11 +53,12 @@ class MetricSnapshotDataFrame(DataFrame):
     def __init__(
         self,
         data=None,
-        ts_values_keep: str = "last",
         index: Optional[Axes] = None,
         columns: Optional[Axes] = None,
         dtype: Optional[Dtype] = None,
         copy: bool = False,
+        ts_values_keep: str = "last",
+        ts_as_datetime: bool = True,
     ):
         """Functions as a constructor for MetricSnapshotDataFrame class."""
         if data is not None:
@@ -79,6 +84,10 @@ def __init__(
             data=data, index=index, columns=columns, dtype=dtype, copy=copy
         )
 
+        # convert to DateTime type instead of Float64
+        if ts_as_datetime:
+            self["timestamp"] = to_datetime(self["timestamp"], unit="s")
+
     @staticmethod
     def _get_nth_ts_value_pair(i: dict, n: int):
         val = i["values"][n] if "values" in i else i["value"]

tests/test_metric_range_df.py

@@ -1,5 +1,6 @@
 """Unit Tests for MetricRangeDataFrame."""
 import unittest
+import pandas as pd
 from prometheus_api_client import MetricRangeDataFrame
 from .test_with_metrics import TestWithMetrics
 
@@ -34,8 +35,9 @@ def test_init_timestamps(self):
         """Test if dataframe contains the correct timestamp indices."""
         # check that the timestamp indices in each series are the same
         for curr_metric_list in self.raw_metrics_list:
+            curr_df = MetricRangeDataFrame(curr_metric_list, ts_as_datetime=False)
             self.assertEqual(
-                set(MetricRangeDataFrame(curr_metric_list).index.values),
+                set(curr_df.index.values),
                 set([v[0] for s in curr_metric_list for v in s["values"]]),
             )
 
@@ -50,6 +52,23 @@ def test_init_columns(self):
                 "incorrect dataframe columns",
             )
 
+    def test_timestamp_dtype_conversion(self):
+        """Test if the timestamp in the dataframe initialized has correct dtype."""
+        for curr_metric_list in self.raw_metrics_list:
+            # index (prometheus metric timestamps) should be datetime type by default
+            curr_df = MetricRangeDataFrame(curr_metric_list)
+            self.assertTrue(
+                isinstance(curr_df.index, pd.DatetimeIndex),
+                "incorrect dtype for timestamp column (expected datetime dtype)",
+            )
+
+            # if explicitly set to false, conversion to dt shouldnt take place
+            curr_df = MetricRangeDataFrame(curr_metric_list, ts_as_datetime=False)
+            self.assertFalse(
+                isinstance(curr_df.index, pd.DatetimeIndex),
+                "incorrect dtype for timestamp column (expected non-datetime dtype)",
+            )
+
     def test_init_single_metric(self):
         """
         Test if dataframe initialized is of correct shape.

tests/test_metric_snapshot_df.py

@@ -3,6 +3,7 @@
 import json
 import os
 from prometheus_api_client import MetricSnapshotDataFrame
+from pandas.api.types import is_datetime64_any_dtype as is_dtype_datetime
 
 
 class TestMetricSnapshotDataFrame(unittest.TestCase):  # noqa D101
@@ -52,6 +53,23 @@ def test_init_columns(self):
                 "incorrect dataframe columns",
             )
 
+    def test_timestamp_dtype_conversion(self):
+        """Test if the timestamp in the dataframe initialized has correct dtype."""
+        for curr_metric_list in self.raw_metrics_list:
+            # timestamp column should be datetime type by default
+            curr_df = MetricSnapshotDataFrame(curr_metric_list,)
+            self.assertTrue(
+                is_dtype_datetime(curr_df["timestamp"]),
+                "incorrect dtype for timestamp column (expected datetime dtype)",
+            )
+
+            # if explicitly set to false, conversion to dt shouldnt take place
+            curr_df = MetricSnapshotDataFrame(curr_metric_list, ts_as_datetime=False,)
+            self.assertFalse(
+                is_dtype_datetime(curr_df["timestamp"]),
+                "incorrect dtype for timestamp column (expected non-datetime dtype)",
+            )
+
     def test_init_single_metric(self):
         """
         Test if dataframe initialized is of correct shape.