|
| 1 | +# Licensed to the Apache Software Foundation (ASF) under one |
| 2 | +# or more contributor license agreements. See the NOTICE file |
| 3 | +# distributed with this work for additional information |
| 4 | +# regarding copyright ownership. The ASF licenses this file |
| 5 | +# to you under the Apache License, Version 2.0 (the |
| 6 | +# "License"); you may not use this file except in compliance |
| 7 | +# with the License. You may obtain a copy of the License at |
| 8 | +# |
| 9 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 10 | +# |
| 11 | +# Unless required by applicable law or agreed to in writing, |
| 12 | +# software distributed under the License is distributed on an |
| 13 | +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| 14 | +# KIND, either express or implied. See the License for the |
| 15 | +# specific language governing permissions and limitations |
| 16 | +# under the License. |
| 17 | + |
| 18 | +"""The driver: split a query into tasks, fan them out, collect the answer. |
| 19 | +
|
| 20 | +The shape is deliberately boring, because the interesting part is not the |
| 21 | +scheduling. What matters is the four things the driver has to get right, each |
| 22 | +of which is a way a real deployment goes wrong: |
| 23 | +
|
| 24 | +1. It serializes the stage **with** its session. ``to_bytes(None)`` uses an |
| 25 | + empty codec chain and cannot encode any library's node. |
| 26 | +2. It ships the codec ids it used, so a worker can refuse a plan it would |
| 27 | + misread rather than decode it with the wrong codec. |
| 28 | +3. It puts the shuffle directory in the session config, not in the message, |
| 29 | + so the directory travels *inside* the encoded plan and the two sides |
| 30 | + cannot disagree. |
| 31 | +4. It waits for every worker before reading, because the stage node decides |
| 32 | + whether to read or recompute by looking at the filesystem. |
| 33 | +""" |
| 34 | + |
| 35 | +from __future__ import annotations |
| 36 | + |
| 37 | +import json |
| 38 | +import pathlib |
| 39 | +import subprocess |
| 40 | +import sys |
| 41 | +from typing import TYPE_CHECKING |
| 42 | + |
| 43 | +from dfx_engine import _internal |
| 44 | +from dfx_engine.session import SessionSpec, build_session |
| 45 | + |
| 46 | +if TYPE_CHECKING: |
| 47 | + import pyarrow as pa |
| 48 | + from datafusion import DataFrame, SessionContext |
| 49 | + from datafusion.plan import ExecutionPlan |
| 50 | + |
| 51 | +__all__ = ["DistributedResult", "find_stage", "run_distributed"] |
| 52 | + |
| 53 | + |
| 54 | +class DistributedResult: |
| 55 | + """What a distributed run produced, and how.""" |
| 56 | + |
| 57 | + def __init__( |
| 58 | + self, |
| 59 | + batches: list[pa.RecordBatch], |
| 60 | + partitions: list[int], |
| 61 | + worker_rows: dict[int, int], |
| 62 | + ) -> None: |
| 63 | + self.batches = batches |
| 64 | + self.partitions = partitions |
| 65 | + """Partition indices that were dispatched, one per worker.""" |
| 66 | + self.worker_rows = worker_rows |
| 67 | + """Rows each worker produced, keyed by partition index.""" |
| 68 | + |
| 69 | + |
| 70 | +STAGE_NODE_NAME = "ShuffleStageExec" |
| 71 | + |
| 72 | + |
| 73 | +def find_stage(plan: ExecutionPlan) -> ExecutionPlan | None: |
| 74 | + """Locate the stage node the planner inserted. |
| 75 | +
|
| 76 | + Matched on the display string because a Python caller has no way to |
| 77 | + downcast a Rust plan node -- there is no ``isinstance`` across an FFI |
| 78 | + boundary. |
| 79 | +
|
| 80 | + Note the *containment* test. The node was built inside this library and |
| 81 | + handed back to the host, so what the host prints is not |
| 82 | + ``ShuffleStageExec: stage=1`` but:: |
| 83 | +
|
| 84 | + FFI_ExecutionPlan: ShuffleStageExec, number_of_children=1 |
| 85 | +
|
| 86 | + A foreign node reports its own name nested inside the wrapper's, which |
| 87 | + makes anchored matches on plan text quietly wrong -- the kind of thing |
| 88 | + that works in a single-library test and fails the moment a real extension |
| 89 | + is involved. |
| 90 | + """ |
| 91 | + if STAGE_NODE_NAME in plan.display(): |
| 92 | + return plan |
| 93 | + for child in plan.children(): |
| 94 | + found = find_stage(child) |
| 95 | + if found is not None: |
| 96 | + return found |
| 97 | + return None |
| 98 | + |
| 99 | + |
| 100 | +def _dispatch( |
| 101 | + envelope: dict, envelope_dir: pathlib.Path, partition: int |
| 102 | +) -> subprocess.Popen[str]: |
| 103 | + """Start one worker for one partition. |
| 104 | +
|
| 105 | + ``sys.executable``, not ``python``: a worker on a different Python minor |
| 106 | + version cannot load a cloudpickled inline UDF, and that failure is far |
| 107 | + from its cause. |
| 108 | + """ |
| 109 | + path = envelope_dir / f"task-{partition}.json" |
| 110 | + path.write_text(json.dumps(envelope)) |
| 111 | + return subprocess.Popen( # noqa: S603 |
| 112 | + [sys.executable, "-m", "dfx_engine.worker", str(path)], |
| 113 | + stdout=subprocess.PIPE, |
| 114 | + stderr=subprocess.PIPE, |
| 115 | + text=True, |
| 116 | + ) |
| 117 | + |
| 118 | + |
| 119 | +def run_distributed(sql: str, spec: SessionSpec) -> DistributedResult: |
| 120 | + """Run `sql`, executing its leaf stage in one worker process per partition. |
| 121 | +
|
| 122 | + Requires ``spec.shuffle_dir``: without it the planner inserts no stage and |
| 123 | + there is nothing to distribute. |
| 124 | + """ |
| 125 | + if not spec.shuffle_dir: |
| 126 | + message = "run_distributed needs a shuffle_dir; build_session got none" |
| 127 | + raise ValueError(message) |
| 128 | + |
| 129 | + ctx, engine, _storage = build_session(spec) |
| 130 | + plan = ctx.sql(sql).execution_plan() |
| 131 | + |
| 132 | + stage = find_stage(plan) |
| 133 | + if stage is None: |
| 134 | + message = ( |
| 135 | + "no ShuffleStageExec in the plan; the engine's planner did not run, " |
| 136 | + "or its config extension was not registered" |
| 137 | + ) |
| 138 | + raise RuntimeError(message) |
| 139 | + |
| 140 | + shuffle_dir = pathlib.Path(spec.shuffle_dir) |
| 141 | + shuffle_dir.mkdir(parents=True, exist_ok=True) |
| 142 | + |
| 143 | + # Encode the stage subtree, through the session that owns the codecs. |
| 144 | + plan_path = shuffle_dir / "stage.plan" |
| 145 | + plan_path.write_bytes(stage.to_bytes(ctx)) |
| 146 | + |
| 147 | + stage_id = _internal.stage_id() |
| 148 | + partitions = list(range(stage.partition_count)) |
| 149 | + envelopes = [ |
| 150 | + { |
| 151 | + "spec": spec.to_json(), |
| 152 | + "plan": str(plan_path), |
| 153 | + "stage_id": stage_id, |
| 154 | + "partition": partition, |
| 155 | + } |
| 156 | + for partition in partitions |
| 157 | + ] |
| 158 | + |
| 159 | + # One process per partition, all in flight together. This is the claim the |
| 160 | + # example is making: each worker reads a different file and writes a |
| 161 | + # different result, so they need no coordination beyond the directory. |
| 162 | + workers = [ |
| 163 | + _dispatch(envelope, shuffle_dir, partition) |
| 164 | + for envelope, partition in zip(envelopes, partitions, strict=True) |
| 165 | + ] |
| 166 | + |
| 167 | + worker_rows: dict[int, int] = {} |
| 168 | + failures = [] |
| 169 | + for partition, worker in zip(partitions, workers, strict=True): |
| 170 | + stdout, stderr = worker.communicate() |
| 171 | + if worker.returncode != 0: |
| 172 | + failures.append(f"partition {partition} failed:\n{stderr}") |
| 173 | + continue |
| 174 | + worker_rows[partition] = json.loads(stdout)["rows"] |
| 175 | + |
| 176 | + if failures: |
| 177 | + raise RuntimeError("\n".join(failures)) |
| 178 | + |
| 179 | + # Now run the whole query here. Every stage partition has a file, so the |
| 180 | + # stage node streams them instead of recomputing -- the driver does only |
| 181 | + # the final merge. |
| 182 | + batches = ctx.sql(sql).collect() |
| 183 | + _ = engine |
| 184 | + return DistributedResult(batches, partitions, worker_rows) |
| 185 | + |
| 186 | + |
| 187 | +def run_local(sql: str, spec: SessionSpec) -> list[pa.RecordBatch]: |
| 188 | + """Run `sql` in this process, for comparison. |
| 189 | +
|
| 190 | + Uses the same session factory with no shuffle directory, so the only |
| 191 | + difference from :func:`run_distributed` is where the work happened. |
| 192 | + """ |
| 193 | + ctx, _engine, _storage = build_session( |
| 194 | + SessionSpec( |
| 195 | + tables=spec.tables, |
| 196 | + shuffle_dir="", |
| 197 | + target_partitions=spec.target_partitions, |
| 198 | + ) |
| 199 | + ) |
| 200 | + return ctx.sql(sql).collect() |
| 201 | + |
| 202 | + |
| 203 | +def dataframe_for(sql: str, spec: SessionSpec) -> tuple[SessionContext, DataFrame]: |
| 204 | + """Session and DataFrame for `sql`, for tests that want to inspect a plan.""" |
| 205 | + ctx, _engine, _storage = build_session(spec) |
| 206 | + return ctx, ctx.sql(sql) |
0 commit comments