Adding unit tests for VSR management, bugfix for VSR close hierarchy

Signed-off-by: Raghuvansh Raj <raghraaj@amazon.com>

Author: raghuvanshraj

modules/parquet-data-format/src/main/java/com/parquet/parquetdataformat/vsr/ManagedVSR.java

@@ -59,11 +59,16 @@ public void setRowCount(int rowCount) {
 
     /**
      * Gets a field vector by name.
+     * Only allowed when VSR is in ACTIVE state.
      *
      * @param fieldName Name of the field
      * @return FieldVector for the field, or null if not found
+     * @throws IllegalStateException if VSR is not in ACTIVE state
      */
     public FieldVector getVector(String fieldName) {
+        if (state != VSRState.ACTIVE) {
+            throw new IllegalStateException("Cannot access vector in VSR state: " + state + ". VSR must be ACTIVE to access vectors.");
+        }
         return vsr.getVector(fieldName);
     }
 

modules/parquet-data-format/src/main/java/com/parquet/parquetdataformat/vsr/VSRManager.java

@@ -12,17 +12,13 @@
 import com.parquet.parquetdataformat.bridge.RustBridge;
 import com.parquet.parquetdataformat.memory.ArrowBufferPool;
 import com.parquet.parquetdataformat.writer.ParquetDocumentInput;
-import org.apache.arrow.vector.FieldVector;
-import org.apache.arrow.vector.types.pojo.Field;
 import org.apache.arrow.vector.types.pojo.Schema;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.opensearch.index.engine.exec.FlushIn;
 import org.opensearch.index.engine.exec.WriteResult;
 
 import java.io.IOException;
-import java.util.HashMap;
-import java.util.Map;
 
 /**
  * Manages VectorSchemaRoot lifecycle with integrated memory management and native call wrappers.
@@ -35,12 +31,11 @@
  *   <li>{@link RustBridge} - Direct JNI calls to Rust backend</li>
  * </ul>
  */
-public class VSRManager {
+public class VSRManager implements AutoCloseable {
 
     private static final Logger logger = LogManager.getLogger(VSRManager.class);
 
     private ManagedVSR managedVSR;
-    private Map<String, FieldVector> fieldVectorMap;
     private final Schema schema;
     private final String fileName;
     private final VSRPool vsrPool;
@@ -54,8 +49,7 @@ public VSRManager(String fileName, Schema schema, ArrowBufferPool arrowBufferPoo
 
         // Get active VSR from pool
         this.managedVSR = vsrPool.getActiveVSR();
-        initializeFieldVectorMap();
-        // Initialize writer lazily to avoid crashes
+
         initializeWriter();
     }
 
@@ -80,7 +74,6 @@ public WriteResult addToManagedVSR(ParquetDocumentInput document) throws IOExcep
             if (managedVSR == null) {
                 throw new IOException("No active VSR available");
             }
-            reinitializeFieldVectorMap();
         }
 
         // Ensure VSR is in ACTIVE state for modifications
@@ -136,6 +129,7 @@ public String flush(FlushIn flushIn) throws IOException {
         }
     }
 
+    @Override
     public void close() {
         try {
             // Direct native calls
@@ -146,12 +140,28 @@ public void close() {
                 logger.warn("Failed to close/flush writer for {}: {}", fileName, e.getMessage(), e);
             }
 
-            // Close VSR Pool
+            // Close VSR Pool - handle IllegalStateException specially
             vsrPool.close();
             managedVSR = null;
 
+        } catch (IllegalStateException e) {
+            // Direct IllegalStateException - re-throw for business logic validation
+            logger.error("Error during close for {}: {}", fileName, e.getMessage(), e);
+            throw e;
+        } catch (RuntimeException e) {
+            // Check if this is a wrapped IllegalStateException from defensive cleanup
+            Throwable cause = e.getCause();
+            if (cause instanceof IllegalStateException) {
+                // Re-throw the original IllegalStateException for business logic validation
+                logger.error("Error during close for {}: {}", fileName, cause.getMessage(), cause);
+                throw (IllegalStateException) cause;
+            }
+            // For other RuntimeExceptions, log and re-throw
+            logger.error("Error during close for {}: {}", fileName, e.getMessage(), e);
+            throw new RuntimeException("Failed to close VSRManager: " + e.getMessage(), e);
         } catch (Exception e) {
             logger.error("Error during close for {}: {}", fileName, e.getMessage(), e);
+            throw new RuntimeException("Failed to close VSRManager: " + e.getMessage(), e);
         }
     }
 
@@ -198,9 +208,6 @@ public void maybeRotateActiveVSR() throws IOException {
                     throw new IOException("No active VSR available after rotation");
                 }
 
-                // Reinitialize field vector map with new VSR
-                reinitializeFieldVectorMap();
-
                 logger.debug("VSR rotation completed for {}, new active VSR: {}, row count: {}",
                     fileName, managedVSR.getId(), managedVSR.getRowCount());
             }
@@ -211,56 +218,20 @@ public void maybeRotateActiveVSR() throws IOException {
     }
 
     /**
-     * Checks if VSR rotation is needed based on row count and memory pressure.
-     * If rotation occurs, updates the managed VSR reference and reinitializes field vectors.
+     * Gets the current active ManagedVSR for document input creation.
      *
-     * @deprecated Use handleVSRRotationAfterAddToManagedVSR() instead for safer rotation after document processing
-     */
-    @Deprecated
-    private void checkAndHandleVSRRotation() throws IOException {
-        // Get active VSR from pool - this will trigger rotation if needed
-        ManagedVSR currentActive = vsrPool.getActiveVSR();
-
-        // Check if we got a different VSR (rotation occurred)
-        if (currentActive != managedVSR) {
-            logger.debug("VSR rotation detected for {}, updating references", fileName);
-
-            // Update the managed VSR reference
-            managedVSR = currentActive;
-
-            // Reinitialize field vector map with new VSR
-            reinitializeFieldVectorMap();
-
-            // Note: Writer initialization is not needed per VSR as it's per file
-            logger.debug("VSR rotation completed for {}, new row count: {}", fileName, managedVSR.getRowCount());
-        }
-    }
-
-    /**
-     * Reinitializes the field vector map with the current managed VSR.
-     * Called after VSR rotation to update vector references.
+     * @return The current managed VSR instance
      */
-    private void reinitializeFieldVectorMap() {
-        fieldVectorMap.clear();
-        initializeFieldVectorMap();
-    }
-
-    private void initializeFieldVectorMap() {
-        fieldVectorMap = new HashMap<>();
-        for (Field field : schema.getFields()) {
-            String fieldName = field.getName();
-            FieldVector fieldVector = managedVSR.getVector(fieldName);
-            // Vector is already properly typed from ManagedVSR.getVector()
-            fieldVectorMap.put(fieldName, fieldVector);
-        }
+    public ManagedVSR getActiveManagedVSR() {
+        return managedVSR;
     }
 
     /**
-     * Gets the current active ManagedVSR for document input creation.
+     * Gets the current frozen VSR for testing purposes.
      *
-     * @return The current managed VSR instance
+     * @return The current frozen VSR instance, or null if none exists
      */
-    public ManagedVSR getActiveManagedVSR() {
-        return managedVSR;
+    public ManagedVSR getFrozenVSR() {
+        return vsrPool.getFrozenVSR();
     }
 }

modules/parquet-data-format/src/main/java/com/parquet/parquetdataformat/vsr/VSRPool.java

@@ -17,7 +17,7 @@
  * in the Project Mustang design. Each ParquetWriter maintains a single ACTIVE VSR
  * for writing and a single FROZEN VSR for Rust handoff.
  */
-public class VSRPool {
+public class VSRPool implements AutoCloseable {
 
     private static final Logger logger = LogManager.getLogger(VSRPool.class);
 
@@ -181,18 +181,45 @@ public void freezeAll() {
 
     /**
      * Closes the pool and cleans up all resources.
+     * Uses defensive cleanup to ensure resources are not orphaned if close operations fail.
      */
+    @Override
     public void close() {
-        // Close active VSR
-        ManagedVSR active = activeVSR.getAndSet(null);
+        // Get references without clearing them yet - defensive cleanup approach
+        ManagedVSR active = activeVSR.get();
+        ManagedVSR frozen = frozenVSR.get();
+
+        Exception firstException = null;
+
+        // Try to close active VSR
         if (active != null) {
-            active.close();
+            try {
+                active.close();
+                activeVSR.set(null); // Only clear if successful
+            } catch (Exception e) {
+                firstException = e;
+                // Don't set to null - leave reference so subsequent close attempts can retry
+            }
         }
 
-        // Close frozen VSR
-        ManagedVSR frozen = frozenVSR.getAndSet(null);
+        // Try to close frozen VSR regardless of active VSR result
         if (frozen != null) {
-            frozen.close();
+            try {
+                frozen.close();
+                frozenVSR.set(null); // Only clear if successful
+            } catch (Exception e) {
+                if (firstException != null) {
+                    firstException.addSuppressed(e);
+                } else {
+                    firstException = e;
+                }
+                // Don't set to null - leave reference so subsequent close attempts can retry
+            }
+        }
+
+        // Throw the most relevant exception after attempting all cleanup
+        if (firstException != null) {
+            throw new RuntimeException("VSRPool cleanup failed", firstException);
         }
     }