spark-test commented on code in PR #59786:
URL: https://github.com/apache/spark/pull/59786#discussion_r4372986249
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowOutput.scala:
##########
@@ -265,47 +265,74 @@ class SliceBytesArrowOutputProcessorImpl(
arrowMaxBytesPerOutputBatch: Int)
extends BaseSliceArrowOutputProcessor(reader, pythonMetrics) {
+ // The number of rows of the last slice output from the current batch, from
which the search
+ // for the next slice 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 of the last slice that fit, since every slice that does not start
at row 0 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)
+ if (batchBytes(rootSlice) <= arrowMaxBytesPerOutputBatch) {
+ sliceRowCount = numRows
+ 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.
+ sliceRowCount = 1
+ 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
+ protected 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,
+ * e.g. a struct vector, which adds up the sizes of its children, such as
the one top-level
+ * struct vector that holds the output of mapInArrow, mapInPandas,
applyInPandas and Python
+ * UDTFs. Other vectors, e.g. `NullVector` and the large variable-width
ones, which cap their
+ * size at Int.MaxValue, fall back to `getBufferSize`.
+ */
+ def bufferSize(vector: ValueVector): Long = vector match {
+ case v: BaseFixedWidthVector if v.getTypeWidth > 0 && v.getValueCount > 0
=>
+ BitVectorHelper.getValidityBufferSizeFromCount(v.getValueCount) +
+ v.getValueCount.toLong * v.getTypeWidth
+ case v: BaseVariableWidthVector if v.getValueCount > 0 =>
+ val offsetWidth = BaseVariableWidthVector.OFFSET_WIDTH
+ BitVectorHelper.getValidityBufferSizeFromCount(v.getValueCount) +
+ (v.getValueCount + 1L) * offsetWidth +
+ v.getOffsetBuffer.getInt(v.getValueCount.toLong * offsetWidth)
+ case v: FieldVector if !v.getChildrenFromFields.isEmpty =>
+ val children = v.getChildrenFromFields.asScala.toSeq
+ // The size of the buffers of the nested vector itself, e.g. the
validity and offsets. Its
+ // `getBufferSize` adds them to the sizes of its children in int, so
subtracting the sizes
+ // of the children in int cancels out any wrap-around of the sum.
+ val ownBytes = v.getBufferSize - children.map(_.getBufferSize).sum
+ ownBytes + children.map(bufferSize).sum
+ case v => v.getBufferSize
Review Comment:
**Nit (P3):** `BaseLargeVariableWidthVector` is not a
`BaseVariableWidthVector`, so LargeVarChar and LargeVarBinary columns fall
through to this arm. In Arrow 19 their `getBufferSizeFor` returns
`capAtMaxInt(validity + offsets + data)`. With the new `<=` checks, that capped
value satisfies a limit of `2147483647`, the largest value `checkValue`
accepts. So with `useLargeVarTypes=true`, a top-level large string or binary
column over 2 GiB is output as one batch over the limit; a scalar pandas UDF
result read by ArrowEvalPythonExec is one example. The merge base's `<`
rejected the capped size and kept shrinking until the measurement was exact.
Struct-wrapped outputs (mapInArrow, applyInPandas) are not affected, because
the struct's own validity bytes push the sum past the limit.
A `BaseLargeVariableWidthVector` arm that mirrors the variable-width one
would close this. It would compute validity + `(valueCount + 1L) * 8` +
`getOffsetBuffer.getLong(valueCount * 8L)` in Long. The scaladoc sentence
saying the large variable-width vectors fall back to `getBufferSize` would need
updating too.
**Verification:**
- **Behavior:** bufferSize matches getBufferSize for large string and binary
vectors, on whole roots and on offset slices below 2 GiB.
- **Regression:** A large variable-width vector whose 64-bit end offset
exceeds Int.MaxValue is measured at its exact Long size rather than
Int.MaxValue, so it does not fit a limit of 2147483647.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowOutput.scala:
##########
@@ -265,47 +265,74 @@ class SliceBytesArrowOutputProcessorImpl(
arrowMaxBytesPerOutputBatch: Int)
extends BaseSliceArrowOutputProcessor(reader, pythonMetrics) {
+ // The number of rows of the last slice output from the current batch, from
which the search
+ // for the next slice 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 of the last slice that fit, since every slice that does not start
at row 0 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)
+ if (batchBytes(rootSlice) <= arrowMaxBytesPerOutputBatch) {
+ sliceRowCount = numRows
+ 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.
+ sliceRowCount = 1
Review Comment:
**Non-blocking (P2):** This records the row count of a slice that did *not*
fit as the start of the next search. `sliceRowCount` is only ever reassigned to
a smaller count that fits, or reseeded when a new Python batch starts. So after
one row larger than the limit, every later search in the same batch probes
exactly one row, and the rest of the batch comes out as one-row batches.
Example: with `maxBytesPerOutputBatch=1m`, take a mapInArrow batch of 10,000
rows where row 0 is 2 MiB and the other rows are 50 bytes. It is now output as
10,000 batches; the merge base restarted from all remaining rows and output 2.
Rows stay correct, but each extra batch pays for a `VectorSchemaRoot.slice`,
new `ArrowColumnVector`s and a `ColumnarBatch`, plus per-batch overhead
downstream. It also contradicts the comment above ("Each search starts from the
number of rows of the last slice that fit").
Dropping `sliceRowCount = 1` makes the next search start from the last slice
that fit, or from 90% of the batch if none has fit yet. The cost is one more
full search per oversized row, which is what the merge base paid for every
output batch. No current test has an oversized row followed by rows that fit:
the `(Seq(2, 0, 5, 1), 3)` case makes every row oversized.
**Verification:**
- **Regression:** A Python output batch of 1,000 rows whose first row alone
exceeds the limit, and whose other 999 rows fit the limit together, is output
as the oversized row alone followed by at most two batches holding the other
999 rows, not 999 one-row batches.
- **Behavior:** Byte slicing of uniform rows, all-oversized rows, empty
batches and a batch exactly at the limit still outputs every row once, in
order, with every multi-row slice within the limit.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonArrowOutputSuite.scala:
##########
@@ -145,7 +152,117 @@ 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("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("slice by bytes: each search starts from the rows of the last slice
that fit") {
+ 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 protected 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")
+ }
+
+ test("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("bufferSize does not wrap around for a struct vector over 2 GiB") {
Review Comment:
**Nit (P3):** This test reaches only the nested arm and the fallback,
because `NullVector` is not a `BaseFixedWidthVector`. The parity test above
runs below 2 GiB, where `bufferSize` equals `getBufferSize` even if the fixed-
or variable-width arm is deleted or loses `.toLong`. If either arm regressed
(e.g. an int64 child of 2^28 rows wrapping to a negative size again), the suite
would still pass. That arm is the flat-column half of the wrap-around fix the
PR describes. It can be tested the same way, without allocating: in Arrow 19,
`BaseFixedWidthVector` has a public constructor and non-final
`getValueCount`/`getTypeWidth`, so a test vector can report a size over
`Int.MaxValue`. A variable-width vector can report a large end offset through
its offset buffer.
**Verification:**
- **Regression:** Dropping the Long widening from either the fixed-width or
the variable-width arm of bufferSize makes a test fail.
- **Behavior:** bufferSize returns the exact Long size for a fixed-width
vector and a variable-width vector whose sizes exceed Int.MaxValue.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonArrowOutputSuite.scala:
##########
@@ -145,7 +152,117 @@ 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("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("slice by bytes: each search starts from the rows of the last slice
that fit") {
+ 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 protected 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")
Review Comment:
**Nit (P3):** These assertions bound the maximum slice size, the minimum
number of slices, and a measurement count that grows with the number of slices.
Nothing fixes how many rows a chosen slice holds. For example, deleting
`sliceRowCount = (rowCount * ...stepRatio).toInt` makes every slice one row.
This test then sees 2000 one-row slices and 2000 measurements, under its 4040
bound, and the other byte-slicing tests still pass. Reverting the in-search
check to `<` is not caught either: the equal-limit test sets the limit to a
whole batch's size, so it reaches only the whole-root check. Asserting the
expected per-slice row counts for a deterministic batch would give the new
search a failure signal for slice sizing.
**Verification:**
- **Regression:** Removing the initial search start for a batch that does
not fit whole, so that every slice has one row, makes a byte-slicing test fail.
- **Regression:** Changing the in-search fit check from at-most to
strictly-less-than makes a byte-slicing test fail.
--
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]
