Store TotalHits and use it to report total in response.

Signed-off-by: Yury-Fridlyand <[email protected]>

Author: Yury-Fridlyand

core/src/main/java/org/opensearch/sql/executor/ExecutionEngine.java

@@ -54,7 +54,7 @@ void execute(PhysicalPlan plan, ExecutionContext context,
   class QueryResponse {
     private final Schema schema;
     private final List<ExprValue> results;
-    
+    private final long total;
     private final Cursor cursor;
   }
 

core/src/main/java/org/opensearch/sql/planner/physical/PhysicalPlan.java

@@ -50,6 +50,10 @@ public ExecutionEngine.Schema schema() {
         + "ProjectOperator, instead of %s", this.getClass().getSimpleName()));
   }
 
+  public long getTotalHits() {
+    return getChild().stream().mapToLong(PhysicalPlan::getTotalHits).max().orElse(0);
+  }
+
   public String toCursor() {
     throw new IllegalStateException(String.format("%s is not compatible with cursor feature",
         this.getClass().getSimpleName()));

core/src/test/java/org/opensearch/sql/executor/QueryServiceTest.java

@@ -133,7 +133,7 @@ Helper executeSuccess(Split split) {
           invocation -> {
             ResponseListener<ExecutionEngine.QueryResponse> listener = invocation.getArgument(2);
             listener.onResponse(
-                new ExecutionEngine.QueryResponse(schema, Collections.emptyList(),
+                new ExecutionEngine.QueryResponse(schema, Collections.emptyList(), 0,
                     Cursor.None));
             return null;
           })

core/src/test/java/org/opensearch/sql/executor/streaming/MicroBatchStreamingExecutionTest.java

@@ -170,7 +170,7 @@ Helper executeSuccess(Long... offsets) {
                 ResponseListener<ExecutionEngine.QueryResponse> listener =
                     invocation.getArgument(2);
                 listener.onResponse(
-                    new ExecutionEngine.QueryResponse(null, Collections.emptyList(),
+                    new ExecutionEngine.QueryResponse(null, Collections.emptyList(), 0,
                         Cursor.None));
 
                 PlanContext planContext = invocation.getArgument(1);

core/src/test/java/org/opensearch/sql/planner/physical/PhysicalPlanTest.java

@@ -5,9 +5,16 @@
 
 package org.opensearch.sql.planner.physical;
 
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.mockito.Mockito.CALLS_REAL_METHODS;
+import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 import java.util.List;
+
+import org.junit.jupiter.api.DisplayNameGeneration;
+import org.junit.jupiter.api.DisplayNameGenerator;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.mockito.Mock;
@@ -16,6 +23,7 @@
 import org.opensearch.sql.storage.split.Split;
 
 @ExtendWith(MockitoExtension.class)
+@DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class)
 class PhysicalPlanTest {
   @Mock
   Split split;
@@ -46,8 +54,25 @@ public List<PhysicalPlan> getChild() {
   };
 
   @Test
-  void addSplitToChildByDefault() {
+  void add_split_to_child_by_default() {
     testPlan.add(split);
     verify(child).add(split);
   }
+
+  @Test
+  void get_total_hits_from_child() {
+    var plan = mock(PhysicalPlan.class);
+    when(child.getTotalHits()).thenReturn(42L);
+    when(plan.getChild()).thenReturn(List.of(child));
+    when(plan.getTotalHits()).then(CALLS_REAL_METHODS);
+    assertEquals(42, plan.getTotalHits());
+    verify(child).getTotalHits();
+  }
+
+  @Test
+  void get_total_hits_uses_default_value() {
+    var plan = mock(PhysicalPlan.class);
+    when(plan.getTotalHits()).then(CALLS_REAL_METHODS);
+    assertEquals(0, plan.getTotalHits());
+  }
 }

core/src/testFixtures/java/org/opensearch/sql/executor/DefaultExecutionEngine.java

@@ -34,7 +34,7 @@ public void execute(
         result.add(plan.next());
       }
       QueryResponse response = new QueryResponse(new Schema(new ArrayList<>()), new ArrayList<>(),
-          Cursor.None);
+          0, Cursor.None);
       listener.onResponse(response);
     } catch (Exception e) {
       listener.onFailure(e);

legacy/src/main/java/org/opensearch/sql/legacy/plugin/RestSQLQueryAction.java

@@ -172,7 +172,8 @@ private ResponseListener<QueryResponse> createQueryResponseListener(
       @Override
       public void onResponse(QueryResponse response) {
         sendResponse(channel, OK,
-            formatter.format(new QueryResult(response.getSchema(), response.getResults(), response.getCursor())));
+            formatter.format(new QueryResult(response.getSchema(), response.getResults(),
+                response.getCursor(), response.getTotal())));
       }
 
       @Override

opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngine.java

@@ -53,7 +53,7 @@ public void execute(PhysicalPlan physicalPlan, ExecutionContext context,
 
             Cursor qc = paginatedPlanCache.convertToCursor(plan);
 
-            QueryResponse response = new QueryResponse(physicalPlan.schema(), result, qc);
+            QueryResponse response = new QueryResponse(physicalPlan.schema(), result, plan.getTotalHits(), qc);
             listener.onResponse(response);
           } catch (Exception e) {
             listener.onFailure(e);

opensearch/src/main/java/org/opensearch/sql/opensearch/executor/protector/ResourceMonitorPlan.java

@@ -83,6 +83,11 @@ public ExprValue next() {
     return delegate.next();
   }
 
+  @Override
+  public long getTotalHits() {
+    return delegate.getTotalHits();
+  }
+
   @Override
   public String toCursor() {
     return delegate.toCursor();

opensearch/src/main/java/org/opensearch/sql/opensearch/response/OpenSearchResponse.java

@@ -71,7 +71,10 @@ public OpenSearchResponse(SearchHits hits, OpenSearchExprValueFactory exprValueF
    */
   public boolean isEmpty() {
     return (hits.getHits() == null) || (hits.getHits().length == 0) && aggregations == null;
-    // TODO             TBD       ^    ^
+  }
+
+  public long getTotalHits() {
+    return hits.getTotalHits().value;
   }
 
   public boolean isAggregationResponse() {