Dispose of Parquet RecordBatches when creating a new filtered batch (#25)

Author: adamreeve

ParquetSharp.Dataset.Benchmark/DatasetRead.cs

@@ -76,8 +76,9 @@ public async Task<long> ReadAllFilesDirectly()
         {
             using var reader = new FileReader(filePath);
             using var batchReader = reader.GetRecordBatchReader();
-            while (await batchReader.ReadNextRecordBatchAsync() is { } batch)
+            while (await batchReader.ReadNextRecordBatchAsync() is { } batch_)
             {
+                using var batch = batch_;
                 rowsRead += batch.Length;
             }
         }
@@ -91,8 +92,9 @@ public async Task<long> ReadAllData()
         var dataset = new DatasetReader(_datasetDirectory, new HivePartitioning.Factory());
         using var reader = dataset.ToBatches();
         long rowsRead = 0;
-        while (await reader.ReadNextRecordBatchAsync() is { } batch)
+        while (await reader.ReadNextRecordBatchAsync() is { } batch_)
         {
+            using var batch = batch_;
             rowsRead += batch.Length;
         }
 
@@ -106,8 +108,9 @@ public async Task<long> FilterPartitions()
         var filter = Col.Named("group").IsEqualTo("group-2").And(Col.Named("day").IsEqualTo(2));
         using var reader = dataset.ToBatches(filter);
         long rowsRead = 0;
-        while (await reader.ReadNextRecordBatchAsync() is { } batch)
+        while (await reader.ReadNextRecordBatchAsync() is { } batch_)
         {
+            using var batch = batch_;
             rowsRead += batch.Length;
         }
 
@@ -121,8 +124,9 @@ public async Task<long> FilterFileData()
         var filter = Col.Named("id").IsEqualTo(5);
         using var reader = dataset.ToBatches(filter);
         long rowsRead = 0;
-        while (await reader.ReadNextRecordBatchAsync() is { } batch)
+        while (await reader.ReadNextRecordBatchAsync() is { } batch_)
         {
+            using var batch = batch_;
             rowsRead += batch.Length;
         }
 
@@ -135,8 +139,9 @@ public async Task<long> FilterToFileColumns()
         var dataset = new DatasetReader(_datasetDirectory, new HivePartitioning.Factory());
         using var reader = dataset.ToBatches(columns: new[] { "id", "ints", "doubles" });
         long rowsRead = 0;
-        while (await reader.ReadNextRecordBatchAsync() is { } batch)
+        while (await reader.ReadNextRecordBatchAsync() is { } batch_)
         {
+            using var batch = batch_;
             rowsRead += batch.Length;
         }
 
@@ -149,8 +154,9 @@ public async Task<long> FilterToSingleColumn()
         var dataset = new DatasetReader(_datasetDirectory, new HivePartitioning.Factory());
         using var reader = dataset.ToBatches(columns: new[] { "ints" });
         long rowsRead = 0;
-        while (await reader.ReadNextRecordBatchAsync() is { } batch)
+        while (await reader.ReadNextRecordBatchAsync() is { } batch_)
         {
+            using var batch = batch_;
             rowsRead += batch.Length;
         }
 

ParquetSharp.Dataset.Test/TestEncryption.cs

@@ -47,8 +47,9 @@ public static async Task TestReadEncryptedData()
         using var reader = dataset.ToBatches();
 
         var rowsRead = 0;
-        while (await reader.ReadNextRecordBatchAsync() is { } batch)
+        while (await reader.ReadNextRecordBatchAsync() is { } batch_)
         {
+            using var batch = batch_;
             rowsRead += batch.Length;
         }
 

ParquetSharp.Dataset.Test/TestFilterParquet.cs

@@ -98,8 +98,9 @@ private static async Task TestIntegerFilter(IFilter filter, Func<int, bool> incl
         var dataset = new DatasetReader(datasetDir.DirectoryPath);
         using var reader = dataset.ToBatches(filter);
         var valuesRead = new List<int?>();
-        while (await reader.ReadNextRecordBatchAsync() is { } batch)
+        while (await reader.ReadNextRecordBatchAsync() is { } batch_)
         {
+            using var batch = batch_;
             var filteredBatchValues = batch.Column(0) as Int32Array;
             foreach (var value in filteredBatchValues!)
             {

ParquetSharp.Dataset/DatasetStreamReader.cs

@@ -1,3 +1,4 @@
+using System;
 using System.Collections.Generic;
 using System.Linq;
 using System.Threading;
@@ -48,7 +49,7 @@ public DatasetStreamReader(
             var nextBatch = await _currentFragmentReader.ReadNextRecordBatchAsync(cancellationToken);
             if (nextBatch != null)
             {
-                var filtered = FilterBatch(nextBatch);
+                var filtered = FilterBatch(ref nextBatch);
                 if (filtered == null)
                 {
                     // All rows excluded
@@ -70,9 +71,16 @@ public DatasetStreamReader(
     /// <summary>
     /// Return a record batch with rows filtered out using the current filter.
     /// Returns null if all rows are excluded.
+    /// Takes ownership of the input record batch and will dispose it if it
+    /// is filtered.
     /// </summary>
-    private RecordBatch? FilterBatch(RecordBatch recordBatch)
+    private RecordBatch? FilterBatch(ref RecordBatch? recordBatch)
     {
+        if (recordBatch == null)
+        {
+            throw new ArgumentNullException(nameof(recordBatch));
+        }
+
         if (_filter == null)
         {
             return recordBatch;
@@ -97,7 +105,13 @@ public DatasetStreamReader(
             arrays.Add(filterApplier.MaskedArray);
         }
 
-        return new RecordBatch(recordBatch.Schema, arrays, filterMask.IncludedCount);
+        var filteredBatch = new RecordBatch(recordBatch.Schema, arrays, filterMask.IncludedCount);
+
+        // Dispose input record batch so its memory can be immediately freed
+        recordBatch.Dispose();
+        recordBatch = null;
+
+        return filteredBatch;
     }
 
     public void Dispose()

ParquetSharp.Dataset/ParquetSharp.Dataset.csproj

@@ -5,7 +5,7 @@
     <LangVersion>10.0</LangVersion>
     <Nullable>enable</Nullable>
     <TreatWarningsAsErrors>true</TreatWarningsAsErrors>
-    <VersionPrefix>0.3.0</VersionPrefix>
+    <VersionPrefix>0.3.1</VersionPrefix>
     <Company>G-Research</Company>
     <Authors>G-Research</Authors>
     <Product>ParquetSharp.Dataset</Product>