Update CAGRA serialization #1755
Conversation
This changes the serialization format of saved CAGRA instances by: * The dtype will now be written in the first 4 bytes of the index, to match the IVF methods and to make it easier to deduce the dtype from python (rapidsai#1729) * Writing out the dataset with the index is now optional. Since many use cases will already have the dataset written out separately, this gives us the option to save disk space by not writing out an extra copy of the input dataset. If the include_dataset=false option is given, you will have to call `index.update_dataset` to set the dataset yourself after loading
benfred
added
improvement
github-actions
bot
added
cpp
python
and removed
improvement
benfred
added
improvement
| @@ -706,6 +781,8 @@ def save(filename, Index index, handle=None): | |||
| Name of the file. | |||
| index : Index | |||
| Trained CAGRA index. | |||
| include_dataset : bool | |||
| Whether or not to write out the dataset along with the index | |||
There was a problem hiding this comment.
It might be useful to mention the implication here just to make it more obvious for the uninformed- like a warning that a dataset can get quite large so it's advisable to set this to false to shrink the size of the serialized index.
| @@ -258,6 +258,13 @@ struct index : ann::index { | |||
| dataset.data_handle(), dataset.extent(0), dataset.extent(1), dataset.extent(1)); | |||
| } | |||
| } | |||
| void update_dataset(raft::resources const& res, | |||
There was a problem hiding this comment.
We probably want to keep these const mdspans. If this is because of python, can we use make_const_mdspan() in that layer?
There was a problem hiding this comment.
I agree that automatically discarding const would be bad - but this is doing the opposite and is automatically adding it (like this is converting a non-const mdspan to a const msdpan), which I feel like is something that should be allowed with our API's.
The issue I have is that Cython kinda sucks with respecting const identifiers, which is why all our Cython api's use non-const mdspans right now. Like if I try to add a get_const_hmv_float (to parallel the non-const get_hmv_float we have now) - I get an error message from Cython, where it doesn't recognize const float as a type inside template parameters:
Error compiling Cython file:
------------------------------------------------------------
...
if cai.dtype != np.float32:
raise TypeError("dtype %s not supported" % cai.dtype)
if check_shape and len(cai.shape) != 2:
raise ValueError("Expected a 2D array, got %d D" % len(cai.shape))
shape = (cai.shape[0], cai.shape[1] if len(cai.shape) == 2 else 1)
return make_host_matrix_view[const float, int64_t, row_major](
^
------------------------------------------------------------
/home/ben/code/raft/python/pylibraft/pylibraft/common/mdspan.pyx:232:39: Expected ']', found 'float'
I can get around this by adding a Cython typedef (like ctypedef const float const_float) - but that introduces the need for other hacks later on (like cython will treat const_float and const float as separate types - meaning that when we define the update_dataset for Cython in c_cagra.pxd I can't just go const T as the type, and have to introduce a new template param =(. I've done this in the last commit - let me know what you think
|
|
||
| if dataset_ai.from_cai: | ||
| self.index[0].update_dataset(deref(handle_), | ||
| get_dmv_float(dataset_ai, |
There was a problem hiding this comment.
Here is where we could use make_const_mdspan. It would simplify things so that we don't need to make non-const functions everywhere (which kind of circumvents the const functions).
|
|
||
| if dataset_ai.from_cai: | ||
| self.index[0].update_dataset(deref(handle_), | ||
| get_dmv_int8(dataset_ai, |
| _check_input_array(dataset_ai, [np.dtype("ubyte")]) | ||
|
|
||
| if dataset_ai.from_cai: | ||
| self.index[0].update_dataset(deref(handle_), |
|
/merge |
This changes the serialization format of saved CAGRA indices by:
index.update_datasetto set the dataset yourself after loading