如何在查询Cosmos时传入多个分区键？最佳实践解析

我有这样的业务场景：会收到一个包含多个ID的逗号分隔字符串，这些ID是Cosmos容器的分区键。查询单个ID时，我通过.WithPartitionKey传入分区键实现查询，代码如下：

// 查询单个ID
public async Task<MyTable> GetExistingDocumentsAsync(string myPartitionKey, bool withNoTracking, CancellationToken cancellationToken)
{
    // 假设 myPartitionKey = "123" 
    var query = withNoTracking ? _dbContext.MyTable.AsNoTracking() : _dbContext.MyTable;
    
    return await query.WithPartitionKey(myPartitionKey).FirstOrDefaultAsync(x => x.myPartitionKey.Equals(myPartitionKey), cancellationToken);
}

但收到多个ID时，我尝试了以下代码实现查询，但不确定是否高效正确：

// 尝试查询多个ID
public async Task<List<MyTable>> GetExistingDocumentsAsync(string myPartitionKey, bool withNoTracking, CancellationToken cancellationToken)
{
    // 假设 myPartitionKey = "1,2,3,4"  
    var ids = myPartitionKey.Split(",").Select(s => s.Trim()).ToArray();
    var records = new List<MyTable>();

    using var feedIterator = _dbContext.Database
        .GetCosmosClient()
        .GetContainer(DbGlobals.SchemaName.ToLower(),  nameof(MyTable).ToLower())
        .GetItemLinqQueryable<MyTable>()
        // 疑问：这里如何传入多个分区键？
        .Where(x => ids.Contains(x.myPartitionKey))
        .ToFeedIterator();
    
    while (feedIterator.HasMoreResults)
    {
        foreach (var document in await feedIterator.ReadNextAsync())
        {
            records.Add(document);
        }
    }
   
    return records; 
}

我有两个问题：

1. 这种查询N个ID的方式是否高效且正确？
2. 如果不正确，对应的最佳实践是什么？

1. 当前方式是否高效且正确？ #

这种方式既不高效也不正确。原因如下：

• 没有指定WithPartitionKey，Cosmos会将该查询视为跨分区查询，需要扫描容器内的所有分区才能匹配数据，RU（请求单位）消耗会大幅飙升，查询速度显著变慢，当容器分区多、数据量大时，还可能出现超时、RU限流等稳定性问题。
• 代码中返回值类型定义错误（原代码写的Task<MyTable>实际返回List<MyTable>），属于语法层面的小问题，但也会导致运行报错。

2. 最佳实践 #

针对多个分区键的查询，推荐以下两种高效方案：

方案一：并行发起单个分区键定向查询 #

对每个分区键单独发起带WithPartitionKey的定向查询，通过Task.WhenAll实现并行异步请求，既保证每个查询都精准命中对应分区（RU消耗低、速度快），又能通过并行提升整体处理效率。

示例代码：

public async Task<List<MyTable>> GetExistingDocumentsAsync(string myPartitionKey, bool withNoTracking, CancellationToken cancellationToken)
{
    var ids = myPartitionKey.Split(",").Select(s => s.Trim()).ToArray();
    // 为每个分区键创建单独的定向查询任务
    var queryTasks = ids.Select(id => 
    {
        var query = withNoTracking ? _dbContext.MyTable.AsNoTracking() : _dbContext.MyTable;
        return query.WithPartitionKey(id).ToListAsync(cancellationToken);
    });

    // 并行执行所有查询并合并结果
    var results = await Task.WhenAll(queryTasks);
    return results.SelectMany(r => r).ToList();
}

方案二：批量封装分区查询（适用于大量ID场景） #

如果需要查询的分区键数量超过100个，可以使用Cosmos的批量操作API，将多个单个分区查询打包成批量请求，减少网络往返开销。本质依然是每个请求定向到对应分区，避免跨分区扫描。

注意事项 #

• 避免一次性发起过多并行请求（比如超过100个），否则可能触发Cosmos的限流机制，建议根据实际情况拆分批次处理。
• 始终确保每个查询都显式指定WithPartitionKey，杜绝隐式跨分区查询。

内容的提问来源于stack exchange，提问作者GThree