def add_operation(op, func_name, db):
"""
Registers the given operation within the Ibis SQL translation toolchain
Parameters
----------
op: operator class
name: used in issuing statements to SQL engine
database: database the relevant operator is registered to
"""
full_name = f'{db}.{func_name}'
# TODO
# if op.input_type is rlz.listof:
# translator = comp.varargs(full_name)
# else:
arity = len(op.__signature__.parameters)
translator = fixed_arity(full_name, arity)
ImpalaExprTranslator._registry[op] = translator
ops.ExtractMonth:
_extract_field("month"),
ops.ExtractDay:
_extract_field("day"),
ops.ExtractHour:
_extract_field("hour"),
ops.ExtractMinute:
_extract_field("minute"),
ops.ExtractSecond:
_extract_field("second"),
ops.ExtractMillisecond:
_extract_field("millisecond"),
ops.Hash:
_hash,
ops.StringReplace:
fixed_arity("REPLACE", 3),
ops.StringSplit:
fixed_arity("SPLIT", 2),
ops.StringConcat:
_string_concat,
ops.StringJoin:
_string_join,
ops.StringAscii:
_string_ascii,
ops.StringFind:
_string_find,
ops.Substring:
_string_substring,
ops.StrRight:
_string_right,
ops.Repeat:
_operation_registry = {
**operation_registry,
}
_operation_registry.update(
{
ops.ExtractYear: _extract_field("year"),
ops.ExtractMonth: _extract_field("month"),
ops.ExtractDay: _extract_field("day"),
ops.ExtractHour: _extract_field("hour"),
ops.ExtractMinute: _extract_field("minute"),
ops.ExtractSecond: _extract_field("second"),
ops.ExtractMillisecond: _extract_field("millisecond"),
ops.Hash: _hash,
ops.StringReplace: fixed_arity("REPLACE", 3),
ops.StringSplit: fixed_arity("SPLIT", 2),
ops.StringConcat: _string_concat,
ops.StringJoin: _string_join,
ops.StringAscii: _string_ascii,
ops.StringFind: _string_find,
ops.Substring: _string_substring,
ops.StrRight: _string_right,
ops.Repeat: fixed_arity("REPEAT", 2),
ops.RegexSearch: _regex_search,
ops.RegexExtract: _regex_extract,
ops.RegexReplace: _regex_replace,
ops.GroupConcat: reduction("STRING_AGG"),
ops.IfNull: fixed_arity("IFNULL", 2),
ops.Cast: _cast,
ops.StructField: _struct_field,
raise com.UnsupportedOperationError(
'{!r} value not supported for arbitrary in Spark SQL'.format(how))
def _day_of_week_name(translator, expr):
arg = expr.op().arg
return 'date_format({}, {!r})'.format(translator.translate(arg), 'EEEE')
operation_registry = {**operation_registry}
operation_registry.update({
ops.IsNan:
unary('isnan'),
ops.IfNull:
fixed_arity('ifnull', 2),
ops.StructField:
_struct_field,
ops.MapValueForKey:
_map_value_for_key,
ops.ArrayLength:
unary('size'),
ops.Round:
_round,
ops.HLLCardinality:
reduction('approx_count_distinct'),
ops.StrRight:
fixed_arity('right', 2),
ops.StringSplit:
fixed_arity('SPLIT', 2),
ops.RegexSearch: