Skip to content

Commit 263624c

Browse files
authored
Speed up decoding of annotation files in rdann (#627)
`proc_ann_bytes` and its helpers decode an annotation file by indexing the Nx2 NumPy byte array one element at a time (`filebytes[bpi, 1]`, `np.int64(...)`, `.astype("i1")`), and build each aux note with `"".join(chr(c) for c in ...)`. Each of those makes a NumPy scalar, which is slow in a Python loop. This converts the byte pairs to a list once with `.tolist()` and does the rest with plain ints. Signed `subtype`/`num` bytes are converted with a small helper, and aux notes are decoded with `bytes(...).decode("latin-1")`, which gives the same characters as `chr()` for 0-255. Since no NumPy scalars are involved any more, the decoding no longer depends on NumPy's scalar type promotion rules, which caused the `OverflowError` in #493 and #522 (already worked around on main with `np.int64`). The output is unchanged: I compared every field (`sample`, `symbol`, `subtype`, `chan`, `num`, `aux_note`, `fs`, `label_store`, custom labels, and dtypes) for all 15 annotation files in `sample-data`, with and without `sampto`, against main. Benchmark (median, Apple Silicon, alternating runs of main and this branch; the machine was busy, so the absolute numbers are noisy): | file | main | this PR | speedup | | --- | --- | --- | --- | | `rdann` 100.atr | 4.10 ms | 2.01 ms | 2.0x | | `rdann` 1003.atr | 6.24 ms | 3.31 ms | 1.9x | | `rdann` 12726.wabp (many aux notes) | 16.5 ms | 6.4 ms | 2.6x | | `rdann` 03700181.gqrsl | 2.99 ms | 1.64 ms | 1.8x | | `rdann`, 100,000 beats with some aux notes (written with `wrann`) | 176 ms | 76 ms | 2.3x | | `proc_ann_bytes` only, 12726.wabp | 22.4 ms | 7.6 ms | 2.9x | Tests: added a round-trip test for `subtype`, `chan` and `num` values (including negative subtypes). `uv run pytest -m "not network"` passes, `tests/test_annotation.py` passes with network tests, `ruff format`/`ruff check` are clean, and the docs build without warnings. Added a changelog entry. Open PRs #491 and #498 also touch `annotation.py`, but only how the file is opened and read into bytes, not the decoding, so they shouldn't conflict beyond nearby lines.
2 parents 21d0eeb + 24d3690 commit 263624c

3 files changed

Lines changed: 58 additions & 28 deletions

File tree

‎changelog.d/rdann-speed.rst‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
**Speed up reading annotation files**
2+
``rdann`` now decodes annotation files using plain Python integers instead of indexing a NumPy array one byte at a time. Reading annotations is about 2 to 3 times faster (a 100,000-beat file now takes about 75 ms instead of 175). The annotations read are unchanged.

‎src/wfdb/io/annotation.py‎

Lines changed: 35 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -2209,6 +2209,8 @@ def proc_ann_bytes(filebytes, sampto):
22092209
# Total number of sample from beginning of record. Annotation bytes
22102210
# only store sample_diff
22112211
sample_total = 0
2212+
# Python ints are much faster to index one at a time than NumPy scalars.
2213+
pairs = filebytes.tolist()
22122214
# Byte pair index
22132215
bpi = 0
22142216

@@ -2218,9 +2220,9 @@ def proc_ann_bytes(filebytes, sampto):
22182220
# - samp + sym pair
22192221
# - other pairs (if any)
22202222
# The last byte pair of the file is 0 indicating eof.
2221-
while bpi < filebytes.shape[0] - 1:
2223+
while bpi < len(pairs) - 1:
22222224
# Get the sample and label_store fields of the current annotation
2223-
sample_diff, current_label_store, bpi = proc_core_fields(filebytes, bpi)
2225+
sample_diff, current_label_store, bpi = proc_core_fields(pairs, bpi)
22242226
sample_total = sample_total + sample_diff
22252227
sample.append(sample_total)
22262228
label_store.append(current_label_store)
@@ -2231,12 +2233,12 @@ def proc_ann_bytes(filebytes, sampto):
22312233
update = {"subtype": True, "chan": True, "num": True, "aux_note": True}
22322234
# Get the next label store value - it may indicate additional
22332235
# fields for this annotation, or the values of the next annotation.
2234-
current_label_store = int(filebytes[bpi, 1]) >> 2
2236+
current_label_store = pairs[bpi][1] >> 2
22352237

22362238
while current_label_store > 59:
22372239
subtype, chan, num, aux_note, update, bpi = proc_extra_field(
22382240
current_label_store,
2239-
filebytes,
2241+
pairs,
22402242
bpi,
22412243
subtype,
22422244
chan,
@@ -2245,7 +2247,7 @@ def proc_ann_bytes(filebytes, sampto):
22452247
update,
22462248
)
22472249

2248-
current_label_store = int(filebytes[bpi, 1]) >> 2
2250+
current_label_store = pairs[bpi][1] >> 2
22492251

22502252
# Set defaults or carry over previous values if necessary
22512253
subtype, chan, num, aux_note = update_extra_fields(
@@ -2261,13 +2263,18 @@ def proc_ann_bytes(filebytes, sampto):
22612263
return sample, label_store, subtype, chan, num, aux_note
22622264

22632265

2266+
def _to_signed_char(value):
2267+
"""Interpret a byte value (0-255) as a signed char."""
2268+
return value - 256 if value > 127 else value
2269+
2270+
22642271
def proc_core_fields(filebytes, bpi):
22652272
"""Get the sample difference and store fields of the current annotation.
22662273
22672274
Parameters
22682275
----------
2269-
filebytes : ndarray
2270-
The input filestream converted to an Nx2 array of unsigned bytes.
2276+
filebytes : list
2277+
The input filestream as a list of N pairs of unsigned bytes.
22712278
bpi : int
22722279
The index to start the conversion.
22732280
@@ -2287,13 +2294,13 @@ def proc_core_fields(filebytes, bpi):
22872294

22882295
# The current byte pair will contain either the actual d_sample + annotation store
22892296
# value, or 0 + SKIP.
2290-
while int(filebytes[bpi, 1]) >> 2 == 59:
2297+
while filebytes[bpi][1] >> 2 == 59:
22912298
# 4 bytes storing dt
22922299
skip_diff = (
2293-
(int(filebytes[bpi + 1, 0]) << 16)
2294-
+ (int(filebytes[bpi + 1, 1]) << 24)
2295-
+ (int(filebytes[bpi + 2, 0]) << 0)
2296-
+ (int(filebytes[bpi + 2, 1]) << 8)
2300+
(filebytes[bpi + 1][0] << 16)
2301+
+ (filebytes[bpi + 1][1] << 24)
2302+
+ filebytes[bpi + 2][0]
2303+
+ (filebytes[bpi + 2][1] << 8)
22972304
)
22982305

22992306
# Data type is long integer (stored in two's complement).
@@ -2305,8 +2312,9 @@ def proc_core_fields(filebytes, bpi):
23052312
bpi = bpi + 3
23062313

23072314
# Not a skip - it is the actual sample number + annotation type store value
2308-
label_store = int(filebytes[bpi, 1]) >> 2
2309-
sample_diff += np.int64(filebytes[bpi, 0]) + 256 * (np.int64(filebytes[bpi, 1]) & 3)
2315+
byte0, byte1 = filebytes[bpi]
2316+
label_store = byte1 >> 2
2317+
sample_diff += byte0 + 256 * (byte1 & 3)
23102318
bpi = bpi + 1
23112319

23122320
return sample_diff, label_store, bpi
@@ -2323,8 +2331,8 @@ def proc_extra_field(label_store, filebytes, bpi, subtype, chan, num, aux_note,
23232331
A numpy array containing the integer values used to store the
23242332
annotation labels. If this field is present, `symbol` must not be
23252333
present.
2326-
filebytes : str
2327-
The input filestream converted to bytes.
2334+
filebytes : list
2335+
The input filestream as a list of N pairs of unsigned bytes.
23282336
bpi : int
23292337
The index to start the conversion.
23302338
subtype : ndarray
@@ -2366,35 +2374,34 @@ def proc_extra_field(label_store, filebytes, bpi, subtype, chan, num, aux_note,
23662374
# SUB
23672375
if label_store == 61:
23682376
# sub is interpreted as signed char.
2369-
subtype.append(filebytes[bpi, 0].astype("i1"))
2377+
subtype.append(_to_signed_char(filebytes[bpi][0]))
23702378
update["subtype"] = False
23712379
bpi = bpi + 1
23722380
# CHAN
23732381
elif label_store == 62:
23742382
# chan is interpreted as un_signed char
2375-
chan.append(filebytes[bpi, 0])
2383+
chan.append(filebytes[bpi][0])
23762384
update["chan"] = False
23772385
bpi = bpi + 1
23782386
# NUM
23792387
elif label_store == 60:
23802388
# num is interpreted as signed char
2381-
num.append(filebytes[bpi, 0].astype("i1"))
2389+
num.append(_to_signed_char(filebytes[bpi][0]))
23822390
update["num"] = False
23832391
bpi = bpi + 1
23842392
# aux_note
23852393
elif label_store == 63:
23862394
# length of aux_note string. Max 256? No need to check other bits of
23872395
# second byte?
2388-
aux_notelen = filebytes[bpi, 0]
2389-
aux_notebytes = filebytes[
2390-
bpi + 1 : bpi + 1 + int(np.ceil(aux_notelen / 2.0)), :
2391-
].flatten()
2392-
if int(aux_notelen) & 1:
2393-
aux_notebytes = aux_notebytes[:-1]
2394-
# The aux_note string
2395-
aux_note.append("".join([chr(char) for char in aux_notebytes]))
2396+
aux_notelen = filebytes[bpi][0]
2397+
npairs = (aux_notelen + 1) // 2
2398+
aux_notebytes = bytes(
2399+
b for pair in filebytes[bpi + 1 : bpi + 1 + npairs] for b in pair
2400+
)
2401+
# The aux_note string, one character per byte
2402+
aux_note.append(aux_notebytes[:aux_notelen].decode("latin-1"))
23962403
update["aux_note"] = False
2397-
bpi = bpi + 1 + int(np.ceil(aux_notelen / 2.0))
2404+
bpi = bpi + 1 + npairs
23982405

23992406
return subtype, chan, num, aux_note, update, bpi
24002407

‎tests/test_annotation.py‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -423,3 +423,24 @@ def test_rdann_path():
423423
assert annotation.record_name == "100"
424424
np.testing.assert_array_equal(annotation.sample, expected.sample)
425425
assert annotation.symbol == expected.symbol
426+
427+
428+
def test_subtype_chan_num_round_trip(temp_path):
429+
"""Subtype (signed), chan and num values are read back unchanged."""
430+
subtype = np.array([-128, -1, 0, 127])
431+
chan = np.array([0, 255, 255, 3])
432+
num = np.array([5, 5, 100, 127])
433+
wfdb.wrann(
434+
"fields",
435+
"atr",
436+
np.array([10, 20, 30, 40]),
437+
symbol=["N", "V", "N", "A"],
438+
subtype=subtype,
439+
chan=chan,
440+
num=num,
441+
write_dir=temp_path,
442+
)
443+
ann = wfdb.rdann(os.path.join(temp_path, "fields"), "atr")
444+
assert ann.subtype.tolist() == subtype.tolist()
445+
assert ann.chan.tolist() == chan.tolist()
446+
assert ann.num.tolist() == num.tolist()

0 commit comments

Comments
 (0)