diff --git a/examples/extensions/benchmarks/e0_cell.py b/examples/extensions/benchmarks/e0_cell.py new file mode 100644 index 00000000..6c086132 --- /dev/null +++ b/examples/extensions/benchmarks/e0_cell.py @@ -0,0 +1,200 @@ +"""One native Pregel propagation cell; the outer owner manages Sail and its budget. + +Reads LDBC Parquet, writes the complete answer and pre-write relation plans, +and never validates/collects graph rows. Run e0_oracle.py after the cell and +the owned engine have exited. Plan recording is an explicit diagnostic cost. +""" + +from __future__ import annotations + +import argparse +import json +import time +from dataclasses import asdict, dataclass +from datetime import datetime, timezone +from pathlib import Path +from typing import Literal, cast + +from pyspark.sql.connect import functions as F +from pyspark.sql.connect.client.retries import DefaultPolicy +from pyspark.sql.connect.session import SparkSession +from pyspark_pecan import GraphAlgorithms, IterationEvent + +Program = Literal["pagerank", "sssp", "landmarks"] + + +@dataclass(frozen=True, slots=True) +class Arguments: + endpoint: str + vertices: Path + edges: Path + output: Path + program: Program + source: int + landmarks: tuple[int, ...] + partitions: int + iterations: int + tolerance: float + reset: float + normalized: bool + vote_to_halt: bool + reversed_edges: bool + input_weights: bool + record_plans: bool + undirected: bool = False + + +@dataclass(slots=True) +class Receipt: + observed_utc: str + status: str + arguments: Arguments + output: str + pipeline_seconds: float | None = None + algorithm_and_export_seconds: float | None = None + iterations: int | None = None + converged: bool | None = None + error: str | None = None + session_stopped: bool = False + + +def write_receipt(path: Path, receipt: Receipt) -> None: + path.write_text(json.dumps(asdict(receipt), default=str, indent=2) + "\n") + + +def execute(arguments: Arguments) -> None: + arguments.output.mkdir(parents=True, exist_ok=False) + receipt = Receipt(datetime.now(timezone.utc).isoformat(), "running", arguments, str(arguments.output / "result")) + write_receipt(arguments.output / "receipt.json", receipt) + events: list[IterationEvent] = [] + + def observe(event: IterationEvent) -> None: + if event.plan is not None: + (arguments.output / f"pre-write-relation-step-{event.iteration:05d}.txt").write_text(event.plan) + events.append(event.model_copy(update={"plan": None})) + + spark = SparkSession.builder.remote(arguments.endpoint).create() + # Do not retry a benchmark action behind the owner's receipt. + spark.client.set_retry_policies([DefaultPolicy(max_retries=0)]) + try: + graph = GraphAlgorithms( + spark, + snapshot_inputs=False, + repartition_checkpoints=False, + record_plans=arguments.record_plans, + observer=observe, + ) + started = time.perf_counter() + vertices = spark.read.parquet(str(arguments.vertices)).select("id") + original_edges = spark.read.parquet(str(arguments.edges)) + edges = original_edges.select(F.col("source").alias("src"), F.col("target").alias("dst")) + if arguments.program == "sssp": + edges = original_edges.select( + F.col("source").alias("src"), + F.col("target").alias("dst"), + (F.col("weight") if arguments.input_weights else F.lit(1.0)).alias("weight"), + ) + elif arguments.undirected: + edges = edges.unionByName(edges.select(F.col("dst").alias("src"), F.col("src").alias("dst"))) + algorithm_started = time.perf_counter() + if arguments.program == "pagerank": + result = graph.pagerank( + vertices, + edges, + method="pregel_delta", + partitions=arguments.partitions, + max_iterations=arguments.iterations, + tolerance=arguments.tolerance, + reset_probability=arguments.reset, + normalize=arguments.normalized, + vote_to_halt=arguments.vote_to_halt, + ) + elif arguments.program == "sssp": + result = graph.sssp( + vertices, + edges, + method="pregel", + source=arguments.source, + partitions=arguments.partitions, + max_iterations=arguments.iterations, + vote_to_halt=arguments.vote_to_halt, + directed=not arguments.undirected, + ) + else: + result = graph.shortest_paths( + vertices, + edges, + landmarks=arguments.landmarks, + to_landmarks=arguments.reversed_edges, + partitions=arguments.partitions, + max_iterations=arguments.iterations, + vote_to_halt=arguments.vote_to_halt, + ) + with result: + result.write_parquet(str(arguments.output / "result")) + finished = time.perf_counter() + receipt.pipeline_seconds = finished - started + receipt.algorithm_and_export_seconds = finished - algorithm_started + receipt.iterations = result.iterations + receipt.converged = result.converged + receipt.status = "completed_unvalidated" + except BaseException as error: + receipt.status = "error" + receipt.error = f"{type(error).__name__}: {error}" + raise + finally: + (arguments.output / "events.json").write_text( + json.dumps([event.model_dump(mode="json", exclude_none=True) for event in events], indent=2) + "\n" + ) + try: + spark.stop() + receipt.session_stopped = True + finally: + write_receipt(arguments.output / "receipt.json", receipt) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--endpoint", required=True) + parser.add_argument("--vertices", type=Path, required=True) + parser.add_argument("--edges", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--program", choices=("pagerank", "sssp", "landmarks"), required=True) + parser.add_argument("--source", type=int, default=0) + parser.add_argument("--landmarks", type=int, nargs="+", default=[0]) + parser.add_argument("--partitions", type=int, default=16) + parser.add_argument("--iterations", type=int, default=10) + parser.add_argument("--tolerance", type=float, default=0.01) + parser.add_argument("--reset", type=float, default=0.15) + parser.add_argument("--normalized", action="store_true") + parser.add_argument("--vote-to-halt", action="store_true") + parser.add_argument("--reversed-edges", action="store_true") + parser.add_argument("--input-weights", action="store_true") + parser.add_argument("--record-plans", action="store_true") + parser.add_argument("--undirected", action="store_true") + args = parser.parse_args() + execute( + Arguments( + args.endpoint, + args.vertices, + args.edges, + args.output, + cast(Program, args.program), + args.source, + tuple(args.landmarks), + args.partitions, + args.iterations, + args.tolerance, + args.reset, + args.normalized, + args.vote_to_halt, + args.reversed_edges, + args.input_weights, + args.record_plans, + args.undirected, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/examples/extensions/benchmarks/e0_oracle.py b/examples/extensions/benchmarks/e0_oracle.py new file mode 100644 index 00000000..b593ad7e --- /dev/null +++ b/examples/extensions/benchmarks/e0_oracle.py @@ -0,0 +1,331 @@ +"""Full, separate physical-output oracle for e0_cell.py. + +This utility reads all original edges and all result rows after the owned +engine exits. It performs the explicit input validation absent from Pecan. +Fixed-budget paths are compared with bounded Bellman-Ford; a halted weighted +run is compared with independent heap Dijkstra, and halted landmark hops +with queue BFS. PageRank uses GraphX's delta recurrence on NumPy arrays. +""" + +from __future__ import annotations + +import argparse +import hashlib +import heapq +import json +from collections import deque +from dataclasses import asdict, dataclass +from datetime import datetime, timezone +from pathlib import Path + +import numpy as np +import pyarrow as pa +import pyarrow.dataset as ds +from e0_cell import Arguments, Receipt +from numpy.typing import NDArray +from pydantic import TypeAdapter + +MAX_HOPS = (1 << 31) - 1 + + +@dataclass(frozen=True, slots=True) +class FilePin: + path: str + bytes: int + sha256: str + + +@dataclass(frozen=True, slots=True) +class Graph: + ids: NDArray[np.int64] + source: NDArray[np.int64] + target: NDArray[np.int64] + weights: NDArray[np.float64] + + +@dataclass(slots=True) +class OracleReceipt: + observed_utc: str + status: str + producer_receipt: FilePin + before: list[FilePin] + after: list[FilePin] + program: str + vertices: int = 0 + edges: int = 0 + result_rows: int = 0 + mismatch_rows: int = 0 + max_absolute_difference: float = 0.0 + error: str | None = None + official_sssp: FilePin | None = None + official_sssp_passed: bool = False + + +def pin(path: Path) -> FilePin: + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(8 * 1024 * 1024), b""): + digest.update(chunk) + return FilePin(str(path.resolve()), path.stat().st_size, digest.hexdigest()) + + +def parquet_files(path: Path) -> list[Path]: + paths = [path] if path.is_file() else sorted(path.rglob("*.parquet")) + if not paths: + raise ValueError(f"no original Parquet files at {path}") + return paths + + +def int64_column(table: pa.Table, name: str) -> NDArray[np.int64]: + column = table[name] + if column.type != pa.int64() or column.null_count: + raise ValueError(f"{name} must be nonnull physical Int64") + return np.asarray(column.combine_chunks().to_numpy(zero_copy_only=False), dtype=np.int64) + + +def positions(ids: NDArray[np.int64], values: NDArray[np.int64]) -> NDArray[np.int64]: + found = np.searchsorted(ids, values) + if np.any(found >= len(ids)) or np.any(ids[found] != values): + raise ValueError("an endpoint or traversal source does not name a vertex") + return found + + +def read_graph(arguments: Arguments) -> Graph: + vertex_table = ds.dataset(str(arguments.vertices), format="parquet").to_table(columns=["id"]) + ids = np.sort(int64_column(vertex_table, "id")) + if len(np.unique(ids)) != len(ids): + raise ValueError("duplicate input vertex IDs") + names = ["source", "target", *(["weight"] if arguments.input_weights else [])] + edge_table = ds.dataset(str(arguments.edges), format="parquet").to_table(columns=names) + source = positions(ids, int64_column(edge_table, "source")) + target = positions(ids, int64_column(edge_table, "target")) + weights = np.ones(len(source), dtype=np.float64) + if arguments.input_weights: + column = edge_table["weight"] + if column.type != pa.float64() or column.null_count: + raise ValueError("weight must be nonnull physical Float64") + weights = np.asarray(column.combine_chunks().to_numpy(zero_copy_only=False), dtype=np.float64) + if np.any(~np.isfinite(weights)) or np.any(weights < 0): + raise ValueError("weights must be finite and nonnegative") + if arguments.undirected: + source, target = np.concatenate((source, target)), np.concatenate((target, source)) + weights = np.concatenate((weights, weights)) + return Graph(ids, source, target, weights) + + +def bounded_distances( + graph: Graph, source: int, iterations: int, *, reverse: bool = False, weighted: bool = True +) -> NDArray[np.float64]: + """Synchronous all-edge Bellman-Ford, independent of active-set Pregel.""" + start, end = (graph.target, graph.source) if reverse else (graph.source, graph.target) + distance = np.full(len(graph.ids), np.inf, dtype=np.float64) + distance[int(positions(graph.ids, np.array([source], dtype=np.int64))[0])] = 0.0 + weights = graph.weights if weighted else np.ones(len(start), dtype=np.float64) + for _ in range(iterations): + updated = distance.copy() + np.minimum.at(updated, end, distance[start] + weights) + distance = updated + return distance + + +def adjacency(graph: Graph, reverse: bool) -> tuple[NDArray[np.int64], NDArray[np.int64], NDArray[np.float64]]: + start, end = (graph.target, graph.source) if reverse else (graph.source, graph.target) + order = np.argsort(start, kind="stable") + offsets = np.zeros(len(graph.ids) + 1, dtype=np.int64) + offsets[1:] = np.cumsum(np.bincount(start, minlength=len(graph.ids))) + return offsets, end[order], graph.weights[order] + + +def complete_distances( + graph: Graph, source: int, *, reverse: bool = False, weighted: bool = True +) -> NDArray[np.float64]: + offsets, targets, weights = adjacency(graph, reverse) + root = int(positions(graph.ids, np.array([source], dtype=np.int64))[0]) + distance = np.full(len(graph.ids), np.inf, dtype=np.float64) + distance[root] = 0.0 + if not weighted: + queue = deque([root]) + while queue: + vertex = queue.popleft() + for target in targets[offsets[vertex] : offsets[vertex + 1]]: + if np.isinf(distance[target]): + distance[target] = distance[vertex] + 1.0 + queue.append(int(target)) + return distance + pending = [(0.0, root)] + while pending: + cost, vertex = heapq.heappop(pending) + if cost != distance[vertex]: + continue + for index in range(int(offsets[vertex]), int(offsets[vertex + 1])): + target = int(targets[index]) + candidate = cost + float(weights[index]) + if candidate < distance[target]: + distance[target] = candidate + heapq.heappush(pending, (candidate, target)) + return distance + + +def pagerank(graph: Graph, arguments: Arguments, iterations: int) -> tuple[NDArray[np.float64], bool]: + degree = np.bincount(graph.source, minlength=len(graph.ids)) + rank = np.full(len(graph.ids), arguments.reset, dtype=np.float64) + delta = rank.copy() + for step in range(iterations): + active = np.ones(len(rank), dtype=np.bool_) if step == 0 else delta > arguments.tolerance + contribution = np.where(active[graph.source], delta[graph.source] / degree[graph.source], 0.0) + received = np.bincount(graph.target, weights=contribution, minlength=len(rank)) + delta = (1.0 - arguments.reset) * received + rank += delta + if arguments.normalized and rank.sum() != 0: + rank /= rank.sum() + return rank, bool(np.all(delta <= arguments.tolerance)) + + +def compare(actual: NDArray[np.float64], expected: NDArray[np.float64], *, exact: bool) -> tuple[int, float]: + finite = np.isfinite(expected) & np.isfinite(actual) + matches = (actual == expected) if exact else np.isclose(actual, expected, rtol=1e-12, atol=1e-12) + mismatches = int(np.count_nonzero(~matches)) + maximum = float(np.max(np.abs(actual[finite] - expected[finite]), initial=0.0)) + return mismatches, maximum + + +def qualify(producer: Receipt, graph: Graph, table: pa.Table, receipt: OracleReceipt) -> None: + args = producer.arguments + if producer.status != "completed_unvalidated" or not producer.session_stopped or producer.iterations is None: + raise ValueError("producer did not complete and close its session") + if not args.vote_to_halt and producer.iterations != args.iterations: + raise ValueError("fixed-budget producer did not run its entire declared budget") + if not args.vote_to_halt and producer.converged is not None: + raise ValueError("fixed-budget producer must not claim a convergence result") + if args.vote_to_halt and producer.converged is not True: + raise ValueError("vote-to-halt producer did not converge") + required = ["id", "pagerank"] if args.program == "pagerank" else ["id", "distance"] + if args.program == "landmarks": + required = ["id", *[f"dist_{key}" for key in sorted(args.landmarks)]] + if table.column_names != required: + raise ValueError("physical output columns differ from the declared program") + output_ids = int64_column(table, "id") + order = np.argsort(output_ids) + if len(output_ids) != len(graph.ids) or not np.array_equal(output_ids[order], graph.ids): + raise ValueError("physical output must contain every original vertex exactly once") + receipt.vertices, receipt.edges, receipt.result_rows = len(graph.ids), len(graph.source), table.num_rows + for name in required[1:]: + column = table[name] + if args.program == "landmarks": + if column.type != pa.int32() or column.null_count: + raise ValueError("landmark distances must be nonnull physical Int32") + actual = np.asarray(column.combine_chunks().to_numpy(zero_copy_only=False), dtype=np.float64)[order] + landmark = int(name.removeprefix("dist_")) + expected = ( + complete_distances(graph, landmark, reverse=args.reversed_edges, weighted=False) + if args.vote_to_halt + else bounded_distances( + graph, landmark, producer.iterations, reverse=args.reversed_edges, weighted=False + ) + ) + expected[np.isinf(expected)] = MAX_HOPS + mismatches, maximum = compare(actual, expected, exact=True) + else: + if column.type != pa.float64(): + raise ValueError("distance/rank must have physical Float64 type") + actual = np.asarray(column.combine_chunks().to_numpy(zero_copy_only=False), dtype=np.float64)[order] + if args.program == "sssp": + expected = ( + complete_distances(graph, args.source) + if args.vote_to_halt + else bounded_distances(graph, args.source, producer.iterations) + ) + if not np.array_equal( + np.asarray(column.is_null().to_numpy(zero_copy_only=False))[order], np.isinf(expected) + ): + raise ValueError("only unreachable SSSP rows may be physically null") + actual[np.isnan(actual)] = np.inf + else: + if column.null_count: + raise ValueError("PageRank must be nonnull") + expected, converged = pagerank(graph, args, producer.iterations) + if args.vote_to_halt and not converged: + raise ValueError("PageRank halt disagrees with the independent active-set recurrence") + mismatches, maximum = compare(actual, expected, exact=False) + receipt.mismatch_rows += mismatches + receipt.max_absolute_difference = max(receipt.max_absolute_difference, maximum) + if receipt.mismatch_rows: + raise ValueError(f"{receipt.mismatch_rows} mismatched physical values") + + +def qualify_official_sssp( + table: pa.Table, graph: Graph, arguments: Arguments, reference: Path, receipt: OracleReceipt +) -> None: + """Compare every weighted SSSP value with an explicitly supplied Graphalytics reference.""" + if arguments.program != "sssp" or not arguments.vote_to_halt or not arguments.input_weights: + raise ValueError("official SSSP comparison requires a completed weighted run to the halt") + records = np.loadtxt(reference, dtype=[("id", np.int64), ("distance", np.float64)], ndmin=1) + reference_ids = np.asarray(records["id"], dtype=np.int64) + expected = np.asarray(records["distance"], dtype=np.float64)[np.argsort(reference_ids)] + if len(reference_ids) != len(graph.ids) or not np.array_equal(np.sort(reference_ids), graph.ids): + raise ValueError("official SSSP reference must cover the full original unique vertex domain") + if np.any(np.isnan(expected)) or np.any(expected < 0): + raise ValueError("official SSSP reference contains invalid distances") + # Graphalytics uses the largest DOUBLE for unreachable; permit explicit +inf too. + expected[expected == np.finfo(np.float64).max] = np.inf + actual = np.asarray(table["distance"].combine_chunks().to_numpy(zero_copy_only=False), dtype=np.float64) + actual = actual[np.argsort(int64_column(table, "id"))] + actual[np.isnan(actual)] = np.inf + mismatches, maximum = compare(actual, expected, exact=False) + receipt.mismatch_rows += mismatches + receipt.max_absolute_difference = max(receipt.max_absolute_difference, maximum) + if mismatches: + raise ValueError(f"{mismatches} mismatches against the full official SSSP reference") + receipt.official_sssp_passed = True + + +def execute(cell: Path, output: Path, *, official_sssp: Path | None = None) -> None: + producer_path = cell / "receipt.json" + producer = TypeAdapter(Receipt).validate_json(producer_path.read_bytes()) + original_paths = [ + producer_path, + *parquet_files(producer.arguments.vertices), + *parquet_files(producer.arguments.edges), + *parquet_files(cell / "result"), + ] + if official_sssp is not None: + original_paths.append(official_sssp) + receipt = OracleReceipt( + datetime.now(timezone.utc).isoformat(), + "running", + pin(producer_path), + [pin(path) for path in original_paths], + [], + producer.arguments.program, + ) + try: + graph = read_graph(producer.arguments) + table = ds.dataset(str(cell / "result"), format="parquet").to_table() + qualify(producer, graph, table, receipt) + if official_sssp is not None: + receipt.official_sssp = pin(official_sssp) + qualify_official_sssp(table, graph, producer.arguments, official_sssp, receipt) + receipt.after = [pin(path) for path in original_paths] + if receipt.before != receipt.after: + raise ValueError("an original input, producer receipt or raw output changed during qualification") + receipt.status = "passed_full_physical_oracle" + except BaseException as error: + receipt.status = "error" + receipt.error = f"{type(error).__name__}: {error}" + raise + finally: + with output.open("x") as stream: + stream.write(json.dumps(asdict(receipt), indent=2, allow_nan=False) + "\n") + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--cell", required=True, type=Path) + parser.add_argument("--output", required=True, type=Path) + parser.add_argument("--official-sssp", type=Path) + args = parser.parse_args() + execute(args.cell, args.output, official_sssp=args.official_sssp) + + +if __name__ == "__main__": + main() diff --git a/examples/extensions/benchmarks/test_e0_oracle.py b/examples/extensions/benchmarks/test_e0_oracle.py new file mode 100644 index 00000000..c1ad4183 --- /dev/null +++ b/examples/extensions/benchmarks/test_e0_oracle.py @@ -0,0 +1,225 @@ +"""Negative physical-output controls and independent algorithm expectations.""" + +from __future__ import annotations + +import json +from dataclasses import asdict, replace +from pathlib import Path + +import e0_oracle +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq +import pytest +from e0_cell import Arguments, Receipt +from e0_oracle import ( + MAX_HOPS, + FilePin, + Graph, + OracleReceipt, + bounded_distances, + complete_distances, + pagerank, + qualify, + qualify_official_sssp, +) + + +def graph() -> Graph: + return Graph( + np.array([-9, 0, 2, 4, 8], dtype=np.int64), + np.array([0, 1, 1, 2, 2], dtype=np.int64), + np.array([1, 2, 2, 1, 3], dtype=np.int64), + np.array([1.0, 2.0, 0.0, 0.0, 0.5], dtype=np.float64), + ) + + +def arguments() -> Arguments: + return Arguments( + "sc://127.0.0.1:1", + Path("vertices"), + Path("edges"), + Path("out"), + "sssp", + -9, + (-9,), + 1, + 10, + 0.01, + 0.15, + False, + True, + False, + True, + False, + ) + + +def oracle() -> OracleReceipt: + return OracleReceipt("2026-10-03T00:00:00+00:00", "running", FilePin("receipt", 0, "unused"), [], [], "sssp") + + +def producer(args: Arguments | None = None, *, iterations: int = 4) -> Receipt: + return Receipt( + "2026-10-03T00:00:00+00:00", + "completed_unvalidated", + args or arguments(), + "out", + iterations=iterations, + converged=None if args is not None and not args.vote_to_halt else True, + session_stopped=True, + ) + + +def output() -> pa.Table: + return pa.table( + { + "id": pa.array([-9, 0, 2, 4, 8], type=pa.int64()), + "distance": pa.array([0.0, 1.0, 1.0, 1.5, None], type=pa.float64()), + } + ) + + +def test_independent_heap_and_bounded_all_edge_recurrence() -> None: + np.testing.assert_array_equal(complete_distances(graph(), -9), [0.0, 1.0, 1.0, 1.5, np.inf]) + np.testing.assert_array_equal(bounded_distances(graph(), -9, 1), [0.0, 1.0, np.inf, np.inf, np.inf]) + np.testing.assert_array_equal(complete_distances(graph(), 4, reverse=True, weighted=False), [3, 2, 1, 0, np.inf]) + + +def test_full_output_domain_order_independence_and_null_unreachable() -> None: + receipt = oracle() + qualify(producer(), graph(), output().take(pa.array([4, 2, 0, 3, 1])), receipt) + assert receipt.result_rows == 5 and receipt.mismatch_rows == 0 + + +@pytest.mark.parametrize( + "distances", + [[0.0, 1.0, 2.0, 1.5, None], [0.0, None, 1.0, 1.5, None], [0.0, 1.0, 1.0, 1.5, 0.0], [0.0, 1.0, np.inf, 1.5, None]], +) +def test_wrong_distance_or_null_classification_is_rejected(distances: list[float | None]) -> None: + table = output().set_column(1, "distance", pa.array(distances, type=pa.float64())) + with pytest.raises(ValueError): + qualify(producer(), graph(), table, oracle()) + + +@pytest.mark.parametrize("ids", [[-9, 0, 2, 4, 4], [-9, 0, 2, 4, 99]]) +def test_duplicate_missing_or_foreign_id_is_rejected(ids: list[int]) -> None: + table = output().set_column(0, "id", pa.array(ids, type=pa.int64())) + with pytest.raises(ValueError, match="every original vertex exactly once"): + qualify(producer(), graph(), table, oracle()) + + +def test_wrong_physical_output_type_is_rejected() -> None: + table = output().set_column(1, "distance", pa.array([0.0, 1.0, 1.0, 1.5, None], type=pa.float32())) + with pytest.raises(ValueError, match="Float64"): + qualify(producer(), graph(), table, oracle()) + + +def test_fixed_budget_is_not_mistaken_for_complete_shortest_paths() -> None: + args = replace(arguments(), vote_to_halt=False, iterations=1) + table = output().set_column(1, "distance", pa.array([0.0, 1.0, None, None, None], type=pa.float64())) + qualify(producer(args, iterations=1), graph(), table, oracle()) + with pytest.raises(ValueError): + qualify(producer(args, iterations=1), graph(), output(), oracle()) + + +def test_landmark_int32_max_and_direction_are_checked() -> None: + args = replace(arguments(), program="landmarks", landmarks=(4,), reversed_edges=True) + table = pa.table( + {"id": pa.array([-9, 0, 2, 4, 8], type=pa.int64()), "dist_4": pa.array([3, 2, 1, 0, MAX_HOPS], type=pa.int32())} + ) + qualify(producer(args), graph(), table, oracle()) + with pytest.raises(ValueError): + qualify(producer(replace(args, reversed_edges=False)), graph(), table, oracle()) + with pytest.raises(ValueError, match="Int32"): + qualify(producer(args), graph(), table.set_column(1, "dist_4", table["dist_4"].cast(pa.int64())), oracle()) + + +def test_graphx_sink_recurrence_and_normalization() -> None: + sink = Graph( + np.array([1, 2, 3, 4], dtype=np.int64), + np.array([0, 0, 1, 2], dtype=np.int64), + np.array([1, 2, 3, 3], dtype=np.int64), + np.ones(4, dtype=np.float64), + ) + args = replace(arguments(), program="pagerank", input_weights=False) + ranks, converged = pagerank(sink, args, 3) + np.testing.assert_allclose(ranks, [0.15, 0.21375, 0.21375, 0.513375], rtol=0, atol=1e-15) + assert converged + normalized, _ = pagerank(sink, replace(args, normalized=True), 3) + assert normalized.sum() == pytest.approx(1.0) + + +def prepare_cell(tmp_path: Path) -> Path: + cell = tmp_path / "cell" + cell.mkdir() + (cell / "result").mkdir() + original = graph() + vertices, edges = tmp_path / "vertices.parquet", tmp_path / "edges.parquet" + pq.write_table(pa.table({"id": pa.array(original.ids, type=pa.int64())}), vertices) + pq.write_table( + pa.table( + { + "source": pa.array(original.ids[original.source], type=pa.int64()), + "target": pa.array(original.ids[original.target], type=pa.int64()), + "weight": pa.array(original.weights, type=pa.float64()), + } + ), + edges, + ) + pq.write_table(output(), cell / "result" / "part-0.parquet") + args = replace(arguments(), vertices=vertices, edges=edges, output=cell) + (cell / "receipt.json").write_text(json.dumps(asdict(producer(args)), default=str)) + return cell + + +def test_complete_disk_oracle_binds_every_original_and_result(tmp_path: Path) -> None: + cell = prepare_cell(tmp_path) + receipt = tmp_path / "oracle.json" + e0_oracle.execute(cell, receipt) + saved = json.loads(receipt.read_text()) + assert saved["status"] == "passed_full_physical_oracle" + assert len(saved["before"]) == 4 and saved["before"] == saved["after"] + assert saved["result_rows"] == 5 and saved["mismatch_rows"] == 0 + + +def test_mutated_original_during_qualification_keeps_failed_receipt( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + cell = prepare_cell(tmp_path) + original_qualify = e0_oracle.qualify + + def change_source(producer: Receipt, graph: Graph, table: pa.Table, receipt: OracleReceipt) -> None: + original_qualify(producer, graph, table, receipt) + pq.write_table(pa.table({"id": pa.array([-9, 0, 2, 4, 8, 99], type=pa.int64())}), producer.arguments.vertices) + + monkeypatch.setattr(e0_oracle, "qualify", change_source) + receipt = tmp_path / "failed-oracle.json" + with pytest.raises(ValueError, match="changed during qualification"): + e0_oracle.execute(cell, receipt) + assert json.loads(receipt.read_text())["status"] == "error" + + +def test_full_official_reference_and_unreachable_sentinel(tmp_path: Path) -> None: + reference = tmp_path / "official-SSSP" + reference.write_text("-9 0.0\n0 1.0\n2 1.0\n4 1.5\n8 1.7976931348623157e308\n") + receipt = oracle() + qualify_official_sssp(output(), graph(), arguments(), reference, receipt) + assert receipt.official_sssp_passed + reference.write_text("-9 0.0\n0 1.0\n2 9.0\n4 1.5\n8 1.7976931348623157e308\n") + with pytest.raises(ValueError, match="official SSSP reference"): + qualify_official_sssp(output(), graph(), arguments(), reference, oracle()) + + +def test_undirected_input_expansion_keeps_actual_weights(tmp_path: Path) -> None: + cell = prepare_cell(tmp_path) + args = replace( + arguments(), + vertices=tmp_path / "vertices.parquet", + edges=tmp_path / "edges.parquet", + output=cell, + undirected=True, + ) + expanded = e0_oracle.read_graph(args) + assert len(expanded.source) == len(expanded.weights) == 10 + np.testing.assert_array_equal(complete_distances(expanded, 4), [1.5, 0.5, 0.5, 0.0, np.inf]) diff --git a/examples/extensions/graph-algorithms/src/pyspark_pecan/__init__.py b/examples/extensions/graph-algorithms/src/pyspark_pecan/__init__.py index 897a8bd5..fd2d44c8 100644 --- a/examples/extensions/graph-algorithms/src/pyspark_pecan/__init__.py +++ b/examples/extensions/graph-algorithms/src/pyspark_pecan/__init__.py @@ -7,11 +7,20 @@ from .algorithms import ConvergenceError, GraphAlgorithms, Observer from .lifecycle import CancellationToken, GraphCancelledError, GraphResult from .pregel import Pregel -from .types import ContractionStep, GraphOptions, IterationEvent, PageRankOptions, TraversalOptions, WccOptions +from .types import ( + ContractionStep, + GraphOptions, + IterationEvent, + PageRankOptions, + PregelSsspOptions, + ShortestPathsOptions, + TraversalOptions, + WccOptions, +) from .utils import CapabilityError, GraphUtils __all__ = [ "CancellationToken", "CapabilityError", "ContractionStep", "ConvergenceError", "GraphAlgorithms", "GraphCancelledError", "GraphOptions", "GraphResult", "GraphUtils", "IterationEvent", "Observer", "PageRankOptions", - "Pregel", "TraversalOptions", "WccOptions", "pregel", + "Pregel", "PregelSsspOptions", "ShortestPathsOptions", "TraversalOptions", "WccOptions", "pregel", ] diff --git a/examples/extensions/graph-algorithms/src/pyspark_pecan/algorithms.py b/examples/extensions/graph-algorithms/src/pyspark_pecan/algorithms.py index 811ea1a6..7c511e58 100644 --- a/examples/extensions/graph-algorithms/src/pyspark_pecan/algorithms.py +++ b/examples/extensions/graph-algorithms/src/pyspark_pecan/algorithms.py @@ -9,13 +9,13 @@ from __future__ import annotations import warnings -from collections.abc import Callable +from collections.abc import Callable, Sequence from typing import TYPE_CHECKING, Any, Protocol, cast from pyspark.sql.connect import functions as F from pyspark.sql.types import LongType -from . import _contracts, pagerank_delta, pagerank_pregel_delta, pregel, traversal, wcc_randomized +from . import _contracts, pagerank_delta, pagerank_pregel_delta, pregel, shortest_pregel, traversal, wcc_randomized from .lifecycle import CancellationToken, GraphCancelledError, GraphResult from .staging import StagingRun from .types import ( @@ -24,6 +24,8 @@ IterationEvent, PageRankMethod, PageRankOptions, + PregelSsspOptions, + ShortestPathsOptions, TraversalMethod, TraversalOptions, WccMethod, @@ -393,13 +395,14 @@ def bfs(self, vertices: DataFrame, edges: DataFrame, *, source: int, method: Tra """ options = TraversalOptions(source=source, method=method, directed=directed, max_iterations=max_iterations, partitions=partitions) - if options.method == "delta_star": + if options.method in ("delta_star", "pregel"): raise ValueError("unsupported traversal method") return traversal.execute(self, vertices, edges, options=options, weighted=False, cancellation=cancellation) def sssp(self, vertices: DataFrame, edges: DataFrame, *, source: int, method: TraversalMethod = "frontier", directed: bool = True, max_iterations: int = 1000, partitions: int = 4, - cancellation: CancellationToken | None = None, delta: float = 1.0) -> GraphResult: + cancellation: CancellationToken | None = None, delta: float = 1.0, + vote_to_halt: bool = True) -> GraphResult: """Single-source shortest distances for finite nonnegative DOUBLE weights. Edges require a `weight` column. Reference is synchronous Bellman–Ford; @@ -411,9 +414,41 @@ def sssp(self, vertices: DataFrame, edges: DataFrame, *, source: int, method: Tr Weights are assumed finite and non-negative and distance sums are assumed finite (the valid graph contract); nothing checks them. Both methods require an explicit convergence certificate. + + method="pregel" runs a distance-only active-source Pregel program with + one DOUBLE message, MIN, and one state write per superstep. Its output + is id BIGINT, distance DOUBLE? (no parent/hops). vote_to_halt=True counts + active vertices after each step and raises if the cap precedes a halt. + False runs exactly max_iterations supersteps and returns that possibly + incomplete state with converged=None; it runs no activity count. """ + if method == "pregel": + pregel_options = PregelSsspOptions(source=source, directed=directed, max_iterations=max_iterations, + partitions=partitions, vote_to_halt=vote_to_halt) + return shortest_pregel.execute_sssp(self, vertices, edges, options=pregel_options, + cancellation=cancellation) + if not vote_to_halt: + raise ValueError("vote_to_halt=False applies only to method='pregel'") options = TraversalOptions(source=source, method=method, directed=directed, max_iterations=max_iterations, partitions=partitions, delta=delta) if options.method == "push_pull": raise ValueError("unsupported traversal method") return traversal.execute(self, vertices, edges, options=options, weighted=True, cancellation=cancellation) + + def shortest_paths(self, vertices: DataFrame, edges: DataFrame, *, landmarks: Sequence[int], + to_landmarks: bool = False, max_iterations: int = 1000, partitions: int = 4, + vote_to_halt: bool = True, cancellation: CancellationToken | None = None) -> GraphResult: + """Unweighted per-landmark hops, as graphframes-rs's shortest_paths. + + Edges are directed: default distances run from each landmark to the + vertex; to_landmarks=True reverses edges. Output is id BIGINT and one + INT32 dist_ column per unique landmark, in sorted landmark + order. Unreachable is 2^31-1. Landmarks must be graph vertices, assumed + by contract without a membership job. With vote_to_halt, the cap + includes the final no-change step and raises if still active. False + returns the fixed-budget state and converged=None, counting nothing. + """ + options = ShortestPathsOptions(landmarks=tuple(landmarks), to_landmarks=to_landmarks, + max_iterations=max_iterations, partitions=partitions, + vote_to_halt=vote_to_halt) + return shortest_pregel.execute_landmarks(self, vertices, edges, options=options, cancellation=cancellation) diff --git a/examples/extensions/graph-algorithms/src/pyspark_pecan/pregel.py b/examples/extensions/graph-algorithms/src/pyspark_pecan/pregel.py index ad48ff23..c370b236 100644 --- a/examples/extensions/graph-algorithms/src/pyspark_pecan/pregel.py +++ b/examples/extensions/graph-algorithms/src/pyspark_pecan/pregel.py @@ -172,7 +172,6 @@ def loop(self, run: StagingRun, vertices: DataFrame, edges: DataFrame) -> Outcom while (self.limit is None or iteration < self.limit) and not converged: iteration += 1 run.cancellation.check() - self.graph._observe(run, self.algorithm, iteration, "iteration_start") sources = state if not self.destination_state and participation is not None: sources = sources.where(F.col(participation.name)) @@ -189,6 +188,7 @@ def loop(self, run: StagingRun, vertices: DataFrame, edges: DataFrame) -> Outcom aggregated = messages.groupBy("id").agg(*self.aggregates) if self.aggregates else messages received = aggregated.withColumnRenamed("id", "__pregel_to") updated = state.join(received, F.col("id") == F.col("__pregel_to"), "left").select(*updates) + self.graph._observe(run, self.algorithm, iteration, "iteration_start", plan_of=updated) next_path, next_state = run.materialize(updated) run.remove(path) path, state = next_path, next_state diff --git a/examples/extensions/graph-algorithms/src/pyspark_pecan/shortest_pregel.py b/examples/extensions/graph-algorithms/src/pyspark_pecan/shortest_pregel.py new file mode 100644 index 00000000..4b1039c6 --- /dev/null +++ b/examples/extensions/graph-algorithms/src/pyspark_pecan/shortest_pregel.py @@ -0,0 +1,150 @@ +"""Shortest-distance programs on the common Pregel loop. + +Weighted SSSP sends one DOUBLE distance and aggregates MIN. Landmark hops +follow graphframes-rs's `connectivity/shortest_paths.rs` at b4da56d: one INT +column and named MIN message per landmark, MAX_INT as unreachable, forward +edges by default, and reversed edges for distances to landmarks. Neither +program inspects graph validity. Only participating sources join adjacency. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from pyspark.sql.connect import functions as F +from pyspark.sql.types import DoubleType + +from ._contracts import ConvergenceError +from .pregel import Pregel, edge, msg, src + +if TYPE_CHECKING: + from pyspark.sql import DataFrame + + from .algorithms import GraphAlgorithms + from .lifecycle import CancellationToken, GraphResult + from .staging import StagingRun + from .types import PregelSsspOptions, ShortestPathsOptions + +MAX_HOPS = (1 << 31) - 1 + + +def sssp_program(graph: GraphAlgorithms, options: PregelSsspOptions) -> Pregel: + """Nonnegative weighted SSSP, with source distance 0 and null unreachable.""" + changed = msg().isNotNull() & (F.col("distance").isNull() | (msg() < F.col("distance"))) + source = F.col("id") == F.lit(options.source) + program = ( + Pregel(graph, algorithm="sssp-pregel") + .vertex_column( + "distance", + F.when(source, F.lit(0.0)).otherwise(F.lit(None).cast("double")), + F.least(F.col("distance"), msg()), + ) + .participation("participates", source, changed) + .edge_column("weight") + .message(src("distance") + edge("weight"), "src_to_dst") + .aggregate(F.min(msg())) + .skip_destination_state() + .max_iterations(options.max_iterations) + ) + if options.vote_to_halt: + program.vote_to_halt("active", changed) + return program + + +def landmarks_program(graph: GraphAlgorithms, options: ShortestPathsOptions) -> Pregel: + """graphframes-rs's per-landmark hop program, on the selected edge direction.""" + changed = F.lit(False) + initial_participation = F.lit(False) + program = Pregel(graph, algorithm="shortest-paths-pregel") + for landmark in sorted(options.landmarks): + name = str(landmark) + column = "dist_" + name + distance = F.col(column) + received = msg(name) + initial_participation = initial_participation | (F.col("id") == F.lit(landmark)) + changed = changed | (distance > received) + sent = F.when(src(column) < F.lit(MAX_HOPS), src(column) + F.lit(1)).otherwise(F.lit(MAX_HOPS)) + program.vertex_column( + column, + F.when(F.col("id") == F.lit(landmark), F.lit(0)).otherwise(F.lit(MAX_HOPS)), + F.least(distance, received), + ) + program.message(sent, "src_to_dst", name=name).aggregate(F.min(received), name=name) + program.participation("participates", initial_participation, changed) + program.skip_destination_state().max_iterations(options.max_iterations) + if options.vote_to_halt: + program.vote_to_halt("active", changed) + return program + + +def execute_sssp( + graph: GraphAlgorithms, + vertices: DataFrame, + edges: DataFrame, + *, + options: PregelSsspOptions, + cancellation: CancellationToken | None, +) -> GraphResult: + """Execute weighted SSSP; output exactly `id: BIGINT, distance: DOUBLE?`.""" + if "weight" not in edges.columns or not isinstance(edges.schema["weight"].dataType, DoubleType): + raise ValueError("SSSP requires a DOUBLE weight column") + program = sssp_program(graph, options) + + def body(run: StagingRun, vertices: DataFrame, edges: DataFrame, size: int | None) -> GraphResult: + adjacency = ( + edges + if options.directed + else edges.unionByName(edges.select(F.col("dst").alias("src"), F.col("src").alias("dst"), "weight")) + ) + outcome = program.loop(run, vertices, adjacency) + if outcome.converged is False: + raise ConvergenceError(f"sssp-pregel still had active vertices after {options.max_iterations} iterations") + path, result = run.materialize(outcome.state.select("id", "distance")) + handle = run.finish( + path, result, algorithm="sssp-pregel", iterations=outcome.iterations, converged=outcome.converged + ) + handle.method = "pregel" + return handle + + return graph._run( + vertices, + edges, + options.partitions, + cancellation, + body, + edge_columns=("src", "dst", "weight"), + count_vertices=False, + ) + + +def execute_landmarks( + graph: GraphAlgorithms, + vertices: DataFrame, + edges: DataFrame, + *, + options: ShortestPathsOptions, + cancellation: CancellationToken | None, +) -> GraphResult: + """Execute per-landmark hops; an absent route is INT32 MAX, never null.""" + if len(set(options.landmarks)) != len(options.landmarks): + raise ValueError("landmarks must be unique") + program = landmarks_program(graph, options) + + def body(run: StagingRun, vertices: DataFrame, edges: DataFrame, size: int | None) -> GraphResult: + adjacency = ( + edges.select(F.col("dst").alias("src"), F.col("src").alias("dst")) if options.to_landmarks else edges + ) + outcome = program.loop(run, vertices, adjacency) + if outcome.converged is False: + raise ConvergenceError( + f"shortest-paths-pregel still had active vertices after {options.max_iterations} iterations" + ) + columns = ["id", *["dist_" + str(landmark) for landmark in sorted(options.landmarks)]] + path, result = run.materialize(outcome.state.select(*columns)) + handle = run.finish( + path, result, algorithm="shortest-paths-pregel", iterations=outcome.iterations, converged=outcome.converged + ) + handle.method = "pregel" + return handle + + return graph._run(vertices, edges, options.partitions, cancellation, body, count_vertices=False) diff --git a/examples/extensions/graph-algorithms/src/pyspark_pecan/types.py b/examples/extensions/graph-algorithms/src/pyspark_pecan/types.py index 50e1cbf1..555c563b 100644 --- a/examples/extensions/graph-algorithms/src/pyspark_pecan/types.py +++ b/examples/extensions/graph-algorithms/src/pyspark_pecan/types.py @@ -7,7 +7,7 @@ from __future__ import annotations -from typing import Literal +from typing import Annotated, Literal from pydantic import BaseModel, ConfigDict, Field @@ -19,7 +19,8 @@ Direction = Literal["push", "pull"] PageRankMethod = Literal["power", "delta", "pregel", "pregel_delta"] WccMethod = Literal["min_label", "randomized", "randomized_fused"] -TraversalMethod = Literal["reference", "frontier", "push_pull", "delta_star"] +TraversalMethod = Literal["reference", "frontier", "push_pull", "delta_star", "pregel"] +Int64Id = Annotated[int, Field(ge=INT64_MIN, le=INT64_MAX)] class IterationEvent(BaseModel): @@ -114,6 +115,30 @@ class TraversalOptions(BaseModel): delta: float = Field(default=1.0, gt=0.0, allow_inf_nan=False) +class PregelSsspOptions(BaseModel): + """Distance-only weighted Pregel, with an optional fixed superstep budget.""" + + model_config = ConfigDict(frozen=True, extra="forbid", strict=True) + + source: Int64Id + directed: bool = True + max_iterations: int = Field(default=1000, ge=1) + partitions: int = Field(default=4, ge=1) + vote_to_halt: bool = True + + +class ShortestPathsOptions(BaseModel): + """Per-landmark unweighted hops; landmark membership is a caller contract.""" + + model_config = ConfigDict(frozen=True, extra="forbid", strict=True) + + landmarks: tuple[Int64Id, ...] = Field(min_length=1) + to_landmarks: bool = False + max_iterations: int = Field(default=1000, ge=0) + partitions: int = Field(default=4, ge=1) + vote_to_halt: bool = True + + class MassResidual(BaseModel): """Scalar statistics of a delta PageRank state.""" diff --git a/examples/extensions/graph-algorithms/tests/test_shortest_pregel.py b/examples/extensions/graph-algorithms/tests/test_shortest_pregel.py new file mode 100644 index 00000000..8978bacf --- /dev/null +++ b/examples/extensions/graph-algorithms/tests/test_shortest_pregel.py @@ -0,0 +1,189 @@ +"""Full-output independent distance oracles for the common Pregel programs.""" + +from __future__ import annotations + +import heapq +from collections import deque +from typing import TYPE_CHECKING + +import pytest +from pyspark.sql.types import DoubleType, IntegerType, LongType +from pyspark_pecan import CancellationToken, ConvergenceError, GraphAlgorithms, GraphCancelledError +from pyspark_pecan.shortest_pregel import MAX_HOPS + +if TYPE_CHECKING: + from pyspark.sql import DataFrame + from pyspark.sql.connect.session import SparkSession + +pytestmark = pytest.mark.integration + +IDS = [-(1 << 63), -9, 0, 2, 3, 5, (1 << 63) - 1] +EDGES = [ + (-(1 << 63), -9, 1.0), + (-(1 << 63), 2, 5.0), + (-9, 0, 0.0), + (0, -9, 0.0), + (0, 2, 1.0), + (0, 2, 3.0), + (2, 3, 0.5), + (3, 3, 0.0), + (5, (1 << 63) - 1, 2.0), +] + + +def frames(spark: SparkSession, ids: list[int], edges: list[tuple[int, int, float]]) -> tuple[DataFrame, DataFrame]: + return ( + spark.createDataFrame([(key,) for key in ids], "id long"), + spark.createDataFrame(edges, "src long, dst long, weight double"), + ) + + +def dijkstra( + ids: list[int], edges: list[tuple[int, int, float]], source: int, directed: bool +) -> dict[int, float | None]: + """Heap Dijkstra, independent of the Pregel superstep recurrence.""" + adjacency: dict[int, list[tuple[int, float]]] = {key: [] for key in ids} + for start, end, weight in edges: + adjacency[start].append((end, weight)) + if not directed: + adjacency[end].append((start, weight)) + found = {source: 0.0} + queue = [(0.0, source)] + while queue: + distance, vertex = heapq.heappop(queue) + if distance != found[vertex]: + continue + for target, weight in adjacency[vertex]: + candidate = distance + weight + if target not in found or candidate < found[target]: + found[target] = candidate + heapq.heappush(queue, (candidate, target)) + return {key: found.get(key) for key in ids} + + +def hop_distances(ids: list[int], edges: list[tuple[int, int, float]], source: int, reverse: bool) -> dict[int, int]: + """Queue BFS from one landmark, independently of the column program.""" + adjacency: dict[int, list[int]] = {key: [] for key in ids} + for start, end, _ in edges: + if reverse: + start, end = end, start + adjacency[start].append(end) + found = {source: 0} + queue = deque([source]) + while queue: + vertex = queue.popleft() + for target in adjacency[vertex]: + if target not in found: + found[target] = found[vertex] + 1 + queue.append(target) + return {key: found.get(key, MAX_HOPS) for key in ids} + + +@pytest.mark.parametrize("directed", [True, False]) +@pytest.mark.parametrize("partitions", [1, 3]) +def test_sssp_full_distance_domain_signed_ids_cycles_and_parallel_edges( + spark: SparkSession, directed: bool, partitions: int +) -> None: + with GraphAlgorithms(spark, snapshot_inputs=False).sssp( + *frames(spark, IDS, EDGES), source=IDS[0], method="pregel", directed=directed, partitions=partitions + ) as result: + rows = result.frame.collect() + assert result.frame.columns == ["id", "distance"] + assert isinstance(result.frame.schema["id"].dataType, LongType) + assert isinstance(result.frame.schema["distance"].dataType, DoubleType) + assert len(rows) == len(IDS) + assert {row.id for row in rows} == set(IDS) + assert {row.id: row.distance for row in rows} == dijkstra(IDS, EDGES, IDS[0], directed) + assert result.algorithm == "sssp-pregel" and result.method == "pregel" + assert result.converged is True + + +def test_sssp_fixed_budget_returns_bounded_state_without_count( + spark: SparkSession, monkeypatch: pytest.MonkeyPatch +) -> None: + vertices, edges = frames(spark, [1, 2, 3], [(1, 2, 1.0), (2, 3, 1.0)]) + + def unexpected_count(frame: DataFrame) -> int: + raise AssertionError("fixed-budget Pregel must not count") + + monkeypatch.setattr(type(vertices), "count", unexpected_count) + with GraphAlgorithms(spark).sssp( + vertices, edges, source=1, method="pregel", max_iterations=1, vote_to_halt=False + ) as result: + assert {row.id: row.distance for row in result.frame.collect()} == {1: 0.0, 2: 1.0, 3: None} + assert result.iterations == 1 and result.converged is None + + +def test_sssp_cap_is_not_reported_as_convergence(spark: SparkSession) -> None: + with pytest.raises(ConvergenceError, match="active vertices after 1 iterations"): + GraphAlgorithms(spark).sssp( + *frames(spark, [1, 2, 3], [(1, 2, 1.0), (2, 3, 1.0)]), source=1, method="pregel", max_iterations=1 + ) + + +def test_sssp_isolated_source_and_empty_adjacency(spark: SparkSession) -> None: + with GraphAlgorithms(spark).sssp(*frames(spark, [0, -9, 8], []), source=-9, method="pregel") as result: + assert {row.id: row.distance for row in result.frame.collect()} == {-9: 0.0, 0: None, 8: None} + assert result.iterations == 1 and result.converged is True + + +@pytest.mark.parametrize("reverse", [True, False]) +def test_landmarks_have_independent_forward_and_backward_hops(spark: SparkSession, reverse: bool) -> None: + landmarks = [5, IDS[0]] + with GraphAlgorithms(spark, snapshot_inputs=False).shortest_paths( + *frames(spark, IDS, EDGES), landmarks=landmarks, to_landmarks=reverse, partitions=3 + ) as result: + rows = result.frame.collect() + assert len(rows) == len(IDS) and {row.id for row in rows} == set(IDS) + assert result.frame.columns == ["id", *[f"dist_{landmark}" for landmark in sorted(landmarks)]] + assert isinstance(result.frame.schema["id"].dataType, LongType) + for landmark in landmarks: + column = f"dist_{landmark}" + assert isinstance(result.frame.schema[column].dataType, IntegerType) + assert {row.id: row[column] for row in rows} == hop_distances(IDS, EDGES, landmark, reverse) + assert result.algorithm == "shortest-paths-pregel" and result.converged is True + + +def test_landmarks_match_graphframes_rs_small_graph(spark: SparkSession) -> None: + edges = [(1, 2, 1.0), (2, 3, 1.0), (2, 4, 1.0), (3, 4, 1.0), (4, 1, 1.0), (4, 2, 1.0), (2, 1, 1.0), (3, 2, 1.0)] + with GraphAlgorithms(spark).shortest_paths( + *frames(spark, [1, 2, 3, 4], edges), landmarks=[4, 1], to_landmarks=True + ) as result: + got = {row.id: (row.dist_1, row.dist_4) for row in result.frame.collect()} + assert got == {1: (0, 2), 2: (1, 1), 3: (2, 1), 4: (1, 0)} + + +def test_landmarks_fixed_zero_budget_and_disconnected_isolates( + spark: SparkSession, monkeypatch: pytest.MonkeyPatch +) -> None: + vertices, edges = frames(spark, [-9, 0, 8], []) + + def unexpected_count(frame: DataFrame) -> int: + raise AssertionError("fixed-budget Pregel must not count") + + monkeypatch.setattr(type(vertices), "count", unexpected_count) + with GraphAlgorithms(spark).shortest_paths( + vertices, edges, landmarks=[-9], max_iterations=0, vote_to_halt=False + ) as result: + assert {row.id: row["dist_-9"] for row in result.frame.collect()} == {-9: 0, 0: MAX_HOPS, 8: MAX_HOPS} + assert result.iterations == 0 and result.converged is None + + +def test_landmark_cap_and_duplicate_arguments(spark: SparkSession) -> None: + vertices, edges = frames(spark, [1, 2, 3], [(1, 2, 1.0), (2, 3, 1.0)]) + with pytest.raises(ConvergenceError, match="active vertices after 1 iterations"): + GraphAlgorithms(spark).shortest_paths(vertices, edges, landmarks=[1], max_iterations=1) + with pytest.raises(ValueError, match="landmarks must be unique"): + GraphAlgorithms(spark).shortest_paths(vertices, edges, landmarks=[1, 1]) + + +def test_cancellation_uses_the_common_owned_loop(spark: SparkSession) -> None: + cancellation = CancellationToken() + + def cancel(event: object) -> None: + cancellation.cancel() + + with pytest.raises(GraphCancelledError): + GraphAlgorithms(spark, observer=cancel).sssp( + *frames(spark, IDS, EDGES), source=IDS[0], method="pregel", cancellation=cancellation + ) diff --git a/examples/extensions/graph-algorithms/tests/test_shortest_pregel_options.py b/examples/extensions/graph-algorithms/tests/test_shortest_pregel_options.py new file mode 100644 index 00000000..a610f5d0 --- /dev/null +++ b/examples/extensions/graph-algorithms/tests/test_shortest_pregel_options.py @@ -0,0 +1,43 @@ +"""Free argument-domain admission for the new Pregel distance programs.""" + +from __future__ import annotations + +from typing import cast + +import pytest +from pyspark_pecan import GraphAlgorithms, PregelSsspOptions, ShortestPathsOptions +from pyspark_pecan.shortest_pregel import landmarks_program, sssp_program +from pyspark_pecan.types import INT64_MAX, INT64_MIN + + +@pytest.mark.parametrize("source", [INT64_MIN, INT64_MAX]) +def test_source_full_signed_domain_and_scalar_message(source: int) -> None: + graph = cast(GraphAlgorithms, object()) + options = PregelSsspOptions(source=source) + program = sssp_program(graph, options) + assert [column.name for column in program.columns] == ["distance"] + assert program.edge_columns == ["src", "dst", "weight"] + assert program.destination_state is False and program.vote is not None + assert len(program.messages) == len(program.aggregates) == 1 + + +def test_fixed_budget_and_sorted_per_landmark_columns() -> None: + graph = cast(GraphAlgorithms, object()) + options = ShortestPathsOptions(landmarks=(INT64_MAX, INT64_MIN), vote_to_halt=False) + program = landmarks_program(graph, options) + assert [column.name for column in program.columns] == [f"dist_{INT64_MIN}", f"dist_{INT64_MAX}"] + assert program.destination_state is False and program.vote is None + assert len(program.messages) == len(program.aggregates) == 2 + assert sssp_program(graph, PregelSsspOptions(source=0, vote_to_halt=False)).vote is None + + +@pytest.mark.parametrize("source", [INT64_MIN - 1, INT64_MAX + 1, True]) +def test_invalid_source_arguments_need_no_data_operation(source: int) -> None: + with pytest.raises(ValueError): + PregelSsspOptions(source=source) + + +@pytest.mark.parametrize("landmarks", [(), (INT64_MIN - 1,), (INT64_MAX + 1,), (True,)]) +def test_invalid_landmark_arguments_need_no_data_operation(landmarks: tuple[int, ...]) -> None: + with pytest.raises(ValueError): + ShortestPathsOptions(landmarks=landmarks)