feat: set query options in Paginator

ObjectBuilder.partial is a new util for composing ObjectBuilder-style functions.

Author: bevzzz

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

@@ -11,6 +11,7 @@
 
 import io.weaviate.ConcurrentTest;
 import io.weaviate.client6.v1.api.WeaviateClient;
+import io.weaviate.client6.v1.api.collections.Property;
 import io.weaviate.client6.v1.api.collections.WeaviateMetadata;
 import io.weaviate.client6.v1.api.collections.WeaviateObject;
 import io.weaviate.containers.Container;
@@ -89,4 +90,31 @@ public void testResumePagination() throws IOException {
     // Assert
     Assertions.assertThat(remaining).isEqualTo(5);
   }
+
+  @Test
+  public void testWithQueryOptions() throws IOException {
+    // Arrange
+    var nsThings = ns("Things");
+    var count = 10;
+
+    client.collections.create(nsThings,
+        c -> c.properties(
+            Property.text("fetch_me"),
+            Property.integer("dont_fetch")));
+
+    var things = client.collections.use(nsThings);
+    var inserted = new ArrayList<String>();
+    for (var i = 0; i < count; i++) {
+      var object = things.data.insert(Collections.emptyMap());
+      inserted.add(object.metadata().uuid());
+    }
+
+    // Act / Assert
+    var withSomeProperties = things.paginate(p -> p.returnProperties("fetch_me"));
+    for (var thing : withSomeProperties) {
+      Assertions.assertThat(thing.properties())
+          .as("uuid=" + thing.metadata().uuid())
+          .doesNotContainKey("dont_fetch");
+    }
+  }
 }

src/main/java/io/weaviate/client6/v1/api/collections/pagination/Paginator.java

@@ -1,21 +1,26 @@
 package io.weaviate.client6.v1.api.collections.pagination;
 
 import java.util.Iterator;
+import java.util.List;
 import java.util.Spliterator;
 import java.util.Spliterators;
 import java.util.function.Function;
 import java.util.stream.Stream;
 import java.util.stream.StreamSupport;
 
 import io.weaviate.client6.v1.api.collections.WeaviateObject;
+import io.weaviate.client6.v1.api.collections.query.FetchObjects;
+import io.weaviate.client6.v1.api.collections.query.Metadata;
 import io.weaviate.client6.v1.api.collections.query.QueryMetadata;
+import io.weaviate.client6.v1.api.collections.query.QueryReference;
 import io.weaviate.client6.v1.api.collections.query.WeaviateQueryClient;
 import io.weaviate.client6.v1.internal.ObjectBuilder;
 
 public class Paginator<T> implements Iterable<WeaviateObject<T, Object, QueryMetadata>> {
   private static final int DEFAULT_PAGE_SIZE = 100;
 
   private final WeaviateQueryClient<T> query;
+  private final Function<FetchObjects.Builder, ObjectBuilder<FetchObjects>> queryOptions;
   private final int pageSize;
   private final String cursor;
 
@@ -30,7 +35,10 @@ public Stream<WeaviateObject<T, Object, QueryMetadata>> stream() {
 
   public Spliterator<WeaviateObject<T, Object, QueryMetadata>> spliterator() {
     return new CursorSpliterator<T>(cursor, pageSize,
-        (after, limit) -> query.fetchObjects(q -> q.after(after).limit(limit)).objects());
+        (after, limit) -> {
+          var fn = ObjectBuilder.partial(queryOptions, q -> q.after(after).limit(limit));
+          return query.fetchObjects(fn).objects();
+        });
   }
 
   public static <T> Paginator<T> of(WeaviateQueryClient<T> query) {
@@ -44,20 +52,24 @@ public static <T> Paginator<T> of(WeaviateQueryClient<T> query,
 
   Paginator(Builder<T> builder) {
     this.query = builder.query;
+    this.queryOptions = builder.queryOptions;
     this.cursor = builder.cursor;
     this.pageSize = builder.pageSize;
   }
 
   public static class Builder<T> implements ObjectBuilder<Paginator<T>> {
     private final WeaviateQueryClient<T> query;
 
-    int pageSize = DEFAULT_PAGE_SIZE;
-    String cursor;
+    private Function<FetchObjects.Builder, ObjectBuilder<FetchObjects>> queryOptions = ObjectBuilder.identity();
+    private int pageSize = DEFAULT_PAGE_SIZE;
+    private String cursor;
 
     public Builder(WeaviateQueryClient<T> query) {
       this.query = query;
     }
 
+    // Pagination options -----------------------------------------------------
+
     public Builder<T> pageSize(int pageSize) {
       this.pageSize = pageSize;
       return this;
@@ -68,6 +80,37 @@ public Builder<T> resumeFrom(String uuid) {
       return this;
     }
 
+    // Query options ----------------------------------------------------------
+
+    public final Builder<T> returnProperties(String... properties) {
+      return applyQueryOption(q -> q.returnProperties(properties));
+    }
+
+    public final Builder<T> returnProperties(List<String> properties) {
+      return applyQueryOption(q -> q.returnProperties(properties));
+    }
+
+    public final Builder<T> returnReferences(QueryReference... references) {
+      return applyQueryOption(q -> q.returnReferences(references));
+    }
+
+    public final Builder<T> returnReferences(List<QueryReference> references) {
+      return applyQueryOption(q -> q.returnReferences(references));
+    }
+
+    public final Builder<T> returnMetadata(Metadata... metadata) {
+      return applyQueryOption(q -> q.returnMetadata(metadata));
+    }
+
+    public final Builder<T> returnMetadata(List<Metadata> metadata) {
+      return applyQueryOption(q -> q.returnMetadata(metadata));
+    }
+
+    private final Builder<T> applyQueryOption(Function<FetchObjects.Builder, FetchObjects.Builder> options) {
+      this.queryOptions = ObjectBuilder.partial(this.queryOptions, options);
+      return this;
+    }
+
     @Override
     public Paginator<T> build() {
       return new Paginator<>(this);

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

@@ -86,17 +86,29 @@ public final SELF where(Where where) {
     }
 
     public final SELF returnProperties(String... properties) {
-      this.returnProperties = Arrays.asList(properties);
+      return returnProperties(Arrays.asList(properties));
+    }
+
+    public final SELF returnProperties(List<String> properties) {
+      this.returnProperties.addAll(properties);
       return (SELF) this;
     }
 
     public final SELF returnReferences(QueryReference... references) {
-      this.returnReferences = Arrays.asList(references);
+      return returnReferences(Arrays.asList(references));
+    }
+
+    public final SELF returnReferences(List<QueryReference> references) {
+      this.returnReferences.addAll(references);
       return (SELF) this;
     }
 
     public final SELF returnMetadata(Metadata... metadata) {
-      this.returnMetadata = Arrays.asList(metadata);
+      return returnMetadata(metadata);
+    }
+
+    public final SELF returnMetadata(List<Metadata> metadata) {
+      this.returnMetadata.addAll(metadata);
       return (SELF) this;
     }
 

src/main/java/io/weaviate/client6/v1/internal/ObjectBuilder.java

@@ -8,4 +8,9 @@ public interface ObjectBuilder<T> {
   static <B extends ObjectBuilder<T>, T> Function<B, ObjectBuilder<T>> identity() {
     return builder -> builder;
   }
+
+  static <B extends ObjectBuilder<T>, T> Function<B, ObjectBuilder<T>> partial(Function<B, ObjectBuilder<T>> fn,
+      Function<B, B> partialFn) {
+    return partialFn.andThen(fn);
+  }
 }