viirya commented on code in PR #59639:
URL: https://github.com/apache/spark/pull/59639#discussion_r4333331832
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:
##########
@@ -2027,8 +2030,8 @@ object SQLConf {
"false, the reader can only skip whole row groups, and it reads the
key columns twice in " +
"the others. With a page index, it can be slower when the matching
rows are spread " +
"across most pages. The matching key values of a row group are
buffered while it is " +
- s"read, up to
${PARQUET_STORAGE_FILTER_PUSHDOWN_MAX_SPLICED_ROW_GROUP_BYTES.key}, " +
- "outside Spark's memory manager.")
+ s"read, up to
${PARQUET_STORAGE_FILTER_PUSHDOWN_MAX_SPLICED_ROW_GROUP_BYTES.key}, in " +
+ "the task's execution memory unless column vectors are off heap.")
Review Comment:
"in the task's execution memory unless column vectors are off heap" reads as
if nothing is accounted off heap. The row ranges are accounted off heap too,
and with scattered survivors they are the larger part. Maybe: "The matching key
values and row ranges of a row group are buffered while it is read, up to ...,
and counted against the task's execution memory, except key values held in
off-heap column vectors."
##########
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);
+ this.reader = reader;
+ }
+
+ @Override
+ public long spill(long size, MemoryConsumer trigger) {
Review Comment:
I suggested a spillable consumer on #58895, but seeing what it takes here,
I'm not sure the spill earns its place in this PR. It frees memory only when
all of these hold:
- the column vectors are on heap;
- phase 2 reads the row group in part;
- the spill comes from the reading thread (so never under a pipelined Python
UDF, whose writer thread drives the scan);
- the rows' keys are still queued, not yet emitted.
It never frees the ranges. In exchange it brings the store pre-build (item
3), `replaceColumnReader`, the `loaded*` / `rowsEmittedInRowGroup` mirrors, the
thread confinement and the key re-read, and the re-read can allocate more than
the spill freed (item 1). Accounting plus denial already gives the memory
manager what it needs: other consumers spill before this one is refused, and a
refused reader backs off.
Would you consider having `spill()` return 0 here, and adding the spill in a
follow-up once several readers per task exist (the SPJ case in the
description), where its benefit can be measured? That would also resolve items
1, 3, 4, 7, 9, 13 and 14 for this PR.
##########
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:
Agreed. Wrapping the two reservations the way `throwIfKilled` wraps a kill,
in a checked `SparkException` around the cause, looks like a cheap fix.
`FileScanRDD` then can't take another consumer's spill failure for a corrupt
file under `ignoreCorruptFiles`. A kill that lands during that spill is already
safe: `TaskMemoryManager` rethrows it as a `RuntimeException`, which the same
wrap would cover, and the task is still killed afterwards anyway.
##########
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:
This reproduces as you predicted. In "the survivor buffer asks for little
more memory than it holds", with `memoryManager.limit(held + 10000L)` the other
consumer keeps `1000000`. With `memoryManager.limit(held + 300L)`, which is
still more than the 220 bytes the buffer ever holds, the rows are still right,
but `other.getUsed` drops to `0`: the whole other consumer is spilled. So the
test should also cover a free amount close to what the buffer holds, and the PR
description's "a buffer that fits in the free memory spills no other consumer"
doesn't hold yet. Since the ask is not clamped to the cap either, near the
default 64 MiB the reservation can briefly reach about twice
`maxSplicedRowGroupBytes`, until L629 trims it.
##########
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:
One consideration on keeping the filter without reserving: with scattered
survivors, the ranges are the larger part of this buffer. Each costs
`ROW_RANGE_BYTES` (40 B), against 9 B for a long key. Continuing unaccounted
after a range denial would leave most of the buffer outside the memory manager,
which is what this PR sets out to fix. So I'd lean toward keeping the give-up,
but logging it once per split and counting it in a metric, so that a scan
losing its filter under memory pressure is visible. Your point about the
repeated acquire off heap stands either way.
##########
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:
I ran the spill at other points against this commit, using a variant of
`readWithSpill` that spills at the n-th batch instead of the first:
- over `twoRowGroupsOf20` with `k % 2 == 0` and batch size 4, so both row
groups are read in part: at each of the six batches, which includes the 2nd and
3rd batches of block 0 and every batch of block 1;
- over `writePagedFile` with `k % 3 == 0` and batch size 16, so the
remaining rows start on later pages: at batches 1, 2 and 5.
Every case returned exactly the expected keys, all marked checked. A spill
at a row group's last batch freed nothing, as it should, since the queue is
empty there. So the current bookkeeping looks right, and these cases would pin
the mutations you list. Could they be added? One spill over the two-key fixture
would also cover replacing every key's reader, not only `keyColumns[0]`'s.
--
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]
