feat(bigquery): add JS UDF support

Author: krzysztof-kwitt

ibis/backends/bigquery/datatypes.py

@@ -20,6 +20,8 @@ class UDFContext(TypeTranslationContext):
     __slots__ = ()
 
 
+UDF_CONTEXT = UDFContext()
+
 ibis_type_to_bigquery_type = Dispatcher("ibis_type_to_bigquery_type")
 
 

ibis/backends/bigquery/tests/system/udf/test_udf_execute.py

@@ -25,7 +25,11 @@ def df(alltypes):
 
 
 def test_udf(alltypes, df):
-    @udf(input_type=[dt.double, dt.double], output_type=dt.double)
+    @udf(
+        input_type=[dt.double, dt.double],
+        output_type=dt.double,
+        determinism=True,
+    )
     def my_add(a, b):
         return a + b
 

ibis/backends/bigquery/tests/unit/test_datatypes.py

@@ -4,8 +4,7 @@
 
 import ibis.expr.datatypes as dt
 from ibis.backends.bigquery.datatypes import (
-    TypeTranslationContext,
-    UDFContext,
+    UDF_CONTEXT,
     ibis_type_to_bigquery_type,
 )
 
@@ -54,8 +53,7 @@ def test_no_ambiguities():
     ],
 )
 def test_simple(datatype, expected):
-    context = TypeTranslationContext()
-    assert ibis_type_to_bigquery_type(datatype, context) == expected
+    assert ibis_type_to_bigquery_type(datatype) == expected
 
 
 @pytest.mark.parametrize("datatype", [dt.uint64, dt.Decimal(8, 3)])
@@ -81,5 +79,4 @@ def test_simple_failure_mode(datatype):
     ],
 )
 def test_ibis_type_to_bigquery_type_udf(type, expected):
-    context = UDFContext()
-    assert ibis_type_to_bigquery_type(type, context) == expected
+    assert ibis_type_to_bigquery_type(type, UDF_CONTEXT) == expected

ibis/backends/bigquery/tests/unit/udf/snapshots/test_usage/test_multiple_calls_redefinition/out.sql

@@ -0,0 +1,21 @@
+CREATE TEMPORARY FUNCTION my_len_0(s STRING)
+RETURNS FLOAT64
+LANGUAGE js AS """
+'use strict';
+function my_len(s) {
+    return s.length;
+}
+return my_len(s);
+""";
+
+CREATE TEMPORARY FUNCTION my_len_1(s STRING)
+RETURNS FLOAT64
+LANGUAGE js AS """
+'use strict';
+function my_len(s) {
+    return (s.length + 1);
+}
+return my_len(s);
+""";
+
+SELECT (my_len_0('abcd') + my_len_0('abcd')) + my_len_1('abcd') AS `tmp`

ibis/backends/bigquery/tests/unit/udf/snapshots/test_usage/test_udf_determinism/False/out.sql

@@ -0,0 +1,12 @@
+CREATE TEMPORARY FUNCTION my_len_0(s STRING)
+RETURNS FLOAT64
+NOT DETERMINISTIC
+LANGUAGE js AS """
+'use strict';
+function my_len(s) {
+    return s.length;
+}
+return my_len(s);
+""";
+
+SELECT my_len_0('abcd') AS `tmp`

ibis/backends/bigquery/tests/unit/udf/snapshots/test_usage/test_udf_determinism/None/out.sql

@@ -0,0 +1,11 @@
+CREATE TEMPORARY FUNCTION my_len_0(s STRING)
+RETURNS FLOAT64
+LANGUAGE js AS """
+'use strict';
+function my_len(s) {
+    return s.length;
+}
+return my_len(s);
+""";
+
+SELECT my_len_0('abcd') AS `tmp`

ibis/backends/bigquery/tests/unit/udf/snapshots/test_usage/test_udf_determinism/True/out.sql

@@ -0,0 +1,12 @@
+CREATE TEMPORARY FUNCTION my_len_0(s STRING)
+RETURNS FLOAT64
+DETERMINISTIC
+LANGUAGE js AS """
+'use strict';
+function my_len(s) {
+    return s.length;
+}
+return my_len(s);
+""";
+
+SELECT my_len_0('abcd') AS `tmp`

ibis/backends/bigquery/tests/unit/udf/test_usage.py

@@ -4,46 +4,49 @@
 import ibis
 import ibis.expr.datatypes as dt
 from ibis.backends.bigquery import udf
+from ibis.backends.bigquery.udf import _udf_name_cache
 
 
-def test_multiple_calls_redefinition():
-    @udf([dt.string], dt.double)
+def test_multiple_calls_redefinition(snapshot):
+    _udf_name_cache.clear()
+
+    @udf.python([dt.string], dt.double)
     def my_len(s):
         return s.length
 
     s = ibis.literal("abcd")
     expr = my_len(s) + my_len(s)
 
-    @udf([dt.string], dt.double)
+    @udf.python([dt.string], dt.double)
     def my_len(s):
         return s.length + 1
 
     expr = expr + my_len(s)
 
     sql = ibis.bigquery.compile(expr)
-    expected = '''\
-CREATE TEMPORARY FUNCTION my_len_0(s STRING)
-RETURNS FLOAT64
-LANGUAGE js AS """
-'use strict';
-function my_len(s) {
-    return s.length;
-}
-return my_len(s);
-""";
+    snapshot.assert_match(sql, "out.sql")
+
+
+@pytest.mark.parametrize(
+    ("determinism",),
+    [
+        param(True),
+        param(False),
+        param(None),
+    ],
+)
+def test_udf_determinism(snapshot, determinism):
+    _udf_name_cache.clear()
+
+    @udf.python([dt.string], dt.double, determinism=determinism)
+    def my_len(s):
+        return s.length
 
-CREATE TEMPORARY FUNCTION my_len_1(s STRING)
-RETURNS FLOAT64
-LANGUAGE js AS """
-'use strict';
-function my_len(s) {
-    return (s.length + 1);
-}
-return my_len(s);
-""";
+    s = ibis.literal("abcd")
+    expr = my_len(s)
 
-SELECT (my_len_0('abcd') + my_len_0('abcd')) + my_len_1('abcd') AS `tmp`'''
-    assert sql == expected
+    sql = ibis.bigquery.compile(expr)
+    snapshot.assert_match(sql, "out.sql")
 
 
 @pytest.mark.parametrize(
@@ -93,6 +96,6 @@ def my_len(s):
 )
 def test_udf_int64(argument_type, return_type):
     # invalid argument type, valid return type
-    @udf([argument_type], return_type)
+    @udf.python([argument_type], return_type)
     def my_int64_add(x):
         return 1.0