This repository was archived by the owner on Apr 1, 2026. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 67
Expand file tree
/
Copy path_function_session.py
More file actions
924 lines (824 loc) · 43.7 KB
/
_function_session.py
File metadata and controls
924 lines (824 loc) · 43.7 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
# Copyright 2023 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.
from __future__ import annotations
import collections.abc
import functools
import inspect
import sys
import threading
from typing import (
Any,
cast,
Dict,
Literal,
Mapping,
Optional,
Sequence,
TYPE_CHECKING,
Union,
)
import warnings
import google.api_core.exceptions
from google.cloud import (
bigquery,
bigquery_connection_v1,
functions_v2,
resourcemanager_v3,
)
from bigframes import clients
import bigframes.exceptions as bfe
import bigframes.formatting_helpers as bf_formatting
from bigframes.functions import function as bq_functions
from bigframes.functions import udf_def
if TYPE_CHECKING:
from bigframes.session import Session
from bigframes.functions import _function_client, _utils
class FunctionSession:
"""Session to manage bigframes functions."""
def __init__(self):
# Session level mapping of function artifacts
self._temp_artifacts: Dict[str, str] = dict()
# Lock to synchronize the update of the session artifacts
self._artifacts_lock = threading.Lock()
def _resolve_session(self, session: Optional[Session]) -> Session:
"""Resolves the BigFrames session."""
import bigframes.pandas as bpd
import bigframes.session
# Using the global session if none is provided.
return cast(bigframes.session.Session, session or bpd.get_global_session())
def _resolve_bigquery_client(
self, session: Session, bigquery_client: Optional[bigquery.Client]
) -> bigquery.Client:
"""Resolves the BigQuery client."""
if not bigquery_client:
bigquery_client = session.bqclient
if not bigquery_client:
raise bf_formatting.create_exception_with_feedback_link(
ValueError,
"A bigquery client must be provided, either directly or via "
"session.",
)
return bigquery_client
def _resolve_bigquery_connection_client(
self,
session: Session,
bigquery_connection_client: Optional[
bigquery_connection_v1.ConnectionServiceClient
],
) -> bigquery_connection_v1.ConnectionServiceClient:
"""Resolves the BigQuery connection client."""
if not bigquery_connection_client:
bigquery_connection_client = session.bqconnectionclient
if not bigquery_connection_client:
raise bf_formatting.create_exception_with_feedback_link(
ValueError,
"A bigquery connection client must be provided, either "
"directly or via session.",
)
return bigquery_connection_client
def _resolve_resource_manager_client(
self,
session: Session,
resource_manager_client: Optional[resourcemanager_v3.ProjectsClient],
) -> resourcemanager_v3.ProjectsClient:
"""Resolves the resource manager client."""
if not resource_manager_client:
resource_manager_client = session.resourcemanagerclient
if not resource_manager_client:
raise bf_formatting.create_exception_with_feedback_link(
ValueError,
"A resource manager client must be provided, either directly "
"or via session.",
)
return resource_manager_client
def _resolve_dataset_reference(
self,
session: Session,
bigquery_client: bigquery.Client,
dataset: Optional[str],
) -> bigquery.DatasetReference:
"""Resolves the dataset reference for the bigframes function."""
if dataset:
dataset_ref = bigquery.DatasetReference.from_string(
dataset, default_project=bigquery_client.project
)
else:
dataset_ref = session._anonymous_dataset
return dataset_ref
def _resolve_cloud_functions_client(
self,
session: Session,
cloud_functions_client: Optional[functions_v2.FunctionServiceClient],
) -> Optional[functions_v2.FunctionServiceClient]:
"""Resolves the Cloud Functions client."""
if not cloud_functions_client:
cloud_functions_client = session.cloudfunctionsclient
if not cloud_functions_client:
raise bf_formatting.create_exception_with_feedback_link(
ValueError,
"A cloud functions client must be provided, either directly "
"or via session.",
)
return cloud_functions_client
def _resolve_bigquery_connection_id(
self,
session: Session,
dataset_ref: bigquery.DatasetReference,
bq_location: str,
bigquery_connection: Optional[str] = None,
) -> str:
"""Resolves BigQuery connection id."""
if not bigquery_connection:
bigquery_connection = session.bq_connection # type: ignore
bigquery_connection = clients.get_canonical_bq_connection_id(
bigquery_connection,
default_project=dataset_ref.project,
default_location=bq_location,
)
# Guaranteed to be the form of <project>.<location>.<connection_id>
(
gcp_project_id,
bq_connection_location,
bq_connection_id,
) = bigquery_connection.split(".")
if gcp_project_id.casefold() != dataset_ref.project.casefold():
raise bf_formatting.create_exception_with_feedback_link(
ValueError,
"The project_id does not match BigQuery connection "
f"gcp_project_id: {dataset_ref.project}.",
)
if bq_connection_location.casefold() != bq_location.casefold():
raise bf_formatting.create_exception_with_feedback_link(
ValueError,
"The location does not match BigQuery connection location: "
f"{bq_location}.",
)
return bq_connection_id
def _update_temp_artifacts(self, bqrf_routine: str, gcf_path: str):
"""Update function artifacts in the current session."""
with self._artifacts_lock:
self._temp_artifacts[bqrf_routine] = gcf_path
def clean_up(
self,
bqclient: bigquery.Client,
gcfclient: functions_v2.FunctionServiceClient,
session_id: str,
):
"""Delete function artifacts in the current session."""
with self._artifacts_lock:
for bqrf_routine, gcf_path in self._temp_artifacts.items():
# Let's accept the possibility that the function may have been
# deleted directly by the user
bqclient.delete_routine(bqrf_routine, not_found_ok=True)
if gcf_path:
# Let's accept the possibility that the cloud function may
# have been deleted directly by the user
try:
gcfclient.delete_function(name=gcf_path)
except google.api_core.exceptions.NotFound:
pass
self._temp_artifacts.clear()
# Inspired by @udf decorator implemented in ibis-bigquery package
# https://github.com/ibis-project/ibis-bigquery/blob/main/ibis_bigquery/udf/__init__.py
# which has moved as @js to the ibis package
# https://github.com/ibis-project/ibis/blob/master/ibis/backends/bigquery/udf/__init__.py
def remote_function(
self,
*,
input_types: Union[None, type, Sequence[type]] = None,
output_type: Optional[type] = None,
session: Optional[Session] = None,
bigquery_client: Optional[bigquery.Client] = None,
bigquery_connection_client: Optional[
bigquery_connection_v1.ConnectionServiceClient
] = None,
cloud_functions_client: Optional[functions_v2.FunctionServiceClient] = None,
resource_manager_client: Optional[resourcemanager_v3.ProjectsClient] = None,
dataset: Optional[str] = None,
bigquery_connection: Optional[str] = None,
reuse: bool = True,
name: Optional[str] = None,
packages: Optional[Sequence[str]] = None,
cloud_function_service_account: str,
cloud_function_kms_key_name: Optional[str] = None,
cloud_function_docker_repository: Optional[str] = None,
max_batching_rows: Optional[int] = None,
cloud_function_timeout: Optional[int] = 600,
cloud_function_max_instances: Optional[int] = None,
cloud_function_vpc_connector: Optional[str] = None,
cloud_function_vpc_connector_egress_settings: Optional[
Literal["all", "private-ranges-only", "unspecified"]
] = None,
cloud_function_memory_mib: Optional[int] = None,
cloud_function_cpus: Optional[float] = None,
cloud_function_ingress_settings: Literal[
"all", "internal-only", "internal-and-gclb"
] = "internal-only",
cloud_build_service_account: Optional[str] = None,
):
"""Decorator to turn a user defined function into a BigQuery remote function.
.. deprecated:: 0.0.1
This is an internal method. Please use :func:`bigframes.pandas.remote_function` instead.
.. warning::
To use remote functions with Bigframes 2.0 and onwards, please (preferred)
set an explicit user-managed ``cloud_function_service_account`` or (discouraged)
set ``cloud_function_service_account`` to use the Compute Engine service account
by setting it to `"default"`.
See, https://cloud.google.com/functions/docs/securing/function-identity.
.. note::
Please make sure following is setup before using this API:
1. Have the below APIs enabled for your project:
* BigQuery Connection API
* Cloud Functions API
* Cloud Run API
* Cloud Build API
* Artifact Registry API
* Cloud Resource Manager API
This can be done from the cloud console (change `PROJECT_ID` to yours):
https://console.cloud.google.com/apis/enableflow?apiid=bigqueryconnection.googleapis.com,cloudfunctions.googleapis.com,run.googleapis.com,cloudbuild.googleapis.com,artifactregistry.googleapis.com,cloudresourcemanager.googleapis.com&project=PROJECT_ID
Or from the gcloud CLI:
`$ gcloud services enable bigqueryconnection.googleapis.com cloudfunctions.googleapis.com run.googleapis.com cloudbuild.googleapis.com artifactregistry.googleapis.com cloudresourcemanager.googleapis.com`
2. Have following IAM roles enabled for you:
* BigQuery Data Editor (roles/bigquery.dataEditor)
* BigQuery Connection Admin (roles/bigquery.connectionAdmin)
* Cloud Functions Developer (roles/cloudfunctions.developer)
* Service Account User (roles/iam.serviceAccountUser) on the service account `PROJECT_NUMBER-compute@developer.gserviceaccount.com`
* Storage Object Viewer (roles/storage.objectViewer)
* Project IAM Admin (roles/resourcemanager.projectIamAdmin) (Only required if the bigquery connection being used is not pre-created and is created dynamically with user credentials.)
3. Either the user has setIamPolicy privilege on the project, or a BigQuery connection is pre-created with necessary IAM role set:
1. To create a connection, follow https://cloud.google.com/bigquery/docs/reference/standard-sql/remote-functions#create_a_connection
2. To set up IAM, follow https://cloud.google.com/bigquery/docs/reference/standard-sql/remote-functions#grant_permission_on_function
Alternatively, the IAM could also be setup via the gcloud CLI:
`$ gcloud projects add-iam-policy-binding PROJECT_ID --member="serviceAccount:CONNECTION_SERVICE_ACCOUNT_ID" --role="roles/run.invoker"`.
Args:
input_types (type or sequence(type), Optional):
For scalar user defined function it should be the input type or
sequence of input types. The supported scalar input types are
`bool`, `bytes`, `float`, `int`, `str`. For row processing user
defined function (i.e. functions that receive a single input
representing a row in form of a Series), type `Series` should be
specified.
output_type (type, Optional):
Data type of the output in the user defined function. If the
user defined function returns an array, then `list[type]` should
be specified. The supported output types are `bool`, `bytes`,
`float`, `int`, `str`, `list[bool]`, `list[float]`, `list[int]`
and `list[str]`.
session (bigframes.Session, Optional):
BigQuery DataFrames session to use for getting default project,
dataset and BigQuery connection.
bigquery_client (google.cloud.bigquery.Client, Optional):
Client to use for BigQuery operations. If this param is not provided
then bigquery client from the session would be used.
bigquery_connection_client (google.cloud.bigquery_connection_v1.ConnectionServiceClient, Optional):
Client to use for BigQuery connection operations. If this param is
not provided then bigquery connection client from the session would
be used.
cloud_functions_client (google.cloud.functions_v2.FunctionServiceClient, Optional):
Client to use for cloud functions operations. If this param is not
provided then the functions client from the session would be used.
resource_manager_client (google.cloud.resourcemanager_v3.ProjectsClient, Optional):
Client to use for cloud resource management operations, e.g. for
getting and setting IAM roles on cloud resources. If this param is
not provided then resource manager client from the session would be
used.
dataset (str, Optional):
Dataset in which to create a BigQuery remote function. It should be in
`<project_id>.<dataset_name>` or `<dataset_name>` format. If this
parameter is not provided then session dataset id is used.
bigquery_connection (str, Optional):
Name of the BigQuery connection in the form of `CONNECTION_ID` or
`LOCATION.CONNECTION_ID` or `PROJECT_ID.LOCATION.CONNECTION_ID`.
If this param is not provided then the bigquery connection from the session
would be used. If it is pre created in the same location as the
`bigquery_client.location` then it would be used, otherwise it is created
dynamically using the `bigquery_connection_client` assuming the user has necessary
priviliges. The PROJECT_ID should be the same as the BigQuery connection project.
reuse (bool, Optional):
Reuse the remote function if already exists.
`True` by default, which will result in reusing an existing remote
function and corresponding cloud function that was previously
created (if any) for the same udf.
Please note that for an unnamed (i.e. created without an explicit
`name` argument) remote function, the BigQuery DataFrames
session id is attached in the cloud artifacts names. So for the
effective reuse across the sessions it is recommended to create
the remote function with an explicit `name`.
Setting it to `False` would force creating a unique remote function.
If the required remote function does not exist then it would be
created irrespective of this param.
name (str, Optional):
Explicit name of the persisted BigQuery remote function. Use it with
caution, because two users working in the same project and dataset
could overwrite each other's remote functions if they use the same
persistent name. When an explicit name is provided, any session
specific clean up (``bigframes.session.Session.close``/
``bigframes.pandas.close_session``/
``bigframes.pandas.reset_session``/
``bigframes.pandas.clean_up_by_session_id``) does not clean up
the function, and leaves it for the user to manage the function
and the associated cloud function directly.
packages (str[], Optional):
Explicit name of the external package dependencies. Each dependency
is added to the `requirements.txt` as is, and can be of the form
supported in https://pip.pypa.io/en/stable/reference/requirements-file-format/.
cloud_function_service_account (str):
Service account to use for the cloud functions. If "default" provided then
the default service account would be used. See
https://cloud.google.com/functions/docs/securing/function-identity
for more details. Please make sure the service account has the
necessary IAM permissions configured as described in
https://cloud.google.com/functions/docs/reference/iam/roles#additional-configuration.
cloud_function_kms_key_name (str, Optional):
Customer managed encryption key to protect cloud functions and
related data at rest. This is of the format
projects/PROJECT_ID/locations/LOCATION/keyRings/KEYRING/cryptoKeys/KEY.
Read https://cloud.google.com/functions/docs/securing/cmek for
more details including granting necessary service accounts
access to the key.
cloud_function_docker_repository (str, Optional):
Docker repository created with the same encryption key as
`cloud_function_kms_key_name` to store encrypted artifacts
created to support the cloud function. This is of the format
projects/PROJECT_ID/locations/LOCATION/repositories/REPOSITORY_NAME.
For more details see
https://cloud.google.com/functions/docs/securing/cmek#before_you_begin.
max_batching_rows (int, Optional):
The maximum number of rows to be batched for processing in the
BQ remote function. Default value is 1000. A lower number can be
passed to avoid timeouts in case the user code is too complex to
process large number of rows fast enough. A higher number can be
used to increase throughput in case the user code is fast enough.
`None` can be passed to let BQ remote functions service apply
default batching. See for more details
https://cloud.google.com/bigquery/docs/remote-functions#limiting_number_of_rows_in_a_batch_request.
cloud_function_timeout (int, Optional):
The maximum amount of time (in seconds) BigQuery should wait for
the cloud function to return a response. See for more details
https://cloud.google.com/functions/docs/configuring/timeout.
Please note that even though the cloud function (2nd gen) itself
allows seeting up to 60 minutes of timeout, BigQuery remote
function can wait only up to 20 minutes, see for more details
https://cloud.google.com/bigquery/quotas#remote_function_limits.
By default BigQuery DataFrames uses a 10 minute timeout. `None`
can be passed to let the cloud functions default timeout take effect.
cloud_function_max_instances (int, Optional):
The maximumm instance count for the cloud function created. This
can be used to control how many cloud function instances can be
active at max at any given point of time. Lower setting can help
control the spike in the billing. Higher setting can help
support processing larger scale data. When not specified, cloud
function's default setting applies. For more details see
https://cloud.google.com/functions/docs/configuring/max-instances.
cloud_function_vpc_connector (str, Optional):
The VPC connector you would like to configure for your cloud
function. This is useful if your code needs access to data or
service(s) that are on a VPC network. See for more details
https://cloud.google.com/functions/docs/networking/connecting-vpc.
cloud_function_vpc_connector_egress_settings (str, Optional):
Egress settings for the VPC connector, controlling what outbound
traffic is routed through the VPC connector.
Options are: `all`, `private-ranges-only`, or `unspecified`.
If not specified, `private-ranges-only` is used by default.
See for more details
https://cloud.google.com/run/docs/configuring/vpc-connectors#egress-job.
cloud_function_memory_mib (int, Optional):
The amounts of memory (in mebibytes) to allocate for the cloud
function (2nd gen) created. This also dictates a corresponding
amount of allocated CPU for the function. By default a memory of
1024 MiB is set for the cloud functions created to support
BigQuery DataFrames remote function. If you want to let the
default memory of cloud functions be allocated, pass `None`. See
for more details
https://cloud.google.com/functions/docs/configuring/memory.
cloud_function_cpus (float, Optional):
The number of cpus to allocate for the cloud
function (2nd gen) created.
https://docs.cloud.google.com/run/docs/configuring/services/cpu.
cloud_function_ingress_settings (str, Optional):
Ingress settings controls dictating what traffic can reach the
function. Options are: `all`, `internal-only`, or `internal-and-gclb`.
If no setting is provided, `internal-only` will be used by default.
See for more details
https://cloud.google.com/functions/docs/networking/network-settings#ingress_settings.
cloud_build_service_account (str, Optional):
Service account in the fully qualified format
`projects/PROJECT_ID/serviceAccounts/SERVICE_ACCOUNT_EMAIL`, or
just the SERVICE_ACCOUNT_EMAIL. The latter would be interpreted
as belonging to the BigQuery DataFrames session project. This is
to be used by Cloud Build to build the function source code into
a deployable artifact. If not provided, the default Cloud Build
service account is used. See
https://cloud.google.com/build/docs/cloud-build-service-account
for more details.
"""
# Some defaults may be used from the session if not provided otherwise.
session = self._resolve_session(session)
# If the user forces the cloud function service argument to None, throw
# an exception
if cloud_function_service_account is None:
raise ValueError(
'You must provide a user managed cloud_function_service_account, or "default" if you would like to let the default service account be used.'
)
# A BigQuery client is required to perform BQ operations.
bigquery_client = self._resolve_bigquery_client(session, bigquery_client)
# A BigQuery connection client is required for BQ connection operations.
bigquery_connection_client = self._resolve_bigquery_connection_client(
session, bigquery_connection_client
)
# A resource manager client is required to get/set IAM operations.
resource_manager_client = self._resolve_resource_manager_client(
session, resource_manager_client
)
# BQ remote function must be persisted, for which we need a dataset.
# https://cloud.google.com/bigquery/docs/reference/standard-sql/remote-functions#:~:text=You%20cannot%20create%20temporary%20remote%20functions.
dataset_ref = self._resolve_dataset_reference(session, bigquery_client, dataset)
# A cloud functions client is required for cloud functions operations.
cloud_functions_client = self._resolve_cloud_functions_client(
session, cloud_functions_client
)
bq_location, cloud_function_region = _utils.get_remote_function_locations(
bigquery_client.location
)
# A connection is required for BQ remote function.
# https://cloud.google.com/bigquery/docs/reference/standard-sql/remote-functions#create_a_remote_function
bq_connection_id = self._resolve_bigquery_connection_id(
session, dataset_ref, bq_location, bigquery_connection
)
# If any CMEK is intended then check that a docker repository is also specified.
if (
cloud_function_kms_key_name is not None
and cloud_function_docker_repository is None
):
raise bf_formatting.create_exception_with_feedback_link(
ValueError,
"cloud_function_docker_repository must be specified with cloud_function_kms_key_name."
" For more details see https://cloud.google.com/functions/docs/securing/cmek#before_you_begin.",
)
# A VPC connector is required to specify VPC egress settings.
if (
cloud_function_vpc_connector_egress_settings is not None
and cloud_function_vpc_connector is None
):
raise bf_formatting.create_exception_with_feedback_link(
ValueError,
"cloud_function_vpc_connector must be specified before cloud_function_vpc_connector_egress_settings.",
)
if cloud_function_ingress_settings is None:
cloud_function_ingress_settings = "internal-only"
msg = bfe.format_message(
"The `cloud_function_ingress_settings` is being set to 'internal-only' by default."
)
warnings.warn(msg, category=UserWarning, stacklevel=2)
bq_connection_manager = session.bqconnectionmanager
def wrapper(func):
nonlocal input_types, output_type
if not callable(func):
raise bf_formatting.create_exception_with_feedback_link(
TypeError, f"func must be a callable, got {func}"
)
if sys.version_info >= (3, 10):
# Add `eval_str = True` so that deferred annotations are turned into their
# corresponding type objects. Need Python 3.10 for eval_str parameter.
# https://docs.python.org/3/library/inspect.html#inspect.signature
signature_kwargs: Mapping[str, Any] = {"eval_str": True}
else:
signature_kwargs = {} # type: ignore
py_sig = inspect.signature(
func,
**signature_kwargs,
)
py_sig = _resolve_signature(py_sig, input_types, output_type)
remote_function_client = _function_client.FunctionClient(
dataset_ref.project,
bq_location,
dataset_ref.dataset_id,
bigquery_client,
bq_connection_manager,
cloud_function_region,
cloud_functions_client,
None
if cloud_function_service_account == "default"
else cloud_function_service_account,
cloud_function_kms_key_name,
cloud_function_docker_repository,
cloud_build_service_account=cloud_build_service_account,
session=session, # type: ignore
)
udf_sig = udf_def.UdfSignature.from_py_signature(
py_sig
).to_remote_function_compatible()
(
rf_name,
cf_name,
created_new,
) = remote_function_client.provision_bq_remote_function(
func,
func_signature=udf_sig,
reuse=reuse or False,
name=name,
package_requirements=tuple(packages) if packages else tuple(),
max_batching_rows=max_batching_rows or 1000,
cloud_function_timeout=cloud_function_timeout,
cloud_function_max_instance_count=cloud_function_max_instances,
cloud_function_vpc_connector=cloud_function_vpc_connector,
cloud_function_vpc_connector_egress_settings=cloud_function_vpc_connector_egress_settings,
cloud_function_memory_mib=cloud_function_memory_mib,
cloud_function_cpus=cloud_function_cpus,
cloud_function_ingress_settings=cloud_function_ingress_settings,
bq_connection_id=bq_connection_id,
)
bigframes_cloud_function = (
remote_function_client.get_cloud_function_fully_qualified_name(cf_name)
)
bigframes_bigquery_function = (
remote_function_client.get_remote_function_fully_qualilfied_name(
rf_name
)
)
# If a new remote function was created, update the cloud artifacts
# created in the session. This would be used to clean up any
# resources in the session. Note that we need to do this only for
# the case where an explicit name was not provided by the user and
# we used an internal name. For the cases where the user provided an
# explicit name, we are assuming that the user wants to persist them
# with that name and would directly manage their lifecycle.
if created_new and (not name):
self._update_temp_artifacts(
bigframes_bigquery_function, bigframes_cloud_function
)
udf_definition = udf_def.BigqueryUdf(
routine_ref=bigquery.RoutineReference.from_string(
bigframes_bigquery_function
),
signature=udf_sig,
)
decorator = functools.wraps(func)
if udf_sig.is_row_processor:
msg = bfe.format_message("input_types=Series is in preview.")
warnings.warn(msg, stacklevel=1, category=bfe.PreviewWarning)
return decorator(
bq_functions.BigqueryCallableRoutine(
udf_definition,
session,
cloud_function_ref=bigframes_cloud_function,
local_func=func,
is_managed=False,
)
)
return wrapper
def deploy_remote_function(
self,
func,
**kwargs,
):
"""Orchestrates the creation of a BigQuery remote function that deploys immediately.
This method ensures that the remote function is created and available for
use in BigQuery as soon as this call is made.
Args:
kwargs:
All arguments are passed directly to
:meth:`~bigframes.session.Session.remote_function`. Please see
its docstring for parameter details.
Returns:
A wrapped remote function, usable in
:meth:`~bigframes.series.Series.apply`.
"""
# TODO(tswast): If we update remote_function to defer deployment, update
# this method to deploy immediately.
return self.remote_function(**kwargs)(func)
def udf(
self,
input_types: Union[None, type, Sequence[type]] = None,
output_type: Optional[type] = None,
session: Optional[Session] = None,
bigquery_client: Optional[bigquery.Client] = None,
dataset: Optional[str] = None,
bigquery_connection: Optional[str] = None,
name: Optional[str] = None,
packages: Optional[Sequence[str]] = None,
max_batching_rows: Optional[int] = None,
container_cpu: Optional[float] = None,
container_memory: Optional[str] = None,
):
"""Decorator to turn a Python user defined function (udf) into a
BigQuery managed function.
.. note::
This feature is in preview. The code in the udf must be
(1) self-contained, i.e. it must not contain any
references to an import or variable defined outside the function
body, and
(2) Python 3.11 compatible, as that is the environment
in which the code is executed in the cloud.
.. note::
Please have following IAM roles enabled for you:
* BigQuery Data Editor (roles/bigquery.dataEditor)
Args:
input_types (type or sequence(type), Optional):
For scalar user defined function it should be the input type or
sequence of input types. The supported scalar input types are
`bool`, `bytes`, `float`, `int`, `str`.
output_type (type, Optional):
Data type of the output in the user defined function. If the
user defined function returns an array, then `list[type]` should
be specified. The supported output types are `bool`, `bytes`,
`float`, `int`, `str`, `list[bool]`, `list[float]`, `list[int]`
and `list[str]`.
session (bigframes.Session, Optional):
BigQuery DataFrames session to use for getting default project,
dataset and BigQuery connection.
bigquery_client (google.cloud.bigquery.Client, Optional):
Client to use for BigQuery operations. If this param is not
provided, then bigquery client from the session would be used.
dataset (str, Optional):
Dataset in which to create a BigQuery managed function. It
should be in `<project_id>.<dataset_name>` or `<dataset_name>`
format. If this parameter is not provided then session dataset
id is used.
bigquery_connection (str, Optional):
Name of the BigQuery connection. It is used to provide an
identity to the serverless instances running the user code. It
helps BigQuery manage and track the resources used by the udf.
This connection is required for internet access and for
interacting with other GCP services. To access GCP services, the
appropriate IAM permissions must also be granted to the
connection's Service Account. When it defaults to None, the udf
will be created without any connection. A udf without a
connection has no internet access and no access to other GCP
services.
name (str, Optional):
Explicit name of the persisted BigQuery managed function. Use it
with caution, because more than one users working in the same
project and dataset could overwrite each other's managed
functions if they use the same persistent name. When an explicit
name is provided, any session specific clean up (
``bigframes.session.Session.close``/
``bigframes.pandas.close_session``/
``bigframes.pandas.reset_session``/
``bigframes.pandas.clean_up_by_session_id``) does not clean up
the function, and leaves it for the user to manage the function
directly.
packages (str[], Optional):
Explicit name of the external package dependencies. Each
dependency is added to the `requirements.txt` as is, and can be
of the form supported in
https://pip.pypa.io/en/stable/reference/requirements-file-format/.
max_batching_rows (int, Optional):
The maximum number of rows in each batch. If you specify
max_batching_rows, BigQuery determines the number of rows in a
batch, up to the max_batching_rows limit. If max_batching_rows
is not specified, the number of rows to batch is determined
automatically.
container_cpu (float, Optional):
The CPU limits for containers that run Python UDFs. By default,
the CPU allocated is 0.33 vCPU. See details at
https://cloud.google.com/bigquery/docs/user-defined-functions-python#configure-container-limits.
container_memory (str, Optional):
The memory limits for containers that run Python UDFs. By
default, the memory allocated to each container instance is
512 MiB. See details at
https://cloud.google.com/bigquery/docs/user-defined-functions-python#configure-container-limits.
"""
warnings.warn("udf is in preview.", category=bfe.PreviewWarning, stacklevel=5)
# Some defaults may be used from the session if not provided otherwise.
session = self._resolve_session(session)
# A BigQuery client is required to perform BQ operations.
bigquery_client = self._resolve_bigquery_client(session, bigquery_client)
# BQ managed function must be persisted, for which we need a dataset.
dataset_ref = self._resolve_dataset_reference(session, bigquery_client, dataset)
bq_location, _ = _utils.get_remote_function_locations(bigquery_client.location)
# A connection is optional for BQ managed function.
bq_connection_id = (
self._resolve_bigquery_connection_id(
session, dataset_ref, bq_location, bigquery_connection
)
if bigquery_connection
else None
)
bq_connection_manager = session.bqconnectionmanager
# TODO(b/399129906): Write a method for the repeated part in the wrapper
# for both managed function and remote function.
def wrapper(func):
nonlocal input_types, output_type
if not callable(func):
raise bf_formatting.create_exception_with_feedback_link(
TypeError, f"func must be a callable, got {func}"
)
if sys.version_info >= (3, 10):
# Add `eval_str = True` so that deferred annotations are turned into their
# corresponding type objects. Need Python 3.10 for eval_str parameter.
# https://docs.python.org/3/library/inspect.html#inspect.signature
signature_kwargs: Mapping[str, Any] = {"eval_str": True}
else:
signature_kwargs = {} # type: ignore
py_sig = inspect.signature(
func,
**signature_kwargs,
)
py_sig = _resolve_signature(py_sig, input_types, output_type)
# The function will actually be receiving a pandas Series, but allow
# both BigQuery DataFrames and pandas object types for compatibility.
udf_sig = udf_def.UdfSignature.from_py_signature(py_sig)
managed_function_client = _function_client.FunctionClient(
dataset_ref.project,
bq_location,
dataset_ref.dataset_id,
bigquery_client,
bq_connection_manager,
session=session, # type: ignore
)
code_def = udf_def.CodeDef.from_func(func)
config = udf_def.ManagedFunctionConfig(
code=code_def,
signature=udf_sig,
max_batching_rows=max_batching_rows,
container_cpu=container_cpu,
container_memory=container_memory,
bq_connection_id=bq_connection_id,
capture_references=False,
)
bq_function_name = managed_function_client.provision_bq_managed_function(
name=name,
config=config,
)
full_rf_name = (
managed_function_client.get_remote_function_fully_qualilfied_name(
bq_function_name
)
)
udf_definition = udf_def.BigqueryUdf(
routine_ref=bigquery.RoutineReference.from_string(full_rf_name),
signature=udf_sig,
)
if udf_sig.is_row_processor:
msg = bfe.format_message("input_types=Series is in preview.")
warnings.warn(msg, stacklevel=1, category=bfe.PreviewWarning)
if not name: # session-managed resource
self._update_temp_artifacts(full_rf_name, "")
return bq_functions.UdfRoutine(func=func, _udf_def=udf_definition)
# user-managed permanent resource
else:
return bq_functions.BigqueryCallableRoutine(
udf_definition, session, local_func=func, is_managed=True
)
return wrapper
def deploy_udf(
self,
func,
**kwargs,
):
"""Orchestrates the creation of a BigQuery UDF that deploys immediately.
This method ensures that the UDF is created and available for
use in BigQuery as soon as this call is made.
Args:
func:
Function to deploy.
kwargs:
All arguments are passed directly to
:meth:`~bigframes.session.Session.udf`. Please see
its docstring for parameter details.
Returns:
A wrapped Python user defined function, usable in
:meth:`~bigframes.series.Series.apply`.
"""
# TODO(tswast): If we update udf to defer deployment, update this method
# to deploy immediately.
return self.udf(**kwargs)(func)
def _resolve_signature(
py_sig: inspect.Signature,
input_types: Union[None, type, Sequence[type]] = None,
output_type: Optional[type] = None,
) -> inspect.Signature:
if input_types is not None:
if not isinstance(input_types, collections.abc.Sequence):
input_types = [input_types]
if _utils.has_conflict_input_type(py_sig, input_types):
msg = bfe.format_message(
"Conflicting input types detected, using the one from the decorator."
)
warnings.warn(msg, category=bfe.FunctionConflictTypeHintWarning)
py_sig = py_sig.replace(
parameters=[
par.replace(annotation=itype)
for par, itype in zip(py_sig.parameters.values(), input_types)
]
)
if output_type:
if _utils.has_conflict_output_type(py_sig, output_type):
msg = bfe.format_message(
"Conflicting return type detected, using the one from the decorator."
)
warnings.warn(msg, category=bfe.FunctionConflictTypeHintWarning)
py_sig = py_sig.replace(return_annotation=output_type)
return py_sig