[SPARK-60294][SQL] Disable Arrow-backed columnar input to Python UDFs by default

docs/sql-migration-guide.md

@@ -28,6 +28,7 @@ license: |
 
 ## Upgrading from Spark SQL 4.3 to 4.4
 
+- 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, 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.
 - Since Spark 4.4, when `spark.sql.preserveCharVarcharTypeInfo` is true and `spark.sql.charVarchar.standardSemantics.enabled` is false, ORC reads that apply a CHAR/VARCHAR schema over STRING storage return the stored values without ORC truncation, matching Parquet. Previously the ORC reader requested `char(n)`/`varchar(n)` and truncated STRING-stored values to `n`. Read-side length checks (`EXCEED_LIMIT_LENGTH`) apply only when `spark.sql.charVarchar.standardSemantics.enabled` is true.
 - Since Spark 4.4, when `spark.sql.charVarchar.standardSemantics.enabled` is true, XML names used as `MAP<CHAR(n), _>` or `MAP<VARCHAR(n), _>` keys in `from_xml` and the XML datasource are length-checked without padding or trimming. The checked names are element names, `attributePrefix` plus attribute names (default prefix `_`), and the `valueTag` (default `_VALUE`) when mixed text is present. A `CHAR(n)` key must already be exactly `n` characters, and a `VARCHAR(n)` key must already be at most `n` characters. Mismatched keys fail with `UNSUPPORTED_XML_CHAR_VARCHAR_MAP_KEY` (SQLSTATE `0A000`) and follow the XML parse mode (`PERMISSIVE` or `FAILFAST`). Exact repeated names last-win; `spark.sql.mapKeyDedupPolicy` is not applied. XML CHAR/VARCHAR values use the same pad and `EXCEED_LIMIT_LENGTH` checks as other parsed text. Empty, attribute-only, and text-only map elements stay SQL NULL or a malformed record; they do not become an empty map or a `valueTag` entry. In Spark 4.3, `from_xml` and XML reader schemas rejected CHAR/VARCHAR with `UNSUPPORTED_CHAR_OR_VARCHAR_AS_STRING`, or replaced them with STRING when `spark.sql.legacy.charVarcharAsString` was true, so CHAR/VARCHAR map keys did not reach the parser. When the standard-semantics flag is false and `spark.sql.preserveCharVarcharTypeInfo` is true, short CHAR keys are padded and over-length CHAR or VARCHAR keys raise `EXCEED_LIMIT_LENGTH`.
@@ -80,6 +81,7 @@ license: |
 
 ## Upgrading from Spark SQL 4.2.0 to 4.2.1
 
+- Since Spark 4.2.1, `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 in Spark 4.1. 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`.
 - Spark 4.2.0 contains a correctness bug in the Data Source V2 `PartitionPredicate`. For scans, `PartitionPredicate` is the additional predicate type Spark sends during the second pass of iterative pushdown. When a data source opts into iterative pushdown by returning `true` from `supportsIterativePushdown()`, Spark may call its pushdown API again with a `PartitionPredicate` that evaluates a filter against the full partition key described by `Table.partitioning()`. Spark also uses `PartitionPredicate` separately for metadata-only `DELETE`. On either path, once the data source accepts the predicate, there is no later evaluator for the condition. In Spark 4.2.0, `PartitionPredicate.eval(InternalRow)` incorrectly returns `true` when evaluation fails, for example when an ANSI cast of a partition value fails, or when the partition key width does not match `Table.partitioning()`. This can produce incorrect results in two ways:
   - For a `SELECT` over partitions `dep = 'hr'` and `dep = '1'` with the filter `CAST(dep AS INT) = 1`, evaluating the predicate for `dep = 'hr'` fails. Returning `true` keeps that partition, and because Spark removes the post-scan filter, the query incorrectly returns its rows.
   - For a metadata-only `DELETE` with the same condition, returning `true` can delete the `dep = 'hr'` partition even though the condition never matches it.

python/pyspark/sql/tests/arrow/test_arrow_python_udf_cached.py

@@ -64,11 +64,14 @@ def conf(cls):
     def setUpClass(cls):
         super().setUpClass()
         cls.spark.conf.set("spark.sql.execution.pythonUDF.arrow.enabled", "true")
+        # The columnar input path, which this suite tests, is disabled by default.
+        cls.spark.conf.set("spark.sql.execution.arrow.pythonUDF.columnarInput.enabled", "true")
 
     @classmethod
     def tearDownClass(cls):
         try:
             cls.spark.conf.unset("spark.sql.execution.pythonUDF.arrow.enabled")
+            cls.spark.conf.unset("spark.sql.execution.arrow.pythonUDF.columnarInput.enabled")
         finally:
             super().tearDownClass()
 

sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala

@@ -5514,11 +5514,14 @@ object SQLConf {
         "Arrow-backed ColumnarBatch (e.g., DataSource V2 connectors), " +
         "bypassing the ColumnarToRow and ArrowWriter conversion. " +
         "This optimization reduces data transfer overhead between " +
-        "the JVM and Python worker processes.")
+        "the JVM and Python worker processes. It keeps the columns that the UDFs pass " +
+        "through until the results of their batch come back from Python, so it requires " +
+        "an upstream operator whose batches stay valid until then, which a columnar " +
+        "reader does not promise.")
       .version("4.2.0")
       .withBindingPolicy(ConfigBindingPolicy.SESSION)
       .booleanConf
-      .createWithDefault(true)
+      .createWithDefault(false)
 
   val ARROW_CACHE_PREFETCH_ENABLED =
     buildConf("spark.sql.execution.arrow.cache.prefetch.enabled")

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

@@ -106,7 +106,10 @@ case class ArrowEvalPythonExec(
   override def supportsRowBased: Boolean = true
 
   override protected def doExecute(): RDD[InternalRow] = {
-    if (child.supportsColumnar) {
+    // 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

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

@@ -45,9 +45,12 @@ import org.apache.spark.util.Utils
  * 1. '''Arrow columnar path''' (UDF inputs are simple column refs AND
  *    columns are [[ArrowColumnVector]]): Arrow FieldVectors are extracted
  *    directly and serialized to IPC. Pass-through columns are kept as
- *    [[ColumnVector]] references (safe because Arrow vectors are
- *    independently allocated per batch). Output is produced by columnar
- *    combining: passThruCols ++ resultCols -> ColumnarBatch.
+ *    [[ColumnVector]] references, which is only correct if the batches stay
+ *    valid until their results come back from Python: a columnar reader may
+ *    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.
  *
  * 2. '''Non-Arrow columnar path''' (UDF inputs are simple column refs
  *    BUT columns are NOT [[ArrowColumnVector]]): Non-Arrow columnar
@@ -175,9 +178,10 @@ private[python] class ColumnarArrowEvalPythonEvaluatorFactory(
     }
 
     /**
-     * Path 1: Arrow columnar. ColumnVector references are safe
-     * because Arrow vectors are independently allocated per batch.
-     * Combines pass-through + UDF result columns into ColumnarBatch.
+     * Path 1: Arrow columnar. Keeps references to the pass-through
+     * ColumnVectors, which requires the input batches to stay valid until
+     * their results come back from Python, and combines them with the UDF
+     * result columns into ColumnarBatch.
      */
     private def evalArrowColumnar(
         inputIter: Iterator[ColumnarBatch],

sql/core/src/test/scala/org/apache/spark/sql/execution/python/ArrowBackedDataSourceV2.scala

@@ -16,8 +16,11 @@
  */
 package org.apache.spark.sql.execution.python
 
+import java.nio.charset.StandardCharsets.UTF_8
 import java.util
 
+import scala.jdk.CollectionConverters._
+
 import org.apache.arrow.memory.BufferAllocator
 import org.apache.arrow.vector._
 
@@ -119,9 +122,13 @@ private class ArrowBackedScanBuilder(options: CaseInsensitiveStringMap)
 
   private val columnar =
     options.getBoolean("columnar", true)
+  private val reuseVectors =
+    options.getBoolean("reuseVectors", false)
+  private val batchSize =
+    options.getInt("batchSize", ArrowBackedDataSourceV2.BATCH_SIZE)
 
   override def createReaderFactory(): PartitionReaderFactory = {
-    new ArrowBackedReaderFactory(columnar)
+    new ArrowBackedReaderFactory(columnar, reuseVectors, batchSize)
   }
 }
 
@@ -137,16 +144,19 @@ private case class ArrowPartition(start: Int, end: Int)
  *   This allows benchmarking the same data source with and without
  *   the columnar optimization.
  */
-private class ArrowBackedReaderFactory(columnar: Boolean)
-    extends PartitionReaderFactory {
+private class ArrowBackedReaderFactory(
+    columnar: Boolean,
+    reuseVectors: Boolean,
+    batchSize: Int)
+  extends PartitionReaderFactory {
 
   override def supportColumnarReads(
       partition: InputPartition): Boolean = columnar
 
   override def createReader(
       partition: InputPartition): PartitionReader[InternalRow] = {
     val ArrowPartition(start, end) = partition
-    val inner = new ArrowBackedPartitionReader(start, end)
+    val inner = new ArrowBackedPartitionReader(start, end, batchSize)
     // Wrap columnar reader as row reader.
     new PartitionReader[InternalRow] {
       private var rowIter: java.util.Iterator[InternalRow] = _
@@ -164,7 +174,58 @@ private class ArrowBackedReaderFactory(columnar: Boolean)
   override def createColumnarReader(
       partition: InputPartition): PartitionReader[ColumnarBatch] = {
     val ArrowPartition(start, end) = partition
-    new ArrowBackedPartitionReader(start, end)
+    if (reuseVectors) {
+      new ReusingArrowPartitionReader(start, end, batchSize)
+    } else {
+      new ArrowBackedPartitionReader(start, end, batchSize)
+    }
+  }
+}
+
+/**
+ * Reads the same rows as [[ArrowBackedPartitionReader]], with short values in `data`, but refills
+ * the same vectors for every batch and closes them once the reader is closed, as columnar readers
+ * may: a batch is only valid until the next one.
+ */
+private class ReusingArrowPartitionReader(start: Int, end: Int, batchSize: Int)
+    extends PartitionReader[ColumnarBatch] {
+  private var current = start
+  private val allocator = ArrowUtils.rootAllocator.newChildAllocator(
+    s"reusing-arrow-test-reader-$start", 0, Long.MaxValue)
+  private val root = VectorSchemaRoot.create(
+    ArrowUtils.toArrowSchema(ArrowBackedDataSourceV2.SCHEMA, null,
+      errorOnDuplicatedFieldNames = true, largeVarTypes = false),
+    allocator)
+  private val batch = new ColumnarBatch(
+    root.getFieldVectors.asScala.map(new ArrowColumnVector(_): ColumnVector).toArray)
+
+  override def next(): Boolean = {
+    if (current >= end) return false
+    val count = math.min(batchSize, end - current)
+    val id = root.getVector("id").asInstanceOf[IntVector]
+    val name = root.getVector("name").asInstanceOf[VarCharVector]
+    val value = root.getVector("value").asInstanceOf[Float8Vector]
+    val data = root.getVector("data").asInstanceOf[VarCharVector]
+    // Overwrites the values of the previous batch in place.
+    root.getFieldVectors.forEach(_.reset())
+    (0 until count).foreach { i =>
+      val row = current + i
+      id.setSafe(i, row)
+      name.setSafe(i, s"row_$row".getBytes(UTF_8))
+      value.setSafe(i, row * 0.1)
+      data.setSafe(i, s"data_$row".getBytes(UTF_8))
+    }
+    root.setRowCount(count)
+    batch.setNumRows(count)
+    current += count
+    true
+  }
+
+  override def get(): ColumnarBatch = batch
+
+  override def close(): Unit = {
+    root.close()
+    allocator.close()
   }
 }
 
@@ -176,10 +237,9 @@ private class ArrowBackedReaderFactory(columnar: Boolean)
  * - name (STRING): "row_<id>"
  * - value (DOUBLE): id * 0.1
  */
-private class ArrowBackedPartitionReader(start: Int, end: Int)
+private class ArrowBackedPartitionReader(start: Int, end: Int, batchSize: Int)
     extends PartitionReader[ColumnarBatch] {
 
-  private val batchSize = ArrowBackedDataSourceV2.BATCH_SIZE
   private var current = start
   private var currentBatch: ColumnarBatch = _
 

sql/core/src/test/scala/org/apache/spark/sql/execution/python/ArrowCachedPythonUDFSuite.scala

@@ -0,0 +1,85 @@
+/*
+ * 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.python
+
+import org.apache.spark.SparkConf
+import org.apache.spark.sql.IntegratedUDFTestUtils
+import org.apache.spark.sql.execution.columnar.{ArrowCachedBatchSerializer, InMemoryRelation}
+import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf}
+import org.apache.spark.sql.test.SharedSparkSession
+import org.apache.spark.sql.types.StringType
+
+/**
+ * Tests Arrow Python UDFs over the Arrow cache, whose scan returns rows or Arrow-backed batches,
+ * and closes the root of the previous batch when it returns the next one.
+ */
+class ArrowCachedPythonUDFSuite extends SharedSparkSession {
+
+  import IntegratedUDFTestUtils._
+
+  override protected def sparkConf: SparkConf = super.sparkConf
+    .set(StaticSQLConf.SPARK_CACHE_SERIALIZER.key, classOf[ArrowCachedBatchSerializer].getName)
+
+  // The serializer of the cache is kept for the JVM once it is used.
+  override def beforeAll(): Unit = {
+    InMemoryRelation.clearSerializer()
+    super.beforeAll()
+  }
+
+  override def afterAll(): Unit = {
+    try {
+      super.afterAll()
+    } finally {
+      InMemoryRelation.clearSerializer()
+    }
+  }
+
+  test("Arrow Python UDFs over the Arrow cache with the default configuration") {
+    assume(shouldTestPandasUDFs)
+    // The scan of the Arrow cache also returns rows, so no ColumnarToRowExec is added below
+    // ArrowEvalPythonExec, whose child still supports columnar: it must still read rows when
+    // columnar input is disabled, as it is by default.
+    withSQLConf(SQLConf.ARROW_EXECUTION_MAX_RECORDS_PER_BATCH.key -> "1000") {
+      val pandasUDF = TestScalarPandasUDF(
+        name = "cached_str_udf", returnType = Some(StringType))
+      registerTestUDF(pandasUDF, spark)
+
+      val cached = spark.range(0, 200000, 1, 1).selectExpr("id", "CAST(id AS STRING) AS s")
+        .cache()
+      try {
+        // The test UDF casts its input to string, which is removed on a string column.
+        val result = cached.selectExpr("id", "s", "cached_str_udf(s) AS u")
+        val arrowExec = result.queryExecution.executedPlan.collect {
+          case exec: ArrowEvalPythonExec => exec
+        }
+        assert(arrowExec.nonEmpty && arrowExec.head.child.supportsColumnar &&
+          !arrowExec.head.supportsColumnar)
+
+        val rows = result.collect()
+        assert(rows.map(_.getLong(0)).toSeq == (0L until 200000L))
+        rows.foreach { row =>
+          val id = row.getLong(0)
+          assert(row.getString(1) == id.toString, s"pass-through mismatch for id $id")
+          assert(row.getString(2) == id.toString, s"UDF result mismatch for id $id")
+        }
+      } finally {
+        cached.unpersist()
+      }
+    }
+  }
+}