Skip to content

Commit 7462cbd

Browse files
ldematteldematte
authored
[Java] Using functions for mapping (#1007)
This PR replaces "array" based ID mappings with a more free-form mapping based on functions. This should be more efficient both when the number of IDs is huge (you don't need to actually allocate big list for all IDs to map) and in general (as it won't need to perform boxing/unboxing of integers). It should also be more flexible; as an example, I have introduced a simple helper to wrap a list, so that the previous approach can still be used. The same can be done for maps etc. Relates to #699 Authors: - Lorenzo Dematté (https://github.com/ldematte) - Corey J. Nolet (https://github.com/cjnolet) - MithunR (https://github.com/mythrocks) Approvers: - Ishan Chattopadhyaya (https://github.com/chatman) - MithunR (https://github.com/mythrocks) URL: #1007
1 parent f02ee48 commit 7462cbd

13 files changed

Lines changed: 433 additions & 256 deletions

‎java/cuvs-java/src/main/java/com/nvidia/cuvs/BruteForceQuery.java‎

Lines changed: 17 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717

1818
import java.util.Arrays;
1919
import java.util.BitSet;
20-
import java.util.List;
20+
import java.util.function.LongToIntFunction;
2121

2222
/**
2323
* BruteForceQuery holds the query vectors to be used while invoking search.
@@ -26,26 +26,30 @@
2626
*/
2727
public class BruteForceQuery {
2828

29-
private List<Integer> mapping;
30-
private float[][] queryVectors;
31-
private BitSet[] prefilters;
29+
private final LongToIntFunction mapping;
30+
private final float[][] queryVectors;
31+
private final BitSet[] prefilters;
3232
private int numDocs = -1;
33-
private int topK;
33+
private final int topK;
3434

3535
/**
3636
* Constructs an instance of {@link BruteForceQuery} using queryVectors,
3737
* mapping, and topK.
3838
*
3939
* @param queryVectors 2D float query vector array
40-
* @param mapping an instance of ID mapping
40+
* @param mapping a function mapping ordinals (neighbor IDs) to custom user IDs
4141
* @param topK the top k results to return
4242
* @param prefilters the prefilters data to use while searching the BRUTEFORCE
4343
* index
4444
* @param numDocs Maximum of bits in each prefilter, representing number of documents in this index.
4545
* Used only when prefilter(s) is/are passed.
4646
*/
4747
public BruteForceQuery(
48-
float[][] queryVectors, List<Integer> mapping, int topK, BitSet[] prefilters, int numDocs) {
48+
float[][] queryVectors,
49+
LongToIntFunction mapping,
50+
int topK,
51+
BitSet[] prefilters,
52+
int numDocs) {
4953
this.queryVectors = queryVectors;
5054
this.mapping = mapping;
5155
this.topK = topK;
@@ -63,11 +67,9 @@ public float[][] getQueryVectors() {
6367
}
6468

6569
/**
66-
* Gets the passed map instance.
67-
*
68-
* @return a map of ID mappings
70+
* Gets the function mapping ordinals (neighbor IDs) to custom user IDs
6971
*/
70-
public List<Integer> getMapping() {
72+
public LongToIntFunction getMapping() {
7173
return mapping;
7274
}
7375

@@ -119,7 +121,7 @@ public static class Builder {
119121
private float[][] queryVectors;
120122
private BitSet[] prefilters;
121123
private int numDocs;
122-
private List<Integer> mapping;
124+
private LongToIntFunction mapping = SearchResults.IDENTITY_MAPPING;
123125
private int topK = 2;
124126

125127
/**
@@ -134,12 +136,12 @@ public Builder withQueryVectors(float[][] queryVectors) {
134136
}
135137

136138
/**
137-
* Sets the instance of mapping to be used for ID mapping.
139+
* Sets the function used to map ordinals (neighbor IDs) to custom user IDs
138140
*
139-
* @param mapping the ID mapping instance
141+
* @param mapping a function mapping ordinals (neighbor IDs) to custom user IDs
140142
* @return an instance of this Builder
141143
*/
142-
public Builder withMapping(List<Integer> mapping) {
144+
public Builder withMapping(LongToIntFunction mapping) {
143145
this.mapping = mapping;
144146
return this;
145147
}

‎java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraQuery.java‎

Lines changed: 15 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717

1818
import java.util.Arrays;
1919
import java.util.BitSet;
20-
import java.util.List;
20+
import java.util.function.LongToIntFunction;
2121

2222
/**
2323
* CagraQuery holds the CagraSearchParams and the query vectors to be used while
@@ -27,12 +27,12 @@
2727
*/
2828
public class CagraQuery {
2929

30-
private CagraSearchParams cagraSearchParameters;
31-
private List<Integer> mapping;
32-
private float[][] queryVectors;
33-
private int topK;
34-
private BitSet prefilter;
35-
private int numDocs;
30+
private final CagraSearchParams cagraSearchParameters;
31+
private final LongToIntFunction mapping;
32+
private final float[][] queryVectors;
33+
private final int topK;
34+
private final BitSet prefilter;
35+
private final int numDocs;
3636

3737
/**
3838
* Constructs an instance of {@link CagraQuery} using cagraSearchParameters,
@@ -41,15 +41,15 @@ public class CagraQuery {
4141
* @param cagraSearchParameters an instance of {@link CagraSearchParams} holding
4242
* the search parameters
4343
* @param queryVectors 2D float query vector array
44-
* @param mapping an instance of ID mapping
44+
* @param mapping a function mapping ordinals (neighbor IDs) to custom user IDs
4545
* @param topK the top k results to return
4646
* @param prefilter A single BitSet to use as filter while searching the CAGRA index
4747
* @param numDocs Total number of dataset vectors; used to align the prefilter correctly
4848
*/
4949
public CagraQuery(
5050
CagraSearchParams cagraSearchParameters,
5151
float[][] queryVectors,
52-
List<Integer> mapping,
52+
LongToIntFunction mapping,
5353
int topK,
5454
BitSet prefilter,
5555
int numDocs) {
@@ -81,11 +81,9 @@ public float[][] getQueryVectors() {
8181
}
8282

8383
/**
84-
* Gets the passed map instance.
85-
*
86-
* @return a map of ID mappings
84+
* Gets the function mapping ordinals (neighbor IDs) to custom user IDs
8785
*/
88-
public List<Integer> getMapping() {
86+
public LongToIntFunction getMapping() {
8987
return mapping;
9088
}
9189

@@ -136,7 +134,7 @@ public static class Builder {
136134

137135
private CagraSearchParams cagraSearchParams;
138136
private float[][] queryVectors;
139-
private List<Integer> mapping;
137+
private LongToIntFunction mapping = SearchResults.IDENTITY_MAPPING;
140138
private int topK = 2;
141139
private BitSet prefilter;
142140
private int numDocs;
@@ -170,12 +168,12 @@ public Builder withQueryVectors(float[][] queryVectors) {
170168
}
171169

172170
/**
173-
* Sets the instance of mapping to be used for ID mapping.
171+
* Sets the function used to map ordinals (neighbor IDs) to custom user IDs
174172
*
175-
* @param mapping the ID mapping instance
173+
* @param mapping a function mapping ordinals (neighbor IDs) to custom user IDs
176174
* @return an instance of this Builder
177175
*/
178-
public Builder withMapping(List<Integer> mapping) {
176+
public Builder withMapping(LongToIntFunction mapping) {
179177
this.mapping = mapping;
180178
return this;
181179
}

‎java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswQuery.java‎

Lines changed: 16 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616
package com.nvidia.cuvs;
1717

1818
import java.util.Arrays;
19-
import java.util.List;
19+
import java.util.function.LongToIntFunction;
2020

2121
/**
2222
* HnswQuery holds the query vectors to be used while invoking search on the
@@ -26,22 +26,25 @@
2626
*/
2727
public class HnswQuery {
2828

29-
private HnswSearchParams hnswSearchParams;
30-
private List<Integer> mapping;
31-
private float[][] queryVectors;
32-
private int topK;
29+
private final HnswSearchParams hnswSearchParams;
30+
private final LongToIntFunction mapping;
31+
private final float[][] queryVectors;
32+
private final int topK;
3333

3434
/**
3535
* Constructs an instance of {@link HnswQuery} using queryVectors, mapping, and
3636
* topK.
3737
*
3838
* @param hnswSearchParams the search parameters to use
3939
* @param queryVectors 2D float query vector array
40-
* @param mapping an instance of ID mapping
40+
* @param mapping a function mapping ordinals (neighbor IDs) to custom user IDs
4141
* @param topK the top k results to return
4242
*/
4343
private HnswQuery(
44-
HnswSearchParams hnswSearchParams, float[][] queryVectors, List<Integer> mapping, int topK) {
44+
HnswSearchParams hnswSearchParams,
45+
float[][] queryVectors,
46+
LongToIntFunction mapping,
47+
int topK) {
4548
this.hnswSearchParams = hnswSearchParams;
4649
this.queryVectors = queryVectors;
4750
this.mapping = mapping;
@@ -67,11 +70,9 @@ public float[][] getQueryVectors() {
6770
}
6871

6972
/**
70-
* Gets the passed map instance.
71-
*
72-
* @return a map of ID mappings
73+
* Gets the function mapping ordinals (neighbor IDs) to custom user IDs
7374
*/
74-
public List<Integer> getMapping() {
75+
public LongToIntFunction getMapping() {
7576
return mapping;
7677
}
7778

@@ -102,7 +103,7 @@ public static class Builder {
102103

103104
private HnswSearchParams hnswSearchParams;
104105
private float[][] queryVectors;
105-
private List<Integer> mapping;
106+
private LongToIntFunction mapping = SearchResults.IDENTITY_MAPPING;
106107
private int topK = 2;
107108

108109
/**
@@ -129,12 +130,12 @@ public Builder withQueryVectors(float[][] queryVectors) {
129130
}
130131

131132
/**
132-
* Sets the instance of mapping to be used for ID mapping.
133+
* Sets the function used to map ordinals (neighbor IDs) to custom user IDs
133134
*
134-
* @param mapping the ID mapping instance
135+
* @param mapping a function mapping ordinals (neighbor IDs) to custom user IDs
135136
* @return an instance of this Builder
136137
*/
137-
public Builder withMapping(List<Integer> mapping) {
138+
public Builder withMapping(LongToIntFunction mapping) {
138139
this.mapping = mapping;
139140
return this;
140141
}

‎java/cuvs-java/src/main/java/com/nvidia/cuvs/SearchResults.java‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,9 +17,24 @@
1717

1818
import java.util.List;
1919
import java.util.Map;
20+
import java.util.function.LongToIntFunction;
2021

2122
public interface SearchResults {
2223

24+
/**
25+
* The default identity function mapping neighbours IDs to user-defined IDs
26+
*/
27+
LongToIntFunction IDENTITY_MAPPING = l -> (int) l;
28+
29+
/**
30+
* Creates a mapping function from a list lookup of custom user IDs
31+
* @param mappingAsList a positional list of custom user IDs
32+
* @return a function that maps the input ordinal to a custom user IDs, using the input as an index in the list
33+
*/
34+
static LongToIntFunction mappingsFromList(List<Integer> mappingAsList) {
35+
return l -> mappingAsList.get((int) l);
36+
}
37+
2338
/**
2439
* Gets a list results as a map of neighbor IDs to distances.
2540
*

‎java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceSearchResults.java‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,8 @@
1919
import java.lang.foreign.MemorySegment;
2020
import java.lang.foreign.SequenceLayout;
2121
import java.util.LinkedHashMap;
22-
import java.util.List;
2322
import java.util.Map;
23+
import java.util.function.LongToIntFunction;
2424

2525
/**
2626
* SearchResult encapsulates the logic for reading and holding search results.
@@ -35,7 +35,7 @@ protected BruteForceSearchResults(
3535
MemorySegment neighboursMemorySegment,
3636
MemorySegment distancesMemorySegment,
3737
int topK,
38-
List<Integer> mapping,
38+
LongToIntFunction mapping,
3939
long numberOfQueries) {
4040
super(
4141
neighboursSequenceLayout,
@@ -58,7 +58,7 @@ protected void readResultMemorySegments() {
5858
for (long i = 0; i < topK * numberOfQueries; i++) {
5959
long id = (long) neighboursVarHandle.get(neighboursMemorySegment, 0L, i);
6060
float dst = (float) distancesVarHandle.get(distancesMemorySegment, 0L, i);
61-
intermediateResultMap.put(mapping != null ? mapping.get((int) id) : (int) id, dst);
61+
intermediateResultMap.put(mapping.applyAsInt(id), dst);
6262
count += 1;
6363
if (count == topK) {
6464
results.add(intermediateResultMap);

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

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -239,10 +239,7 @@ private static MemorySegment createCagraIndex() {
239239
public SearchResults search(CagraQuery query) throws Throwable {
240240
try (var localArena = Arena.ofConfined()) {
241241
checkNotDestroyed();
242-
int topK =
243-
query.getMapping() != null
244-
? Math.min(query.getMapping().size(), query.getTopK())
245-
: query.getTopK();
242+
int topK = query.getTopK();
246243
long numQueries = query.getQueryVectors().length;
247244
long numBlocks = topK * numQueries;
248245
int vectorDimension = numQueries > 0 ? query.getQueryVectors()[0].length : 0;

‎java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraSearchResults.java‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,8 @@
1919
import java.lang.foreign.MemorySegment;
2020
import java.lang.foreign.SequenceLayout;
2121
import java.util.LinkedHashMap;
22-
import java.util.List;
2322
import java.util.Map;
23+
import java.util.function.LongToIntFunction;
2424

2525
/**
2626
* SearchResult encapsulates the logic for reading and holding search results.
@@ -35,7 +35,7 @@ protected CagraSearchResults(
3535
MemorySegment neighboursMemorySegment,
3636
MemorySegment distancesMemorySegment,
3737
int topK,
38-
List<Integer> mapping,
38+
LongToIntFunction mapping,
3939
long numberOfQueries) {
4040
super(
4141
neighboursSequenceLayout,
@@ -56,10 +56,10 @@ protected void readResultMemorySegments() {
5656
Map<Integer, Float> intermediateResultMap = new LinkedHashMap<Integer, Float>();
5757
int count = 0;
5858
for (long i = 0; i < topK * numberOfQueries; i++) {
59-
int id = (int) neighboursVarHandle.get(neighboursMemorySegment, 0L, i);
59+
long id = (long) neighboursVarHandle.get(neighboursMemorySegment, 0L, i);
6060
float dst = (float) distancesVarHandle.get(distancesMemorySegment, 0L, i);
6161
if (id != Integer.MAX_VALUE) {
62-
intermediateResultMap.put(mapping != null ? mapping.get(id) : id, dst);
62+
intermediateResultMap.put(mapping.applyAsInt(id), dst);
6363
}
6464
count += 1;
6565
if (count == topK) {

‎java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/HnswSearchResults.java‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,8 @@
1919
import java.lang.foreign.MemorySegment;
2020
import java.lang.foreign.SequenceLayout;
2121
import java.util.LinkedHashMap;
22-
import java.util.List;
2322
import java.util.Map;
23+
import java.util.function.LongToIntFunction;
2424

2525
/**
2626
* SearchResult encapsulates the logic for reading and holding search results.
@@ -35,7 +35,7 @@ protected HnswSearchResults(
3535
MemorySegment neighboursMemorySegment,
3636
MemorySegment distancesMemorySegment,
3737
int topK,
38-
List<Integer> mapping,
38+
LongToIntFunction mapping,
3939
long numberOfQueries) {
4040
super(
4141
neighboursSequenceLayout,
@@ -58,7 +58,7 @@ protected void readResultMemorySegments() {
5858
for (long i = 0; i < topK * numberOfQueries; i++) {
5959
long id = (long) neighboursVarHandle.get(neighboursMemorySegment, 0L, i);
6060
float dst = (float) distancesVarHandle.get(distancesMemorySegment, 0L, i);
61-
intermediateResultMap.put(mapping != null ? mapping.get((int) id) : (int) id, dst);
61+
intermediateResultMap.put(mapping.applyAsInt(id), dst);
6262
count += 1;
6363
if (count == topK) {
6464
results.add(intermediateResultMap);

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

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -23,11 +23,12 @@
2323
import java.util.LinkedList;
2424
import java.util.List;
2525
import java.util.Map;
26+
import java.util.function.LongToIntFunction;
2627

2728
public abstract class SearchResultsImpl implements SearchResults {
2829

2930
protected final List<Map<Integer, Float>> results;
30-
protected final List<Integer> mapping; // TODO: Is this performant in a user application?
31+
protected final LongToIntFunction mapping;
3132

3233
protected final MemorySegment neighboursMemorySegment;
3334
protected final MemorySegment distancesMemorySegment;
@@ -42,7 +43,7 @@ protected SearchResultsImpl(
4243
MemorySegment neighboursMemorySegment,
4344
MemorySegment distancesMemorySegment,
4445
int topK,
45-
List<Integer> mapping,
46+
LongToIntFunction mapping,
4647
long numberOfQueries) {
4748
this.topK = topK;
4849
this.numberOfQueries = numberOfQueries;

0 commit comments

Comments
 (0)