Don't yield in jitted function

This might have caused some caching issues.
otsaloma · Jun 14, 2023 · c7ea8a9 · c7ea8a9
1 parent 908b303
commit c7ea8a9
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/NEWS.md b/NEWS.md
@@ -3,6 +3,7 @@ PENDING: Dataiter 0.45
 
 * `USE_NUMBA_CACHE`: New option, read from environment variable
   `DATAITER_USE_NUMBA_CACHE` if exists, defauls to `True`
+* Fix a possible issue with Numba caching
 
 2023-06-13: Dataiter 0.44
 =========================

diff --git a/dataiter/aggregate.py b/dataiter/aggregate.py
@@ -659,10 +659,12 @@ def yield_groups_numba(x, group, drop_na):
     # Groups must be contiguous for this to work!
     i = 0
     n = len(x)
+    out = []
     for j in range(1, n + 1):
         if j < n and group[j] == group[i]: continue
         xij = x[i:j]
         if drop_na:
             xij = xij[~is_na_numba(xij)]
-        yield xij
+        out.append(xij)
         i = j
+    return out