Coverage for python/lsst/daf/butler/queries/tree/_column_reference.py: 64%
75 statements
« prev ^ index » next coverage.py v7.5.0, created at 2024-04-27 03:00 -0700
« prev ^ index » next coverage.py v7.5.0, created at 2024-04-27 03:00 -0700
1# This file is part of daf_butler.
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 software is dual licensed under the GNU General Public License and also
10# under a 3-clause BSD license. Recipients may choose which of these licenses
11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt,
12# respectively. If you choose the GPL option then the following text applies
13# (but note that there is still no warranty even if you opt for BSD instead):
14#
15# This program is free software: you can redistribute it and/or modify
16# it under the terms of the GNU General Public License as published by
17# the Free Software Foundation, either version 3 of the License, or
18# (at your option) any later version.
19#
20# This program is distributed in the hope that it will be useful,
21# but WITHOUT ANY WARRANTY; without even the implied warranty of
22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
23# GNU General Public License for more details.
24#
25# You should have received a copy of the GNU General Public License
26# along with this program. If not, see <http://www.gnu.org/licenses/>.
28from __future__ import annotations
30__all__ = ("ColumnReference", "DimensionKeyReference", "DimensionFieldReference", "DatasetFieldReference")
32from typing import TYPE_CHECKING, ClassVar, Literal, TypeAlias, TypeVar, Union, final
34import pydantic
36from ..._exceptions import InvalidQueryError
37from ...column_spec import ColumnType
38from ...dimensions import Dimension, DimensionElement
39from ._base import ColumnExpressionBase, DatasetFieldName
41if TYPE_CHECKING:
42 from ..visitors import ColumnExpressionVisitor
43 from ._column_set import ColumnSet
46_T = TypeVar("_T")
49@final
50class DimensionKeyReference(ColumnExpressionBase):
51 """A column expression that references a dimension primary key column."""
53 expression_type: Literal["dimension_key"] = "dimension_key"
55 is_column_reference: ClassVar[bool] = True
57 dimension: Dimension
58 """Definition and name of this dimension."""
60 def gather_required_columns(self, columns: ColumnSet) -> None:
61 # Docstring inherited.
62 columns.update_dimensions(self.dimension.minimal_group)
64 @property
65 def column_type(self) -> ColumnType:
66 # Docstring inherited.
67 return self.dimension.primary_key.type
69 def __str__(self) -> str:
70 return self.dimension.name
72 def visit(self, visitor: ColumnExpressionVisitor[_T]) -> _T:
73 # Docstring inherited.
74 return visitor.visit_dimension_key_reference(self)
77@final
78class DimensionFieldReference(ColumnExpressionBase):
79 """A column expression that references a dimension record column that is
80 not a primary key.
81 """
83 expression_type: Literal["dimension_field"] = "dimension_field"
85 is_column_reference: ClassVar[bool] = True
87 element: DimensionElement
88 """Definition and name of the dimension element."""
90 field: str
91 """Name of the field (i.e. column) in the element's logical table."""
93 def gather_required_columns(self, columns: ColumnSet) -> None:
94 # Docstring inherited.
95 columns.update_dimensions(self.element.minimal_group)
96 columns.dimension_fields[self.element.name].add(self.field)
98 @property
99 def column_type(self) -> ColumnType:
100 # Docstring inherited.
101 return self.element.schema.remainder[self.field].type
103 def __str__(self) -> str:
104 return f"{self.element}.{self.field}"
106 def visit(self, visitor: ColumnExpressionVisitor[_T]) -> _T:
107 # Docstring inherited.
108 return visitor.visit_dimension_field_reference(self)
110 @pydantic.model_validator(mode="after")
111 def _validate_field(self) -> DimensionFieldReference:
112 if self.field not in self.element.schema.remainder.names:
113 raise InvalidQueryError(f"Dimension field {self.element.name}.{self.field} does not exist.")
114 return self
117@final
118class DatasetFieldReference(ColumnExpressionBase):
119 """A column expression that references a column associated with a dataset
120 type.
121 """
123 expression_type: Literal["dataset_field"] = "dataset_field"
125 is_column_reference: ClassVar[bool] = True
127 dataset_type: str
128 """Name of the dataset type."""
130 field: DatasetFieldName
131 """Name of the field (i.e. column) in the dataset's logical table."""
133 def gather_required_columns(self, columns: ColumnSet) -> None:
134 # Docstring inherited.
135 columns.dataset_fields[self.dataset_type].add(self.field)
137 @property
138 def column_type(self) -> ColumnType:
139 # Docstring inherited.
140 match self.field:
141 case "dataset_id":
142 return "uuid"
143 case "ingest_date":
144 return "datetime"
145 case "run":
146 return "string"
147 case "collection":
148 return "string"
149 case "timespan":
150 return "timespan"
151 raise AssertionError(f"Invalid field {self.field!r} for dataset.")
153 def __str__(self) -> str:
154 return f"{self.dataset_type}.{self.field}"
156 def visit(self, visitor: ColumnExpressionVisitor[_T]) -> _T:
157 # Docstring inherited.
158 return visitor.visit_dataset_field_reference(self)
161ColumnReference: TypeAlias = Union[
162 DimensionKeyReference,
163 DimensionFieldReference,
164 DatasetFieldReference,
165]