microecs.query_result

 1"""query_result.py - A list of pools seen as a contiguous view. Implements array interface to look like numpy"""
 2import numpy as np
 3
 4from .utils import Shape, EntityId
 5from .pool import Pool
 6from .qr_field import QRField
 7
 8# Note: if QueryResult gets new fields, add them here! Otherwise the user code may overwrite them e.g. qr._data=xxx
 9_QR_INTERNAL_ATTRS = {"pool_list", "fields", "_field_shapes", "_field_dtypes", "_data",
10                      "_cache", "_entity_ids", "_len", "_pool_ids"}
11
12class _QRArray(np.ndarray):
13    """small shim array so we don't instantiate QRField which is more expensive (contiguous view for >=2 pools)"""
14    def numpy(self) -> np.ndarray:
15        """for compatibility with QRField.numpy()"""
16        return np.asarray(self)
17
18    @property
19    def parts(self) -> list[np.ndarray]:
20        """for compatibility with QRField.parts"""
21        return [np.asarray(self)]
22
23class QueryResult:
24    """A query result containing entities. Fields (e.g. qr.position) implement array interface to look like numpy"""
25    def __init__(self, pool_list: list[Pool], field_shapes: dict[str, Shape], field_dtypes: dict[str, np.dtype],
26                 pool_ids: dict[Pool, list[EntityId]]):
27        self.pool_list = pool_list
28        self.fields = list(field_shapes)
29        self._field_shapes = field_shapes
30        self._field_dtypes = field_dtypes
31        self._pool_ids = pool_ids
32
33        self._data: dict[str, list[np.ndarray]] = {f: [p.data[f][0:len(p)] for p in pool_list] for f in field_shapes}
34        self._cache: dict[str, QRField | _QRArray] = {}
35        self._entity_ids: np.ndarray | None = None
36        self._len: int | None = None
37
38    @property
39    def entity_ids(self) -> np.ndarray:
40        """The entity ids of this query result"""
41        if self._entity_ids is None:
42            self._entity_ids = np.array(sum((self._pool_ids[p] for p in self.pool_list), []), dtype="int64")
43        return self._entity_ids
44
45    def __getattr__(self, name):
46        data: dict[str, list[np.ndarray]]
47        # .get and not self._data: on an instance built without __init__ (copy/deepcopy/pickle.loads probe
48        # hasattr(__setstate__)), _data is MISSING -- and self._data would re-enter __getattr__ -> RecursionError
49        if name not in (data := self.__dict__.get("_data", {})):
50            raise AttributeError(f"'{name}' not part of {self.__dict__.get('fields')}")
51
52        if name not in self._cache:
53            if len(parts := data[name]) in (0, 1): # optimized path for a single pool -> return an actual np array
54                # the if/else part is in case no pools match the query so we create a (0, k) array for that field.
55                arr = parts[0] if parts else np.empty((0, *self._field_shapes[name]), self._field_dtypes[name])
56                self._cache[name] = arr.view(_QRArray)
57            else:
58                self._cache[name] = QRField(parts)
59
60        return self._cache[name]
61
62    def __setattr__(self, name, value):
63        if name in _QR_INTERNAL_ATTRS:
64            super().__setattr__(name, value)
65            return
66
67        # When is _data None and __setattr__ called? On any instance built without __init__ (copy/deepcopy/pickle.loads)
68        if name not in self.__dict__.get("_data", {}): # note: self._data bounces to getattr (recursion).
69            raise AttributeError(f"Attribute '{name}' not in query result fields: {self.__dict__.get('fields')}")
70
71        # (!46) Weird optimization trick with [:]: `qr.f += x`: the ufunc already wrote in place; `col[:] = col` skipped
72        if (col := getattr(self, name)) is value:
73            return
74
75        col[:] = value
76
77    def __iter__(self):
78        raise TypeError(("QueryResult is not iterable. Use `qr.field = ..` that applies to all items at once.\n"
79                         "Common pattern: `for e in world.query(..): e.attr = X` -> `qr=world.query(..); qr.attr = X`"))
80
81    def __len__(self):
82        if self._len is None:
83            self._len = sum(len(p) for p in self.pool_list)
84        return self._len
85
86    def __repr__(self):
87        return (f"[QueryResult]\n- Entities: {len(self.entity_ids)} (pools: {len(self.pool_list)})"
88                f"\n- Fields: {self.fields}"
89                f"\n- Shapes: {list(self._field_shapes.values())}\n- Dtypes: {list(self._field_dtypes.values())}")
90
91QUERY_RESULT_RESERVED_NAMES = _QR_INTERNAL_ATTRS | {n for n in vars(QueryResult) if not n.startswith("__")}