jorisvandenbossche commented on a change in pull request #11185: URL: https://github.com/apache/arrow/pull/11185#discussion_r726188672
########## File path: python/pyarrow/tests/test_array.py ########## @@ -2745,6 +2745,33 @@ def test_array_masked(): mask=np.array([False, True, False, True])) assert arr.type == pa.int64() + # ARROW-13883 + arr = pa.array([4, None, 4, 3], + mask=pa.array([False, True, False, True])) + assert arr.to_pylist() == [4, None, 4, None] + + arr = pa.array([4, None, 4, 3], + mask=[False, True, False, True]) Review comment: Can you slightly change the input data so there is also a None in the data that is not covered by a True in the mask? (eg `[None, 4, 4, 3]` instead, and the result will then be `[None, None, 4, None]`) ########## File path: python/pyarrow/tests/test_array.py ########## @@ -2745,6 +2745,33 @@ def test_array_masked(): mask=np.array([False, True, False, True])) assert arr.type == pa.int64() + # ARROW-13883 Review comment: Can you maybe start a new test function here? (the above is explicitly about the type of values being ignored if they are masked (so a float being present does not force it to become a floating array if it is masked an all other values are integers) ########## File path: cpp/src/arrow/python/iterators.h ########## @@ -97,31 +99,64 @@ 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"); - } + 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"); + } - PyArrayObject* mask = reinterpret_cast<PyArrayObject*>(mo); - if (PyArray_NDIM(mask) != 1) { - return Status::Invalid("Mask must be 1D array"); - } + const int dtype = fix_numpy_type_num(PyArray_DESCR(mask)->type_num); + if (dtype == NPY_BOOL) { + Ndarray1DIndexer<uint8_t> mask_values(mask); - 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"); - } + 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"); Review comment: Can you also add tests that cover those checks? -- 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