ENH: Add max_results argument to read_gbq. (#286)

* ENH: Add max_results argument to read_gbq.

This argument allows you to set a maximum number of rows
in the resulting dataframe. Set to 0 to ignore any results.

* Call get_table on destination so listing works on older versions of the client library.

* Add max_results to changelog.

* Fix unit tests for extra call to get_table.

* Add unit tests for max_results.

* Adjust fail-under percent

Author: tswast

docs/source/changelog.rst

@@ -1,6 +1,16 @@
 Changelog
 =========
 
+.. _changelog-0.12.0:
+
+0.12.0 / TBD
+------------
+
+- Add ``max_results`` argument to :func:`~pandas_gbq.read_gbq()`. Use this
+  argument to limit the number of rows in the results DataFrame. Set
+  ``max_results`` to 0 to ignore query outputs, such as for DML or DDL
+  queries. (:issue:`102`)
+
 .. _changelog-0.11.0:
 
 0.11.0 / 2019-07-29

noxfile.py

@@ -48,7 +48,7 @@ def unit(session):
 @nox.session
 def cover(session, python=latest_python):
     session.install("coverage", "pytest-cov")
-    session.run("coverage", "report", "--show-missing", "--fail-under=40")
+    session.run("coverage", "report", "--show-missing", "--fail-under=74")
     session.run("coverage", "erase")
 
 

pandas_gbq/gbq.py

@@ -436,7 +436,7 @@ def process_http_error(ex):
 
         raise GenericGBQException("Reason: {0}".format(ex))
 
-    def run_query(self, query, **kwargs):
+    def run_query(self, query, max_results=None, **kwargs):
         from concurrent.futures import TimeoutError
         from google.auth.exceptions import RefreshError
         from google.cloud import bigquery
@@ -526,15 +526,33 @@ def run_query(self, query, **kwargs):
                 )
             )
 
+        return self._download_results(query_reply, max_results=max_results)
+
+    def _download_results(self, query_job, max_results=None):
+        # No results are desired, so don't bother downloading anything.
+        if max_results == 0:
+            return None
+
+        if max_results is None:
+            # Only use the BigQuery Storage API if the full result set is requested.
+            bqstorage_client = self.bqstorage_client
+        else:
+            bqstorage_client = None
+
         try:
-            rows_iter = query_reply.result()
+            query_job.result()
+            # Get the table schema, so that we can list rows.
+            destination = self.client.get_table(query_job.destination)
+            rows_iter = self.client.list_rows(
+                destination, max_results=max_results
+            )
         except self.http_error as ex:
             self.process_http_error(ex)
 
         schema_fields = [field.to_api_repr() for field in rows_iter.schema]
         nullsafe_dtypes = _bqschema_to_nullsafe_dtypes(schema_fields)
         df = rows_iter.to_dataframe(
-            dtypes=nullsafe_dtypes, bqstorage_client=self.bqstorage_client
+            dtypes=nullsafe_dtypes, bqstorage_client=bqstorage_client
         )
 
         if df.empty:
@@ -812,6 +830,7 @@ def read_gbq(
     configuration=None,
     credentials=None,
     use_bqstorage_api=False,
+    max_results=None,
     verbose=None,
     private_key=None,
 ):
@@ -907,9 +926,16 @@ def read_gbq(
         ``configuration`` dictionary.
 
         This feature requires the ``google-cloud-bigquery-storage`` and
-        ``fastavro`` packages.
+        ``pyarrow`` packages.
+
+        This value is ignored if ``max_results`` is set.
 
         .. versionadded:: 0.10.0
+    max_results : int, optional
+        If set, limit the maximum number of rows to fetch from the query
+        results.
+
+        .. versionadded:: 0.12.0
     verbose : None, deprecated
         Deprecated in Pandas-GBQ 0.4.0. Use the `logging module
         to adjust verbosity instead
@@ -969,7 +995,9 @@ def read_gbq(
         use_bqstorage_api=use_bqstorage_api,
     )
 
-    final_df = connector.run_query(query, configuration=configuration)
+    final_df = connector.run_query(
+        query, configuration=configuration, max_results=max_results
+    )
 
     # Reindex the DataFrame on the provided column
     if index_col is not None:

tests/system/test_gbq.py

@@ -584,6 +584,30 @@ def test_download_dataset_larger_than_200k_rows(self, project_id):
         )
         assert len(df.drop_duplicates()) == test_size
 
+    def test_ddl(self, random_dataset, project_id):
+        # Bug fix for https://github.com/pydata/pandas-gbq/issues/45
+        df = gbq.read_gbq(
+            "CREATE OR REPLACE TABLE {}.test_ddl (x INT64)".format(
+                random_dataset.dataset_id
+            )
+        )
+        assert len(df) == 0
+
+    def test_ddl_w_max_results(self, random_dataset, project_id):
+        df = gbq.read_gbq(
+            "CREATE OR REPLACE TABLE {}.test_ddl (x INT64)".format(
+                random_dataset.dataset_id
+            ),
+            max_results=0,
+        )
+        assert df is None
+
+    def test_max_results(self, random_dataset, project_id):
+        df = gbq.read_gbq(
+            "SELECT * FROM UNNEST(GENERATE_ARRAY(1, 100))", max_results=10
+        )
+        assert len(df) == 10
+
     def test_zero_rows(self, project_id):
         # Bug fix for https://github.com/pandas-dev/pandas/issues/10273
         df = gbq.read_gbq(

tests/unit/conftest.py

@@ -15,7 +15,6 @@ def reset_context():
 
 @pytest.fixture(autouse=True)
 def mock_bigquery_client(monkeypatch):
-    from google.api_core.exceptions import NotFound
     import google.cloud.bigquery
     import google.cloud.bigquery.table
 
@@ -35,7 +34,6 @@ def mock_bigquery_client(monkeypatch):
     mock_query.result.return_value = mock_rows
     mock_client.query.return_value = mock_query
     # Mock table creation.
-    mock_client.get_table.side_effect = NotFound("nope")
     monkeypatch.setattr(google.cloud.bigquery, "Client", mock_client)
     mock_client.reset_mock()
     return mock_client

tests/unit/test_gbq.py

@@ -34,10 +34,6 @@ def min_bq_version():
     return pkg_resources.parse_version("1.9.0")
 
 
-def mock_none_credentials(*args, **kwargs):
-    return None, None
-
-
 def mock_get_credentials_no_project(*args, **kwargs):
     import google.auth.credentials
 
@@ -76,16 +72,6 @@ def mock_compute_engine_credentials():
     return mock_credentials
 
 
-@pytest.fixture
-def mock_get_user_credentials(*args, **kwargs):
-    import google.auth.credentials
-
-    mock_credentials = mock.create_autospec(
-        google.auth.credentials.Credentials
-    )
-    return mock_credentials
-
-
 @pytest.fixture(autouse=True)
 def no_auth(monkeypatch):
     import pydata_google_auth
@@ -351,6 +337,17 @@ def test_read_gbq_without_inferred_project_id_from_compute_engine_credentials(
         )
 
 
+def test_read_gbq_with_max_results_zero(monkeypatch):
+    df = gbq.read_gbq("SELECT 1", dialect="standard", max_results=0)
+    assert df is None
+
+
+def test_read_gbq_with_max_results_ten(monkeypatch, mock_bigquery_client):
+    df = gbq.read_gbq("SELECT 1", dialect="standard", max_results=10)
+    assert df is not None
+    mock_bigquery_client.list_rows.assert_called_with(mock.ANY, max_results=10)
+
+
 def test_read_gbq_with_invalid_private_key_json_should_fail():
     with pytest.raises(pandas_gbq.exceptions.InvalidPrivateKeyFormat):
         gbq.read_gbq(
@@ -511,8 +508,12 @@ def test_generate_bq_schema_deprecated():
         gbq.generate_bq_schema(df)
 
 
-def test_load_does_not_modify_schema_arg():
-    # Test of Issue # 277
+def test_load_does_not_modify_schema_arg(mock_bigquery_client):
+    """Test of Issue # 277."""
+    from google.api_core.exceptions import NotFound
+
+    # Create table with new schema.
+    mock_bigquery_client.get_table.side_effect = NotFound("nope")
     df = DataFrame(
         {
             "field1": ["a", "b"],