Skip to content

Commit 675aa92

Browse files
punAhujalowener
authored andcommitted
[Java]Fixed TieredIndexParams creation/destruction, calling cuvsTieredIndexParamsCreate (rapidsai#1147)
Using `cuvsTieredIndexParamsCreate` and `cuvsTieredIndexParamsDestroy` now instead of allocating arena in Java. -Used CloseableHandle and CuvsParamsHelper as used in rapidsai#1109 and rapidsai#1110 for consistency This fixes rapidsai#1138 Authors: - Puneet Ahuja (https://github.com/punAhuja) Approvers: - MithunR (https://github.com/mythrocks) URL: rapidsai#1147
1 parent 3b07201 commit 675aa92

File tree

2 files changed

+37
-17
lines changed

2 files changed

+37
-17
lines changed

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

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -151,4 +151,23 @@ public void close() {
151151
};
152152
}
153153
}
154+
155+
static CloseableHandle createTieredIndexParams() {
156+
try (var localArena = Arena.ofConfined()) {
157+
var paramsPtrPtr = localArena.allocate(cuvsTieredIndexParams_t);
158+
checkCuVSError(cuvsTieredIndexParamsCreate(paramsPtrPtr), "cuvsTieredIndexParamsCreate");
159+
var paramsPtr = paramsPtrPtr.get(cuvsTieredIndexParams_t, 0L);
160+
return new CloseableHandle() {
161+
@Override
162+
public MemorySegment handle() {
163+
return paramsPtr;
164+
}
165+
166+
@Override
167+
public void close() {
168+
checkCuVSError(cuvsTieredIndexParamsDestroy(paramsPtr), "cuvsTieredIndexParamsDestroy");
169+
}
170+
};
171+
}
172+
}
154173
}

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

Lines changed: 18 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
*/
1616
package com.nvidia.cuvs.internal;
1717

18+
import static com.nvidia.cuvs.internal.CuVSParamsHelper.createTieredIndexParams;
1819
import static com.nvidia.cuvs.internal.common.LinkerHelper.C_FLOAT;
1920
import static com.nvidia.cuvs.internal.common.LinkerHelper.C_FLOAT_BYTE_SIZE;
2021
import static com.nvidia.cuvs.internal.common.LinkerHelper.C_INT_BYTE_SIZE;
@@ -32,6 +33,7 @@
3233

3334
import com.nvidia.cuvs.*;
3435
import com.nvidia.cuvs.CuVSMatrix;
36+
import com.nvidia.cuvs.internal.common.CloseableHandle;
3537
import com.nvidia.cuvs.internal.common.Util;
3638
import com.nvidia.cuvs.internal.panama.cuvsCagraIndexParams;
3739
import com.nvidia.cuvs.internal.panama.cuvsCagraSearchParams;
@@ -118,15 +120,16 @@ private IndexReference build() {
118120
long rows = dataset.size();
119121
long cols = dataset.columns();
120122

121-
MemorySegment indexParamsMemorySegment =
122-
tieredIndexParameters != null
123-
? segmentFromIndexParams(localArena, tieredIndexParameters)
124-
: MemorySegment.NULL;
125-
126123
// Get host data
127124
MemorySegment hostDataSeg = ((CuVSHostMatrixImpl) dataset).memorySegment();
128125

129-
try (var resourceAccess = resources.access()) {
126+
try (var resourceAccess = resources.access();
127+
var indexParamsHandle =
128+
tieredIndexParameters != null
129+
? segmentFromIndexParams(localArena, tieredIndexParameters)
130+
: CloseableHandle.NULL) {
131+
132+
MemorySegment indexParamsMemorySegment = indexParamsHandle.handle();
130133
long cuvsRes = resourceAccess.handle();
131134

132135
// TieredIndex REQUIRES device memory - allocate it
@@ -389,16 +392,20 @@ public void execute() {
389392
/**
390393
* Allocates the configured index parameters in the MemorySegment.
391394
*/
392-
private static MemorySegment segmentFromIndexParams(Arena arena, TieredIndexParams params) {
393-
MemorySegment seg = cuvsTieredIndexParams.allocate(arena);
395+
396+
/**
397+
* Allocates the configured index parameters in a MemorySegment and returns a CloseableHandle
398+
* for safe resource management.
399+
*/
400+
private static CloseableHandle segmentFromIndexParams(Arena arena, TieredIndexParams params) {
401+
CloseableHandle paramsHandle = createTieredIndexParams();
402+
MemorySegment seg = paramsHandle.handle();
394403

395404
// Get the metric from CagraParams if available, otherwise use TieredIndex metric
396405
int metric;
397406
if (params.getCagraParams() != null) {
398-
// Use the metric from CagraParams to ensure consistency
399407
metric = params.getCagraParams().getCuvsDistanceType().value;
400408
} else {
401-
// Fallback to TieredIndex metric
402409
metric =
403410
switch (params.getMetric()) {
404411
case L2 -> 0;
@@ -407,34 +414,28 @@ private static MemorySegment segmentFromIndexParams(Arena arena, TieredIndexPara
407414
throw new IllegalArgumentException("Unsupported metric: " + params.getMetric());
408415
};
409416
}
410-
411417
cuvsTieredIndexParams.metric(seg, metric);
412-
413418
int algo = 0; // CUVS_TIERED_INDEX_ALGO_CAGRA
414419
cuvsTieredIndexParams.algo(seg, algo);
415-
416420
cuvsTieredIndexParams.min_ann_rows(seg, params.getMinAnnRows());
417421
cuvsTieredIndexParams.create_ann_index_on_extend(seg, params.isCreateAnnIndexOnExtend());
418422

419423
CagraIndexParams cagraParams = params.getCagraParams();
420424
if (cagraParams != null) {
421425
MemorySegment cagraParamsSeg = cuvsCagraIndexParams.allocate(arena);
422-
423426
cuvsCagraIndexParams.intermediate_graph_degree(
424427
cagraParamsSeg, cagraParams.getIntermediateGraphDegree());
425428
cuvsCagraIndexParams.graph_degree(cagraParamsSeg, cagraParams.getGraphDegree());
426429
cuvsCagraIndexParams.build_algo(cagraParamsSeg, cagraParams.getCagraGraphBuildAlgo().value);
427430
cuvsCagraIndexParams.nn_descent_niter(
428431
cagraParamsSeg, cagraParams.getNNDescentNumIterations());
429432
cuvsCagraIndexParams.metric(cagraParamsSeg, metric);
430-
431433
cuvsTieredIndexParams.cagra_params(seg, cagraParamsSeg);
432434
}
433-
434435
cuvsTieredIndexParams.ivf_flat_params(seg, MemorySegment.NULL);
435436
cuvsTieredIndexParams.ivf_pq_params(seg, MemorySegment.NULL);
436437

437-
return seg;
438+
return paramsHandle;
438439
}
439440

440441
/**

0 commit comments

Comments
 (0)