diff --git a/build.sbt b/build.sbt index 035edbb68cc..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)) + } } ) @@ -239,6 +245,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..097ca127301 --- /dev/null +++ b/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/PyOpExecHarness.scala @@ -0,0 +1,366 @@ +/* + * 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. + * + * 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 { + + 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 + OpExecHarness.validateInputCoverage(externalInputs, inputs.keySet) + + val inputSchemas: Map[PortIdentity, Schema] = + inputs.map { case (portId, path) => portId -> TupleIO.readSchemaSidecar(path) } + + // 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 + // 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() + } + +} + +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 + )