Destroy chunk after it is no longer needed.

Giorgi · Sep 13, 2023 · 462cbc4 · 462cbc4
1 parent 17747a5
commit 462cbc4
Show file tree

Hide file tree

Showing 4 changed files with 31 additions and 11 deletions.
diff --git a/DuckDB.NET.Bindings/DuckDBWrapperObjects.cs b/DuckDB.NET.Bindings/DuckDBWrapperObjects.cs
@@ -93,4 +93,17 @@ protected override bool ReleaseHandle()
             return true;
         }
     }
+
+    public class DuckDBDataChunk : SafeHandleZeroOrMinusOneIsInvalid
+    {
+        public DuckDBDataChunk() : base(true)
+        {
+        }
+
+        protected override bool ReleaseHandle()
+        {
+            NativeMethods.DataChunks.DuckDBDestroyDataChunk(out handle);
+            return true;
+        }
+    }
 }
diff --git a/DuckDB.NET.Bindings/NativeMethods/NativeMethods.DataChunks.cs b/DuckDB.NET.Bindings/NativeMethods/NativeMethods.DataChunks.cs
@@ -11,10 +11,10 @@ public static class DataChunks
         public static extern long DuckDBDataChunkGetColumnCount(IntPtr chunk);
 
         [DllImport(DuckDbLibrary, CallingConvention = CallingConvention.Cdecl, EntryPoint = "duckdb_data_chunk_get_vector")]
-        public static extern IntPtr DuckDBDataChunkGetVector(IntPtr chunk, long columnIndex);
+        public static extern IntPtr DuckDBDataChunkGetVector(DuckDBDataChunk chunk, long columnIndex);
 
         [DllImport(DuckDbLibrary, CallingConvention = CallingConvention.Cdecl, EntryPoint = "duckdb_data_chunk_get_size")]
-        public static extern long DuckDBDataChunkGetSize(IntPtr chunk);
+        public static extern long DuckDBDataChunkGetSize(DuckDBDataChunk chunk);
 
         [DllImport(DuckDbLibrary, CallingConvention = CallingConvention.Cdecl, EntryPoint = "duckdb_vector_get_column_type")]
         public static extern IntPtr DuckDBVectorGetColumnType(IntPtr vector);
@@ -33,5 +33,8 @@ public static class DataChunks
 
         [DllImport(DuckDbLibrary, CallingConvention = CallingConvention.Cdecl, EntryPoint = "duckdb_struct_vector_get_child")]
         public static extern long DuckDBStructVectorGetChild(IntPtr vector, long index);
+
+        [DllImport(DuckDbLibrary, CallingConvention = CallingConvention.Cdecl, EntryPoint = "duckdb_destroy_data_chunk")]
+        public static extern void DuckDBDestroyDataChunk(out IntPtr chunk);
     }
 }
diff --git a/DuckDB.NET.Bindings/NativeMethods/NativeMethods.Types.cs b/DuckDB.NET.Bindings/NativeMethods/NativeMethods.Types.cs
@@ -62,7 +62,7 @@ public static class Types
         public static extern DuckDBTimestampStruct DuckDBValueTimestamp([In, Out] DuckDBResult result, long col, long row);
 
         [DllImport(DuckDbLibrary, CallingConvention = CallingConvention.Cdecl, EntryPoint = "duckdb_result_get_chunk")]
-        public static extern IntPtr DuckDBResultGetChunk([In, Out] DuckDBResultStruct result, long chunkIndex);
+        public static extern DuckDBDataChunk DuckDBResultGetChunk([In, Out] DuckDBResultStruct result, long chunkIndex);
 
         [DllImport(DuckDbLibrary, CallingConvention = CallingConvention.Cdecl, EntryPoint = "duckdb_result_chunk_count")]
         public static extern long DuckDBResultChunkCount([In, Out] DuckDBResultStruct result);

diff --git a/DuckDB.NET.Data/DuckDBDataReader.cs b/DuckDB.NET.Data/DuckDBDataReader.cs
@@ -16,6 +16,7 @@ public class DuckDBDataReader : DbDataReader
         private readonly CommandBehavior behavior;
 
         private DuckDBResult currentResult;
+        private DuckDBDataChunk currentChunk;
         private readonly List<DuckDBResult> queryResults;
 
         private bool closed;
@@ -56,18 +57,18 @@ private void InitReaderData()
 
             InitChunkData();
 
-            // recordsAffected = (int)NativeMethods.Query.DuckDBRowsChanged(currentResult);
+            //recordsAffected = (int)NativeMethods.Query.DuckDBRowsChanged(currentResult);
         }
 
         private void InitChunkData()
         {
-            var chunk = NativeMethods.Types.DuckDBResultGetChunk(currentResult, currentChunkIndex);
-
-            currentChunkRowCount = NativeMethods.DataChunks.DuckDBDataChunkGetSize(chunk);
+            currentChunk?.Dispose();
+            currentChunk = NativeMethods.Types.DuckDBResultGetChunk(currentResult, currentChunkIndex);
+            currentChunkRowCount = NativeMethods.DataChunks.DuckDBDataChunkGetSize(currentChunk);
 
             for (int i = 0; i < fieldCount; i++)
             {
-                var vector = NativeMethods.DataChunks.DuckDBDataChunkGetVector(chunk, i);
+                var vector = NativeMethods.DataChunks.DuckDBDataChunkGetVector(currentChunk, i);
 
                 vectors[i] = NativeMethods.DataChunks.DuckDBVectorGetData(vector);
                 vectorValidityMask[i] = NativeMethods.DataChunks.DuckDBVectorGetValidity(vector);
@@ -427,7 +428,7 @@ public override IEnumerator GetEnumerator()
 
         public override DataTable GetSchemaTable()
         {
-            DataTable table = new DataTable
+            var table = new DataTable
             {
                 Columns =
                 {
@@ -438,8 +439,10 @@ public override DataTable GetSchemaTable()
                      { "AllowDBNull", typeof(bool) }
                 }
             };
-            object[] rowData = new object[5];
-            for (int i = 0; i < FieldCount; i++)
+
+            var rowData = new object[5];
+
+            for (var i = 0; i < FieldCount; i++)
             {
                 rowData[0] = i;
                 rowData[1] = GetName(i);
@@ -456,6 +459,7 @@ public override void Close()
         {
             if (closed) return;
 
+            currentChunk?.Dispose();
             foreach (var result in queryResults)
             {
                 result.Dispose();