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'),
),
