feat: implement the graph API and core algorithms

.pre-commit-config.yaml

@@ -38,7 +38,14 @@ repos:
   rev: v2.1.0
   hooks:
   - id: mypy                           # Static type checker for Python
-    additional_dependencies: [daft>=0.7.5]
+    additional_dependencies: [daft>=0.7.5, typing_extensions]
+    # Scope to the library, matching [tool.mypy] files in pyproject.toml. This
+    # hook passes filenames explicitly, which overrides that setting, so the
+    # exclude has to be repeated here. Tests build graphs from daft expressions
+    # (e.g. `col("id") != 3`), which daft's stubs type as `bool` rather than
+    # `Expression`, so type-checking them reports dozens of false positives that
+    # would drown real errors.
+    exclude: ^tests/
 
   # Check that uv.lock is up to date with pyproject.toml
 - repo: https://github.com/astral-sh/uv-pre-commit

.ruff.toml

@@ -46,7 +46,7 @@ ignore = [
 preview = true
 
 [lint.isort]
-known-first-party = ["daft_ext_template"]
+known-first-party = ["daft_graph"]
 known-third-party = ["daft"]
 required-imports = ["from __future__ import annotations"]
 

README.md

@@ -1,3 +1,100 @@
 # daft-graph
 
 Graph processing algorithms using Daft under the hood.
+
+`daft-graph` brings GraphFrames style graph operations to the [Daft](https://docs.daft.ai)
+DataFrame engine. A graph is two DataFrames, a vertex table keyed by `id` and an
+edge table keyed by `src`/`dst`, and every algorithm returns an ordinary DataFrame
+that composes with the rest of a Daft pipeline. It runs on the native runner locally
+and on Ray when distributed.
+
+## Install
+
+```bash
+uv sync
+```
+
+Requires Python 3.10 to 3.13. The core install depends on `daft` alone. The single
+node solves (connected components `strategy="local"` and `svd_plus_plus`) use numpy
+and scipy from the optional `local` extra:
+
+```bash
+uv sync --extra local          # or: pip install 'daft-graph[local]'
+```
+
+## Quick start
+
+```python
+import daft
+from daft_graph import UndirectedGraph, connected_components
+
+edges = daft.from_pydict({"src": [0, 1, 3], "dst": [1, 2, 4]})
+graph = UndirectedGraph(edges)          # vertices derived from the endpoints
+
+connected_components(graph).show()
+```
+
+Graphs come in two flavors and the type carries the direction semantics, so an
+algorithm declares which one it needs. `pagerank` takes a `DirectedGraph`;
+`connected_components` accepts either.
+
+```python
+from daft_graph import DirectedGraph, pagerank
+
+g = DirectedGraph(edges)
+pagerank(g).show()                      # directed
+connected_components(g).show()          # undirected semantics either way
+pagerank(g.reverse()).show()            # every edge flipped
+```
+
+Convert between flavors with `g.as_undirected()` and `g.as_directed()`. Traversal
+follows the graph's own semantics, so there is no `directed` keyword; pass an
+`UndirectedGraph` (or call `as_undirected()`) to walk edges both ways.
+
+## API
+
+### Graph model
+
+- `DirectedGraph(edges, vertices=None, *, src_col="src", dst_col="dst", id_col="id", validate=False)`
+- `UndirectedGraph(edges, vertices=None, *, src_col="src", dst_col="dst", id_col="id", validate=False)`
+- `Graph` - the abstract base, for type annotations and isinstance checks
+
+Base methods: `degrees`, `triplets`, `filter_vertices`, `filter_edges`,
+`drop_isolated_vertices`, `degree_by_type`, `num_vertices`, `num_edges`, `bfs_paths`.
+`DirectedGraph` adds `in_degrees`, `out_degrees`, `reverse`, `as_undirected`, `find`.
+`UndirectedGraph` adds `as_directed`. Transforms return the caller's flavor.
+
+### Algorithms
+
+| Group | Functions |
+|---|---|
+| Connected components | `connected_components`, `strongly_connected_components` |
+| Centrality | `pagerank` (+ personalized), `parallel_personalized_pagerank` |
+| Traversal | `bfs`, `bfs_paths`, `shortest_paths`, `all_shortest_paths`, `all_paths` |
+| Community | `label_propagation`, `power_iteration_clustering` |
+| Motif | `find` (GraphFrames style DSL) |
+| Message passing | `aggregate_messages`, `pregel` |
+| Other | `triangle_count`, `k_core`, `has_cycle`, `vertices_on_cycles`, `maximal_independent_set`, `random_walks`, `svd_plus_plus`, `hyper_anf` |
+| Id indexing | `reindex`, `restore_ids` for arbitrary (non int) ids |
+| Edge utils | `canonicalize`, `symmetrize`, `dedupe_edges`, `drop_self_loops`, `to_edges`, `validate_edges` |
+
+Which flavor each algorithm takes, and full examples, are in [`docs/usage.md`](docs/usage.md).
+
+## Design notes
+
+- Iterative algorithms run on `iterate_to_fixed_point`, which materializes the state
+  between rounds to truncate the Daft logical plan. This is the analog of GraphFrames
+  checkpointing and is what keeps the iterative joins from growing an unbounded plan.
+- Connected components ports the large star and small star contraction algorithm
+  (Kiveris et al. 2014), the same method Spark GraphFrames uses by default, with an
+  optional `scipy.sparse.csgraph` single node solve for small edge sets.
+- Correctness is validated against igraph (connected components) and networkx
+  (PageRank and others), which are test only dependencies.
+
+## Development
+
+```bash
+uv sync                      # install with dev group
+uv run pytest tests/ -v      # run the suite
+uv run pre-commit run --all-files   # ruff + mypy style checks
+```

benchmarks/bench_cc.py

@@ -0,0 +1,91 @@
+"""Benchmark connected components on synthetic random graphs.
+
+Reports the star contraction round count and wall time on the active Daft
+runner. Uses a few internal helpers (prefixed ``_``) so it can report the round
+count, which the public ``connected_components`` does not expose.
+
+Requires the optional ``local`` extra for numpy (``uv sync --extra local`` or
+``pip install 'daft-graph[local]'``).
+
+Examples:
+    uv run python benchmarks/bench_cc.py --edges 1000000
+    DAFT_RUNNER=ray uv run python benchmarks/bench_cc.py --edges 10000000
+"""
+
+from __future__ import annotations
+
+import argparse
+import time
+
+import daft
+import numpy as np
+
+from daft_graph.algorithms.connected_components import (
+    _assign_components,
+    _attach_isolated,
+    _canonical_equal,
+    _propagate_min_labels,
+    _star_step,
+)
+from daft_graph.edges import canonicalize
+from daft_graph.graph import UndirectedGraph
+from daft_graph.iterate import iterate_to_fixed_point
+from daft_graph.schema import COMPONENT, DST, SRC
+
+
+def _generate_edges(n_nodes: int, n_edges: int, seed: int) -> daft.DataFrame:
+    rng = np.random.default_rng(seed)
+    src = rng.integers(0, n_nodes, size=n_edges)
+    dst = rng.integers(0, n_nodes, size=n_edges)
+    return daft.from_pydict({SRC: src.tolist(), DST: dst.tolist()})
+
+
+def _parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="daft-graph connected components benchmark")
+    parser.add_argument("--edges", type=int, default=1_000_000)
+    parser.add_argument("--nodes", type=int, default=None)
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--max-iters", type=int, default=30)
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = _parse_args()
+    n_nodes = args.nodes if args.nodes is not None else max(2, args.edges // 5)
+
+    graph = UndirectedGraph(_generate_edges(n_nodes, args.edges, args.seed))
+
+    start = time.perf_counter()
+    edges = canonicalize(graph.edges)
+    final_edges, star_rounds = iterate_to_fixed_point(edges, _star_step, _canonical_equal, max_iters=args.max_iters)
+    after_star = time.perf_counter()
+
+    assignments = _assign_components(final_edges)
+    assignments = _propagate_min_labels(
+        final_edges, assignments, max_iters=args.max_iters, materialize_every=1, checkpoint_dir=None
+    )
+    result = _attach_isolated(graph.vertices, assignments).collect()
+    end = time.perf_counter()
+
+    n_vertices = result.count_rows()
+    n_components = result.select(COMPONENT).distinct().count_rows()
+
+    rows = [
+        ("requested edges", f"{args.edges:,}"),
+        ("requested nodes", f"{n_nodes:,}"),
+        ("vertices", f"{n_vertices:,}"),
+        ("components", f"{n_components:,}"),
+        ("star rounds", f"{star_rounds:,}"),
+        ("star loop seconds", f"{after_star - start:.2f}"),
+        ("assign + label seconds", f"{end - after_star:.2f}"),
+        ("total seconds", f"{end - start:.2f}"),
+    ]
+    print("daft-graph connected components benchmark")
+    print(f"{'metric':<26}{'value':>16}")
+    print("-" * 42)
+    for name, value in rows:
+        print(f"{name:<26}{value:>16}")
+
+
+if __name__ == "__main__":
+    main()

daft_graph/__init__.py

@@ -1,6 +1,88 @@
-from __future__ import annotations
+"""daft-graph: distributed graph operations for the Daft DataFrame engine on Ray.
 
-import daft
+Construct a :class:`DirectedGraph` or an :class:`UndirectedGraph`. ``Graph`` is
+the abstract base, exported for type annotations and isinstance checks.
 
+Example:
+    >>> import daft
+    >>> from daft_graph import UndirectedGraph, connected_components
+    >>> edges = daft.from_pydict({"src": [0, 1, 3], "dst": [1, 2, 4]})
+    >>> components = connected_components(UndirectedGraph(edges))
+"""
 
-__all__ = ["greet"]
+from importlib.metadata import PackageNotFoundError, version
+
+from daft_graph.algorithms.all_paths import all_paths
+from daft_graph.algorithms.bfs import all_shortest_paths, bfs, bfs_paths
+from daft_graph.algorithms.connected_components import connected_components
+from daft_graph.algorithms.cycles import has_cycle, vertices_on_cycles
+from daft_graph.algorithms.hyper_anf import hyper_anf
+from daft_graph.algorithms.k_core import k_core
+from daft_graph.algorithms.label_propagation import label_propagation
+from daft_graph.algorithms.maximal_independent_set import maximal_independent_set
+from daft_graph.algorithms.pagerank import pagerank, parallel_personalized_pagerank
+from daft_graph.algorithms.power_iteration_clustering import power_iteration_clustering
+from daft_graph.algorithms.random_walks import random_walks
+from daft_graph.algorithms.shortest_paths import shortest_paths
+from daft_graph.algorithms.strongly_connected_components import (
+    strongly_connected_components,
+)
+from daft_graph.algorithms.svd_plus_plus import SvdPlusPlusResult, svd_plus_plus
+from daft_graph.algorithms.triangle_count import triangle_count
+from daft_graph.edges import (
+    canonicalize,
+    dedupe_edges,
+    drop_self_loops,
+    symmetrize,
+    to_edges,
+    validate_edges,
+)
+from daft_graph.graph import DirectedGraph, Graph, UndirectedGraph
+from daft_graph.indexing import ORIGINAL, IndexedGraph, reindex, restore_ids
+from daft_graph.message_passing import aggregate_messages, pregel
+from daft_graph.motif import find
+
+try:
+    __version__ = version("daft-graph")
+except PackageNotFoundError:  # pragma: no cover
+    __version__ = "0.0.0"
+
+__all__ = [
+    "ORIGINAL",
+    "DirectedGraph",
+    "Graph",
+    "IndexedGraph",
+    "SvdPlusPlusResult",
+    "UndirectedGraph",
+    "__version__",
+    "aggregate_messages",
+    "all_paths",
+    "all_shortest_paths",
+    "bfs",
+    "bfs_paths",
+    "canonicalize",
+    "connected_components",
+    "dedupe_edges",
+    "drop_self_loops",
+    "find",
+    "has_cycle",
+    "hyper_anf",
+    "k_core",
+    "label_propagation",
+    "maximal_independent_set",
+    "pagerank",
+    "parallel_personalized_pagerank",
+    "power_iteration_clustering",
+    "pregel",
+    "random_walks",
+    "reindex",
+    "restore_ids",
+    "shortest_paths",
+    "strongly_connected_components",
+    "svd_plus_plus",
+    "symmetrize",
+    "to_edges",
+    "triangle_count",
+    "validate_edges",
+    "vertices_on_cycles",
+]

daft_graph/_compare.py

@@ -0,0 +1,24 @@
+"""Internal helper for comparing DataFrame row sets during iteration."""
+
+from __future__ import annotations
+
+from daft import DataFrame, Expression
+
+from daft_graph.iterate import bound_partitions
+
+
+def rows_equal(a: DataFrame, b: DataFrame, on: list[str | Expression]) -> bool:
+    """True when ``a`` and ``b`` hold the same rows over the columns ``on``.
+
+    Uses anti join counts so the comparison stays inside the Daft engine instead
+    of materializing rows into Python. Both sides are capped with
+    :func:`daft_graph.iterate.bound_partitions` first, a plan rewrite that costs
+    no execution: this runs every iteration, so collecting here would add a
+    distributed round trip per round, while leaving the inputs uncapped would let
+    the anti join inherit a large partition count from the caller's shuffles.
+    """
+    a2 = bound_partitions(a)
+    b2 = bound_partitions(b)
+    left = a2.join(b2, on=on, how="anti").count_rows()
+    right = b2.join(a2, on=on, how="anti").count_rows()
+    return left == 0 and right == 0