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

64 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2023-01-07 02:01 -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 

47 expression: ColumnExpression 

48 ascending: bool = True 

49 

50 def __str__(self) -> str: 

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

52 

53 

54@final 

55@dataclasses.dataclass(frozen=True) 

56class Sort(Reordering): 

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

58 column expressions. 

59 """ 

60 

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

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

63 

64 @property 

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

66 # Docstring inherited. 

67 result: set[ColumnTag] = set() 

68 for term in self.terms: 

69 result.update(term.expression.columns_required) 

70 return result 

71 

72 def __str__(self) -> str: 

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

74 

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

76 # Docstring inherited. 

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

78 

79 def _begin_apply( 

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

81 ) -> tuple[UnaryOperation, Engine]: 

82 # Docstring inherited. 

83 if not self.terms: 

84 return Identity(), target.engine 

85 for term in self.terms: 

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

87 raise ColumnError( 

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

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

90 ) 

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

92 

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

94 # Docstring inherited. 

95 if not self.terms: 

96 return target 

97 return super()._finish_apply(target) 

98 

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

100 """Compose this sort with another one. 

101 

102 Parameters 

103 ---------- 

104 next : `Sort` 

105 Sort that acts after ``self``. 

106 

107 Returns 

108 ------- 

109 composition : `Sort` 

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

111 back-to-back. 

112 """ 

113 new_terms = list(next.terms) 

114 for term in self.terms: 

115 if term not in new_terms: 

116 new_terms.append(term) 

117 return Sort(tuple(new_terms)) 

118 

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

120 # Docstring inherited. 

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

122 return UnaryCommutator( 

123 first=None, 

124 second=current.operation, 

125 done=False, 

126 messages=( 

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

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

129 ), 

130 ) 

131 if current.operation.is_order_dependent: 

132 return UnaryCommutator( 

133 first=None, 

134 second=current.operation, 

135 done=False, 

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

137 ) 

138 return UnaryCommutator(self, current.operation) 

139 

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

141 # Docstring inherited. 

142 if not self.terms: 

143 return upstream 

144 match upstream: 

145 case Sort(): 

146 return upstream.then(self) 

147 return None