Support hierarchical partition keys in Cosmos DB grain storage

[cmeyertons]

Summary

The current IPartitionKeyProvider interface returns a string, which limits Cosmos DB grain storage to single-level partition keys. With hierarchical partition keys now GA in Cosmos DB, it would be valuable to support multi-level partitioning for Orleans grain state.

Background

I'm returning to Orleans after a hiatus and starting a new project that would benefit from hierarchical partitioning. The current implementation wraps the string result in new PartitionKey(partitionKey):

// Current implementation in CosmosGrainStorage.cs
private ValueTask<string> BuildPartitionKey(string grainType, GrainId grainId) =>
    _partitionKeyProvider.GetPartitionKey(grainType, grainId);

// Later used as:
var pk = new PartitionKey(partitionKey);

Proposal

Add a virtual default interface method (DIM) to IPartitionKeyProvider that returns a PartitionKey directly, allowing implementations to construct hierarchical keys:

public interface IPartitionKeyProvider
{
    // Existing method (unchanged for backwards compatibility)
    ValueTask<string> GetPartitionKey(string grainType, GrainId grainId);

    // New DIM for hierarchical partition key support
    ValueTask<PartitionKey> GetPartitionKeyValue(string grainType, GrainId grainId)
        => new(new PartitionKey(GetPartitionKey(grainType, grainId).Result));
}

Then update CosmosGrainStorage to call GetPartitionKeyValue() instead of wrapping the string manually.

Example Usage

public class TenantAwarePartitionKeyProvider : IPartitionKeyProvider
{
    public ValueTask<string> GetPartitionKey(string grainType, GrainId grainId)
        => new(grainType); // Legacy fallback

    public ValueTask<PartitionKey> GetPartitionKeyValue(string grainType, GrainId grainId)
    {
        var tenantId = ExtractTenantFromGrainId(grainId);
        // Hierarchical: TenantId -> GrainType -> Region
        return new(new PartitionKeyBuilder()
            .Add(tenantId)
            .Add(grainType)
            .Build());
    }
}

Benefits

1. Multi-tenant scenarios: Partition first by tenant, then by grain type
2. Better data distribution: Avoid hot partitions in high-volume scenarios
3. Backwards compatible: Existing implementations continue to work unchanged via the DIM
4. Cosmos DB best practices: Aligns with Microsoft's guidance on partition key design

Considerations

• CosmosStorageOptions.PartitionKeyPath would need to support hierarchical paths (e.g., /TenantId/GrainType) (we should be able to define GrainStateEntity as an object w/ nested properties rather than a string)
• Container creation logic may need updates for hierarchical key configuration
• Documentation updates for multi-level partitioning patterns

Environment

• Orleans version: 10.x
• Target: Orleans.Persistence.Cosmos

Happy to submit a PR if this direction makes sense.

[ScottArbeit]

I need this as well, and for now I'll probably vibe-code an implementation to move forward, but would love to see this land officially in Orleans. ❤️ I spent a couple of hours last night having Codex explore the idea of switching my actor storage backend from Cosmos to other supported (and not-yet-supported) data stores, but CosmosDB with hierarchical partition keys is the best way forward for my case.

[ReubenBond]

If we can make sure it's backward compatible out-of-the box, then I would gladly accept a PR here, and it would be much appreciated!