Skip to content

Commit 8c7b839

Browse files
plamutplamut
authored
feat: add support for user defined Table View Functions (#724)
* Add auxiliary classes for TVF routines * Add return_table_type property to Routine * Add system test for TVF routines * Use the generated StandardSqlTableType class * Update docs with new changes * Add missing space in misc. Sphinx directives
1 parent b8b5433 commit 8c7b839

File tree

8 files changed

+270
-14
lines changed

8 files changed

+270
-14
lines changed

‎docs/reference.rst

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -118,6 +118,7 @@ Routine
118118
routine.Routine
119119
routine.RoutineArgument
120120
routine.RoutineReference
121+
routine.RoutineType
121122

122123
Schema
123124
======

‎google/cloud/bigquery/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,7 @@
8585
from google.cloud.bigquery.routine import Routine
8686
from google.cloud.bigquery.routine import RoutineArgument
8787
from google.cloud.bigquery.routine import RoutineReference
88+
from google.cloud.bigquery.routine import RoutineType
8889
from google.cloud.bigquery.schema import SchemaField
8990
from google.cloud.bigquery.table import PartitionRange
9091
from google.cloud.bigquery.table import RangePartitioning
@@ -162,6 +163,7 @@
162163
"KeyResultStatementKind",
163164
"OperationType",
164165
"QueryPriority",
166+
"RoutineType",
165167
"SchemaUpdateOption",
166168
"SourceFormat",
167169
"SqlTypeNames",

‎google/cloud/bigquery/job/query.py

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1386,12 +1386,12 @@ def to_arrow(
13861386
13871387
This argument does nothing if ``bqstorage_client`` is supplied.
13881388
1389-
..versionadded:: 1.24.0
1389+
.. versionadded:: 1.24.0
13901390
13911391
max_results (Optional[int]):
13921392
Maximum number of rows to include in the result. No limit by default.
13931393
1394-
..versionadded:: 2.21.0
1394+
.. versionadded:: 2.21.0
13951395
13961396
Returns:
13971397
pyarrow.Table
@@ -1403,7 +1403,7 @@ def to_arrow(
14031403
ValueError:
14041404
If the :mod:`pyarrow` library cannot be imported.
14051405
1406-
..versionadded:: 1.17.0
1406+
.. versionadded:: 1.17.0
14071407
"""
14081408
query_result = wait_for_query(self, progress_bar_type, max_results=max_results)
14091409
return query_result.to_arrow(
@@ -1452,7 +1452,7 @@ def to_dataframe(
14521452
:func:`~google.cloud.bigquery.table.RowIterator.to_dataframe`
14531453
for details.
14541454
1455-
..versionadded:: 1.11.0
1455+
.. versionadded:: 1.11.0
14561456
create_bqstorage_client (Optional[bool]):
14571457
If ``True`` (default), create a BigQuery Storage API client
14581458
using the default API settings. The BigQuery Storage API
@@ -1461,18 +1461,18 @@ def to_dataframe(
14611461
14621462
This argument does nothing if ``bqstorage_client`` is supplied.
14631463
1464-
..versionadded:: 1.24.0
1464+
.. versionadded:: 1.24.0
14651465
14661466
date_as_object (Optional[bool]):
14671467
If ``True`` (default), cast dates to objects. If ``False``, convert
14681468
to datetime64[ns] dtype.
14691469
1470-
..versionadded:: 1.26.0
1470+
.. versionadded:: 1.26.0
14711471
14721472
max_results (Optional[int]):
14731473
Maximum number of rows to include in the result. No limit by default.
14741474
1475-
..versionadded:: 2.21.0
1475+
.. versionadded:: 2.21.0
14761476
14771477
Returns:
14781478
A :class:`~pandas.DataFrame` populated with row data and column

‎google/cloud/bigquery/routine/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,11 +19,13 @@
1919
from google.cloud.bigquery.routine.routine import Routine
2020
from google.cloud.bigquery.routine.routine import RoutineArgument
2121
from google.cloud.bigquery.routine.routine import RoutineReference
22+
from google.cloud.bigquery.routine.routine import RoutineType
2223

2324

2425
__all__ = (
2526
"DeterminismLevel",
2627
"Routine",
2728
"RoutineArgument",
2829
"RoutineReference",
30+
"RoutineType",
2931
)

‎google/cloud/bigquery/routine/routine.py

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,21 @@
2121
import google.cloud._helpers
2222
from google.cloud.bigquery import _helpers
2323
import google.cloud.bigquery_v2.types
24+
from google.cloud.bigquery_v2.types import StandardSqlTableType
25+
26+
27+
class RoutineType:
28+
"""The fine-grained type of the routine.
29+
30+
https://cloud.google.com/bigquery/docs/reference/rest/v2/routines#routinetype
31+
32+
.. versionadded:: 2.22.0
33+
"""
34+
35+
ROUTINE_TYPE_UNSPECIFIED = "ROUTINE_TYPE_UNSPECIFIED"
36+
SCALAR_FUNCTION = "SCALAR_FUNCTION"
37+
PROCEDURE = "PROCEDURE"
38+
TABLE_VALUED_FUNCTION = "TABLE_VALUED_FUNCTION"
2439

2540

2641
class Routine(object):
@@ -48,6 +63,7 @@ class Routine(object):
4863
"modified": "lastModifiedTime",
4964
"reference": "routineReference",
5065
"return_type": "returnType",
66+
"return_table_type": "returnTableType",
5167
"type_": "routineType",
5268
"description": "description",
5369
"determinism_level": "determinismLevel",
@@ -204,6 +220,35 @@ def return_type(self, value):
204220
resource = None
205221
self._properties[self._PROPERTY_TO_API_FIELD["return_type"]] = resource
206222

223+
@property
224+
def return_table_type(self) -> StandardSqlTableType:
225+
"""The return type of a Table Valued Function (TVF) routine.
226+
227+
.. versionadded:: 2.22.0
228+
"""
229+
resource = self._properties.get(
230+
self._PROPERTY_TO_API_FIELD["return_table_type"]
231+
)
232+
if not resource:
233+
return resource
234+
235+
output = google.cloud.bigquery_v2.types.StandardSqlTableType()
236+
raw_protobuf = json_format.ParseDict(
237+
resource, output._pb, ignore_unknown_fields=True
238+
)
239+
return type(output).wrap(raw_protobuf)
240+
241+
@return_table_type.setter
242+
def return_table_type(self, value):
243+
if not value:
244+
resource = None
245+
else:
246+
resource = {
247+
"columns": [json_format.MessageToDict(col._pb) for col in value.columns]
248+
}
249+
250+
self._properties[self._PROPERTY_TO_API_FIELD["return_table_type"]] = resource
251+
207252
@property
208253
def imported_libraries(self):
209254
"""List[str]: The path of the imported JavaScript libraries.

‎google/cloud/bigquery/table.py

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1684,7 +1684,7 @@ def to_arrow(
16841684
16851685
This argument does nothing if ``bqstorage_client`` is supplied.
16861686
1687-
..versionadded:: 1.24.0
1687+
.. versionadded:: 1.24.0
16881688
16891689
Returns:
16901690
pyarrow.Table
@@ -1695,7 +1695,7 @@ def to_arrow(
16951695
Raises:
16961696
ValueError: If the :mod:`pyarrow` library cannot be imported.
16971697
1698-
..versionadded:: 1.17.0
1698+
.. versionadded:: 1.17.0
16991699
"""
17001700
if pyarrow is None:
17011701
raise ValueError(_NO_PYARROW_ERROR)
@@ -1775,7 +1775,7 @@ def to_dataframe_iterable(
17751775
created by the server. If ``max_queue_size`` is :data:`None`, the queue
17761776
size is infinite.
17771777
1778-
..versionadded:: 2.14.0
1778+
.. versionadded:: 2.14.0
17791779
17801780
Returns:
17811781
pandas.DataFrame:
@@ -1861,7 +1861,7 @@ def to_dataframe(
18611861
Use the :func:`tqdm.tqdm_gui` function to display a
18621862
progress bar as a graphical dialog box.
18631863
1864-
..versionadded:: 1.11.0
1864+
.. versionadded:: 1.11.0
18651865
create_bqstorage_client (Optional[bool]):
18661866
If ``True`` (default), create a BigQuery Storage API client
18671867
using the default API settings. The BigQuery Storage API
@@ -1870,13 +1870,13 @@ def to_dataframe(
18701870
18711871
This argument does nothing if ``bqstorage_client`` is supplied.
18721872
1873-
..versionadded:: 1.24.0
1873+
.. versionadded:: 1.24.0
18741874
18751875
date_as_object (Optional[bool]):
18761876
If ``True`` (default), cast dates to objects. If ``False``, convert
18771877
to datetime64[ns] dtype.
18781878
1879-
..versionadded:: 1.26.0
1879+
.. versionadded:: 1.26.0
18801880
18811881
Returns:
18821882
pandas.DataFrame:
@@ -2010,7 +2010,7 @@ def to_dataframe_iterable(
20102010
) -> Iterator["pandas.DataFrame"]:
20112011
"""Create an iterable of pandas DataFrames, to process the table as a stream.
20122012
2013-
..versionadded:: 2.21.0
2013+
.. versionadded:: 2.21.0
20142014
20152015
Args:
20162016
bqstorage_client:

‎tests/system/test_client.py

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2228,6 +2228,85 @@ def test_create_routine(self):
22282228
assert len(rows) == 1
22292229
assert rows[0].max_value == 100.0
22302230

2231+
def test_create_tvf_routine(self):
2232+
from google.cloud.bigquery import Routine, RoutineArgument, RoutineType
2233+
2234+
StandardSqlDataType = bigquery_v2.types.StandardSqlDataType
2235+
StandardSqlField = bigquery_v2.types.StandardSqlField
2236+
StandardSqlTableType = bigquery_v2.types.StandardSqlTableType
2237+
2238+
INT64 = StandardSqlDataType.TypeKind.INT64
2239+
STRING = StandardSqlDataType.TypeKind.STRING
2240+
2241+
client = Config.CLIENT
2242+
2243+
dataset = self.temp_dataset(_make_dataset_id("create_tvf_routine"))
2244+
routine_ref = dataset.routine("test_tvf_routine")
2245+
2246+
routine_body = """
2247+
SELECT int_col, str_col
2248+
FROM (
2249+
UNNEST([1, 2, 3]) int_col
2250+
JOIN
2251+
(SELECT str_col FROM UNNEST(["one", "two", "three"]) str_col)
2252+
ON TRUE
2253+
)
2254+
WHERE int_col > threshold
2255+
"""
2256+
2257+
return_table_type = StandardSqlTableType(
2258+
columns=[
2259+
StandardSqlField(
2260+
name="int_col", type=StandardSqlDataType(type_kind=INT64),
2261+
),
2262+
StandardSqlField(
2263+
name="str_col", type=StandardSqlDataType(type_kind=STRING),
2264+
),
2265+
]
2266+
)
2267+
2268+
routine_args = [
2269+
RoutineArgument(
2270+
name="threshold", data_type=StandardSqlDataType(type_kind=INT64),
2271+
)
2272+
]
2273+
2274+
routine_def = Routine(
2275+
routine_ref,
2276+
type_=RoutineType.TABLE_VALUED_FUNCTION,
2277+
arguments=routine_args,
2278+
return_table_type=return_table_type,
2279+
body=routine_body,
2280+
)
2281+
2282+
# Create TVF routine.
2283+
client.delete_routine(routine_ref, not_found_ok=True)
2284+
routine = client.create_routine(routine_def)
2285+
2286+
assert routine.body == routine_body
2287+
assert routine.return_table_type == return_table_type
2288+
assert routine.arguments == routine_args
2289+
2290+
# Execute the routine to see if it's working as expected.
2291+
query_job = client.query(
2292+
f"""
2293+
SELECT int_col, str_col
2294+
FROM `{routine.reference}`(1)
2295+
ORDER BY int_col, str_col ASC
2296+
"""
2297+
)
2298+
2299+
result_rows = [tuple(row) for row in query_job.result()]
2300+
expected = [
2301+
(2, "one"),
2302+
(2, "three"),
2303+
(2, "two"),
2304+
(3, "one"),
2305+
(3, "three"),
2306+
(3, "two"),
2307+
]
2308+
assert result_rows == expected
2309+
22312310
def test_create_table_rows_fetch_nested_schema(self):
22322311
table_name = "test_table"
22332312
dataset = self.temp_dataset(_make_dataset_id("create_table_nested_schema"))

0 commit comments

Comments
 (0)