GraphRAG on Lakebase

Knowledge-graph-augmented RAG served entirely from Lakebase (managed Postgres). Classic RAG retrieves text chunks by vector similarity and misses the relationships between them. GraphRAG adds a knowledge graph so retrieval can traverse relationships, not just match text — then answers from the expanded context.

Architecture

Question
   │
   ▼
pgvector HNSW (cosine)  ──►  seed nodes (semantic entry points)
   │
   ▼
recursive CTE over edges ──►  k-hop expansion (bounded by max_hops)
   │
   ▼
blended rank (seed_similarity × 0.5^hop) ──►  context ──►  model-serving answer

Everything lives in Lakebase: nodes + edges relational tables (typed, JSONB props) as the graph store, a recursive CTE (WITH RECURSIVE) for traversal, and pgvector for the semantic seed. A generic Databricks model-serving endpoint supplies embeddings and answer synthesis — swap in whichever provider your workspace has registered.

The graph schema

Three tables under a graph schema hold everything (sql/schema.sql):

Table	Key columns	Notes
`graph.nodes`	`node_id TEXT PK`, `node_type`, `name`, `props JSONB`, `path LTREE`	Business-key ids (`product:122`, `supplier:S2`); `props` holds type-specific attributes; optional `ltree` path for hierarchy rollups (`us.northeast.boston`). Trigram + GiST indexes for fuzzy name and path lookups.
`graph.edges`	PK `(src_id, dst_id, rel)`, `props JSONB`	The composite key dedupes relationships. Relations: `SUPPLIED_BY`, `LOCATED_IN`, `BELONGS_TO`, `SUBSTITUTE_FOR`, `SURGES_IN`. Indexed both directions (`edges_src_rel_idx`, `edges_dst_rel_idx`).
`graph.node_embeddings`	`node_id TEXT PK`, `embedding VECTOR(1024)`	1024-dim to match `databricks-gte-large-en`. HNSW index: `vector_cosine_ops` with `m = 16, ef_construction = 64`.

Cross-stack interop (`gold_triplets`)

A knowledge graph is portable between engines when they agree on one row shape. sql/gold_triplets_mapping.sql projects graph.nodes / graph.edges into a portable 8-column triplet contract — subject_id, subject_type, predicate, object_id, object_type, confidence, source_method, source_agent — so the same graph can be emitted to or consumed from another graph stack without migrating tables. It is a column contract, not a dependency: the underlying tables are unchanged and nothing new is imported. Provenance in an edge’s props wins where present; otherwise source_method is derived from the relation and confidence defaults to 1.0, with the numeric cast guarded so a malformed props value becomes NULL rather than 1.0, so a consumer filtering on high confidence cannot silently ingest malformed rows as gold. Undirected relations (SUBSTITUTE_FOR) are stored once with sorted endpoints and made bidirectional only at query time, so the view emits both directions for them — a directed consumer would otherwise miss the reverse.

How retrieval works

The retrieval SQL is three stages in one statement:

Semantic seed — an HNSW ANN scan finds the entry nodes closest to the question embedding, ordered by embedding <=> :query_embedding (pgvector’s <=> is cosine distance, so similarity = 1 - distance). Cosine is used because embeddings are direction-, not magnitude-, meaningful.
Graph expansion — a WITH RECURSIVE walk follows edges out from each seed up to :max_hops (typically 2 — “two degrees of separation”). Edges are first materialized in both directions (UNION ALL of forward and reverse) so an undirected relation like SUBSTITUTE_FOR traverses either way. A per-path visited array guards against cycles.
Blended ranking — each reachable node is scored
```
graph_score = MAX(seed_similarity × 0.5 ^ hop)
```
The 0.5 ^ hop decay halves a node’s contribution per hop, so a direct neighbor of a strong seed outranks a distant one; MAX means a node reachable from several seeds keeps its strongest path.
Authority — a node marked as curated in props ({"source_method": "uc_certified"}) has its score multiplied by :authority_boost, so a certified definition wins the tie against an equally-relevant inferred one. The query returns source_class so an answer can cite what resolved it, and the top 25 by the authority-weighted score become the LLM context.

A multiplier rather than a sort key, deliberately: ordering by (authority, score) would let a barely-relevant certified node outrank a highly relevant inferred one, which is worse than no authority signal at all. Three properties worth knowing:
- Only positive scores are boosted. seed_similarity spans [-1, 1] and the default seed_floor of -1.0 admits negatives, so multiplying would make a negative score more negative and demote the very node the boost exists to promote.
- source_class is passed through verbatim, not normalised to two values, so any producer’s label survives. In practice it is uc_certified or inferred: the certified-core builder below produces the former, though it writes nothing — it returns rows and you load them. Only uc_certified is boosted; everything else carries weight 1.0. Branch on equality with uc_certified, never on inequality with inferred.
- A boost below 1.0 is clamped up to 1.0, for the same reason: a fractional or negative multiplier would demote the node it is meant to promote.
Populating those certified nodes is a separate concern from ranking them: anything that can write the props key participates, which keeps the ranking independent of any one producer. The example ships one such producer — see the certified core — and marking a node by hand works just as well.

This is the GraphRAG win: a supplier or substitute that shares no keywords with the question still surfaces because it’s one hop from a semantically-matched node — something flat vector RAG never retrieves.

The certified core (Unity Catalog metric views)

The authority stage needs some node to be certified. graph_certified.py produces those from Unity Catalog metric views: each view becomes a node, each field becomes a node, and a directional HAS_MEASURE / HAS_DIMENSION edge joins them — every row tagged source_method='uc_certified'.

Fields are separate nodes, not properties on the view, because retrieval seeds on node embeddings and expands along edges. A question about “revenue” has to match the measure and then reach its view; a measure inside a JSON blob is unreachable, and authority ranking only ever re-orders nodes retrieval already found.

Two things worth knowing if you adapt this:

The read path needs two queries. information_schema.columns gives names, types and comments but cannot distinguish a measure from a dimension — both data_type and full_data_type report a measure as plain bigint. Only DESCRIBE EXTENDED marks it, with a measure type suffix. That, plus display_name, synonyms, Owner and Created Time, is why DESCRIBE is not optional.
information_schema.semantic_* is not populated for native metric views. semantic_views, semantic_dimensions and friends look like the structured source you want. In one workspace survey every catalog exposing them was a Snowflake federation connection, and the native catalog holding a real Databricks metric view exposed none — so they appear to be populated for federated Snowflake semantic views rather than native ones. Whether that is permanent or just not implemented yet is unknown; either way they find nothing for a Databricks metric view today.

Only metric views are required, so the prerequisite stays “you have a metric view”. Domains, governed Pages and certification signals are the other three parts of Unity Catalog business semantics and are left as future enrichment.

Building the graph from documents

The example above starts from structured dims. When the source is documents, graph_upstream.py covers the indexing phase that comes first: parse with ai_parse_document, chunk into passages with stable ids, extract typed triples with ai_query under a constrained schema, resolve entities, and emit the gold_triplets contract. The two steps needing a workspace live in sql/upstream_ai_functions.sql.

Two lessons from running it against a live workspace are worth carrying into your own build:

Entity resolution is the load-bearing step, and Jaccard alone is not enough. The same organization appears as “Acme”, “Acme Foods” and “ACME Foods Inc.”; each spelling would otherwise become its own node and fragment the graph precisely where multi-hop retrieval needs it joined. Jaccard trigram similarity scores “Acme Foods” against “Acme Foods Incorporated” at just 0.42 — below any sane threshold — because it penalizes length difference, and legal suffixes are the most common way one entity is written two ways. resolve_entities() therefore also applies a containment measure, the shape of pg_trgm’s word_similarity(), which scores that pair 0.92. Every merge is written back as a SAME_AS edge, so the resolution is auditable in the graph.

Pin the relation vocabulary, not just the entity types. Left free, the model invents a predicate per sentence: a live run over two short documents produced SUPPLIES_TO_RETAILERS_ACROSS, OPERATES_DISTRIBUTION_CENTER_IN, IS_LOCATED_IN and BELONGS_TO_CATEGORY — the last two near-misses for this schema’s LOCATED_IN and BELONGS_TO. Each spelling becomes a distinct edge label, and typed-path traversal degrades toward an untyped walk. Constrain the predicate list in the extraction prompt and pass allowed_predicates=DEFAULT_PREDICATES, on_unknown="drop" as a backstop.

Building the graph safely

assemble_graph() in graph_build.py is a pure function that turns rows + LLM enrichment into (nodes, edges). Its add_edge() only adds an edge if both endpoints exist — so hallucinated substitute ids or orphaned supply rows are dropped and logged rather than creating dangling references. Undirected SUBSTITUTE_FOR edges are stored once (endpoints sorted) to avoid duplicates.

The retrieval and build logic is validated entirely offline by smoketest/graphrag_logic_smoketest.py — 134 assertions (on DuckDB, no Lakebase or model endpoint needed) covering the semantic seed, graph expansion surfacing context flat RAG misses, the dangling-edge guard, 0.5^hop score decay, the max_hops depth bound, and the seed_floor distractor guard.

Deploy with Asset Bundles

Prerequisites: a Databricks workspace with Lakebase (Autoscaling) enabled, a model-serving embeddings endpoint and a chat endpoint registered, plus the databricks CLI and uv.

cd agents/graphrag
databricks bundle deploy -t dev \
  --var lakebase_database="projects/<project>/branches/<branch>/databases/<id>"
databricks bundle run graphrag_build -t dev

The bundle deploys notebooks/graphrag_build_and_query.py as a job that assembles a small example supply-chain graph, embeds its nodes, writes to Lakebase, and queries it. The two Lakebase I/O cells are scaffolding you complete (the Postgres connection is workspace-specific); the retrieval logic itself is fully validated offline by the smoke test:

cd agents/graphrag
uv run --python 3.11 --with duckdb --with numpy smoketest/graphrag_logic_smoketest.py

Configuration and tuning

Variable / setting	Purpose
`lakebase_database`	Full Lakebase database resource path (required).
`max_hops`	Traversal depth from the seeds. `2` is the sweet spot; higher pulls in more distant (and lower-scored) context at the cost of a wider recursive walk.
`seed_floor`	Minimum cosine similarity a semantic seed must clear before it enters the graph walk. Cosine ranges `[-1, 1]`, so the default `-1.0` keeps every seed (identical to no floor). Raise it on distractor-heavy corpora, where weak seeds bridge into unrelated subgraphs and dilute precision. Calibrate to your embedding model’s observed range rather than a fixed constant: measured live with `databricks-gte-large-en`, this graph’s similarities span `0.38`-`0.71`, so `0.3` is a no-op and `0.555` is where the distractors drop. Required bind — Postgres has no server-side default for a named parameter, so SQL callers must pass it (`-1.0` for the old behavior); the Python twin defaults it. A floor above every similarity empties the seed CTE and returns zero rows, so the caller should retry at `-1.0` or decline to answer.
`authority_boost`	Multiplier applied to a node whose `props` carry `source_method='uc_certified'`, so a curated definition outranks an equally-relevant inferred one. The default `1.0` leaves every weight at `1.0`, so scores are unchanged and distinctly-scored rows keep their order — though tied rows now sort deterministically by `node_id` rather than arbitrarily, which measured live is common (two tied pairs in one top-ten). Calibrate once a certified core exists — around `1.2`-`2.0` is a reasonable start, but the useful value depends on the score gaps your embedding model produces, not a portable constant. Only scores above zero are boosted, since multiplying a negative score would demote the node. Required bind, exactly like `seed_floor`: adding it changes the SQL’s parameter arity, so an existing caller that omits it raises a `bind message supplies N parameters, but prepared statement requires N+1` error. Bind `1.0` to keep the previous ranking. The Python twin defaults it.
`VECTOR(1024)`	Embedding dimension — must match your embeddings endpoint (1024 for `databricks-gte-large-en`).
HNSW `m` / `ef_construction`	Index build quality vs. speed (`16` / `64` here). Raise for higher recall on larger graphs.
Embeddings endpoint	Model-serving endpoint used to embed nodes and questions.
Chat endpoint	Model-serving endpoint used to synthesize the final answer.

Edit this page →