perf: faster reading of bytes,bytearray,np.ndarray,mmap index

Author: william-silversmith

mapbuffer/intmap.py

@@ -67,9 +67,20 @@ def index(self):
       return self._index
 
     N = len(self)
-    index_length = 2 * N * 8
-    index = self.buffer[self.HEADER_LENGTH:index_length+self.HEADER_LENGTH]
-    self._index = np.frombuffer(index, dtype=np.uint64).reshape((N,2))
+    index_length = 2 * N
+
+    if isinstance(self.buffer, (bytes,bytearray,np.ndarray,mmap.mmap)):
+      self._index = np.frombuffer(
+        self.buffer,
+        offset=self.HEADER_LENGTH,
+        count=index_length,
+        dtype=np.uint64,
+      ).reshape((N,2))
+    else:
+      index_length *= 8
+      index = self.buffer[self.HEADER_LENGTH:index_length+self.HEADER_LENGTH]
+      self._index = np.frombuffer(index, dtype=np.uint64).reshape((N,2))
+
     return self._index
 
   def keys(self):

mapbuffer/mapbuffer.py

@@ -105,9 +105,20 @@ def index(self):
       return self._index
 
     N = len(self)
-    index_length = 2 * N * 8
-    index = self.buffer[HEADER_LENGTH:index_length+HEADER_LENGTH]
-    self._index = np.frombuffer(index, dtype=np.uint64).reshape((N,2))
+    index_length = 2 * N
+
+    if isinstance(self.buffer, (bytes,bytearray,np.ndarray,mmap.mmap)):
+      self._index = np.frombuffer(
+        self.buffer,
+        offset=HEADER_LENGTH,
+        count=index_length,
+        dtype=np.uint64,
+      ).reshape((N,2))
+    else:
+      index_length *= 8
+      index = self.buffer[HEADER_LENGTH:index_length+HEADER_LENGTH]
+      self._index = np.frombuffer(index, dtype=np.uint64).reshape((N,2))
+    
     return self._index
 
   def keys(self):