diff --git a/astra-db-java/src/main/java/com/datastax/astra/client/collections/commands/cursor/CollectionFindAndRerankCursor.java b/astra-db-java/src/main/java/com/datastax/astra/client/collections/commands/cursor/CollectionFindAndRerankCursor.java index 64e8996e..d0aa5d18 100644 --- a/astra-db-java/src/main/java/com/datastax/astra/client/collections/commands/cursor/CollectionFindAndRerankCursor.java +++ b/astra-db-java/src/main/java/com/datastax/astra/client/collections/commands/cursor/CollectionFindAndRerankCursor.java @@ -26,6 +26,7 @@ import com.datastax.astra.client.core.query.Filter; import com.datastax.astra.client.core.query.Projection; import com.datastax.astra.client.core.query.Sort; +import com.datastax.astra.client.core.rerank.RerankServiceOptions; import com.datastax.astra.client.core.rerank.RerankedResult; import com.datastax.astra.client.core.vector.DataAPIVector; import com.datastax.astra.internal.command.AbstractCursor; @@ -155,6 +156,20 @@ public CollectionFindAndRerankCursor project(Projection... newProjection) return newCursor; } + /** + * Creates a new {@link CollectionFindAndRerankCursor} with an updated reranking services + * + * @param rerankService + * the new projection to apply + * @return a new {@link CollectionFindAndRerankCursor} instance with the specified reranking service + */ + public CollectionFindAndRerankCursor rerankService(RerankServiceOptions rerankService) { + checkIdleState(); + CollectionFindAndRerankCursor newCursor = this.clone(); + newCursor.options.rerankService(rerankService); + return newCursor; + } + /** * Creates a new {@link CollectionFindAndRerankCursor} with a specified sort order. * diff --git a/astra-db-java/src/main/java/com/datastax/astra/client/collections/commands/options/CollectionFindAndRerankOptions.java b/astra-db-java/src/main/java/com/datastax/astra/client/collections/commands/options/CollectionFindAndRerankOptions.java index 6e188be6..1c7d96b9 100644 --- a/astra-db-java/src/main/java/com/datastax/astra/client/collections/commands/options/CollectionFindAndRerankOptions.java +++ b/astra-db-java/src/main/java/com/datastax/astra/client/collections/commands/options/CollectionFindAndRerankOptions.java @@ -26,6 +26,8 @@ import com.datastax.astra.client.core.options.BaseOptions; import com.datastax.astra.client.core.query.Projection; import com.datastax.astra.client.core.query.Sort; +import com.datastax.astra.client.core.rerank.RerankProvider; +import com.datastax.astra.client.core.rerank.RerankServiceOptions; import lombok.Getter; import lombok.Setter; import lombok.experimental.Accessors; @@ -71,6 +73,11 @@ public class CollectionFindAndRerankOptions extends BaseOptions parameters) { + return vectorize(provider, modeName, null, parameters); + } + /** * Enable Vectorization within the collection. * @@ -328,6 +340,8 @@ public CollectionDefinition vectorize(String provider, String modeName, String s return this; } + + // --------------------- // Lexical options // --------------------- diff --git a/astra-db-java/src/test/java/com/datastax/astra/test/integration/AbstractCollectionFindAndRerankIT.java b/astra-db-java/src/test/java/com/datastax/astra/test/integration/AbstractCollectionFindAndRerankIT.java index 37d11e7c..c68eed41 100644 --- a/astra-db-java/src/test/java/com/datastax/astra/test/integration/AbstractCollectionFindAndRerankIT.java +++ b/astra-db-java/src/test/java/com/datastax/astra/test/integration/AbstractCollectionFindAndRerankIT.java @@ -47,6 +47,7 @@ import com.datastax.astra.client.core.vector.VectorOptions; import com.datastax.astra.client.collections.commands.cursor.CollectionFindAndRerankCursor; import com.datastax.astra.client.databases.commands.results.FindRerankingProvidersResult; +import com.datastax.astra.client.exceptions.DataAPIException; import lombok.Data; import lombok.extern.slf4j.Slf4j; import org.junit.jupiter.api.*; @@ -65,6 +66,7 @@ import static com.datastax.astra.client.core.lexical.AnalyzerTypes.STANDARD; import static com.datastax.astra.client.core.lexical.AnalyzerTypes.WHITESPACE; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.fail; /** * Abstract integration tests for findAndRerank, lexical search, and hybrid search. @@ -651,4 +653,89 @@ void should_replaceOne_withVectorizeSort() { assertThat(doc).isPresent(); assertThat(doc.get().getString("quote")).isEqualTo("Math is beautiful"); } + + // ========== Reranker Service Mutation ========== + + @Test + @Order(26) + void should_mutateRerankService_andVerifyChanges() { + if (skipIfNoRerankingKey()) return; + + String collectionName = "c_rerank_mutation"; + getDatabase().dropCollection(collectionName); + + // Create collection with initial rerank service configuration + RerankServiceOptions initialRerankService = new RerankServiceOptions() + .modelName("nvidia/llama-3.2-nv-rerankqa-1b-v2") + .provider("nvidia"); + + CollectionRerankOptions initialRerankOptions = new CollectionRerankOptions() + .enabled(true) + .service(initialRerankService); + + VectorServiceOptions vectorService = new VectorServiceOptions() + .provider("nvidia") + .modelName("NV-Embed-QA"); + + VectorOptions vectorOptions = new VectorOptions() + .dimension(1024) + .metric(SimilarityMetric.COSINE.getValue()) + .service(vectorService); + + CollectionDefinition initialDef = new CollectionDefinition() + .vector(vectorOptions) + .lexical(new LexicalOptions().enabled(true).analyzer(new Analyzer(STANDARD))) + .rerank(initialRerankOptions); + + Collection col = getDatabase().createCollection(collectionName, initialDef); + assertThat(col).isNotNull(); + assertThat(col.getDefinition().getRerank()).isNotNull(); + assertThat(col.getDefinition().getRerank().getService().getModelName()) + .isEqualTo("nvidia/llama-3.2-nv-rerankqa-1b-v2"); + + // Insert test documents + col.insertMany(List.of( + new Document().id("m1").put("text", "Artificial intelligence transforms technology").vectorize("Artificial intelligence transforms technology").lexical("Artificial intelligence transforms technology"), + new Document().id("m2").put("text", "Machine learning enables predictions").vectorize("Machine learning enables predictions").lexical("Machine learning enables predictions"), + new Document().id("m3").put("text", "Deep learning powers neural networks").vectorize("Deep learning powers neural networks").lexical("Deep learning powers neural networks")), + new CollectionInsertManyOptions().chunkSize(3)); + + // Perform findAndRerank with initial configuration + CollectionFindAndRerankOptions options = baseFindAndRerankOptions() + .sort(Sort.hybrid(new Hybrid("artificial intelligence and machine learning"))) + .hybridLimits(10) + .limit(3); + + List> initialResults = col.findAndRerank(options).toList(); + assertThat(initialResults).isNotEmpty(); + log.info("Initial rerank results count: {}", initialResults.size()); + + // Mutate the rerank service by creating a new collection definition + // Note: In practice, mutation would involve updating collection settings if supported + // For this test, we verify that different rerank configurations can be applied + RerankServiceOptions validUpdatedRerankService = new RerankServiceOptions() + .modelName("nvidia/llama-3.2-nv-rerankqa-1b-v2") + .provider("nvidia"); // Add custom parameters + + CollectionFindAndRerankOptions mutatedRerankOptions = new CollectionFindAndRerankOptions() + .rerankService(validUpdatedRerankService); + + List> updatedResults = col.findAndRerank(mutatedRerankOptions).toList(); + assertThat(updatedResults).isNotEmpty(); + log.info("Updated rerank results count: {}", initialResults.size()); + + CollectionFindAndRerankOptions mutatedInvalidRerankOptions = new CollectionFindAndRerankOptions() + .rerankService(new RerankServiceOptions().modelName("bla").provider("ble")); + + try { + col.findAndRerank(mutatedInvalidRerankOptions).toList(); + fail(); + } catch (DataAPIException dataAPIException) { + // Should fail + } + + // Cleanup + getDatabase().dropCollection(collectionName); + } + } diff --git a/tools/data-api-tools/src/main/java/com/datastax/astra/tool/copy/CollectionCloneSettings.java b/tools/data-api-tools/src/main/java/com/datastax/astra/tool/copy/CollectionCloneSettings.java index b114c3db..af0a8be4 100644 --- a/tools/data-api-tools/src/main/java/com/datastax/astra/tool/copy/CollectionCloneSettings.java +++ b/tools/data-api-tools/src/main/java/com/datastax/astra/tool/copy/CollectionCloneSettings.java @@ -29,6 +29,14 @@ public class CollectionCloneSettings { @Builder.Default private final int insertThreadPoolSize = 10; + /** + * Number of parallel threads for reading from source collection. + * More threads = faster reading through parallel skip/limit queries. + * Default: 5 + */ + @Builder.Default + private final int readThreadPoolSize = 5; + /** * Maximum time in seconds to wait for all insertions to complete. * Default: 300 seconds (5 minutes)