Add LIMIT > k rejection in top-k vector search mode

VectorSearchQueryBuilder now accepts options map and rejects
pushDownLimit when LIMIT exceeds k. Radial modes (max_distance,
min_score) have no LIMIT restriction.

Signed-off-by: Eric Wei <mengwei.eric@gmail.com>

Author: mengweieric

opensearch/src/main/java/org/opensearch/sql/opensearch/storage/VectorSearchIndex.java

@@ -48,7 +48,7 @@ public TableScanBuilder createScanBuilder() {
 
     // Use VectorSearchQueryBuilder to keep knn in must (scoring) context.
     // WHERE filters will be placed in filter (non-scoring) context.
-    var queryBuilder = new VectorSearchQueryBuilder(requestBuilder, buildKnnQuery());
+    var queryBuilder = new VectorSearchQueryBuilder(requestBuilder, buildKnnQuery(), options);
     requestBuilder.pushDownTrackedScore(true);
 
     // Top-k mode: default size to k so queries without LIMIT return k results

opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/VectorSearchQueryBuilder.java

@@ -5,14 +5,18 @@
 
 package org.opensearch.sql.opensearch.storage.scan;
 
+import java.util.Map;
 import org.opensearch.index.query.BoolQueryBuilder;
 import org.opensearch.index.query.QueryBuilder;
 import org.opensearch.index.query.QueryBuilders;
+import org.opensearch.sql.exception.ExpressionEvaluationException;
 import org.opensearch.sql.expression.Expression;
 import org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder;
 import org.opensearch.sql.opensearch.storage.script.filter.FilterQueryBuilder;
 import org.opensearch.sql.opensearch.storage.serde.DefaultExpressionSerializer;
 import org.opensearch.sql.planner.logical.LogicalFilter;
+import org.opensearch.sql.planner.logical.LogicalLimit;
+import org.opensearch.sql.planner.logical.LogicalSort;
 
 /**
  * Query builder for vector search that keeps the knn query in a scoring (must) context and puts
@@ -25,12 +29,14 @@
 public class VectorSearchQueryBuilder extends OpenSearchIndexScanQueryBuilder {
 
   private final QueryBuilder knnQuery;
+  private final Map<String, String> options;
 
-  public VectorSearchQueryBuilder(OpenSearchRequestBuilder requestBuilder, QueryBuilder knnQuery) {
+  public VectorSearchQueryBuilder(
+      OpenSearchRequestBuilder requestBuilder, QueryBuilder knnQuery, Map<String, String> options) {
     super(requestBuilder);
-    // Set knn as the initial query (scoring context)
     requestBuilder.getSourceBuilder().query(knnQuery);
     this.knnQuery = knnQuery;
+    this.options = options;
   }
 
   @Override
@@ -44,4 +50,25 @@ public boolean pushDownFilter(LogicalFilter filter) {
     requestBuilder.getSourceBuilder().query(combined);
     return true;
   }
+
+  @Override
+  public boolean pushDownLimit(LogicalLimit limit) {
+    if (options.containsKey("k")) {
+      int k = Integer.parseInt(options.get("k"));
+      if (limit.getLimit() > k) {
+        throw new ExpressionEvaluationException(
+            String.format("LIMIT %d exceeds k=%d in top-k vector search", limit.getLimit(), k));
+      }
+    }
+    return super.pushDownLimit(limit);
+  }
+
+  @Override
+  public boolean pushDownSort(LogicalSort sort) {
+    // Vector search returns results sorted by _score DESC by default.
+    // Reject non-trivial sort pushdowns — only _score DESC is meaningful.
+    // For now, let the parent handle it; unsupported sort rejection is
+    // deferred until we can inspect the sort expression for _score references.
+    return super.pushDownSort(sort);
+  }
 }

opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/VectorSearchQueryBuilderTest.java

@@ -6,21 +6,25 @@
 package org.opensearch.sql.opensearch.storage.scan;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.Mockito.mock;
 import static org.opensearch.sql.data.type.ExprCoreType.STRING;
 
 import java.util.Collections;
+import java.util.Map;
 import org.junit.jupiter.api.Test;
 import org.opensearch.index.query.BoolQueryBuilder;
 import org.opensearch.index.query.QueryBuilder;
 import org.opensearch.index.query.WrapperQueryBuilder;
 import org.opensearch.sql.common.setting.Settings;
+import org.opensearch.sql.exception.ExpressionEvaluationException;
 import org.opensearch.sql.expression.DSL;
 import org.opensearch.sql.expression.ReferenceExpression;
 import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory;
 import org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder;
 import org.opensearch.sql.planner.logical.LogicalFilter;
+import org.opensearch.sql.planner.logical.LogicalLimit;
 import org.opensearch.sql.planner.logical.LogicalValues;
 
 class VectorSearchQueryBuilderTest {
@@ -30,7 +34,7 @@ void knnQuerySetAsScoringQuery() {
     var requestBuilder = createRequestBuilder();
     var knnQuery = new WrapperQueryBuilder("{\"knn\":{}}");
 
-    new VectorSearchQueryBuilder(requestBuilder, knnQuery);
+    new VectorSearchQueryBuilder(requestBuilder, knnQuery, Map.of("k", "5"));
 
     QueryBuilder query = requestBuilder.getSourceBuilder().query();
     assertTrue(
@@ -42,7 +46,7 @@ void knnQuerySetAsScoringQuery() {
   void pushDownFilterKeepsKnnInScoringContext() {
     var requestBuilder = createRequestBuilder();
     var knnQuery = new WrapperQueryBuilder("{\"knn\":{}}");
-    var builder = new VectorSearchQueryBuilder(requestBuilder, knnQuery);
+    var builder = new VectorSearchQueryBuilder(requestBuilder, knnQuery, Map.of("k", "5"));
 
     // Simulate WHERE name = 'John'
     var condition = DSL.equal(new ReferenceExpression("name", STRING), DSL.literal("John"));
@@ -62,6 +66,60 @@ void pushDownFilterKeepsKnnInScoringContext() {
         "must clause should contain the original knn WrapperQueryBuilder");
   }
 
+  @Test
+  void pushDownLimitWithinKSucceeds() {
+    var requestBuilder = createRequestBuilder();
+    var knnQuery = new WrapperQueryBuilder("{\"knn\":{}}");
+    var builder = new VectorSearchQueryBuilder(requestBuilder, knnQuery, Map.of("k", "5"));
+
+    var dummyChild = new LogicalValues(Collections.emptyList());
+    var limit = new LogicalLimit(dummyChild, 3, 0);
+
+    boolean pushed = builder.pushDownLimit(limit);
+    assertTrue(pushed, "LIMIT within k should succeed");
+  }
+
+  @Test
+  void pushDownLimitExceedingKThrows() {
+    var requestBuilder = createRequestBuilder();
+    var knnQuery = new WrapperQueryBuilder("{\"knn\":{}}");
+    var builder = new VectorSearchQueryBuilder(requestBuilder, knnQuery, Map.of("k", "5"));
+
+    var dummyChild = new LogicalValues(Collections.emptyList());
+    var limit = new LogicalLimit(dummyChild, 10, 0);
+
+    ExpressionEvaluationException ex =
+        assertThrows(ExpressionEvaluationException.class, () -> builder.pushDownLimit(limit));
+    assertTrue(ex.getMessage().contains("LIMIT 10 exceeds k=5"));
+  }
+
+  @Test
+  void pushDownLimitEqualToKSucceeds() {
+    var requestBuilder = createRequestBuilder();
+    var knnQuery = new WrapperQueryBuilder("{\"knn\":{}}");
+    var builder = new VectorSearchQueryBuilder(requestBuilder, knnQuery, Map.of("k", "5"));
+
+    var dummyChild = new LogicalValues(Collections.emptyList());
+    var limit = new LogicalLimit(dummyChild, 5, 0);
+
+    boolean pushed = builder.pushDownLimit(limit);
+    assertTrue(pushed, "LIMIT equal to k should succeed");
+  }
+
+  @Test
+  void pushDownLimitRadialModeNoRestriction() {
+    var requestBuilder = createRequestBuilder();
+    var knnQuery = new WrapperQueryBuilder("{\"knn\":{}}");
+    var builder =
+        new VectorSearchQueryBuilder(requestBuilder, knnQuery, Map.of("max_distance", "10.0"));
+
+    var dummyChild = new LogicalValues(Collections.emptyList());
+    var limit = new LogicalLimit(dummyChild, 100, 0);
+
+    boolean pushed = builder.pushDownLimit(limit);
+    assertTrue(pushed, "Radial mode should not restrict LIMIT");
+  }
+
   private OpenSearchRequestBuilder createRequestBuilder() {
     return new OpenSearchRequestBuilder(
         mock(OpenSearchExprValueFactory.class), 10000, mock(Settings.class));