Skip to content

Commit 90e1ff1

Browse files
bevzzzbevzzz
committed
feat: add multiple references in a batch
1 parent cd3dc70 commit 90e1ff1

9 files changed

Lines changed: 308 additions & 8 deletions

File tree

‎src/it/java/io/weaviate/integration/DataITest.java‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
import io.weaviate.client6.v1.api.collections.Property;
1414
import io.weaviate.client6.v1.api.collections.Vectors;
1515
import io.weaviate.client6.v1.api.collections.WeaviateObject;
16+
import io.weaviate.client6.v1.api.collections.data.BatchReference;
1617
import io.weaviate.client6.v1.api.collections.data.DeleteManyResponse;
1718
import io.weaviate.client6.v1.api.collections.data.Reference;
1819
import io.weaviate.client6.v1.api.collections.query.Metadata;
@@ -346,4 +347,44 @@ public void testInsertMany() throws IOException {
346347
.as("collection has 5 objects")
347348
.isEqualTo(5);
348349
}
350+
351+
@Test
352+
public void testReferenceAddMany() throws IOException {
353+
// Arrange
354+
var nsCities = ns("Cities");
355+
var nsAirports = ns("Airports");
356+
357+
client.collections.create(nsAirports);
358+
client.collections.create(nsCities, c -> c
359+
.references(Property.reference("hasAirports", nsAirports)));
360+
361+
var airports = client.collections.use(nsAirports);
362+
var cities = client.collections.use(nsCities);
363+
364+
var alpha = airports.data.insert(Map.of()).uuid();
365+
var goodburg = cities.data.insert(Map.of(), city -> city
366+
.reference("hasAirports", Reference.uuids(alpha)));
367+
368+
// Act
369+
var newAirports = airports.data.insertMany(Map.of(), Map.of());
370+
var bravo = newAirports.responses().get(0).uuid();
371+
var charlie = newAirports.responses().get(1).uuid();
372+
373+
var response = cities.data.referenceAddMany(BatchReference.uuids(goodburg, "hasAirports", bravo, charlie));
374+
375+
// Assert
376+
Assertions.assertThat(response.errors()).isEmpty();
377+
378+
var goodburgAirports = cities.query.byId(goodburg.metadata().uuid(),
379+
city -> city.returnReferences(
380+
QueryReference.single("hasAirports",
381+
airport -> airport.returnMetadata(Metadata.ID))));
382+
383+
Assertions.assertThat(goodburgAirports).get()
384+
.as("Goodburg has 3 airports")
385+
.extracting(WeaviateObject::references)
386+
.extracting(references -> references.get("hasAirports"), InstanceOfAssertFactories.list(WeaviateObject.class))
387+
.extracting(WeaviateObject::uuid)
388+
.contains(alpha, bravo, charlie);
389+
}
349390
}

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

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,16 @@ public record WeaviateObject<P, R, M extends WeaviateMetadata>(
2727
Map<String, List<R>> references,
2828
M metadata) {
2929

30+
/** Shorthand for accesing objects's UUID from metadata. */
31+
public String uuid() {
32+
return metadata.uuid();
33+
}
34+
35+
/** Shorthand for accesing objects's vectors from metadata. */
36+
public Vectors vectors() {
37+
return metadata.vectors();
38+
}
39+
3040
public static <P, R, M extends WeaviateMetadata> WeaviateObject<P, R, M> of(
3141
Function<Builder<P, R, M>, ObjectBuilder<WeaviateObject<P, R, M>>> fn) {
3242
return fn.apply(new Builder<>()).build();

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

Lines changed: 108 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,108 @@
1+
package io.weaviate.client6.v1.api.collections.data;
2+
3+
import java.io.IOException;
4+
import java.util.Arrays;
5+
6+
import com.google.gson.TypeAdapter;
7+
import com.google.gson.stream.JsonReader;
8+
import com.google.gson.stream.JsonWriter;
9+
10+
import io.weaviate.client6.v1.api.collections.WeaviateObject;
11+
12+
public record BatchReference(String fromCollection, String fromProperty, String fromUuid, Reference reference) {
13+
14+
public static BatchReference[] objects(WeaviateObject<?, ?, ?> fromObject, String fromProperty,
15+
WeaviateObject<?, ?, ?>... toObjects) {
16+
return Arrays.stream(toObjects)
17+
.map(to -> new BatchReference(
18+
fromObject.collection(), fromProperty, fromObject.metadata().uuid(),
19+
Reference.object(to)))
20+
.toArray(BatchReference[]::new);
21+
}
22+
23+
public static BatchReference[] uuids(WeaviateObject<?, ?, ?> fromObject, String fromProperty,
24+
String... toUuids) {
25+
return Arrays.stream(toUuids)
26+
.map(to -> new BatchReference(
27+
fromObject.collection(), fromProperty, fromObject.metadata().uuid(),
28+
Reference.uuids(to)))
29+
.toArray(BatchReference[]::new);
30+
}
31+
32+
public static final TypeAdapter<BatchReference> TYPE_ADAPTER = new TypeAdapter<BatchReference>() {
33+
34+
@Override
35+
public void write(JsonWriter out, BatchReference value) throws IOException {
36+
out.beginObject();
37+
38+
out.name("from");
39+
out.value(Reference.toBeacon(value.fromCollection, value.fromProperty, value.fromUuid));
40+
41+
out.name("to");
42+
out.value(Reference.toBeacon(value.reference.collection(), value.reference.uuids().get(0)));
43+
44+
// TODO: add tenant
45+
46+
out.endObject();
47+
}
48+
49+
@Override
50+
public BatchReference read(JsonReader in) throws IOException {
51+
String fromCollection = null;
52+
String fromProperty = null;
53+
String fromUuid = null;
54+
Reference toReference = null;
55+
56+
in.beginObject();
57+
while (in.hasNext()) {
58+
switch (in.nextName()) {
59+
60+
case "from": {
61+
var beacon = in.nextString();
62+
beacon = beacon.replaceFirst("weaviate://localhost/", "");
63+
64+
var parts = beacon.split("/");
65+
fromCollection = parts[0];
66+
fromUuid = parts[1];
67+
fromProperty = parts[2];
68+
break;
69+
}
70+
71+
case "to": {
72+
String collection = null;
73+
String id = null;
74+
75+
var beacon = in.nextString();
76+
beacon = beacon.replaceFirst("weaviate://localhost/", "");
77+
if (beacon.contains("/")) {
78+
var parts = beacon.split("/");
79+
collection = parts[0];
80+
id = parts[1];
81+
} else {
82+
id = beacon;
83+
}
84+
toReference = new Reference(collection, id);
85+
break;
86+
}
87+
88+
// case "tenant":
89+
// switch (in.peek()) {
90+
// case STRING:
91+
// in.nextString();
92+
// case NULL:
93+
// in.nextNull();
94+
// default:
95+
// // We don't expect anything else
96+
// }
97+
// System.out.println("processed tenant");
98+
// break;
99+
// default:
100+
// in.skipValue();
101+
}
102+
}
103+
in.endObject();
104+
105+
return new BatchReference(fromCollection, fromProperty, fromUuid, toReference);
106+
}
107+
}.nullSafe();
108+
}

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

Lines changed: 17 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -44,21 +44,30 @@ public static Reference collection(String collection, String... uuids) {
4444
return new Reference(collection, Arrays.asList(uuids));
4545
}
4646

47+
public static String toBeacon(String collection, String uuid) {
48+
return toBeacon(collection, null, uuid);
49+
}
50+
51+
public static String toBeacon(String collection, String property, String uuid) {
52+
var beacon = "weaviate://localhost";
53+
if (collection != null) {
54+
beacon += "/" + collection;
55+
}
56+
beacon += "/" + uuid;
57+
if (property != null) {
58+
beacon += "/" + property;
59+
}
60+
return beacon;
61+
}
62+
4763
public static final TypeAdapter<Reference> TYPE_ADAPTER = new TypeAdapter<Reference>() {
4864

4965
@Override
5066
public void write(JsonWriter out, Reference value) throws IOException {
5167
for (var uuid : value.uuids()) {
5268
out.beginObject();
5369
out.name("beacon");
54-
55-
var beacon = "weaviate://localhost";
56-
if (value.collection() != null) {
57-
beacon += "/" + value.collection();
58-
}
59-
beacon += "/" + uuid;
60-
61-
out.value(beacon);
70+
out.value(toBeacon(value.collection(), uuid));
6271
out.endObject();
6372
}
6473
}

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

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
package io.weaviate.client6.v1.api.collections.data;
2+
3+
import java.util.ArrayList;
4+
import java.util.Collections;
5+
import java.util.List;
6+
7+
import org.apache.hc.core5.http.HttpStatus;
8+
9+
import io.weaviate.client6.v1.internal.json.JSON;
10+
import io.weaviate.client6.v1.internal.rest.Endpoint;
11+
12+
public record ReferenceAddManyRequest(List<BatchReference> references) {
13+
14+
public static final Endpoint<ReferenceAddManyRequest, ReferenceAddManyResponse> endpoint(
15+
List<BatchReference> references) {
16+
return Endpoint.of(
17+
request -> "POST",
18+
request -> "/batch/references",
19+
(gson, request) -> JSON.serialize(request.references),
20+
request -> Collections.emptyMap(),
21+
code -> code != HttpStatus.SC_SUCCESS,
22+
(gson, response) -> {
23+
var result = JSON.deserialize(response, ReferenceAddManyResponse.class);
24+
var errors = new ArrayList<ReferenceAddManyResponse.BatchError>();
25+
26+
for (var err : result.errors()) {
27+
errors.add(new ReferenceAddManyResponse.BatchError(
28+
err.message(),
29+
references.get(err.referenceIndex()),
30+
err.referenceIndex()));
31+
}
32+
return new ReferenceAddManyResponse(errors);
33+
});
34+
}
35+
36+
}

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

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
package io.weaviate.client6.v1.api.collections.data;
2+
3+
import java.lang.reflect.Type;
4+
import java.util.ArrayList;
5+
import java.util.List;
6+
7+
import com.google.gson.JsonDeserializationContext;
8+
import com.google.gson.JsonDeserializer;
9+
import com.google.gson.JsonElement;
10+
import com.google.gson.JsonParseException;
11+
12+
public record ReferenceAddManyResponse(List<BatchError> errors) {
13+
public record BatchError(String message, BatchReference reference, int referenceIndex) {
14+
}
15+
16+
public static enum CustomJsonDeserializer implements JsonDeserializer<ReferenceAddManyResponse> {
17+
INSTANCE;
18+
19+
@Override
20+
public ReferenceAddManyResponse deserialize(JsonElement json, Type typeOfT, JsonDeserializationContext context)
21+
throws JsonParseException {
22+
23+
var errors = new ArrayList<BatchError>();
24+
int i = 0;
25+
for (var el : json.getAsJsonArray()) {
26+
var result = el.getAsJsonObject().get("result").getAsJsonObject();
27+
if (result.get("status").getAsString().equals("FAILED")) {
28+
var errorMsg = result
29+
.get("errors").getAsJsonObject()
30+
.get("error").getAsJsonArray()
31+
.get(0).getAsString();
32+
33+
var batchErr = new BatchError(errorMsg, null, i);
34+
errors.add(batchErr);
35+
}
36+
i++;
37+
}
38+
return new ReferenceAddManyResponse(errors);
39+
}
40+
}
41+
}

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

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -108,6 +108,15 @@ public void referenceAdd(String fromUuid, String fromProperty, Reference referen
108108
}
109109
}
110110

111+
public ReferenceAddManyResponse referenceAddMany(BatchReference... references) throws IOException {
112+
return referenceAddMany(Arrays.asList(references));
113+
}
114+
115+
public ReferenceAddManyResponse referenceAddMany(List<BatchReference> references) throws IOException {
116+
return this.restTransport.performRequest(new ReferenceAddManyRequest(references),
117+
ReferenceAddManyRequest.endpoint(references));
118+
}
119+
111120
public void referenceDelete(String fromUuid, String fromProperty, Reference reference) throws IOException {
112121
for (var uuid : reference.uuids()) {
113122
var singleRef = new Reference(reference.collection(), uuid);

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

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,8 @@ public final class JSON {
99

1010
static {
1111
var gsonBuilder = new GsonBuilder();
12+
13+
// TypeAdapterFactories ---------------------------------------------------
1214
gsonBuilder.registerTypeAdapterFactory(
1315
io.weaviate.client6.v1.api.collections.WeaviateObject.CustomTypeAdapterFactory.INSTANCE);
1416
gsonBuilder.registerTypeAdapterFactory(
@@ -24,12 +26,21 @@ public final class JSON {
2426
gsonBuilder.registerTypeAdapterFactory(
2527
io.weaviate.client6.v1.api.collections.Generative.CustomTypeAdapterFactory.INSTANCE);
2628

29+
// TypeAdapters -----------------------------------------------------------
2730
gsonBuilder.registerTypeAdapter(
2831
io.weaviate.client6.v1.api.collections.vectorizers.NoneVectorizer.class,
2932
io.weaviate.client6.v1.api.collections.vectorizers.NoneVectorizer.TYPE_ADAPTER);
3033
gsonBuilder.registerTypeAdapter(
3134
io.weaviate.client6.v1.api.collections.data.Reference.class,
3235
io.weaviate.client6.v1.api.collections.data.Reference.TYPE_ADAPTER);
36+
gsonBuilder.registerTypeAdapter(
37+
io.weaviate.client6.v1.api.collections.data.BatchReference.class,
38+
io.weaviate.client6.v1.api.collections.data.BatchReference.TYPE_ADAPTER);
39+
40+
// Deserilizers -----------------------------------------------------------
41+
gsonBuilder.registerTypeAdapter(
42+
io.weaviate.client6.v1.api.collections.data.ReferenceAddManyResponse.class,
43+
io.weaviate.client6.v1.api.collections.data.ReferenceAddManyResponse.CustomJsonDeserializer.INSTANCE);
3344
gson = gsonBuilder.create();
3445
}
3546

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

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,9 @@
2323
import io.weaviate.client6.v1.api.collections.Vectorizer;
2424
import io.weaviate.client6.v1.api.collections.Vectors;
2525
import io.weaviate.client6.v1.api.collections.WeaviateObject;
26+
import io.weaviate.client6.v1.api.collections.data.BatchReference;
2627
import io.weaviate.client6.v1.api.collections.data.Reference;
28+
import io.weaviate.client6.v1.api.collections.data.ReferenceAddManyResponse;
2729
import io.weaviate.client6.v1.api.collections.rerankers.CohereReranker;
2830
import io.weaviate.client6.v1.api.collections.vectorindex.Distance;
2931
import io.weaviate.client6.v1.api.collections.vectorindex.Flat;
@@ -289,6 +291,17 @@ public static Object[][] testCases() {
289291
}
290292
""",
291293
},
294+
{
295+
BatchReference.class,
296+
new BatchReference("FromCollection", "fromProperty", "from-uuid",
297+
Reference.collection("ToCollection", "to-uuid")),
298+
"""
299+
{
300+
"from": "weaviate://localhost/FromCollection/from-uuid/fromProperty",
301+
"to": "weaviate://localhost/ToCollection/to-uuid"
302+
}
303+
""",
304+
},
292305
};
293306
}
294307

@@ -346,4 +359,26 @@ private static void compareVectors(Object got, Object want) {
346359
.withEqualsForType(Arrays::deepEquals, Float[][].class)
347360
.isEqualTo(want);
348361
}
362+
363+
@Test
364+
public void test_ReferenceAddManyResponse_CustomDeserializer() {
365+
var json = """
366+
[
367+
{
368+
"result": { "status": "SUCCESS", "errors": {} }
369+
},
370+
{
371+
"result": { "status": "FAILED", "errors": {
372+
"error": [ "oops" ]
373+
}}
374+
}
375+
]
376+
""";
377+
378+
var got = JSON.deserialize(json, ReferenceAddManyResponse.class);
379+
380+
Assertions.assertThat(got.errors())
381+
.as("response contains 1 error")
382+
.hasSize(1);
383+
}
349384
}

0 commit comments

Comments
 (0)