feat: convert StreamedResultSet to Pandas Dataframe

google/cloud/spanner_v1/_pandas_helpers.py

@@ -0,0 +1,76 @@
+import logging
+import six
+
+class pandas_df():
+
+    def pd_dataframe(self):
+        """Returns the response in a pandas dataframe of the StreamedResultSet object"""
+        try :
+            from pandas import DataFrame
+        except ImportError:
+            logging.error("Pandas module not found. It is needed for converting query result to Dataframe.\n Try running 'pip3 install pandas'")
+
+        code_to_spanner_dtype_dict = {
+                                        1 : 'BOOL',
+                                        2 : 'INT64',
+                                        3 : 'FLOAT64',
+                                        4 : 'TIMESTAMP',
+                                        5 : 'DATE',
+                                        6 : 'STRING',
+                                        7 : 'BYTES',
+                                        8 : 'ARRAY',
+                                        10 : 'NUMERIC'
+                                    }
+        response = six.next(self._response_iterator)
+        if self._metadata is None:  # first response
+            metadata = self._metadata = response.metadata
+
+        #Creating dictionary to store column name maping of spanner to pandas dataframe
+        columns_dict={}
+        try :
+            for item in metadata.row_type.fields :
+                columns_dict[item.name]=code_to_spanner_dtype_dict[item.type_.code]
+        except :
+            logging.warning("Not able to create spanner to pandas fields mapping")
+
+        #Creating list of columns to be mapped with the data
+        column_list=[k for k,v in columns_dict.items()]
+
+        #Creating list of data values to be converted to dataframe
+        values = list(response.values)
+        if self._pending_chunk is not None:
+            values[0] = self._merge_chunk(values[0])
+        if response.chunked_value:
+            self._pending_chunk = values.pop()
+        self._merge_values(values)
+
+        width = len(column_list)
+
+        # list to store each row as a sub-list
+        data=[] 
+        while len(values)/width > 0 :
+            data.append(values[:width])
+            values=values[width:]
+
+        #Creating dataframe using column headers and list of data rows
+        df = DataFrame(data,columns=column_list)
+
+        #Mapping dictionary to map every spanner datatype to a pandas compatible datatype
+        mapping_dict={
+                    'INT64':'int64',
+                    'STRING':'object',
+                    'BOOL':'bool',
+                    'BYTES':'object', 
+                    'ARRAY':'object',
+                    'DATE':'datetime64[ns, UTC]',
+                    'FLOAT64':'float64', 
+                    'NUMERIC':'object', 
+                    'TIMESTAMP':'datetime64[ns, UTC]'
+                    }
+        for k,v in columns_dict.items() :
+            try:
+                df[k]= df[k].astype(mapping_dict[v])
+            except KeyError:
+                print("Spanner Datatype not present in datatype mapping dictionary")
+
+        return df

google/cloud/spanner_v1/streamed.py

@@ -17,9 +17,11 @@
 from google.cloud import exceptions
 from google.cloud.spanner_v1 import TypeCode
 import six
+import logging
 
 # pylint: disable=ungrouped-imports
 from google.cloud.spanner_v1._helpers import _parse_value
+from google.cloud.spanner_v1._pandas_helpers import pandas_df
 
 # pylint: enable=ungrouped-imports
 
@@ -143,6 +145,11 @@ def __iter__(self):
             while iter_rows:
                 yield iter_rows.pop(0)
 
+    def to_dataframe(self):
+        """Returns the response in a pandas dataframe of the StreamedResultSet object"""
+        df = pandas_df.pd_dataframe(self)
+        return df
+
     def one(self):
         """Return exactly one result, or raise an exception.
 
@@ -302,4 +309,4 @@ def _merge_struct(lhs, rhs, type_):
 def _merge_by_type(lhs, rhs, type_):
     """Helper for '_merge_chunk'."""
     merger = _MERGE_BY_TYPE[type_.code]
-    return merger(lhs, rhs, type_)
+    return merger(lhs, rhs, type_)