dongjoon-hyun commented on code in PR #59786:
URL: https://github.com/apache/spark/pull/59786#discussion_r4383082511
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowOutput.scala:
##########
@@ -265,47 +265,85 @@ 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
+
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 that would fit at the bytes per row of the last slice that fit,
rather than from all
+ // the remaining rows, since a slice that does not start at row 0 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)
Review Comment:
**9. [Medium] A fixed-width column over 2 GiB now makes this `root.slice`
throw, where the merge base output the batch whole**
With the Long measurement, a Python output batch whose fixed-width leaf
holds more than 2 GiB now reaches this `root.slice`, including the fixed-width
element vector of an array or map column. In Arrow 19, and still on arrow-java
main, `BaseFixedWidthVector.splitAndTransferValueBuffer` computes `startIndex *
typeWidth` and `length * typeWidth` in int, so the slice throws
`IndexOutOfBoundsException` and the task fails.
For example, take `maxBytesPerOutputBatch=256m` and one applyInPandas group
or mapInArrow batch of 300,000,000 rows with one bigint column. At the merge
base, the int `getBufferSize` of the child wrapped to -1,857,467,296 and that
of the struct to -1,819,967,296, so the whole root went out. Here `bufferSize`
gives 2,475,000,000, L278 sets `sliceRowCount` to 270,000,000, and
`root.slice(0, 270000000)` computes 270,000,000 * 8 = -2,134,967,296 in int, so
`ArrowBuf.slice` throws on every attempt. The half-built slice is never closed
either, so `allocator.close()` in the completion listener (L80-88) also reports
a leak.
An int64 column fails at the first probe from 298,261,618 rows, and from
268,435,457 rows once a slice starts at row 2^28 or later. The same holds for
16-byte decimals over 2^27 values and int32 over 2^29 values. So "Output
batches over 2 GiB are now sliced" in the description holds only while every
fixed-width leaf is under 2 GiB, as in its 3 x 2^27 example.
Could L274 also output the whole root when it cannot be sliced, e.g. `if
(batchBytes(root) <= arrowMaxBytesPerOutputBatch || !canSlice(root))`?
`canSlice` would walk struct children and the element vectors of lists and
maps, and return false if any `BaseFixedWidthVector` has `getValueCount.toLong
* getTypeWidth > Int.MaxValue`. That also covers sizes the merge base already
failed on, e.g. two int64 columns of 300M rows. The scaladoc could name the
gap, and it may be worth an arrow-java issue.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowOutput.scala:
##########
@@ -265,47 +265,85 @@ 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
+
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 that would fit at the bytes per row of the last slice that fit,
rather than from all
+ // the remaining rows, since a slice that does not start at row 0 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
}
- 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. The next
+ // search starts from where this one did, since the next rows may be much
smaller.
Review Comment:
**10. [Low] Correction to my item 1: with the estimate in place,
`sliceRowCount = 1` no longer needs to go**
In item 1, I suggested the estimate "plus dropping `sliceRowCount = 1`", and
I gave a +1 to [the earlier
point](https://github.com/apache/spark/pull/59786#discussion_r4372986243) that
asked for the same. Its premise, that `sliceRowCount = 1` pins every later
slice of the batch to one row, held only before the estimate existed. Now the
one-row fallback leaves `sliceRowCount` unchanged, so in a run of rows that are
each over the limit, or while nothing in the batch has fit yet, every row
reruns the whole 10% search from up to `0.9 * rowCount` rows.
For example, a Python batch of 10,000 rows of 2,000 bytes with a 1,000-byte
limit, the "even one record is still larger" case of the config doc, takes
620,916 measurements and builds slices of about 492M rows in total. That is
about the same as the merge base (620,979 and 497M), versus 10,071 measurements
and 110K rows at a52b968. Array and map probes always copy (n + 1) offsets, and
string probes copy them at a non-zero start offset, so this is the c = 1 case
of the O(R^2/c) cost that the description says is removed.
With the estimate, restoring `sliceRowCount = 1` would not pin later slices,
because the next row, probed alone, seeds the estimate at L290. The case at
`PythonArrowOutputSuite.scala:220` would become `Seq(1, 1) ++ Seq.fill(9)(100)
++ Seq(98)`, still 12 batches with 21 fewer probes, and the 3-batch example in
item 1 does not change. It is not free, though: the row after each isolated
oversized row would go out alone, one more batch for no saving. Could the
fallback instead probe the next row alone, output it directly if it is also
over the limit, and otherwise close it and search from the unchanged
`sliceRowCount`? That needs one field and about 8 lines, and has neither cost.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowOutput.scala:
##########
@@ -265,47 +265,85 @@ 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
+
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
Review Comment:
**12. [Low] The first search of each batch starts blind, although L274 just
measured the root**
When a Python batch does not fit, its first search starts at `(rowCount *
0.9).toInt` and shrinks by 10% per probe, while the size from L274 is thrown
away. With 10,000 rows of 1,000 bytes and a 1 MiB limit, that is 32
measurements for 10 batches, 22 of them in the first search, over 80,095 rows.
These probes start at row 0, so flat columns share their buffers, but in Arrow
19 `ListVector` and `MapVector` allocate and rebase (length + 1) offsets on
every slice, even from row 0 (`ListVector.java` L562-568). That is about 320 KB
of offsets per array or map column per batch, versus about 40 KB for the output
slices.
Could the first search start from the estimate of L290, applied to the root?
```scala
val rootBytes = batchBytes(root)
if (rootBytes <= arrowMaxBytesPerOutputBatch) {
currentRowIdx = rowCount
return root
}
val estimate = math.max(rowCount.toLong * arrowMaxBytesPerOutputBatch
/ rootBytes, 1L)
sliceRowCount =
math.min(estimate, (rowCount *
SliceBytesArrowOutputProcessorImpl.stepRatio).toLong).toInt
```
The example then takes 11 measurements for the same 10 batches, and the
first search never takes more probes than now. The `max(..., 1L)` is needed,
since an estimate of 0 would skip the loop and output every row alone. The
trade-off is that a batch whose leading rows are smaller than its average gets
one more output batch, e.g. 17 instead of 16 in the first model of my item 1,
and two expectations in `PythonArrowOutputSuite.scala` change, e.g. L218
becomes `Seq.fill(5)(10)`. This cost predates the PR (SPARK-51931), so it could
also be a follow-up.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowOutput.scala:
##########
@@ -265,47 +265,85 @@ 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
+
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 that would fit at the bytes per row of the last slice that fit,
rather than from all
+ // the remaining rows, since a slice that does not start at row 0 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
}
- 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. The next
+ // search starts from where this one did, since the next rows may be much
smaller.
+ val rootSlice = root.slice(currentRowIdx, 1)
Review Comment:
**13. [Nit] The one-row fallback still rebuilds the one-row slice that the
loop just closed**
When a search never fits, its last probe is always the one-row slice, since
`(n * 0.9).toInt` strictly decreases and reaches 1 for every n >= 2. L295
closes that slice and L301 builds it again, so every oversized row pays for one
more `VectorSchemaRoot.slice`. That slice copies validity bits when the start
is not a multiple of 8, the offsets of array and map columns always, and those
of string columns at a non-zero start offset. The all-oversized case at
`PythonArrowOutputSuite.scala:226` builds 12 slices where 7 would do. This is
the second half of [my comment on
#59616](https://github.com/apache/spark/pull/59616#discussion_r4331550202)
("the one-row fallback rebuilds the 1-row slice that the loop just closed"),
which you [added to
SPARK-60395](https://github.com/apache/spark/pull/59616#discussion_r4331794478).
The first half, the `slice(0, rowCount)` probe, is fixed.
Could the loop return that slice itself, e.g. before L295?
```scala
if (numRows == 1) {
// Even a single row is larger than the limit, so we output a batch
with that one row.
currentRowIdx += 1
return rootSlice
}
```
L301 would then remain only for a one-row batch whose root does not fit,
where the loop does not run, so [the `root.slice(currentRowIdx, 1)`
form](https://github.com/apache/spark/pull/59616#discussion_r4331550214) stays.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonArrowOutputSuite.scala:
##########
@@ -76,7 +81,8 @@ 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
+ * the rows and the buffer size in bytes, as
`SliceBytesArrowOutputProcessorImpl` measures it,
+ * of every output batch. Rows are copied before the
Review Comment:
**17. [Nit] The `readAll` scaladoc has a short line left from the edit**
Inserting "as `SliceBytesArrowOutputProcessorImpl` measures it," left L85 at
54 characters, between lines of 96 and 89. Could L85-86 be re-flowed, with the
same text?
```scala
* of every output batch. Rows are copied before the next batch is
produced, since the
* processor releases the previous output batch then.
```
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowOutput.scala:
##########
@@ -265,47 +265,85 @@ 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
+
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 that would fit at the bytes per row of the last slice that fit,
rather than from all
+ // the remaining rows, since a slice that does not start at row 0 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
}
- 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. The next
+ // search starts from where this one did, since the next rows may be much
smaller.
+ 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 = {
+
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,
Review Comment:
**14. [Nit] The scaladoc no longer says that large variable-width vectors
cap their size**
At a52b968, this scaladoc said that the large variable-width vectors "cap
their size at Int.MaxValue". That sentence went away with the new arm, so the
first paragraph now says that `getBufferSize` "computes the size in int, so it
wraps around" for every vector. But
`BaseLargeVariableWidthVector.getBufferSizeFor` computes it in long and caps it
with `capAtMaxInt`, and a limit of `Int.MaxValue` accepts the capped size under
`<=`, which is why the arm at L334 exists. Only the test comment at
`PythonArrowOutputSuite.scala:323` says so now. Could the paragraph keep it?
For example:
```scala
* 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.
```
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonArrowOutputSuite.scala:
##########
@@ -145,7 +152,194 @@ 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)
Review Comment:
**16. [Nit] `newSliceBytesProcessor` could be placeholder syntax**
L165 already builds a factory with `new ArrowOutputProcessorImpl(_, _)`, and
the parameter type of `readAll` gives the types here too. Could L159-162 go,
with these at L168 and L172?
```scala
val output = readAll(batchSizes, new
SliceBytesArrowOutputProcessorImpl(_, _, maxBytes))
val sliced = readAll(batchSizes, new
SliceBytesArrowOutputProcessorImpl(_, _, maxBytes - 1))
```
The suite would then build these processors in two ways instead of three.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonArrowOutputSuite.scala:
##########
@@ -145,7 +152,194 @@ 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("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, 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("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")
+ }
+
+ /**
+ * 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.
+ */
+ private def sliceRowCounts(batchSizes: Seq[Int], maxBytes: Int)(rowBytes:
Long => Long)
+ : Seq[Int] = {
+ val output = readAll(batchSizes, (reader, metrics) =>
+ new SliceBytesArrowOutputProcessorImpl(reader, metrics, maxBytes) {
+ override private[python] def batchBytes(root: VectorSchemaRoot): Long
= {
+ 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)
+ }
+
+ test("SPARK-60395: slice by bytes: number of rows of each slice") {
+ // The first search goes 45, 40, 36, ..., 11, 9 rows, and later searches
start from the rows
+ // that would fit at the bytes per row of the last slice, here 10, which
fits exactly.
+ assert(sliceRowCounts(Seq(50), 1000)(_ => 100) === Seq(9, 10, 10, 10, 10,
1))
+ // Row 0 alone is larger than the limit, which does not keep the other
slices at one row.
+ assert(sliceRowCounts(Seq(1000), 1000)(i => if (i == 0) 5000 else 10) ===
+ Seq(1, 94) ++ Seq.fill(9)(100) ++ Seq(5))
+ // 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: 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: NullVector): Unit = putChild(name,
child)
+ }
+ val struct = new TestStructVector
+ try {
+ // Children that report 1 GiB each, like three int64 columns of 2^27
rows.
Review Comment:
**11. [Low] This test still passes if the nested arm stops recursing**
The children are `NullVector`s, whose `bufferSize` is just their overridden
`getBufferSize`. So if `PythonArrowOutput.scala:345` became `ownBytes +
children.map(_.getBufferSize.toLong).sum`, this test would still get `3L * (1
<< 30) + 1`. The parity test above is an identity below 2 GiB, and the flat
test below measures only top-level vectors, so every test would pass. Yet with
that change, a struct with a bigint child of 2^28 rows would measure
-2,080,374,784 again and go out unsliced, the bug this PR fixes, and a struct
with a child over 2 GiB is the shape of every mapInArrow, mapInPandas and
applyInPandas output. [The earlier
point](https://github.com/apache/spark/pull/59786#discussion_r4372986253)
covered only the top-level flat arms.
Could `addChild` take a `FieldVector`, and the test add a fourth child that
reports 2^28 values without allocating them?
```scala
struct.addChild("c3", new FixedSizeBinaryVector("c3", allocator, 8) {
override def getValueCount: Int = 1 << 28
override def setValueCount(valueCount: Int): Unit = {}
})
```
The expected size would become `3L * (1 << 30) + (1L << 31) + (1 << 25) +
1`, while the non-recursive variant gives the old `3L * (1 << 30) + 1` and
fails.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonArrowOutputSuite.scala:
##########
@@ -145,7 +152,194 @@ 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("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, 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("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")
+ }
+
+ /**
+ * 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.
+ */
+ private def sliceRowCounts(batchSizes: Seq[Int], maxBytes: Int)(rowBytes:
Long => Long)
+ : Seq[Int] = {
+ val output = readAll(batchSizes, (reader, metrics) =>
+ new SliceBytesArrowOutputProcessorImpl(reader, metrics, maxBytes) {
+ override private[python] def batchBytes(root: VectorSchemaRoot): Long
= {
+ 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)
+ }
+
+ test("SPARK-60395: slice by bytes: number of rows of each slice") {
+ // The first search goes 45, 40, 36, ..., 11, 9 rows, and later searches
start from the rows
+ // that would fit at the bytes per row of the last slice, here 10, which
fits exactly.
+ assert(sliceRowCounts(Seq(50), 1000)(_ => 100) === Seq(9, 10, 10, 10, 10,
1))
+ // Row 0 alone is larger than the limit, which does not keep the other
slices at one row.
+ assert(sliceRowCounts(Seq(1000), 1000)(i => if (i == 0) 5000 else 10) ===
+ Seq(1, 94) ++ Seq.fill(9)(100) ++ Seq(5))
+ // 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: 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: NullVector): 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
+ })
+ }
+ struct.setValueCount(1)
+ // `StructVector.getBufferSize` wraps around to a negative size.
+ assert(struct.getBufferSize < 0)
+ assert(SliceBytesArrowOutputProcessorImpl.bufferSize(struct) === 3L * (1
<< 30) + 1)
+ } 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)
+ try {
+ // An int64 column of 2^28 rows, without allocating it.
+ val fixed = new FixedSizeBinaryVector("f", allocator, 8) {
+ override def getValueCount: Int = 1 << 28
+ }
+ try {
Review Comment:
**15. [Nit] This test could use one try/finally, like the struct test above**
In Arrow 19, the constructors of `FixedSizeBinaryVector`, `VarCharVector`
and `LargeVarCharVector` allocate nothing, and `fixed` is never allocated. So
the try/finally at L297-301 guards nothing, and the one at L306 and L327 only
orders closes that one `finally` would order the same way. Could the test
create `fixed`, `varChar` and `largeVarChar` right after the allocator, and
close them all in one `try { ... } finally { fixed.close(); varChar.close();
largeVarChar.close(); allocator.close() }`? That is about 39 lines at one level
instead of 45 at three, in the same shape as the test above.
--
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]
