fix: data types for querying into DataFrame (influxdata#380)

bednar · web-flow · commit afee531fd1dc · 2021-12-14T14:31:37.000+01:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -4,6 +4,7 @@
 1. [#375](https://github.com/influxdata/influxdb-client-python/pull/375): Construct `InfluxDBError` without HTTP response 
 1. [#378](https://github.com/influxdata/influxdb-client-python/pull/378): Correct serialization DataFrame with nan values [DataFrame]
 1. [#384](https://github.com/influxdata/influxdb-client-python/pull/384): Timeout can be specified as a `float`
+1. [#380](https://github.com/influxdata/influxdb-client-python/pull/380): Correct data types for querying [DataFrame]
 
 ### CI
 1. [#370](https://github.com/influxdata/influxdb-client-python/pull/370): Add Python 3.10 to CI builds
diff --git a/influxdb_client/client/flux_csv_parser.py b/influxdb_client/client/flux_csv_parser.py
@@ -9,9 +9,8 @@
 
 from urllib3 import HTTPResponse
 
-from influxdb_client.client.util.date_utils import get_date_helper
 from influxdb_client.client.flux_table import FluxTable, FluxColumn, FluxRecord
-
+from influxdb_client.client.util.date_utils import get_date_helper
 
 ANNOTATION_DEFAULT = "#default"
 ANNOTATION_GROUP = "#group"
@@ -187,7 +186,7 @@ def _prepare_data_frame(self):
             _temp_df = _temp_df.set_index(self._data_frame_index)
 
         # Append data
-        return self._data_frame.append(_temp_df)
+        return self._data_frame.astype(_temp_df.dtypes).append(_temp_df)
 
     def parse_record(self, table_index, table, csv):
         """Parse one record."""
diff --git a/tests/test_FluxCSVParser.py b/tests/test_FluxCSVParser.py
@@ -235,6 +235,29 @@ def test_pandas_lot_of_columns(self):
         _dataFrames = list(parser.generator())
         self.assertEqual(1, _dataFrames.__len__())
 
+    def test_pandas_column_datatype(self):
+        data = "#datatype,string,long,dateTime:RFC3339,dateTime:RFC3339,string,string,string,string,long,unsignedLong,string,boolean,double\n" \
+               "#group,false,false,true,true,true,true,true,true,false,false,false,false,false\n" \
+               "#default,_result,,,,,,,,,,,,\n" \
+               ",result,table,_start,_stop,_field,_measurement,host,region,value1,value2,value3,value4,value5\n" \
+               ",,0,1977-09-21T00:12:43.145224192Z,2018-07-16T11:21:02.547596934Z,free,mem,A,west,121,11,test,true,6.56\n"
+        parser = self._parse(data=data, serialization_mode=FluxSerializationMode.dataFrame)
+        df = list(parser.generator())[0]
+        self.assertEqual(13, df.dtypes.__len__())
+        self.assertEqual('object', df.dtypes['result'].name)
+        self.assertEqual('int64', df.dtypes['table'].name)
+        self.assertIn('datetime64[ns,', df.dtypes['_start'].name)
+        self.assertIn('datetime64[ns,', df.dtypes['_stop'].name)
+        self.assertEqual('object', df.dtypes['_field'].name)
+        self.assertEqual('object', df.dtypes['_measurement'].name)
+        self.assertEqual('object', df.dtypes['host'].name)
+        self.assertEqual('object', df.dtypes['region'].name)
+        self.assertEqual('int64', df.dtypes['value1'].name)
+        self.assertEqual('int64', df.dtypes['value2'].name)
+        self.assertEqual('object', df.dtypes['value3'].name)
+        self.assertEqual('bool', df.dtypes['value4'].name)
+        self.assertEqual('float64', df.dtypes['value5'].name)
+
     @staticmethod
     def _parse_to_tables(data: str, serialization_mode=FluxSerializationMode.tables):
         _parser = FluxCsvParserTest._parse(data, serialization_mode)
