viirya commented on code in PR #59786:
URL: https://github.com/apache/spark/pull/59786#discussion_r4386262306
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowOutput.scala:
##########
@@ -265,47 +265,85 @@ class SliceBytesArrowOutputProcessorImpl(
arrowMaxBytesPerOutputBatch: Int)
extends BaseSliceArrowOutputProcessor(reader, pythonMetrics) {
+ // The number of rows from which the search for the next slice of the
current batch starts.
+ private var sliceRowCount = 0
+
protected override def getRoot: VectorSchemaRoot = {
- // Output the whole root only if no row of this batch has been output yet.
- if (currentRowIdx == 0 && getBatchBytes(root) <
arrowMaxBytesPerOutputBatch) {
- currentRowIdx = rowCount
- root
- } else {
- // We iteratively slice the batch by 90% until we reach a batch with
required byte size.
- // If we cannot have it even with one single row, we will output a batch
with one row.
- var rootSlice = root
- val stepRatio = 0.9
- var sliceBatchCount = rowCount - currentRowIdx
-
- while (sliceBatchCount > 0) {
- rootSlice = root.slice(currentRowIdx, sliceBatchCount)
- val batchBytes = getBatchBytes(rootSlice)
- if (batchBytes < arrowMaxBytesPerOutputBatch) {
- currentRowIdx = currentRowIdx + rootSlice.getRowCount
-
- return rootSlice
- } else {
- rootSlice.close()
- sliceBatchCount = (sliceBatchCount * stepRatio).toInt
- }
+ if (currentRowIdx == 0) {
+ // Output the whole root if it fits, which happens only if no row of it
has been output yet.
+ if (batchBytes(root) <= arrowMaxBytesPerOutputBatch) {
+ currentRowIdx = rowCount
+ return root
}
- // If the batch size is still too large, we need to slice it down to a
single row.
- // Normally, this should not happen, but it is possible if the data is
very large.
- if (sliceBatchCount == 0) {
- sliceBatchCount = 1
+ sliceRowCount = (rowCount *
SliceBytesArrowOutputProcessorImpl.stepRatio).toInt
+ }
+ // We iteratively shrink the slice by 10% until it fits. Each search
starts from the number of
+ // rows that would fit at the bytes per row of the last slice that fit,
rather than from all
+ // the remaining rows, since a slice that does not start at row 0 usually
copies its offset
+ // buffers, and restarting from all the remaining rows would copy them
again and again.
+ var numRows = math.min(sliceRowCount, rowCount - currentRowIdx)
+ while (numRows > 0) {
+ val rootSlice = root.slice(currentRowIdx, numRows)
+ val bytes = batchBytes(rootSlice)
+ if (bytes <= arrowMaxBytesPerOutputBatch) {
+ // Never below `numRows`, so the next slice can grow again after a run
of large rows.
+ val estimate = numRows.toLong * arrowMaxBytesPerOutputBatch /
math.max(bytes, 1L)
+ sliceRowCount = math.min(estimate, Int.MaxValue).toInt
+ currentRowIdx += numRows
+ return rootSlice
}
- rootSlice = root.slice(currentRowIdx, sliceBatchCount)
- currentRowIdx = currentRowIdx + rootSlice.getRowCount
- rootSlice
+ rootSlice.close()
+ numRows = (numRows * SliceBytesArrowOutputProcessorImpl.stepRatio).toInt
}
+ // Even a single row is larger than the limit, so we output a batch with
that one row.
+ // Normally, this should not happen, but it is possible if the data is
very large. The next
+ // search starts from where this one did, since the next rows may be much
smaller.
+ val rootSlice = root.slice(currentRowIdx, 1)
+ currentRowIdx += 1
+ rootSlice
}
- private def getBatchBytes(root: VectorSchemaRoot): Long = {
- var batchBytes = 0L
- root.getFieldVectors.asScala.foreach { vector =>
- batchBytes += vector.getBufferSize
- }
- batchBytes
+ private[python] def batchBytes(root: VectorSchemaRoot): Long = {
+
root.getFieldVectors.asScala.map(SliceBytesArrowOutputProcessorImpl.bufferSize).sum
+ }
+}
+
+object SliceBytesArrowOutputProcessorImpl {
+ private val stepRatio = 0.9
+
+ /**
+ * Returns the size in bytes of the buffers of `vector`, like
`ValueVector.getBufferSize` but as
+ * a Long. `getBufferSize` computes the size in int, so it wraps around for
a vector over 2 GiB,
Review Comment:
Added "A large variable-width vector computes it in Long but caps it at
Int.MaxValue instead."
##########
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)
+ } finally {
+ struct.close()
+ allocator.close()
+ }
+ }
+
+ test("SPARK-60395: bufferSize does not wrap around for flat vectors over 2
GiB") {
+ val allocator = ArrowUtils.rootAllocator.newChildAllocator("flat", 0,
Long.MaxValue)
+ try {
+ // An int64 column of 2^28 rows, without allocating it.
+ val fixed = new FixedSizeBinaryVector("f", allocator, 8) {
+ override def getValueCount: Int = 1 << 28
+ }
+ try {
Review Comment:
Done.
##########
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)
Review Comment:
Done, `newSliceBytesProcessor` is gone.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonArrowOutputSuite.scala:
##########
@@ -76,7 +81,8 @@ class PythonArrowOutputSuite extends SparkFunSuite {
/**
* Drives `processor` the way `PythonArrowOutput.newReaderIterator.read()`
does, and returns
- * the rows and the buffer size in bytes of every output batch. Rows are
copied before the
+ * the rows and the buffer size in bytes, as
`SliceBytesArrowOutputProcessorImpl` measures it,
+ * of every output batch. Rows are copied before the
Review Comment:
Re-flowed.
--
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]
