dongjoon-hyun commented on code in PR #59639:
URL: https://github.com/apache/spark/pull/59639#discussion_r4340886143
##########
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:
**An `IOException` from evicting a cached block inside this reservation gets
past this catch, so under `ignoreCorruptFiles` the read skips the rest of a
healthy file.** `acquireMemory` can grow the execution pool by evicting
storage, and eviction writes a block with a disk level to disk on this thread:
- `maybeGrowPool(numBytes - memoryFree)` (ExecutionMemoryPool.scala:120) and
`freeSpaceToShrinkPool` (UnifiedMemoryManager.scala:171);
- `evictBlocksToFreeSpace` (StorageMemoryPool.scala:127), whose loop has a
`finally` and no `catch` (MemoryStore.scala:548-550);
- `dropFromMemory`, then `diskStore.put` or `putBytes`
(BlockManager.scala:2273 and 2283), then `DiskStore.put`
(DiskStore.scala:86-89).
`TaskMemoryManager` converts only an `IOException` from a consumer's
`spill()` (TaskMemoryManager.java:331), not one from
`memoryManager.acquireExecutionMemory` at TaskMemoryManager.java:206 and 301.
So a full or failing local disk sends `IOException: No space left on device`
out of `nextBatch()`, `DataSourceUtils.shouldIgnoreCorruptFileException`
accepts it (DataSourceUtils.scala:273), and FileScanRDD.scala:284-289 logs
"Skipped the rest of the content in the corrupted file" while the task
succeeds. If `openForWrite` fails (DiskStore.scala:159), the
`FileNotFoundException` is skipped as a missing file under `ignoreMissingFiles`
(FileScanRDD.scala:277-280), or otherwise fails the query as
FAILED_READ_FILE.FILE_NOT_EXIST naming the Parquet file
(FileDataSourceV2.scala:139-140). If the error comes from the retry at
TaskMemoryManager.java:301, the bytes granted at TaskMemoryManager.java:206
also belong to no consumer, so the task that now succeeds reports "Managed
memory leak detecte
d". This needs only a storage pool that has borrowed execution memory and
holds blocks of a disk level, such as `Dataset.cache()`'s default
MEMORY_AND_DISK, and even a row group's first reservation, 49 bytes for a long
key, can trigger it. Before this PR the read never acquired execution memory
inside FileScanRDD's per-file read, and the PR description says the reader
"rethrows in a form `FileScanRDD` cannot take for a corrupt file under
`ignoreCorruptFiles`". I traced this by reading the code and did not run it.
Could this catch `Exception | InternalError`, which compiles although nothing
here declares a checked exception, with `memoryReservationFailed` wrapping any
`NonFatal` exception in the checked `SparkException` (which leaves an
`InterruptedException` from the 1/2N wait as it is), and its doc and the
description naming the eviction? A test could pass a `TaskMemoryManager`
subclass whose `acquireExecutionMemory` throws an `IOException`, since
`TestMemoryManager`'s is `private[
memory]`.
##########
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:
**This counts a splicing give-up before phase 2 can still give the filter
up, so a row group whose offset index parquet cannot read is counted as one
that kept the filter.** When phase 2's `readFilteredRowGroup` throws
`MissingOffsetIndexException` (L577-595), L588 sets `filterGivenUp` and the row
group comes back whole and unchecked, but SPLICING_ON_CAP or SPLICING_ON_MEMORY
has already been counted here, and the INFO line ending "The storage filter is
still applied" may have gone out for it. `StorageFilterMetrics` documents those
counters as "row groups that kept the filter but read their key columns a
second time" (ParquetStorageFilter.scala:55-56), and the PR description says
"Each row group counts once, under the last thing given up and what caused it".
The same splicing give-up followed by a decoding or evaluation error in phase 1
counts nothing, since those paths return before L817. For example, take the
damaged file of "the checked column is false for a row group read with
out its offset index" (ParquetStorageFilterSuite.scala:2082-2083) and read it
with `maxSplicedRowGroupBytes=2100b`. Each block's 50 scattered even keys pass
the cap at survivor 43 (43 x 49 = 2,107 B), while their 2,000 B of ranges fit.
So all three blocks count here, block 1 then gives the filter up in phase 2,
and `budgetGiveUps` reads (3, 0, 0, 0) where (2, 0, 0, 0) is right. The rows
are right. I traced this by reading the code and did not run it. Could
`evaluateStorageFilter` hand the give-up back, or keep it in a field, so that
`loadNextRowGroup` records it inside `if (!filterGivenUp)` at L598, with a test
over that fixture?
##########
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:
Minor: Nothing carries a denial over to the next row group: L536 starts each
one with `spliceCurrentRowGroup = canNarrowRowGroup`, and L510 releases the
reservation to 0. When the free memory stays between what the ranges need and
what the keys need, as in the denial test's `memoryTask(105)`, every later row
group of the split repeats the same work. It allocates its accumulator sets
(36,864 B per set for a long key at capacity 4096), copies each survivor's key,
makes about log2(F / 49) + F / 1 MiB reservations through the executor-wide
memory manager for F bytes free, ends in a denied one that runs the task's
spill pass, frees it all, and reads the key column again in phase 2.
ParquetStorageFilterSuite.scala:2295 pins that repeat with `denied == Seq(0L,
2L, 0L, 0L)`. This is the denial counterpart of my item 9, which went with the
spill. Pressure can be transient, so giving splicing up for the rest of the
file could cost re-reads that a later grant would have saved. Could a denial
make the reader skip splicing for the next row group only, counted under
`rowGroupsSplicingGivenUpByMemory`, while it keeps the filter and still
reserves the ranges?
##########
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:
Minor: This copies two cases of `rethrowIfMustPropagate` (L271-282) instead
of calling it, and leaves out its search of the cause chain. A
`RuntimeException` from another consumer's `spill()` with a fatal error in its
cause chain goes out here as `SparkException(RuntimeException(...))`, and
`FileScanRDD` wraps that in FAILED_READ_FILE (FileDataSourceV2.scala:141-142).
So the fatal error sits two levels deeper than when the same spill error comes
out of that consumer's own acquisition. At the default
`spark.executor.killOnFatalError.depth` of 5 (Executor.scala:1144 and
1654-1661), a fatal error at depth 4 or 5 of the spill's exception is then
missed, and the task fails and is retried on a JVM that ran out of memory
instead of the executor exiting. With that internal conf at 2 or 3, one or two
wrappers are enough. I found no in-tree spill that builds such a chain, so this
is about robustness. Also, the doc's "and anything else, such as the
`SparkOutOfMemoryError` of a failed spill,
as it is, as in [[rethrowIfMustPropagate]]" describes input that the only
caller never passes, since its `catch (RuntimeException | InternalError e)`
(LateMaterializationParquetRecordReader.java:1052) does not catch a
`SparkOutOfMemoryError`. Could this be `rethrowIfMustPropagate(e, fromRead =
false)` followed by `throw new SparkException("A storage filter failed to
reserve memory", e)`? That keeps both outcomes of my item 2 and the spill-error
test, and it fits the wider catch of item 16.
##########
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:
**No test reaches a partial grant, `0 < granted < missing`, or checks what a
denied read leaves reserved, so five one-line mutations here and around should
pass the whole suite.** Every grant the suite makes is 0, the full ask, or at
least `missing`. In "memory the task's memory manager denies gives splicing up,
and then the filter", `rows(105)` is granted 49, 49, then 7 against a missing
of 5, then 0. `rows(0)`, `bytes(0)` and the off-heap cases get 0 or all of the
ask, the spill-error test throws before any grant, and "the survivor buffer
asks for little more memory than it holds" always has room. The `memoryTask`
contexts are never completed either, so no leak check runs. These should still
pass:
- dropping `if (granted > 0) survivorMemory.freeMemory(granted);`, which
leaks the partial grant: it stays in the consumer's `used` but never in
`reservedBytes`, so `close()` never frees it, and the executor reports "Managed
memory leak detected";
- accepting only `granted >= ask` at L1057, since `rows(105)` is then denied
at survivor 7 instead of 8 with the same bytes, metrics and log, so the comment
at ParquetStorageFilterSuite.scala:2280 is not pinned;
- asking for exactly `missing`, since no other consumer is spilled either
way, so "a growing buffer asks only a few times" is not pinned;
- dropping L1013 in `abandonSplicing`, or L509-510, since L557 frees the
same before the first sample.
I checked these by reading the tests and did not run them. Could the denial
test's helpers check that the task holds nothing after each read, and could the
test add `memoryTask(50)`, where survivor 2 asks for 49 with 1 byte free and is
granted 1 against a missing of 9, and `bytes(130L)`, whose block 0 keeps
splicing on a grant of 32 against a missing of 5 and which expects `Seq(0L, 1L,
0L, 0L)`?
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetStorageFilterSuite.scala:
##########
@@ -1179,6 +1207,8 @@ class ParquetStorageFilterSuite extends QueryTest with
SharedSparkSession
s"rows 1-100 and 201-300 must splice and rows 101-200 must not; got
$splicedByRowGroup")
assert(m.rowGroupsSkipped.value == 1,
s"expected the row group with no survivor skipped; got
${m.rowGroupsSkipped.value}")
+ assert(m.rowGroupsSplicingGivenUpByCap.value == 1,
Review Comment:
Minor: This and L1936 check only the counter each test expects, while the
other cap tests check all four through `budgetGiveUps` (L1119-1121, L1902 and
the denial test). A row group that is also counted under another cause still
passes them. For example, if `splicingGivenUp` were kept in a field that only
the normal return resets, "a row group that gives the filter up does not take
the next one with it" would count its first row group's splicing give-up again
at the second, giving `Seq(1L, 0L, 1L, 0L)`, and every test in the suite should
still pass. The PR description says "the existing cap tests now count each row
group under what it gave up, once". Could both build their metrics with
`metricMap()` and assert `budgetGiveUps(metrics) == Seq(1L, 0L, 0L, 0L)` here
and `Seq(0L, 0L, 1L, 0L)` at L1936?
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetTestAccess.scala:
##########
@@ -236,7 +236,8 @@ object ParquetTestAccess {
def queuedSurvivorSets(
reader: LateMaterializationParquetRecordReader):
Seq[Seq[WritableColumnVector]] =
survivorBatchesField.get(reader)
-
.asInstanceOf[java.util.ArrayDeque[Array[WritableColumnVector]]].asScala.map(_.toSeq).toSeq
+
.asInstanceOf[java.util.ArrayDeque[LateMaterializationParquetRecordReader.SurvivorBatch]]
Review Comment:
Nit: This names the reader's `SurvivorBatch` in a cast and, through
`_.vectors()`, in the lambda, and the reader widened the record from `private`
to package-private for it
(LateMaterializationParquetRecordReader.java:1111-1115). L36-39 here say "Under
`spark-submit --jars`, test and main classes load from different classloaders,
blocking package-private access despite the matching package name. ... So a
package-private class is looked up by name here, never named in a `classOf` or
a cast.", and every other member does that. The benchmarks that use this object
run under `spark-submit --jars` (.github/workflows/benchmark.yml:192-194), so
one that called `queuedSurvivorSets` would fail with an `IllegalAccessError`.
Only ParquetStorageFilterSuite.scala:3710 calls it today, so this is latent.
Could the record go back to `private`, with the by-name lookup 0c670c9 had,
held in a lazy val like `getLeavesMethod` (L251)?
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetStorageFilterSuite.scala:
##########
@@ -2119,6 +2163,206 @@ class ParquetStorageFilterSuite extends QueryTest with
SharedSparkSession
}
}
+ // A task whose memory manager grants at most `limit` bytes.
+ private def memoryTask(limit: Long = Long.MaxValue): TaskContextImpl = {
+ val memoryManager = new TestMemoryManager(spark.sparkContext.conf)
+ memoryManager.limit(limit)
+ new TaskContextImpl(stageId = 0, stageAttemptNumber = 0, partitionId = 0,
taskAttemptId = 0,
+ attemptNumber = 0, numPartitions = 1,
+ taskMemoryManager = new TaskMemoryManager(memoryManager, 0),
+ localProperties = new Properties, metricsSystem = null)
+ }
+
+ // Reads `path` under `k >= 10` in batches of `batchSize` and a cap of
`cap`, as a task of
+ // `context`. Returns the rows, and what the task holds as each batch is
read. The confs are set
+ // outside the task, which sees a read-only one.
+ private def readInTask(
Review Comment:
Nit: `readInTask` sets its three confs with `withSQLConf` and then calls
`readChecked` inside `asTask`, which `readUnderCap` (L203-219) now does through
`context: Option[TaskContext] = None`. Could `readChecked` take the same
parameter, so `readInTask` passes its confs through `readChecked`'s `conf`?
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/LateMaterializationParquetRecordReader.java:
##########
@@ -202,6 +221,18 @@ public class LateMaterializationParquetRecordReader
extends VectorizedParquetRec
/** The array the emitted batch is built over, whose key slots each batch
rewrites. */
private ColumnVector[] spliceBatchColumns;
+ /** What the survivor ranges of the row group being read or emitted cost to
hold. */
+ private long rangeCharge;
+ /** The survivor buffer's account with the task's memory manager, or null
outside a task. */
+ private SurvivorMemory survivorMemory;
+ /**
+ * What {@link #survivorMemory} holds. Only this reader changes it, so it is
kept here rather
+ * than read off the consumer's atomic counter per survivor.
+ */
+ private long reservedBytes;
+ /** Whether the key values are accounted too, which needs them on heap. */
+ private boolean accountKeys;
Review Comment:
Nit: `accountKeys` is `MEMORY_MODE == MemoryMode.ON_HEAP` wherever it has an
effect. Without a task memory manager every use of `heldCharge()` is already a
no-op: `reserveSurvivorMemory` returns at L1043, and `reservedBytes` stays 0,
so the releases and the trim return at once. L800-802 also rebuild
`heldCharge()` by hand, where `keyCharge > 0L` is `held > rangeCharge`, and
L805's `rangeCharge` is `heldCharge()` once `abandonSplicing()` has run. Could
the field go, with `heldCharge()` testing `MEMORY_MODE` and L800-805 reading
`long held = heldCharge(); boolean denied = !reserveSurvivorMemory(held); if
(denied && held > rangeCharge) { ... denied =
!reserveSurvivorMemory(heldCharge()); }`? A term added to `heldCharge()` would
then also reach the reservation.
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/LateMaterializationParquetRecordReader.java:
##########
@@ -663,11 +718,12 @@ private RowRanges evaluateStorageFilter(
PrimitiveIterator.OfLong rowIndexIter = pushedFilterRanges.iterator();
RowRanges.Builder finalRangesBuilder = RowRanges.builder();
- // The budget weighs both the bytes buffered for splicing and the
survivors' ranges.
- long splicedBytes = 0L;
- long rangeBytes = 0L;
+ // The budget weighs both the key values buffered for splicing and the
survivors' ranges.
+ rangeCharge = 0L;
Review Comment:
Nit: This store never changes the value. L509 zeroes `rangeCharge` at the
start of every `loadNextRowGroup`, L527's `continue` comes before phase 1, the
one at L567 is reached only when nothing survived, since a range is charged
only for a survivor (L770-771), and a row group whose filter was given up
zeroes it at L552 and returns at L615. L545 is the only caller. Could it go,
with the comment at L721 moving to `cap`?
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetStorageFilterSuite.scala:
##########
@@ -2119,6 +2163,206 @@ class ParquetStorageFilterSuite extends QueryTest with
SharedSparkSession
}
}
+ // A task whose memory manager grants at most `limit` bytes.
+ private def memoryTask(limit: Long = Long.MaxValue): TaskContextImpl = {
+ val memoryManager = new TestMemoryManager(spark.sparkContext.conf)
+ memoryManager.limit(limit)
+ new TaskContextImpl(stageId = 0, stageAttemptNumber = 0, partitionId = 0,
taskAttemptId = 0,
+ attemptNumber = 0, numPartitions = 1,
+ taskMemoryManager = new TaskMemoryManager(memoryManager, 0),
+ localProperties = new Properties, metricsSystem = null)
+ }
+
+ // Reads `path` under `k >= 10` in batches of `batchSize` and a cap of
`cap`, as a task of
+ // `context`. Returns the rows, and what the task holds as each batch is
read. The confs are set
+ // outside the task, which sees a read-only one.
+ private def readInTask(
+ path: String,
+ context: TaskContext,
+ offHeap: Boolean = false,
+ batchSize: Int = 4,
+ cap: String = "64MB"): (Seq[(Long, Boolean)], Seq[Long]) = withSQLConf(
+ SQLConf.PARQUET_VECTORIZED_READER_BATCH_SIZE.key -> batchSize.toString,
+ SQLConf.COLUMN_VECTOR_OFFHEAP_ENABLED.key -> offHeap.toString,
+ SQLConf.PARQUET_STORAGE_FILTER_PUSHDOWN_MAX_SPLICED_ROW_GROUP_BYTES.key
-> cap) {
+ var held = Seq.empty[Long]
+ val rows = asTask(context) {
+ readChecked(path, checkedSchema(), Seq(keyAtLeast(10L)), keyOf = (b, i)
=> {
+ if (i == 0) held :+=
context.taskMemoryManager().getMemoryConsumptionForThisTask
+ b.column(0).getLong(i)
+ })
+ }
+ (rows, held)
+ }
+
+ // One row group of 400,000 rows of a long key and a string value. Under `k
>= 10` its 399,990
+ // survivors fall in one range and buffer 3,599,950 bytes, several steps of
the reservation.
+ private def oneRowGroupOf400000(dir: File): String = {
Review Comment:
Nit: "emitting a row group releases its memory in steps" and "the survivor
buffer asks for little more memory than it holds" each write this 400,000-row
file through `repartition(1)` and collect about 400,000 `(Long, Boolean)` rows.
In [this commit's CI
run](https://github.com/peter-toth/spark/actions/runs/39632367694) they are the
suite's slowest and fourth-slowest tests, 1.126 s and 0.601 s of its 26.1 s. A
wider key would need fewer rows, but it would change the steps the first test
pins, and at a few thousand rows the doubling would stop inside the second
test's margin, so that test would no longer catch an unclamped ask. Could the
second test's 400,000-row case move into the first test, reading the same file
twice, so it is written once?
##########
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:
Nit: Each reader creates its own `SurvivorMemory` here,
`TaskMemoryManager.acquireExecutionMemory` adds it to the task's `consumers` at
its first reservation (TaskMemoryManager.java:256), and only
`cleanUpAllAllocatedMemory` clears that set (TaskMemoryManager.java:916). So a
task keeps one per file it reserved for, about 80 B each. Every contended
acquisition by any consumer of the task, every page-allocation recovery and
every memory-usage snapshot walks them under the task memory manager's lock,
although a closed one costs only a `getUsed()` check. At the default 128 MB / 4
MB packing that is about 33 per scan per task, so it is negligible today. It
grows with a lowered `openCostInBytes`, with a bucketed scan over many files,
and with the several readers per task that the description expects once storage
filters reach storage-partitioned joins. The round-1 javadoc said "The task's
memory manager keeps every consumer until the task ends", and the new one does
not. Could the `Surv
ivorMemory` javadoc keep that, so that the change for several readers per task
can share one consumer per task and scan?
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetStorageFilterSuite.scala:
##########
@@ -2119,6 +2163,206 @@ class ParquetStorageFilterSuite extends QueryTest with
SharedSparkSession
}
}
+ // A task whose memory manager grants at most `limit` bytes.
+ private def memoryTask(limit: Long = Long.MaxValue): TaskContextImpl = {
+ val memoryManager = new TestMemoryManager(spark.sparkContext.conf)
+ memoryManager.limit(limit)
+ new TaskContextImpl(stageId = 0, stageAttemptNumber = 0, partitionId = 0,
taskAttemptId = 0,
+ attemptNumber = 0, numPartitions = 1,
+ taskMemoryManager = new TaskMemoryManager(memoryManager, 0),
+ localProperties = new Properties, metricsSystem = null)
+ }
+
+ // Reads `path` under `k >= 10` in batches of `batchSize` and a cap of
`cap`, as a task of
+ // `context`. Returns the rows, and what the task holds as each batch is
read. The confs are set
+ // outside the task, which sees a read-only one.
+ private def readInTask(
+ path: String,
+ context: TaskContext,
+ offHeap: Boolean = false,
+ batchSize: Int = 4,
+ cap: String = "64MB"): (Seq[(Long, Boolean)], Seq[Long]) = withSQLConf(
+ SQLConf.PARQUET_VECTORIZED_READER_BATCH_SIZE.key -> batchSize.toString,
+ SQLConf.COLUMN_VECTOR_OFFHEAP_ENABLED.key -> offHeap.toString,
+ SQLConf.PARQUET_STORAGE_FILTER_PUSHDOWN_MAX_SPLICED_ROW_GROUP_BYTES.key
-> cap) {
+ var held = Seq.empty[Long]
+ val rows = asTask(context) {
+ readChecked(path, checkedSchema(), Seq(keyAtLeast(10L)), keyOf = (b, i)
=> {
+ if (i == 0) held :+=
context.taskMemoryManager().getMemoryConsumptionForThisTask
+ b.column(0).getLong(i)
+ })
+ }
+ (rows, held)
+ }
+
+ // One row group of 400,000 rows of a long key and a string value. Under `k
>= 10` its 399,990
+ // survivors fall in one range and buffer 3,599,950 bytes, several steps of
the reservation.
+ private def oneRowGroupOf400000(dir: File): String = {
+ val path = writeSingleParquetFile(dir,
+ spark.range(400000L).selectExpr("id AS k", "CAST(id AS STRING) AS v"),
+ rowGroupSize = ParquetWriter.DEFAULT_BLOCK_SIZE)
+ assert(footerOf(path).getBlocks.size == 1, "the fixture must be one row
group")
+ path
+ }
+
+ test("the survivor buffer is accounted to the task's execution memory") {
+ // What a row group buffers is held while it is emitted, and released when
the next one loads.
+ // Phase 1 reserves a little ahead of what it buffers, and gives that back
before phase 2. Off
+ // heap, only the ranges are accounted, since the vectors' native memory
is not the memory
+ // manager's.
+ withTempDir { dir =>
+ val path = twoRowGroupsOf20(dir)
+ // Past a cap of 100 bytes, at the seventh survivor of each row group,
the keys go, a full
+ // batch of them already queued, and only the range is held.
+ Seq((false, "64MB", Seq(130L, 130L, 130L, 220L, 220L, 220L, 220L, 220L)),
+ (true, "64MB", Seq.fill(8)(40L)),
+ (false, "100b", Seq.fill(8)(40L))).foreach { case (offHeap, cap,
expected) =>
+ val context = memoryTask()
+ val (rows, held) = readInTask(path, context, offHeap, cap = cap)
+ val name = s"off heap $offHeap under a cap of $cap"
+ assert(rows == marked(10L until 40L, checked = true), s"$name: got
$rows")
+ assert(held == expected, s"$name: each batch must hold what is still
buffered; got $held")
Review Comment:
Nit: "each batch must hold what is still buffered" fits round 1's per-batch
release, where this assertion expected `Seq(130L, 94L, 58L, 220L, 184L, 148L,
112L, 76L)` at 0c670c9. Now each batch holds its row group's whole reservation
until the next one loads, e.g. 130 at block 0's second batch, which buffers 94.
Could the message say something like "each batch must hold its row group's
reservation until the next row group loads"?
--
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]
