Support proper numpy integration for ~100x performance boost

Author: VeaaC

flatdata-py/README.md

@@ -18,6 +18,37 @@ Once you have [created a flatdata schema file](../README.md#creating-a-schema),
 flatdata-generator --gen py --schema locations.flatdata --output-file locations.py
 ```
 
+## Performance tips
+
+`flatdata-py` supports two data access patterns with very different performance characteristics on large archives.
+
+Iterating over a vector yields one Python object per element. Each field access unpacks bits from the underlying memory-mapped data. This is fine for accessing individual elements or small ranges, but has significant per-element overhead for bulk operations:
+
+```python
+count = sum(1 for x in archive.links if x.speed_limit > 100)
+```
+
+For bulk operations, use the vectorized access methods that read fields directly into NumPy arrays:
+
+```python
+# single column access, returns a pandas DataFrame
+df = archive.links.speed_limit
+count = len(df[df['speed_limit'] > 100])
+
+# full NumPy structured array with all fields
+arr = archive.links.to_numpy()
+count = int(np.sum(arr['speed_limit'] > 100))
+
+# slices work too
+arr = archive.links[1000:2000].to_numpy()
+df = archive.links[::10].to_data_frame()
+```
+
+* Use `vector.field_name` (column access) when you only need one or a few fields.
+* Use `vector.to_numpy()` or `vector.to_data_frame()` when you need all fields at once.
+* Use `vector[i].field` for random access to individual elements.
+* The underlying data is memory-mapped; the OS pages it from disk on demand. Vectorized results are materialized as NumPy arrays in RAM.
+
 ## Using the inspector
 
 `flatdata-py` comes with a handy tool called the `flatdata-inspector` to inspect the contents of an archive:

flatdata-py/flatdata/lib/archive.py

@@ -39,9 +39,9 @@ def __init__(self, resource_storage):
             self.__getattr__(name)
 
     def __getattr__(self, name):
-        if name not in list(self._RESOURCES.keys()):
+        if name not in self._RESOURCES:
             raise AttributeError("Resource %s not defined in archive." % name)
-        if name not in list(self._loaded_resources.keys()):
+        if name not in self._loaded_resources:
             self._loaded_resources[name] = self._open_resource(name)
         return self._loaded_resources[name]
 

flatdata-py/flatdata/lib/data_access.py

@@ -3,6 +3,8 @@
  See the LICENSE file in the root of this project for license details.
 '''
 
+import numpy as np
+
 # Sign bits cache for the value reading.
 _SIGN_BITS = [0] + [(1 << (bits - 1)) for bits in range(1, 65)]
 
@@ -62,3 +64,32 @@ def write_value(data, offset_bits, num_bits, is_signed, value):
         surrounding_bits = data[offset_bytes + byte_idx] & ~((1 << offset_bits) - 1)
         data[offset_bytes + byte_idx] = value_in_little_endian[byte_idx] & ((1 << (8 - (bits_written % 8))) - 1)
         data[offset_bytes + byte_idx] |= surrounding_bits
+
+
+def read_field_vectorized(raw_bytes_2d, field_offset_bits, field_width_bits, is_signed):
+    """Read a bit-packed field from all elements at once, returning a numpy array.
+
+    :param raw_bytes_2d: numpy uint8 array shaped (num_elements, struct_size_bytes)
+    :param field_offset_bits: bit offset of the field within each element
+    :param field_width_bits: width of the field in bits (max 64)
+    :param is_signed: whether to sign-extend the result
+    :return: numpy array of field values
+    """
+    byte_start = field_offset_bits // 8
+    bit_shift = field_offset_bits % 8
+    bytes_needed = (bit_shift + field_width_bits + 7) // 8
+
+    result = np.zeros(raw_bytes_2d.shape[0], dtype=np.uint64)
+    for b in range(bytes_needed):
+        result |= raw_bytes_2d[:, byte_start + b].astype(np.uint64) << np.uint64(b * 8)
+    result >>= np.uint64(bit_shift)
+
+    if field_width_bits < 64:
+        result &= np.uint64((1 << field_width_bits) - 1)
+
+    if is_signed:
+        sign_bit = np.uint64(1 << (field_width_bits - 1))
+        signed = result.astype(np.int64) - np.int64(1 << field_width_bits)
+        result = np.where(result & sign_bit, signed, result.astype(np.int64))
+
+    return result

flatdata-py/flatdata/lib/resources.py

@@ -8,7 +8,7 @@
 import pandas as pd
 import numpy as np
 
-from .data_access import read_value
+from .data_access import read_value, read_field_vectorized
 from .errors import CorruptResourceError
 
 SIZE_OFFSET_IN_BITS = 64
@@ -24,6 +24,7 @@ def __init__(self, mem, element_type):
         self._element_type = element_type
         self._element_types = [element_type]
         self._type_size_in_bytes = self._element_type._SIZE_IN_BYTES if self._element_type else 1
+        self._raw_numpy_2d = None
 
     def size_in_bytes(self):
         return len(self._mem)
@@ -35,6 +36,20 @@ def _get_item(self, index):
         offset = self._item_offset(index)
         return self._element_type(self._mem, offset)
 
+    def _as_numpy_2d(self):
+        """Return the raw data as a 2D numpy uint8 array of shape (n, struct_size).
+        Zero-copy via np.frombuffer on the mmap'd memory. Cached after first call.
+        """
+        if self._raw_numpy_2d is None:
+            n = len(self)
+            struct_size = self._type_size_in_bytes
+            raw = np.frombuffer(
+                self._mem[SIZE_OFFSET_IN_BYTES:SIZE_OFFSET_IN_BYTES + n * struct_size],
+                dtype=np.uint8,
+            )
+            self._raw_numpy_2d = raw.reshape(n, struct_size)
+        return self._raw_numpy_2d
+
     def _repr_attributes(self):
         return {
             "container_type": self.__class__.__name__,
@@ -60,14 +75,19 @@ def __init__(self, s, sequence):
         self._sequence = sequence
 
     def to_numpy(self, limit=None):
+        raw_2d = self._sequence._as_numpy_2d()
         indices = self._slice.indices(len(self._sequence))
-        num_items = len(range(*indices)) if not limit else limit
-        result = np.empty(
-            shape=num_items,
-            dtype=self._sequence._element_type.dtype()
-        )
-        for index, item in enumerate(self):
-            result[index] = item.as_tuple()
+        sliced = raw_2d[self._slice]
+        if limit is not None:
+            sliced = sliced[:limit]
+
+        fields = self._sequence._element_type._FIELDS
+        dtype = self._sequence._element_type.dtype()
+        result = np.empty(sliced.shape[0], dtype=dtype)
+        for name, field in fields.items():
+            result[name] = read_field_vectorized(
+                sliced, field.offset, field.width, field.is_signed
+            )
         return result
 
     def to_data_frame(self, limit=None):
@@ -78,7 +98,10 @@ def __iter__(self):
             yield self._sequence[i]
 
     def __getattr__(self, name):
-        return pd.DataFrame(data=[[getattr(item, name)] for item in self], columns=[name])
+        raw_2d = self._sequence._as_numpy_2d()[self._slice]
+        field = self._sequence._element_type._FIELDS[name]
+        values = read_field_vectorized(raw_2d, field.offset, field.width, field.is_signed)
+        return pd.DataFrame(data=values, columns=[name])
 
     def __repr__(self):
         return "Displaying first 100 records:\n" + self.to_data_frame(limit=100).__repr__()
@@ -92,8 +115,20 @@ def __init__(self, mem, element_type):
         assert rem == 0, "Malformed vector"
         self._size = size
 
+    def to_numpy(self):
+        """Convert entire vector to a numpy structured array (vectorized)."""
+        raw_2d = self._as_numpy_2d()
+        fields = self._element_type._FIELDS
+        dtype = self._element_type.dtype()
+        result = np.empty(self._size, dtype=dtype)
+        for name, field in fields.items():
+            result[name] = read_field_vectorized(
+                raw_2d, field.offset, field.width, field.is_signed
+            )
+        return result
+
     def to_data_frame(self):
-        return self[:].to_data_frame()
+        return pd.DataFrame(data=self.to_numpy())
 
     def __getitem__(self, index):
         if isinstance(index, slice):
@@ -106,11 +141,17 @@ def __getitem__(self, index):
         return self._get_item(index)
 
     def __iter__(self):
-        for i in range(len(self)):
-            yield self._get_item(i)
+        mem = self._mem
+        element_type = self._element_type
+        size_bytes = self._type_size_in_bytes
+        for i in range(self._size):
+            yield element_type(mem, SIZE_OFFSET_IN_BYTES + size_bytes * i)
 
     def __getattr__(self, name):
-        return pd.DataFrame(data=[[getattr(item, name)] for item in self], columns=[name])
+        raw_2d = self._as_numpy_2d()
+        field = self._element_type._FIELDS[name]
+        values = read_field_vectorized(raw_2d, field.offset, field.width, field.is_signed)
+        return pd.DataFrame(data=values, columns=[name])
 
     def __len__(self):
         return self._size

flatdata-py/flatdata/lib/structure.py

@@ -9,6 +9,8 @@
 
 
 class Structure:
+    __slots__ = ('_mem', '_pos')
+
     def __init__(self, mem, pos):
         self._mem = mem
         self._pos = pos

flatdata-py/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "flatdata-py"
-version = "0.4.10"
+version = "0.4.11"
 description = "Python 3 implementation of Flatdata"
 readme = "README.md"
 authors = [