dongjoon-hyun commented on code in PR #59639:
URL: https://github.com/apache/spark/pull/59639#discussion_r4327382584
##########
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();
Review Comment:
**A spill can leave the reader holding more heap than before, while it tells
the task memory manager that memory was freed.** This returns only the queued
key bytes, and `TaskMemoryManager` hands them straight to the consumer that
asked. The next `nextBatch()` then allocates two things that `heldCharge()`
does not count, since L1050 counts `rangeCharge` for `loadedRanges` only:
- `remainingRanges` (L1032-1036), a second `RowRanges` with one `Range` per
remaining run, while `loadedRanges` stays referenced by the field, by the
non-key readers' `RangeListState` (ParquetReadState.java:107) and by
`dataPages`;
- `rereadKeyPages` (L1038), the compressed key pages of every remaining row,
which parquet reads into heap buffers that stay until the next row group.
For example, take a 1M-row row group with a long key and 100k scattered
survivors (one range each), spilled after its first batch of 4096 rows. The
spill frees 95,904 x 9 B, about 0.86 MB. The re-read then allocates about 3.6
MB of new ranges (3.8 MB at `ROW_RANGE_BYTES`) and reads the key chunk's
remaining pages, about 7.7 MB for plain longs. So the reader's real footprint
grows by several MB under the memory pressure that caused the spill, and
scattered survivors are what a selective bloom filter produces. Two things make
it more likely that this consumer is the one spilled: `getUsed()` includes the
unspillable `rangeCharge`, so `ceilingEntry(required - got)`
(TaskMemoryManager.java:242-249) can pick it when it can free little, and
`spill()` ignores `size`, so a request for one page drops the whole queue. I
traced this by reading the code and did not run it. Could
`spillQueuedSurvivors` return 0 when the re-read would cost at least what it
frees, e.g. the remaining ranges at `RO
W_RANGE_BYTES` plus `blockAccounting.compressedBytesForRowRanges` over the key
paths (no I/O, since L579 built the store), or else charge the new ranges?
##########
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
Review Comment:
Minor: The javadoc says a spill on the reading thread "can come only between
two batches", but nothing here checks it. `spillableRowGroup` is already true
from L575 on, including while `super.nextBatch()` reads each batch. If a
consumer ever acquired task memory on the reading thread inside `nextBatch()`,
this would run between the re-read check at L376 and the poll at L381.
`spliceCurrentRowGroup` would turn false, the poll would be skipped, and
L394-395 would publish the persistent key vectors, which the base reset but did
not fill because `installRowGroup` cleared their readers. Those rows are marked
checked, so the post-scan `Filter` would skip the storage filter on them. I
found nothing in the parquet, vectorized or `FileScanRDD` read paths that
acquires task memory there today, so this is latent. Could a flag set on entry
to `nextBatch()` and cleared in a `finally` make `spillQueuedSurvivors` return
0 inside it?
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/LateMaterializationParquetRecordReader.java:
##########
@@ -513,6 +568,16 @@ protected boolean loadNextRowGroup() throws IOException {
continue;
}
+ // A spill has the rest of a spliced row group read its keys again, in
part. That needs their
+ // offset indexes. Parquet builds a block's store once, over the columns
requested then, so it
+ // is built here over every projected column. Only a row group phase 2
reads in part can be
+ // spilled. Phase 2 reads its store anyway, and gives the filter up on
an unreadable index.
+ spillableRowGroup =
Review Comment:
**`finalRowCount < blockRowCount` leaves out the row groups that buffer the
most, those in which every row survives, and the reasons in the PR description
do not apply to them.** For such a row group the key re-read covers the whole
block, and parquet 1.18.1 reads a whole block through `internalReadRowGroup`,
with no `ColumnIndexStore` and no offset index (`readFilteredRowGroup` returns
it when `rowCount == block.getRowCount()`). So there is no extra read and no
index that can be damaged. The rows already emitted can be skipped by decoding
and dropping them through `keyScratchVectors`, the way phase 1 decodes at
L749-750, which keeps the key readers on the same row as the non-key ones. With
clustered or sorted keys, where storage filters skip the most, most surviving
row groups are whole, like the 220-byte second row group of the PR's fixture.
Their buffers, up to the 64 MB cap, are accounted, but `spill()` returns 0, so
the memory manager spills a sort or an aggregation to disk i
nstead, or a consumer that cannot spill fails, while the reader keeps them.
The test "a spill leaves a row group phase 2 read whole, whose offset index may
not parse" damages column 1, the value column `v`
(ParquetStorageFilterSuite.scala:2300), which a key-only re-read never reads,
so it pins the current choice rather than a hazard. Could spillable be
`spliceCurrentRowGroup && accountKeys`, with a whole-block re-read and
decode-and-drop for such row groups, keeping the remaining-ranges re-read for
the partial ones?
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/LateMaterializationParquetRecordReader.java:
##########
@@ -730,12 +800,18 @@ private RowRanges evaluateStorageFilter(
splicedBytes = 0L;
}
}
- if (splicedBytes + rangeBytes > cap) {
+ if (splicedBytes + rangeCharge > cap) {
// The buffer goes first, which is a no-op once splicing is given up.
abandonSplicing();
splicedBytes = 0L;
// The ranges would be thrown away, so stop.
- if (rangeBytes > cap) return null;
+ if (rangeCharge > cap) return null;
+ }
+ if (!reserveSurvivorMemory((accountKeys ? splicedBytes : 0L) +
rangeCharge)) {
Review Comment:
**This is the first task-memory acquisition inside `FileScanRDD`'s per-file
read, so an unchecked error from a spill it triggers now looks like a corrupt
file.** The two reservations (L810 and L814) are outside every catch of the
reader. When the pool cannot grant them, `TaskMemoryManager.spillConsumer`
spills other consumers of the task and converts only interrupts and
`IOException` (TaskMemoryManager.java:326-341), so any other `RuntimeException`
or `InternalError` from their `spill()` propagates unchanged. Consumers that
accept a spill triggered by another consumer include `UnsafeExternalSorter`
while it is being read, `HybridRowQueue`, `Spillable` collections and another
late-materialization reader of the same task. With
`spark.sql.files.ignoreCorruptFiles=true`, such an error now leaves
`nextBatch()` and reaches FileScanRDD.scala:284-289, which skips the rest of
the file. The task then succeeds with a truncated scan, and the other consumer
may be left half-spilled. Before thi
s PR, the same error could only come from that consumer's own allocation
outside `FileScanRDD`, and it failed the task. This is the concern for which
`throwIfKilled` wraps a kill in a checked `SparkException`
(ParquetStorageFilter.scala:276-280). I found no in-tree spill path that throws
such an error today, apart from invariant checks, so this is about robustness.
Could the two reservations rethrow a foreign `RuntimeException` or
`InternalError` in a form that `ignoreCorruptFiles` does not swallow?
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/LateMaterializationParquetRecordReader.java:
##########
@@ -730,12 +800,18 @@ private RowRanges evaluateStorageFilter(
splicedBytes = 0L;
}
}
- if (splicedBytes + rangeBytes > cap) {
+ if (splicedBytes + rangeCharge > cap) {
// The buffer goes first, which is a no-op once splicing is given up.
abandonSplicing();
splicedBytes = 0L;
// The ranges would be thrown away, so stop.
- if (rangeBytes > cap) return null;
+ if (rangeCharge > cap) return null;
+ }
+ if (!reserveSurvivorMemory((accountKeys ? splicedBytes : 0L) +
rangeCharge)) {
+ // A reservation the task's memory manager denies is handled as
passing the cap is.
+ abandonSplicing();
+ splicedBytes = 0L;
+ if (!reserveSurvivorMemory(rangeCharge)) return null;
Review Comment:
**A denied range reservation, as small as the first 40 bytes, gives the
filter up for the whole row group, which costs far more than it saves.** Phase
2 then reads every projected column over `pushedFilterRanges`, key columns
included, and marks every row unchecked, so the post-scan `Filter` evaluates
them all: a plain read plus a second read of the keys, with more unaccounted
page memory than the filtered read. The cap-only reader before this PR kept the
filter in the same state. It is reachable whenever the task is at its 1/N share
and that memory is held by consumers that refuse a spill triggered by another
consumer, e.g. an inserting `ShuffleExternalSorter`
(ShuffleExternalSorter.java:302) after `numActiveTasks` grew:
`ExecutionMemoryPool` grants 0 without waiting, and this consumer's self-spill
returns 0. Every row group with a survivor that loads in that window loses the
filter, and nothing logs or counts it. Off heap, the first request carries no
key bytes, so `abandonSplic
ing()` frees nothing accounted and L814 repeats the same acquire and spill
pass, and the comment at L811 does not hold there, since the cap counts the
off-heap keys. Could a range denial keep the filter and stop reserving for the
rest of the row group, which the cap still bounds as before this PR, or at
least log it once and count it? The conf doc ("A denied reservation is handled
the same way") and ParquetStorageFilterSuite.scala:2253-2255 would change with
it. Also, could L814 be skipped when the denied request had no key bytes?
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/LateMaterializationParquetRecordReader.java:
##########
@@ -513,6 +568,16 @@ protected boolean loadNextRowGroup() throws IOException {
continue;
}
+ // A spill has the rest of a spliced row group read its keys again, in
part. That needs their
+ // offset indexes. Parquet builds a block's store once, over the columns
requested then, so it
+ // is built here over every projected column. Only a row group phase 2
reads in part can be
+ // spilled. Phase 2 reads its store anyway, and gives the filter up on
an unreadable index.
+ spillableRowGroup =
+ spliceCurrentRowGroup && accountKeys && finalRowCount <
blockRowCount;
+ if (spillableRowGroup) {
Review Comment:
Minor: When no parquet filter is pushed, this is the block's first store
build, so it also reads the key columns' offset indexes, and "Phase 2 reads its
store anyway" at L574 holds only for the non-key ones. With a pushed filter,
usually the inferred `isnotnull(key)`, `initialize` has already built every
block's store over the whole projection through `getFilteredRecordCount`, so
nothing changes. Without one (`spark.sql.parquet.filterPushdown=false`, an
INT96 timestamp key that `ParquetFilters` cannot convert, a `coalesce(...)`
join key, or the suite's direct readers), the base built the store in phase 2
over `nonKeyColumns` only. Now every partly read spliced row group in an
on-heap task:
- reads one more offset index per key column chunk, whether or not a spill
comes, and the metrics do not count it;
- gives the filter up when a key column's offset index is unreadable,
because `ColumnIndexStoreImpl.create` then returns its empty store and phase 2
throws `MissingOffsetIndexException`. L596 then logs `OFFSET_INDEX_UNREADABLE`
naming a non-key column whose index is fine. The same file read outside a task,
or off heap, keeps the filter. For example, "the checked column is false for a
row group read without its offset index" with `breakOffsetIndex(damaged, 1,
columnIdx = 0)` under `readInTask(memoryTask(...))` should return that row
group unchecked (I did not run it);
- fails the read on an unchecked error decoding that index, such as an
encrypted file's `TagVerificationException`, where a plain read never reads it.
No test covers the in-task, no-filter, partly read case. Could you note this
in the comment at L571-574 and add such a test?
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/LateMaterializationParquetRecordReader.java:
##########
@@ -920,6 +1095,35 @@ private static void closeAll(ColumnVector[] vectors) {
}
}
+ /**
+ * The survivor buffer's account with the task's memory manager, so other
consumers see what it
+ * holds. A spill can free the survivors not yet emitted, see {@link
#spillQueuedSurvivors}.
+ * ON_HEAP whatever the reader's memory mode, since off-heap vectors
allocate native memory
+ * outside the Tungsten off-heap pool. With them only the row ranges, which
are always on heap,
+ * are accounted, and a spill frees nothing. The task's memory manager keeps
every consumer until
+ * the task ends, so this lets go of its reader when the reader closes.
+ */
+ private static final class SurvivorMemory extends MemoryConsumer {
+ private volatile LateMaterializationParquetRecordReader reader;
+
+ SurvivorMemory(
+ TaskMemoryManager taskMemoryManager,
LateMaterializationParquetRecordReader reader) {
+ super(taskMemoryManager, MemoryMode.ON_HEAP);
Review Comment:
Minor: A storage-filter scan now holds on-heap execution memory through each
row group that can be read in part, at least the 40-byte range charges, even
with off-heap vectors. Scan tasks that held none before, e.g. one feeding
`BypassMergeSortShuffleWriter` (the default for up to 200 shuffle partitions)
or a result stage, now enter `ExecutionMemoryPool.memoryForTask` and count in
`numActiveTasks` (ExecutionMemoryPool.scala:114). That lowers every other
task's cap to `maxPoolSize / N` and its floor to `poolSize / 2N`
(ExecutionMemoryPool.scala:128-129). On an executor shared by concurrent stages
or jobs, one sort task next to seven such scans sees its cap drop from the
whole pool to 1/8 while the scans hold a few hundred bytes each, so it spills
much more. The release to 0 at each row group's load (L513-514) also takes the
task out of the pool and back in, and a 40-byte re-entry into a full pool below
1/2N can wait in `lock.wait()`, which a plain scan never does. Nothing changes w
hen every slot runs the same scan stage. Could the reader skip the reservation
when nothing in it can be spilled, i.e. off heap, and could the conf doc
mention that an on-heap scan now counts as an active task?
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/LateMaterializationParquetRecordReader.java:
##########
@@ -335,19 +368,26 @@ private void decideAllKeysMissingFile() {
*/
@Override
public boolean nextBatch() throws IOException {
+ readerThread = Thread.currentThread();
closeAll(publishedKeyVectors);
publishedKeyVectors = null;
+ publishedKeyCharge = 0L;
+ releaseSurvivorMemoryTo(heldCharge());
Review Comment:
Minor: This walks the whole remaining queue in `heldCharge()` on every
batch, then frees one batch's share through
`MemoryManager.releaseExecutionMemory`, which is `synchronized` on the
executor-wide memory manager (MemoryManager.scala:138) and ends in
`lock.notifyAll()` (ExecutionMemoryPool.scala:171), waking every task that
waits for execution memory. At the 64 MiB cap with a long key and the default
batch size, a row group has 1,821 batches, so emitting it visits about 1.66M
queue entries and takes the global lock 1,821 times. With an int key the walk
is 5.4M, and with a batch size of 1024 it is 26.5M. Before this PR the emit
path made no memory-manager calls. Phase 1 likewise reads an `AtomicLong` per
survivor (`survivorMemory.getUsed()` at L1066), against the file's own rule at
L741, for about 20 actual changes per row group. Could the reader keep a
running queued-charge total and a plain mirror of the reservation (only the
reading thread changes it), and release in larger st
eps or at the end of the row group?
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/LateMaterializationParquetRecordReader.java:
##########
@@ -460,6 +507,11 @@ protected boolean loadNextRowGroup() throws IOException {
// assignment comes only after the next row group's phase 1, whose budget
does not count them.
closeDataPages();
releaseRowGroupReaders();
+ loadedRanges = null;
+ spillableRowGroup = false;
Review Comment:
Minor: Nothing records that a spill happened, so the next row group splices
again (`spliceCurrentRowGroup = canNarrowRowGroup` at L540). If a downstream
consumer on the same thread falls short once per row group, e.g. an aggregation
that has fallen back to `UnsafeKVExternalSorter` or a sort that spills itself,
every row group pays for phase 1's copy of all its survivors (about 245
accumulator sets for 1M long keys) and for its reservations. It is then spilled
at the first shortfall, since inserting sorters and maps refuse a spill
triggered by another consumer and so this consumer is picked before the
requester spills itself, and then it reads its keys again. That is more work
than not splicing, apart from the batches emitted before the spill. Could the
reader stop splicing for the rest of the file, or for the next few row groups,
after a spill?
##########
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:
**Asking for `max(missing, used)` makes the memory manager treat the whole
amount as required, so a few missing bytes can spill other consumers, evict
cached blocks, or block the task.** `TaskMemoryManager.acquireExecutionMemory`
spills other consumers until `got >= required`
(TaskMemoryManager.java:240-249), `maybeGrowPool(numBytes - memoryFree)` can
evict storage blocks (ExecutionMemoryPool.scala:120), and a task below 1/2N of
the pool waits in `lock.wait()` while `toGrant < numBytes`
(ExecutionMemoryPool.scala:139-142). All of this can happen when `missing`
alone was free. The ask is not clamped to the cap either, so at the default 64
MiB the last step asks about 49 MiB for a 9-byte shortfall, and the excess
stays reserved through the store build and phase 2's I/O until L629. The test
"the survivor buffer asks for little more memory than it holds" passes only
because 10,000 bytes are free against at most 392 reserved. With
`memoryManager.limit(held + 300L)` at ParquetStorageFil
terSuite.scala:2266, which is still more than the 220 bytes the buffer ever
holds, I expect the assertion at ParquetStorageFilterSuite.scala:2272 to fail.
At survivor 18 of the second row group the buffer holds 202 bytes and has 196
reserved, so it needs 6 more. It asks for 196, gets the 104 that are free, and
`TaskMemoryManager` spills `other` for the rest. So "a buffer that fits in the
free memory spills no other consumer" in the PR description holds only for a
buffer under about half the free memory. I traced this by reading the code and
did not run it. Could the ask be bounded, e.g. `max(missing, min(used, 1 MiB,
cap - used))`, and the excess released right after phase 1, with the test and
the description narrowed to match?
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetStorageFilterSuite.scala:
##########
@@ -1977,19 +1978,26 @@ class ParquetStorageFilterSuite extends QueryTest with
SharedSparkSession
private def marked(keys: Seq[Long], checked: Boolean): Seq[(Long, Boolean)] =
keys.map(_ -> checked)
+ // 40 rows of a long key and a string value in two row groups of 20. A
filter of `k >= 10` keeps
+ // one range of 10 rows in the first and all 20 rows of the second. A
buffered long key costs 9
+ // bytes and a range 40, so the first row group holds 130 bytes and the
second 220.
+ private def twoRowGroupsOf20(dir: File): String = {
Review Comment:
Nit: The same fixture is still inlined at L1630-1635, in "a key vector that
cannot grow, or a codec's InternalError, gives the row group up", which could
call this and keep `val rows` for its expected `rows.take(20) ++ rows.drop(30)`
at L1665. Two of the new helpers could also reuse existing ones.
`readWithSpill` could deny the next acquisition with
`memoryManager.markExecutionAsOutOfMemoryOnce()`, as
UnsafeExternalSorterSuite.java:257 does, instead of `limit(0L)` and the
`limit(1L << 40)` restore, and the 76 and 130 assertions should still hold. The
two local `bytes` helpers (L2233-2240 and L2313-2318) repeat `readUnderCap`
with `asTask` around the read. Could you reuse them?
##########
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();
Review Comment:
Nit: This walks `loadedRanges` one row at a time: `nextLong()` per emitted
row, then `addSelectedRow` per remaining row. Parquet 1.18.1's public
`RowRanges.getRanges()` and `Builder.addSelectedRange(from, to)` build the same
ranges per range, and ParquetReadState.java:107 already reads `getRanges()`:
`long skip = rowsEmittedInRowGroup; for (RowRanges.Range r :
loadedRanges.getRanges()) { long n = r.to - r.from + 1; if (skip >= n) { skip
-= n; continue; } remaining.addSelectedRange(r.from + skip, r.to); skip = 0;
}`. Also, a spill comes only after a batch was emitted, so `remainingRows` is
below the block's row count and `readerRangesFor` at L1041 always returns
`remainingRanges`, and the queue is non-empty only while rows remain, so the
guard at L1027 never fires. `loadedBlockIndex` and `loadedRowCount` also equal
`nextBlockIndex - 1` and `loadedRanges.rowCount()` here. Could you simplify
this?
##########
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
Review Comment:
Minor: This now says that the re-read after a spill reads the key columns,
but the `spliceCurrentRowGroup` javadoc at L168-169 still says "False means
phase 2 read every projected column, so the emit path takes the key columns
from the persistent batch". After a spill that flag is false, but phase 2 read
only the non-key columns and the keys come from `rereadKeyPages`. Also,
`spillableRowGroup` stays true after the `MissingOffsetIndexException` fallback
abandons splicing at L601, which is harmless only because
`spillQueuedSurvivors` also checks `spliceCurrentRowGroup` and the queue. Could
the field javadoc say something like "False means the key columns are read from
the file, by phase 2 or by the re-read after a spill", and could the fallback
clear `spillableRowGroup`?
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetStorageFilterSuite.scala:
##########
@@ -2119,6 +2127,201 @@ class ParquetStorageFilterSuite extends QueryTest with
SharedSparkSession
}
}
+ // A task whose memory manager draws on `memoryManager`, which a test can
limit.
+ private def memoryTask(memoryManager: TestMemoryManager): TaskContextImpl =
+ 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 `filter` in batches of `batchSize`, as a task of
`context`. `onRow` gets
+ // each row's index in its batch before its key 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,
+ filter: Expression = keyAtLeast(10L),
+ schema: StructType = checkedSchema(),
+ batchSize: Int = 4)(onRow: Int => Unit): Seq[(Long, Boolean)] =
withSQLConf(
+ SQLConf.PARQUET_VECTORIZED_READER_BATCH_SIZE.key -> batchSize.toString,
+ SQLConf.COLUMN_VECTOR_OFFHEAP_ENABLED.key -> offHeap.toString) {
+ asTask(context) {
+ readChecked(path, schema, Seq(filter),
+ keyOf = (b, i) => { onRow(i); b.column(0).getLong(i) })
+ }
+ }
+
+ // Reads `path` as `readInTask` does. At the first batch, another consumer
asks for memory the
+ // manager has no more of, so the reader is spilled. Returns the rows, with
what the task held
+ // before and after the spill.
+ private def readWithSpill(
+ path: String,
+ otherThread: Boolean = false,
+ filter: Expression = keyAtLeast(10L),
+ schema: StructType = checkedSchema(),
+ batchSize: Int = 4): (Seq[(Long, Boolean)], Long, Long) = {
+ val memoryManager = new TestMemoryManager(spark.sparkContext.conf)
+ val context = memoryTask(memoryManager)
+ val memory = context.taskMemoryManager
+ val other = new TestMemoryConsumer(memory)
+ var before, after = -1L
+ val rows = readInTask(path, context, filter = filter, schema = schema,
+ batchSize = batchSize) { i =>
+ if (before < 0 && i == 0) {
+ before = memory.getMemoryConsumptionForThisTask
+ memoryManager.limit(0L)
+ if (otherThread) ThreadUtils.runInNewThread("spill")(other.use(before))
+ else other.use(before)
+ after = memory.getMemoryConsumptionForThisTask - other.getUsed
+ other.free(other.getUsed)
+ // Not Long.MaxValue, which the releases still to come would overflow.
+ memoryManager.limit(1L << 40)
+ }
+ }
+ (rows, before, after)
+ }
+
+ test("the survivor buffer is accounted to the task's execution memory") {
+ // What is buffered is held while a spliced row group is emitted, each
batch's share released as
+ // it goes, and the rest at the end. Off heap, only the ranges are
accounted, since the vectors'
+ // native memory is not the memory manager's.
+ withTempDir { dir =>
+ val path = twoRowGroupsOf20(dir)
+ Seq(false -> Seq(130L, 94L, 58L, 220L, 184L, 148L, 112L, 76L),
+ true -> Seq.fill(8)(40L)).foreach { case (offHeap, expected) =>
+ val context = memoryTask(new
TestMemoryManager(spark.sparkContext.conf))
+ val memory = context.taskMemoryManager
+ var held = Seq.empty[Long]
+ val rows = readInTask(path, context, offHeap) { i =>
+ if (i == 0) held :+= memory.getMemoryConsumptionForThisTask
+ }
+ assert(rows == marked(10L until 40L, checked = true), s"off heap
$offHeap: got $rows")
+ assert(held == expected,
+ s"off heap $offHeap: each batch must hold what is still buffered;
got $held")
+ assert(memory.getMemoryConsumptionForThisTask == 0, s"off heap
$offHeap: not released")
+ }
+ // A reader closed after its first batch, as under a limit, releases
what it still holds.
+ val context = memoryTask(new TestMemoryManager(spark.sparkContext.conf))
+ val file = PartitionedFile(
+ InternalRow.empty, SparkPath.fromPathString(path), 0, new
File(path).length())
+ val heldAtFirst =
withSQLConf(SQLConf.PARQUET_VECTORIZED_READER_BATCH_SIZE.key -> "4") {
+ asTask(context) {
+ drain(formatReader(checkedSchema(), Nil,
Seq(keyAtLeast(10L)))(file)) { it =>
+ it.next()
+ context.taskMemoryManager.getMemoryConsumptionForThisTask
+ }
+ }
+ }
+ assert(heldAtFirst == 130L,
+ s"the first batch must hold the first row group; got $heldAtFirst")
+ assert(context.taskMemoryManager.getMemoryConsumptionForThisTask == 0,
+ "a reader closed early must release what it holds")
+ }
+ }
+
+ test("memory the task's memory manager denies gives splicing up, and then
the filter") {
+ withTempDir { dir =>
+ val path = twoRowGroupsOf20(dir)
+ def limitedTask(available: Long): TaskContext = {
+ val memoryManager = new TestMemoryManager(spark.sparkContext.conf)
+ memoryManager.limit(available)
+ memoryTask(memoryManager)
+ }
+ def rows(available: Long): Seq[(Long, Boolean)] =
+ readInTask(path, limitedTask(available))(_ => ())
+ def bytes(available: Long, cap: String = "64MB"): Long = withSQLConf(
+ SQLConf.PARQUET_VECTORIZED_READER_BATCH_SIZE.key -> "4",
+
SQLConf.PARQUET_STORAGE_FILTER_PUSHDOWN_MAX_SPLICED_ROW_GROUP_BYTES.key -> cap)
{
+ asTask(limitedTask(available)) {
+ readCounting(path, formatReader(
+ checkedSchema(), Nil, Seq(keyAtLeast(10L)),
countingHadoopConf()))._2
+ }
+ }
+ val plain = rows(Long.MaxValue)
+ assert(plain == marked(10L until 40L, checked = true), s"got $plain")
+ // 105 bytes hold the 40 of the first row group's range and 7 of its
keys, but not the 8th.
+ // So splicing is given up there, as a cap of 105 bytes gives it up. The
read must cost
+ // exactly what the cap does, the key column read again, while the
filter stays.
+ val keysDenied = rows(105L)
+ assert(keysDenied == plain, s"denied keys must keep the filter; got
$keysDenied")
+ val cappedBytes = bytes(Long.MaxValue, cap = "105b")
+ assert(cappedBytes > bytes(Long.MaxValue), "a cap of 105 bytes must give
splicing up")
+ val keysDeniedBytes = bytes(105L)
+ assert(keysDeniedBytes == cappedBytes,
+ s"denied keys must read what the cap reads; got $keysDeniedBytes
against $cappedBytes")
+ // Nothing at all does not hold the ranges either, so the filter is
given up.
+ assert(rows(0L) == marked(0L until 40L, checked = false),
+ "denied ranges must give the filter up")
+ }
+ }
+
+ test("the survivor buffer asks for little more memory than it holds") {
+ // Another consumer holds all but 10,000 bytes, and the buffer never holds
more than 220. A
+ // reservation asking for more than is free would spill that consumer.
+ withTempDir { dir =>
+ val path = twoRowGroupsOf20(dir)
+ val memoryManager = new TestMemoryManager(spark.sparkContext.conf)
+ val held = 1000000L
+ memoryManager.limit(held + 10000L)
+ val context = memoryTask(memoryManager)
+ val other = new TestMemoryConsumer(context.taskMemoryManager)
+ other.use(held)
+ val rows = readInTask(path, context)(_ => ())
+ assert(rows == marked(10L until 40L, checked = true), s"got $rows")
+ assert(other.getUsed == held, "the read must not spill another consumer")
+ }
+ }
+
+ test("a spill while a spliced row group is emitted frees its survivors, and
its keys are read " +
Review Comment:
**The re-read bookkeeping is not pinned: this is the only test that drives
it, and it spills once, after the first batch of block 0, over one range, one
page and one long key.** Each of these one-line mutations should still pass the
whole suite:
- dropping `rowsEmittedInRowGroup = 0L;`
(LateMaterializationParquetRecordReader.java:512), since the field starts at 0
for block 0, and block 1 here survives whole, so it is never spilled;
- `rowsEmittedInRowGroup = columnarBatch.numRows()` instead of `+=`
(LateMaterializationParquetRecordReader.java:378), since the two agree after
the first batch;
- dropping `loadedBlockIndex = blockIdx;`
(LateMaterializationParquetRecordReader.java:626), since the default 0 is right
for block 0;
- replacing only `keyColumns[0]`'s reader
(LateMaterializationParquetRecordReader.java:1042-1045), since there is one key.
In production a spill can come at any batch of any row group, so each of
these would put keys on other rows, marked checked, which the post-scan
`Filter` trusts. `requireRowCount` catches it only when the row counts
disagree. I checked these by reading the tests and did not run them. Could you
add spills after the second batch and at the first batch of block 1, e.g. over
`twoRowGroupsOf20` with the even-keys filter so that both row groups are read
in part, one over `writePagedFile` whose remaining rows start on a later page,
and one over the two-key fixture of "two key columns are both charged"?
Deriving `loadedBlockIndex` from `nextBlockIndex - 1`, and the rows left from
the base's `totalCountLoadedSoFar - rowsReturned` through a protected accessor,
would also remove two of these mirrors.
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/LateMaterializationParquetRecordReader.java:
##########
@@ -183,20 +188,44 @@ public class LateMaterializationParquetRecordReader
extends VectorizedParquetRec
private WritableColumnVector checkedVector;
/**
- * Splicing state. {@link #survivorBatches} holds one array of survivor
vectors per batch to emit,
- * indexed like {@link #keyColumns}, and {@link #currentKeyAccumulators} the
array still filling.
- * Together they hold every survivor of one row group, outside any
MemoryConsumer.
+ * Splicing state. {@link #survivorBatches} holds one {@link SurvivorBatch}
per batch to emit, and
+ * {@link #currentKeyAccumulators} the vectors still filling, indexed like
{@link #keyColumns}.
+ * Together they hold the survivors of one row group not yet emitted.
*/
- private final ArrayDeque<WritableColumnVector[]> survivorBatches = new
ArrayDeque<>();
+ private final ArrayDeque<SurvivorBatch> survivorBatches = new ArrayDeque<>();
/** The survivor vectors still filling, or null when there are none. */
private WritableColumnVector[] currentKeyAccumulators;
- /** Row count of {@link #currentKeyAccumulators}. */
+ /** Row count and charge of {@link #currentKeyAccumulators}. */
private int currentKeyAccumulatorRowCount;
+ private long currentKeyAccumulatorCharge;
/**
- * The survivor vectors the current batch is built on. The next batch frees
them, or
- * {@link #close()} when no next batch comes.
+ * The survivor vectors the current batch is built on, and their charge. The
next batch frees
+ * them, or {@link #close()} when no next batch comes.
*/
private WritableColumnVector[] publishedKeyVectors;
+ private long publishedKeyCharge;
+ /** 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;
+ /** Whether the key values are accounted too, which needs them on heap. */
+ private boolean accountKeys;
+ /** The thread that reads, the only one whose spill may touch this reader's
state. */
+ private volatile Thread readerThread;
+
+ /** The block, final ranges and row count of the row group being emitted. */
+ private int loadedBlockIndex;
+ private RowRanges loadedRanges;
+ private long loadedRowCount;
+ /** Whether that row group can be spilled, which {@link #loadNextRowGroup}
decides. */
+ private boolean spillableRowGroup;
+ /** The rows of that row group already emitted. */
+ private long rowsEmittedInRowGroup;
+ /** Set by a spill, so the next batch reads the keys of the rest of the row
group again. */
+ private boolean rereadRemainingKeys;
+ /** The key pages read again after a spill, closed with the row group. */
+ private PageReadStore rereadKeyPages;
Review Comment:
Nit: These two new groups are inserted inside the "Splicing state." block,
before its last two members, so `persistentBatchColumns` and
`spliceBatchColumns` (L229-232) now sit under "The block, final ranges and row
count of the row group being emitted.", although they are per-file vectors that
`takeOverBatch` sets. AGENTS.md asks a new member to "follow the sectioning the
file already uses and put the new member with the code it belongs with". Could
you move L209-228 below `spliceBatchColumns`?
--
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]
