From 5cfc67ba970ce41060fed3432953f08e75632799 Mon Sep 17 00:00:00 2001
From: Mohit Godwani <mgodwan@amazon.com>
Date: Thu, 4 Dec 2025 18:18:26 +0530
Subject: [PATCH] Fix cleanup due to bug in Datafusion Reader not releasing
 catalog snapshot

Signed-off-by: Mohit Godwani <mgodwan@amazon.com>
---
 .../opensearch/datafusion/jni/handle/ReaderHandle.java    | 8 +++++++-
 .../opensearch/datafusion/search/DatafusionReader.java    | 5 ++++-
 2 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/jni/handle/ReaderHandle.java b/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/jni/handle/ReaderHandle.java
index 2df7a4b11bf70..dd1c93f74b7b0 100644
--- a/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/jni/handle/ReaderHandle.java
+++ b/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/jni/handle/ReaderHandle.java
@@ -11,17 +11,23 @@
 import org.opensearch.datafusion.jni.NativeBridge;
 import org.opensearch.vectorized.execution.jni.RefCountedNativeHandle;
 
+import java.io.Closeable;
+
 /**
  * Reference-counted handle for native reader.
  */
 public final class ReaderHandle extends RefCountedNativeHandle {
 
-    public ReaderHandle(String path, String[] files) {
+    private final Runnable onClose;
+
+    public ReaderHandle(String path, String[] files, Runnable onClose) {
         super(NativeBridge.createDatafusionReader(path, files));
+        this.onClose = onClose;
     }
 
     @Override
     protected void doClose() {
         NativeBridge.closeDatafusionReader(ptr);
+        onClose.run();
     }
 }
diff --git a/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/search/DatafusionReader.java b/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/search/DatafusionReader.java
index 7f74e337155e6..ed3a3f15ddfce 100644
--- a/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/search/DatafusionReader.java
+++ b/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/search/DatafusionReader.java
@@ -55,7 +55,7 @@ public DatafusionReader(String directoryPath, CompositeEngine.ReleasableRef<Cata
         }
         System.out.println("File names: " + Arrays.toString(fileNames));
         System.out.println("Directory path: " + directoryPath);
-        this.readerHandle = new ReaderHandle(directoryPath, fileNames);
+        this.readerHandle = new ReaderHandle(directoryPath, fileNames, this::releaseCatalogSnapshot);
     }
 
     /**
@@ -91,6 +91,9 @@ public int getRefCount() {
     @Override
     public void close() {
         readerHandle.close();
+    }
+
+    private void releaseCatalogSnapshot() {
         try {
             if (catalogSnapshotRef != null)
                 catalogSnapshotRef.close();