def eval_operator(self, op: str, value: ValueType) -> Set[str]:
"""
Evaluates a given operator on the index for a given value and returns all
partition labels allowed by this index.
Parameters
----------
op: str
A string representation of the operator to be evaluated. Supported are
"==", "<=", ">=", "<", ">", "in"
For details, see documentation of kartothek.serialization
value: object
The value to be evaluated
Returns
-------
set
Allowed partition labels
"""
result = set()
index_dct = self.index_dct
index_type = self.dtype
index_arr = None
if index_type is not None and index_type and not pa.types.is_date(
index_type):
index_type = index_type.to_pandas_dtype()
try:
index_arr = np.fromiter(index_dct.keys(), dtype=index_type)
except ValueError:
pass
if index_arr is None:
index_arr = np.array(list(index_dct.keys()))
index = filter_array_like(index_arr,
op,
value,
strict_date_types=True,
column_name=self.column)
allowed_values = index_arr[index]
# Need to determine allowed values to include predicates like `in`
for value in allowed_values:
result.update(set(self.index_dct[value]))
return result
def test_raise_on_type(value, filter_value, op):
array_like = pd.Series([value])
with pytest.raises(TypeError, match="Unexpected type encountered."):
filter_array_like(array_like, op, filter_value, strict_date_types=True)
def test_filter_array_like_lower_eq(array_like):
ix = 3
value = array_like[ix]
res = filter_array_like(array_like, "<=", value)
assert all(array_like[res] == array_like[:ix + 1])
def test_filter_array_like_categoricals(op, expected, cat_type):
ser = pd.Series(["B", "C", "A"]).astype(cat_type)
res = filter_array_like(ser, op, "B")
assert all(res == expected)
def test_filter_array_like_larger_eq(array_like):
ix = 3
value = array_like[ix]
res = filter_array_like(array_like, ">=", value)
assert all(array_like[res] == array_like[ix:])
