viirya commented on code in PR #58978: URL: https://github.com/apache/spark/pull/58978#discussion_r4333611732
########## sql/core/src/main/scala/org/apache/spark/sql/execution/python/InProcessPythonRuntime.scala: ########## @@ -0,0 +1,432 @@ +/* + * 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 java.io.File +import java.util.concurrent.{Callable, ExecutionException, Executors, ThreadFactory, TimeoutException, TimeUnit} +import java.util.concurrent.atomic.AtomicInteger + +import scala.collection.mutable +import scala.jdk.CollectionConverters._ + +import jep.{JepConfig, JepException, MainInterpreter, NamingConventionClassEnquirer, PyConfig, SharedInterpreter} +import org.apache.arrow.c.{ArrowSchema, Data} +import org.apache.arrow.vector.types.pojo.Field + +import org.apache.spark.{TaskContext, TaskKilledException} +import org.apache.spark.api.python.{PythonException, PythonUtils} +import org.apache.spark.internal.Logging +import org.apache.spark.internal.config.Python +import org.apache.spark.sql.util.ArrowUtils +import org.apache.spark.util.Utils + +/** Owns one interpreter generation per executor plugin lifecycle. */ +private[python] object InProcessPythonRuntime extends Logging { + private val TRACEBACK_SENTINEL = "__INPROCESS_UDF_TRACEBACK__:" + private var active: InterpreterSession = _ + private var mainConfigured = false + @volatile private var sharedConfigured = false + @volatile private var bootstrappedSitePackages: Option[Seq[String]] = None + // Every sitePackages given to JEP. JEP adds them to sys.path before it imports its package, + // so a failed configuration leaves them there, and the bootstrap of a retry removes them. + // Guarded by the object's monitor. + private var configuredSitePackages = Seq.empty[String] + + private[python] class LifecycleException(message: String) extends IllegalStateException(message) + + // Keep JEP references out of the singleton's verifier so currentSession can report + // an uninitialized runtime even when the provided JEP JAR is absent. + private[python] object InterpreterConfiguration { + def configure(sitePackages: Seq[String]): Unit = { + if (!mainConfigured) { + // Like Python workers, use a stable default hash seed on every executor. This must + // happen before JEP creates its process-wide main interpreter, including on restarts. + MainInterpreter.setInitParams( + PyConfig.isolated().setUseEnvironment(false).setHashSeed(0).setUseHashSeed(true)) + mainConfigured = true + } + if (!sharedConfigured) { + // JEP imports its Python package during construction, before our bootstrap runs. + SharedInterpreter.setConfig(interpreterConfig(sitePackages)) + configuredSitePackages = (configuredSitePackages ++ sitePackages).distinct + } + } + + def interpreterConfig(sitePackages: Seq[String]): JepConfig = { + require(sitePackages.forall(Python.isValidInProcessPath), Python.IN_PROCESS_PATH_RULE) + val config = new JepConfig().setClassEnquirer(new NamingConventionClassEnquirer(false)) + // Calling addIncludePaths with no arguments adds the working directory in JEP. + if (sitePackages.nonEmpty) config.addIncludePaths(sitePackages: _*) + config + } + } + + private class ManagedSharedInterpreter extends SharedInterpreter { + override protected def configureInterpreter(config: JepConfig): Unit = { + // JEP invokes this hook after native initialization, from its constructor. Close + // here if configuration fails, before the caller can receive an interpreter handle. + try { + super.configureInterpreter(config) + sharedConfigured = true + } catch { + case t: Throwable => Utils.tryWithSafeFinally { throw t } { close() } + } + } + } + + private[python] def bootstrapScript(script: String): String = { + "try:\n" + script.linesIterator.map(" " + _).mkString("\n") + + "\nexcept BaseException as _bootstrap_error:\n" + + " raise RuntimeError('In-process Python bootstrap failed: ' + " + + "ascii(type(_bootstrap_error).__name__ + ': ' + str(_bootstrap_error))) from None\n" + } + + def initialize(sitePackages: Seq[String] = Seq.empty): Unit = synchronized { + bootstrappedSitePackages.foreach { paths => + if (paths != sitePackages) { + throw new LifecycleException("In-process Python has already configured different " + + "sitePackages. Restart the executor process before changing interpreter configuration.") + } + } + if (active != null && !active.isTerminated) { + // The sitePackages were checked above: an active session bootstrapped them. + active.requireRunning() + } else { + InterpreterConfiguration.configure(sitePackages) + val candidate = new InterpreterSession( + sitePackages, configuredSitePackages.filterNot(sitePackages.contains)) + try { + candidate.initialize() + active = candidate + } catch { + case t: Throwable => Utils.tryWithSafeFinally { throw t } { candidate.shutdown() } + } + } + } + + def currentSession: InterpreterSession = synchronized { + checkState(active != null) + // `shutdown` keeps the stopped session, so a session that is not running was stopped. + checkState(active.isRunning, StoppedMessage) + active + } + + def shutdown(): Unit = { + val session = synchronized { active } + if (session != null) session.shutdown() + } + + private def checkState(running: Boolean): Unit = { + checkState(running, "In-process Python is not running; initialize the executor plugin first") + } + + private val StoppedMessage = + "In-process Python has been stopped (executor or SparkContext shutdown)" + + private def checkState(running: Boolean, message: String): Unit = { + if (!running) throw new IllegalStateException(message) + } + + /** + * Tasks retain this generation, so stale tasks cannot enter a later SparkContext's interpreter. + * Lifecycle operations only hold the monitor while enqueueing work, never while running Python. + */ + /** + * @param stalePaths include paths of earlier, failed JEP configurations, which are removed + * from sys.path unless `sitePackages` configures them again + */ + private[python] class InterpreterSession( + val sitePackages: Seq[String] = Seq.empty, + stalePaths: Seq[String] = Seq.empty) { + // CPython native calls need more stack than the usual JVM thread default. This is a + // platform-dependent size request, not protection against arbitrary native crashes. + private val executor = Executors.newSingleThreadExecutor(new ThreadFactory { + override def newThread(runnable: Runnable): Thread = { + val thread = new Thread(null, runnable, "inprocess-python", 8L * 1024 * 1024) + thread.setDaemon(true) + thread + } + }) + @volatile private var running = true + // Calls submitted to the interpreter thread that have not finished or been cancelled. + private val pendingCalls = new AtomicInteger() + // Accessed only on the owning thread. + private var interp: SharedInterpreter = _ + // Guarded by this session's monitor. Shutdown must keep Python-owned result buffers + // pinned until their tasks have released the JVM CDI references. + private val registeredHandles = mutable.Set.empty[String] + + def isRunning: Boolean = running + def isTerminated: Boolean = executor.isTerminated + + def requireRunning(): Unit = { + if (!isRunning) { + throw new LifecycleException("In-process Python is still stopping. Wait for outstanding " + + "native work to finish or replace the executor process before starting a new context.") + } + } + + private[python] def onInterpreterThread[T](body: => T): T = { + val context = Option(TaskContext.get()) + context.foreach(_.killTaskIfInterrupted()) + val gate = new Object + var started = false + var cancelled = false + val future = synchronized { + checkRunning() + pendingCalls.incrementAndGet() + executor.submit(new Callable[T] { + override def call(): T = { + gate.synchronized { + if (cancelled) throw new TaskKilledException("Cancelled before Python invocation") + started = true + } + try body finally pendingCalls.decrementAndGet() + } + }) + } + var interrupted = false + try { + while (true) { + val taskCancelled = context.exists(_.isInterrupted()) + if (interrupted || taskCancelled) { + val cancelledBeforeStart = gate.synchronized { + if (started) false else { + cancelled = true + future.cancel(false) + pendingCalls.decrementAndGet() + true + } + } + if (cancelledBeforeStart) { + context.foreach(_.killTaskIfInterrupted()) + throw new InterruptedException("Cancelled before Python invocation") + } + } + try { + val result = future.get(100, TimeUnit.MILLISECONDS) + context.foreach(_.killTaskIfInterrupted()) + return result + } catch { + case _: TimeoutException => + case _: InterruptedException => interrupted = true + case e: ExecutionException => throw e.getCause + } + } + throw new IllegalStateException("Unreachable") + } finally { + // Once native work starts, wait for it even after cancellation: the caller still owns + // CDI structs that Python may use. Pending work, however, is safe to cancel immediately. + if (interrupted) Thread.currentThread().interrupt() + } + } + + def initialize(): Unit = onInterpreterThread { + val candidate = new ManagedSharedInterpreter() + // SharedInterpreter keeps sys.modules and sys.path for the JVM lifetime, even when + // the following bootstrap fails. A new context cannot switch Python environments. + bootstrappedSitePackages = Some(sitePackages) + try { + candidate.set("_site_packages", sitePackages.asJava) + candidate.set("_stale_paths", stalePaths.asJava) + val sparkPaths = PythonUtils.mergePythonPaths( + PythonUtils.sparkPythonPath, sys.env.getOrElse("PYTHONPATH", "")) + .split(File.pathSeparator).filter(_.nonEmpty) + candidate.set("_spark_paths", sparkPaths.toSeq.asJava) + candidate.exec(bootstrapScript( + """import os, site, sys + |_configured = [os.path.abspath(p) for p in _site_packages] + |_stale = {os.path.abspath(p) for p in _stale_paths} - set(_configured) + |sys.path[:] = [p for p in sys.path if os.path.abspath(p) not in _stale] Review Comment: Thanks, done in 961ebc6 as suggested: `ManagedSharedInterpreter.configureInterpreter` saves `sys.path` before `super.configureInterpreter(config)` and restores it before `close()` when that fails, so only JEP's appended copies go away, and `configuredSitePackages`, `stalePaths` and the bootstrap filter are gone. The fresh-JVM test now also lists the stdlib directory, a default entry, in the failed configuration and checks that the retried interpreter still has it; it also still checks that a module only in the failed configuration's directory is not importable, and fails without the restore. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/python/InProcessPythonRuntime.scala: ########## @@ -0,0 +1,432 @@ +/* + * 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 java.io.File +import java.util.concurrent.{Callable, ExecutionException, Executors, ThreadFactory, TimeoutException, TimeUnit} +import java.util.concurrent.atomic.AtomicInteger + +import scala.collection.mutable +import scala.jdk.CollectionConverters._ + +import jep.{JepConfig, JepException, MainInterpreter, NamingConventionClassEnquirer, PyConfig, SharedInterpreter} +import org.apache.arrow.c.{ArrowSchema, Data} +import org.apache.arrow.vector.types.pojo.Field + +import org.apache.spark.{TaskContext, TaskKilledException} +import org.apache.spark.api.python.{PythonException, PythonUtils} +import org.apache.spark.internal.Logging +import org.apache.spark.internal.config.Python +import org.apache.spark.sql.util.ArrowUtils +import org.apache.spark.util.Utils + +/** Owns one interpreter generation per executor plugin lifecycle. */ +private[python] object InProcessPythonRuntime extends Logging { + private val TRACEBACK_SENTINEL = "__INPROCESS_UDF_TRACEBACK__:" + private var active: InterpreterSession = _ + private var mainConfigured = false + @volatile private var sharedConfigured = false + @volatile private var bootstrappedSitePackages: Option[Seq[String]] = None + // Every sitePackages given to JEP. JEP adds them to sys.path before it imports its package, + // so a failed configuration leaves them there, and the bootstrap of a retry removes them. + // Guarded by the object's monitor. + private var configuredSitePackages = Seq.empty[String] + + private[python] class LifecycleException(message: String) extends IllegalStateException(message) + + // Keep JEP references out of the singleton's verifier so currentSession can report + // an uninitialized runtime even when the provided JEP JAR is absent. + private[python] object InterpreterConfiguration { + def configure(sitePackages: Seq[String]): Unit = { + if (!mainConfigured) { + // Like Python workers, use a stable default hash seed on every executor. This must + // happen before JEP creates its process-wide main interpreter, including on restarts. + MainInterpreter.setInitParams( + PyConfig.isolated().setUseEnvironment(false).setHashSeed(0).setUseHashSeed(true)) + mainConfigured = true + } + if (!sharedConfigured) { + // JEP imports its Python package during construction, before our bootstrap runs. + SharedInterpreter.setConfig(interpreterConfig(sitePackages)) + configuredSitePackages = (configuredSitePackages ++ sitePackages).distinct + } + } + + def interpreterConfig(sitePackages: Seq[String]): JepConfig = { + require(sitePackages.forall(Python.isValidInProcessPath), Python.IN_PROCESS_PATH_RULE) + val config = new JepConfig().setClassEnquirer(new NamingConventionClassEnquirer(false)) + // Calling addIncludePaths with no arguments adds the working directory in JEP. + if (sitePackages.nonEmpty) config.addIncludePaths(sitePackages: _*) + config + } + } + + private class ManagedSharedInterpreter extends SharedInterpreter { + override protected def configureInterpreter(config: JepConfig): Unit = { + // JEP invokes this hook after native initialization, from its constructor. Close + // here if configuration fails, before the caller can receive an interpreter handle. + try { + super.configureInterpreter(config) Review Comment: Done in 961ebc6: the same hook calls `importlib.invalidate_caches()` before JEP's configuration, while `sharedConfigured` is false. The fresh-JVM test now lists a directory that does not exist when the first attempt fails, creates it as a link to JEP's directory, and retries with it; without the call, the retry fails with `No module named 'jep'`. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/python/InProcessArrowEvalPythonEvaluatorFactory.scala: ########## @@ -0,0 +1,598 @@ +/* + * 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 java.io.File +import java.nio.file.Files +import java.util.UUID +import java.util.concurrent.TimeUnit +import java.util.concurrent.atomic.{AtomicBoolean, AtomicInteger} +import java.util.concurrent.locks.ReentrantLock + +import scala.collection.mutable.ArrayBuffer +import scala.jdk.CollectionConverters._ + +import com.google.common.util.concurrent.Uninterruptibles +import org.apache.arrow.c.{ArrowArray, ArrowSchema, BaseStruct} +import org.apache.arrow.util.AutoCloseables +import org.apache.arrow.vector.{BaseLargeVariableWidthVector, BaseVariableWidthVector, DensityAwareVector, VectorSchemaRoot} + +import org.apache.spark.{SparkEnv, SparkException, TaskContext} +import org.apache.spark.api.python.ChainedPythonFunctions +import org.apache.spark.memory.MemoryConsumer +import org.apache.spark.serializer.SerializerManager +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, JoinedRow, PythonUDF, UnsafeProjection, UnsafeRow} +import org.apache.spark.sql.execution.arrow.ArrowWriter +import org.apache.spark.sql.execution.metric.SQLMetric +import org.apache.spark.sql.execution.python.EvalPythonExec.ArgumentMetadata +import org.apache.spark.sql.types._ +import org.apache.spark.sql.util.ArrowUtils +import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch, ColumnVector} +import org.apache.spark.util.Utils + +/** + * Evaluates scalar Python UDFs using Arrow CDI in the executor process. Only UDF arguments + * are converted to Arrow. Original rows are buffered in a spillable queue and joined with + * the results, unless all of them are UDF arguments that read back from Arrow unchanged. + * Each batch owns its Arrow buffers so Python can safely retain input arrays. + * + * The evaluator owns its queue, so that cleanup at task completion is coordinated with a + * consumer on another thread, such as a pipelined Python writer or a TRANSFORM feed thread. + */ +class InProcessArrowEvalPythonEvaluatorFactory( + childOutput: Seq[Attribute], + udfs: Seq[PythonUDF], + output: Seq[Attribute], + batchSize: Int, + maxBytes: Long, + timeZoneId: String, + largeVarTypes: Boolean, + hideTraceback: Boolean, + simplifiedTraceback: Boolean, + tracebackWithLocals: Boolean, + fullValidation: Boolean, + metrics: Map[String, SQLMetric]) + extends EvalPythonEvaluatorFactory(childOutput, udfs, output) { + + private[python] def runtimeSession: InProcessPythonRuntime.InterpreterSession = + InProcessPythonRuntime.currentSession + + /** Unused: `evaluateJoined` always evaluates the UDFs. */ + override protected def evaluate( + funcs: Seq[(ChainedPythonFunctions, Long)], + argMetas: Array[Array[ArgumentMetadata]], + rows: Iterator[InternalRow], + inputSchema: StructType, + context: TaskContext): Iterator[InternalRow] = + throw SparkException.internalError("In-process UDFs are evaluated with their input rows") + + override protected def evaluateJoined( + funcs: Seq[(ChainedPythonFunctions, Long)], + argMetas: Array[Array[ArgumentMetadata]], + rows: Iterator[InternalRow], + inputs: Seq[Expression], + inputSchema: StructType, + context: TaskContext): Option[Iterator[InternalRow]] = { + import InProcessArrowEvalPythonEvaluatorFactory.{Buffered, ReadBack, readsBack} + val inputColumns = inputs.length == childOutput.length && inputs.zip(childOutput).forall { + case (a: Attribute, c) => a.exprId == c.exprId + case _ => false + } + // If all input columns are UDF arguments, they are written to Arrow regardless. Read them + // back from the exported input vectors instead of buffering every input row, if their + // values read back from Arrow exactly as written and as fast as an unsafe row copy. + val joinInput = if (inputColumns && inputSchema.forall(f => readsBack(f.dataType))) { + ReadBack + } else if (inputColumns) { + Buffered(None) + } else { + // Each projected row is written to Arrow before the next input row is pulled, so the + // arguments go into a reused buffer rather than being copied value by value. + val projection = UnsafeProjection.create(inputs, childOutput) + projection.initialize(context.partitionId()) + Buffered(Some(projection)) + } + Some(evaluateBatches(funcs, argMetas, rows, inputSchema, context, joinInput)) + } + + private[python] def evaluateBatches( + funcs: Seq[(ChainedPythonFunctions, Long)], + argMetas: Array[Array[ArgumentMetadata]], + rows: Iterator[InternalRow], + inputSchema: StructType, + context: TaskContext, + joinInput: InProcessArrowEvalPythonEvaluatorFactory.JoinInput): Iterator[InternalRow] = { + import InProcessArrowEvalPythonEvaluatorFactory.{Buffered, ReadBack} + ArrowUtils.failDuplicatedFieldNames(inputSchema) + val functions = funcs.map { case (chain, _) => + if (chain.funcs.size != 1) { + throw SparkException.internalError( + "In-process UDF chains must use separate evaluation nodes") + } + chain.funcs.head + } + val inputOrdinals = argMetas.map(_.map(_.offset)) + def checkCancellation(): Unit = context.killTaskIfInterrupted() + + val expectedFields = udfs.map { udf => + ArrowUtils.toArrowField("result", udf.dataType, true, timeZoneId, largeVarTypes) + } + val arrowSchema = ArrowUtils.toArrowSchema(inputSchema, timeZoneId, largeVarTypes) + // Capture before consuming input: an old task must never join a later context's session. + val runtime = runtimeSession + // Rows are copied out of the queue and Arrow vectors before they are returned, so they + // remain valid after task completion releases those, on whichever thread consumes them. + val resultProj = UnsafeProjection.create(output, output) + // Spill files go into a directory of the queue's own, created with the first disk queue, so + // that task completion can delete them when it cannot close the queue. + @volatile var spillDir: File = null + // Guarded by the queue's monitor. + var queueAbandoned = false + // Only the consumer holding the iterator's lock adds and removes rows. + class InProcessRowQueue(localDir: File, serializerManager: SerializerManager) + extends HybridRowQueue(context.taskMemoryManager(), localDir, childOutput.length, + serializerManager, lockFree = true) { + + override protected def createDiskQueue(): RowQueue = synchronized { + if (spillDir == null) { + spillDir = Files.createTempDirectory(localDir.toPath, "inprocess-udf-").toFile + } + DiskRowQueue(Files.createTempFile(spillDir.toPath, "buffer", "").toFile, + childOutput.length, serializerManager) + } + + // Once task completion leaves the queue to the executor, it must not spill for other + // consumers into a directory that nothing deletes. + override def spill(size: Long, trigger: MemoryConsumer): Long = synchronized { + if (queueAbandoned) 0L else super.spill(size, trigger) + } + + /** + * Leaves the queue to the executor: closes its disk queues, which hold no task memory + * but open streams, and deletes their directory. The consumer no longer uses the + * queue, and its last changes are visible, since it marked a read afterwards. + */ + def abandon(): Unit = synchronized { + queueAbandoned = true + Utils.tryWithSafeFinally { + (queues.asScala ++ Option(reading)).filterNot(isInMemoryQueue).foreach { q => + Utils.tryLogNonFatalError(q.close()) + } + } { + Utils.deleteQuietly(spillDir) + } + } + } + val (queue, projection) = joinInput match { + case Buffered(projection) => + val queue = new InProcessRowQueue( + new File(Utils.getLocalDir(SparkEnv.get.conf)), SparkEnv.get.serializerManager) + (queue, projection.orNull) + case ReadBack => (null, null) + } + val joined = new JoinedRow + val handles = functions.map(_ => UUID.randomUUID().toString) + var registered = false + var writer: ArrowWriter = null + val results = ArrayBuffer.empty[ArrowColumnVector] + + def closeBatch(): Unit = { + val resources = ArrayBuffer.empty[AutoCloseable] + resources ++= results + results.clear() + if (writer != null) { + resources += writer.root + writer = null + } + AutoCloseables.close(resources.asJava) + } + + val resources = new InProcessArrowEvalPythonEvaluatorFactory.IteratorResources( + // Closing the queue deletes the spill files it tracks; deleteQuietly also removes any + // other, without starting a process or throwing, also on an interrupted thread. + releaseTaskMemory = () => if (queue != null) { + Utils.tryWithSafeFinally(queue.close())(Utils.deleteQuietly(spillDir)) + }, + abandonTaskMemory = () => if (queue != null) queue.abandon(), + recordTime = nanos => metrics("pythonTotalTime") += nanos, + releaseOthers = () => { + Utils.tryWithSafeFinally { + closeBatch() + } { + if (registered) runtime.release(handles) + } + }) + + context.addTaskCompletionListener[Unit](_ => resources.close()) + + new Iterator[InternalRow] { + private var batchIter: Iterator[InternalRow] = Iterator.empty + + private def endOfInput: Nothing = + throw new NoSuchElementException("End of in-process UDF input") + + // Releases the resources on failure without replacing its exception. + private def fail(t: Throwable): Nothing = + Utils.tryWithSafeFinally { throw t } { resources.close() } + + // Called with the lock held. + private def hasNextLocked: Boolean = { + resources.start() + checkCancellation() + val available = batchIter.hasNext || { + var open = false + val more = try { + resources.startReadingInput() && rows.hasNext + } finally { + open = resources.endReadingInput() + } + more && open + } + if (!available) resources.close() + available + } + + // Each call takes the lock without allocating a closure per row. Within a batch, only + // the consumer advances `batchIter`, whose `hasNext` compares row indexes. + override def hasNext: Boolean = { + if (batchIter.hasNext && !resources.isClosed) return true + if (!resources.enter()) return false + val available = try { + try hasNextLocked catch { case t: Throwable => fail(t) } + } finally { + resources.exit() + } + // Task completion may have closed the iterator while the input was being read. + available && !resources.isClosed + } + + override def next(): InternalRow = { + if (!resources.enter()) endOfInput + try { + try { + if (!hasNextLocked) endOfInput + if (!batchIter.hasNext) nextBatch() + val result = batchIter.next() + resultProj(if (queue != null) joined(queue.remove(), result) else result) + } catch { + case t: Throwable => fail(t) + } + } finally { + resources.exit() + } + } + + // Runs Python without the lock unless task completion already happened, adds the + // nanoseconds it took to `metric`, and ends the input if task completion happened + // meanwhile. + private def python(metric: SQLMetric)(body: => Long): Unit = { + if (resources.isClosed) endOfInput + metric += resources.withoutLock(body) + if (resources.isClosed) endOfInput + } + + /** + * Writes the next input row to the batch, returning false at the end of input or once + * task completion happened. If it happens while the row is read, the row is dropped. + */ + private def pullRow(): Boolean = { + var open = false + val row = try { + // Checked again after `hasNext`, which may wait for input, not to read another row. + if (resources.startReadingInput() && rows.hasNext && !resources.isClosed) { + rows.next() + } else { + null + } + } finally { + open = resources.endReadingInput() + } + if (!open) endOfInput + if (row == null) return false + if (queue != null) queue.add(row.asInstanceOf[UnsafeRow]) + writer.write(if (projection != null) projection(row) else row) + true + } + + // Called with the lock held. + private def nextBatch(): Unit = { + // Each batch needs buffers of its own, which Python may retain, but sizing them like + // the previous batch's avoids growing them from Arrow's defaults row by row. + val capacities = if (writer == null) Seq.empty else { + writer.root.getFieldVectors.asScala.map { v => + val density = v match { + case v: BaseVariableWidthVector => v.getDensity + case v: BaseLargeVariableWidthVector => v.getDensity + case _ => 0.0 + } + (v.getValueCount, density) + }.toSeq + } + closeBatch() + val root = VectorSchemaRoot.create(arrowSchema, ArrowUtils.rootAllocator) + capacities.zip(root.getFieldVectors.asScala).foreach { case ((count, density), v) => + if (count > 0) v match { + case v: DensityAwareVector if density > 0 => v.setInitialCapacity(count, density) + case v => v.setInitialCapacity(count) Review Comment: Done in 961ebc6 with the narrower option: only `BaseFixedWidthVector` and the two variable-width bases are sized from the previous batch, and struct, list and map vectors keep Arrow's defaults, as before round 21. With no record limit and 40M empty `array<string>` values, the peak off-heap allocation went from 2039 MiB at b33b115 to 260 MiB, at the same time per query; `array<int>` with 0 to 3 elements at 100,000 rows per batch is unchanged (0.438 vs 0.445 s). ########## sql/core/src/main/scala/org/apache/spark/sql/execution/python/InProcessArrowEvalPythonEvaluatorFactory.scala: ########## @@ -0,0 +1,598 @@ +/* + * 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 java.io.File +import java.nio.file.Files +import java.util.UUID +import java.util.concurrent.TimeUnit +import java.util.concurrent.atomic.{AtomicBoolean, AtomicInteger} +import java.util.concurrent.locks.ReentrantLock + +import scala.collection.mutable.ArrayBuffer +import scala.jdk.CollectionConverters._ + +import com.google.common.util.concurrent.Uninterruptibles +import org.apache.arrow.c.{ArrowArray, ArrowSchema, BaseStruct} +import org.apache.arrow.util.AutoCloseables +import org.apache.arrow.vector.{BaseLargeVariableWidthVector, BaseVariableWidthVector, DensityAwareVector, VectorSchemaRoot} + +import org.apache.spark.{SparkEnv, SparkException, TaskContext} +import org.apache.spark.api.python.ChainedPythonFunctions +import org.apache.spark.memory.MemoryConsumer +import org.apache.spark.serializer.SerializerManager +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, JoinedRow, PythonUDF, UnsafeProjection, UnsafeRow} +import org.apache.spark.sql.execution.arrow.ArrowWriter +import org.apache.spark.sql.execution.metric.SQLMetric +import org.apache.spark.sql.execution.python.EvalPythonExec.ArgumentMetadata +import org.apache.spark.sql.types._ +import org.apache.spark.sql.util.ArrowUtils +import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch, ColumnVector} +import org.apache.spark.util.Utils + +/** + * Evaluates scalar Python UDFs using Arrow CDI in the executor process. Only UDF arguments + * are converted to Arrow. Original rows are buffered in a spillable queue and joined with + * the results, unless all of them are UDF arguments that read back from Arrow unchanged. + * Each batch owns its Arrow buffers so Python can safely retain input arrays. + * + * The evaluator owns its queue, so that cleanup at task completion is coordinated with a + * consumer on another thread, such as a pipelined Python writer or a TRANSFORM feed thread. + */ +class InProcessArrowEvalPythonEvaluatorFactory( + childOutput: Seq[Attribute], + udfs: Seq[PythonUDF], + output: Seq[Attribute], + batchSize: Int, + maxBytes: Long, + timeZoneId: String, + largeVarTypes: Boolean, + hideTraceback: Boolean, + simplifiedTraceback: Boolean, + tracebackWithLocals: Boolean, + fullValidation: Boolean, + metrics: Map[String, SQLMetric]) + extends EvalPythonEvaluatorFactory(childOutput, udfs, output) { + + private[python] def runtimeSession: InProcessPythonRuntime.InterpreterSession = + InProcessPythonRuntime.currentSession + + /** Unused: `evaluateJoined` always evaluates the UDFs. */ + override protected def evaluate( + funcs: Seq[(ChainedPythonFunctions, Long)], + argMetas: Array[Array[ArgumentMetadata]], + rows: Iterator[InternalRow], + inputSchema: StructType, + context: TaskContext): Iterator[InternalRow] = + throw SparkException.internalError("In-process UDFs are evaluated with their input rows") + + override protected def evaluateJoined( + funcs: Seq[(ChainedPythonFunctions, Long)], + argMetas: Array[Array[ArgumentMetadata]], + rows: Iterator[InternalRow], + inputs: Seq[Expression], + inputSchema: StructType, + context: TaskContext): Option[Iterator[InternalRow]] = { + import InProcessArrowEvalPythonEvaluatorFactory.{Buffered, ReadBack, readsBack} + val inputColumns = inputs.length == childOutput.length && inputs.zip(childOutput).forall { + case (a: Attribute, c) => a.exprId == c.exprId + case _ => false + } + // If all input columns are UDF arguments, they are written to Arrow regardless. Read them + // back from the exported input vectors instead of buffering every input row, if their + // values read back from Arrow exactly as written and as fast as an unsafe row copy. + val joinInput = if (inputColumns && inputSchema.forall(f => readsBack(f.dataType))) { + ReadBack + } else if (inputColumns) { + Buffered(None) + } else { + // Each projected row is written to Arrow before the next input row is pulled, so the + // arguments go into a reused buffer rather than being copied value by value. + val projection = UnsafeProjection.create(inputs, childOutput) + projection.initialize(context.partitionId()) + Buffered(Some(projection)) + } + Some(evaluateBatches(funcs, argMetas, rows, inputSchema, context, joinInput)) + } + + private[python] def evaluateBatches( + funcs: Seq[(ChainedPythonFunctions, Long)], + argMetas: Array[Array[ArgumentMetadata]], + rows: Iterator[InternalRow], + inputSchema: StructType, + context: TaskContext, + joinInput: InProcessArrowEvalPythonEvaluatorFactory.JoinInput): Iterator[InternalRow] = { + import InProcessArrowEvalPythonEvaluatorFactory.{Buffered, ReadBack} + ArrowUtils.failDuplicatedFieldNames(inputSchema) + val functions = funcs.map { case (chain, _) => + if (chain.funcs.size != 1) { + throw SparkException.internalError( + "In-process UDF chains must use separate evaluation nodes") + } + chain.funcs.head + } + val inputOrdinals = argMetas.map(_.map(_.offset)) + def checkCancellation(): Unit = context.killTaskIfInterrupted() + + val expectedFields = udfs.map { udf => + ArrowUtils.toArrowField("result", udf.dataType, true, timeZoneId, largeVarTypes) + } + val arrowSchema = ArrowUtils.toArrowSchema(inputSchema, timeZoneId, largeVarTypes) + // Capture before consuming input: an old task must never join a later context's session. + val runtime = runtimeSession + // Rows are copied out of the queue and Arrow vectors before they are returned, so they + // remain valid after task completion releases those, on whichever thread consumes them. + val resultProj = UnsafeProjection.create(output, output) + // Spill files go into a directory of the queue's own, created with the first disk queue, so + // that task completion can delete them when it cannot close the queue. + @volatile var spillDir: File = null + // Guarded by the queue's monitor. + var queueAbandoned = false + // Only the consumer holding the iterator's lock adds and removes rows. + class InProcessRowQueue(localDir: File, serializerManager: SerializerManager) + extends HybridRowQueue(context.taskMemoryManager(), localDir, childOutput.length, + serializerManager, lockFree = true) { + + override protected def createDiskQueue(): RowQueue = synchronized { + if (spillDir == null) { + spillDir = Files.createTempDirectory(localDir.toPath, "inprocess-udf-").toFile + } + DiskRowQueue(Files.createTempFile(spillDir.toPath, "buffer", "").toFile, + childOutput.length, serializerManager) + } + + // Once task completion leaves the queue to the executor, it must not spill for other + // consumers into a directory that nothing deletes. + override def spill(size: Long, trigger: MemoryConsumer): Long = synchronized { + if (queueAbandoned) 0L else super.spill(size, trigger) + } + + /** + * Leaves the queue to the executor: closes its disk queues, which hold no task memory + * but open streams, and deletes their directory. The consumer no longer uses the + * queue, and its last changes are visible, since it marked a read afterwards. + */ + def abandon(): Unit = synchronized { + queueAbandoned = true + Utils.tryWithSafeFinally { + (queues.asScala ++ Option(reading)).filterNot(isInMemoryQueue).foreach { q => + Utils.tryLogNonFatalError(q.close()) + } + } { + Utils.deleteQuietly(spillDir) + } + } + } + val (queue, projection) = joinInput match { + case Buffered(projection) => + val queue = new InProcessRowQueue( + new File(Utils.getLocalDir(SparkEnv.get.conf)), SparkEnv.get.serializerManager) + (queue, projection.orNull) + case ReadBack => (null, null) + } + val joined = new JoinedRow + val handles = functions.map(_ => UUID.randomUUID().toString) + var registered = false + var writer: ArrowWriter = null + val results = ArrayBuffer.empty[ArrowColumnVector] + + def closeBatch(): Unit = { + val resources = ArrayBuffer.empty[AutoCloseable] + resources ++= results + results.clear() + if (writer != null) { + resources += writer.root + writer = null + } + AutoCloseables.close(resources.asJava) + } + + val resources = new InProcessArrowEvalPythonEvaluatorFactory.IteratorResources( + // Closing the queue deletes the spill files it tracks; deleteQuietly also removes any + // other, without starting a process or throwing, also on an interrupted thread. + releaseTaskMemory = () => if (queue != null) { + Utils.tryWithSafeFinally(queue.close())(Utils.deleteQuietly(spillDir)) + }, + abandonTaskMemory = () => if (queue != null) queue.abandon(), + recordTime = nanos => metrics("pythonTotalTime") += nanos, + releaseOthers = () => { + Utils.tryWithSafeFinally { + closeBatch() + } { + if (registered) runtime.release(handles) + } + }) + + context.addTaskCompletionListener[Unit](_ => resources.close()) + + new Iterator[InternalRow] { + private var batchIter: Iterator[InternalRow] = Iterator.empty + + private def endOfInput: Nothing = + throw new NoSuchElementException("End of in-process UDF input") + + // Releases the resources on failure without replacing its exception. + private def fail(t: Throwable): Nothing = + Utils.tryWithSafeFinally { throw t } { resources.close() } + + // Called with the lock held. + private def hasNextLocked: Boolean = { + resources.start() + checkCancellation() + val available = batchIter.hasNext || { + var open = false + val more = try { + resources.startReadingInput() && rows.hasNext + } finally { + open = resources.endReadingInput() + } + more && open + } + if (!available) resources.close() + available + } + + // Each call takes the lock without allocating a closure per row. Within a batch, only + // the consumer advances `batchIter`, whose `hasNext` compares row indexes. + override def hasNext: Boolean = { + if (batchIter.hasNext && !resources.isClosed) return true + if (!resources.enter()) return false + val available = try { + try hasNextLocked catch { case t: Throwable => fail(t) } + } finally { + resources.exit() + } + // Task completion may have closed the iterator while the input was being read. + available && !resources.isClosed + } + + override def next(): InternalRow = { + if (!resources.enter()) endOfInput + try { + try { + if (!hasNextLocked) endOfInput + if (!batchIter.hasNext) nextBatch() + val result = batchIter.next() + resultProj(if (queue != null) joined(queue.remove(), result) else result) + } catch { + case t: Throwable => fail(t) + } + } finally { + resources.exit() + } + } + + // Runs Python without the lock unless task completion already happened, adds the + // nanoseconds it took to `metric`, and ends the input if task completion happened + // meanwhile. + private def python(metric: SQLMetric)(body: => Long): Unit = { + if (resources.isClosed) endOfInput + metric += resources.withoutLock(body) + if (resources.isClosed) endOfInput + } + + /** + * Writes the next input row to the batch, returning false at the end of input or once + * task completion happened. If it happens while the row is read, the row is dropped. + */ + private def pullRow(): Boolean = { + var open = false + val row = try { + // Checked again after `hasNext`, which may wait for input, not to read another row. + if (resources.startReadingInput() && rows.hasNext && !resources.isClosed) { + rows.next() + } else { + null + } + } finally { + open = resources.endReadingInput() + } + if (!open) endOfInput + if (row == null) return false + if (queue != null) queue.add(row.asInstanceOf[UnsafeRow]) + writer.write(if (projection != null) projection(row) else row) + true + } + + // Called with the lock held. + private def nextBatch(): Unit = { + // Each batch needs buffers of its own, which Python may retain, but sizing them like + // the previous batch's avoids growing them from Arrow's defaults row by row. + val capacities = if (writer == null) Seq.empty else { + writer.root.getFieldVectors.asScala.map { v => + val density = v match { + case v: BaseVariableWidthVector => v.getDensity + case v: BaseLargeVariableWidthVector => v.getDensity + case _ => 0.0 + } + (v.getValueCount, density) + }.toSeq + } + closeBatch() + val root = VectorSchemaRoot.create(arrowSchema, ArrowUtils.rootAllocator) + capacities.zip(root.getFieldVectors.asScala).foreach { case ((count, density), v) => + if (count > 0) v match { + case v: DensityAwareVector if density > 0 => v.setInitialCapacity(count, density) Review Comment: Done in 961ebc6: the data request is rounded up to the next power of two, capped at `Int.MaxValue` for 32-bit offsets. Since `setInitialCapacity` truncates `count * density`, the density carries an extra half byte; "string data is pre-sized to the next power of two of the previous batch" checks the resulting sizes, including the cap. With strings of 1,000 to 2,999 bytes, the peak off-heap allocation went from 256 MiB to 192 MiB against b33b115 (0.508 vs 0.526 s per query), and the fixed-length long-string, short-string and bigint benchmarks are unchanged. ########## python/pyspark/sql/tests/test_inprocess_udf.py: ########## @@ -0,0 +1,1992 @@ +# +# 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 tests for in-process Python UDFs. + +Run with python/run-tests like other SQL tests. JEP paths are discovered from the +selected Python environment before the Spark JVM starts. ARROW_C_DATA_JAR must +point to the provided Arrow CDI JAR. Set INPROCESS_TESTS=1 +to require the suite (missing dependencies then fail), or 0 to disable it. +Otherwise, the suite runs when JEP, PyArrow and the CDI JAR are available. +""" + +import os +import shutil +import tempfile +import unittest +import zipfile +from importlib.util import find_spec +from pathlib import Path +from unittest.mock import patch + +from pyspark.testing.sqlutils import ReusedSQLTestCase +from pyspark.testing.utils import eventually + +_jep_spec = find_spec("jep") +_cdi_jar = os.environ.get("ARROW_C_DATA_JAR") +_test_mode = os.environ.get("INPROCESS_TESTS") +_run_inprocess = _test_mode == "1" or ( + _test_mode != "0" + and _jep_spec is not None + and find_spec("pyarrow") is not None + and _cdi_jar is not None + and Path(_cdi_jar).is_file() +) + + [email protected](_run_inprocess, "Requires JEP, PyArrow and ARROW_C_DATA_JAR") +class InProcessUDFTests(ReusedSQLTestCase): + """ + End-to-end tests for @inprocess_udf that require jep + CPython + PyArrow. + + The plugin initializes JEP before any task starts. Calls from task threads and + shutdown must use the same dedicated interpreter thread. + """ + + @classmethod + def master(cls): + return "local[2]" + + @classmethod + def conf(cls): + return ( + super() + .conf() + .set("spark.task.cpus", "0.5") + .set("spark.driver.extraClassPath", os.pathsep.join([str(cls.jep_jar), cls.cdi_jar])) + .set("spark.driver.extraLibraryPath", str(cls.jep_dir)) + .set( + "spark.python.inProcess.sitePackages", + ",".join([cls.site_packages, str(cls.jep_dir.parent)]), + ) + .set("spark.plugins", "org.apache.spark.sql.execution.python.InProcessPythonPlugin") + ) + + @classmethod + def setUpClass(cls): + if _jep_spec is None: + raise RuntimeError("INPROCESS_TESTS=1 requires JEP in the selected Python environment") + # Do not import jep: it can only be imported by an embedded interpreter. + cls.jep_dir = Path(_jep_spec.origin).parent + jars = list(cls.jep_dir.glob("jep-*.jar")) + if len(jars) != 1: + raise RuntimeError(f"Expected one JEP JAR in {cls.jep_dir}, found {len(jars)}") + cls.jep_jar = jars[0] + if not _cdi_jar or not Path(_cdi_jar).is_file(): + raise RuntimeError("Set ARROW_C_DATA_JAR to the provided Arrow CDI JAR") + cls.cdi_jar = str(Path(_cdi_jar).resolve()) + cls.site_packages = tempfile.mkdtemp() + helper_dir = os.path.join(cls.site_packages, "extra") + system_dir = os.path.join(cls.site_packages, "system") + os.mkdir(helper_dir) + os.mkdir(system_dir) + with open(os.path.join(system_dir, "_inprocess_process_helper.py"), "w") as f: + f.write("MAGIC = -1\n") + with open(os.path.join(helper_dir, "_inprocess_test_helper.py"), "w") as f: + f.write("MAGIC = 99\n") + with open(os.path.join(cls.site_packages, "helper.pth"), "w") as f: + f.write("extra\n") + for name in ["spire", "redis"]: + package = Path(cls.site_packages) / name + package.mkdir() + (package / "__init__.py").write_text("PYTHON_PACKAGE = True\n") + shadow = os.path.join(cls.site_packages, "pyspark") + os.mkdir(shadow) + with open(os.path.join(shadow, "__init__.py"), "w") as f: + f.write("raise RuntimeError('site-packages must not shadow Spark PySpark')\n") + try: + # CI extracts compiled targets without building pyspark.zip. Keep the source + # tree available, while still requiring the configured paths to provide JEP. + cls.python_source = Path(__file__).resolve().parents[3] + archive = cls.python_source / "lib" / "pyspark.zip" + if archive.is_file(): + with zipfile.ZipFile(archive) as packaged: + for module in (cls.python_source / "pyspark" / "inprocess").glob("*.py"): + name = module.relative_to(cls.python_source).as_posix() + if ( + name not in packaged.namelist() + or packaged.read(name) != module.read_bytes() + ): + raise RuntimeError("Rebuild or remove stale python/lib/pyspark.zip") + python_path = os.pathsep.join([str(cls.python_source), system_dir]) + with patch.dict(os.environ, {"PYTHONPATH": python_path}): + super().setUpClass() + except Exception: + shutil.rmtree(cls.site_packages) + raise + + @classmethod + def tearDownClass(cls): + try: + super().tearDownClass() + finally: + shutil.rmtree(cls.site_packages) + + def test_driver_defined_udt_return_type(self): + import sys + + from pyspark.inprocess import inprocess_udf + from pyspark.sql.types import LongType, UserDefinedType + + class DriverUDT(UserDefinedType): + @classmethod + def sqlType(cls): + return LongType() + + @classmethod + def module(cls): + return "__main__" + + def serialize(self, value): + return value + + def deserialize(self, value): + return value + + DriverUDT.__module__ = "__main__" + with patch.object(sys.modules["__main__"], "DriverUDT", DriverUDT, create=True): + identity = inprocess_udf(DriverUDT())(lambda x: x) + result = self.spark.range(3).select(identity("id")) + self.assertEqual([r[0] for r in result.collect()], [0, 1, 2]) + + def test_unsupported_ddl_return_type_fails_on_driver(self): + from pyspark.errors import PySparkNotImplementedError + from pyspark.inprocess import inprocess_udf + + wrapper = inprocess_udf("interval year to month")(lambda x: x) + with patch("pyspark.inprocess.udf._serialize_udf") as serialize: + with self.assertRaises(PySparkNotImplementedError): + wrapper("id") + serialize.assert_not_called() + self.assertIsNone(wrapper._jfunc) + + def test_worker_environment_is_rejected(self): + from pyspark.inprocess import inprocess_udf + + identity = inprocess_udf("long")(lambda x: x) + column = identity("id") + with self.sql_conf({"spark.pythonWorkerEnv.INPROCESS_TEST_VALUE": "value"}): + # Columns are session-agnostic; validate using the query's session at execution. + identity("id") + with self.assertRaisesRegex(Exception, "UNSUPPORTED_IN_PROCESS_PYTHON_UDF"): + self.spark.range(1).select(column).collect() + + def test_profiler_and_memory_settings_are_rejected(self): + from pyspark.inprocess import inprocess_udf + + identity = inprocess_udf("long")(lambda x: x) + column = identity("id") + with self.sql_conf({"spark.sql.pyspark.udf.profiler": "perf"}): + # Columns are session-agnostic; validate using the query's session at execution. + identity("id") + with self.assertRaisesRegex(Exception, "UNSUPPORTED_IN_PROCESS_PYTHON_UDF"): + self.spark.range(1).select(column).collect() + conf = self.spark.sparkContext._jvm.org.apache.spark.SparkEnv.get().conf() + key = "spark.executor.pyspark.memory" + try: + conf.set(key, "128m") + # Columns are session-agnostic; validate using the query's session at execution. + identity("id") + with self.assertRaisesRegex(Exception, "UNSUPPORTED_IN_PROCESS_PYTHON_UDF"): + self.spark.range(1).select(column).collect() + finally: + conf.remove(key) + + def test_configuration_uses_the_query_session(self): + from pyspark.inprocess import inprocess_udf + + identity = inprocess_udf("long")(lambda x: x) + second = self.spark.newSession() + with self.sql_conf({"spark.sql.pyspark.udf.profiler": "perf"}): + column = identity("id") + self.assertEqual([r[0] for r in second.range(2).select(column).collect()], [0, 1]) + with self.assertRaisesRegex(Exception, "UNSUPPORTED_IN_PROCESS_PYTHON_UDF"): + self.spark.range(2).select(column).collect() + conf = self.spark.sparkContext._jvm.org.apache.spark.SparkEnv.get().conf() + try: + conf.set("spark.executor.pyspark.memory", "0") + self.assertEqual(self.spark.range(1).select(identity("id")).first()[0], 0) + finally: + conf.remove("spark.executor.pyspark.memory") + + def test_python_packages_are_not_shadowed_by_java_imports(self): + from pyspark.inprocess import inprocess_udf + + def probe(x): + import sys + + import pyarrow as pa + import redis + import spire + + good = spire.PYTHON_PACKAGE and redis.PYTHON_PACKAGE + good = good and sys.stdout.line_buffering and sys.stdout.write_through + return pa.array([good] * len(x)) + + self.assertTrue( + self.spark.range(1).select(inprocess_udf("boolean")(probe)("id")).first()[0] + ) + + def test_result_schema_adapts_to_session_representation(self): + from pyspark.inprocess import inprocess_udf + + def timestamp(x): + import pyarrow as pa + + return pa.array([0] * len(x), type=pa.timestamp("us", tz="UTC")) + + def nested(x): + import pyarrow as pa + + # Declare the field order: newer PyArrow versions sort inferred struct fields. + return pa.array( + [{"s": ["hello"], "b": b"data"}] * len(x), + pa.struct([("s", pa.list_(pa.string())), ("b", pa.binary())]), + ) + + for zone in ["UTC", "Etc/UTC", "America/Los_Angeles"]: + with self.sql_conf({"spark.sql.session.timeZone": zone}): + result = self.spark.range(1).select(inprocess_udf("timestamp")(timestamp)("id")) + self.assertEqual( + result.toDF("value").selectExpr("unix_micros(value)").first()[0], 0 + ) + with self.sql_conf({"spark.sql.execution.arrow.useLargeVarTypes": "true"}): + result = ( + self.spark.range(1) + .select(inprocess_udf("struct<s:array<string>,b:binary>")(nested)("id")) + .first()[0] + ) + self.assertEqual(result.s, ["hello"]) + self.assertEqual(result.b, b"data") + + def test_kwargs_only_function(self): + from pyspark.inprocess import inprocess_udf + + identity = inprocess_udf("long")(lambda **cols: cols["x"]) + self.assertEqual( + [row[0] for row in self.spark.range(2).select(identity(x="id")).collect()], [0, 1] + ) + + def test_isolated_interpreter_and_explicit_process_pythonpath(self): + from pyspark.inprocess import inprocess_udf + + def flags(x): + import faulthandler + import sys + + import _inprocess_process_helper + import pyarrow as pa + + value = ( + sys.flags.isolated == 1 + and sys.flags.ignore_environment == 1 + and not faulthandler.is_enabled() + and _inprocess_process_helper.MAGIC == -1 + ) + return pa.array([value] * len(x)) + + self.assertTrue( + self.spark.range(1).select(inprocess_udf("boolean")(flags)("id")).first()[0] + ) + + def test_missing_cdi_dependency_fails_plugin_startup(self): + import subprocess + + jvm = self.spark.sparkContext._jvm + java_home = jvm.java.lang.System.getProperty("java.home") + classpath = os.pathsep.join( + p + for p in jvm.java.lang.System.getProperty("java.class.path").split(os.pathsep) + if p != self.cdi_jar + ) + source = """ +import java.lang.reflect.Proxy; +import java.util.Collections; +import org.apache.spark.SparkConf; +import org.apache.spark.api.plugin.PluginContext; +import org.apache.spark.sql.execution.python.InProcessPythonPlugin; + +class MissingCdiProbe { + public static void main(String[] args) { + PluginContext context = (PluginContext) Proxy.newProxyInstance( + PluginContext.class.getClassLoader(), new Class<?>[] {PluginContext.class}, + (proxy, method, values) -> new SparkConf(false)); + try { + new InProcessPythonPlugin().executorPlugin().init(context, Collections.emptyMap()); + throw new AssertionError("Plugin accepted a missing CDI dependency"); + } catch (IllegalStateException expected) { + if (!(expected.getCause() instanceof NoClassDefFoundError)) throw expected; + if (!expected.getMessage().contains("arrow-c-data.jar")) throw expected; + System.out.println("MISSING_CDI_REJECTED"); + } + } +} +""" + with tempfile.TemporaryDirectory() as directory: + source_file = Path(directory) / "MissingCdiProbe.java" + source_file.write_text(source) + result = subprocess.run( + [str(Path(java_home) / "bin" / "java"), "-cp", classpath, str(source_file)], + capture_output=True, + text=True, + timeout=90, + ) + self.assertEqual(result.returncode, 0, result.stdout + result.stderr) + self.assertIn("MISSING_CDI_REJECTED", result.stdout) + + def test_missing_jep_has_an_initialization_hint(self): + import subprocess + + jvm = self.spark.sparkContext._jvm + java_home = jvm.java.lang.System.getProperty("java.home") + classpath = os.pathsep.join( + p + for p in jvm.java.lang.System.getProperty("java.class.path").split(os.pathsep) + if p != str(self.jep_jar) + ) + source = """ +import java.lang.reflect.Proxy; +import java.util.Collections; +import org.apache.spark.SparkConf; +import org.apache.spark.api.plugin.PluginContext; +import org.apache.spark.sql.execution.python.InProcessPythonPlugin; +import org.apache.spark.sql.execution.python.InProcessPythonRuntime; + +class MissingJepProbe { + public static void main(String[] args) { + try { + InProcessPythonRuntime.currentSession(); + throw new AssertionError("Uninitialized runtime was accepted"); + } catch (IllegalStateException expected) { + if (!expected.getMessage().contains("executor plugin")) throw expected; + } + PluginContext context = (PluginContext) Proxy.newProxyInstance( + PluginContext.class.getClassLoader(), new Class<?>[] {PluginContext.class}, + (proxy, method, values) -> new SparkConf(false)); + try { + new InProcessPythonPlugin().executorPlugin().init(context, Collections.emptyMap()); + throw new AssertionError("Plugin accepted a missing JEP dependency"); + } catch (IllegalStateException expected) { + if (!(expected.getCause() instanceof LinkageError)) throw expected; + if (!expected.getMessage().contains("jep.jar")) throw expected; + System.out.println("MISSING_JEP_REJECTED"); + } + } +} +""" + with tempfile.TemporaryDirectory() as directory: + source_file = Path(directory) / "MissingJepProbe.java" + source_file.write_text(source) + result = subprocess.run( + [ + str(Path(java_home) / "bin" / "java"), + "--add-opens=java.base/java.nio=ALL-UNNAMED", + "--add-opens=java.base/sun.nio.ch=ALL-UNNAMED", + "-Dio.netty.tryReflectionSetAccessible=true", + "-cp", + classpath, + str(source_file), + ], + capture_output=True, + text=True, + timeout=90, + ) + self.assertEqual(result.returncode, 0, result.stdout + result.stderr) + self.assertIn("MISSING_JEP_REJECTED", result.stdout) + + def test_fresh_jvm_retries_configuration_without_spark_home(self): + import subprocess + import venv + + from pyspark import cloudpickle + + jvm = self.spark.sparkContext._jvm + java_home = jvm.java.lang.System.getProperty("java.home") + classpath = jvm.java.lang.System.getProperty("java.class.path") + spark_home = jvm.java.lang.System.getenv("SPARK_HOME") + self.assertIsNotNone(spark_home) + python_lib = Path(spark_home) / "python" / "lib" + if not (python_lib / "pyspark.zip").is_file(): + self.skipTest("Packaged bootstrap coverage requires python/lib/pyspark.zip") + env = os.environ.copy() + env.pop("SPARK_HOME", None) + env.pop("VIRTUAL_ENV", None) + env["PATH"] = os.defpath + env["PYTHONPATH"] = os.pathsep.join(str(p) for p in python_lib.glob("*.zip")) + # These flags must be ignored by the embedded interpreter. No signals are raised. + env["PYTHONFAULTHANDLER"] = "1" + env["PYTHONDEVMODE"] = "1" + env["LC_ALL"] = "C" + env["LANG"] = "C" + + class BootstrapCheck: + def __reduce__(self): + return eval, ( + "(__import__('sys').flags.isolated == 1 and " + "__import__('sys').flags.ignore_environment == 1 and " + "not __import__('faulthandler').is_enabled() and " + "'pyspark.zip' in __import__('pyspark').__file__ and " + "__import__('sys').stdout.line_buffering and " + "__import__('sys').stdout.write_through and " + "__import__('importlib.util', fromlist=['find_spec'])" + ".find_spec('only_in_failed_config') is None and " + "(print('SHUTDOWN_FLUSH', end='') or (lambda x: x))) or " + "(_ for _ in ()).throw(AssertionError('unexpected bootstrap state'))", + ) + + source = """ +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.Arrays; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.spark.sql.execution.python.InProcessPythonRuntime; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.util.ArrowUtils; + +class BootstrapProbe { + public static void main(String[] args) throws Exception { + var bad = scala.jdk.javaapi.CollectionConverters.asScala(Arrays.asList(args[0])).toSeq(); + boolean failed = false; + try { + InProcessPythonRuntime.initialize(bad); + } catch (jep.JepException expected) { + failed = true; + } + if (!failed) throw new AssertionError("Expected missing JEP package"); + var good = scala.jdk.javaapi.CollectionConverters + .asScala(Arrays.asList(args[1], args[2])).toSeq(); + try { + InProcessPythonRuntime.initialize(good); + Field field = ArrowUtils.toArrowField("result", DataTypes.LongType, true, "UTC", + false, org.apache.spark.sql.types.Metadata.empty(), false); + InProcessPythonRuntime.currentSession().register("probe", + Files.readAllBytes(Path.of(args[3])), field, args[4], false, false, false, true); + if (ArrowUtils.rootAllocator().getAllocatedMemory() != 0) { + throw new AssertionError("Unreleased registration schema"); + } + System.out.println("BOOTSTRAP_OK"); + } finally { + InProcessPythonRuntime.currentSession().release( + scala.jdk.javaapi.CollectionConverters.asScala(Arrays.asList("probe")).toSeq()); + InProcessPythonRuntime.shutdown(); + } + } +} +""" + with tempfile.TemporaryDirectory() as directory: + # Keep JEP absent initially even when it is installed in the system Python. + clean_python = Path(directory) / "python" + venv.EnvBuilder(with_pip=False).create(clean_python) + env["PATH"] = str(clean_python / "bin") + os.pathsep + os.defpath + # A module that only the failed configuration's directory has, which the retried + # configuration must not import. + (Path(directory) / "only_in_failed_config.py").write_text("") Review Comment: Done in 961ebc6: when `python/lib/pyspark.zip` is missing, the test zips `python/pyspark` into its temporary directory and puts that with the py4j zip on `PYTHONPATH`, so it no longer skips in CI. I ran it both with and without the packaged zip. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
