diff --git a/bigquery/google/cloud/bigquery/client.py b/bigquery/google/cloud/bigquery/client.py index 12c0b57ad641..3f95e55e9290 100644 --- a/bigquery/google/cloud/bigquery/client.py +++ b/bigquery/google/cloud/bigquery/client.py @@ -1717,6 +1717,10 @@ def list_rows( max_results=max_results, page_size=page_size, extra_params=params, + table=table, + # Pass in selected_fields separately from schema so that full + # tables can be fetched without a column filter. + selected_fields=selected_fields, ) return row_iterator diff --git a/bigquery/google/cloud/bigquery/table.py b/bigquery/google/cloud/bigquery/table.py index 868921fe5a4c..8fadbd499bec 100644 --- a/bigquery/google/cloud/bigquery/table.py +++ b/bigquery/google/cloud/bigquery/table.py @@ -18,6 +18,7 @@ import copy import datetime +import json import operator import warnings @@ -1242,6 +1243,17 @@ class RowIterator(HTTPIterator): page_size (int, optional): The number of items to return per page. extra_params (Dict[str, object]): Extra query string parameters for the API call. + table (Union[ \ + :class:`~google.cloud.bigquery.table.Table`, \ + :class:`~google.cloud.bigquery.table.TableReference`, \ + ]): + Optional. The table which these rows belong to, or a reference to + it. Used to call the BigQuery Storage API to fetch rows. + selected_fields (Sequence[ \ + google.cloud.bigquery.schema.SchemaField, \ + ]): + Optional. A subset of columns to select from this table. + """ def __init__( @@ -1254,6 +1266,8 @@ def __init__( max_results=None, page_size=None, extra_params=None, + table=None, + selected_fields=None, ): super(RowIterator, self).__init__( client, @@ -1271,6 +1285,9 @@ def __init__( self._field_to_index = _helpers._field_to_index_mapping(schema) self._total_rows = None self._page_size = page_size + self._table = table + self._selected_fields = selected_fields + self._project = client.project def _get_next_page_response(self): """Requests the next page from the path provided. @@ -1296,9 +1313,81 @@ def total_rows(self): """int: The total number of rows in the table.""" return self._total_rows - def to_dataframe(self): + def _to_dataframe_tabledata_list(self): + """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) + + def _to_dataframe_bqstorage(self, bqstorage_client): + """Use (faster, but billable) BQ Storage API to construct DataFrame.""" + import concurrent.futures + from google.cloud import bigquery_storage_v1beta1 + + if "$" in self._table.table_id: + raise ValueError( + "Reading from a specific partition is not currently supported." + ) + if "@" in self._table.table_id: + raise ValueError( + "Reading from a specific snapshot is not currently supported." + ) + + read_options = bigquery_storage_v1beta1.types.TableReadOptions() + if self._selected_fields is not None: + for field in self._selected_fields: + read_options.selected_fields.append(field.name) + + session = bqstorage_client.create_read_session( + self._table.to_bqstorage(), + "projects/{}".format(self._project), + read_options=read_options, + ) + + # We need to parse the schema manually so that we can rearrange the + # columns. + schema = json.loads(session.avro_schema.schema) + columns = [field["name"] for field in schema["fields"]] + + # Avoid reading rows from an empty table. pandas.concat will fail on an + # empty list. + if not session.streams: + return pandas.DataFrame(columns=columns) + + def get_dataframe(stream): + position = bigquery_storage_v1beta1.types.StreamPosition(stream=stream) + rowstream = bqstorage_client.read_rows(position) + return rowstream.to_dataframe(session) + + with concurrent.futures.ThreadPoolExecutor() as pool: + frames = pool.map(get_dataframe, session.streams) + + # rowstream.to_dataframe() does not preserve column order. Rearrange at + # the end using manually-parsed schema. + return pandas.concat(frames)[columns] + + def to_dataframe(self, bqstorage_client=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. + + This method requires the ``fastavro`` and + ``google-cloud-bigquery-storage`` libraries. + + Reading from a specific partition or snapshot is not + currently supported by this method. + + **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. + Returns: pandas.DataFrame: A :class:`~pandas.DataFrame` populated with row data and column @@ -1312,11 +1401,10 @@ def to_dataframe(self): if pandas is None: raise ValueError(_NO_PANDAS_ERROR) - 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) + if bqstorage_client is not None: + return self._to_dataframe_bqstorage(bqstorage_client) + else: + return self._to_dataframe_tabledata_list() class _EmptyRowIterator(object): diff --git a/bigquery/tests/system.py b/bigquery/tests/system.py index fe5e3ce3dea0..a77d270c4279 100644 --- a/bigquery/tests/system.py +++ b/bigquery/tests/system.py @@ -28,6 +28,10 @@ import six import pytest +try: + from google.cloud import bigquery_storage_v1beta1 +except ImportError: # pragma: NO COVER + bigquery_storage_v1beta1 = None try: import pandas except ImportError: # pragma: NO COVER @@ -1496,7 +1500,7 @@ def test_query_iter(self): def test_query_results_to_dataframe(self): QUERY = """ SELECT id, author, time_ts, dead - from `bigquery-public-data.hacker_news.comments` + FROM `bigquery-public-data.hacker_news.comments` LIMIT 10 """ @@ -1518,6 +1522,53 @@ def test_query_results_to_dataframe(self): if not row[col] is None: self.assertIsInstance(row[col], exp_datatypes[col]) + @unittest.skipIf(pandas is None, "Requires `pandas`") + @unittest.skipIf( + bigquery_storage_v1beta1 is None, "Requires `google-cloud-bigquery-storage`" + ) + def test_query_results_to_dataframe_w_bqstorage(self): + dest_dataset = self.temp_dataset(_make_dataset_id("bqstorage_to_dataframe_")) + dest_ref = dest_dataset.table("query_results") + + query = """ + SELECT id, author, time_ts, dead + FROM `bigquery-public-data.hacker_news.comments` + LIMIT 10 + """ + + bqstorage_client = bigquery_storage_v1beta1.BigQueryStorageClient( + credentials=Config.CLIENT._credentials + ) + df = ( + Config.CLIENT.query( + query, + # There is a known issue reading small anonymous query result + # tables with the BQ Storage API. Writing to a destination + # table works around this issue. + job_config=bigquery.QueryJobConfig( + destination=dest_ref, write_disposition="WRITE_TRUNCATE" + ), + ) + .result() + .to_dataframe(bqstorage_client) + ) + + self.assertIsInstance(df, pandas.DataFrame) + self.assertEqual(len(df), 10) # verify the number of rows + column_names = ["id", "author", "time_ts", "dead"] + self.assertEqual(list(df), column_names) + exp_datatypes = { + "id": int, + "author": six.text_type, + "time_ts": pandas.Timestamp, + "dead": bool, + } + for index, row in df.iterrows(): + for col in column_names: + # all the schema fields are nullable, so None is acceptable + if not row[col] is None: + self.assertIsInstance(row[col], exp_datatypes[col]) + def test_insert_rows_nested_nested(self): # See #2951 SF = bigquery.SchemaField diff --git a/bigquery/tests/unit/test_table.py b/bigquery/tests/unit/test_table.py index d51a5dfec52e..9ed6eea2a3d0 100644 --- a/bigquery/tests/unit/test_table.py +++ b/bigquery/tests/unit/test_table.py @@ -13,6 +13,7 @@ # limitations under the License. import itertools +import json import unittest import mock @@ -31,6 +32,14 @@ from google.cloud.bigquery.dataset import DatasetReference +def _mock_client(): + from google.cloud.bigquery import client + + mock_client = mock.create_autospec(client.Client) + mock_client.project = "my-project" + return mock_client + + class _SchemaBase(object): def _verify_field(self, field, r_field): self.assertEqual(field.name, r_field["name"]) @@ -1289,7 +1298,7 @@ def test_constructor(self): from google.cloud.bigquery.table import _item_to_row from google.cloud.bigquery.table import _rows_page_start - client = mock.sentinel.client + client = _mock_client() api_request = mock.sentinel.api_request path = "/foo" schema = [] @@ -1322,7 +1331,7 @@ def test_iterate(self): ] path = "/foo" api_request = mock.Mock(return_value={"rows": rows}) - row_iterator = RowIterator(mock.sentinel.client, api_request, path, schema) + row_iterator = RowIterator(_mock_client(), api_request, path, schema) self.assertEqual(row_iterator.num_results, 0) rows_iter = iter(row_iterator) @@ -1356,7 +1365,7 @@ def test_page_size(self): api_request = mock.Mock(return_value={"rows": rows}) row_iterator = RowIterator( - mock.sentinel.client, api_request, path, schema, page_size=4 + _mock_client(), api_request, path, schema, page_size=4 ) row_iterator._get_next_page_response() @@ -1383,7 +1392,7 @@ def test_to_dataframe(self): ] path = "/foo" api_request = mock.Mock(return_value={"rows": rows}) - row_iterator = RowIterator(mock.sentinel.client, api_request, path, schema) + row_iterator = RowIterator(_mock_client(), api_request, path, schema) df = row_iterator.to_dataframe() @@ -1404,7 +1413,7 @@ def test_to_dataframe_w_empty_results(self): ] path = "/foo" api_request = mock.Mock(return_value={"rows": []}) - row_iterator = RowIterator(mock.sentinel.client, api_request, path, schema) + row_iterator = RowIterator(_mock_client(), api_request, path, schema) df = row_iterator.to_dataframe() @@ -1435,7 +1444,7 @@ def test_to_dataframe_w_various_types_nullable(self): 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.sentinel.client, api_request, path, schema) + row_iterator = RowIterator(_mock_client(), api_request, path, schema) df = row_iterator.to_dataframe() @@ -1475,7 +1484,7 @@ def test_to_dataframe_column_dtypes(self): 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.sentinel.client, api_request, path, schema) + row_iterator = RowIterator(_mock_client(), api_request, path, schema) df = row_iterator.to_dataframe() @@ -1506,11 +1515,150 @@ def test_to_dataframe_error_if_pandas_is_none(self): ] path = "/foo" api_request = mock.Mock(return_value={"rows": rows}) - row_iterator = RowIterator(mock.sentinel.client, api_request, path, schema) + row_iterator = RowIterator(_mock_client(), api_request, path, schema) with self.assertRaises(ValueError): row_iterator.to_dataframe() + @unittest.skipIf(pandas is None, "Requires `pandas`") + @unittest.skipIf( + bigquery_storage_v1beta1 is None, "Requires `google-cloud-bigquery-storage`" + ) + def test_to_dataframe_w_bqstorage_empty(self): + from google.cloud.bigquery import schema + from google.cloud.bigquery import table as mut + + bqstorage_client = mock.create_autospec( + bigquery_storage_v1beta1.BigQueryStorageClient + ) + session = bigquery_storage_v1beta1.types.ReadSession() + session.avro_schema.schema = json.dumps( + { + "fields": [ + {"name": "colA"}, + # Not alphabetical to test column order. + {"name": "colC"}, + {"name": "colB"}, + ] + } + ) + bqstorage_client.create_read_session.return_value = session + + row_iterator = mut.RowIterator( + _mock_client(), + None, # api_request: ignored + None, # path: ignored + [ + schema.SchemaField("colA", "IGNORED"), + schema.SchemaField("colC", "IGNORED"), + schema.SchemaField("colB", "IGNORED"), + ], + table=mut.TableReference.from_string("proj.dset.tbl"), + ) + + got = row_iterator.to_dataframe(bqstorage_client) + column_names = ["colA", "colC", "colB"] + self.assertEqual(list(got), column_names) + self.assertTrue(got.empty) + + @unittest.skipIf(pandas is None, "Requires `pandas`") + @unittest.skipIf( + bigquery_storage_v1beta1 is None, "Requires `google-cloud-bigquery-storage`" + ) + def test_to_dataframe_w_bqstorage_nonempty(self): + from google.cloud.bigquery import schema + from google.cloud.bigquery import table as mut + from google.cloud.bigquery_storage_v1beta1 import reader + + mock_rowstream = mock.create_autospec(reader.ReadRowsStream) + mock_rowstream.to_dataframe.return_value = pandas.DataFrame( + [ + {"colA": 1, "colB": "abc", "colC": 2.0}, + {"colA": -1, "colB": "def", "colC": 4.0}, + ] + ) + bqstorage_client = mock.create_autospec( + bigquery_storage_v1beta1.BigQueryStorageClient + ) + session = bigquery_storage_v1beta1.types.ReadSession( + streams=[{"name": "/projects/proj/dataset/dset/tables/tbl/streams/1234"}] + ) + session.avro_schema.schema = json.dumps( + { + "fields": [ + {"name": "colA"}, + # Not alphabetical to test column order. + {"name": "colC"}, + {"name": "colB"}, + ] + } + ) + bqstorage_client.create_read_session.return_value = session + bqstorage_client.read_rows.return_value = mock_rowstream + schema = [ + schema.SchemaField("colA", "IGNORED"), + schema.SchemaField("colC", "IGNORED"), + schema.SchemaField("colB", "IGNORED"), + ] + + row_iterator = mut.RowIterator( + _mock_client(), + None, # api_request: ignored + None, # path: ignored + schema, + table=mut.TableReference.from_string("proj.dset.tbl"), + selected_fields=schema, + ) + + got = row_iterator.to_dataframe(bqstorage_client) + column_names = ["colA", "colC", "colB"] + self.assertEqual(list(got), column_names) + self.assertEqual(len(got.index), 2) + + @unittest.skipIf( + bigquery_storage_v1beta1 is None, "Requires `google-cloud-bigquery-storage`" + ) + def test_to_dataframe_w_bqstorage_partition(self): + from google.cloud.bigquery import schema + from google.cloud.bigquery import table as mut + + bqstorage_client = mock.create_autospec( + bigquery_storage_v1beta1.BigQueryStorageClient + ) + + row_iterator = mut.RowIterator( + _mock_client(), + None, # api_request: ignored + None, # path: ignored + [schema.SchemaField("colA", "IGNORED")], + table=mut.TableReference.from_string("proj.dset.tbl$20181225"), + ) + + with pytest.raises(ValueError): + row_iterator.to_dataframe(bqstorage_client) + + @unittest.skipIf( + bigquery_storage_v1beta1 is None, "Requires `google-cloud-bigquery-storage`" + ) + def test_to_dataframe_w_bqstorage_snapshot(self): + from google.cloud.bigquery import schema + from google.cloud.bigquery import table as mut + + bqstorage_client = mock.create_autospec( + bigquery_storage_v1beta1.BigQueryStorageClient + ) + + row_iterator = mut.RowIterator( + _mock_client(), + None, # api_request: ignored + None, # path: ignored + [schema.SchemaField("colA", "IGNORED")], + table=mut.TableReference.from_string("proj.dset.tbl@1234567890000"), + ) + + with pytest.raises(ValueError): + row_iterator.to_dataframe(bqstorage_client) + class TestTimePartitioning(unittest.TestCase): def _get_target_class(self):