Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

add n-way merge for async enumerable #2126

Draft
wants to merge 2 commits into
base: dev
Choose a base branch
from
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
@@ -0,0 +1,63 @@
// Copyright (c) Microsoft. All rights reserved.
// Licensed under the MIT license. See LICENSE file in the project root for full license information.

using System;
using System.Collections.Generic;
using System.Threading;

namespace Microsoft.Azure.SignalR.Common;

public class NWayMergeAsyncEnumerable<T> : IAsyncEnumerable<T>
{
private readonly IComparer<T> _comparer;
private readonly IAsyncEnumerable<T>[] _sources;

public NWayMergeAsyncEnumerable(params IAsyncEnumerable<T>[] sources)
: this(null, sources)
{
}

public NWayMergeAsyncEnumerable(IComparer<T> comparer, params IAsyncEnumerable<T>[] sources)
{
_comparer = comparer ?? Comparer<T>.Default;
_sources = sources ?? throw new ArgumentNullException(nameof(sources));
foreach (var source in _sources)
{
if (source == null)
{
throw new ArgumentException("Item cannot be null.", nameof(sources));
}
}
}

public async IAsyncEnumerator<T> GetAsyncEnumerator(CancellationToken cancellationToken = default)
{
var sources = Array.ConvertAll(_sources, source => source.GetAsyncEnumerator(cancellationToken));
var hasMore = new bool[sources.Length];
for (int i = 0; i < sources.Length; i++)
{
hasMore[i] = await sources[i].MoveNextAsync();
}
while (Array.IndexOf(hasMore, true) != -1)
{
for (int i = 0; i < sources.Length; i++)
{
if (!hasMore[i])
{
continue;
}
var current = i;
for (int j = 0; j < sources.Length; j++)
{
if (j != current && hasMore[j] && _comparer.Compare(sources[j].Current, sources[current].Current) < 0)
{
current = j;
}
}
yield return sources[current].Current;
hasMore[current] = await sources[current].MoveNextAsync();
break;
}
}
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,101 @@
// Copyright (c) Microsoft. All rights reserved.
// Licensed under the MIT license. See LICENSE file in the project root for full license information.

using System;
using System.Collections.Generic;
using System.Threading.Tasks;

using Xunit;

namespace Microsoft.Azure.SignalR.Common.Tests;

public class NWayMergeAsyncEnumerableTest
{
[Fact]
public async Task TestOneWayMergeAsyncEnumerable()
{
List<int> source = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
var sources = new IAsyncEnumerable<int>[]
{
ToAsyncEnumerable(source),
};
var multiWayMergeAsyncEnumerable = new NWayMergeAsyncEnumerable<int>(sources);
var result = await ToListAsync(multiWayMergeAsyncEnumerable);
Assert.Equal([1, 2, 3, 4, 5, 6, 7, 8, 9, 10], result);
}

[Fact]
public async Task TestEmptyOneWayMergeAsyncEnumerable()
{
List<int> source = [];
var sources = new IAsyncEnumerable<int>[]
{
ToAsyncEnumerable(source),
};
var multiWayMergeAsyncEnumerable = new NWayMergeAsyncEnumerable<int>(sources);
var result = await ToListAsync(multiWayMergeAsyncEnumerable);
Assert.Equal([], result);
}

[Fact]
public async Task TestThreeWayMergeAsyncEnumerable()
{
List<int> source1 = [1, 3, 5, 7, 9];
List<int> source2 = [2, 4, 6, 8, 10];
List<int> source3 = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
var sources = new IAsyncEnumerable<int>[]
{
ToAsyncEnumerable(source1),
ToAsyncEnumerable(source2),
ToAsyncEnumerable(source3),
};
var multiWayMergeAsyncEnumerable = new NWayMergeAsyncEnumerable<int>(sources);
var result = await ToListAsync(multiWayMergeAsyncEnumerable);
Assert.Equal([1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7, 8, 8, 9, 9, 10, 10], result);
}

[Fact]
public async Task TestEmptyThreeWayMergeAsyncEnumerable()
{
List<int> source1 = [];
List<int> source2 = [];
List<int> source3 = [];
var sources = new IAsyncEnumerable<int>[]
{
ToAsyncEnumerable(source1),
ToAsyncEnumerable(source2),
ToAsyncEnumerable(source3),
};
var multiWayMergeAsyncEnumerable = new NWayMergeAsyncEnumerable<int>(sources);
var result = await ToListAsync(multiWayMergeAsyncEnumerable);
Assert.Equal([], result);
}

[Fact]
public void TestNullCases()
{
IAsyncEnumerable<int>[] args = null;
Assert.Throws<ArgumentNullException>(() => new NWayMergeAsyncEnumerable<int>(args));
IAsyncEnumerable<int> arg = null;
Assert.Throws<ArgumentException>(() => new NWayMergeAsyncEnumerable<int>(arg));
}

private static async IAsyncEnumerable<T> ToAsyncEnumerable<T>(IEnumerable<T> source)
{
foreach (var item in source)
{
await Task.Delay(1);
yield return item;
}
}

private static async Task<List<T>> ToListAsync<T>(IAsyncEnumerable<T> source)
{
var list = new List<T>();
await foreach (var item in source)
{
list.Add(item);
}
return list;
}
}
Loading