Merge pull request influxdata#123 from timtroendle/feature-multiple-time-series-to-dataframe

aviau · aviau · commit 3745d57e66b9 · 2015-03-17T23:04:39.000-04:00
Feature multiple time series to dataframe (Thanks @timtroendle!)
diff --git a/influxdb/dataframe_client.py b/influxdb/dataframe_client.py
@@ -77,7 +77,10 @@ def write_points_with_precision(self, data, time_precision='s'):
 
     def query(self, query, time_precision='s', chunked=False):
         """
-        Quering data into a DataFrame.
+        Quering data into DataFrames.
+
+        Returns a DataFrame for a single time series and a map for multiple
+        time series with the time series as value and its name as key.
 
         :param time_precision: [Optional, default 's'] Either 's', 'm', 'ms'
             or 'u'.
@@ -88,10 +91,14 @@ def query(self, query, time_precision='s', chunked=False):
         result = InfluxDBClient.query(self, query=query,
                                       time_precision=time_precision,
                                       chunked=chunked)
-        if len(result) > 0:
+        if len(result) == 0:
+            return result
+        elif len(result) == 1:
             return self._to_dataframe(result[0], time_precision)
         else:
-            return result
+            return {time_series['name']: self._to_dataframe(time_series,
+                                                            time_precision)
+                    for time_series in result}
 
     def _to_dataframe(self, json_result, time_precision):
         dataframe = pd.DataFrame(data=json_result['points'],
diff --git a/tests/influxdb/dataframe_client_test.py b/tests/influxdb/dataframe_client_test.py
@@ -203,6 +203,46 @@ def test_query_into_dataframe(self):
             result = cli.query('select column_one from foo;')
             assert_frame_equal(dataframe, result)
 
+    def test_query_multiple_time_series(self):
+        data = [
+            {
+                "name": "series1",
+                "columns": ["time", "mean", "min", "max", "stddev"],
+                "points": [[0, 323048, 323048, 323048, 0]]
+            },
+            {
+                "name": "series2",
+                "columns": ["time", "mean", "min", "max", "stddev"],
+                "points": [[0, -2.8233, -2.8503, -2.7832, 0.0173]]
+            },
+            {
+                "name": "series3",
+                "columns": ["time", "mean", "min", "max", "stddev"],
+                "points": [[0, -0.01220, -0.01220, -0.01220, 0]]
+            }
+        ]
+        dataframes = {
+            'series1': pd.DataFrame(data=[[323048, 323048, 323048, 0]],
+                                    index=pd.to_datetime([0], unit='s',
+                                                         utc=True),
+                                    columns=['mean', 'min', 'max', 'stddev']),
+            'series2': pd.DataFrame(data=[[-2.8233, -2.8503, -2.7832, 0.0173]],
+                                    index=pd.to_datetime([0], unit='s',
+                                                         utc=True),
+                                    columns=['mean', 'min', 'max', 'stddev']),
+            'series3': pd.DataFrame(data=[[-0.01220, -0.01220, -0.01220, 0]],
+                                    index=pd.to_datetime([0], unit='s',
+                                                         utc=True),
+                                    columns=['mean', 'min', 'max', 'stddev'])
+        }
+        with _mocked_session('get', 200, data):
+            cli = DataFrameClient('host', 8086, 'username', 'password', 'db')
+            result = cli.query("""select mean(value), min(value), max(value),
+                stddev(value) from series1, series2, series3""")
+            assert dataframes.keys() == result.keys()
+            for key in dataframes.keys():
+                assert_frame_equal(dataframes[key], result[key])
+
     def test_query_with_empty_result(self):
         with _mocked_session('get', 200, []):
             cli = DataFrameClient('host', 8086, 'username', 'password', 'db')
