kszucs commented on a change in pull request #11185:
URL: https://github.com/apache/arrow/pull/11185#discussion_r712110860
##########
File path: cpp/src/arrow/python/iterators.h
##########
@@ -97,31 +98,53 @@ inline Status VisitSequence(PyObject* obj, int64_t offset,
VisitorFunc&& func) {
template <class VisitorFunc>
inline Status VisitSequenceMasked(PyObject* obj, PyObject* mo, int64_t offset,
VisitorFunc&& func) {
- if (mo == nullptr || !PyArray_Check(mo)) {
- return Status::Invalid("Null mask must be NumPy array");
- }
-
- PyArrayObject* mask = reinterpret_cast<PyArrayObject*>(mo);
- if (PyArray_NDIM(mask) != 1) {
- return Status::Invalid("Mask must be 1D array");
- }
+ if (PyArray_Check(mo)) {
+ PyArrayObject* mask = reinterpret_cast<PyArrayObject*>(mo);
+ if (PyArray_NDIM(mask) != 1) {
+ return Status::Invalid("Mask must be 1D array");
+ }
+ if (PyArray_SIZE(mask) != static_cast<int64_t>(PySequence_Size(obj))) {
+ return Status::Invalid("Mask was a different length from sequence being
converted");
+ }
- const Py_ssize_t obj_size = PySequence_Size(obj);
- if (PyArray_SIZE(mask) != static_cast<int64_t>(obj_size)) {
- return Status::Invalid("Mask was a different length from sequence being
converted");
- }
+ const int dtype = fix_numpy_type_num(PyArray_DESCR(mask)->type_num);
+ if (dtype == NPY_BOOL) {
+ Ndarray1DIndexer<uint8_t> mask_values(mask);
- const int dtype = fix_numpy_type_num(PyArray_DESCR(mask)->type_num);
- if (dtype == NPY_BOOL) {
- Ndarray1DIndexer<uint8_t> mask_values(mask);
+ return VisitSequenceGeneric(
+ obj, offset,
+ [&func, &mask_values](PyObject* value, int64_t i, bool* keep_going) {
+ return func(value, mask_values[i], keep_going);
+ });
+ } else {
+ return Status::Invalid("Mask must be boolean dtype");
+ }
+ } else if (PySequence_Check(mo)) {
+ if (PySequence_Size(mo) != static_cast<int64_t>(PySequence_Size(obj))) {
+ return Status::Invalid("Mask was a different length from sequence being
converted");
+ }
return VisitSequenceGeneric(
- obj, offset, [&func, &mask_values](PyObject* value, int64_t i, bool*
keep_going) {
- return func(value, mask_values[i], keep_going);
+ obj, offset, [&func, &mo](PyObject* value, int64_t i, bool*
keep_going) {
+ OwnedRef value_ref(PySequence_ITEM(mo, i));
+ PyObject* value_obj = value_ref.obj();
+ auto scalar = unwrap_scalar(value_obj);
+ if (scalar.ok()) {
+ BooleanScalar* z =
dynamic_cast<BooleanScalar*>(scalar.ValueOrDie().get());
+ if (z == nullptr)
+ return Status::Invalid("Mask must be made of boolean values");
+ return func(value, z->value, keep_going);
+ } else if (PyBool_Check(value_obj)) {
+ return func(value, value_obj == Py_True, keep_going);
+ } else {
+ return Status::Invalid("Mask must be made of boolean values");
+ }
});
} else {
- return Status::Invalid("Mask must be boolean dtype");
+ return Status::Invalid("Null mask must be a NumPy or Arrow array");
Review comment:
According to the tests below generic boolean list seems to be supported.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscr...@arrow.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
