| /* |
| * 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 "kll_sketch.hpp" |
| |
| #include <pybind11/pybind11.h> |
| #include <pybind11/stl.h> |
| #include <pybind11/numpy.h> |
| #include <sstream> |
| #include <vector> |
| |
| namespace py = pybind11; |
| |
| namespace datasketches { |
| |
| // Wrapper class for Numpy compatibility |
| template <typename T, typename C = std::less<T>, typename S = serde<T>> |
| class vector_of_kll_sketches { |
| public: |
| static const uint32_t DEFAULT_K = kll_sketch<T, C, S>::DEFAULT_K; |
| static const uint32_t DEFAULT_D = 1; |
| |
| explicit vector_of_kll_sketches(uint32_t k = DEFAULT_K, uint32_t d = 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,S>& operator=(const vector_of_kll_sketches& other); |
| vector_of_kll_sketches<T,C,S>& operator=(vector_of_kll_sketches&& other); |
| |
| // container parameters |
| inline uint32_t get_k() const; |
| inline uint32_t get_d() const; |
| |
| // sketch updates/merges |
| void update(const py::array_t<T>& items); |
| void merge(const vector_of_kll_sketches<T>& other); |
| |
| // returns a single sketch combining all data in the array |
| kll_sketch<T,C,S> collapse(const py::array_t<int>& isk) const; |
| |
| // sketch queries returning an array of results |
| py::array is_empty() const; |
| py::array get_n() const; |
| py::array is_estimation_mode() const; |
| py::array get_min_values() const; |
| py::array get_max_values() const; |
| py::array get_num_retained() const; |
| py::array get_quantiles(const py::array_t<double>& fractions, const py::array_t<int>& isk) const; |
| py::array get_ranks(const py::array_t<T>& values, const py::array_t<int>& isk) const; |
| py::array get_pmf(const py::array_t<T>& split_points, const py::array_t<int>& isk) const; |
| py::array get_cdf(const py::array_t<T>& split_points, const py::array_t<int>& isk) const; |
| |
| // human-readable output |
| std::string to_string(bool print_levels = false, bool print_items = false) const; |
| |
| // binary output/input |
| py::list serialize(py::array_t<uint32_t>& isk); |
| // note: deserialize() replaces the sketch at the specified |
| // index. Not a static method. |
| void deserialize(const py::bytes& sk_bytes, uint32_t idx); |
| |
| private: |
| std::vector<uint32_t> get_indices(const py::array_t<int>& isk) 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,S>> sketches_; |
| }; |
| |
| template<typename T, typename C, typename S> |
| vector_of_kll_sketches<T,C,S>::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, typename S> |
| vector_of_kll_sketches<T,C,S>::vector_of_kll_sketches(const vector_of_kll_sketches& other) : |
| k_(other.k_), |
| d_(other.d_), |
| sketches_(other.sketches_) |
| {} |
| |
| template<typename T, typename C, typename S> |
| vector_of_kll_sketches<T,C,S>::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, typename S> |
| vector_of_kll_sketches<T,C,S>& vector_of_kll_sketches<T,C,S>::operator=(const vector_of_kll_sketches& other) { |
| vector_of_kll_sketches<T,C,S> copy(other); |
| k_ = copy.k_; |
| d_ = copy.d_; |
| std::swap(sketches_, copy.sketches_); |
| return *this; |
| } |
| |
| template<typename T, typename C, typename S> |
| vector_of_kll_sketches<T,C,S>& vector_of_kll_sketches<T,C,S>::operator=(vector_of_kll_sketches&& other) { |
| k_ = other.k_; |
| d_ = other.d_; |
| std::swap(sketches_, other.sketches_); |
| return *this; |
| } |
| |
| template<typename T, typename C, typename S> |
| uint32_t vector_of_kll_sketches<T,C,S>::get_k() const { |
| return k_; |
| } |
| |
| template<typename T, typename C, typename S> |
| uint32_t vector_of_kll_sketches<T,C,S>::get_d() const { |
| return d_; |
| } |
| |
| template<typename T, typename C, typename S> |
| std::vector<uint32_t> vector_of_kll_sketches<T,C,S>::get_indices(const py::array_t<int>& isk) const { |
| std::vector<uint32_t> indices; |
| if (isk.size() == 1) { |
| auto data = isk.unchecked(); |
| if (data(0) == -1) { |
| indices.reserve(d_); |
| for (uint32_t i = 0; i < d_; ++i) { |
| indices.push_back(i); |
| } |
| } else { |
| indices.push_back(static_cast<uint32_t>(data(0))); |
| } |
| } else { |
| auto data = isk.unchecked<1>(); |
| indices.reserve(isk.size()); |
| for (uint32_t i = 0; i < isk.size(); ++i) { |
| const uint32_t idx = static_cast<uint32_t>(data(i)); |
| if (idx < d_) { |
| indices.push_back(idx); |
| } else { |
| throw std::invalid_argument("request for invalid dimenions >= d (" |
| + std::to_string(d_) +"): "+ std::to_string(idx)); |
| } |
| } |
| } |
| return indices; |
| } |
| |
| // Checks if each sketch is empty or not |
| template<typename T, typename C, typename S> |
| py::array vector_of_kll_sketches<T,C,S>::is_empty() const { |
| std::vector<bool> vals(d_); |
| for (uint32_t i = 0; i < d_; ++i) { |
| vals[i] = sketches_[i].is_empty(); |
| } |
| |
| return py::cast(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, typename S> |
| void vector_of_kll_sketches<T,C,S>::update(const py::array_t<T>& items) { |
| |
| 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 |
| auto data = items.template unchecked<1>(); |
| 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 |
| auto data = items.template unchecked<2>(); |
| if (items.flags() & py::array::f_style) { |
| for (uint32_t j = 0; j < d_; ++j) { |
| for (uint32_t i = 0; i < items.shape(0); ++i) { |
| sketches_[j].update(data(i,j)); |
| } |
| } |
| } else { // py::array::c_style or py::array::forcecast |
| for (uint32_t i = 0; i < items.shape(0); ++i) { |
| 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, typename S> |
| void vector_of_kll_sketches<T,C,S>::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, typename S> |
| kll_sketch<T,C,S> vector_of_kll_sketches<T,C,S>::collapse(const py::array_t<int>& isk) const { |
| std::vector<uint32_t> inds = get_indices(isk); |
| |
| kll_sketch<T,C,S> result(k_); |
| for (auto& idx : inds) { |
| result.merge(sketches_[idx]); |
| } |
| return result; |
| } |
| |
| // Number of updates for each sketch |
| template<typename T, typename C, typename S> |
| py::array vector_of_kll_sketches<T,C,S>::get_n() const { |
| std::vector<uint64_t> vals(d_); |
| for (uint32_t i = 0; i < d_; ++i) { |
| vals[i] = sketches_[i].get_n(); |
| } |
| return py::cast(vals); |
| } |
| |
| // Number of retained values for each sketch |
| template<typename T, typename C, typename S> |
| py::array vector_of_kll_sketches<T,C,S>::get_num_retained() const { |
| std::vector<uint32_t> vals(d_); |
| for (uint32_t i = 0; i < d_; ++i) { |
| vals[i] = sketches_[i].get_num_retained(); |
| } |
| return py::cast(vals); |
| } |
| |
| // Gets the minimum value of each sketch |
| // TODO: allow subsets of sketches |
| template<typename T, typename C, typename S> |
| py::array vector_of_kll_sketches<T,C,S>::get_min_values() const { |
| std::vector<T> vals(d_); |
| for (uint32_t i = 0; i < d_; ++i) { |
| vals[i] = sketches_[i].get_min_value(); |
| } |
| return py::cast(vals); |
| } |
| |
| // Gets the maximum value of each sketch |
| // TODO: allow subsets of sketches |
| template<typename T, typename C, typename S> |
| py::array vector_of_kll_sketches<T,C,S>::get_max_values() const { |
| std::vector<T> vals(d_); |
| for (uint32_t i = 0; i < d_; ++i) { |
| vals[i] = sketches_[i].get_max_value(); |
| } |
| return py::cast(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, typename S> |
| std::string vector_of_kll_sketches<T,C,S>::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, typename S> |
| py::array vector_of_kll_sketches<T,C,S>::is_estimation_mode() const { |
| std::vector<bool> vals(d_); |
| for (uint32_t i = 0; i < d_; ++i) { |
| vals[i] = sketches_[i].is_estimation_mode(); |
| } |
| return py::cast(vals); |
| } |
| |
| // Value of sketch(es) corresponding to some quantile(s) |
| template<typename T, typename C, typename S> |
| py::array vector_of_kll_sketches<T,C,S>::get_quantiles(const py::array_t<double>& fractions, |
| const py::array_t<int>& isk) const { |
| std::vector<uint32_t> inds = get_indices(isk); |
| size_t num_sketches = inds.size(); |
| size_t num_quantiles = fractions.size(); |
| |
| std::vector<std::vector<T>> quants(num_sketches, std::vector<T>(num_quantiles)); |
| for (uint32_t i = 0; i < num_sketches; ++i) { |
| auto quant = sketches_[inds[i]].get_quantiles(fractions.data(), num_quantiles); |
| for (size_t j = 0; j < num_quantiles; ++j) { |
| quants[i][j] = quant[j]; |
| } |
| } |
| |
| return py::cast(quants); |
| } |
| |
| // Value of sketch(es) corresponding to some rank(s) |
| template<typename T, typename C, typename S> |
| py::array vector_of_kll_sketches<T,C,S>::get_ranks(const py::array_t<T>& values, |
| const py::array_t<int>& isk) const { |
| std::vector<uint32_t> inds = get_indices(isk); |
| size_t num_sketches = inds.size(); |
| size_t num_ranks = values.size(); |
| auto vals = values.data(); |
| |
| std::vector<std::vector<float>> ranks(num_sketches, std::vector<float>(num_ranks)); |
| for (uint32_t i = 0; i < num_sketches; ++i) { |
| for (size_t j = 0; j < num_ranks; ++j) { |
| ranks[i][j] = sketches_[inds[i]].get_rank(vals[j]); |
| } |
| } |
| |
| return py::cast(ranks); |
| } |
| |
| // PMF(s) of sketch(es) |
| template<typename T, typename C, typename S> |
| py::array vector_of_kll_sketches<T,C,S>::get_pmf(const py::array_t<T>& split_points, |
| const py::array_t<int>& isk) const { |
| std::vector<uint32_t> inds = get_indices(isk); |
| size_t num_sketches = inds.size(); |
| size_t num_splits = split_points.size(); |
| |
| std::vector<std::vector<T>> pmfs(num_sketches, std::vector<T>(num_splits + 1)); |
| for (uint32_t i = 0; i < num_sketches; ++i) { |
| auto pmf = sketches_[inds[i]].get_PMF(split_points.data(), num_splits); |
| for (size_t j = 0; j <= num_splits; ++j) { |
| pmfs[i][j] = pmf[j]; |
| } |
| } |
| |
| return py::cast(pmfs); |
| } |
| |
| // CDF(s) of sketch(es) |
| template<typename T, typename C, typename S> |
| py::array vector_of_kll_sketches<T,C,S>::get_cdf(const py::array_t<T>& split_points, |
| const py::array_t<int>& isk) const { |
| std::vector<uint32_t> inds = get_indices(isk); |
| size_t num_sketches = inds.size(); |
| size_t num_splits = split_points.size(); |
| |
| std::vector<std::vector<T>> cdfs(num_sketches, std::vector<T>(num_splits + 1)); |
| for (uint32_t i = 0; i < num_sketches; ++i) { |
| auto cdf = sketches_[inds[i]].get_CDF(split_points.data(), num_splits); |
| for (size_t j = 0; j <= num_splits; ++j) { |
| cdfs[i][j] = cdf[j]; |
| } |
| } |
| |
| return py::cast(cdfs); |
| } |
| |
| template<typename T, typename C, typename S> |
| void vector_of_kll_sketches<T,C,S>::deserialize(const py::bytes& sk_bytes, |
| uint32_t idx) { |
| if (idx >= d_) { |
| throw std::invalid_argument("request for invalid dimenions >= d (" |
| + std::to_string(d_) +"): "+ std::to_string(idx)); |
| } |
| std::string skStr = sk_bytes; // implicit cast |
| // load the sketch into the proper index |
| sketches_[idx] = std::move(kll_sketch<T>::deserialize(skStr.c_str(), skStr.length())); |
| } |
| |
| template<typename T, typename C, typename S> |
| py::list vector_of_kll_sketches<T,C,S>::serialize(py::array_t<uint32_t>& isk) { |
| std::vector<uint32_t> inds = get_indices(isk); |
| const size_t num_sketches = inds.size(); |
| |
| py::list list(num_sketches); |
| for (uint32_t i = 0; i < num_sketches; ++i) { |
| auto serResult = sketches_[inds[i]].serialize(); |
| list[i] = py::bytes((char*)serResult.data(), serResult.size()); |
| } |
| |
| return list; |
| } |
| |
| namespace python { |
| template<typename T> |
| double kll_sketch_generic_normalized_rank_error(uint16_t k, bool pmf) { |
| return kll_sketch<T>::get_normalized_rank_error(k, pmf); |
| } |
| |
| } // namespace datasketches::python |
| |
| } // namespace datasketches |
| |
| namespace dspy = datasketches::python; |
| |
| template<typename T> |
| void bind_vector_of_kll_sketches(py::module &m, const char* name) { |
| using namespace datasketches; |
| |
| py::class_<vector_of_kll_sketches<T>>(m, name) |
| .def(py::init<uint32_t, uint32_t>(), py::arg("k")=vector_of_kll_sketches<T>::DEFAULT_K, |
| py::arg("d")=vector_of_kll_sketches<T>::DEFAULT_D) |
| .def(py::init<const vector_of_kll_sketches<T>&>()) |
| // allow user to retrieve k or d, in case it's instantiated w/ defaults |
| .def("get_k", &vector_of_kll_sketches<T>::get_k, |
| "Returns the value of `k` of the sketch(es)") |
| .def("get_d", &vector_of_kll_sketches<T>::get_d, |
| "Returns the number of sketches") |
| .def("update", &vector_of_kll_sketches<T>::update, py::arg("items"), |
| "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") |
| .def("__str__", &vector_of_kll_sketches<T>::to_string, py::arg("print_levels")=false, py::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, py::arg("print_levels")=false, |
| py::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, py::arg("fractions"), |
| py::arg("isk")=-1, |
| "Returns the value(s) associated with the specified quantile(s) for the specified sketch(es). `fractions` 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, py::arg("values"), |
| py::arg("isk")=-1, |
| "Returns the value(s) associated with the specified ranks(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, py::arg("split_points"), py::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, py::arg("split_points"), py::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", &dspy::kll_sketch_generic_normalized_rank_error<T>, |
| py::arg("k"), py::arg("as_pmf"), "Returns the normalized rank error") |
| .def("serialize", &vector_of_kll_sketches<T>::serialize, py::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, py::arg("skBytes"), py::arg("isk"), |
| "Deserializes the specified sketch. `isk` must be an int.") |
| .def("merge", &vector_of_kll_sketches<T>::merge, py::arg("array_of_sketches"), |
| "Merges the input array of KLL sketches into the existing array.") |
| .def("collapse", &vector_of_kll_sketches<T>::collapse, py::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(py::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"); |
| } |