This is an automated email from the ASF dual-hosted git repository.
dongjoon-hyun pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new e3f4817f1ad4 [SPARK-60395][SQL] Fix the batch size computation and
limit check of maxBytesPerOutputBatch in sliced Arrow output
e3f4817f1ad4 is described below
commit e3f4817f1ad4f879793cf1093e0c6a3e4cd3a644
Author: Liang-Chi Hsieh <[email protected]>
AuthorDate: Thu Oct 29 06:56:44 2026 -0700
[SPARK-60395][SQL] Fix the batch size computation and limit check of
maxBytesPerOutputBatch in sliced Arrow output
### What changes were proposed in this pull request?
This PR fixes how `SliceBytesArrowOutputProcessorImpl`, used when the
internal config `spark.sql.execution.arrow.maxBytesPerOutputBatch` is set,
measures an output batch and compares it with the limit. It addresses the six
items in SPARK-60395, raised in the review of #59616:
1. Add `SliceBytesArrowOutputProcessorImpl.bufferSize(vector): Long` and
use it to measure batches. `ValueVector.getBufferSize` computes the size in
int, so it wraps around over 2 GiB: a struct vector sums its children and
returns `(int) buffer`, and fixed- and variable-width vectors compute
`valueCount * typeWidth` and their offsets in int. Large variable-width vectors
cap it at `Int.MaxValue`, which a limit of `Int.MaxValue` would accept. The new
function computes fixed-width and (l [...]
2. Both checks now use `<=`. This matches the config doc ("equal or less
than this value") and the records processor.
3. The shrink search no longer restarts from all the remaining rows for
every output batch. When the whole batch does not fit, the first search starts
from the number of rows that would fit at the root's bytes per row, at most
`0.9 * rowCount`, which skips the `slice(0, rowCount)` probe that always
measured the same as the root. After a slice fits, the next search starts from
the number of rows that would fit at that slice's bytes per row, capped by the
remaining rows. This is never b [...]
4. A row larger than the limit is output as the one-row slice that the
search measured last, rather than a slice rebuilt after closing it, without the
always-true `if (sliceBatchCount == 0)` and the unused initial `rootSlice`.
Only a one-row batch that does not fit still builds `root.slice(currentRowIdx,
1)`.
5. `SQLConf.arrowMaxBytesPerOutputBatch` now returns `math.max(value, 0L)`,
still as a Long, before `PythonArrowOutput` narrows it with `.toInt`. Before, a
value below `Int.MinValue` could wrap around to a positive limit.
6. Drop the redundant `rowCount > 0 &&` in
`BaseSliceArrowOutputProcessor.loadBatch`.
7. Output a batch whole when Arrow cannot slice it. Arrow 19 computes some
byte positions of a slice in int:
`BaseFixedWidthVector.splitAndTransferValueBuffer` computes `startIndex *
typeWidth` and `length * typeWidth`, `ListVector` (and so `MapVector`) the
positions of its offsets, and the view vectors `startIndex * 16`. So
`root.slice` throws for a fixed-width vector of 2 GiB or more, a list vector
with 2^29 - 1 rows or more or a view vector with 2^27 rows or more, at any
level of t [...]
### Why are the changes needed?
- The outputs of mapInArrow, mapInPandas and applyInPandas are one
top-level struct vector. A batch over 2 GiB could measure negative and be
output without slicing. For example, three int64 columns of 2^27 rows measure
as -1,006,632,960 bytes.
- A batch whose size equals the limit was sliced, although the doc says it
is allowed.
- A probe usually copies its offset buffers: string and binary columns when
it does not start at row 0, and array and map columns always. Restarting the
search from all the remaining rows made splitting one batch of R rows into
chunks of c rows copy O(R^2/c) offsets.
- A negative value, which the doc describes as no limit, could wrap around
to a positive limit, e.g. `-3g` became 1 GiB and `-4294967295` became 1 byte.
If `maxRecordsPerOutputBatch` was also set, the check that both limits are not
set at once then threw.
### Does this PR introduce _any_ user-facing change?
Only for the internal config
`spark.sql.execution.arrow.maxBytesPerOutputBatch`:
- Output batches over 2 GiB are now sliced, unless Arrow cannot slice them
(item 7), in which case they are output whole as before.
- A batch whose size equals the limit is no longer sliced.
- Values below `Int.MinValue` now mean no limit, as the doc says.
### How was this patch tested?
New tests in `PythonArrowOutputSuite`:
- A batch whose size equals the limit is not sliced, and one byte less
slices it.
- The number of size measurements stays linear in the number of slices.
Splitting 2000 rows into about 50 slices took 1513 measurements before this
change.
- The exact row counts of the slices, with a modeled size per row: uniform
rows, including a slice that fits the limit exactly; one row larger than the
limit followed by small rows; a few large rows followed by small rows; and rows
that are all larger than the limit.
- A run of rows larger than the limit takes about one measurement per row,
after earlier slices have fit.
- A batch that Arrow cannot slice is output whole by the bytes slicing, and
`canSlice` rejects a struct vector that holds an int64 vector of 2^28 values.
SPARK-60627 already tests `canSlice` at the boundaries of flat vectors.
- `bufferSize` matches `getBufferSize` for nested and leaf vectors
(boolean, decimal, binary, array of struct, map of array) on roots and slices
below 2 GiB.
- `bufferSize` and `batchBytes` do not wrap around for a struct vector
whose children report 3 GiB in total and include an int64 vector of 2^28
values, which the struct measures only if it measures its children in Long;
`bufferSize` does not wrap around for a fixed-width vector of 2^28 int64
values, a variable-width vector whose data is close to 2 GiB, and a large
variable-width vector of 3 GiB. These vectors report a large size or end
offset, so nothing over 2 GiB is allocated.
- `maxBytesPerOutputBatch` values of `-3g`, `-4294967295`, `-1` and `0`
mean no limit.
The tests for the equal-limit batch, the measurement counts, the slice row
counts, the batch output whole, the struct over 2 GiB and the negative values
fail without the fix. Also ran `pyspark.sql.tests.arrow.test_arrow_map`,
`pyspark.sql.tests.pandas.test_pandas_grouped_map_with_state`,
`pyspark.sql.tests.arrow.test_arrow_python_udf_cached` and
`pyspark.sql.tests.test_udtf`.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Claude Opus 5.5)
This pull request and its description were written by Isaac.
Closes #59786 from viirya/SPARK-60395.
Authored-by: Liang-Chi Hsieh <[email protected]>
Signed-off-by: Dongjoon Hyun <[email protected]>
---
.../org/apache/spark/sql/internal/SQLConf.scala | 11 +-
.../sql/execution/python/PythonArrowOutput.scala | 142 +++++++++----
.../execution/python/PythonArrowOutputSuite.scala | 228 ++++++++++++++++++++-
3 files changed, 336 insertions(+), 45 deletions(-)
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
index 550536cf2a95..4a16f2618f5a 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
@@ -5582,8 +5582,10 @@ object SQLConf {
"the columnar-based operator implemented by third-party libraries.
Spark will try to " +
"create batches with the size equal or less than this value. Normally
it should not " +
"happen, but if in extreme case that even one record is still larger
than this value, " +
- "Spark will create a batch with one record. applyInPandasWithState,
Arrow Python " +
- "UDTFs and Arrow Python UDFs that take columnar input directly ignore
this config.")
+ "Spark will create a batch with one record. A batch that Arrow cannot
slice, e.g. one " +
+ "with a fixed-width column of 2 GiB or more, is sent as a whole. " +
+ "applyInPandasWithState, Arrow Python UDTFs and Arrow Python UDFs that
" +
+ "take columnar input directly ignore this config.")
.version("4.1.0")
.internal()
.bytesConf(ByteUnit.BYTE)
@@ -10114,7 +10116,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 to an int.
+ def arrowMaxBytesPerOutputBatch: Long =
+ math.max(getConf(ARROW_EXECUTION_MAX_BYTES_PER_OUTPUT_BATCH), 0L)
def arrowMaxBytesPerBatch: Long =
getConf(ARROW_EXECUTION_MAX_BYTES_PER_BATCH)
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowOutput.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowOutput.scala
index 07a0939736e9..aa6c0f2a9bf9 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowOutput.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowOutput.scala
@@ -21,7 +21,7 @@ import java.util.concurrent.atomic.AtomicBoolean
import scala.jdk.CollectionConverters._
-import org.apache.arrow.vector.{BaseFixedWidthVector,
BaseVariableWidthViewVector, FieldVector, VectorSchemaRoot}
+import org.apache.arrow.vector.{BaseFixedWidthVector,
BaseLargeVariableWidthVector, BaseVariableWidthVector,
BaseVariableWidthViewVector, BitVectorHelper, FieldVector, ValueVector,
VectorSchemaRoot}
import org.apache.arrow.vector.complex.{BaseRepeatedValueVector, ListVector}
import org.apache.arrow.vector.ipc.ArrowStreamReader
@@ -224,7 +224,7 @@ abstract class BaseSliceArrowOutputProcessor(
protected var prevRoot: VectorSchemaRoot = root
override def loadBatch(): Boolean = {
- if (rowCount > 0 && currentRowIdx < rowCount) {
+ if (currentRowIdx < rowCount) {
true
} else {
val loaded = super.loadBatch()
@@ -303,47 +303,115 @@ 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
}
- // 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
+ // 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 = root.slice(currentRowIdx, sliceBatchCount)
- currentRowIdx = currentRowIdx + rootSlice.getRowCount
- 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, or of the
+ // root while none has, 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 (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.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 def getBatchBytes(root: VectorSchemaRoot): Long = {
- var batchBytes = 0L
- root.getFieldVectors.asScala.foreach { vector =>
- batchBytes += vector.getBufferSize
- }
- batchBytes
+ private[python] def batchBytes(root: VectorSchemaRoot): Long =
+ SliceBytesArrowOutputProcessorImpl.batchBytes(root)
+}
+
+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 reach 2 GiB, e.g. the offsets of a list
vector with 2^29 - 1 rows
+ * or more. 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
}
}
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonArrowOutputSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonArrowOutputSuite.scala
index 735b5b8f50ed..bceec11ca8b6 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonArrowOutputSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonArrowOutputSuite.scala
@@ -22,15 +22,17 @@ import java.io.{ByteArrayInputStream, ByteArrayOutputStream}
import scala.collection.mutable.ArrayBuffer
import scala.jdk.CollectionConverters._
-import org.apache.arrow.vector.{BigIntVector, FieldVector,
FixedSizeBinaryVector, VectorSchemaRoot}
+import org.apache.arrow.vector.{BigIntVector, FieldVector,
FixedSizeBinaryVector, LargeVarCharVector, NullVector, VarCharVector,
VectorSchemaRoot}
import org.apache.arrow.vector.complex.{ListVector, StructVector}
import org.apache.arrow.vector.ipc.{ArrowStreamReader, ArrowStreamWriter}
import org.apache.arrow.vector.types.pojo.{ArrowType, FieldType}
import org.apache.spark.SparkFunSuite
import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.util.{ArrayBasedMapData, GenericArrayData}
import org.apache.spark.sql.execution.arrow.ArrowWriter
import org.apache.spark.sql.execution.metric.SQLMetric
+import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.types._
import org.apache.spark.sql.util.ArrowUtils
import org.apache.spark.sql.vectorized.ArrowColumnVector
@@ -79,8 +81,9 @@ 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
- * next batch is produced, since the processor releases the previous output
batch then.
+ * the rows and the buffer size in bytes, as
`SliceBytesArrowOutputProcessorImpl` measures it,
+ * of every output batch. Rows are copied before the next batch is produced,
since the
+ * processor releases the previous output batch then.
*/
private def readAll(
batchSizes: Seq[Int],
@@ -96,7 +99,8 @@ class PythonArrowOutputSuite extends SparkFunSuite {
while (processor.loadBatch()) {
val batch = processor.produceBatch()
val bytes = (0 until batch.numCols()).map { i =>
-
batch.column(i).asInstanceOf[ArrowColumnVector].getValueVector.getBufferSize.toLong
+ SliceBytesArrowOutputProcessorImpl.bufferSize(
+ batch.column(i).asInstanceOf[ArrowColumnVector].getValueVector)
}.sum
output += ((batch.rowIterator().asScala.map(_.copy()).toSeq, bytes))
}
@@ -216,7 +220,221 @@ 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 a few 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 <= 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 and the number of size
measurements.
+ */
+ private def sliceRowCountsAndMeasurements(batchSizes: Seq[Int], maxBytes:
Int)(
+ rowBytes: Long => Long): (Seq[Int], Int) = {
+ var measured = 0
+ val output = readAll(batchSizes, (reader, metrics) =>
+ new SliceBytesArrowOutputProcessorImpl(reader, metrics, maxBytes) {
+ override private[python] def batchBytes(root: VectorSchemaRoot): Long
= {
+ measured += 1
+ 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), measured)
+ }
+
+ private def sliceRowCounts(batchSizes: Seq[Int], maxBytes: Int)(rowBytes:
Long => Long)
+ : Seq[Int] = sliceRowCountsAndMeasurements(batchSizes,
maxBytes)(rowBytes)._1
+
+ test("SPARK-60395: slice by bytes: number of rows of each slice") {
+ // The first search starts from the rows that would fit at the bytes per
row of the root, and
+ // later searches from those of the last slice, here 10 rows, which fit
exactly.
+ assert(sliceRowCounts(Seq(50), 1000)(_ => 100) === Seq.fill(5)(10))
+ // Row 0 alone is larger than the limit, which does not keep the other
slices at one row. The
+ // first search goes 66, 59, 53, ..., 2, 1 rows, and the next search
starts from 66 rows.
+ assert(sliceRowCounts(Seq(1000), 1000)(i => if (i == 0) 5000 else 10) ===
+ Seq(1, 66) ++ Seq.fill(9)(100) ++ Seq(33))
+ // 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: slice by bytes: rows larger than the limit are measured
once each") {
+ // 100 small rows make the search start from about 100 rows, then 900 rows
are each larger
+ // than the limit. Searching from there for each of them would take about
45 measurements.
+ val (rowCounts, measured) = sliceRowCountsAndMeasurements(Seq(1000), 1000)
{ i =>
+ if (i < 100) 10 else 2000
+ }
+ assert(rowCounts.takeRight(900) === Seq.fill(900)(1))
+ assert(measured <= 2 * rowCounts.size, s"$measured measurements for
$rowCounts")
+ }
+
+ test("SPARK-60395: slice by bytes: a batch that Arrow cannot slice is output
whole") {
+ val batchSizes = Seq(10, 20)
+ val output = readAll(batchSizes, (reader, metrics) =>
+ new SliceBytesArrowOutputProcessorImpl(reader, metrics, 1) {
+ override private[python] def canSlice(root: VectorSchemaRoot): Boolean
= false
+ })
+ checkAllRowsOnceInOrder(batchSizes, output)
+ assert(output.map(_._1.size) === batchSizes)
+ }
+
+ 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: FieldVector): 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
+ })
+ }
+ // An int64 column of 2^28 rows, without allocating it, which the struct
vector measures
+ // only if it measures its children in Long too.
+ struct.addChild("c3", new FixedSizeBinaryVector("c3", allocator, 8) {
+ override def getValueCount: Int = 1 << 28
+ override def setValueCount(valueCount: Int): Unit = {}
+ })
+ struct.setValueCount(1)
+ // `StructVector.getBufferSize` wraps around to a negative size.
+ assert(struct.getBufferSize < 0)
+ val expected = 3L * (1 << 30) + (1L << 31) + (1 << 25) + 1
+ assert(SliceBytesArrowOutputProcessorImpl.bufferSize(struct) ===
expected)
+ val root = new VectorSchemaRoot(Seq[FieldVector](struct).asJava)
+ assert(SliceBytesArrowOutputProcessorImpl.batchBytes(root) === expected)
+ // Arrow cannot slice the int64 child of 2 GiB, which only the check of
the children sees.
+ assert(!BaseSliceArrowOutputProcessor.canSlice(struct))
+ } 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)
+ // An int64 column of 2^28 rows, without allocating it.
+ val fixed = new FixedSizeBinaryVector("f", allocator, 8) {
+ override def getValueCount: Int = 1 << 28
+ }
+ val varChar = new VarCharVector("v", allocator)
+ val largeVarChar = new LargeVarCharVector("l", allocator)
+ try {
+ assert(SliceBytesArrowOutputProcessorImpl.bufferSize(fixed) === (1L <<
31) + (1 << 25))
+
+ // Variable-width vectors of 3 values whose end offset is then set to a
large value.
+ varChar.allocateNew()
+ largeVarChar.allocateNew()
+ Seq("a", "bb", "ccc").zipWithIndex.foreach { case (str, i) =>
+ varChar.setSafe(i, str.getBytes("UTF-8"))
+ largeVarChar.setSafe(i, str.getBytes("UTF-8"))
+ }
+ Seq(varChar, largeVarChar).foreach { v =>
+ v.setValueCount(3)
+ assert(SliceBytesArrowOutputProcessorImpl.bufferSize(v) ===
v.getBufferSize)
+ }
+ varChar.getOffsetBuffer.setInt(3 * 4, Int.MaxValue - 10)
+ // `getBufferSize` wraps around to a negative size.
+ assert(varChar.getBufferSize < 0)
+ assert(SliceBytesArrowOutputProcessorImpl.bufferSize(varChar) ===
+ 1 + 4 * 4 + (Int.MaxValue - 10L))
+ largeVarChar.getOffsetBuffer.setLong(3 * 8, 3L << 30)
+ // `getBufferSize` caps the size at Int.MaxValue.
+ assert(largeVarChar.getBufferSize === Int.MaxValue)
+ assert(SliceBytesArrowOutputProcessorImpl.bufferSize(largeVarChar) ===
+ 1 + 4 * 8 + (3L << 30))
+ } finally {
+ fixed.close()
+ varChar.close()
+ largeVarChar.close()
+ allocator.close()
+ }
+ }
+
+ test("SPARK-60395: maxBytesPerOutputBatch below Int.MinValue means no
limit") {
+ val conf = new SQLConf
+ Seq("-3g", "-4294967295", "-1", "0").foreach { value =>
+
conf.setConfString(SQLConf.ARROW_EXECUTION_MAX_BYTES_PER_OUTPUT_BATCH.key,
value)
+ assert(conf.arrowMaxBytesPerOutputBatch === 0, value)
+ }
+ conf.setConfString(SQLConf.ARROW_EXECUTION_MAX_BYTES_PER_OUTPUT_BATCH.key,
"256m")
+ assert(conf.arrowMaxBytesPerOutputBatch === 256 * 1024 * 1024)
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
