Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Support New Set Membership Query Condition #1837

Merged
merged 4 commits into from
Oct 14, 2023
Merged

Support New Set Membership Query Condition #1837

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
efe919d
Support New Set Membership Query Condition
nguyenv Sep 29, 2023
2253b3d
Correctly Handle Values for Python 3.7
nguyenv Oct 2, 2023
51b4e56
Support `Not In` Condition
nguyenv Oct 4, 2023
426b3de
Format
nguyenv Oct 4, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
85 changes: 78 additions & 7 deletions tiledb/query_condition.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -65,19 +65,34 @@ class PyQueryCondition {
use_enumeration);
}

template <typename T>
static PyQueryCondition
create(py::object pyctx, const std::string &field_name,
const std::vector<T> &values, tiledb_query_condition_op_t op) {
auto pyqc = PyQueryCondition(pyctx);

const Context ctx = std::as_const(pyqc.ctx_);

auto set_membership_qc =
QueryConditionExperimental::create(ctx, field_name, values, op);

pyqc.qc_ = std::make_shared<QueryCondition>(std::move(set_membership_qc));

return pyqc;
}

PyQueryCondition
combine(PyQueryCondition rhs,
combine(PyQueryCondition qc,
tiledb_query_condition_combination_op_t combination_op) const {

auto pyqc = PyQueryCondition(nullptr, ctx_.ptr().get());

tiledb_query_condition_t *combined_qc = nullptr;
ctx_.handle_error(
tiledb_query_condition_alloc(ctx_.ptr().get(), &combined_qc));

ctx_.handle_error(tiledb_query_condition_combine(
ctx_.ptr().get(), qc_->ptr().get(), rhs.qc_->ptr().get(),
combination_op, &combined_qc));
ctx_.ptr().get(), qc_->ptr().get(), qc.qc_->ptr().get(), combination_op,
&combined_qc));

pyqc.qc_ = std::shared_ptr<QueryCondition>(
new QueryCondition(pyqc.ctx_, combined_qc));
Expand Down Expand Up @@ -112,7 +127,6 @@ void init_query_condition(py::module &m) {
static_cast<void (PyQueryCondition::*)(
const string &, const string &, tiledb_query_condition_op_t)>(
&PyQueryCondition::init))

.def("init_uint64",
static_cast<void (PyQueryCondition::*)(const string &, uint64_t,
tiledb_query_condition_op_t)>(
Expand Down Expand Up @@ -145,7 +159,6 @@ void init_query_condition(py::module &m) {
static_cast<void (PyQueryCondition::*)(const string &, int8_t,
tiledb_query_condition_op_t)>(
&PyQueryCondition::init))

.def("init_float32",
static_cast<void (PyQueryCondition::*)(const string &, float,
tiledb_query_condition_op_t)>(
Expand All @@ -157,7 +170,63 @@ void init_query_condition(py::module &m) {

.def("__capsule__", &PyQueryCondition::__capsule__)

.def("combine", &PyQueryCondition::combine);
.def("combine", &PyQueryCondition::combine)

.def_static(
"create_string",
static_cast<PyQueryCondition (*)(
py::object, const std::string &, const std::vector<std::string> &,
tiledb_query_condition_op_t)>(&PyQueryCondition::create))
.def_static(
"create_uint64",
static_cast<PyQueryCondition (*)(
py::object, const std::string &, const std::vector<uint64_t> &,
tiledb_query_condition_op_t)>(&PyQueryCondition::create))
.def_static(
"create_int64",
static_cast<PyQueryCondition (*)(
py::object, const std::string &, const std::vector<int64_t> &,
tiledb_query_condition_op_t)>(&PyQueryCondition::create))
.def_static(
"create_uint32",
static_cast<PyQueryCondition (*)(
py::object, const std::string &, const std::vector<uint32_t> &,
tiledb_query_condition_op_t)>(&PyQueryCondition::create))
.def_static(
"create_int32",
static_cast<PyQueryCondition (*)(
py::object, const std::string &, const std::vector<int32_t> &,
tiledb_query_condition_op_t)>(&PyQueryCondition::create))
.def_static(
"create_uint16",
static_cast<PyQueryCondition (*)(
py::object, const std::string &, const std::vector<uint16_t> &,
tiledb_query_condition_op_t)>(&PyQueryCondition::create))
.def_static(
"create_int8",
static_cast<PyQueryCondition (*)(
py::object, const std::string &, const std::vector<int8_t> &,
tiledb_query_condition_op_t)>(&PyQueryCondition::create))
.def_static(
"create_uint16",
static_cast<PyQueryCondition (*)(
py::object, const std::string &, const std::vector<uint16_t> &,
tiledb_query_condition_op_t)>(&PyQueryCondition::create))
.def_static(
"create_int8",
static_cast<PyQueryCondition (*)(
py::object, const std::string &, const std::vector<int8_t> &,
tiledb_query_condition_op_t)>(&PyQueryCondition::create))
.def_static(
"create_float32",
static_cast<PyQueryCondition (*)(
py::object, const std::string &, const std::vector<float> &,
tiledb_query_condition_op_t)>(&PyQueryCondition::create))
.def_static(
"create_float64",
static_cast<PyQueryCondition (*)(
py::object, const std::string &, const std::vector<double> &,
tiledb_query_condition_op_t)>(&PyQueryCondition::create));

py::enum_<tiledb_query_condition_op_t>(m, "tiledb_query_condition_op_t",
py::arithmetic())
Expand All @@ -167,6 +236,8 @@ void init_query_condition(py::module &m) {
.value("TILEDB_GE", TILEDB_GE)
.value("TILEDB_EQ", TILEDB_EQ)
.value("TILEDB_NE", TILEDB_NE)
.value("TILEDB_IN", TILEDB_IN)
.value("TILEDB_NOT_IN", TILEDB_NOT_IN)
.export_values();

py::enum_<tiledb_query_condition_combination_op_t>(
Expand Down
40 changes: 32 additions & 8 deletions tiledb/query_condition.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -180,6 +180,9 @@ def visit_NotEq(self, node):
def visit_In(self, node):
return node

def visit_NotIn(self, node):
return node

def visit_List(self, node):
return list(node.elts)

Expand Down Expand Up @@ -208,21 +211,28 @@ def visit_Compare(self, node: Type[ast.Compare]) -> qc.PyQueryCondition:
self.visit(lhs), self.visit(op), self.visit(rhs)
)
result = result.combine(value, qc.TILEDB_AND)
elif isinstance(operator, ast.In):
elif isinstance(operator, (ast.In, ast.NotIn)):
rhs = node.comparators[0]
if not isinstance(rhs, ast.List):
raise TileDBError(
"`in` operator syntax must be written as `variable in ['l', 'i', 's', 't']`"
)

consts = self.visit(rhs)
result = self.aux_visit_Compare(
self.visit(node.left), qc.TILEDB_EQ, consts[0]
)
variable = node.left.id
values = [self.get_value_from_node(val) for val in self.visit(rhs)]

for val in consts[1:]:
value = self.aux_visit_Compare(self.visit(node.left), qc.TILEDB_EQ, val)
result = result.combine(value, qc.TILEDB_OR)
if self.array.schema.has_attr(variable):
enum_label = self.array.attr(variable).enum_label
if enum_label is not None:
dt = self.array.enum(enum_label).dtype
else:
dt = self.array.attr(variable).dtype
else:
dt = self.array.schema.attr_or_dim_dtype(variable)

dtype = "string" if dt.kind in "SUa" else dt.name
op = qc.TILEDB_IN if isinstance(operator, ast.In) else qc.TILEDB_NOT_IN
result = self.create_pyqc(dtype)(self.ctx, node.left.id, values, op)

return result

Expand Down Expand Up @@ -401,6 +411,20 @@ def init_pyqc(self, pyqc: qc.PyQueryCondition, dtype: str) -> Callable:

return getattr(pyqc, init_fn_name)

def create_pyqc(self, dtype: str) -> Callable:
if dtype != "string":
if np.issubdtype(dtype, np.datetime64):
dtype = "int64"
elif np.issubdtype(dtype, bool):
dtype = "uint8"

create_fn_name = f"create_{dtype}"

if not hasattr(qc.PyQueryCondition, create_fn_name):
raise TileDBError(f"PyQueryCondition.{create_fn_name}() not found.")

return getattr(qc.PyQueryCondition, create_fn_name)

def visit_BinOp(self, node: ast.BinOp) -> qc.PyQueryCondition:
try:
op = self.visit(node.op)
Expand Down
8 changes: 8 additions & 0 deletions tiledb/tests/test_query_condition.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -547,6 +547,10 @@ def test_in_operator_sparse(self):
result = A.query(cond="S in ['8']")[:]
assert len(result["S"]) == 0

result = A.query(cond="U not in [5, 6, 7]")[:]
for val in result["U"]:
assert val not in [5, 6, 7]

def test_in_operator_dense(self):
with tiledb.open(self.create_input_array_UIDSA(sparse=False)) as A:
U_mask = A.attr("U").fill
Expand Down Expand Up @@ -574,6 +578,10 @@ def test_in_operator_dense(self):
result = A.query(cond="S in ['8']")[:]
assert len(self.filter_dense(result["S"], S_mask)) == 0

result = A.query(cond="U not in [5, 6, 7]")[:]
for val in self.filter_dense(result["U"], U_mask):
assert val not in [5, 6, 7]

@pytest.mark.skipif(not has_pandas(), reason="pandas not installed")
def test_dense_datetime(self):
import pandas as pd
Expand Down