feat(bigquery): implement ops.Typeof

ci/schema/snowflake.sql

@@ -1,4 +1,3 @@
-USE WAREHOUSE ibis_testing;
 DROP DATABASE IF EXISTS ibis_testing;
 CREATE DATABASE IF NOT EXISTS ibis_testing;
 CREATE SCHEMA IF NOT EXISTS ibis_testing.ibis_testing;

ibis/backends/base/sql/registry/literal.py

@@ -91,7 +91,7 @@ def literal(translator, op):
     elif dtype.is_set():
         typeclass = 'set'
     else:
-        raise NotImplementedError
+        raise NotImplementedError(f'Unsupported type: {dtype!r}')
 
     return literal_formatters[typeclass](translator, op)
 

ibis/backends/bigquery/compiler.py

@@ -52,6 +52,8 @@ def keyword(cls, distinct):
 
 def find_bigquery_udf(op):
     """Filter which includes only UDFs from expression tree."""
+    if type(op) in BigQueryExprTranslator._rewrites:
+        op = BigQueryExprTranslator._rewrites[type(op)](op)
     if isinstance(op, operations.BigQueryUDFNode):
         result = op
     else:
@@ -118,3 +120,7 @@ def _generate_setup_queries(expr, context):
         # UDFs are uniquely identified by the name of the Node subclass we
         # generate.
         return list(toolz.unique(queries, key=lambda x: type(x.expr.op()).__name__))
+
+
+# Register custom UDFs
+import ibis.backends.bigquery.custom_udfs  # noqa:  F401, E402

ibis/backends/bigquery/custom_udfs.py

@@ -0,0 +1,39 @@
+import ibis.expr.datatypes as dt
+import ibis.expr.operations as ops
+from ibis.backends.bigquery.compiler import BigQueryExprTranslator
+from ibis.backends.bigquery.udf import udf
+
+# Based on:
+# https://github.com/GoogleCloudPlatform/bigquery-utils/blob/45e1ac51367ab6209f68e04b1660d5b00258c131/udfs/community/typeof.sqlx#L1
+typeof_ = udf.sql(
+    name="typeof",
+    params={"input": 'ANY TYPE'},
+    output_type=dt.str,
+    sql_expression=r"""
+    (
+        SELECT
+            CASE
+                -- Process NUMERIC, DATE, DATETIME, TIME, TIMESTAMP,
+                WHEN REGEXP_CONTAINS(literal, r'^[A-Z]+ "') THEN REGEXP_EXTRACT(literal, r'^([A-Z]+) "')
+                WHEN REGEXP_CONTAINS(literal, r'^-?[0-9]*$') THEN 'INT64'
+                WHEN
+                    REGEXP_CONTAINS(literal, r'^(-?[0-9]+[.e].*|CAST\("([^"]*)" AS FLOAT64\))$')
+                THEN
+                    'FLOAT64'
+                WHEN literal IN ('true', 'false') THEN 'BOOL'
+                WHEN literal LIKE '"%' OR literal LIKE "'%" THEN 'STRING'
+                WHEN literal LIKE 'b"%' THEN 'BYTES'
+                WHEN literal LIKE '[%' THEN 'ARRAY'
+                WHEN REGEXP_CONTAINS(literal, r'^(STRUCT)?\(') THEN 'STRUCT'
+                WHEN literal LIKE 'ST_%' THEN 'GEOGRAPHY'
+                WHEN literal = 'NULL' THEN 'NULL'
+            ELSE
+                'UNKNOWN'
+            END
+        FROM
+            UNNEST([FORMAT('%T', input)]) AS literal
+    )
+    """,
+)
+
+BigQueryExprTranslator.rewrites(ops.TypeOf)(lambda op: typeof_(op.arg).op())

ibis/backends/pandas/execution/decimal.py

@@ -56,13 +56,18 @@ def execute_decimal_log2(op, data, **kwargs):
 # exactly
 @execute_node.register((ops.Unary, ops.Negate), decimal.Decimal)
 def execute_decimal_unary(op, data, **kwargs):
-    operation_name = type(op).__name__.lower()
-    math_function = getattr(math, operation_name, None)
+    op_type = type(op)
+    operation_name = op_type.__name__.lower()
     function = getattr(
         decimal.Decimal,
         operation_name,
-        lambda x: decimal.Decimal(math_function(x)),
+        None,
     )
+    if function is None:
+        math_function = getattr(math, operation_name, None)
+        if math_function is None:
+            raise NotImplementedError(f'{op_type.__name__} not supported')
+        function = lambda x: decimal.Decimal(math_function(x))
     try:
         return function(data)
     except decimal.InvalidOperation:

ibis/backends/pandas/execution/generic.py

@@ -200,6 +200,8 @@ def execute_sort_key_series(op, data, _, **kwargs):
 def call_numpy_ufunc(func, op, data, **kwargs):
     if getattr(data, "dtype", None) == np.dtype(np.object_):
         return data.apply(functools.partial(execute_node, op, **kwargs))
+    if func is None:
+        raise NotImplementedError(f'{type(op).__name__} not supported')
     return func(data)
 
 

ibis/backends/polars/compiler.py

@@ -408,7 +408,10 @@ def string_length(op):
 @translate.register(ops.StringUnary)
 def string_unary(op):
     arg = translate(op.arg)
-    func = _string_unary[type(op)]
+    func = _string_unary.get(type(op))
+    if func is None:
+        raise NotImplementedError(f'{type(op).__name__} not supported')
+
     method = getattr(arg.str, func)
     return method()
 
@@ -869,7 +872,9 @@ def day_of_week_name(op):
 @translate.register(ops.Unary)
 def unary(op):
     arg = translate(op.arg)
-    func = _unary[type(op)]
+    func = _unary.get(type(op))
+    if func is None:
+        raise NotImplementedError(f'{type(op).__name__} not supported')
     return func(arg)
 
 
@@ -887,7 +892,9 @@ def unary(op):
 def comparison(op):
     left = translate(op.left)
     right = translate(op.right)
-    func = _comparisons[type(op)]
+    func = _comparisons.get(type(op))
+    if func is None:
+        raise NotImplementedError(f'{type(op).__name__} not supported')
     return func(left, right)
 
 
@@ -910,7 +917,9 @@ def between(op):
 
 @translate.register(ops.BitwiseBinary)
 def bitwise_binops(op):
-    ufunc = _bitwise_binops[type(op)]
+    ufunc = _bitwise_binops.get(type(op))
+    if ufunc is None:
+        raise NotImplementedError(f'{type(op).__name__} not supported')
     left = translate(op.left)
     right = translate(op.right)
 
@@ -948,7 +957,9 @@ def bitwise_not(op):
 def binop(op):
     left = translate(op.left)
     right = translate(op.right)
-    func = _binops[type(op)]
+    func = _binops.get(type(op))
+    if func is None:
+        raise NotImplementedError(f'{type(op).__name__} not supported')
     return func(left, right)