Fix explain action on query rewrite (opensearch-project#17286)

Signed-off-by: Fen Qin <[email protected]>

Author: fen-qin

CHANGELOG-3.0.md

@@ -71,6 +71,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Fix swapped field formats in nodes API where `total_indexing_buffer_in_bytes` and `total_indexing_buffer` values were reversed ([#17070](https://github.com/opensearch-project/OpenSearch/pull/17070))
 - Add HTTP/2 protocol support to HttpRequest.HttpVersion ([#17248](https://github.com/opensearch-project/OpenSearch/pull/17248))
 - Fix missing bucket in terms aggregation with missing value ([#17418](https://github.com/opensearch-project/OpenSearch/pull/17418))
+- Fix explain action on query rewrite ([#17286](https://github.com/opensearch-project/OpenSearch/pull/17286))
 
 ### Security
 

server/src/internalClusterTest/java/org/opensearch/explain/ExplainActionIT.java

@@ -40,6 +40,8 @@
 import org.opensearch.core.common.io.stream.InputStreamStreamInput;
 import org.opensearch.core.common.io.stream.OutputStreamStreamOutput;
 import org.opensearch.index.query.QueryBuilders;
+import org.opensearch.index.query.TermsQueryBuilder;
+import org.opensearch.indices.TermsLookup;
 import org.opensearch.test.OpenSearchIntegTestCase;
 
 import java.io.ByteArrayInputStream;
@@ -52,6 +54,7 @@
 import java.util.Set;
 
 import static java.util.Collections.singleton;
+import static org.opensearch.cluster.metadata.IndexMetadata.SETTING_NUMBER_OF_SHARDS;
 import static org.opensearch.common.xcontent.XContentFactory.jsonBuilder;
 import static org.opensearch.index.query.QueryBuilders.queryStringQuery;
 import static org.opensearch.test.hamcrest.OpenSearchAssertions.assertAcked;
@@ -305,4 +308,25 @@ public void testStreamExplain() throws Exception {
         result = Lucene.readExplanation(esBuffer);
         assertThat(exp.toString(), equalTo(result.toString()));
     }
+
+    public void testQueryRewrite() {
+        client().admin()
+            .indices()
+            .prepareCreate("twitter")
+            .setMapping("user", "type=integer", "followers", "type=integer")
+            .setSettings(Settings.builder().put(SETTING_NUMBER_OF_SHARDS, 2).put("index.number_of_routing_shards", 2))
+            .get();
+        client().prepareIndex("twitter").setId("1").setSource("followers", new int[] { 1, 2, 3 }).get();
+        refresh();
+
+        TermsQueryBuilder termsLookupQuery = QueryBuilders.termsLookupQuery("user", new TermsLookup("twitter", "1", "followers"));
+        ExplainResponse response = client().prepareExplain("twitter", "1").setQuery(termsLookupQuery).get();
+        assertNotNull(response);
+        assertTrue(response.isExists());
+        assertFalse(response.isMatch());
+        assertThat(response.getIndex(), equalTo("twitter"));
+        assertThat(response.getId(), equalTo("1"));
+        assertNotNull(response.getExplanation());
+        assertFalse(response.getExplanation().isMatch());
+    }
 }

server/src/main/java/org/opensearch/action/explain/TransportExplainAction.java

@@ -52,6 +52,8 @@
 import org.opensearch.index.get.GetResult;
 import org.opensearch.index.mapper.IdFieldMapper;
 import org.opensearch.index.mapper.Uid;
+import org.opensearch.index.query.QueryBuilder;
+import org.opensearch.index.query.Rewriteable;
 import org.opensearch.index.shard.IndexShard;
 import org.opensearch.search.SearchService;
 import org.opensearch.search.internal.AliasFilter;
@@ -65,6 +67,7 @@
 
 import java.io.IOException;
 import java.util.Set;
+import java.util.function.LongSupplier;
 
 /**
  * Explain transport action. Computes the explain on the targeted shard.
@@ -101,7 +104,14 @@ public TransportExplainAction(
     @Override
     protected void doExecute(Task task, ExplainRequest request, ActionListener<ExplainResponse> listener) {
         request.nowInMillis = System.currentTimeMillis();
-        super.doExecute(task, request, listener);
+        ActionListener<QueryBuilder> rewriteListener = ActionListener.wrap(rewrittenQuery -> {
+            request.query(rewrittenQuery);
+            super.doExecute(task, request, listener);
+        }, listener::onFailure);
+
+        assert request.query() != null;
+        LongSupplier timeProvider = () -> request.nowInMillis;
+        Rewriteable.rewriteAndFetch(request.query(), searchService.getIndicesService().getRewriteContext(timeProvider), rewriteListener);
     }
 
     @Override