ParquetDataFormat plugin - Fix document corruption bug due to dirty writes from rolled back document in VSR

[rayshrey]

Description

[Describe what this change achieves]

Related Issues

Resolves #22417

Check List

• Functionality includes testing.
• API changes companion pull request created, if applicable.
• Public documentation issue/PR created, if applicable.

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check here.

[github-actions]

PR Reviewer Guide 🔍

(Review updated until commit 77e23fc)

Here are some key observations to aid the review process:

🧪 PR contains tests

🔒 No security concerns identified

📝 TODO sections sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/vsr/VSRManager.java [235]: Remove this check after validation on mapping update is in place.

🔀 No multiple PR themes

⚡ Recommended focus areas for review Incomplete Scrub for Multi-Vector Fields scrubPartialRow clears only the top-level FieldVector returned by activeVSR.getVector(name) for each written field. For composite Arrow types produced by some ParquetField implementations (e.g. list/struct vectors, or fields that populate additional child/aux vectors beyond the parent), parquetField.createField may write into child vectors whose validity bits are not cleared by calling setNull on the parent. This can leave stale data in child vectors when the reused row is later written by a document that doesn't set that field. Confidence is limited without visibility into all ParquetField implementations, but the potential impact (silent data corruption — the exact bug this PR fixes) warrants verification that only simple fixed/variable-width vectors are produced, or extending the scrub to walk child vectors. private void scrubPartialRow(ParquetDocumentInput doc, ManagedVSR activeVSR, int rowIndex, int writtenFields, boolean rowIdWritten) { List<FieldValuePair> fields = doc.getFinalInput(); for (int i = 0; i < writtenFields; i++) { scrubVector(activeVSR.getVector(fields.get(i).getFieldType().name()), rowIndex); } if (rowIdWritten) { scrubVector(activeVSR.getVector(DocumentInput.ROW_ID_FIELD), rowIndex); } } Fallback May Not Clear Value Data The default branch in setNull only calls BitVectorHelper.unsetBit on the validity buffer. For any future non-fixed/non-variable-width vector type (e.g. union, list, struct, view vectors), this clears the parent validity bit but leaves value/offset/child buffers with stale content. If such a vector's isNull(index) implementation checks anything beyond the top-level validity bit, or if downstream readers read the value buffer directly, the stale value could still leak. Consider either explicitly enumerating supported vector types and throwing on unknown, or delegating to the vector's own setNull where available. private static void setNull(FieldVector vector, int index) { switch (vector) { case BaseFixedWidthVector fixed -> fixed.setNull(index); case BaseVariableWidthVector variable -> variable.setNull(index); case BaseLargeVariableWidthVector large -> large.setNull(index); default -> BitVectorHelper.unsetBit(vector.getValidityBuffer(), index); } }

[github-actions]

PR Code Suggestions ✨

Latest suggestions up to 77e23fc

Explore these optional code suggestions:

Category	Suggestion	Impact
Possible issue	Broaden catch to include Errors/OOM Catching Exception will also catch InterruptedException and swallow its interrupt semantics, and it silently narrows the declared throws IOException. Consider catching Throwable (or at minimum RuntimeException | IOException | Error) to also handle OutOfMemoryError/OutOfMemoryException from Arrow, which is the very scenario this fix targets. Otherwise an Arrow OOM (an Error or unchecked type outside Exception) will bypass the scrub and still leak partial state. sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/vsr/VSRManager.java [263-270] -} catch (Exception e) { - // Any failure between the first field write and acceptance leaves an uncounted partial - // row. Scrub it (best-effort, never throws) so no stale value can leak into the next doc - // that reuses this row index, then rethrow the original failure unchanged. Precise - // rethrow keeps addDocument's throws clause unchanged. +} catch (Throwable e) { scrubPartialRow(doc, activeVSR, rowIndex, writtenFields, rowIdWritten); throw e; } Suggestion importance[1-10]: 7 __ Why: Valid concern: Arrow's OutOfMemoryException extends RuntimeException (so is caught by Exception), but true Errors like OutOfMemoryError would bypass the scrub. Broadening to Throwable would make the fix more robust, though the primary target (Arrow OOM) is already caught.	Medium
General	Guard row index invariant rowIndex is captured before the loop but the rowId vector write and setRowCount use rowIndex + 1 implicitly via activeVSR.getRowCount() elsewhere; confirm the scrub uses the same rowIndex the writes actually targeted. If any parquetField.createField internally advances or reads activeVSR.getRowCount(), the pre-loop capture may diverge from the actual slot written. Consider asserting the invariant or having createField accept an explicit rowIndex. sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/vsr/VSRManager.java [219] final int rowIndex = activeVSR.getRowCount(); +assert rowIndex == acceptedRows : "rowIndex/acceptedRows drift before addDocument row write"; Suggestion importance[1-10]: 2 __ Why: The suggestion is speculative; createField writes via setSafe at a specific index without advancing row count, and the invariant is already checked earlier. Adding an assert provides marginal value.	Low

---

Previous suggestions

Suggestions up to commit e04ff29

Category	Suggestion	Impact
Possible issue	Track vector before write to ensure scrubbing writtenVectors.add(vector) runs after parquetField.createField(...), so if createField throws mid-write the vector that was being written is not tracked and therefore not scrubbed, leaving its partial value at rowIndex. Add the vector to writtenVectors before invoking createField so any partial write is guaranteed to be cleared on failure. sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/vsr/VSRManager.java [250-251] for (FieldValuePair pair : doc.getFinalInput()) { MappedFieldType fieldType = pair.getFieldType(); ParquetField parquetField = ArrowFieldRegistry.getParquetField(fieldType.typeName()); if (parquetField == null) { ... } FieldVector vector = activeVSR.getVector(fieldType.name()); if (vector == null) { ... } + writtenVectors.add(vector); parquetField.createField(fieldType, activeVSR, pair.getValue()); - writtenVectors.add(vector); } Suggestion importance[1-10]: 8 __ Why: Valid correctness concern: if createField throws partway through writing a vector's row, that vector is not added to writtenVectors and thus not scrubbed, potentially leaving a partial write — exactly the bug this PR aims to fix. Adding the vector before the write ensures scrubbing covers this case.	Medium

Suggestions up to commit 7a2e797

Category	Suggestion	Impact
General	Track rowId vector before setSafe call setSafe on the rowId vector can itself trigger a reallocation and throw an Arrow OOM. If that happens after the field loop succeeded, the field vectors have already been written for rowIndex but writtenVectors has been populated, so the scrub path will still run — good — but ensure rowIdVector is added to writtenVectors before calling setSafe so a failure inside setSafe still scrubs the rowId slot's validity bit if it was partially set. sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/vsr/VSRManager.java [253-257] BigIntVector rowIdVector = (BigIntVector) activeVSR.getVector(DocumentInput.ROW_ID_FIELD); if (rowIdVector != null) { + writtenVectors.add(rowIdVector); rowIdVector.setSafe(rowIndex, doc.getRowId()); - writtenVectors.add(rowIdVector); } activeVSR.setRowCount(rowIndex + 1); acceptedRows++; Suggestion importance[1-10]: 5 __ Why: Reordering writtenVectors.add(rowIdVector) before setSafe is a reasonable defensive change to ensure scrub covers the rowId slot if setSafe itself OOMs mid-allocation. Minor but valid improvement.	Low
Possible issue	Ensure scrub covers all written vectors The vector obtained via activeVSR.getVector(fieldType.name()) may not be the same instance that parquetField.createField actually writes to (e.g., if createField looks up the vector by a different name/alias or writes to multiple vectors). Consider tracking the actual vector(s) written by createField (or scrubbing all vectors in the VSR at rowIndex) to ensure the scrub covers every partially-written column. sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/vsr/VSRManager.java [236-251] FieldVector vector = activeVSR.getVector(fieldType.name()); if (vector == null) { logger.error( "[Gen: {}] VSR schema mismatch: field [{}] not in active VSR. VSR schema fields: {}", writerGeneration, fieldType.name(), activeVSR.getSchema().getFields().stream().map(f -> f.getName()).collect(java.util.stream.Collectors.joining(", ")) ); throw new MismatchedInputException( "Active VSR has no vector for field [" + fieldType.name() + "] — schema reconciliation must run via updateMappingVersion before addDocument" ); } parquetField.createField(fieldType, activeVSR, pair.getValue()); +// Track all vectors in the VSR to ensure scrub covers any vector createField may have touched. writtenVectors.add(vector); Suggestion importance[1-10]: 3 __ Why: The suggestion raises a theoretical concern but the improved_code is identical to the existing_code (only a comment is added), providing no actual code change. The concern about createField writing to different vectors is speculative without evidence.	Low

[github-actions]

✅ Gradle check result for 7a2e797: SUCCESS

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 71.53%. Comparing base (599785a) to head (77e23fc).
⚠️ Report is 8 commits behind head on main.

Additional details and impacted files

@@             Coverage Diff              @@
##               main   #22482      +/-   ##
============================================
- Coverage     71.54%   71.53%   -0.02%     
+ Complexity    77023    77022       -1     
============================================
  Files          6153     6156       +3     
  Lines        358354   358410      +56     
  Branches      52237    52243       +6     
============================================
- Hits         256399   256383      -16     
- Misses        81586    81636      +50     
- Partials      20369    20391      +22     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[github-actions]

Persistent review updated to latest commit e04ff29

[github-actions]

✅ Gradle check result for e04ff29: SUCCESS

[mgodwan]

instead of always collecting writtenVectors, can we not just not re iterate only on failure?

[github-actions]

Persistent review updated to latest commit 77e23fc

[github-actions]

✅ Gradle check result for 77e23fc: SUCCESS