add filter test

Author: martindurant

fsspec/caching.py

@@ -652,11 +652,28 @@ def __init__(
         else:
             self.data = {}
 
+    @property
+    def size(self):
+        return sum(_[1] - _[0] for _ in self.data)
+
+    @size.setter
+    def size(self, value):
+        pass
+
+    @property
+    def nblocks(self):
+        return len(self.data)
+
+    @nblocks.setter
+    def nblocks(self, value):
+        pass
+
     def _fetch(self, start: int | None, stop: int | None) -> bytes:
         if start is None:
             start = 0
         if stop is None:
             stop = self.size
+        self.total_requested_bytes += stop - start
 
         out = b""
         started = False
@@ -665,24 +682,34 @@ def _fetch(self, start: int | None, stop: int | None) -> bytes:
             if (loc0 <= start < loc1) and (loc0 <= stop <= loc1):
                 # entirely within the block
                 off = start - loc0
+                self.hit_count += 1
                 return self.data[(loc0, loc1)][off : off + stop - start]
+            if stop <= loc0:
+                break
             if started and loc0 > loc_old:
                 # a gap where we need data
+                self.miss_count += 1
                 if self.strict:
                     raise ValueError
                 out += b"\x00" * (loc0 - loc_old)
             if loc0 <= start < loc1:
                 # found the start
+                self.hit_count += 1
                 off = start - loc0
                 out = self.data[(loc0, loc1)][off : off + stop - start]
                 started = True
             elif start < loc0 and stop > loc1:
                 # the whole block
+                self.hit_count += 1
                 out += self.data[(loc0, loc1)]
             elif loc0 <= stop <= loc1:
                 # end block
+                self.hit_count += 1
                 return out + self.data[(loc0, loc1)][: stop - loc0]
             loc_old = loc1
+        self.miss_count += 1
+        if started and not self.strict:
+            return out + b"\x00" * (stop - loc_old)
         raise ValueError
 
 

fsspec/tests/test_caches.py

@@ -225,19 +225,37 @@ def test_cache_basic(Cache_imp, blocksize, size_requests):
         assert result == expected
 
 
+@pytest.mark.parametrize("strict", [True, False])
 @pytest.mark.parametrize("sort", [True, False])
-def test_known(sort):
-    parts = {(10, 20): b"1" * 10, (20, 30): b"2" * 10, (0, 10): b"0" * 10}
+def test_known(strict, sort):
+    parts = {
+        (10, 20): b"1" * 10,
+        (20, 30): b"2" * 10,
+        (0, 10): b"0" * 10,
+        (40, 50): b"3" * 10,
+    }
     if sort:
         parts = dict(sorted(parts.items()))
-    c = caches["parts"](None, None, 100, parts)
+    c = caches["parts"](None, None, 100, parts, strict=strict)
+    assert c.size == 40
     assert (0, 30) in c.data  # got consolidated
+    assert c.nblocks == 2
+
     assert c._fetch(5, 15) == b"0" * 5 + b"1" * 5
     assert c._fetch(15, 25) == b"1" * 5 + b"2" * 5
-    # Over-read will raise error
-    with pytest.raises(ValueError):
-        # tries to call None fetcher
-        c._fetch(25, 35)
+    assert c.hit_count
+    assert not c.miss_count
+
+    if strict:
+        # Over-read will raise error
+        with pytest.raises(ValueError):
+            c._fetch(25, 35)
+        with pytest.raises(ValueError):
+            c._fetch(25, 45)
+    else:
+        assert c._fetch(25, 35) == b"2" * 5 + b"\x00" * 5
+        assert c._fetch(25, 45) == b"2" * 5 + b"\x00" * 10 + b"3" * 5
+    assert c.miss_count
 
 
 def test_background(server, monkeypatch):

fsspec/tests/test_parquet.py

@@ -21,7 +21,7 @@
 
 @pytest.fixture(
     params=[
-        # pytest.param("fastparquet", marks=FASTPARQUET_MARK),
+        pytest.param("fastparquet", marks=FASTPARQUET_MARK),
         pytest.param("pyarrow", marks=PYARROW_MARK),
     ]
 )
@@ -145,3 +145,32 @@ def test_open_parquet_file(
                 max_block=max_block,
                 footer_sample_size=footer_sample_size,
             )
+
+
+@FASTPARQUET_MARK
+def test_with_filter(tmpdir):
+    import pandas as pd
+
+    df = pd.DataFrame(
+        {
+            "a": [10, 1, 2, 3, 7, 8, 9],
+            "b": ["a", "a", "a", "b", "b", "b", "b"],
+        }
+    )
+    fn = os.path.join(str(tmpdir), "test.parquet")
+    df.to_parquet(fn, engine="fastparquet", row_group_offsets=[0, 3], stats=True)
+
+    expect = pd.read_parquet(fn, engine="fastparquet", filters=[["b", "==", "b"]])
+    f = open_parquet_file(
+        fn,
+        engine="fastparquet",
+        filters=[["b", "==", "b"]],
+        max_gap=1,
+        max_block=1,
+        footer_sample_size=8,
+    )
+    assert (0, 4) in f.cache.data
+    assert f.cache.size < os.path.getsize(fn)
+
+    result = pd.read_parquet(f, engine="fastparquet", filters=[["b", "==", "b"]])
+    pd.testing.assert_frame_equal(expect, result)