peter-toth commented on code in PR #59639:
URL: https://github.com/apache/spark/pull/59639#discussion_r4350558389
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/LateMaterializationParquetRecordReader.java:
##########
@@ -912,6 +1024,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 MEMORY_MODE == MemoryMode.ON_HEAP ? 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 (Exception | InternalError e) {
+ // Another consumer's spill and the eviction of cached blocks run
inside, so this can raise
Review Comment:
Done in
[09b8b2186e3](https://github.com/apache/spark/pull/59639/commits/09b8b2186e31b49f1d5d864601739f989294f582).
The comment at the catch says so, and the description's Errors bullet has the
matching sub-bullet.
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/LateMaterializationParquetRecordReader.java:
##########
@@ -118,6 +127,13 @@ public class LateMaterializationParquetRecordReader
extends VectorizedParquetRec
*/
private static final int DEFAULT_CHILD_BYTES_PER_ROW = 4;
+ /**
+ * The most the survivor buffer reserves beyond what it holds, see
+ * {@link #reserveSurvivorMemory}. The task's memory manager spills other
consumers to grant a
Review Comment:
Done in
[09b8b2186e3](https://github.com/apache/spark/pull/59639/commits/09b8b2186e31b49f1d5d864601739f989294f582),
with your wording in this javadoc and in the description's No spill bullet.
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/LateMaterializationParquetRecordReader.java:
##########
@@ -500,7 +551,13 @@ protected boolean loadNextRowGroup() throws IOException {
// Null means phase 1 gave the filter up, and phase 2 then reads every
row the pushed filter
// kept.
boolean filterGivenUp = survivors == null;
- if (filterGivenUp) abandonSplicing();
+ if (filterGivenUp) {
+ rangeCharge = 0L;
Review Comment:
Done in
[09b8b2186e3](https://github.com/apache/spark/pull/59639/commits/09b8b2186e31b49f1d5d864601739f989294f582).
Both give-up sites call a new `abandonFilter()`, which drops the ranges'
charge. The new test "a row group whose filter phase 1 gives up holds no memory
while it is read whole" reads your two cases. It expects `(0, 0, 2, 0)` and
`(0, 0, 0, 2)`, every row unchecked, and nothing held at any batch. Dropping
the reset now fails it, and the phase-2 test of item 17.
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/LateMaterializationParquetRecordReader.java:
##########
@@ -920,6 +1089,51 @@ 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 frees nothing. The buffer is bounded by {@code
maxSplicedRowGroupBytes} and
+ * drains as its batches are emitted, and freeing it would have the rest of
the row group read
+ * its key pages again, which the memory manager does not see either. A
denied reservation is
+ * handled where it is made, as passing that bound is. 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. The task's
memory manager keeps
+ * every consumer until the task ends, so a task holds one per file it
reserved for.
+ */
+ private static final class SurvivorMemory extends MemoryConsumer {
+ SurvivorMemory(TaskMemoryManager taskMemoryManager) {
+ super(taskMemoryManager, MemoryMode.ON_HEAP);
Review Comment:
Done in
[09b8b2186e3](https://github.com/apache/spark/pull/59639/commits/09b8b2186e31b49f1d5d864601739f989294f582).
`ObservedTaskMemoryManager` passes `consumer.getMode`, and each case of the
accounted test expects only `ON_HEAP` asks. For the record, "off-heap column
vectors through the planner" already failed under that mutation at 4b28a3f. The
real memory manager has no off-heap pool there, so the reader gave the filter
up and returned 100 rows instead of 51. The new check names the cause.
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/LateMaterializationParquetRecordReader.java:
##########
@@ -727,15 +785,30 @@ 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) splicingGiveUp =
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 held = heldCharge();
+ boolean denied = !reserveSurvivorMemory(held);
+ if (denied && held > rangeCharge) {
Review Comment:
Done in
[09b8b2186e3](https://github.com/apache/spark/pull/59639/commits/09b8b2186e31b49f1d5d864601739f989294f582).
The denial test now expects one 40-byte ask per row group off heap under
`memoryTask(0L)`, and `(0, 0, 0, 2)` under `keyAtLeast(19L)` with nothing free.
Each of the two mutations fails it.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetStorageFilterSuite.scala:
##########
@@ -2119,6 +2199,270 @@ class ParquetStorageFilterSuite extends QueryTest with
SharedSparkSession
}
}
+ // A task memory manager that calls `onAcquire` before each acquisition but
a test consumer's,
+ // with what is asked for and what the task holds then.
+ private class ObservedTaskMemoryManager(
+ memoryManager: MemoryManager,
+ onAcquire: (Long, Long) => Unit) extends
TaskMemoryManager(memoryManager, 0) {
+ override def acquireExecutionMemory(required: Long, consumer:
MemoryConsumer): Long = {
+ if (!consumer.isInstanceOf[TestMemoryConsumer]) {
+ onAcquire(required, getMemoryConsumptionForThisTask)
+ }
+ super.acquireExecutionMemory(required, consumer)
+ }
+ }
+
+ // A task whose memory manager grants at most `limit` bytes, and calls
`onAcquire` as
+ // `ObservedTaskMemoryManager` does.
+ private def memoryTask(
+ limit: Long = Long.MaxValue,
+ onAcquire: (Long, Long) => Unit = (_, _) => ()): 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 ObservedTaskMemoryManager(memoryManager,
onAcquire),
+ localProperties = new Properties, metricsSystem = null)
+ }
+
+ // Reads `path` under `filter` 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.
+ private def readInTask(
+ path: String,
+ context: TaskContext,
+ offHeap: Boolean = false,
+ batchSize: Int = 4,
+ cap: String = "64MB",
+ filter: Expression = keyAtLeast(10L),
+ schema: StructType = checkedSchema()): (Seq[(Long, Boolean)], Seq[Long])
= {
+ var held = Seq.empty[Long]
+ val rows = readChecked(path, schema, Seq(filter), conf = Map(
+ 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),
+ keyOf = (b, i) => {
+ if (i == 0) held :+=
context.taskMemoryManager().getMemoryConsumptionForThisTask
+ b.column(0).getLong(i)
+ },
+ context = Some(context))
+ (rows, held)
+ }
+
+ // Reads `path` under `k >= 10` while another consumer holds all of the
task's memory but
+ // `margin`, and checks that the read neither drops rows nor spills that
consumer.
+ private def readBesideOther(
+ path: String,
+ rows: Long,
+ margin: Long,
+ cap: String = "64MB"): Unit = {
+ val held = 1000000L
+ val context = memoryTask(held + margin)
+ val other = new TestMemoryConsumer(context.taskMemoryManager)
+ other.use(held)
+ val (read, _) = readInTask(path, context, batchSize = 4096, cap = cap)
+ val name = s"$rows rows with a margin of $margin under a cap of $cap"
+ assert(read == marked(10L until rows, checked = true), s"$name: got
${read.size} rows")
+ assert(other.getUsed == held, s"$name: the read must not spill another
consumer")
+ }
+
+ // 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 its row group's reservation until the
next row group " +
+ s"loads; got $held")
+ assert(context.taskMemoryManager.getMemoryConsumptionForThisTask == 0,
+ s"$name: not released")
+ }
+ // On heap, each row group asks first for its first survivor's 49 bytes,
with nothing held
+ // since the previous one's release, and then for as much again as it
holds, three and four
+ // times for the 130 and 220 bytes it reaches.
+ var asked = Seq.empty[(Long, Long)]
+ readInTask(path, memoryTask(onAcquire = (size, held) => asked :+=
((size, held))))
+ assert(asked == Seq((49L, 0L), (49L, 49L), (98L, 98L), (49L, 0L), (49L,
49L), (98L, 98L),
+ (196L, 196L)), s"got $asked")
+ // A reader closed after its first batch, as under a limit, releases
what it still holds.
+ val context = memoryTask()
+ 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("emitting a row group releases its memory in steps") {
+ // Each batch emitted frees its survivors, but the reservation is cut only
once it is a step
+ // past what is still buffered, so a large row group gives its memory back
a few times rather
+ // than per batch.
+ withTempDir { dir =>
+ val path = oneRowGroupOf400000(dir)
+ var asked = Seq.empty[Long]
+ val context = memoryTask(onAcquire = (size, _) => asked :+= size)
+ val (rows, reserved) = readInTask(path, context, batchSize = 4096)
+ assert(rows.size == 399990, s"got ${rows.size} rows")
+ // What batch `b` still holds, its own survivors and those after it at 9
bytes each, and the
+ // range. A batch frees 36,864 bytes, so the reservation steps down
every 29 batches.
+ def held(b: Int): Long = 40L + 9L * (399990L - 4096L * b)
+ val expected = (0 until 98).map(b => held(b - b % 29))
+ assert(reserved == expected, s"expected the reservation to step down at
batches 29, 58 " +
+ s"and 87; got ${reserved.distinct}")
+ assert(context.taskMemoryManager.getMemoryConsumptionForThisTask == 0,
"not released")
+ // A growing buffer asks a few times. First for the first survivor, then
for as much again as
+ // it holds up to 802,816 bytes, and then 1 MiB at a time.
+ assert(asked == Seq(49L) ++ (0 to 14).map(49L << _) ++ Seq.fill(2)(1L <<
20),
+ s"got ${asked.size} asks: ${asked.take(20)}")
+ // The 3,599,950 bytes it ends up buffering are not asked for twice
over. Asking for as much
+ // again as it holds would spill the other consumer.
+ readBesideOther(path, rows = 400000L, margin = 3599950L + (1L << 20))
+ }
+ }
+
+ test("memory the task's memory manager denies gives splicing up, and then
the filter") {
+ withTempDir { dir =>
+ val path = twoRowGroupsOf20(dir)
+ // Each read must leave the task holding nothing.
+ def released(context: TaskContext, available: Long): Unit =
+ assert(context.taskMemoryManager().getMemoryConsumptionForThisTask ==
0,
+ s"a read with $available bytes free must release what it reserved")
+ def rows(available: Long, offHeap: Boolean = false): Seq[(Long,
Boolean)] = {
+ val context = memoryTask(available)
+ val (read, _) = readInTask(path, context, offHeap)
+ released(context, available)
+ read
+ }
+ // The bytes a read transfers, with the row groups the budget limited,
see `budgetGiveUps`.
+ def bytes(available: Long, cap: String = "64MB"): (Long, Seq[Long]) = {
+ val metrics = metricMap()
+ val context = memoryTask(available)
+ val (_, read) = readUnderCap(path, checkedSchema(),
Seq(keyAtLeast(10L)), cap,
+ batchSize = 4, metrics, Some(context))
+ released(context, available)
+ (read, budgetGiveUps(metrics))
+ }
+ 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, and in the second row group, 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. The metrics tell the two causes apart.
+ val keysDenied = rows(105L)
+ assert(keysDenied == plain, s"denied keys must keep the filter; got
$keysDenied")
+ // 50 bytes take the first survivor's 49, and then grant 1 of the 9 the
second one misses,
+ // which goes back.
+ assert(rows(50L) == plain, "a partly granted reservation must keep the
filter")
+ // 130 bytes grant the seventh survivor 32 of the 98 it asks for, which
covers the 5 it
+ // misses, so the first row group keeps splicing. The second is denied
at its eleventh.
+ val (_, partlyGranted) = bytes(130L)
+ assert(partlyGranted == Seq(0L, 1L, 0L, 0L),
+ s"only the second row group must be denied; got $partlyGranted")
+ val (uncappedBytes, uncapped) = bytes(Long.MaxValue)
+ assert(uncapped == Seq(0L, 0L, 0L, 0L), s"got $uncapped")
+ val (cappedBytes, capped) = bytes(Long.MaxValue, cap = "105b")
+ assert(cappedBytes > uncappedBytes, "a cap of 105 bytes must give
splicing up")
+ assert(capped == Seq(2L, 0L, 0L, 0L), s"both row groups pass the cap;
got $capped")
+ val ((keysDeniedBytes, denied), keysLogged) =
+ withLogged(splicingGivenUpOnMemory)(bytes(105L))
+ assert(keysDeniedBytes == cappedBytes,
+ s"denied keys must read what the cap reads; got $keysDeniedBytes
against $cappedBytes")
+ assert(denied == Seq(0L, 2L, 0L, 0L), s"both row groups are denied; got
$denied")
+ assert(keysLogged == 1, s"which is logged once per split; got
$keysLogged")
+ // Nothing at all does not hold the ranges either, so the filter is
given up. Off heap only
+ // the ranges are asked for, so 40 bytes keep the filter and nothing
gives it up.
+ assert(rows(0L) == marked(0L until 40L, checked = false),
+ "denied ranges must give the filter up")
+ val ((_, rangesDenied), rangesLogged) =
withLogged(filterGivenUpOnMemory)(bytes(0L))
+ assert(rangesDenied == Seq(0L, 0L, 0L, 2L),
+ s"both row groups count as giving the filter up, not splicing; got
$rangesDenied")
+ assert(rangesLogged == 1, s"which is logged once per split; got
$rangesLogged")
+ assert(rows(40L, offHeap = true) == plain, "off heap, ranges that fit
must keep the filter")
+ assert(rows(0L, offHeap = true) == marked(0L until 40L, checked = false),
+ "off heap, denied ranges must give the filter up")
+ }
+ }
+
+ test("an error raised while reserving fails the read in a form
ignoreCorruptFiles keeps") {
+ // Reserving the buffer can have the task's memory manager spill another
consumer, which can
+ // raise anything. A `RuntimeException`, which a kill's interrupt during a
spill also becomes,
+ // goes out checked, and an `InternalError` in a fatal `UnknownError`.
Under
+ // `ignoreCorruptFiles` either would otherwise be taken for a corrupt
file, and the rest of the
+ // file dropped.
+ withTempDir { dir =>
+ val path = twoRowGroupsOf20(dir)
+ Seq(new RuntimeException("spill"), new InternalError("spill")).foreach {
error =>
Review Comment:
Done in
[09b8b2186e3](https://github.com/apache/spark/pull/59639/commits/09b8b2186e31b49f1d5d864601739f989294f582).
The spill also throws `new RuntimeException("spill", new
OutOfMemoryError("deep"))`, and the test expects that `OutOfMemoryError`
itself. A two-case `memoryReservationFailed` fails it.
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/LateMaterializationParquetRecordReader.java:
##########
@@ -541,6 +599,7 @@ protected boolean loadNextRowGroup() throws IOException {
requireRowCount(dataPages, finalRowCount, blockIdx);
// A row group whose filter was given up reports nothing (see
`StorageFilterMetrics`).
Review Comment:
Done in
[09b8b2186e3](https://github.com/apache/spark/pull/59639/commits/09b8b2186e31b49f1d5d864601739f989294f582),
with your wording here, and "at most once" in `ParquetStorageFilter` and the
description.
--
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]
