viirya commented on code in PR #59786:
URL: https://github.com/apache/spark/pull/59786#discussion_r4404175333
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowOutput.scala:
##########
@@ -265,47 +266,138 @@ 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
+ // Whether the last output batch is one row larger than the limit.
+ private var lastRowTooLarge = false
+
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) {
+ lastRowTooLarge = false
+ // Output the whole root if it fits, which happens only if no row of it
has been output yet,
+ // or if Arrow cannot slice it.
+ val rootBytes = batchBytes(root)
+ if (rootBytes <= arrowMaxBytesPerOutputBatch || !canSlice(root)) {
+ currentRowIdx = rowCount
+ return root
+ }
+ // The first search starts from the rows that would fit at the bytes per
row of the root.
+ val estimate = math.max(rowCount.toLong * arrowMaxBytesPerOutputBatch /
rootBytes, 1L)
+ sliceRowCount = math.min(
+ estimate, (rowCount *
SliceBytesArrowOutputProcessorImpl.stepRatio).toLong).toInt
+ }
+ if (lastRowTooLarge) {
+ // Probe the row after a row larger than the limit alone, and output it
if it is larger than
+ // the limit too, so that a run of such rows does not search from
`sliceRowCount` every row.
+ val rootSlice = root.slice(currentRowIdx, 1)
+ if (batchBytes(rootSlice) > arrowMaxBytesPerOutputBatch) {
+ currentRowIdx += 1
+ return rootSlice
+ }
+ rootSlice.close()
+ lastRowTooLarge = false
+ }
+ // We iteratively shrink the slice by 10% until it fits. Each later 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 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
}
- // 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
+ if (numRows == 1) {
+ // 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.
+ lastRowTooLarge = true
+ currentRowIdx += 1
+ return rootSlice
}
- rootSlice = root.slice(currentRowIdx, sliceBatchCount)
- currentRowIdx = currentRowIdx + rootSlice.getRowCount
- rootSlice
+ rootSlice.close()
+ numRows = (numRows * SliceBytesArrowOutputProcessorImpl.stepRatio).toInt
}
+ // A batch of one row that is larger than the limit, where the search does
not run.
+ val rootSlice = root.slice(currentRowIdx, 1)
+ currentRowIdx += 1
+ rootSlice
+ }
+
+ private[python] def batchBytes(root: VectorSchemaRoot): Long =
+ SliceBytesArrowOutputProcessorImpl.batchBytes(root)
+
+ private[python] def canSlice(root: VectorSchemaRoot): Boolean =
+
root.getFieldVectors.asScala.forall(SliceBytesArrowOutputProcessorImpl.canSlice)
+}
+
+object SliceBytesArrowOutputProcessorImpl {
+ private val stepRatio = 0.9
+
+ /** Returns the size in bytes of the buffers of `root`, as the sum of
[[bufferSize]]. */
+ private[python] def batchBytes(root: VectorSchemaRoot): Long =
+ root.getFieldVectors.asScala.map(bufferSize).sum
+
+ /**
+ * 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 and
applyInPandas. A large
+ * variable-width vector computes it in Long but caps it at Int.MaxValue
instead.
+ *
+ * Fixed-width and (large) variable-width vectors are measured in Long. A
nested vector adds up
+ * its children in Long, but its own buffers come from its int
`getBufferSize`, so they are
+ * still wrong when they alone exceed 2 GiB, e.g. the offsets of a list
vector with more than
+ * 2^29 rows. Other vectors, e.g. the view vectors, fall back to
`getBufferSize`.
+ */
+ private[python] 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: BaseLargeVariableWidthVector if v.getValueCount > 0 =>
+ val offsetWidth = BaseLargeVariableWidthVector.OFFSET_WIDTH
+ BitVectorHelper.getValidityBufferSizeFromCount(v.getValueCount) +
+ (v.getValueCount + 1L) * offsetWidth +
+ v.getOffsetBuffer.getLong(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
}
- private def getBatchBytes(root: VectorSchemaRoot): Long = {
- var batchBytes = 0L
- root.getFieldVectors.asScala.foreach { vector =>
- batchBytes += vector.getBufferSize
+ /**
+ * Returns whether `VectorSchemaRoot.slice` can slice `vector` and its
children. Arrow computes
+ * some byte positions of a slice in int, so it throws, and leaks the
vectors sliced so far, for
+ * a fixed-width vector over 2 GiB, a view vector with more than 2^27 rows,
or a list vector with
+ * more than 2^29 rows. A batch that holds one is output whole instead.
+ */
+ private[python] def canSlice(vector: ValueVector): Boolean = {
+ val fits = vector match {
+ case v: BaseFixedWidthVector => v.getValueCount.toLong * v.getTypeWidth
<= Int.MaxValue
+ case v: BaseVariableWidthViewVector =>
+ v.getValueCount.toLong * BaseVariableWidthViewVector.ELEMENT_SIZE <=
Int.MaxValue
+ case v: ListVector =>
+ (v.getValueCount + 1L) * BaseRepeatedValueVector.OFFSET_WIDTH <=
Int.MaxValue
+ case _ => true
}
- batchBytes
+ fits && (vector match {
+ case v: FieldVector => v.getChildrenFromFields.asScala.forall(canSlice)
Review Comment:
Added the assertion to the struct test. Replacing the check of the children
with `true` now fails it with "canSlice(struct) was true", where every test
passed before.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonArrowOutputSuite.scala:
##########
@@ -145,7 +152,240 @@ 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 })
}
}
+
+ 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, new SliceBytesArrowOutputProcessorImpl(_,
_, maxBytes))
+ checkAllRowsOnceInOrder(batchSizes, output)
+ assert(output.map(_._1.size) === batchSizes)
+ // One byte less, the larger batch is sliced.
+ val sliced = readAll(batchSizes, new SliceBytesArrowOutputProcessorImpl(_,
_, 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")
Review Comment:
Took the suggestion. It measures 58 times for 52 slices, as your model says.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowOutput.scala:
##########
@@ -265,47 +266,138 @@ 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
+ // Whether the last output batch is one row larger than the limit.
+ private var lastRowTooLarge = false
+
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) {
+ lastRowTooLarge = false
+ // Output the whole root if it fits, which happens only if no row of it
has been output yet,
+ // or if Arrow cannot slice it.
+ val rootBytes = batchBytes(root)
+ if (rootBytes <= arrowMaxBytesPerOutputBatch || !canSlice(root)) {
+ currentRowIdx = rowCount
+ return root
+ }
+ // The first search starts from the rows that would fit at the bytes per
row of the root.
+ val estimate = math.max(rowCount.toLong * arrowMaxBytesPerOutputBatch /
rootBytes, 1L)
+ sliceRowCount = math.min(
+ estimate, (rowCount *
SliceBytesArrowOutputProcessorImpl.stepRatio).toLong).toInt
+ }
+ if (lastRowTooLarge) {
+ // Probe the row after a row larger than the limit alone, and output it
if it is larger than
+ // the limit too, so that a run of such rows does not search from
`sliceRowCount` every row.
+ val rootSlice = root.slice(currentRowIdx, 1)
+ if (batchBytes(rootSlice) > arrowMaxBytesPerOutputBatch) {
+ currentRowIdx += 1
+ return rootSlice
+ }
+ rootSlice.close()
+ lastRowTooLarge = false
+ }
+ // We iteratively shrink the slice by 10% until it fits. Each later 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 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
}
- // 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
+ if (numRows == 1) {
+ // 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.
+ lastRowTooLarge = true
+ currentRowIdx += 1
+ return rootSlice
}
- rootSlice = root.slice(currentRowIdx, sliceBatchCount)
- currentRowIdx = currentRowIdx + rootSlice.getRowCount
- rootSlice
+ rootSlice.close()
+ numRows = (numRows * SliceBytesArrowOutputProcessorImpl.stepRatio).toInt
}
+ // A batch of one row that is larger than the limit, where the search does
not run.
+ val rootSlice = root.slice(currentRowIdx, 1)
+ currentRowIdx += 1
+ rootSlice
+ }
+
+ private[python] def batchBytes(root: VectorSchemaRoot): Long =
+ SliceBytesArrowOutputProcessorImpl.batchBytes(root)
+
+ private[python] def canSlice(root: VectorSchemaRoot): Boolean =
+
root.getFieldVectors.asScala.forall(SliceBytesArrowOutputProcessorImpl.canSlice)
+}
+
+object SliceBytesArrowOutputProcessorImpl {
+ private val stepRatio = 0.9
+
+ /** Returns the size in bytes of the buffers of `root`, as the sum of
[[bufferSize]]. */
+ private[python] def batchBytes(root: VectorSchemaRoot): Long =
+ root.getFieldVectors.asScala.map(bufferSize).sum
+
+ /**
+ * 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 and
applyInPandas. A large
+ * variable-width vector computes it in Long but caps it at Int.MaxValue
instead.
+ *
+ * Fixed-width and (large) variable-width vectors are measured in Long. A
nested vector adds up
+ * its children in Long, but its own buffers come from its int
`getBufferSize`, so they are
+ * still wrong when they alone exceed 2 GiB, e.g. the offsets of a list
vector with more than
+ * 2^29 rows. Other vectors, e.g. the view vectors, fall back to
`getBufferSize`.
+ */
+ private[python] 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: BaseLargeVariableWidthVector if v.getValueCount > 0 =>
+ val offsetWidth = BaseLargeVariableWidthVector.OFFSET_WIDTH
+ BitVectorHelper.getValidityBufferSizeFromCount(v.getValueCount) +
+ (v.getValueCount + 1L) * offsetWidth +
+ v.getOffsetBuffer.getLong(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
}
- private def getBatchBytes(root: VectorSchemaRoot): Long = {
- var batchBytes = 0L
- root.getFieldVectors.asScala.foreach { vector =>
- batchBytes += vector.getBufferSize
+ /**
+ * Returns whether `VectorSchemaRoot.slice` can slice `vector` and its
children. Arrow computes
+ * some byte positions of a slice in int, so it throws, and leaks the
vectors sliced so far, for
+ * a fixed-width vector over 2 GiB, a view vector with more than 2^27 rows,
or a list vector with
+ * more than 2^29 rows. A batch that holds one is output whole instead.
Review Comment:
Took the suggestion, and updated items 1 and 7 of the description to match.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowOutput.scala:
##########
@@ -265,47 +266,138 @@ 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
+ // Whether the last output batch is one row larger than the limit.
+ private var lastRowTooLarge = false
+
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) {
+ lastRowTooLarge = false
+ // Output the whole root if it fits, which happens only if no row of it
has been output yet,
+ // or if Arrow cannot slice it.
+ val rootBytes = batchBytes(root)
+ if (rootBytes <= arrowMaxBytesPerOutputBatch || !canSlice(root)) {
+ currentRowIdx = rowCount
+ return root
+ }
+ // The first search starts from the rows that would fit at the bytes per
row of the root.
+ val estimate = math.max(rowCount.toLong * arrowMaxBytesPerOutputBatch /
rootBytes, 1L)
+ sliceRowCount = math.min(
+ estimate, (rowCount *
SliceBytesArrowOutputProcessorImpl.stepRatio).toLong).toInt
+ }
+ if (lastRowTooLarge) {
+ // Probe the row after a row larger than the limit alone, and output it
if it is larger than
+ // the limit too, so that a run of such rows does not search from
`sliceRowCount` every row.
+ val rootSlice = root.slice(currentRowIdx, 1)
+ if (batchBytes(rootSlice) > arrowMaxBytesPerOutputBatch) {
+ currentRowIdx += 1
+ return rootSlice
+ }
+ rootSlice.close()
+ lastRowTooLarge = false
+ }
+ // We iteratively shrink the slice by 10% until it fits. Each later 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 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
}
- // 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
+ if (numRows == 1) {
+ // 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.
+ lastRowTooLarge = true
+ currentRowIdx += 1
+ return rootSlice
}
- rootSlice = root.slice(currentRowIdx, sliceBatchCount)
- currentRowIdx = currentRowIdx + rootSlice.getRowCount
- rootSlice
+ rootSlice.close()
+ numRows = (numRows * SliceBytesArrowOutputProcessorImpl.stepRatio).toInt
}
+ // A batch of one row that is larger than the limit, where the search does
not run.
+ val rootSlice = root.slice(currentRowIdx, 1)
+ currentRowIdx += 1
+ rootSlice
+ }
+
+ private[python] def batchBytes(root: VectorSchemaRoot): Long =
+ SliceBytesArrowOutputProcessorImpl.batchBytes(root)
+
+ private[python] def canSlice(root: VectorSchemaRoot): Boolean =
+
root.getFieldVectors.asScala.forall(SliceBytesArrowOutputProcessorImpl.canSlice)
+}
+
+object SliceBytesArrowOutputProcessorImpl {
+ private val stepRatio = 0.9
+
+ /** Returns the size in bytes of the buffers of `root`, as the sum of
[[bufferSize]]. */
+ private[python] def batchBytes(root: VectorSchemaRoot): Long =
+ root.getFieldVectors.asScala.map(bufferSize).sum
+
+ /**
+ * 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 and
applyInPandas. A large
+ * variable-width vector computes it in Long but caps it at Int.MaxValue
instead.
+ *
+ * Fixed-width and (large) variable-width vectors are measured in Long. A
nested vector adds up
+ * its children in Long, but its own buffers come from its int
`getBufferSize`, so they are
+ * still wrong when they alone exceed 2 GiB, e.g. the offsets of a list
vector with more than
+ * 2^29 rows. Other vectors, e.g. the view vectors, fall back to
`getBufferSize`.
+ */
+ private[python] 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: BaseLargeVariableWidthVector if v.getValueCount > 0 =>
+ val offsetWidth = BaseLargeVariableWidthVector.OFFSET_WIDTH
+ BitVectorHelper.getValidityBufferSizeFromCount(v.getValueCount) +
+ (v.getValueCount + 1L) * offsetWidth +
+ v.getOffsetBuffer.getLong(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
}
- private def getBatchBytes(root: VectorSchemaRoot): Long = {
- var batchBytes = 0L
- root.getFieldVectors.asScala.foreach { vector =>
- batchBytes += vector.getBufferSize
+ /**
+ * Returns whether `VectorSchemaRoot.slice` can slice `vector` and its
children. Arrow computes
+ * some byte positions of a slice in int, so it throws, and leaks the
vectors sliced so far, for
+ * a fixed-width vector over 2 GiB, a view vector with more than 2^27 rows,
or a list vector with
+ * more than 2^29 rows. A batch that holds one is output whole instead.
+ */
+ private[python] def canSlice(vector: ValueVector): Boolean = {
Review Comment:
Done.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowOutput.scala:
##########
@@ -265,47 +266,138 @@ 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
+ // Whether the last output batch is one row larger than the limit.
+ private var lastRowTooLarge = false
+
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) {
+ lastRowTooLarge = false
+ // Output the whole root if it fits, which happens only if no row of it
has been output yet,
+ // or if Arrow cannot slice it.
+ val rootBytes = batchBytes(root)
+ if (rootBytes <= arrowMaxBytesPerOutputBatch || !canSlice(root)) {
+ currentRowIdx = rowCount
+ return root
+ }
+ // The first search starts from the rows that would fit at the bytes per
row of the root.
+ val estimate = math.max(rowCount.toLong * arrowMaxBytesPerOutputBatch /
rootBytes, 1L)
+ sliceRowCount = math.min(
+ estimate, (rowCount *
SliceBytesArrowOutputProcessorImpl.stepRatio).toLong).toInt
+ }
+ if (lastRowTooLarge) {
+ // Probe the row after a row larger than the limit alone, and output it
if it is larger than
+ // the limit too, so that a run of such rows does not search from
`sliceRowCount` every row.
+ val rootSlice = root.slice(currentRowIdx, 1)
+ if (batchBytes(rootSlice) > arrowMaxBytesPerOutputBatch) {
+ currentRowIdx += 1
+ return rootSlice
+ }
+ rootSlice.close()
+ lastRowTooLarge = false
+ }
+ // We iteratively shrink the slice by 10% until it fits. Each later 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 usually copies its offset
buffers, and restarting
+ // from all the remaining rows would copy them again and again.
Review Comment:
Took the suggestion.
--
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]
