From 91f80d2da4c1eb65674a32714ebe24c3faa958ee Mon Sep 17 00:00:00 2001 From: kary zheng Date: Wed, 2 Sep 2026 16:30:44 -0700 Subject: [PATCH 01/19] 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 02/19] 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 03/19] 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 04/19] 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 { From 9f51cd5e14aac5107cd32486751f12765b3ef53c Mon Sep 17 00:00:00 2001 From: kary zheng Date: Fri, 4 Sep 2026 13:22:26 -0700 Subject: [PATCH 05/19] fix(verify): drop an import nothing uses `PhysicalPlan` is named only by `getPhysicalPlan` and by a comment, so scalafix reports the import and the check fails before any test runs. Co-Authored-By: Claude Opus 5 (1M context) --- .../apache/texera/amber/translator/verify/PyOpExecHarness.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 097ca127301..ac51cfd0308 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 @@ -28,7 +28,7 @@ import org.apache.texera.amber.core.virtualidentity.{ PhysicalOpIdentity, WorkflowIdentity } -import org.apache.texera.amber.core.workflow.{PhysicalPlan, PortIdentity} +import org.apache.texera.amber.core.workflow.PortIdentity import org.apache.texera.amber.operator.LogicalOp import org.apache.texera.amber.util.JSONUtils.objectMapper import org.apache.texera.amber.util.python.PythonWorkerPool From c2c89a49e0b510ae8e519c31ef7a373f8dfb0ec4 Mon Sep 17 00:00:00 2001 From: kary zheng Date: Wed, 9 Sep 2026 15:51:41 -0700 Subject: [PATCH 06/19] test(verify): finalize each emitted tuple against the output schema The driver projected the operator's output onto the schema's attribute names and coerced each value with int/float/str. That is looser than the engine, which hands every yielded value to all_output_to_tuple and finalizes the result: an operator putting 1.5 in an INTEGER field, or a row with a field the schema does not have, or one missing a field, all raise there and none of the three raised here. A reference path that accepts output the engine rejects cannot say what the engine would do, so call the engine's own two functions instead of re-deriving them. finalize pickles a BINARY field's object itself and tags it, where the standalone side pickles the same object untagged, so drop the tag on the way out and both sides carry the same payload. Co-Authored-By: Claude Opus 5 (1M context) --- .../src/test/resources/python/py_op_driver.py | 63 ++++++++++--------- 1 file changed, 34 insertions(+), 29 deletions(-) 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 index 1ff2a9000da..b698a07f0b0 100644 --- a/workflow-compiling-service/src/test/resources/python/py_op_driver.py +++ b/workflow-compiling-service/src/test/resources/python/py_op_driver.py @@ -56,7 +56,6 @@ import base64 import inspect import json -import pickle import sys import traceback from pathlib import Path @@ -88,6 +87,7 @@ ) from core.models.schema.schema import Schema as TexeraSchema from core.models.schema.attribute_type import AttributeType, RAW_TYPE_MAPPING + from core.models.table import all_output_to_tuple except ImportError as exc: sys.stderr.write( "py_op_driver.py: failed to import pytexera/pyamber. The harness must " @@ -167,6 +167,12 @@ def _write_schema_sidecar(data_path: Path, schema: TexeraSchema) -> None: # -------------------------------------------------------------------------- # Tuple I/O. JSONL with sidecar — same on-disk shape as TupleIO on the JVM. # -------------------------------------------------------------------------- +# Prefix Tuple.cast_to_schema writes in front of an object it pickles into a +# BINARY field. Nothing in the repo reads it back, so it is a tag, not part of +# the value. +_CAST_PICKLE_MARKER = b"pickle " + + def _coerce_field(raw: Any, attr_type: AttributeType) -> Any: """Coerce a JSON-decoded field to the type the schema expects.""" if raw is None: @@ -225,31 +231,23 @@ def _emit_as_dicts( ) -> 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) + output schema's attribute names. + + Both steps are the engine's own: ``DataProcessor`` hands each yielded + value to ``all_output_to_tuple`` (a DataFrame, a Series, a dict, a list + of those, a Tuple, or None) and finalizes every tuple that comes back + against the output port's schema. Running the same two calls here is what + makes this path reject what the engine rejects. A row carrying a field the + schema does not have, a row missing one, and a value the cast cannot bring + to its declared type all raise here exactly as they would in the worker, + rather than being projected or coerced away. """ - 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}" - ) + for tup in all_output_to_tuple(item): + if tup is None: + continue + tup.finalize(schema) + yield dict(tup.as_key_value_pairs()) def _jsonify(value: Any, attr_type: AttributeType) -> Any: @@ -279,11 +277,18 @@ def _jsonify(value: Any, attr_type: AttributeType) -> Any: 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) + # Trained-model / object columns: base64 so the value survives the + # 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. + # + # The finalize in _emit_as_dicts has already pickled anything that was + # not bytes, prefixing it with the marker the engine's cast writes. The + # standalone path pickles the same object without one, so strip it and + # both sides carry the same payload. + raw = value + if raw.startswith(_CAST_PICKLE_MARKER): + raw = raw[len(_CAST_PICKLE_MARKER) :] return base64.b64encode(raw).decode("ascii") if attr_type == AttributeType.TIMESTAMP: # Emit the same JDBC string java.sql.Timestamp.toString produces (>=1 From e59d0d824bb475d4e61ff43cadda41c244ea84ea Mon Sep 17 00:00:00 2001 From: kary zheng Date: Thu, 10 Sep 2026 13:02:13 -0700 Subject: [PATCH 07/19] test(verify): finalize each yield where it is yielded The driver collected everything the operator yielded and converted the list afterwards. `DataProcessor` reads a yield where it happens, turning it into tuples and finalizing them on the spot, so a UDF that yields one dict twice and changes it in between is recorded as the two values it produced. The driver recorded the last one twice. The operator now takes a callback and hands each value over at the yield, before the generator is advanced. Co-Authored-By: Claude Opus 5 (1M context) --- .../src/test/resources/python/py_op_driver.py | 57 ++++++----- .../verify/PyOpExecHarnessSpec.scala | 99 +++++++++++++++++++ 2 files changed, 133 insertions(+), 23 deletions(-) create mode 100644 workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/PyOpExecHarnessSpec.scala 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 index b698a07f0b0..c8eda83735a 100644 --- a/workflow-compiling-service/src/test/resources/python/py_op_driver.py +++ b/workflow-compiling-service/src/test/resources/python/py_op_driver.py @@ -59,7 +59,7 @@ import sys import traceback from pathlib import Path -from typing import Any, Iterable, Iterator, List, Mapping, Sequence +from typing import Any, Callable, Iterable, Iterator, List, Mapping, Sequence import pandas as pd @@ -379,16 +379,20 @@ def _run_operator( 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. + consume: "Callable[[Any], None]", +) -> None: """ - emitted: List[Any] = [] + Drive the operator's lifecycle, handing each emitted value to ``consume``. + ``consume`` is called AT the yield, before the generator is advanced, + because that is where ``DataProcessor`` reads one: it turns each yielded + value into tuples and finalizes them on the spot. Collecting the yields + into a list first records a UDF that yields one dict twice, mutating it in + between, as the same value twice rather than as the two it produced. + + 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. + """ op.open() try: if is_source: @@ -396,22 +400,20 @@ def _run_operator( # 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 + consume(item) + return 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) + consume(item) for item in op.on_finish(port): if item is not None: - emitted.append(item) + consume(item) finally: op.close() - return emitted - # -------------------------------------------------------------------------- # Entry point. @@ -458,20 +460,29 @@ def run_config(config: Mapping[str, Any]) -> None: 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) + + if not outputs: + _run_operator(op_instance, is_source, port_order, inputs_by_port, lambda _item: None) + return + + out_entry = outputs[0] + out_path = Path(out_entry["dataPath"]) + out_schema = _schema_from_dict(out_entry["schema"]) + rows: List["dict[str, Any]"] = [] + _run_operator( + op_instance, + is_source, + port_order, + inputs_by_port, + lambda item: rows.extend(_emit_as_dicts([item], out_schema)), + ) + _write_tuples(out_path, rows, out_schema) def main(argv: Sequence[str]) -> int: diff --git a/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/PyOpExecHarnessSpec.scala b/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/PyOpExecHarnessSpec.scala new file mode 100644 index 00000000000..1843c3cf21a --- /dev/null +++ b/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/PyOpExecHarnessSpec.scala @@ -0,0 +1,99 @@ +/* + * 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 org.apache.texera.amber.core.tuple.{Attribute, AttributeType, Schema} +import org.apache.texera.amber.core.workflow.{InputPort, OutputPort, PortIdentity} +import org.apache.texera.amber.operator.PythonOperatorDescriptor +import org.apache.texera.amber.operator.metadata.{OperatorGroupConstants, OperatorInfo} +import org.apache.texera.amber.translator.verify.tags.IntegrationTest +import org.scalatest.flatspec.AnyFlatSpec +import org.scalatest.matchers.should.Matchers + +import java.nio.file.Files + +/** The reference path's own behaviour, where it can differ from the worker's. + * + * Tagged @IntegrationTest: the harness forks Python. + */ +@IntegrationTest +class PyOpExecHarnessSpec extends AnyFlatSpec with Matchers { + + /** Yields one dict twice, changing it in between. `DataProcessor` reads a + * yield where it happens, so the engine records 1 and then 2; a driver that + * collects the yields first and converts them afterwards sees the same + * object twice and records 2 twice. + */ + private class MutatingYieldOpDesc extends PythonOperatorDescriptor { + override def operatorInfo: OperatorInfo = + OperatorInfo( + userFriendlyName = "Mutating Yield", + operatorDescription = "yields one dict twice, changing it in between", + operatorGroupName = OperatorGroupConstants.UTILITY_GROUP, + inputPorts = List(InputPort()), + outputPorts = List(OutputPort()) + ) + + override def getOutputSchemas( + inputSchemas: Map[PortIdentity, Schema] + ): Map[PortIdentity, Schema] = + Map( + operatorInfo.outputPorts.head.id -> Schema().add(new Attribute("x", AttributeType.INTEGER)) + ) + + override def generatePythonCode(): String = + """from pytexera import * + | + |class ProcessTableOperator(UDFTableOperator): + | + | @overrides + | def process_table(self, table: Table, port: int) -> Iterator[Optional[TableLike]]: + | row = {"x": 1} + | yield row + | row["x"] = 2 + | yield row + |""".stripMargin + } + + "PyOpExecHarness" should "record each yield as it was yielded" in { + val dir = Files.createTempDirectory("py-op-harness-") + val inputSchema = Schema().add(new Attribute("seed", AttributeType.INTEGER)) + val input = CuratedHandlers.writeFixture( + dir.resolve("input_port_0.jsonl"), + Seq(("seed", AttributeType.INTEGER)), + Seq(Seq[Any](1)) + ) + + val result = PyOpExecHarness.execute( + new MutatingYieldOpDesc, + inputs = Map(PortIdentity(0) -> input), + outputDir = dir.resolve("actual") + ) + + val out = result.outputs(PortIdentity(0)) + val values = TupleIO + .readTuples(out, TupleIO.readSchemaSidecar(out)) + .map(_.getField[Integer]("x")) + .toSeq + withClue(s"schema was $inputSchema\n") { + values shouldBe Seq(1, 2) + } + } +} From cf3544a69e0fe43f025bf774543481b51368f892 Mon Sep 17 00:00:00 2001 From: kary zheng Date: Thu, 10 Sep 2026 15:57:06 -0700 Subject: [PATCH 08/19] test(verify): tell the operator what an empty port was carrying DataProcessor writes each input port's schema into the executor before calling on_finish, so a TableOperator handed no rows can still say what its columns were. The driver stands in for DataProcessor here, so it has to do the same or the reference path answers on a frame of no columns where the engine answers on the declared ones. Co-Authored-By: Claude Opus 5 (1M context) --- .../src/test/resources/python/py_op_driver.py | 6 ++++++ 1 file changed, 6 insertions(+) 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 index c8eda83735a..7568956ffcd 100644 --- a/workflow-compiling-service/src/test/resources/python/py_op_driver.py +++ b/workflow-compiling-service/src/test/resources/python/py_op_driver.py @@ -445,10 +445,12 @@ def run_config(config: Mapping[str, Any]) -> None: port_order: Sequence[int] = list(config.get("portOrder", [])) inputs_by_port: "dict[int, List[Tuple]]" = {} + schemas_by_port: "dict[int, TexeraSchema]" = {} for entry in config.get("inputs", []): port = int(entry["portIndex"]) data_path = Path(entry["dataPath"]) schema = _read_schema_sidecar(data_path) + schemas_by_port[port] = schema inputs_by_port[port] = _read_tuples(data_path, schema) # Default port order: sorted by index. Matches OpExecHarness's fallback @@ -460,6 +462,10 @@ def run_config(config: Mapping[str, Any]) -> None: op_class = _discover_operator_class(namespace) op_instance = op_class() + # What DataProcessor hands the executor at EndChannel: a port that carried + # no rows still says what its columns were. + op_instance.input_schemas.update(schemas_by_port) + outputs = config.get("outputs", []) if len(outputs) > 1: raise NotImplementedError( From 0c9c6f4202bca784afe7d569e436950e8056ab73 Mon Sep 17 00:00:00 2001 From: kary zheng Date: Thu, 10 Sep 2026 16:45:49 -0700 Subject: [PATCH 09/19] test(verify): cut the doc to what the code cannot say It restated its own signature and listed a codec table that lives in TupleIO. The reasons a reader cannot derive stay. Co-Authored-By: Claude Opus 5 (1M context) --- .../translator/verify/PyOpExecHarness.scala | 16 ++++++---------- 1 file changed, 6 insertions(+), 10 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 ac51cfd0308..dda2c9acf78 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 @@ -45,17 +45,13 @@ import scala.sys.process._ * 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. + * It returns the same Result shape, so the rest of the pipeline does not care + * which harness ran. * - * 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. + * One PhysicalOp with one output port is all it drives, which is all a + * `PythonOperatorDescriptor` emits: the UDF traits yield a TupleLike without + * naming a port. TIMESTAMP and the binaries are out until both [[TupleIO]] and + * the driver carry a codec for them. */ object PyOpExecHarness extends LazyLogging { From 76e3e822fe81f7b381a353694f6d339e0d11fede Mon Sep 17 00:00:00 2001 From: kary zheng Date: Thu, 10 Sep 2026 21:22:52 -0700 Subject: [PATCH 10/19] test(verify): write this spec's one row without the curated handlers The handlers are written nine steps later in the set, and this borrowed a fixture writer from them for a single one-column row. Writing it with TupleIO, which the spec already reads with, leaves the order free of the one edge that pointed backwards. Co-Authored-By: Claude Opus 5 (1M context) --- .../amber/translator/verify/PyOpExecHarnessSpec.scala | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/PyOpExecHarnessSpec.scala b/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/PyOpExecHarnessSpec.scala index 1843c3cf21a..c6d2db20368 100644 --- a/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/PyOpExecHarnessSpec.scala +++ b/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/PyOpExecHarnessSpec.scala @@ -19,7 +19,7 @@ package org.apache.texera.amber.translator.verify -import org.apache.texera.amber.core.tuple.{Attribute, AttributeType, Schema} +import org.apache.texera.amber.core.tuple.{Attribute, AttributeType, Schema, Tuple} import org.apache.texera.amber.core.workflow.{InputPort, OutputPort, PortIdentity} import org.apache.texera.amber.operator.PythonOperatorDescriptor import org.apache.texera.amber.operator.metadata.{OperatorGroupConstants, OperatorInfo} @@ -75,11 +75,9 @@ class PyOpExecHarnessSpec extends AnyFlatSpec with Matchers { "PyOpExecHarness" should "record each yield as it was yielded" in { val dir = Files.createTempDirectory("py-op-harness-") val inputSchema = Schema().add(new Attribute("seed", AttributeType.INTEGER)) - val input = CuratedHandlers.writeFixture( - dir.resolve("input_port_0.jsonl"), - Seq(("seed", AttributeType.INTEGER)), - Seq(Seq[Any](1)) - ) + val input = dir.resolve("input_port_0.jsonl") + val seed = Tuple.builder(inputSchema).add("seed", AttributeType.INTEGER, Int.box(1)).build() + TupleIO.writeTuples(input, Iterator(seed), inputSchema) val result = PyOpExecHarness.execute( new MutatingYieldOpDesc, From f4cdf6a37ce8e8615bb0eed48dfd6a5591f94b3f Mon Sep 17 00:00:00 2001 From: kary zheng Date: Fri, 18 Sep 2026 13:16:56 -0700 Subject: [PATCH 11/19] test(verify): carry a large binary column through the Python harness The schema sidecar names large_binary and the driver's type map reads it, but neither the input coercion nor the output conversion knew the type, so an operator handed such a column could not be run at all. A large binary is a reference: the bytes live in S3 and the field is the s3:// URI that points at them. Both harnesses now carry that URI, which is what the worker sends and what LargeBinary holds on the other side, and building one touches no storage. Co-Authored-By: Claude Opus 5 (1M context) --- .../src/test/resources/python/py_op_driver.py | 17 +++- .../translator/verify/OpExecHarness.scala | 8 +- .../translator/verify/PyOpExecHarness.scala | 3 +- .../verify/PyOpExecHarnessSpec.scala | 98 ++++++++++++++++++- 4 files changed, 120 insertions(+), 6 deletions(-) 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 index 7568956ffcd..62c39e20ab4 100644 --- a/workflow-compiling-service/src/test/resources/python/py_op_driver.py +++ b/workflow-compiling-service/src/test/resources/python/py_op_driver.py @@ -82,6 +82,7 @@ UDFSourceOperator, UDFTableOperator, Union as PyUnion, # noqa: F401 + largebinary, logger as pytexera_logger, # noqa: F401 overrides, # noqa: F401 ) @@ -189,13 +190,19 @@ def _coerce_field(raw: Any, attr_type: AttributeType) -> Any: return bool(raw) if attr_type == AttributeType.BINARY: return base64.b64decode(raw) + if attr_type == AttributeType.LARGE_BINARY: + # The bytes live in S3; the field is the reference to them, which the + # worker hands the operator as a largebinary. Building one from a URI + # touches no storage, so a run that only carries the reference through + # needs none either. + return largebinary(str(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. + # Every type the schema can name is read above. A new one fails loud rather + # than passing a string through as though it had been read. raise NotImplementedError( f"py_op_driver: reading attribute type {attr_type!r} from JSONL is " f"not implemented yet" @@ -290,6 +297,12 @@ def _jsonify(value: Any, attr_type: AttributeType) -> Any: if raw.startswith(_CAST_PICKLE_MARKER): raw = raw[len(_CAST_PICKLE_MARKER) :] return base64.b64encode(raw).decode("ascii") + if attr_type == AttributeType.LARGE_BINARY: + # The reference, not the object it points at: the URI is the whole field + # on the wire, which is what the worker sends and what LargeBinary holds + # on the other side. An operator that wrote the column itself may hand + # back the URI as plain text. + return value.uri if isinstance(value, largebinary) else str(value) 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 diff --git a/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/OpExecHarness.scala b/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/OpExecHarness.scala index 01e48e40da2..cbfbcc34021 100644 --- a/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/OpExecHarness.scala +++ b/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/OpExecHarness.scala @@ -22,7 +22,7 @@ package org.apache.texera.amber.translator.verify import com.fasterxml.jackson.databind.node.ObjectNode import com.typesafe.scalalogging.LazyLogging import org.apache.texera.amber.core.executor.{ExecFactory, OpExecWithClassName, OperatorExecutor} -import org.apache.texera.amber.core.tuple.{AttributeType, Schema, Tuple, TupleLike} +import org.apache.texera.amber.core.tuple.{AttributeType, LargeBinary, Schema, Tuple, TupleLike} import org.apache.texera.amber.core.virtualidentity.{ ExecutionIdentity, PhysicalOpIdentity, @@ -383,6 +383,10 @@ object TupleIO { case AttributeType.BOOLEAN => Boolean.box(fieldNode.asBoolean()) case AttributeType.BINARY => Base64.getDecoder.decode(fieldNode.asText()) + // The bytes live in S3 and the field is the reference to them, so + // the s3:// URI is the whole value on either side of the file. + case AttributeType.LARGE_BINARY => + new LargeBinary(fieldNode.asText()) // Timestamps round-trip through the JDBC string form // ("yyyy-mm-dd hh:mm:ss[.f]"), the exact inverse of Timestamp.toString // below — timezone-free, so no shift across write/read. The Python @@ -425,6 +429,8 @@ object TupleIO { attr.getName, Base64.getEncoder.encodeToString(v.asInstanceOf[Array[Byte]]) ) + case AttributeType.LARGE_BINARY => + node.put(attr.getName, v.asInstanceOf[LargeBinary].getUri) case AttributeType.TIMESTAMP => node.put(attr.getName, v.asInstanceOf[Timestamp].toString) case other => 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 dda2c9acf78..96ecaca0ab2 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 @@ -50,8 +50,7 @@ import scala.sys.process._ * * One PhysicalOp with one output port is all it drives, which is all a * `PythonOperatorDescriptor` emits: the UDF traits yield a TupleLike without - * naming a port. TIMESTAMP and the binaries are out until both [[TupleIO]] and - * the driver carry a codec for them. + * naming a port. */ object PyOpExecHarness extends LazyLogging { diff --git a/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/PyOpExecHarnessSpec.scala b/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/PyOpExecHarnessSpec.scala index c6d2db20368..77e6cffab8d 100644 --- a/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/PyOpExecHarnessSpec.scala +++ b/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/PyOpExecHarnessSpec.scala @@ -19,7 +19,7 @@ package org.apache.texera.amber.translator.verify -import org.apache.texera.amber.core.tuple.{Attribute, AttributeType, Schema, Tuple} +import org.apache.texera.amber.core.tuple.{Attribute, AttributeType, LargeBinary, Schema, Tuple} import org.apache.texera.amber.core.workflow.{InputPort, OutputPort, PortIdentity} import org.apache.texera.amber.operator.PythonOperatorDescriptor import org.apache.texera.amber.operator.metadata.{OperatorGroupConstants, OperatorInfo} @@ -72,6 +72,102 @@ class PyOpExecHarnessSpec extends AnyFlatSpec with Matchers { |""".stripMargin } + /** Reads a LARGE_BINARY column and writes out the reference it holds. + * + * The bytes of a large binary live in S3 and the field is the s3:// URI that + * points at them, so an operator can be handed one, and can hand one back, + * without any storage being reachable. + */ + private class LargeBinaryOpDesc(readsInput: Boolean) extends PythonOperatorDescriptor { + override def operatorInfo: OperatorInfo = + OperatorInfo( + userFriendlyName = "Large Binary", + operatorDescription = "reads or writes a reference to a large binary", + operatorGroupName = OperatorGroupConstants.UTILITY_GROUP, + inputPorts = List(InputPort()), + outputPorts = List(OutputPort()) + ) + + override def getOutputSchemas( + inputSchemas: Map[PortIdentity, Schema] + ): Map[PortIdentity, Schema] = + Map( + operatorInfo.outputPorts.head.id -> + (if (readsInput) Schema().add(new Attribute("uri", AttributeType.STRING)) + else Schema().add(new Attribute("blob", AttributeType.LARGE_BINARY))) + ) + + override def generatePythonCode(): String = + if (readsInput) + """from pytexera import * + | + |class ProcessTupleOperator(UDFOperatorV2): + | + | @overrides + | def process_tuple(self, tuple_: Tuple, port: int) -> Iterator[Optional[TupleLike]]: + | yield {"uri": tuple_["blob"].uri} + |""".stripMargin + else + """from pytexera import * + | + |class ProcessTupleOperator(UDFOperatorV2): + | + | @overrides + | def process_tuple(self, tuple_: Tuple, port: int) -> Iterator[Optional[TupleLike]]: + | yield {"blob": largebinary(tuple_["uri"])} + |""".stripMargin + } + + private val someObject = "s3://a-bucket/a/large/object" + + // The schema sidecar names large_binary, so the harness has to carry it in as + // well as out. The operator reads the reference itself, which a column handed + // over as plain text could not answer for. + "PyOpExecHarness" should "hand a large binary column to the operator as a reference" in { + val dir = Files.createTempDirectory("py-op-harness-large-binary-in-") + val inputSchema = Schema().add(new Attribute("blob", AttributeType.LARGE_BINARY)) + val input = dir.resolve("input_port_0.jsonl") + val row = Tuple + .builder(inputSchema) + .add("blob", AttributeType.LARGE_BINARY, new LargeBinary(someObject)) + .build() + TupleIO.writeTuples(input, Iterator(row), inputSchema) + + val result = PyOpExecHarness.execute( + new LargeBinaryOpDesc(readsInput = true), + inputs = Map(PortIdentity(0) -> input), + outputDir = dir.resolve("actual") + ) + + val out = result.outputs(PortIdentity(0)) + TupleIO + .readTuples(out, TupleIO.readSchemaSidecar(out)) + .map(_.getField[String]("uri")) + .toSeq shouldBe Seq(someObject) + } + + it should "write back a large binary the operator made" in { + val dir = Files.createTempDirectory("py-op-harness-large-binary-out-") + val inputSchema = Schema().add(new Attribute("uri", AttributeType.STRING)) + val input = dir.resolve("input_port_0.jsonl") + val row = Tuple.builder(inputSchema).add("uri", AttributeType.STRING, someObject).build() + TupleIO.writeTuples(input, Iterator(row), inputSchema) + + val result = PyOpExecHarness.execute( + new LargeBinaryOpDesc(readsInput = false), + inputs = Map(PortIdentity(0) -> input), + outputDir = dir.resolve("actual") + ) + + val out = result.outputs(PortIdentity(0)) + // Read back as a LargeBinary and not as the text of one: the column is + // declared large_binary, and this is the type the rest of the pipeline gets. + TupleIO + .readTuples(out, TupleIO.readSchemaSidecar(out)) + .map(_.getField[LargeBinary]("blob").getUri) + .toSeq shouldBe Seq(someObject) + } + "PyOpExecHarness" should "record each yield as it was yielded" in { val dir = Files.createTempDirectory("py-op-harness-") val inputSchema = Schema().add(new Attribute("seed", AttributeType.INTEGER)) From 153db9c9f6b04c4c4f635710b610838e14a156c0 Mon Sep 17 00:00:00 2001 From: kary zheng Date: Tue, 22 Sep 2026 15:53:14 -0700 Subject: [PATCH 12/19] test(verify): stop handing the driver each port's declared schema apache/texera#8488 is closed: a table with no rows carries no schema by design, so there is nothing for the driver to hand the operator and no engine behavior here to match. Co-Authored-By: Claude Opus 5 (1M context) --- .../src/test/resources/python/py_op_driver.py | 9 +-------- 1 file changed, 1 insertion(+), 8 deletions(-) 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 index 62c39e20ab4..103567d1396 100644 --- a/workflow-compiling-service/src/test/resources/python/py_op_driver.py +++ b/workflow-compiling-service/src/test/resources/python/py_op_driver.py @@ -458,13 +458,10 @@ def run_config(config: Mapping[str, Any]) -> None: port_order: Sequence[int] = list(config.get("portOrder", [])) inputs_by_port: "dict[int, List[Tuple]]" = {} - schemas_by_port: "dict[int, TexeraSchema]" = {} for entry in config.get("inputs", []): port = int(entry["portIndex"]) data_path = Path(entry["dataPath"]) - schema = _read_schema_sidecar(data_path) - schemas_by_port[port] = schema - inputs_by_port[port] = _read_tuples(data_path, schema) + inputs_by_port[port] = _read_tuples(data_path, _read_schema_sidecar(data_path)) # Default port order: sorted by index. Matches OpExecHarness's fallback # when getInputPortDependencyPairs is empty. @@ -475,10 +472,6 @@ def run_config(config: Mapping[str, Any]) -> None: op_class = _discover_operator_class(namespace) op_instance = op_class() - # What DataProcessor hands the executor at EndChannel: a port that carried - # no rows still says what its columns were. - op_instance.input_schemas.update(schemas_by_port) - outputs = config.get("outputs", []) if len(outputs) > 1: raise NotImplementedError( From b210ba197146687baea524bdd8912272f4eb7156 Mon Sep 17 00:00:00 2001 From: kary zheng Date: Wed, 23 Sep 2026 16:16:47 -0700 Subject: [PATCH 13/19] test(verify): hand a pickled binary cell to the operator as the object The worker's ArrowTableTupleProvider unpickles a BINARY cell that starts with the cast's pickle marker, so a model column reaches the next operator as a model. The driver handed the operator the bytes instead, and a predict on them raised AttributeError. Decoding the cell alone was not enough. The driver then finalized each input tuple, and the cast in finalize pickles any non-bytes value back into bytes. The worker's InputManager attaches the schema without finalizing, so the driver now builds input tuples the same way. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../src/test/resources/python/py_op_driver.py | 19 ++++-- .../verify/PyOpExecHarnessSpec.scala | 64 +++++++++++++++++++ 2 files changed, 77 insertions(+), 6 deletions(-) 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 index 103567d1396..eed16e1ee10 100644 --- a/workflow-compiling-service/src/test/resources/python/py_op_driver.py +++ b/workflow-compiling-service/src/test/resources/python/py_op_driver.py @@ -56,6 +56,7 @@ import base64 import inspect import json +import pickle import sys import traceback from pathlib import Path @@ -169,8 +170,8 @@ def _write_schema_sidecar(data_path: Path, schema: TexeraSchema) -> None: # Tuple I/O. JSONL with sidecar — same on-disk shape as TupleIO on the JVM. # -------------------------------------------------------------------------- # Prefix Tuple.cast_to_schema writes in front of an object it pickles into a -# BINARY field. Nothing in the repo reads it back, so it is a tag, not part of -# the value. +# BINARY field. The worker's ArrowTableTupleProvider unpickles a cell that +# carries it, which is how a model column reaches the next operator as a model. _CAST_PICKLE_MARKER = b"pickle " @@ -189,7 +190,12 @@ def _coerce_field(raw: Any, attr_type: AttributeType) -> Any: if attr_type == AttributeType.BOOL: return bool(raw) if attr_type == AttributeType.BINARY: - return base64.b64decode(raw) + # Same test and slice as ArrowTableTupleProvider, so an operator gets + # the object the worker would give it and any other bytes stay bytes. + value = base64.b64decode(raw) + if value[:6] == b"pickle": + return pickle.loads(value[len(_CAST_PICKLE_MARKER) :]) + return value if attr_type == AttributeType.LARGE_BINARY: # The bytes live in S3; the field is the reference to them, which the # worker hands the operator as a largebinary. Building one from a URI @@ -227,9 +233,10 @@ def _read_tuples(data_path: Path, schema: TexeraSchema) -> List[Tuple]: 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) + # Built the way the worker's InputManager builds it: the schema is + # attached, not finalized against. A finalize would cast, and the + # cast pickles a decoded model straight back into bytes. + rows.append(Tuple(field_data, schema=schema)) return rows diff --git a/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/PyOpExecHarnessSpec.scala b/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/PyOpExecHarnessSpec.scala index 77e6cffab8d..c8de9924e2a 100644 --- a/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/PyOpExecHarnessSpec.scala +++ b/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/PyOpExecHarnessSpec.scala @@ -118,6 +118,70 @@ class PyOpExecHarnessSpec extends AnyFlatSpec with Matchers { |""".stripMargin } + /** Reports the Python type of the BINARY cell it is handed. */ + private class BinaryTypeOpDesc extends PythonOperatorDescriptor { + override def operatorInfo: OperatorInfo = + OperatorInfo( + userFriendlyName = "Binary Type", + operatorDescription = "reports the type of a binary cell", + operatorGroupName = OperatorGroupConstants.UTILITY_GROUP, + inputPorts = List(InputPort()), + outputPorts = List(OutputPort()) + ) + + override def getOutputSchemas( + inputSchemas: Map[PortIdentity, Schema] + ): Map[PortIdentity, Schema] = + Map( + operatorInfo.outputPorts.head.id -> Schema().add( + new Attribute("kind", AttributeType.STRING) + ) + ) + + override def generatePythonCode(): String = + """from pytexera import * + | + |class ProcessTupleOperator(UDFOperatorV2): + | + | @overrides + | def process_tuple(self, tuple_: Tuple, port: int) -> Iterator[Optional[TupleLike]]: + | yield {"kind": type(tuple_["blob"]).__name__} + |""".stripMargin + } + + private def binaryCellKind(cell: Array[Byte]): String = { + val dir = Files.createTempDirectory("py-op-harness-binary-") + val inputSchema = Schema().add(new Attribute("blob", AttributeType.BINARY)) + val input = dir.resolve("input_port_0.jsonl") + val row = Tuple.builder(inputSchema).add("blob", AttributeType.BINARY, cell).build() + TupleIO.writeTuples(input, Iterator(row), inputSchema) + + val result = PyOpExecHarness.execute( + new BinaryTypeOpDesc, + inputs = Map(PortIdentity(0) -> input), + outputDir = dir.resolve("actual") + ) + + val out = result.outputs(PortIdentity(0)) + TupleIO + .readTuples(out, TupleIO.readSchemaSidecar(out)) + .map(_.getField[String]("kind")) + .toSeq + .head + } + + // A model column arrives as the marker followed by the pickle, and the worker + // unpickles it before the operator sees it. This is pickle.dumps(["a"], + // protocol=0), which stands in for a fitted estimator. + "PyOpExecHarness" should "hand a pickled binary cell to the operator as the object" in { + val pickled = "pickle ".getBytes("US-ASCII") ++ "(lp0\nVa\np1\na.".getBytes("US-ASCII") + binaryCellKind(pickled) shouldBe "list" + } + + it should "hand any other binary cell to the operator as bytes" in { + binaryCellKind(Array[Byte](0, 1, 2)) shouldBe "bytes" + } + private val someObject = "s3://a-bucket/a/large/object" // The schema sidecar names large_binary, so the harness has to carry it in as From 1f0a0ef021964592fc9c2bc0c5647e955cd50da8 Mon Sep 17 00:00:00 2001 From: kary zheng Date: Wed, 23 Sep 2026 23:48:53 -0700 Subject: [PATCH 14/19] test(verify): hand a real fitted model to the operator in the harness spec The pickled-cell test used a pickled list in place of a model. It now pickles a fitted DecisionTreeClassifier with the interpreter the harness runs and has the operator call predict on it, which is the call that raised AttributeError when the driver handed over bytes. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../translator/verify/PyOpExecHarness.scala | 2 +- .../verify/PyOpExecHarnessSpec.scala | 43 ++++++++++++++++--- 2 files changed, 37 insertions(+), 8 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 96ecaca0ab2..4dc669c75e8 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 @@ -299,7 +299,7 @@ object PyOpExecHarness extends LazyLogging { // -------------------------------------------------------------------------- // Resolution helpers. // -------------------------------------------------------------------------- - private def resolvePython(): String = + private[verify] def resolvePython(): String = sys.env.get("UDF_PYTHON_PATH").filter(_.nonEmpty).getOrElse("python3.12") /** diff --git a/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/PyOpExecHarnessSpec.scala b/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/PyOpExecHarnessSpec.scala index c8de9924e2a..2072330bb1d 100644 --- a/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/PyOpExecHarnessSpec.scala +++ b/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/PyOpExecHarnessSpec.scala @@ -28,6 +28,8 @@ import org.scalatest.flatspec.AnyFlatSpec import org.scalatest.matchers.should.Matchers import java.nio.file.Files +import java.util.Base64 +import scala.sys.process._ /** The reference path's own behaviour, where it can differ from the worker's. * @@ -149,7 +151,36 @@ class PyOpExecHarnessSpec extends AnyFlatSpec with Matchers { |""".stripMargin } - private def binaryCellKind(cell: Array[Byte]): String = { + /** Calls predict on the BINARY cell it is handed, as a model's consumer does. */ + private class ModelPredictOpDesc extends BinaryTypeOpDesc { + override def generatePythonCode(): String = + """from pytexera import * + | + |class ProcessTupleOperator(UDFOperatorV2): + | + | @overrides + | def process_tuple(self, tuple_: Tuple, port: int) -> Iterator[Optional[TupleLike]]: + | yield {"kind": str(tuple_["blob"].predict([[1.0]])[0])} + |""".stripMargin + } + + // A fitted DecisionTreeClassifier behind the marker the cast writes, pickled + // by the interpreter the harness runs so the two share a sklearn. + private def fittedTreeCell(): Array[Byte] = { + val script = + """import base64, pickle, sys + |from sklearn.tree import DecisionTreeClassifier + |model = DecisionTreeClassifier(random_state=0).fit([[0.0], [1.0]], [0, 1]) + |sys.stdout.write(base64.b64encode(b"pickle " + pickle.dumps(model)).decode("ascii")) + |""".stripMargin + val encoded = Process(Seq(PyOpExecHarness.resolvePython(), "-c", script)).!! + Base64.getDecoder.decode(encoded.trim) + } + + private def binaryCellKind( + cell: Array[Byte], + opDesc: PythonOperatorDescriptor = new BinaryTypeOpDesc + ): String = { val dir = Files.createTempDirectory("py-op-harness-binary-") val inputSchema = Schema().add(new Attribute("blob", AttributeType.BINARY)) val input = dir.resolve("input_port_0.jsonl") @@ -157,7 +188,7 @@ class PyOpExecHarnessSpec extends AnyFlatSpec with Matchers { TupleIO.writeTuples(input, Iterator(row), inputSchema) val result = PyOpExecHarness.execute( - new BinaryTypeOpDesc, + opDesc, inputs = Map(PortIdentity(0) -> input), outputDir = dir.resolve("actual") ) @@ -171,11 +202,9 @@ class PyOpExecHarnessSpec extends AnyFlatSpec with Matchers { } // A model column arrives as the marker followed by the pickle, and the worker - // unpickles it before the operator sees it. This is pickle.dumps(["a"], - // protocol=0), which stands in for a fitted estimator. - "PyOpExecHarness" should "hand a pickled binary cell to the operator as the object" in { - val pickled = "pickle ".getBytes("US-ASCII") ++ "(lp0\nVa\np1\na.".getBytes("US-ASCII") - binaryCellKind(pickled) shouldBe "list" + // unpickles it before the operator sees it, so predict works on it. + "PyOpExecHarness" should "hand a pickled model cell to the operator as the model" in { + binaryCellKind(fittedTreeCell(), new ModelPredictOpDesc) shouldBe "1" } it should "hand any other binary cell to the operator as bytes" in { From a2f966878105ea335254cfb4c1ab331b2a307abe Mon Sep 17 00:00:00 2001 From: kary zheng Date: Sat, 26 Sep 2026 20:09:37 -0700 Subject: [PATCH 15/19] test(verify): hand a Python operator the timestamp the engine sends, and refuse a row the schema does not describe Co-Authored-By: Claude Opus 5.5 (1M context) --- .../src/test/resources/python/py_op_driver.py | 14 ++- .../verify/PyOpExecHarnessSpec.scala | 94 +++++++++++++++++++ 2 files changed, 104 insertions(+), 4 deletions(-) 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 index eed16e1ee10..bdffe300f10 100644 --- a/workflow-compiling-service/src/test/resources/python/py_op_driver.py +++ b/workflow-compiling-service/src/test/resources/python/py_op_driver.py @@ -54,6 +54,7 @@ from __future__ import annotations import base64 +import datetime import inspect import json import pickle @@ -203,10 +204,15 @@ def _coerce_field(raw: Any, attr_type: AttributeType) -> Any: # needs none either. return largebinary(str(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() + # TupleIO writes java.sql.Timestamp.toString ("YYYY-MM-DD HH:MM:SS[.f]"). + # The engine hands a worker the wall clock as a zoneless Arrow + # microsecond timestamp, which reaches the operator as a naive datetime + # cut to the microsecond. Read here without pandas, whose nanoseconds + # end in 2262 where the engine's years go on. + date, _, fraction = str(raw).partition(".") + return datetime.datetime.strptime(date, "%Y-%m-%d %H:%M:%S").replace( + microsecond=int((fraction + "000000")[:6]) + ) # Every type the schema can name is read above. A new one fails loud rather # than passing a string through as though it had been read. raise NotImplementedError( diff --git a/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/PyOpExecHarnessSpec.scala b/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/PyOpExecHarnessSpec.scala index 2072330bb1d..d88f1a65d6a 100644 --- a/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/PyOpExecHarnessSpec.scala +++ b/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/PyOpExecHarnessSpec.scala @@ -28,6 +28,7 @@ import org.scalatest.flatspec.AnyFlatSpec import org.scalatest.matchers.should.Matchers import java.nio.file.Files +import java.sql.Timestamp import java.util.Base64 import scala.sys.process._ @@ -74,6 +75,48 @@ class PyOpExecHarnessSpec extends AnyFlatSpec with Matchers { |""".stripMargin } + /** Reports the timestamp it was handed, as Python prints it. */ + private class TimestampTextOpDesc extends BinaryTypeOpDesc { + override def generatePythonCode(): String = + """from pytexera import * + | + |class ProcessTupleOperator(UDFOperatorV2): + | + | @overrides + | def process_tuple(self, tuple_: Tuple, port: int) -> Iterator[Optional[TupleLike]]: + | yield {"kind": str(tuple_["ts"])} + |""".stripMargin + } + + /** Declares one INTEGER column `x` and yields `row`, a Python dict literal. */ + private class YieldOpDesc(row: String) extends PythonOperatorDescriptor { + override def operatorInfo: OperatorInfo = + OperatorInfo( + userFriendlyName = "Yield", + operatorDescription = "yields one row as given", + operatorGroupName = OperatorGroupConstants.UTILITY_GROUP, + inputPorts = List(InputPort()), + outputPorts = List(OutputPort()) + ) + + override def getOutputSchemas( + inputSchemas: Map[PortIdentity, Schema] + ): Map[PortIdentity, Schema] = + Map( + operatorInfo.outputPorts.head.id -> Schema().add(new Attribute("x", AttributeType.INTEGER)) + ) + + override def generatePythonCode(): String = + s"""from pytexera import * + | + |class ProcessTupleOperator(UDFOperatorV2): + | + | @overrides + | def process_tuple(self, tuple_: Tuple, port: int) -> Iterator[Optional[TupleLike]]: + | yield $row + |""".stripMargin + } + /** Reads a LARGE_BINARY column and writes out the reference it holds. * * The bytes of a large binary live in S3 and the field is the s3:// URI that @@ -211,6 +254,32 @@ class PyOpExecHarnessSpec extends AnyFlatSpec with Matchers { binaryCellKind(Array[Byte](0, 1, 2)) shouldBe "bytes" } + // The worker gets the wall clock as a zoneless Arrow microsecond timestamp, + // so the operator holds a naive datetime cut to the microsecond, and a year + // past pandas' nanosecond range reaches it all the same. + it should "hand a timestamp to the operator as the worker hands it" in { + val dir = Files.createTempDirectory("py-op-harness-timestamp-") + val inputSchema = Schema().add(new Attribute("ts", AttributeType.TIMESTAMP)) + val input = dir.resolve("input_port_0.jsonl") + val row = Tuple + .builder(inputSchema) + .add("ts", AttributeType.TIMESTAMP, Timestamp.valueOf("2500-01-01 00:00:00.123456789")) + .build() + TupleIO.writeTuples(input, Iterator(row), inputSchema) + + val result = PyOpExecHarness.execute( + new TimestampTextOpDesc, + inputs = Map(PortIdentity(0) -> input), + outputDir = dir.resolve("actual") + ) + + val out = result.outputs(PortIdentity(0)) + TupleIO + .readTuples(out, TupleIO.readSchemaSidecar(out)) + .map(_.getField[String]("kind")) + .toSeq shouldBe Seq("2500-01-01 00:00:00.123456") + } + private val someObject = "s3://a-bucket/a/large/object" // The schema sidecar names large_binary, so the harness has to carry it in as @@ -261,6 +330,31 @@ class PyOpExecHarnessSpec extends AnyFlatSpec with Matchers { .toSeq shouldBe Seq(someObject) } + // The worker finalizes each yielded row against the output schema, and so + // does the driver: a value the cast cannot bring to the declared type, and a + // field the schema does not have or lacks, end the run in both. + "PyOpExecHarness" should "refuse a row the output schema does not describe" in { + val dir = Files.createTempDirectory("py-op-harness-finalize-") + val inputSchema = Schema().add(new Attribute("seed", AttributeType.INTEGER)) + val input = dir.resolve("input_port_0.jsonl") + val seed = Tuple.builder(inputSchema).add("seed", AttributeType.INTEGER, Int.box(1)).build() + TupleIO.writeTuples(input, Iterator(seed), inputSchema) + + def run(row: String): Unit = + PyOpExecHarness.execute( + new YieldOpDesc(row), + inputs = Map(PortIdentity(0) -> input), + outputDir = dir.resolve("actual") + ) + + noException should be thrownBy run("""{"x": 1}""") + Seq("""{"x": 1.5}""", """{"x": 1, "y": 2}""", "{}").foreach { row => + withClue(s"yielding $row: ") { + a[PyOpDriverException] should be thrownBy run(row) + } + } + } + "PyOpExecHarness" should "record each yield as it was yielded" in { val dir = Files.createTempDirectory("py-op-harness-") val inputSchema = Schema().add(new Attribute("seed", AttributeType.INTEGER)) From 135b9be9a848b5af10cbccd98dd83427edd4186f Mon Sep 17 00:00:00 2001 From: kary zheng Date: Tue, 29 Sep 2026 23:54:18 -0700 Subject: [PATCH 16/19] test(verify): hand a port that carried no rows its declared columns Co-Authored-By: Claude Opus 5.5 --- .../src/test/resources/python/py_op_driver.py | 19 ++++++- .../verify/PyOpExecHarnessSpec.scala | 54 +++++++++++++++++++ 2 files changed, 71 insertions(+), 2 deletions(-) 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 index bdffe300f10..384fa0637bc 100644 --- a/workflow-compiling-service/src/test/resources/python/py_op_driver.py +++ b/workflow-compiling-service/src/test/resources/python/py_op_driver.py @@ -405,6 +405,7 @@ def _run_operator( is_source: bool, port_order: Sequence[int], inputs_by_port: Mapping[int, Sequence[Tuple]], + schemas_by_port: Mapping[int, TexeraSchema], consume: "Callable[[Any], None]", ) -> None: """ @@ -434,6 +435,10 @@ def _run_operator( for item in op.process_tuple(tup, port): if item is not None: consume(item) + # DataProcessor declares the finishing port's schema right before + # on_finish, so a port that carried no rows still has its columns. + if port in schemas_by_port: + op.input_schemas[port] = schemas_by_port[port] for item in op.on_finish(port): if item is not None: consume(item) @@ -471,10 +476,12 @@ def run_config(config: Mapping[str, Any]) -> None: port_order: Sequence[int] = list(config.get("portOrder", [])) inputs_by_port: "dict[int, List[Tuple]]" = {} + schemas_by_port: "dict[int, TexeraSchema]" = {} for entry in config.get("inputs", []): port = int(entry["portIndex"]) data_path = Path(entry["dataPath"]) - inputs_by_port[port] = _read_tuples(data_path, _read_schema_sidecar(data_path)) + schemas_by_port[port] = _read_schema_sidecar(data_path) + inputs_by_port[port] = _read_tuples(data_path, schemas_by_port[port]) # Default port order: sorted by index. Matches OpExecHarness's fallback # when getInputPortDependencyPairs is empty. @@ -493,7 +500,14 @@ def run_config(config: Mapping[str, Any]) -> None: ) if not outputs: - _run_operator(op_instance, is_source, port_order, inputs_by_port, lambda _item: None) + _run_operator( + op_instance, + is_source, + port_order, + inputs_by_port, + schemas_by_port, + lambda _item: None, + ) return out_entry = outputs[0] @@ -505,6 +519,7 @@ def run_config(config: Mapping[str, Any]) -> None: is_source, port_order, inputs_by_port, + schemas_by_port, lambda item: rows.extend(_emit_as_dicts([item], out_schema)), ) _write_tuples(out_path, rows, out_schema) diff --git a/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/PyOpExecHarnessSpec.scala b/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/PyOpExecHarnessSpec.scala index d88f1a65d6a..72835bc3049 100644 --- a/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/PyOpExecHarnessSpec.scala +++ b/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/PyOpExecHarnessSpec.scala @@ -75,6 +75,37 @@ class PyOpExecHarnessSpec extends AnyFlatSpec with Matchers { |""".stripMargin } + /** Reports the column names of the table it was handed, as one row. */ + private class TableColumnsOpDesc extends PythonOperatorDescriptor { + override def operatorInfo: OperatorInfo = + OperatorInfo( + userFriendlyName = "Table Columns", + operatorDescription = "yields the column names of its input table", + operatorGroupName = OperatorGroupConstants.UTILITY_GROUP, + inputPorts = List(InputPort()), + outputPorts = List(OutputPort()) + ) + + override def getOutputSchemas( + inputSchemas: Map[PortIdentity, Schema] + ): Map[PortIdentity, Schema] = + Map( + operatorInfo.outputPorts.head.id -> Schema().add( + new Attribute("columns", AttributeType.STRING) + ) + ) + + override def generatePythonCode(): String = + """from pytexera import * + | + |class ProcessTableOperator(UDFTableOperator): + | + | @overrides + | def process_table(self, table: Table, port: int) -> Iterator[Optional[TableLike]]: + | yield {"columns": ",".join(table.columns)} + |""".stripMargin + } + /** Reports the timestamp it was handed, as Python prints it. */ private class TimestampTextOpDesc extends BinaryTypeOpDesc { override def generatePythonCode(): String = @@ -377,4 +408,27 @@ class PyOpExecHarnessSpec extends AnyFlatSpec with Matchers { values shouldBe Seq(1, 2) } } + + // A port that carried no rows gives the table no tuples to read column names + // off, so the worker hands the operator the port's declared schema instead. + "PyOpExecHarness" should "hand a port that carried no rows its declared columns" in { + val dir = Files.createTempDirectory("py-op-harness-empty-") + val inputSchema = Schema() + .add(new Attribute("name", AttributeType.STRING)) + .add(new Attribute("n", AttributeType.INTEGER)) + val input = dir.resolve("input_port_0.jsonl") + TupleIO.writeTuples(input, Iterator.empty, inputSchema) + + val result = PyOpExecHarness.execute( + new TableColumnsOpDesc, + inputs = Map(PortIdentity(0) -> input), + outputDir = dir.resolve("actual") + ) + + val out = result.outputs(PortIdentity(0)) + TupleIO + .readTuples(out, TupleIO.readSchemaSidecar(out)) + .map(_.getField[String]("columns")) + .toSeq shouldBe Seq("name,n") + } } From b581ae06c5bb3a14c7776ee95530cfce26d38414 Mon Sep 17 00:00:00 2001 From: kary zheng Date: Mon, 12 Oct 2026 19:22:33 -0700 Subject: [PATCH 17/19] test(verify): pickle the harness spec's model from a script file Windows drops the quotes from a multiline `-c` argument, so the bytes literal reached Python as `bpickle` and the fixture failed before the harness ran. Write the script to a temporary .py file and run that. Co-Authored-By: Claude Opus 5.5 --- .../amber/translator/verify/PyOpExecHarnessSpec.scala | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/PyOpExecHarnessSpec.scala b/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/PyOpExecHarnessSpec.scala index 72835bc3049..2f917bccd8c 100644 --- a/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/PyOpExecHarnessSpec.scala +++ b/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/PyOpExecHarnessSpec.scala @@ -239,15 +239,20 @@ class PyOpExecHarnessSpec extends AnyFlatSpec with Matchers { } // A fitted DecisionTreeClassifier behind the marker the cast writes, pickled - // by the interpreter the harness runs so the two share a sklearn. + // by the interpreter the harness runs so the two share a sklearn. The script + // goes through a file because Windows drops the quotes from a `-c` argument. private def fittedTreeCell(): Array[Byte] = { - val script = + val script = Files.createTempFile("py-op-harness-model-", ".py") + script.toFile.deleteOnExit() + Files.writeString( + script, """import base64, pickle, sys |from sklearn.tree import DecisionTreeClassifier |model = DecisionTreeClassifier(random_state=0).fit([[0.0], [1.0]], [0, 1]) |sys.stdout.write(base64.b64encode(b"pickle " + pickle.dumps(model)).decode("ascii")) |""".stripMargin - val encoded = Process(Seq(PyOpExecHarness.resolvePython(), "-c", script)).!! + ) + val encoded = Process(Seq(PyOpExecHarness.resolvePython(), script.toString)).!! Base64.getDecoder.decode(encoded.trim) } From a739b78e685b949df814ac2568e4d53a4e200cfc Mon Sep 17 00:00:00 2001 From: kary zheng Date: Mon, 12 Oct 2026 22:38:42 -0700 Subject: [PATCH 18/19] test(verify): give the executor each finished port's declared schema As the worker's EndChannelHandler does since #8506: a HashJoin probe that received no build rows names its right columns from the build port's declared schema. Co-Authored-By: Claude Opus 5.5 --- .../apache/texera/amber/translator/verify/OpExecHarness.scala | 3 +++ 1 file changed, 3 insertions(+) diff --git a/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/OpExecHarness.scala b/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/OpExecHarness.scala index cbfbcc34021..56ea492a4f1 100644 --- a/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/OpExecHarness.scala +++ b/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/OpExecHarness.scala @@ -283,6 +283,9 @@ object OpExecHarness extends LazyLogging { tuples.foreach { t => bucket(opExec.processTupleMultiPort(t, portId.id)) } + // As the worker's EndChannelHandler does: a HashJoin probe that received + // no build rows names its columns from this. + phOp.inputPorts(portId)._3.foreach(schema => opExec.inputSchemas(portId.id) = schema) bucket(opExec.onFinishMultiPort(portId.id)) } From 3f7a878cd7f4244a96d562c7b77e2e2a7b640073 Mon Sep 17 00:00:00 2001 From: kary zheng Date: Thu, 15 Oct 2026 21:44:31 -0700 Subject: [PATCH 19/19] test(verify): write a Python operator's infinities as the engine's writer does json.dumps writes a bare Infinity, which is not JSON; the engine's writer spells it as text. A NaN still leaves as null, as the worker casts it. Co-Authored-By: Claude Opus 5.5 --- .../src/test/resources/python/py_op_driver.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) 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 index 384fa0637bc..503abffd9ed 100644 --- a/workflow-compiling-service/src/test/resources/python/py_op_driver.py +++ b/workflow-compiling-service/src/test/resources/python/py_op_driver.py @@ -57,6 +57,7 @@ import datetime import inspect import json +import math import pickle import sys import traceback @@ -293,7 +294,13 @@ def _jsonify(value: Any, attr_type: AttributeType) -> Any: if attr_type in (AttributeType.INT, AttributeType.LONG): return int(value) if attr_type == AttributeType.DOUBLE: - return float(value) + value = float(value) + # The engine's writer spells these as text; json.dumps would write a + # bare Infinity, which is not JSON. A NaN never gets here: the worker + # casts it to null before the row leaves, and so does the branch above. + if value in (math.inf, -math.inf): + return "Infinity" if value > 0 else "-Infinity" + return value if attr_type == AttributeType.BOOL: return bool(value) if attr_type == AttributeType.BINARY: