viirya commented on code in PR #59871:
URL: https://github.com/apache/spark/pull/59871#discussion_r4403928991
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonArrowOutputSuite.scala:
##########
@@ -132,6 +135,72 @@ class PythonArrowOutputSuite extends SparkFunSuite {
}
}
+ test("SPARK-60627: slice by records: a batch that Arrow cannot slice is
output whole") {
+ val allocator = ArrowUtils.rootAllocator.newChildAllocator("unsliceable",
0, Long.MaxValue)
+ // An array column that reports 2^29 + 2^20 rows without allocating them,
so its offsets take
+ // over 2 GiB. Arrow computes the positions of the offsets of a slice in
int, so a slice that
+ // reaches row 2^29 throws and leaks the vectors sliced so far. The
fixed-width vectors that
Review Comment:
Without the fix, this test does fail, but not at the int overflow: only
`getValueCount` of the list is overridden, so the struct's `valueCount` field
is still 0, and the first `root.slice(0, 1 << 20)` throws
`IllegalArgumentException` from the `Preconditions` check in `StructVector`'s
transfer pair. And since no buffer is allocated, `allocator.close()` cannot
catch a leak here. Testing only the dispatch is reasonable, since a real
overflow needs a buffer over 2 GiB, but could this comment, and the PR
description's "It fails without this fix" and "nothing is leaked", say that the
test checks that the batch is output whole instead of sliced?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowOutput.scala:
##########
@@ -267,6 +268,29 @@ abstract class BaseSliceArrowOutputProcessor(
closePrevSlice()
super.close()
}
+
+ private[python] def canSlice(root: VectorSchemaRoot): Boolean =
+ root.getFieldVectors.asScala.forall(BaseSliceArrowOutputProcessor.canSlice)
+}
+
+object BaseSliceArrowOutputProcessor {
+ /**
+ * Returns whether `VectorSchemaRoot.slice` can slice `vector` and its
children. Arrow computes
+ * some byte positions of a slice in int, so it throws, and leaks the
vectors sliced so far, for
+ * a fixed-width vector of 2 GiB or more, a view vector with 2^27 rows or
more, or a list vector
+ * with 2^29 - 1 rows or more. A batch that holds one is output whole
instead.
+ */
+ private[python] def canSlice(vector: FieldVector): Boolean = {
+ val fits = vector match {
+ case v: BaseFixedWidthVector => v.getValueCount.toLong * v.getTypeWidth
<= Int.MaxValue
+ case v: BaseVariableWidthViewVector =>
+ v.getValueCount.toLong * BaseVariableWidthViewVector.ELEMENT_SIZE <=
Int.MaxValue
+ case v: ListVector =>
+ (v.getValueCount + 1L) * BaseRepeatedValueVector.OFFSET_WIDTH <=
Int.MaxValue
Review Comment:
The `+ 1L` is the subtle part of this check. `startIndex * OFFSET_WIDTH` and
`(startIndex + length) * OFFSET_WIDTH` only overflow from 2^29 rows. The `+ 1`
is needed because the transfer pair also allocates `(length + 1) *
OFFSET_WIDTH` in int, and a nested list (e.g. `array<array<...>>`) has its
child sliced over the range covered by the parent's slice, which can be all the
child's rows. Could the doc above say so? Otherwise the `+ 1` reads like a typo
for the plain offset positions.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:
##########
@@ -5562,7 +5562,8 @@ object SQLConf {
"there is no limit. Note that the complete ArrowRecordBatch is
actually created but " +
"the number of records is limited when sending it to the downstream
operator. This is " +
"used to avoid large batches being sent to the downstream operator
including " +
- "the columnar-based operator implemented by third-party libraries. " +
+ "the columnar-based operator implemented by third-party libraries. A
batch that " +
+ "Arrow cannot slice is sent as a whole. " +
Review Comment:
nit: a reader of the doc can't tell when Arrow cannot slice a batch. Maybe
"A batch that Arrow cannot slice, e.g. one with a column of 2 GiB or more, is
sent as a whole."
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonArrowOutputSuite.scala:
##########
@@ -132,6 +135,72 @@ class PythonArrowOutputSuite extends SparkFunSuite {
}
}
+ test("SPARK-60627: slice by records: a batch that Arrow cannot slice is
output whole") {
+ val allocator = ArrowUtils.rootAllocator.newChildAllocator("unsliceable",
0, Long.MaxValue)
+ // An array column that reports 2^29 + 2^20 rows without allocating them,
so its offsets take
+ // over 2 GiB. Arrow computes the positions of the offsets of a slice in
int, so a slice that
+ // reaches row 2^29 throws and leaks the vectors sliced so far. The
fixed-width vectors that
+ // Spark reads, e.g. of an int64 column over 2 GiB, are final classes, so
a list stands in.
+ val rows = (1 << 29) + (1 << 20)
+ val list = new ListVector("l", allocator,
FieldType.nullable(ArrowType.List.INSTANCE), null) {
+ override def getValueCount: Int = rows
+ }
+ list.addOrGetVector[BigIntVector](FieldType.nullable(new ArrowType.Int(64,
true)))
+ // Like the output of mapInArrow, the column is in a struct column.
+ val struct = new StructVector(
+ "st", allocator, FieldType.nullable(ArrowType.Struct.INSTANCE), null) {
+ putChild("l", list)
+ }
+ val root = new VectorSchemaRoot(
+ Seq(struct.getField).asJava, Seq[FieldVector](struct).asJava, rows)
+ val reader = new ArrowStreamReader(
+ new ByteArrayInputStream(Array.emptyByteArray), allocator) {
+ private var loaded = false
+ override def getVectorSchemaRoot: VectorSchemaRoot = root
+ override def loadNextBatch(): Boolean = {
+ val hasNext = !loaded
+ loaded = true
+ hasNext
+ }
+ override def bytesRead(): Long = 0L
+ }
+ val metrics = Map(
+ "pythonNumRowsReceived" -> new SQLMetric("sum"),
+ "pythonDataReceived" -> new SQLMetric("size"))
+ val processor = new SliceRecordsArrowOutputProcessorImpl(reader, metrics,
1 << 20)
+ try {
+ assert(processor.loadBatch())
+ assert(processor.produceBatch().numRows() === rows)
+ assert(!processor.loadBatch())
+ } finally {
+ processor.close()
+ allocator.close()
+ }
+ }
+
+ test("SPARK-60627: canSlice rejects vectors whose slice positions overflow
int") {
+ val allocator = ArrowUtils.rootAllocator.newChildAllocator("canSlice", 0,
Long.MaxValue)
+ // Vectors that report many values without allocating them.
+ def fixed(rows: Int): FixedSizeBinaryVector =
+ new FixedSizeBinaryVector("f", allocator, 8) {
+ override def getValueCount: Int = rows
+ }
+ def list(rows: Int): ListVector = new ListVector(
+ "l", allocator, FieldType.nullable(ArrowType.List.INSTANCE), null) {
+ override def getValueCount: Int = rows
+ }
+ val vectors = Seq(fixed((1 << 28) - 1), fixed(1 << 28), list(1000), list(1
<< 29))
Review Comment:
The fixed-width cases sit right at the boundary, but the list cases (`1000`
and `1 << 29`) don't, so dropping the `+ 1L` from `canSlice` would still pass
this test. How about `list((1 << 29) - 2)` -> true and `list((1 << 29) - 1)` ->
false?
--
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]
