| // Copyright (C) 2020 The Android Open Source Project |
| // |
| // Licensed under the Apache License, Version 2.0 (the "License"); |
| // you may not use this file except in compliance with the License. |
| // You may obtain a copy of the License at |
| // |
| // http://www.apache.org/licenses/LICENSE-2.0 |
| // |
| // Unless required by applicable law or agreed to in writing, software |
| // distributed under the License is distributed on an "AS IS" BASIS, |
| // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| // See the License for the specific language governing permissions and |
| // limitations under the License. |
| #include "query.h" |
| |
| #include "hunk.h" |
| #include "io_spec.h" |
| #include "operator_context.h" |
| #include "pyerrfmt.h" |
| #include "pyobj.h" |
| #include "pyseq.h" |
| #include "qe_call_resize_buffers.h" |
| |
| namespace dctv { |
| |
| static_assert(sizeof (npy_intp) == sizeof (Py_ssize_t)); |
| |
| _Py_IDENTIFIER(_data); |
| _Py_IDENTIFIER(_mask); |
| |
| void |
| QeCall::do_operator_setup(QueryExecution*, |
| OperatorContext* oc, |
| QueryDb*) |
| { |
| throw_pyerr_fmt(PyExc_RuntimeError, |
| "operator %s did not call setup", |
| repr(oc->get_query_action())); |
| } |
| |
| IoSpec |
| QeCall::extract_io_spec() |
| { |
| throw_pyerr_fmt(PyExc_RuntimeError, "not an IO operation"); |
| } |
| |
| void |
| QeCall::refresh_resize_buffers(const OperatorContext* oc, |
| std::unique_ptr<QeCall>* self_ptr) |
| { |
| assume(self_ptr->get() == this); |
| // If we fail here, we'll leave *self_ptr empty, which will make |
| // QueryExecution complain if someone tries to continue execution. |
| std::unique_ptr<QeCall> new_qe_call = |
| std::make_unique<QeCallResizeBuffers>( |
| oc, self_ptr, std::move(*self_ptr)); |
| *self_ptr = std::move(new_qe_call); |
| } |
| |
| QueryClasses QueryClasses::instance; |
| |
| void |
| QueryClasses::ensure() |
| { |
| // TODO(dancol): use less centralized class interning |
| if (!instance.initialized()) { |
| instance.cls_tmpdir = |
| find_pyclass("dctv.util", "TmpDir"); |
| instance.cls_query_node = |
| find_pyclass("dctv.query", "QueryNode"); |
| instance.cls_query_action = |
| find_pyclass("dctv.query", "QueryAction"); |
| instance.cls_invalid_query_exception = |
| find_pyclass("dctv.query", "InvalidQueryException"); |
| instance.np_ma_nomask = |
| getattr(import_pymodule("numpy.ma"), "nomask"); |
| instance. cls_masked_array = |
| find_pyclass("numpy.ma", "masked_array"); |
| assert(instance.initialized()); |
| } |
| } |
| |
| void |
| throw_invalid_query(String msg) // NOLINT: min call site size |
| { |
| throw_pyerr_msg(QueryClasses::get().cls_invalid_query_exception.get(), |
| msg.c_str()); |
| } |
| |
| unique_pyarray |
| collapse_useless_views(unique_pyarray array) |
| { |
| if (npy_ndim(array) != 1) |
| return array; |
| for (;;) { |
| obj_pyref base = npy_base(array); |
| if (!base) |
| return array; |
| if (!is_base_ndarray(base)) |
| return array; |
| pyarray_ref abase = base.as_unsafe<PyArrayObject>(); |
| if (npy_ndim(abase) != 1) |
| return array; |
| if (npy_size1d(abase) != npy_size1d(array)) |
| return array; |
| if (npy_dtype(abase) != npy_dtype(array)) |
| return array; |
| if (npy_data_raw(abase) != npy_data_raw(array)) |
| return array; |
| array = abase.addref(); |
| } |
| } |
| |
| unique_pyarray |
| npy_get_data_of_ma(pyarray_ref ma) |
| { |
| assert(maybe_ma(ma)); |
| unique_pyarray data = getattr(ma, &PyId__data) |
| .addref_as_unsafe<PyArrayObject>(); |
| assert(is_base_ndarray(data)); |
| return collapse_useless_views(std::move(data)); |
| } |
| |
| unique_pyref |
| npy_get_mask_or_nomask_of_ma(pyarray_ref ma) |
| { |
| assert(maybe_ma(ma)); |
| unique_pyref m = getattr(ma, &PyId__mask); |
| assert(m == nomask() || is_base_ndarray(m)); |
| if (m == nomask()) |
| return m; |
| return collapse_useless_views( |
| std::move(m).addref_as_unsafe<PyArrayObject>()); |
| } |
| |
| unique_pyarray |
| npy_get_mask_of_ma(pyarray_ref ma) |
| { |
| unique_pyref m = npy_get_mask_or_nomask_of_ma(ma); |
| if (m == nomask()) |
| return {}; |
| assert(is_base_ndarray(m)); |
| return std::move(m).addref_as_unsafe<PyArrayObject>(); |
| } |
| |
| bool |
| npy_has_mask(pyarray_ref array) |
| { |
| pyarray_ref ma = maybe_ma(array); |
| return ma && npy_get_mask_of_ma(ma); |
| } |
| |
| static |
| unique_pyref |
| py_npy_has_mask(PyObject*, pyref obj) |
| { |
| return make_pybool(npy_has_mask(obj.as<PyArrayObject>())); |
| } |
| |
| static |
| unique_pyref |
| py_npy_get_data(PyObject*, pyref obj) |
| { |
| pyarray_ref ma = maybe_ma(obj); |
| if (ma) |
| return npy_get_data_of_ma(ma); |
| return obj.addref(); |
| } |
| |
| static |
| unique_pyref |
| py_npy_get_mask(PyObject*, pyref obj) |
| { |
| pyarray_ref ma = maybe_ma(obj); |
| if (ma) |
| return npy_get_mask_or_nomask_of_ma(ma); |
| return nomask().addref(); |
| } |
| |
| static |
| unique_pyref |
| py_npy_explode(PyObject*, pyref obj) |
| { |
| pyarray_ref ma = maybe_ma(obj); |
| if (!ma) |
| return pytuple::of(obj, nomask()); |
| return pytuple::of(npy_get_data_of_ma(ma), |
| npy_get_mask_or_nomask_of_ma(ma)); |
| } |
| |
| template<bool JustCheck> |
| static std::conditional_t< |
| JustCheck, |
| bool, |
| unique_pyarray> npy_get_broadcaster_1( |
| pyarray_ref array) noexcept(JustCheck); |
| |
| template<bool JustCheck> |
| static |
| std::conditional_t<JustCheck, bool, unique_pyarray> |
| npy_get_broadcaster_ma_1(pyarray_ref ma) |
| { |
| using Wanted = std::conditional_t<JustCheck, bool, unique_pyarray>; |
| unique_pyarray data = npy_get_data_of_ma(ma); |
| Wanted data_broadcaster = npy_get_broadcaster_1<JustCheck>(data); |
| if (!data_broadcaster) |
| return {}; |
| unique_pyarray mask = npy_get_mask_of_ma(ma); |
| if (!mask) |
| return data_broadcaster; |
| Wanted mask_broadcaster = npy_get_broadcaster_1<JustCheck>(mask); |
| if (!mask_broadcaster) |
| return {}; |
| if constexpr (JustCheck) { |
| return true; |
| } else { // NOLINT |
| return make_masked_array(data_broadcaster, mask_broadcaster); |
| } |
| } |
| |
| pyarray_ref |
| npy_get_broadcaster_raw(pyarray_ref array) noexcept |
| { |
| if (!is_base_ndarray(array)) |
| return {}; |
| if (npy_ndim(array) != 1) |
| return {}; |
| if (npy_strides(array)[0] != 0) |
| return {}; |
| pyref base = npy_base(array); |
| if (!base || !is_base_ndarray(base)) |
| return {}; |
| pyarray_ref abase = base.as_unsafe<PyArrayObject>(); |
| if (npy_dtype(array) != npy_dtype(abase)) |
| return {}; |
| if (npy_data_raw(array) != npy_data_raw(abase)) |
| return {}; |
| if (npy_ndim(abase) == 0 || |
| (npy_ndim(abase) == 1 && npy_dims(abase)[0] == 1)) |
| return abase; |
| return {}; |
| } |
| |
| template<bool JustCheck> |
| static |
| std::conditional_t<JustCheck, bool, unique_pyarray> |
| npy_get_broadcaster_1(pyarray_ref array) noexcept(JustCheck) |
| { |
| pyarray_ref ma = maybe_ma(array); |
| if (ma) |
| return npy_get_broadcaster_ma_1<JustCheck>(ma); |
| pyarray_ref bcaster_raw = npy_get_broadcaster_raw(array); |
| if (!bcaster_raw) |
| return {}; |
| if constexpr (JustCheck) { |
| return true; |
| } else { // NOLINT |
| if (npy_ndim(bcaster_raw) == 0) |
| return npy_ravel(bcaster_raw); |
| return bcaster_raw.addref(); |
| } |
| } |
| |
| unique_pyarray |
| npy_get_broadcaster(pyarray_ref array) |
| { |
| return npy_get_broadcaster_1<false>(array); |
| } |
| |
| bool |
| npy_is_broadcasted(pyarray_ref array) noexcept |
| { |
| return npy_get_broadcaster_1<true>(array); |
| } |
| |
| static |
| unique_pyref |
| py_npy_get_broadcaster(PyObject*, pyref obj) |
| { |
| pyarray_ref array = obj.as<PyArrayObject>(); |
| unique_pyarray broadcaster = npy_get_broadcaster(array); |
| if (!broadcaster) |
| return addref(Py_None); |
| return broadcaster; |
| } |
| |
| static |
| unique_pyref |
| py_npy_is_broadcasted(PyObject*, pyref obj) |
| { |
| pyarray_ref array = obj.as<PyArrayObject>(); |
| return make_pybool(npy_is_broadcasted(array)); |
| } |
| |
| unique_pyarray |
| make_masked_array(pyref data, pyref mask) |
| { |
| return call(QueryClasses::get().cls_masked_array, data, mask) |
| .addref_as_unsafe<PyArrayObject>(); |
| } |
| |
| void |
| maybe_drop_mask(unique_pyarray* mask) |
| { |
| if (!*mask) |
| return; |
| |
| pyarray_ref bcaster = npy_get_broadcaster_raw(*mask); |
| if (bcaster) { |
| assume(npy_ndim(bcaster) <= 1); |
| if (npy_ndim(bcaster)) |
| assert(npy_size1d(bcaster) == 1); |
| assume(npy_dtype(bcaster) == type_descr<bool>()); |
| if (!npy_data<bool>(bcaster)[0]) { |
| mask->reset(); |
| } |
| } else { |
| // TODO(dancol): is the full scan worth it? |
| if (!npy_any(*mask)) |
| mask->reset(); |
| } |
| } |
| |
| |
| |
| static PyMethodDef functions[] = { |
| make_methoddef("npy_has_mask", |
| wraperr<py_npy_has_mask>(), |
| METH_O, |
| "Return whether an object has a mask"), |
| make_methoddef("npy_get_mask", |
| wraperr<py_npy_get_mask>(), |
| METH_O, |
| "Return an object's mask if present or None"), |
| make_methoddef("npy_get_data", |
| wraperr<py_npy_get_data>(), |
| METH_O, |
| "Return data of masked array or array"), |
| make_methoddef("npy_explode", |
| wraperr<py_npy_explode>(), |
| METH_O, |
| "Split an array into value and mask"), |
| make_methoddef("npy_get_broadcaster", |
| wraperr<py_npy_get_broadcaster>(), |
| METH_O, |
| "Return the underlying broadcasted array " |
| "or None if array is not broadcasted"), |
| make_methoddef("npy_is_broadcasted", |
| wraperr<py_npy_is_broadcasted>(), |
| METH_O, |
| "Return the underlying broadcasted array " |
| "or None if array is not broadcasted"), |
| { 0 } |
| }; |
| |
| void |
| init_query(pyref m) |
| { |
| register_functions(m, functions); |
| } |
| |
| } // namespace dctv |