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.

bigquery/google/cloud/bigquery/table.py

@@ -16,6 +16,7 @@
 
 from __future__ import absolute_import
 
+import collections
 import copy
 import datetime
 import json
@@ -1315,14 +1316,24 @@ def total_rows(self):
         """int: The total number of rows in the table."""
         return self._total_rows
 
-    def _to_dataframe_tabledata_list(self):
+    def _to_dataframe_dtypes(self, page, column_names, dtypes):
+        columns = collections.defaultdict(list)
+        for row in page:
+            for column in column_names:
+                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 _to_dataframe_tabledata_list(self, dtypes):
         """Use (slower, but free) tabledata.list to construct a DataFrame."""
-        column_headers = [field.name for field in self.schema]
-        # Use generator, rather than pulling the whole rowset into memory.
-        rows = (row.values() for row in iter(self))
-        return pandas.DataFrame(rows, columns=column_headers)
+        column_names = [field.name for field in self.schema]
+        frames = []
+        for page in iter(self.pages):
+            frames.append(self._to_dataframe_dtypes(page, column_names, dtypes))
+        return pandas.concat(frames)
 
-    def _to_dataframe_bqstorage(self, bqstorage_client):
+    def _to_dataframe_bqstorage(self, bqstorage_client, dtypes):
         """Use (faster, but billable) BQ Storage API to construct DataFrame."""
         import concurrent.futures
         from google.cloud import bigquery_storage_v1beta1
@@ -1360,7 +1371,7 @@ def _to_dataframe_bqstorage(self, bqstorage_client):
         def get_dataframe(stream):
             position = bigquery_storage_v1beta1.types.StreamPosition(stream=stream)
             rowstream = bqstorage_client.read_rows(position)
-            return rowstream.to_dataframe(session)
+            return rowstream.to_dataframe(session, dtypes=dtypes)
 
         with concurrent.futures.ThreadPoolExecutor() as pool:
             frames = pool.map(get_dataframe, session.streams)
@@ -1369,16 +1380,16 @@ def get_dataframe(stream):
         # the end using manually-parsed schema.
         return pandas.concat(frames)[columns]
 
-    def to_dataframe(self, bqstorage_client=None):
+    def to_dataframe(self, bqstorage_client=None, dtypes=None):
         """Create a pandas DataFrame from the query results.
 
         Args:
             bqstorage_client ( \
                 google.cloud.bigquery_storage_v1beta1.BigQueryStorageClient \
             ):
-                Optional. A BigQuery Storage API client. If supplied, use the
-                faster BigQuery Storage API to fetch rows from BigQuery. This
-                API is a billable API.
+                **Alpha Feature** Optional. A BigQuery Storage API client. If
+                supplied, use the faster BigQuery Storage API to fetch rows
+                from BigQuery. This API is a billable API.
 
                 This method requires the ``fastavro`` and
                 ``google-cloud-bigquery-storage`` libraries.
@@ -1389,6 +1400,13 @@ def to_dataframe(self, bqstorage_client=None):
                 **Caution**: There is a known issue reading small anonymous
                 query result tables with the BQ Storage API. Write your query
                 results to a destination table to work around this issue.
+            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:
@@ -1402,11 +1420,13 @@ def to_dataframe(self, bqstorage_client=None):
         """
         if pandas is None:
             raise ValueError(_NO_PANDAS_ERROR)
+        if dtypes is None:
+            dtypes = {}
 
         if bqstorage_client is not None:
-            return self._to_dataframe_bqstorage(bqstorage_client)
+            return self._to_dataframe_bqstorage(bqstorage_client, dtypes)
         else:
-            return self._to_dataframe_tabledata_list()
+            return self._to_dataframe_tabledata_list(dtypes)
 
 
 class _EmptyRowIterator(object):

bigquery/noxfile.py

@@ -22,6 +22,9 @@
 LOCAL_DEPS = (
     os.path.join('..', 'api_core[grpc]'),
     os.path.join('..', 'core'),
+    # TODO: Move bigquery_storage back to dev_install once dtypes feature is
+    #       released. Issue #7049
+    os.path.join('..', 'bigquery_storage[pandas,fastavro]'),
 )
 
 
@@ -40,9 +43,9 @@ def default(session):
 
     # Pyarrow does not support Python 3.7
     if session.python == '3.7':
-        dev_install = '.[bqstorage, pandas]'
+        dev_install = '.[pandas]'
     else:
-        dev_install = '.[bqstorage, pandas, pyarrow]'
+        dev_install = '.[pandas, pyarrow]'
     session.install('-e', dev_install)
 
     # IPython does not support Python 2 after version 5.x

bigquery/setup.py

@@ -34,7 +34,7 @@
     'google-resumable-media >= 0.3.1',
 ]
 extras = {
-    'bqstorage': 'google-cloud-bigquery-storage<=2.0.0dev',
+    'bqstorage': 'google-cloud-bigquery-storage >= 0.2.0dev1, <2.0.0dev',
     'pandas': 'pandas>=0.17.1',
     # Exclude PyArrow dependency from Windows Python 2.7.
     'pyarrow: platform_system != "Windows" or python_version >= "3.4"':

bigquery/tests/system.py

@@ -1733,13 +1733,22 @@ def test_nested_table_to_dataframe(self):
                     ),
                 ],
             ),
+            SF("bigfloat_col", "FLOAT", mode="NULLABLE"),
+            SF("smallfloat_col", "FLOAT", mode="NULLABLE"),
         ]
         record = {
             "nested_string": "another string value",
             "nested_repeated": [0, 1, 2],
             "nested_record": {"nested_nested_string": "some deep insight"},
         }
-        to_insert = [{"string_col": "Some value", "record_col": record}]
+        to_insert = [
+            {
+                "string_col": "Some value",
+                "record_col": record,
+                "bigfloat_col": 3.14,
+                "smallfloat_col": 2.72,
+            }
+        ]
         rows = [json.dumps(row) for row in to_insert]
         body = six.BytesIO("{}\n".format("\n".join(rows)).encode("ascii"))
         table_id = "test_table"
@@ -1753,11 +1762,13 @@ def test_nested_table_to_dataframe(self):
         # Load a table using a local JSON file from memory.
         Config.CLIENT.load_table_from_file(body, table, job_config=job_config).result()
 
-        df = Config.CLIENT.list_rows(table, selected_fields=schema).to_dataframe()
+        df = Config.CLIENT.list_rows(table, selected_fields=schema).to_dataframe(
+            dtypes={"smallfloat_col": "float16"}
+        )
 
         self.assertIsInstance(df, pandas.DataFrame)
         self.assertEqual(len(df), 1)  # verify the number of rows
-        exp_columns = ["string_col", "record_col"]
+        exp_columns = ["string_col", "record_col", "bigfloat_col", "smallfloat_col"]
         self.assertEqual(list(df), exp_columns)  # verify the column names
         row = df.iloc[0]
         # verify the row content
@@ -1769,6 +1780,9 @@ def test_nested_table_to_dataframe(self):
             row["record_col"]["nested_record"]["nested_nested_string"],
             "some deep insight",
         )
+        # verify dtypes
+        self.assertEqual(df.dtypes["bigfloat_col"].name, "float64")
+        self.assertEqual(df.dtypes["smallfloat_col"].name, "float16")
 
     def test_list_rows_empty_table(self):
         from google.cloud.bigquery.table import RowIterator

bigquery/tests/unit/test_table.py

@@ -1472,21 +1472,22 @@ def test_to_dataframe_column_dtypes(self):
             SchemaField("start_timestamp", "TIMESTAMP"),
             SchemaField("seconds", "INT64"),
             SchemaField("miles", "FLOAT64"),
+            SchemaField("km", "FLOAT64"),
             SchemaField("payment_type", "STRING"),
             SchemaField("complete", "BOOL"),
             SchemaField("date", "DATE"),
         ]
         row_data = [
-            ["1.4338368E9", "420", "1.1", "Cash", "true", "1999-12-01"],
-            ["1.3878117E9", "2580", "17.7", "Cash", "false", "1953-06-14"],
-            ["1.3855653E9", "2280", "4.4", "Credit", "true", "1981-11-04"],
+            ["1.4338368E9", "420", "1.1", "1.77", "Cash", "true", "1999-12-01"],
+            ["1.3878117E9", "2580", "17.7", "28.5", "Cash", "false", "1953-06-14"],
+            ["1.3855653E9", "2280", "4.4", "7.1", "Credit", "true", "1981-11-04"],
         ]
         rows = [{"f": [{"v": field} for field in row]} for row in row_data]
         path = "/foo"
         api_request = mock.Mock(return_value={"rows": rows})
         row_iterator = RowIterator(_mock_client(), api_request, path, schema)
 
-        df = row_iterator.to_dataframe()
+        df = row_iterator.to_dataframe(dtypes={"km": "float16"})
 
         self.assertIsInstance(df, pandas.DataFrame)
         self.assertEqual(len(df), 3)  # verify the number of rows
@@ -1496,6 +1497,7 @@ def test_to_dataframe_column_dtypes(self):
         self.assertEqual(df.start_timestamp.dtype.name, "datetime64[ns, UTC]")
         self.assertEqual(df.seconds.dtype.name, "int64")
         self.assertEqual(df.miles.dtype.name, "float64")
+        self.assertEqual(df.km.dtype.name, "float16")
         self.assertEqual(df.payment_type.dtype.name, "object")
         self.assertEqual(df.complete.dtype.name, "bool")
         self.assertEqual(df.date.dtype.name, "object")

bigquery_storage/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):

bigquery_storage/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',

bigquery_storage/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',

bigquery_storage/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()
 
 

bigquery_storage/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