Fix PIT context leaks when a query without cursor/pagination fails during execution

Signed-off-by: Lantao Jin <ltjin@amazon.com>

Author: LantaoJin

opensearch/src/main/java/org/opensearch/sql/opensearch/data/utils/OpenSearchJsonContent.java

@@ -156,7 +156,7 @@ public Pair<Double, Double> geoValue() {
       GeoUtils.parseGeoPoint(parser, point, true);
       return Pair.of(point.getLat(), point.getLon());
     } catch (IOException ex) {
-      throw new OpenSearchParseException("error parsing geo point", ex);
+      throw new OpenSearchParseException(String.format("error parsing geo point %s", value), ex);
     }
   }
 
@@ -175,7 +175,7 @@ private long parseLongValue(JsonNode node) {
       }
       return Numbers.toLong(node.textValue(), true);
     } else {
-      throw new OpenSearchParseException("node must be a number");
+      throw new OpenSearchParseException(String.format("node %s must be a number", node));
     }
   }
 
@@ -189,7 +189,7 @@ private double parseDoubleValue(JsonNode node) {
       }
       return Double.parseDouble(node.textValue());
     } else {
-      throw new OpenSearchParseException("node must be a number");
+      throw new OpenSearchParseException(String.format("node %s must be a number", node));
     }
   }
 
@@ -200,7 +200,7 @@ private boolean parseBooleanValue(JsonNode node) {
     } else if (node.isTextual()) {
       return Boolean.parseBoolean(node.textValue());
     } else {
-      throw new OpenSearchParseException("node must be a boolean");
+      throw new OpenSearchParseException(String.format("node %s must be a boolean", node));
     }
   }
 }

opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java

@@ -44,6 +44,12 @@ public class OpenSearchIndexScan extends TableScanOperator implements Serializab
   /** Number of rows returned. */
   private Integer queryCount;
 
+  /**
+   * Whether the cursor (including PIT) has been serialized for a subsequent page request. When
+   * true, {@link #close()} must preserve the PIT because a future request will resume from it.
+   */
+  private boolean cursorSerialized = false;
+
   /** Search response for current batch. */
   private Iterator<ExprValue> iterator;
 
@@ -108,12 +114,12 @@ private void fetchNextBatch() {
   public void close() {
     super.close();
 
-    if (request.hasAnotherBatch()) {
-      // PIT is still needed for the next page — only clean up the in-memory state
-      // without deleting the PIT from OpenSearch.
+    if (request.hasAnotherBatch() && cursorSerialized) {
+      // PIT has been serialized into a cursor for the next page request.
+      // Only clean up in-memory state; the PIT must survive for the next request.
       client.cleanup(request);
     } else {
-      // No more pages expected (last page consumed, non-paginated query, or error path).
+      // No more pages, or query failed/aborted before cursor was serialized.
       // Force delete the PIT to prevent leaking.
       client.forceCleanup(request);
     }
@@ -186,5 +192,8 @@ public void writeExternal(ObjectOutput out) throws IOException {
     out.write(reqAsBytes, 0, reqOut.size());
 
     out.writeInt(maxResponseSize);
+
+    // Mark that the PIT has been serialized into a cursor, so close() preserves it.
+    cursorSerialized = true;
   }
 }

opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java

@@ -281,6 +281,49 @@ void query_some_results_with_scroll() {
     verify(client).forceCleanup(any());
   }
 
+  /**
+   * When close() is called mid-pagination without cursor serialization (e.g., query failed or
+   * aborted), the PIT should be force-deleted to prevent leaking.
+   */
+  @Test
+  void close_mid_pagination_without_cursor_serialized_should_force_cleanup() {
+    var request = mock(OpenSearchRequest.class);
+    when(request.hasAnotherBatch()).thenReturn(true);
+    var indexScan = new OpenSearchIndexScan(client, request);
+    indexScan.close();
+    verify(client).forceCleanup(request);
+    verify(client, never()).cleanup(any());
+  }
+
+  /**
+   * When close() is called after cursor has been serialized (normal pagination), the PIT should be
+   * preserved via cleanup() (in-memory only, no PIT deletion).
+   */
+  @Test
+  @SneakyThrows
+  void close_mid_pagination_with_cursor_serialized_should_cleanup() {
+    var request = mock(OpenSearchRequest.class);
+    when(request.hasAnotherBatch()).thenReturn(true);
+    var indexScan = new OpenSearchIndexScan(client, request);
+
+    // Simulate successful cursor serialization by calling writeExternal
+    var out = mock(ObjectOutput.class);
+    indexScan.writeExternal(out);
+
+    indexScan.close();
+    verify(client).cleanup(request);
+    verify(client, never()).forceCleanup(any());
+  }
+
+  /** forceClose() should always force-delete the PIT regardless of pagination state. */
+  @Test
+  void forceClose_should_always_force_cleanup() {
+    var request = mock(OpenSearchRequest.class);
+    var indexScan = new OpenSearchIndexScan(client, request);
+    indexScan.forceClose();
+    verify(client).forceCleanup(request);
+  }
+
   static void mockTwoPageResponse(OpenSearchClient client) {
     mockResponse(
         client,