dongjoon-hyun commented on code in PR #59871:
URL: https://github.com/apache/spark/pull/59871#discussion_r4404508626
##########
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:
Thank you. Added a paragraph to the scaladoc: the offset positions of a list
vector overflow only from 2^29 rows, but a slice also allocates `(length + 1) *
4` bytes of offsets in int, and a list vector nested in another one is sliced
over all its rows when the slice of its parent covers them.
##########
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:
Thanks. Changed the list cases to `list((1 << 29) - 2)` -> true and `list((1
<< 29) - 1)` -> false, so dropping the `+ 1L` now fails this test.
##########
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:
Right, it checks only the dispatch. One correction: without the fix, it
fails even earlier, with `ClassCastException` from
`AbstractStructVector.addOrGet`, when `StructVector`'s transfer pair re-creates
the anonymous list child, before the `Preconditions` check. I updated the
comment to say that it checks only that the batch is output whole instead of
sliced, and the PR description no longer says that it fails without this fix or
that nothing is leaked. The manual runs in the description cover the overflow
and the leak.
##########
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:
Thanks, updated. I wrote "a fixed-width column of 2 GiB or more", since a
var-width column over 2 GiB can still be sliced, as Arrow computes its
positions in long.
--
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]
