Add introspection and referenced fields to eval tree

Author: jeromekelleher

tests/test_filter.py

@@ -86,6 +86,33 @@ def test_evaluate(self, expression, data, expected):
         result = fee.evaluate(numpify_values(data))
         nt.assert_array_equal(result, expected)
 
+    @pytest.mark.parametrize(
+        ("expr", "expected"),
+        [
+            ("a == b", {"variant_a", "variant_b"}),
+            ("a == b + c", {"variant_a", "variant_b", "variant_c"}),
+            ("(a + 1) < (b + c) - d / a", {f"variant_{x}" for x in "abcd"}),
+        ],
+    )
+    def test_referenced_fields(self, expr, expected):
+        fe = filter_mod.FilterExpression(include=expr)
+        assert fe.referenced_fields == expected
+
+    @pytest.mark.parametrize(
+        ("expr", "expected"),
+        [
+            ("a == b", "(variant_a)==(variant_b)"),
+            ("a + 1", "(variant_a)+(1)"),
+            ("a + 1 + 2", "(variant_a)+(1)+(2)"),
+            ("a + (1 + 2)", "(variant_a)+((1)+(2))"),
+            ("POS<10", "(variant_position)<(10)"),
+            ('CHROM=="chr1"', "(variant_contig)==('chr1')"),
+        ],
+    )
+    def test_repr(self, expr, expected):
+        fe = filter_mod.FilterExpression(include=expr)
+        assert repr(fe.parse_result[0]) == expected
+
 
 class TestBcftoolsParser:
     @pytest.mark.parametrize(

vcztools/filter.py

@@ -29,15 +29,28 @@ class Constant(EvaluationNode):
     def eval(self, data):
         return self.tokens
 
+    def __repr__(self):
+        return repr(self.tokens)
+
+    def referenced_fields(self):
+        return frozenset()
+
 
 class Identifier(EvaluationNode):
     def __init__(self, mapper, tokens):
         self.field_name = mapper(tokens[0])
         logger.debug(f"Mapped {tokens[0]} to {self.field_name}")
+        # TODO add errors for unsupported things like call_ fields etc.
 
     def eval(self, data):
         return data[self.field_name]
 
+    def __repr__(self):
+        return self.field_name
+
+    def referenced_fields(self):
+        return frozenset([self.field_name])
+
 
 class BinaryOperator(EvaluationNode):
     op_map = {
@@ -55,19 +68,31 @@ class BinaryOperator(EvaluationNode):
     }
 
     def eval(self, data):
-        # start by eval()'ing the first operand
-        ret = self.tokens[0].eval(data)
-
-        # get following operators and operands in pairs
+        # get the  operators and operands in pairs
+        operands = self.tokens[0::2]
         ops = self.tokens[1::2]
-        operands = self.tokens[2::2]
-        for op, operand in zip(ops, operands):
-            # print(f"Eval {op}, {ret}, {operand}")
-            # update cumulative value by add/subtract/mult/divide the next operand
+        # start by eval()'ing the first operand
+        ret = operands[0].eval(data)
+        for op, operand in zip(ops, operands[1:]):
             arith_fn = self.op_map[op]
             ret = arith_fn(ret, operand.eval(data))
         return ret
 
+    def __repr__(self):
+        ops = self.tokens[1::2]
+        operands = self.tokens[0::2]
+        ret = f"({repr(operands[0])})"
+        for op, operand in zip(ops, operands[1:]):
+            ret += f"{op}({repr(operand)})"
+        return ret
+
+    def referenced_fields(self):
+        operands = self.tokens[0::2]
+        ret = operands[0].referenced_fields()
+        for operand in operands[1:]:
+            ret |= operand.referenced_fields()
+        return ret
+
 
 class ComparisonOperator(EvaluationNode):
     op_map = {
@@ -85,6 +110,14 @@ def eval(self, data):
         comparison_fn = self.op_map[op]
         return comparison_fn(op1.eval(data), op2.eval(data))
 
+    def __repr__(self):
+        op1, op, op2 = self.tokens
+        return f"({repr(op1)}){op}({repr(op2)})"
+
+    def referenced_fields(self):
+        op1, _, op2 = self.tokens
+        return op1.referenced_fields() | op2.referenced_fields()
+
 
 def _identity(x):
     return x
@@ -110,6 +143,7 @@ def make_bcftools_filter_parser(all_fields=None, map_vcf_identifiers=True):
     filter_expression = pp.infix_notation(
         constant | identifier,
         [
+            # FIXME Does bcftools support unary minus?
             # ("-", 1, pp.OpAssoc.RIGHT, ),
             (pp.one_of("* /"), 2, pp.OpAssoc.LEFT, BinaryOperator),
             (pp.one_of("+ -"), 2, pp.OpAssoc.LEFT, BinaryOperator),
@@ -128,6 +162,7 @@ def __init__(self, *, field_names=None, include=None, exclude=None):
         if field_names is None:
             field_names = set()
         self.parse_result = None
+        self.referenced_fields = set()
         self.invert = False
         expr = None
         if include is not None and exclude is not None:
@@ -144,9 +179,8 @@ def __init__(self, *, field_names=None, include=None, exclude=None):
         if expr is not None:
             parser = make_bcftools_filter_parser(field_names)
             self.parse_result = parser.parse_string(expr, parse_all=True)
-
-        # Setting to None for now so that we retrieve all fields
-        self.referenced_fields = None
+            # This isn't a very good pattern, fix
+            self.referenced_fields = self.parse_result[0].referenced_fields()
 
     def evaluate(self, chunk_data):
         if self.parse_result is None:

vcztools/query.py

@@ -272,9 +272,7 @@ def generate(root):
             # NOTE: this should be done at the top-level when we've
             # figured out what fields need to be retrieved from both
             # the parsed query and filter expressions.
-            reader = retrieval.VariantChunkReader(
-                root, fields=filter_expr.referenced_fields
-            )
+            reader = retrieval.VariantChunkReader(root)
             for v_chunk in range(root["variant_position"].cdata_shape[0]):
                 # print("Read v_chunk", v_chunk)
                 chunk_data = reader[v_chunk]