Merge branch 'gusutabopb-dataframeclient_chunked_queries'

sebito91 · sebito91 · commit e07dafccf486 · 2017-08-15T19:36:07.000-05:00
diff --git a/influxdb/_dataframe_client.py b/influxdb/_dataframe_client.py
@@ -7,6 +7,7 @@
 from __future__ import unicode_literals
 
 import math
+from collections import defaultdict
 
 import pandas as pd
 
@@ -137,23 +138,54 @@ def write_points(self,
 
         return True
 
-    def query(self, query, chunked=False, database=None):
-        """Quering data into a DataFrame.
-
-        :param chunked: [Optional, default=False] True if the data shall be
-            retrieved in chunks, False otherwise.
+    def query(self,
+              query,
+              params=None,
+              epoch=None,
+              expected_response_code=200,
+              database=None,
+              raise_errors=True,
+              chunked=False,
+              chunk_size=0,
+              dropna=True):
+        """
+        Quering data into a DataFrame.
+
+        :param query: the actual query string
+        :param params: additional parameters for the request, defaults to {}
+        :param epoch: response timestamps to be in epoch format either 'h',
+            'm', 's', 'ms', 'u', or 'ns',defaults to `None` which is
+            RFC3339 UTC format with nanosecond precision
+        :param expected_response_code: the expected status code of response,
+            defaults to 200
+        :param database: database to query, defaults to None
+        :param raise_errors: Whether or not to raise exceptions when InfluxDB
+            returns errors, defaults to True
+        :param chunked: Enable to use chunked responses from InfluxDB.
+            With ``chunked`` enabled, one ResultSet is returned per chunk
+            containing all results within that chunk
+        :param chunk_size: Size of each chunk to tell InfluxDB to use.
+        :param dropna: drop columns where all values are missing
+        :returns: the queried data
+        :rtype: :class:`~.ResultSet`
         """
-        results = super(DataFrameClient, self).query(query, database=database)
+        query_args = dict(params=params,
+                          epoch=epoch,
+                          expected_response_code=expected_response_code,
+                          raise_errors=raise_errors,
+                          chunked=chunked,
+                          chunk_size=chunk_size)
+        results = super(DataFrameClient, self).query(query, **query_args)
         if query.strip().upper().startswith("SELECT"):
             if len(results) > 0:
-                return self._to_dataframe(results)
+                return self._to_dataframe(results, dropna)
             else:
                 return {}
         else:
             return results
 
-    def _to_dataframe(self, rs):
-        result = {}
+    def _to_dataframe(self, rs, dropna=True):
+        result = defaultdict(list)
         if isinstance(rs, list):
             return map(self._to_dataframe, rs)
 
@@ -168,6 +200,11 @@ def _to_dataframe(self, rs):
             df.set_index('time', inplace=True)
             df.index = df.index.tz_localize('UTC')
             df.index.name = None
+            result[key].append(df)
+        for key, data in result.items():
+            df = pd.concat(data).sort_index()
+            if dropna:
+                df.dropna(how='all', axis=1, inplace=True)
             result[key] = df
 
         return result
