From c6981d31a37da3727b6f02377c5ef4c86d143e93 Mon Sep 17 00:00:00 2001
From: zhichao-aws <zhichaog@amazon.com>
Date: Thu, 25 Jul 2024 14:23:25 +0800
Subject: [PATCH] Fix the visit of inner query for NestedQueryBuilder(#14967)

Signed-off-by: zhichao-aws <zhichaog@amazon.com>
---
 CHANGELOG.md                                          |  1 +
 .../opensearch/index/query/NestedQueryBuilder.java    |  9 +++++++++
 .../index/query/NestedQueryBuilderTests.java          | 11 +++++++++++
 3 files changed, 21 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2f42d27ab34d1..728de32e413d8 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -97,6 +97,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Create new IndexInput for multi part upload ([#14888](https://github.com/opensearch-project/OpenSearch/pull/14888))
 - Fix searchable snapshot failure with scripted fields ([#14411](https://github.com/opensearch-project/OpenSearch/pull/14411))
 - Fix constant_keyword field type not working when creating index ([#14807](https://github.com/opensearch-project/OpenSearch/pull/14807))
+- Fix the visit of inner query for NestedQueryBuilder ([#14739](https://github.com/opensearch-project/OpenSearch/pull/14739))
 
 ### Security
 
diff --git a/server/src/main/java/org/opensearch/index/query/NestedQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/NestedQueryBuilder.java
index b5ba79632b622..5908882472ce7 100644
--- a/server/src/main/java/org/opensearch/index/query/NestedQueryBuilder.java
+++ b/server/src/main/java/org/opensearch/index/query/NestedQueryBuilder.java
@@ -34,6 +34,7 @@
 
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.ReaderUtil;
+import org.apache.lucene.search.BooleanClause;
 import org.apache.lucene.search.MatchNoDocsQuery;
 import org.apache.lucene.search.MultiCollector;
 import org.apache.lucene.search.Query;
@@ -505,4 +506,12 @@ public TopDocsAndMaxScore topDocs(SearchHit hit) throws IOException {
             }
         }
     }
+
+    @Override
+    public void visit(QueryBuilderVisitor visitor) {
+        visitor.accept(this);
+        if (query != null) {
+            visitor.getChildVisitor(BooleanClause.Occur.MUST).accept(query);
+        }
+    }
 }
diff --git a/server/src/test/java/org/opensearch/index/query/NestedQueryBuilderTests.java b/server/src/test/java/org/opensearch/index/query/NestedQueryBuilderTests.java
index 1cb595c87675c..6957d9908069f 100644
--- a/server/src/test/java/org/opensearch/index/query/NestedQueryBuilderTests.java
+++ b/server/src/test/java/org/opensearch/index/query/NestedQueryBuilderTests.java
@@ -59,8 +59,10 @@
 import org.hamcrest.Matchers;
 
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 
@@ -565,4 +567,13 @@ void doWithDepth(int depth, ThrowingConsumer<QueryShardContext> test) throws Exc
                 );
         }
     }
+
+    public void testVisit() {
+        NestedQueryBuilder builder = new NestedQueryBuilder("path", new MatchAllQueryBuilder(), ScoreMode.None);
+
+        List<QueryBuilder> visitedQueries = new ArrayList<>();
+        builder.visit(createTestVisitor(visitedQueries));
+
+        assertEquals(2, visitedQueries.size());
+    }
 }