Coverage for tests/test_calculation.py: 19%

58 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2023-03-31 02:35 -0700

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 Calculation, 

28 ColumnError, 

29 ColumnExpression, 

30 EngineError, 

31 SortTerm, 

32 UnaryOperationRelation, 

33 iteration, 

34 tests, 

35) 

36 

37 

38class CalculationTestCase(tests.RelationTestCase): 

39 """Tests for the Calculation operation and relations based on it.""" 

40 

41 def setUp(self) -> None: 

42 self.a = tests.ColumnTag("a") 

43 self.b = tests.ColumnTag("b") 

44 self.expression = ColumnExpression.function( 

45 "__add__", 

46 ColumnExpression.reference(self.a), 

47 ColumnExpression.literal(10), 

48 ) 

49 self.engine = iteration.Engine(name="preferred") 

50 self.leaf = self.engine.make_leaf( 

51 {self.a}, payload=iteration.RowSequence([{self.a: 0}, {self.a: 1}]), name="leaf" 

52 ) 

53 

54 def test_attributes(self) -> None: 

55 """Check that all UnaryOperation and Relation attributes have the 

56 expected values. 

57 """ 

58 relation = self.leaf.with_calculated_column(self.b, self.expression) 

59 assert isinstance(relation, UnaryOperationRelation) 

60 self.assertEqual(relation.columns, {self.a, self.b}) 

61 self.assertEqual(relation.engine, self.engine) 

62 self.assertEqual(relation.min_rows, self.leaf.min_rows) 

63 self.assertEqual(relation.max_rows, self.leaf.max_rows) 

64 operation = relation.operation 

65 assert isinstance(operation, Calculation) 

66 self.assertEqual(operation.expression, self.expression) 

67 self.assertEqual(operation.columns_required, {self.a}) 

68 self.assertTrue(operation.is_empty_invariant) 

69 self.assertTrue(operation.is_count_invariant) 

70 self.assertFalse(operation.is_order_dependent) 

71 self.assertFalse(operation.is_count_dependent) 

72 

73 def test_apply_failures(self) -> None: 

74 """Test failure modes of constructing and applying Calculations.""" 

75 # Calculations must depend on existing columns. 

76 with self.assertRaises(ColumnError): 

77 self.leaf.with_calculated_column(self.b, ColumnExpression.literal(10)) 

78 # Required columns must be present. 

79 with self.assertRaises(ColumnError): 

80 self.leaf.with_calculated_column(self.b, ColumnExpression.reference(tests.ColumnTag("c"))) 

81 # New column must not already be present. 

82 with self.assertRaises(ColumnError): 

83 self.leaf.with_calculated_column(self.a, self.expression) 

84 

85 def test_backtracking_apply(self) -> None: 

86 """Test apply logic that involves reordering operations in the existing 

87 tree to perform the new operation in a preferred engine. 

88 """ 

89 new_engine = iteration.Engine(name="downstream") 

90 predicate = ColumnExpression.reference(self.b).lt(ColumnExpression.literal(20)) 

91 # Apply a bunch of operations in a new engine that a Calculation should 

92 # commute with. 

93 target = ( 

94 self.leaf.transferred_to(new_engine) 

95 .with_calculated_column(self.b, self.expression) 

96 .with_rows_satisfying(predicate) 

97 .without_duplicates() 

98 .with_only_columns({self.a}) 

99 .sorted([SortTerm(ColumnExpression.reference(self.a))]) 

100 )[:2] 

101 # Apply a new Calculation with backtracking and see that it appears 

102 # before the transfer to the new engine, with adjustments as needed 

103 # downstream (to the Projection and Chain, in this case). 

104 tag = tests.ColumnTag("tag") 

105 relation = target.with_calculated_column( 

106 tag, self.expression, preferred_engine=self.engine, require_preferred_engine=True 

107 ) 

108 self.assert_relations_equal( 

109 relation, 

110 ( 

111 self.leaf.with_calculated_column(tag, self.expression) 

112 .transferred_to(new_engine) 

113 .with_calculated_column(self.b, self.expression) 

114 .with_rows_satisfying(predicate) 

115 .without_duplicates() 

116 .with_only_columns({self.a, tag}) 

117 .sorted([SortTerm(ColumnExpression.reference(self.a))]) 

118 )[:2], 

119 ) 

120 

121 def test_no_backtracking(self) -> None: 

122 """Test apply logic that handles preferred engines without reordering 

123 operations in the existing tree. 

124 """ 

125 new_engine = iteration.Engine(name="downstream") 

126 # Construct a relation tree we can't reorder when inserting a 

127 # Calculation, for various reasons. 

128 leaf = self.engine.make_leaf( 

129 {self.a}, 

130 payload=iteration.RowSequence([{self.a: 100}, {self.a: 200}]), 

131 name="leaf", 

132 ) 

133 # Can't insert after leaf because Materialization is locked. 

134 target = leaf.transferred_to(new_engine).materialized("lock") 

135 # Preferred engine is ignored if we can't backtrack and don't enable 

136 # anything else. 

137 self.assert_relations_equal( 

138 target.with_calculated_column(self.b, self.expression, preferred_engine=self.engine), 

139 target.with_calculated_column(self.b, self.expression), 

140 ) 

141 # We can force this to be an error. 

142 with self.assertRaises(EngineError): 

143 target.with_calculated_column( 

144 self.b, self.expression, preferred_engine=self.engine, require_preferred_engine=True 

145 ) 

146 # We can also automatically transfer (back) to the preferred engine. 

147 self.assert_relations_equal( 

148 target.with_calculated_column( 

149 self.b, self.expression, preferred_engine=self.engine, transfer=True 

150 ), 

151 target.transferred_to(self.engine).with_calculated_column(self.b, self.expression), 

152 ) 

153 # Can't commute past an existing Calculation that provides a required 

154 # column. 

155 target = leaf.transferred_to(new_engine).with_calculated_column(self.b, self.expression) 

156 with self.assertRaises(EngineError): 

157 target.with_calculated_column( 

158 tests.ColumnTag("c"), 

159 ColumnExpression.reference(self.b), 

160 preferred_engine=self.engine, 

161 require_preferred_engine=True, 

162 ) 

163 

164 def test_iteration(self) -> None: 

165 """Test Calculation execution in the iteration engine.""" 

166 relation = self.leaf.with_calculated_column(self.b, self.expression) 

167 self.assertEqual( 

168 list(self.engine.execute(relation)), 

169 [{self.a: 0, self.b: 10}, {self.a: 1, self.b: 11}], 

170 ) 

171 

172 def test_str(self) -> None: 

173 """Test str(Calculation) and 

174 str(UnaryOperationRelation[Calculation]). 

175 """ 

176 relation = self.leaf.with_calculated_column(self.b, self.expression) 

177 self.assertEqual(str(relation), "+[b=__add__(a, 10)](leaf)") 

178 

179 

180if __name__ == "__main__": 180 ↛ 181line 180 didn't jump to line 181, because the condition on line 180 was never true

181 unittest.main()