-
Notifications
You must be signed in to change notification settings - Fork 8
Expand file tree
/
Copy pathvo_wrapper.cpp
More file actions
173 lines (152 loc) · 6.9 KB
/
Copy pathvo_wrapper.cpp
File metadata and controls
173 lines (152 loc) · 6.9 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
/*
* 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 "var_opt_sketch.hpp"
#include "var_opt_union.hpp"
#include "py_serde.hpp"
#include <pybind11/pybind11.h>
namespace py = pybind11;
namespace datasketches {
namespace python {
template<typename T>
var_opt_sketch<T> vo_sketch_deserialize(py::bytes& skBytes, py_object_serde& sd) {
std::string skStr = skBytes; // implicit cast
return var_opt_sketch<T>::deserialize(skStr.c_str(), skStr.length(), sd);
}
template<typename T>
py::object vo_sketch_serialize(const var_opt_sketch<T>& sk, py_object_serde& sd) {
auto serResult = sk.serialize(0, sd);
return py::bytes((char*)serResult.data(), serResult.size());
}
template<typename T>
size_t vo_sketch_size_bytes(const var_opt_sketch<T>& sk, py_object_serde& sd) {
return sk.get_serialized_size_bytes(sd);
}
template<typename T>
var_opt_union<T> vo_union_deserialize(py::bytes& uBytes, py_object_serde& sd) {
std::string uStr = uBytes; // implicit cast
return var_opt_union<T>::deserialize(uStr.c_str(), uStr.length(), sd);
}
template<typename T>
py::object vo_union_serialize(const var_opt_union<T>& u, py_object_serde& sd) {
auto serResult = u.serialize(0, sd);
return py::bytes((char*)serResult.data(), serResult.size());
}
template<typename T>
size_t vo_union_size_bytes(const var_opt_union<T>& u, py_object_serde& sd) {
return u.get_serialized_size_bytes(sd);
}
template<typename T>
py::list vo_sketch_get_samples(const var_opt_sketch<T>& sk) {
py::list list;
for (auto item : sk) {
py::tuple t = py::make_tuple(item.first, item.second);
list.append(t);
}
return list;
}
template<typename T>
py::dict vo_sketch_estimate_subset_sum(const var_opt_sketch<T>& sk, const std::function<bool(T)> func) {
subset_summary summary = sk.estimate_subset_sum(func);
py::dict d;
d["estimate"] = summary.estimate;
d["lower_bound"] = summary.lower_bound;
d["upper_bound"] = summary.upper_bound;
d["total_sketch_weight"] = summary.total_sketch_weight;
return d;
}
template<typename T>
std::string vo_sketch_to_string(const var_opt_sketch<T>& sk, bool print_items) {
if (print_items) {
std::ostringstream ss;
ss << sk.to_string();
ss << "### VarOpt Sketch Items" << std::endl;
int i = 0;
for (auto item : sk) {
// item.second is always a double
// item.first is an arbitrary py::object, so get the value by
// using internal str() method then casting to C++ std::string
py::str item_pystr(item.first);
std::string item_str = py::cast<std::string>(item_pystr);
ss << i++ << ": " << item_str << "\twt = " << item.second << std::endl;
}
return ss.str();
} else {
return sk.to_string();
}
}
}
}
namespace dspy = datasketches::python;
template<typename T>
void bind_vo_sketch(py::module &m, const char* name) {
using namespace datasketches;
py::class_<var_opt_sketch<T>>(m, name)
.def(py::init<uint32_t>(), py::arg("k"))
.def("__str__", &dspy::vo_sketch_to_string<T>, py::arg("print_items")=false,
"Produces a string summary of the sketch")
.def("to_string", &dspy::vo_sketch_to_string<T>, py::arg("print_items")=false,
"Produces a string summary of the sketch")
.def("update", (void (var_opt_sketch<T>::*)(const T&, double)) &var_opt_sketch<T>::update, py::arg("item"), py::arg("weight")=1.0,
"Updates the sketch with the given value and weight")
.def_property_readonly("k", &var_opt_sketch<T>::get_k,
"Returns the sketch's maximum configured sample size")
.def_property_readonly("n", &var_opt_sketch<T>::get_n,
"Returns the total stream length")
.def_property_readonly("num_samples", &var_opt_sketch<T>::get_num_samples,
"Returns the number of samples currently in the sketch")
.def("get_samples", &dspy::vo_sketch_get_samples<T>,
"Returns the set of samples in the sketch")
.def("is_empty", &var_opt_sketch<T>::is_empty,
"Returns True if the sketch is empty, otherwise False")
.def("estimate_subset_sum", &dspy::vo_sketch_estimate_subset_sum<T>,
"Applies a provided predicate to the sketch and returns the estimated total weight matching the predicate, as well "
"as upper and lower bounds on the estimate and the total weight processed by the sketch")
.def("get_serialized_size_bytes", &dspy::vo_sketch_size_bytes<T>, py::arg("serde"),
"Computes the size in bytes needed to serialize the current sketch")
.def("serialize", &dspy::vo_sketch_serialize<T>, py::arg("serde"), "Serialize the var opt sketch using the provided serde")
.def_static("deserialize", &dspy::vo_sketch_deserialize<T>, py::arg("bytes"), py::arg("serde"),
"Constructs a var opt sketch from the given bytes using the provided serde")
.def("__iter__", [](const var_opt_sketch<T>& sk) { return py::make_iterator(sk.begin(), sk.end()); });
}
template<typename T>
void bind_vo_union(py::module &m, const char* name) {
using namespace datasketches;
py::class_<var_opt_union<T>>(m, name)
.def(py::init<uint32_t>(), py::arg("max_k"))
.def("__str__", &var_opt_union<T>::to_string,
"Produces a string summary of the sketch")
.def("to_string", &var_opt_union<T>::to_string,
"Produces a string summary of the sketch")
.def("update", (void (var_opt_union<T>::*)(const var_opt_sketch<T>& sk)) &var_opt_union<T>::update, py::arg("sketch"),
"Updates the union with the given sketch")
.def("get_result", &var_opt_union<T>::get_result,
"Returns a sketch corresponding to the union result")
.def("reset", &var_opt_union<T>::reset,
"Resets the union to the empty state")
.def("get_serialized_size_bytes", &dspy::vo_union_size_bytes<T>, py::arg("serde"),
"Computes the size in bytes needed to serialize the current sketch")
.def("serialize", &dspy::vo_union_serialize<T>, py::arg("serde"), "Serialize the var opt union using the provided serde")
.def_static("deserialize", &dspy::vo_union_deserialize<T>, py::arg("bytes"), py::arg("serde"),
"Constructs a var opt union from the given bytes using the provided serde")
;
}
void init_vo(py::module &m) {
bind_vo_sketch<py::object>(m, "var_opt_sketch");
bind_vo_union<py::object>(m, "var_opt_union");
}