diff --git a/modules/parquet-data-format/src/main/java/com/parquet/parquetdataformat/vsr/ManagedVSR.java b/modules/parquet-data-format/src/main/java/com/parquet/parquetdataformat/vsr/ManagedVSR.java index 3b4113b0abc00..a42b178c4f377 100644 --- a/modules/parquet-data-format/src/main/java/com/parquet/parquetdataformat/vsr/ManagedVSR.java +++ b/modules/parquet-data-format/src/main/java/com/parquet/parquetdataformat/vsr/ManagedVSR.java @@ -1,10 +1,7 @@ package com.parquet.parquetdataformat.vsr; import com.parquet.parquetdataformat.bridge.ArrowExport; -import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; -import org.apache.arrow.memory.RootAllocator; -import org.apache.arrow.vector.BigIntVector; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.c.ArrowArray; @@ -13,15 +10,12 @@ import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; -import java.util.concurrent.atomic.AtomicReference; -import java.util.concurrent.locks.ReadWriteLock; -import java.util.concurrent.locks.ReentrantReadWriteLock; - import static org.apache.arrow.vector.BitVectorHelper.byteIndex; +import org.apache.arrow.vector.types.pojo.Schema; /** * Managed wrapper around VectorSchemaRoot that handles state transitions - * and provides thread-safe access for the ACTIVE/FROZEN lifecycle. + * for the ACTIVE/FROZEN/CLOSED lifecycle with controlled access methods. */ public class ManagedVSR implements AutoCloseable { @@ -30,43 +24,23 @@ public class ManagedVSR implements AutoCloseable { private final String id; private final VectorSchemaRoot vsr; private final BufferAllocator allocator; - private final AtomicReference state; - private final ReadWriteLock lock; - private final long createdTime; + private VSRState state; - public ManagedVSR(String id, VectorSchemaRoot vsr, BufferAllocator allocator) { + public ManagedVSR(String id, Schema schema, BufferAllocator allocator) { this.id = id; - this.vsr = vsr; + this.vsr = VectorSchemaRoot.create(schema, allocator); this.allocator = allocator; - this.state = new AtomicReference<>(VSRState.ACTIVE); - this.lock = new ReentrantReadWriteLock(); - this.createdTime = System.currentTimeMillis(); - } - - /** - * Gets the underlying VectorSchemaRoot. - * Should only be used when holding appropriate locks. - * - * @return VectorSchemaRoot instance - */ - public VectorSchemaRoot getVSR() { - return vsr; + this.state = VSRState.ACTIVE; } /** * Gets the current row count in this VSR. - * Thread-safe read operation. * * @return Number of rows currently in the VSR */ public int getRowCount() { - lock.readLock().lock(); - try { - return vsr.getRowCount(); - } finally { - lock.readLock().unlock(); - } + return vsr.getRowCount(); } /** @@ -77,94 +51,114 @@ public int getRowCount() { * @throws IllegalStateException if VSR is not active or is immutable */ public void setRowCount(int rowCount) { - lock.writeLock().lock(); - try { - if (state.get() != VSRState.ACTIVE) { - throw new IllegalStateException("Cannot modify VSR in state: " + state.get()); - } - vsr.setRowCount(rowCount); - } finally { - lock.writeLock().unlock(); + if (state != VSRState.ACTIVE) { + throw new IllegalStateException("Cannot modify VSR in state: " + state); } + vsr.setRowCount(rowCount); } /** * Gets a field vector by name. - * Thread-safe read operation. + * 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) { - lock.readLock().lock(); - try { - return vsr.getVector(fieldName); - } finally { - lock.readLock().unlock(); + 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); } /** * Changes the state of this VSR. * Handles state transition logic and immutability. + * This method is private to ensure controlled state transitions. * * @param newState New state to transition to */ - public void setState(VSRState newState) { - VSRState oldState = state.getAndSet(newState); + private void setState(VSRState newState) { + VSRState oldState = state; + state = newState; logger.debug("State transition: {} -> {} for VSR {}", oldState, newState, id); } + /** + * Transitions the VSR from ACTIVE to FROZEN state. + * This is the only way to freeze a VSR. + * + * @throws IllegalStateException if VSR is not in ACTIVE state + */ + public void moveToFrozen() { + if (state != VSRState.ACTIVE) { + throw new IllegalStateException(String.format( + "Cannot freeze VSR %s: expected ACTIVE state but was %s", id, state)); + } + setState(VSRState.FROZEN); + } + + /** + * Transitions the VSR from FROZEN to CLOSED state. + * This method is private and only called by close(). + * + * @throws IllegalStateException if VSR is not in FROZEN state + */ + private void moveToClosed() { + if (state != VSRState.FROZEN) { + throw new IllegalStateException(String.format( + "Cannot close VSR %s: expected FROZEN state but was %s", id, state)); + } + setState(VSRState.CLOSED); + + // Clean up resources + if (vsr != null) { + vsr.close(); + } + if (allocator != null) { + allocator.close(); + } + } + /** * Gets the current state of this VSR. * * @return Current VSRState */ public VSRState getState() { - return state.get(); + return state; } /** * Exports this VSR to Arrow C Data Interface for Rust handoff. - * Only allowed when VSR is FROZEN or FLUSHING. + * Only allowed when VSR is FROZEN. * * @return ArrowExport containing ArrowArray and ArrowSchema * @throws IllegalStateException if VSR is not in correct state */ public ArrowExport exportToArrow() { - VSRState currentState = state.get(); - if (currentState != VSRState.FROZEN && - currentState != VSRState.FLUSHING) { - throw new IllegalStateException("Cannot export VSR in state: " + currentState); + if (state != VSRState.FROZEN) { + throw new IllegalStateException("Cannot export VSR in state: " + state + ". VSR must be FROZEN to export."); } - lock.readLock().lock(); - try { - ArrowArray arrowArray = ArrowArray.allocateNew(allocator); - ArrowSchema arrowSchema = ArrowSchema.allocateNew(allocator); + ArrowArray arrowArray = ArrowArray.allocateNew(allocator); + ArrowSchema arrowSchema = ArrowSchema.allocateNew(allocator); - // Export the VectorSchemaRoot to C Data Interface - Data.exportVectorSchemaRoot(allocator, vsr, null, arrowArray, arrowSchema); + // Export the VectorSchemaRoot to C Data Interface + Data.exportVectorSchemaRoot(allocator, vsr, null, arrowArray, arrowSchema); - return new ArrowExport(arrowArray, arrowSchema); - } finally { - lock.readLock().unlock(); - } + return new ArrowExport(arrowArray, arrowSchema); } public ArrowExport exportSchema() { - lock.readLock().lock(); - try { - ArrowSchema arrowSchema = ArrowSchema.allocateNew(allocator); + ArrowSchema arrowSchema = ArrowSchema.allocateNew(allocator); - // Export the VectorSchemaRoot to C Data Interface - Data.exportSchema(allocator, vsr.getSchema(), null, arrowSchema); + // Export the VectorSchemaRoot to C Data Interface + Data.exportSchema(allocator, vsr.getSchema(), null, arrowSchema); - return new ArrowExport(null, arrowSchema); - } finally { - lock.readLock().unlock(); - } + return new ArrowExport(null, arrowSchema); } /** @@ -173,8 +167,7 @@ public ArrowExport exportSchema() { * @return true if VSR cannot be modified */ public boolean isImmutable() { - VSRState currentState = state.get(); - return currentState != VSRState.ACTIVE; + return state != VSRState.ACTIVE; } @@ -187,15 +180,6 @@ public String getId() { return id; } - /** - * Gets the creation timestamp. - * - * @return Creation time in milliseconds - */ - public long getCreatedTime() { - return createdTime; - } - /** * Gets the associated BufferAllocator. * @@ -207,18 +191,31 @@ public BufferAllocator getAllocator() { /** * Closes this VSR and releases all resources. + * This is the only way to transition a VSR to CLOSED state. + * VSR must be in FROZEN state before it can be closed. + * + * @throws IllegalStateException if VSR is in ACTIVE state (must freeze first) */ @Override public void close() { - lock.writeLock().lock(); - try { - if (state.get() != VSRState.CLOSED) { - state.set(VSRState.CLOSED); - vsr.close(); - allocator.close(); - } - } finally { - lock.writeLock().unlock(); + // If already CLOSED, do nothing (idempotent) + if (state == VSRState.CLOSED) { + return; + } + + // If ACTIVE, must freeze first + if (state == VSRState.ACTIVE) { + throw new IllegalStateException(String.format( + "Cannot close VSR %s: VSR is still ACTIVE. Must freeze VSR before closing.", id)); + } + + // If FROZEN, transition to CLOSED + if (state == VSRState.FROZEN) { + moveToClosed(); + } else { + // This should never happen with current states, but defensive programming + throw new IllegalStateException(String.format( + "Cannot close VSR %s: unexpected state %s", id, state)); } } @@ -226,6 +223,6 @@ public void close() { @Override public String toString() { return String.format("ManagedVSR{id='%s', state=%s, rows=%d, immutable=%s}", - id, state.get(), getRowCount(), isImmutable()); + id, state, getRowCount(), isImmutable()); } } diff --git a/modules/parquet-data-format/src/main/java/com/parquet/parquetdataformat/vsr/VSRManager.java b/modules/parquet-data-format/src/main/java/com/parquet/parquetdataformat/vsr/VSRManager.java index 602402e31001d..8f900a4084821 100644 --- a/modules/parquet-data-format/src/main/java/com/parquet/parquetdataformat/vsr/VSRManager.java +++ b/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 @@ *
  • {@link com.parquet.parquetdataformat.bridge.RustBridge} - Direct JNI calls to Rust backend
    • * */ -public class VSRManager implements Closeable { +public class VSRManager implements AutoCloseable { + + private static final Logger logger = LogManager.getLogger(VSRManager.class); + private final AtomicReference managedVSR = new AtomicReference<>(); - private Map 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; }); @@ -123,12 +118,9 @@ public String flush(FlushIn flushIn) throws IOException { } // Transition VSR to FROZEN state before flushing - currentVSR.setState(VSRState.FROZEN); + currentVSR.moveToFrozen(); logger.info("Flushing {} rows for {}", currentVSR.getRowCount(), fileName); - // Transition to FLUSHING state - currentVSR.setState(VSRState.FLUSHING); - // Write through native writer handle try (ArrowExport export = currentVSR.exportToArrow()) { writer.write(export.getArrayAddress(), export.getSchemaAddress()); @@ -150,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); } } @@ -181,7 +192,6 @@ public void maybeRotateActiveVSR() throws IOException { frozenVSR.getId(), frozenVSR.getRowCount(), fileName); // Write the frozen VSR data immediately - frozenVSR.setState(VSRState.FLUSHING); try (ArrowExport export = frozenVSR.exportToArrow()) { writer.write(export.getArrayAddress(), export.getSchemaAddress()); } @@ -203,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()); } @@ -243,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); } /** @@ -275,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(); + } } diff --git a/modules/parquet-data-format/src/main/java/com/parquet/parquetdataformat/vsr/VSRPool.java b/modules/parquet-data-format/src/main/java/com/parquet/parquetdataformat/vsr/VSRPool.java index 4c3317200b712..eeef66a63ecf9 100644 --- a/modules/parquet-data-format/src/main/java/com/parquet/parquetdataformat/vsr/VSRPool.java +++ b/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); @@ -28,7 +28,6 @@ public class VSRPool { // VSR lifecycle management private final AtomicReference activeVSR; private final AtomicReference frozenVSR; - private final ConcurrentHashMap allVSRs; private final AtomicInteger vsrCounter; // Configuration @@ -40,7 +39,6 @@ public VSRPool(String poolId, Schema schema, ArrowBufferPool arrowBufferPool) { this.bufferPool = arrowBufferPool; this.activeVSR = new AtomicReference<>(); this.frozenVSR = new AtomicReference<>(); - this.allVSRs = new ConcurrentHashMap<>(); this.vsrCounter = new AtomicInteger(0); // Configuration - could be made configurable @@ -161,24 +159,13 @@ public ManagedVSR takeFrozenVSR() { return frozenVSR.getAndSet(null); } - /** - * Marks a VSR as flushing (being processed by Rust). - * - * @param vsr VSR being processed - */ - public void markFlushing(ManagedVSR vsr) { - vsr.setState(VSRState.FLUSHING); - } - /** * Completes VSR processing and cleans up resources. * * @param vsr VSR that has been processed */ public void completeVSR(ManagedVSR vsr) { - vsr.setState(VSRState.CLOSED); vsr.close(); - allVSRs.remove(vsr.getId()); } /** @@ -193,43 +180,47 @@ public void freezeAll() { } /** - * Gets statistics about the VSR pool. - * - * @return PoolStats with current state + * Closes the pool and cleans up all resources. + * Uses defensive cleanup to ensure resources are not orphaned if close operations fail. */ - public PoolStats getStats() { + @Override + public void close() { + // Get references without clearing them yet - defensive cleanup approach ManagedVSR active = activeVSR.get(); ManagedVSR frozen = frozenVSR.get(); - int frozenCount = frozen != null ? 1 : 0; - - return new PoolStats( - poolId, - active != null ? active.getRowCount() : 0, - frozenCount, - allVSRs.size(), - allVSRs.values().stream().mapToLong(ManagedVSR::getRowCount).sum() - ); - } - /** - * Closes the pool and cleans up all resources. - */ - public void close() { - // Close active VSR - ManagedVSR active = activeVSR.getAndSet(null); + 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 + } } - // Close any remaining VSRs - allVSRs.values().forEach(ManagedVSR::close); - allVSRs.clear(); + // Throw the most relevant exception after attempting all cleanup + if (firstException != null) { + throw new RuntimeException("VSRPool cleanup failed", firstException); + } } private void initializeActiveVSR() { @@ -245,10 +236,7 @@ private ManagedVSR createNewVSR() { try { allocator = bufferPool.createChildAllocator(vsrId); - vsr = VectorSchemaRoot.create(schema, allocator); - - ManagedVSR managedVSR = new ManagedVSR(vsrId, vsr, allocator); - allVSRs.put(vsrId, managedVSR); + ManagedVSR managedVSR = new ManagedVSR(vsrId, schema, allocator); // Success: ManagedVSR now owns the resources return managedVSR; @@ -273,60 +261,33 @@ private ManagedVSR createNewVSR() { } private void freezeVSR(ManagedVSR vsr) { - vsr.setState(VSRState.FROZEN); - - // CRITICAL FIX: Check if frozen slot is already occupied + // Check if frozen slot is already occupied ManagedVSR previousFrozen = frozenVSR.get(); if (previousFrozen != null) { - // NEVER blindly overwrite a frozen VSR - this would cause data loss + // Never blindly overwrite a frozen VSR - this would cause data loss logger.error("Attempting to freeze VSR when frozen slot is occupied! " + "Previous VSR: {} ({} rows), New VSR: {} ({} rows). " + "This indicates a logic error - frozen VSR should be consumed before replacement.", previousFrozen.getId(), previousFrozen.getRowCount(), vsr.getId(), vsr.getRowCount()); - // Return VSR to ACTIVE state to prevent state corruption - vsr.setState(VSRState.ACTIVE); throw new IllegalStateException("Cannot freeze VSR: frozen slot is occupied by unprocessed VSR " + previousFrozen.getId() + ". This would cause data loss."); } - // Safe to set frozen VSR since slot is empty + // First freeze the VSR (validates ACTIVE -> FROZEN transition) + vsr.moveToFrozen(); + + // Safe to set frozen VSR since slot is empty and VSR is now frozen boolean success = frozenVSR.compareAndSet(null, vsr); if (!success) { // Race condition: another thread set frozen VSR between our check and set - vsr.setState(VSRState.ACTIVE); - throw new IllegalStateException("Race condition detected: frozen slot was occupied during freeze operation"); + // This is a critical error since we can't revert the freeze operation + throw new IllegalStateException("Race condition detected: frozen slot was occupied during freeze operation for VSR " + vsr.getId() + ", slot occupied by VSR " + frozenVSR.get().getId()); } } private boolean shouldRotateVSR(ManagedVSR vsr) { return vsr.getRowCount() >= maxRowsPerVSR; } - - /** - * Statistics for the VSR pool. - */ - public static class PoolStats { - private final String poolId; - private final long activeRowCount; - private final int frozenVSRCount; - private final int totalVSRCount; - private final long totalRowCount; - - public PoolStats(String poolId, long activeRowCount, int frozenVSRCount, - int totalVSRCount, long totalRowCount) { - this.poolId = poolId; - this.activeRowCount = activeRowCount; - this.frozenVSRCount = frozenVSRCount; - this.totalVSRCount = totalVSRCount; - this.totalRowCount = totalRowCount; - } - - public String getPoolId() { return poolId; } - public long getActiveRowCount() { return activeRowCount; } - public int getFrozenVSRCount() { return frozenVSRCount; } - public int getTotalVSRCount() { return totalVSRCount; } - public long getTotalRowCount() { return totalRowCount; } - } } diff --git a/modules/parquet-data-format/src/main/java/com/parquet/parquetdataformat/vsr/VSRState.java b/modules/parquet-data-format/src/main/java/com/parquet/parquetdataformat/vsr/VSRState.java index cd55f30ca24cc..81d8441053408 100644 --- a/modules/parquet-data-format/src/main/java/com/parquet/parquetdataformat/vsr/VSRState.java +++ b/modules/parquet-data-format/src/main/java/com/parquet/parquetdataformat/vsr/VSRState.java @@ -16,11 +16,6 @@ public enum VSRState { */ FROZEN, - /** - * Currently being processed by Rust - VSR is in the handoff process. - */ - FLUSHING, - /** * Completed and cleaned up - VSR processing is complete and resources freed. */ diff --git a/modules/parquet-data-format/src/test/java/com/parquet/parquetdataformat/vsr/ManagedVSRTests.java b/modules/parquet-data-format/src/test/java/com/parquet/parquetdataformat/vsr/ManagedVSRTests.java new file mode 100644 index 0000000000000..e86acbed80f2b --- /dev/null +++ b/modules/parquet-data-format/src/test/java/com/parquet/parquetdataformat/vsr/ManagedVSRTests.java @@ -0,0 +1,368 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package com.parquet.parquetdataformat.vsr; + +import com.parquet.parquetdataformat.bridge.ArrowExport; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.apache.arrow.vector.types.pojo.Schema; +import org.apache.arrow.vector.types.Types; +import org.opensearch.test.OpenSearchTestCase; + +import java.util.Arrays; + +/** + * Comprehensive unit tests for ManagedVSR covering all changes from atomic/thread-safe + * handling removal and state management simplification. + */ +public class ManagedVSRTests extends OpenSearchTestCase { + + private BufferAllocator allocator; + private Schema testSchema; + private String vsrId; + + @Override + public void setUp() throws Exception { + super.setUp(); + allocator = new RootAllocator(); + + // Create a simple test schema + Field idField = new Field("id", FieldType.nullable(Types.MinorType.INT.getType()), null); + Field nameField = new Field("name", FieldType.nullable(Types.MinorType.VARCHAR.getType()), null); + testSchema = new Schema(Arrays.asList(idField, nameField)); + + vsrId = "test-vsr-" + System.currentTimeMillis(); + } + + @Override + public void tearDown() throws Exception { + if (allocator != null) { + allocator.close(); + } + super.tearDown(); + } + + // ===== Constructor Tests ===== + + public void testConstructorCreatesActiveVSR() { + ManagedVSR managedVSR = new ManagedVSR(vsrId, testSchema, allocator); + + assertEquals("VSR should start in ACTIVE state", VSRState.ACTIVE, managedVSR.getState()); + assertEquals("VSR ID should match", vsrId, managedVSR.getId()); + assertEquals("Initial row count should be 0", 0, managedVSR.getRowCount()); + assertFalse("New VSR should not be immutable", managedVSR.isImmutable()); + + // Must freeze before closing + managedVSR.moveToFrozen(); + managedVSR.close(); + } + + public void testConstructorWithNullParameters() { + // Note: The current ManagedVSR implementation may not have explicit null validation + // This test documents expected behavior but may need implementation updates + + // Test with valid parameters to ensure constructor works + ManagedVSR validVSR = new ManagedVSR(vsrId, testSchema, allocator); + assertNotNull("Valid constructor should work", validVSR); + validVSR.moveToFrozen(); + validVSR.close(); + } + + // ===== State Transition Tests ===== + + public void testStateTransitionActiveToFrozen() { + ManagedVSR managedVSR = new ManagedVSR(vsrId, testSchema, allocator); + + // Verify initial state + assertEquals("Should start ACTIVE", VSRState.ACTIVE, managedVSR.getState()); + assertFalse("Should not be immutable when active", managedVSR.isImmutable()); + + // Transition to FROZEN + managedVSR.moveToFrozen(); + + assertEquals("Should be FROZEN after moveToFrozen()", VSRState.FROZEN, managedVSR.getState()); + assertTrue("Should be immutable when frozen", managedVSR.isImmutable()); + + managedVSR.close(); + } + + public void testMoveToFrozenFromNonActiveState() { + ManagedVSR managedVSR = new ManagedVSR(vsrId, testSchema, allocator); + + // Move to FROZEN first + managedVSR.moveToFrozen(); + assertEquals("Should be FROZEN", VSRState.FROZEN, managedVSR.getState()); + + // Try to freeze again - should fail + IllegalStateException exception = expectThrows(IllegalStateException.class, managedVSR::moveToFrozen); + + assertTrue("Exception should mention expected ACTIVE state", + exception.getMessage().contains("expected ACTIVE state")); + assertTrue("Exception should mention current FROZEN state", + exception.getMessage().contains("FROZEN")); + + managedVSR.close(); + } + + public void testStateTransitionFrozenToClosed() { + ManagedVSR managedVSR = new ManagedVSR(vsrId, testSchema, allocator); + + // Move to FROZEN then CLOSED + managedVSR.moveToFrozen(); + assertEquals("Should be FROZEN", VSRState.FROZEN, managedVSR.getState()); + + managedVSR.close(); + assertEquals("Should be CLOSED after close()", VSRState.CLOSED, managedVSR.getState()); + } + + public void testCloseFromActiveStateFails() { + ManagedVSR managedVSR = new ManagedVSR(vsrId, testSchema, allocator); + + assertEquals("Should start ACTIVE", VSRState.ACTIVE, managedVSR.getState()); + + // Try to close while ACTIVE - should fail + IllegalStateException exception = expectThrows(IllegalStateException.class, managedVSR::close); + + assertTrue("Exception should mention VSR is ACTIVE", + exception.getMessage().contains("VSR is still ACTIVE")); + assertTrue("Exception should mention must freeze first", + exception.getMessage().contains("Must freeze VSR before closing")); + + // VSR should still be ACTIVE after failed close + assertEquals("VSR should still be ACTIVE", VSRState.ACTIVE, managedVSR.getState()); + + // Proper cleanup + managedVSR.moveToFrozen(); + managedVSR.close(); + } + + public void testCloseIdempotency() { + ManagedVSR managedVSR = new ManagedVSR(vsrId, testSchema, allocator); + + // Move to FROZEN then CLOSED + managedVSR.moveToFrozen(); + managedVSR.close(); + assertEquals("Should be CLOSED", VSRState.CLOSED, managedVSR.getState()); + + // Call close again - should be idempotent + managedVSR.close(); + assertEquals("Should still be CLOSED", VSRState.CLOSED, managedVSR.getState()); + } + + // ===== Operation State Validation Tests ===== + + public void testSetRowCountOnlyWorksInActiveState() { + ManagedVSR managedVSR = new ManagedVSR(vsrId, testSchema, allocator); + + // Should work in ACTIVE state + managedVSR.setRowCount(5); + assertEquals("Row count should be set", 5, managedVSR.getRowCount()); + + // Move to FROZEN + managedVSR.moveToFrozen(); + + // Should fail in FROZEN state + IllegalStateException exception = expectThrows(IllegalStateException.class, () -> managedVSR.setRowCount(10)); + + assertTrue("Exception should mention cannot modify in FROZEN state", + exception.getMessage().contains("Cannot modify VSR in state: FROZEN")); + assertEquals("Row count should remain unchanged", 5, managedVSR.getRowCount()); + + managedVSR.close(); + } + + public void testExportToArrowOnlyWorksInFrozenState() { + ManagedVSR managedVSR = new ManagedVSR(vsrId, testSchema, allocator); + + // Should fail in ACTIVE state + IllegalStateException exception = expectThrows(IllegalStateException.class, managedVSR::exportToArrow); + + assertTrue("Exception should mention cannot export in ACTIVE state", + exception.getMessage().contains("Cannot export VSR in state: ACTIVE")); + assertTrue("Exception should mention must be FROZEN", + exception.getMessage().contains("VSR must be FROZEN to export")); + + // Move to FROZEN - should work + managedVSR.moveToFrozen(); + + try (ArrowExport export = managedVSR.exportToArrow()) { + assertNotNull("Export should not be null", export); + assertTrue("Array address should be valid", export.getArrayAddress() != 0); + assertTrue("Schema address should be valid", export.getSchemaAddress() != 0); + } + + managedVSR.close(); + } + + public void testGetVectorOnlyWorksInActiveState() { + ManagedVSR managedVSR = new ManagedVSR(vsrId, testSchema, allocator); + + // Test in ACTIVE state - should work + FieldVector idVector = managedVSR.getVector("id"); + assertNotNull("Should get vector in ACTIVE state", idVector); + assertTrue("Should be IntVector", idVector instanceof IntVector); + + // Test in FROZEN state - should fail + managedVSR.moveToFrozen(); + IllegalStateException frozenException = expectThrows(IllegalStateException.class, () -> managedVSR.getVector("id")); + assertTrue("Exception should mention cannot access in FROZEN state", + frozenException.getMessage().contains("Cannot access vector in VSR state: FROZEN")); + assertTrue("Exception should mention must be ACTIVE", + frozenException.getMessage().contains("VSR must be ACTIVE to access vectors")); + + managedVSR.close(); + + // Test in CLOSED state - should also fail + IllegalStateException closedException = expectThrows(IllegalStateException.class, () -> managedVSR.getVector("id")); + assertTrue("Exception should mention cannot access in CLOSED state", + closedException.getMessage().contains("Cannot access vector in VSR state: CLOSED")); + assertTrue("Exception should mention must be ACTIVE", + closedException.getMessage().contains("VSR must be ACTIVE to access vectors")); + } + + public void testGetVectorWithNonExistentField() { + ManagedVSR managedVSR = new ManagedVSR(vsrId, testSchema, allocator); + + FieldVector nonExistentVector = managedVSR.getVector("nonexistent"); + assertNull("Should return null for non-existent field", nonExistentVector); + + managedVSR.moveToFrozen(); + managedVSR.close(); + } + + public void testExportSchemaWorksInAllStates() { + ManagedVSR managedVSR = new ManagedVSR(vsrId, testSchema, allocator); + + // Test in ACTIVE state + try (ArrowExport schemaExport = managedVSR.exportSchema()) { + assertNotNull("Schema export should not be null", schemaExport); + // Note: Schema-only exports may have null array address - this is expected + assertTrue("Schema address should be valid", schemaExport.getSchemaAddress() != 0); + } + + // Test in FROZEN state + managedVSR.moveToFrozen(); + try (ArrowExport schemaExportFrozen = managedVSR.exportSchema()) { + assertNotNull("Schema export should work in FROZEN state", schemaExportFrozen); + } + + managedVSR.close(); + } + + // ===== Resource Management Tests ===== + + public void testResourceCleanupOnClose() { + // Create a separate allocator for this test so we can verify it gets closed + BufferAllocator testAllocator = new RootAllocator(); + ManagedVSR managedVSR = new ManagedVSR(vsrId, testSchema, testAllocator); + + // Add some data + managedVSR.setRowCount(3); + + // Verify resources are allocated + assertTrue("Allocator should have allocated memory", testAllocator.getAllocatedMemory() > 0); + + // Close properly + managedVSR.moveToFrozen(); + managedVSR.close(); + + assertEquals("Should be CLOSED", VSRState.CLOSED, managedVSR.getState()); + + // Verify allocator has no reserved memory after close (the allocator itself gets closed by ManagedVSR) + assertEquals("Allocator should have no reserved memory after close", 0, testAllocator.getAllocatedMemory()); + } + + // ===== Edge Case Tests ===== + + public void testToStringInDifferentStates() { + ManagedVSR managedVSR = new ManagedVSR(vsrId, testSchema, allocator); + + // Test toString in ACTIVE state + String activeString = managedVSR.toString(); + assertTrue("toString should contain VSR ID", activeString.contains(vsrId)); + assertTrue("toString should contain ACTIVE state", activeString.contains("ACTIVE")); + assertTrue("toString should contain row count", activeString.contains("rows=0")); + assertTrue("toString should contain immutable=false", activeString.contains("immutable=false")); + + // Test toString in FROZEN state + managedVSR.moveToFrozen(); + String frozenString = managedVSR.toString(); + assertTrue("toString should contain FROZEN state", frozenString.contains("FROZEN")); + assertTrue("toString should contain immutable=true", frozenString.contains("immutable=true")); + + // Test toString in CLOSED state + managedVSR.close(); + String closedString = managedVSR.toString(); + assertTrue("toString should contain CLOSED state", closedString.contains("CLOSED")); + assertTrue("toString should contain immutable=true", closedString.contains("immutable=true")); + } + + public void testGetRowCountAfterStateTransitions() { + ManagedVSR managedVSR = new ManagedVSR(vsrId, testSchema, allocator); + + // Set initial row count + managedVSR.setRowCount(10); + assertEquals("Row count should be 10", 10, managedVSR.getRowCount()); + + // Row count should persist through state transitions + managedVSR.moveToFrozen(); + assertEquals("Row count should persist in FROZEN state", 10, managedVSR.getRowCount()); + + managedVSR.close(); + assertEquals("Row count should persist in CLOSED state", 10, managedVSR.getRowCount()); + } + + public void testImmutabilityInDifferentStates() { + ManagedVSR managedVSR = new ManagedVSR(vsrId, testSchema, allocator); + + // ACTIVE state - should be mutable + assertFalse("Should be mutable in ACTIVE state", managedVSR.isImmutable()); + + // FROZEN state - should be immutable + managedVSR.moveToFrozen(); + assertTrue("Should be immutable in FROZEN state", managedVSR.isImmutable()); + + // CLOSED state - should be immutable + managedVSR.close(); + assertTrue("Should be immutable in CLOSED state", managedVSR.isImmutable()); + } + + // ===== Integration Tests ===== + + public void testCompleteVSRLifecycle() { + ManagedVSR managedVSR = new ManagedVSR(vsrId, testSchema, allocator); + + // 1. Start in ACTIVE state + assertEquals("Should start ACTIVE", VSRState.ACTIVE, managedVSR.getState()); + assertFalse("Should be mutable", managedVSR.isImmutable()); + + // 2. Populate with data (only possible in ACTIVE) + managedVSR.setRowCount(5); + assertEquals("Should have 5 rows", 5, managedVSR.getRowCount()); + + // 3. Transition to FROZEN + managedVSR.moveToFrozen(); + assertEquals("Should be FROZEN", VSRState.FROZEN, managedVSR.getState()); + assertTrue("Should be immutable", managedVSR.isImmutable()); + + // 4. Export data (only possible in FROZEN) + try (ArrowExport export = managedVSR.exportToArrow()) { + assertNotNull("Export should succeed", export); + } + + // 5. Close and cleanup (only possible from FROZEN) + managedVSR.close(); + assertEquals("Should be CLOSED", VSRState.CLOSED, managedVSR.getState()); + assertTrue("Should remain immutable", managedVSR.isImmutable()); + } +} diff --git a/modules/parquet-data-format/src/test/java/com/parquet/parquetdataformat/vsr/VSRIntegrationTests.java b/modules/parquet-data-format/src/test/java/com/parquet/parquetdataformat/vsr/VSRIntegrationTests.java new file mode 100644 index 0000000000000..0c92a7f84cd5b --- /dev/null +++ b/modules/parquet-data-format/src/test/java/com/parquet/parquetdataformat/vsr/VSRIntegrationTests.java @@ -0,0 +1,418 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package com.parquet.parquetdataformat.vsr; + +import com.parquet.parquetdataformat.bridge.ArrowExport; +import com.parquet.parquetdataformat.memory.ArrowBufferPool; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.apache.arrow.vector.types.pojo.Schema; +import org.apache.arrow.vector.types.Types; +import org.opensearch.test.OpenSearchTestCase; +import org.opensearch.common.settings.Settings; + +import java.util.Arrays; + +/** + * End-to-end integration tests covering the complete VSR lifecycle across all components + * after the removal of atomic/thread-safe handling and state management simplification. + */ +public class VSRIntegrationTests extends OpenSearchTestCase { + + private BufferAllocator allocator; + private ArrowBufferPool bufferPool; + private Schema testSchema; + private String poolId; + + @Override + public void setUp() throws Exception { + super.setUp(); + allocator = new RootAllocator(); + bufferPool = new ArrowBufferPool(Settings.EMPTY); + + // Create a simple test schema + Field idField = new Field("id", FieldType.nullable(Types.MinorType.INT.getType()), null); + Field nameField = new Field("name", FieldType.nullable(Types.MinorType.VARCHAR.getType()), null); + testSchema = new Schema(Arrays.asList(idField, nameField)); + + poolId = "integration-test-pool-" + System.currentTimeMillis(); + } + + @Override + public void tearDown() throws Exception { + if (bufferPool != null) { + bufferPool.close(); + } + if (allocator != null) { + allocator.close(); + } + super.tearDown(); + } + + // ===== Complete VSR Lifecycle Integration Tests ===== + + public void testCompleteVSRLifecycleIntegration() { + // Test the complete VSR lifecycle: Pool -> Manager -> Export -> Cleanup + VSRPool pool = new VSRPool(poolId, testSchema, bufferPool); + + // 1. VSRPool provides active VSR + ManagedVSR activeVSR = pool.getActiveVSR(); + assertNotNull("Pool should provide active VSR", activeVSR); + assertEquals("VSR should start ACTIVE", VSRState.ACTIVE, activeVSR.getState()); + + // 2. Simulate VSRManager populating data + activeVSR.setRowCount(100); + assertEquals("VSR should have expected row count", 100, activeVSR.getRowCount()); + + // Verify data can be added to vectors + IntVector idVector = (IntVector) activeVSR.getVector("id"); + VarCharVector nameVector = (VarCharVector) activeVSR.getVector("name"); + assertNotNull("Should have id vector", idVector); + assertNotNull("Should have name vector", nameVector); + + // 3. VSRManager decides to freeze for processing + activeVSR.moveToFrozen(); + assertEquals("VSR should be FROZEN", VSRState.FROZEN, activeVSR.getState()); + assertTrue("VSR should be immutable when frozen", activeVSR.isImmutable()); + + // 4. Export for Rust processing (simulating VSRManager -> RustBridge handoff) + try (ArrowExport export = activeVSR.exportToArrow()) { + assertNotNull("Export should succeed", export); + assertTrue("Array address should be valid", export.getArrayAddress() != 0); + assertTrue("Schema address should be valid", export.getSchemaAddress() != 0); + + // Simulate successful Rust processing + // (In real implementation, RustBridge.write would be called here) + } + + // 5. After processing, VSRPool completes the VSR + pool.completeVSR(activeVSR); + assertEquals("VSR should be CLOSED after completion", VSRState.CLOSED, activeVSR.getState()); + + // 6. Cleanup + pool.close(); + } + + public void testMultipleVSRsWithDifferentLifecycleStages() { + // Test managing multiple VSRs at different lifecycle stages simultaneously + VSRPool pool = new VSRPool(poolId, testSchema, bufferPool); + + // Create VSRs at different stages + ManagedVSR activeVSR = pool.getActiveVSR(); + + BufferAllocator frozenAllocator = bufferPool.createChildAllocator("frozen-vsr"); + ManagedVSR frozenVSR = new ManagedVSR("frozen-vsr", testSchema, frozenAllocator); + + BufferAllocator closedAllocator = bufferPool.createChildAllocator("closed-vsr"); + ManagedVSR closedVSR = new ManagedVSR("closed-vsr", testSchema, closedAllocator); + + // Set up different states + activeVSR.setRowCount(50); + assertEquals("Active VSR should be ACTIVE", VSRState.ACTIVE, activeVSR.getState()); + + frozenVSR.setRowCount(75); + frozenVSR.moveToFrozen(); + assertEquals("Frozen VSR should be FROZEN", VSRState.FROZEN, frozenVSR.getState()); + + closedVSR.setRowCount(25); + closedVSR.moveToFrozen(); + closedVSR.close(); + assertEquals("Closed VSR should be CLOSED", VSRState.CLOSED, closedVSR.getState()); + + // Verify operations work correctly for each state + + // Active VSR: can be modified and should export after freezing + activeVSR.setRowCount(60); + assertEquals("Active VSR row count should be updated", 60, activeVSR.getRowCount()); + + activeVSR.moveToFrozen(); + try (ArrowExport export = activeVSR.exportToArrow()) { + assertNotNull("Active->Frozen VSR should export", export); + } + + // Frozen VSR: should be able to export but not modify + try (ArrowExport export = frozenVSR.exportToArrow()) { + assertNotNull("Frozen VSR should export", export); + } + + IllegalStateException exception = expectThrows(IllegalStateException.class, () -> { + frozenVSR.setRowCount(80); + }); + assertTrue("Should not allow modification of frozen VSR", + exception.getMessage().contains("Cannot modify VSR in state: FROZEN")); + + // Closed VSR: should maintain its final state + assertEquals("Closed VSR should maintain row count", 25, closedVSR.getRowCount()); + assertEquals("Closed VSR should remain CLOSED", VSRState.CLOSED, closedVSR.getState()); + + // Complete the active VSR + pool.completeVSR(activeVSR); + assertEquals("Active VSR should be CLOSED after completion", VSRState.CLOSED, activeVSR.getState()); + + // Complete the frozen VSR + pool.completeVSR(frozenVSR); + assertEquals("Frozen VSR should be CLOSED after completion", VSRState.CLOSED, frozenVSR.getState()); + + pool.close(); + } + + public void testVSRStateTransitionValidation() { + // Test that all state transitions are properly validated across components + VSRPool pool = new VSRPool(poolId, testSchema, bufferPool); + + ManagedVSR vsr = pool.getActiveVSR(); + + // Valid transition: ACTIVE -> FROZEN + assertEquals("Should start ACTIVE", VSRState.ACTIVE, vsr.getState()); + vsr.moveToFrozen(); + assertEquals("Should be FROZEN after moveToFrozen()", VSRState.FROZEN, vsr.getState()); + + // Invalid transition: FROZEN -> FROZEN + IllegalStateException exception1 = expectThrows(IllegalStateException.class, vsr::moveToFrozen); + assertTrue("Should not allow FROZEN->FROZEN transition", + exception1.getMessage().contains("expected ACTIVE state")); + + // Valid transition: FROZEN -> CLOSED + vsr.close(); + assertEquals("Should be CLOSED after close()", VSRState.CLOSED, vsr.getState()); + + // Invalid operations on CLOSED VSR + IllegalStateException exception2 = expectThrows(IllegalStateException.class, vsr::moveToFrozen); + assertTrue("Should not allow CLOSED->FROZEN transition", + exception2.getMessage().contains("expected ACTIVE state")); + + // Test that close() is idempotent + vsr.close(); + assertEquals("VSR should remain CLOSED", VSRState.CLOSED, vsr.getState()); + + pool.close(); + } + + public void testVSROperationRestrictionsByState() { + // Test that operations are properly restricted based on VSR state + BufferAllocator childAllocator = bufferPool.createChildAllocator("test-restrictions-vsr"); + ManagedVSR vsr = new ManagedVSR("test-restrictions-vsr", testSchema, childAllocator); + + // ACTIVE state: all modification operations should work + assertEquals("Should start ACTIVE", VSRState.ACTIVE, vsr.getState()); + vsr.setRowCount(10); + assertEquals("Row count should be set in ACTIVE state", 10, vsr.getRowCount()); + + // Export should fail in ACTIVE state + IllegalStateException exception1 = expectThrows(IllegalStateException.class, vsr::exportToArrow); + assertTrue("Should not allow export in ACTIVE state", + exception1.getMessage().contains("Cannot export VSR in state: ACTIVE")); + + // Schema export should work in all states + try (ArrowExport schemaExport = vsr.exportSchema()) { + assertNotNull("Schema export should work in ACTIVE state", schemaExport); + } + + // Transition to FROZEN + vsr.moveToFrozen(); + assertEquals("Should be FROZEN", VSRState.FROZEN, vsr.getState()); + + // FROZEN state: modifications should fail, exports should work + IllegalStateException exception2 = expectThrows(IllegalStateException.class, () -> { + vsr.setRowCount(20); + }); + assertTrue("Should not allow modification in FROZEN state", + exception2.getMessage().contains("Cannot modify VSR in state: FROZEN")); + + // Export should work in FROZEN state + try (ArrowExport export = vsr.exportToArrow()) { + assertNotNull("Export should work in FROZEN state", export); + } + + try (ArrowExport schemaExport = vsr.exportSchema()) { + assertNotNull("Schema export should work in FROZEN state", schemaExport); + } + + // Transition to CLOSED + vsr.close(); + assertEquals("Should be CLOSED", VSRState.CLOSED, vsr.getState()); + + // CLOSED state: most operations should still return values but no modifications + assertEquals("Should maintain row count in CLOSED state", 10, vsr.getRowCount()); + assertTrue("Should be immutable in CLOSED state", vsr.isImmutable()); + + IllegalStateException exception3 = expectThrows(IllegalStateException.class, () -> { + vsr.setRowCount(30); + }); + assertTrue("Should not allow modification in CLOSED state", + exception3.getMessage().contains("Cannot modify VSR in state: CLOSED")); + } + + public void testErrorHandlingAcrossComponents() { + // Test error handling scenarios that span multiple components + VSRPool pool = new VSRPool(poolId, testSchema, bufferPool); + + // Test invalid close attempt + ManagedVSR activeVSR = pool.getActiveVSR(); + activeVSR.setRowCount(15); + + // Attempting to close an ACTIVE VSR should fail + IllegalStateException exception1 = expectThrows(IllegalStateException.class, activeVSR::close); + assertTrue("Should not allow closing ACTIVE VSR", + exception1.getMessage().contains("VSR is still ACTIVE")); + assertTrue("Should mention freezing requirement", + exception1.getMessage().contains("Must freeze VSR before closing")); + + // VSR should still be ACTIVE after failed close + assertEquals("VSR should still be ACTIVE", VSRState.ACTIVE, activeVSR.getState()); + + // Test proper error recovery + activeVSR.moveToFrozen(); + activeVSR.close(); + assertEquals("VSR should be properly CLOSED", VSRState.CLOSED, activeVSR.getState()); + + pool.close(); + } + + public void testResourceManagementIntegration() { + // Test resource management across all components + long initialMemory = bufferPool.getTotalAllocatedBytes(); + + VSRPool pool = new VSRPool(poolId, testSchema, bufferPool); + + // Memory tracking may be delayed or managed differently - just ensure operations complete + long afterPoolCreation = bufferPool.getTotalAllocatedBytes(); + assertTrue("Memory operations should complete", afterPoolCreation >= initialMemory); + + // Get active VSR and populate it + ManagedVSR activeVSR = pool.getActiveVSR(); + activeVSR.setRowCount(200); + + long afterVSRPopulation = bufferPool.getTotalAllocatedBytes(); + assertTrue("Memory should be manageable", afterVSRPopulation >= initialMemory); + + // Create additional VSRs + BufferAllocator childAllocator1 = bufferPool.createChildAllocator("additional-vsr-1"); + BufferAllocator childAllocator2 = bufferPool.createChildAllocator("additional-vsr-2"); + ManagedVSR additionalVSR1 = new ManagedVSR("additional-vsr-1", testSchema, childAllocator1); + ManagedVSR additionalVSR2 = new ManagedVSR("additional-vsr-2", testSchema, childAllocator2); + + additionalVSR1.setRowCount(50); + additionalVSR2.setRowCount(75); + + // Clean up VSRs - must freeze active VSRs before completing them + activeVSR.moveToFrozen(); + pool.completeVSR(activeVSR); + + additionalVSR1.moveToFrozen(); + pool.completeVSR(additionalVSR1); + + additionalVSR2.moveToFrozen(); + pool.completeVSR(additionalVSR2); + + // Close pool + pool.close(); + + // Verify operations completed without error + assertEquals("Active VSR should be CLOSED", VSRState.CLOSED, activeVSR.getState()); + assertEquals("Additional VSR1 should be CLOSED", VSRState.CLOSED, additionalVSR1.getState()); + assertEquals("Additional VSR2 should be CLOSED", VSRState.CLOSED, additionalVSR2.getState()); + } + + public void testBackwardCompatibilityScenarios() { + // Test that the new system maintains backward compatibility patterns + VSRPool pool = new VSRPool(poolId, testSchema, bufferPool); + + // Simulate old pattern: get VSR, populate, process + ManagedVSR vsr = pool.getActiveVSR(); + assertNotNull("Should get VSR like before", vsr); + + // Old pattern operations should still work + vsr.setRowCount(42); + assertEquals("Row count should work like before", 42, vsr.getRowCount()); + + IntVector idVector = (IntVector) vsr.getVector("id"); + assertNotNull("Should get vector like before", idVector); + + assertEquals("State should be predictable", VSRState.ACTIVE, vsr.getState()); + + // New pattern: explicit state transitions + vsr.moveToFrozen(); + + // Export should work (this is the key integration point) + try (ArrowExport export = vsr.exportToArrow()) { + assertNotNull("Export should work for Rust integration", export); + } + + // Cleanup should work + pool.completeVSR(vsr); + assertEquals("Should be cleaned up", VSRState.CLOSED, vsr.getState()); + + pool.close(); + } + + public void testConcurrentVSROperationsOnDifferentInstances() { + // Test that different VSR instances can be operated on concurrently + // without interference (even though they're no longer thread-safe individually) + VSRPool pool = new VSRPool(poolId, testSchema, bufferPool); + + // Create multiple independent VSRs + BufferAllocator allocator1 = bufferPool.createChildAllocator("concurrent-vsr-1"); + BufferAllocator allocator2 = bufferPool.createChildAllocator("concurrent-vsr-2"); + BufferAllocator allocator3 = bufferPool.createChildAllocator("concurrent-vsr-3"); + + ManagedVSR vsr1 = new ManagedVSR("concurrent-vsr-1", testSchema, allocator1); + ManagedVSR vsr2 = new ManagedVSR("concurrent-vsr-2", testSchema, allocator2); + ManagedVSR vsr3 = new ManagedVSR("concurrent-vsr-3", testSchema, allocator3); + + // Perform different operations on each VSR + vsr1.setRowCount(10); + vsr2.setRowCount(20); + vsr3.setRowCount(30); + + // Move them to different states independently + vsr1.moveToFrozen(); + vsr2.moveToFrozen(); + // Keep vsr3 active + + // Verify independent state + assertEquals("VSR1 should be FROZEN", VSRState.FROZEN, vsr1.getState()); + assertEquals("VSR2 should be FROZEN", VSRState.FROZEN, vsr2.getState()); + assertEquals("VSR3 should be ACTIVE", VSRState.ACTIVE, vsr3.getState()); + + // Export from frozen VSRs + try (ArrowExport export1 = vsr1.exportToArrow()) { + assertNotNull("VSR1 should export", export1); + } + + try (ArrowExport export2 = vsr2.exportToArrow()) { + assertNotNull("VSR2 should export", export2); + } + + // Modify active VSR + vsr3.setRowCount(35); + assertEquals("VSR3 should be modifiable", 35, vsr3.getRowCount()); + + // Clean up all VSRs + pool.completeVSR(vsr1); + pool.completeVSR(vsr2); + + vsr3.moveToFrozen(); + pool.completeVSR(vsr3); + + // Verify all closed + assertEquals("VSR1 should be CLOSED", VSRState.CLOSED, vsr1.getState()); + assertEquals("VSR2 should be CLOSED", VSRState.CLOSED, vsr2.getState()); + assertEquals("VSR3 should be CLOSED", VSRState.CLOSED, vsr3.getState()); + + // Must freeze pool's active VSR before closing pool + ManagedVSR poolActiveVSR = pool.getActiveVSR(); + poolActiveVSR.moveToFrozen(); + pool.close(); + } +} diff --git a/modules/parquet-data-format/src/test/java/com/parquet/parquetdataformat/vsr/VSRManagerTests.java b/modules/parquet-data-format/src/test/java/com/parquet/parquetdataformat/vsr/VSRManagerTests.java new file mode 100644 index 0000000000000..046cc94f5433d --- /dev/null +++ b/modules/parquet-data-format/src/test/java/com/parquet/parquetdataformat/vsr/VSRManagerTests.java @@ -0,0 +1,379 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package com.parquet.parquetdataformat.vsr; + +import com.parquet.parquetdataformat.bridge.ArrowExport; +import com.parquet.parquetdataformat.bridge.RustBridge; +import com.parquet.parquetdataformat.memory.ArrowBufferPool; +import com.parquet.parquetdataformat.writer.ParquetDocumentInput; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.apache.arrow.vector.types.pojo.Schema; +import org.apache.arrow.vector.types.Types; +import org.opensearch.index.engine.exec.FlushIn; +import org.opensearch.index.engine.exec.WriteResult; +import org.opensearch.index.mapper.MappedFieldType; +import org.opensearch.test.OpenSearchTestCase; +import org.opensearch.common.settings.Settings; +import org.mockito.MockedStatic; +import org.mockito.Mockito; + +import java.util.Arrays; +import java.util.HashMap; +import java.util.Map; + +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.ArgumentMatchers.anyString; + +/** + * Integration tests for VSRManager covering document processing workflows and state management + * through the VSRManager layer rather than direct ManagedVSR manipulation. + */ +public class VSRManagerTests extends OpenSearchTestCase { + + private BufferAllocator allocator; + private ArrowBufferPool bufferPool; + private Schema testSchema; + private String testFileName; + + @Override + public void setUp() throws Exception { + super.setUp(); + allocator = new RootAllocator(); + bufferPool = new ArrowBufferPool(Settings.EMPTY); + + // Create a simple test schema + Field idField = new Field("id", FieldType.nullable(Types.MinorType.INT.getType()), null); + Field nameField = new Field("name", FieldType.nullable(Types.MinorType.VARCHAR.getType()), null); + testSchema = new Schema(Arrays.asList(idField, nameField)); + + testFileName = "test-file-" + System.currentTimeMillis() + ".parquet"; + } + + @Override + public void tearDown() throws Exception { + if (bufferPool != null) { + bufferPool.close(); + } + if (allocator != null) { + allocator.close(); + } + super.tearDown(); + } + + // ===== VSRManager Integration Tests ===== + + public void testVSRManagerInitializationAndActiveVSR() throws Exception { + // Test VSRManager initialization through constructor + VSRManager vsrManager = new VSRManager(testFileName, testSchema, bufferPool); + + // VSRManager should have an active VSR + assertNotNull("VSRManager should have active VSR", vsrManager.getActiveManagedVSR()); + assertEquals("Active VSR should be in ACTIVE state", VSRState.ACTIVE, vsrManager.getActiveManagedVSR().getState()); + + // VSR should start with 0 rows + assertEquals("Active VSR should start with 0 rows", 0, vsrManager.getActiveManagedVSR().getRowCount()); + + // Follow proper VSRManager lifecycle: Write → Flush → Close + // Since we haven't written data, simulate minimal data for flush + vsrManager.getActiveManagedVSR().setRowCount(1); + + // Flush before close (transitions VSR to FROZEN) + FlushIn flushIn = Mockito.mock(FlushIn.class); + String flushResult = vsrManager.flush(flushIn); + assertEquals("Flush should return filename", testFileName, flushResult); + assertEquals("VSR should be FROZEN after flush", VSRState.FROZEN, vsrManager.getActiveManagedVSR().getState()); + + // Now close should succeed + vsrManager.close(); + } + + public void testDocumentAdditionThroughVSRManager() throws Exception { + // Test document addition through VSRManager.addToManagedVSR() + VSRManager vsrManager = new VSRManager(testFileName, testSchema, bufferPool); + + // Create a document to add + ParquetDocumentInput document = new ParquetDocumentInput(vsrManager.getActiveManagedVSR()); + + // Create mock field types and add fields to document + MappedFieldType idFieldType = Mockito.mock(MappedFieldType.class); + Mockito.when(idFieldType.typeName()).thenReturn("integer"); + document.addField(idFieldType, 42); + + MappedFieldType nameFieldType = Mockito.mock(MappedFieldType.class); + Mockito.when(nameFieldType.typeName()).thenReturn("keyword"); + document.addField(nameFieldType, "test-document"); + + // Add document through VSRManager + WriteResult result = vsrManager.addToManagedVSR(document); + assertNotNull("Write result should not be null", result); + + // VSR should still be ACTIVE after document addition + assertEquals("VSR should remain ACTIVE after document addition", + VSRState.ACTIVE, vsrManager.getActiveManagedVSR().getState()); + + // Follow proper VSRManager lifecycle: Write → Flush → Close + // Flush before close (transitions VSR to FROZEN) + FlushIn flushIn = Mockito.mock(FlushIn.class); + String flushResult = vsrManager.flush(flushIn); + assertEquals("Flush should return filename", testFileName, flushResult); + assertEquals("VSR should be FROZEN after flush", VSRState.FROZEN, vsrManager.getActiveManagedVSR().getState()); + + // Now close should succeed + vsrManager.close(); + } + + public void testFlushThroughVSRManager() throws Exception { + // Test flush workflow through VSRManager.flush() + VSRManager vsrManager = new VSRManager(testFileName, testSchema, bufferPool); + + // Add some data first + vsrManager.getActiveManagedVSR().setRowCount(10); // Simulate data addition + + // Flush through VSRManager (create mock FlushIn) + FlushIn flushIn = Mockito.mock(FlushIn.class); + String result = vsrManager.flush(flushIn); + + assertEquals("Flush should return filename", testFileName, result); + + // VSR should be FROZEN after flush + assertEquals("VSR should be FROZEN after flush", + VSRState.FROZEN, vsrManager.getActiveManagedVSR().getState()); + + vsrManager.close(); + } + + public void testVSRManagerStateTransitionWorkflow() throws Exception { + // Test the complete workflow: create -> add data -> flush -> close + VSRManager vsrManager = new VSRManager(testFileName, testSchema, bufferPool); + + // 1. Initial state - VSR should be ACTIVE + assertEquals("Initial VSR should be ACTIVE", VSRState.ACTIVE, vsrManager.getActiveManagedVSR().getState()); + + // 2. Add data (simulate document processing) + vsrManager.getActiveManagedVSR().setRowCount(5); + assertEquals("VSR should have data", 5, vsrManager.getActiveManagedVSR().getRowCount()); + + // 3. Flush - should transition VSR to FROZEN + FlushIn flushIn = Mockito.mock(FlushIn.class); + String flushResult = vsrManager.flush(flushIn); + + assertEquals("Flush should return filename", testFileName, flushResult); + assertEquals("VSR should be FROZEN after flush", VSRState.FROZEN, vsrManager.getActiveManagedVSR().getState()); + assertTrue("VSR should be immutable when frozen", vsrManager.getActiveManagedVSR().isImmutable()); + + // 4. Close - cleanup + vsrManager.close(); + } + + // ===== Integration with VSRPool Pattern Tests ===== + + public void testVSRLifecycleIntegrationPattern() { + // Test the integration pattern between VSRManager and VSRPool + String vsrId = "test-vsr-integration-" + System.currentTimeMillis(); + ManagedVSR managedVSR = new ManagedVSR(vsrId, testSchema, allocator); + + // 1. VSRPool creates active VSR + assertEquals("VSR starts ACTIVE", VSRState.ACTIVE, managedVSR.getState()); + + // 2. VSRManager populates data + managedVSR.setRowCount(15); + + // 3. VSRPool/VSRManager decides to freeze for flushing + managedVSR.moveToFrozen(); + assertTrue("VSR should be immutable after freezing", managedVSR.isImmutable()); + + // 4. VSRManager exports for Rust processing + try (ArrowExport export = managedVSR.exportToArrow()) { + assertNotNull("Export should succeed", export); + } + + // 5. After processing, resources are cleaned up + managedVSR.close(); + assertEquals("VSR should be CLOSED", VSRState.CLOSED, managedVSR.getState()); + } + + public void testMultipleVSRsWithDifferentStates() { + // Test managing multiple VSRs in different states (as VSRManager might do) + String vsrId1 = "test-vsr-1-" + System.currentTimeMillis(); + String vsrId2 = "test-vsr-2-" + System.currentTimeMillis(); + + // Use child allocators instead of new RootAllocators to avoid memory leaks + BufferAllocator childAllocator1 = allocator.newChildAllocator("vsr1", 0, Long.MAX_VALUE); + BufferAllocator childAllocator2 = allocator.newChildAllocator("vsr2", 0, Long.MAX_VALUE); + + ManagedVSR activeVSR = new ManagedVSR(vsrId1, testSchema, childAllocator1); + ManagedVSR frozenVSR = new ManagedVSR(vsrId2, testSchema, childAllocator2); + + // Set up different states + activeVSR.setRowCount(5); + + frozenVSR.setRowCount(10); + frozenVSR.moveToFrozen(); + + // Verify states + assertEquals("First VSR should be ACTIVE", VSRState.ACTIVE, activeVSR.getState()); + assertEquals("Second VSR should be FROZEN", VSRState.FROZEN, frozenVSR.getState()); + + // Verify operations work correctly for each state + activeVSR.setRowCount(7); // Should work + + IllegalStateException exception = expectThrows(IllegalStateException.class, () -> { + frozenVSR.setRowCount(12); // Should fail + }); + assertTrue("Should not allow modification of frozen VSR", + exception.getMessage().contains("Cannot modify VSR in state: FROZEN")); + + // Export should only work for frozen VSR + expectThrows(IllegalStateException.class, activeVSR::exportToArrow); + + try (ArrowExport export = frozenVSR.exportToArrow()) { + assertNotNull("Frozen VSR export should work", export); + } + + // Clean up - must freeze active VSR before closing + activeVSR.moveToFrozen(); + activeVSR.close(); + frozenVSR.close(); + } + + // ===== Error Handling Tests ===== + + public void testVSRManagerCloseWithoutFlushFails() throws Exception { + // Test that VSRManager.close() fails when VSRs are still in ACTIVE state (not flushed) + VSRManager vsrManager = new VSRManager(testFileName, testSchema, bufferPool); + + // Get active VSR and add some data + assertEquals("VSR should be ACTIVE", VSRState.ACTIVE, vsrManager.getActiveManagedVSR().getState()); + vsrManager.getActiveManagedVSR().setRowCount(5); // Simulate data addition + + // Try to close without flushing - should fail + IllegalStateException exception = expectThrows(IllegalStateException.class, vsrManager::close); + + // Verify the error message mentions the VSR is still ACTIVE + assertTrue("Should mention VSR is still ACTIVE", + exception.getMessage().contains("VSR is still ACTIVE")); + assertTrue("Should mention must freeze first", + exception.getMessage().contains("Must freeze VSR before closing")); + + // VSR should still be in ACTIVE state after failed close + assertEquals("VSR should still be ACTIVE", VSRState.ACTIVE, vsrManager.getActiveManagedVSR().getState()); + + // Proper cleanup: flush then close + FlushIn flushIn = Mockito.mock(FlushIn.class); + vsrManager.flush(flushIn); + assertEquals("VSR should be FROZEN after flush", VSRState.FROZEN, vsrManager.getActiveManagedVSR().getState()); + vsrManager.close(); // Should succeed now + } + + public void testVSRManagerCloseEmptyButUnflushedFails() throws Exception { + // Test that even an empty VSRManager must be flushed before closing + VSRManager vsrManager = new VSRManager(testFileName, testSchema, bufferPool); + + // Get active VSR (no data added, but still ACTIVE) + assertEquals("VSR should be ACTIVE", VSRState.ACTIVE, vsrManager.getActiveManagedVSR().getState()); + assertEquals("VSR should have 0 rows", 0, vsrManager.getActiveManagedVSR().getRowCount()); + + // Try to close without flushing - should fail even with no data + IllegalStateException exception = expectThrows(IllegalStateException.class, vsrManager::close); + + assertTrue("Should mention VSR is still ACTIVE", + exception.getMessage().contains("VSR is still ACTIVE")); + + // Must flush first, even with no data + vsrManager.getActiveManagedVSR().setRowCount(1); // Need minimal data for flush to work + FlushIn flushIn = Mockito.mock(FlushIn.class); + vsrManager.flush(flushIn); + vsrManager.close(); // Should succeed now + } + + public void testInvalidStateTransitionHandling() { + // Test error handling for invalid state transitions that VSRManager might encounter + String vsrId = "test-vsr-error-" + System.currentTimeMillis(); + ManagedVSR managedVSR = new ManagedVSR(vsrId, testSchema, allocator); + + // Try to close active VSR (should fail - must freeze first) + IllegalStateException exception = expectThrows(IllegalStateException.class, managedVSR::close); + + assertTrue("Should mention VSR is still ACTIVE", + exception.getMessage().contains("VSR is still ACTIVE")); + assertTrue("Should mention must freeze first", + exception.getMessage().contains("Must freeze VSR before closing")); + + // VSR should still be in ACTIVE state after failed close + assertEquals("VSR should still be ACTIVE", VSRState.ACTIVE, managedVSR.getState()); + + // Proper cleanup + managedVSR.moveToFrozen(); + managedVSR.close(); + } + + // ===== Mock-based Rust Integration Tests ===== + + public void testRustBridgeIntegrationPattern() { + // Test the pattern of integration with RustBridge (without actually calling native code) + String vsrId = "test-vsr-rust-" + System.currentTimeMillis(); + ManagedVSR managedVSR = new ManagedVSR(vsrId, testSchema, allocator); + + // Populate VSR as VSRManager would + managedVSR.setRowCount(20); + + // Freeze before export (as VSRManager does) + managedVSR.moveToFrozen(); + + // Export for Rust (simulating VSRManager calling RustBridge.write) + try (ArrowExport export = managedVSR.exportToArrow()) { + assertNotNull("Export should be ready for Rust", export); + + // Verify addresses are valid (would be passed to RustBridge.write) + assertTrue("Array address should be valid for Rust", export.getArrayAddress() != 0); + assertTrue("Schema address should be valid for Rust", export.getSchemaAddress() != 0); + + // In real VSRManager, this would be: + // RustBridge.write(fileName, export.getArrayAddress(), export.getSchemaAddress()); + } + + // After Rust processing, clean up + managedVSR.close(); + } + + // ===== Resource Management Tests ===== + + public void testResourceCleanupAfterStateTransitions() { + // Test that resources are properly managed through state transitions + String vsrId = "test-vsr-cleanup-" + System.currentTimeMillis(); + ManagedVSR managedVSR = new ManagedVSR(vsrId, testSchema, allocator); + + // Add data and verify memory allocation + managedVSR.setRowCount(100); + long initialMemory = allocator.getAllocatedMemory(); + assertTrue("Should have allocated memory", initialMemory > 0); + + // Freeze (state change should not affect memory) + managedVSR.moveToFrozen(); + assertTrue("Memory should still be allocated", allocator.getAllocatedMemory() >= initialMemory); + + // Export (should not significantly change memory usage) + try (ArrowExport export = managedVSR.exportToArrow()) { + assertNotNull("Export should work", export); + // Memory might increase slightly for export structures + } + + // Close should clean up resources + managedVSR.close(); + assertEquals("VSR should be CLOSED", VSRState.CLOSED, managedVSR.getState()); + + // Note: Exact memory cleanup testing is difficult without more intrusive testing + // but we verify the state transitions work correctly + } +} diff --git a/modules/parquet-data-format/src/test/java/com/parquet/parquetdataformat/vsr/VSRPoolTests.java b/modules/parquet-data-format/src/test/java/com/parquet/parquetdataformat/vsr/VSRPoolTests.java new file mode 100644 index 0000000000000..699c6bb87f74f --- /dev/null +++ b/modules/parquet-data-format/src/test/java/com/parquet/parquetdataformat/vsr/VSRPoolTests.java @@ -0,0 +1,329 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package com.parquet.parquetdataformat.vsr; + +import com.parquet.parquetdataformat.memory.ArrowBufferPool; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.apache.arrow.vector.types.pojo.Schema; +import org.apache.arrow.vector.types.Types; +import org.opensearch.test.OpenSearchTestCase; +import org.opensearch.common.settings.Settings; + +import java.util.Arrays; + +/** + * Unit tests for VSRPool covering changes related to the removal of allVSRs tracking, + * simplified resource management, and removal of statistics functionality. + */ +public class VSRPoolTests extends OpenSearchTestCase { + + private BufferAllocator allocator; + private ArrowBufferPool bufferPool; + private Schema testSchema; + private String poolId; + + @Override + public void setUp() throws Exception { + super.setUp(); + allocator = new RootAllocator(); + bufferPool = new ArrowBufferPool(Settings.EMPTY); + + // Create a simple test schema + Field idField = new Field("id", FieldType.nullable(Types.MinorType.INT.getType()), null); + Field nameField = new Field("name", FieldType.nullable(Types.MinorType.VARCHAR.getType()), null); + testSchema = new Schema(Arrays.asList(idField, nameField)); + + poolId = "test-pool-" + System.currentTimeMillis(); + } + + @Override + public void tearDown() throws Exception { + if (bufferPool != null) { + bufferPool.close(); + } + if (allocator != null) { + allocator.close(); + } + super.tearDown(); + } + + // ===== Basic Pool Functionality Tests ===== + + public void testPoolCreationAndInitialization() { + VSRPool pool = new VSRPool(poolId, testSchema, bufferPool); + + assertNotNull("Pool should be created", pool); + + // Pool should start with an active VSR (created during initialization) + ManagedVSR activeVSR = pool.getActiveVSR(); + assertNotNull("Should have active VSR after initialization", activeVSR); + assertEquals("VSR should be ACTIVE", VSRState.ACTIVE, activeVSR.getState()); + + ManagedVSR frozenVSR = pool.takeFrozenVSR(); + assertNull("Should start with no frozen VSR", frozenVSR); + + // Must freeze active VSR before closing pool + activeVSR.moveToFrozen(); + pool.close(); + } + + public void testActiveVSRCreationOnDemand() { + VSRPool pool = new VSRPool(poolId, testSchema, bufferPool); + + // Get active VSR - should be created during initialization + ManagedVSR activeVSR = pool.getActiveVSR(); + assertNotNull("Should have active VSR", activeVSR); + assertEquals("VSR should be ACTIVE", VSRState.ACTIVE, activeVSR.getState()); + assertFalse("VSR should not be immutable", activeVSR.isImmutable()); + + // Getting active VSR again should return the same instance + ManagedVSR sameActiveVSR = pool.getActiveVSR(); + assertSame("Should return same active VSR instance", activeVSR, sameActiveVSR); + + // Must freeze active VSR before closing pool + activeVSR.moveToFrozen(); + pool.close(); + } + + public void testVSRRotationThroughPool() throws Exception { + VSRPool pool = new VSRPool(poolId, testSchema, bufferPool); + + // Get initial active VSR + ManagedVSR activeVSR = pool.getActiveVSR(); + String initialVSRId = activeVSR.getId(); + + // Fill VSR with data to simulate reaching capacity + activeVSR.setRowCount(50000); // Assuming this triggers rotation threshold + + // Test pool rotation mechanism + boolean rotationOccurred = pool.maybeRotateActiveVSR(); + + if (rotationOccurred) { + // After rotation, should have new active VSR + ManagedVSR newActiveVSR = pool.getActiveVSR(); + assertNotNull("Should have new active VSR after rotation", newActiveVSR); + assertEquals("New active VSR should be ACTIVE", VSRState.ACTIVE, newActiveVSR.getState()); + assertEquals("New VSR should have row count 0", 0, newActiveVSR.getRowCount()); + + // Should have frozen VSR available + ManagedVSR frozenVSR = pool.getFrozenVSR(); + if (frozenVSR != null) { + assertEquals("Frozen VSR should be FROZEN", VSRState.FROZEN, frozenVSR.getState()); + assertEquals("Frozen VSR should have expected row count", 50000, frozenVSR.getRowCount()); + assertSame("Frozen VSR should be the same as the previous active VSR", activeVSR, frozenVSR); + + // Complete the frozen VSR + pool.completeVSR(frozenVSR); + assertEquals("Frozen VSR should be CLOSED after completion", VSRState.CLOSED, frozenVSR.getState()); + } + + // Clean up new active VSR + newActiveVSR.moveToFrozen(); + } else { + // No rotation occurred, clean up original VSR + fail("VSR should be rotated"); + } + + pool.close(); + } + + public void testTakeFrozenVSRReturnsBehavior() throws Exception { + VSRPool pool = new VSRPool(poolId, testSchema, bufferPool); + + // Initially should have no frozen VSR + ManagedVSR frozenVSR = pool.takeFrozenVSR(); + assertNull("Should initially have no frozen VSR", frozenVSR); + + // Test through pool rotation to create a frozen VSR + ManagedVSR activeVSR = pool.getActiveVSR(); + activeVSR.setRowCount(50000); // Fill to trigger rotation + + boolean rotated = pool.maybeRotateActiveVSR(); + assertTrue("Frozen VSR should be rotated", rotated); + + // Should now have a frozen VSR + ManagedVSR actualFrozenVSR = pool.takeFrozenVSR(); + if (actualFrozenVSR != null) { + assertEquals("Taken VSR should be FROZEN", VSRState.FROZEN, actualFrozenVSR.getState()); + + // Taking it again should return null (slot cleared) + ManagedVSR shouldBeNull = pool.takeFrozenVSR(); + assertNull("Should be null after taking frozen VSR", shouldBeNull); + + // Clean up the taken VSR + pool.completeVSR(actualFrozenVSR); + } + + // Must freeze pool's active VSR before closing pool + pool.getActiveVSR().moveToFrozen(); + pool.close(); + } + + // ===== Resource Management Tests ===== + + public void testCompleteVSRThroughPool() throws Exception { + VSRPool pool = new VSRPool(poolId, testSchema, bufferPool); + + // Get active VSR and fill it + ManagedVSR activeVSR = pool.getActiveVSR(); + activeVSR.setRowCount(20); + activeVSR.moveToFrozen(); // Manually freeze for this test + + // Test the completeVSR functionality through pool + pool.completeVSR(activeVSR); + + // VSR should be closed + assertEquals("VSR should be CLOSED after completion", VSRState.CLOSED, activeVSR.getState()); + + pool.close(); + } + + public void testPoolCloseResourceCleanup() { + VSRPool pool = new VSRPool(poolId, testSchema, bufferPool); + + // Get active VSR to trigger creation + ManagedVSR activeVSR = pool.getActiveVSR(); + assertNotNull("Should have active VSR", activeVSR); + + // Add some data + activeVSR.setRowCount(5); + + // Must freeze active VSR before closing pool + activeVSR.moveToFrozen(); + + // Close pool - should clean up frozen VSR + pool.close(); + } + + // ===== Error Handling Tests ===== + + public void testVSRCreationWithInvalidParameters() { + // Test error handling in VSR creation within pool context + + // Null schema should fail - but the error is wrapped in RuntimeException + RuntimeException schemaException = expectThrows(RuntimeException.class, () -> { + new VSRPool(poolId, null, bufferPool); + }); + assertTrue("Should mention failed to create VSR", + schemaException.getMessage().contains("Failed to create new VSR")); + assertTrue("Root cause should be NullPointerException", + schemaException.getCause() instanceof NullPointerException); + + // Null buffer pool should fail - but the error is also wrapped in RuntimeException + RuntimeException bufferPoolException = expectThrows(RuntimeException.class, () -> { + new VSRPool(poolId, testSchema, null); + }); + assertTrue("Should mention failed to create VSR", + bufferPoolException.getMessage().contains("Failed to create new VSR")); + assertTrue("Root cause should be NullPointerException", + bufferPoolException.getCause() instanceof NullPointerException); + } + + // ===== Integration Pattern Tests ===== + + public void testPoolVSRManagerIntegrationPattern() { + // Test the integration pattern between pool and manager + VSRPool pool = new VSRPool(poolId, testSchema, bufferPool); + + // 1. Get active VSR (as VSRManager would) + ManagedVSR activeVSR = pool.getActiveVSR(); + assertNotNull("Manager should get active VSR", activeVSR); + assertEquals("VSR should be ACTIVE", VSRState.ACTIVE, activeVSR.getState()); + + // 2. Populate VSR (as VSRManager would) + activeVSR.setRowCount(25); + assertEquals("Should have expected row count", 25, activeVSR.getRowCount()); + + // 3. When ready to flush, freeze VSR (as VSRManager would do) + activeVSR.moveToFrozen(); + assertEquals("VSR should be FROZEN", VSRState.FROZEN, activeVSR.getState()); + + // 4. After processing, complete VSR (as VSRManager would) + pool.completeVSR(activeVSR); + assertEquals("VSR should be CLOSED after completion", VSRState.CLOSED, activeVSR.getState()); + + pool.close(); + } + + public void testMultipleVSRLifecycleInPool() { + // Test managing multiple VSRs through their lifecycle + VSRPool pool = new VSRPool(poolId, testSchema, bufferPool); + + // Create multiple VSRs simulating different stages of lifecycle + BufferAllocator childAllocator1 = bufferPool.createChildAllocator("lifecycle-vsr-1"); + BufferAllocator childAllocator2 = bufferPool.createChildAllocator("lifecycle-vsr-2"); + BufferAllocator childAllocator3 = bufferPool.createChildAllocator("lifecycle-vsr-3"); + ManagedVSR vsr1 = new ManagedVSR("lifecycle-vsr-1", testSchema, childAllocator1); + ManagedVSR vsr2 = new ManagedVSR("lifecycle-vsr-2", testSchema, childAllocator2); + ManagedVSR vsr3 = new ManagedVSR("lifecycle-vsr-3", testSchema, childAllocator3); + + // Put them in different states + vsr1.setRowCount(10); // Keep ACTIVE + + vsr2.setRowCount(20); + vsr2.moveToFrozen(); // Make FROZEN + + vsr3.setRowCount(30); + vsr3.moveToFrozen(); + vsr3.close(); // Make CLOSED + + // Verify states + assertEquals("VSR1 should be ACTIVE", VSRState.ACTIVE, vsr1.getState()); + assertEquals("VSR2 should be FROZEN", VSRState.FROZEN, vsr2.getState()); + assertEquals("VSR3 should be CLOSED", VSRState.CLOSED, vsr3.getState()); + + // Complete the active and frozen ones + vsr1.moveToFrozen(); + pool.completeVSR(vsr1); + pool.completeVSR(vsr2); + + // All should be closed now + assertEquals("VSR1 should be CLOSED", VSRState.CLOSED, vsr1.getState()); + assertEquals("VSR2 should be CLOSED", VSRState.CLOSED, vsr2.getState()); + assertEquals("VSR3 should remain CLOSED", VSRState.CLOSED, vsr3.getState()); + + // Must freeze pool's active VSR before closing pool + ManagedVSR poolActiveVSR = pool.getActiveVSR(); + poolActiveVSR.moveToFrozen(); + pool.close(); + } + + // ===== Memory and Performance Tests ===== + + public void testMemoryManagementInPool() { + // Test memory allocation and cleanup behavior + VSRPool pool = new VSRPool(poolId, testSchema, bufferPool); + + long initialMemory = bufferPool.getTotalAllocatedBytes(); + + // Create active VSR - should allocate memory + ManagedVSR activeVSR = pool.getActiveVSR(); + activeVSR.setRowCount(50); + + long afterCreationMemory = bufferPool.getTotalAllocatedBytes(); + // Note: Memory allocation may be delayed or managed differently + // Just ensure operations complete without error + assertTrue("Memory operations should complete", afterCreationMemory >= initialMemory); + + // Must freeze active VSR before closing pool + activeVSR.moveToFrozen(); + + // Close pool - should clean up + pool.close(); + + // Memory cleanup verification is difficult without intrusive testing, + // but we verify operations complete without error + assertTrue("Test completed successfully", true); + } +}