Skip to content

Commit

Permalink
Added udf_resource_uri parameter to read_gbq
Browse files Browse the repository at this point in the history
Now more complicated queries could be processed.
  • Loading branch information
necnec committed Nov 25, 2016
1 parent 22d982a commit 55bf05c
Show file tree
Hide file tree
Showing 2 changed files with 18 additions and 3 deletions.
1 change: 1 addition & 0 deletions doc/source/whatsnew/v0.19.0.txt
Original file line number Diff line number Diff line change
Expand Up @@ -379,6 +379,7 @@ Google BigQuery Enhancements

- The :func:`read_gbq` method has gained the ``dialect`` argument to allow users to specify whether to use BigQuery's legacy SQL or BigQuery's standard SQL. See the :ref:`docs <io.bigquery_reader>` for more details (:issue:`13615`).
- The :func:`~DataFrame.to_gbq` method now allows the DataFrame column order to differ from the destination table schema (:issue:`11359`).
- The :func:`read_gbq` method now allows passing user defined functions

.. _whatsnew_0190.errstate:

Expand Down
20 changes: 17 additions & 3 deletions pandas/io/gbq.py
Original file line number Diff line number Diff line change
Expand Up @@ -375,7 +375,7 @@ def process_insert_errors(self, insert_errors):

raise StreamingInsertError

def run_query(self, query):
def run_query(self, query, udf_resource_uri=None):
try:
from googleapiclient.errors import HttpError
except:
Expand All @@ -395,6 +395,14 @@ def run_query(self, query):
}
}
}

if udf_resource_uri is not None:
if not isinstance(udf_resource_uri, list):
udf_resource_uri = [udf_resource_uri]

job_data['configuration']['query']['userDefinedFunctionResources'] = \
[{'resourceUri': uri} for uri in udf_resource_uri]


self._start_timer()
try:
Expand Down Expand Up @@ -622,7 +630,7 @@ def _parse_entry(field_value, field_type):


def read_gbq(query, project_id=None, index_col=None, col_order=None,
reauth=False, verbose=True, private_key=None, dialect='legacy'):
reauth=False, verbose=True, private_key=None, dialect='legacy', udf_resource_uri=None):
"""Load data from Google BigQuery.
THIS IS AN EXPERIMENTAL LIBRARY
Expand Down Expand Up @@ -682,6 +690,12 @@ def read_gbq(query, project_id=None, index_col=None, col_order=None,
.. versionadded:: 0.19.0
udf_resource_uri : list(str) or str (optional)
A code resource to load from a Google Cloud Storage URI.
Describes user-defined function resources used in the query.
.. versionadded:: 0.19.0
Returns
-------
df: DataFrame
Expand All @@ -698,7 +712,7 @@ def read_gbq(query, project_id=None, index_col=None, col_order=None,
connector = GbqConnector(project_id, reauth=reauth, verbose=verbose,
private_key=private_key,
dialect=dialect)
schema, pages = connector.run_query(query)
schema, pages = connector.run_query(query, udf_resource_uri)
dataframe_list = []
while len(pages) > 0:
page = pages.pop()
Expand Down

0 comments on commit 55bf05c

Please sign in to comment.