如何對 IAsyncEnumerable 進行惰性分區？

Question

我有一個 IAsyncEnumerable，它返回本質上是一系列Key/IEnumerable<Value>對。 我有代碼使用這個和其他類似的枚舉，假設它將接收一個唯一的鍵集合。 但是我的一個數據源不遵守此約束。 但是，它確實會將重復的鍵組合在一起。 （您不會看到 [ k1 , k2 , k1 ]。）

使用按鍵對數據進行分區並連接值的包裝器來解決這個問題應該相當簡單，只是我在System.Linq.Async不到任何可用的分區運算符。 有GroupBy和ToLookup ，但這兩個都是急切的運算符，會立即消耗整個序列。 由於涉及大量數據，這不適合我的目的。

是否有任何簡單的方法來划分類似於GroupBy的IAsyncEnumerable ，根據鍵選擇器對輸入進行分組，但保持其行為完全惰性並在鍵更改時按需生成新分組？

編輯：我查看了 MoreLINQ 是否有類似的東西，並找到GroupAdjacent ，但代碼顯示，雖然它不會急切地消耗整個輸入序列，但在開始一個新組時它仍然會急切地消耗整個組。 我正在尋找一種將在其分組中返回惰性可枚舉的方法。 這比聽起來更棘手！

Answer 1

這是一個用於異步序列的GroupAdjacent運算符，類似於MoreLinq package 的同義運算符，不同之處在於它不緩沖發出的分組的元素。 分組應該以正確的順序完全枚舉，一次一個分組，否則將拋出InvalidOperationException 。

此實現需要 package System.Linq.Async ，因為它發出實現IAsyncGrouping<out TKey, out TElement>接口的分組。

/// <summary>
/// Groups the adjacent elements of a sequence according to a specified
/// key selector function.
/// </summary>
/// <remarks>
/// The groups don't contain buffered elements.
/// Enumerating the groups in the correct order is mandatory.
/// </remarks>
public static IAsyncEnumerable<IAsyncGrouping<TKey, TSource>>
    GroupAdjacent<TSource, TKey>(
        this IAsyncEnumerable<TSource> source,
        Func<TSource, TKey> keySelector,
        IEqualityComparer<TKey> keyComparer = null)
{
    ArgumentNullException.ThrowIfNull(source);
    ArgumentNullException.ThrowIfNull(keySelector);
    keyComparer ??= EqualityComparer<TKey>.Default;
    return Implementation();

    async IAsyncEnumerable<IAsyncGrouping<TKey, TSource>> Implementation(
        [EnumeratorCancellation] CancellationToken cancellationToken = default)
    {
        Tuple<TSource, TKey, bool> sharedState = null;
        var enumerator = source.GetAsyncEnumerator(cancellationToken);
        try
        {
            if (!await enumerator.MoveNextAsync().ConfigureAwait(false))
                yield break;
            var firstItem = enumerator.Current;
            var firstKey = keySelector(firstItem);
            sharedState = new(firstItem, firstKey, true);

            Tuple<TSource, TKey, bool> previousState = null;
            while (true)
            {
                var state = Volatile.Read(ref sharedState);
                if (ReferenceEquals(state, previousState))
                    throw new InvalidOperationException("Out of order enumeration.");
                var (item, key, exists) = state;
                if (!exists) yield break;
                previousState = state;
                yield return new AsyncGrouping<TKey, TSource>(key, GetAdjacent(state));
            }
        }
        finally { await enumerator.DisposeAsync().ConfigureAwait(false); }

        async IAsyncEnumerable<TSource> GetAdjacent(Tuple<TSource, TKey, bool> state)
        {
            if (!ReferenceEquals(Volatile.Read(ref sharedState), state))
                throw new InvalidOperationException("Out of order enumeration.");
            var (stateItem, stateKey, stateExists) = state;
            Debug.Assert(stateExists);
            yield return stateItem;
            Tuple<TSource, TKey, bool> nextState;
            while (true)
            {
                if (!ReferenceEquals(Volatile.Read(ref sharedState), state))
                    throw new InvalidOperationException("Out of order enumeration.");
                if (!await enumerator.MoveNextAsync().ConfigureAwait(false))
                {
                    nextState = new(default, default, false);
                    break;
                }
                var item = enumerator.Current;
                var key = keySelector(item);
                if (!keyComparer.Equals(key, stateKey))
                {
                    nextState = new(item, key, true);
                    break;
                }
                yield return item;
            }
            if (!ReferenceEquals(Interlocked.CompareExchange(
                ref sharedState, nextState, state), state))
                throw new InvalidOperationException("Out of order enumeration.");
        }
    }
}

private class AsyncGrouping<TKey, TElement> : IAsyncGrouping<TKey, TElement>
{
    private readonly TKey _key;
    private readonly IAsyncEnumerable<TElement> _sequence;

    public AsyncGrouping(TKey key, IAsyncEnumerable<TElement> sequence)
    {
        _key = key;
        _sequence = sequence;
    }

    public TKey Key => _key;

    public IAsyncEnumerator<TElement> GetAsyncEnumerator(
        CancellationToken cancellationToken = default)
    {
        return _sequence.GetAsyncEnumerator(cancellationToken);
    }
}

使用示例：

IAsyncEnumerable<IGrouping<string, double>> source = //...

IAsyncEnumerable<IAsyncGrouping<string, double>> merged = source
    .GroupAdjacent(g => g.Key)
    .Select(gg => new AsyncGrouping<string, double>(
        gg.Key, gg.Select(g => g.ToAsyncEnumerable()).Concat()));

此示例從包含分組的序列開始，目標是將具有相同鍵的任何相鄰分組組合成包含其所有元素的單個異步分組。 應用GroupAdjacent(g => g.Key)運算符后，我們得到這種類型：

IAsyncEnumerable<IAsyncGrouping<string, IGrouping<string, double>>>

所以在這個階段，每個異步分組都包含內部分組，而不是單個元素。 我們需要Concat這個嵌套結構以獲得我們想要的。 Concat運算符存在於System.Interactive.Async package 中，它具有以下簽名：

public static IAsyncEnumerable<TSource> Concat<TSource>(
    this IAsyncEnumerable<IAsyncEnumerable<TSource>> sources);

ToAsyncEnumerable運算符 (System.Linq.Async) 附加到同步內部分組，以滿足此簽名。