feat: add supports for custom precision for index specified as number [DataFrame] (influxdata#335)

bednar · web-flow · commit 9b20924e7c7f · 2021-10-12T10:56:24.000+02:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,7 +1,8 @@
 ## 1.22.0 [unreleased]
 
 ### Features
-1. [#330](https://github.com/influxdata/influxdb-client-python/pull/330): Add supports for write structured data - `NamedTuple`, `Data Classes`
+1. [#330](https://github.com/influxdata/influxdb-client-python/pull/330): Add support for write structured data - `NamedTuple`, `Data Classes`
+1. [#335](https://github.com/influxdata/influxdb-client-python/pull/335): Add support for custom precision for index specified as number [DataFrame] 
 
 ### Documentation
 1. [#331](https://github.com/influxdata/influxdb-client-python/pull/331): Add [Migration Guide](MIGRATION_GUIDE.rst)
diff --git a/influxdb_client/client/write/dataframe_serializer.py b/influxdb_client/client/write/dataframe_serializer.py
@@ -101,7 +101,7 @@ def __init__(self, data_frame, point_settings, precision=DEFAULT_WRITE_PRECISION
             # Instead, it would probably be better to leave
             # out the timestamp unless a time column is explicitly
             # enabled.
-            data_frame.index = pd.to_datetime(data_frame.index)
+            data_frame.index = pd.to_datetime(data_frame.index, unit=precision)
 
         if data_frame.index.tzinfo is None:
             data_frame.index = data_frame.index.tz_localize('UTC')
diff --git a/influxdb_client/client/write_api.py b/influxdb_client/client/write_api.py
@@ -252,7 +252,7 @@ def write(self, bucket: str, org: str = None,
         :param WritePrecision write_precision: specifies the precision for the unix timestamps within
                                                the body line-protocol. The precision specified on a Point has precedes
                                                and is use for write.
-        :param record: Point, Line Protocol, Dictionary, NamedTuple,  Data Classes, Pandas DataFrame or
+        :param record: Point, Line Protocol, Dictionary, NamedTuple, Data Classes, Pandas DataFrame or
                        RxPY Observable to write
         :key data_frame_measurement_name: name of measurement for writing Pandas DataFrame - ``DataFrame``
         :key data_frame_tag_columns: list of DataFrame columns which are tags,
@@ -284,7 +284,26 @@ def write(self, bucket: str, org: str = None,
                 point = Point("h2o_feet").tag("location", "us-west").field("level", 125).time(1)
                 write_api.write("my-bucket", "my-org", point)
 
-        """
+        DataFrame:
+            The index of `Pandas DataFrame <https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.html>`_
+            is used as a ``timestamp`` for written data. The index should be `PeriodIndex <https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.PeriodIndex.html#pandas.PeriodIndex>`_
+            or its must be transformable to ``datetime`` by
+            `pandas.to_datetime <https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.to_datetime.html#pandas.to_datetime>`_.
+
+            If you would like to transform a column to ``PeriodIndex``, you can use something like:
+
+            .. code-block:: python
+
+                import pandas as pd
+
+                # DataFrame
+                data_frame = ...
+                # Set column as Index
+                data_frame.set_index('column_name', inplace=True)
+                # Transform index to PeriodIndex
+                data_frame.index = pd.to_datetime(data_frame.index, unit='s')
+
+        """  # noqa: E501
         org = get_org_query_param(org=org, client=self._influxdb_client)
 
         if self._point_settings.defaultTags and record is not None:
diff --git a/tests/test_WriteApiDataFrame.py b/tests/test_WriteApiDataFrame.py
@@ -354,6 +354,26 @@ def test_write_precision(self):
             self.assertEqual(1, len(points))
             self.assertEqual(f"h2o level=15i {precision[1]}", points[0])
 
+    def test_index_not_periodIndex_respect_write_precision(self):
+        from influxdb_client.extras import pd
+
+        precisions = [
+            (WritePrecision.NS, 1586044800000000000),
+            (WritePrecision.US, 1586044800000000),
+            (WritePrecision.MS, 1586044800000),
+            (WritePrecision.S, 1586044800),
+            (None, 1586044800000000000)
+        ]
+
+        for precision in precisions:
+            data_frame = pd.DataFrame([15], index=[precision[1]], columns=['level'])
+            points = data_frame_to_list_of_points(data_frame=data_frame,
+                                                  data_frame_measurement_name='h2o',
+                                                  point_settings=PointSettings(),
+                                                  precision=precision[0])
+            self.assertEqual(1, len(points))
+            self.assertEqual(f"h2o level=15i {precision[1]}", points[0])
+
 
 class DataSerializerChunksTest(unittest.TestCase):
     def test_chunks(self):
