peter-toth commented on code in PR #59639:
URL: https://github.com/apache/spark/pull/59639#discussion_r4350560773
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/LateMaterializationParquetRecordReader.java:
##########
@@ -337,17 +376,21 @@ private void decideAllKeysMissingFile() {
public boolean nextBatch() throws IOException {
closeAll(publishedKeyVectors);
publishedKeyVectors = null;
+ trimSurvivorMemory();
if (!super.nextBatch()) return false;
if (keyColumns != null) {
if (spliceCurrentRowGroup) {
- publishedKeyVectors = survivorBatches.pollFirst();
- if (publishedKeyVectors == null) {
+ SurvivorBatch next = survivorBatches.pollFirst();
+ if (next == null) {
// Unreachable, since the queue holds exactly the survivors the emit
loop counts. Without
// this check the key slots would show stale values.
throw ParquetStorageFilter.internalError(String.format(
"Storage-filter survivor queue of row group %d in %s ran out
with %d rows still "
+ "to emit", nextBlockIndex - 1, fileReader.getFile(),
columnarBatch.numRows()));
}
+ publishedKeyVectors = next.vectors();
+ // Still reserved, until the next batch frees these vectors and trims
the reservation.
Review Comment:
Done in
[09b8b2186e3](https://github.com/apache/spark/pull/59639/commits/09b8b2186e31b49f1d5d864601739f989294f582),
with your wording.
##########
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 =>
+ val context = memoryTask(1000L)
+ val failing = new TestMemoryConsumer(context.taskMemoryManager) {
+ override def spill(size: Long, trigger: MemoryConsumer): Long =
throw error
+ }
+ failing.use(1000L)
+ val thrown = intercept[Throwable](readInTask(path, context))
+ val expected = if (error.isInstanceOf[InternalError]) {
+ classOf[UnknownError]
+ } else {
+ classOf[SparkException]
+ }
+ assert(thrown.getClass == expected && (thrown.getCause eq error),
+ s"expected $error inside ${expected.getSimpleName}; got $thrown")
+ assert(!DataSourceUtils.shouldIgnoreCorruptFileException(thrown),
s"got $thrown")
+ }
+ // Evicting a cached block to disk can raise an `IOException` there too,
which the task's
+ // memory manager passes on as it is.
+ val full = new IOException("No space left on device")
+ val evicting = memoryTask(onAcquire = (_, _) => throw full)
+ val thrown = intercept[Throwable](readInTask(path, evicting))
+ assert(thrown.getClass == classOf[SparkException] && (thrown.getCause eq
full),
+ s"expected $full inside a SparkException; got $thrown")
+ assert(!DataSourceUtils.shouldIgnoreCorruptFileException(thrown), s"got
$thrown")
+ }
+ }
+
+ test("the survivor buffer asks for little more memory than it holds") {
+ // Another consumer holds all but what the buffer needs plus a margin. The
task's memory manager
+ // spills that consumer to grant a reservation past what is free, so the
buffer must not ask
+ // for more than its margin: less than as much again as it holds while
small, less than 1 MiB
+ // more once it is past that, see "emitting a row group releases its
memory in steps", and
+ // never past the cap.
+ withTempDir { dir =>
+ // The buffer never holds more than 220 bytes, so it asks for less than
440.
+ val path = twoRowGroupsOf20(dir)
+ readBesideOther(path, rows = 40L, margin = 440L)
+ // Doubling from 196 would ask for 392, which a cap of 300 bytes stops
at 300.
Review Comment:
Done in
[09b8b2186e3](https://github.com/apache/spark/pull/59639/commits/09b8b2186e31b49f1d5d864601739f989294f582),
with your wording.
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/LateMaterializationParquetRecordReader.java:
##########
@@ -202,6 +221,22 @@ 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. */
Review Comment:
Done in
[09b8b2186e3](https://github.com/apache/spark/pull/59639/commits/09b8b2186e31b49f1d5d864601739f989294f582).
##########
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
Review Comment:
Done in
[09b8b2186e3](https://github.com/apache/spark/pull/59639/commits/09b8b2186e31b49f1d5d864601739f989294f582).
The javadoc says that up to a batch per key column is held beyond the charge.
##########
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;
Review Comment:
Done in
[09b8b2186e3](https://github.com/apache/spark/pull/59639/commits/09b8b2186e31b49f1d5d864601739f989294f582).
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/LateMaterializationParquetRecordReader.java:
##########
@@ -243,9 +278,13 @@ public void close() throws IOException {
try {
closeSplicingState();
} finally {
- if (checkedVector != null) {
- checkedVector.close();
- checkedVector = null;
+ try {
+ if (checkedVector != null) {
+ checkedVector.close();
+ checkedVector = null;
+ }
+ } finally {
+ releaseSurvivorMemoryTo(0L);
Review Comment:
Done in
[09b8b2186e3](https://github.com/apache/spark/pull/59639/commits/09b8b2186e31b49f1d5d864601739f989294f582).
`close()` zeroes `rangeCharge` first, so `abandonSplicing()` releases
everything in one trip.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetStorageFilterSuite.scala:
##########
@@ -2048,6 +2099,35 @@ class ParquetStorageFilterSuite extends QueryTest with
SharedSparkSession
}
}
+ test("a row group whose filter phase 2 gives up does not count as giving
splicing up") {
Review Comment:
Done in
[09b8b2186e3](https://github.com/apache/spark/pull/59639/commits/09b8b2186e31b49f1d5d864601739f989294f582).
Both tests use `damagedMiddleRowGroup`, beside `twoRowGroupsOf20`, over
`threeRowGroupsOf100` and `evenKey`.
--
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]
