Coverage for python/lsst/dax/apdb/cassandra/partitioner.py: 91%

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 

22from __future__ import annotations 

23 

24__all__ = ["Partitioner"] 

25 

26from typing import Any, overload 

27 

28import astropy.time 

29 

30from lsst import sphgeom 

31 

32from ..apdbSchema import ApdbTables 

33from ..pixelization import Pixelization 

34from .config import ApdbCassandraConfig, ApdbCassandraTimePartitionRange 

35from .queries import Column as C # noqa: N817 

36from .queries import QExpr 

37 

38 

39class Partitioner: 

40 """Logic for temporal and spatial partitioning of APDB tables. 

41 

42 Parameters 

43 ---------- 

44 config : `ApdbCassandraConfig` 

45 Configuration object. 

46 """ 

47 

48 partition_zero_epoch = astropy.time.Time(0, format="unix_tai") 

49 """Start time for partition 0, this should never be changed.""" 

50 

51 def __init__(self, config: ApdbCassandraConfig): 

52 self._config = config 

53 self.pixelization = Pixelization( 

54 config.partitioning.part_pixelization, 

55 config.partitioning.part_pix_level, 

56 config.partitioning.part_pix_max_ranges, 

57 ) 

58 self._epoch = float(self.partition_zero_epoch.mjd) 

59 

60 @overload 

61 def pixel(self, direction: sphgeom.UnitVector3d, /) -> int: ... 61 ↛ exitline 61 didn't return from function 'pixel' because

62 

63 @overload 

64 def pixel(self, ra: float, dec: float, /) -> int: ... 64 ↛ exitline 64 didn't return from function 'pixel' because

65 

66 def pixel(self, *args: Any) -> int: 

67 """Compute the index of the pixel for given direction. 

68 

69 Parameters 

70 ---------- 

71 args 

72 The method can take either a single `lsst.sphgeom.UnitVector3d` or 

73 a pair of floating point numbers (or values convertible to floats) 

74 representing RA and Dec in degrees. 

75 

76 Returns 

77 ------- 

78 pixel : `int` 

79 Pixel index. 

80 """ 

81 return self.pixelization.pixel(*args) 

82 

83 def time_partition(self, time: float | astropy.time.Time) -> int: 

84 """Calculate time partition number for a given time. 

85 

86 Parameters 

87 ---------- 

88 time : `float` or `astropy.time.Time` 

89 Time for which to calculate partition number. Can be float to mean 

90 MJD or `astropy.time.Time` 

91 

92 Returns 

93 ------- 

94 partition : `int` 

95 Partition number for a given time. 

96 """ 

97 if isinstance(time, astropy.time.Time): 

98 mjd = float(time.mjd) 

99 else: 

100 mjd = time 

101 days_since_epoch = mjd - self._epoch 

102 partition = int(days_since_epoch) // self._config.partitioning.time_partition_days 

103 return partition 

104 

105 def partition_period(self, time_partition: int) -> tuple[astropy.time.Time, astropy.time.Time]: 

106 """Return time period for specified time partition. 

107 

108 Parameters 

109 ---------- 

110 time_partition : `int` 

111 Time partition. 

112 

113 Returns 

114 ------- 

115 start : `astropy.time.Time` 

116 Start of the period, inclusive boundary. 

117 end : `astropy.time.Time` 

118 Start of the period, exclusive boundary. 

119 """ 

120 partition_days = self._config.partitioning.time_partition_days 

121 start_mjd = self._epoch + partition_days * time_partition 

122 end_mjd = self._epoch + partition_days * (time_partition + 1) 

123 start = astropy.time.Time(start_mjd, format="mjd", scale="tai") 

124 end = astropy.time.Time(end_mjd, format="mjd", scale="tai") 

125 return (start, end) 

126 

127 def spatial_where( 

128 self, region: sphgeom.Region | None, *, use_ranges: bool = False 

129 ) -> tuple[list[QExpr], int]: 

130 """Generate expressions for spatial part of WHERE clause. 

131 

132 Parameters 

133 ---------- 

134 region : `sphgeom.Region` 

135 Spatial region for query results. 

136 use_ranges : `bool`, optional 

137 If True then use pixel ranges ("apdb_part >= p1 AND apdb_part <= 

138 p2") instead of exact list of pixels. Should be set to True for 

139 large regions covering very many pixels. 

140 

141 Returns 

142 ------- 

143 expressions : `list` [ `QExpr` ] 

144 Empty list is returned if ``region`` is `None`, otherwise a list 

145 of one or more `QExpr`. 

146 partition_count : `int` 

147 Number of spatial partitions in the result. 

148 """ 

149 if region is None: 149 ↛ 150line 149 didn't jump to line 150 because the condition on line 149 was never true

150 return [], 0 

151 

152 count = 0 

153 expressions: list[QExpr] = [] 

154 if use_ranges: 

155 pixel_ranges = self.pixelization.envelope(region) 

156 for lower, upper in pixel_ranges: 

157 upper -= 1 

158 if lower == upper: 158 ↛ 159line 158 didn't jump to line 159 because the condition on line 158 was never true

159 expressions.append(C("apdb_part") == lower) 

160 count += 1 

161 elif lower + 1 == upper: 161 ↛ 162line 161 didn't jump to line 162 because the condition on line 161 was never true

162 expressions.append(C("apdb_part") == lower) 

163 expressions.append(C("apdb_part") == upper) 

164 count += 2 

165 else: 

166 count += upper - lower + 1 

167 expressions.append((C("apdb_part") >= lower) & (C("apdb_part") <= upper)) 

168 else: 

169 pixels = self.pixelization.pixels(region) 

170 count = len(pixels) 

171 if self._config.partitioning.query_per_spatial_part: 

172 expressions.extend((C("apdb_part") == pixel) for pixel in pixels) 

173 else: 

174 # If the are many pixels then don't prepare statements. 

175 can_prepare = len(pixels) <= 3 

176 expressions.append(C("apdb_part").in_(pixels, can_prepare=can_prepare)) 

177 

178 return expressions, count 

179 

180 def temporal_where( 

181 self, 

182 table: ApdbTables, 

183 start_time: float | astropy.time.Time, 

184 end_time: float | astropy.time.Time, 

185 *, 

186 query_per_time_part: bool | None = None, 

187 partitons_range: ApdbCassandraTimePartitionRange | None = None, 

188 ) -> tuple[list[str], list[QExpr]]: 

189 """Generate table names and expressions for temporal part of WHERE 

190 clauses. 

191 

192 Parameters 

193 ---------- 

194 table : `ApdbTables` 

195 Table to select from. 

196 start_time : `astropy.time.Time` or `float` 

197 Starting Datetime of MJD value of the time range. 

198 end_time : `astropy.time.Time` or `float` 

199 Starting Datetime of MJD value of the time range. 

200 query_per_time_part : `bool`, optional 

201 If None then use ``query_per_time_part`` from configuration. 

202 partitons_range : `ApdbCassandraTimePartitionRange` or `None` 

203 Partitions range to further restrict time range. 

204 

205 Returns 

206 ------- 

207 tables : `list` [ `str` ] 

208 List of the table names to query. Empty list is returned when time 

209 range does not overlap ``partitons_range``. 

210 expressions : `list` [ `QExpr` ] 

211 A list of zero or more `QExpr` instances. 

212 """ 

213 tables: list[str] 

214 temporal_where: list[QExpr] = [] 

215 # First and last partition. 

216 time_part_start = self.time_partition(start_time) 

217 time_part_end = self.time_partition(end_time) 

218 if partitons_range: 

219 # Check for non-overlapping ranges. 

220 if time_part_start > partitons_range.end or time_part_end < partitons_range.start: 

221 return [], [] 

222 if time_part_start < partitons_range.start: 

223 time_part_start = partitons_range.start 

224 if time_part_end > partitons_range.end: 

225 time_part_end = partitons_range.end 

226 # Inclusive range. 

227 time_parts = list(range(time_part_start, time_part_end + 1)) 

228 if self._config.partitioning.time_partition_tables: 

229 tables = [table.table_name(self._config.prefix, part) for part in time_parts] 

230 else: 

231 tables = [table.table_name(self._config.prefix)] 

232 if query_per_time_part is None: 

233 query_per_time_part = self._config.partitioning.query_per_time_part 

234 if query_per_time_part: 

235 temporal_where = [QExpr('"apdb_time_part" = {}', (time_part,)) for time_part in time_parts] 

236 else: 

237 # If the are many partitions then don't prepare statements. 

238 can_prepare = len(time_parts) <= 3 

239 temporal_where = [C("apdb_time_part").in_(time_parts, can_prepare=can_prepare)] 

240 

241 return tables, temporal_where