Apply filters and group-by logic to a metadata iterator

[huddlej]

Description of proposed changes

This PR rewrites the logic for filtering records to avoid loading all metadata into memory at once. To achieve this goal, we filter metadata in one or two passes through the full data, iterating through a fixed number of chunks in memory at any given time. In the best case scenario (filtering only or grouping by columns with predefined number of sequences per group), we only need one pass through the metadata. During this pass, we can stream records that pass all filters to disk or store the highest priority records in a (fixed size) priority queue to be written later.

In the worst case scenario (grouping by columns with a maximum number of total sequences), we need the first pass to calculate the number of sequences per group and a second pass to select the highest priority sequences in each group.

In addition to this major change in how the filter logic works, this PR also makes the following changes:

• Adds an --output-log argument that accepts a tab-delimited file that will contain each filtered strain and its reason for being filtered (or force-included)
• Adds a --metadata-chunksize argument to control the number of metadata records loaded into memory at a time. This argument allows users to control memory use by Augur to match their compute environment.
• Defaults to non-probabilistic subsampling when possible and only switches to probabilistic sampling with --probabilistic-sampling is True and the non-probabilistic number of sequences per group cannot be calculated.
• Adds support for indexing VCF files with augur index and from the Augur Python API with a new index_vcf function. These interfaces allow VCFs to be treated the same as FASTA files when filtering by sequence presence/absence.
• Removes a check for strains actually present in input sequences (when provided) and relies entirely on the sequence index to filter metadata and strain outputs (this is a breaking change from previous behavior but only affects the use case when a provided sequence index is out of date relative to the provided sequences).
• Removes ordering of strain names prior to output via --output-strains.

Related issue(s)

Fixes #699
Fixes #424

TODO

• Support subsampling by priority
• Fix pandas read_csv iteration in Python 3.6
• Support probabilistic sampling using subsample seed argument
• Log strains filtered by subsampling
• Report exclude/include files separately in the log output
• Rename include_by_query to include_by_include_where
• Rewrite heap logic as a class to support counter/random number generator state

Maybe for later PRs:

• Clean up and refactor code in run that could be in its own functions (e.g., setup grouping, write metadata, etc.)
• Standardize naming of "strains", "sequences", and "records"

Steps for later PRs:

• Throw an error when any requested filter or group-by columns don't exist #754
• Replace construct_filters with factory functions that generate filters and integrate with argparse types

Testing

• Tested by CI
• Tests for probabilistic sampling
• Tests for subsampling by priority
• Test with full ncov metadata
• Test with full Nextstrain build

[codecov]

Codecov Report

Merging #750 (bc076d4) into master (653079e) will increase coverage by 1.93%.
The diff coverage is 58.52%.

@@            Coverage Diff             @@
##           master     #750      +/-   ##
==========================================
+ Coverage   31.52%   33.46%   +1.93%     
==========================================
  Files          41       41              
  Lines        5754     5893     +139     
  Branches     1389     1431      +42     
==========================================
+ Hits         1814     1972     +158     
+ Misses       3865     3837      -28     
- Partials       75       84       +9     

Impacted Files	Coverage Δ
augur/index.py	64.86% <10.00%> (+12.23%)	⬆️
augur/filter.py	67.72% <58.63%> (+15.64%)	⬆️
augur/utils.py	41.59% <91.66%> (+1.41%)	⬆️
augur/io.py	100.00% <100.00%> (ø)
augur/export_v2.py	11.04% <0.00%> (-0.04%)	⬇️

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 653079e...bc076d4. Read the comment docs.

[huddlej]

@tsibley This PR supersedes #749 but adds the remainder of the work required to fully reduce memory usage. I tried to group commits into logical chunks (rebasing), but a couple are unavoidably large. Nearly all commits could be squashed after review, but I left a couple unsquashed that might be contentious or at least worth discussing separately.

It might be easiest to walk through this in person like we did with previous PRs, but if you prefer to scan it on your own that's cool, too.

I'll also post some benchmarking results here later when I've finalized those.

[huddlej]

Some thoughts from chatting with Tom today:

• The current implementation of two passes through metadata assumes the input can be re-read, but this will not be true for streaming input. To support streaming input, we should stream the metadata to a temporary local file in the first pass and read from this file in the second pass.
  • Since we don't immediately need to support streaming inputs, we can keep the current implementation and add that support in a later PR.
  • The sanitize metadata script in the ncov workflow will need to use a similar spool-to-disk approach because it transforms and deduplicates records. I will test out a generic implementation of the required logic in that script and then port it into augur filter when we've solidified the implementation in ncov. (This has been our pattern for most new functionality in augur lately anyway.)
• We should consider what interface (i.e., argument/flag names) we want for Augur and then map these names to underlying technology like pandas instead of echoing the names used by pandas. This gives us flexibility to swap out pandas for something else eventually (at least, in some ways; we still refer to .loc, .index.values, etc. everywhere) and also allows us to name the interface in a way that matches our users' preferred nomenclature. For example, instead of referring to "metadata index columns" we could refer to "metadata unique id columns".
• We do not need to add support for indexing VCFs right now, so the current approach in this PR is more than what's necessary to address the original memory issue. However, we should discuss as a group whether we want to support a "sequence index" for VCFs (something that is technically possible, given the original reference FASTA) because current users would benefit from it or we might in the future (e.g., if we move the ncov workflow to VCF from FASTA MSAs someday).

[tsibley]

New changes look better, I agree. I had one small comment about an error message, but instead of leaving a comment, I pushed a commit (0be8869) with my proposal. I also pushed a tangential change (562b0e4) that I noticed while doing the error message work.

I'll also note what I already noted in Slack: One thing that jumped out is that I'd have broken e975fef up into three commits to make each distinct change smaller. One tell I use when thinking about how to chunk up commits is if I'm tempted to use "also" when describing the changes. That commit uses "also" twice, to introduce two additional changes. :-)