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

@@ -13,7 +13,6 @@
 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;
@@ -37,16 +36,16 @@
  *   <li>{@link com.parquet.parquetdataformat.bridge.RustBridge} - Direct JNI calls to Rust backend</li>
  * </ul>
  */
-public class VSRManager implements Closeable {
+public class VSRManager implements AutoCloseable {
+
+    private static final Logger logger = LogManager.getLogger(VSRManager.class);
+
     private final AtomicReference<ManagedVSR> managedVSR = new AtomicReference<>();
-    private Map<String, FieldVector> fieldVectorMap;
     private final Schema schema;
     private final String fileName;
     private final VSRPool vsrPool;
     private NativeParquetWriter writer;
 
-    private static final Logger logger = LogManager.getLogger(VSRManager.class);
-
 
     public VSRManager(String fileName, Schema schema, ArrowBufferPool arrowBufferPool) {
         this.fileName = fileName;
@@ -57,7 +56,7 @@ public VSRManager(String fileName, Schema schema, ArrowBufferPool arrowBufferPoo
 
         // Get active VSR from pool
         this.managedVSR.set(vsrPool.getActiveVSR());
-        initializeFieldVectorMap();
+
         // Initialize writer lazily to avoid crashes
         initializeWriter();
     }
@@ -75,11 +74,7 @@ private void initializeWriter() {
     public WriteResult addToManagedVSR(ParquetDocumentInput document) throws IOException {
         ManagedVSR currentVSR = managedVSR.updateAndGet(vsr -> {
             if (vsr == null) {
-                ManagedVSR newVSR = vsrPool.getActiveVSR();
-                if (newVSR != null) {
-                    reinitializeFieldVectorMap();
-                }
-                return newVSR;
+                return vsrPool.getActiveVSR();
             }
             return vsr;
         });
@@ -147,10 +142,29 @@ public void close() {
                 writer.flush();
                 writer.close();
             }
+
+            // Close VSR Pool - handle IllegalStateException specially
             vsrPool.close();
             managedVSR.set(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);
         }
     }
 
@@ -199,9 +213,6 @@ public void maybeRotateActiveVSR() throws IOException {
                 }
                 updateVSRAndReinitialize(oldVSR, newVSR);
 
-                // Reinitialize field vector map with new VSR
-                reinitializeFieldVectorMap();
-
                 logger.debug("VSR rotation completed for {}, new active VSR: {}, row count: {}",
                     fileName, newVSR.getId(), newVSR.getRowCount());
             }
@@ -239,28 +250,7 @@ private void checkAndHandleVSRRotation() throws IOException {
      * Atomically updates managedVSR and reinitializes field vector map.
      */
     private void updateVSRAndReinitialize(ManagedVSR oldVSR, ManagedVSR newVSR) {
-        if (managedVSR.compareAndSet(oldVSR, newVSR)) {
-            reinitializeFieldVectorMap();
-        }
-    }
-
-    /**
-     * Reinitializes the field vector map with the current managed VSR.
-     * Called after VSR rotation to update vector references.
-     */
-    private void reinitializeFieldVectorMap() {
-        fieldVectorMap.clear();
-        initializeFieldVectorMap();
-    }
-
-    private void initializeFieldVectorMap() {
-        fieldVectorMap = new HashMap<>();
-        for (Field field : schema.getFields()) {
-            String fieldName = field.getName();
-            FieldVector fieldVector = managedVSR.get().getVector(fieldName);
-            // Vector is already properly typed from ManagedVSR.getVector()
-            fieldVectorMap.put(fieldName, fieldVector);
-        }
+        managedVSR.compareAndSet(oldVSR, newVSR);
     }
 
     /**
@@ -271,4 +261,13 @@ private void initializeFieldVectorMap() {
     public ManagedVSR getActiveManagedVSR() {
         return managedVSR.get();
     }
+
+    /**
+     * Gets the current frozen VSR for testing purposes.
+     *
+     * @return The current frozen VSR instance, or null if none exists
+     */
+    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);
         }
     }