Skip to content

Commit e0bf5b9

Browse files
bevzzzbevzzz
committed
feat: add quantization
1 parent 91144a9 commit e0bf5b9

15 files changed

Lines changed: 618 additions & 19 deletions

File tree

‎src/main/java/io/weaviate/client6/v1/api/collections/MultiTenancy.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
import io.weaviate.client6.v1.internal.ObjectBuilder;
88

99
public record MultiTenancy(
10-
@SerializedName("enabled") Boolean enabled,
10+
@SerializedName("enabled") boolean enabled,
1111
@SerializedName("autoTenantCreation") Boolean createAutomatically,
1212
@SerializedName("autoTenantActivation") Boolean activateAutomatically) {
1313

@@ -23,7 +23,7 @@ public MultiTenancy(Builder builder) {
2323
}
2424

2525
public static class Builder implements ObjectBuilder<MultiTenancy> {
26-
private Boolean enabled = true;
26+
private boolean enabled = true;
2727
private Boolean createAutomatically;
2828
private Boolean activateAutomatically;
2929

‎src/main/java/io/weaviate/client6/v1/api/collections/Quantization.java‎

Lines changed: 156 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,156 @@
1+
package io.weaviate.client6.v1.api.collections;
2+
3+
import java.io.IOException;
4+
import java.util.EnumMap;
5+
import java.util.Map;
6+
import java.util.function.Function;
7+
8+
import com.google.gson.Gson;
9+
import com.google.gson.JsonParser;
10+
import com.google.gson.TypeAdapter;
11+
import com.google.gson.TypeAdapterFactory;
12+
import com.google.gson.reflect.TypeToken;
13+
import com.google.gson.stream.JsonReader;
14+
import com.google.gson.stream.JsonWriter;
15+
16+
import io.weaviate.client6.v1.api.collections.quantizers.BQ;
17+
import io.weaviate.client6.v1.api.collections.quantizers.PQ;
18+
import io.weaviate.client6.v1.api.collections.quantizers.RQ;
19+
import io.weaviate.client6.v1.api.collections.quantizers.SQ;
20+
import io.weaviate.client6.v1.api.collections.quantizers.Uncompressed;
21+
import io.weaviate.client6.v1.internal.ObjectBuilder;
22+
import io.weaviate.client6.v1.internal.json.JsonEnum;
23+
24+
public interface Quantization {
25+
26+
public enum Kind implements JsonEnum<Kind> {
27+
UNCOMPRESSED("skipDefaultQuantization"),
28+
RQ("rq"),
29+
BQ("bq"),
30+
PQ("pq"),
31+
SQ("sq");
32+
33+
private static final Map<String, Kind> jsonValueMap = JsonEnum.collectNames(Kind.values());
34+
private final String jsonValue;
35+
36+
private Kind(String jsonValue) {
37+
this.jsonValue = jsonValue;
38+
}
39+
40+
@Override
41+
public String jsonValue() {
42+
return this.jsonValue;
43+
}
44+
45+
public static Kind valueOfJson(String jsonValue) {
46+
return JsonEnum.valueOfJson(jsonValue, jsonValueMap, Kind.class);
47+
}
48+
}
49+
50+
Kind _kind();
51+
52+
Object _self();
53+
54+
public static Quantization uncompressed() {
55+
return Uncompressed.of();
56+
}
57+
58+
public static Quantization bq() {
59+
return BQ.of();
60+
}
61+
62+
public static Quantization bq(Function<BQ.Builder, ObjectBuilder<BQ>> fn) {
63+
return BQ.of(fn);
64+
}
65+
66+
public static Quantization pq() {
67+
return PQ.of();
68+
}
69+
70+
public static Quantization pq(Function<PQ.Builder, ObjectBuilder<PQ>> fn) {
71+
return PQ.of(fn);
72+
}
73+
74+
public static Quantization sq() {
75+
return SQ.of();
76+
}
77+
78+
public static Quantization sq(Function<SQ.Builder, ObjectBuilder<SQ>> fn) {
79+
return SQ.of(fn);
80+
}
81+
82+
public static Quantization rq() {
83+
return RQ.of();
84+
}
85+
86+
public static Quantization rq(Function<RQ.Builder, ObjectBuilder<RQ>> fn) {
87+
return RQ.of(fn);
88+
}
89+
90+
public static enum CustomTypeAdapterFactory implements TypeAdapterFactory {
91+
INSTANCE;
92+
93+
private static final EnumMap<Quantization.Kind, TypeAdapter<? extends Quantization>> delegateAdapters = new EnumMap<>(
94+
Quantization.Kind.class);
95+
96+
private final void addAdapter(Gson gson, Quantization.Kind kind, Class<? extends Quantization> cls) {
97+
delegateAdapters.put(kind,
98+
(TypeAdapter<? extends Quantization>) gson.getDelegateAdapter(this, TypeToken.get(cls)));
99+
}
100+
101+
private final void init(Gson gson) {
102+
addAdapter(gson, Quantization.Kind.UNCOMPRESSED, Uncompressed.class);
103+
addAdapter(gson, Quantization.Kind.BQ, BQ.class);
104+
addAdapter(gson, Quantization.Kind.RQ, RQ.class);
105+
addAdapter(gson, Quantization.Kind.SQ, SQ.class);
106+
addAdapter(gson, Quantization.Kind.PQ, PQ.class);
107+
}
108+
109+
@SuppressWarnings("unchecked")
110+
@Override
111+
public <T> TypeAdapter<T> create(Gson gson, TypeToken<T> type) {
112+
final var rawType = type.getRawType();
113+
if (!Quantization.class.isAssignableFrom(rawType)) {
114+
return null;
115+
}
116+
117+
if (delegateAdapters.isEmpty()) {
118+
init(gson);
119+
}
120+
121+
return (TypeAdapter<T>) new TypeAdapter<Quantization>() {
122+
123+
@Override
124+
public void write(JsonWriter out, Quantization value) throws IOException {
125+
if (value._kind() == Quantization.Kind.UNCOMPRESSED) {
126+
// out.name(value._kind().jsonValue());
127+
out.value(true);
128+
return;
129+
}
130+
TypeAdapter<T> adapter = (TypeAdapter<T>) delegateAdapters.get(value._kind());
131+
adapter.write(out, (T) value._self());
132+
}
133+
134+
@Override
135+
public Quantization read(JsonReader in) throws IOException {
136+
var quantizerObject = JsonParser.parseReader(in).getAsJsonObject();
137+
var quantizationName = quantizerObject.keySet().iterator().next();
138+
Quantization.Kind kind;
139+
try {
140+
kind = Quantization.Kind.valueOfJson(quantizationName);
141+
} catch (IllegalArgumentException e) {
142+
return null;
143+
}
144+
145+
if (kind == Quantization.Kind.UNCOMPRESSED) {
146+
return new Uncompressed();
147+
}
148+
149+
var adapter = delegateAdapters.get(kind);
150+
var concreteQuantizer = quantizerObject.get(quantizationName).getAsJsonObject();
151+
return adapter.fromJsonTree(concreteQuantizer);
152+
}
153+
}.nullSafe();
154+
}
155+
}
156+
}

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

Lines changed: 36 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,8 @@ public static Kind valueOfJson(String jsonValue) {
5454

5555
VectorIndex vectorIndex();
5656

57+
Quantization quantization();
58+
5759
/** Create a bring-your-own-vector vector index. */
5860
public static Map.Entry<String, VectorConfig> selfProvided() {
5961
return selfProvided(VectorIndex.DEFAULT_VECTOR_NAME);
@@ -269,29 +271,50 @@ public void write(JsonWriter out, VectorConfig value) throws IOException {
269271
TypeAdapter<T> adapter = (TypeAdapter<T>) delegateAdapters.get(value._kind());
270272

271273
// Serialize vectorizer config as { "vectorizer-kind": { ... } }
272-
// and remove "vectorIndex" object which every vectorizer has.
274+
// and remove "vectorIndex" and quantization objects which every vectorizer has.
273275
var vectorizer = new JsonObject();
274276
var config = adapter.toJsonTree((T) value._self());
275277

276278
// This will create { "vectorIndexType": "", "vectorIndexConfig": { ... } }
277-
// to which we just need to add "vectorizer": { ... } key.
279+
// to which we just need to add "vectorizer": { ... } key
280+
// and "bq"/"pg"/"sq"/"rq": { ... } (quantizer) key.
278281
var vectorIndex = config.getAsJsonObject().remove("vectorIndex");
279282

280283
vectorizer.add(value._kind().jsonValue(), config);
281284
vectorIndex.getAsJsonObject().add("vectorizer", vectorizer);
282285

286+
if (value.quantization() != null) {
287+
vectorIndex.getAsJsonObject()
288+
.get("vectorIndexConfig").getAsJsonObject()
289+
.add(value.quantization()._kind().jsonValue(), config.getAsJsonObject().remove("quantization"));
290+
}
291+
283292
Streams.write(vectorIndex, out);
284293
}
285294

286295
@Override
287296
public VectorConfig read(JsonReader in) throws IOException {
288297
var jsonObject = JsonParser.parseReader(in).getAsJsonObject();
298+
var vectorIndexConfig = jsonObject.get("vectorIndexConfig").getAsJsonObject();
299+
300+
String quantizationKind = null;
301+
if (vectorIndexConfig.has(Quantization.Kind.BQ.jsonValue())) {
302+
quantizationKind = Quantization.Kind.BQ.jsonValue();
303+
} else if (vectorIndexConfig.has(Quantization.Kind.PQ.jsonValue())) {
304+
quantizationKind = Quantization.Kind.PQ.jsonValue();
305+
} else if (vectorIndexConfig.has(Quantization.Kind.SQ.jsonValue())) {
306+
quantizationKind = Quantization.Kind.SQ.jsonValue();
307+
} else if (vectorIndexConfig.has(Quantization.Kind.RQ.jsonValue())) {
308+
quantizationKind = Quantization.Kind.RQ.jsonValue();
309+
} else {
310+
quantizationKind = Quantization.Kind.UNCOMPRESSED.jsonValue();
311+
}
289312

290313
// VectorIndex.CustomTypeAdapterFactory expects keys
291314
// ["vectorIndexType", "vectorIndexConfig"].
292315
var vectorIndex = new JsonObject();
293316
vectorIndex.add("vectorIndexType", jsonObject.get("vectorIndexType"));
294-
vectorIndex.add("vectorIndexConfig", jsonObject.get("vectorIndexConfig"));
317+
vectorIndex.add("vectorIndexConfig", vectorIndexConfig);
295318

296319
var vectorizerObject = jsonObject.get("vectorizer").getAsJsonObject();
297320
var vectorizerName = vectorizerObject.keySet().iterator().next();
@@ -309,6 +332,16 @@ public VectorConfig read(JsonReader in) throws IOException {
309332
// Each individual vectorizer has a `VectorIndex vectorIndex` field.
310333
concreteVectorizer.add("vectorIndex", vectorIndex);
311334

335+
// Each individual vectorizer has a `Quantization quantization` field.
336+
// We need to specify the kind in order for
337+
// Quantization.CustomTypeAdapterFactory to be able to find the right adapter.
338+
if (vectorIndexConfig.has(quantizationKind)) {
339+
JsonObject quantization = new JsonObject();
340+
quantization.add(quantizationKind, vectorIndexConfig.get(quantizationKind));
341+
concreteVectorizer.add("quantization", quantization);
342+
} else {
343+
concreteVectorizer.add("quantization", null);
344+
}
312345
return adapter.fromJsonTree(concreteVectorizer);
313346
}
314347
}.nullSafe();

‎src/main/java/io/weaviate/client6/v1/api/collections/quantizers/BQ.java‎

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
package io.weaviate.client6.v1.api.collections.quantizers;
2+
3+
import java.util.function.Function;
4+
5+
import com.google.gson.annotations.SerializedName;
6+
7+
import io.weaviate.client6.v1.api.collections.Quantization;
8+
import io.weaviate.client6.v1.internal.ObjectBuilder;
9+
10+
public record BQ(
11+
@SerializedName("enabled") boolean enabled,
12+
@SerializedName("rescore_limit") Integer rescoreLimit,
13+
@SerializedName("cache") Boolean cache) implements Quantization {
14+
15+
@Override
16+
public Quantization.Kind _kind() {
17+
return Quantization.Kind.BQ;
18+
}
19+
20+
@Override
21+
public Object _self() {
22+
return this;
23+
}
24+
25+
public static BQ of() {
26+
return of(ObjectBuilder.identity());
27+
}
28+
29+
public static BQ of(Function<Builder, ObjectBuilder<BQ>> fn) {
30+
return fn.apply(new Builder()).build();
31+
}
32+
33+
public BQ(Builder builder) {
34+
this(builder.enabled, builder.rescoreLimit, builder.cache);
35+
}
36+
37+
public static class Builder implements ObjectBuilder<BQ> {
38+
private boolean enabled = true;
39+
private Integer rescoreLimit;
40+
private Boolean cache;
41+
42+
public Builder enabled(boolean enabled) {
43+
this.enabled = enabled;
44+
return this;
45+
}
46+
47+
public Builder rescoreLimit(int rescoreLimit) {
48+
this.rescoreLimit = rescoreLimit;
49+
return this;
50+
}
51+
52+
public Builder cache(boolean enabled) {
53+
this.cache = enabled;
54+
return this;
55+
}
56+
57+
@Override
58+
public BQ build() {
59+
return new BQ(this);
60+
}
61+
}
62+
}

0 commit comments

Comments
 (0)