Coverage for tests/test_cassandraPartitioner.py: 100%

1# This file is part of dax_apdb. 

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 

22import math 

23import unittest 

24from typing import Any 

25 

26from astropy.time import Time 

27 

28from lsst.dax.apdb import ApdbTables 

29from lsst.dax.apdb.cassandra import ( 

30 ApdbCassandraConfig, 

31 ApdbCassandraPartitioningConfig, 

32 ApdbCassandraTimePartitionRange, 

33) 

34from lsst.dax.apdb.cassandra.partitioner import Partitioner 

35from lsst.dax.apdb.cassandra.queries import QExpr 

36from lsst.sphgeom import Box, UnitVector3d 

37 

38 

39class CassandraPartitionerTestCase(unittest.TestCase): 

40 """A test case for ApdbCassandra class""" 

41 

42 def make_partitioner(self, **kwargs: Any) -> Partitioner: 

43 """Make Partitioner instance, keyword arguments are passed to 

44 ApdbCassandraPartitioningConfig. 

45 """ 

46 params = { 

47 "part_pixelization": "mq3c", 

48 "part_pix_level": 10, 

49 "query_per_spatial_part": False, 

50 } 

51 params.update(kwargs) 

52 

53 partitioning = ApdbCassandraPartitioningConfig(**params) # type: ignore[arg-type] 

54 config = ApdbCassandraConfig(partitioning=partitioning) 

55 return Partitioner(config) 

56 

57 def test_pixel(self) -> None: 

58 """Test pixel() method.""" 

59 partitioner = self.make_partitioner() 

60 

61 # Avoid getting on the boundary between pixels with tiny deltas. 

62 self.assertEqual(partitioner.pixel(UnitVector3d(1.0001, 1.0, 1.0001)), 0xD00000) 

63 self.assertEqual(partitioner.pixel(UnitVector3d(-1.0001, 1.0, 0.9999)), 0xE00000) 

64 self.assertEqual(partitioner.pixel(UnitVector3d(0.9999, -1.0, 0.9999)), 0xF55555) 

65 self.assertEqual(partitioner.pixel(UnitVector3d(1.0001, 1.0, -1.0001)), 0xAAAAAA) 

66 self.assertEqual(partitioner.pixel(UnitVector3d(-1.0001, -1.0, 0.9999)), 0xEFFFFF) 

67 self.assertEqual(partitioner.pixel(UnitVector3d(0.9999, -1.0, -0.9999)), 0xFAAAAA) 

68 

69 # Same with ra/dec in degrees. 

70 dec = math.atan(1 / math.sqrt(2.0)) / math.pi * 180 

71 self.assertEqual(partitioner.pixel(45.01, dec + 0.01), 0xD00000) 

72 self.assertEqual(partitioner.pixel(135.01, dec - 0.01), 0xE00000) 

73 self.assertEqual(partitioner.pixel(314.99, dec - 0.01), 0xF55555) 

74 self.assertEqual(partitioner.pixel(45.01, -dec - 0.01), 0xAAAAAA) 

75 self.assertEqual(partitioner.pixel(224.99, dec - 0.01), 0xEFFFFF) 

76 self.assertEqual(partitioner.pixel(314.99, -dec + 0.01), 0xFAAAAA) 

77 

78 def test_time_partition(self) -> None: 

79 """Test time_partition() method.""" 

80 partitioner = self.make_partitioner() 

81 

82 astrotime = Time("2025-06-01T00:00:00", format="isot", scale="tai") 

83 self.assertEqual(partitioner.time_partition(astrotime), 674) 

84 astrotime = Time("2025-07-01T00:00:00", format="isot", scale="tai") 

85 self.assertEqual(partitioner.time_partition(astrotime), 675) 

86 

87 self.assertEqual(partitioner.time_partition(60827.0), 674) 

88 self.assertEqual(partitioner.time_partition(60857.0), 675) 

89 

90 def test_partition_period(self) -> None: 

91 """Test partition_period() method.""" 

92 partitioner = self.make_partitioner() 

93 

94 self.assertEqual( 

95 partitioner.partition_period(674), 

96 ( 

97 Time("2025-05-12T00:00:00", format="isot", scale="tai"), 

98 Time("2025-06-11T00:00:00", format="isot", scale="tai"), 

99 ), 

100 ) 

101 self.assertEqual( 

102 partitioner.partition_period(800), 

103 ( 

104 Time("2035-09-17T00:00:00", format="isot", scale="tai"), 

105 Time("2035-10-17T00:00:00", format="isot", scale="tai"), 

106 ), 

107 ) 

108 

109 def test_spatial_where(self) -> None: 

110 """Test spatial_where() method.""" 

111 region = Box.fromDegrees(-0.05, 0.05, 0.05, 0.15) 

112 

113 partitioner = self.make_partitioner() 

114 result, count = partitioner.spatial_where(region) 

115 self.assertEqual(count, 4) 

116 self.assertEqual( 

117 result, 

118 [ 

119 QExpr( 

120 "apdb_part IN ({},{},{},{})", 

121 (12058622, 12058623, 12058624, 12058625), 

122 can_prepare=False, 

123 ) 

124 ], 

125 ) 

126 result, count = partitioner.spatial_where(region, use_ranges=True) 

127 self.assertEqual(count, 4) 

128 self.assertEqual(result, [QExpr("apdb_part >= {} AND apdb_part <= {}", (12058622, 12058625))]) 

129 

130 partitioner = self.make_partitioner(query_per_spatial_part=True) 

131 result, count = partitioner.spatial_where(region) 

132 self.assertEqual(count, 4) 

133 self.assertIn(QExpr("apdb_part = {}", (12058622,)), result) 

134 self.assertIn(QExpr("apdb_part = {}", (12058623,)), result) 

135 self.assertIn(QExpr("apdb_part = {}", (12058624,)), result) 

136 self.assertIn(QExpr("apdb_part = {}", (12058625,)), result) 

137 

138 result, count = partitioner.spatial_where(region, use_ranges=True) 

139 self.assertEqual(count, 4) 

140 self.assertEqual(result, [QExpr("apdb_part >= {} AND apdb_part <= {}", (12058622, 12058625))]) 

141 

142 def _check_temporal_where( 

143 self, 

144 tables: list[str], 

145 where: list[QExpr], 

146 part_start: int, 

147 part_end: int, 

148 *, 

149 time_partition_tables: bool = False, 

150 query_per_time_part: bool = False, 

151 ) -> None: 

152 if part_start > part_end: 

153 self.assertEqual(tables, []) 

154 self.assertEqual(where, []) 

155 elif time_partition_tables: 

156 expect_tables = [f"DiaSource_{part}" for part in range(part_start, part_end + 1)] 

157 self.assertEqual(tables, expect_tables) 

158 self.assertEqual(where, []) 

159 elif query_per_time_part: 

160 where_str = '"apdb_time_part" = {}' 

161 expect_where = [QExpr(where_str, (part,)) for part in range(part_start, part_end + 1)] 

162 self.assertEqual(tables, ["DiaSource"]) 

163 self.assertEqual(where, expect_where) 

164 else: 

165 num_part = part_end + 1 - part_start 

166 placeholders = ",".join(["{}"] * num_part) 

167 self.assertEqual(tables, ["DiaSource"]) 

168 can_prepare = num_part <= 3 

169 self.assertEqual( 

170 where, 

171 [ 

172 QExpr( 

173 f"apdb_time_part IN ({placeholders})", 

174 tuple(range(part_start, part_end + 1)), 

175 can_prepare=can_prepare, 

176 ) 

177 ], 

178 ) 

179 

180 def test_temporal_where(self) -> None: 

181 """Test temporal_where() method.""" 

182 start_time = Time("2025-01-01T00:00:00", format="isot", scale="tai") 

183 end_time = Time("2025-06-01T00:00:00", format="isot", scale="tai") 

184 

185 partitioner = self.make_partitioner() 

186 

187 tables, where = partitioner.temporal_where(ApdbTables.DiaSource, start_time, end_time) 

188 self._check_temporal_where(tables, where, 669, 674) 

189 

190 tables, where = partitioner.temporal_where( 

191 ApdbTables.DiaSource, start_time, end_time, query_per_time_part=True 

192 ) 

193 self._check_temporal_where(tables, where, 669, 674, query_per_time_part=True) 

194 

195 tables, where = partitioner.temporal_where( 

196 ApdbTables.DiaSource, 

197 start_time, 

198 end_time, 

199 query_per_time_part=True, 

200 ) 

201 self._check_temporal_where(tables, where, 669, 674, query_per_time_part=True) 

202 

203 partitioner = self.make_partitioner(query_per_time_part=True) 

204 

205 tables, where = partitioner.temporal_where(ApdbTables.DiaSource, start_time, end_time) 

206 self._check_temporal_where(tables, where, 669, 674, query_per_time_part=True) 

207 

208 partitioner = self.make_partitioner(time_partition_tables=True) 

209 tables, where = partitioner.temporal_where(ApdbTables.DiaSource, start_time, end_time) 

210 self._check_temporal_where(tables, where, 669, 674, time_partition_tables=True) 

211 

212 # Check additional partition range constraint. 

213 ranges = [ 

214 ((0, 1000), (669, 674)), 

215 ((0, 1), (0, -1)), 

216 ((600, 670), (669, 670)), 

217 ((670, 770), (670, 674)), 

218 ((671, 672), (671, 672)), 

219 ] 

220 

221 partitioner = self.make_partitioner() 

222 for (range_start, range_end), (result_start, result_end) in ranges: 

223 part_range = ApdbCassandraTimePartitionRange(start=range_start, end=range_end) 

224 tables, where = partitioner.temporal_where( 

225 ApdbTables.DiaSource, start_time, end_time, partitons_range=part_range 

226 ) 

227 self._check_temporal_where(tables, where, result_start, result_end) 

228 

229 partitioner = self.make_partitioner(time_partition_tables=True) 

230 for (range_start, range_end), (result_start, result_end) in ranges: 

231 part_range = ApdbCassandraTimePartitionRange(start=range_start, end=range_end) 

232 tables, where = partitioner.temporal_where( 

233 ApdbTables.DiaSource, start_time, end_time, partitons_range=part_range 

234 ) 

235 self._check_temporal_where(tables, where, result_start, result_end, time_partition_tables=True) 

236 

237 

238if __name__ == "__main__": 

239 unittest.main()