Refactored IVF-PQ Python API as a thin wrapper around the cpp objects

rapidsai · Nov 9, 2022 · d1c0f1b · d1c0f1b
1 parent 94e1adf
commit d1c0f1b
Show file tree

Hide file tree

Showing 4 changed files with 386 additions and 243 deletions.
diff --git a/python/pylibraft/pylibraft/neighbors/__init__.py b/python/pylibraft/pylibraft/neighbors/__init__.py
@@ -12,4 +12,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-from .ivf_pq import IvfPq
+# from .ivf_pq import IvfPq
diff --git a/python/pylibraft/pylibraft/neighbors/c_ivf_pq.pxd b/python/pylibraft/pylibraft/neighbors/c_ivf_pq.pxd
@@ -78,6 +78,18 @@ cdef extern from "raft/neighbors/ivf_pq_types.hpp" \
               uint32_t pq_dim,
               uint32_t n_nonempty_lists)
 
+        IdxT size()
+        uint32_t dim()
+        uint32_t pq_dim()
+        uint32_t pq_len()
+        uint32_t pq_bits()
+        DistanceType metric()
+        uint32_t n_lists()
+        uint32_t rot_dim()
+
+
+
+
     cpdef cppclass search_params(ann_search_params):
         uint32_t n_probes
         cudaDataType_t lut_dtype
@@ -90,19 +102,19 @@ cdef extern from "raft/neighbors/specializations/ivf_pq_specialization.hpp" \
              const index_params& params,     
              const float* dataset,               
              uint64_t n_rows,                    
-             uint32_t dim) # except +
+             uint32_t dim) #except +
 
     cdef index[uint64_t] build(const handle_t& handle,      
              const index_params& params,     
              const int8_t* dataset,               
              uint64_t n_rows,                    
-             uint32_t dim) # except +
+             uint32_t dim) #except +
 
     cdef index[uint64_t] build(const handle_t& handle,      
              const index_params& params,     
              const uint8_t* dataset,               
              uint64_t n_rows,                    
-             uint32_t dim) # except +
+             uint32_t dim) #except +
 
     cdef index[uint64_t] extend(const handle_t& handle,        
               const index[uint64_t]& orig_index,