viirya commented on code in PR #59786:
URL: https://github.com/apache/spark/pull/59786#discussion_r4381040030
##########
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:
Fixed in 643cfbac1fb: the one-row fallback no longer touches
`sliceRowCount`, so the next search starts from where the failed one did. The
new "number of rows of each slice" test covers a 5000-byte row 0 followed by
10-byte rows: it comes out as 1, 94, 100 x 9, 5.
##########
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:
Added a `BaseLargeVariableWidthVector` arm that reads the end offset with
`getLong` and computes the size in Long, and updated the scaladoc. Tested with
a LargeVarChar vector whose end offset is 3 GiB: `getBufferSize` returns
`Int.MaxValue`, `bufferSize` the exact size.
##########
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:
Added "bufferSize does not wrap around for flat vectors over 2 GiB": a
`FixedSizeBinaryVector` reporting 2^28 values of 8 bytes, a VarChar vector
whose end offset is `Int.MaxValue - 10` (its `getBufferSize` is negative), and
the LargeVarChar case above. Nothing over 2 GiB is allocated.
##########
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:
Added "number of rows of each slice", which models a size per row through
the `id` column and asserts the exact row counts. The uniform case has a slice
that fits the limit exactly, so going back to `<` in the search, dropping the
initial search start, or pinning the search to one row each makes it fail
(checked by reverting the search to the previous commit: `List(9, 9, 9, 9, 9,
5)`).
##########
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)
Review Comment:
Thanks, took the estimate as suggested: after a slice fits, `sliceRowCount =
numRows * limit / max(bytes, 1)`, capped at Int.MaxValue and then by the
remaining rows. The test's third case (10 rows of 500 bytes, then 10-byte rows,
limit 1000) now outputs 6 x 2, 100 x 9, 88, i.e. 16 batches.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:
##########
@@ -10084,7 +10084,10 @@ class SQLConf extends Serializable with Logging with
SqlApiConf {
def arrowMaxRecordsPerOutputBatch: Int =
getConf(ARROW_EXECUTION_MAX_RECORDS_PER_OUTPUT_BATCH)
- def arrowMaxBytesPerOutputBatch: Long =
getConf(ARROW_EXECUTION_MAX_BYTES_PER_OUTPUT_BATCH)
+ // Zero or negative means no limit; clamp it so that a value below
Int.MinValue does not wrap
+ // around to a positive limit when narrowed.
+ def arrowMaxBytesPerOutputBatch: Int =
Review Comment:
Good point. Kept `Long` with the clamp, and `PythonArrowOutput` narrows it
with `.toInt` again.
--
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]
