Skip to content

Commit 7c2e850

Browse files
authored
Merge pull request #488 from weaviate/v6-dynamic-index
v6: Dynamic vector index configuration
2 parents 03d711f + c1de067 commit 7c2e850

7 files changed

Lines changed: 188 additions & 13 deletions

File tree

src/main/java/io/weaviate/client6/v1/api/collections/VectorConfig.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -285,7 +285,7 @@ public void write(JsonWriter out, VectorConfig value) throws IOException {
285285
vectorizer.add(value._kind().jsonValue(), config);
286286
vectorIndex.getAsJsonObject().add("vectorizer", vectorizer);
287287

288-
if (value.quantization() != null) {
288+
if (value.quantization() != null && !config.getAsJsonObject().get("quantization").isJsonNull()) {
289289
vectorIndex.getAsJsonObject()
290290
.get("vectorIndexConfig").getAsJsonObject()
291291
.add(value.quantization()._kind().jsonValue(), config.getAsJsonObject().remove("quantization"));

src/main/java/io/weaviate/client6/v1/api/collections/VectorIndex.java

Lines changed: 33 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -13,15 +13,17 @@
1313
import com.google.gson.stream.JsonReader;
1414
import com.google.gson.stream.JsonWriter;
1515

16+
import io.weaviate.client6.v1.api.collections.vectorindex.Dynamic;
1617
import io.weaviate.client6.v1.api.collections.vectorindex.Flat;
1718
import io.weaviate.client6.v1.api.collections.vectorindex.Hnsw;
19+
import io.weaviate.client6.v1.internal.TaggedUnion;
1820
import io.weaviate.client6.v1.internal.json.JsonEnum;
1921

20-
public interface VectorIndex {
22+
public interface VectorIndex extends TaggedUnion<VectorIndex.Kind, Object> {
2123
static final String DEFAULT_VECTOR_NAME = "default";
2224
static final VectorIndex DEFAULT_VECTOR_INDEX = Hnsw.of();
2325

24-
public enum Kind implements JsonEnum<Kind> {
26+
enum Kind implements JsonEnum<Kind> {
2527
HNSW("hnsw"),
2628
FLAT("flat"),
2729
DYNAMIC("dynamic");
@@ -43,17 +45,37 @@ public static Kind valueOfJson(String jsonValue) {
4345
}
4446
}
4547

46-
VectorIndex.Kind _kind();
48+
/** Is this vector index of type HNSW? */
49+
default Hnsw isHnsw() {
50+
return _as(VectorIndex.Kind.HNSW);
51+
}
52+
53+
/** Get as {@link Hnsw} instance. */
54+
default Hnsw asHnsw() {
55+
return _as(VectorIndex.Kind.HNSW);
56+
}
57+
58+
/** Is this vector index of type FLAT? */
59+
default Flat isFlat() {
60+
return _as(VectorIndex.Kind.FLAT);
61+
}
62+
63+
/** Get as {@link Flat} instance. */
64+
default Flat asFlat() {
65+
return _as(VectorIndex.Kind.FLAT);
66+
}
4767

48-
/** Returns the on-the-wire name of the vector index type. */
49-
default String type() {
50-
return _kind().jsonValue();
68+
/** Is this vector index of type DYNAMIC? */
69+
default Dynamic isDynamic() {
70+
return _as(VectorIndex.Kind.DYNAMIC);
5171
}
5272

53-
/** Get the concrete vector index configuration object. */
54-
Object config();
73+
/** Get as {@link Dynamic} instance. */
74+
default Dynamic asDynamic() {
75+
return _as(VectorIndex.Kind.DYNAMIC);
76+
}
5577

56-
public static enum CustomTypeAdapterFactory implements TypeAdapterFactory {
78+
static enum CustomTypeAdapterFactory implements TypeAdapterFactory {
5779
INSTANCE;
5880

5981
private static final EnumMap<VectorIndex.Kind, TypeAdapter<? extends VectorIndex>> readAdapters = new EnumMap<>(
@@ -66,6 +88,7 @@ private final void addAdapter(Gson gson, VectorIndex.Kind kind, Class<? extends
6688
private final void init(Gson gson) {
6789
addAdapter(gson, VectorIndex.Kind.HNSW, Hnsw.class);
6890
addAdapter(gson, VectorIndex.Kind.FLAT, Flat.class);
91+
addAdapter(gson, VectorIndex.Kind.DYNAMIC, Dynamic.class);
6992
}
7093

7194
@SuppressWarnings("unchecked")
@@ -90,7 +113,7 @@ public void write(JsonWriter out, VectorIndex value) throws IOException {
90113
out.value(value._kind().jsonValue());
91114

92115
out.name("vectorIndexConfig");
93-
var config = writeAdapter.toJsonTree((T) value.config());
116+
var config = writeAdapter.toJsonTree((T) value._self());
94117
config.getAsJsonObject().remove("name");
95118
Streams.write(config, out);
96119

Lines changed: 120 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,120 @@
1+
package io.weaviate.client6.v1.api.collections.vectorindex;
2+
3+
import java.io.IOException;
4+
import java.util.function.Function;
5+
6+
import com.google.gson.Gson;
7+
import com.google.gson.JsonObject;
8+
import com.google.gson.JsonParser;
9+
import com.google.gson.TypeAdapter;
10+
import com.google.gson.TypeAdapterFactory;
11+
import com.google.gson.annotations.SerializedName;
12+
import com.google.gson.internal.Streams;
13+
import com.google.gson.reflect.TypeToken;
14+
import com.google.gson.stream.JsonReader;
15+
import com.google.gson.stream.JsonWriter;
16+
17+
import io.weaviate.client6.v1.api.collections.VectorIndex;
18+
import io.weaviate.client6.v1.internal.ObjectBuilder;
19+
20+
public record Dynamic(
21+
@SerializedName("hnsw") Hnsw hnsw,
22+
@SerializedName("flat") Flat flat,
23+
@SerializedName("threshold") Long threshold)
24+
implements VectorIndex {
25+
26+
@Override
27+
public VectorIndex.Kind _kind() {
28+
return VectorIndex.Kind.DYNAMIC;
29+
}
30+
31+
@Override
32+
public Object _self() {
33+
return this;
34+
}
35+
36+
public static Dynamic of() {
37+
return of(ObjectBuilder.identity());
38+
}
39+
40+
public static Dynamic of(Function<Builder, ObjectBuilder<Dynamic>> fn) {
41+
return fn.apply(new Builder()).build();
42+
}
43+
44+
public Dynamic(Builder builder) {
45+
this(
46+
builder.hnsw,
47+
builder.flat,
48+
builder.threshold);
49+
}
50+
51+
public static class Builder implements ObjectBuilder<Dynamic> {
52+
53+
private Hnsw hnsw;
54+
private Flat flat;
55+
private Long threshold;
56+
57+
public Builder hnsw(Hnsw hnsw) {
58+
this.hnsw = hnsw;
59+
return this;
60+
}
61+
62+
public Builder flat(Flat flat) {
63+
this.flat = flat;
64+
return this;
65+
}
66+
67+
public Builder threshold(long threshold) {
68+
this.threshold = threshold;
69+
return this;
70+
}
71+
72+
@Override
73+
public Dynamic build() {
74+
return new Dynamic(this);
75+
}
76+
}
77+
78+
public static enum CustomTypeAdapterFactory implements TypeAdapterFactory {
79+
INSTANCE;
80+
81+
@SuppressWarnings("unchecked")
82+
@Override
83+
public <T> TypeAdapter<T> create(Gson gson, TypeToken<T> type) {
84+
var rawType = type.getRawType();
85+
if (!Dynamic.class.isAssignableFrom(rawType)) {
86+
return null;
87+
}
88+
89+
final var hnswAdapter = gson.getDelegateAdapter(VectorIndex.CustomTypeAdapterFactory.INSTANCE,
90+
TypeToken.get(Hnsw.class));
91+
final var flatAdapter = gson.getDelegateAdapter(VectorIndex.CustomTypeAdapterFactory.INSTANCE,
92+
TypeToken.get(Flat.class));
93+
94+
return (TypeAdapter<T>) new TypeAdapter<Dynamic>() {
95+
96+
@Override
97+
public void write(JsonWriter out, Dynamic value) throws IOException {
98+
99+
var dynamic = new JsonObject();
100+
101+
dynamic.addProperty("threshold", value.threshold);
102+
dynamic.add("hnsw", hnswAdapter.toJsonTree(value.hnsw));
103+
dynamic.add("flat", flatAdapter.toJsonTree(value.flat));
104+
105+
Streams.write(dynamic, out);
106+
}
107+
108+
@Override
109+
public Dynamic read(JsonReader in) throws IOException {
110+
var jsonObject = JsonParser.parseReader(in).getAsJsonObject();
111+
112+
var hnsw = hnswAdapter.fromJsonTree(jsonObject.get("hnsw"));
113+
var flat = flatAdapter.fromJsonTree(jsonObject.get("flat"));
114+
var threshold = jsonObject.get("threshold").getAsLong();
115+
return new Dynamic(hnsw, flat, threshold);
116+
}
117+
}.nullSafe();
118+
}
119+
}
120+
}

src/main/java/io/weaviate/client6/v1/api/collections/vectorindex/Flat.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ public VectorIndex.Kind _kind() {
1616
}
1717

1818
@Override
19-
public Object config() {
19+
public Object _self() {
2020
return this;
2121
}
2222

src/main/java/io/weaviate/client6/v1/api/collections/vectorindex/Hnsw.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@ public VectorIndex.Kind _kind() {
2929
}
3030

3131
@Override
32-
public Object config() {
32+
public Object _self() {
3333
return this;
3434
}
3535

src/main/java/io/weaviate/client6/v1/internal/json/JSON.java

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,8 +28,17 @@ public final class JSON {
2828
io.weaviate.client6.v1.api.collections.Vectors.CustomTypeAdapterFactory.INSTANCE);
2929
gsonBuilder.registerTypeAdapterFactory(
3030
io.weaviate.client6.v1.api.collections.VectorConfig.CustomTypeAdapterFactory.INSTANCE);
31+
32+
// These 2 adapters need to be registered in this exact order: Dynamic
33+
// (narrower), VectorIndex (broader).
34+
// When searching for an adapter, Gson will pick the first adapter factory that
35+
// can process the class, and it's important that Dynamic.class is processed by
36+
// this factory.
37+
gsonBuilder.registerTypeAdapterFactory(
38+
io.weaviate.client6.v1.api.collections.vectorindex.Dynamic.CustomTypeAdapterFactory.INSTANCE);
3139
gsonBuilder.registerTypeAdapterFactory(
3240
io.weaviate.client6.v1.api.collections.VectorIndex.CustomTypeAdapterFactory.INSTANCE);
41+
3342
gsonBuilder.registerTypeAdapterFactory(
3443
io.weaviate.client6.v1.api.collections.Reranker.CustomTypeAdapterFactory.INSTANCE);
3544
gsonBuilder.registerTypeAdapterFactory(

src/test/java/io/weaviate/client6/v1/internal/json/JSONTest.java

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
import io.weaviate.client6.v1.api.collections.quantizers.PQ;
3333
import io.weaviate.client6.v1.api.collections.rerankers.CohereReranker;
3434
import io.weaviate.client6.v1.api.collections.vectorindex.Distance;
35+
import io.weaviate.client6.v1.api.collections.vectorindex.Dynamic;
3536
import io.weaviate.client6.v1.api.collections.vectorindex.Flat;
3637
import io.weaviate.client6.v1.api.collections.vectorindex.Hnsw;
3738
import io.weaviate.client6.v1.api.collections.vectorindex.MultiVector;
@@ -166,6 +167,28 @@ public static Object[][] testCases() {
166167
}
167168
""",
168169
},
170+
{
171+
VectorConfig.class,
172+
SelfProvidedVectorizer.of(none -> none
173+
.vectorIndex(Dynamic.of(idx -> idx
174+
.hnsw(Hnsw.of(hnsw -> hnsw
175+
.ef(1)
176+
.efConstruction(2)))
177+
.flat(Flat.of(flat -> flat
178+
.vectorCacheMaxObjects(100)))
179+
.threshold(5)))),
180+
"""
181+
{
182+
"vectorIndexType": "dynamic",
183+
"vectorizer": {"none": {}},
184+
"vectorIndexConfig": {
185+
"flat": {"vectorCacheMaxObjects": 100},
186+
"hnsw": {"ef": 1, "efConstruction": 2},
187+
"threshold": 5
188+
}
189+
}
190+
""",
191+
},
169192
{
170193
VectorConfig.class,
171194
SelfProvidedVectorizer.of(none -> none

0 commit comments

Comments
 (0)