From 6db0f11d7fe1cea32ebd1d46fc5984f8329a05db Mon Sep 17 00:00:00 2001 From: Marcin Antas Date: Wed, 6 Dec 2023 11:48:13 +0100 Subject: [PATCH] Add support for Flat index config --- .../client/v1/misc/model/BQConfig.java | 19 ++++++++ .../client/v1/misc/model/PQConfig.java | 5 ++- .../v1/misc/model/VectorIndexConfig.java | 1 + .../v1/schema/model/WeaviateClassTest.java | 25 +++++++++++ .../integration/client/WeaviateVersion.java | 2 +- .../batch/ClientBatchGrpcCreateTest.java | 45 ++++++++++++++++--- src/test/resources/docker-compose-azure.yaml | 2 +- .../resources/docker-compose-cluster.yaml | 4 +- .../resources/docker-compose-okta-cc.yaml | 2 +- .../resources/docker-compose-okta-users.yaml | 2 +- src/test/resources/docker-compose-proxy.yaml | 2 +- src/test/resources/docker-compose-test.yaml | 2 +- src/test/resources/docker-compose-wcs.yaml | 2 +- 13 files changed, 96 insertions(+), 17 deletions(-) create mode 100644 src/main/java/io/weaviate/client/v1/misc/model/BQConfig.java diff --git a/src/main/java/io/weaviate/client/v1/misc/model/BQConfig.java b/src/main/java/io/weaviate/client/v1/misc/model/BQConfig.java new file mode 100644 index 00000000..4e7c5417 --- /dev/null +++ b/src/main/java/io/weaviate/client/v1/misc/model/BQConfig.java @@ -0,0 +1,19 @@ +package io.weaviate.client.v1.misc.model; + +import lombok.AccessLevel; +import lombok.Builder; +import lombok.EqualsAndHashCode; +import lombok.Getter; +import lombok.ToString; +import lombok.experimental.FieldDefaults; + +@Getter +@Builder +@ToString +@EqualsAndHashCode +@FieldDefaults(makeFinal = true, level = AccessLevel.PRIVATE) +public class BQConfig { + Boolean enabled; + Long rescoreLimit; + Boolean cache; +} diff --git a/src/main/java/io/weaviate/client/v1/misc/model/PQConfig.java b/src/main/java/io/weaviate/client/v1/misc/model/PQConfig.java index 3e8cadf4..4a8ab038 100644 --- a/src/main/java/io/weaviate/client/v1/misc/model/PQConfig.java +++ b/src/main/java/io/weaviate/client/v1/misc/model/PQConfig.java @@ -13,13 +13,16 @@ @EqualsAndHashCode @FieldDefaults(makeFinal = true, level = AccessLevel.PRIVATE) public class PQConfig { - Boolean enabled; + // HNSW index specific settings Boolean bitCompression; Integer segments; Integer centroids; Integer trainingLimit; Encoder encoder; + // Flat index specific settings + Long rescoreLimit; + Boolean cache; @Getter diff --git a/src/main/java/io/weaviate/client/v1/misc/model/VectorIndexConfig.java b/src/main/java/io/weaviate/client/v1/misc/model/VectorIndexConfig.java index 59e93e6c..96108838 100644 --- a/src/main/java/io/weaviate/client/v1/misc/model/VectorIndexConfig.java +++ b/src/main/java/io/weaviate/client/v1/misc/model/VectorIndexConfig.java @@ -25,4 +25,5 @@ public class VectorIndexConfig { Integer cleanupIntervalSeconds; Boolean skip; PQConfig pq; + BQConfig bq; } diff --git a/src/test/java/io/weaviate/client/v1/schema/model/WeaviateClassTest.java b/src/test/java/io/weaviate/client/v1/schema/model/WeaviateClassTest.java index 50c3ed23..833cd3b1 100644 --- a/src/test/java/io/weaviate/client/v1/schema/model/WeaviateClassTest.java +++ b/src/test/java/io/weaviate/client/v1/schema/model/WeaviateClassTest.java @@ -1,6 +1,8 @@ package io.weaviate.client.v1.schema.model; import com.google.gson.GsonBuilder; +import io.weaviate.client.v1.misc.model.BQConfig; +import io.weaviate.client.v1.misc.model.VectorIndexConfig; import org.junit.Test; import java.util.HashMap; @@ -55,6 +57,29 @@ public void shouldSerializeClass() { } + @Test + public void shouldSerializeClassWithFlatIndexType() { + BQConfig bqConfig = BQConfig.builder().enabled(true).rescoreLimit(100l).build(); + VectorIndexConfig vectorIndexConfig = VectorIndexConfig.builder() + .bq(bqConfig) + .build(); + WeaviateClass clazz = WeaviateClass.builder() + .moduleConfig(createDummyModuleConfig()) + .className("Band") + .description("Band that plays and produces music") + .vectorIndexType("flat") + .vectorIndexConfig(vectorIndexConfig) + .vectorizer("text2vec-contextionary") + .build(); + + String result = new GsonBuilder().create().toJson(clazz); + + assertThat(result).isEqualTo("{\"class\":\"Band\",\"description\":\"Band that plays and produces music\"," + + "\"moduleConfig\":{\"text2vec-contextionary\":{\"vectorizeClassName\":false}}," + + "\"vectorIndexConfig\":{\"bq\":{\"enabled\":true,\"rescoreLimit\":100}},\"vectorIndexType\":\"flat\"," + + "\"vectorizer\":\"text2vec-contextionary\"}"); + } + private Object createDummyModuleConfig() { Map text2vecContextionary = new HashMap<>(); text2vecContextionary.put("vectorizeClassName", false); diff --git a/src/test/java/io/weaviate/integration/client/WeaviateVersion.java b/src/test/java/io/weaviate/integration/client/WeaviateVersion.java index e00a8fdd..18993a4e 100644 --- a/src/test/java/io/weaviate/integration/client/WeaviateVersion.java +++ b/src/test/java/io/weaviate/integration/client/WeaviateVersion.java @@ -5,7 +5,7 @@ public class WeaviateVersion { // to be set according to weaviate docker image public static final String EXPECTED_WEAVIATE_VERSION = "1.22.6"; // to be set according to weaviate docker image - public static final String EXPECTED_WEAVIATE_GIT_HASH = "0cc2d22"; + public static final String EXPECTED_WEAVIATE_GIT_HASH = "ea13956"; private WeaviateVersion() {} } diff --git a/src/test/java/io/weaviate/integration/client/batch/ClientBatchGrpcCreateTest.java b/src/test/java/io/weaviate/integration/client/batch/ClientBatchGrpcCreateTest.java index 62baae0e..90db3025 100644 --- a/src/test/java/io/weaviate/integration/client/batch/ClientBatchGrpcCreateTest.java +++ b/src/test/java/io/weaviate/integration/client/batch/ClientBatchGrpcCreateTest.java @@ -5,6 +5,8 @@ import io.weaviate.client.base.Result; import io.weaviate.client.v1.batch.model.ObjectGetResponse; import io.weaviate.client.v1.data.model.WeaviateObject; +import io.weaviate.client.v1.misc.model.BQConfig; +import io.weaviate.client.v1.misc.model.VectorIndexConfig; import io.weaviate.client.v1.schema.model.Property; import io.weaviate.client.v1.schema.model.WeaviateClass; import io.weaviate.integration.client.WeaviateTestGenerics; @@ -16,6 +18,7 @@ import org.junit.Before; import org.junit.ClassRule; import org.junit.Test; +import org.junit.platform.commons.util.StringUtils; import org.testcontainers.containers.ComposeContainer; import org.testcontainers.containers.wait.strategy.Wait; @@ -77,6 +80,11 @@ public void shouldCreateBatchWithMultiCrossReferencesWithNestedPropertiesUsingGR testCreateBatchWithMultiReferenceWithNested(true); } + @Test + public void shouldCreateBatchUsingGRPCWithFlatBQConfig() { + testCreateBatchWithFlatVectorIndex(true); + } + @Test public void shouldCreateBatchUsingRest() { testCreateBatch(false); @@ -112,6 +120,11 @@ public void shouldCreateBatchWithMultiCrossReferencesWithNestedPropertiesUsingRe testCreateBatchWithMultiReferenceWithNested(false); } + @Test + public void shouldCreateBatchUsingRestWithFlatBQConfig() { + testCreateBatchWithFlatVectorIndex(false); + } + private void testCreateBatchWithReferenceWithoutNested(Boolean useGRPC) { WeaviateClient client = createClient(useGRPC); WeaviateTestGenerics.AllPropertiesSchema testData = new WeaviateTestGenerics.AllPropertiesSchema(); @@ -168,13 +181,24 @@ private void testCreateBatchWithMultiReferenceWithNested(Boolean useGRPC) { testData.deleteRefClasses(client); } - private void testCreateBatch(Boolean useGRPC) { + private void testCreateBatch(Boolean useGRPC, String vectorIndexType, VectorIndexConfig vectorIndexConfig) { WeaviateClient client = createClient(useGRPC); WeaviateTestGenerics.AllPropertiesSchema testData = new WeaviateTestGenerics.AllPropertiesSchema(); String className = testData.CLASS_NAME; List properties = testData.properties(); WeaviateObject[] objects = testData.objects(); - testCreateBatch(client, className, properties, objects); + testCreateBatch(client, className, properties, objects, vectorIndexType, vectorIndexConfig); + } + + private void testCreateBatchWithFlatVectorIndex(Boolean useGRPC) { + VectorIndexConfig vectorIndexConfig = VectorIndexConfig.builder() + .bq(BQConfig.builder().enabled(true).build()) + .build(); + testCreateBatch(useGRPC, "flat", vectorIndexConfig); + } + + private void testCreateBatch(Boolean useGRPC) { + testCreateBatch(useGRPC, null, null); } private void testCreateBatchWithNested(Boolean useGRPC) { @@ -205,13 +229,20 @@ private WeaviateClient createClient(Boolean useGRPC) { } private void testCreateBatch(WeaviateClient client, String className, List properties, WeaviateObject[] objects) { + testCreateBatch(client, className, properties, objects, null, null); + } + + private void testCreateBatch(WeaviateClient client, String className, List properties, WeaviateObject[] objects, + String vectorIndexType, VectorIndexConfig vectorIndexConfig) { // create schema + WeaviateClass.WeaviateClassBuilder weaviateClassBuilder = WeaviateClass.builder() + .className(className) + .properties(properties); + if (StringUtils.isNotBlank(vectorIndexType) && vectorIndexConfig != null) { + weaviateClassBuilder.vectorIndexType(vectorIndexType).vectorIndexConfig(vectorIndexConfig); + } Result createResult = client.schema().classCreator() - .withClass(WeaviateClass.builder() - .className(className) - .properties(properties) - .build() - ) + .withClass(weaviateClassBuilder.build()) .run(); assertThat(createResult).isNotNull() .returns(false, Result::hasErrors) diff --git a/src/test/resources/docker-compose-azure.yaml b/src/test/resources/docker-compose-azure.yaml index 2bb2b27b..8867e606 100644 --- a/src/test/resources/docker-compose-azure.yaml +++ b/src/test/resources/docker-compose-azure.yaml @@ -10,7 +10,7 @@ services: - --scheme - http - --write-timeout=600s - image: semitechnologies/weaviate:1.22.6 + image: semitechnologies/weaviate:preview-hnsw-bq-ea13956 restart: on-failure:0 environment: PERSISTENCE_DATA_PATH: '/var/lib/weaviate' diff --git a/src/test/resources/docker-compose-cluster.yaml b/src/test/resources/docker-compose-cluster.yaml index 7b3c3e72..27022d52 100644 --- a/src/test/resources/docker-compose-cluster.yaml +++ b/src/test/resources/docker-compose-cluster.yaml @@ -10,7 +10,7 @@ services: - --scheme - http - --write-timeout=600s - image: semitechnologies/weaviate:1.22.6 + image: semitechnologies/weaviate:preview-hnsw-bq-ea13956 restart: on-failure:0 environment: LOG_LEVEL: 'debug' @@ -41,7 +41,7 @@ services: - '8088' - --scheme - http - image: semitechnologies/weaviate:1.22.6 + image: semitechnologies/weaviate:preview-hnsw-bq-ea13956 restart: on-failure:0 environment: LOG_LEVEL: 'debug' diff --git a/src/test/resources/docker-compose-okta-cc.yaml b/src/test/resources/docker-compose-okta-cc.yaml index 59c9c59c..0ec840a2 100644 --- a/src/test/resources/docker-compose-okta-cc.yaml +++ b/src/test/resources/docker-compose-okta-cc.yaml @@ -10,7 +10,7 @@ services: - --scheme - http - --write-timeout=600s - image: semitechnologies/weaviate:1.22.6 + image: semitechnologies/weaviate:preview-hnsw-bq-ea13956 restart: on-failure:0 environment: PERSISTENCE_DATA_PATH: '/var/lib/weaviate' diff --git a/src/test/resources/docker-compose-okta-users.yaml b/src/test/resources/docker-compose-okta-users.yaml index 62a92eb7..2867a147 100644 --- a/src/test/resources/docker-compose-okta-users.yaml +++ b/src/test/resources/docker-compose-okta-users.yaml @@ -10,7 +10,7 @@ services: - --scheme - http - --write-timeout=600s - image: semitechnologies/weaviate:1.22.6 + image: semitechnologies/weaviate:preview-hnsw-bq-ea13956 restart: on-failure:0 environment: PERSISTENCE_DATA_PATH: '/var/lib/weaviate' diff --git a/src/test/resources/docker-compose-proxy.yaml b/src/test/resources/docker-compose-proxy.yaml index a0bb3260..79aecd57 100644 --- a/src/test/resources/docker-compose-proxy.yaml +++ b/src/test/resources/docker-compose-proxy.yaml @@ -9,7 +9,7 @@ services: - '8080' - --scheme - http - image: semitechnologies/weaviate:1.22.6 + image: semitechnologies/weaviate:preview-hnsw-bq-ea13956 restart: on-failure:0 environment: LOG_LEVEL: "debug" diff --git a/src/test/resources/docker-compose-test.yaml b/src/test/resources/docker-compose-test.yaml index 477ea855..be74c8df 100644 --- a/src/test/resources/docker-compose-test.yaml +++ b/src/test/resources/docker-compose-test.yaml @@ -9,7 +9,7 @@ services: - '8080' - --scheme - http - image: semitechnologies/weaviate:1.22.6 + image: semitechnologies/weaviate:preview-hnsw-bq-ea13956 links: - "contextionary:contextionary" restart: on-failure:0 diff --git a/src/test/resources/docker-compose-wcs.yaml b/src/test/resources/docker-compose-wcs.yaml index 7bc0667c..31f47fea 100644 --- a/src/test/resources/docker-compose-wcs.yaml +++ b/src/test/resources/docker-compose-wcs.yaml @@ -10,7 +10,7 @@ services: - --scheme - http - --write-timeout=600s - image: semitechnologies/weaviate:1.22.6 + image: semitechnologies/weaviate:preview-hnsw-bq-ea13956 restart: on-failure:0 environment: PERSISTENCE_DATA_PATH: '/var/lib/weaviate'