class SearchedCase(Value): cases = rlz.value_list_of(rlz.boolean) results = rlz.value_list_of(rlz.any) default = rlz.any output_shape = rlz.shape_like("cases") def __init__(self, cases, results, default): assert len(cases) == len(results) super().__init__(cases=cases, results=results, default=default) def root_tables(self): return distinct_roots(*self.flat_args()) @immutable_property def output_dtype(self): exprs = [*self.results, self.default] return rlz.highest_precedence_dtype(exprs)
class Contains(Value): value = rlz.any options = rlz.one_of([ rlz.value_list_of(rlz.any), rlz.set_, rlz.column(rlz.any), rlz.array_of(rlz.any), ]) output_dtype = dt.boolean output_shape = rlz.shape_like("args")
class SimpleCase(Value): base = rlz.any cases = rlz.value_list_of(rlz.any) results = rlz.value_list_of(rlz.any) default = rlz.any output_shape = rlz.shape_like("base") def __init__(self, cases, results, **kwargs): assert len(cases) == len(results) super().__init__(cases=cases, results=results, **kwargs) def root_tables(self): return distinct_roots(*self.flat_args()) @immutable_property def output_dtype(self): # TODO(kszucs): we could extend the functionality of # rlz.shape_like to support varargs with .flat_args() # to define a subset of input arguments values = [*self.results, self.default] return rlz.highest_precedence_dtype(values)
class ArrayColumn(Value): cols = rlz.value_list_of(rlz.column(rlz.any), min_length=1) output_shape = rlz.Shape.COLUMNAR def __init__(self, cols): if len({col.type() for col in cols}) > 1: raise com.IbisTypeError( f'The types of all input columns must match exactly in a ' f'{type(self).__name__} operation.') super().__init__(cols=cols) @immutable_property def output_dtype(self): first_dtype = self.cols[0].type() return dt.Array(first_dtype)
class CoalesceLike(Value): # According to Impala documentation: # Return type: same as the initial argument value, except that integer # values are promoted to BIGINT and floating-point values are promoted to # DOUBLE; use CAST() when inserting into a smaller numeric column arg = rlz.value_list_of(rlz.any) output_shape = rlz.shape_like('arg') @immutable_property def output_dtype(self): # filter out null types non_null_exprs = [arg for arg in self.arg if arg.type() != dt.null] if non_null_exprs: return rlz.highest_precedence_dtype(non_null_exprs) else: return dt.null
[ (Foo, 'c', IbisTypeError), (Bar, 'c', IbisTypeError), (Baz(3), 'b', IbisTypeError), ], ) def test_invalid_member_of(obj, value, expected): with pytest.raises(expected): rlz.member_of(obj, value) @pytest.mark.parametrize( ('validator', 'values', 'expected'), [ param( rlz.value_list_of(identity), (3, 2), ibis.sequence([3, 2]), id="list_identity", ), param( rlz.value_list_of(rlz.integer), (3, 2), ibis.sequence([3, 2]), id="list_int", ), param( rlz.value_list_of(rlz.integer), (3, None), ibis.sequence([3, ibis.NA]), id="list_int_null",
class StringConcat(Value): arg = rlz.value_list_of(rlz.string) output_shape = rlz.shape_like("arg") output_dtype = dt.string
class StringJoin(Value): sep = rlz.string arg = rlz.value_list_of(rlz.string, min_length=1) output_dtype = dt.string output_shape = rlz.shape_like("arg")
class FindInSet(Value): needle = rlz.string values = rlz.value_list_of(rlz.string, min_length=1) output_shape = rlz.shape_like("needle") output_dtype = dt.int64
('obj', 'value', 'expected'), [ (Foo, 'c', IbisTypeError), (Bar, 'c', IbisTypeError), (Baz(3), 'b', IbisTypeError), ], ) def test_invalid_member_of(obj, value, expected): with pytest.raises(expected): rlz.member_of(obj, value) @pytest.mark.parametrize( ('validator', 'values', 'expected'), [ (rlz.value_list_of(identity), (3, 2), ibis.sequence([3, 2])), (rlz.value_list_of(rlz.integer), (3, 2), ibis.sequence([3, 2])), ( rlz.value_list_of(rlz.integer), (3, None), ibis.sequence([3, ibis.NA]), ), (rlz.value_list_of(rlz.string), ('a', ), ibis.sequence(['a'])), (rlz.value_list_of(rlz.string), ['a', 'b'], ibis.sequence(['a', 'b'])), pytest.param( rlz.value_list_of(rlz.value_list_of(rlz.string)), [[], ['a']], ibis.sequence([[], ['a']]), marks=pytest.mark.xfail(raises=ValueError, reason='Not yet implemented'), ),