Fix stream

Signed-off-by: Harsha Vamsi Kalluri <[email protected]>
harshavamsi · Jan 30, 2025 · 85b334e · 85b334e
1 parent b484146
commit 85b334e
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 5 deletions.
diff --git a/server/src/main/java/org/opensearch/action/search/SearchPhaseController.java b/server/src/main/java/org/opensearch/action/search/SearchPhaseController.java
@@ -804,8 +804,6 @@ public ReducedQueryPhase reducedFromStream(
             } catch (InterruptedException | ExecutionException e) {
                 Thread.currentThread().interrupt();
                 throw new RuntimeException("Error processing tickets in parallel", e);
-            } finally {
-                streamManager.close();
             }
 
             TotalHits totalHits = new TotalHits(totalRows, Relation.EQUAL_TO);

diff --git a/server/src/main/java/org/opensearch/search/aggregations/support/StreamingAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/support/StreamingAggregator.java
@@ -125,6 +125,7 @@ public void finish() throws IOException {
                 if (currentRow[0] > 0) {
                     flushBatch();
                 }
+                root.close();
             }
         };
     }

diff --git a/server/src/main/java/org/opensearch/search/query/StreamSearchPhase.java b/server/src/main/java/org/opensearch/search/query/StreamSearchPhase.java
@@ -111,6 +111,7 @@ private boolean searchWithCollector(
                 throw new RuntimeException("StreamManager not setup");
             }
             final boolean[] isCancelled = { false };
+            final Schema[] schema = { null };
             StreamTicket ticket = streamManager.registerStream(new StreamProducer() {
 
                 @Override
@@ -174,7 +175,7 @@ public void onCancel() {
 
                         @Override
                         public boolean isCancelled() {
-                            return searchContext.isCancelled() || isCancelled();
+                            return searchContext.isCancelled() || isCancelled[0];
                         }
                     };
                 }
@@ -185,8 +186,8 @@ public VectorSchemaRoot createRoot(BufferAllocator allocator) {
                     Field countField = new Field("count", FieldType.nullable(new ArrowType.Int(64, false)), null);
                     arrowFields.put("count", countField);
                     arrowFields.put("ord", new Field("ord", FieldType.nullable(new ArrowType.Utf8()), null));
-                    Schema schema = new Schema(arrowFields.values());
-                    return VectorSchemaRoot.create(schema, allocator);
+                    schema[0] = new Schema(arrowFields.values());
+                    return VectorSchemaRoot.create(schema[0], allocator);
                 }
 
                 @Override
-Original file line number
+Diff line change
@@ Expand Up / @@ -125,6 +125,7 @@ public void finish() throws IOException { @@
                     if (currentRow[0] > 0) {
                         flushBatch();
                     }
+                    root.close();
                 }
             };
         }
@@ Expand Down @@