Fix reading all row groups when none are selected (#13)

Author: adamreeve

ParquetSharp.Dataset.Test/InstrumentedFilter.cs

@@ -0,0 +1,44 @@
+using System.Collections.Generic;
+using Apache.Arrow;
+using ParquetSharp.Dataset.Filter;
+
+namespace ParquetSharp.Dataset.Test;
+
+internal sealed class InstrumentedFilter : IFilter
+{
+    public InstrumentedFilter(IFilter filter)
+    {
+        _filter = filter;
+    }
+
+    public int IncludePartitionCallCount { get; private set; }
+
+    public int IncludeRowGroupCallCount { get; private set; }
+
+    public int ComputeMaskRowCount { get; private set; }
+
+    public bool IncludePartition(PartitionInformation partitionInformation)
+    {
+        ++IncludePartitionCallCount;
+        return _filter.IncludePartition(partitionInformation);
+    }
+
+    public bool IncludeRowGroup(IReadOnlyDictionary<string, LogicalStatistics> columnStatistics)
+    {
+        ++IncludeRowGroupCallCount;
+        return _filter.IncludeRowGroup(columnStatistics);
+    }
+
+    public FilterMask? ComputeMask(RecordBatch dataBatch)
+    {
+        ComputeMaskRowCount += dataBatch.Length;
+        return _filter.ComputeMask(dataBatch);
+    }
+
+    public IEnumerable<string> Columns()
+    {
+        return _filter.Columns();
+    }
+
+    private readonly IFilter _filter;
+}

ParquetSharp.Dataset.Test/TestDatasetReader.cs

@@ -97,12 +97,18 @@ await VerifyData(
             new Dictionary<string, int> { { "a", 20 }, { "b", 20 } });
 
         // Read filtered on partition
-        var filter = Col.Named("part").IsEqualTo("b");
+        var filter = new InstrumentedFilter(Col.Named("part").IsEqualTo("b"));
         using var filteredReader = dataset.ToBatches(filter);
         await VerifyData(
             filteredReader,
             new Dictionary<int, int> { { 2, 10 }, { 3, 10 } },
             new Dictionary<string, int> { { "b", 20 } });
+
+        // IncludePartition is called for the root directory and 2 subdirectories
+        Assert.That(filter.IncludePartitionCallCount, Is.EqualTo(3));
+        // No data files are used in the filter, so row group filtering is not used
+        Assert.That(filter.IncludeRowGroupCallCount, Is.EqualTo(0));
+        Assert.That(filter.ComputeMaskRowCount, Is.EqualTo(20));
     }
 
     [Test]
@@ -322,10 +328,15 @@ public async Task TestFilterOnFileColumn()
             new NoPartitioning(),
             schema: schema);
 
-        var filter = Col.Named("id").IsInRange(0, 1);
+        var filter = new InstrumentedFilter(Col.Named("id").IsInRange(0, 1));
         using var reader = dataset.ToBatches(filter);
 
         await VerifyData(reader, new Dictionary<int, int> { { 0, 10 }, { 1, 10 } });
+
+        Assert.That(filter.IncludePartitionCallCount, Is.EqualTo(1));
+        Assert.That(filter.IncludeRowGroupCallCount, Is.EqualTo(4));
+        // Should only read data from one row group (10 rows) from each file
+        Assert.That(filter.ComputeMaskRowCount, Is.EqualTo(20));
     }
 
     [Test]
@@ -348,10 +359,15 @@ public async Task TestAllRowGroupsInFileExcluded()
             new NoPartitioning(),
             schema: schema);
 
-        var filter = Col.Named("id").IsEqualTo(2);
+        var filter = new InstrumentedFilter(Col.Named("id").IsEqualTo(2));
         using var reader = dataset.ToBatches(filter);
 
         await VerifyData(reader, new Dictionary<int, int> { { 2, 10 } });
+
+        Assert.That(filter.IncludePartitionCallCount, Is.EqualTo(1));
+        Assert.That(filter.IncludeRowGroupCallCount, Is.EqualTo(4));
+        // Should only read data from one row group (10 rows) from the second file
+        Assert.That(filter.ComputeMaskRowCount, Is.EqualTo(10));
     }
 
     [Test]

ParquetSharp.Dataset/DatasetStreamReader.cs

@@ -111,22 +111,28 @@ private void GetNextReader()
         _currentFragmentReader?.Dispose();
         _currentFileReader?.Dispose();
 
-        if (_fragmentEnumerator.MoveNext())
+        while (_fragmentEnumerator.MoveNext())
         {
             _currentFileReader = new FileReader(
                 _fragmentEnumerator.Current.FilePath, _readerProperties, _arrowReaderProperties);
 
             var rowGroups = _rowGroupSelector?.GetRequiredRowGroups(_currentFileReader);
+            if (rowGroups != null && rowGroups.Length == 0)
+            {
+                _currentFileReader.Dispose();
+                continue;
+            }
+
             var columnIndices = GetFileColumnIndices(_currentFileReader);
 
             _currentFragmentReader = _currentFileReader.GetRecordBatchReader(
                 rowGroups: rowGroups, columns: columnIndices);
+
+            return;
         }
-        else
-        {
-            _currentFileReader = null;
-            _currentFragmentReader = null;
-        }
+
+        _currentFileReader = null;
+        _currentFragmentReader = null;
     }
 
     private int[] GetFileColumnIndices(FileReader fileReader)

ParquetSharp.Dataset/Filter/DateRangeEvaluator.cs

@@ -26,7 +26,7 @@ public void Visit(Date32Array array)
             var endNumber = _end.DayNumber - ArrowEpoch.DayNumber;
             if (inputArray.NullCount == 0)
             {
-                var values = array.Values;
+                var values = inputArray.Values;
                 for (var i = 0; i < inputArray.Length; ++i)
                 {
                     var value = values[i];
@@ -38,7 +38,7 @@ public void Visit(Date32Array array)
             {
                 for (var i = 0; i < inputArray.Length; ++i)
                 {
-                    var value = array.GetValue(i);
+                    var value = inputArray.GetValue(i);
                     var isInRange = value.HasValue && value.Value >= startNumber && value.Value <= endNumber;
                     BitUtility.SetBit(mask, i, isInRange);
                 }
@@ -54,7 +54,7 @@ public void Visit(Date64Array array)
         {
             for (var i = 0; i < inputArray.Length; ++i)
             {
-                var value = array.GetDateOnly(i);
+                var value = inputArray.GetDateOnly(i);
                 var isInRange = value.HasValue && value.Value >= _start && value.Value <= _end;
                 BitUtility.SetBit(mask, i, isInRange);
             }