viirya commented on code in PR #59786:
URL: https://github.com/apache/spark/pull/59786#discussion_r4381042342
##########
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
Review Comment:
Agreed that they are latent here. The scaladoc now says that the own buffers
of a nested vector are still measured in int, e.g. the offsets of a list vector
with more than 2^29 rows, and that view vectors fall back to `getBufferSize`.
##########
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 {
Review Comment:
I'd keep it next to its only caller in this PR, and move it to `ArrowUtils`
in the follow-up that fixes the other int sums you listed, so that the move
comes with its second user.
##########
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
Review Comment:
Right that each level walks its subtree again. A recursion that returns the
int and Long sums together still needs `v.getBufferSize` at every level to get
the nested vector's own buffers, so it would go from about 2d to d walks rather
than to one. Since the schemas here are shallow and each measurement follows a
`VectorSchemaRoot.slice`, which costs more, I left it as is.
--
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]
