DO NOT MERGE — POC: entity-level routing with composite tablet keys and TabletIndex

[mlwelles]

⚠️ POC ONLY — NOT FOR PRODUCTION MERGE
Stacked on #9574 (sharding-poc). Exists to validate entity-level routing before any production work.

What sharding-poc (#9574) provides

The base branch implements schema-level label routing: a predicate declared with @label(secret) is pinned to the Alpha group running with --label=secret. One predicate → one group. It also adds program-based auth via dgraph.programs facets.

What this branch adds

Entity-level routing — the same predicate can now live on multiple groups simultaneously, one per label. Instead of the schema pinning an entire predicate to a group, individual entities are pinned via a dgraph.label value on their UID:

# No @label directives in schema — routing is per-entity, not per-predicate:
Document.name: string @index(term) .
Document.text: string @index(term) .

_:doc1 <dgraph.label> "secret" .
_:doc1 <Document.name> "Secret.pdf" .       # → routed to "secret" group
_:doc1 <Document.text> "Classified" .       # → routed to "secret" group

_:doc2 <Document.name> "Public.pdf" .       # → routed to unlabeled group

Zero's tablet state now contains composite keys for each (predicate, label) pair:

0-Document.name             → group 1 (unlabeled)
0-Document.name@secret      → group 2
0-Document.name@top_secret  → group 3

A query like has(Document.name) fans out to all three groups and merges the results.

Implementation breakdown

Composite tablet keys (protos/pb/labeled.go)

• TabletKey(pred, label) → "pred@label" (or bare "pred" when unlabeled)
• ParseTabletKey(key) → splits back into (pred, label)
• These are the encoding/decoding for the flat proto map<string, Tablet>

dgraph.label as a reserved predicate (schema/schema.go, x/keys.go)

• Registered in initialSchemaInternal with string type + exact index
• Added to starAllPredicateMap so * queries include it

Entity label cache (worker/entity_label_cache.go — new file)

• Concurrency-safe UID → label map on Alpha (1M entries, ~16MB)
• Avoids querying group 1 for every mutation to resolve an entity's label
• Cleared on DropAll (worker/draft.go)

Two-phase mutation routing (worker/mutation.go)

• Phase 1: Scan mutation edges for dgraph.label to build entity → label batch map
• Phase 2: Route each edge using resolveLabel(uid, batchLabels) — checks batch → cache → group-1 read
• Reserved predicates (dgraph.type, dgraph.label, ACL) still use predicate-level routing

Zero state machine (dgraph/cmd/zero/raft.go, zero.go)

• handleTablet stores/looks up tablets using composite keys
• Duplicate detection uses servingTablet(pred, label) — multiple groups CAN serve the same predicate with different labels
• ShouldServe updated: when a group requests a predicate it already serves (under a different label), return the matching label tablet

Query fan-out (worker/task.go, worker/sort.go)

• AllTablets(pred) returns all label variants of a predicate from the tablet cache
• ProcessTaskOverNetwork / SortOverNetwork: if multiple label tablets exist, dispatch to all groups in parallel and merge
• mergeResults / mergeSortResults: concatenate UID matrices, value matrices, counts; sort merged UIDs for downstream binary search

Oracle commit validation (dgraph/cmd/zero/oracle.go)

• ServingTablets(pred) returns all groups serving a predicate
• Commit validation checks if the committing group is any of them (not just the single tablet owner)

Proposal validation (worker/proposal.go)

• checkTablet switched from Tablet(pred, label) to BelongsToReadOnly(pred, 0) which uses GetForGroup — avoids depending on schema.State().GetLabel() which returns a single label and may not match the receiving group's label tablet

TabletIndex data structure (protos/pb/labeled.go)

• Nested map[pred]map[label]*Tablet for O(1) lookups
• Replaces three O(n) flat-map patterns introduced by the composite key design:
  • ServingTablet fallback loop scanning all tablets via ParseTabletKey
  • AllTablets walking the entire flat map filtering by predicate prefix
  • applyState two-pass group ordering to guarantee bare-key → own-group mapping
• Read cache rebuilt from authoritative flat proto maps after every write

New files

File	Purpose
protos/pb/tablet_index_test.go	TabletIndex unit tests (7 tests)
protos/pb/labeled_test.go	TabletKey/ParseTabletKey unit tests
worker/entity_label_cache.go	UID → label cache
worker/entity_label_cache_test.go	Entity label cache tests

Test results

• 7/7 TabletIndex unit tests pass
• 27/27 systest/label integration tests pass, including:
  • TestEntityLevelRouting — 3 entities, 3 groups, fan-out query returns all
  • TestLabelAuthNoLevel — regression test for fallback-free ServingTablet
  • TestLabeledPredicateCannotBeMoved — labeled predicate stays pinned

Diff stats

20 files changed, 2569 insertions(+), 112 deletions(-)

[xqqp]

I appreciate the progress made on sharding predicates among alphas. However, I feel that both pinning predicates and entity-level routing are merely workarounds for issues that Dgraph should resolve independently. Ideally, Dgraph would automatically detect when a predicate has grown too large on an alpha and then move a portion of that predicate to another alpha, without requiring intervention from a database administrator.

[matthewmcneely]

@xqqp Note this branch is a proof of concept. You should not take anything happening in this branch as future direction.

Ideally, Dgraph would automatically detect when a predicate has grown too large on an alpha and then move a portion of that predicate to another alpha, without requiring intervention from a database administrator.

What you describe here is, in essence, the current functionality of group rebalancing in Dgraph and has been for years.

[xqqp]

@matthewmcneely

Note this branch is a proof of concept. You should not take anything happening in this branch as future direction.

I'm aware of that. Just thought I give some feedback.

What you describe here is, in essence, the current functionality of group rebalancing in Dgraph and has been for years.

In which release was that implemented? If I'm reading this correctly at least in December 2024 predicates couldn't be split: https://discuss.dgraph.io/t/does-dgraph-have-a-scalability-problems-with-graphs-having-single-heavy-predicate/19647/2

[matthewmcneely]

Right, sharding is still per predicate. When I said "in essense", I was referring to the rebalancing based on disk consumption per-predicate.

[xqqp]

Alright, then we are on the same page.

[mlwelles]

Closing to prevent confusion, as it's not intended as anything but a POC of feasibility.