WIP: fix pushDownFilter to preserve knn scores with WHERE clause

Introduce VectorSearchQueryBuilder that keeps the knn query in a
bool.must (scoring) context while placing WHERE filters in bool.filter
(non-scoring) context. Without this fix, both knn and WHERE queries
were wrapped in bool.filter, which destroyed knn relevance scores.

Verified: knn + WHERE now returns correct scores matching DSL baseline.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: mengweieric

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

@@ -15,6 +15,7 @@
 import org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder;
 import org.opensearch.sql.opensearch.storage.scan.OpenSearchIndexScan;
 import org.opensearch.sql.opensearch.storage.scan.OpenSearchIndexScanBuilder;
+import org.opensearch.sql.opensearch.storage.scan.VectorSearchQueryBuilder;
 import org.opensearch.sql.storage.read.TableScanBuilder;
 
 /**
@@ -45,20 +46,21 @@ public VectorSearchIndex(
   public TableScanBuilder createScanBuilder() {
     final TimeValue cursorKeepAlive =
         getSettings().getSettingValue(Settings.Key.SQL_CURSOR_KEEP_ALIVE);
-    var builder = createRequestBuilder();
+    var requestBuilder = createRequestBuilder();
 
-    // Inject knn query
-    builder.pushDownFilter(buildKnnQuery());
-    builder.pushDownTrackedScore(true);
+    // 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());
+    requestBuilder.pushDownTrackedScore(true);
 
     Function<OpenSearchRequestBuilder, OpenSearchIndexScan> createScanOperator =
-        requestBuilder ->
+        rb ->
             new OpenSearchIndexScan(
                 getClient(),
-                requestBuilder.getMaxResponseSize(),
-                requestBuilder.build(
+                rb.getMaxResponseSize(),
+                rb.build(
                     getIndexName(), cursorKeepAlive, getClient(), getFieldTypes().isEmpty()));
-    return new OpenSearchIndexScanBuilder(builder, createScanOperator);
+    return new OpenSearchIndexScanBuilder(queryBuilder, createScanOperator);
   }
 
   private QueryBuilder buildKnnQuery() {

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

@@ -45,8 +45,8 @@ public OpenSearchIndexScanBuilder(
     this.scanFactory = scanFactory;
   }
 
-  /** Constructor used for unit tests. */
-  protected OpenSearchIndexScanBuilder(
+  /** Constructor that accepts a custom PushDownQueryBuilder delegate. */
+  public OpenSearchIndexScanBuilder(
       PushDownQueryBuilder translator,
       Function<OpenSearchRequestBuilder, OpenSearchIndexScan> scanFactory) {
     this.delegate = translator;

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

@@ -0,0 +1,50 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.opensearch.storage.scan;
+
+import org.opensearch.index.query.BoolQueryBuilder;
+import org.opensearch.index.query.QueryBuilder;
+import org.opensearch.index.query.QueryBuilders;
+import org.opensearch.sql.expression.Expression;
+import org.opensearch.sql.expression.FunctionExpression;
+import org.opensearch.sql.expression.function.OpenSearchFunctions;
+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;
+
+/**
+ * Query builder for vector search that keeps the knn query in a scoring (must) context and puts
+ * WHERE filters in a non-scoring (filter) context. This prevents the knn relevance scores from
+ * being destroyed when a WHERE clause is pushed down.
+ *
+ * <p>Without this, the default pushDownFilter wraps both queries into bool.filter, which is a
+ * non-scoring context.
+ */
+public class VectorSearchQueryBuilder extends OpenSearchIndexScanQueryBuilder {
+
+  private final QueryBuilder knnQuery;
+
+  public VectorSearchQueryBuilder(OpenSearchRequestBuilder requestBuilder, QueryBuilder knnQuery) {
+    super(requestBuilder);
+    // Set knn as the initial query (scoring context)
+    requestBuilder.getSourceBuilder().query(knnQuery);
+    this.knnQuery = knnQuery;
+  }
+
+  @Override
+  public boolean pushDownFilter(LogicalFilter filter) {
+    FilterQueryBuilder queryBuilder = new FilterQueryBuilder(new DefaultExpressionSerializer());
+    Expression queryCondition = filter.getCondition();
+    QueryBuilder whereQuery = queryBuilder.build(queryCondition);
+
+    // Combine: knn in must (scores), WHERE in filter (no scoring impact)
+    BoolQueryBuilder combined =
+        QueryBuilders.boolQuery().must(knnQuery).filter(whereQuery);
+    requestBuilder.getSourceBuilder().query(combined);
+    return true;
+  }
+}

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

@@ -0,0 +1,48 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.opensearch.storage.scan;
+
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.Mockito.mock;
+
+import org.junit.jupiter.api.Test;
+import org.opensearch.index.query.QueryBuilder;
+import org.opensearch.index.query.WrapperQueryBuilder;
+import org.opensearch.sql.common.setting.Settings;
+import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory;
+import org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder;
+
+class VectorSearchQueryBuilderTest {
+
+  @Test
+  void knnQuerySetAsScoringQuery() {
+    var requestBuilder = createRequestBuilder();
+    var knnQuery = new WrapperQueryBuilder("eyJrbm4iOnt9fQ==");
+
+    new VectorSearchQueryBuilder(requestBuilder, knnQuery);
+
+    QueryBuilder query = requestBuilder.getSourceBuilder().query();
+    assertTrue(query instanceof WrapperQueryBuilder,
+        "knn query should be set directly as top-level query (scoring context)");
+  }
+
+  @Test
+  void knnQueryNotWrappedInFilterWhenNoWhere() {
+    var requestBuilder = createRequestBuilder();
+    var knnQuery = new WrapperQueryBuilder("eyJrbm4iOnt9fQ==");
+
+    new VectorSearchQueryBuilder(requestBuilder, knnQuery);
+
+    QueryBuilder query = requestBuilder.getSourceBuilder().query();
+    assertTrue(query instanceof WrapperQueryBuilder,
+        "Without WHERE clause, knn query should NOT be wrapped in bool.filter");
+  }
+
+  private OpenSearchRequestBuilder createRequestBuilder() {
+    return new OpenSearchRequestBuilder(
+        mock(OpenSearchExprValueFactory.class), 10000, mock(Settings.class));
+  }
+}