peter-toth commented on code in PR #59639:
URL: https://github.com/apache/spark/pull/59639#discussion_r4339163205
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/LateMaterializationParquetRecordReader.java:
##########
@@ -888,30 +972,121 @@ private void reportOnce(Throwable t, String message) {
}
/**
- * Gives up splicing for the row group being evaluated and frees its
buffered survivors. Phase 2
- * then reads every projected column, so the rows are unaffected.
+ * Gives up splicing for the row group being read and frees its buffered
survivors, but not the
+ * published ones. Phase 2, or the key re-read after a spill, then reads the
key columns, so the
+ * rows are unaffected.
*/
private void abandonSplicing() {
- for (WritableColumnVector[] batch : survivorBatches) {
- closeAll(batch);
+ for (SurvivorBatch batch : survivorBatches) {
+ closeAll(batch.vectors());
}
survivorBatches.clear();
closeAll(currentKeyAccumulators);
currentKeyAccumulators = null;
currentKeyAccumulatorRowCount = 0;
+ currentKeyAccumulatorCharge = 0L;
spliceCurrentRowGroup = false;
+ releaseSurvivorMemoryTo(heldCharge());
}
/** Closes every survivor vector this reader still holds, the published ones
included. */
private void closeSplicingState() {
try {
closeAll(publishedKeyVectors);
publishedKeyVectors = null;
+ publishedKeyCharge = 0L;
} finally {
abandonSplicing();
}
}
+ /**
+ * Frees the queued survivors of the row group being emitted, and returns
the accounted bytes this
+ * released. The next batch then reads the rest of the row group's keys
again. It acts only on the
+ * reading thread, since the reader takes no lock. There it can come only
between two batches.
+ */
+ private long spillQueuedSurvivors() {
+ if (Thread.currentThread() != readerThread || !spillableRowGroup ||
!spliceCurrentRowGroup
+ || survivorBatches.isEmpty()) {
+ return 0L;
+ }
+ long before = survivorMemory.getUsed();
+ abandonSplicing();
+ rereadRemainingKeys = true;
+ return before - survivorMemory.getUsed();
+ }
+
+ /**
+ * Reads the keys of the rest of the loaded row group again, after a spill
freed its queued
+ * survivors, and points the key slots' column readers at them. The non-key
readers go on where
+ * they are, so both read the same rows.
+ */
+ private void rereadRemainingKeys() throws IOException {
+ rereadRemainingKeys = false;
+ long remainingRows = loadedRowCount - rowsEmittedInRowGroup;
+ if (remainingRows <= 0) return;
+ PrimitiveIterator.OfLong rows = loadedRanges.iterator();
+ for (long i = 0; i < rowsEmittedInRowGroup; i++) {
+ rows.nextLong();
+ }
+ RowRanges.Builder remaining = RowRanges.builder();
+ while (rows.hasNext()) {
+ remaining.addSelectedRow(rows.nextLong());
+ }
+ RowRanges remainingRanges = remaining.build();
+ fileReader.setRequestedSchema(keyOnlyColumns);
+ rereadKeyPages = fileReader.readFilteredRowGroup(loadedBlockIndex,
remainingRanges);
+ requireRowCount(rereadKeyPages, remainingRows, loadedBlockIndex);
+ long blockRowCount =
fileReader.getRowGroups().get(loadedBlockIndex).getRowCount();
+ RowRanges readerRanges = readerRangesFor(remainingRanges, remainingRows,
blockRowCount);
+ for (KeyColumn key : keyColumns) {
+ replaceColumnReader(key.batchSlot(),
+ newColumnReader(key.descriptor(), key.required(), rereadKeyPages,
readerRanges));
+ }
+ }
+
+ /** The accounted bytes the survivor buffer holds, see {@link
SurvivorMemory}. */
+ private long heldCharge() {
+ long held = rangeCharge;
+ if (accountKeys) {
+ held += publishedKeyCharge + currentKeyAccumulatorCharge;
+ for (SurvivorBatch batch : survivorBatches) {
+ held += batch.charge();
+ }
+ }
+ return held;
+ }
+
+ /**
+ * Grows the reservation to cover {@code held} bytes, asking for at least
what it already holds,
+ * so a growing buffer asks only a few times. False when it is denied.
+ */
+ private boolean reserveSurvivorMemory(long held) {
+ if (survivorMemory == null) return true;
+ long used = survivorMemory.getUsed();
+ long missing = held - used;
+ if (missing <= 0) return true;
+ long granted = survivorMemory.acquireMemory(Math.max(missing, used));
Review Comment:
Done in
[778c14954a9](https://github.com/apache/spark/pull/59639/commits/778c14954a9ede1af0fc6357b0b5546322e69a23),
with your formula. What phase 1 reserved ahead is released before phase 2
reads.
The test now leaves 440 bytes free, twice the 220 the buffer holds. A new
case with a 300-byte cap and 300 bytes free fails without the clamp to the cap.
The description now says the reservation stays below twice what is held while
small, and below 1 MiB more once larger.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
