refactor: make reranking candidate-read budgets explicit

src/main/java/io/stargate/sgv2/jsonapi/api/model/command/CommandContext.java

@@ -2,7 +2,6 @@
 
 import com.google.common.base.Preconditions;
 import io.micrometer.core.instrument.MeterRegistry;
-import io.stargate.sgv2.jsonapi.api.model.command.impl.FindAndRerankCommand;
 import io.stargate.sgv2.jsonapi.api.model.command.tracing.DefaultRequestTracing;
 import io.stargate.sgv2.jsonapi.api.model.command.tracing.RequestTracing;
 import io.stargate.sgv2.jsonapi.api.request.RequestContext;
@@ -67,9 +66,6 @@ public class CommandContext<SchemaT extends SchemaObject> implements LoggingMDCC
   // add to this list in the ctor. See {@link #addToMDC()} and {@link #removeFromMDC()}
   private final List<LoggingMDCContext> loggingMDCContexts = new ArrayList<>();
 
-  // see accessors
-  private FindAndRerankCommand.HybridLimits hybridLimits;
-
   // used to track the features used in the command
   private final CommandFeatures commandFeatures;
 
@@ -110,22 +106,6 @@ public static BuilderSupplier builderSupplier() {
     return new BuilderSupplier();
   }
 
-  /**
-   * HACK: for https://github.com/stargate/data-api/issues/1961 This is a temporary work around for
-   * needing to pass the page size to the FindCollectionOperation when doing the inner finds for
-   * findAndRerank because they will only run the command once, and not multiple times to exhaust
-   * the cursor.
-   *
-   * @return
-   */
-  public FindAndRerankCommand.HybridLimits getHybridLimits() {
-    return hybridLimits;
-  }
-
-  public void setHybridLimits(FindAndRerankCommand.HybridLimits hybridLimits) {
-    this.hybridLimits = hybridLimits;
-  }
-
   public SchemaT schemaObject() {
     return schemaObject;
   }

src/main/java/io/stargate/sgv2/jsonapi/service/operation/reranking/IntermediateCollectionReadTask.java

@@ -72,7 +72,7 @@ protected IntermediateReadResultSupplier buildResultSupplier(
     }
 
     Operation<CollectionSchemaObject> findOperation =
-        findCommandResolver.resolveCommand(commandContext, findCommand);
+        findCommandResolver.resolveCollectionCandidateCommand(commandContext, findCommand);
     return new IntermediateReadResultSupplier(
         () -> {
           commandContext

src/main/java/io/stargate/sgv2/jsonapi/service/resolver/FindAndRerankOperationBuilder.java

@@ -181,15 +181,6 @@ private TaskGroupAndDeferrables<IntermediateCollectionReadTask, CollectionSchema
     TaskGroup<IntermediateCollectionReadTask, CollectionSchemaObject> taskGroup =
         new TaskGroup<>(false);
 
-    // Hack: See https://github.com/stargate/data-api/issues/1961
-    // copying the hybrid limits on the command context so the find command resolver can pick it up
-    // when the command runs later, so we can set the page size to be the same as the limit
-    commandContext.setHybridLimits(
-        getOrDefault(
-            command.options(),
-            FindAndRerankCommand.Options::hybridLimits,
-            FindAndRerankCommand.HybridLimits.DEFAULT));
-
     var includeScores =
         getOrDefault(command.options(), FindAndRerankCommand.Options::includeScores, false);
     var includeSortVector =

src/main/java/io/stargate/sgv2/jsonapi/service/resolver/FindCommandResolver.java

@@ -1,12 +1,12 @@
 package io.stargate.sgv2.jsonapi.service.resolver;
 
 import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.base.Preconditions;
 import io.micrometer.core.instrument.MeterRegistry;
 import io.stargate.sgv2.jsonapi.api.model.command.CommandContext;
 import io.stargate.sgv2.jsonapi.api.model.command.clause.sort.SortClause;
 import io.stargate.sgv2.jsonapi.api.model.command.clause.sort.SortExpression;
 import io.stargate.sgv2.jsonapi.api.model.command.impl.FindCommand;
-import io.stargate.sgv2.jsonapi.api.model.command.impl.FindOneCommand;
 import io.stargate.sgv2.jsonapi.api.v1.metrics.JsonApiMetricsConfig;
 import io.stargate.sgv2.jsonapi.config.OperationsConfig;
 import io.stargate.sgv2.jsonapi.exception.SortException;
@@ -24,7 +24,7 @@
 import java.util.List;
 import java.util.Map;
 
-/** Resolves the {@link FindOneCommand } */
+/** Resolves collection and table {@link FindCommand}s. */
 @ApplicationScoped
 public class FindCommandResolver implements CommandResolver<FindCommand> {
 
@@ -72,6 +72,34 @@ public Operation<TableSchemaObject> resolveTableCommand(
   @Override
   public Operation<CollectionSchemaObject> resolveCollectionCommand(
       CommandContext<CollectionSchemaObject> commandContext, FindCommand command) {
+    int pageSize = operationsConfig.defaultPageSize();
+    return resolveCollectionRead(
+        commandContext, command, new CollectionReadLimits(pageSize, pageSize, pageSize));
+  }
+
+  /**
+   * Resolves an internal reranking read. The planner validates user-supplied hybridLimits before
+   * building the inner find, so a missing or non-positive limit here is an internal state error.
+   */
+  public Operation<CollectionSchemaObject> resolveCollectionCandidateCommand(
+      CommandContext<CollectionSchemaObject> commandContext, FindCommand command) {
+    var options = command.options();
+    Preconditions.checkState(
+        options != null && options.limit() != null && options.limit() > 0,
+        "Candidate reads require an explicit positive options.limit");
+    return resolveCollectionRead(
+        commandContext,
+        command,
+        new CollectionReadLimits(
+            Math.min(options.limit(), operationsConfig.maxVectorSearchLimit()),
+            options.limit(),
+            operationsConfig.defaultPageSize()));
+  }
+
+  private FindCollectionOperation resolveCollectionRead(
+      CommandContext<CollectionSchemaObject> commandContext,
+      FindCommand command,
+      CollectionReadLimits readLimits) {
 
     var resolvedDbLogicalExpression =
         collectionFilterResolver.resolve(commandContext, command).target();
@@ -128,18 +156,13 @@ public Operation<CollectionSchemaObject> resolveCollectionCommand(
           Math.min(
               limit, operationsConfig.maxVectorSearchLimit()); // Max vector search support is 1000
 
-      // Hack: See https://github.com/stargate/data-api/issues/1961
-      int pageSize =
-          commandContext.getHybridLimits() == null
-              ? operationsConfig.defaultPageSize()
-              : commandContext.getHybridLimits().vectorLimit();
       return FindCollectionOperation.vsearch(
           commandContext,
           resolvedDbLogicalExpression,
           command.buildProjector(includeSimilarity),
           pageState,
           limit,
-          pageSize,
+          readLimits.vectorPageSize(),
           CollectionReadType.DOCUMENT,
           objectMapper,
           vector,
@@ -149,18 +172,13 @@ public Operation<CollectionSchemaObject> resolveCollectionCommand(
     // BM25 search / sort?
     SortExpression bm25Expr = SortClauseUtil.resolveBM25Search(sortClause);
     if (bm25Expr != null) {
-      // Hack: See https://github.com/stargate/data-api/issues/1961
-      int pageSize =
-          commandContext.getHybridLimits() == null
-              ? operationsConfig.defaultPageSize()
-              : commandContext.getHybridLimits().lexicalLimit();
       return FindCollectionOperation.bm25Multi(
           commandContext,
           resolvedDbLogicalExpression,
           command.buildProjector(),
           pageState,
           limit,
-          pageSize,
+          readLimits.lexicalPageSize(),
           CollectionReadType.DOCUMENT,
           objectMapper,
           bm25Expr);
@@ -169,16 +187,12 @@ public Operation<CollectionSchemaObject> resolveCollectionCommand(
     List<FindCollectionOperation.OrderBy> orderBy = SortClauseUtil.resolveOrderBy(sortClause);
     // if orderBy present
     if (orderBy != null) {
-      // Hack: See https://github.com/stargate/data-api/issues/1961
-      // not commandContext.getHybridLimits() becuase there is no limit for a non ANN or BM25 query
       return FindCollectionOperation.sorted(
           commandContext,
           resolvedDbLogicalExpression,
           command.buildProjector(),
           pageState,
-          // For in-memory sorting if no limit provided in the request will use
-          // documentConfig.defaultPageSize() as limit
-          Math.min(limit, operationsConfig.defaultPageSize()),
+          Math.min(limit, readLimits.sortedLimit()),
           // For in-memory sorting we read more data than needed, so defaultSortPageSize like 100
           operationsConfig.defaultSortPageSize(),
           CollectionReadType.SORTED_DOCUMENT,
@@ -188,8 +202,6 @@ public Operation<CollectionSchemaObject> resolveCollectionCommand(
           operationsConfig.maxDocumentSortCount(),
           includeSortVector);
     }
-    // Hack: See https://github.com/stargate/data-api/issues/1961
-    // not commandContext.getHybridLimits() becuase there is no limit for a non ANN or BM25 query
     return FindCollectionOperation.unsorted(
         commandContext,
         resolvedDbLogicalExpression,
@@ -201,4 +213,7 @@ public Operation<CollectionSchemaObject> resolveCollectionCommand(
         objectMapper,
         includeSortVector);
   }
+
+  /** Concrete execution limits supplied by the caller, independent of query resolution. */
+  private record CollectionReadLimits(int vectorPageSize, int lexicalPageSize, int sortedLimit) {}
 }

src/test/java/io/stargate/sgv2/jsonapi/service/operation/reranking/FindAndRerankPlannerTest.java

@@ -11,6 +11,7 @@
 import com.fasterxml.jackson.core.json.JsonReadFeature;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import io.stargate.sgv2.jsonapi.TestConstants;
+import io.stargate.sgv2.jsonapi.api.model.command.CommandErrorFactory;
 import io.stargate.sgv2.jsonapi.api.model.command.clause.sort.FindAndRerankSort;
 import io.stargate.sgv2.jsonapi.api.model.command.clause.sort.LegMode;
 import io.stargate.sgv2.jsonapi.api.model.command.impl.FindAndRerankCommand;
@@ -30,6 +31,7 @@
 import io.stargate.sgv2.jsonapi.service.reranking.configuration.RerankingProvidersConfig;
 import io.stargate.sgv2.jsonapi.service.schema.collections.CollectionRerankDef;
 import io.stargate.sgv2.jsonapi.service.schema.collections.CollectionSchemaObject;
+import java.util.List;
 import java.util.stream.Stream;
 import org.junit.jupiter.api.Nested;
 import org.junit.jupiter.api.Test;
@@ -38,6 +40,7 @@
 import org.junit.jupiter.params.provider.Arguments;
 import org.junit.jupiter.params.provider.EnumSource;
 import org.junit.jupiter.params.provider.MethodSource;
+import org.junit.jupiter.params.provider.ValueSource;
 
 /**
  * Tests for {@link FindAndRerankPlanner}, plain JUnit without Quarkus.
@@ -320,6 +323,29 @@ void failVectorizeSortWithoutVectorizeService() {
           () -> plan(hybridText("cheese"), RERANK_OPTIONS, VECTOR_LEXICAL));
     }
 
+    @ParameterizedTest
+    @ValueSource(ints = {0, -1})
+    void nonPositiveUserLimitsReturnRequestErrors(int limit) {
+      for (var limitsJson :
+          List.of(
+              String.valueOf(limit),
+              "{'$vector': %d, '$lexical': 10}".formatted(limit),
+              "{'$vector': 10, '$lexical': %d}".formatted(limit))) {
+        var error =
+            assertPlanFails(
+                RequestException.class,
+                RequestException.Code.COMMAND_FIELD_VALUE_INVALID,
+                () ->
+                    plan(
+                        hybridObject(null, "cows", VECTOR),
+                        "{'rerankQuery': 'q', 'rerankOn': 'body', 'hybridLimits': %s}"
+                            .formatted(limitsJson),
+                        VECTOR_LEXICAL));
+        assertThat(CommandErrorFactory.create(error).errorCode())
+            .isEqualTo(RequestException.Code.COMMAND_FIELD_VALUE_INVALID.name());
+      }
+    }
+
     @Test
     void failHybridLimitsOutOfBounds() {
       var expectedBounds =

src/test/java/io/stargate/sgv2/jsonapi/service/resolver/FindAndRerankOperationBuilderBehaviorTest.java

@@ -43,8 +43,7 @@
 /**
  * Behavior baseline for {@link FindAndRerankOperationBuilder#build()}: pins the outcome for many
  * request shapes and collection configs. A row either builds an operation, and then pins the reads
- * it runs (rank source and limit, in order) and the hybrid limits on the command context, or it
- * pins the error code.
+ * it runs (rank source and limit, in order), or it pins the error code.
  *
  * <p>Some rows pin behaviour that is easy to change by accident:
  *
@@ -294,13 +293,13 @@ static Stream<Row> rows() {
             "{'$hybrid': {'$vector': " + VECTOR_JSON + ", '$lexical': 'cows'}}",
             "{'rerankQuery': 'q', 'rerankOn': 'body', 'hybridLimits': {'$vector': 100, '$lexical': 25}}",
             Collection.VECTOR_LEXICAL,
-            success(100, 25, lexicalRead(25), vectorRead(100))),
+            success(lexicalRead(25), vectorRead(100))),
         row(
             "G23",
             "{'$hybrid': {'$vector': " + VECTOR_JSON + ", '$lexical': 'cows'}}",
             "{'rerankQuery': 'q', 'rerankOn': 'body', 'hybridLimits': 30}",
             Collection.VECTOR_LEXICAL,
-            success(30, 30, lexicalRead(30), vectorRead(30))),
+            success(lexicalRead(30), vectorRead(30))),
 
         // ---- lexical-only $hybrid ----
         row(
@@ -464,7 +463,7 @@ static Stream<Row> rows() {
             "{'$hybrid': {'$vectorize': 'cheese'}}",
             "{'hybridLimits': 30}",
             Collection.VECTORIZE_NO_LEXICAL,
-            success(30, 30, vectorRead(30))),
+            success(vectorRead(30))),
         row(
             "R42",
             "{'$hybrid': {'$vector': " + VECTOR_JSON + ", '$lexical': 'cows'}}",
@@ -686,12 +685,6 @@ void buildOutcome(Row row) throws Exception {
                     .toList())
             .as(row.id() + " reads, in order")
             .containsExactlyElementsOf(success.reads());
-        assertThat(commandContext.getHybridLimits().vectorLimit())
-            .as(row.id() + " hybridLimits.$vector on the command context")
-            .isEqualTo(success.vectorLimit());
-        assertThat(commandContext.getHybridLimits().lexicalLimit())
-            .as(row.id() + " hybridLimits.$lexical on the command context")
-            .isEqualTo(success.lexicalLimit());
       }
       case Failure failure -> {
         var ex =
@@ -774,15 +767,9 @@ private static Row row(
     return new Row(id, sort, options, filter, projection, collection, expected);
   }
 
-  /** build() returns an operation with these reads, and the default hybrid limits. */
+  /** build() returns an operation with these reads and their candidate limits. */
   private static Expected success(ExpectedRead... reads) {
-    var defaults = FindAndRerankCommand.HybridLimits.DEFAULT;
-    return new Success(defaults.vectorLimit(), defaults.lexicalLimit(), List.of(reads));
-  }
-
-  /** build() returns an operation with these reads, and these hybrid limits. */
-  private static Expected success(int vectorLimit, int lexicalLimit, ExpectedRead... reads) {
-    return new Success(vectorLimit, lexicalLimit, List.of(reads));
+    return new Success(List.of(reads));
   }
 
   private static <T extends APIException> Expected failure(
@@ -857,12 +844,8 @@ public String toString() {
 
   private sealed interface Expected permits Success, Failure {}
 
-  /**
-   * build() returns an operation with these reads, and copies these hybrid limits onto the command
-   * context.
-   */
-  private record Success(int vectorLimit, int lexicalLimit, List<ExpectedRead> reads)
-      implements Expected {}
+  /** build() returns an operation with these reads and their candidate limits. */
+  private record Success(List<ExpectedRead> reads) implements Expected {}
 
   /**
    * build() throws exactly this API exception class with this code, and optionally with this