diff --git a/bigquery/google/cloud/bigquery/job.py b/bigquery/google/cloud/bigquery/job.py index 442420a7191b..5cd90ada9637 100644 --- a/bigquery/google/cloud/bigquery/job.py +++ b/bigquery/google/cloud/bigquery/job.py @@ -2896,7 +2896,9 @@ def result(self, timeout=None, page_size=None, retry=DEFAULT_RETRY): rows._preserve_order = _contains_order_by(self.query) return rows - def to_arrow(self, progress_bar_type=None): + # If changing the signature of this method, make sure to apply the same + # changes to table.RowIterator.to_arrow() + def to_arrow(self, progress_bar_type=None, bqstorage_client=None): """[Beta] Create a class:`pyarrow.Table` by loading all pages of a table or query. @@ -2919,6 +2921,18 @@ def to_arrow(self, progress_bar_type=None): ``'tqdm_gui'`` Use the :func:`tqdm.tqdm_gui` function to display a progress bar as a graphical dialog box. + bqstorage_client ( \ + google.cloud.bigquery_storage_v1beta1.BigQueryStorageClient \ + ): + **Beta 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 ``pyarrow`` and + ``google-cloud-bigquery-storage`` libraries. + + Reading from a specific partition or snapshot is not + currently supported by this method. Returns: pyarrow.Table @@ -2932,8 +2946,12 @@ def to_arrow(self, progress_bar_type=None): ..versionadded:: 1.17.0 """ - return self.result().to_arrow(progress_bar_type=progress_bar_type) + return self.result().to_arrow( + progress_bar_type=progress_bar_type, bqstorage_client=bqstorage_client + ) + # If changing the signature of this method, make sure to apply the same + # changes to table.RowIterator.to_dataframe() def to_dataframe(self, bqstorage_client=None, dtypes=None, progress_bar_type=None): """Return a pandas DataFrame from a QueryJob diff --git a/bigquery/google/cloud/bigquery/table.py b/bigquery/google/cloud/bigquery/table.py index 3c8ce6cc39d7..5b1de345c068 100644 --- a/bigquery/google/cloud/bigquery/table.py +++ b/bigquery/google/cloud/bigquery/table.py @@ -1449,6 +1449,8 @@ def _to_arrow_iterable(self, bqstorage_client=None): bqstorage_client=bqstorage_client, ) + # If changing the signature of this method, make sure to apply the same + # changes to job.QueryJob.to_arrow() def to_arrow(self, progress_bar_type=None, bqstorage_client=None): """[Beta] Create a class:`pyarrow.Table` by loading all pages of a table or query. @@ -1552,6 +1554,8 @@ def _to_dataframe_iterable(self, bqstorage_client=None, dtypes=None): bqstorage_client=bqstorage_client, ) + # If changing the signature of this method, make sure to apply the same + # changes to job.QueryJob.to_dataframe() def to_dataframe(self, bqstorage_client=None, dtypes=None, progress_bar_type=None): """Create a pandas DataFrame by loading all pages of a query. diff --git a/bigquery/tests/unit/test_signature_compatibility.py b/bigquery/tests/unit/test_signature_compatibility.py new file mode 100644 index 000000000000..6002ae3e87c9 --- /dev/null +++ b/bigquery/tests/unit/test_signature_compatibility.py @@ -0,0 +1,51 @@ +# Copyright 2019 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import inspect + +import pytest + + +@pytest.fixture +def query_job_class(): + from google.cloud.bigquery.job import QueryJob + + return QueryJob + + +@pytest.fixture +def row_iterator_class(): + from google.cloud.bigquery.table import RowIterator + + return RowIterator + + +@pytest.mark.skipif( + not hasattr(inspect, "signature"), + reason="inspect.signature() is not availalbe in older Python versions", +) +def test_to_arrow_method_signatures_match(query_job_class, row_iterator_class): + sig = inspect.signature(query_job_class.to_arrow) + sig2 = inspect.signature(row_iterator_class.to_arrow) + assert sig == sig2 + + +@pytest.mark.skipif( + not hasattr(inspect, "signature"), + reason="inspect.signature() is not availalbe in older Python versions", +) +def test_to_dataframe_method_signatures_match(query_job_class, row_iterator_class): + sig = inspect.signature(query_job_class.to_dataframe) + sig2 = inspect.signature(row_iterator_class.to_dataframe) + assert sig == sig2