🚧 Use tsv-utils for --output-metadata

tsv-join is much faster than the other implementation here (18x faster -
12s vs. 3m43s on the current SARS-CoV-2 GISAID dataset containing 16
million rows).

Author: victorlin

augur/filter/_run.py

@@ -419,14 +419,13 @@ def run(args):
             write_vcf(args.sequences, args.output_sequences, dropped_samps)
         else:
             subset_fasta(args.sequences, args.output_sequences, strains_file, args.nthreads)
-            if not args.output_strains:
-                os.remove(strains_file)
 
     if args.output_metadata:
         print_debug(f"Reading metadata from {args.metadata!r} and writing to {args.output_metadata!r}…")
-        write_output_metadata(args.metadata, args.metadata_delimiters,
-                              args.metadata_id_columns, args.output_metadata,
-                              valid_strains)
+        write_output_metadata(args.metadata, metadata_object.id_column, args.output_metadata, strains_file)
+
+    if not args.output_strains:
+        os.remove(strains_file)
 
     # Calculate the number of strains that don't exist in either metadata or
     # sequences.

augur/filter/io.py

@@ -1,18 +1,20 @@
 import argparse
-import csv
 from argparse import Namespace
 import os
 import re
+from shlex import quote as shquote
+from shutil import which
 from textwrap import dedent
-from typing import Sequence, Set
+from typing import Sequence
 import numpy as np
 from collections import defaultdict
-from xopen import xopen
 
 from augur.errors import AugurError
 from augur.io.file import open_file
-from augur.io.metadata import Metadata, METADATA_DATE_COLUMN
+from augur.io.metadata import METADATA_DATE_COLUMN
 from augur.io.print import print_err
+from augur.io.shell_command_runner import run_shell_command
+from augur.utils import augur
 from .constants import GROUP_BY_GENERATED_COLUMNS
 from .include_exclude_rules import extract_variables, parse_filter_query
 
@@ -96,25 +98,29 @@ def constant_factory(value):
         raise AugurError(f"missing or malformed priority scores file {fname}")
 
 
-def write_output_metadata(input_metadata_path: str, delimiters: Sequence[str],
-                          id_columns: Sequence[str], output_metadata_path: str,
-                          ids_to_write: Set[str]):
+def write_output_metadata(input_filename: str, id_column: str, output_filename: str, ids_file: str):
     """
-    Write output metadata file given input metadata information and a set of IDs
-    to write.
+    Write output metadata file given input metadata information and a file
+    containing ids to write.
     """
-    input_metadata = Metadata(input_metadata_path, delimiters, id_columns)
-
-    with xopen(output_metadata_path, "w", newline="") as output_metadata_handle:
-        output_metadata = csv.DictWriter(output_metadata_handle, fieldnames=input_metadata.columns,
-                                         delimiter="\t", lineterminator=os.linesep)
-        output_metadata.writeheader()
-
-        # Write outputs based on rows in the original metadata.
-        for row in input_metadata.rows():
-            row_id = row[input_metadata.id_column]
-            if row_id in ids_to_write:
-                output_metadata.writerow(row)
+    # FIXME: make this a function like augur() and seqkit()
+    tsv_join = which("tsv-join")
+
+    command = f"""
+        {augur()} read-file {shquote(input_filename)} |
+        {tsv_join} -H --filter-file {ids_file} --key-fields {id_column} |
+        {augur()} write-file {shquote(output_filename)}
+    """
+
+    try:
+        run_shell_command(command, raise_errors=True)
+    except Exception:
+        if os.path.isfile(output_filename):
+            # Remove the partial output file.
+            os.remove(output_filename)
+            raise AugurError(f"Metadata output failed, see error(s) above.")
+        else:
+            raise AugurError(f"Metadata output failed, see error(s) above. The command may have already written data to stdout. You may want to clean up any partial outputs.")
 
 
 # These are the types accepted in the following function.

tests/functional/filter/cram/filter-output-metadata-header.t

@@ -7,6 +7,8 @@ the default quotechar, any column names with that character may be altered.
 
 Quoted columns containing the tab delimiter are left unchanged.
 
+# FIXME: tsv-join has different behavior here. Test both?
+
   $ cat >metadata.tsv <<~~
   > strain	"col	1"
   > SEQ_1	a