src/dctv/block.cpp

// 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 "block.h"

#include "automethod.h"
#include "fmt.h"
#include "hunk.h"
#include "npy.h"
#include "pyobj.h"
#include "pyobj.h"
#include "pyparsetuple.h"
#include "pyparsetuplenpy.h"
#include "pyseq.h"
#include "query.h"

namespace dctv {

unique_obj_pyref<Block>
Block::make_empty(QueryCache* qc, unique_dtype dtype)
{
  return make_pyobj<Block>(Hunk::make_empty(qc, std::move(dtype)),
                           unique_hunk());
}

int
Block::py_traverse(visitproc visit, void* arg) const noexcept
{
  Py_VISIT(this->data.get());  // NOLINT
  Py_VISIT(this->mask.get());  // NOLINT
  return 0;
}

static
unique_pyref
Block_get_data(PyObject* self, void*)
{
  return pyref(self).as_unsafe<Block>()->get_data().notnull();
}

static
unique_pyref
Block_get_mask(PyObject* self, void*)
{
  unique_pyref mask =
      addref(pyref(self).as_unsafe<Block>()->get_mask());
  if (!mask)
    mask = nomask().notnull().addref();
  return mask;
}

static
unique_pyref
Block_has_mask(PyObject* self, void*)
{
  return make_pybool(pyref(self).as_unsafe<Block>()->has_mask());
}

static
unique_pyref
Block_dtype(PyObject* self, void*)
{
  return pyref(self).as_unsafe<Block>()
      ->get_dtype().notnull().addref();
}

unique_pyarray
Block::as_array(unique_dtype dtype) const
{
  unique_pyarray data = this->get_data();
  assert(is_base_ndarray(data));
  assert(npy_ndim(data) == 1);
  if (dtype && npy_dtype(data) != dtype) {
    data = as_base_1d_pyarray(data, dtype.addref());
    assume(npy_dtype(data) == dtype.get());
  }
  return this->has_mask()
      ? make_masked_array(std::move(data), this->get_mask())
      : std::move(data);
}

Block::operator String() const
{
  return fmt("<Block len=%s dtype=%s%s>",
             this->get_size(),
             repr(this->get_dtype()),
             this->has_mask() ? " masked" : "");
}

unique_obj_pyref<Block>
Block::from_arrays(QueryCache* qc,
                   unique_pyarray data,
                   unique_pyarray mask /* nullable */)
{
  return make_pyobj<Block>(
      Hunk::convert_pyarray(qc, std::move(data.notnull())),
      (mask
       ? Hunk::convert_pyarray(qc, std::move(mask))
       : unique_hunk()));
}

PyMethodDef Block::pymethods[] = {
  AUTOMETHOD(&Block::as_array,
             ("Make a numpy array out of this block "
              "(masked_array if masked)"),
             (OPTIONAL_ARGS_FOLLOW)
             (unique_dtype, dtype, unique_dtype(), convert_dtype)
  ),
  AUTOMETHOD(&Block::as_array,
             "__array__",
             ("Make a numpy array out of this block "
              "(masked_array if masked)"),
             (OPTIONAL_ARGS_FOLLOW)
             (unique_dtype, dtype, unique_dtype(), convert_dtype)
  ),
  { 0 },
};

PyGetSetDef Block::pygetset[] = {
  make_getset("data",
              "Get numpy array of block data",
              wraperr<Block_get_data>()),
  make_getset("mask",
              "Get numpy array of block mask or False if no mask",
              wraperr<Block_get_mask>()),
  make_getset("has_mask",
              "Whether this block has a mask.  May be more "
              "efficient than testing the mask property directly, "
              "since the mask property has to inflate a mask if "
              "present",
              wraperr<Block_has_mask>()),
  make_getset("dtype",
              "dtype of this block",
              wraperr<Block_dtype>()),
  { 0 },
};

PySequenceMethods Block::pyseq = {
  .sq_length = wraperr<&Block::get_size>(),
};

PyTypeObject Block::pytype =
    make_py_type<Block>(
        "dctv._native.Block",
        "Block of query data",
        [](PyTypeObject* t){
          t->tp_methods = Block::pymethods;
          t->tp_getset = Block::pygetset;
          t->tp_as_sequence = &Block::pyseq;
        });

void
init_block(pyref m)
{
  register_type(m, &Block::pytype);
}

}  // namespace dctv