This is an automated email from the ASF dual-hosted git repository.
viirya pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new d33352c5e40a [SPARK-56642][SQL] Add pipelined JVM-Python UDF data
transfer
d33352c5e40a is described below
commit d33352c5e40a9358636cbf7d5cef040edaf10ac8
Author: Liang-Chi Hsieh <[email protected]>
AuthorDate: Thu Jun 25 14:38:36 2026 -0700
[SPARK-56642][SQL] Add pipelined JVM-Python UDF data transfer
### What changes were proposed in this pull request?
Add an opt-in pipelined execution mode for Python UDFs
(spark.python.udf.pipelined.enabled). When enabled, a dedicated writer thread
serializes input and writes directly to the Python worker socket in blocking
mode, while the task main thread reads output from the same socket
concurrently. Socket full-duplex allows both directions to overlap, achieving
pipeline parallelism.
Key changes:
- PipelinedWriterRunnable in PythonRunner.scala: serializes input +
writes to socket in a pool thread
- pipelined_process() in worker.py: background reader thread pre-fetches
input batches; lazy iterators from grouped/aggregate serializers are eagerly
materialized to avoid cross-thread socket conflicts
- InMemoryRowQueue gains a lockFree parameter: when true (pipelined mode
only), synchronized is skipped since blocking socket I/O already provides
happens-before guarantees
- Proper propagation of TaskContext and InputFileBlockHolder to the
writer pool thread
- Socket idle timeout handling via SO_TIMEOUT + SocketTimeoutException
wrapper
Benchmark: `python/pyspark/sql/tests/pandas/bench_pipelined_udf.py`
Environment: local[*] (16 cores), 5 iterations, 2 warmup
| Scenario | Sync (ms) | Pipelined (ms) | Speedup |
|------------------------------|-----------|----------------|-----------|
| Light UDF (1M rows) | 100 | 90 | 1.11x |
| CPU-bound UDF (1M rows) | 143 | 146 | 0.98x |
| Heavy UDF (10ms sleep/batch) | 1180 | 1176 | 1.00x |
| Large data (5M rows) | 338 | 300 | **1.13x** |
| Multi-UDF (3 columns, 1M) | 123 | 100 | **1.24x** |
ASV Benchmark: `python/benchmarks/bench_pipelined_udf.py`
**ScalarUDFTimeBench** (scalar UDF `x + 1`)
| pipelined | 100K rows | 1M rows |
|-----------|-----------|---------|
| False | 122±0ms | 200±0ms |
| True | 86.5±0ms | 164±0ms |
**LargeDataUDFTimeBench** (5M rows, scalar UDF `x + 1`)
| pipelined | time | peak memory |
|-----------|------|-------------|
| False | 526±0ms | 116M |
| True | 496±0ms | 110M |
**MultiUDFTimeBench** (3 UDF columns)
| pipelined | 100K rows | 1M rows |
|-----------|-----------|---------|
| False | 157±0ms | 305±0ms |
| True | 123±0ms | 269±0ms |
Memory usage is the same for both modes (~110M).
#### Relationship to SPARK-44705
SPARK-44705 made PythonRunner single-threaded by removing the original
WriterThread. That WriterThread shared a single blocking java.net.Socket with
the main reader thread, where Thread.interrupt() does not unblock
OutputStream.write. This led to two recurring bugs:
- SPARK-33277 — race between writer-thread access to off-heap rows and
the FileScan task-completion listener freeing those rows; segfault.
- SPARK-38677 — three-way deadlock (task → writer → python → task)
because interrupt() couldn't unblock the writer's blocking socket write.
SPARK-44705 mitigated both by collapsing to one thread + SocketChannel
selector and required a few monitor threads to keep the design tractable. This
PR is not a revert — the existing single-threaded selector path remains the
default and is unchanged. The new pipelined path is opt-in
(spark.python.udf.pipelined.enabled=false by default) and is structured so the
two original bugs do not return.
| | Pre-44705 `WriterThread` | This PR's pipelined path |
| --- | --- | --- |
| Socket API | `java.net.Socket.OutputStream` — `Thread.interrupt()` does
**not** unblock `write` | `java.nio.channels.SocketChannel` —
`Thread.interrupt()` closes the channel and throws `ClosedByInterruptException`
|
| Cancel + cleanup | `interrupt() + join()`; could deadlock (SPARK-38677),
needed a `WriterMonitorThread` to break it | `Future.cancel(true) +
Future.get()`; `cancel(true)` reliably unblocks the writer, `get()` then
provides the same SPARK-33277 ordering guarantee as `join()`, no monitor thread
required |
A regression test mirroring the SPARK-33277 reproducer
(`test_offheap_reader_with_head_does_not_segfault`) is included.
### Why are the changes needed?
The current single-threaded NIO selector model serializes input and reads
output in the same thread alternately. For multi-column UDFs or compute-heavy
UDFs, this leaves the JVM idle while waiting for Python output. Pipelined mode
overlaps serialization with output reading, improving throughput by up to 22%
for multi-UDF queries and 10% for large data workloads.
### Does this PR introduce _any_ user-facing change?
No
### How was this patch tested?
- New dedicated test suite test_pipelined_udf.py with 12 tests running
with spark.python.udf.pipelined.enabled=true via SparkConf, covering scalar
UDF, string UDF, multi-column UDF, chained UDF, null handling, UDAF, empty
partitions, multiple partitions, large data, batched UDF, and exception
propagation
- All PySpark UDF test suites pass with pipelined=true as the
compile-time default: test_pandas_udf_scalar, test_pandas_udf_grouped_agg,
test_pandas_udf_window, test_udf, test_udtf
- JVM test suites pass: PythonUDFSuite, BatchEvalPythonExecSuite,
ArrowColumnarPythonUDFSuite
- Benchmark (bench_pipelined_udf.py) across 5 scenarios confirms no
regression vs sync mode and up to 22% speedup for multi-UDF workloads
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code
Closes #55552 from viirya/pipelined-python-udf.
Authored-by: Liang-Chi Hsieh <[email protected]>
Signed-off-by: Liang-Chi Hsieh <[email protected]>
(cherry picked from commit 59d4fa29acc11361b374965e1f9bbc4cad691d31)
Signed-off-by: Liang-Chi Hsieh <[email protected]>
---
.../org/apache/spark/api/python/PythonRunner.scala | 257 ++++++++++++++++-
.../org/apache/spark/internal/config/Python.scala | 25 ++
dev/sparktestsupport/modules.py | 1 +
python/benchmarks/bench_pipelined_udf.py | 220 +++++++++++++++
python/pyspark/sql/pandas/serializers.py | 7 +-
.../sql/tests/pandas/bench_pipelined_udf.py | 305 +++++++++++++++++++++
.../pyspark/sql/tests/pandas/test_pipelined_udf.py | 285 +++++++++++++++++++
python/pyspark/worker.py | 101 ++++++-
.../python/ArrowAggregatePythonExec.scala | 7 +-
.../python/ArrowWindowPythonEvaluatorFactory.scala | 6 +-
.../ColumnarArrowEvalPythonEvaluatorFactory.scala | 7 +-
.../python/EvalPythonEvaluatorFactory.scala | 8 +-
.../sql/execution/python/EvalPythonUDTFExec.scala | 6 +-
.../spark/sql/execution/python/RowQueue.scala | 66 ++++-
14 files changed, 1274 insertions(+), 27 deletions(-)
diff --git a/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala
b/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala
index 896ab6e3b19d..ee15bd5e46ef 100644
--- a/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala
+++ b/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala
@@ -23,7 +23,7 @@ import java.nio.ByteBuffer
import java.nio.channels.{AsynchronousCloseException, Channels, SelectionKey,
ServerSocketChannel, SocketChannel}
import java.nio.file.{Files => JavaFiles, Path}
import java.util.UUID
-import java.util.concurrent.{ConcurrentHashMap, TimeUnit}
+import java.util.concurrent.{CancellationException, ConcurrentHashMap,
ExecutionException, TimeUnit}
import java.util.concurrent.atomic.AtomicBoolean
import scala.jdk.CollectionConverters._
@@ -121,6 +121,21 @@ private[spark] object PythonEvalType {
private[spark] object BasePythonRunner extends Logging {
+ /**
+ * Shared thread pool for pipelined writer tasks. Using a cached thread pool
ensures that
+ * writer threads are reused across tasks, which keeps JIT-compiled code,
branch prediction
+ * history, and CPU caches warm.
+ * Bounded by executor cores since each task uses at most one writer thread.
+ */
+ private[python] lazy val pipelinedWriterThreadPool = {
+ // Each concurrent task uses at most one writer thread. Bound the pool by
available
+ // processors, which is the natural upper limit for concurrent tasks on
this executor.
+ // Using availableProcessors() instead of EXECUTOR_CORES because the
latter defaults
+ // to 1 in local[*] mode even though multiple tasks run concurrently.
+ val maxThreads = Runtime.getRuntime.availableProcessors()
+ ThreadUtils.newDaemonCachedThreadPool("python-udf-pipelined-writer",
maxThreads)
+ }
+
private[spark] lazy val faultHandlerLogDir = Utils.createTempDir(namePrefix
= "faulthandler")
private[spark] def faultHandlerLogPath(pid: Int): Path = {
@@ -211,6 +226,8 @@ private[spark] abstract class BasePythonRunner[IN, OUT](
conf.get(PYTHON_WORKER_TRACEBACK_DUMP_INTERVAL_SECONDS)
protected val killWorkerOnFlushFailure: Boolean =
conf.get(PYTHON_DAEMON_KILL_WORKER_ON_FLUSH_FAILURE)
+ protected val pipelinedEnabled: Boolean =
conf.get(PYTHON_UDF_PIPELINED_EXECUTION)
+ protected val pipelinedQueueDepth: Int =
conf.get(PYTHON_UDF_PIPELINED_QUEUE_DEPTH)
protected val hideTraceback: Boolean = false
protected val simplifiedTraceback: Boolean = false
protected val tracebackWithLocals: Boolean = false
@@ -329,6 +346,12 @@ private[spark] abstract class BasePythonRunner[IN, OUT](
envVars.put("SPARK_JOB_ARTIFACT_UUID",
jobArtifactUUID.getOrElse("default"))
envVars.put("SPARK_PYTHON_RUNTIME", "PYTHON_WORKER")
+ // Pipelined mode is only for UDF eval types, not NON_UDF
(mapPartitions/RDD path).
+ val usePipelined = pipelinedEnabled && evalType != PythonEvalType.NON_UDF
+ if (usePipelined) {
+ envVars.put("SPARK_PIPELINED_UDF", "1")
+ envVars.put("SPARK_PIPELINED_UDF_QUEUE_DEPTH",
pipelinedQueueDepth.toString)
+ }
val (worker: PythonWorker, handle: Option[ProcessHandle]) =
env.createPythonWorker(
pythonExec, workerModule, daemonModule, envVars.asScala.toMap, useDaemon)
@@ -363,15 +386,128 @@ private[spark] abstract class BasePythonRunner[IN, OUT](
}
// Return an iterator that read lines from the process's stdout
- val dataIn = new DataInputStream(new BufferedInputStream(
- new ReaderInputStream(worker, writer, handle,
- faultHandlerEnabled, idleTimeoutSeconds, killOnIdleTimeout, context),
- bufferSize))
+ val dataIn = if (usePipelined) {
+ createPipelinedDataIn(worker, writer, handle, context)
+ } else {
+ new DataInputStream(new BufferedInputStream(
+ new ReaderInputStream(worker, writer, handle,
+ faultHandlerEnabled, idleTimeoutSeconds, killOnIdleTimeout, context),
+ bufferSize))
+ }
val stdoutIterator = newReaderIterator(
dataIn, writer, startTime, env, worker, handle.map(_.pid.toInt),
releasedOrClosed, context)
new InterruptibleIterator(context, stdoutIterator)
}
+ /**
+ * Sets up pipelined mode: switches the socket to blocking mode, starts the
writer
+ * thread, configures idle timeout, and returns a DataInputStream for
reading output.
+ */
+ private def createPipelinedDataIn(
+ worker: PythonWorker,
+ writer: Writer,
+ handle: Option[ProcessHandle],
+ context: TaskContext): DataInputStream = {
+ // Switch the channel to blocking mode for true full-duplex I/O.
+ // Must close the selector first because configureBlocking() fails
+ // if the channel is registered with a selector
(IllegalBlockingModeException).
+ if (worker.selectionKey != null) {
+ worker.selectionKey.cancel()
+ }
+ if (worker.selector != null) {
+ worker.selector.close()
+ }
+ worker.channel.configureBlocking(true)
+ worker.refresh() // re-initializes (no selector in blocking mode)
+
+ val writerRunnable = new PipelinedWriterRunnable(worker, writer,
bufferSize, context)
+ val writerFuture =
BasePythonRunner.pipelinedWriterThreadPool.submit(writerRunnable)
+
+ // Wait for the writer to actually exit before letting subsequent task
completion listeners
+ // run. Subsequent listeners (registered earlier, executed later under
LIFO) free off-heap
+ // memory backing the input rows; if the writer is still serializing such
a row when free
+ // happens, we get a use-after-free segfault (SPARK-33277). cancel(true)
is enough to unblock
+ // a writer stuck on channel.write (JDK closes the SocketChannel and throws
+ // ClosedByInterruptException on interrupt), so this get() returns
promptly in normal cases;
+ // the worst case is a bounded wait for the writer to finish serializing
the current row or
+ // batch and observe the interrupt flag at the top of its loop.
+ context.addTaskCompletionListener[Unit] { _ =>
+ writerFuture.cancel(true)
+ try {
+ writerFuture.get()
+ } catch {
+ case _: CancellationException | _: ExecutionException | _:
InterruptedException =>
+ // Expected: cancel(true) raced ahead, or writer exited via
_exception path.
+ }
+ }
+
+ // Set socket read timeout for idle timeout detection in pipelined mode.
+ // Always set explicitly (including 0 = no timeout) because reused workers
may
+ // retain a stale SO_TIMEOUT from a previous task that had a different
setting.
+ worker.channel.socket().setSoTimeout(
+ if (idleTimeoutSeconds > 0) idleTimeoutSeconds.toInt * 1000 else 0)
+
+ // Wrap the socket InputStream to handle idle timeout, matching sync mode
behavior:
+ // - Log warning on each timeout
+ // - If killOnIdleTimeout=true: kill worker, then throw
PythonWorkerException
+ // - If killOnIdleTimeout=false: log only, retry read (continue waiting)
+ val socketInput = new InputStream {
+ private val inner = worker.channel.socket().getInputStream
+ private var pythonWorkerKilled = false
+ override def read(): Int = doRead(() => inner.read())
+ override def read(b: Array[Byte], off: Int, len: Int): Int =
+ doRead(() => inner.read(b, off, len))
+ private def doRead(op: () => Int): Int = {
+ var result = 0
+ var retry = true
+ while (retry) {
+ try {
+ result = op()
+ retry = false
+ } catch {
+ case _: java.net.SocketTimeoutException =>
+ if (pythonWorkerKilled) {
+ logWarning(
+ log"Waiting for Python worker process to terminate after
idle timeout: " +
+ pythonWorkerStatusMessageWithContext(
+ handle, worker, hasInputs = true))
+ } else {
+ logWarning(
+ log"Idle timeout reached for Python worker (timeout: " +
+ log"${MDC(PYTHON_WORKER_IDLE_TIMEOUT, idleTimeoutSeconds)}
seconds). " +
+ log"No data received from the worker process - " +
+ pythonWorkerStatusMessageWithContext(
+ handle, worker, hasInputs = true) +
+ log" - ${MDC(TASK_NAME, taskIdentifier(context))}")
+ if (killOnIdleTimeout) {
+ handle.foreach { h =>
+ if (h.isAlive) {
+ logWarning(
+ log"Terminating Python worker process due to idle
timeout " +
+ log"(timeout: " +
+ log"${MDC(PYTHON_WORKER_IDLE_TIMEOUT,
idleTimeoutSeconds)} " +
+ log"seconds) - ${MDC(TASK_NAME,
taskIdentifier(context))}")
+ pythonWorkerKilled = h.destroy()
+ }
+ }
+ }
+ }
+ }
+ }
+ if (result == -1 && pythonWorkerKilled) {
+ val base = "Python worker process terminated due to idle timeout " +
+ s"(timeout: $idleTimeoutSeconds seconds)"
+ val msg = tryReadFaultHandlerLog(faultHandlerEnabled,
handle.map(_.pid.toInt))
+ .map(error => s"$base: $error")
+ .getOrElse(base)
+ throw new PythonWorkerException(msg)
+ }
+ result
+ }
+ }
+ new DataInputStream(new BufferedInputStream(socketInput, bufferSize))
+ }
+
protected def newWriter(
env: SparkEnv,
worker: PythonWorker,
@@ -408,6 +544,11 @@ private[spark] abstract class BasePythonRunner[IN, OUT](
/** Contains the throwable thrown while writing the parent iterator to the
Python process. */
def exception: Option[Throwable] = Option(_exception)
+ /** Records a throwable observed by an external collaborator (e.g. the
pipelined writer). */
+ private[python] def setException(t: Throwable): Unit = {
+ _exception = t
+ }
+
/**
* Writes a command section to the stream connected to the Python worker.
*/
@@ -1000,6 +1141,112 @@ private[spark] abstract class BasePythonRunner[IN, OUT](
}
}
+ /**
+ * A dedicated thread that serializes input data and writes it directly to
the Python worker
+ * socket in blocking mode. The task main thread simultaneously reads output
from the same
+ * socket. TCP sockets are full-duplex, so concurrent read() and write()
from different
+ * threads is safe -- they operate on independent OS-level buffers.
+ *
+ * This design achieves true pipeline parallelism without any inter-thread
queues or locks:
+ * Writer Thread: serialize batch N -> channel.write(batch N)
[blocking]
+ * Reader Thread: channel.read(output N-1)
[blocking]
+ * Python: read batch N-1 -> compute -> write output -> read
batch N
+ *
+ * Deadlock safety: Python's UDF loop is "read input -> compute -> write
output -> repeat".
+ * As long as the reader thread is consuming Python's output (freeing
Python's send buffer),
+ * Python will eventually consume input from the socket (freeing the JVM's
send buffer for
+ * the writer thread). The reader thread is always actively reading because
the task's
+ * downstream operators pull output on demand.
+ *
+ * Unlike the old WriterThread (removed in SPARK-44705), this design uses a
blocking socket
+ * in full-duplex mode rather than two threads competing on the same
blocking socket with
+ * shared mutable state. The old design's deadlocks were caused by complex
interactions
+ * with vectorized readers and monitor threads, not by the fundamental
read/write split.
+ */
+ class PipelinedWriterRunnable(
+ worker: PythonWorker,
+ writer: Writer,
+ bufferSize: Int,
+ context: TaskContext)
+ extends Runnable {
+
+ // Capture InputFileBlockHolder from the task thread so we can propagate it
+ // to the writer pool thread. This is needed because upstream scan
operators
+ // set InputFileBlockHolder via InheritableThreadLocal, but pool threads
+ // don't inherit from the task thread.
+ private val parentInputFileBlockHolder =
InputFileBlockHolder.getThreadLocalValue()
+
+ override def run(): Unit = {
+ // Propagate TaskContext and InputFileBlockHolder to the pool thread so
that
+ // upstream operators work correctly.
+ TaskContext.setTaskContext(context)
+ InputFileBlockHolder.setThreadLocalValue(parentInputFileBlockHolder)
+ val bufferStream = new DirectByteBufferOutputStream(bufferSize)
+ val dataOut = new DataOutputStream(bufferStream)
+ try {
+ // Write command/metadata (partition index, task context, broadcasts,
UDF definition).
+ writer.open(dataOut)
+ flushToSocket(bufferStream)
+
+ // Write input data in a loop, batching into buffers of ~bufferSize.
+ var hasInput = true
+ while (hasInput && !Thread.currentThread().isInterrupted) {
+ hasInput = writer.writeNextInputToStream(dataOut)
+ if (bufferStream.size() >= bufferSize || !hasInput) {
+ if (!hasInput) {
+ writer.close(dataOut)
+ }
+ flushToSocket(bufferStream)
+ }
+ }
+ } catch {
+ case _: InterruptedException =>
+ // Task cancelled via Future.cancel(true)
+ Thread.currentThread().interrupt()
+ case _: java.nio.channels.ClosedByInterruptException =>
+ // Task cancelled while blocked in channel.write(). The channel is
+ // automatically closed by the JVM, which will cause Python to
receive
+ // EOF and the reader thread to get IOException.
+ Thread.currentThread().interrupt()
+ case NonFatal(t) =>
+ // InterruptedException and ClosedByInterruptException are matched
above; what
+ // remains here is genuine non-fatal failure that needs to be
propagated to the
+ // reader through writer.exception + a socket EOF.
+ writer.setException(t)
+ // Shut down the socket output so Python receives EOF and terminates.
+ // This unblocks the reader thread which is waiting on socket input:
+ // Python will exit, closing its end of the socket, causing the
reader's
+ // read() to return -1. The ReaderIterator will then check
writer.exception
+ // and propagate the failure.
+ if (worker.channel.isConnected) {
+ Utils.tryLog(worker.channel.shutdownOutput())
+ }
+ } finally {
+ TaskContext.unset()
+ InputFileBlockHolder.unset()
+ try {
+ bufferStream.close()
+ } catch {
+ case _: Exception => // ignore
+ }
+ }
+ }
+
+ /**
+ * Writes all buffered data to the socket and resets the buffer for reuse.
+ * Uses the DirectByteBufferOutputStream's direct buffer view for zero-copy
+ * socket writes. The write() call is blocking -- it will wait until the OS
+ * socket send buffer has room, which provides natural backpressure.
+ */
+ private def flushToSocket(bufferStream: DirectByteBufferOutputStream):
Unit = {
+ val buf = bufferStream.toByteBuffer
+ while (buf.hasRemaining) {
+ worker.channel.write(buf) // blocking write
+ }
+ bufferStream.reset()
+ }
+ }
+
}
private[spark] object PythonRunner {
diff --git a/core/src/main/scala/org/apache/spark/internal/config/Python.scala
b/core/src/main/scala/org/apache/spark/internal/config/Python.scala
index dc16d1ff255d..1f5829226d88 100644
--- a/core/src/main/scala/org/apache/spark/internal/config/Python.scala
+++ b/core/src/main/scala/org/apache/spark/internal/config/Python.scala
@@ -150,4 +150,29 @@ private[spark] object Python {
.version("4.1.0")
.booleanConf
.createWithDefault(true)
+
+ val PYTHON_UDF_PIPELINED_EXECUTION =
+ ConfigBuilder("spark.python.udf.pipelined.enabled")
+ .doc("When true, enables pipelined (asynchronous) data transfer between
JVM and Python " +
+ "UDF workers. In pipelined mode, input serialization runs in a
separate writer thread " +
+ "while the main task thread reads output from the Python worker,
allowing the two " +
+ "directions to overlap. This can improve throughput for some workloads
" +
+ "(e.g., multi-column UDFs or compute-heavy UDFs like ML inference);
for light, " +
+ "single-column UDFs the overhead of the extra thread may offset the
gain.")
+ .version("4.3.0")
+ .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE)
+ .booleanConf
+ .createWithDefault(false)
+
+ val PYTHON_UDF_PIPELINED_QUEUE_DEPTH =
+ ConfigBuilder("spark.python.udf.pipelined.queueDepth")
+ .doc("The maximum number of input batches the Python worker's background
reader thread " +
+ "can pre-fetch ahead of UDF computation. A higher value allows more
overlap between " +
+ "input reading and UDF processing, at the cost of increased memory
usage. " +
+ "Only effective when spark.python.udf.pipelined.enabled is true.")
+ .version("4.3.0")
+ .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE)
+ .intConf
+ .checkValue(_ > 0, "Queue depth must be positive.")
+ .createWithDefault(2)
}
diff --git a/dev/sparktestsupport/modules.py b/dev/sparktestsupport/modules.py
index c474b87c8b36..860158b941f6 100644
--- a/dev/sparktestsupport/modules.py
+++ b/dev/sparktestsupport/modules.py
@@ -605,6 +605,7 @@ pyspark_sql = Module(
"pyspark.sql.tests.pandas.test_pandas_udf",
"pyspark.sql.tests.pandas.test_pandas_udf_grouped_agg",
"pyspark.sql.tests.pandas.test_pandas_udf_scalar",
+ "pyspark.sql.tests.pandas.test_pipelined_udf",
"pyspark.sql.tests.pandas.test_pandas_udf_typehints",
"pyspark.sql.tests.pandas.test_pandas_udf_typehints_with_future_annotations",
"pyspark.sql.tests.pandas.test_pandas_udf_window",
diff --git a/python/benchmarks/bench_pipelined_udf.py
b/python/benchmarks/bench_pipelined_udf.py
new file mode 100644
index 000000000000..31eb9d0b16e3
--- /dev/null
+++ b/python/benchmarks/bench_pipelined_udf.py
@@ -0,0 +1,220 @@
+#
+# 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.
+#
+
+"""
+End-to-end benchmarks for pipelined vs synchronous Python UDF execution.
+
+Unlike the microbenchmarks in bench_eval_type.py (which test the Python worker
+in isolation), these benchmarks run full Spark queries through a real
+SparkSession to measure the JVM-Python socket I/O pipeline overlap.
+"""
+
+import pandas as pd
+
+from pyspark import SparkConf
+from pyspark.sql import SparkSession
+from pyspark.sql.functions import col, pandas_udf
+from pyspark.sql.types import LongType, StringType
+
+
+class _PipelinedUDFBenchBase:
+ """Base class for pipelined UDF benchmarks.
+
+ Each benchmark parameterizes over pipelined=true/false to compare
+ the two execution modes. SparkSession is created in setup() and
+ stopped in teardown() because spark.python.udf.pipelined.enabled
+ is a SparkConf-level config.
+ """
+
+ # Subclasses must define timeout (seconds per benchmark iteration).
+ timeout = 120
+
+ def _spark_conf(self, pipelined):
+ return (
+ SparkConf()
+ .setMaster("local[1]")
+ .setAppName("PipelinedUDFBench")
+ .set("spark.sql.execution.arrow.pyspark.enabled", "true")
+ .set("spark.python.worker.reuse", "true")
+ .set("spark.ui.enabled", "false")
+ .set("spark.sql.shuffle.partitions", "1")
+ .set("spark.python.udf.pipelined.enabled", str(pipelined).lower())
+ )
+
+ def _setup_spark(self, pipelined):
+ conf = self._spark_conf(pipelined)
+ self.spark = SparkSession.builder.config(conf=conf).getOrCreate()
+
+ @pandas_udf(LongType())
+ def _add_one(x: pd.Series) -> pd.Series:
+ return x + 1
+
+ self._add_one = _add_one
+
+ # Warmup: start Python worker, JIT
+
self.spark.range(100).select(_add_one(col("id"))).write.format("noop").mode(
+ "overwrite"
+ ).save()
+
+ def _teardown_spark(self):
+ if hasattr(self, "spark"):
+ self.spark.stop()
+ # Clear the active session so the next setup() creates a fresh one
+ SparkSession.builder._options = {}
+
+
+class ScalarUDFTimeBench(_PipelinedUDFBenchBase):
+ """Benchmark scalar Arrow UDF with light computation (x + 1)."""
+
+ params = [[False, True], [100000, 1000000]]
+ param_names = ["pipelined", "n_rows"]
+
+ def setup(self, pipelined, n_rows):
+ self._setup_spark(pipelined)
+
+ def teardown(self, pipelined, n_rows):
+ self._teardown_spark()
+
+ def time_scalar_udf(self, pipelined, n_rows):
+
self.spark.range(n_rows).select(self._add_one(col("id")).alias("result")).write.format(
+ "noop"
+ ).mode("overwrite").save()
+
+ def peakmem_scalar_udf(self, pipelined, n_rows):
+
self.spark.range(n_rows).select(self._add_one(col("id")).alias("result")).write.format(
+ "noop"
+ ).mode("overwrite").save()
+
+
+class MultiUDFTimeBench(_PipelinedUDFBenchBase):
+ """Benchmark multiple UDF columns in a single query."""
+
+ params = [[False, True], [100000, 1000000]]
+ param_names = ["pipelined", "n_rows"]
+
+ def setup(self, pipelined, n_rows):
+ self._setup_spark(pipelined)
+
+ @pandas_udf(LongType())
+ def _mul_two(x: pd.Series) -> pd.Series:
+ return x * 2
+
+ @pandas_udf(LongType())
+ def _sub_one(x: pd.Series) -> pd.Series:
+ return x - 1
+
+ self._mul_two = _mul_two
+ self._sub_one = _sub_one
+
+ def teardown(self, pipelined, n_rows):
+ self._teardown_spark()
+
+ def time_multi_udf(self, pipelined, n_rows):
+ self.spark.range(n_rows).select(
+ col("id"),
+ self._add_one(col("id")).alias("a"),
+ self._mul_two(col("id")).alias("b"),
+ self._sub_one(col("id")).alias("c"),
+ ).write.format("noop").mode("overwrite").save()
+
+ def peakmem_multi_udf(self, pipelined, n_rows):
+ self.spark.range(n_rows).select(
+ col("id"),
+ self._add_one(col("id")).alias("a"),
+ self._mul_two(col("id")).alias("b"),
+ self._sub_one(col("id")).alias("c"),
+ ).write.format("noop").mode("overwrite").save()
+
+
+class LargeDataUDFTimeBench(_PipelinedUDFBenchBase):
+ """Benchmark scalar UDF with large data to exercise throughput."""
+
+ params = [[False, True]]
+ param_names = ["pipelined"]
+
+ def setup(self, pipelined):
+ self._setup_spark(pipelined)
+
+ def teardown(self, pipelined):
+ self._teardown_spark()
+
+ def time_large_data(self, pipelined):
+
self.spark.range(5000000).select(self._add_one(col("id")).alias("result")).write.format(
+ "noop"
+ ).mode("overwrite").save()
+
+ def peakmem_large_data(self, pipelined):
+
self.spark.range(5000000).select(self._add_one(col("id")).alias("result")).write.format(
+ "noop"
+ ).mode("overwrite").save()
+
+
+class WideRowUDFTimeBench(_PipelinedUDFBenchBase):
+ """Benchmark scalar UDF with larger per-batch in-memory size.
+
+ Each row carries a wide string payload and the Arrow batch size is bumped
so
+ one batch is ~10-50 MB rather than ~80 KB. This exercises the regime that
+ Yicong-Huang asked about in the SPARK-56642 review: how does pipelined mode
+ behave when each batch is large enough that the queue's memory overhead is
+ no longer negligible?
+ """
+
+ # (pipelined, n_rows, payload_chars, records_per_batch)
+ # 50_000 rows * 1024 chars = ~50 MB raw per dataset; with records_per_batch
+ # = 10_000 that's ~10 MB per Arrow batch.
+ params = [
+ [False, True],
+ [(50_000, 1024, 10_000), (50_000, 4096, 5_000)],
+ ]
+ param_names = ["pipelined", "shape"]
+
+ def setup(self, pipelined, shape):
+ n_rows, payload_chars, records_per_batch = shape
+ self._setup_spark(pipelined)
+ self.spark.conf.set("spark.sql.execution.arrow.maxRecordsPerBatch",
str(records_per_batch))
+ self._n_rows = n_rows
+ self._payload_chars = payload_chars
+
+ @pandas_udf(StringType())
+ def _wide_passthrough(s: pd.Series) -> pd.Series:
+ # Non-trivial work proportional to row width so the UDF actually
+ # holds a batch for a measurable amount of time.
+ return s.str.upper()
+
+ self._wide_udf = _wide_passthrough
+
+ def teardown(self, pipelined, shape):
+ self._teardown_spark()
+
+ def _make_df(self):
+ chars = self._payload_chars
+
+ @pandas_udf(StringType())
+ def _make_payload(x: pd.Series) -> pd.Series:
+ return pd.Series(["x" * chars] * len(x))
+
+ return
self.spark.range(self._n_rows).select(_make_payload(col("id")).alias("payload"))
+
+ def time_wide_row_udf(self, pipelined, shape):
+
self._make_df().select(self._wide_udf(col("payload")).alias("result")).write.format(
+ "noop"
+ ).mode("overwrite").save()
+
+ def peakmem_wide_row_udf(self, pipelined, shape):
+
self._make_df().select(self._wide_udf(col("payload")).alias("result")).write.format(
+ "noop"
+ ).mode("overwrite").save()
diff --git a/python/pyspark/sql/pandas/serializers.py
b/python/pyspark/sql/pandas/serializers.py
index d6bf3b7d7416..20e1f80a99b4 100644
--- a/python/pyspark/sql/pandas/serializers.py
+++ b/python/pyspark/sql/pandas/serializers.py
@@ -127,9 +127,10 @@ class ArrowStreamSerializer(Serializer):
output batch. Default False.
"""
- def __init__(self, write_start_stream: bool = False) -> None:
+ def __init__(self, write_start_stream: bool = False, flush_per_batch: bool
= False) -> None:
super().__init__()
self._write_start_stream: bool = write_start_stream
+ self._flush_per_batch: bool = flush_per_batch
def dump_stream(self, iterator: Iterable["pa.RecordBatch"], stream:
IO[bytes]) -> None:
"""Optionally prepend START_ARROW_STREAM, then write batches."""
@@ -144,6 +145,10 @@ class ArrowStreamSerializer(Serializer):
if writer is None:
writer = pa.RecordBatchStreamWriter(stream, batch.schema)
writer.write_batch(batch)
+ # In pipelined mode, flush after each batch so the JVM can
read output
+ # while still sending input, rather than buffering all output.
+ if self._flush_per_batch:
+ stream.flush()
finally:
if writer is not None:
writer.close()
diff --git a/python/pyspark/sql/tests/pandas/bench_pipelined_udf.py
b/python/pyspark/sql/tests/pandas/bench_pipelined_udf.py
new file mode 100644
index 000000000000..5deed766cbce
--- /dev/null
+++ b/python/pyspark/sql/tests/pandas/bench_pipelined_udf.py
@@ -0,0 +1,305 @@
+#
+# 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.
+#
+
+"""
+Benchmark: Pipelined vs synchronous JVM-Python UDF data transfer.
+
+Compares end-to-end execution time of Python UDFs with
+spark.python.udf.pipelined.enabled = true vs false.
+
+Because spark.python.udf.pipelined.enabled is a SparkConf-level config (read at
+SparkContext startup), each benchmark scenario runs in a separate subprocess
with
+its own SparkSession to ensure the config takes effect.
+
+Note: In local[1] mode (single core), pipelined mode may show overhead because
+the writer thread and selector thread compete for the same CPU. The benefit of
+pipeline parallelism is expected on multi-core executors where serialization
can
+overlap with output reading.
+
+Usage:
+ cd $SPARK_HOME
+ # Build Spark first (needed for PySpark to find JVM jars):
+ # build/sbt -Phive package
+ # cd python && zip -r lib/pyspark.zip pyspark && cd ..
+ python python/pyspark/sql/tests/pandas/bench_pipelined_udf.py \
+ [--rows N] [--iterations N] [--partitions N] [--sleep-ms N]
+"""
+
+import argparse
+import json
+import os
+import subprocess
+import sys
+
+
+SPARK_HOME = os.path.join(os.path.dirname(os.path.abspath(__file__)),
"../../../../..")
+PIPELINED_CONF = "spark.python.udf.pipelined.enabled"
+QUEUE_DEPTH_CONF = "spark.python.udf.pipelined.queueDepth"
+
+
+# ---- Subprocess worker script template ----
+# Each benchmark scenario is run in a fresh Python process to get a fresh
SparkContext.
+WORKER_TEMPLATE = """
+import os, sys, time, json
+sys.path.insert(0, "{spark_home}")
+
+import pandas as pd
+from pyspark.sql import SparkSession
+from pyspark.sql.functions import pandas_udf, col
+from pyspark.sql.types import LongType
+
+spark = (
+ SparkSession.builder.master("{master}")
+ .appName("PipelinedUDFBench")
+ .config("spark.sql.execution.arrow.pyspark.enabled", "true")
+ .config("spark.python.worker.reuse", "true")
+ .config("spark.ui.enabled", "false")
+ .config("spark.sql.shuffle.partitions", "1")
+ .config("{pipelined_conf}", "{pipelined}")
+ .config("{queue_depth_conf}", "{queue_depth}")
+ .getOrCreate()
+)
+
+{udf_code}
+
+df = {make_df_code}
+
+# Warmup
+for _ in range({warmup}):
+ df.write.format("noop").mode("overwrite").save()
+
+# Timed runs
+times = []
+for _ in range({iterations}):
+ start = time.perf_counter()
+ df.write.format("noop").mode("overwrite").save()
+ elapsed = time.perf_counter() - start
+ times.append(elapsed)
+
+# Output results as JSON to stdout
+print("BENCH_RESULT:" + json.dumps(times))
+spark.stop()
+"""
+
+
+def run_subprocess(pipelined, udf_code, make_df_code, args):
+ """Run a benchmark in a fresh subprocess, return list of timing results."""
+ script = WORKER_TEMPLATE.format(
+ spark_home=os.path.abspath(SPARK_HOME),
+ master=args.master,
+ pipelined_conf=PIPELINED_CONF,
+ pipelined="true" if pipelined else "false",
+ queue_depth_conf=QUEUE_DEPTH_CONF,
+ queue_depth=args.queue_depth,
+ udf_code=udf_code,
+ make_df_code=make_df_code,
+ warmup=args.warmup,
+ iterations=args.iterations,
+ )
+ env = os.environ.copy()
+ env["SPARK_HOME"] = os.path.abspath(SPARK_HOME)
+ py4j_zip = os.path.join(os.path.abspath(SPARK_HOME),
"python/lib/py4j-0.10.9.9-src.zip")
+ pyspark_path = os.path.join(os.path.abspath(SPARK_HOME), "python")
+ env["PYTHONPATH"] = f"{pyspark_path}:{py4j_zip}:" + env.get("PYTHONPATH",
"")
+
+ result = subprocess.run(
+ [sys.executable, "-c", script], capture_output=True, text=True,
env=env, timeout=600
+ )
+
+ for line in result.stdout.splitlines():
+ if line.startswith("BENCH_RESULT:"):
+ return json.loads(line[len("BENCH_RESULT:") :])
+
+ print(" ERROR: no BENCH_RESULT in output")
+ print(" STDERR (last 500 chars):", result.stderr[-500:] if result.stderr
else "<empty>")
+ return None
+
+
+def print_stats(label, times):
+ if not times:
+ print(f" {label:40s} FAILED")
+ return 0.0
+ avg = sum(times) / len(times)
+ mn = min(times)
+ mx = max(times)
+ print(
+ f" {label:40s} "
+ f"avg = {avg * 1000:8.1f} ms "
+ f"min = {mn * 1000:8.1f} ms "
+ f"max = {mx * 1000:8.1f} ms "
+ f"({len(times)} iters)"
+ )
+ return avg
+
+
+def run_benchmark(label, udf_code, make_df_code, args):
+ """Run sync and pipelined in separate subprocesses, print comparison."""
+ print(f" [{label}]")
+
+ sync_times = run_subprocess(False, udf_code, make_df_code, args)
+ sync_avg = print_stats("sync (pipelined=false)", sync_times)
+
+ pipe_times = run_subprocess(True, udf_code, make_df_code, args)
+ pipe_avg = print_stats("pipelined (pipelined=true)", pipe_times)
+
+ if pipe_avg > 0 and sync_avg > 0:
+ speedup = sync_avg / pipe_avg
+ diff_ms = (sync_avg - pipe_avg) * 1000
+ marker = "faster" if speedup > 1.0 else "slower"
+ print(f" --> pipelined is {speedup:.2f}x {marker} ({diff_ms:+.1f}
ms)")
+ print()
+ return sync_avg, pipe_avg
+
+
+def main():
+ parser = argparse.ArgumentParser(
+ description="Benchmark pipelined vs synchronous Python UDF data
transfer"
+ )
+ parser.add_argument(
+ "--rows",
+ type=int,
+ default=1_000_000,
+ help="Rows for standard benchmarks (default: 1000000)",
+ )
+ parser.add_argument(
+ "--large-rows",
+ type=int,
+ default=5_000_000,
+ help="Rows for large data benchmark (default: 5000000)",
+ )
+ parser.add_argument(
+ "--iterations", type=int, default=5, help="Timed iterations per
scenario (default: 5)"
+ )
+ parser.add_argument("--warmup", type=int, default=2, help="Warmup
iterations (default: 2)")
+ parser.add_argument(
+ "--partitions", type=int, default=1, help="Number of partitions
(default: 1)"
+ )
+ parser.add_argument(
+ "--sleep-ms",
+ type=float,
+ default=10.0,
+ help="Sleep time in ms per batch for heavy UDF (default: 10.0)",
+ )
+ parser.add_argument(
+ "--queue-depth", type=int, default=2, help="Pipelined queue depth
(default: 2)"
+ )
+ parser.add_argument(
+ "--master", type=str, default="local[1]", help="Spark master URL
(default: local[1])"
+ )
+ args = parser.parse_args()
+
+ nparts = args.partitions
+
+ print("=" * 78)
+ print(" Pipelined vs Synchronous Python UDF Data Transfer Benchmark")
+ print("=" * 78)
+ print(
+ f" master={args.master} rows={args.rows}
large_rows={args.large_rows} "
+ f"partitions={nparts}"
+ )
+ print(
+ f" iterations={args.iterations} warmup={args.warmup} "
+ f"sleep_ms={args.sleep_ms} queue_depth={args.queue_depth}"
+ )
+ print()
+
+ # --- Benchmark 1: Light UDF ---
+ run_benchmark(
+ "Light UDF (x + 1)",
+ udf_code="""
+@pandas_udf(LongType())
+def bench_udf(x: pd.Series) -> pd.Series:
+ return x + 1
+""",
+ make_df_code=f"spark.range({args.rows}, numPartitions={nparts})"
+ f'.select(col("id"), bench_udf(col("id")).alias("result"))',
+ args=args,
+ )
+
+ # --- Benchmark 2: CPU-bound UDF ---
+ run_benchmark(
+ "CPU-bound UDF (iterative computation)",
+ udf_code="""
+@pandas_udf(LongType())
+def bench_udf(x: pd.Series) -> pd.Series:
+ result = x + 1
+ for _ in range(20):
+ result = result + (x % 7) - 3
+ return result
+""",
+ make_df_code=f"spark.range({args.rows}, numPartitions={nparts})"
+ f'.select(col("id"), bench_udf(col("id")).alias("result"))',
+ args=args,
+ )
+
+ # --- Benchmark 3: Heavy UDF (sleep) ---
+ run_benchmark(
+ f"Heavy UDF ({args.sleep_ms}ms sleep/batch)",
+ udf_code=f"""
+import time as _time
+@pandas_udf(LongType())
+def bench_udf(x: pd.Series) -> pd.Series:
+ _time.sleep({args.sleep_ms / 1000.0})
+ return x + 1
+""",
+ make_df_code=f"spark.range({args.rows}, numPartitions={nparts})"
+ f'.select(col("id"), bench_udf(col("id")).alias("result"))',
+ args=args,
+ )
+
+ # --- Benchmark 4: Large data ---
+ run_benchmark(
+ f"Large data ({args.large_rows} rows, x + 1)",
+ udf_code="""
+@pandas_udf(LongType())
+def bench_udf(x: pd.Series) -> pd.Series:
+ return x + 1
+""",
+ make_df_code=f"spark.range({args.large_rows}, numPartitions={nparts})"
+ f'.select(col("id"), bench_udf(col("id")).alias("result"))',
+ args=args,
+ )
+
+ # --- Benchmark 5: Multiple UDF columns ---
+ run_benchmark(
+ "Multi-UDF (3 UDF columns)",
+ udf_code="""
+@pandas_udf(LongType())
+def udf_a(x: pd.Series) -> pd.Series:
+ return x + 1
+
+@pandas_udf(LongType())
+def udf_b(x: pd.Series) -> pd.Series:
+ return x * 2
+
+@pandas_udf(LongType())
+def udf_c(x: pd.Series) -> pd.Series:
+ return x - 1
+""",
+ make_df_code=f"spark.range({args.rows}, numPartitions={nparts})"
+ f'.select(col("id"), udf_a(col("id")).alias("a"), '
+ f'udf_b(col("id")).alias("b"), udf_c(col("id")).alias("c"))',
+ args=args,
+ )
+
+ print("=" * 78)
+ print(" Benchmark complete.")
+ print("=" * 78)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/python/pyspark/sql/tests/pandas/test_pipelined_udf.py
b/python/pyspark/sql/tests/pandas/test_pipelined_udf.py
new file mode 100644
index 000000000000..1e133f6e219b
--- /dev/null
+++ b/python/pyspark/sql/tests/pandas/test_pipelined_udf.py
@@ -0,0 +1,285 @@
+#
+# 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.
+#
+
+"""
+Tests for pipelined Python UDF execution mode.
+
+These tests run with spark.python.udf.pipelined.enabled=true to verify
+correctness of the pipelined data transfer path for various UDF types.
+"""
+
+import os
+import unittest
+
+from pyspark import SparkConf
+from pyspark.sql.functions import col, pandas_udf, udf
+from pyspark.sql.types import (
+ DoubleType,
+ LongType,
+ StringType,
+ StructType,
+ StructField,
+)
+from pyspark.testing.sqlutils import ReusedSQLTestCase
+from pyspark.testing.utils import (
+ have_pandas,
+ have_pyarrow,
+ pandas_requirement_message,
+ pyarrow_requirement_message,
+)
+
+if have_pandas:
+ import pandas as pd
+
+
[email protected](
+ not have_pandas or not have_pyarrow,
+ pandas_requirement_message or pyarrow_requirement_message,
+)
+class PipelinedUDFTests(ReusedSQLTestCase):
+ """Tests that run with pipelined mode enabled."""
+
+ @classmethod
+ def conf(cls):
+ return (
+ SparkConf()
+ .set("spark.python.udf.pipelined.enabled", "true")
+ .set("spark.sql.execution.arrow.pyspark.enabled", "true")
+ )
+
+ def test_pipelined_mode_is_active(self):
+ """Verify the pipelined code path is actually being used."""
+
+ @pandas_udf(StringType())
+ def check_env(x: pd.Series) -> pd.Series:
+ # SPARK_PIPELINED_UDF is set by the JVM when pipelined mode is
enabled.
+ flag = os.environ.get("SPARK_PIPELINED_UDF", "not_set")
+ return pd.Series([flag] * len(x))
+
+ result = self.spark.range(1).select(check_env(col("id"))).first()[0]
+ self.assertEqual(result, "1", "JVM should set SPARK_PIPELINED_UDF=1")
+
+ def test_scalar_arrow_udf(self):
+ """Basic scalar Arrow UDF with pipelined mode."""
+
+ @pandas_udf(LongType())
+ def add_one(x: pd.Series) -> pd.Series:
+ return x + 1
+
+ result = (
+ self.spark.range(100).select(col("id"),
add_one(col("id")).alias("result")).collect()
+ )
+ for row in result:
+ self.assertEqual(row.result, row.id + 1)
+
+ def test_scalar_arrow_udf_string(self):
+ """Scalar Arrow UDF with string type."""
+
+ @pandas_udf(StringType())
+ def to_str(x: pd.Series) -> pd.Series:
+ return x.astype(str) + "_val"
+
+ result = self.spark.range(50).select(col("id"),
to_str(col("id")).alias("s")).collect()
+ for row in result:
+ self.assertEqual(row.s, f"{row.id}_val")
+
+ def test_multiple_udf_columns(self):
+ """Multiple UDF columns in a single query."""
+
+ @pandas_udf(LongType())
+ def udf_a(x: pd.Series) -> pd.Series:
+ return x + 1
+
+ @pandas_udf(LongType())
+ def udf_b(x: pd.Series) -> pd.Series:
+ return x * 2
+
+ @pandas_udf(LongType())
+ def udf_c(x: pd.Series) -> pd.Series:
+ return x - 1
+
+ result = (
+ self.spark.range(100)
+ .select(
+ col("id"),
+ udf_a(col("id")).alias("a"),
+ udf_b(col("id")).alias("b"),
+ udf_c(col("id")).alias("c"),
+ )
+ .collect()
+ )
+ for row in result:
+ self.assertEqual(row.a, row.id + 1)
+ self.assertEqual(row.b, row.id * 2)
+ self.assertEqual(row.c, row.id - 1)
+
+ def test_multiple_partitions(self):
+ """UDF across multiple partitions."""
+
+ @pandas_udf(LongType())
+ def add_one(x: pd.Series) -> pd.Series:
+ return x + 1
+
+ result = (
+ self.spark.range(1000, numPartitions=4)
+ .select(col("id"), add_one(col("id")).alias("result"))
+ .collect()
+ )
+ self.assertEqual(len(result), 1000)
+ for row in result:
+ self.assertEqual(row.result, row.id + 1)
+
+ def test_empty_partition(self):
+ """UDF with some empty partitions."""
+
+ @pandas_udf(LongType())
+ def add_one(x: pd.Series) -> pd.Series:
+ return x + 1
+
+ # 2 rows across 4 partitions means some partitions are empty
+ result = (
+ self.spark.range(2, numPartitions=4)
+ .select(col("id"), add_one(col("id")).alias("result"))
+ .collect()
+ )
+ self.assertEqual(len(result), 2)
+ for row in result:
+ self.assertEqual(row.result, row.id + 1)
+
+ def test_chained_udf(self):
+ """Chained UDF calls."""
+
+ @pandas_udf(LongType())
+ def add_one(x: pd.Series) -> pd.Series:
+ return x + 1
+
+ @pandas_udf(LongType())
+ def double_it(x: pd.Series) -> pd.Series:
+ return x * 2
+
+ result = (
+
self.spark.range(50).select(double_it(add_one(col("id"))).alias("result")).collect()
+ )
+ for row in result:
+ self.assertEqual(row.result, (row.result // 2) * 2) # even number
+
+ def test_udf_with_null(self):
+ """UDF handling null values."""
+
+ @pandas_udf(LongType())
+ def add_one(x: pd.Series) -> pd.Series:
+ return x + 1
+
+ df = self.spark.createDataFrame(
+ [(1,), (None,), (3,), (None,), (5,)],
+ schema=StructType([StructField("v", LongType(), True)]),
+ )
+ result = df.select(col("v"),
add_one(col("v")).alias("result")).collect()
+ expected = [(1, 2), (None, None), (3, 4), (None, None), (5, 6)]
+ for row, (v, r) in zip(result, expected):
+ self.assertEqual(row.v, v)
+ self.assertEqual(row.result, r)
+
+ def test_grouped_agg_udf(self):
+ """Grouped aggregation UDF (UDAF) with pipelined mode."""
+
+ @pandas_udf(DoubleType())
+ def mean_udf(x: pd.Series) -> float:
+ return float(x.mean())
+
+ df = self.spark.range(100).selectExpr("id", "id % 5 as grp")
+ result =
df.groupBy("grp").agg(mean_udf(col("id")).alias("avg")).orderBy("grp").collect()
+
+ self.assertEqual(len(result), 5)
+ # grp=0: mean of 0,5,10,...,95 = 47.5
+ self.assertAlmostEqual(result[0].avg, 47.5)
+
+ def test_scalar_udf_large_data(self):
+ """Scalar UDF with large data to exercise backpressure."""
+
+ @pandas_udf(LongType())
+ def add_one(x: pd.Series) -> pd.Series:
+ return x + 1
+
+ result = (
+ self.spark.range(500000)
+ .select(add_one(col("id")).alias("result"))
+ .agg({"result": "sum"})
+ .collect()
+ )
+ # sum of (1..500000) = 500000 * 500001 / 2 = 125000250000
+ self.assertEqual(result[0][0], 125000250000)
+
+ def test_batched_udf(self):
+ """Non-Arrow batched UDF (pickle serialization)."""
+
+ @udf(StringType())
+ def simple_udf(x):
+ return str(x) + "_done"
+
+ result = self.spark.range(50).select(col("id"),
simple_udf(col("id")).alias("s")).collect()
+ for row in result:
+ self.assertEqual(row.s, f"{row.id}_done")
+
+ def test_udf_exception_propagation(self):
+ """UDF that raises an exception should propagate correctly."""
+
+ @pandas_udf(LongType())
+ def bad_udf(x: pd.Series) -> pd.Series:
+ raise ValueError("intentional error")
+
+ with self.assertRaisesRegex(Exception, "intentional error"):
+ self.spark.range(10).select(bad_udf(col("id"))).collect()
+
+ def test_offheap_reader_with_head_does_not_segfault(self):
+ """Regression test: SPARK-33277 reproducer adapted for pipelined mode.
+
+ With the off-heap vectorized Parquet reader, head() triggers a
+ driver-side cancel of the still-running task. The task completion
+ listener for the writer thread must wait for the writer to actually
+ exit before subsequent listeners free off-heap memory backing the
+ rows the writer is still serializing. Otherwise the writer reads
+ invalidated off-heap memory and the executor segfaults.
+
+ The race is probabilistic (the SPARK-33277 PR reported ~3% failure
+ rate without the fix), so we repeat the head() many times to make
+ the test reliably crash without the fix while still finishing fast.
+ """
+ import shutil
+ import tempfile
+
+ path = tempfile.mkdtemp()
+ shutil.rmtree(path)
+ try:
+ self.spark.range(0, 200000, 1, 1).write.parquet(path)
+
+ @pandas_udf(LongType())
+ def to_zero(x: pd.Series) -> pd.Series:
+ return pd.Series([0] * len(x))
+
+ with self.sql_conf({"spark.sql.columnVector.offheap.enabled":
"true"}):
+ for _ in range(100):
+ row =
self.spark.read.parquet(path).select(to_zero(col("id"))).head()
+ self.assertEqual(row[0], 0)
+ finally:
+ shutil.rmtree(path, ignore_errors=True)
+
+
+if __name__ == "__main__":
+ from pyspark.testing import main
+
+ main()
diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py
index 5a82d11ebe9b..5ceebca15338 100644
--- a/python/pyspark/worker.py
+++ b/python/pyspark/worker.py
@@ -26,6 +26,7 @@ import time
import inspect
import itertools
import json
+import warnings
from collections.abc import Iterator
from typing import (
Any,
@@ -3731,12 +3732,108 @@ def invoke_udf(message_receiver: SparkMessageReceiver,
outfile: BinaryIO):
if hasattr(out_iter, "close"):
out_iter.close()
+ def pipelined_process():
+ """
+ Pipelined variant of process() that pre-fetches input batches in a
background
+ reader thread while the main thread computes the UDF and writes
output.
+ This allows input deserialization to overlap with UDF computation.
+ """
+ import queue
+ import threading
+
+ queue_depth =
int(os.environ.get("SPARK_PIPELINED_UDF_QUEUE_DEPTH", "2"))
+ _SENTINEL = object()
+ input_queue = queue.Queue(maxsize=queue_depth)
+ reader_error = [None]
+ # Event to signal the reader thread to stop (set by main thread on
+ # exception or completion). The reader checks this after each
failed
+ # put attempt instead of polling with a timeout.
+ stop_event = threading.Event()
+
+ def _reader_thread():
+ try:
+ for batch in deserializer.load_stream(input_data_stream):
+ # Some serializers (e.g., ArrowStreamGroupSerializer,
+ # ArrowStreamAggPandasUDFSerializer) yield lazy
iterators
+ # that still read from the input stream. Materialize
them here so
+ # the main thread can consume them without touching
the stream.
+ if hasattr(batch, "__next__"):
+ batch = list(batch)
+ # Block on put, but wake up when stop_event is set.
+ # stop_event.wait() returns immediately if already set.
+ while not stop_event.is_set():
+ try:
+ input_queue.put(batch, timeout=0.1)
+ break
+ except queue.Full:
+ continue
+ if stop_event.is_set():
+ return
+ except Exception as e:
+ reader_error[0] = e
+ finally:
+ # Enqueue sentinel so the consumer knows we're done.
+ while not stop_event.is_set():
+ try:
+ input_queue.put(_SENTINEL, timeout=0.1)
+ break
+ except queue.Full:
+ continue
+
+ t = threading.Thread(
+ target=_reader_thread, name="pyspark-pipelined-reader",
daemon=True
+ )
+ t.start()
+
+ def _queued_iter():
+ while True:
+ item = input_queue.get()
+ if item is _SENTINEL:
+ if reader_error[0] is not None:
+ raise reader_error[0]
+ return
+ yield item
+
+ out_iter = func(init_info.split_index, _queued_iter())
+ try:
+ serializer.dump_stream(out_iter, outfile)
+ finally:
+ if hasattr(out_iter, "close"):
+ out_iter.close()
+ # Signal reader thread to stop, drain the queue so it can
unblock,
+ # then wait for it to finish.
+ stop_event.set()
+ try:
+ while not input_queue.empty():
+ input_queue.get_nowait()
+ except Exception:
+ pass
+ # If the reader is still blocked in input_data_stream.read(),
the stop_event
+ # check only fires between put attempts -- it cannot interrupt
a syscall.
+ # Force-closing the stream here would break worker reuse (the
next task uses
+ # the same socket fd), so we settle for a bounded join and a
loud warning
+ # so an undetected leak shows up in the worker log.
+ t.join(timeout=5)
+ if t.is_alive():
+ warnings.warn(
+ "pipelined reader thread did not exit within 5s; "
+ "it may still be blocked in input_data_stream.read()
and could "
+ "read data intended for a subsequent reused-worker
task. "
+ "Consider disabling spark.python.worker.reuse if this
recurs.",
+ RuntimeWarning,
+ )
+
+ is_pipelined = os.environ.get("SPARK_PIPELINED_UDF") == "1"
+ if is_pipelined and hasattr(serializer, "_flush_per_batch"):
+ serializer._flush_per_batch = True
+ run_process = pipelined_process if is_pipelined else process
+
processing_start_time = time.time()
with capture_outputs():
if profiler:
- profiler.profile(process)
+ profiler.profile(run_process)
else:
- process()
+ run_process()
processing_time_ms = int(1000 * (time.time() - processing_start_time))
# Cleanup
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowAggregatePythonExec.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowAggregatePythonExec.scala
index b265d1de54b6..09bfe6538420 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowAggregatePythonExec.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowAggregatePythonExec.scala
@@ -23,6 +23,7 @@ import scala.collection.mutable.ArrayBuffer
import org.apache.spark.{JobArtifactSet, SparkEnv, SparkException, TaskContext}
import org.apache.spark.api.python.{ChainedPythonFunctions, PythonEvalType}
+import org.apache.spark.internal.config.Python.PYTHON_UDF_PIPELINED_EXECUTION
import org.apache.spark.rdd.RDD
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions._
@@ -176,8 +177,12 @@ case class ArrowAggregatePythonExec(
// The queue used to buffer input rows so we can drain it to
// combine input with output from Python.
+ // In pipelined mode the queue's add() runs in the writer thread and
remove() runs in
+ // the task thread; use lock-free mode to skip per-row synchronization.
+ val pipelined = SparkEnv.get.conf.get(PYTHON_UDF_PIPELINED_EXECUTION)
val queue = HybridRowQueue(context.taskMemoryManager(),
- new File(Utils.getLocalDir(SparkEnv.get.conf)),
groupingExpressions.length)
+ new File(Utils.getLocalDir(SparkEnv.get.conf)),
groupingExpressions.length,
+ lockFree = pipelined)
context.addTaskCompletionListener[Unit] { _ =>
queue.close()
}
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowWindowPythonEvaluatorFactory.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowWindowPythonEvaluatorFactory.scala
index 5ba47db7fc94..ab9671c022f9 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowWindowPythonEvaluatorFactory.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowWindowPythonEvaluatorFactory.scala
@@ -24,6 +24,7 @@ import scala.jdk.CollectionConverters._
import org.apache.spark.{JobArtifactSet, PartitionEvaluator,
PartitionEvaluatorFactory, SparkEnv, TaskContext}
import org.apache.spark.api.python.ChainedPythonFunctions
+import org.apache.spark.internal.config.Python.PYTHON_UDF_PIPELINED_EXECUTION
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions.{Attribute,
AttributeReference, BoundReference, EmptyRow, Expression, JoinedRow,
NamedArgumentExpression, NamedExpression, PythonFuncExpression, PythonUDAF,
SortOrder, SpecificInternalRow, UnsafeProjection, UnsafeRow, WindowExpression}
import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression
@@ -255,8 +256,11 @@ class ArrowWindowPythonEvaluatorFactory(
// The queue used to buffer input rows so we can drain it to
// combine input with output from Python.
+ // In pipelined mode the queue's add() runs in the writer thread and
remove() runs in
+ // the task thread; use lock-free mode to skip per-row synchronization.
+ val pipelined = SparkEnv.get.conf.get(PYTHON_UDF_PIPELINED_EXECUTION)
val queue = HybridRowQueue(context.taskMemoryManager(),
- new File(Utils.getLocalDir(SparkEnv.get.conf)), childOutput.length)
+ new File(Utils.getLocalDir(SparkEnv.get.conf)), childOutput.length,
lockFree = pipelined)
context.addTaskCompletionListener[Unit] { _ =>
queue.close()
}
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala
index 4e699f975fec..bad609919a4b 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala
@@ -24,6 +24,7 @@ import scala.jdk.CollectionConverters._
import org.apache.spark.{PartitionEvaluator, PartitionEvaluatorFactory,
SparkEnv, TaskContext}
import org.apache.spark.api.python.ChainedPythonFunctions
+import org.apache.spark.internal.config.Python.PYTHON_UDF_PIPELINED_EXECUTION
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions._
import org.apache.spark.sql.errors.QueryExecutionErrors
@@ -235,10 +236,14 @@ private[python] class
ColumnarArrowEvalPythonEvaluatorFactory(
inputColumnIndices: Option[Array[Int]]
): Iterator[ColumnarBatch] = {
+ // In pipelined mode the queue's add() runs in the writer thread and
remove() runs in
+ // the task thread; use lock-free mode to skip per-row synchronization.
+ val pipelined = SparkEnv.get.conf.get(PYTHON_UDF_PIPELINED_EXECUTION)
val queue = HybridRowQueue(
context.taskMemoryManager(),
new File(Utils.getLocalDir(SparkEnv.get.conf)),
- childOutput.length)
+ childOutput.length,
+ lockFree = pipelined)
context.addTaskCompletionListener[Unit] { _ => queue.close() }
val unsafeProj = UnsafeProjection.create(
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/EvalPythonEvaluatorFactory.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/EvalPythonEvaluatorFactory.scala
index 34f9be0aa633..c0a983c60afa 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/EvalPythonEvaluatorFactory.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/EvalPythonEvaluatorFactory.scala
@@ -23,6 +23,7 @@ import scala.collection.mutable.ArrayBuffer
import org.apache.spark.{PartitionEvaluator, PartitionEvaluatorFactory,
SparkEnv, TaskContext}
import org.apache.spark.api.python.ChainedPythonFunctions
+import org.apache.spark.internal.config.Python.PYTHON_UDF_PIPELINED_EXECUTION
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions._
import org.apache.spark.sql.execution.python.EvalPythonExec.ArgumentMetadata
@@ -67,10 +68,15 @@ abstract class EvalPythonEvaluatorFactory(
// The queue used to buffer input rows so we can drain it to
// combine input with output from Python.
+ // In pipelined mode, add() runs in the writer thread and remove() in
the task thread.
+ // Use lock-free mode to avoid synchronized overhead (memory visibility
is guaranteed
+ // by the blocking socket I/O between the two threads).
+ val pipelined = SparkEnv.get.conf.get(PYTHON_UDF_PIPELINED_EXECUTION)
val queue = HybridRowQueue(
context.taskMemoryManager(),
new File(Utils.getLocalDir(SparkEnv.get.conf)),
- childOutput.length)
+ childOutput.length,
+ lockFree = pipelined)
context.addTaskCompletionListener[Unit] { ctx =>
queue.close()
}
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/EvalPythonUDTFExec.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/EvalPythonUDTFExec.scala
index 3cb9431fed6f..c5bad567add6 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/EvalPythonUDTFExec.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/EvalPythonUDTFExec.scala
@@ -22,6 +22,7 @@ import java.io.File
import scala.collection.mutable.ArrayBuffer
import org.apache.spark.{SparkEnv, TaskContext}
+import org.apache.spark.internal.config.Python.PYTHON_UDF_PIPELINED_EXECUTION
import org.apache.spark.rdd.RDD
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions._
@@ -59,8 +60,11 @@ trait EvalPythonUDTFExec extends UnaryExecNode {
// The queue used to buffer input rows so we can drain it to
// combine input with output from Python.
+ // In pipelined mode the queue's add() runs in the writer thread and
remove() runs in
+ // the task thread; use lock-free mode to skip per-row synchronization.
+ val pipelined = SparkEnv.get.conf.get(PYTHON_UDF_PIPELINED_EXECUTION)
val queue = HybridRowQueue(context.taskMemoryManager(),
- new File(Utils.getLocalDir(SparkEnv.get.conf)), child.output.length)
+ new File(Utils.getLocalDir(SparkEnv.get.conf)), child.output.length,
lockFree = pipelined)
context.addTaskCompletionListener[Unit] { ctx =>
queue.close()
}
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/RowQueue.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/RowQueue.scala
index d2008cfa1309..9f980be51d51 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/RowQueue.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/RowQueue.scala
@@ -43,40 +43,63 @@ trait RowQueue extends Queue[UnsafeRow]
* A RowQueue that is based on in-memory page. UnsafeRows are appended into it
until it's full.
* Another thread could read from it at the same time (behind the writer).
*
+ * When `lockFree` is false (default), add() and remove() use synchronized for
thread safety.
+ * When `lockFree` is true (pipelined Python UDF mode), synchronized is
replaced by a
+ * volatile `writeOffset` using SPSC release-acquire semantics:
+ * - add() performs a volatile store on writeOffset after writing row data
(release fence),
+ * ensuring all prior Platform.putInt/copyMemory writes are visible before
the offset update.
+ * - remove() performs a volatile load on writeOffset (acquire fence) to see
the latest data.
+ * - readOffset does not need to be volatile because the writer never reads
it.
+ *
* The format of UnsafeRow in page:
* [4 bytes to hold length of record (N)] [N bytes to hold record] [...]
*
* -1 length means end of page.
*/
-private[python] abstract class InMemoryRowQueue(val page: MemoryBlock,
numFields: Int)
+private[python] abstract class InMemoryRowQueue(
+ val page: MemoryBlock, numFields: Int, lockFree: Boolean = false)
extends RowQueue {
private val base: AnyRef = page.getBaseObject
private val endOfPage: Long = page.getBaseOffset + page.size
// the first location where a new row would be written
- private var writeOffset = page.getBaseOffset
- // points to the start of the next row to read
+ // When lockFree=true, this is accessed via volatile read/write for SPSC
visibility.
+ // When lockFree=false, synchronized provides the memory barrier.
+ @volatile private var writeOffset = page.getBaseOffset
+ // points to the start of the next row to read (only updated by consumer)
private var readOffset = page.getBaseOffset
private val resultRow = new UnsafeRow(numFields)
- def add(row: UnsafeRow): Boolean = synchronized {
+ private def doAdd(row: UnsafeRow): Boolean = {
+ // Cache writeOffset in a local var to avoid repeated volatile reads in
lockFree mode.
+ val curOffset = writeOffset
val size = row.getSizeInBytes
- if (writeOffset + 4 + size > endOfPage) {
+ if (curOffset + 4 + size > endOfPage) {
// if there is not enough space in this page to hold the new record
- if (writeOffset + 4 <= endOfPage) {
+ if (curOffset + 4 <= endOfPage) {
// if there's extra space at the end of the page, store a special
"end-of-page" length (-1)
- Platform.putInt(base, writeOffset, -1)
+ Platform.putInt(base, curOffset, -1)
+ // Volatile store to publish the end-of-page marker. The reader relies
on seeing
+ // -1 to know this page is exhausted and switch to the next queue.
+ writeOffset = curOffset
}
false
} else {
- Platform.putInt(base, writeOffset, size)
- Platform.copyMemory(row.getBaseObject, row.getBaseOffset, base,
writeOffset + 4, size)
- writeOffset += 4 + size
+ Platform.putInt(base, curOffset, size)
+ Platform.copyMemory(row.getBaseObject, row.getBaseOffset, base,
curOffset + 4, size)
+ // Volatile store acts as a release fence: all prior writes (row data)
are visible
+ // to any thread that subsequently reads this writeOffset via volatile
load.
+ writeOffset = curOffset + 4 + size
true
}
}
- def remove(): UnsafeRow = synchronized {
- assert(readOffset <= writeOffset, "reader should not go beyond writer")
+ private def doRemove(): UnsafeRow = {
+ // Volatile load acts as an acquire fence: ensures all row data written by
the
+ // producer (before its volatile store of writeOffset) is visible to this
thread.
+ // Read unconditionally into a local val so the acquire fence is not
dependent on
+ // assert being enabled.
+ val curWriteOffset = writeOffset
+ assert(readOffset <= curWriteOffset, "reader should not go beyond writer")
if (readOffset + 4 > endOfPage || Platform.getInt(base, readOffset) < 0) {
null
} else {
@@ -86,6 +109,12 @@ private[python] abstract class InMemoryRowQueue(val page:
MemoryBlock, numFields
resultRow
}
}
+
+ def add(row: UnsafeRow): Boolean =
+ if (lockFree) doAdd(row) else synchronized { doAdd(row) }
+
+ def remove(): UnsafeRow =
+ if (lockFree) doRemove() else synchronized { doRemove() }
}
/**
@@ -156,7 +185,8 @@ case class HybridRowQueue(
memManager: TaskMemoryManager,
tempDir: File,
numFields: Int,
- serMgr: SerializerManager)
+ serMgr: SerializerManager,
+ lockFree: Boolean = false)
extends HybridQueue[UnsafeRow, RowQueue](memManager, tempDir, serMgr) {
override protected def createDiskQueue(): RowQueue = {
@@ -164,7 +194,7 @@ case class HybridRowQueue(
}
override protected def createInMemoryQueue(page: MemoryBlock): RowQueue = {
- new InMemoryRowQueue(page, numFields) {
+ new InMemoryRowQueue(page, numFields, lockFree) {
override def close(): Unit = {
freePage(this.page)
}
@@ -185,6 +215,14 @@ object HybridRowQueue {
HybridRowQueue(taskMemoryMgr, file, fields, SparkEnv.get.serializerManager)
}
+ def apply(
+ taskMemoryMgr: TaskMemoryManager,
+ file: File,
+ fields: Int,
+ lockFree: Boolean): HybridRowQueue = {
+ HybridRowQueue(taskMemoryMgr, file, fields,
SparkEnv.get.serializerManager, lockFree)
+ }
+
def apply(taskMemoryMgr: TaskMemoryManager, fields: Int): HybridRowQueue = {
apply(taskMemoryMgr, new File(Utils.getLocalDir(SparkEnv.get.conf)),
fields)
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
