Skip to content

Type resolution reads the whole entity table once per subject #514

Description

@roycegeo

resolve_batch walks its subjects one at a time, and each iteration goes to the database twice:

// crates/utopia-server/src/type_resolution.rs:218
for (i, s) in subjects.iter().enumerate() {
    let descendants = match s.coarse_id {
        Some(c) => utopia_store::resolution::descendants_of(&state.pool, kb_id, c).await?...,
        None => HashSet::new(),
    };
    ...
    let raw = utopia_store::resolution::nearest_typed_entities(&state.pool, kb_id, s.id, NEIGHBOURS).await?;

BATCH is 60 and MAX_ROUNDS is 10, so a resolution run is up to 600 iterations of this, taken in order.

nearest_typed_entities is not a cheap query. It reads the whole entity table:

-- crates/utopia-store/src/resolution.rs:2744
ORDER BY e.profile_embedding <=> (SELECT v FROM me)
LIMIT $3

entities.profile_embedding has no index either, so each call computes cosine distance against every typed, unmerged entity in the knowledge base. Sixty subjects means sixty full passes over the same table, back to back, per round.

The cheap half

descendants_of is keyed on s.coarse_id and nothing else. The coarse class comes from extraction and is drawn from a small vocabulary — person, organization, product and a handful more — so across sixty subjects the same recursive CTE is issued over and over for the same few arguments. A HashMap<Uuid, HashSet<Uuid>> held for the batch removes almost all of them and cannot change an answer: same input, same query, same result, within one batch where the ontology is not moving.

This is the part to do first. It is small, it is provably behaviour-preserving, and it is measurable on its own.

The expensive half

The sixty nearest_typed_entities calls are independent of each other — nothing in iteration i reads anything iteration i-1 wrote. They are serial only because the loop is.

The pattern is already established twenty lines up. type_resolution.rs:184 runs the two ontology retrieval routes concurrently:

let (profile_hits, name_hits) = tokio::join!(...);

So the awareness is there; the per-subject loop just did not get it. Bounded concurrency over the retrieval half of the loop — gather the neighbours for all sixty subjects, then run the per-subject reasoning against the gathered results — turns sixty serial scans into a handful of concurrent ones.

The bound matters. DEFAULT_MAX_CONNECTIONS is 32 and the comment at db.rs:5 is explicit that the pool is sized for "how many short queries run at once", with a recorded history of an undersized pool showing up as slow requests and timeouts rather than as anything saying the pool is too small. Sixty concurrent full-table scans would be exactly that mistake at a new site. Pick a ceiling well under the pool, and say what it is under and why.

There is an ordering question to settle honestly: the loop currently appends to out in subject order, and adjudication downstream may depend on that order. Gathering concurrently and then reasoning in the original order preserves it; reasoning concurrently does not. Do the first.

The real fix underneath

Both halves are palliative. nearest_typed_entities scans because entities.profile_embedding has no ANN index, for the same dimensionless-column reason as chunks.embedding. If that index lands, this loop stops being quadratic and the concurrency question becomes a much smaller one. These two issues should be sequenced together, and the entity-side index may belong in the same migration.

Not this

Caching nearest_typed_entities across subjects. Every subject has a different query vector; there is nothing to share.

Raising the pool to cover sixty concurrent scans. That treats the symptom by moving the load onto Postgres, and db.rs already argues against sizing the pool to the worker count rather than to the concurrent short queries.

Unit tests

The SME rule is that type resolution's answer is a function of the ledger, not of the order or timing in which candidates were fetched. A concurrency change here is only safe if the same knowledge base yields the same types.

Database-gated, in crates/utopia-store/tests/ and beside type_resolution.rs:

  • a_batch_resolves_to_the_same_types_however_it_is_scheduled — the anchor test. Seed a fixed knowledge base, run resolution, record the assignments; run it again with the concurrency ceiling set to 1 and then to its production value, and assert the results are identical. If they are not, the change is wrong regardless of how much faster it is.
  • the_order_of_the_batch_survives_concurrency — subjects come out in the order they went in, because adjudication downstream reads that order.
  • a_memoised_descendant_set_is_the_set_the_query_returns — for every distinct coarse_id in a batch, the cached set equals what descendants_of returns when called directly. Cheap, and it is the whole correctness argument for the memoisation.
  • a_subject_with_no_coarse_class_considers_the_whole_table — the None arm, which memoisation must not accidentally fold in with a real class. type_resolution.rs:231 records that an entity with no class has no descendants axis and the whole class table is a candidate (0009); a HashMap keyed on Option<Uuid> would quietly break this.
  • a_signature_is_guidance_and_not_a_gate — the rule at type_resolution.rs:222, measured at four correct answers lost out of seventeen when it was a hard gate. Non-descendant classes remain reachable candidates after the refactor.
  • an_entity_is_never_its_own_neighbour and an_untyped_entity_is_never_offered_as_evidence — the two gates written into the SQL at resolution.rs:2738, restated as tests so a rewrite of the query cannot drop them.
  • a_neighbour_from_the_same_document_is_marked_as_suchsame_document is documented as a known weakness the caller must be able to see; it has to survive.
  • the_database_pool_is_never_exhausted_by_one_batch — run a full batch against a deliberately small pool and assert it completes rather than hitting the 10-second acquire timeout.

Acceptance testing

  1. On a knowledge base with a substantial entity count, time a full resolution run before and after. Report both, and report them separately for the memoisation and the concurrency change so the two are attributable.
  2. Count the queries. log_min_duration_statement = 0 on the dev database, run one batch, and compare the number of descendants_of and nearest_typed_entities statements before and after. The first should collapse to roughly the number of distinct coarse classes; the second should stay at sixty but overlap.
  3. Run resolution twice on the same untouched knowledge base and diff the resulting type assignments. They must be identical — this is the acceptance form of the anchor test and it should be done against real data, not a fixture.
  4. Watch the Review page during a run and confirm the queue drains and nothing sits locked. Concurrency changes are where a lock gets held across an await.
  5. Watch Postgres connection count during a run; it must stay clear of the pool ceiling.
  6. cargo clippy --workspace --all-targets -- -D warnings and cargo test --workspace with UTOPIA_TEST_REQUIRE_DB=1.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions