Coverage for tests/test_selection.py: 18%
60 statements
« prev ^ index » next coverage.py v7.3.2, created at 2023-12-12 11:35 +0000
« prev ^ index » next coverage.py v7.3.2, created at 2023-12-12 11:35 +0000
1# This file is part of daf_relation.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (http://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# This program is free software: you can redistribute it and/or modify
10# it under the terms of the GNU General Public License as published by
11# the Free Software Foundation, either version 3 of the License, or
12# (at your option) any later version.
13#
14# This program is distributed in the hope that it will be useful,
15# but WITHOUT ANY WARRANTY; without even the implied warranty of
16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
17# GNU General Public License for more details.
18#
19# You should have received a copy of the GNU General Public License
20# along with this program. If not, see <http://www.gnu.org/licenses/>.
22from __future__ import annotations
24import unittest
26from lsst.daf.relation import (
27 ColumnError,
28 ColumnExpression,
29 EngineError,
30 Predicate,
31 Selection,
32 SortTerm,
33 UnaryOperationRelation,
34 iteration,
35 tests,
36)
39class SelectionTestCase(tests.RelationTestCase):
40 """Tests for the Selection operation and relations based on it."""
42 def setUp(self) -> None:
43 self.a = tests.ColumnTag("a")
44 self.predicate = ColumnExpression.reference(self.a).gt(ColumnExpression.literal(0))
45 self.engine = iteration.Engine(name="preferred")
46 self.leaf = self.engine.make_leaf(
47 {self.a}, payload=iteration.RowSequence([{self.a: 0}, {self.a: 1}]), name="leaf"
48 )
50 def test_attributes(self) -> None:
51 """Check that all UnaryOperation and Relation attributes have the
52 expected values.
53 """
54 relation = self.leaf.with_rows_satisfying(self.predicate)
55 assert isinstance(relation, UnaryOperationRelation)
56 self.assertEqual(relation.columns, {self.a})
57 self.assertEqual(relation.engine, self.engine)
58 self.assertEqual(relation.min_rows, 0)
59 self.assertEqual(relation.max_rows, self.leaf.max_rows)
60 operation = relation.operation
61 assert isinstance(operation, Selection)
62 self.assertEqual(operation.predicate, self.predicate)
63 self.assertEqual(operation.columns_required, {self.a})
64 self.assertFalse(operation.is_empty_invariant)
65 self.assertFalse(operation.is_count_invariant)
66 self.assertFalse(operation.is_order_dependent)
67 self.assertFalse(operation.is_count_dependent)
69 def test_apply_failures(self) -> None:
70 """Test failure modes of constructing and applying Selections."""
71 # Required columns must be present.
72 with self.assertRaises(ColumnError):
73 self.leaf.with_rows_satisfying(
74 ColumnExpression.reference(tests.ColumnTag("c")).lt(ColumnExpression.literal(0))
75 )
77 def test_backtracking_apply(self) -> None:
78 """Test apply logic that involves reordering operations in the existing
79 tree to perform the new operation in a preferred engine.
80 """
81 new_engine = iteration.Engine(name="downstream")
82 b = tests.ColumnTag("b")
83 expression = ColumnExpression.function(
84 "__add__", ColumnExpression.reference(self.a), ColumnExpression.literal(5)
85 )
86 sort_terms = [SortTerm(ColumnExpression.reference(self.a))]
87 other_predicate = ColumnExpression.reference(b).gt(ColumnExpression.literal(0))
88 # Apply a bunch of operations in a new engine that a Selection should
89 # commute with.
90 target = (
91 self.leaf.transferred_to(new_engine)
92 .with_calculated_column(b, expression)
93 .with_rows_satisfying(other_predicate)
94 .without_duplicates()
95 .with_only_columns({self.a})
96 .sorted(sort_terms)
97 )
98 # Apply a new Selection with backtracking and see that it appears
99 # before the transfer to the new engine, with adjustments as needed.
100 relation = target.with_rows_satisfying(
101 self.predicate, preferred_engine=self.engine, require_preferred_engine=True
102 )
103 self.assert_relations_equal(
104 relation,
105 (
106 self.leaf.with_rows_satisfying(self.predicate)
107 .transferred_to(new_engine)
108 .with_calculated_column(b, expression)
109 .with_rows_satisfying(other_predicate)
110 .without_duplicates()
111 .with_only_columns({self.a})
112 .sorted(sort_terms)
113 ),
114 )
116 def test_no_backtracking(self) -> None:
117 """Test apply logic that handles preferred engines without reordering
118 operations in the existing tree.
119 """
120 new_engine = iteration.Engine(name="downstream")
121 # Construct a relation tree we can't reorder when inserting a
122 # Selection, because there is a locked Materialization in the way.
123 target = self.leaf.transferred_to(new_engine).materialized("lock")
124 # Preferred engine is ignored if we can't backtrack and don't enable
125 # anything else.
126 self.assert_relations_equal(
127 target.with_rows_satisfying(self.predicate, preferred_engine=self.engine),
128 target.with_rows_satisfying(self.predicate),
129 )
130 # We can force this to be an error.
131 with self.assertRaises(EngineError):
132 target.with_rows_satisfying(
133 self.predicate, preferred_engine=self.engine, require_preferred_engine=True
134 )
135 # We can also automatically transfer (back) to the preferred engine.
136 self.assert_relations_equal(
137 target.with_rows_satisfying(self.predicate, preferred_engine=self.engine, transfer=True),
138 target.transferred_to(self.engine).with_rows_satisfying(self.predicate),
139 )
140 # Can't backtrack through a Calculation that provides required columns.
141 # In the future, we could make this possible by subsuming the
142 # calculated columns into the predicate.
143 b = tests.ColumnTag("b")
144 target = self.leaf.transferred_to(new_engine).with_calculated_column(
145 b, ColumnExpression.reference(self.a)
146 )
147 with self.assertRaises(EngineError):
148 target.with_rows_satisfying(
149 ColumnExpression.reference(b).gt(ColumnExpression.literal(0)),
150 preferred_engine=self.engine,
151 require_preferred_engine=True,
152 )
153 # Can't backtrack through a slice.
154 target = self.leaf.transferred_to(new_engine)[1:]
155 with self.assertRaises(EngineError):
156 target.with_rows_satisfying(
157 self.predicate,
158 preferred_engine=self.engine,
159 require_preferred_engine=True,
160 )
162 def test_apply_simplify(self) -> None:
163 """Test simplification logic in Selection.apply."""
164 self.assertEqual(self.leaf.with_rows_satisfying(Predicate.literal(True)), self.leaf)
165 new_predicate = ColumnExpression.reference(self.a).lt(ColumnExpression.literal(5))
166 self.assertEqual(
167 self.leaf.with_rows_satisfying(self.predicate).with_rows_satisfying(new_predicate),
168 self.leaf.with_rows_satisfying(self.predicate.logical_and(new_predicate)),
169 )
171 def test_iteration(self) -> None:
172 """Test Selection execution in the iteration engine."""
173 relation = self.leaf.with_rows_satisfying(self.predicate)
174 self.assertEqual(list(self.engine.execute(relation)), [{self.a: 1}])
176 def test_str(self) -> None:
177 """Test str(Selection) and
178 str(UnaryOperationRelation[Selection]).
179 """
180 relation = self.leaf.with_rows_satisfying(self.predicate)
181 self.assertEqual(str(relation), "σ[a>0](leaf)")
184if __name__ == "__main__":
185 unittest.main()