Add max_results kwarg to read_gbq (#34639) (#34641)

Since max_results is a new kwarg (added in pandas-gbq 0.12.0), it
is handled and tested in the same way as use_bqstorage_api,
using the "new kwargs" mechanism to maintain backwards
compatibility with older pandas-gbq versions.

Author: JohnPaton

doc/source/whatsnew/v1.1.0.rst

@@ -292,6 +292,7 @@ Other enhancements
 - :meth:`groupby.transform` now allows ``func`` to be ``pad``, ``backfill`` and ``cumcount`` (:issue:`31269`).
 - :meth:`~pandas.io.json.read_json` now accepts `nrows` parameter. (:issue:`33916`).
 - :meth `~pandas.io.gbq.read_gbq` now allows to disable progress bar (:issue:`33360`).
+- :meth:`~pandas.io.gbq.read_gbq` now supports the ``max_results`` kwarg from ``pandas-gbq`` (:issue:`34639`).
 
 .. ---------------------------------------------------------------------------
 

pandas/io/gbq.py

@@ -30,6 +30,7 @@ def read_gbq(
     configuration: Optional[Dict[str, Any]] = None,
     credentials=None,
     use_bqstorage_api: Optional[bool] = None,
+    max_results: Optional[int] = None,
     private_key=None,
     verbose=None,
     progress_bar_type: Optional[str] = None,
@@ -125,6 +126,13 @@ def read_gbq(
         ``fastavro`` packages.
 
         .. versionadded:: 0.25.0
+    max_results : int, optional
+        If set, limit the maximum number of rows to fetch from the query
+        results.
+
+        *New in version 0.12.0 of pandas-gbq*.
+
+        .. versionadded:: 1.1.0
     progress_bar_type : Optional, str
         If set, use the `tqdm <https://tqdm.github.io/>`__ library to
         display a progress bar while the data downloads. Install the
@@ -162,11 +170,13 @@ def read_gbq(
     """
     pandas_gbq = _try_import()
 
-    kwargs: Dict[str, Union[str, bool, None]] = {}
+    kwargs: Dict[str, Union[str, bool, int, None]] = {}
 
     # START: new kwargs.  Don't populate unless explicitly set.
     if use_bqstorage_api is not None:
         kwargs["use_bqstorage_api"] = use_bqstorage_api
+    if max_results is not None:
+        kwargs["max_results"] = max_results
 
     kwargs["progress_bar_type"] = progress_bar_type
     # END: new kwargs

pandas/tests/io/test_gbq.py

@@ -113,9 +113,10 @@ def mock_read_gbq(sql, **kwargs):
         return DataFrame([[1.0]])
 
     monkeypatch.setattr("pandas_gbq.read_gbq", mock_read_gbq)
-    pd.read_gbq("SELECT 1", use_bqstorage_api=True)
+    pd.read_gbq("SELECT 1", use_bqstorage_api=True, max_results=1)
 
     assert captured_kwargs["use_bqstorage_api"]
+    assert captured_kwargs["max_results"]
 
 
 def test_read_gbq_without_new_kwargs(monkeypatch):
@@ -129,6 +130,7 @@ def mock_read_gbq(sql, **kwargs):
     pd.read_gbq("SELECT 1")
 
     assert "use_bqstorage_api" not in captured_kwargs
+    assert "max_results" not in captured_kwargs
 
 
 @pytest.mark.parametrize("progress_bar", [None, "foo"])