Coverage for python/lsst/daf/relation/_operations/_sort.py: 35%
64 statements
« prev ^ index » next coverage.py v6.5.0, created at 2023-01-10 02:26 -0800
« prev ^ index » next coverage.py v6.5.0, created at 2023-01-10 02:26 -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/>.
22from __future__ import annotations
24__all__ = (
25 "SortTerm",
26 "Sort",
27)
29import dataclasses
30from collections.abc import Set
31from typing import TYPE_CHECKING, final
33from .._columns import ColumnTag
34from .._exceptions import ColumnError
35from .._operation_relations import UnaryOperationRelation
36from .._unary_operation import Identity, Reordering, UnaryCommutator, UnaryOperation
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
44@dataclasses.dataclass
45class SortTerm:
47 expression: ColumnExpression
48 ascending: bool = True
50 def __str__(self) -> str:
51 return f"{'' if self.ascending else '-'}{self.expression}"
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 """
61 terms: tuple[SortTerm, ...] = ()
62 """Criteria for sorting rows (`Sequence` [ `SortTerm` ])."""
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
72 def __str__(self) -> str:
73 return f"sort[{', '.join(str(term) for term in self.terms)}]"
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)
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)
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)
99 def then(self, next: Sort) -> Sort:
100 """Compose this sort with another one.
102 Parameters
103 ----------
104 next : `Sort`
105 Sort that acts after ``self``.
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))
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)
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