perf: make robust to network access patterns by caching header

Author: william-silversmith

automated_test.py

@@ -2,6 +2,7 @@
 import numpy as np
 from mapbuffer import MapBuffer, HEADER_LENGTH
 import random
+import mmap
 
 @pytest.mark.parametrize("compress", (None, "gzip", "br", "zstd", "lzma"))
 def test_empty(compress):
@@ -49,4 +50,46 @@ def test_full(compress):
 
   mbuf.validate()
 
-  assert len(mbuf.buffer) > HEADER_LENGTH
+  assert len(mbuf.buffer) > HEADER_LENGTH
+
+@pytest.mark.parametrize("compress", (None, "gzip", "br", "zstd"))
+def test_mmap_access(compress):
+  data = { 
+    1: b"hello",
+    2: b"world",
+  }
+  mbuf = MapBuffer(data, compress=compress)
+  with open("test_mmap.mb", "wb") as f:
+    f.write(mbuf.tobytes())
+
+  with open("test_mmap.mb", "rb") as f:
+    mb = MapBuffer(f)
+
+    assert mb[1] == b"hello"
+    assert mb[2] == b"world"
+
+@pytest.mark.parametrize("compress", (None, "gzip", "br", "zstd"))
+def test_object_access(compress):
+  data = { 
+    1: b"hello",
+    2: b"world",
+  }
+  mbuf = MapBuffer(data, compress=compress)
+
+  class Reader:
+    def __init__(self):
+      self.lst = mbuf.tobytes()
+    def __getitem__(self, slc):
+      return self.lst[slc]
+
+  mbuf2 = MapBuffer(Reader())
+  assert mbuf2[1] == b"hello"
+  assert mbuf2[2] == b"world"
+
+
+
+
+
+
+
+

mapbuffer/mapbuffer.py

@@ -17,7 +17,8 @@ class MapBuffer:
   """Represents a usable int->bytes dictionary as a byte string."""
   __slots__ = (
     "data", "tobytesfn", "frombytesfn", 
-    "dtype", "buffer", "_index", "_compress"
+    "dtype", "buffer", "_header", 
+    "_index", "_compress"
   )
   def __init__(
     self, data=None, compress=None,
@@ -41,6 +42,7 @@ def __init__(
     self.dtype = np.uint64
     self.buffer = None
 
+    self._header = None
     self._index = None
     self._compress = None
 
@@ -50,26 +52,31 @@ def __init__(
       self.buffer = mmap.mmap(data.fileno(), 0, prot=mmap.PROT_READ)
     elif isinstance(data, (bytes, mmap.mmap)):
       self.buffer = data
+    elif hasattr(data, "__getitem__"):
+      self.buffer = data
     else:
-      raise TypeError("data must be a dict, bytes, file, or mmap. Got: " + str(type(data)))
+      raise TypeError(
+        f"data must be a dict, bytes, file, mmap, or otherwise support "
+        f"__getitem__ with slice support for byte ranges. Got: {type(data)}"
+      )
 
   def __len__(self):
     """Returns number of keys."""
-    return int.from_bytes(self.buffer[12:16], byteorder="little", signed=False)
+    return int.from_bytes(self.header[12:16], byteorder="little", signed=False)
 
   @property
   def compress(self):
     if self._compress:
       return self._compress
 
     self._compress = compression.normalize_encoding(
-      self.buffer[8:12]
+      self.header[8:12]
     )
     return self._compress
 
   @property
   def format_version(self):
-    return self.buffer[len(MAGIC_NUMBERS)]
+    return self.header[len(MAGIC_NUMBERS)]
 
   def __iter__(self):
     yield from self.keys()
@@ -78,6 +85,17 @@ def datasize(self):
     """Returns size of data region in bytes."""
     return len(self.buffer) - HEADER_LENGTH - len(self) * 2 * 8
 
+  @property
+  def header(self):
+    """Get the header bytes."""
+    if self._header is not None:
+      return self._header
+
+    # seems dumb, buf if self.buffer is an object that
+    # requires network access, this is a valuable cache
+    self._header = self.buffer[:HEADER_LENGTH]
+    return self._header
+
   def index(self):
     """Get an Nx2 numpy array representing the index."""
     if self._index is not None:
@@ -217,11 +235,12 @@ def validate(self):
   @staticmethod
   def validate_buffer(buf):
     mapbuf = MapBuffer(buf)
+    header = mapbuf.header
     index = mapbuf.index()
     if len(index) != len(mapbuf):
       raise ValidationError(f"Index size doesn't match. len(mapbuf): {len(mapbuf)}")
 
-    magic = buf[:len(MAGIC_NUMBERS)]
+    magic = header[:len(MAGIC_NUMBERS)]
     if magic != MAGIC_NUMBERS:
       raise ValidationError(f"Magic number mismatch. Expected: {MAGIC_NUMBERS} Got: {magic}")