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:])