1515 */
1616package com .nvidia .cuvs .internal ;
1717
18+ import static com .nvidia .cuvs .internal .CuVSParamsHelper .createTieredIndexParams ;
1819import static com .nvidia .cuvs .internal .common .LinkerHelper .C_FLOAT ;
1920import static com .nvidia .cuvs .internal .common .LinkerHelper .C_FLOAT_BYTE_SIZE ;
2021import static com .nvidia .cuvs .internal .common .LinkerHelper .C_INT_BYTE_SIZE ;
3233
3334import com .nvidia .cuvs .*;
3435import com .nvidia .cuvs .CuVSMatrix ;
36+ import com .nvidia .cuvs .internal .common .CloseableHandle ;
3537import com .nvidia .cuvs .internal .common .Util ;
3638import com .nvidia .cuvs .internal .panama .cuvsCagraIndexParams ;
3739import 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