viirya commented on code in PR #59786:
URL: https://github.com/apache/spark/pull/59786#discussion_r4386261306
##########
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:
Good catch, confirmed in Arrow 19: `splitAndTransferValueBuffer` computes
`startIndex * typeWidth` and `length * typeWidth` in int, and `ArrowBuf.slice`
then fails `checkPositionIndex`. Added `canSlice`, and a batch that fails it is
output whole, as at the merge base. Besides a fixed-width vector over 2 GiB, it
also rejects a list (or map) vector with more than 2^29 rows, since
`ListVector`'s transfer pair computes the positions of its offsets in int too,
and a view vector with more than 2^27 rows, which computes `startIndex * 16` in
int. It checks every level of the schema. The scaladoc names the gap, and the
description says so in items 7 and the user-facing change.
##########
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:
Took the variant you suggested: after a row larger than the limit is output
alone, the next row is probed alone and output if it is also larger than the
limit; otherwise that slice is closed and the search starts from the unchanged
`sliceRowCount`. The new test "rows larger than the limit are measured once
each" has 100 small rows followed by 900 rows over the limit; without the probe
it took 24452 measurements, now fewer than 2 per output batch.
##########
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:
Added the 2^28-value int64 child as suggested (`addChild` takes a
`FieldVector`), so the expected size is `3L * (1 << 30) + (1L << 31) + (1 <<
25) + 1`. Replacing the recursion with each child's `getBufferSize.toLong` now
fails the test with 3221225473.
##########
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)
Review Comment:
`batchBytes` now lives in the companion, and the processor's method calls
it. The struct test also measures a root holding that struct through it, so
summing the field vectors' `getBufferSize` fails the test with -1073741823.
##########
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:
Done as suggested: the first search starts from `max(rowCount * limit /
rootBytes, 1)`, at most `0.9 * rowCount`. The uniform case is now
`Seq.fill(5)(10)`, and the row-0 case `Seq(1, 66) ++ Seq.fill(9)(100) ++
Seq(33)`.
##########
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:
Done: the loop returns the one-row slice it measured, and
`root.slice(currentRowIdx, 1)` after the loop is left only for a one-row batch
that does not fit.
--
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]
