Fixed batching

Author: GuillaumeVadeBe

Extensions/AzureTableAPI/Cosmos.DataTransfer.AzureTableAPIExtension/AzureTableAPIDataSinkExtension.cs

@@ -23,18 +23,41 @@ public async Task WriteAsync(IAsyncEnumerable<IDataItem> dataItems, IConfigurati
 
             await tableClient.CreateIfNotExistsAsync(cancellationToken);
 
+            await foreach (var batch in GetBatches(dataItems, settings).WithCancellation(cancellationToken))
+            {
+                await InnerWriteAsync(batch, tableClient, logger, cancellationToken);
+            }
+        }
+
+        private static async IAsyncEnumerable<List<TableEntity>> GetBatches(IAsyncEnumerable<IDataItem> dataItems, AzureTableAPIDataSinkSettings settings)
+        {
             var entities = new List<TableEntity>();
+            var first = true;
+            var partitionKey = string.Empty;
 
-            await foreach (var item in dataItems.WithCancellation(cancellationToken))
+            await foreach (var item in dataItems)
             {
-                var entity = item.ToTableEntity(settings.PartitionKeyFieldName, settings.RowKeyFieldName);
-                entities.Add(entity);
+                var tableEntity = item.ToTableEntity(settings.PartitionKeyFieldName, settings.RowKeyFieldName);
+
+                if (first)
+                {
+                    partitionKey = tableEntity.PartitionKey;
+                    first = false;
+                }
 
-                if (entities.Count == 100)
+                if (!tableEntity.PartitionKey.Equals(partitionKey) || entities.Count == 100)
                 {
-                    await InnerWriteAsync(entities, tableClient, logger, cancellationToken);
-                    entities.Clear();
+                    yield return entities;
+                    entities = new List<TableEntity>();
+                    partitionKey = tableEntity.PartitionKey;
                 }
+
+                entities.Add(tableEntity);
+            }
+
+            if (entities.Count > 0)
+            {
+                yield return entities;
             }
         }