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

60 statements  

« prev     ^ index     » next       coverage.py v7.2.7, created at 2023-07-21 09:39 +0000

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: 

39 from .._columns import ColumnExpression 

40 from .._engine import Engine 

41 from .._relation import Relation 

42 

43 

44@dataclasses.dataclass 

45class SortTerm: 

46 """Sort expression and indication of sort direction.""" 

47 

48 expression: ColumnExpression 

49 ascending: bool = True 

50 

51 def __str__(self) -> str: 

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

53 

54 

55@final 

56@dataclasses.dataclass(frozen=True) 

57class Sort(Reordering): 

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

59 column expressions. 

60 """ 

61 

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

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

64 

65 @property 

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

67 # Docstring inherited. 

68 result: set[ColumnTag] = set() 

69 for term in self.terms: 

70 result.update(term.expression.columns_required) 

71 return result 

72 

73 def __str__(self) -> str: 

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

75 

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

77 # Docstring inherited. 

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

79 

80 def _begin_apply( 

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

82 ) -> tuple[UnaryOperation, Engine]: 

83 # Docstring inherited. 

84 if not self.terms: 

85 return Identity(), target.engine 

86 for term in self.terms: 

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

88 raise ColumnError( 

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

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

91 ) 

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

93 

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

95 # Docstring inherited. 

96 if not self.terms: 

97 return target 

98 return super()._finish_apply(target) 

99 

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

101 """Compose this sort with another one. 

102 

103 Parameters 

104 ---------- 

105 next : `Sort` 

106 Sort that acts after ``self``. 

107 

108 Returns 

109 ------- 

110 composition : `Sort` 

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

112 back-to-back. 

113 """ 

114 new_terms = list(next.terms) 

115 for term in self.terms: 

116 if term not in new_terms: 

117 new_terms.append(term) 

118 return Sort(tuple(new_terms)) 

119 

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

121 # Docstring inherited. 

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

123 return UnaryCommutator( 

124 first=None, 

125 second=current.operation, 

126 done=False, 

127 messages=( 

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

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

130 ), 

131 ) 

132 if current.operation.is_order_dependent: 

133 return UnaryCommutator( 

134 first=None, 

135 second=current.operation, 

136 done=False, 

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

138 ) 

139 return UnaryCommutator(self, current.operation) 

140 

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

142 # Docstring inherited. 

143 if not self.terms: 

144 return upstream 

145 match upstream: 

146 case Sort(): 

147 return upstream.then(self) 

148 return None