spark-test commented on code in PR #59786:
URL: https://github.com/apache/spark/pull/59786#discussion_r4383559974
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonArrowOutputSuite.scala:
##########
@@ -145,7 +152,194 @@ class PythonArrowOutputSuite extends SparkFunSuite {
checkAllRowsOnceInOrder(batchSizes, output)
// Guards the choice of cases: each has a batch that does not fit, which
must be sliced.
assert(output.size > batchSizes.size)
- assert(output.forall { case (rows, bytes) => rows.size <= 1 || bytes <
maxBytes })
+ assert(output.forall { case (rows, bytes) => rows.size <= 1 || bytes <=
maxBytes })
}
}
+
+ private def newSliceBytesProcessor(maxBytes: Int)
+ : (ArrowStreamReader, Map[String, SQLMetric]) => ArrowOutputProcessor =
+ (reader, metrics) => new SliceBytesArrowOutputProcessorImpl(reader,
metrics, maxBytes)
+
+ test("SPARK-60395: slice by bytes: a batch whose size equals the limit is
not sliced") {
+ val batchSizes = Seq(50, 50)
+ val batchBytes = readAll(batchSizes, new ArrowOutputProcessorImpl(_,
_)).map(_._2)
+ assert(batchBytes.distinct.size === 2)
+ val maxBytes = batchBytes.max.toInt
+ val output = readAll(batchSizes, newSliceBytesProcessor(maxBytes))
+ checkAllRowsOnceInOrder(batchSizes, output)
+ assert(output.map(_._1.size) === batchSizes)
+ // One byte less, the larger batch is sliced.
+ val sliced = readAll(batchSizes, newSliceBytesProcessor(maxBytes - 1))
+ checkAllRowsOnceInOrder(batchSizes, sliced)
+ assert(sliced.size > batchSizes.size)
+ assert(sliced.forall(_._2 <= maxBytes - 1))
+ }
+
+ test("SPARK-60395: slice by bytes: each search does not restart from all
remaining rows") {
+ val batchSizes = Seq(2000)
+ val maxBytes = (readAll(batchSizes, new ArrowOutputProcessorImpl(_,
_)).head._2 / 50).toInt
+ var measured = 0
+ val output = readAll(batchSizes, (reader, metrics) =>
+ new SliceBytesArrowOutputProcessorImpl(reader, metrics, maxBytes) {
+ override private[python] def batchBytes(root: VectorSchemaRoot): Long
= {
+ measured += 1
+ super.batchBytes(root)
+ }
+ })
+ checkAllRowsOnceInOrder(batchSizes, output)
+ assert(output.forall { case (rows, bytes) => rows.size <= 1 || bytes <=
maxBytes })
+ assert(output.size >= 50)
+ // The whole batch and the first search take about log(50) / log(1 / 0.9)
= 38 measurements,
+ // and each later slice one, or a few more where its rows are larger.
Restarting every search
+ // from all the remaining rows would take hundreds.
+ assert(measured <= 40 + 2 * output.size, s"$measured measurements for
${output.size} slices")
+ }
+
+ /**
+ * Slices `batchSizes` by bytes, where a row whose `id` is i takes
`rowBytes(i)` bytes, and
+ * returns the number of rows of every output batch.
+ */
+ private def sliceRowCounts(batchSizes: Seq[Int], maxBytes: Int)(rowBytes:
Long => Long)
+ : Seq[Int] = {
+ val output = readAll(batchSizes, (reader, metrics) =>
+ new SliceBytesArrowOutputProcessorImpl(reader, metrics, maxBytes) {
+ override private[python] def batchBytes(root: VectorSchemaRoot): Long
= {
+ val ids = root.getVector("id").asInstanceOf[BigIntVector]
+ (0 until root.getRowCount).map(i => rowBytes(ids.get(i))).sum
+ }
+ })
+ checkAllRowsOnceInOrder(batchSizes, output)
+ output.map(_._1.size)
+ }
+
+ test("SPARK-60395: slice by bytes: number of rows of each slice") {
+ // The first search goes 45, 40, 36, ..., 11, 9 rows, and later searches
start from the rows
+ // that would fit at the bytes per row of the last slice, here 10, which
fits exactly.
+ assert(sliceRowCounts(Seq(50), 1000)(_ => 100) === Seq(9, 10, 10, 10, 10,
1))
+ // Row 0 alone is larger than the limit, which does not keep the other
slices at one row.
+ assert(sliceRowCounts(Seq(1000), 1000)(i => if (i == 0) 5000 else 10) ===
+ Seq(1, 94) ++ Seq.fill(9)(100) ++ Seq(5))
+ // Slices grow again after the large rows.
+ assert(sliceRowCounts(Seq(1000), 1000)(i => if (i < 10) 500 else 10) ===
+ Seq.fill(6)(2) ++ Seq.fill(9)(100) ++ Seq(88))
+ // Every row is larger than the limit.
+ assert(sliceRowCounts(Seq(3, 0, 2), 1000)(_ => 2000) === Seq(1, 1, 1, 0,
1, 1))
+ }
+
+ test("SPARK-60395: bufferSize matches getBufferSize of vectors below 2 GiB")
{
+ val nestedSchema = new StructType()
+ .add("b", BooleanType)
+ .add("d", DecimalType(20, 2))
+ .add("bin", BinaryType)
+ .add("arr", ArrayType(new StructType().add("x", LongType).add("y",
StringType)))
+ .add("m", MapType(StringType, ArrayType(IntegerType)))
+ val writer = ArrowWriter.create(nestedSchema, "UTC")
+ try {
+ (0 until 20).foreach { i =>
+ val arr = new GenericArrayData((0 until i % 4).map { j =>
+ InternalRow(j.toLong, if (j % 2 == 0) null else
UTF8String.fromString(s"y_$j"))
+ })
+ val m = ArrayBasedMapData(
+ Array[Any](UTF8String.fromString(s"k_$i")),
+ Array[Any](new GenericArrayData((0 until i % 3).toArray[Any])))
+ writer.write(InternalRow(i % 2 == 0, Decimal(i * 100L, 20, 2),
+ Array.fill[Byte](i)(1), if (i % 5 == 0) null else arr, m))
+ }
+ writer.finish()
+ val root = writer.root
+ Seq((0, 20), (3, 7), (5, 1), (19, 1), (0, 13)).foreach { case (start,
length) =>
+ val slice = root.slice(start, length)
+ try {
+ (slice.getFieldVectors.asScala ++
root.getFieldVectors.asScala).foreach { v =>
+ assert(SliceBytesArrowOutputProcessorImpl.bufferSize(v) ===
v.getBufferSize,
+ s"${v.getField} in slice ($start, $length)")
+ }
+ } finally {
+ slice.close()
+ }
+ }
+ } finally {
+ writer.root.close()
+ }
+ }
+
+ test("SPARK-60395: bufferSize does not wrap around for a struct vector over
2 GiB") {
+ val allocator = ArrowUtils.rootAllocator.newChildAllocator("struct", 0,
Long.MaxValue)
+ class TestStructVector extends StructVector(
+ "st", allocator, FieldType.nullable(ArrowType.Struct.INSTANCE), null) {
+ def addChild(name: String, child: NullVector): Unit = putChild(name,
child)
+ }
+ val struct = new TestStructVector
+ try {
+ // Children that report 1 GiB each, like three int64 columns of 2^27
rows.
+ (0 until 3).foreach { i =>
+ struct.addChild(s"c$i", new NullVector(s"c$i") {
+ override def getBufferSize: Int = 1 << 30
+ })
+ }
+ struct.setValueCount(1)
+ // `StructVector.getBufferSize` wraps around to a negative size.
+ assert(struct.getBufferSize < 0)
+ assert(SliceBytesArrowOutputProcessorImpl.bufferSize(struct) === 3L * (1
<< 30) + 1)
Review Comment:
**Nit (P3):** This assertion holds whether or not the nested arm of
`bufferSize` recurses, and nothing else in the suite ties the Long measurement
to the processor's decision. The three children are `NullVector`s, so they
reach the `case v => v.getBufferSize` fallback (`PythonArrowOutput.scala:346`),
and `ownBytes + children.map(_.getBufferSize.toLong).sum` at L345 would give
the same `3L * (1 << 30) + 1`. Item 11 of the latest review covers that half.
The other half is `batchBytes`. If L307 went back to summing
`getBufferSize.toLong`, every test would still pass. The equal-limit,
measurement-count and parameterized tests stay far below 2 GiB, where the two
sums agree, and `sliceRowCounts` overrides `batchBytes`. Either change brings
back the bug this PR fixes: the description's struct of three int64 columns of
2^27 rows would again measure -1,006,632,960 bytes and go out unsliced.
A child that reports 2^28 values of 8 bytes through an overridden
`getValueCount`, like `fixed` in the flat test, needs no allocation (with
`addChild` taking a `FieldVector`). Its `getBufferSize` reads the `valueCount`
field that `struct.setValueCount(1)` sets, so it reports 9 bytes, and the
struct's expected size becomes `3L * (1 << 30) + (1L << 31) + (1 << 25) + 1`.
Measuring a root that holds this struct through a processor's `batchBytes`
would cover the wiring the same way.
This extends an [existing
discussion](https://github.com/apache/spark/pull/59786#discussion_r4383082539).
**Verification:**
- **Regression:** Replacing batchBytes' sum with the field vectors' int
getBufferSize, or the nested arm's recursion with each child's getBufferSize as
a Long, makes a test fail.
- **Behavior:** The measured size of a struct that holds a child over 2 GiB
equals the exact sum of its children's sizes plus its validity buffer, with
nothing over 2 GiB allocated.
--
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]
