googleapis · shobsi · Sep 6, 2024 · Sep 3, 2024 · Sep 3, 2024 · Sep 3, 2024
@@ -14,6 +14,7 @@

 from __future__ import annotations

+import inspect
 import logging
 from typing import cast, Optional, TYPE_CHECKING
 import warnings
@@ -149,6 +150,13 @@ def func(*ignored_args, **ignored_kwargs):
        expr = node(*ignored_args, **ignored_kwargs)  # type: ignore
        return ibis_client.execute(expr)

+    func.__signature__ = inspect.signature(func).replace(  # type: ignore
+        parameters=[
+            inspect.Parameter(name, inspect.Parameter.POSITIONAL_OR_KEYWORD)
+            for name in ibis_signature.parameter_names
+        ]
+    )
+
    # TODO: Move ibis logic to compiler step

    func.__name__ = routine_ref.routine_id

@@ -1239,12 +1239,22 @@ def read_gbq_function(

        **Examples:**

-        Use the ``cw_lower_case_ascii_only`` function from Community UDFs.
-        (https://github.com/GoogleCloudPlatform/bigquery-utils/blob/master/udfs/community/cw_lower_case_ascii_only.sqlx)
-
            >>> import bigframes.pandas as bpd
            >>> bpd.options.display.progress_bar = None

+        Use the [cw_lower_case_ascii_only](https://github.com/GoogleCloudPlatform/bigquery-utils/blob/master/udfs/community/README.md#cw_lower_case_ascii_onlystr-string)
+        function from Community UDFs.
+
+            >>> func = bpd.read_gbq_function("bqutil.fn.cw_lower_case_ascii_only")
+
+        You can run it on scalar input. Usually you would do so to verify that
+        it works as expected before applying to all values in a Series.
+
+            >>> func('AURÉLIE')
+            'aurÉlie'
+
+        You can apply it to a BigQuery DataFrame Series.
+
            >>> df = bpd.DataFrame({'id': [1, 2, 3], 'name': ['AURÉLIE', 'CÉLESTINE', 'DAPHNÉ']})
            >>> df
               id       name
@@ -1254,7 +1264,6 @@ def read_gbq_function(
            <BLANKLINE>
            [3 rows x 2 columns]

-            >>> func = bpd.read_gbq_function("bqutil.fn.cw_lower_case_ascii_only")
            >>> df1 = df.assign(new_name=df['name'].apply(func))
            >>> df1
               id       name   new_name
@@ -1264,9 +1273,17 @@ def read_gbq_function(
            <BLANKLINE>
            [3 rows x 3 columns]

+        You can even use a function with multiple inputs. For example, let's use
+        [cw_instr4](https://github.com/GoogleCloudPlatform/bigquery-utils/blob/master/udfs/community/README.md#cw_instr4source-string-search-string-position-int64-ocurrence-int64)
+        from Community UDFs.
+
+            >>> func = bpd.read_gbq_function("bqutil.fn.cw_instr4")
+            >>> func('TestStr123456Str', 'Str', 1, 2)
+            14
+
        Args:
            function_name (str):
-                the function's name in BigQuery in the format
+                The function's name in BigQuery in the format
                `project_id.dataset_id.function_name`, or
                `dataset_id.function_name` to load from the default project, or
                `function_name` to load from the default project and the dataset

@@ -671,12 +671,19 @@ def square1(x):


 @pytest.mark.flaky(retries=2, delay=120)
-def test_read_gbq_function_runs_existing_udf(session, bigquery_client, dataset_id):
+def test_read_gbq_function_runs_existing_udf(session):
    func = session.read_gbq_function("bqutil.fn.cw_lower_case_ascii_only")
    got = func("AURÉLIE")
    assert got == "aurÉlie"


+@pytest.mark.flaky(retries=2, delay=120)
+def test_read_gbq_function_runs_existing_udf_4_params(session):
+    func = session.read_gbq_function("bqutil.fn.cw_instr4")
+    got = func("TestStr123456Str", "Str", 1, 2)
+    assert got == 14
+
+
 @pytest.mark.flaky(retries=2, delay=120)
 def test_read_gbq_function_reads_udfs(session, bigquery_client, dataset_id):
    dataset_ref = bigquery.DatasetReference.from_string(dataset_id)