viirya commented on code in PR #59581:
URL: https://github.com/apache/spark/pull/59581#discussion_r4381033411
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/ColumnarEvaluatorFactory.scala:
##########
@@ -44,10 +45,14 @@ class ColumnarToRowEvaluatorFactory(
inputs: Iterator[ColumnarBatch]*): Iterator[InternalRow] = {
assert(inputs.length == 1)
val toUnsafe = UnsafeProjection.create(childOutput, childOutput)
+ // Releases each batch once all its rows have been returned, before the
next batch is
+ // fetched, as the generated code of ColumnarToRowExec does in
advanceBatch: flatMap checks
+ // the iterator of the current batch, which runs its completion, before
the next batch.
inputs.head.flatMap { input =>
numInputBatches += 1
numOutputRows += input.numRows()
- input.rowIterator().asScala.map(toUnsafe)
+ CompletionIterator[InternalRow, Iterator[InternalRow]](
Review Comment:
Thanks, you're right that this holds only where `ColumnarToRowExec` is
compiled. I added a line to the 4.4 section of the SQL migration guide in
dd1c0fefaab: the conversion of columnar batches to rows now releases each batch
with `closeIfFreeable` also when it is not whole-stage code-generated, e.g.
with more than `spark.sql.codegen.maxFields` fields or with
`spark.sql.codegen.wholeStage=false`, as the generated code has done since
3.5.4 and 4.0, so a `ColumnVector` that a Data Source V2 source reuses across
batches must override `closeIfFreeable` to do nothing. The user-facing section
of the description says so too.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowEvalPythonExec.scala:
##########
@@ -117,28 +117,13 @@ case class ArrowEvalPythonExec(
// accept columnar input to avoid the ColumnarToRow -> ArrowWriter
round-trip.
// The Arrow FieldVectors are extracted directly from ArrowColumnVector and
// serialized to IPC, bypassing the row-based ArrowWriter conversion.
+ // This node then outputs only columnar batches, so the planner inserts a
+ // ColumnarToRowExec above it for a row-based parent, which releases each
Review Comment:
Thanks, added to the class doc in dd1c0fefaab: the consumer of the output
releases the pass-through columns with `closeIfFreeable` on the task thread,
and with `spark.python.udf.pipelined.enabled` the Python writer reads the next
input batches on its own thread meanwhile, so the upstream operator must allow
its vectors to be released while it returns the next batches.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala:
##########
@@ -255,12 +258,12 @@ private[python] class
ColumnarArrowEvalPythonEvaluatorFactory(
val resultIter = if (inputColumnIndices.isDefined) {
// Path 2: simple column refs, non-Arrow vectors.
- val bufferedIter = inputIter.map { batch =>
+ val bufferedIter = new WrittenBatchReleasingIterator(inputIter, {
batch =>
Review Comment:
Thanks, added in dd1c0fefaab: the test source has an `arrowIdColumn` option
that wraps `id` in a release-counting `ColumnVector` that delegates to the
Arrow vector but is not an `ArrowColumnVector`, and the UDF test runs
`arrow_release_udf(name)` over it with the same value and counter assertions,
with whole-stage codegen on and off. It also checks the path from the output
batches, which carry the input columns only on the Arrow columnar path. If this
iterator stops releasing, the test fails with 0 released batches out of 20, and
if it releases a batch before the writer serializes it, the UDF results come
back as `null`.
##########
sql/core/src/main/java/org/apache/spark/sql/execution/vectorized/ReusedArrowColumnVector.java:
##########
@@ -0,0 +1,43 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark.sql.execution.vectorized;
+
+import org.apache.arrow.vector.ValueVector;
+
+import org.apache.spark.sql.vectorized.ArrowColumnVector;
+
+/**
+ * An {@link ArrowColumnVector} whose producer may release or overwrite it
after returning it, for
+ * example by loading every batch into the same Arrow vectors. A consumer must
not hold it past the
+ * next batch.
+ *
+ * <p>Like the vectors that {@link WritableColumnVector} and {@link
ConstantColumnVector} reuse
+ * across batches, the producer releases it: {@link #closeIfFreeable()}, which
a consumer calls
+ * once it is done with a batch, is a no-op, and the producer frees the vector
with
+ * {@link #close()}.
+ */
+public class ReusedArrowColumnVector extends ArrowColumnVector {
+
+ public ReusedArrowColumnVector(ValueVector vector) {
+ super(vector);
+ }
+
+ @Override
+ public void closeIfFreeable() {
Review Comment:
Thanks, I made the Javadoc say that it covers the vector itself, and that
the vectors `getChild` returns for the fields of a struct are plain
`ArrowColumnVector`s, which consumers do not release on their own. Returning
reused child vectors would need a change to `ArrowColumnVector`, which builds
them in package-private code.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/ArrowBackedDataSourceV2.scala:
##########
@@ -269,24 +304,42 @@ private class ArrowBackedPartitionReader(start: Int, end:
Int, batchSize: Int)
// Wrap in ArrowColumnVector
val columns: Array[ColumnVector] = Array(
- new ArrowColumnVector(idVector),
+ new ReleaseCountingArrowColumnVector(idVector),
new ArrowColumnVector(nameVector),
new ArrowColumnVector(valueVector),
new ArrowColumnVector(dataVector))
currentBatch = new ColumnarBatch(columns, count)
- allBatches.add(currentBatch)
+ allBatches.synchronized {
+ allBatches.add(currentBatch)
+ ArrowBackedDataSourceV2.returnedBatches.incrementAndGet()
+ ArrowBackedDataSourceV2.maxUnreleasedBatches.accumulateAndGet(
+ allBatches.size - numReleased, (a, b) => math.max(a, b))
+ }
current += count
true
}
override def get(): ColumnarBatch = currentBatch
+ // The consumer may still hold the returned batches, so they are closed when
the task completes.
+ // The listener is registered here rather than in the constructor: the
reader may be created by
+ // the writer thread of a pipelined Python UDF after the listener that waits
for that writer to
+ // stop, and the listeners run in reverse order, so the batches could be
freed while the writer
Review Comment:
Thanks, I hadn't noticed that SPARK-60113 made this moot. The registration
is back in the constructor in dd1c0fefaab, without `closed`, and the comment
says that the pipelined writer is stopped by a pre-completion listener, which
runs before it.
--
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]
