def wrapper(f):
if not callable(f):
raise TypeError('f must be callable, got {}'.format(f))
signature = inspect.signature(f)
parameter_names = signature.parameters.keys()
udf_node_fields = collections.OrderedDict(
[(name, Arg(rlz.value(type)))
for name, type in zip(parameter_names, input_type)] + [
('output_type', lambda self, output_type=output_type: rlz.
shape_like(self.args, dtype=output_type)),
('__slots__', ('js', )),
])
udf_node = create_udf_node(f.__name__, udf_node_fields)
@compiles(udf_node)
def compiles_udf_node(t, expr):
return '{}({})'.format(udf_node.__name__,
', '.join(map(t.translate,
expr.op().args)))
type_translation_context = UDFContext()
return_type = ibis_type_to_bigquery_type(dt.dtype(output_type),
type_translation_context)
bigquery_signature = ', '.join(
'{name} {type}'.format(name=name,
type=ibis_type_to_bigquery_type(
dt.dtype(type),
type_translation_context))
for name, type in zip(parameter_names, input_type))
source = PythonToJavaScriptTranslator(f).compile()
js = '''\
CREATE TEMPORARY FUNCTION {external_name}({signature})
RETURNS {return_type}
LANGUAGE js AS """
{strict}{source}
return {internal_name}({args});
"""{libraries};'''.format(external_name=udf_node.__name__,
internal_name=f.__name__,
return_type=return_type,
source=source,
signature=bigquery_signature,
strict=repr('use strict') + ';\n' if strict else '',
args=', '.join(parameter_names),
libraries=('\nOPTIONS (\n library={}\n)'.format(
repr(list(libraries))) if libraries else ''))
@functools.wraps(f)
def wrapped(*args, **kwargs):
node = udf_node(*args, **kwargs)
node.js = js
return node.to_expr()
wrapped.__signature__ = signature
wrapped.js = js
return wrapped
def bq_param_array(param, value):
param_type = param.type()
assert isinstance(param_type, dt.Array), str(param_type)
try:
bigquery_type = ibis_type_to_bigquery_type(param_type.value_type)
except NotImplementedError:
raise com.UnsupportedBackendType(param_type)
else:
if isinstance(param_type.value_type, dt.Struct):
query_value = [
bigquery_param(param[i].name('element_{:d}'.format(i)), struct)
for i, struct in enumerate(value)
]
bigquery_type = 'STRUCT'
elif isinstance(param_type.value_type, dt.Array):
raise TypeError('ARRAY<ARRAY<T>> is not supported in BigQuery')
else:
query_value = value
result = bq.ArrayQueryParameter(
param.get_name(), bigquery_type, query_value)
return result
def compiles_floor(t, e):
bigquery_type = ibis_type_to_bigquery_type(e.type())
arg = e.op().arg
return 'CAST(FLOOR({}) AS {})'.format(t.translate(arg), bigquery_type)
def bigquery_cast_generate(compiled_arg, from_, to):
sql_type = ibis_type_to_bigquery_type(to)
return 'CAST({} AS {})'.format(compiled_arg, sql_type)
def test_ibis_type_to_bigquery_type_udf(type, expected):
context = UDFContext()
assert ibis_type_to_bigquery_type(type, context) == expected
def test_simple_failure_mode(datatype):
with pytest.raises(TypeError):
ibis_type_to_bigquery_type(datatype)
def test_simple(datatype, expected):
context = TypeTranslationContext()
assert ibis_type_to_bigquery_type(datatype, context) == expected
def test_ibis_type_to_bigquery_type(type, expected):
result = ibis_type_to_bigquery_type(dt.dtype(type))
assert result == expected
def test_ibis_type_to_bigquery_type_udf(type, expected):
context = UDFContext()
assert ibis_type_to_bigquery_type(type, context) == expected
def test_simple_failure_mode(datatype):
with pytest.raises(TypeError):
ibis_type_to_bigquery_type(datatype)
def test_simple(datatype, expected):
context = TypeTranslationContext()
assert ibis_type_to_bigquery_type(datatype, context) == expected
def compiles_floor(t, e):
bigquery_type = ibis_type_to_bigquery_type(e.type())
arg = e.op().arg
return 'CAST(FLOOR({}) AS {})'.format(t.translate(arg), bigquery_type)
def bigquery_cast_generate(compiled_arg, from_, to):
sql_type = ibis_type_to_bigquery_type(to)
return 'CAST({} AS {})'.format(compiled_arg, sql_type)
def wrapper(f):
if not callable(f):
raise TypeError('f must be callable, got {}'.format(f))
signature = inspect.signature(f)
parameter_names = signature.parameters.keys()
udf_node_fields = collections.OrderedDict(
[
(name, Arg(rlz.value(type)))
for name, type in zip(parameter_names, input_type)
]
+ [
(
'output_type',
lambda self, output_type=output_type: rlz.shape_like(
self.args, dtype=output_type
),
),
('__slots__', ('js',)),
]
)
udf_node = create_udf_node(f.__name__, udf_node_fields)
@compiles(udf_node)
def compiles_udf_node(t, expr):
return '{}({})'.format(
udf_node.__name__, ', '.join(map(t.translate, expr.op().args))
)
type_translation_context = UDFContext()
return_type = ibis_type_to_bigquery_type(
dt.dtype(output_type), type_translation_context
)
bigquery_signature = ', '.join(
'{name} {type}'.format(
name=name,
type=ibis_type_to_bigquery_type(
dt.dtype(type), type_translation_context
),
)
for name, type in zip(parameter_names, input_type)
)
source = PythonToJavaScriptTranslator(f).compile()
js = '''\
CREATE TEMPORARY FUNCTION {external_name}({signature})
RETURNS {return_type}
LANGUAGE js AS """
{strict}{source}
return {internal_name}({args});
"""{libraries};'''.format(
external_name=udf_node.__name__,
internal_name=f.__name__,
return_type=return_type,
source=source,
signature=bigquery_signature,
strict=repr('use strict') + ';\n' if strict else '',
args=', '.join(parameter_names),
libraries=(
'\nOPTIONS (\n library={}\n)'.format(repr(list(libraries)))
if libraries
else ''
),
)
@functools.wraps(f)
def wrapped(*args, **kwargs):
node = udf_node(*args, **kwargs)
node.js = js
return node.to_expr()
wrapped.__signature__ = signature
wrapped.js = js
return wrapped
