Use int64_t for getters in the cagra/ivf_flat c-api

cpp/include/cuvs/neighbors/cagra.h

@@ -356,7 +356,7 @@ cuvsError_t cuvsCagraIndexDestroy(cuvsCagraIndex_t index);
  * @param[out] dim return dimension of the index
  * @return cuvsError_t
  */
-cuvsError_t cuvsCagraIndexGetDims(cuvsCagraIndex_t index, int* dim);
+cuvsError_t cuvsCagraIndexGetDims(cuvsCagraIndex_t index, int64_t* dim);
 
 /**
  * @brief Get size of the CAGRA index
@@ -365,7 +365,7 @@ cuvsError_t cuvsCagraIndexGetDims(cuvsCagraIndex_t index, int* dim);
  * @param[out] size return number of vectors in the index
  * @return cuvsError_t
  */
-cuvsError_t cuvsCagraIndexGetSize(cuvsCagraIndex_t index, uint32_t* size);
+cuvsError_t cuvsCagraIndexGetSize(cuvsCagraIndex_t index, int64_t* size);
 
 /**
  * @brief Get graph degree of the CAGRA index
@@ -374,7 +374,7 @@ cuvsError_t cuvsCagraIndexGetSize(cuvsCagraIndex_t index, uint32_t* size);
  * @param[out] graph_degree return graph degree
  * @return cuvsError_t
  */
-cuvsError_t cuvsCagraIndexGetGraphDegree(cuvsCagraIndex_t index, uint32_t* graph_degree);
+cuvsError_t cuvsCagraIndexGetGraphDegree(cuvsCagraIndex_t index, int64_t* graph_degree);
 
 /**
  * @brief Returns a view of the CAGRA dataset

cpp/include/cuvs/neighbors/ivf_flat.h

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2024, NVIDIA CORPORATION.
+ * Copyright (c) 2024-2025, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -166,10 +166,10 @@ cuvsError_t cuvsIvfFlatIndexCreate(cuvsIvfFlatIndex_t* index);
 cuvsError_t cuvsIvfFlatIndexDestroy(cuvsIvfFlatIndex_t index);
 
 /** Get the number of clusters/inverted lists */
-uint32_t cuvsIvfFlatIndexGetNLists(cuvsIvfFlatIndex_t index);
+cuvsError_t cuvsIvfFlatIndexGetNLists(cuvsIvfFlatIndex_t index, int64_t* n_lists);
 
 /** Get the dimensionality of the data */
-uint32_t cuvsIvfFlatIndexGetDim(cuvsIvfFlatIndex_t index);
+cuvsError_t cuvsIvfFlatIndexGetDim(cuvsIvfFlatIndex_t index, int64_t* dim);
 
 /**
  * @brief Get the cluster centers corresponding to the lists [n_lists, dim]

cpp/src/neighbors/cagra_c.cpp

@@ -426,23 +426,23 @@ extern "C" cuvsError_t cuvsCagraIndexDestroy(cuvsCagraIndex_t index_c_ptr)
   });
 }
 
-extern "C" cuvsError_t cuvsCagraIndexGetDims(cuvsCagraIndex_t index, int* dim)
+extern "C" cuvsError_t cuvsCagraIndexGetDims(cuvsCagraIndex_t index, int64_t* dim)
 {
   return cuvs::core::translate_exceptions([=] {
     auto index_ptr = reinterpret_cast<cuvs::neighbors::cagra::index<float, uint32_t>*>(index->addr);
     *dim           = index_ptr->dim();
   });
 }
 
-extern "C" cuvsError_t cuvsCagraIndexGetSize(cuvsCagraIndex_t index, uint32_t* size)
+extern "C" cuvsError_t cuvsCagraIndexGetSize(cuvsCagraIndex_t index, int64_t* size)
 {
   return cuvs::core::translate_exceptions([=] {
     auto index_ptr = reinterpret_cast<cuvs::neighbors::cagra::index<float, uint32_t>*>(index->addr);
     *size          = index_ptr->size();
   });
 }
 
-extern "C" cuvsError_t cuvsCagraIndexGetGraphDegree(cuvsCagraIndex_t index, uint32_t* graph_degree)
+extern "C" cuvsError_t cuvsCagraIndexGetGraphDegree(cuvsCagraIndex_t index, int64_t* graph_degree)
 {
   return cuvs::core::translate_exceptions([=] {
     auto index_ptr = reinterpret_cast<cuvs::neighbors::cagra::index<float, uint32_t>*>(index->addr);

cpp/src/neighbors/ivf_flat_c.cpp

@@ -1,6 +1,6 @@
 
 /*
- * Copyright (c) 2024, NVIDIA CORPORATION.
+ * Copyright (c) 2024-2025, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -395,50 +395,54 @@ extern "C" cuvsError_t cuvsIvfFlatExtend(cuvsResources_t res,
   });
 }
 
-extern "C" uint32_t cuvsIvfFlatIndexGetNLists(cuvsIvfFlatIndex_t index)
+extern "C" cuvsError_t cuvsIvfFlatIndexGetNLists(cuvsIvfFlatIndex_t index, int64_t* n_lists)
 {
-  if (index->dtype.code == kDLFloat && index->dtype.bits == 32) {
-    auto index_ptr =
-      reinterpret_cast<cuvs::neighbors::ivf_flat::index<float, int64_t>*>(index->addr);
-    return index_ptr->n_lists();
-  } else if (index->dtype.code == kDLFloat && index->dtype.bits == 16) {
-    auto index_ptr =
-      reinterpret_cast<cuvs::neighbors::ivf_flat::index<half, int64_t>*>(index->addr);
-    return index_ptr->n_lists();
-  } else if (index->dtype.code == kDLInt && index->dtype.bits == 8) {
-    auto index_ptr =
-      reinterpret_cast<cuvs::neighbors::ivf_flat::index<int8_t, int64_t>*>(index->addr);
-    return index_ptr->n_lists();
-  } else if (index->dtype.code == kDLUInt && index->dtype.bits == 8) {
-    auto index_ptr =
-      reinterpret_cast<cuvs::neighbors::ivf_flat::index<uint8_t, int64_t>*>(index->addr);
-    return index_ptr->n_lists();
-  } else {
-    return 0;
-  }
+  return cuvs::core::translate_exceptions([=] {
+    if (index->dtype.code == kDLFloat && index->dtype.bits == 32) {
+      auto index_ptr =
+        reinterpret_cast<cuvs::neighbors::ivf_flat::index<float, int64_t>*>(index->addr);
+      *n_lists = index_ptr->n_lists();
+    } else if (index->dtype.code == kDLFloat && index->dtype.bits == 16) {
+      auto index_ptr =
+        reinterpret_cast<cuvs::neighbors::ivf_flat::index<half, int64_t>*>(index->addr);
+      *n_lists = index_ptr->n_lists();
+    } else if (index->dtype.code == kDLInt && index->dtype.bits == 8) {
+      auto index_ptr =
+        reinterpret_cast<cuvs::neighbors::ivf_flat::index<int8_t, int64_t>*>(index->addr);
+      *n_lists = index_ptr->n_lists();
+    } else if (index->dtype.code == kDLUInt && index->dtype.bits == 8) {
+      auto index_ptr =
+        reinterpret_cast<cuvs::neighbors::ivf_flat::index<uint8_t, int64_t>*>(index->addr);
+      *n_lists = index_ptr->n_lists();
+    } else {
+      RAFT_FAIL("Unsupported index dtype: %d and bits: %d", index->dtype.code, index->dtype.bits);
+    }
+  });
 }
 
-extern "C" uint32_t cuvsIvfFlatIndexGetDim(cuvsIvfFlatIndex_t index)
+extern "C" cuvsError_t cuvsIvfFlatIndexGetDim(cuvsIvfFlatIndex_t index, int64_t* dim)
 {
-  if (index->dtype.code == kDLFloat && index->dtype.bits == 32) {
-    auto index_ptr =
-      reinterpret_cast<cuvs::neighbors::ivf_flat::index<float, int64_t>*>(index->addr);
-    return index_ptr->dim();
-  } else if (index->dtype.code == kDLFloat && index->dtype.bits == 16) {
-    auto index_ptr =
-      reinterpret_cast<cuvs::neighbors::ivf_flat::index<half, int64_t>*>(index->addr);
-    return index_ptr->dim();
-  } else if (index->dtype.code == kDLInt && index->dtype.bits == 8) {
-    auto index_ptr =
-      reinterpret_cast<cuvs::neighbors::ivf_flat::index<int8_t, int64_t>*>(index->addr);
-    return index_ptr->dim();
-  } else if (index->dtype.code == kDLUInt && index->dtype.bits == 8) {
-    auto index_ptr =
-      reinterpret_cast<cuvs::neighbors::ivf_flat::index<uint8_t, int64_t>*>(index->addr);
-    return index_ptr->dim();
-  } else {
-    return 0;
-  }
+  return cuvs::core::translate_exceptions([=] {
+    if (index->dtype.code == kDLFloat && index->dtype.bits == 32) {
+      auto index_ptr =
+        reinterpret_cast<cuvs::neighbors::ivf_flat::index<float, int64_t>*>(index->addr);
+      *dim = index_ptr->dim();
+    } else if (index->dtype.code == kDLFloat && index->dtype.bits == 16) {
+      auto index_ptr =
+        reinterpret_cast<cuvs::neighbors::ivf_flat::index<half, int64_t>*>(index->addr);
+      *dim = index_ptr->dim();
+    } else if (index->dtype.code == kDLInt && index->dtype.bits == 8) {
+      auto index_ptr =
+        reinterpret_cast<cuvs::neighbors::ivf_flat::index<int8_t, int64_t>*>(index->addr);
+      *dim = index_ptr->dim();
+    } else if (index->dtype.code == kDLUInt && index->dtype.bits == 8) {
+      auto index_ptr =
+        reinterpret_cast<cuvs::neighbors::ivf_flat::index<uint8_t, int64_t>*>(index->addr);
+      *dim = index_ptr->dim();
+    } else {
+      RAFT_FAIL("Unsupported index dtype: %d and bits: %d", index->dtype.code, index->dtype.bits);
+    }
+  });
 }
 
 extern "C" cuvsError_t cuvsIvfFlatIndexGetCenters(cuvsResources_t res,

cpp/tests/neighbors/ann_cagra_c.cu

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2023-2024, NVIDIA CORPORATION.
+ * Copyright (c) 2023-2025, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -541,7 +541,7 @@ TEST(CagraC, BuildMergeSearch)
   cuvsCagraIndex_t index_array[2] = {index_main, index_add};
   ASSERT_EQ(cuvsCagraMerge(res, merge_params, index_array, 2, index_merged), CUVS_SUCCESS);
 
-  int merged_dim = -1;
+  int64_t merged_dim = -1;
   ASSERT_EQ(cuvsCagraIndexGetDims(index_merged, &merged_dim), CUVS_SUCCESS);
   EXPECT_EQ(merged_dim, 2);
 

java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java

@@ -45,7 +45,6 @@
 import com.nvidia.cuvs.internal.common.CloseableHandle;
 import com.nvidia.cuvs.internal.common.CloseableRMMAllocation;
 import com.nvidia.cuvs.internal.common.CompositeCloseableHandle;
-import com.nvidia.cuvs.internal.common.Util;
 import com.nvidia.cuvs.internal.panama.cuvsCagraCompressionParams;
 import com.nvidia.cuvs.internal.panama.cuvsCagraIndexParams;
 import com.nvidia.cuvs.internal.panama.cuvsCagraMergeParams;
@@ -413,16 +412,16 @@ public void serialize(OutputStream outputStream, Path tempFile, int bufferLength
   @Override
   public CuVSMatrix getGraph() {
     try (var localArena = Arena.ofConfined()) {
-      var outPtr = localArena.allocate(uint32_t);
+      var outPtr = localArena.allocate(C_LONG);
       checkCuVSError(
           cuvsCagraIndexGetGraphDegree(cagraIndexReference.getMemorySegment(), outPtr),
           "cuvsCagraIndexGetGraphDegree");
-      long graphDegree = Util.dereferenceUnsignedInt(outPtr);
+      long graphDegree = outPtr.get(C_LONG, 0);
 
       checkCuVSError(
           cuvsCagraIndexGetSize(cagraIndexReference.getMemorySegment(), outPtr),
           "cuvsCagraIndexGetSize");
-      long size = Util.dereferenceUnsignedInt(outPtr);
+      long size = outPtr.get(C_LONG, 0);
 
       // TODO: use a "device" graph + tensor, avoid (defer) copy
       var graph = new CuVSHostMatrixArenaImpl(size, graphDegree, CuVSMatrix.DataType.UINT);

java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/Util.java

@@ -69,13 +69,6 @@ public static void checkCudaError(int value, String caller) {
     }
   }
 
-  private static final long UNSIGNED_INT_MASK = 0xFFFFFFFFL;
-
-  public static long dereferenceUnsignedInt(MemorySegment ptr) {
-    assert ptr.byteSize() == 4;
-    return ptr.get(uint32_t, 0) & UNSIGNED_INT_MASK;
-  }
-
   /**
    * Java analog to CUDA's cudaMemcpyKind, used for cudaMemcpy() calls.
    * @see <a href="https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__TYPES.html">CUDA Runtime API</a>

python/cuvs/cuvs/neighbors/cagra/cagra.pxd

@@ -134,10 +134,10 @@ cdef extern from "cuvs/neighbors/cagra.h" nogil:
 
     cuvsError_t cuvsCagraIndexDestroy(cuvsCagraIndex_t index)
 
-    cuvsError_t cuvsCagraIndexGetDims(cuvsCagraIndex_t index, int32_t* dim)
-    cuvsError_t cuvsCagraIndexGetSize(cuvsCagraIndex_t index, uint32_t* size)
+    cuvsError_t cuvsCagraIndexGetDims(cuvsCagraIndex_t index, int64_t* dim)
+    cuvsError_t cuvsCagraIndexGetSize(cuvsCagraIndex_t index, int64_t* size)
     cuvsError_t cuvsCagraIndexGetGraphDegree(cuvsCagraIndex_t index,
-                                             uint32_t* degree)
+                                             int64_t* degree)
     cuvsError_t cuvsCagraIndexGetGraph(cuvsCagraIndex_t index,
                                        DLManagedTensor * graph)
     cuvsError_t cuvsCagraIndexGetDataset(cuvsCagraIndex_t index,

python/cuvs/cuvs/neighbors/cagra/cagra.pyx

@@ -271,18 +271,18 @@ cdef class Index:
 
     @property
     def dim(self):
-        cdef int32_t dim
+        cdef int64_t dim
         check_cuvs(cuvsCagraIndexGetDims(self.index, &dim))
         return dim
 
     @property
     def graph_degree(self):
-        cdef uint32_t degree
+        cdef int64_t degree
         check_cuvs(cuvsCagraIndexGetGraphDegree(self.index, &degree))
         return degree
 
     def __len__(self):
-        cdef uint32_t size
+        cdef int64_t size
         check_cuvs(cuvsCagraIndexGetSize(self.index, &size))
         return size
 

python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pxd

@@ -1,5 +1,5 @@
 #
-# Copyright (c) 2024, NVIDIA CORPORATION.
+# Copyright (c) 2024-2025, NVIDIA CORPORATION.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -15,7 +15,7 @@
 #
 # cython: language_level=3
 
-from libc.stdint cimport uint32_t, uintptr_t
+from libc.stdint cimport int64_t, uint32_t, uintptr_t
 from libcpp cimport bool
 
 from cuvs.common.c_api cimport cuvsError_t, cuvsResources_t
@@ -62,9 +62,10 @@ cdef extern from "cuvs/neighbors/ivf_flat.h" nogil:
 
     cuvsError_t cuvsIvfFlatIndexDestroy(cuvsIvfFlatIndex_t index)
 
-    uint32_t cuvsIvfFlatIndexGetNLists(cuvsIvfFlatIndex_t index)
+    cuvsError_t cuvsIvfFlatIndexGetNLists(cuvsIvfFlatIndex_t index,
+                                          int64_t * n_lists)
 
-    uint32_t cuvsIvfFlatIndexGetDim(cuvsIvfFlatIndex_t index)
+    cuvsError_t cuvsIvfFlatIndexGetDim(cuvsIvfFlatIndex_t index, int64_t * dim)
 
     cuvsError_t cuvsIvfFlatIndexGetCenters(cuvsResources_t res,
                                            cuvsIvfFlatIndex_t index,

python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pyx

@@ -1,5 +1,5 @@
 #
-# Copyright (c) 2024, NVIDIA CORPORATION.
+# Copyright (c) 2024-2025, NVIDIA CORPORATION.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -185,12 +185,16 @@ cdef class Index:
     @property
     def n_lists(self):
         """ The number of inverted lists (clusters) """
-        return cuvsIvfFlatIndexGetNLists(self.index)
+        cdef int64_t n_lists = 0
+        cuvsIvfFlatIndexGetNLists(self.index, &n_lists)
+        return n_lists
 
     @property
     def dim(self):
         """ dimensionality of the cluster centers """
-        return cuvsIvfFlatIndexGetDim(self.index)
+        cdef int64_t dim = 0
+        cuvsIvfFlatIndexGetDim(self.index, &dim)
+        return dim
 
     @property
     def centers(self):