Merge pull request #246 from weaviate/add-support-for-flat-index

Add support for Flat index config
weaviate · Dec 6, 2023 · fb0a48c · fb0a48c
2 parents 15fb8c8 + 6db0f11
commit fb0a48c
Show file tree

Hide file tree

Showing 13 changed files with 96 additions and 17 deletions.
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
@@ -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
@@ -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
@@ -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
@@ -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<String, Object> 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
@@ -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
@@ -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<Property> 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<Property> properties, WeaviateObject[] objects) {
+    testCreateBatch(client, className, properties, objects, null, null);
+  }
+
+  private void testCreateBatch(WeaviateClient client, String className, List<Property> 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<Boolean> 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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'