Skip Dictionary allocation when all LogEvents in single bucket

Author: snakefoot

src/NLog.Extensions.AzureDataTables/DataTablesTarget.cs

@@ -35,8 +35,8 @@ struct TablePartitionKey : IEquatable<TablePartitionKey>
 
             public TablePartitionKey(string tableName, string partitionKey)
             {
-                TableName = tableName;
-                PartitionKey = partitionKey;
+                TableName = tableName ?? string.Empty;
+                PartitionKey = partitionKey ?? string.Empty;
             }
 
             public bool Equals(TablePartitionKey other)

src/NLog.Extensions.AzureEventHub/EventHubTarget.cs

@@ -362,7 +362,7 @@ protected override Task WriteAsyncTask(IList<LogEventInfo> logEvents, Cancellati
                 }
             }
 
-            var partitionBuckets = PartitionKey != null ? SortHelpers.BucketSort(logEvents, _getEventHubPartitionKeyDelegate) : new Dictionary<string, IList<LogEventInfo>>() { { string.Empty, logEvents } };
+            var partitionBuckets = PartitionKey != null ? SortHelpers.BucketSort(logEvents, _getEventHubPartitionKeyDelegate) : new[] { new KeyValuePair<string, IList<LogEventInfo>>(string.Empty, logEvents) };
             IList<Task> multipleTasks = partitionBuckets.Count > 1 ? new List<Task>(partitionBuckets.Count) : null;
             foreach (var partitionBucket in partitionBuckets)
             {

src/NLog.Extensions.AzureServiceBus/ServiceBusTarget.cs

@@ -425,7 +425,7 @@ protected override Task WriteAsyncTask(IList<LogEventInfo> logEvents, Cancellati
                 }
             }
 
-            var partitionBuckets = (SessionId != null || PartitionKey != null) ? SortHelpers.BucketSort(logEvents, _getMessagePartitionKeyDelegate) : new Dictionary<string, IList<LogEventInfo>>() { { string.Empty, logEvents } };
+            var partitionBuckets = (SessionId != null || PartitionKey != null) ? SortHelpers.BucketSort(logEvents, _getMessagePartitionKeyDelegate) : new [] { new KeyValuePair<string, IList<LogEventInfo>>(string.Empty, logEvents) };
             IList<Task> multipleTasks = partitionBuckets.Count > 1 ? new List<Task>(partitionBuckets.Count) : null;
             foreach (var partitionBucket in partitionBuckets)
             {

src/NLog.Extensions.AzureStorage/SortHelpers.cs

@@ -22,24 +22,59 @@ internal static class SortHelpers
         /// <param name="inputs">The inputs.</param>
         /// <param name="keySelector">The key selector.</param>
         /// <returns></returns>
-        internal static Dictionary<TKey, IList<TValue>> BucketSort<TValue, TKey>(IList<TValue> inputs, KeySelector<TValue, TKey> keySelector) where TKey : IEquatable<TKey>
+        internal static ICollection<KeyValuePair<TKey, IList<TValue>>> BucketSort<TValue, TKey>(IList<TValue> inputs, KeySelector<TValue, TKey> keySelector) where TKey : IEquatable<TKey>
         {
-            var retVal = new Dictionary<TKey, IList<TValue>>();
+            if (inputs.Count == 0)
+                return Array.Empty<KeyValuePair<TKey, IList<TValue>>>();
 
-            for (int i = 0; i < inputs.Count; ++i)
+            Dictionary<TKey, IList<TValue>> buckets = null;
+            TKey firstBucketKey = keySelector(inputs[0]);
+            for (int i = 1; i < inputs.Count; ++i)
             {
-                var input = inputs[i];
-                var keyValue = keySelector(input);
-                if (!retVal.TryGetValue(keyValue, out var eventsInBucket))
+                TKey keyValue = keySelector(inputs[i]);
+                if (buckets is null)
                 {
-                    eventsInBucket = new List<TValue>(inputs.Count - i);
-                    retVal.Add(keyValue, eventsInBucket);
+                    if (!firstBucketKey.Equals(keyValue))
+                    {
+                        // Multiple buckets needed, allocate full dictionary
+                        buckets = CreateBucketDictionaryWithValue(inputs, i, firstBucketKey, keyValue);
+                    }
                 }
+                else
+                {
+                    if (!buckets.TryGetValue(keyValue, out var eventsInBucket))
+                    {
+                        eventsInBucket = new List<TValue>();
+                        buckets.Add(keyValue, eventsInBucket);
+                    }
+                    eventsInBucket.Add(inputs[i]);
+                }
+            }
 
-                eventsInBucket.Add(input);
+            if (buckets is null)
+            {
+                // All inputs belong to the same bucket
+                return new[] { new KeyValuePair<TKey, IList<TValue>>(firstBucketKey, inputs) };
+            }
+            else
+            {
+                return buckets;
+            }
+        }
+
+        private static Dictionary<TKey, IList<TValue>> CreateBucketDictionaryWithValue<TValue, TKey>(IList<TValue> inputs, int currentIndex, TKey firstBucketKey, TKey nextBucketKey)
+        {
+            var buckets = new Dictionary<TKey, IList<TValue>>();
+            var firstBucket = new List<TValue>(currentIndex);
+            for (int i = 0; i < currentIndex; i++)
+            {
+                firstBucket.Add(inputs[i]);
             }
+            buckets[firstBucketKey] = firstBucket;
 
-            return retVal;
+            var nextBucket = new List<TValue> { inputs[currentIndex] };
+            buckets[nextBucketKey] = nextBucket;
+            return buckets;
         }
     }
 }