Enable customization of saves based on Annoy-related types.

Author: LTLA

include/knncolle_annoy/Annoy.hpp

@@ -336,6 +336,11 @@ class AnnoyPrebuilt final : public knncolle::Prebuilt<Index_, Data_, Distance_>
         const auto dname = get_distance_name<AnnoyDistance_>();
         knncolle::quick_save(prefix + "distance", dname, std::strlen(dname));
 
+        auto& cust = customize_save_for_annoy_types<AnnoyDistance_, AnnoyIndex_, AnnoyData_>();
+        if (cust) {
+            cust(prefix);
+        }
+
         // Not bothering to save anything else; the RNG and thread policy
         // should not be relevant once the index is built.
 

include/knncolle_annoy/utils.hpp

@@ -7,6 +7,7 @@
 #include <cstddef>
 #include <stdexcept>
 #include <type_traits>
+#include <functional>
 
 #include "knncolle/knncolle.hpp"
 
@@ -47,6 +48,8 @@ struct has_name<AnnoyDistance_, decltype(AnnoyDistance_::name(), 0)> {
  * @tparam AnnoyDistance_ An **Annoy**-compatible class to compute the distance between vectors, as used in `AnnoyBuilder()`.
  * @return Name of the distance metric, e.g., `"euclidean"`, `"manhattan"`.
  * This is taken from `AnnoyDistance_::name()` if such a method exists, otherwise `"unknown"` is returned.
+ *
+ * For unknown distances, consider using `customize_save_for_annoy_types()` to add more information to the on-disk representation during a `knncolle::Prebuilt::save()` call.
  */
 template<typename AnnoyDistance_> 
 const char* get_distance_name() { 
@@ -79,6 +82,8 @@ enum class NumericType : char {
 /**
  * @tparam Type_ Some integer or floating-point type, typically used as `AnnoyIndex_` or `AnnoyData_` in `AnnoyBuilder()`.
  * @return Identity of the numeric type.
+ *
+ * For unknown types, consider using `customize_save_for_annoy_types()` to add more information to the on-disk representation during a `knncolle::Prebuilt::save()` call.
  */
 template<typename Type_>
 NumericType get_numeric_type() {
@@ -217,14 +222,29 @@ inline PrebuiltSaveConfig scan_prebuilt_save_config(const std::string& prefix) {
     PrebuiltSaveConfig config;
     config.index = types[0];
     config.data = types[1];
-
-    const std::string distpath = prefix + "distance";
-    std::ifstream input(distpath);
-    config.distance.insert(config.distance.end(), (std::istreambuf_iterator<char>(input)), (std::istreambuf_iterator<char>()) );
+    config.distance = knncolle::quick_load_as_string(prefix + "distance");
 
     return config;
 }
 
+/**
+ * Define a customized saving function to preserve type information from the Annoy index in `knncolle::Prebuilt::save()`.
+ * Users can provide their own function here, to handle types that are unknown to `get_numeric_type()` or `get_distance_name()`.
+ * Any modifications to this function are not thread-safe and should be done in a serial section. 
+ *
+ * @tparam AnnoyDistance_ An **Annoy**-compatible class to compute the distance between vectors. 
+ * @tparam AnnoyIndex_ Integer type for the observation indices in the Annoy index.
+ * @tparam AnnoyData_ Floating-point type for data in the Annoy index.
+ *
+ * @return A global function for saving information about `AnnoyDistance_`, `AnnoyIndex_` and `AnnoyData_`.
+ * If set, this will be called by the `knncolle::Prebuilt::save()` method for the Annoy `Prebuilt` subclass.
+ */
+template<class AnnoyDistance_, typename AnnoyIndex_, typename AnnoyData_>
+std::function<void(const std::string&)>& customize_save_for_annoy_types() {
+    static std::function<void(const std::string&)> fun;
+    return fun;
+}
+
 }
 
 #endif

tests/src/load_annoy_prebuilt.cpp

@@ -74,6 +74,37 @@ TEST_F(AnnoyLoadPrebuiltTest, Manhattan) {
     }
 }
 
+TEST_F(AnnoyLoadPrebuiltTest, Custom) {
+    auto& cust = knncolle_annoy::customize_save_for_annoy_types<Annoy::Euclidean, int, float>();
+    cust = [](const std::string& prefix) -> void {
+        knncolle::quick_save(prefix + "FOO", "bar", 3);
+    };
+
+    knncolle_annoy::AnnoyBuilder<int, double, double, Annoy::Euclidean> ab;
+    auto aptr = ab.build_unique(knncolle::SimpleMatrix<int, double>(ndim, nobs, data.data()));
+
+    const auto prefix = (savedir / "custom_").string();
+    aptr->save(prefix);
+
+    // Custom function is respected.
+    EXPECT_EQ(knncolle::quick_load_as_string(prefix + "FOO"), "bar");
+    cust = nullptr;
+
+    // Everything else is still fine.
+    auto reloaded = knncolle::load_prebuilt_shared<int, double, double>(prefix);
+    std::vector<int> output_i, output_i2;
+    std::vector<double> output_d, output_d2;
+
+    auto searcher = aptr->initialize();
+    auto researcher = reloaded->initialize();
+    for (int x = 0; x < nobs; ++x) {
+        searcher->search(x, 5, &output_i, &output_d);
+        researcher->search(x, 5, &output_i2, &output_d2);
+        EXPECT_EQ(output_i, output_i2);
+        EXPECT_EQ(output_d, output_d2);
+    }
+}
+
 TEST_F(AnnoyLoadPrebuiltTest, Error) {
     knncolle_annoy::AnnoyBuilder<int, double, double, Annoy::Manhattan> ab;
     auto aptr = ab.build_unique(knncolle::SimpleMatrix<int, double>(ndim, nobs, data.data()));