viirya commented on code in PR #59582:
URL: https://github.com/apache/spark/pull/59582#discussion_r4461161282
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatchSerializer.scala:
##########
@@ -1348,19 +1350,24 @@ private class ArrowCachedBatchToColumnarBatchIterator(
deserializeToRoot(batchIter.next().asInstanceOf[ArrowCachedBatch])
}
- previousRoot = root
-
// When projected on load, the root already holds only the selected
columns in output order,
// so wrap its vectors directly. Otherwise it holds all cache columns and
must be selected.
- val selectedColumns = if (canProjectOnLoad) {
- root.getFieldVectors.asScala.map(v => new
ArrowColumnVector(v)).toArray[ColumnVector]
+ val selectedVectors = if (canProjectOnLoad) {
+ root.getFieldVectors.asScala.toArray
+ } else {
+ columnIndices.map(root.getFieldVectors.get(_))
+ }
+ val numRows = root.getRowCount
+ val batch = if (selectedVectors.isEmpty) {
+ // No column can release the root, and it holds no buffer the batch
needs.
+ root.close()
+ new ColumnarBatch(Array.empty[ColumnVector], numRows)
} else {
- val allColumns = root.getFieldVectors.asScala.map { vector =>
- new ArrowColumnVector(vector)
- }.toArray[ColumnVector]
- columnIndices.map(allColumns(_))
+ unreleasedRoots.synchronized(unreleasedRoots.add(root))
Review Comment:
Fixed in cfce388a0f5 by reverting `ArrowCachedBatchSerializer` to master: it
closes the previous root in `next()` again, so the row fallback releases each
batch as before. Its vectors are not retainable, so cache + UDF takes the row
queue path.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala:
##########
@@ -200,15 +214,22 @@ private[python] class
ColumnarArrowEvalPythonEvaluatorFactory(
outputTypes: Seq[DataType],
columnIndices: Array[Int]): Iterator[ColumnarBatch] = {
+ // In pipelined mode the batches are added in the writer thread and
+ // polled in the task thread.
val passThruQueue =
- new ArrayDeque[(Array[ColumnVector], Int)]()
+ new ConcurrentLinkedQueue[(Array[ColumnVector], Int)]()
- val bufferedIter = inputIter.map { batch =>
+ val bufferedIter = new WrittenBatchReleasingIterator(inputIter, { batch
=>
val passThruCols = childOutput.indices.map(
i => batch.column(i)).toArray
- passThruQueue.add((passThruCols, batch.numRows()))
- batch
- }
+ if (passThruCols.exists(_.isInstanceOf[ReusedColumnVector])) {
Review Comment:
Fixed in cfce388a0f5: both checks now use the same predicate. The first
batch takes Path 1 only if every column is an `ArrowColumnVector` whose
`isRetainable` is true, and a later batch is held only if all its pass-through
columns are retainable, otherwise they are copied. An `OrcColumnVector` is not
retainable, so it is copied.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala:
##########
@@ -200,15 +214,22 @@ private[python] class
ColumnarArrowEvalPythonEvaluatorFactory(
outputTypes: Seq[DataType],
columnIndices: Array[Int]): Iterator[ColumnarBatch] = {
+ // In pipelined mode the batches are added in the writer thread and
+ // polled in the task thread.
val passThruQueue =
- new ArrayDeque[(Array[ColumnVector], Int)]()
+ new ConcurrentLinkedQueue[(Array[ColumnVector], Int)]()
- val bufferedIter = inputIter.map { batch =>
+ val bufferedIter = new WrittenBatchReleasingIterator(inputIter, { batch
=>
val passThruCols = childOutput.indices.map(
i => batch.column(i)).toArray
- passThruQueue.add((passThruCols, batch.numRows()))
- batch
- }
+ if (passThruCols.exists(_.isInstanceOf[ReusedColumnVector])) {
+ passThruQueue.add((copyColumns(batch), batch.numRows()))
+ true
+ } else {
+ passThruQueue.add((passThruCols, batch.numRows()))
Review Comment:
Not fixed, but now part of the contract: the doc of
`ColumnVector.isRetainable` says that a consumer that keeps such vectors does
not account for their memory and cannot spill them, so a producer opts in
knowing that. The conf doc says the same. Since nothing in the tree opts in,
this only applies to a source that chooses to.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala:
##########
@@ -155,11 +158,17 @@ private[python] class
ColumnarArrowEvalPythonEvaluatorFactory(
val inputColumnIndices = resolveColumnIndices(allInputs.toSeq)
- // Peek at first batch to check if Arrow-backed.
+ // Peek at the first batch to check whether its columns are Arrow vectors
+ // that may be held until the result of the batch comes back.
val peekIter = new PeekableIterator(inputIter)
val isArrow = peekIter.peek().exists { batch =>
- batch.numCols() > 0 &&
- batch.column(0).isInstanceOf[ArrowColumnVector]
+ batch.numCols() > 0 && (0 until batch.numCols()).forall { i =>
+ batch.column(i) match {
+ case _: ReusedColumnVector => false
+ case _: ArrowColumnVector => true
Review Comment:
Agreed, and fixed in cfce388a0f5 as you suggest: `ReusedColumnVector` is
gone, and a producer opts in with the public `ColumnVector.isRetainable()`,
which is false by default. A source that does not know about it is not held.
See my comment on the PR for the contract and the alternatives I considered.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala:
##########
@@ -314,6 +335,15 @@ private[python] class
ColumnarArrowEvalPythonEvaluatorFactory(
rowsToColumnarBatches(rowIter, context)
}
+ /** Copies the pass-through columns of `batch` into on-heap vectors. */
+ private def copyColumns(batch: ColumnarBatch): Array[ColumnVector] = {
+ val vectors = OnHeapColumnVector.allocateColumns(
+ math.max(batch.numRows(), 1),
childSchema).toArray[WritableColumnVector]
+ val converter = new RowToColumnConverter(childSchema)
Review Comment:
Agreed. This also fails with columnar input enabled on 4.2 and 4.3, where it
is the default, so I'll file it as a bug and fix it separately. This PR no
longer changes the default on master.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala:
##########
@@ -314,6 +335,15 @@ private[python] class
ColumnarArrowEvalPythonEvaluatorFactory(
rowsToColumnarBatches(rowIter, context)
Review Comment:
Agreed. Like item 5, it happens with the default configuration of 4.2 and
4.3, so I'll file it as a bug and fix it separately.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala:
##########
@@ -314,6 +335,15 @@ private[python] class
ColumnarArrowEvalPythonEvaluatorFactory(
rowsToColumnarBatches(rowIter, context)
Review Comment:
Agreed. Like item 5, it happens with the default configuration of 4.2 and
4.3, so I'll file it as a bug and fix it separately.
--
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]
