From 91f80d2da4c1eb65674a32714ebe24c3faa958ee Mon Sep 17 00:00:00 2001 From: kary zheng Date: Wed, 2 Sep 2026 16:30:44 -0700 Subject: [PATCH 1/4] test(verify): run a Python operator the way the engine runs it A Python operator is not called; it is handed to an interpreter that imports the generated module and drives it through the same open, process and close the engine uses. `PyOpExecHarness` writes that module, starts the driver, and reads back what the operator emitted. The driver is the engine's side of the contract written out plainly: it is what makes the answer this side produces the engine's answer rather than an approximation of it. Co-Authored-By: Claude Opus 5 (1M context) --- build.sbt | 1 + .../src/test/resources/python/py_op_driver.py | 531 ++++++++++++++++++ .../translator/verify/PyOpExecHarness.scala | 406 +++++++++++++ 3 files changed, 938 insertions(+) create mode 100644 workflow-compiling-service/src/test/resources/python/py_op_driver.py create mode 100644 workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/PyOpExecHarness.scala diff --git a/build.sbt b/build.sbt index 035edbb68cc..94a5af5606c 100644 --- a/build.sbt +++ b/build.sbt @@ -239,6 +239,7 @@ lazy val WorkflowCompiler = (project in file("common/workflow-compiler")) .dependsOn(WorkflowOperator) lazy val WorkflowCompilingService = (project in file("workflow-compiling-service")) .dependsOn(WorkflowCompiler, Auth, Config, Resource) + .dependsOn(WorkflowOperator % "test->test") // reuse PythonWorkerPool in verify tests .settings(commonModuleSettings) .settings( dependencyOverrides ++= Seq( diff --git a/workflow-compiling-service/src/test/resources/python/py_op_driver.py b/workflow-compiling-service/src/test/resources/python/py_op_driver.py new file mode 100644 index 00000000000..1ff2a9000da --- /dev/null +++ b/workflow-compiling-service/src/test/resources/python/py_op_driver.py @@ -0,0 +1,531 @@ +#!/usr/bin/env python3 +# +# 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. +""" +Driver that runs a Texera Python-native operator without spinning up the +Pekko/Arrow worker stack. + +Symmetric to ``OpExecHarness`` on the JVM side: take an OpDesc's +``generatePythonCode()`` output (which defines a ``UDFOperatorV2`` / +``UDFTableOperator`` / ``UDFBatchOperator`` / ``UDFSourceOperator`` subclass), +load JSONL+sidecar inputs into ``Tuple`` instances, drive +``open -> process_tuple/on_finish per port -> close``, and write the emitted +tuples back as JSONL+sidecar in the same format ``TupleIO`` reads. + +The harness invokes us as:: + + python3 py_op_driver.py + +with ``PYTHONPATH`` pointing at ``amber/src/main/python`` so ``pytexera`` / +``pyamber`` import cleanly. + +Config schema (all paths absolute):: + + { + "operatorCode": "", + "isSource": false, + "portOrder": [0, 1], # input-port dependency order + "inputs": [{"portIndex": 0, "dataPath": "...", "schemaPath": "..."}], + "outputs": [{"portIndex": 0, "dataPath": "...", "schema": + {"attributes": [{"attributeName": "...", + "attributeType": "..."}]}}] + } + +Output schemas come from the JVM side (``PhysicalOp.propagateSchema``) so +this driver never has to infer them. The driver writes the schema back as a +``.jsonl.schema.json`` sidecar next to each ``dataPath``, matching +``TupleIO.writeTuples``. +""" +from __future__ import annotations + +import base64 +import inspect +import json +import pickle +import sys +import traceback +from pathlib import Path +from typing import Any, Iterable, Iterator, List, Mapping, Sequence + +import pandas as pd + +# pytexera re-exports the operator base classes and the Tuple/Table types. +# The Scala side prepends `amber/src/main/python` to PYTHONPATH so these +# resolve. If they don't, raise a clean error rather than a cryptic +# ImportError deep in user code. +try: + from pytexera import ( # noqa: F401 (used dynamically in user code's globals) + Batch, + BatchLike, + Iterator as PyIterator, # noqa: F401 + Optional as PyOptional, # noqa: F401 + Table, + TableLike, + Tuple, + TupleLike, + UDFBatchOperator, + UDFOperatorV2, + UDFSourceOperator, + UDFTableOperator, + Union as PyUnion, # noqa: F401 + logger as pytexera_logger, # noqa: F401 + overrides, # noqa: F401 + ) + from core.models.schema.schema import Schema as TexeraSchema + from core.models.schema.attribute_type import AttributeType, RAW_TYPE_MAPPING +except ImportError as exc: + sys.stderr.write( + "py_op_driver.py: failed to import pytexera/pyamber. The harness must " + "set PYTHONPATH to `amber/src/main/python` and the venv must have all " + "amber Python deps installed (see amber/requirements.txt).\n" + f"Underlying error: {exc!r}\n" + ) + raise + + +# -------------------------------------------------------------------------- +# Schema sidecar I/O. +# -------------------------------------------------------------------------- +# The JVM writes attributes using AttributeType's Jackson @JsonValue ("string", +# "integer", "long", "double", "boolean", "timestamp", "binary", +# "large_binary"). The Python Schema's RAW_TYPE_MAPPING uses uppercase keys +# ("STRING", "INTEGER", ...). Translate at the boundary; keep the rest of +# the pipeline using Python's AttributeType enum. +_SCALA_TO_PY_TYPE: Mapping[str, str] = { + "string": "STRING", + "integer": "INTEGER", + "long": "LONG", + "double": "DOUBLE", + "boolean": "BOOLEAN", + "timestamp": "TIMESTAMP", + "binary": "BINARY", + "large_binary": "LARGE_BINARY", +} + +_PY_TO_SCALA_TYPE: Mapping[AttributeType, str] = { + AttributeType.STRING: "string", + AttributeType.INT: "integer", + AttributeType.LONG: "long", + AttributeType.DOUBLE: "double", + AttributeType.BOOL: "boolean", + AttributeType.TIMESTAMP: "timestamp", + AttributeType.BINARY: "binary", + AttributeType.LARGE_BINARY: "large_binary", +} + + +def _schema_from_dict(payload: Mapping[str, Any]) -> TexeraSchema: + raw: "dict[str, str]" = {} + for attr in payload["attributes"]: + raw_name = attr["attributeName"] + raw_type = attr["attributeType"].lower() + if raw_type not in _SCALA_TO_PY_TYPE: + raise ValueError( + f"py_op_driver: unknown attributeType {attr['attributeType']!r} " + f"for attribute {raw_name!r}" + ) + raw[raw_name] = _SCALA_TO_PY_TYPE[raw_type] + return TexeraSchema(raw_schema=raw) + + +def _schema_to_dict(schema: TexeraSchema) -> "dict[str, Any]": + return { + "attributes": [ + {"attributeName": name, "attributeType": _PY_TO_SCALA_TYPE[attr_type]} + for name, attr_type in schema.as_key_value_pairs() + ] + } + + +def _read_schema_sidecar(data_path: Path) -> TexeraSchema: + sidecar = data_path.with_name(data_path.name + ".schema.json") + with sidecar.open("r", encoding="utf-8") as fh: + return _schema_from_dict(json.load(fh)) + + +def _write_schema_sidecar(data_path: Path, schema: TexeraSchema) -> None: + sidecar = data_path.with_name(data_path.name + ".schema.json") + with sidecar.open("w", encoding="utf-8") as fh: + json.dump(_schema_to_dict(schema), fh) + + +# -------------------------------------------------------------------------- +# Tuple I/O. JSONL with sidecar — same on-disk shape as TupleIO on the JVM. +# -------------------------------------------------------------------------- +def _coerce_field(raw: Any, attr_type: AttributeType) -> Any: + """Coerce a JSON-decoded field to the type the schema expects.""" + if raw is None: + return None + if attr_type == AttributeType.STRING: + return str(raw) + if attr_type == AttributeType.INT: + return int(raw) + if attr_type == AttributeType.LONG: + return int(raw) + if attr_type == AttributeType.DOUBLE: + return float(raw) + if attr_type == AttributeType.BOOL: + return bool(raw) + if attr_type == AttributeType.BINARY: + return base64.b64decode(raw) + if attr_type == AttributeType.TIMESTAMP: + # TupleIO writes java.sql.Timestamp.toString ("YYYY-MM-DD HH:MM:SS[.f]"); + # the native path's schema maps TIMESTAMP -> datetime.datetime, and + # pandas parses the JDBC form robustly. + return pd.Timestamp(raw).to_pydatetime() + # LARGE_BINARY: defer until an operator actually exercises it. Failing loud + # beats silently passing a string through. + raise NotImplementedError( + f"py_op_driver: reading attribute type {attr_type!r} from JSONL is " + f"not implemented yet" + ) + + +def _read_tuples(data_path: Path, schema: TexeraSchema) -> List[Tuple]: + rows: List[Tuple] = [] + if not data_path.exists(): + return rows + with data_path.open("r", encoding="utf-8") as fh: + for line_num, raw_line in enumerate(fh, 1): + line = raw_line.strip() + if not line: + continue + try: + obj = json.loads(line) + except json.JSONDecodeError as exc: + raise ValueError( + f"py_op_driver: invalid JSON on line {line_num} of {data_path}: {exc}" + ) from exc + field_data: "dict[str, Any]" = {} + for name, attr_type in schema.as_key_value_pairs(): + field_data[name] = _coerce_field(obj.get(name), attr_type) + tup = Tuple(field_data) + tup.finalize(schema) + rows.append(tup) + return rows + + +def _emit_as_dicts( + emitted: Iterable[Any], schema: TexeraSchema +) -> Iterator["dict[str, Any]"]: + """ + Flatten whatever the operator yields into per-row dicts keyed by the + output schema's attribute names. The operator may yield: + * pandas.DataFrame (UDFTableOperator's process_table return) + * pandas.Series (single row) + * dict / OrderedDict (e.g. BarChart yields {'html-content': html}) + * Tuple + * None (skip — matches the engine's behavior) + """ + attr_names = schema.get_attr_names() + for item in emitted: + if item is None: + continue + if isinstance(item, pd.DataFrame): + for _, row in item.iterrows(): + yield {col: row[col] for col in attr_names if col in row.index} + elif isinstance(item, pd.Series): + yield {col: item[col] for col in attr_names if col in item.index} + elif isinstance(item, Tuple): + yield {name: item[name] for name in attr_names} + elif isinstance(item, Mapping): + yield {name: item.get(name) for name in attr_names} + else: + raise TypeError( + f"py_op_driver: cannot serialize emitted value of type " + f"{type(item).__name__}: {item!r}" + ) + + +def _jsonify(value: Any, attr_type: AttributeType) -> Any: + """Convert a Python value into something json.dumps will accept.""" + if value is None: + return None + # A missing cell reaches pandas as NaN or NaT, not None, and every branch + # below assumes a real value: the timestamp one formats NaT's float + # microsecond with a "d" code and raises. Emitting null matches the + # standalone path, whose to_json writes NaN and NaT that way. Scalars only — + # an object column can hold a list or an array, where isna answers + # element-wise and the result is not a truth value. + if pd.api.types.is_scalar(value) and pd.isna(value): + return None + # pandas often hands us numpy scalars; .item() collapses them to native. + if hasattr(value, "item") and not isinstance(value, (str, bytes)): + try: + value = value.item() + except (ValueError, AttributeError): + pass + if attr_type == AttributeType.STRING: + return str(value) + if attr_type in (AttributeType.INT, AttributeType.LONG): + return int(value) + if attr_type == AttributeType.DOUBLE: + return float(value) + if attr_type == AttributeType.BOOL: + return bool(value) + if attr_type == AttributeType.BINARY: + # Trained-model / object columns: pickle then base64 so the value + # survives JSONL round-trip. Mirrors the BINARY read path in + # _coerce_field. For deterministic estimators the pickle is byte-stable + # across processes, so the two verification paths compare equal. + raw = value if isinstance(value, (bytes, bytearray)) else pickle.dumps(value) + return base64.b64encode(raw).decode("ascii") + if attr_type == AttributeType.TIMESTAMP: + # Emit the same JDBC string java.sql.Timestamp.toString produces (>=1 + # fractional digit), so a passed-through timestamp column matches the + # standalone path, which carries it as that exact string. + ts = pd.Timestamp(value) + frac = f"{ts.microsecond:06d}".rstrip("0") or "0" + return ts.strftime("%Y-%m-%d %H:%M:%S") + "." + frac + raise NotImplementedError( + f"py_op_driver: writing attribute type {attr_type!r} to JSONL is " + f"not implemented yet" + ) + + +def _write_tuples( + data_path: Path, rows: Iterable["dict[str, Any]"], schema: TexeraSchema +) -> None: + _write_schema_sidecar(data_path, schema) + with data_path.open("w", encoding="utf-8") as fh: + for row in rows: + serialized: "dict[str, Any]" = {} + for name, attr_type in schema.as_key_value_pairs(): + serialized[name] = _jsonify(row.get(name), attr_type) + fh.write(json.dumps(serialized)) + fh.write("\n") + + +# -------------------------------------------------------------------------- +# Operator discovery + lifecycle. +# -------------------------------------------------------------------------- +_OPERATOR_BASES = ( + UDFOperatorV2, + UDFTableOperator, + UDFBatchOperator, + UDFSourceOperator, +) + + +def _exec_user_code(code: str) -> "dict[str, Any]": + """ + Execute the operator code in a fresh namespace seeded with the pytexera + re-exports, the way the real Texera Python worker does it (see + ``InitializeExecutorHandler``). Returning the namespace lets us pick the + user's operator class out of it. + """ + namespace: "dict[str, Any]" = { + "__name__": "__texera_user_op__", + "__builtins__": __builtins__, + } + # pytexera does `from pyamber import *` itself, so this single import is + # equivalent to what the generated code's `from pytexera import *` brings + # into scope. + exec("from pytexera import *", namespace) + try: + exec(code, namespace) + except Exception: + sys.stderr.write("py_op_driver: error executing operator code:\n") + traceback.print_exc() + raise + return namespace + + +def _discover_operator_class(namespace: Mapping[str, Any]) -> type: + candidates: List[type] = [] + for name, obj in namespace.items(): + if not inspect.isclass(obj): + continue + if obj in _OPERATOR_BASES: + continue # the base classes themselves come in via the import + if any(issubclass(obj, base) for base in _OPERATOR_BASES): + candidates.append(obj) + if not candidates: + raise RuntimeError( + "py_op_driver: operator code did not define a subclass of " + "UDFOperatorV2 / UDFTableOperator / UDFBatchOperator / UDFSourceOperator" + ) + if len(candidates) > 1: + names = ", ".join(c.__name__ for c in candidates) + raise RuntimeError( + f"py_op_driver: operator code defined multiple UDF subclasses " + f"({names}); expected exactly one" + ) + return candidates[0] + + +def _run_operator( + op: Any, + is_source: bool, + port_order: Sequence[int], + inputs_by_port: Mapping[int, Sequence[Tuple]], +) -> List[Any]: + """ + Drive the operator's lifecycle. Returns the flat list of emitted values + (anything not-None yielded by process_tuple / on_finish, in emission + order). UDF operators don't expose multi-output ports today, so we don't + bucket by output port — same convention as ``OpExecHarness`` when port + is unset. + """ + emitted: List[Any] = [] + + op.open() + try: + if is_source: + # Source ops: SourceOperator.on_finish iterates produce() and + # yields Tuples. Single synthetic port 0 — see OpExecHarness. + for item in op.on_finish(0): + if item is not None: + emitted.append(item) + return emitted + + for port in port_order: + for tup in inputs_by_port.get(port, ()): # type: ignore[arg-type] + for item in op.process_tuple(tup, port): + if item is not None: + emitted.append(item) + for item in op.on_finish(port): + if item is not None: + emitted.append(item) + finally: + op.close() + + return emitted + + +# -------------------------------------------------------------------------- +# Entry point. +# -------------------------------------------------------------------------- +def run_config(config: Mapping[str, Any]) -> None: + """Run one operator to completion from a parsed config dict. Writes the + output JSONL+sidecar as a side effect. Raises on any failure. Shared by the + CLI (main) and the persistent server (serve) so both behave identically. + + Each call execs the user code in a FRESH namespace and constructs a FRESH + operator instance, so operators don't share Python-level state across jobs + when run through the server (the isolation the per-process CLI gave for + free).""" + # Both paths seed numpy's global RNG with the same value before running, so + # an estimator built without random_state (sklearn reads the global RNG for + # that) draws the same samples on each and the two models come out + # identical. Without it a stochastic estimator makes the parity check + # inconclusive in both directions: a difference could be the translation or + # could be the draw, and a match could be either. Per call, not per process: + # the worker pool reuses this process, so a job that inherited the previous + # job's RNG position would not line up with a fresh standalone one. Keep in + # step with StandaloneRunner.VerifySeed. + import numpy as _texera_np + + _texera_np.random.seed(20260811) + + operator_code: str = config["operatorCode"] + is_source: bool = bool(config.get("isSource", False)) + port_order: Sequence[int] = list(config.get("portOrder", [])) + + inputs_by_port: "dict[int, List[Tuple]]" = {} + for entry in config.get("inputs", []): + port = int(entry["portIndex"]) + data_path = Path(entry["dataPath"]) + schema = _read_schema_sidecar(data_path) + inputs_by_port[port] = _read_tuples(data_path, schema) + + # Default port order: sorted by index. Matches OpExecHarness's fallback + # when getInputPortDependencyPairs is empty. + if not port_order: + port_order = sorted(inputs_by_port.keys()) + + namespace = _exec_user_code(operator_code) + op_class = _discover_operator_class(namespace) + op_instance = op_class() + + emitted = _run_operator(op_instance, is_source, port_order, inputs_by_port) + + outputs = config.get("outputs", []) + if len(outputs) > 1: + raise NotImplementedError( + "py_op_driver: multi-output Python operators are not supported " + "yet (no UDF base class exposes per-port emission)" + ) + if outputs: + out_entry = outputs[0] + out_path = Path(out_entry["dataPath"]) + out_schema = _schema_from_dict(out_entry["schema"]) + rows = list(_emit_as_dicts(emitted, out_schema)) + _write_tuples(out_path, rows, out_schema) + + +def main(argv: Sequence[str]) -> int: + if len(argv) != 2: + sys.stderr.write(f"usage: {argv[0]} \n") + return 2 + config_path = Path(argv[1]) + with config_path.open("r", encoding="utf-8") as fh: + config = json.load(fh) + run_config(config) + return 0 + + +def serve() -> int: + """Persistent driver: import pyamber once, then run many operators. + + pytexera/pyamber import at module load (~300 ms) is the dominant per-call + cost; paying it once here instead of per operator is the whole point. Reads + one JSON job per line on stdin, writes one JSON result per line on stdout: + + request {"configPath": ""}\n + response {"exit": 0|1, "stdout": "...", "stderr": "..."}\n + + exit=1 with the traceback on stderr mirrors a nonzero CLI exit, so the + Scala side's PyOpDriverException path is unchanged. An operator error never + kills the server; only closing stdin (EOF) ends it. The executed script's + stdout/stderr are captured so they can't corrupt the protocol channel. + """ + import io + from contextlib import redirect_stderr, redirect_stdout + + sys.stdout.write(json.dumps({"ready": True}) + "\n") + sys.stdout.flush() + + for line in sys.stdin: + line = line.strip() + if not line: + continue + out_buf, err_buf = io.StringIO(), io.StringIO() + try: + job = json.loads(line) + with Path(job["configPath"]).open("r", encoding="utf-8") as fh: + config = json.load(fh) + with redirect_stdout(out_buf), redirect_stderr(err_buf): + run_config(config) + resp = {"exit": 0, "stdout": out_buf.getvalue(), "stderr": err_buf.getvalue()} + except BaseException: # noqa: BLE001 — a bad job must not kill the server + resp = { + "exit": 1, + "stdout": out_buf.getvalue(), + "stderr": err_buf.getvalue() + traceback.format_exc(), + } + sys.stdout.write(json.dumps(resp) + "\n") + sys.stdout.flush() + return 0 + + +if __name__ == "__main__": + if len(sys.argv) > 1 and sys.argv[1] == "--serve": + sys.exit(serve()) + else: + sys.exit(main(sys.argv)) diff --git a/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/PyOpExecHarness.scala b/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/PyOpExecHarness.scala new file mode 100644 index 00000000000..a6ae7b3b373 --- /dev/null +++ b/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/PyOpExecHarness.scala @@ -0,0 +1,406 @@ +/* + * 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.texera.amber.translator.verify + +import com.fasterxml.jackson.databind.node.{ArrayNode, ObjectNode} +import com.typesafe.scalalogging.LazyLogging +import org.apache.texera.amber.core.executor.OpExecWithCode +import org.apache.texera.amber.core.tuple.Schema +import org.apache.texera.amber.core.virtualidentity.{ + ExecutionIdentity, + PhysicalOpIdentity, + WorkflowIdentity +} +import org.apache.texera.amber.core.workflow.{PhysicalPlan, PortIdentity} +import org.apache.texera.amber.operator.LogicalOp +import org.apache.texera.amber.util.JSONUtils.objectMapper +import org.apache.texera.amber.util.python.PythonWorkerPool + +import java.io.File +import java.nio.charset.StandardCharsets +import java.nio.file.{Files, Path, Paths, StandardCopyOption} +import scala.collection.mutable.ArrayBuffer +import scala.sys.process._ + +/** + * Counterpart to [[OpExecHarness]] for Python-native operators + * ([[OpExecWithCode]] with language="python"). Drives the operator's + * generatePythonCode() output through a thin subprocess driver rather than + * spinning up the Pekko/Arrow worker stack. + * + * Same Result(outputs, outputSchemas) shape as OpExecHarness so the rest of + * the verify pipeline (Comparator, category runners) is harness-agnostic. + * + * Scope (MVP, mirrors OpExecHarness's MVP): + * - Single-PhysicalOp plans only. PythonOperatorDescriptor only emits + * either a `sourcePhysicalOp` or a `oneToOnePhysicalOp`, so multi-op + * plans don't exist for Python-native ops today. If that changes, add + * topo-order driving here the way OpExecHarness does. + * - Single output port. UDFOperatorV2 / UDFTableOperator / UDFBatchOperator + * / UDFSourceOperator all yield TupleLike without specifying a port — + * same convention OpExecHarness uses when port is unset. + * - JSONL types: STRING / INTEGER / LONG / DOUBLE / BOOLEAN. TIMESTAMP / + * BINARY / LARGE_BINARY require explicit codecs in both [[TupleIO]] and + * the driver — add when the first operator needs them. + */ +object PyOpExecHarness extends LazyLogging { + + private val TestWorkflowId = WorkflowIdentity(0L) + private val TestExecutionId = ExecutionIdentity(0L) + + // Same Result shape as OpExecHarness so callers can swap harnesses + // transparently. + final case class Result( + outputs: Map[PortIdentity, Path], + outputSchemas: Map[PortIdentity, Schema] + ) + + // Driver script lives on the test classpath at /python/py_op_driver.py + // (sibling to compare.py). Extracted to a temp file at runtime so it works + // whether the test resources are loose files or sealed in a jar. + private val DriverResourcePath = "/python/py_op_driver.py" + + def execute( + opDesc: LogicalOp, + inputs: Map[PortIdentity, Path], + outputDir: Path, + pythonExe: String = resolvePython(), + amberPythonHome: Path = resolveAmberPythonHome() + ): Result = { + Files.createDirectories(outputDir) + + val plan = opDesc.getPhysicalPlan(TestWorkflowId, TestExecutionId) + + // PythonOperatorDescriptor builds single-op plans; bail loudly if some + // future Python op produces a multi-stage plan (need to extend the driver + // and the per-PhysicalOp config the way OpExecHarness does). + require( + plan.operators.size == 1, + s"PyOpExecHarness only supports single-PhysicalOp plans for now, got " + + s"${plan.operators.size} PhysicalOps" + ) + val phOp = plan.operators.head + + val (pythonCode, language) = phOp.opExecInitInfo match { + case OpExecWithCode(code, lang) => (code, lang) + case other => + throw new UnsupportedOperationException( + s"PyOpExecHarness only supports OpExecWithCode; got ${other.getClass.getSimpleName}. " + + "For OpExecWithClassName, use OpExecHarness." + ) + } + require( + language == "python", + s"""PyOpExecHarness only supports language="python", got "$language".""" + ) + + // External input ports = same definition as OpExecHarness. For a + // single-op plan that's just every input port the op declares. + val externalInputs: Set[(PhysicalOpIdentity, PortIdentity)] = + phOp.inputPorts.keys.map(portId => (phOp.id, portId)).toSet + validateInputCoverage(externalInputs, inputs.keySet) + + val inputSchemas: Map[PortIdentity, Schema] = + inputs.map { case (portId, path) => portId -> TupleIO.readSchemaSidecar(path) } + + val planWithSchemas = propagateExternalSchemas(plan, externalInputs, inputSchemas) + val phOpWithSchemas = planWithSchemas.operators.head + + // Output port schemas come from PhysicalPlan.propagateSchema — same + // ground truth OpExecHarness writes to its own outputs. + val outputPortSchemas: Map[PortIdentity, Schema] = + phOpWithSchemas.outputPorts.map { + case (portId, (_, _, schemaOrErr)) => + portId -> schemaOrErr.toOption.getOrElse( + throw new IllegalStateException( + s"Output schema for ($portId) was not propagated" + ) + ) + } + + require( + outputPortSchemas.size == 1, + s"PyOpExecHarness only supports single-output-port operators, got " + + s"${outputPortSchemas.size} output ports" + ) + val (outputPortId, outputSchema) = outputPortSchemas.head + val outputPath = outputDir.resolve(s"output_port_${outputPortId.id}.jsonl") + + // Port ordering for multi-input ops: respect declared dependencies + // (matches OpExecHarness — e.g. HashJoin probe processes build-side + // first). Default = sorted by port id when no dependencies declared. + val portOrder: Seq[Int] = + if (phOpWithSchemas.getInputPortDependencyPairs.nonEmpty) + phOpWithSchemas.getInputPortDependencyPairs.map(_.id) + else phOpWithSchemas.inputPorts.keys.toList.map(_.id).sorted + + val config = buildConfig( + pythonCode = pythonCode, + isSource = phOpWithSchemas.isSourceOperator, + portOrder = portOrder, + inputs = inputs, + outputPath = outputPath, + outputSchema = outputSchema + ) + + val configPath = outputDir.resolve("py_op_driver_config.json") + Files.write(configPath, config.getBytes(StandardCharsets.UTF_8)) + + val driverPath = extractDriverScript() + runDriver(driverPath, configPath, outputDir, pythonExe, amberPythonHome) + + Result( + outputs = Map(outputPortId -> outputPath), + outputSchemas = Map(outputPortId -> outputSchema) + ) + } + + // -------------------------------------------------------------------------- + // Config serialization. Matches the driver's expected schema (see + // py_op_driver.py's module docstring). + // -------------------------------------------------------------------------- + private def buildConfig( + pythonCode: String, + isSource: Boolean, + portOrder: Seq[Int], + inputs: Map[PortIdentity, Path], + outputPath: Path, + outputSchema: Schema + ): String = { + val root: ObjectNode = objectMapper.createObjectNode() + root.put("operatorCode", pythonCode) + root.put("isSource", isSource) + + val portOrderArr: ArrayNode = root.putArray("portOrder") + portOrder.foreach(portOrderArr.add) + + val inputsArr: ArrayNode = root.putArray("inputs") + inputs.toSeq.sortBy(_._1.id).foreach { + case (portId, dataPath) => + val entry: ObjectNode = inputsArr.addObject() + entry.put("portIndex", portId.id) + entry.put("dataPath", dataPath.toAbsolutePath.toString) + // schemaPath is implicit (data_path + ".schema.json") — the driver + // resolves it the same way TupleIO does. + } + + val outputsArr: ArrayNode = root.putArray("outputs") + val outEntry: ObjectNode = outputsArr.addObject() + outEntry.put("dataPath", outputPath.toAbsolutePath.toString) + // Embed the schema directly. We can't just write the sidecar ahead of + // time and have the driver read it, because writing a sidecar before + // outputs exist would leave a stale sidecar on partial failures. + outEntry.set[ObjectNode]( + "schema", + objectMapper.valueToTree[ObjectNode](outputSchema) + ) + + objectMapper.writeValueAsString(root) + } + + // -------------------------------------------------------------------------- + // Subprocess invocation. + // -------------------------------------------------------------------------- + private def runDriver( + driverPath: Path, + configPath: Path, + cwd: Path, + pythonExe: String, + amberPythonHome: Path + ): Unit = { + // Prepend amber's Python source to PYTHONPATH so `import pytexera` + // resolves. Existing PYTHONPATH (if any) is preserved as the lower- + // priority suffix. + val existing = sys.env.getOrElse("PYTHONPATH", "") + val newPyPath = + if (existing.isEmpty) amberPythonHome.toAbsolutePath.toString + else s"${amberPythonHome.toAbsolutePath}${File.pathSeparator}$existing" + + val (exit, stdout, stderr) = execDriver(driverPath, configPath, cwd, pythonExe, newPyPath) + if (exit != 0) { + throw new PyOpDriverException( + exitCode = exit, + driverPath = driverPath, + configPath = configPath, + stdout = stdout, + stderr = stderr + ) + } + } + + // Prefer a pooled persistent worker (imports pytexera/pyamber once via + // `py_op_driver.py --serve`, the ~300 ms cost that dominates a per-Python-op + // run — see PythonWorkerPool). A rare hard worker crash falls back to a + // one-shot subprocess so behavior is never worse than the original path. Both + // paths use absolute config paths, so cwd only matters to the subprocess + // form; the worker constructs a fresh operator per job for isolation. + private def execDriver( + driverPath: Path, + configPath: Path, + cwd: Path, + pythonExe: String, + pythonPath: String + ): (Int, String, String) = { + if (PythonWorkerPool.enabled) { + try { + val req = objectMapper.createObjectNode() + req.put("configPath", configPath.toAbsolutePath.toString) + val o = PythonWorkerPool.run( + DriverResourcePath, + Seq("--serve"), + pythonExe, + req, + env = Map("PYTHONPATH" -> pythonPath) + ) + return (o.exit, o.stdout, o.stderr) + } catch { + case e: PythonWorkerPool.WorkerDiedException => + logger.warn( + s"py_op_driver worker unavailable; falling back to one-shot subprocess " + + s"for $configPath: ${e.getMessage}" + ) + } + } + runDriverSubprocess(driverPath, configPath, cwd, pythonExe, pythonPath) + } + + // Original one-process-per-operator path. Retained as the fallback and as the + // behavior selected by TEXERA_TEST_PYTHON_WORKER=0. + private def runDriverSubprocess( + driverPath: Path, + configPath: Path, + cwd: Path, + pythonExe: String, + pythonPath: String + ): (Int, String, String) = { + val outBuf = ArrayBuffer.empty[String] + val errBuf = ArrayBuffer.empty[String] + val procLogger = ProcessLogger(line => outBuf += line, line => errBuf += line) + val exit = Process( + Seq(pythonExe, driverPath.toString, configPath.toString), + Some(cwd.toFile), + "PYTHONPATH" -> pythonPath + ).!(procLogger) + (exit, outBuf.mkString("\n"), errBuf.mkString("\n")) + } + + // -------------------------------------------------------------------------- + // Resolution helpers. + // -------------------------------------------------------------------------- + private def resolvePython(): String = + sys.env.get("UDF_PYTHON_PATH").filter(_.nonEmpty).getOrElse("python3.12") + + /** + * Locate `amber/src/main/python`. Resolution chain: + * 1. Env var TEXERA_AMBER_PYTHON_HOME (set by CI / dev shell). + * 2. Walk up from cwd looking for `amber/src/main/python`. + * sbt runs tests with cwd = the subproject dir (`workflow-compiling-service/`), + * so the walk-up is two levels at most for the normal layout. + */ + private def resolveAmberPythonHome(): Path = { + sys.env.get("TEXERA_AMBER_PYTHON_HOME").filter(_.nonEmpty).map(Paths.get(_)).getOrElse { + val cwd = Paths.get(".").toAbsolutePath.normalize() + val maxDepth = 5 + var current: Path = cwd + var depth = 0 + while (current != null && depth <= maxDepth) { + val candidate = current.resolve("amber/src/main/python") + if (Files.isDirectory(candidate)) return candidate.toAbsolutePath + current = current.getParent + depth += 1 + } + throw new RuntimeException( + s"PyOpExecHarness: could not locate amber/src/main/python from cwd $cwd. " + + "Set TEXERA_AMBER_PYTHON_HOME to the absolute path." + ) + } + } + + private def extractDriverScript(): Path = { + val stream = getClass.getResourceAsStream(DriverResourcePath) + require( + stream != null, + s"py_op_driver.py not found on classpath at $DriverResourcePath" + ) + try { + val tmp = Files.createTempFile("py_op_driver-", ".py") + Files.copy(stream, tmp, StandardCopyOption.REPLACE_EXISTING) + tmp.toFile.deleteOnExit() + tmp + } finally stream.close() + } + + // -------------------------------------------------------------------------- + // Schema propagation — mirrors OpExecHarness.propagateExternalSchemas. + // Kept inline (rather than shared) so the two harnesses stay independently + // readable; consolidate if a third harness shows up. + // -------------------------------------------------------------------------- + private def propagateExternalSchemas( + plan: PhysicalPlan, + externalPorts: Set[(PhysicalOpIdentity, PortIdentity)], + schemas: Map[PortIdentity, Schema] + ): PhysicalPlan = { + var acc = PhysicalPlan(operators = Set.empty, links = Set.empty) + plan.topologicalIterator().map(plan.getOperator).foreach { phOp => + val updated = phOp.inputPorts.keys.foldLeft(phOp) { (op, portId) => + if (externalPorts.contains((phOp.id, portId)) && schemas.contains(portId)) { + op.propagateSchema(Some((portId, schemas(portId)))) + } else op + } + acc = acc.addOperator(updated.propagateSchema()) + plan.getUpstreamPhysicalLinks(phOp.id).foreach { link => + acc = acc.addLink(link) + } + } + acc + } + + private def validateInputCoverage( + external: Set[(PhysicalOpIdentity, PortIdentity)], + provided: Set[PortIdentity] + ): Unit = { + val expected = external.map(_._2) + val missing = expected -- provided + val extra = provided -- expected + require( + missing.isEmpty, + s"Missing input fixtures for external ports: $missing (expected $expected)" + ) + if (extra.nonEmpty) { + logger.warn(s"Input fixtures provided for non-external ports (ignored): $extra") + } + } +} + +final class PyOpDriverException( + val exitCode: Int, + val driverPath: Path, + val configPath: Path, + val stdout: String, + val stderr: String +) extends RuntimeException( + s"""py_op_driver.py exited with code $exitCode. + |Driver: $driverPath + |Config: $configPath + |--- stdout --- + |$stdout + |--- stderr --- + |$stderr""".stripMargin + ) From d2e8495a7a713c09bc2dee744d0f8dec332ab48b Mon Sep 17 00:00:00 2001 From: kary zheng Date: Wed, 2 Sep 2026 17:11:45 -0700 Subject: [PATCH 2/4] build: read the fork options once, outside the grouping loop Eight lines this change is already in the file for. sbt lifts a `.value` written inside a lambda to the top of the task, so the options were read once rather than per suite either way; written where they were, they said the other thing, and sbt warned on it in every build. Co-Authored-By: Claude Opus 5 (1M context) --- build.sbt | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/build.sbt b/build.sbt index 94a5af5606c..6fca342b849 100644 --- a/build.sbt +++ b/build.sbt @@ -227,8 +227,14 @@ lazy val FileService = (project in file("file-service")) Test / fork := true, Test / forkOptions := (Test / forkOptions).value .withWorkingDirectory((ThisBuild / baseDirectory).value), - Test / testGrouping := (Test / definedTests).value.map { suite => - Tests.Group(suite.name, Seq(suite), Tests.SubProcess((Test / forkOptions).value)) + Test / testGrouping := { + // Read once, outside the loop: sbt lifts a `.value` written inside a lambda to + // the top of the task anyway, so leaving it there reads as one lookup per suite + // when it is not, and sbt warns about exactly that. + val forked = (Test / forkOptions).value + (Test / definedTests).value.map { suite => + Tests.Group(suite.name, Seq(suite), Tests.SubProcess(forked)) + } } ) From 4659852dfa474d75e1b71a50528da46ac142861e Mon Sep 17 00:00:00 2001 From: kary zheng Date: Thu, 3 Sep 2026 15:11:24 -0700 Subject: [PATCH 3/4] test(verify): call the plan preparation next door instead of copying it The two helpers here were a second copy of OpExecHarness's, carrying a note that they were kept inline so each harness read on its own and would be consolidated if a third arrived. Two copies of the same twenty lines is already the cost that note was deferring: they have to be changed together, and nothing says so at either site. Preparing the plan does not vary with the executor, so it happens once now. What differs between the harnesses stays here: how the prepared op is run. Co-Authored-By: Claude Opus 5 (1M context) --- .../translator/verify/PyOpExecHarness.scala | 47 ++----------------- 1 file changed, 5 insertions(+), 42 deletions(-) diff --git a/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/PyOpExecHarness.scala b/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/PyOpExecHarness.scala index a6ae7b3b373..fe4f3565c83 100644 --- a/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/PyOpExecHarness.scala +++ b/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/PyOpExecHarness.scala @@ -115,12 +115,15 @@ object PyOpExecHarness extends LazyLogging { // single-op plan that's just every input port the op declares. val externalInputs: Set[(PhysicalOpIdentity, PortIdentity)] = phOp.inputPorts.keys.map(portId => (phOp.id, portId)).toSet - validateInputCoverage(externalInputs, inputs.keySet) + OpExecHarness.validateInputCoverage(externalInputs, inputs.keySet) val inputSchemas: Map[PortIdentity, Schema] = inputs.map { case (portId, path) => portId -> TupleIO.readSchemaSidecar(path) } - val planWithSchemas = propagateExternalSchemas(plan, externalInputs, inputSchemas) + // Preparing the plan is the same work for either executor, so it is done in + // one place; only what runs the prepared op differs between the harnesses. + val planWithSchemas = + OpExecHarness.propagateExternalSchemas(plan, externalInputs, inputSchemas) val phOpWithSchemas = planWithSchemas.operators.head // Output port schemas come from PhysicalPlan.propagateSchema — same @@ -347,46 +350,6 @@ object PyOpExecHarness extends LazyLogging { } finally stream.close() } - // -------------------------------------------------------------------------- - // Schema propagation — mirrors OpExecHarness.propagateExternalSchemas. - // Kept inline (rather than shared) so the two harnesses stay independently - // readable; consolidate if a third harness shows up. - // -------------------------------------------------------------------------- - private def propagateExternalSchemas( - plan: PhysicalPlan, - externalPorts: Set[(PhysicalOpIdentity, PortIdentity)], - schemas: Map[PortIdentity, Schema] - ): PhysicalPlan = { - var acc = PhysicalPlan(operators = Set.empty, links = Set.empty) - plan.topologicalIterator().map(plan.getOperator).foreach { phOp => - val updated = phOp.inputPorts.keys.foldLeft(phOp) { (op, portId) => - if (externalPorts.contains((phOp.id, portId)) && schemas.contains(portId)) { - op.propagateSchema(Some((portId, schemas(portId)))) - } else op - } - acc = acc.addOperator(updated.propagateSchema()) - plan.getUpstreamPhysicalLinks(phOp.id).foreach { link => - acc = acc.addLink(link) - } - } - acc - } - - private def validateInputCoverage( - external: Set[(PhysicalOpIdentity, PortIdentity)], - provided: Set[PortIdentity] - ): Unit = { - val expected = external.map(_._2) - val missing = expected -- provided - val extra = provided -- expected - require( - missing.isEmpty, - s"Missing input fixtures for external ports: $missing (expected $expected)" - ) - if (extra.nonEmpty) { - logger.warn(s"Input fixtures provided for non-external ports (ignored): $extra") - } - } } final class PyOpDriverException( From 1fb3581a5680d52d99e97b712691e99fff3f6eb3 Mon Sep 17 00:00:00 2001 From: kary zheng Date: Thu, 3 Sep 2026 15:46:05 -0700 Subject: [PATCH 4/4] test(verify): say it once, in the shape the code does not already give Three kinds of comment came out. A drawing of the string the code below assembles. A restatement of a branch the reader can see. And the word MVP, which dated the scope to a moment rather than stating it. What replaces them says the same thing shorter, or says what the code cannot: which cases the harness does not drive and why none of them has an operator asking for it. Co-Authored-By: Claude Opus 5 (1M context) --- .../translator/verify/PyOpExecHarness.scala | 19 ++++++++----------- 1 file changed, 8 insertions(+), 11 deletions(-) diff --git a/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/PyOpExecHarness.scala b/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/PyOpExecHarness.scala index fe4f3565c83..097ca127301 100644 --- a/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/PyOpExecHarness.scala +++ b/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/PyOpExecHarness.scala @@ -48,17 +48,14 @@ import scala.sys.process._ * Same Result(outputs, outputSchemas) shape as OpExecHarness so the rest of * the verify pipeline (Comparator, category runners) is harness-agnostic. * - * Scope (MVP, mirrors OpExecHarness's MVP): - * - Single-PhysicalOp plans only. PythonOperatorDescriptor only emits - * either a `sourcePhysicalOp` or a `oneToOnePhysicalOp`, so multi-op - * plans don't exist for Python-native ops today. If that changes, add - * topo-order driving here the way OpExecHarness does. - * - Single output port. UDFOperatorV2 / UDFTableOperator / UDFBatchOperator - * / UDFSourceOperator all yield TupleLike without specifying a port — - * same convention OpExecHarness uses when port is unset. - * - JSONL types: STRING / INTEGER / LONG / DOUBLE / BOOLEAN. TIMESTAMP / - * BINARY / LARGE_BINARY require explicit codecs in both [[TupleIO]] and - * the driver — add when the first operator needs them. + * What it does not drive, each because no Python-native operator asks for it: + * - Multi-PhysicalOp plans. PythonOperatorDescriptor emits either a + * `sourcePhysicalOp` or a `oneToOnePhysicalOp`, so topo-order driving, + * which OpExecHarness has, would have nothing to order. + * - More than one output port. The UDF operator traits yield TupleLike + * without naming a port, the same convention OpExecHarness reads as port 0. + * - JSONL types beyond STRING / INTEGER / LONG / DOUBLE / BOOLEAN. TIMESTAMP + * and the binaries need a codec in both [[TupleIO]] and the driver. */ object PyOpExecHarness extends LazyLogging {