feat(query): match GTDB taxonomy for --species and --genus

CHANGELOG.md

@@ -4,6 +4,19 @@ All notable changes to `atb-cli` are documented here. Format follows [Keep a Cha
 
 ## [Unreleased]
 
+### Changed
+
+- `--species` and `--genus` now match GTDB taxonomy. ATB stores GTDB names,
+  which split NCBI genera and species into alphabetic clades
+  (`Enterococcus_A faecium`, `Enterococcus_B faecium`,
+  `Campylobacter jejunii_A`), so an NCBI-style query such as
+  `atb mlst --species "Enterococcus faecium"` previously returned nothing. The
+  GTDB suffix is now stripped from the stored name during comparison, so an
+  unsuffixed query matches every clade while an explicit clade name
+  (`Campylobacter_D`) still selects only that clade. This applies to `query`,
+  `mlst`, and `amr` on every backend (SQLite index, parquet scan, and AMR genus
+  partitions).
+
 ### Fixed
 
 - `atb <unknown-command>` no longer printed `Run 'atb --help' for usage.`, a

internal/amr/amr.go

@@ -87,6 +87,12 @@ func Query(dataDir string, filters Filters) ([]Result, error) {
 		return nil, err
 	}
 
+	// Expand each requested genus to the GTDB clade partitions that hold its
+	// rows (e.g. "Enterococcus" -> "Enterococcus_A", "Enterococcus_B") so a
+	// split genus is read from every clade file rather than falling back to a
+	// full monolithic scan.
+	filters.Genera = expandGeneraToPartitions(dataDir, filters.Genera)
+
 	// Try SQLite indexes first for each genus.
 	if len(filters.Genera) > 0 {
 		return queryWithIndexes(dataDir, filters)
@@ -351,9 +357,13 @@ func matchesFilters(row pq.AMRRow, f Filters) bool {
 	return true
 }
 
+// matchesAny reports whether value matches any candidate taxon name, treating
+// GTDB alphabetic suffixes as equivalent so an NCBI-style query
+// ("Enterococcus faecium") matches the GTDB-split names stored in the data
+// ("Enterococcus_A faecium").
 func matchesAny(value string, candidates []string) bool {
 	for _, c := range candidates {
-		if strings.EqualFold(value, c) {
+		if match.SpeciesMatches(c, value) {
 			return true
 		}
 	}

internal/amr/amr_gtdb_test.go

@@ -0,0 +1,93 @@
+package amr_test
+
+import (
+	"fmt"
+	"os"
+	"path/filepath"
+	"testing"
+
+	parquetgo "github.com/parquet-go/parquet-go"
+
+	"github.com/allthebacteria/atb-cli/internal/amr"
+	pq "github.com/allthebacteria/atb-cli/internal/parquet"
+)
+
+// writeAMRParquetAt writes the fixture rows to an explicit parquet path.
+func writeAMRParquetAt(t *testing.T, path string, rows []amrFixtureRow) {
+	t.Helper()
+	if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
+		t.Fatalf("mkdir: %v", err)
+	}
+	f, err := os.Create(path)
+	if err != nil {
+		t.Fatalf("create parquet: %v", err)
+	}
+	defer f.Close()
+
+	w := parquetgo.NewGenericWriter[pq.AMRRow](f)
+	var idx int
+	for _, r := range rows {
+		for i := 0; i < r.count; i++ {
+			row := pq.AMRRow{
+				Name:        fmt.Sprintf("SAMN%08d", idx),
+				GeneSymbol:  fmt.Sprintf("gene_%d", idx),
+				ElementType: "AMR",
+				Coverage:    100,
+				Identity:    100,
+				Method:      "EXACT",
+				Class:       "BETA-LACTAM",
+				Species:     r.species,
+				Genus:       r.genus,
+			}
+			if _, err := w.Write([]pq.AMRRow{row}); err != nil {
+				t.Fatalf("write row: %v", err)
+			}
+			idx++
+		}
+	}
+	if err := w.Close(); err != nil {
+		t.Fatalf("close writer: %v", err)
+	}
+}
+
+// TestQueryAMRSpeciesGTDBSuffix verifies that an NCBI-style --species query
+// matches GTDB-split rows on the monolithic scan path.
+func TestQueryAMRSpeciesGTDBSuffix(t *testing.T) {
+	dir := t.TempDir()
+	writeAMRFixture(t, dir, []amrFixtureRow{
+		{species: "Enterococcus_A faecium", genus: "Enterococcus_A", count: 2},
+		{species: "Enterococcus_B faecium", genus: "Enterococcus_B", count: 3},
+		{species: "Escherichia coli", genus: "Escherichia", count: 4},
+	})
+
+	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, 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) {
+	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}})
+
+	results, err := amr.Query(dir, amr.Filters{
+		Genera: []string{"Enterococcus"},
+	})
+	if err != nil {
+		t.Fatalf("Query: %v", err)
+	}
+	if len(results) != 5 {
+		t.Fatalf("expected 5 rows across both clade partitions, got %d", len(results))
+	}
+}

internal/amr/partition.go

@@ -11,6 +11,7 @@ import (
 
 	parquetgo "github.com/parquet-go/parquet-go"
 
+	"github.com/allthebacteria/atb-cli/internal/match"
 	pq "github.com/allthebacteria/atb-cli/internal/parquet"
 )
 
@@ -135,6 +136,60 @@ func BuildPartitions(dataDir string, logFn func(string, ...any)) error {
 	return nil
 }
 
+// expandGeneraToPartitions replaces each requested genus with the on-disk
+// partition genera whose canonical (GTDB-suffix-stripped) name matches it, so
+// an NCBI-style genus such as "Enterococcus" resolves to the clade partitions
+// "Enterococcus_A" and "Enterococcus_B". A genus with no matching partition is
+// kept unchanged so the caller can fall back to the monolithic scan. When no
+// partition directory exists, the input is returned as-is.
+func expandGeneraToPartitions(dataDir string, genera []string) []string {
+	if len(genera) == 0 {
+		return genera
+	}
+
+	entries, err := os.ReadDir(filepath.Join(dataDir, PartitionDir))
+	if err != nil {
+		return genera
+	}
+	var partitions []string
+	for _, e := range entries {
+		if e.IsDir() || !strings.HasSuffix(e.Name(), ".parquet") {
+			continue
+		}
+		name := strings.TrimSuffix(e.Name(), ".parquet")
+		if name == otherPartition {
+			continue
+		}
+		partitions = append(partitions, name)
+	}
+
+	seen := make(map[string]struct{}, len(genera))
+	var out []string
+	add := func(g string) {
+		if _, ok := seen[g]; ok {
+			return
+		}
+		seen[g] = struct{}{}
+		out = append(out, g)
+	}
+	for _, g := range genera {
+		var matched []string
+		for _, p := range partitions {
+			if match.SpeciesMatches(g, p) {
+				matched = append(matched, p)
+			}
+		}
+		if len(matched) > 0 {
+			for _, m := range matched {
+				add(m)
+			}
+			continue
+		}
+		add(g)
+	}
+	return out
+}
+
 // PartitionPath returns the path to a genus partition file if it exists.
 // Returns empty string if the partition doesn't exist. Lookup is case-
 // insensitive so that GTDB letter clades (e.g. Legionella_C) match files

internal/cli/amr_cmd.go

@@ -188,26 +188,18 @@ Run 'atb fetch' to download the data before querying.`,
 			if hqOnly {
 				fmt.Fprintf(os.Stderr, "Loading HQ sample set...\n")
 				assemblyPath := filepath.Join(dir, "assembly.parquet")
-				lowerGenera := make(map[string]bool, len(genera))
-				for _, g := range genera {
-					lowerGenera[strings.ToLower(g)] = true
-				}
-				lowerSpecies := make(map[string]bool, len(speciesList))
-				for _, s := range speciesList {
-					lowerSpecies[strings.ToLower(s)] = true
-				}
 				hqRows, hqErr := pq.ReadStreamFiltered[pq.AssemblyRow](assemblyPath, func(r pq.AssemblyRow) bool {
 					if r.HQFilter != "PASS" {
 						return false
 					}
 					if speciesLike != "" && !match.Like(r.SylphSpecies, speciesLike) {
 						return false
 					}
-					if len(lowerSpecies) > 0 {
-						return lowerSpecies[strings.ToLower(r.SylphSpecies)]
+					if len(speciesList) > 0 {
+						return matchesAnySpecies(speciesList, r.SylphSpecies)
 					}
-					if len(lowerGenera) > 0 {
-						return lowerGenera[strings.ToLower(pq.GenusFromSpecies(r.SylphSpecies))]
+					if len(genera) > 0 {
+						return matchesAnySpecies(genera, pq.GenusFromSpecies(r.SylphSpecies))
 					}
 					return true
 				}, 0)
@@ -489,6 +481,18 @@ func intersectSampleSets(a, b map[string]struct{}) map[string]struct{} {
 	return out
 }
 
+// matchesAnySpecies reports whether the stored GTDB species or genus matches
+// any of the user-supplied queries, treating GTDB alphabetic suffixes as
+// equivalent (see match.SpeciesMatches).
+func matchesAnySpecies(queries []string, stored string) bool {
+	for _, q := range queries {
+		if match.SpeciesMatches(q, stored) {
+			return true
+		}
+	}
+	return false
+}
+
 // amrColumns returns the fixed column order for AMR output. Headers match the
 // AMRFinderPlus v4.2.5 TSV verbatim so downstream tooling sees the same names
 // regardless of source. When withENA is true, country/collection_date/

internal/index/query.go

@@ -11,6 +11,7 @@ import (
 	_ "modernc.org/sqlite"
 
 	"github.com/allthebacteria/atb-cli/internal/match"
+	pq "github.com/allthebacteria/atb-cli/internal/parquet"
 )
 
 // DB wraps a read-only SQLite connection to the index.
@@ -160,22 +161,76 @@ func (d *DB) MLSTForSample(sampleAccession string) (map[string]string, error) {
 	return result, err
 }
 
+// speciesMatching returns the distinct stored sylph_species values for which
+// keep reports true. It is used to resolve a species or genus filter to the
+// concrete GTDB-suffixed names present in the index, so matching can strip
+// GTDB suffixes in Go where SQLite has no regex support.
+func (d *DB) speciesMatching(keep func(stored string) bool) ([]string, error) {
+	rows, err := d.db.Query("SELECT DISTINCT sylph_species FROM samples")
+	if err != nil {
+		return nil, fmt.Errorf("distinct species query: %w", err)
+	}
+	defer rows.Close()
+
+	var names []string
+	for rows.Next() {
+		var s string
+		if err := rows.Scan(&s); err != nil {
+			return nil, fmt.Errorf("scanning species: %w", err)
+		}
+		if keep(s) {
+			names = append(names, s)
+		}
+	}
+	return names, rows.Err()
+}
+
+// inClause builds a "column IN (?, ?, ...)" condition and its arguments for the
+// given values. An empty list yields a condition that matches no rows, so a
+// filter that resolves to nothing returns zero results rather than everything.
+func inClause(column string, values []string) (string, []any) {
+	if len(values) == 0 {
+		return "1 = 0", nil
+	}
+	placeholders := make([]string, len(values))
+	args := make([]any, len(values))
+	for i, v := range values {
+		placeholders[i] = "?"
+		args[i] = v
+	}
+	return fmt.Sprintf("%s IN (%s)", column, strings.Join(placeholders, ", ")), args
+}
+
 // Query runs a filtered query returning result rows as map[string]string.
 func (d *DB) Query(params QueryParams) ([]map[string]string, error) {
 	var conditions []string
 	var args []any
 
 	if params.Species != "" {
-		conditions = append(conditions, "lower(sylph_species) = lower(?)")
-		args = append(args, params.Species)
+		names, err := d.speciesMatching(func(stored string) bool {
+			return match.SpeciesMatches(params.Species, stored)
+		})
+		if err != nil {
+			return nil, err
+		}
+		cond, inArgs := inClause("sylph_species", names)
+		conditions = append(conditions, cond)
+		args = append(args, inArgs...)
 	}
 	if params.SpeciesLike != "" {
 		conditions = append(conditions, "lower(sylph_species) LIKE ? ESCAPE '\\'")
 		args = append(args, match.ToSQLLike(params.SpeciesLike))
 	}
 	if params.Genus != "" {
-		conditions = append(conditions, "lower(substr(sylph_species, 1, instr(sylph_species, ' ') - 1)) = lower(?)")
-		args = append(args, params.Genus)
+		names, err := d.speciesMatching(func(stored string) bool {
+			return match.SpeciesMatches(params.Genus, pq.GenusFromSpecies(stored))
+		})
+		if err != nil {
+			return nil, err
+		}
+		cond, inArgs := inClause("sylph_species", names)
+		conditions = append(conditions, cond)
+		args = append(args, inArgs...)
 	}
 	if params.HQOnly {
 		conditions = append(conditions, "hq_filter = 'PASS'")