void DefineKmvSketchClass()

in tfx_bsl/cc/sketches/sketches_submodule.cc [37:108]

• 69 lines of code
• 5 McCabe index (conditional complexity)

void DefineKmvSketchClass(py::module sketch_module) {
  py::class_<KmvSketch>(sketch_module, "KmvSketch")
      .def(py::init<const int&>())
      .def(
          "AddValues",
          [](KmvSketch& sketch, const std::shared_ptr<arrow::Array>& array) {
            absl::Status s = sketch.AddValues(*array);
            if (!s.ok()) {
              throw std::runtime_error(s.ToString());
            }
          },
          py::doc("Updates the sketch with an Arrow array of values."),
          py::call_guard<py::gil_scoped_release>())
      .def(
          "Merge",
          [](KmvSketch& sketch, KmvSketch& other) {
            absl::Status s = sketch.Merge(other);
            if (!s.ok()) {
              throw std::runtime_error(s.ToString());
            }
          },
          py::doc("Merges another KMV sketch into this sketch. Returns error "
                  "if the other sketch has a different number of buckets than "
                  "this sketch."),
          py::call_guard<py::gil_scoped_release>())
      .def("Estimate", &KmvSketch::Estimate,
           py::doc("Estimates the number of distinct elements."),
           py::call_guard<py::gil_scoped_release>())
      .def(
          "Serialize",
          [](KmvSketch& sketch) {
            std::string serialized;
            {
              // Release the GIL during the call to Serialize
              py::gil_scoped_release release_gil;
              serialized = sketch.Serialize();
            }
            return py::bytes(serialized);
          },
          py::doc("Serializes the sketch as a string."))
      .def_static(
          "Deserialize",
          [](absl::string_view byte_string) {
            std::unique_ptr<KmvSketch> result;
            absl::Status s = KmvSketch::Deserialize(byte_string, &result);
            if (!s.ok()) throw std::runtime_error(s.ToString());
            return result;
          },
          py::doc("Deserializes the string to a KmvSketch object."),
          py::call_guard<py::gil_scoped_release>())
      // Pickle support
      .def(py::pickle(
          [](KmvSketch& sketch) {  // __getstate__
            std::string serialized;
            {
              // Release the GIL during the call to Serialize
              py::gil_scoped_release release_gil;
              serialized = sketch.Serialize();
            }
            return py::bytes(serialized);
          },
          [](py::bytes byte_string) {  // __setstate__
            char* data;
            Py_ssize_t size;
            PyBytes_AsStringAndSize(byte_string.ptr(), &data, &size);
            std::unique_ptr<KmvSketch> result;
            absl::Status s =
                KmvSketch::Deserialize(absl::string_view(data, size), &result);
            if (!s.ok()) throw std::runtime_error(s.ToString());
            return result;
          }));
}