|
1 | 1 | import logging
|
2 |
| -import os |
3 | 2 | import time
|
4 | 3 | import warnings
|
5 |
| -from collections import OrderedDict |
6 | 4 | from datetime import datetime
|
7 | 5 |
|
8 | 6 | import numpy as np
|
9 |
| -from pandas import DataFrame |
10 | 7 |
|
11 | 8 | from pandas_gbq.exceptions import AccessDenied
|
12 | 9 |
|
@@ -69,7 +66,7 @@ def _check_google_client_version():
|
69 | 66 | def _test_google_api_imports():
|
70 | 67 |
|
71 | 68 | try:
|
72 |
| - import pydata_google_auth |
| 69 | + import pydata_google_auth # noqa |
73 | 70 | except ImportError as ex:
|
74 | 71 | raise ImportError(
|
75 | 72 | "pandas-gbq requires pydata-google-auth: {0}".format(ex)
|
@@ -483,15 +480,9 @@ def run_query(self, query, **kwargs):
|
483 | 480 | rows_iter = query_reply.result()
|
484 | 481 | except self.http_error as ex:
|
485 | 482 | self.process_http_error(ex)
|
486 |
| - result_rows = list(rows_iter) |
487 |
| - total_rows = rows_iter.total_rows |
488 |
| - schema = { |
489 |
| - "fields": [field.to_api_repr() for field in rows_iter.schema] |
490 |
| - } |
491 |
| - |
492 |
| - logger.debug("Got {} rows.\n".format(total_rows)) |
493 |
| - |
494 |
| - return schema, result_rows |
| 483 | + df = rows_iter.to_dataframe() |
| 484 | + logger.debug("Got {} rows.\n".format(rows_iter.total_rows)) |
| 485 | + return df |
495 | 486 |
|
496 | 487 | def load_data(
|
497 | 488 | self,
|
@@ -662,25 +653,6 @@ def _parse_schema(schema_fields):
|
662 | 653 | yield name, dtype
|
663 | 654 |
|
664 | 655 |
|
665 |
| -def _parse_data(schema, rows): |
666 |
| - |
667 |
| - column_dtypes = OrderedDict(_parse_schema(schema["fields"])) |
668 |
| - df = DataFrame(data=(iter(r) for r in rows), columns=column_dtypes.keys()) |
669 |
| - |
670 |
| - for column in df: |
671 |
| - dtype = column_dtypes[column] |
672 |
| - null_safe = ( |
673 |
| - df[column].notnull().all() |
674 |
| - or dtype == float |
675 |
| - or dtype == "datetime64[ns]" |
676 |
| - ) |
677 |
| - if dtype and null_safe: |
678 |
| - df[column] = df[column].astype( |
679 |
| - column_dtypes[column], errors="ignore" |
680 |
| - ) |
681 |
| - return df |
682 |
| - |
683 |
| - |
684 | 656 | def read_gbq(
|
685 | 657 | query,
|
686 | 658 | project_id=None,
|
@@ -833,8 +805,8 @@ def read_gbq(
|
833 | 805 | credentials=credentials,
|
834 | 806 | private_key=private_key,
|
835 | 807 | )
|
836 |
| - schema, rows = connector.run_query(query, configuration=configuration) |
837 |
| - final_df = _parse_data(schema, rows) |
| 808 | + |
| 809 | + final_df = connector.run_query(query, configuration=configuration) |
838 | 810 |
|
839 | 811 | # Reindex the DataFrame on the provided column
|
840 | 812 | if index_col is not None:
|
|
0 commit comments