Refactor Bulk Insert to Use Memory<object> for Enhanced Flexibility and Efficiency

ClickHouse.Client.Tests/BulkCopy/BulkCopyTests.cs

@@ -1,6 +1,7 @@
 using System;
 using System.Buffers;
 using System.Collections.Generic;
+using System.Data;
 using System.Diagnostics;
 using System.Linq;
 using System.Runtime.CompilerServices;
@@ -451,7 +452,7 @@ public async Task ShouldThrowExceptionOnInnerException(double fraction)
     [Test]
     public async Task ShouldNotAffectSharedArrayPool()
     {
-        var targetTable = "test." + SanitizeTableName($"array_pool");
+        var targetTable = "test." + SanitizeTableName($"Array_pool");
 
         await connection.ExecuteStatementAsync($"DROP TABLE IF EXISTS {targetTable}");
         await connection.ExecuteStatementAsync($"CREATE TABLE IF NOT EXISTS {targetTable} (int Int32, str String, dt DateTime) ENGINE Null");
@@ -466,9 +467,9 @@ public async Task ShouldNotAffectSharedArrayPool()
         await bulkCopy.InitAsync();
         await bulkCopy.WriteToServerAsync(Enumerable.Repeat(new object[] { 0, "a", DateTime.Now }, 100));
 
-        var rentedArray = ArrayPool<object>.Shared.Rent(poolSize);
-        Assert.DoesNotThrow(() => { rentedArray[0] = 1; });
-        ArrayPool<object>.Shared.Return(rentedArray);
+        var rentedArray = ArrayPool<Memory<object>>.Shared.Rent(poolSize);
+        Assert.DoesNotThrow(() => { rentedArray[0] = new Memory<object>(); });
+        ArrayPool<Memory<object>>.Shared.Return(rentedArray);
     }
 
     [Test]
@@ -496,5 +497,77 @@ public async Task ShouldInsertJson()
             Assert.That(reader.GetValue(0), Is.EqualTo(jsonObject).UsingPropertiesComparer());
         }
     }
+
+    [Test]
+    public async Task ShouldInsertDataReader()
+    {
+        var targetTable = "test." + SanitizeTableName($"bulk_datareader");
+        await connection.ExecuteStatementAsync($"DROP TABLE IF EXISTS {targetTable}");
+        await connection.ExecuteStatementAsync($"CREATE TABLE IF NOT EXISTS {targetTable} (int Int32, str String) ENGINE Memory");
+
+        using var bulkCopy = new ClickHouseBulkCopy(connection)
+        {
+            DestinationTableName = targetTable,
+        };
+
+        var dataTable = new DataTable();
+        dataTable.Columns.Add(new DataColumn("int", typeof(int)));
+        dataTable.Columns.Add(new DataColumn("str", typeof(string)));
+        dataTable.Rows.Add(new object[] { 1, "a" });
+        dataTable.Rows.Add(new object[] { 2, "b" });
+        dataTable.Rows.Add(new object[] { 3, "c" });
+
+
+        await bulkCopy.InitAsync();
+        await bulkCopy.WriteToServerAsync(dataTable.CreateDataReader());
+
+        Assert.That(bulkCopy.RowsWritten, Is.EqualTo(3));
+        using var reader = await connection.ExecuteReaderAsync($"SELECT * from {targetTable}");
+        Assert.That(reader.Read(), Is.True);
+        Assert.That(reader.GetInt32(0), Is.EqualTo(dataTable.Rows[0]["int"]));
+        Assert.That(reader.GetString(1), Is.EqualTo(dataTable.Rows[0]["str"]));
+        Assert.That(reader.Read(), Is.True);
+        Assert.That(reader.GetInt32(0), Is.EqualTo(dataTable.Rows[1]["int"]));
+        Assert.That(reader.GetString(1), Is.EqualTo(dataTable.Rows[1]["str"]));
+        Assert.That(reader.Read(), Is.True);
+        Assert.That(reader.GetInt32(0), Is.EqualTo(dataTable.Rows[2]["int"]));
+        Assert.That(reader.GetString(1), Is.EqualTo(dataTable.Rows[2]["str"]));
+    }
+
+    [Test]
+    public async Task ShouldInsertDataTable()
+    {
+        var targetTable = "test." + SanitizeTableName($"bulk_dataTable");
+        await connection.ExecuteStatementAsync($"DROP TABLE IF EXISTS {targetTable}");
+        await connection.ExecuteStatementAsync($"CREATE TABLE IF NOT EXISTS {targetTable} (int Int32, str String) ENGINE Memory");
+
+        using var bulkCopy = new ClickHouseBulkCopy(connection)
+        {
+            DestinationTableName = targetTable,
+        };
+
+        var dataTable = new DataTable();
+        dataTable.Columns.Add(new DataColumn("int", typeof(int)));
+        dataTable.Columns.Add(new DataColumn("str", typeof(string)));
+        dataTable.Rows.Add(new object[] { 1, "a" });
+        dataTable.Rows.Add(new object[] { 2, "b" });
+        dataTable.Rows.Add(new object[] { 3, "c" });
+
+
+        await bulkCopy.InitAsync();
+        await bulkCopy.WriteToServerAsync(dataTable, CancellationToken.None);
+
+        Assert.That(bulkCopy.RowsWritten, Is.EqualTo(3));
+        using var reader = await connection.ExecuteReaderAsync($"SELECT * from {targetTable}");
+        Assert.That(reader.Read(), Is.True);
+        Assert.That(reader.GetInt32(0), Is.EqualTo(dataTable.Rows[0]["int"]));
+        Assert.That(reader.GetString(1), Is.EqualTo(dataTable.Rows[0]["str"]));
+        Assert.That(reader.Read(), Is.True);
+        Assert.That(reader.GetInt32(0), Is.EqualTo(dataTable.Rows[1]["int"]));
+        Assert.That(reader.GetString(1), Is.EqualTo(dataTable.Rows[1]["str"]));
+        Assert.That(reader.Read(), Is.True);
+        Assert.That(reader.GetInt32(0), Is.EqualTo(dataTable.Rows[2]["int"]));
+        Assert.That(reader.GetString(1), Is.EqualTo(dataTable.Rows[2]["str"]));
+    }
 }
 

ClickHouse.Client/Copy/Batch.cs

@@ -7,7 +7,7 @@ namespace ClickHouse.Client.Copy;
 // Convenience argument collection
 internal struct Batch : IDisposable
 {
-    public object[][] Rows;
+    public Memory<object>[] Rows;
     public int Size;
     public string Query;
     public ClickHouseType[] Types;
@@ -16,7 +16,7 @@ public void Dispose()
     {
         if (Rows != null)
         {
-            ArrayPool<object[]>.Shared.Return(Rows, true);
+            ArrayPool<Memory<object>>.Shared.Return(Rows, true);
             Rows = null;
         }
     }

ClickHouse.Client/Copy/ClickHouseBulkCopy.cs

@@ -117,21 +117,25 @@ public Task WriteToServerAsync(IDataReader reader, CancellationToken token)
         if (reader is null)
             throw new ArgumentNullException(nameof(reader));
 
-        return WriteToServerAsync(reader.AsEnumerable(), token);
+        return WriteToServerAsync(reader.AsEnumerable().Select(r => new Memory<object>(r)), token);
     }
 
     public Task WriteToServerAsync(DataTable table, CancellationToken token)
     {
         if (table is null)
             throw new ArgumentNullException(nameof(table));
 
-        var rows = table.Rows.Cast<DataRow>().Select(r => r.ItemArray);
+        var rows = table.Rows.Cast<DataRow>().Select(r => new Memory<object>(r.ItemArray));
         return WriteToServerAsync(rows, token);
     }
 
-    public Task WriteToServerAsync(IEnumerable<object[]> rows) => WriteToServerAsync(rows, CancellationToken.None);
+    public Task WriteToServerAsync(IEnumerable<object[]> rows) =>
+        WriteToServerAsync(rows.Select(r => new Memory<object>(r)), CancellationToken.None);
 
-    public async Task WriteToServerAsync(IEnumerable<object[]> rows, CancellationToken token)
+    public Task WriteToServerAsync(IEnumerable<object[]> rows, CancellationToken token) =>
+        WriteToServerAsync(rows.Select(r => new Memory<object>(r)), token);
+
+    public async Task WriteToServerAsync(IEnumerable<Memory<object>> rows, CancellationToken token)
     {
         if (rows is null)
             throw new ArgumentNullException(nameof(rows));
@@ -201,7 +205,7 @@ public void Dispose()
 
     private static string GetColumnsExpression(IReadOnlyCollection<string> columns) => columns == null || columns.Count == 0 ? "*" : string.Join(",", columns);
 
-    private IEnumerable<Batch> IntoBatches(IEnumerable<object[]> rows, string query, ClickHouseType[] types)
+    private IEnumerable<Batch> IntoBatches(IEnumerable<Memory<object>> rows, string query, ClickHouseType[] types)
     {
         foreach (var (batch, size) in rows.BatchRented(BatchSize))
         {

ClickHouse.Client/Copy/Serializer/BatchSerializer.cs

@@ -35,24 +35,19 @@ public void Serialize(Batch batch, Stream stream)
 
         using var writer = new ExtendedBinaryWriter(gzipStream);
 
-        object[] row = null;
-        int counter = 0;
-        var enumerator = batch.Rows.GetEnumerator();
+        Memory<object> row = null;
+        var enumerator = batch.Rows.AsSpan(0, batch.Size).GetEnumerator();
         try
         {
             while (enumerator.MoveNext())
             {
-                row = (object[])enumerator.Current;
-                rowSerializer.Serialize(row, batch.Types, writer);
-
-                counter++;
-                if (counter >= batch.Size)
-                    break; // We've reached the batch size
+                row = enumerator.Current;
+                rowSerializer.Serialize(row.Span, batch.Types, writer);
             }
         }
         catch (Exception e)
         {
-            throw new ClickHouseBulkCopySerializationException(row, e);
+            throw new ClickHouseBulkCopySerializationException(row.ToArray(), e);
         }
     }
 }

ClickHouse.Client/Copy/Serializer/IRowSerializer.cs

@@ -1,9 +1,10 @@
+using System;
 using ClickHouse.Client.Formats;
 using ClickHouse.Client.Types;
 
 namespace ClickHouse.Client.Copy.Serializer;
 
 internal interface IRowSerializer
 {
-    void Serialize(object[] row, ClickHouseType[] types, ExtendedBinaryWriter writer);
+    void Serialize(Span<object> row, ClickHouseType[] types, ExtendedBinaryWriter writer);
 }

ClickHouse.Client/Copy/Serializer/RowBinarySerializer.cs

@@ -1,11 +1,12 @@
+using System;
 using ClickHouse.Client.Formats;
 using ClickHouse.Client.Types;
 
 namespace ClickHouse.Client.Copy.Serializer;
 
 internal class RowBinarySerializer : IRowSerializer
 {
-    public void Serialize(object[] row, ClickHouseType[] types, ExtendedBinaryWriter writer)
+    public void Serialize(Span<object> row, ClickHouseType[] types, ExtendedBinaryWriter writer)
     {
         for (int col = 0; col < row.Length; col++)
         {

ClickHouse.Client/Copy/Serializer/RowBinaryWithDefaultsSerializer.cs

@@ -1,3 +1,4 @@
+using System;
 using ClickHouse.Client.Constraints;
 using ClickHouse.Client.Formats;
 using ClickHouse.Client.Types;
@@ -7,7 +8,7 @@ namespace ClickHouse.Client.Copy.Serializer;
 // https://clickhouse.com/docs/en/interfaces/formats#rowbinarywithdefaults
 internal class RowBinaryWithDefaultsSerializer : IRowSerializer
 {
-    public void Serialize(object[] row, ClickHouseType[] types, ExtendedBinaryWriter writer)
+    public void Serialize(Span<object> row, ClickHouseType[] types, ExtendedBinaryWriter writer)
     {
         for (int col = 0; col < row.Length; col++)
         {