peter-toth commented on code in PR #59639:
URL: https://github.com/apache/spark/pull/59639#discussion_r4346739571
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/LateMaterializationParquetRecordReader.java:
##########
@@ -912,6 +1023,63 @@ private void closeSplicingState() {
}
}
+ /** The bytes of key values buffered for splicing, which the cap weighs even
off heap. */
+ private long bufferedKeyCharge() {
+ return queuedKeyCharge + currentKeyAccumulatorCharge;
+ }
+
+ /** The accounted bytes the survivor buffer holds, see {@link
SurvivorMemory}. */
+ private long heldCharge() {
+ return accountKeys ? rangeCharge + bufferedKeyCharge() : rangeCharge;
+ }
+
+ /**
+ * Grows the reservation to cover {@code held} bytes, which the caller has
kept within the cap.
+ * It asks for at least what it already holds, up to {@link
#SURVIVOR_MEMORY_STEP_BYTES} and the
+ * cap, so a growing buffer asks only a few times and the reservation never
passes the cap. False
+ * when it is denied.
+ */
+ private boolean reserveSurvivorMemory(long held) {
+ if (survivorMemory == null) return true;
+ long missing = held - reservedBytes;
+ if (missing <= 0) return true;
+ long headroom = storageFilter.maxSplicedRowGroupBytes() - reservedBytes;
+ long ask = Math.max(missing,
+ Math.min(Math.min(reservedBytes, SURVIVOR_MEMORY_STEP_BYTES),
headroom));
+ long granted;
+ try {
+ granted = survivorMemory.acquireMemory(ask);
+ } catch (RuntimeException | InternalError e) {
Review Comment:
Done in
[4b28a3f1a8f](https://github.com/apache/spark/pull/59639/commits/4b28a3f1a8fd03ff47b17333e2349f070a23af69).
The catch is now `Exception | InternalError`, and `memoryReservationFailed`
wraps whatever `rethrowIfMustPropagate` does not rethrow in a checked
`SparkException`. A new case in the error test has the task's memory manager
throw an `IOException`, through a `TaskMemoryManager` subclass, and checks that
`shouldIgnoreCorruptFileException` rejects what comes out. The doc and the
description name the eviction.
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/LateMaterializationParquetRecordReader.java:
##########
@@ -727,22 +783,38 @@ private RowRanges evaluateStorageFilter(
// where a plain read fails.
reportOnce(e, SPLICING_GIVEN_UP_ON_ERROR);
abandonSplicing();
- splicedBytes = 0L;
}
}
- if (splicedBytes + rangeBytes > cap) {
+ if (bufferedKeyCharge() + rangeCharge > cap) {
// The buffer goes first, which is a no-op once splicing is given up.
+ if (spliceCurrentRowGroup) splicingGivenUp =
BudgetGiveUp.SPLICING_ON_CAP;
abandonSplicing();
- splicedBytes = 0L;
// The ranges would be thrown away, so stop.
- if (rangeBytes > cap) return null;
+ if (rangeCharge > cap) {
+ recordBudgetGiveUp(BudgetGiveUp.FILTER_ON_CAP);
+ return null;
+ }
+ }
+ // A reservation the task's memory manager denies is handled as
passing the cap is. With no
+ // key values in it, giving splicing up frees nothing, so the ranges
are denied too.
+ long keyCharge = accountKeys ? bufferedKeyCharge() : 0L;
+ boolean denied = !reserveSurvivorMemory(keyCharge + rangeCharge);
+ if (denied && keyCharge > 0L) {
+ splicingGivenUp = BudgetGiveUp.SPLICING_ON_MEMORY;
+ abandonSplicing();
+ denied = !reserveSurvivorMemory(rangeCharge);
+ }
+ if (denied) {
+ recordBudgetGiveUp(BudgetGiveUp.FILTER_ON_MEMORY);
+ return null;
}
}
remaining -= num;
}
// A partly filled set of accumulators goes to the queue as this row
group's last batch.
if (currentKeyAccumulators != null) pushAccumulators();
+ if (splicingGivenUp != null) recordBudgetGiveUp(splicingGivenUp);
Review Comment:
Done in
[4b28a3f1a8f](https://github.com/apache/spark/pull/59639/commits/4b28a3f1a8fd03ff47b17333e2349f070a23af69).
Phase 1 leaves the give-up in a field, and `loadNextRowGroup` records it only
once phase 2 keeps the filter. A new test reads your example, the damaged file
at a cap of 2,100 bytes, and expects `(2, 0, 0, 0)`. It also checks that the
row group read whole holds nothing while it is emitted.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetStorageFilter.scala:
##########
@@ -295,6 +328,22 @@ object ParquetStorageFilter {
}
}
+ /**
+ * Throws what reserving the survivor buffer's memory raised, in a form
`FileScanRDD` cannot take
+ * for a corrupt file under `ignoreCorruptFiles`. The task's memory manager
runs other consumers'
+ * spills inside the reservation, so it can raise anything. A
`RuntimeException`, which is also
+ * what a kill's interrupt during a spill becomes, goes out wrapped in a
checked
+ * [[SparkException]], as in [[throwIfKilled]]. An `InternalError` goes out
in an `UnknownError`,
+ * and anything else, such as the `SparkOutOfMemoryError` of a failed spill,
as it is, as in
+ * [[rethrowIfMustPropagate]]. Typed as returning a `RuntimeException` so
Java code can write
+ * `throw memoryReservationFailed(e)`.
+ */
+ def memoryReservationFailed(e: Throwable): RuntimeException = e match {
Review Comment:
Done in
[4b28a3f1a8f](https://github.com/apache/spark/pull/59639/commits/4b28a3f1a8fd03ff47b17333e2349f070a23af69),
as you suggested: `rethrowIfMustPropagate(e, fromRead = false)`, then the
`SparkException`. The doc no longer describes input the caller never passes.
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/LateMaterializationParquetRecordReader.java:
##########
@@ -727,22 +783,38 @@ private RowRanges evaluateStorageFilter(
// where a plain read fails.
reportOnce(e, SPLICING_GIVEN_UP_ON_ERROR);
abandonSplicing();
- splicedBytes = 0L;
}
}
- if (splicedBytes + rangeBytes > cap) {
+ if (bufferedKeyCharge() + rangeCharge > cap) {
// The buffer goes first, which is a no-op once splicing is given up.
+ if (spliceCurrentRowGroup) splicingGivenUp =
BudgetGiveUp.SPLICING_ON_CAP;
abandonSplicing();
- splicedBytes = 0L;
// The ranges would be thrown away, so stop.
- if (rangeBytes > cap) return null;
+ if (rangeCharge > cap) {
+ recordBudgetGiveUp(BudgetGiveUp.FILTER_ON_CAP);
+ return null;
+ }
+ }
+ // A reservation the task's memory manager denies is handled as
passing the cap is. With no
+ // key values in it, giving splicing up frees nothing, so the ranges
are denied too.
+ long keyCharge = accountKeys ? bufferedKeyCharge() : 0L;
+ boolean denied = !reserveSurvivorMemory(keyCharge + rangeCharge);
+ if (denied && keyCharge > 0L) {
+ splicingGivenUp = BudgetGiveUp.SPLICING_ON_MEMORY;
Review Comment:
I'd rather keep each row group deciding on its own. Memory pressure can pass
between two row groups, and skipping splicing then costs a second read of the
key column that the memory would have saved. What a denied row group does in
vain, copying its survivors and a few reservations, is small next to that read.
So I'll leave this one as it is.
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/LateMaterializationParquetRecordReader.java:
##########
@@ -445,6 +484,13 @@ private void initializeLateMaterialization() throws
IOException {
}
keyFixedBytesPerRow = fixedBytesPerRow;
keyOnlyColumns = Arrays.stream(keys).map(KeyColumn::descriptor).toList();
+ // A test's task can have no memory manager.
+ TaskContext context = TaskContext.get();
+ TaskMemoryManager taskMemoryManager = context == null ? null :
context.taskMemoryManager();
+ if (taskMemoryManager != null) {
+ survivorMemory = new SurvivorMemory(taskMemoryManager);
Review Comment:
Done in
[4b28a3f1a8f](https://github.com/apache/spark/pull/59639/commits/4b28a3f1a8fd03ff47b17333e2349f070a23af69).
The `SurvivorMemory` javadoc says so again.
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/LateMaterializationParquetRecordReader.java:
##########
@@ -912,6 +1023,63 @@ private void closeSplicingState() {
}
}
+ /** The bytes of key values buffered for splicing, which the cap weighs even
off heap. */
+ private long bufferedKeyCharge() {
+ return queuedKeyCharge + currentKeyAccumulatorCharge;
+ }
+
+ /** The accounted bytes the survivor buffer holds, see {@link
SurvivorMemory}. */
+ private long heldCharge() {
+ return accountKeys ? rangeCharge + bufferedKeyCharge() : rangeCharge;
+ }
+
+ /**
+ * Grows the reservation to cover {@code held} bytes, which the caller has
kept within the cap.
+ * It asks for at least what it already holds, up to {@link
#SURVIVOR_MEMORY_STEP_BYTES} and the
+ * cap, so a growing buffer asks only a few times and the reservation never
passes the cap. False
+ * when it is denied.
+ */
+ private boolean reserveSurvivorMemory(long held) {
+ if (survivorMemory == null) return true;
+ long missing = held - reservedBytes;
+ if (missing <= 0) return true;
+ long headroom = storageFilter.maxSplicedRowGroupBytes() - reservedBytes;
+ long ask = Math.max(missing,
+ Math.min(Math.min(reservedBytes, SURVIVOR_MEMORY_STEP_BYTES),
headroom));
+ long granted;
+ try {
+ granted = survivorMemory.acquireMemory(ask);
+ } catch (RuntimeException | InternalError e) {
+ // Another consumer's spill runs inside, so this can raise anything, see
+ // `ParquetStorageFilter.memoryReservationFailed`.
+ throw ParquetStorageFilter.memoryReservationFailed(e);
+ }
+ if (granted < missing) {
+ if (granted > 0) survivorMemory.freeMemory(granted);
Review Comment:
Done in
[4b28a3f1a8f](https://github.com/apache/spark/pull/59639/commits/4b28a3f1a8fd03ff47b17333e2349f070a23af69):
- the denial test's helpers check that each read leaves the task holding
nothing;
- `memoryTask(50)` and `bytes(130L)` cover a partial grant, denied and kept,
the latter expecting `(0, 1, 0, 0)`;
- the step test pins the exact asks of the 400,000-row read, and the
accounted test those of the small one, including that each row group's first
ask finds nothing held.
Each of the five mutations you list now fails the suite.
--
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]
