Skip to content

Commit 9ad79eb

Browse files
committed
Split up the customized saves into three separate functions.
One per type, so users don't have to define all possible combinations if they're just interested in customizing the save for one of the types.
1 parent 86e9a41 commit 9ad79eb

4 files changed

Lines changed: 73 additions & 17 deletions

File tree

include/knncolle_annoy/Annoy.hpp

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -336,9 +336,20 @@ class AnnoyPrebuilt final : public knncolle::Prebuilt<Index_, Data_, Distance_>
336336
const auto dname = get_distance_name<AnnoyDistance_>();
337337
knncolle::quick_save(prefix + "distance", dname, std::strlen(dname));
338338

339-
auto& cust = customize_save_for_annoy_types<AnnoyDistance_, AnnoyIndex_, AnnoyData_>();
340-
if (cust) {
341-
cust(prefix);
339+
// Handling custom Annoy types.
340+
auto& icust = custom_save_for_annoy_index<AnnoyIndex_>();
341+
if (icust) {
342+
icust(prefix);
343+
}
344+
345+
auto& dcust = custom_save_for_annoy_data<AnnoyData_>();
346+
if (dcust) {
347+
dcust(prefix);
348+
}
349+
350+
auto& dscust = custom_save_for_annoy_distance<AnnoyDistance_>();
351+
if (dscust) {
352+
dscust(prefix);
342353
}
343354

344355
// Not bothering to save anything else; the RNG and thread policy

include/knncolle_annoy/load_annoy_prebuilt.hpp

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -62,15 +62,17 @@ inline AnnoyPrebuiltTypes load_annoy_prebuilt_types(const std::string& prefix) {
6262
* Helper function to define a `knncolle::LoadPrebuiltFunction` for Annoy in `knncolle::load_prebuilt()`.
6363
*
6464
* To load an Annoy index from disk, users are expected to define and register an Annoy-specific `knncolle::LoadPrebuiltFunction`.
65-
* In this function, users should first call `load_annoy_prebuilt_types()` to figure out the saved index's `AnnoyDistance_`, `AnnoyIndex` and `AnnoyData_`.
65+
* In this function, users should call `load_annoy_prebuilt_types()` to figure out the saved index's `AnnoyDistance_`, `AnnoyIndex` and `AnnoyData_`.
6666
* Then, they should call `load_annoy_prebuilt()` with the appropriate types to return a pointer to a `knncolle::Prebuilt` object.
6767
* This user-defined function should be registered in `load_prebuilt_registry()` with the key in `knncolle_annoy::save_name`.
6868
*
6969
* We do not define a default function for loading Annoy indices as there are too many possible combinations of types.
7070
* Instead, the user is responsible for deciding which combinations of types should be handled.
71-
* This avoids binary bloat from repeated instantiations of the Annoy template classes, if an application only deals with a certain subset of combinations.
72-
* For types or distances that are unknown to `knncolle::get_numeric_type()` or `get_distance_name()`, respectively,
73-
* users can store additional information on disk via `customize_save_for_annoy_types()` for use in loading.
71+
* This avoids binary bloat from repeated instantiations of the Annoy template classes, if the user's application only deals with a certain subset of combinations.
72+
*
73+
* For unknown types or distances, users can set `custom_save_for_annoy_index()`, `custom_save_for_annoy_data()` and/or `custom_save_for_annoy_distance()`.
74+
* Each custom function saves additional information about its type to disk during a `knncolle::Prebuilt::save()` call.
75+
* That information can then be parsed in the user-defined `knncolle::LoadPrebuiltFunction` to recover an Annoy index with the appropriate template types.
7476
*
7577
* @tparam Index_ Integer type for the observation indices.
7678
* @tparam Data_ Numeric type for the input and query data.

include/knncolle_annoy/utils.hpp

Lines changed: 40 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -60,19 +60,52 @@ const char* get_distance_name() {
6060
}
6161

6262
/**
63-
* Define a customized saving function to preserve type information from the Annoy index in `knncolle::Prebuilt::save()`.
64-
* Users can provide their own function here, to handle types that are unknown to `knncolle::get_numeric_type()` or `get_distance_name()`.
63+
* Define a customized saving function to preserve `AnnoyIndex_` type information when saving a prebuilt Annoy index in `knncolle::Prebuilt::save()`.
64+
* Users can provide their own function here to handle an `AnnoyIndex_` type that is unknown to `knncolle::get_numeric_type()`.
6565
* Any modifications to this function are not thread-safe and should be done in a serial section.
6666
*
67-
* @tparam AnnoyDistance_ An **Annoy**-compatible class to compute the distance between vectors.
6867
* @tparam AnnoyIndex_ Integer type for the observation indices in the Annoy index.
68+
*
69+
* @return A global function for saving information about `AnnoyIndex_`.
70+
* By default, no function is provided.
71+
* If set, the function will be called by the `knncolle::Prebuilt::save()` method for the Annoy `Prebuilt` subclass.
72+
*/
73+
template<class AnnoyIndex_>
74+
std::function<void(const std::string&)>& custom_save_for_annoy_index() {
75+
static std::function<void(const std::string&)> fun;
76+
return fun;
77+
}
78+
79+
/**
80+
* Define a customd saving function to preserve `AnnoyData_` type information when saving a prebuilt Annoy index in `knncolle::Prebuilt::save()`.
81+
* Users can provide their own function here to handle an `AnnoyData_` type that is unknown to `knncolle::get_numeric_type()`.
82+
* Any modifications to this function are not thread-safe and should be done in a serial section.
83+
*
6984
* @tparam AnnoyData_ Floating-point type for data in the Annoy index.
7085
*
71-
* @return A global function for saving information about `AnnoyDistance_`, `AnnoyIndex_` and `AnnoyData_`.
72-
* If set, this will be called by the `knncolle::Prebuilt::save()` method for the Annoy `Prebuilt` subclass.
86+
* @return A global function for saving information about `AnnoyData_`.
87+
* By default, no function is provided.
88+
* If set, the function will be called by the `knncolle::Prebuilt::save()` method for the Annoy `Prebuilt` subclass.
89+
*/
90+
template<class AnnoyData_>
91+
std::function<void(const std::string&)>& custom_save_for_annoy_data() {
92+
static std::function<void(const std::string&)> fun;
93+
return fun;
94+
}
95+
96+
/**
97+
* Define a customd saving function to preserve `AnnoyDistance_` type information when saving a prebuilt Annoy index in `knncolle::Prebuilt::save()`.
98+
* Users can provide their own function here to handle an `AnnoyDistance_` type that is unknown to `get_distance_name()`.
99+
* Any modifications to this function are not thread-safe and should be done in a serial section.
100+
*
101+
* @tparam AnnoyDistance_ An **Annoy**-compatible class to compute the distance between vectors.
102+
*
103+
* @return A global function for saving information about `AnnoyDistance_`.
104+
* By default, no function is provided.
105+
* If set, the function will be called by the `knncolle::Prebuilt::save()` method for the Annoy `Prebuilt` subclass.
73106
*/
74-
template<class AnnoyDistance_, typename AnnoyIndex_, typename AnnoyData_>
75-
std::function<void(const std::string&)>& customize_save_for_annoy_types() {
107+
template<class AnnoyDistance_>
108+
std::function<void(const std::string&)>& custom_save_for_annoy_distance() {
76109
static std::function<void(const std::string&)> fun;
77110
return fun;
78111
}

tests/src/load_annoy_prebuilt.cpp

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -75,10 +75,15 @@ TEST_F(AnnoyLoadPrebuiltTest, Manhattan) {
7575
}
7676

7777
TEST_F(AnnoyLoadPrebuiltTest, Custom) {
78-
auto& cust = knncolle_annoy::customize_save_for_annoy_types<Annoy::Euclidean, int, float>();
79-
cust = [](const std::string& prefix) -> void {
78+
knncolle_annoy::custom_save_for_annoy_index<int>() = [](const std::string& prefix) -> void {
8079
knncolle::quick_save(prefix + "FOO", "bar", 3);
8180
};
81+
knncolle_annoy::custom_save_for_annoy_data<float>() = [](const std::string& prefix) -> void {
82+
knncolle::quick_save(prefix + "WHEE", "stuff", 5);
83+
};
84+
knncolle_annoy::custom_save_for_annoy_distance<Annoy::Euclidean>() = [](const std::string& prefix) -> void {
85+
knncolle::quick_save(prefix + "YAY", "blah", 4);
86+
};
8287

8388
knncolle_annoy::AnnoyBuilder<int, double, double, Annoy::Euclidean> ab;
8489
auto aptr = ab.build_unique(knncolle::SimpleMatrix<int, double>(ndim, nobs, data.data()));
@@ -88,7 +93,12 @@ TEST_F(AnnoyLoadPrebuiltTest, Custom) {
8893

8994
// Custom function is respected.
9095
EXPECT_EQ(knncolle::quick_load_as_string(prefix + "FOO"), "bar");
91-
cust = nullptr;
96+
EXPECT_EQ(knncolle::quick_load_as_string(prefix + "WHEE"), "stuff");
97+
EXPECT_EQ(knncolle::quick_load_as_string(prefix + "YAY"), "blah");
98+
99+
knncolle_annoy::custom_save_for_annoy_index<int>() = nullptr;
100+
knncolle_annoy::custom_save_for_annoy_data<float>() = nullptr;
101+
knncolle_annoy::custom_save_for_annoy_distance<Annoy::Euclidean>() = nullptr;
92102

93103
// Everything else is still fine.
94104
auto reloaded = knncolle::load_prebuilt_shared<int, double, double>(prefix);

0 commit comments

Comments
 (0)