Coverage for python/lsst/daf/relation/_operations/_sort.py: 39%
60 statements
« prev ^ index » next coverage.py v7.4.4, created at 2024-04-06 10:42 +0000
« prev ^ index » next coverage.py v7.4.4, created at 2024-04-06 10:42 +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/>.
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:
39 from .._columns import ColumnExpression
40 from .._engine import Engine
41 from .._relation import Relation
44@dataclasses.dataclass
45class SortTerm:
46 """Sort expression and indication of sort direction."""
48 expression: ColumnExpression
49 ascending: bool = True
51 def __str__(self) -> str:
52 return f"{'' if self.ascending else '-'}{self.expression}"
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 """
62 terms: tuple[SortTerm, ...] = ()
63 """Criteria for sorting rows (`Sequence` [ `SortTerm` ])."""
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
73 def __str__(self) -> str:
74 return f"sort[{', '.join(str(term) for term in self.terms)}]"
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)
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)
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)
100 def then(self, next: Sort) -> Sort:
101 """Compose this sort with another one.
103 Parameters
104 ----------
105 next : `Sort`
106 Sort that acts after ``self``.
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))
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)
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