peter-toth commented on code in PR #59639:
URL: https://github.com/apache/spark/pull/59639#discussion_r4346741305
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetStorageFilterSuite.scala:
##########
@@ -1179,6 +1207,8 @@ class ParquetStorageFilterSuite extends QueryTest with
SharedSparkSession
s"rows 1-100 and 201-300 must splice and rows 101-200 must not; got
$splicedByRowGroup")
assert(m.rowGroupsSkipped.value == 1,
s"expected the row group with no survivor skipped; got
${m.rowGroupsSkipped.value}")
+ assert(m.rowGroupsSplicingGivenUpByCap.value == 1,
Review Comment:
Done in
[4b28a3f1a8f](https://github.com/apache/spark/pull/59639/commits/4b28a3f1a8fd03ff47b17333e2349f070a23af69).
Both check all four through `budgetGiveUps`.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetTestAccess.scala:
##########
@@ -236,7 +236,8 @@ object ParquetTestAccess {
def queuedSurvivorSets(
reader: LateMaterializationParquetRecordReader):
Seq[Seq[WritableColumnVector]] =
survivorBatchesField.get(reader)
-
.asInstanceOf[java.util.ArrayDeque[Array[WritableColumnVector]]].asScala.map(_.toSeq).toSeq
+
.asInstanceOf[java.util.ArrayDeque[LateMaterializationParquetRecordReader.SurvivorBatch]]
Review Comment:
Done in
[4b28a3f1a8f](https://github.com/apache/spark/pull/59639/commits/4b28a3f1a8fd03ff47b17333e2349f070a23af69).
The record is private again, and `ParquetTestAccess` looks up `vectors()` by
name, held in a lazy val.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetStorageFilterSuite.scala:
##########
@@ -2119,6 +2163,206 @@ class ParquetStorageFilterSuite extends QueryTest with
SharedSparkSession
}
}
+ // A task whose memory manager grants at most `limit` bytes.
+ private def memoryTask(limit: Long = Long.MaxValue): TaskContextImpl = {
+ val memoryManager = new TestMemoryManager(spark.sparkContext.conf)
+ memoryManager.limit(limit)
+ new TaskContextImpl(stageId = 0, stageAttemptNumber = 0, partitionId = 0,
taskAttemptId = 0,
+ attemptNumber = 0, numPartitions = 1,
+ taskMemoryManager = new TaskMemoryManager(memoryManager, 0),
+ localProperties = new Properties, metricsSystem = null)
+ }
+
+ // Reads `path` under `k >= 10` in batches of `batchSize` and a cap of
`cap`, as a task of
+ // `context`. Returns the rows, and what the task holds as each batch is
read. The confs are set
+ // outside the task, which sees a read-only one.
+ private def readInTask(
+ path: String,
+ context: TaskContext,
+ offHeap: Boolean = false,
+ batchSize: Int = 4,
+ cap: String = "64MB"): (Seq[(Long, Boolean)], Seq[Long]) = withSQLConf(
+ SQLConf.PARQUET_VECTORIZED_READER_BATCH_SIZE.key -> batchSize.toString,
+ SQLConf.COLUMN_VECTOR_OFFHEAP_ENABLED.key -> offHeap.toString,
+ SQLConf.PARQUET_STORAGE_FILTER_PUSHDOWN_MAX_SPLICED_ROW_GROUP_BYTES.key
-> cap) {
+ var held = Seq.empty[Long]
+ val rows = asTask(context) {
+ readChecked(path, checkedSchema(), Seq(keyAtLeast(10L)), keyOf = (b, i)
=> {
+ if (i == 0) held :+=
context.taskMemoryManager().getMemoryConsumptionForThisTask
+ b.column(0).getLong(i)
+ })
+ }
+ (rows, held)
+ }
+
+ // One row group of 400,000 rows of a long key and a string value. Under `k
>= 10` its 399,990
+ // survivors fall in one range and buffer 3,599,950 bytes, several steps of
the reservation.
+ private def oneRowGroupOf400000(dir: File): String = {
Review Comment:
Done in
[4b28a3f1a8f](https://github.com/apache/spark/pull/59639/commits/4b28a3f1a8fd03ff47b17333e2349f070a23af69).
The step test reads the same file again for the margin check, so it is written
once.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetStorageFilterSuite.scala:
##########
@@ -2119,6 +2163,206 @@ class ParquetStorageFilterSuite extends QueryTest with
SharedSparkSession
}
}
+ // A task whose memory manager grants at most `limit` bytes.
+ private def memoryTask(limit: Long = Long.MaxValue): TaskContextImpl = {
+ val memoryManager = new TestMemoryManager(spark.sparkContext.conf)
+ memoryManager.limit(limit)
+ new TaskContextImpl(stageId = 0, stageAttemptNumber = 0, partitionId = 0,
taskAttemptId = 0,
+ attemptNumber = 0, numPartitions = 1,
+ taskMemoryManager = new TaskMemoryManager(memoryManager, 0),
+ localProperties = new Properties, metricsSystem = null)
+ }
+
+ // Reads `path` under `k >= 10` in batches of `batchSize` and a cap of
`cap`, as a task of
+ // `context`. Returns the rows, and what the task holds as each batch is
read. The confs are set
+ // outside the task, which sees a read-only one.
+ private def readInTask(
+ path: String,
+ context: TaskContext,
+ offHeap: Boolean = false,
+ batchSize: Int = 4,
+ cap: String = "64MB"): (Seq[(Long, Boolean)], Seq[Long]) = withSQLConf(
+ SQLConf.PARQUET_VECTORIZED_READER_BATCH_SIZE.key -> batchSize.toString,
+ SQLConf.COLUMN_VECTOR_OFFHEAP_ENABLED.key -> offHeap.toString,
+ SQLConf.PARQUET_STORAGE_FILTER_PUSHDOWN_MAX_SPLICED_ROW_GROUP_BYTES.key
-> cap) {
+ var held = Seq.empty[Long]
+ val rows = asTask(context) {
+ readChecked(path, checkedSchema(), Seq(keyAtLeast(10L)), keyOf = (b, i)
=> {
+ if (i == 0) held :+=
context.taskMemoryManager().getMemoryConsumptionForThisTask
+ b.column(0).getLong(i)
+ })
+ }
+ (rows, held)
+ }
+
+ // One row group of 400,000 rows of a long key and a string value. Under `k
>= 10` its 399,990
+ // survivors fall in one range and buffer 3,599,950 bytes, several steps of
the reservation.
+ private def oneRowGroupOf400000(dir: File): String = {
+ val path = writeSingleParquetFile(dir,
+ spark.range(400000L).selectExpr("id AS k", "CAST(id AS STRING) AS v"),
+ rowGroupSize = ParquetWriter.DEFAULT_BLOCK_SIZE)
+ assert(footerOf(path).getBlocks.size == 1, "the fixture must be one row
group")
+ path
+ }
+
+ test("the survivor buffer is accounted to the task's execution memory") {
+ // What a row group buffers is held while it is emitted, and released when
the next one loads.
+ // Phase 1 reserves a little ahead of what it buffers, and gives that back
before phase 2. Off
+ // heap, only the ranges are accounted, since the vectors' native memory
is not the memory
+ // manager's.
+ withTempDir { dir =>
+ val path = twoRowGroupsOf20(dir)
+ // Past a cap of 100 bytes, at the seventh survivor of each row group,
the keys go, a full
+ // batch of them already queued, and only the range is held.
+ Seq((false, "64MB", Seq(130L, 130L, 130L, 220L, 220L, 220L, 220L, 220L)),
+ (true, "64MB", Seq.fill(8)(40L)),
+ (false, "100b", Seq.fill(8)(40L))).foreach { case (offHeap, cap,
expected) =>
+ val context = memoryTask()
+ val (rows, held) = readInTask(path, context, offHeap, cap = cap)
+ val name = s"off heap $offHeap under a cap of $cap"
+ assert(rows == marked(10L until 40L, checked = true), s"$name: got
$rows")
+ assert(held == expected, s"$name: each batch must hold what is still
buffered; got $held")
Review Comment:
Done in
[4b28a3f1a8f](https://github.com/apache/spark/pull/59639/commits/4b28a3f1a8fd03ff47b17333e2349f070a23af69),
with your wording.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetStorageFilterSuite.scala:
##########
@@ -2119,6 +2163,206 @@ class ParquetStorageFilterSuite extends QueryTest with
SharedSparkSession
}
}
+ // A task whose memory manager grants at most `limit` bytes.
+ private def memoryTask(limit: Long = Long.MaxValue): TaskContextImpl = {
+ val memoryManager = new TestMemoryManager(spark.sparkContext.conf)
+ memoryManager.limit(limit)
+ new TaskContextImpl(stageId = 0, stageAttemptNumber = 0, partitionId = 0,
taskAttemptId = 0,
+ attemptNumber = 0, numPartitions = 1,
+ taskMemoryManager = new TaskMemoryManager(memoryManager, 0),
+ localProperties = new Properties, metricsSystem = null)
+ }
+
+ // Reads `path` under `k >= 10` in batches of `batchSize` and a cap of
`cap`, as a task of
+ // `context`. Returns the rows, and what the task holds as each batch is
read. The confs are set
+ // outside the task, which sees a read-only one.
+ private def readInTask(
Review Comment:
Done in
[4b28a3f1a8f](https://github.com/apache/spark/pull/59639/commits/4b28a3f1a8fd03ff47b17333e2349f070a23af69).
`readChecked` takes the task, and `readInTask` passes its confs through `conf`.
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/LateMaterializationParquetRecordReader.java:
##########
@@ -202,6 +221,18 @@ public class LateMaterializationParquetRecordReader
extends VectorizedParquetRec
/** The array the emitted batch is built over, whose key slots each batch
rewrites. */
private ColumnVector[] spliceBatchColumns;
+ /** What the survivor ranges of the row group being read or emitted cost to
hold. */
+ private long rangeCharge;
+ /** The survivor buffer's account with the task's memory manager, or null
outside a task. */
+ private SurvivorMemory survivorMemory;
+ /**
+ * What {@link #survivorMemory} holds. Only this reader changes it, so it is
kept here rather
+ * than read off the consumer's atomic counter per survivor.
+ */
+ private long reservedBytes;
+ /** Whether the key values are accounted too, which needs them on heap. */
+ private boolean accountKeys;
Review Comment:
Done in
[4b28a3f1a8f](https://github.com/apache/spark/pull/59639/commits/4b28a3f1a8fd03ff47b17333e2349f070a23af69),
as you suggested.
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/LateMaterializationParquetRecordReader.java:
##########
@@ -663,11 +718,12 @@ private RowRanges evaluateStorageFilter(
PrimitiveIterator.OfLong rowIndexIter = pushedFilterRanges.iterator();
RowRanges.Builder finalRangesBuilder = RowRanges.builder();
- // The budget weighs both the bytes buffered for splicing and the
survivors' ranges.
- long splicedBytes = 0L;
- long rangeBytes = 0L;
+ // The budget weighs both the key values buffered for splicing and the
survivors' ranges.
+ rangeCharge = 0L;
Review Comment:
Done in
[4b28a3f1a8f](https://github.com/apache/spark/pull/59639/commits/4b28a3f1a8fd03ff47b17333e2349f070a23af69).
--
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]
