/*

* Licensed to the Apache Software Foundation (ASF) under one

* or more contributor license agreements. See the NOTICE file

* distributed with this work for additional information

* regarding copyright ownership. The ASF licenses this file

* to you 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 <functional>

#include <string>

#include <sstream>

#include <stdexcept>

#include <variant>

#include <vector>

#include <nanobind/nanobind.h>

#include <nanobind/ndarray.h>

#include <nanobind/stl/variant.h>

#include <nanobind/stl/vector.h>

#include <nanobind/stl/string.h>

#include "kll_sketch.hpp"

namespace nb = nanobind;

namespace datasketches {

namespace vector_of_kll_constants {

static const uint32_t DEFAULT_K = kll_constants::DEFAULT_K;

static const uint32_t DEFAULT_D = 1;

}

// Wrapper class for Numpy compatibility

template <typename T, typename C = std::less<T>>

class vector_of_kll_sketches {

public:

explicit vector_of_kll_sketches(uint32_t k = vector_of_kll_constants::DEFAULT_K, uint32_t d = vector_of_kll_constants::DEFAULT_D);

vector_of_kll_sketches(const vector_of_kll_sketches& other);

vector_of_kll_sketches(vector_of_kll_sketches&& other) noexcept;

vector_of_kll_sketches<T, C>& operator=(const vector_of_kll_sketches& other);

vector_of_kll_sketches<T, C>& operator=(vector_of_kll_sketches&& other);

// container parameters

inline uint32_t get_k() const;

inline uint32_t get_d() const;

template<typename V>

using Array1D = nb::ndarray<V, nb::numpy, nb::ndim<1>>;

template<typename V>

using Array2D = nb::ndarray<V, nb::numpy, nb::ndim<2>, nb::c_contig>;

// sketch updates/merges

void update(nb::ndarray<T>& items, char order);

void merge(const vector_of_kll_sketches<T>& other);

template<typename V>

using ArrInputType = std::variant<nb::ndarray<>, nb::list, V>;

// returns a single sketch combining all data in the array

kll_sketch<T, C> collapse(ArrInputType<int>& isk) const;

// sketch queries returning an array of results

Array1D<bool> is_empty() const;

Array1D<uint64_t> get_n() const;

Array1D<bool> is_estimation_mode() const;

Array1D<T> get_min_values() const;

Array1D<T> get_max_values() const;

Array1D<uint32_t> get_num_retained() const;

Array2D<T> get_quantiles(ArrInputType<double>& ranks, ArrInputType<int>& isk) const;

Array2D<double> get_ranks(ArrInputType<T>& values, ArrInputType<int>& isk) const;

Array2D<double> get_pmf(ArrInputType<T>& split_points, ArrInputType<int>& isk) const;

Array2D<double> get_cdf(ArrInputType<T>& split_points, ArrInputType<int>& isk) const;

// human-readable output

std::string to_string(bool print_levels = false, bool print_items = false) const;

// binary output/input

nb::list serialize(ArrInputType<int>& isk);

// note: deserialize() replaces the sketch at the specified

// index. Not a static method.

void deserialize(const nb::bytes& sk_bytes, uint32_t idx);

private:

template<typename TT>

Array1D<TT> input_to_vec(ArrInputType<TT>& input) const;

Array1D<uint32_t> get_indices(Array1D<int>& isk) const;

template<typename TT>

Array1D<TT> make_ndarray(size_t size) const;

template<typename TT>

Array2D<TT> make_ndarray(size_t rows, size_t cols) const;

const uint32_t k_; // kll sketch k parameter

const uint32_t d_; // number of dimensions (here: sketches) to hold

std::vector<kll_sketch<T, C>> sketches_;

};

template<typename T, typename C>

vector_of_kll_sketches<T, C>::vector_of_kll_sketches(uint32_t k, uint32_t d):

k_(k),

d_(d)

{

// check d is valid (k is checked by kll_sketch)

if (d < 1) {

throw std::invalid_argument("D must be >= 1: " + std::to_string(d));

}

sketches_.reserve(d);

// spawn the sketches

for (uint32_t i = 0; i < d; i++) {

sketches_.emplace_back(k);

}

}

template<typename T, typename C>

vector_of_kll_sketches<T, C>::vector_of_kll_sketches(const vector_of_kll_sketches& other) :

k_(other.k_),

d_(other.d_),

sketches_(other.sketches_)

{}

template<typename T, typename C>

vector_of_kll_sketches<T, C>::vector_of_kll_sketches(vector_of_kll_sketches&& other) noexcept :

k_(other.k_),

d_(other.d_),

sketches_(std::move(other.sketches_))

{}

template<typename T, typename C>

vector_of_kll_sketches<T, C>& vector_of_kll_sketches<T, C>::operator=(const vector_of_kll_sketches& other) {

vector_of_kll_sketches<T, C> copy(other);

k_ = copy.k_;

d_ = copy.d_;

std::swap(sketches_, copy.sketches_);

return *this;

}

template<typename T, typename C>

vector_of_kll_sketches<T, C>& vector_of_kll_sketches<T, C>::operator=(vector_of_kll_sketches&& other) {

k_ = other.k_;

d_ = other.d_;

std::swap(sketches_, other.sketches_);

return *this;

}

template<typename T, typename C>

uint32_t vector_of_kll_sketches<T, C>::get_k() const {

return k_;

}

template<typename T, typename C>

uint32_t vector_of_kll_sketches<T, C>::get_d() const {

return d_;

}

template<typename T, typename C>

template<typename TT>

auto vector_of_kll_sketches<T, C>::make_ndarray(size_t size) const -> Array1D<TT> {

TT* data = new TT[size];

nb::capsule owner(data, [](void *p) noexcept {

delete[] static_cast<TT*>(p);

});

return Array1D<TT>(data, {size}, owner);

}

template<typename T, typename C>

template<typename TT>

auto vector_of_kll_sketches<T, C>::make_ndarray(size_t rows, size_t cols) const -> Array2D<TT> {

TT* data = new TT[rows * cols];

nb::capsule owner(data, [](void *p) noexcept {

delete[] static_cast<TT*>(p);

});

return Array2D<TT>(data, {rows, cols}, owner);

}

template<typename T, typename C>

template<typename TT>

auto vector_of_kll_sketches<T, C>::input_to_vec(ArrInputType<TT>& input) const -> Array1D<TT> {

if (std::holds_alternative<nb::ndarray<>>(input)) {

nb::ndarray<> arr = std::get<nb::ndarray<>>(input);

return Array1D<TT>(arr);

} else if (std::holds_alternative<nb::list>(input)) {

nb::list list = std::get<nb::list>(input);

Array1D<TT> output = make_ndarray<TT>(list.size());

for (size_t i = 0; i < list.size(); ++i)

output(i) = nb::cast<TT>(list[i]);

return output;

} else if (std::holds_alternative<TT>(input)) {

Array1D<TT> output = make_ndarray<TT>(1);

output(0) = std::get<TT>(input);

return output;

} else {

nb::raise_type_error("Unknown input type");

}

}

template<typename T, typename C>

auto vector_of_kll_sketches<T, C>::get_indices(Array1D<int>& isk) const -> Array1D<uint32_t> {

auto input = isk.view<nb::ndim<1>>();

size_t num_input = input.shape(0);

Array1D<uint32_t> output;

if (num_input == 1) {

if (input(0) == -1) {

output = make_ndarray<uint32_t>(d_);

for (uint32_t i = 0; i < d_; ++i) {

output(i) = i;

}

} else {

output = make_ndarray<uint32_t>(1);

output(0) = static_cast<uint32_t>(input(0));

}

} else {

output = make_ndarray<uint32_t>(num_input);

for (uint32_t i = 0; i < num_input; ++i) {

const uint32_t idx = static_cast<uint32_t>(input(i));

if (idx < d_) {

output(i) = idx;

} else {

throw std::invalid_argument("request for invalid dimensions >= d ("

+ std::to_string(d_) +"): "+ std::to_string(idx));

}

}

}

return output;

}

// Checks if each sketch is empty or not

template<typename T, typename C>

auto vector_of_kll_sketches<T, C>::is_empty() const -> Array1D<bool> {

auto vals = make_ndarray<bool>(d_);

auto view = vals.view();

for (uint32_t i = 0; i < d_; ++i) {

view(i) = sketches_[i].is_empty();

}

return vals;

}

// Updates each sketch with values

// Currently: all values must be present

// TODO: allow subsets of sketches to be updated

template<typename T, typename C>

void vector_of_kll_sketches<T, C>::update(nb::ndarray<T>& items, char order) {

size_t ndim = items.ndim();

if (items.shape(ndim-1) != d_) {

throw std::invalid_argument("input data must have rows with " + std::to_string(d_)

+ " elements. Found: " + std::to_string(items.shape(ndim-1)));

}

if (ndim == 1) {

// 1D case: single value to update per sketch

const T* data = items.data();

for (uint32_t i = 0; i < d_; ++i) {

sketches_[i].update(data[i]);

}

}

else if (ndim == 2) {

// 2D case: multiple values to update per sketch

// We could speedthis up by using raw array access and pre-computing an offset for the

// row/column, but if we use the wrong ordering the computation would be incorrect.

// By using a view and dereferencing by (row, column) each time we ensure correct

// processing at the cost of an extra multiply each derference. Using a mismatched

// ordering versus the actual data storage will be potentially slower but will still

// produce correct output.

auto data = items.template view<nb::ndim<2>>();

if (order == 'F' || order == 'f') { // Fortran-style (column-major) order

for (uint32_t j = 0; j < d_; ++j) {

const size_t offset = j * d_;

for (uint32_t i = 0; i < items.shape(0); ++i) {

sketches_[j].update(data(i, j));

}

}

} else { // nb::c_contig or nb::any_contig

for (uint32_t i = 0; i < items.shape(0); ++i) {

const size_t offset = i * items.shape(0);

for (uint32_t j = 0; j < d_; ++j) {

sketches_[j].update(data(i, j));

}

}

}

}

else {

throw std::invalid_argument("Update input must be 2 or fewer dimensions : " + std::to_string(ndim));

}

}

// Merges two arrays of sketches

// Currently: all values must be present

template<typename T, typename C>

void vector_of_kll_sketches<T, C>::merge(const vector_of_kll_sketches<T>& other) {

if (d_ != other.get_d()) {

throw std::invalid_argument("Must have same number of dimensions to merge: " + std::to_string(d_)

+ " vs " + std::to_string(other.d_));

} else {

for (uint32_t i = 0; i < d_; ++i) {

sketches_[i].merge(other.sketches_[i]);

}

}

}

template<typename T, typename C>

kll_sketch<T, C> vector_of_kll_sketches<T, C>::collapse(ArrInputType<int>& isk) const {

Array1D<int> indices = input_to_vec<int>(isk);

Array1D<uint32_t> index_arr = get_indices(indices);

auto inds = index_arr.view();

kll_sketch<T, C> result(k_);

for (size_t idx = 0; idx < inds.shape(0); ++idx) {

result.merge(sketches_[inds(idx)]);

}

return result;

}

// Number of updates for each sketch

template<typename T, typename C>

auto vector_of_kll_sketches<T, C>::get_n() const -> Array1D<uint64_t> {

auto vals = make_ndarray<uint64_t>(d_);

auto view = vals.view();

for (uint32_t i = 0; i < d_; ++i) {

view(i) = sketches_[i].get_n();

}

return vals;

}

// Number of retained values for each sketch

template<typename T, typename C>

auto vector_of_kll_sketches<T, C>::get_num_retained() const -> Array1D<uint32_t> {

auto vals = make_ndarray<uint32_t>(d_);

auto view = vals.view();

for (uint32_t i = 0; i < d_; ++i) {

view(i) = sketches_[i].get_num_retained();

}

return vals;

}

// Gets the minimum value of each sketch

// TODO: allow subsets of sketches

template<typename T, typename C>

auto vector_of_kll_sketches<T, C>::get_min_values() const -> Array1D<T> {

//std::vector<T> vals(d_);

auto vals = make_ndarray<T>(d_);

auto view = vals.view();

for (uint32_t i = 0; i < d_; ++i) {

view(i) = sketches_[i].get_min_item();

}

return vals;

}

// Gets the maximum value of each sketch

// TODO: allow subsets of sketches

template<typename T, typename C>

auto vector_of_kll_sketches<T, C>::get_max_values() const -> Array1D<T> {

auto vals = make_ndarray<T>(d_);

auto view = vals.view();

for (uint32_t i = 0; i < d_; ++i) {

view(i) = sketches_[i].get_max_item();

}

return vals;

}

// Summary of each sketch as one long string

// Users should use .split('\n\n') when calling it to build a list of each

// sketch's summary

template<typename T, typename C>

std::string vector_of_kll_sketches<T, C>::to_string(bool print_levels, bool print_items) const {

std::ostringstream ss;

for (uint32_t i = 0; i < d_; ++i) {

// all streams into 1 string, for compatibility with Python's str() behavior

// users will need to split by \n\n, e.g., str(kll).split('\n\n')

if (i > 0) ss << "\n";

ss << sketches_[i].to_string(print_levels, print_items);

}

return ss.str();

}

template<typename T, typename C>

auto vector_of_kll_sketches<T, C>::is_estimation_mode() const -> Array1D<bool> {

auto vals = make_ndarray<bool>(d_);

auto view = vals.view();

for (uint32_t i = 0; i < d_; ++i) {

view(i) = sketches_[i].is_estimation_mode();

}

return vals;

}

// Value of sketch(es) corresponding to some quantile(s)

template<typename T, typename C>

auto vector_of_kll_sketches<T, C>::get_quantiles(ArrInputType<double>& ranks,

ArrInputType<int>& isk) const -> Array2D<T> {

Array1D<int> indices = input_to_vec<int>(isk);

Array1D<uint32_t> inds = get_indices(indices);

size_t num_sketches = inds.size();

Array1D<double> ranks_arr = input_to_vec<double>(ranks);

size_t num_quantiles = ranks_arr.size();

auto quants = make_ndarray<T>(num_sketches, num_quantiles);

auto view = quants.view();

auto ranks_view = ranks_arr.view();

for (uint32_t i = 0; i < num_sketches; ++i) {

for (size_t j = 0; j < num_quantiles; ++j) {

view(i, j) = sketches_[inds(i)].get_quantile(ranks_view(j));

}

}

return quants;

}

// Value of sketch(es) corresponding to some rank(s)

template<typename T, typename C>

auto vector_of_kll_sketches<T, C>::get_ranks(ArrInputType<T>& values,

ArrInputType<int>& isk) const -> Array2D<double> {

Array1D<int> indices = input_to_vec<int>(isk);

Array1D<uint32_t> inds = get_indices(indices);

size_t num_sketches = inds.size();

Array1D<T> values_arr = input_to_vec<T>(values);

size_t num_ranks = values_arr.size();

auto vals = values_arr.view();

auto ranks = make_ndarray<double>(num_sketches, num_ranks);

auto view = ranks.view();

for (uint32_t i = 0; i < num_sketches; ++i) {

const size_t offset = i * num_ranks;

for (size_t j = 0; j < num_ranks; ++j) {

view(i, j) = sketches_[inds(i)].get_rank(vals(j));

}

}

return ranks;

}

// PMF(s) of sketch(es)

template<typename T, typename C>

auto vector_of_kll_sketches<T, C>::get_pmf(ArrInputType<T>& split_points,

ArrInputType<int>& isk) const -> Array2D<double> {

Array1D<int> indices = input_to_vec<int>(isk);

Array1D<uint32_t> inds = get_indices(indices);

size_t num_sketches = inds.size();

Array1D<T> splits_arr = input_to_vec<T>(split_points);

size_t num_splits = splits_arr.size();

auto pmfs = make_ndarray<double>(num_sketches, num_splits + 1);

auto view = pmfs.view();

for (uint32_t i = 0; i < num_sketches; ++i) {

auto pmf = sketches_[inds(i)].get_PMF(splits_arr.data(), num_splits);

for (size_t j = 0; j <= num_splits; ++j) {

view(i, j) = pmf[j];

}

}

return pmfs;

}

// CDF(s) of sketch(es)

template<typename T, typename C>

auto vector_of_kll_sketches<T, C>::get_cdf(ArrInputType<T>& split_points,

ArrInputType<int>& isk) const -> Array2D<double> {

Array1D<int> indices = input_to_vec<int>(isk);

Array1D<uint32_t> inds = get_indices(indices);

size_t num_sketches = inds.size();

Array1D<T> splits_arr = input_to_vec<T>(split_points);

size_t num_splits = splits_arr.size();

auto cdfs = make_ndarray<double>(num_sketches, num_splits + 1);

auto view = cdfs.view();

for (uint32_t i = 0; i < num_sketches; ++i) {

auto cdf = sketches_[inds(i)].get_CDF(splits_arr.data(), num_splits);

for (size_t j = 0; j <= num_splits; ++j) {

view(i, j) = cdf[j];

}

}

return cdfs;

}

template<typename T, typename C>

void vector_of_kll_sketches<T, C>::deserialize(const nb::bytes& sk_bytes,

uint32_t idx) {

if (idx >= d_) {

throw std::invalid_argument("request for invalid dimensions >= d ("

+ std::to_string(d_) +"): "+ std::to_string(idx));

}

// load the sketch into the proper index

sketches_[idx] = std::move(kll_sketch<T>::deserialize(sk_bytes.c_str(), sk_bytes.size()));

}

template<typename T, typename C>

nb::list vector_of_kll_sketches<T, C>::serialize(ArrInputType<int>& isk) {

Array1D<int> indices = input_to_vec<int>(isk);

Array1D<uint32_t> inds = get_indices(indices);

const size_t num_sketches = inds.size();

nb::list list;

for (uint32_t i = 0; i < num_sketches; ++i) {

auto serResult = sketches_[inds(i)].serialize();

list.append(nb::bytes((char*)serResult.data(), serResult.size()));

}

return list;

}

} // namespace datasketches

template<typename T>

void bind_vector_of_kll_sketches(nb::module_ &m, const char* name) {

using namespace datasketches;

nb::class_<vector_of_kll_sketches<T>>(m, name)

.def(nb::init<uint32_t, uint32_t>(), nb::arg("k")=vector_of_kll_constants::DEFAULT_K,

nb::arg("d")=vector_of_kll_constants::DEFAULT_D,

"Creates a new Vector of KLL Sketches instance with the given values of k and d.\n\n"

":param k: The value of k for every sketch in the vector\n:type k: int\n"

":param d: The number of sketches in the vector\n:type d: int"

)

.def("__copy__", [](const vector_of_kll_sketches<T>& sk){ return vector_of_kll_sketches<T>(sk); })

// allow user to retrieve k or d, in case it's instantiated w/ defaults

.def_prop_ro("k", &vector_of_kll_sketches<T>::get_k,

"The value of `k` of the sketch(es)")

.def_prop_ro("d", &vector_of_kll_sketches<T>::get_d,

"The number of sketches")

.def("update", &vector_of_kll_sketches<T>::update, nb::arg("items"), nb::arg("order") = "C",

"Updates the sketch(es) with value(s). Must be a 1D array of size equal to the number of sketches. Can also be 2D array of shape (n_updates, n_sketches). If a sketch does not have a value to update, use np.nan. "

" Order 'F' specifies a column-major (Fortran style) matrix; any other value assumes row-major (C style) matrix.")

.def("__str__", &vector_of_kll_sketches<T>::to_string, nb::arg("print_levels")=false, nb::arg("print_items")=false,

"Produces a string summary of all sketches. Users should split the returned string by '\n\n'")

.def("to_string", &vector_of_kll_sketches<T>::to_string, nb::arg("print_levels")=false,

nb::arg("print_items")=false,

"Produces a string summary of all sketches. Users should split the returned string by '\n\n'")

.def("is_empty", &vector_of_kll_sketches<T>::is_empty,

"Returns whether the sketch(es) is(are) empty of not")

.def("get_n", &vector_of_kll_sketches<T>::get_n,

"Returns the number of values seen by the sketch(es)")

.def("get_num_retained", &vector_of_kll_sketches<T>::get_num_retained,

"Returns the number of values retained by the sketch(es)")

.def("is_estimation_mode", &vector_of_kll_sketches<T>::is_estimation_mode,

"Returns whether the sketch(es) is(are) in estimation mode")

.def("get_min_values", &vector_of_kll_sketches<T>::get_min_values,

"Returns the minimum value(s) of the sketch(es)")

.def("get_max_values", &vector_of_kll_sketches<T>::get_max_values,

"Returns the maximum value(s) of the sketch(es)")

.def("get_quantiles", &vector_of_kll_sketches<T>::get_quantiles, nb::arg("ranks"),

nb::arg("isk")=-1,

"Returns the value(s) associated with the specified quantile(s) for the specified sketch(es). `ranks` can be a float between 0 and 1 (inclusive), or a list/array of values. `isk` specifies which sketch(es) to return the value(s) for (default: all sketches)")

.def("get_ranks", &vector_of_kll_sketches<T>::get_ranks,

nb::arg("value"), nb::arg("isk")=-1,

"Returns the value(s) associated with the specified rank(s) for the specified sketch(es). `values` can be an int between 0 and the number of values retained, or a list/array of values. `isk` specifies which sketch(es) to return the value(s) for (default: all sketches)")

.def("get_pmf", &vector_of_kll_sketches<T>::get_pmf, nb::arg("split_points"), nb::arg("isk")=-1,

"Returns the probability mass function (PMF) at `split_points` of the specified sketch(es). `split_points` should be a list/array of floats between 0 and 1 (inclusive). `isk` specifies which sketch(es) to return the PMF for (default: all sketches)")

.def("get_cdf", &vector_of_kll_sketches<T>::get_cdf, nb::arg("split_points"), nb::arg("isk")=-1,

"Returns the cumulative distribution function (CDF) at `split_points` of the specified sketch(es). `split_points` should be a list/array of floats between 0 and 1 (inclusive). `isk` specifies which sketch(es) to return the CDF for (default: all sketches)")

.def_static("get_normalized_rank_error",

[](uint16_t k, bool pmf) { return kll_sketch<T>::get_normalized_rank_error(k, pmf); },

nb::arg("k"), nb::arg("as_pmf"), "Returns the normalized rank error")

.def("serialize", &vector_of_kll_sketches<T>::serialize, nb::arg("isk")=-1,

"Serializes the specified sketch(es). `isk` can be an int or a list/array of ints (default: all sketches)")

.def("deserialize", &vector_of_kll_sketches<T>::deserialize, nb::arg("skBytes"), nb::arg("isk"),

"Deserializes the specified sketch. `isk` must be an int.")

.def("merge", &vector_of_kll_sketches<T>::merge, nb::arg("array_of_sketches"),

"Merges the input array of KLL sketches into the existing array.")

.def("collapse", &vector_of_kll_sketches<T>::collapse, nb::arg("isk")=-1,

"Returns the result of collapsing all sketches in the array into a single sketch. 'isk' can be an int or a list/array of ints (default: all sketches)")

;

}

void init_vector_of_kll(nb::module_ &m) {

bind_vector_of_kll_sketches<int>(m, "vector_of_kll_ints_sketches");

bind_vector_of_kll_sketches<float>(m, "vector_of_kll_floats_sketches");

}