[SPARK-60325][SQL] Re-enable Arrow columnar input to Python UDFs, holding only vectors that are not reused

sql/core/src/main/java/org/apache/spark/sql/execution/vectorized/ConstantColumnVector.java

@@ -34,8 +34,11 @@
  * which will set the exact same value to all rows.
  *
  * Capacity: The vector stores only one copy of the data.
+ *
+ * Readers reuse a ConstantColumnVector across batches and free it when they are closed, so it is a
+ * {@link ReusedColumnVector}.
  */
-public class ConstantColumnVector extends ColumnVector {
+public class ConstantColumnVector extends ColumnVector implements ReusedColumnVector {
 
   // The data stored in this ConstantColumnVector, the vector stores only one copy of the data.
   private byte nullData;

sql/core/src/main/java/org/apache/spark/sql/execution/vectorized/ReusedArrowColumnVector.java

@@ -0,0 +1,40 @@
+/*
+ * 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. See {@link ReusedColumnVector}.
+ *
+ * <p>The producer releases the vector, so {@link #closeIfFreeable()} is a no-op, and the
+ * producer frees it with {@link #close()}.
+ */
+public class ReusedArrowColumnVector extends ArrowColumnVector implements ReusedColumnVector {
+
+  public ReusedArrowColumnVector(ValueVector vector) {
+    super(vector);
+  }
+
+  @Override
+  public void closeIfFreeable() {
+    // no-op: the producer owns the vector and frees it in `close`.
+  }
+}

sql/core/src/main/java/org/apache/spark/sql/execution/vectorized/ReusedColumnVector.java

@@ -0,0 +1,37 @@
+/*
+ * 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;
+
+/**
+ * A marker for a {@link org.apache.spark.sql.vectorized.ColumnVector} that its producer may
+ * release or overwrite after returning it in a batch: for example, a reader that refills the same
+ * vectors for every batch, or frees the vectors of a batch once the next one is read or once it is
+ * closed. A consumer may only read such a vector until it asks the producer for the next batch,
+ * and must copy what it needs past that point.
+ *
+ * <p>There is no other agreement between the producer and the consumer of a batch, so the type of
+ * its vectors is the protocol, as for {@code ColumnVector.closeIfFreeable}: the producer owns the
+ * lifetime of a vector carrying this marker, and its {@code closeIfFreeable} is a no-op. A vector
+ * without it may be held by the consumer past the next batch without a copy, and the consumer
+ * releases it with {@code closeIfFreeable} once done with it.
+ *
+ * <p>The marker covers the whole vector, so it is only checked on the top-level columns of a
+ * batch. The child vectors of a struct, array or map column belong to that column and share its
+ * lifetime.
+ */
+public interface ReusedColumnVector {
+}

sql/core/src/main/java/org/apache/spark/sql/execution/vectorized/WritableColumnVector.java

@@ -50,9 +50,9 @@
  * A WritableColumnVector should be considered immutable once originally created. In other words,
  * it is not valid to call put APIs after reads until reset() is called.
  *
- * WritableColumnVector are intended to be reused.
+ * WritableColumnVector are intended to be reused, so they are {@link ReusedColumnVector}s.
  */
-public abstract class WritableColumnVector extends ColumnVector {
+public abstract class WritableColumnVector extends ColumnVector implements ReusedColumnVector {
   private final byte[] byte8 = new byte[8];
 
   // Lookup table that "expands" each of the 8 bits in a byte into a separate byte within a long.

sql/core/src/main/scala/org/apache/spark/sql/execution/ColumnarEvaluatorFactory.scala

@@ -17,8 +17,6 @@
 
 package org.apache.spark.sql.execution
 
-import scala.jdk.CollectionConverters._
-
 import org.apache.spark.{PartitionEvaluator, PartitionEvaluatorFactory, TaskContext}
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions.{Attribute, UnsafeProjection}
@@ -44,12 +42,52 @@ class ColumnarToRowEvaluatorFactory(
         inputs: Iterator[ColumnarBatch]*): Iterator[InternalRow] = {
       assert(inputs.length == 1)
       val toUnsafe = UnsafeProjection.create(childOutput, childOutput)
-      inputs.head.flatMap { input =>
+      new ColumnarToRowIterator(inputs.head, toUnsafe, { batch =>
         numInputBatches += 1
-        numOutputRows += input.numRows()
-        input.rowIterator().asScala.map(toUnsafe)
+        numOutputRows += batch.numRows()
+      })
+    }
+  }
+}
+
+/**
+ * Flattens columnar batches into `UnsafeRow`s. A batch is released with closeIfFreeable once all
+ * its rows have been returned, before the next batch is fetched, as the generated code of
+ * [[ColumnarToRowExec]] does in [[ColumnarToRowExec.advanceBatch]]. The rows are projected into a
+ * buffer of their own, so the last row of a batch stays valid after the batch is released.
+ *
+ * @param onNextBatch called on each batch once it is fetched
+ */
+private[sql] class ColumnarToRowIterator(
+    input: Iterator[ColumnarBatch],
+    toUnsafe: UnsafeProjection,
+    onNextBatch: ColumnarBatch => Unit = _ => ())
+  extends Iterator[InternalRow] {
+
+  private var batch: ColumnarBatch = _
+  private var rows: java.util.Iterator[InternalRow] = java.util.Collections.emptyIterator()
+
+  override def hasNext: Boolean = {
+    while (!rows.hasNext) {
+      if (batch != null) {
+        batch.closeIfFreeable()
+        batch = null
+      }
+      if (!input.hasNext) {
+        return false
       }
+      batch = input.next()
+      onNextBatch(batch)
+      rows = batch.rowIterator()
+    }
+    true
+  }
+
+  override def next(): InternalRow = {
+    if (!hasNext) {
+      throw new NoSuchElementException("End of the columnar batches")
     }
+    toUnsafe(rows.next())
   }
 }
 

sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatchSerializer.scala

@@ -1257,6 +1257,12 @@ private class ColumnarBatchToArrowCachedBatchIterator(
 
 /**
  * Iterator that converts ArrowCachedBatch to ColumnarBatch.
+ *
+ * Every batch is loaded into its own root, and its columns are plain [[ArrowColumnVector]]s, which
+ * a consumer may hold past the next batch: the consumer releases a batch with closeIfFreeable once
+ * it is done with it (see [[org.apache.spark.sql.execution.vectorized.ReusedColumnVector]]), which
+ * closes its root once all its columns are closed. The roots of the batches that are not released
+ * that way, e.g. when the consumer stops early, are closed when the task completes.
  */
 private class ArrowCachedBatchToColumnarBatchIterator(
     batchIter: Iterator[CachedBatch],
@@ -1291,14 +1297,16 @@ private class ArrowCachedBatchToColumnarBatchIterator(
       arrowSchema
     }
 
-  // Track only the previous root to close it when next batch is produced
-  private var previousRoot: VectorSchemaRoot = null
+  // The roots of the returned batches that the consumer has not released yet. A consumer may
+  // release a batch from another thread than the one calling next(), e.g. the writer thread of a
+  // pipelined Python UDF, so it is guarded by its own lock.
+  private val unreleasedRoots = java.util.Collections.newSetFromMap(
+    new java.util.IdentityHashMap[VectorSchemaRoot, java.lang.Boolean]())
 
   // Prefetch support: deserialize the next batch into its own root in a background thread while
   // the current batch is being consumed. Only the deserialization (IPC read + decompression +
-  // loading into a fresh root) happens off-thread; closing the previous root stays on the
-  // consumer thread in next(), so the vectors backing a returned ColumnarBatch are never released
-  // while the consumer may still read them.
+  // loading into a fresh root) happens off-thread; the vectors backing a returned ColumnarBatch are
+  // only released by its consumer, so never while the consumer may still read them.
   private val prefetchExecutor: ExecutorService = if (prefetchEnabled) {
     Executors.newSingleThreadExecutor(r => {
       val t = new Thread(r, "arrow-cache-prefetch")
@@ -1310,7 +1318,7 @@ private class ArrowCachedBatchToColumnarBatchIterator(
   }
   private var prefetchFuture: Future[VectorSchemaRoot] = _
 
-  // Register cleanup - close remaining root and allocator when task completes
+  // Register cleanup - close the unreleased roots and the allocator when the task completes
   Option(TaskContext.get()).foreach { tc =>
     tc.addTaskCompletionListener[Unit] { _ =>
       // Stop the worker and close any root it already produced before closing the allocator.
@@ -1319,9 +1327,9 @@ private class ArrowCachedBatchToColumnarBatchIterator(
       // would then fail with "Memory was leaked by query".
       prefetchFuture = ArrowCachedBatchSerializer.drainAndClosePrefetch(
         prefetchExecutor, prefetchFuture)
-      if (previousRoot != null) {
-        previousRoot.close()
-        previousRoot = null
+      unreleasedRoots.synchronized {
+        unreleasedRoots.forEach(_.close())
+        unreleasedRoots.clear()
       }
       allocator.close()
     }
@@ -1330,12 +1338,6 @@ private class ArrowCachedBatchToColumnarBatchIterator(
   override def hasNext: Boolean = prefetchFuture != null || batchIter.hasNext
 
   override def next(): ColumnarBatch = {
-    // Close the previous root since the consumer has moved on from the batch it backed.
-    if (previousRoot != null) {
-      previousRoot.close()
-      previousRoot = null
-    }
-
     val root = if (prefetchFuture != null) {
       val r = try {
         prefetchFuture.get()
@@ -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))
+      val release = new RootRelease(root, selectedVectors.length)
+      new ColumnarBatch(
+        selectedVectors.map(new ReleasingArrowColumnVector(_, release): ColumnVector), numRows)
     }
-    val batch = new ColumnarBatch(selectedColumns, root.getRowCount)
 
     // Start prefetching the next batch while this one is being consumed.
     submitPrefetch()
@@ -1400,6 +1407,42 @@ private class ArrowCachedBatchToColumnarBatchIterator(
     }
   }
 
+  /** Closes `root` once all `numColumns` columns of the batch over it have been closed. */
+  private class RootRelease(root: VectorSchemaRoot, numColumns: Int) {
+    private var numOpen = numColumns
+
+    /**
+     * Closes a column of the batch with `closeColumn`, and the root after the last one. A no-op
+     * once the root is closed, e.g. by the task completion listener.
+     */
+    def closeColumn(closeColumn: () => Unit): Unit = unreleasedRoots.synchronized {
+      if (unreleasedRoots.contains(root)) {
+        closeColumn()
+        numOpen -= 1
+        if (numOpen == 0) {
+          unreleasedRoots.remove(root)
+          root.close()
+        }
+      }
+    }
+  }
+
+  /** A column of a returned batch, which releases the batch's root once all are closed. */
+  private class ReleasingArrowColumnVector(
+      vector: org.apache.arrow.vector.FieldVector,
+      release: RootRelease) extends ArrowColumnVector(vector) {
+    private var closed = false
+
+    private def closeVector(): Unit = super.close()
+
+    override def close(): Unit = unreleasedRoots.synchronized {
+      if (!closed) {
+        closed = true
+        release.closeColumn(() => closeVector())
+      }
+    }
+  }
+
   /** Submit deserialization of the next batch to the background thread, if prefetch is enabled. */
   private def submitPrefetch(): Unit = {
     if (prefetchEnabled && batchIter.hasNext) {

sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowEvalPythonExec.scala

@@ -25,7 +25,7 @@ import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.errors.QueryExecutionErrors
-import org.apache.spark.sql.execution.SparkPlan
+import org.apache.spark.sql.execution.{ColumnarToRowIterator, SparkPlan}
 import org.apache.spark.sql.execution.metric.SQLMetric
 import org.apache.spark.sql.execution.python.EvalPythonExec.ArgumentMetadata
 import org.apache.spark.sql.types.{StructType, UserDefinedType}
@@ -110,10 +110,11 @@ case class ArrowEvalPythonExec(
       // Columnar path: delegate to doExecuteColumnar, flatten to
       // UnsafeRow. ColumnarBatchRow from rowIterator() is NOT
       // UnsafeRow, and downstream operators (e.g., outer
-      // EvalPythonExec) may cast to UnsafeRow.
+      // EvalPythonExec) may cast to UnsafeRow. Like ColumnarToRowExec,
+      // this releases each batch once its rows are consumed: the
+      // batches of the Arrow columnar path carry the input columns.
       doExecuteColumnar().mapPartitionsInternal { batchIter =>
-        val toUnsafe = UnsafeProjection.create(schema)
-        batchIter.flatMap(_.rowIterator().asScala.map(toUnsafe))
+        new ColumnarToRowIterator(batchIter, UnsafeProjection.create(schema))
       }
     } else {
       // Row-based path: unchanged.