[SPARK-60306][SQL] Release consumed columnar batches in non-codegen ColumnarToRow and Arrow Python UDFs

docs/sql-migration-guide.md

@@ -30,6 +30,7 @@ license: |
 
 - Since Spark 4.4, quoted Variant path keys decode `\\`, `\/`, `\b`, `\f`, `\n`, `\r`, `\t`, `\uXXXX` (including valid surrogate pairs), and escapes of the enclosing quote (`\'` or `\"`). Previously, backslashes were literal, so existing paths can match different keys or become invalid. For example, use `$['a\\']` instead of `$['a\']`. Unknown escapes and unpaired surrogate escapes stay literal. Dot-notation keys are unchanged. To keep the old behavior, set `spark.sql.legacy.variantPathBackslashAsLiteral` to `true` before creating the query.
 - Since Spark 4.4, `spark.sql.execution.arrow.pythonUDF.columnarInput.enabled` is `false` by default: Arrow-based Python UDFs read rows from an Arrow-backed upstream operator, such as a Data Source V2 scan or the Arrow cache, as before Spark 4.2. With columnar input, they kept the columns that they pass through until the results of their batch came back from Python, and failed or returned wrong results when the operator reused or released the vectors of the batch in the meantime. To enable columnar input again, set `spark.sql.execution.arrow.pythonUDF.columnarInput.enabled` to `true`.
+- Since Spark 4.4, the conversion of columnar batches to rows releases each batch with `ColumnVector.closeIfFreeable` once its rows are consumed also when it is not whole-stage code-generated, e.g. when the query has more than `spark.sql.codegen.maxFields` fields or `spark.sql.codegen.wholeStage` is `false`, as the generated code has done since Spark 3.5.4 and 4.0. A `ColumnVector` that a Data Source V2 source reuses across batches, such as an `ArrowColumnVector` over a `VectorSchemaRoot` that is refilled for every batch, must override `closeIfFreeable` to do nothing.
 - Since Spark 4.4, `spark.sql.charVarchar.standardSemantics.enabled` (default `false`) is a preview feature that may change in future releases. When true, it enables first-class `CHAR(n)` and `VARCHAR(n)` semantics throughout schemas and query plans: `CHAR` -> `VARCHAR` -> `STRING` form a type-precedence family, explicit casts can return constrained types, least-common-type resolution for `COALESCE`/`CASE`/`UNION` can return `CHAR` or `VARCHAR`, and string expressions that transform their input return `STRING`. Assignment and reads enforce the declared length, padding `CHAR` to `n` characters and raising `EXCEED_LIMIT_LENGTH` when non-space content does not fit. Comparing a CHAR column with a string literal no longer pads the literal, so `WHERE char_col = 'ab'` does not match a stored `'ab  '`. The default remains the Spark 3.1 table-schema overlay on `STRING`. This opt-in mode is preferred over the experimental `spark.sql.preserveCharVarcharTypeInfo`, which instead keeps constrained types through transforming expressions. See [Character String Types](sql-ref-character-string-types.html) for the preview semantics; the entries below note area-specific behavior that changed. Later CHAR/VARCHAR bullets that do not name `preserveCharVarcharTypeInfo` apply when the preview flag is true, except `from_avro` value assignment, which also runs when only `preserveCharVarcharTypeInfo` is true.
 - Since Spark 4.4, when `spark.sql.charVarchar.standardSemantics.enabled` is true, `from_json` length-checks JSON object names used as `MAP<CHAR(n), _>` or `MAP<VARCHAR(n), _>` keys without padding or trimming them. A `CHAR(n)` key must already be exactly `n` characters, and a `VARCHAR(n)` key must already be at most `n` characters; duplicate names are kept, exactly as for STRING keys, and `spark.sql.mapKeyDedupPolicy` is not applied. A key that fails the check makes the enclosing map a bad record: in `PERMISSIVE` mode the nearest enclosing struct field becomes `null` (its sibling fields are preserved), and if the map is inside an array or another map, that whole array or map becomes `null`; in `FAILFAST` mode parsing fails with `UNSUPPORTED_JSON_CHAR_VARCHAR_MAP_KEY` (SQLSTATE `0A000`) as the cause of `MALFORMED_RECORD_IN_PARSING`. The same parser-level check runs for any JSON scan whose reader schema carries a CHAR/VARCHAR map key, including catalog JSON tables, since under standard semantics those keys stay CHAR/VARCHAR all the way to the scan: an off-width key becomes a bad record there too, rather than being padded or raising `EXCEED_LIMIT_LENGTH`. Unlike `from_json`, keys that pass the check are then rebuilt read-side, so on the scan path `spark.sql.mapKeyDedupPolicy` does apply and a duplicate name can raise `DUPLICATED_MAP_KEY`. JSON file scans, including scans inside persisted views, read the flag from the executing session. The optimizer no longer folds `from_json(to_json(m))` to `m` when that schema has a CHAR/VARCHAR map key, so the length check still runs on the round trip.
 - Since Spark 4.4, `schema_of_json`, `schema_of_csv`, and `schema_of_xml` accept CHAR/VARCHAR document arguments when CHAR/VARCHAR are first-class. In Spark 4.3, with `spark.sql.preserveCharVarcharTypeInfo` on, they rejected those arguments because the child type was not STRING.

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

@@ -0,0 +1,45 @@
+/*
+ * 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()}. This covers the vector itself: the vectors that {@link #getChild(int)}
+ * returns for the fields of a struct are plain {@link ArrowColumnVector}s, which consumers do not
+ * release on their own.
+ */
+public class ReusedArrowColumnVector extends ArrowColumnVector {
+
+  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/scala/org/apache/spark/sql/execution/ColumnarEvaluatorFactory.scala

@@ -27,6 +27,7 @@ import org.apache.spark.sql.execution.vectorized.{OffHeapColumnVector, OnHeapCol
 import org.apache.spark.sql.types.StructType
 import org.apache.spark.sql.vectorized.ColumnarBatch
 import org.apache.spark.util.ArrayImplicits._
+import org.apache.spark.util.CompletionIterator
 
 class ColumnarToRowEvaluatorFactory(
     childOutput: Seq[Attribute],
@@ -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]](
+          input.rowIterator().asScala.map(toUnsafe), input.closeIfFreeable())
       }
     }
   }

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
+  // batch once its rows are consumed. Otherwise it reads rows from its child,
+  // even from one that also supports columnar, such as the scan of the Arrow
+  // cache. Either way, the choice is fixed when the plan is built.
   override def supportsColumnar: Boolean =
     child.supportsColumnar && conf.arrowPySparkUDFColumnarInputEnabled
-  override def supportsRowBased: Boolean = true
-
-  override protected def doExecute(): RDD[InternalRow] = {
-    // Checks the configuration too, through `supportsColumnar`: a child that also returns rows,
-    // such as the scan of the Arrow cache, gets no ColumnarToRowExec when columnar input is
-    // disabled, and still supports columnar.
-    if (supportsColumnar) {
-      // 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.
-      doExecuteColumnar().mapPartitionsInternal { batchIter =>
-        val toUnsafe = UnsafeProjection.create(schema)
-        batchIter.flatMap(_.rowIterator().asScala.map(toUnsafe))
-      }
-    } else {
-      // Row-based path: unchanged.
-      super.doExecute()
-    }
-  }
 
   override protected def doExecuteColumnar(): RDD[ColumnarBatch] = {
     val inputRDD = child.executeColumnar()

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

@@ -50,7 +50,12 @@ import org.apache.spark.util.Utils
  *    reuse or release the vectors of a batch once it is asked for the next
  *    one. So `spark.sql.execution.arrow.pythonUDF.columnarInput.enabled`
  *    is disabled by default. Output is produced by columnar combining:
- *    passThruCols ++ resultCols -> ColumnarBatch.
+ *    passThruCols ++ resultCols -> ColumnarBatch, so the consumer of the
+ *    output releases the pass-through columns with closeIfFreeable on the
+ *    task thread. 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.
  *
  * 2. '''Non-Arrow columnar path''' (UDF inputs are simple column refs
  *    BUT columns are NOT [[ArrowColumnVector]]): Non-Arrow columnar
@@ -63,6 +68,9 @@ import org.apache.spark.util.Utils
  *    [[BasicPythonArrowInput]]. Uses [[HybridRowQueue]] for pass-through
  *    and converts joined rows back to [[ColumnarBatch]].
  *
+ * Paths 2 and 3 copy what they need from an input batch, so they release
+ * it with closeIfFreeable once it has been written to Python.
+ *
  * TODO: Add a physical plan rule that inserts a ProjectExec before
  *   ArrowEvalPythonExec to pre-evaluate complex UDF input expressions
  *   into simple column references. This would eliminate path 3.
@@ -255,12 +263,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 =>
           batch.rowIterator().asScala.foreach { row =>
             queue.add(unsafeProj(row).copy())
           }
-          batch
-        }
+          true
+        })
         val pyRunner =
           new ColumnarArrowPythonWithNamedArgumentRunner(
             pyFuncs, evalType, argMetas, udfInputSchema,
@@ -274,7 +282,9 @@ private[python] class ColumnarArrowEvalPythonEvaluatorFactory(
         val projection = MutableProjection.create(
           allInputs, childOutput)
         projection.initialize(context.partitionId())
-        val batchIter = inputIter.map { batch =>
+        // The rows of a batch are written to Python before the writer asks
+        // for the next batch, which releases it.
+        val batchIter = new WrittenBatchReleasingIterator(inputIter, _ => true).map { batch =>
           batch.rowIterator().asScala.map { row =>
             queue.add(unsafeProj(row).copy())
             projection(row)
@@ -343,6 +353,41 @@ private[python] class ColumnarArrowEvalPythonEvaluatorFactory(
   }
 }
 
+/**
+ * Hands the input batches to the Python writer. `enqueue` is called on each
+ * batch before the writer gets it, and returns whether the batch should be
+ * released with closeIfFreeable once it has been written: the writer is done
+ * with a batch when it asks for the next one or reaches the end.
+ */
+private[python] class WrittenBatchReleasingIterator(
+    input: Iterator[ColumnarBatch],
+    enqueue: ColumnarBatch => Boolean)
+  extends Iterator[ColumnarBatch] {
+
+  private var written: ColumnarBatch = _
+
+  private def releaseWritten(): Unit = {
+    if (written != null) {
+      written.closeIfFreeable()
+      written = null
+    }
+  }
+
+  override def hasNext: Boolean = {
+    releaseWritten()
+    input.hasNext
+  }
+
+  override def next(): ColumnarBatch = {
+    releaseWritten()
+    val batch = input.next()
+    if (enqueue(batch)) {
+      written = batch
+    }
+    batch
+  }
+}
+
 /** Iterator wrapper that allows peeking at the first element. */
 private[python] class PeekableIterator[T](underlying: Iterator[T])
     extends Iterator[T] {

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

@@ -27,10 +27,11 @@ import org.apache.arrow.vector.ipc.ArrowStreamReader
 import org.apache.spark.{SparkEnv, TaskContext}
 import org.apache.spark.api.python.{BasePythonRunner, PythonWorker, PythonWorkerHandle, SpecialLengths}
 import org.apache.spark.sql.execution.metric.SQLMetric
+import org.apache.spark.sql.execution.vectorized.ReusedArrowColumnVector
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types.StructType
 import org.apache.spark.sql.util.ArrowUtils
-import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch, ColumnVector}
+import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector}
 
 /**
  * A trait that can be mixed-in with [[BasePythonRunner]]. It implements the logic from
@@ -170,12 +171,18 @@ trait ArrowOutputProcessor {
   def close(): Unit
 }
 
+/**
+ * Loads every batch from Python into the reader's root and returns batches over the same vectors,
+ * so the vectors of a batch are only valid until the next one is loaded. They are
+ * [[ReusedArrowColumnVector]]s: a consumer does not hold them past the next batch, and its
+ * `closeIfFreeable` leaves them to this processor, which closes them in `close`.
+ */
 class ArrowOutputProcessorImpl(reader: ArrowStreamReader, pythonMetrics: Map[String, SQLMetric])
     extends ArrowOutputProcessor {
   protected val root = reader.getVectorSchemaRoot()
   protected val schema: StructType = ArrowUtils.fromArrowSchema(root.getSchema())
   private val vectors: Array[ColumnVector] = root.getFieldVectors().asScala.map { vector =>
-    new ArrowColumnVector(vector)
+    new ReusedArrowColumnVector(vector)
   }.toArray[ColumnVector]
 
   protected var rowCount = -1
@@ -235,12 +242,15 @@ abstract class BaseSliceArrowOutputProcessor(
     batch
   }
 
+  // The vectors are either those of the reader's root, into which the next batch is loaded, or
+  // those of a slice of it, which is closed once the next one is produced, so both are reused
+  // vectors.
   protected override def getVectors(root: VectorSchemaRoot): Array[ColumnVector] = {
     if (root eq this.root) {
       super.getVectors(root)
     } else {
       root.getFieldVectors.asScala.map { vector =>
-        new ArrowColumnVector(vector)
+        new ReusedArrowColumnVector(vector)
       }.toArray[ColumnVector]
     }
   }