Refactor integ-tests data for PIT and fix unit tests

Signed-off-by: Manasvini B S <manasvis@amazon.com>
manasvinibs · Sep 13, 2024 · 1d27eac · 1d27eac
1 parent ae80476
commit 1d27eac
Show file tree

Hide file tree

Showing 30 changed files with 293 additions and 367 deletions.
diff --git a/integ-test/src/test/java/org/opensearch/sql/legacy/CursorIT.java b/integ-test/src/test/java/org/opensearch/sql/legacy/CursorIT.java
@@ -182,7 +182,7 @@ public void validTotalResultWithAndWithoutPaginationOrderBy() throws IOException
     String selectQuery =
         StringUtils.format(
             "SELECT firstname, state FROM %s ORDER BY balance DESC ", TEST_INDEX_ACCOUNT);
-    verifyWithAndWithoutPaginationResponse(selectQuery + " LIMIT 2000", selectQuery, 26, false);
+    verifyWithAndWithoutPaginationResponse(selectQuery + " LIMIT 2000", selectQuery, 25, false);
   }
 
   @Test

diff --git a/integ-test/src/test/java/org/opensearch/sql/ppl/StandaloneIT.java b/integ-test/src/test/java/org/opensearch/sql/ppl/StandaloneIT.java
@@ -149,8 +149,11 @@ public void onFailure(Exception e) {
 
   private Settings defaultSettings() {
     return new Settings() {
-      private final Map<Key, Integer> defaultSettings =
-          new ImmutableMap.Builder<Key, Integer>().put(Key.QUERY_SIZE_LIMIT, 200).build();
+      private final Map<Key, Object> defaultSettings =
+          new ImmutableMap.Builder<Key, Object>()
+                  .put(Key.QUERY_SIZE_LIMIT, 200)
+                  .put(Key.SQL_PAGINATION_API_SEARCH_AFTER, true)
+                  .build();
 
       @Override
       public <T> T getSettingValue(Key key) {

diff --git a/integ-test/src/test/java/org/opensearch/sql/sql/PaginationFilterIT.java b/integ-test/src/test/java/org/opensearch/sql/sql/PaginationFilterIT.java
@@ -36,8 +36,8 @@ public class PaginationFilterIT extends SQLIntegTestCase {
       Map.of(
           "SELECT * FROM " + TestsConstants.TEST_INDEX_ACCOUNT, 1000,
           "SELECT * FROM " + TestsConstants.TEST_INDEX_ACCOUNT + " WHERE match(address, 'street')",
-              385,
-          "SELECT * FROM "
+             385,
+         "SELECT * FROM "
                   + TestsConstants.TEST_INDEX_ACCOUNT
                   + " WHERE match(address, 'street') AND match(city, 'Ola')",
               1,
@@ -134,6 +134,9 @@ public void test_pagination_with_where() {
 
         paginatedResponse = executeCursorQuery(cursor);
 
+        System.out.println("**** MANA HERE: paginatedResponse -  " + paginatedResponse);
+        System.out.println("**** MANA HERE: paginatedResponse -size -  " + paginatedResponse.getInt("size"));
+
         this.logger.info(
             testReportPrefix
                 + String.format(
@@ -142,6 +145,8 @@ public void test_pagination_with_where() {
         break;
       }
     } while (true);
+    System.out.println("**** MANA HERE: last paginatedResponse -  " + paginatedResponse);
+
     // last page expected results:
     assertEquals(
         testReportPrefix + "Last page", totalHits % pageSize, paginatedResponse.getInt("size"));

diff --git a/integ-test/src/test/java/org/opensearch/sql/sql/StandalonePaginationIT.java b/integ-test/src/test/java/org/opensearch/sql/sql/StandalonePaginationIT.java
@@ -23,6 +23,7 @@
 import org.opensearch.client.Request;
 import org.opensearch.client.ResponseException;
 import org.opensearch.client.RestHighLevelClient;
+import org.opensearch.client.node.NodeClient;
 import org.opensearch.common.inject.Injector;
 import org.opensearch.common.inject.ModulesBuilder;
 import org.opensearch.common.unit.TimeValue;
@@ -166,6 +167,7 @@ private Settings defaultSettings() {
           new ImmutableMap.Builder<Key, Object>()
               .put(Key.QUERY_SIZE_LIMIT, 200)
               .put(Key.SQL_CURSOR_KEEP_ALIVE, TimeValue.timeValueMinutes(1))
+                  .put(Key.SQL_PAGINATION_API_SEARCH_AFTER, true)
               .build();
 
       @Override

diff --git a/integ-test/src/test/resources/expectedOutput/ppl/explain_filter_agg_push.json b/integ-test/src/test/resources/expectedOutput/ppl/explain_filter_agg_push.json
@@ -8,7 +8,7 @@
       {
         "name": "OpenSearchIndexScan",
         "description": {
-          "request": "OpenSearchQueryRequest(indexName\u003dopensearch-sql_test_index_account, sourceBuilder\u003d{\"from\":0,\"size\":10000,\"timeout\":\"1m\",\"query\":{\"range\":{\"age\":{\"from\":30,\"to\":null,\"include_lower\":false,\"include_upper\":true,\"boost\":1.0}}},\"sort\":[{\"_doc\":{\"order\":\"asc\"}}],\"aggregations\":{\"composite_buckets\":{\"composite\":{\"size\":1000,\"sources\":[{\"state\":{\"terms\":{\"field\":\"state.keyword\",\"missing_bucket\":true,\"missing_order\":\"first\",\"order\":\"asc\"}}},{\"city\":{\"terms\":{\"field\":\"city.keyword\",\"missing_bucket\":true,\"missing_order\":\"first\",\"order\":\"asc\"}}}]},\"aggregations\":{\"avg_age\":{\"avg\":{\"field\":\"age\"}}}}}}, searchDone\u003dfalse)"
+          "request": "OpenSearchQueryRequest(indexName\u003dopensearch-sql_test_index_account, sourceBuilder\u003d{\"from\":0,\"size\":10000,\"timeout\":\"1m\",\"query\":{\"range\":{\"age\":{\"from\":30,\"to\":null,\"include_lower\":false,\"include_upper\":true,\"boost\":1.0}}},\"sort\":[{\"_doc\":{\"order\":\"asc\"}}],\"aggregations\":{\"composite_buckets\":{\"composite\":{\"size\":1000,\"sources\":[{\"state\":{\"terms\":{\"field\":\"state.keyword\",\"missing_bucket\":true,\"missing_order\":\"first\",\"order\":\"asc\"}}},{\"city\":{\"terms\":{\"field\":\"city.keyword\",\"missing_bucket\":true,\"missing_order\":\"first\",\"order\":\"asc\"}}}]},\"aggregations\":{\"avg_age\":{\"avg\":{\"field\":\"age\"}}}}}}, needClean\u003dtrue, searchDone\u003dfalse, pitId\u003dnull, client\u003dnull, cursorKeepAlive\u003dnull, searchAfter\u003dnull, searchResponse\u003dnull)"
         },
         "children": []
       }

diff --git a/integ-test/src/test/resources/expectedOutput/ppl/explain_filter_push.json b/integ-test/src/test/resources/expectedOutput/ppl/explain_filter_push.json
@@ -8,7 +8,7 @@
       {
         "name": "OpenSearchIndexScan",
         "description": {
-          "request": "OpenSearchQueryRequest(indexName\u003dopensearch-sql_test_index_account, sourceBuilder\u003d{\"from\":0,\"size\":10000,\"timeout\":\"1m\",\"query\":{\"bool\":{\"filter\":[{\"bool\":{\"filter\":[{\"range\":{\"balance\":{\"from\":10000,\"to\":null,\"include_lower\":false,\"include_upper\":true,\"boost\":1.0}}},{\"range\":{\"age\":{\"from\":null,\"to\":40,\"include_lower\":true,\"include_upper\":false,\"boost\":1.0}}}],\"adjust_pure_negative\":true,\"boost\":1.0}},{\"range\":{\"age\":{\"from\":30,\"to\":null,\"include_lower\":false,\"include_upper\":true,\"boost\":1.0}}}],\"adjust_pure_negative\":true,\"boost\":1.0}},\"_source\":{\"includes\":[\"age\"],\"excludes\":[]},\"sort\":[{\"_doc\":{\"order\":\"asc\"}}]}, searchDone\u003dfalse)"
+          "request": "OpenSearchQueryRequest(indexName\u003dopensearch-sql_test_index_account, sourceBuilder\u003d{\"from\":0,\"size\":10000,\"timeout\":\"1m\",\"query\":{\"bool\":{\"filter\":[{\"bool\":{\"filter\":[{\"range\":{\"balance\":{\"from\":10000,\"to\":null,\"include_lower\":false,\"include_upper\":true,\"boost\":1.0}}},{\"range\":{\"age\":{\"from\":null,\"to\":40,\"include_lower\":true,\"include_upper\":false,\"boost\":1.0}}}],\"adjust_pure_negative\":true,\"boost\":1.0}},{\"range\":{\"age\":{\"from\":30,\"to\":null,\"include_lower\":false,\"include_upper\":true,\"boost\":1.0}}}],\"adjust_pure_negative\":true,\"boost\":1.0}},\"_source\":{\"includes\":[\"age\"],\"excludes\":[]},\"sort\":[{\"_doc\":{\"order\":\"asc\"}}]}, needClean\u003dtrue, searchDone\u003dfalse, pitId\u003dnull, client\u003dnull, cursorKeepAlive\u003dnull, searchAfter\u003dnull, searchResponse\u003dnull)"
         },
         "children": []
       }

diff --git a/integ-test/src/test/resources/expectedOutput/ppl/explain_limit_push.json b/integ-test/src/test/resources/expectedOutput/ppl/explain_limit_push.json
@@ -16,7 +16,7 @@
           {
             "name": "OpenSearchIndexScan",
             "description": {
-              "request": "OpenSearchQueryRequest(indexName=opensearch-sql_test_index_account, sourceBuilder={\"from\":0,\"size\":5,\"timeout\":\"1m\"}, searchDone=false)"
+              "request": "OpenSearchQueryRequest(indexName=opensearch-sql_test_index_account, sourceBuilder={\"from\":0,\"size\":5,\"timeout\":\"1m\"}, needClean\u003dtrue, searchDone\u003dfalse, pitId\u003dnull, client\u003dnull, cursorKeepAlive\u003dnull, searchAfter\u003dnull, searchResponse\u003dnull)"
             },
             "children": []
           }

diff --git a/integ-test/src/test/resources/expectedOutput/ppl/explain_output.json b/integ-test/src/test/resources/expectedOutput/ppl/explain_output.json
@@ -31,7 +31,7 @@
                   {
                     "name": "OpenSearchIndexScan",
                     "description": {
-                      "request": "OpenSearchQueryRequest(indexName\u003dopensearch-sql_test_index_account, sourceBuilder\u003d{\"from\":0,\"size\":10000,\"timeout\":\"1m\",\"query\":{\"range\":{\"age\":{\"from\":30,\"to\":null,\"include_lower\":false,\"include_upper\":true,\"boost\":1.0}}},\"sort\":[{\"_doc\":{\"order\":\"asc\"}}],\"aggregations\":{\"composite_buckets\":{\"composite\":{\"size\":1000,\"sources\":[{\"state\":{\"terms\":{\"field\":\"state.keyword\",\"missing_bucket\":true,\"missing_order\":\"first\",\"order\":\"asc\"}}},{\"city\":{\"terms\":{\"field\":\"city.keyword\",\"missing_bucket\":true,\"missing_order\":\"first\",\"order\":\"asc\"}}}]},\"aggregations\":{\"avg_age\":{\"avg\":{\"field\":\"age\"}}}}}}, searchDone\u003dfalse)"
+                      "request": "OpenSearchQueryRequest(indexName\u003dopensearch-sql_test_index_account, sourceBuilder\u003d{\"from\":0,\"size\":10000,\"timeout\":\"1m\",\"query\":{\"range\":{\"age\":{\"from\":30,\"to\":null,\"include_lower\":false,\"include_upper\":true,\"boost\":1.0}}},\"sort\":[{\"_doc\":{\"order\":\"asc\"}}],\"aggregations\":{\"composite_buckets\":{\"composite\":{\"size\":1000,\"sources\":[{\"state\":{\"terms\":{\"field\":\"state.keyword\",\"missing_bucket\":true,\"missing_order\":\"first\",\"order\":\"asc\"}}},{\"city\":{\"terms\":{\"field\":\"city.keyword\",\"missing_bucket\":true,\"missing_order\":\"first\",\"order\":\"asc\"}}}]},\"aggregations\":{\"avg_age\":{\"avg\":{\"field\":\"age\"}}}}}}, needClean\u003dtrue, searchDone\u003dfalse, pitId\u003dnull, client\u003dnull, cursorKeepAlive\u003dnull, searchAfter\u003dnull, searchResponse\u003dnull)"
                     },
                     "children": []
                   }

diff --git a/integ-test/src/test/resources/expectedOutput/ppl/explain_sort_push.json b/integ-test/src/test/resources/expectedOutput/ppl/explain_sort_push.json
@@ -8,7 +8,7 @@
       {
         "name": "OpenSearchIndexScan",
         "description": {
-          "request": "OpenSearchQueryRequest(indexName\u003dopensearch-sql_test_index_account, sourceBuilder\u003d{\"from\":0,\"size\":10000,\"timeout\":\"1m\",\"query\":{\"range\":{\"age\":{\"from\":30,\"to\":null,\"include_lower\":false,\"include_upper\":true,\"boost\":1.0}}},\"_source\":{\"includes\":[\"age\"],\"excludes\":[]},\"sort\":[{\"age\":{\"order\":\"asc\",\"missing\":\"_first\"}}]}, searchDone\u003dfalse)"
+          "request": "OpenSearchQueryRequest(indexName\u003dopensearch-sql_test_index_account, sourceBuilder\u003d{\"from\":0,\"size\":10000,\"timeout\":\"1m\",\"query\":{\"range\":{\"age\":{\"from\":30,\"to\":null,\"include_lower\":false,\"include_upper\":true,\"boost\":1.0}}},\"_source\":{\"includes\":[\"age\"],\"excludes\":[]},\"sort\":[{\"age\":{\"order\":\"asc\",\"missing\":\"_first\"}}]}, needClean\u003dtrue, searchDone\u003dfalse, pitId\u003dnull, client\u003dnull, cursorKeepAlive\u003dnull, searchAfter\u003dnull, searchResponse\u003dnull)"
         },
         "children": []
       }

diff --git a/legacy/src/main/java/org/opensearch/sql/legacy/executor/format/SelectResultSet.java b/legacy/src/main/java/org/opensearch/sql/legacy/executor/format/SelectResultSet.java
@@ -567,7 +567,7 @@ private void populateDefaultCursor(DefaultCursor cursor) {
     Integer limit = cursor.getLimit();
     long rowsLeft = rowsLeft(cursor.getFetchSize(), cursor.getLimit());
     if (rowsLeft <= 0) {
-      // close the cursor
+      // Delete Point In Time ID
       if (LocalClusterState.state().getSettingValue(Settings.Key.SQL_PAGINATION_API_SEARCH_AFTER)) {
         String pitId = cursor.getPitId();
         PointInTimeHandler pit = new PointInTimeHandlerImpl(client, pitId);

diff --git a/...er/physical/node/scroll/SearchHitRow.java → ...y/planner/physical/node/SearchHitRow.java b/...er/physical/node/scroll/SearchHitRow.java → ...y/planner/physical/node/SearchHitRow.java
@@ -3,16 +3,17 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-package org.opensearch.sql.legacy.query.planner.physical.node.scroll;
+package org.opensearch.sql.legacy.query.planner.physical.node;
 
 import com.google.common.base.Strings;
-import java.util.HashMap;
-import java.util.Map;
 import org.opensearch.common.document.DocumentField;
 import org.opensearch.index.mapper.MapperService;
 import org.opensearch.search.SearchHit;
 import org.opensearch.sql.legacy.query.planner.physical.Row;
 
+import java.util.HashMap;
+import java.util.Map;
+
 /**
  *
  *
@@ -36,7 +37,7 @@
  * ----------------------------------------------------------------------------------------------------------------------
  * </pre>
  */
-class SearchHitRow implements Row<SearchHit> {
+public class SearchHitRow implements Row<SearchHit> {
 
   /** Native OpenSearch data object for each row */
   private final SearchHit hit;
@@ -47,7 +48,7 @@ class SearchHitRow implements Row<SearchHit> {
   /** Table alias owned the row. Empty if this row comes from combination of two other rows */
   private final String tableAlias;
 
-  SearchHitRow(SearchHit hit, String tableAlias) {
+  public SearchHitRow(SearchHit hit, String tableAlias) {
     this.hit = hit;
     this.source = hit.getSourceAsMap();
     this.tableAlias = tableAlias;

diff --git a/...n/java/org/opensearch/sql/legacy/query/planner/physical/node/pointInTime/PointInTime.java b/...n/java/org/opensearch/sql/legacy/query/planner/physical/node/pointInTime/PointInTime.java
@@ -25,6 +25,7 @@
 import org.opensearch.sql.legacy.query.planner.physical.estimation.Cost;
 import org.opensearch.sql.legacy.query.planner.physical.node.BatchPhysicalOperator;
 import org.opensearch.sql.legacy.query.planner.resource.ResourceManager;
+import org.opensearch.sql.legacy.query.planner.physical.node.SearchHitRow;
 
 /** OpenSearch Search API with Point in time as physical implementation of TableScan */
 public class PointInTime extends BatchPhysicalOperator<SearchHit> {

diff --git a/.../java/org/opensearch/sql/legacy/query/planner/physical/node/pointInTime/SearchHitRow.java b/.../java/org/opensearch/sql/legacy/query/planner/physical/node/pointInTime/SearchHitRow.java