diff --git a/aborist/sources/wikipedia.py b/aborist/sources/wikipedia.py index b202fd0..9721fa1 100644 --- a/aborist/sources/wikipedia.py +++ b/aborist/sources/wikipedia.py @@ -142,23 +142,33 @@ def _split_values_tuples(payload: str) -> list[list[str | None]]: def _iter_insert_statements(file_obj: IO[str]) -> Iterator[str]: """Yield complete SQL statements using buffered string find. - Reads the file in 512 KB chunks, accumulates pending text, and uses - str.find for ';' and "'" to scan. String state crosses chunk boundaries - via the `in_string` flag. Avoids the per-char Python loop that dominated - the hot path before. + State persists across read() chunks via three indices: + yield_from position in `pending` where the next yielded statement starts + scan_pos position to resume the scanner from (do NOT restart at 0 + across chunks — would re-find the already-processed open `'`) + in_string are we inside a single-quoted string + + The previous version reset the scanner to position 0 on every chunk, + which silently flipped in_string=False as soon as the first `'` in the + new buffer was re-encountered (the title-field open quote of the + article we were already deep inside). Persisting scan_pos fixes it. """ pending = "" in_string = False + yield_from = 0 # start of current statement, in `pending` coords + scan_pos = 0 # next position to scan, in `pending` coords - def flush_complete(): - """Yield as many complete statements as exist in `pending`.""" - nonlocal pending, in_string - i = 0 - last_yield = 0 + while True: + chunk = file_obj.read(1 << 19) + if not chunk: + break + pending += chunk n = len(pending) + + i = scan_pos while i < n: if in_string: - # Find string terminator (handle escapes). + # Find string terminator (handle escape parity). end = i while True: pos = pending.find("'", end) @@ -166,8 +176,10 @@ def _iter_insert_statements(file_obj: IO[str]) -> Iterator[str]: i = n # need more data break bs = 0 + # Don't walk back past the start of the current statement; + # chars before yield_from belong to a previous statement. p = pos - 1 - while p >= 0 and pending[p] == "\\": + while p >= yield_from and pending[p] == "\\": bs += 1 p -= 1 if bs % 2 == 0: @@ -184,23 +196,23 @@ def _iter_insert_statements(file_obj: IO[str]) -> Iterator[str]: i = n break if qt < 0 or (sc >= 0 and sc < qt): - yield pending[last_yield : sc + 1] - last_yield = sc + 1 + yield pending[yield_from : sc + 1] + yield_from = sc + 1 i = sc + 1 else: in_string = True i = qt + 1 - if last_yield > 0: - pending = pending[last_yield:] - while True: - chunk = file_obj.read(1 << 19) - if not chunk: - break - pending += chunk - yield from flush_complete() - if pending.strip(): - yield pending + scan_pos = i + # Release yielded bytes; remap our two indices into the new buffer. + if yield_from > 0: + pending = pending[yield_from:] + scan_pos -= yield_from + yield_from = 0 + + tail = pending[yield_from:] + if tail.strip(): + yield tail def _extract_wikilinks(text: str, base_uri: str) -> list[Edge]: