Merge pull request #592 from weaviate/feat/boost

feat(query): add Boost parameter to all query types

Author: bevzzz

.github/workflows/test.yaml

@@ -92,7 +92,15 @@ jobs:
       fail-fast: false
       matrix:
         WEAVIATE_VERSION:
-          ["1.32.24", "1.33.11", "1.34.7", "1.35.2", "1.36.9", "1.37.7"]
+          [
+            "1.32.24",
+            "1.33.11",
+            "1.34.7",
+            "1.35.2",
+            "1.36.9",
+            "1.37.7",
+            "1.38.2",
+          ]
     steps:
       - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6
 

src/it/java/io/weaviate/ConcurrentTest.java

@@ -126,6 +126,21 @@ public static void requireAtLeast(Weaviate.Version required) {
         .isGreaterThanOrEqualTo(required.semver);
   }
 
+  /**
+   * Skip the test if the version that the {@link Weaviate}
+   * container is running is newer than the required one.
+   *
+   * <br>
+   * This is useful for guarding tests from server versions
+   * which the client does not fully support yet.
+   */
+  public static void requireAtMost(Weaviate.Version required) {
+    var actual = SemanticVersion.of(Weaviate.VERSION);
+    Assumptions.assumeThat(actual)
+        .as("requires at most %s, but running %s", required.semver, actual)
+        .isLessThanOrEqualTo(required.semver);
+  }
+
   @FunctionalInterface
   public interface ThrowingRunnable {
     void run() throws Exception;

src/it/java/io/weaviate/containers/Weaviate.java

@@ -45,7 +45,8 @@ public enum Version {
     V134(1, 34, 7),
     V135(1, 35, 2),
     V136(1, 36, 9),
-    V137(1, 37, 7);
+    V137(1, 37, 7),
+    V138(1, 38, 2);
 
     public final SemanticVersion semver;
 

src/it/java/io/weaviate/integration/RbacITest.java

@@ -6,6 +6,7 @@
 
 import org.assertj.core.api.Assertions;
 import org.assertj.core.api.InstanceOfAssertFactories;
+import org.junit.BeforeClass;
 import org.junit.Test;
 
 import io.weaviate.ConcurrentTest;
@@ -14,10 +15,10 @@
 import io.weaviate.client6.v1.api.rbac.AliasesPermission;
 import io.weaviate.client6.v1.api.rbac.BackupsPermission;
 import io.weaviate.client6.v1.api.rbac.ClusterPermission;
-import io.weaviate.client6.v1.api.rbac.McpPermission;
 import io.weaviate.client6.v1.api.rbac.CollectionsPermission;
 import io.weaviate.client6.v1.api.rbac.DataPermission;
 import io.weaviate.client6.v1.api.rbac.GroupsPermission;
+import io.weaviate.client6.v1.api.rbac.McpPermission;
 import io.weaviate.client6.v1.api.rbac.NodesPermission;
 import io.weaviate.client6.v1.api.rbac.Permission;
 import io.weaviate.client6.v1.api.rbac.ReplicatePermission;
@@ -56,6 +57,12 @@ public class RbacITest extends ConcurrentTest {
   private static final WeaviateClient client = container
       .getClient(fn -> fn.authentication(Authentication.apiKey(API_KEY)));
 
+  @BeforeClass
+  public static void __() {
+    // TODO(dyma): remove once namespace permissions are supported (v1.38 feature)
+    ConcurrentTest.requireAtMost(Weaviate.Version.V137);
+  }
+
   @Test
   public void test_roles_Lifecycle() throws IOException {
     // Arrange

src/it/java/io/weaviate/integration/SearchITest.java

@@ -17,6 +17,11 @@
 import org.junit.Ignore;
 import org.junit.Test;
 import org.junit.rules.TestRule;
+import org.junit.runner.RunWith;
+
+import com.jparams.junit4.JParamsTestRunner;
+import com.jparams.junit4.data.DataMethod;
+import com.jparams.junit4.description.Name;
 
 import io.weaviate.ConcurrentTest;
 import io.weaviate.client6.v1.api.WeaviateApiException;
@@ -32,6 +37,7 @@
 import io.weaviate.client6.v1.api.collections.generate.GenerativeObject;
 import io.weaviate.client6.v1.api.collections.generate.TaskOutput;
 import io.weaviate.client6.v1.api.collections.generative.DummyGenerative;
+import io.weaviate.client6.v1.api.collections.query.Boost;
 import io.weaviate.client6.v1.api.collections.query.Diversity;
 import io.weaviate.client6.v1.api.collections.query.FetchObjectById;
 import io.weaviate.client6.v1.api.collections.query.Filter;
@@ -54,6 +60,7 @@
 import io.weaviate.containers.Weaviate;
 import io.weaviate.containers.Weaviate.Version;
 
+@RunWith(JParamsTestRunner.class)
 public class SearchITest extends ConcurrentTest {
   private static final ContainerGroup compose = Container.compose(
       Weaviate.custom()
@@ -132,7 +139,10 @@ private static Map<String, float[]> populateTest(int n) throws IOException {
     for (int i = 0; i < n; i++) {
       var vector = randomVector(10, -.01f, .001f);
       var object = things.data.insert(
-          Map.of("category", CATEGORIES.get(i % CATEGORIES.size())),
+          Map.of(
+              "category", CATEGORIES.get(i % CATEGORIES.size()),
+              "created_at", OffsetDateTime.now(),
+              "position", i),
           metadata -> metadata
               .uuid(randomUUID())
               .vectors(Vectors.of(VECTOR_INDEX, vector)));
@@ -150,7 +160,10 @@ private static Map<String, float[]> populateTest(int n) throws IOException {
    */
   private static void createTestCollection() throws IOException {
     client.collections.create(COLLECTION, cfg -> cfg
-        .properties(Property.text("category"))
+        .properties(
+            Property.text("category"),
+            Property.date("created_at"),
+            Property.integer("position"))
         .vectorConfig(VectorConfig.selfProvided(VECTOR_INDEX)));
   }
 
@@ -886,6 +899,61 @@ public void testQueryProfile_groupBy() throws Exception {
             .isNotEmpty());
   }
 
+  public static Object[][] boostCases() {
+    return new Object[][] {
+        { "filter", Boost.filter(Filter.property("category").eq("red")), },
+        { "timeDecay", Boost.timeDecay("created_at", "365d",
+            time -> time
+                .origin("2024-01-01T00:00:00Z")
+                .curve(Boost.Curve.EXPONENTIAL)
+                .decay(.3f)
+                .weight(1f)),
+        },
+        {
+            "numericDecay", Boost.numericDecay("position", 1f, 3f,
+                num -> num
+                    .curve(Boost.Curve.LINEAR)
+                    .decay(0.2f)
+                    .weight(1f)),
+        },
+        {
+            "numericProperty", Boost.numericProperty("position",
+                prop -> prop
+                    .modifier(Boost.Modifier.LOG1P)
+                    .weight(1f)),
+        },
+    };
+  }
+
+  @Test
+  @DataMethod(source = SearchITest.class, method = "boostCases")
+  @Name("{0}")
+  public void testBoost(String __, Object boost) throws Exception {
+    Version.V138.orSkip();
+
+    // Boosting reranks query results. To verify the boost parameter
+    // made it to request, check that boosted results arrive in the
+    // different order from those in a plain vector search.
+    //
+    // Because boosting is a common parameter for all query types,
+    // testing one (e.g. nearVector) is sufficient.
+
+    // Arrange
+    var things = client.collections.use(COLLECTION);
+    var baseline = things.query.nearVector(searchVector,
+        opt -> opt.limit(5))
+        .objects().stream().map(WeaviateObject::uuid).toList();
+
+    // Act
+    var got = things.query.nearVector(searchVector,
+        opt -> opt.limit(5).boost((Boost) boost))
+        .objects().stream().map(WeaviateObject::uuid).toList();
+
+    // Assert
+    Assertions.assertThat(got).hasSameSizeAs(baseline);
+    Assertions.assertThat(got).doesNotContainSequence(baseline);
+  }
+  
   @Test
   public void testDiversity() throws Exception {
     Version.V137.orSkip();

src/main/java/io/weaviate/client6/v1/api/collections/query/BaseQueryOptions.java

@@ -20,6 +20,7 @@ public record BaseQueryOptions(
     String after,
     ConsistencyLevel consistencyLevel,
     Filter filters,
+    Boost boost,
     GenerativeSearch generativeSearch,
     List<String> returnProperties,
     List<QueryReference> returnReferences,
@@ -38,6 +39,7 @@ private <T extends Object> BaseQueryOptions(Builder<? extends Builder<?, T>, T>
         builder.after,
         builder.consistencyLevel,
         builder.filter,
+        builder.boost,
         builder.generativeSearch,
         builder.returnProperties,
         builder.returnReferences,
@@ -54,6 +56,7 @@ public static abstract class Builder<SelfT extends Builder<SelfT, T>, T extends
     private String after;
     private ConsistencyLevel consistencyLevel;
     private Filter filter;
+    private Boost boost;
     private GenerativeSearch generativeSearch;
     private List<String> returnProperties = new ArrayList<>();
     private List<QueryReference> returnReferences = new ArrayList<>();
@@ -143,6 +146,12 @@ public final SelfT filters(Filter... filters) {
       return (SelfT) this;
     }
 
+    /** Boost search results. */
+    public final SelfT boost(Boost boost) {
+      this.boost = boost;
+      return (SelfT) this;
+    }
+
     /** Select properties to include in the query result. */
     public final SelfT returnProperties(String... properties) {
       return returnProperties(Arrays.asList(properties));
@@ -230,6 +239,10 @@ final void appendTo(WeaviateProtoSearchGet.SearchRequest.Builder req) {
       req.setFilters(filter);
     }
 
+    if (boost != null) {
+      req.setBoost(boost.toProto());
+    }
+
     if (generativeSearch != null) {
       var generative = WeaviateProtoGenerative.GenerativeSearch.newBuilder();
       generativeSearch.appendTo(generative);