Add option to choose dtypes by column in to_dataframe. (#7126)

* Add option to choose dtypes by column in to_dataframe.

This allows pandas users to select different sized floats for
performance at the expense of accuracy. With pandas 0.24, it will also
allow pandas users to use the new pandas.Int64Dtype() for nullable
integer columns.

* Adjust deps for testing. Blacken.

Author: tswast

google/cloud/bigquery_storage_v1beta1/reader.py

@@ -14,6 +14,7 @@
 
 from __future__ import absolute_import
 
+import collections
 import itertools
 import json
 
@@ -155,11 +156,11 @@ def rows(self, read_session):
         if fastavro is None:
             raise ImportError(_FASTAVRO_REQUIRED)
 
-        avro_schema = _avro_schema(read_session)
+        avro_schema, _ = _avro_schema(read_session)
         blocks = (_avro_rows(block, avro_schema) for block in self)
         return itertools.chain.from_iterable(blocks)
 
-    def to_dataframe(self, read_session):
+    def to_dataframe(self, read_session, dtypes=None):
         """Create a :class:`pandas.DataFrame` of all rows in the stream.
 
         This method requires the pandas libary to create a data frame and the
@@ -176,6 +177,13 @@ def to_dataframe(self, read_session):
                 The read session associated with this read rows stream. This
                 contains the schema, which is required to parse the data
                 blocks.
+            dtypes ( \
+                Map[str, Union[str, pandas.Series.dtype]] \
+            ):
+                Optional. A dictionary of column names pandas ``dtype``s. The
+                provided ``dtype`` is used when constructing the series for
+                the column specified. Otherwise, the default pandas behavior
+                is used.
 
         Returns:
             pandas.DataFrame:
@@ -186,14 +194,29 @@ def to_dataframe(self, read_session):
         if pandas is None:
             raise ImportError("pandas is required to create a DataFrame")
 
-        avro_schema = _avro_schema(read_session)
+        if dtypes is None:
+            dtypes = {}
+
+        avro_schema, column_names = _avro_schema(read_session)
         frames = []
         for block in self:
-            dataframe = pandas.DataFrame(list(_avro_rows(block, avro_schema)))
+            dataframe = _to_dataframe_with_dtypes(
+                _avro_rows(block, avro_schema), column_names, dtypes
+            )
             frames.append(dataframe)
         return pandas.concat(frames)
 
 
+def _to_dataframe_with_dtypes(rows, column_names, dtypes):
+    columns = collections.defaultdict(list)
+    for row in rows:
+        for column in row:
+            columns[column].append(row[column])
+    for column in dtypes:
+        columns[column] = pandas.Series(columns[column], dtype=dtypes[column])
+    return pandas.DataFrame(columns, columns=column_names)
+
+
 def _avro_schema(read_session):
     """Extract and parse Avro schema from a read session.
 
@@ -206,10 +229,13 @@ def _avro_schema(read_session):
             blocks.
 
     Returns:
-        A parsed Avro schema, using :func:`fastavro.schema.parse_schema`.
+        Tuple[fastavro.schema, Tuple[str]]:
+            A parsed Avro schema, using :func:`fastavro.schema.parse_schema`
+            and the column names for a read session.
     """
     json_schema = json.loads(read_session.avro_schema.schema)
-    return fastavro.parse_schema(json_schema)
+    column_names = tuple((field["name"] for field in json_schema["fields"]))
+    return fastavro.parse_schema(json_schema), column_names
 
 
 def _avro_rows(block, avro_schema):

noxfile.py

@@ -43,7 +43,6 @@ def default(session):
     session.run(
         'py.test',
         '--quiet',
-        '--cov=google.cloud.bigquery_storage',
         '--cov=google.cloud.bigquery_storage_v1beta1',
         '--cov=tests.unit',
         '--cov-append',

setup.py

@@ -21,7 +21,7 @@
 
 name = 'google-cloud-bigquery-storage'
 description = 'BigQuery Storage API API client library'
-version = '0.1.1'
+version = '0.2.0dev1'
 release_status = 'Development Status :: 3 - Alpha'
 dependencies = [
     'google-api-core[grpc] >= 1.6.0, < 2.0.0dev',

tests/system/test_system.py

@@ -17,6 +17,7 @@
 
 import os
 
+import numpy
 import pytest
 
 from google.cloud import bigquery_storage_v1beta1
@@ -78,11 +79,15 @@ def test_read_rows_to_dataframe(client, project_id):
         stream=session.streams[0]
     )
 
-    frame = client.read_rows(stream_pos).to_dataframe(session)
+    frame = client.read_rows(stream_pos).to_dataframe(
+        session, dtypes={"latitude": numpy.float16}
+    )
 
     # Station ID is a required field (no nulls), so the datatype should always
     # be integer.
     assert frame.station_id.dtype.name == "int64"
+    assert frame.latitude.dtype.name == "float16"
+    assert frame.longitude.dtype.name == "float64"
     assert frame["name"].str.startswith("Central Park").any()
 
 

tests/unit/test_reader.py

@@ -55,6 +55,7 @@
     {"name": "time_col", "type": "time"},
     {"name": "ts_col", "type": "timestamp"},
 ]
+SCALAR_COLUMN_NAMES = [field["name"] for field in SCALAR_COLUMNS]
 SCALAR_BLOCKS = [
     [
         {
@@ -281,7 +282,9 @@ def test_to_dataframe_w_scalars(class_under_test):
     )
     got = reader.to_dataframe(read_session)
 
-    expected = pandas.DataFrame(list(itertools.chain.from_iterable(SCALAR_BLOCKS)))
+    expected = pandas.DataFrame(
+        list(itertools.chain.from_iterable(SCALAR_BLOCKS)), columns=SCALAR_COLUMN_NAMES
+    )
     # fastavro provides its own UTC definition, so
     # compare the timestamp columns separately.
     got_ts = got["ts_col"]
@@ -301,6 +304,39 @@ def test_to_dataframe_w_scalars(class_under_test):
     )
 
 
+def test_to_dataframe_w_dtypes(class_under_test):
+    # TODOTODOTODOTODO
+    avro_schema = _bq_to_avro_schema(
+        [
+            {"name": "bigfloat", "type": "float64"},
+            {"name": "lilfloat", "type": "float64"},
+        ]
+    )
+    read_session = _generate_read_session(avro_schema)
+    blocks = [
+        [{"bigfloat": 1.25, "lilfloat": 30.5}, {"bigfloat": 2.5, "lilfloat": 21.125}],
+        [{"bigfloat": 3.75, "lilfloat": 11.0}],
+    ]
+    avro_blocks = _bq_to_avro_blocks(blocks, avro_schema)
+
+    reader = class_under_test(
+        avro_blocks, mock_client, bigquery_storage_v1beta1.types.StreamPosition(), {}
+    )
+    got = reader.to_dataframe(read_session, dtypes={"lilfloat": "float16"})
+
+    expected = pandas.DataFrame(
+        {
+            "bigfloat": [1.25, 2.5, 3.75],
+            "lilfloat": pandas.Series([30.5, 21.125, 11.0], dtype="float16"),
+        },
+        columns=["bigfloat", "lilfloat"],
+    )
+    pandas.testing.assert_frame_equal(
+        got.reset_index(drop=True),  # reset_index to ignore row labels
+        expected.reset_index(drop=True),
+    )
+
+
 def test_copy_stream_position(mut):
     read_position = bigquery_storage_v1beta1.types.StreamPosition(
         stream={"name": "test"}, offset=41