ENH: add dtypes argument to read_gbq (#333)

tswast · web-flow · commit 46c579ac2187 · 2020-10-02T13:09:23.000-05:00
Use this argument to override the default ``dtype`` for a particular column in
the query results. For example, this can be used to select nullable integer
columns as the ``Int64`` nullable integer pandas extension type.

    df = gbq.read_gbq(
        "SELECT CAST(NULL AS INT64) AS null_integer",
        dtypes={"null_integer": "Int64"},
    )
diff --git a/docs/source/changelog.rst b/docs/source/changelog.rst
@@ -6,6 +6,19 @@ Changelog
 0.14.0 / TBD
 ------------
 
+- Add ``dtypes`` argument to ``read_gbq``. Use this argument to override the
+  default ``dtype`` for a particular column in the query results. For
+  example, this can be used to select nullable integer columns as the
+  ``Int64`` nullable integer pandas extension type. (:issue:`242`,
+  :issue:`332`)
+
+.. code-block:: python
+
+   df = gbq.read_gbq(
+       "SELECT CAST(NULL AS INT64) AS null_integer",
+       dtypes={"null_integer": "Int64"},
+   )
+
 Dependency updates
 ~~~~~~~~~~~~~~~~~~
 
@@ -15,7 +28,7 @@ Dependency updates
 Internal changes
 ~~~~~~~~~~~~~~~~
 
-- Update tests to run against for Python 3.8. (:issue:`331`)
+- Update tests to run against Python 3.8. (:issue:`331`)
 
 
 .. _changelog-0.13.3:
diff --git a/pandas_gbq/gbq.py b/pandas_gbq/gbq.py
@@ -526,19 +526,28 @@ def run_query(
                 )
             )
 
+        dtypes = kwargs.get("dtypes")
         return self._download_results(
             query_reply,
             max_results=max_results,
             progress_bar_type=progress_bar_type,
+            user_dtypes=dtypes,
         )
 
     def _download_results(
-        self, query_job, max_results=None, progress_bar_type=None
+        self,
+        query_job,
+        max_results=None,
+        progress_bar_type=None,
+        user_dtypes=None,
     ):
         # No results are desired, so don't bother downloading anything.
         if max_results == 0:
             return None
 
+        if user_dtypes is None:
+            user_dtypes = {}
+
         try:
             bqstorage_client = None
             if max_results is None:
@@ -555,9 +564,10 @@ def _download_results(
             )
 
             schema_fields = [field.to_api_repr() for field in rows_iter.schema]
-            nullsafe_dtypes = _bqschema_to_nullsafe_dtypes(schema_fields)
+            conversion_dtypes = _bqschema_to_nullsafe_dtypes(schema_fields)
+            conversion_dtypes.update(user_dtypes)
             df = rows_iter.to_dataframe(
-                dtypes=nullsafe_dtypes,
+                dtypes=conversion_dtypes,
                 bqstorage_client=bqstorage_client,
                 progress_bar_type=progress_bar_type,
             )
@@ -790,6 +800,7 @@ def read_gbq(
     verbose=None,
     private_key=None,
     progress_bar_type="tqdm",
+    dtypes=None,
 ):
     r"""Load data from Google BigQuery using google-cloud-python
 
@@ -910,6 +921,10 @@ def read_gbq(
         ``'tqdm_gui'``
             Use the :func:`tqdm.tqdm_gui` function to display a
             progress bar as a graphical dialog box.
+    dtypes : dict, optional
+        A dictionary of column names to pandas ``dtype``. The provided
+        ``dtype`` is used when constructing the series for the column
+        specified. Otherwise, a default ``dtype`` is used.
     verbose : None, deprecated
         Deprecated in Pandas-GBQ 0.4.0. Use the `logging module
         to adjust verbosity instead
@@ -965,6 +980,7 @@ def read_gbq(
         configuration=configuration,
         max_results=max_results,
         progress_bar_type=progress_bar_type,
+        dtypes=dtypes,
     )
 
     # Reindex the DataFrame on the provided column
diff --git a/tests/system/test_gbq.py b/tests/system/test_gbq.py
@@ -8,6 +8,11 @@
 import pandas.api.types
 import pandas.util.testing as tm
 from pandas import DataFrame, NaT
+
+try:
+    import pkg_resources  # noqa
+except ImportError:
+    raise ImportError("Could not import pkg_resources (setuptools).")
 import pytest
 import pytz
 
@@ -16,14 +21,14 @@
 
 
 TABLE_ID = "new_test"
+PANDAS_VERSION = pkg_resources.parse_version(pandas.__version__)
+NULLABLE_INT_PANDAS_VERSION = pkg_resources.parse_version("0.24.0")
+NULLABLE_INT_MESSAGE = (
+    "Require pandas 0.24+ in order to use nullable integer type."
+)
 
 
 def test_imports():
-    try:
-        import pkg_resources  # noqa
-    except ImportError:
-        raise ImportError("Could not import pkg_resources (setuptools).")
-
     gbq._test_google_api_imports()
 
 
@@ -87,62 +92,92 @@ def test_should_properly_handle_null_strings(self, project_id):
         tm.assert_frame_equal(df, DataFrame({"null_string": [None]}))
 
     def test_should_properly_handle_valid_integers(self, project_id):
-        query = "SELECT INTEGER(3) AS valid_integer"
+        query = "SELECT CAST(3 AS INT64) AS valid_integer"
         df = gbq.read_gbq(
             query,
             project_id=project_id,
             credentials=self.credentials,
-            dialect="legacy",
+            dialect="standard",
         )
         tm.assert_frame_equal(df, DataFrame({"valid_integer": [3]}))
 
     def test_should_properly_handle_nullable_integers(self, project_id):
+        if PANDAS_VERSION < NULLABLE_INT_PANDAS_VERSION:
+            pytest.skip(msg=NULLABLE_INT_MESSAGE)
+
         query = """SELECT * FROM
-                    (SELECT 1 AS nullable_integer),
-                    (SELECT NULL AS nullable_integer)"""
+                    UNNEST([1, NULL]) AS nullable_integer
+                """
         df = gbq.read_gbq(
             query,
             project_id=project_id,
             credentials=self.credentials,
-            dialect="legacy",
+            dialect="standard",
+            dtypes={"nullable_integer": "Int64"},
+        )
+        tm.assert_frame_equal(
+            df,
+            DataFrame(
+                {
+                    "nullable_integer": pandas.Series(
+                        [1, pandas.NA], dtype="Int64"
+                    )
+                }
+            ),
         )
-        tm.assert_frame_equal(df, DataFrame({"nullable_integer": [1, None]}))
 
     def test_should_properly_handle_valid_longs(self, project_id):
         query = "SELECT 1 << 62 AS valid_long"
         df = gbq.read_gbq(
             query,
             project_id=project_id,
             credentials=self.credentials,
-            dialect="legacy",
+            dialect="standard",
         )
         tm.assert_frame_equal(df, DataFrame({"valid_long": [1 << 62]}))
 
     def test_should_properly_handle_nullable_longs(self, project_id):
+        if PANDAS_VERSION < NULLABLE_INT_PANDAS_VERSION:
+            pytest.skip(msg=NULLABLE_INT_MESSAGE)
+
         query = """SELECT * FROM
-                    (SELECT 1 << 62 AS nullable_long),
-                    (SELECT NULL AS nullable_long)"""
+                    UNNEST([1 << 62, NULL]) AS nullable_long
+                """
         df = gbq.read_gbq(
             query,
             project_id=project_id,
             credentials=self.credentials,
-            dialect="legacy",
+            dialect="standard",
+            dtypes={"nullable_long": "Int64"},
         )
         tm.assert_frame_equal(
-            df, DataFrame({"nullable_long": [1 << 62, None]})
+            df,
+            DataFrame(
+                {
+                    "nullable_long": pandas.Series(
+                        [1 << 62, pandas.NA], dtype="Int64"
+                    )
+                }
+            ),
         )
 
     def test_should_properly_handle_null_integers(self, project_id):
-        query = "SELECT INTEGER(NULL) AS null_integer"
+        if PANDAS_VERSION < NULLABLE_INT_PANDAS_VERSION:
+            pytest.skip(msg=NULLABLE_INT_MESSAGE)
+
+        query = "SELECT CAST(NULL AS INT64) AS null_integer"
         df = gbq.read_gbq(
             query,
             project_id=project_id,
             credentials=self.credentials,
-            dialect="legacy",
+            dialect="standard",
+            dtypes={"null_integer": "Int64"},
         )
         tm.assert_frame_equal(
             df,
-            DataFrame({"null_integer": pandas.Series([None], dtype="object")}),
+            DataFrame(
+                {"null_integer": pandas.Series([pandas.NA], dtype="Int64")}
+            ),
         )
 
     def test_should_properly_handle_valid_floats(self, project_id):
diff --git a/tests/unit/test_gbq.py b/tests/unit/test_gbq.py
@@ -476,6 +476,27 @@ def test_load_does_not_modify_schema_arg(mock_bigquery_client):
     assert original_schema == original_schema_cp
 
 
+def test_read_gbq_passes_dtypes(
+    mock_bigquery_client, mock_service_account_credentials
+):
+    mock_service_account_credentials.project_id = "service_account_project_id"
+    df = gbq.read_gbq(
+        "SELECT 1 AS int_col",
+        dialect="standard",
+        credentials=mock_service_account_credentials,
+        dtypes={"int_col": "my-custom-dtype"},
+    )
+    assert df is not None
+
+    mock_list_rows = mock_bigquery_client.list_rows("dest", max_results=100)
+
+    mock_list_rows.to_dataframe.assert_called_once_with(
+        dtypes={"int_col": "my-custom-dtype"},
+        bqstorage_client=mock.ANY,
+        progress_bar_type=mock.ANY,
+    )
+
+
 def test_read_gbq_calls_tqdm(
     mock_bigquery_client, mock_service_account_credentials
 ):
