From c88812c2f8a2962b70fa371ab789bd14b9737bf1 Mon Sep 17 00:00:00 2001 From: Thanh Lee Date: Wed, 2 Sep 2026 15:07:21 +0000 Subject: [PATCH] fix(amr): apply GTDB species matching on the SQLite index path The AMR genus SQLite index matched --species with an exact LOWER(species) IN (...) clause against the raw query string, so an unsuffixed NCBI name such as "Enterococcus faecium" returned 0 results when an index existed. v0.22.0's GTDB-suffix stripping reached the parquet scan but not the index path. QueryIndex now resolves the requested species against the stored names with match.SpeciesMatches before building the IN clause, mirroring the internal/index resolver. Explicit clade names still select one clade. Closes #49 --- CHANGELOG.md | 10 +++++++ internal/amr/amr_gtdb_test.go | 51 +++++++++++++++++++++++++++++++++++ internal/amr/indexer.go | 38 ++++++++++++++++++++++++++ 3 files changed, 99 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8965f29..8499883 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,16 @@ All notable changes to `atb-cli` are documented here. Format follows [Keep a Cha ## [Unreleased] +### Fixed + +- `atb amr --species "Enterococcus faecium"` (an unsuffixed NCBI name) returned + 0 results when a genus SQLite index was present. The AMR index path matched the + raw query with an exact `LOWER(species) IN (...)` clause and never applied the + GTDB-suffix stripping used everywhere else, so v0.22.0's GTDB matching reached + the parquet scan but not the SQLite index. The index path now resolves the + requested species against the stored names with the same GTDB-aware matcher + before querying. Explicit clade names still select a single clade. + ## [v0.22.0](https://github.com/allthebacteria/atb-cli/releases/tag/v0.22.0) - 2026-09-02 ### Changed diff --git a/internal/amr/amr_gtdb_test.go b/internal/amr/amr_gtdb_test.go index 0f30ad2..c2eb02e 100644 --- a/internal/amr/amr_gtdb_test.go +++ b/internal/amr/amr_gtdb_test.go @@ -72,6 +72,57 @@ func TestQueryAMRSpeciesGTDBSuffix(t *testing.T) { } } +// TestQueryAMRSpeciesGTDBSuffixIndexed verifies that an NCBI-style --species +// query matches GTDB-split rows through the SQLite index path (not just the +// parquet scan). Building indexes makes Query route via QueryIndex/buildSQL. +func TestQueryAMRSpeciesGTDBSuffixIndexed(t *testing.T) { + dir := t.TempDir() + writeAMRParquetAt(t, filepath.Join(dir, amr.PartitionDir, "Enterococcus_A.parquet"), + []amrFixtureRow{{species: "Enterococcus_A faecium", genus: "Enterococcus_A", count: 2}}) + writeAMRParquetAt(t, filepath.Join(dir, amr.PartitionDir, "Enterococcus_B.parquet"), + []amrFixtureRow{{species: "Enterococcus_B faecium", genus: "Enterococcus_B", count: 3}}) + + if err := amr.BuildIndexes(dir, nil); err != nil { + t.Fatalf("BuildIndexes: %v", err) + } + + results, err := amr.Query(dir, amr.Filters{ + Genera: []string{"Enterococcus"}, + Species: []string{"Enterococcus faecium"}, + }) + if err != nil { + t.Fatalf("Query: %v", err) + } + if len(results) != 5 { + t.Fatalf("expected 5 GTDB-split rows via index, got %d", len(results)) + } +} + +// TestQueryAMRSpeciesCladeSelectIndexed verifies that an explicit clade name +// through the SQLite index path selects only that clade. +func TestQueryAMRSpeciesCladeSelectIndexed(t *testing.T) { + dir := t.TempDir() + writeAMRParquetAt(t, filepath.Join(dir, amr.PartitionDir, "Enterococcus_A.parquet"), + []amrFixtureRow{{species: "Enterococcus_A faecium", genus: "Enterococcus_A", count: 2}}) + writeAMRParquetAt(t, filepath.Join(dir, amr.PartitionDir, "Enterococcus_B.parquet"), + []amrFixtureRow{{species: "Enterococcus_B faecium", genus: "Enterococcus_B", count: 3}}) + + if err := amr.BuildIndexes(dir, nil); err != nil { + t.Fatalf("BuildIndexes: %v", err) + } + + results, err := amr.Query(dir, amr.Filters{ + Genera: []string{"Enterococcus"}, + Species: []string{"Enterococcus_B faecium"}, + }) + if err != nil { + t.Fatalf("Query: %v", err) + } + if len(results) != 3 { + t.Fatalf("expected 3 rows for explicit _B clade, got %d", len(results)) + } +} + // TestQueryAMRGenusGTDBPartitions verifies that an NCBI-style genus query reads // every on-disk GTDB clade partition, even when no monolithic file is present. func TestQueryAMRGenusGTDBPartitions(t *testing.T) { diff --git a/internal/amr/indexer.go b/internal/amr/indexer.go index 295a401..b123934 100644 --- a/internal/amr/indexer.go +++ b/internal/amr/indexer.go @@ -300,6 +300,17 @@ func QueryIndex(dbPath string, filters Filters) ([]Result, error) { } } + if len(filters.Species) > 0 { + resolved, err := resolveIndexSpecies(db, filters.Species) + if err != nil { + return nil, wrapStaleIndexError(dbPath, err) + } + if len(resolved) == 0 { + return nil, nil + } + filters.Species = resolved + } + query, args := buildSQL(filters) rows, err := db.Query(query, args...) if err != nil { @@ -329,6 +340,33 @@ func QueryIndex(dbPath string, filters Filters) ([]Result, error) { return results, rows.Err() } +// resolveIndexSpecies reads the distinct species stored in the index and returns +// the concrete values that match any requested query under GTDB-aware matching. +// SQLite cannot evaluate the GTDB suffix regex, so matching happens in Go and the +// resolved names feed buildSQL's exact IN clause. +func resolveIndexSpecies(db *sql.DB, queries []string) ([]string, error) { + rows, err := db.Query("SELECT DISTINCT species FROM amr") + if err != nil { + return nil, err + } + defer rows.Close() + + var resolved []string + for rows.Next() { + var stored string + if err := rows.Scan(&stored); err != nil { + return nil, err + } + for _, q := range queries { + if match.SpeciesMatches(q, stored) { + resolved = append(resolved, stored) + break + } + } + } + return resolved, rows.Err() +} + func buildSQL(f Filters) (string, []any) { var clauses []string var args []any