Expose neural filters to query visitors (#1992)

* Expose neural filters to query visitors

Signed-off-by: Cédric Pelvet <cedric.pelvet@gmail.com>

* Add changelog entry for query visitor fix

Signed-off-by: Cédric Pelvet <cedric.pelvet@gmail.com>

---------

Signed-off-by: Cédric Pelvet <cedric.pelvet@gmail.com>

Author: sharp-pixel

CHANGELOG.md

@@ -11,6 +11,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Add `model_selection` (language_option/model_type) parameter to semantic field to resolve the model id from cluster settings ([#1918](https://github.com/opensearch-project/neural-search/issues/1918))
 
 ### Bug Fixes
+* [Neural Query] Expose embedded filters to `QueryBuilderVisitor` traversal ([#1992](https://github.com/opensearch-project/neural-search/pull/1992))
 * [Hybrid Query] Fix NoSuchElementException in hybrid query with sort/search_after when a shard returns no results ([#1939](https://github.com/opensearch-project/neural-search/pull/1939))
 * [SemanticHighlighter] Fix SemanticHighlighterExtBuilder.toXContent ([#1906](https://github.com/opensearch-project/neural-search/issues/1906)) (query-insights [#651](https://github.com/opensearch-project/query-insights/issues/651))
 * [Sparse ANN] Fold sparse vector tokens into the signed-short range (modulus 32768) so folded tokens are never sign-extended to a negative value when stored in short[] ([#1926](https://github.com/opensearch-project/neural-search/pull/1926))

src/main/java/org/opensearch/neuralsearch/query/NeuralQueryBuilder.java

@@ -57,6 +57,7 @@
 import lombok.experimental.Accessors;
 import org.apache.commons.lang3.StringUtils;
 import org.apache.commons.lang3.builder.EqualsBuilder;
+import org.apache.lucene.search.BooleanClause;
 import org.apache.lucene.search.MatchNoDocsQuery;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.join.ScoreMode;
@@ -77,6 +78,7 @@
 import org.opensearch.index.mapper.RankFeaturesFieldMapper;
 import org.opensearch.index.query.NestedQueryBuilder;
 import org.opensearch.index.query.QueryBuilder;
+import org.opensearch.index.query.QueryBuilderVisitor;
 import org.opensearch.index.query.QueryCoordinatorContext;
 import org.opensearch.index.query.WithFieldName;
 import org.opensearch.index.query.QueryRewriteContext;
@@ -525,6 +527,19 @@ public QueryBuilder filter(QueryBuilder filterToBeAdded) {
 
     }
 
+    /**
+     * Exposes the embedded query filter as a {@link BooleanClause.Occur#FILTER} child. The default
+     * {@link QueryBuilder#visit(QueryBuilderVisitor)} implementation visits only this builder, so composite query
+     * builders must override it to make their complete query tree available to visitors.
+     */
+    @Override
+    public void visit(QueryBuilderVisitor visitor) {
+        visitor.accept(this);
+        if (queryfilter != null) {
+            queryfilter.visit(visitor.getChildVisitor(BooleanClause.Occur.FILTER));
+        }
+    }
+
     @Override
     protected void doXContent(XContentBuilder xContentBuilder, Params params) throws IOException {
         xContentBuilder.startObject(NAME);

src/test/java/org/opensearch/neuralsearch/query/NeuralQueryBuilderTests.java

@@ -6,6 +6,7 @@
 
 import lombok.SneakyThrows;
 import org.apache.commons.lang3.tuple.Pair;
+import org.apache.lucene.search.BooleanClause;
 import org.apache.lucene.search.MatchNoDocsQuery;
 import org.apache.lucene.search.Query;
 import org.junit.Before;
@@ -28,6 +29,7 @@
 import org.opensearch.index.query.MatchAllQueryBuilder;
 import org.opensearch.index.query.MatchNoneQueryBuilder;
 import org.opensearch.index.query.QueryBuilder;
+import org.opensearch.index.query.QueryBuilderVisitor;
 import org.opensearch.index.query.QueryShardContext;
 import org.opensearch.index.query.TermQueryBuilder;
 import org.opensearch.knn.index.query.KNNQueryBuilder;
@@ -39,6 +41,7 @@
 import org.opensearch.test.OpenSearchTestCase;
 
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 import java.util.function.Supplier;
@@ -826,6 +829,57 @@ public void testFilter_whenAddBoolQueryBuilderToNeuralQueryBuilder_thenFilterSuc
         assertEquals(TEST_FILTER, neuralQueryBuilder.queryfilter());
     }
 
+    public void testVisit_whenFilterPresent_thenVisitsFilterClause() {
+        NeuralQueryBuilder neuralQueryBuilder = getBaselineNeuralQueryBuilder();
+        List<QueryBuilder> visitedQueries = new ArrayList<>();
+        List<BooleanClause.Occur> visitedOccurrences = new ArrayList<>();
+        QueryBuilderVisitor visitor = new QueryBuilderVisitor() {
+            @Override
+            public void accept(QueryBuilder queryBuilder) {
+                visitedQueries.add(queryBuilder);
+            }
+
+            @Override
+            public QueryBuilderVisitor getChildVisitor(BooleanClause.Occur occur) {
+                visitedOccurrences.add(occur);
+                return this;
+            }
+        };
+
+        neuralQueryBuilder.visit(visitor);
+
+        assertEquals(List.of(neuralQueryBuilder, TEST_FILTER), visitedQueries);
+        assertEquals(List.of(BooleanClause.Occur.FILTER), visitedOccurrences);
+    }
+
+    public void testVisit_whenFilterAbsent_thenVisitsOnlyNeuralQuery() {
+        NeuralQueryBuilder neuralQueryBuilder = NeuralQueryBuilder.builder()
+            .fieldName(FIELD_NAME)
+            .queryText(QUERY_TEXT)
+            .modelId(MODEL_ID)
+            .k(K)
+            .build();
+        List<QueryBuilder> visitedQueries = new ArrayList<>();
+        List<BooleanClause.Occur> visitedOccurrences = new ArrayList<>();
+        QueryBuilderVisitor visitor = new QueryBuilderVisitor() {
+            @Override
+            public void accept(QueryBuilder queryBuilder) {
+                visitedQueries.add(queryBuilder);
+            }
+
+            @Override
+            public QueryBuilderVisitor getChildVisitor(BooleanClause.Occur occur) {
+                visitedOccurrences.add(occur);
+                return this;
+            }
+        };
+
+        neuralQueryBuilder.visit(visitor);
+
+        assertEquals(List.of(neuralQueryBuilder), visitedQueries);
+        assertTrue(visitedOccurrences.isEmpty());
+    }
+
     public void testQueryCreation_whenCreateQueryWithDoToQuery_thenFail() {
         setUpClusterService();
         NeuralQueryBuilder neuralQueryBuilder = NeuralQueryBuilder.builder()