peter-toth commented on code in PR #59639:
URL: https://github.com/apache/spark/pull/59639#discussion_r4339165212
##########
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:
Gone with the spill in
[778c14954a9](https://github.com/apache/spark/pull/59639/commits/778c14954a9ede1af0fc6357b0b5546322e69a23).
Thanks for running the cases, viirya.
--
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]
