refactor(bigquery): explicite disallow INT64 in JS UDF

Author: krzysztof-kwitt

ibis/backends/bigquery/datatypes.py

@@ -4,24 +4,6 @@
 
 import ibis.expr.datatypes as dt
 
-
-class TypeTranslationContext:
-    """A tag class to alter the way a type is translated.
-
-    This is used to raise an exception when INT64 types are encountered to
-    avoid suprising results due to BigQuery's handling of INT64 types in
-    JavaScript UDFs.
-    """
-
-    __slots__ = ()
-
-
-class UDFContext(TypeTranslationContext):
-    __slots__ = ()
-
-
-UDF_CONTEXT = UDFContext()
-
 ibis_type_to_bigquery_type = Dispatcher("ibis_type_to_bigquery_type")
 
 
@@ -30,81 +12,60 @@ def trans_string_default(datatype):
     return ibis_type_to_bigquery_type(dt.dtype(datatype))
 
 
-@ibis_type_to_bigquery_type.register(dt.DataType)
-def trans_default(t):
-    return ibis_type_to_bigquery_type(t, TypeTranslationContext())
-
-
-@ibis_type_to_bigquery_type.register(str, TypeTranslationContext)
-def trans_string_context(datatype, context):
-    return ibis_type_to_bigquery_type(dt.dtype(datatype), context)
-
-
-@ibis_type_to_bigquery_type.register(dt.Floating, TypeTranslationContext)
-def trans_float64(t, context):
+@ibis_type_to_bigquery_type.register(dt.Floating)
+def trans_float64(t):
     return "FLOAT64"
 
 
-@ibis_type_to_bigquery_type.register(dt.Integer, TypeTranslationContext)
-def trans_integer(t, context):
+@ibis_type_to_bigquery_type.register(dt.Integer)
+def trans_integer(t):
     return "INT64"
 
 
-@ibis_type_to_bigquery_type.register(dt.Binary, TypeTranslationContext)
-def trans_binary(t, context):
+@ibis_type_to_bigquery_type.register(dt.Binary)
+def trans_binary(t):
     return "BYTES"
 
 
-@ibis_type_to_bigquery_type.register(dt.UInt64, (TypeTranslationContext, UDFContext))
-def trans_lossy_integer(t, context):
+@ibis_type_to_bigquery_type.register(dt.UInt64)
+def trans_lossy_integer(t):
     raise TypeError("Conversion from uint64 to BigQuery integer type (int64) is lossy")
 
 
-@ibis_type_to_bigquery_type.register(dt.Array, TypeTranslationContext)
-def trans_array(t, context):
-    return f"ARRAY<{ibis_type_to_bigquery_type(t.value_type, context)}>"
+@ibis_type_to_bigquery_type.register(dt.Array)
+def trans_array(t):
+    return f"ARRAY<{ibis_type_to_bigquery_type(t.value_type)}>"
 
 
-@ibis_type_to_bigquery_type.register(dt.Struct, TypeTranslationContext)
-def trans_struct(t, context):
+@ibis_type_to_bigquery_type.register(dt.Struct)
+def trans_struct(t):
     return "STRUCT<{}>".format(
         ", ".join(
-            f"{name} {ibis_type_to_bigquery_type(dt.dtype(type), context)}"
-            for name, type in zip(t.names, t.types)
+            f"{name} {ibis_type_to_bigquery_type(dt.dtype(type_))}"
+            for name, type_ in zip(t.names, t.types)
         )
     )
 
 
-@ibis_type_to_bigquery_type.register(dt.Date, TypeTranslationContext)
-def trans_date(t, context):
+@ibis_type_to_bigquery_type.register(dt.Date)
+def trans_date(t):
     return "DATE"
 
 
-@ibis_type_to_bigquery_type.register(dt.Timestamp, TypeTranslationContext)
-def trans_timestamp(t, context):
+@ibis_type_to_bigquery_type.register(dt.Timestamp)
+def trans_timestamp(t):
     if t.timezone is not None:
         raise TypeError("BigQuery does not support timestamps with timezones")
     return "TIMESTAMP"
 
 
-@ibis_type_to_bigquery_type.register(dt.DataType, TypeTranslationContext)
-def trans_type(t, context):
+@ibis_type_to_bigquery_type.register(dt.DataType)
+def trans_type(t):
     return str(t).upper()
 
 
-@ibis_type_to_bigquery_type.register(dt.Integer, UDFContext)
-def trans_integer_udf(t, context):
-    # JavaScript does not have integers, only a Number class. BigQuery doesn't
-    # behave as expected with INT64 inputs or outputs
-    raise TypeError(
-        "BigQuery does not support INT64 as an argument type or a return type "
-        "for UDFs. Replace INT64 with FLOAT64 in your UDF signature and "
-        "cast all INT64 inputs to FLOAT64."
-    )
-
-
-@ibis_type_to_bigquery_type.register(dt.Decimal, TypeTranslationContext)
-def trans_numeric(t, context):
+@ibis_type_to_bigquery_type.register(dt.Decimal)
+def trans_numeric(t):
     if (t.precision, t.scale) != (38, 9):
         raise TypeError(
             "BigQuery only supports decimal types with precision of 38 and "
@@ -113,11 +74,22 @@ def trans_numeric(t, context):
     return "NUMERIC"
 
 
-@ibis_type_to_bigquery_type.register(dt.Decimal, UDFContext)
-def trans_numeric_udf(t, context):
-    raise TypeError("Decimal types are not supported in BigQuery UDFs")
+@ibis_type_to_bigquery_type.register(dt.JSON)
+def trans_json(t):
+    return "JSON"
 
 
-@ibis_type_to_bigquery_type.register(dt.JSON, TypeTranslationContext)
-def trans_json(t, context):
-    return "JSON"
+def spread_type(dt: dt.DataType):
+    """Returns a generator that contains all the types in the given type.
+
+    For complex types like set and array, it returns the types of the elements.
+    """
+    if dt.is_array():
+        yield from spread_type(dt.value_type)
+    elif dt.is_struct():
+        for type_ in dt.types:
+            yield from spread_type(type_)
+    elif dt.is_map():
+        yield from spread_type(dt.key_type)
+        yield from spread_type(dt.value_type)
+    yield dt

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

@@ -4,8 +4,8 @@
 
 import ibis.expr.datatypes as dt
 from ibis.backends.bigquery.datatypes import (
-    UDF_CONTEXT,
     ibis_type_to_bigquery_type,
+    spread_type,
 )
 
 
@@ -63,20 +63,25 @@ def test_simple_failure_mode(datatype):
 
 
 @pytest.mark.parametrize(
-    ("type", "expected"),
+    ("type_", "expected"),
     [
-        param(dt.int64, "INT64", marks=pytest.mark.xfail(raises=TypeError)),
+        param(
+            dt.int64,
+            [dt.int64],
+        ),
         param(
             dt.Array(dt.int64),
-            "ARRAY<INT64>",
-            marks=pytest.mark.xfail(raises=TypeError),
+            [dt.int64, dt.Array(value_type=dt.int64)],
         ),
         param(
             dt.Struct.from_tuples([("a", dt.Array(dt.int64))]),
-            "STRUCT<a ARRAY<INT64>>",
-            marks=pytest.mark.xfail(raises=TypeError),
+            [
+                dt.int64,
+                dt.Array(value_type=dt.int64),
+                dt.Struct.from_tuples([('a', dt.Array(value_type=dt.int64))]),
+            ],
         ),
     ],
 )
-def test_ibis_type_to_bigquery_type_udf(type, expected):
-    assert ibis_type_to_bigquery_type(type, UDF_CONTEXT) == expected
+def test_spread_type(type_, expected):
+    assert list(spread_type(type_)) == expected

ibis/backends/bigquery/udf/__init__.py

@@ -7,7 +7,7 @@
 
 import ibis.expr.datatypes as dt
 import ibis.expr.rules as rlz
-from ibis.backends.bigquery.datatypes import UDF_CONTEXT, ibis_type_to_bigquery_type
+from ibis.backends.bigquery.datatypes import ibis_type_to_bigquery_type, spread_type
 from ibis.backends.bigquery.operations import BigQueryUDFNode
 from ibis.backends.bigquery.udf.core import PythonToJavaScriptTranslator
 from ibis.udf.validate import validate_output_type
@@ -254,6 +254,16 @@ def js(
         """;
         '''
         validate_output_type(output_type)
+        if any(
+            type_ == dt.int64
+            for param_type in params.values()
+            for type_ in spread_type(param_type)
+        ) or any(type_ == dt.int64 for type_ in spread_type(output_type)):
+            raise TypeError(
+                "BigQuery does not support INT64 as an argument type or a return type "
+                "for UDFs. Replace INT64 with FLOAT64 in your UDF signature and "
+                "cast all INT64 inputs to FLOAT64."
+            )
 
         if libraries is None:
             libraries = []
@@ -276,11 +286,11 @@ def compiles_udf_node(t, op):
         bigquery_signature = ", ".join(
             "{name} {type}".format(
                 name=name,
-                type=ibis_type_to_bigquery_type(dt.dtype(type_), UDF_CONTEXT),
+                type=ibis_type_to_bigquery_type(dt.dtype(type_)),
             )
             for name, type_ in params.items()
         )
-        return_type = ibis_type_to_bigquery_type(dt.dtype(output_type), UDF_CONTEXT)
+        return_type = ibis_type_to_bigquery_type(dt.dtype(output_type))
         libraries_opts = (
             f"\nOPTIONS (\n    library={repr(list(libraries))}\n)" if libraries else ""
         )