Store TotalHits and use it to report total in response.

Signed-off-by: Yury-Fridlyand <yury.fridlyand@improving.com>
opensearch-project · Mar 11, 2023 · f4ea4ad · f4ea4ad
1 parent 1b5ab7e
commit f4ea4ad
Show file tree

Hide file tree

Showing 20 changed files with 158 additions and 62 deletions.
diff --git a/core/src/main/java/org/opensearch/sql/executor/ExecutionEngine.java b/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;
   }
 

diff --git a/core/src/main/java/org/opensearch/sql/planner/physical/PhysicalPlan.java b/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()));

diff --git a/core/src/test/java/org/opensearch/sql/executor/QueryServiceTest.java b/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;
           })

diff --git a/...src/test/java/org/opensearch/sql/executor/streaming/MicroBatchStreamingExecutionTest.java b/...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);

diff --git a/core/src/test/java/org/opensearch/sql/planner/physical/PhysicalPlanTest.java b/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());
+  }
 }
diff --git a/core/src/testFixtures/java/org/opensearch/sql/executor/DefaultExecutionEngine.java b/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);

diff --git a/legacy/src/main/java/org/opensearch/sql/legacy/plugin/RestSQLQueryAction.java b/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

diff --git a/...earch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngine.java b/...earch/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);

diff --git a/...h/src/main/java/org/opensearch/sql/opensearch/executor/protector/ResourceMonitorPlan.java b/...h/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();

diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/response/OpenSearchResponse.java b/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() {

diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java
@@ -84,6 +84,12 @@ public ExprValue next() {
     return iterator.next();
   }
 
+  @Override
+  public long getTotalHits() {
+    // TODO maybe store totalHits from `response`
+    return queryCount;
+  }
+
   private void fetchNextBatch() {
     OpenSearchResponse response = client.search(request);
     if (!response.isEmpty()) {

diff --git a/...ch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScan.java b/...ch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScan.java
@@ -27,6 +27,7 @@ public class OpenSearchPagedIndexScan extends TableScanOperator {
   private OpenSearchRequest request;
   private Iterator<ExprValue> iterator;
   private boolean needClean = false;
+  private long totalHits = 0;
 
   public OpenSearchPagedIndexScan(OpenSearchClient client,
                                   PagedRequestBuilder requestBuilder) {
@@ -56,6 +57,7 @@ public void open() {
     OpenSearchResponse response = client.search(request);
     if (!response.isEmpty()) {
       iterator = response.iterator();
+      totalHits = response.getTotalHits();
     } else {
       needClean = true;
       iterator = Collections.emptyIterator();
@@ -71,6 +73,11 @@ public void close() {
     }
   }
 
+  @Override
+  public long getTotalHits() {
+    return totalHits;
+  }
+
   @Override
   public String toCursor() {
     // TODO this assumes exactly one index is scanned.

diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ResourceMonitorPlanTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ResourceMonitorPlanTest.java
@@ -107,4 +107,10 @@ void acceptSuccess() {
     monitorPlan.accept(visitor, context);
     verify(plan, times(1)).accept(visitor, context);
   }
+
+  @Test
+  void getTotalHitsSuccess() {
+    monitorPlan.getTotalHits();
+    verify(plan, times(1)).getTotalHits();
+  }
 }
diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/response/OpenSearchResponseTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/response/OpenSearchResponseTest.java
@@ -74,20 +74,29 @@ void isEmpty() {
                 new TotalHits(2L, TotalHits.Relation.EQUAL_TO),
                 1.0F));
 
-    assertFalse(new OpenSearchResponse(searchResponse, factory).isEmpty());
+    var response = new OpenSearchResponse(searchResponse, factory);
+    assertFalse(response.isEmpty());
+    assertEquals(2L, response.getTotalHits());
 
     when(searchResponse.getHits()).thenReturn(SearchHits.empty());
     when(searchResponse.getAggregations()).thenReturn(null);
-    assertTrue(new OpenSearchResponse(searchResponse, factory).isEmpty());
+
+    response = new OpenSearchResponse(searchResponse, factory);
+    assertTrue(response.isEmpty());
+    assertEquals(0L, response.getTotalHits());
 
     when(searchResponse.getHits())
         .thenReturn(new SearchHits(null, new TotalHits(0, TotalHits.Relation.EQUAL_TO), 0));
-    OpenSearchResponse response3 = new OpenSearchResponse(searchResponse, factory);
-    assertTrue(response3.isEmpty());
+    response = new OpenSearchResponse(searchResponse, factory);
+    assertTrue(response.isEmpty());
+    assertEquals(0L, response.getTotalHits());
 
     when(searchResponse.getHits()).thenReturn(SearchHits.empty());
     when(searchResponse.getAggregations()).thenReturn(new Aggregations(emptyList()));
-    assertFalse(new OpenSearchResponse(searchResponse, factory).isEmpty());
+
+    response = new OpenSearchResponse(searchResponse, factory);
+    assertFalse(response.isEmpty());
+    assertEquals(0L, response.getTotalHits());
   }
 
   @Test
@@ -104,7 +113,8 @@ void iterator() {
     when(factory.construct(any())).thenReturn(exprTupleValue1).thenReturn(exprTupleValue2);
 
     int i = 0;
-    for (ExprValue hit : new OpenSearchResponse(searchResponse, factory)) {
+    var response = new OpenSearchResponse(searchResponse, factory);
+    for (ExprValue hit : response) {
       if (i == 0) {
         assertEquals(exprTupleValue1, hit);
       } else if (i == 1) {
@@ -114,6 +124,7 @@ void iterator() {
       }
       i++;
     }
+    assertEquals(2L, response.getTotalHits());
   }
 
   @Test