add new batch overload (bucket factory)

MoreLinq.Test/BatchTest.cs

@@ -26,14 +26,14 @@ public class BatchTest
         [Test]
         public void BatchZeroSize()
         {
-            AssertThrowsArgument.OutOfRangeException("size",() =>
+            AssertThrowsArgument.OutOfRangeException("size", () =>
                 new object[0].Batch(0));
         }
 
         [Test]
         public void BatchNegativeSize()
         {
-            AssertThrowsArgument.OutOfRangeException("size",() =>
+            AssertThrowsArgument.OutOfRangeException("size", () =>
                 new object[0].Batch(-1));
         }
 
@@ -61,6 +61,41 @@ public void BatchUnevenlyDivisibleSequence()
             reader.ReadEnd();
         }
 
+        [Test]
+        public void BatchFactoryUnevenlyDivisibleSequence()
+        {
+            int lastSize = 0;
+            int[] temp = null;
+
+            var result = new[] { 1, 2, 3, 4, 5, 6, 7, 8, 9 }.Batch(4,
+            x => x.Take(lastSize),
+            size =>
+            {
+                lastSize = size;
+
+                return temp ??= new int[size];
+            });
+
+            using var reader = result.Read();
+
+            var first = reader.Read();
+            first.AssertSequenceEqual(1, 2, 3, 4);
+            first.AssertSequenceEqual(temp);
+
+            var second = reader.Read();
+            second.AssertSequenceEqual(5, 6, 7, 8);
+            second.AssertSequenceEqual(temp);
+
+            var third = reader.Read();
+            third.AssertSequenceEqual(9);
+
+            first.AssertSequenceEqual(9, 6, 7, 8);
+            second.AssertSequenceEqual(9, 6, 7, 8);
+            temp.AssertSequenceEqual(9, 6, 7, 8);
+
+            reader.ReadEnd();
+        }
+
         [Test]
         public void BatchSequenceTransformingResult()
         {

MoreLinq/Batch.cs

@@ -80,10 +80,43 @@ public static IEnumerable<IEnumerable<TSource>> Batch<TSource>(this IEnumerable<
 
         public static IEnumerable<TResult> Batch<TSource, TResult>(this IEnumerable<TSource> source, int size,
             Func<IEnumerable<TSource>, TResult> resultSelector)
+        {
+            return Batch(source, size, resultSelector, size => new TSource[size]);
+        }
+
+        /// <summary>
+        /// Batches the source sequence into sized buckets and applies a projection to each bucket.
+        /// </summary>
+        /// <typeparam name="TSource">Type of elements in <paramref name="source"/> sequence.</typeparam>
+        /// <typeparam name="TResult">Type of result returned by <paramref name="resultSelector"/>.</typeparam>
+        /// <param name="source">The source sequence.</param>
+        /// <param name="size">Size of buckets.</param>
+        /// <param name="resultSelector">The projection to apply to each bucket.</param>
+        /// <param name="bucketFactory">A function that receive the requested size for the next bucket and return the array where elements will be placed.</param>
+        /// <returns>A sequence of projections on equally sized buckets containing elements of the source collection.</returns>
+        /// <para>
+        /// This operator uses deferred execution and streams its results
+        /// (buckets are streamed but their content buffered).</para>
+        /// <para>
+        /// <para>
+        /// When more than one bucket is streamed, all buckets except the last
+        /// is guaranteed to have <paramref name="size"/> elements. The last
+        /// bucket may be smaller depending on the remaining elements in the
+        /// <paramref name="source"/> sequence.</para>
+        /// Each bucket is pre-allocated to <paramref name="size"/> elements.
+        /// If <paramref name="size"/> is set to a very large value, e.g.
+        /// <see cref="int.MaxValue"/> to effectively disable batching by just
+        /// hoping for a single bucket, then it can lead to memory exhaustion
+        /// (<see cref="OutOfMemoryException"/>).
+        /// </para>
+
+        public static IEnumerable<TResult> Batch<TSource, TResult>(this IEnumerable<TSource> source, int size,
+            Func<IEnumerable<TSource>, TResult> resultSelector, Func<int, TSource[]> bucketFactory)
         {
             if (source == null) throw new ArgumentNullException(nameof(source));
             if (size <= 0) throw new ArgumentOutOfRangeException(nameof(size));
             if (resultSelector == null) throw new ArgumentNullException(nameof(resultSelector));
+            if (bucketFactory == null) throw new ArgumentNullException(nameof(bucketFactory));
 
             switch (source)
             {
@@ -95,7 +128,7 @@ public static IEnumerable<TResult> Batch<TSource, TResult>(this IEnumerable<TSou
                 {
                     return _(); IEnumerable<TResult> _()
                     {
-                        var bucket = new TSource[collection.Count];
+                        var bucket = bucketFactory(collection.Count);
                         collection.CopyTo(bucket, 0);
                         yield return resultSelector(bucket);
                     }
@@ -108,7 +141,7 @@ public static IEnumerable<TResult> Batch<TSource, TResult>(this IEnumerable<TSou
                 {
                     return _(); IEnumerable<TResult> _()
                     {
-                        var bucket = new TSource[list.Count];
+                        var bucket = bucketFactory(list.Count);
                         for (var i = 0; i < list.Count; i++)
                             bucket[i] = list[i];
                         yield return resultSelector(bucket);
@@ -130,7 +163,7 @@ IEnumerable<TResult> Batch(int size)
 
                     foreach (var item in source)
                     {
-                        bucket ??= new TSource[size];
+                        bucket ??= bucketFactory(size);
                         bucket[count++] = item;
 
                         // The bucket is fully buffered before it's yielded
@@ -146,8 +179,9 @@ IEnumerable<TResult> Batch(int size)
                     // Return the last bucket with all remaining elements
                     if (bucket != null && count > 0)
                     {
-                        Array.Resize(ref bucket, count);
-                        yield return resultSelector(bucket);
+                        var newArray = bucketFactory(count);
+                        Array.Copy(bucket, 0, newArray, 0, count);
+                        yield return resultSelector(newArray);
                     }
                 }
             }

MoreLinq/Extensions.g.cs

@@ -710,6 +710,36 @@ public static IEnumerable<TResult> Batch<TSource, TResult>(this IEnumerable<TSou
             Func<IEnumerable<TSource>, TResult> resultSelector)
             => MoreEnumerable.Batch(source, size, resultSelector);
 
+        /// <summary>
+        /// Batches the source sequence into sized buckets and applies a projection to each bucket.
+        /// </summary>
+        /// <typeparam name="TSource">Type of elements in <paramref name="source"/> sequence.</typeparam>
+        /// <typeparam name="TResult">Type of result returned by <paramref name="resultSelector"/>.</typeparam>
+        /// <param name="source">The source sequence.</param>
+        /// <param name="size">Size of buckets.</param>
+        /// <param name="resultSelector">The projection to apply to each bucket.</param>
+        /// <param name="bucketFactory">A function that receive the requested size for the next bucket and return the array where elements will be placed.</param>
+        /// <returns>A sequence of projections on equally sized buckets containing elements of the source collection.</returns>
+        /// <para>
+        /// This operator uses deferred execution and streams its results
+        /// (buckets are streamed but their content buffered).</para>
+        /// <para>
+        /// <para>
+        /// When more than one bucket is streamed, all buckets except the last
+        /// is guaranteed to have <paramref name="size"/> elements. The last
+        /// bucket may be smaller depending on the remaining elements in the
+        /// <paramref name="source"/> sequence.</para>
+        /// Each bucket is pre-allocated to <paramref name="size"/> elements.
+        /// If <paramref name="size"/> is set to a very large value, e.g.
+        /// <see cref="int.MaxValue"/> to effectively disable batching by just
+        /// hoping for a single bucket, then it can lead to memory exhaustion
+        /// (<see cref="OutOfMemoryException"/>).
+        /// </para>
+
+        public static IEnumerable<TResult> Batch<TSource, TResult>(this IEnumerable<TSource> source, int size,
+            Func<IEnumerable<TSource>, TResult> resultSelector, Func<int, TSource[]> bucketFactory)
+            => MoreEnumerable.Batch(source, size, resultSelector, bucketFactory);
+
     }
 
     /// <summary><c>Cartesian</c> extension.</summary>