Implement efficient filter pushdown branching and build-time validation in VectorSearchQueryBuilder

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

Author: mengweieric

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

@@ -6,6 +6,7 @@
 package org.opensearch.sql.opensearch.storage.scan;
 
 import java.util.Map;
+import java.util.function.Function;
 import org.apache.commons.lang3.tuple.Pair;
 import org.opensearch.index.query.BoolQueryBuilder;
 import org.opensearch.index.query.QueryBuilder;
@@ -16,6 +17,7 @@
 import org.opensearch.sql.expression.Expression;
 import org.opensearch.sql.expression.ReferenceExpression;
 import org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder;
+import org.opensearch.sql.opensearch.storage.FilterType;
 import org.opensearch.sql.opensearch.storage.script.filter.FilterQueryBuilder;
 import org.opensearch.sql.opensearch.storage.serde.DefaultExpressionSerializer;
 import org.opensearch.sql.planner.logical.LogicalFilter;
@@ -27,31 +29,62 @@
  * 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.
+ * <p>Supports two filter placement strategies via {@link FilterType}:
+ *
+ * <ul>
+ *   <li>{@code POST} — WHERE in {@code bool.filter} outside knn (post-filtering, default)
+ *   <li>{@code EFFICIENT} — WHERE inside {@code knn.filter} for pre-filtering during ANN search
+ * </ul>
  */
 public class VectorSearchQueryBuilder extends OpenSearchIndexScanQueryBuilder {
 
   private final QueryBuilder knnQuery;
   private final Map<String, String> options;
+  private final FilterType filterType;
+  private final boolean filterTypeExplicit;
+  private final Function<QueryBuilder, QueryBuilder> rebuildKnnWithFilter;
+  private boolean filterPushed = false;
 
+  /** Full constructor with filter type support. */
   public VectorSearchQueryBuilder(
-      OpenSearchRequestBuilder requestBuilder, QueryBuilder knnQuery, Map<String, String> options) {
+      OpenSearchRequestBuilder requestBuilder,
+      QueryBuilder knnQuery,
+      Map<String, String> options,
+      FilterType filterType,
+      boolean filterTypeExplicit,
+      Function<QueryBuilder, QueryBuilder> rebuildKnnWithFilter) {
     super(requestBuilder);
     requestBuilder.getSourceBuilder().query(knnQuery);
     this.knnQuery = knnQuery;
     this.options = options;
+    this.filterType = filterType != null ? filterType : FilterType.POST;
+    this.filterTypeExplicit = filterTypeExplicit;
+    this.rebuildKnnWithFilter = rebuildKnnWithFilter;
+  }
+
+  /** Backward-compatible constructor — defaults to POST, not explicit. */
+  public VectorSearchQueryBuilder(
+      OpenSearchRequestBuilder requestBuilder,
+      QueryBuilder knnQuery,
+      Map<String, String> options) {
+    this(requestBuilder, knnQuery, options, FilterType.POST, false, null);
   }
 
   @Override
   public boolean pushDownFilter(LogicalFilter filter) {
     FilterQueryBuilder queryBuilder = new FilterQueryBuilder(new DefaultExpressionSerializer());
     Expression queryCondition = filter.getCondition();
     QueryBuilder whereQuery = queryBuilder.build(queryCondition);
+    filterPushed = true;
 
-    // Combine: knn in must (scores), WHERE in filter (no scoring impact)
-    BoolQueryBuilder combined = QueryBuilders.boolQuery().must(knnQuery).filter(whereQuery);
-    requestBuilder.getSourceBuilder().query(combined);
+    if (filterType == FilterType.EFFICIENT) {
+      QueryBuilder rebuiltKnn = rebuildKnnWithFilter.apply(whereQuery);
+      requestBuilder.getSourceBuilder().query(rebuiltKnn);
+    } else {
+      // POST mode: knn in must (scores), WHERE in filter (no scoring impact)
+      BoolQueryBuilder combined = QueryBuilders.boolQuery().must(knnQuery).filter(whereQuery);
+      requestBuilder.getSourceBuilder().query(combined);
+    }
     return true;
   }
 
@@ -94,4 +127,13 @@ public boolean pushDownSort(LogicalSort sort) {
     }
     return true;
   }
+
+  @Override
+  public OpenSearchRequestBuilder build() {
+    if (filterTypeExplicit && !filterPushed) {
+      throw new ExpressionEvaluationException(
+          "filter_type requires a pushdownable WHERE clause");
+    }
+    return super.build();
+  }
 }

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

@@ -14,6 +14,7 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.function.Function;
 import org.junit.jupiter.api.Test;
 import org.opensearch.index.query.BoolQueryBuilder;
 import org.opensearch.index.query.QueryBuilder;
@@ -25,6 +26,7 @@
 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.opensearch.storage.FilterType;
 import org.opensearch.sql.planner.logical.LogicalFilter;
 import org.opensearch.sql.planner.logical.LogicalLimit;
 import org.opensearch.sql.planner.logical.LogicalValues;
@@ -267,6 +269,54 @@ void pushDownFilterCompoundPredicateSurvives() {
     assertEquals(1, boolQuery.filter().size(), "compound WHERE should be in filter (non-scoring)");
   }
 
+  @Test
+  void pushDownFilterEfficientPlacesInsideKnn() {
+    var requestBuilder = createRequestBuilder();
+    var knnQuery = new WrapperQueryBuilder("{\"knn\":{}}");
+    // Callback simulates VectorSearchIndex rebuilding knn with filter
+    Function<QueryBuilder, QueryBuilder> rebuildWithFilter =
+        whereQuery -> new WrapperQueryBuilder("{\"knn\":{\"filter\":\"embedded\"}}");
+    var builder =
+        new VectorSearchQueryBuilder(
+            requestBuilder, knnQuery, Map.of("k", "5"),
+            FilterType.EFFICIENT, true, rebuildWithFilter);
+
+    var condition = DSL.equal(new ReferenceExpression("city", STRING), DSL.literal("Miami"));
+    var dummyChild = new LogicalValues(Collections.emptyList());
+    var filter = new LogicalFilter(dummyChild, condition);
+
+    boolean pushed = builder.pushDownFilter(filter);
+
+    assertTrue(pushed, "pushDownFilter should succeed");
+    QueryBuilder resultQuery = requestBuilder.getSourceBuilder().query();
+    assertTrue(
+        resultQuery instanceof WrapperQueryBuilder,
+        "Efficient filter should produce a WrapperQueryBuilder (rebuilt knn), not BoolQuery");
+  }
+
+  @Test
+  void pushDownFilterExplicitPostProducesBool() {
+    var requestBuilder = createRequestBuilder();
+    var knnQuery = new WrapperQueryBuilder("{\"knn\":{}}");
+    var builder =
+        new VectorSearchQueryBuilder(
+            requestBuilder, knnQuery, Map.of("k", "5"),
+            FilterType.POST, true, null);
+
+    var condition = DSL.equal(new ReferenceExpression("name", STRING), DSL.literal("John"));
+    var dummyChild = new LogicalValues(Collections.emptyList());
+    var filter = new LogicalFilter(dummyChild, condition);
+
+    boolean pushed = builder.pushDownFilter(filter);
+
+    assertTrue(pushed);
+    QueryBuilder resultQuery = requestBuilder.getSourceBuilder().query();
+    assertTrue(resultQuery instanceof BoolQueryBuilder);
+    BoolQueryBuilder boolQuery = (BoolQueryBuilder) resultQuery;
+    assertEquals(1, boolQuery.must().size());
+    assertEquals(1, boolQuery.filter().size());
+  }
+
   private OpenSearchRequestBuilder createRequestBuilder() {
     return new OpenSearchRequestBuilder(
         mock(OpenSearchExprValueFactory.class), 10000, mock(Settings.class));