diff --git a/build.sbt b/build.sbt index cb16bda9528..6a7aa592d12 100644 --- a/build.sbt +++ b/build.sbt @@ -228,8 +228,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)) + } } ) @@ -240,6 +246,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..503abffd9ed --- /dev/null +++ b/workflow-compiling-service/src/test/resources/python/py_op_driver.py @@ -0,0 +1,594 @@ +#!/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 datetime +import inspect +import json +import math +import pickle +import sys +import traceback +from pathlib import Path +from typing import Any, Callable, 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 + largebinary, + 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 + 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 " + "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. +# -------------------------------------------------------------------------- +# Prefix Tuple.cast_to_schema writes in front of an object it pickles into a +# 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 " + + +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: + # 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 + # 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 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( + 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) + # 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 + + +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. + + 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. + """ + for item in emitted: + 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: + """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: + 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: + # 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.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 + # 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]], + schemas_by_port: Mapping[int, TexeraSchema], + consume: "Callable[[Any], None]", +) -> None: + """ + 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: + # 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: + 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: + 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) + finally: + op.close() + + +# -------------------------------------------------------------------------- +# 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]]" = {} + schemas_by_port: "dict[int, TexeraSchema]" = {} + for entry in config.get("inputs", []): + port = int(entry["portIndex"]) + data_path = Path(entry["dataPath"]) + 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. + 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() + + 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 not outputs: + _run_operator( + op_instance, + is_source, + port_order, + inputs_by_port, + schemas_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, + schemas_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: + 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/OpExecHarness.scala b/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/OpExecHarness.scala index 01e48e40da2..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 @@ -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, @@ -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)) } @@ -383,6 +386,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 +432,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 new file mode 100644 index 00000000000..4dc669c75e8 --- /dev/null +++ b/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/PyOpExecHarness.scala @@ -0,0 +1,361 @@ +/* + * 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.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. + * + * It returns the same Result shape, so the rest of the pipeline does not care + * which harness ran. + * + * 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. + */ +object PyOpExecHarness extends LazyLogging { + + private val TestWorkflowId = WorkflowIdentity(0L) + private val TestExecutionId = ExecutionIdentity(0L) + + // Same Result shape as OpExecHarness so callers can swap harnesses + // transparently. + final case class Result( + outputs: Map[PortIdentity, Path], + outputSchemas: Map[PortIdentity, Schema] + ) + + // Driver script lives on the test classpath at /python/py_op_driver.py + // (sibling to compare.py). Extracted to a temp file at runtime so it works + // whether the test resources are loose files or sealed in a jar. + private val DriverResourcePath = "/python/py_op_driver.py" + + def execute( + opDesc: LogicalOp, + inputs: Map[PortIdentity, Path], + outputDir: Path, + pythonExe: String = resolvePython(), + amberPythonHome: Path = resolveAmberPythonHome() + ): Result = { + Files.createDirectories(outputDir) + + val plan = opDesc.getPhysicalPlan(TestWorkflowId, TestExecutionId) + + // PythonOperatorDescriptor builds single-op plans; bail loudly if some + // future Python op produces a multi-stage plan (need to extend the driver + // and the per-PhysicalOp config the way OpExecHarness does). + require( + plan.operators.size == 1, + s"PyOpExecHarness only supports single-PhysicalOp plans for now, got " + + s"${plan.operators.size} PhysicalOps" + ) + val phOp = plan.operators.head + + val (pythonCode, language) = phOp.opExecInitInfo match { + case OpExecWithCode(code, lang) => (code, lang) + case other => + throw new UnsupportedOperationException( + s"PyOpExecHarness only supports OpExecWithCode; got ${other.getClass.getSimpleName}. " + + "For OpExecWithClassName, use OpExecHarness." + ) + } + require( + language == "python", + s"""PyOpExecHarness only supports language="python", got "$language".""" + ) + + // External input ports = same definition as OpExecHarness. For a + // single-op plan that's just every input port the op declares. + val externalInputs: Set[(PhysicalOpIdentity, PortIdentity)] = + phOp.inputPorts.keys.map(portId => (phOp.id, portId)).toSet + OpExecHarness.validateInputCoverage(externalInputs, inputs.keySet) + + val inputSchemas: Map[PortIdentity, Schema] = + inputs.map { case (portId, path) => portId -> TupleIO.readSchemaSidecar(path) } + + // Preparing the plan is the same work for either executor, so it is done in + // one place; only what runs the prepared op differs between the harnesses. + val planWithSchemas = + OpExecHarness.propagateExternalSchemas(plan, externalInputs, inputSchemas) + val phOpWithSchemas = planWithSchemas.operators.head + + // Output port schemas come from PhysicalPlan.propagateSchema — same + // ground truth OpExecHarness writes to its own outputs. + val outputPortSchemas: Map[PortIdentity, Schema] = + phOpWithSchemas.outputPorts.map { + case (portId, (_, _, schemaOrErr)) => + portId -> schemaOrErr.toOption.getOrElse( + throw new IllegalStateException( + s"Output schema for ($portId) was not propagated" + ) + ) + } + + require( + outputPortSchemas.size == 1, + s"PyOpExecHarness only supports single-output-port operators, got " + + s"${outputPortSchemas.size} output ports" + ) + val (outputPortId, outputSchema) = outputPortSchemas.head + val outputPath = outputDir.resolve(s"output_port_${outputPortId.id}.jsonl") + + // Port ordering for multi-input ops: respect declared dependencies + // (matches OpExecHarness — e.g. HashJoin probe processes build-side + // first). Default = sorted by port id when no dependencies declared. + val portOrder: Seq[Int] = + if (phOpWithSchemas.getInputPortDependencyPairs.nonEmpty) + phOpWithSchemas.getInputPortDependencyPairs.map(_.id) + else phOpWithSchemas.inputPorts.keys.toList.map(_.id).sorted + + val config = buildConfig( + pythonCode = pythonCode, + isSource = phOpWithSchemas.isSourceOperator, + portOrder = portOrder, + inputs = inputs, + outputPath = outputPath, + outputSchema = outputSchema + ) + + val configPath = outputDir.resolve("py_op_driver_config.json") + Files.write(configPath, config.getBytes(StandardCharsets.UTF_8)) + + val driverPath = extractDriverScript() + runDriver(driverPath, configPath, outputDir, pythonExe, amberPythonHome) + + Result( + outputs = Map(outputPortId -> outputPath), + outputSchemas = Map(outputPortId -> outputSchema) + ) + } + + // -------------------------------------------------------------------------- + // Config serialization. Matches the driver's expected schema (see + // py_op_driver.py's module docstring). + // -------------------------------------------------------------------------- + private def buildConfig( + pythonCode: String, + isSource: Boolean, + portOrder: Seq[Int], + inputs: Map[PortIdentity, Path], + outputPath: Path, + outputSchema: Schema + ): String = { + val root: ObjectNode = objectMapper.createObjectNode() + root.put("operatorCode", pythonCode) + root.put("isSource", isSource) + + val portOrderArr: ArrayNode = root.putArray("portOrder") + portOrder.foreach(portOrderArr.add) + + val inputsArr: ArrayNode = root.putArray("inputs") + inputs.toSeq.sortBy(_._1.id).foreach { + case (portId, dataPath) => + val entry: ObjectNode = inputsArr.addObject() + entry.put("portIndex", portId.id) + entry.put("dataPath", dataPath.toAbsolutePath.toString) + // schemaPath is implicit (data_path + ".schema.json") — the driver + // resolves it the same way TupleIO does. + } + + val outputsArr: ArrayNode = root.putArray("outputs") + val outEntry: ObjectNode = outputsArr.addObject() + outEntry.put("dataPath", outputPath.toAbsolutePath.toString) + // Embed the schema directly. We can't just write the sidecar ahead of + // time and have the driver read it, because writing a sidecar before + // outputs exist would leave a stale sidecar on partial failures. + outEntry.set[ObjectNode]( + "schema", + objectMapper.valueToTree[ObjectNode](outputSchema) + ) + + objectMapper.writeValueAsString(root) + } + + // -------------------------------------------------------------------------- + // Subprocess invocation. + // -------------------------------------------------------------------------- + private def runDriver( + driverPath: Path, + configPath: Path, + cwd: Path, + pythonExe: String, + amberPythonHome: Path + ): Unit = { + // Prepend amber's Python source to PYTHONPATH so `import pytexera` + // resolves. Existing PYTHONPATH (if any) is preserved as the lower- + // priority suffix. + val existing = sys.env.getOrElse("PYTHONPATH", "") + val newPyPath = + if (existing.isEmpty) amberPythonHome.toAbsolutePath.toString + else s"${amberPythonHome.toAbsolutePath}${File.pathSeparator}$existing" + + val (exit, stdout, stderr) = execDriver(driverPath, configPath, cwd, pythonExe, newPyPath) + if (exit != 0) { + throw new PyOpDriverException( + exitCode = exit, + driverPath = driverPath, + configPath = configPath, + stdout = stdout, + stderr = stderr + ) + } + } + + // Prefer a pooled persistent worker (imports pytexera/pyamber once via + // `py_op_driver.py --serve`, the ~300 ms cost that dominates a per-Python-op + // run — see PythonWorkerPool). A rare hard worker crash falls back to a + // one-shot subprocess so behavior is never worse than the original path. Both + // paths use absolute config paths, so cwd only matters to the subprocess + // form; the worker constructs a fresh operator per job for isolation. + private def execDriver( + driverPath: Path, + configPath: Path, + cwd: Path, + pythonExe: String, + pythonPath: String + ): (Int, String, String) = { + if (PythonWorkerPool.enabled) { + try { + val req = objectMapper.createObjectNode() + req.put("configPath", configPath.toAbsolutePath.toString) + val o = PythonWorkerPool.run( + DriverResourcePath, + Seq("--serve"), + pythonExe, + req, + env = Map("PYTHONPATH" -> pythonPath) + ) + return (o.exit, o.stdout, o.stderr) + } catch { + case e: PythonWorkerPool.WorkerDiedException => + logger.warn( + s"py_op_driver worker unavailable; falling back to one-shot subprocess " + + s"for $configPath: ${e.getMessage}" + ) + } + } + runDriverSubprocess(driverPath, configPath, cwd, pythonExe, pythonPath) + } + + // Original one-process-per-operator path. Retained as the fallback and as the + // behavior selected by TEXERA_TEST_PYTHON_WORKER=0. + private def runDriverSubprocess( + driverPath: Path, + configPath: Path, + cwd: Path, + pythonExe: String, + pythonPath: String + ): (Int, String, String) = { + val outBuf = ArrayBuffer.empty[String] + val errBuf = ArrayBuffer.empty[String] + val procLogger = ProcessLogger(line => outBuf += line, line => errBuf += line) + val exit = Process( + Seq(pythonExe, driverPath.toString, configPath.toString), + Some(cwd.toFile), + "PYTHONPATH" -> pythonPath + ).!(procLogger) + (exit, outBuf.mkString("\n"), errBuf.mkString("\n")) + } + + // -------------------------------------------------------------------------- + // Resolution helpers. + // -------------------------------------------------------------------------- + private[verify] def resolvePython(): String = + sys.env.get("UDF_PYTHON_PATH").filter(_.nonEmpty).getOrElse("python3.12") + + /** + * Locate `amber/src/main/python`. Resolution chain: + * 1. Env var TEXERA_AMBER_PYTHON_HOME (set by CI / dev shell). + * 2. Walk up from cwd looking for `amber/src/main/python`. + * sbt runs tests with cwd = the subproject dir (`workflow-compiling-service/`), + * so the walk-up is two levels at most for the normal layout. + */ + private def resolveAmberPythonHome(): Path = { + sys.env.get("TEXERA_AMBER_PYTHON_HOME").filter(_.nonEmpty).map(Paths.get(_)).getOrElse { + val cwd = Paths.get(".").toAbsolutePath.normalize() + val maxDepth = 5 + var current: Path = cwd + var depth = 0 + while (current != null && depth <= maxDepth) { + val candidate = current.resolve("amber/src/main/python") + if (Files.isDirectory(candidate)) return candidate.toAbsolutePath + current = current.getParent + depth += 1 + } + throw new RuntimeException( + s"PyOpExecHarness: could not locate amber/src/main/python from cwd $cwd. " + + "Set TEXERA_AMBER_PYTHON_HOME to the absolute path." + ) + } + } + + private def extractDriverScript(): Path = { + val stream = getClass.getResourceAsStream(DriverResourcePath) + require( + stream != null, + s"py_op_driver.py not found on classpath at $DriverResourcePath" + ) + try { + val tmp = Files.createTempFile("py_op_driver-", ".py") + Files.copy(stream, tmp, StandardCopyOption.REPLACE_EXISTING) + tmp.toFile.deleteOnExit() + tmp + } finally stream.close() + } + +} + +final class PyOpDriverException( + val exitCode: Int, + val driverPath: Path, + val configPath: Path, + val stdout: String, + val stderr: String +) extends RuntimeException( + s"""py_op_driver.py exited with code $exitCode. + |Driver: $driverPath + |Config: $configPath + |--- stdout --- + |$stdout + |--- stderr --- + |$stderr""".stripMargin + ) 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..2f917bccd8c --- /dev/null +++ b/workflow-compiling-service/src/test/scala/org/apache/texera/amber/translator/verify/PyOpExecHarnessSpec.scala @@ -0,0 +1,439 @@ +/* + * 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, 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} +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 +import java.sql.Timestamp +import java.util.Base64 +import scala.sys.process._ + +/** 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 + } + + /** 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 = + """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 + * 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 + } + + /** 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 + } + + /** 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. The script + // goes through a file because Windows drops the quotes from a `-c` argument. + private def fittedTreeCell(): Array[Byte] = { + 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(), script.toString)).!! + 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") + val row = Tuple.builder(inputSchema).add("blob", AttributeType.BINARY, cell).build() + TupleIO.writeTuples(input, Iterator(row), inputSchema) + + val result = PyOpExecHarness.execute( + opDesc, + 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, 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 { + 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 + // 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) + } + + // 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)) + 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, + 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) + } + } + + // 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") + } +}