udf/worker -- Language-agnostic UDF Worker Framework

Package structure for the UDF worker framework described in SPIP SPARK-55278.

Overview

Spark processes a UDF by obtaining a WorkerDispatcher from a worker specification. The dispatcher manages workers behind the scenes. From the dispatcher, Spark gets a WorkerSession -- one per UDF invocation -- with an Iterator-to-Iterator process API that streams input batches through the worker and returns result batches.

UDFWorkerSpecification   -- how to create and configure workers
    |
    v
WorkerDispatcher      -- manages workers, creates sessions
    |
    v
WorkerSession         -- one UDF execution
    |  1. session.init(Init proto)
    |  2. val results = session.process(inputBatches)
    |  3. session.close()

How workers are created depends on the dispatcher implementation. The framework currently provides direct worker creation (local OS processes) and is designed for future indirect creation (via a provisioning service or daemon).

Sub-packages

udf/worker/
├── proto/                        -- protobuf message classes only (protobuf-java)
│     worker_spec.proto           -- UDFWorkerSpecification protobuf
│     udf_message.proto           -- UDF execution protocol messages (Init, UdfPayload, ...)
│     udf_service.proto           -- UdfWorker gRPC service (Execute, Manage)
│     common.proto                -- shared enums (UDFWorkerDataFormat, etc.)
│
├── core/                         -- abstract interfaces
│     WorkerDispatcher.scala      -- creates sessions, manages worker lifecycle
│     WorkerSession.scala         -- per-UDF init/process/cancel/close
│     WorkerConnection.scala      -- transport channel abstraction
│     WorkerSecurityScope.scala   -- security boundary for worker pooling
│     │
│     └── direct/                 -- "direct" creation: local OS processes
│           DirectWorkerDispatcher.scala  -- spawns processes, env lifecycle
│           DirectWorkerProcess.scala     -- OS process + connection + UDS socket
│           DirectWorkerSession.scala     -- session backed by a direct process
│
└── grpc/                         -- gRPC transport (gRPC runtime confined here)
      (generated)                 -- UdfWorkerGrpc service stubs from proto/udf_service.proto

The core/ package defines abstract interfaces that are independent of how workers are created. The core/direct/ sub-package implements “direct” worker creation where Spark spawns local OS processes. Future packages (e.g., core/indirect/) can implement alternative creation modes such as obtaining workers from a provisioning service or daemon.

The grpc/ module owns the gRPC service-stub generation (from proto/'s udf_service.proto) and the gRPC runtime dependencies. Keeping gRPC here means proto/, core/, and their consumers (core, catalyst, sql/core) carry no gRPC dependency on their classpath.

Wire protocol

Each UDF execution uses a single bidirectional Execute gRPC stream.

Engine -> Worker:  Init -> PayloadChunk* -> (DataRequest)* -> Finish (Cancel)?
                                                            | Cancel
Worker -> Engine:          InitResponse  -> (DataResponse)* -> (ErrorResponse)? -> (FinishResponse | CancelResponse)

See udf/worker/proto/src/main/protobuf/udf_message.proto for the complete message definitions, ordering invariants, and error contract, and udf_service.proto for the gRPC service.

Direct worker creation

DirectWorkerDispatcher spawns worker processes locally. On the first session, it runs the optional environment lifecycle callables from the UDFWorkerSpecification:

environmentVerification -- checks if the environment is ready (exit 0 = ready). When it succeeds, installation is skipped.
installation -- prepares the environment (installs runtime, dependencies, worker binaries). Only runs when verification is absent or fails.
environmentCleanup -- runs after the dispatcher is closed or on JVM shutdown to clean up temporary resources.

Environment setup runs once per dispatcher (not per session). Workers are terminated via SIGTERM/SIGKILL when the dispatcher is closed.

Basic usage (Scala)

import org.apache.spark.udf.worker.{
  DirectWorker, Init, ProcessCallable, UdfPayload,
  UDFProtoCommunicationPattern, UDFWorkerDataFormat, UDFWorkerProperties,
  UDFWorkerSpecification, UnixDomainSocket, WorkerCapabilities,
  WorkerConnectionSpec, WorkerEnvironment}
import org.apache.spark.udf.worker.core._
import com.google.protobuf.ByteString

// 1. Define a worker spec (direct creation mode).
val spec = UDFWorkerSpecification.newBuilder()
  .setEnvironment(WorkerEnvironment.newBuilder()
    .setEnvironmentVerification(ProcessCallable.newBuilder()
      .addCommand("python").addCommand("-c").addCommand("import my_udf_worker").build())
    .setInstallation(ProcessCallable.newBuilder()
      .addCommand("pip").addCommand("install").addCommand("my_udf_worker").build())
    .build())
  .setCapabilities(WorkerCapabilities.newBuilder()
    .addSupportedDataFormats(UDFWorkerDataFormat.ARROW)
    .addSupportedCommunicationPatterns(
      UDFProtoCommunicationPattern.BIDIRECTIONAL_STREAMING)
    .build())
  .setDirect(DirectWorker.newBuilder()
    .setRunner(ProcessCallable.newBuilder()
      .addCommand("python").addCommand("-m").addCommand("my_udf_worker").build())
    .setProperties(UDFWorkerProperties.newBuilder()
      .setConnection(WorkerConnectionSpec.newBuilder()
        .setUnixDomainSocket(UnixDomainSocket.getDefaultInstance).build())
      .build())
    .build())
  .build()

// 2. Create a dispatcher. Use a protocol-specific subclass of
//    DirectWorkerDispatcher (e.g., gRPC over UDS).
val dispatcher: WorkerDispatcher = ...

// 3. Create a session for one UDF execution.
val session = dispatcher.createSession(securityScope = None)
try {
  // 4. Initialize with the serialized function and schemas.
  session.init(Init.newBuilder()
    .setProtocolVersion(1)
    .setUdf(UdfPayload.newBuilder()
      .setPayload(ByteString.copyFrom(serializedFunction))
      .setFormat(payloadFormat)   // worker-recognised tag
      .build())
    .setDataFormat(UDFWorkerDataFormat.ARROW)
    .setInputSchema(ByteString.copyFrom(arrowInputSchema))
    .setOutputSchema(ByteString.copyFrom(arrowOutputSchema))
    .build())

  // 5. Process data -- Iterator in, Iterator out.
  val results: Iterator[Array[Byte]] =
    session.process(inputBatches)

  // Consume results lazily.
  results.foreach(processResultBatch)
} finally {
  session.close()
}

// 6. Shut down all workers.
dispatcher.close()

Build

SBT:

build/sbt "udf-worker-proto/compile" "udf-worker-core/compile" "udf-worker-grpc/compile"

Maven:

build/mvn compile -pl udf/worker/proto,udf/worker/core,udf/worker/grpc -am

Test