Support filtering on FILTER field

Author: tomwhite

lib/vcf_encoder.c

@@ -729,6 +729,7 @@ vcz_variant_encoder_write_filter(const vcz_variant_encoder_t *self, size_t varia
 {
     const vcz_field_t filter_id = self->filter_id;
     bool all_missing = true;
+    bool first = true;
     const int8_t *restrict data = self->filter_data + (variant * filter_id.num_columns);
     const char *filter_id_data = (const char *) self->filter_id.data;
     size_t j, k, source_offset;
@@ -747,7 +748,10 @@ vcz_variant_encoder_write_filter(const vcz_variant_encoder_t *self, size_t varia
         source_offset = 0;
         for (j = 0; j < filter_id.num_columns; j++) {
             if (data[j]) {
-                source_offset = j * filter_id.item_size;
+                if (!first) {
+                    offset = append_char(buf, ';', offset, buflen);
+                }
+                    source_offset = j * filter_id.item_size;
                 for (k = 0; k < filter_id.item_size; k++) {
                     if (filter_id_data[source_offset] == VCZ_STRING_FILL) {
                         break;
@@ -758,6 +762,7 @@ vcz_variant_encoder_write_filter(const vcz_variant_encoder_t *self, size_t varia
                         goto out;
                     }
                     source_offset++;
+                    first = false;
                 }
             }
         }

tests/test_bcftools_validation.py

@@ -51,6 +51,7 @@ def run_vcztools(args: str, expect_error=False) -> tuple[str, str]:
         # ("view --no-version -i 'ID == \"rs6054257\"'", "sample.vcf.gz"),
         ("view --no-version -i 'DB=0'", "sample.vcf.gz"),
         ("view --no-version -i 'DB=1'", "sample.vcf.gz"),
+        ("view --no-version -i 'FILTER=\"PASS\"'", "sample.vcf.gz"),
         ("view --no-version -i 'INFO/DP > 10'", "sample.vcf.gz"),
         ("view --no-version -i 'FMT/DP >= 5'", "sample.vcf.gz"),
         ("view --no-version -i 'FMT/DP >= 5 && FMT/GQ > 10'", "sample.vcf.gz"),
@@ -91,7 +92,7 @@ def run_vcztools(args: str, expect_error=False) -> tuple[str, str]:
         (
             "view --no-version -r '20:1230236-' -i 'FMT/DP>3' -s 'NA00002,NA00003'",
             "sample.vcf.gz"
-        )
+        ),
     ],
     # This is necessary when trying to run individual tests, as the arguments above
     # make for unworkable command lines
@@ -182,6 +183,7 @@ def test_vcf_output_with_output_option(tmp_path, args, vcf_file):
         # (r"query  -f '%AC{1}\n' -i 'AC[1]>10' ", "sample.vcf.gz"),
         # TODO fill-out more of these when supported for more stuff is available
         # in filtering
+        ("query -f '%CHROM %POS %FILTER\n' -i 'FILTER=\"PASS\"'", "sample.vcf.gz"),
         # Per-sample query tests
         (
             r"query -f '[%CHROM %POS %SAMPLE %GT %DP %GQ\n]' -i 'FMT/DP>3'",

tests/test_filter.py

@@ -37,7 +37,6 @@ def test_invalid_expressions(self, parser, expression):
             ("INFO/HAYSTACK ~ 0", filter_mod.UnsupportedRegexError),
             ('CHROM="1"', filter_mod.UnsupportedChromFieldError),
             ('DP="."', filter_mod.UnsupportedMissingDataError),
-            ('FILTER="PASS"', filter_mod.UnsupportedFilterFieldError),
             ("ID!=@~/file", filter_mod.UnsupportedFileReferenceError),
             ("INFO/TAG=@file", filter_mod.UnsupportedFileReferenceError),
             ("INFO/X[0] == 1", filter_mod.UnsupportedArraySubscriptError),
@@ -199,6 +198,38 @@ def test_evaluate(self, expression, data, expected):
         result = fee.evaluate(numpify_values(data))
         nt.assert_array_equal(result, expected)
 
+    @pytest.mark.parametrize(
+        ("expression", "expected"),
+        [
+            ('FILTER="PASS"', [False, True, False, False, False, False]),
+            ('FILTER="."', [True, False, False, False, False, False]),
+            ('FILTER="A"', [False, False, True, False, False, False]),
+            ('FILTER!="A"', [True, True, False, True, True, True]),
+            ('FILTER~"A"', [False, False, True, False, True, True]),
+            ('FILTER="A;B"', [False, False, False, False, True, False]),
+            ('FILTER="B;A"', [False, False, False, False, True, False]),
+            ('FILTER!="A;B"', [True, True, True, True, False, True]),
+            ('FILTER~"A;B"', [False, False, False, False, True, True]),
+            ('FILTER~"B;A"', [False, False, False, False, True, True]),
+            ('FILTER!~"A;B"', [True, True, True, True, False, False]),
+        ],
+    )
+    def test_evaluate_filter_comparison(self, expression, expected):
+        data = {
+            "variant_filter": [
+                [False, False, False, False],
+                [True, False, False, False],
+                [False, True, False, False],
+                [False, False, True, False],
+                [False, True, True, False],
+                [False, True, True, True],
+            ],
+            "filter_id": ["PASS", "A", "B", "C"],
+        }
+        fee = filter_mod.FilterExpression(include=expression)
+        result = fee.evaluate(numpify_values(data))
+        nt.assert_array_equal(result, expected)
+
     @pytest.mark.parametrize(
         ("expr", "expected"),
         [

vcztools/filter.py

@@ -32,11 +32,6 @@ class UnsupportedMissingDataError(UnsupportedFilteringFeatureError):
     feature = "Missing data"
 
 
-class UnsupportedFilterFieldError(UnsupportedFilteringFeatureError):
-    issue = "164"
-    feature = "FILTER field"
-
-
 class UnsupportedGenotypeValuesError(UnsupportedFilteringFeatureError):
     issue = "165"
     feature = "Genotype values"
@@ -131,16 +126,18 @@ def __init__(self, mapper, tokens):
         token = tokens[0]
         if token == "CHROM":
             raise UnsupportedChromFieldError()
-        elif token == "FILTER":
-            raise UnsupportedFilterFieldError()
         elif token == "GT":
             raise UnsupportedGenotypeValuesError()
         self.field_name = mapper(token)
         logger.debug(f"Mapped {token} to {self.field_name}")
 
     def eval(self, data):
         value = np.asarray(data[self.field_name])
-        if not self.field_name.startswith("call_") and len(value.shape) > 1:
+        if (
+            not self.field_name.startswith("call_")
+            and self.field_name != "variant_filter"
+            and len(value.shape) > 1
+        ):
             raise Unsupported2DFieldsError()
         return value
 
@@ -301,6 +298,69 @@ def referenced_fields(self):
         return op1.referenced_fields() | op2.referenced_fields()
 
 
+# FILTER field expressions have special set-like semantics
+# so they are handled by dedicated operators.
+
+
+class FilterString(Constant):
+    def __init__(self, tokens):
+        super().__init__(tokens)
+
+    def eval(self, data):
+        # convert string to a 1D boolean array (one element per filter)
+        if self.tokens == ".":
+            return np.zeros_like(data["filter_id"], dtype=bool)
+        filters = self.tokens.split(";")
+        return np.isin(data["filter_id"], filters)
+
+    def referenced_fields(self):
+        return frozenset(["filter_id"])
+
+
+# 'a' is a 2D boolean array with shape (variants, filters)
+# 'b' is a 1D boolean array with shape (filters)
+
+
+def filter_eq(a, b):
+    return np.all(a == b, axis=1)
+
+
+def filter_ne(a, b):
+    return ~filter_eq(a, b)
+
+
+def filter_subset_match(a, b):
+    return np.all(a[:, b], axis=1)
+
+
+def filter_complement_match(a, b):
+    return ~filter_subset_match(a, b)
+
+
+class FilterFieldOperator(EvaluationNode):
+    op_map = {
+        "=": filter_eq,
+        "==": filter_eq,
+        "!=": filter_ne,
+        "~": filter_subset_match,
+        "!~": filter_complement_match,
+    }
+
+    def __init__(self, tokens):
+        super().__init__(tokens)
+        self.op1, self.op, self.op2 = tokens  # not self.tokens
+        self.comparison_fn = self.op_map[self.op]
+
+    def eval(self, data):
+        return self.comparison_fn(self.op1.eval(data), self.op2.eval(data))
+
+    def __repr__(self):
+        return f"({repr(self.op1)}){self.op}({repr(self.op2)})"
+
+    def referenced_fields(self):
+        return self.op1.referenced_fields() | self.op2.referenced_fields()
+
+
 def _identity(x):
     return x
 
@@ -321,6 +381,18 @@ def make_bcftools_filter_parser(all_fields=None, map_vcf_identifiers=True):
     vcf_prefixes = pp.Literal("INFO/") | pp.Literal("FORMAT/") | pp.Literal("FMT/")
     vcf_identifier = pp.Combine(vcf_prefixes + identifier) | identifier
 
+    name_mapper = _identity
+    if map_vcf_identifiers:
+        name_mapper = functools.partial(vcf_name_to_vcz_name, all_fields)
+
+    filter_field_identifier = pp.Literal("FILTER")
+    filter_field_identifier = filter_field_identifier.set_parse_action(
+        functools.partial(Identifier, name_mapper)
+    )
+    filter_string = pp.QuotedString('"').set_parse_action(FilterString)
+    filter_field_expr = filter_field_identifier + pp.one_of("= != ~ !~") + filter_string
+    filter_field_expr = filter_field_expr.set_parse_action(FilterFieldOperator)
+
     lbracket, rbracket = map(pp.Suppress, "[]")
     # TODO we need to define the indexing grammar more carefully, but
     # this at least let's us match correct strings and raise an informative
@@ -334,9 +406,6 @@ def make_bcftools_filter_parser(all_fields=None, map_vcf_identifiers=True):
     )
     indexed_identifier = pp.Group(vcf_identifier + (lbracket + index_expr + rbracket))
 
-    name_mapper = _identity
-    if map_vcf_identifiers:
-        name_mapper = functools.partial(vcf_name_to_vcz_name, all_fields)
     identifier = vcf_identifier.set_parse_action(
         functools.partial(Identifier, name_mapper)
     )
@@ -350,7 +419,12 @@ def make_bcftools_filter_parser(all_fields=None, map_vcf_identifiers=True):
 
     comp_op = pp.oneOf("< = == > >= <= !=")
     filter_expression = pp.infix_notation(
-        function | constant | indexed_identifier | identifier | file_expr,
+        filter_field_expr
+        | function
+        | constant
+        | indexed_identifier
+        | identifier
+        | file_expr,
         [
             ("-", 1, pp.OpAssoc.RIGHT, UnaryMinus),
             (pp.one_of("* /"), 2, pp.OpAssoc.LEFT, BinaryOperator),

vcztools/query.py

@@ -132,6 +132,7 @@ def generate(chunk_data):
             array = chunk_data[vcz_name]
             for row in array:
                 is_missing = np.any(row == -1)
+                sep = ","
 
                 if tag == "CHROM":
                     row = self.contig_ids[row]
@@ -144,6 +145,7 @@ def generate(chunk_data):
                         row = self.filter_ids[row]
                     else:
                         row = "."
+                    sep = ";"
                 if tag == "QUAL":
                     if math.isnan(row):
                         row = "."
@@ -154,7 +156,7 @@ def generate(chunk_data):
                     and not sample_loop
                     and (isinstance(row, np.ndarray) or isinstance(row, list))
                 ):
-                    row = ",".join(map(str, row))
+                    row = sep.join(map(str, row))
 
                 if sample_loop:
                     if isinstance(row, np.ndarray):