Coverage for tests/test_selection.py: 18%

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/>. 

21 

22from __future__ import annotations 

23 

24import unittest 

25 

26from lsst.daf.relation import ( 

27 ColumnError, 

28 ColumnExpression, 

29 EngineError, 

30 Predicate, 

31 Selection, 

32 SortTerm, 

33 UnaryOperationRelation, 

34 iteration, 

35 tests, 

36) 

37 

38 

39class SelectionTestCase(tests.RelationTestCase): 

40 """Tests for the Selection operation and relations based on it.""" 

41 

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 ) 

49 

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) 

68 

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 ) 

76 

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 ) 

115 

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 ) 

161 

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 ) 

170 

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}]) 

175 

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)") 

182 

183 

184if __name__ == "__main__": 

185 unittest.main()