Saturday, February 14, 2015

Batch and Partition Extensions for .NET

Did you ever need to take a collection of items and break it into a set of batches, or divide those items into partitions? Well here are some simple extension methods to help you do that:

Extension Methods

public static class EnumerableExtensions
{
    public static IList<IList<T>> Partition<T>(
        this IEnumerable<T> items, 
        int partitionCount)
    {
        if (partitionCount == 0)
        {
            throw new ArgumentException(
                "Partition Count must be greater than zero", 
                "partitionCount");
        }
 
        return items
            .Select(
                (v, i) => new
                {
                    Group = i%partitionCount,
                    Value = v
                })
            .GroupBy(k => k.Group, v => v.Value)
            .Select(g => (IList<T>) g.ToList())
            .ToList();
    }
 
    public static IList<IList<T>> Batch<T>(
        this IEnumerable<T> items, 
        int batchSize)
    {
        if (batchSize == 0)
        {
            throw new ArgumentException(
                "Batch Size must be greater than zero", 
                "batchSize");
        }
 
        var batches = new List<IList<T>>();
        var batch = new List<T>();
 
        foreach (var item in items)
        {
            if (batch.Count == batchSize)
            {
                batches.Add(batch);
                batch = new List<T>();
            }
 
            batch.Add(item);
        }
 
        if (batch.Count > 0)
        {
            batches.Add(batch);
        }
 
        return batches;
    }
}

Unit Tests

public class EnumerableExtensionsTests
{
    [Theory]
    [InlineData(0, 5)]
    [InlineData(4, 5)]
    [InlineData(5, 5)]
    [InlineData(6, 5)]
    [InlineData(10, 5)]
    [InlineData(20, 5)]
    public void Partition(int count, int partitionCount)
    {
        var items = Enumerable.Range(1, count);
        var partitions = items.Partition(partitionCount);
 
        if (count < partitionCount)
        {
            Assert.Equal(count, partitions.Count);
        }
        else
        {
            Assert.Equal(partitionCount, partitions.Count);
        }
 
        if (count > 0)
        {
            var firstCount = Math.Ceiling((double)count / partitionCount);
            Assert.Equal(firstCount, partitions.First().Count());
        }
    }
 
    [Theory]
    [InlineData(0, 5)]
    [InlineData(4, 5)]
    [InlineData(5, 5)]
    [InlineData(6, 5)]
    [InlineData(10, 5)]
    [InlineData(20, 5)]
    public void Batch(int count, int batchSize)
    {
        var items = Enumerable.Range(1, count);
        var batches = items.Batch(batchSize);
            
        var batchCount = Math.Ceiling((double)count / batchSize);
        Assert.Equal(batchCount, batches.Count);
    }
}

Enjoy,
Tom

No comments:

Post a Comment

Real Time Web Analytics