Skip to content

fix: off-by-one in fetch_items_iter drops one row from every transform run - #54

Open
vasa-develop wants to merge 2 commits into
LeMaterial:mainfrom
vasa-develop:fix/fetch-items-iter-off-by-one
Open

fix: off-by-one in fetch_items_iter drops one row from every transform run#54
vasa-develop wants to merge 2 commits into
LeMaterial:mainfrom
vasa-develop:fix/fetch-items-iter-off-by-one

Conversation

@vasa-develop

Copy link
Copy Markdown

Summary

StructuresDatabase.fetch_items_iter has an issue with its id-based pagination logic.
When you ask fetch_items_iter to start reading from some offset, it first looks up which row sits at
that offset (using get_id_at_offset, which counts from zero), and then it fetches
rows using the condition WHERE id > start_id. The problem is the >: the row
that sits exactly at the offset gets excluded, because the query only takes rows that
come after it.

Here is what that looks like in practice. Take a table with 11 rows, ids A through
K, and read it in batches of 3 the way the transformers do (this is real output from
running the current code against Postgres):

batch at offset 0: ['A', 'B', 'C']
batch at offset 3: ['E', 'F', 'G']   <- D is lost here
batch at offset 6: ['H', 'I', 'J']
batch at offset 9: ['K']
missing: ['D']

At offset 3, the code looks up the row at position 3 (that's D) and then fetches
WHERE id > 'D', so it starts at E and nobody ever comes back for D. At offset 6
the same exclusion happens to G, but G was already returned by the previous batch
(which had shifted forward by one), so from that point on the off-by-ones cancel out.
The net result is that exactly one row goes missing per read, always the one sitting
at position batch_size, and it doesn't matter how many rows the table has.

This isn't specific to any one data source. The batch loop in BaseTransformer (which
the MP and Alexandria transforms run through) and OQMD's custom transform all paginate
through this method. So as far as I can tell, every transform run to date has quietly
lost one raw structure on its way to the clean table. Since the source tables hold
millions of rows, one missing row was effectively invisible in production. Still, it's
silent data loss in a pipeline whose whole job is faithful ingestion, and there is no
warning anywhere in the logs when it happens.

The fix itself is one character: change the scan to WHERE id >= start_id, so the row
at the offset is included.

How I found it

I'm working on an AFLOW source adapter (PR coming separately), and while validating it
end-to-end against a local Postgres, fetch pulled 180 raw rows but transform produced
179 clean rows, with no skip warning in the logs. The missing row transformed fine on
its own, so the row wasn't the problem. What stood out was where it sat: at position
exactly 50 in the table's id ordering, which happened to be my --batch-size.

To rule out a coincidence, I emptied the destination table and re-ran the same
transform with --batch-size 30. A different row went missing, this time at position
exactly 30. The missing row moves with the batch size, which confirmed it: the row at
position batch_size is always dropped. At production scale (millions of rows) one
missing row is unnoticeable, which is probably why this survived so long.

What's in this PR

  • The one-character fix in fetch_items_iter, plus a short comment in the code
    explaining why the scan has to be inclusive.
  • tests/database/test_postgres.py: regression tests that run against a real
    PostgreSQL server. The most important test builds a small 10-row table and reads it
    in consecutive (offset, batch_size) windows, exactly the way BaseTransformer
    does, and then checks that every row came back exactly once. Before the fix, this
    test fails with the row at position batch_size missing. After the fix, it passes.
    There are also more direct tests checking that fetch_items_iter(offset=k) and
    fetch_items(offset=k) really start at position k, and tests for the edges
    (offset 0, and an offset past the end of the table).
  • A new postgres pytest marker, registered in pyproject.toml. If no PostgreSQL
    server is reachable, these tests skip themselves automatically, so a plain pytest
    run stays green on machines without Postgres. The docstring at the top of the test
    file documents the one-line docker command to run them locally.
  • A separate commit that adds a postgres:16 service to the CI workflow, so that
    these tests actually execute in CI instead of skipping. I kept it as its own commit
    on purpose, but if you would rather not touch the workflow, just drop that commit and the tests will simply skip in CI.

The reason why these new tests need a real Postgres is that the bug lives inside the SQL string itself: a mocked cursor just returns whatever it's told to, so it never really tests the pagination logic (which is why the existing test suite never caught this). Only a real Postgres can actually exercise the pagination behavior.

How to verify the fix

  • The new regression tests fail on main and pass with the fix (5 passed).
  • The full test suite passes after the fix, and ruff check and
    ruff format --check are both clean.
  • I re-ran the AFLOW end-to-end pipeline that surfaced the bug in the first place:
    the transform now produces 180 clean rows from 180 raw rows.

fetch_items_iter resolves the id at the requested offset via
get_id_at_offset and then scanned with WHERE id > start_id, which
excludes the boundary row itself. As a result, every multi-batch read
(BaseTransformer's batch loop, OQMD's custom transform) silently dropped
exactly one row: the one at rank batch_size in the source table's id
ordering.

Fix by making the scan inclusive (WHERE id >= start_id), and add
Postgres-backed regression tests under a new 'postgres' pytest marker,
since the bug lives in the SQL itself and cannot be caught by a mocked
cursor. The tests skip automatically when no server is reachable.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant