Coverage for python/lsst/daf/relation/_operations/_sort.py: 35%

64 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2023-02-22 02:59 -0800

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 

24__all__ = ( 

25 "SortTerm", 

26 "Sort", 

27) 

28 

29import dataclasses 

30from collections.abc import Set 

31from typing import TYPE_CHECKING, final 

32 

33from .._columns import ColumnTag 

34from .._exceptions import ColumnError 

35from .._operation_relations import UnaryOperationRelation 

36from .._unary_operation import Identity, Reordering, UnaryCommutator, UnaryOperation 

37 

38if TYPE_CHECKING: 38 ↛ 39line 38 didn't jump to line 39, because the condition on line 38 was never true

39 from .._columns import ColumnExpression 

40 from .._engine import Engine 

41 from .._relation import Relation 

42 

43 

44@dataclasses.dataclass 

45class SortTerm: 

46 expression: ColumnExpression 

47 ascending: bool = True 

48 

49 def __str__(self) -> str: 

50 return f"{'' if self.ascending else '-'}{self.expression}" 

51 

52 

53@final 

54@dataclasses.dataclass(frozen=True) 

55class Sort(Reordering): 

56 """A relation operation that orders rows according to a sequence of 

57 column expressions. 

58 """ 

59 

60 terms: tuple[SortTerm, ...] = () 

61 """Criteria for sorting rows (`Sequence` [ `SortTerm` ]).""" 

62 

63 @property 

64 def columns_required(self) -> Set[ColumnTag]: 

65 # Docstring inherited. 

66 result: set[ColumnTag] = set() 

67 for term in self.terms: 

68 result.update(term.expression.columns_required) 

69 return result 

70 

71 def __str__(self) -> str: 

72 return f"sort[{', '.join(str(term) for term in self.terms)}]" 

73 

74 def is_supported_by(self, engine: Engine) -> bool: 

75 # Docstring inherited. 

76 return all(term.expression.is_supported_by(engine) for term in self.terms) 

77 

78 def _begin_apply( 

79 self, target: Relation, preferred_engine: Engine | None 

80 ) -> tuple[UnaryOperation, Engine]: 

81 # Docstring inherited. 

82 if not self.terms: 

83 return Identity(), target.engine 

84 for term in self.terms: 

85 if not term.expression.columns_required <= target.columns: 

86 raise ColumnError( 

87 f"Sort term {term} for target relation {target} needs " 

88 f"columns {set(term.expression.columns_required - target.columns)}." 

89 ) 

90 return super()._begin_apply(target, preferred_engine) 

91 

92 def _finish_apply(self, target: Relation) -> Relation: 

93 # Docstring inherited. 

94 if not self.terms: 

95 return target 

96 return super()._finish_apply(target) 

97 

98 def then(self, next: Sort) -> Sort: 

99 """Compose this sort with another one. 

100 

101 Parameters 

102 ---------- 

103 next : `Sort` 

104 Sort that acts after ``self``. 

105 

106 Returns 

107 ------- 

108 composition : `Sort` 

109 Sort that is equivalent to ``self`` and ``next`` being applied 

110 back-to-back. 

111 """ 

112 new_terms = list(next.terms) 

113 for term in self.terms: 

114 if term not in new_terms: 

115 new_terms.append(term) 

116 return Sort(tuple(new_terms)) 

117 

118 def commute(self, current: UnaryOperationRelation) -> UnaryCommutator: 

119 # Docstring inherited. 

120 if not self.columns_required <= current.target.columns: 

121 return UnaryCommutator( 

122 first=None, 

123 second=current.operation, 

124 done=False, 

125 messages=( 

126 f"{current.target} is missing columns " 

127 f"{set(self.columns_required - current.target.columns)}", 

128 ), 

129 ) 

130 if current.operation.is_order_dependent: 

131 return UnaryCommutator( 

132 first=None, 

133 second=current.operation, 

134 done=False, 

135 messages=(f"{current.operation} is order-dependent",), 

136 ) 

137 return UnaryCommutator(self, current.operation) 

138 

139 def simplify(self, upstream: UnaryOperation) -> UnaryOperation | None: 

140 # Docstring inherited. 

141 if not self.terms: 

142 return upstream 

143 match upstream: 

144 case Sort(): 

145 return upstream.then(self) 

146 return None