Coverage for python/lsst/daf/butler/tests/butler_queries.py: 8%
824 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-10-09 03:23 -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__ = ()
32import unittest
33import unittest.mock
34from abc import ABC, abstractmethod
35from collections.abc import Iterable, Sequence
36from operator import attrgetter
37from uuid import UUID
39import astropy.coordinates
40import astropy.time
41from numpy import int64
43from lsst.sphgeom import LonLat, Region
45from .._butler import Butler
46from .._collection_type import CollectionType
47from .._dataset_ref import DatasetRef
48from .._dataset_type import DatasetType
49from .._exceptions import (
50 EmptyQueryResultError,
51 InvalidQueryError,
52 MissingCollectionError,
53 MissingDatasetTypeError,
54)
55from .._timespan import Timespan
56from ..dimensions import DataCoordinate, DimensionRecord
57from ..direct_query_driver import DirectQueryDriver
58from ..queries import DimensionRecordQueryResults, Query
59from ..queries.tree import Predicate
60from ..registry import NoDefaultCollectionError, RegistryDefaults
61from .utils import TestCaseMixin
63# Simplified tuples of the detector records we'll frequently be querying for.
64DETECTOR_TUPLES = {
65 1: ("Cam1", 1, "Aa", "SCIENCE"),
66 2: ("Cam1", 2, "Ab", "SCIENCE"),
67 3: ("Cam1", 3, "Ba", "SCIENCE"),
68 4: ("Cam1", 4, "Bb", "WAVEFRONT"),
69}
72def make_detector_tuples(records: Iterable[DimensionRecord]) -> dict[int, tuple[str, int, str, str]]:
73 """Make tuples with the same entries as DETECTOR_TUPLES from an iterable of
74 detector dimension records.
76 Parameters
77 ----------
78 records : `~collections.abc.Iterable` [ `.dimensions.DimensionRecord` ]
79 Detector dimension records.
81 Returns
82 -------
83 tuples : `dict` [ `int`, `tuple` ]
84 Dictionary mapping detector ID to tuples with the same fields as the
85 ``DETECTOR_TUPLES`` constant in this file.
86 """
87 return {record.id: (record.instrument, record.id, record.full_name, record.purpose) for record in records}
90class ButlerQueryTests(ABC, TestCaseMixin):
91 """Base class for unit tests that test `lsst.daf.butler.Butler.query`
92 implementations.
93 """
95 @abstractmethod
96 def make_butler(self, *args: str) -> Butler:
97 """Make Butler instance populated with data used in the tests below.
99 Parameters
100 ----------
101 *args : str
102 Names of the files to pass to `load_data`.
104 Returns
105 -------
106 butler : `Butler`
107 Butler to use for tests.
108 """
109 raise NotImplementedError()
111 def load_data(self, butler: Butler, filename: str) -> None:
112 """Load registry test data from ``data_dir/<filename>``,
113 which should be a YAML import/export file.
115 This method should be called from implementations of `make_butler`
116 where the Registry should exist.
118 Parameters
119 ----------
120 butler : `~lsst.daf.butler.Butler`
121 The butler to use.
122 filename : `str`
123 Location of test data.
124 """
125 butler.import_(
126 filename=f"resource://lsst.daf.butler/tests/registry_data/{filename}",
127 without_datastore=True,
128 )
130 def check_detector_records(
131 self,
132 results: DimensionRecordQueryResults,
133 ids: Sequence[int] = (1, 2, 3, 4),
134 ordered: bool = False,
135 messages: Iterable[str] = (),
136 doomed: bool = False,
137 has_postprocessing: bool = False,
138 ) -> None:
139 self.assertEqual(results.element.name, "detector")
140 self.assertEqual(results.dimensions, results.dimensions.universe["detector"].minimal_group)
141 if has_postprocessing and not doomed:
142 self.assertEqual(results.count(discard=True), len(ids))
143 self.assertGreaterEqual(results.count(discard=False, exact=False), len(ids))
144 with self.assertRaisesRegex(InvalidQueryError, "^Cannot count query rows"):
145 results.count()
146 else:
147 self.assertEqual(results.count(discard=True), len(ids))
148 self.assertEqual(results.count(discard=False), len(ids))
149 self.assertEqual(results.count(discard=True, exact=False), len(ids))
150 self.assertEqual(results.count(discard=False, exact=False), len(ids))
151 self.assertEqual(results.any(), bool(ids))
152 if not doomed:
153 self.assertTrue(results.any(exact=False, execute=False))
154 with self.assertRaisesRegex(InvalidQueryError, "^Cannot obtain exact"):
155 results.any(exact=True, execute=False)
156 else:
157 self.assertFalse(results.any(exact=False, execute=False))
158 self.assertFalse(results.any(exact=True, execute=False))
159 self.assertCountEqual(results.explain_no_results(), list(messages))
160 self.check_detector_records_returned(list(results), ids=ids, ordered=ordered)
162 def check_detector_records_returned(
163 self,
164 results: list[DimensionRecord],
165 ids: Sequence[int] = (1, 2, 3, 4),
166 ordered: bool = False,
167 ) -> None:
168 expected = [DETECTOR_TUPLES[i] for i in ids]
169 queried = list(make_detector_tuples(results).values())
170 if ordered:
171 self.assertEqual(queried, expected)
172 else:
173 self.assertCountEqual(queried, expected)
175 def test_simple_record_query(self) -> None:
176 """Test query-system basics with simple queries for dimension
177 records.
179 This includes tests for order_by, limit, and where expressions, but
180 only for cases where there are no datasets, dimension projections,
181 or spatial/temporal overlaps.
182 """
183 butler = self.make_butler("base.yaml")
184 with butler.query() as query:
185 _x = query.expression_factory
186 results = query.dimension_records("detector")
187 self.check_detector_records(results)
188 self.check_detector_records_returned(butler.query_dimension_records("detector"))
189 self.assertEqual(len(butler.query_dimension_records("detector", limit=0)), 0)
190 self.check_detector_records(results.order_by("detector"), ordered=True)
191 self.check_detector_records_returned(
192 butler.query_dimension_records("detector", order_by="detector"), ordered=True
193 )
194 self.check_detector_records(
195 results.order_by(_x.detector.full_name.desc), [4, 3, 2, 1], ordered=True
196 )
197 self.check_detector_records_returned(
198 butler.query_dimension_records("detector", order_by="-full_name"),
199 ids=[4, 3, 2, 1],
200 ordered=True,
201 )
202 self.check_detector_records(results.order_by("detector").limit(2), [1, 2], ordered=True)
203 self.check_detector_records_returned(
204 butler.query_dimension_records("detector", limit=2, order_by="detector"),
205 ids=[1, 2],
206 ordered=True,
207 )
208 with self.assertLogs("lsst.daf.butler", level="WARNING") as wcm:
209 self.check_detector_records_returned(
210 butler.query_dimension_records("detector", limit=-2, order_by="-detector"),
211 ids=[4, 3],
212 ordered=True,
213 )
214 self.assertIn("More dimension records are available", wcm.output[0])
215 self.check_detector_records(results.where(_x.detector.raft == "B", instrument="Cam1"), [3, 4])
216 self.check_detector_records_returned(
217 butler.query_dimension_records(
218 "detector", where="detector.raft = :R", bind={"R": "B"}, instrument="Cam1"
219 ),
220 ids=[3, 4],
221 )
222 self.check_detector_records(
223 results.where(_x.detector.full_name.glob("B?"), instrument="Cam1"), [3, 4]
224 )
225 self.check_detector_records(
226 results.where(_x.detector.full_name.glob("*a"), instrument="Cam1"), [1, 3]
227 )
229 # Test incorrect type for glob() parameter.
230 with self.assertRaises(InvalidQueryError):
231 results.where(_x.detector.full_name.glob(1), instrument="Cam1") # type: ignore[arg-type]
233 def test_simple_data_coordinate_query(self) -> None:
234 butler = self.make_butler("base.yaml")
236 expected_detectors = [1, 2, 3, 4]
237 universe = butler.dimensions
238 expected_coordinates = [
239 DataCoordinate.standardize({"instrument": "Cam1", "detector": x}, universe=universe)
240 for x in expected_detectors
241 ]
243 with butler.query() as query:
244 # Test empty query
245 empty = DataCoordinate.make_empty(butler.dimensions)
246 self.assertCountEqual(list(query.data_ids([])), [empty])
247 self.assertCountEqual(butler.query_data_ids([]), [empty])
249 # Test query for a single dimension
250 results = query.data_ids(["detector"])
251 self.assertCountEqual(list(results), expected_coordinates)
253 # Limit.
254 results = query.data_ids(["detector"]).order_by("-detector").limit(2)
255 self.assertCountEqual(list(results), expected_coordinates[2:])
257 data_ids = butler.query_data_ids("detector")
258 self.assertCountEqual(data_ids, expected_coordinates)
260 data_ids = butler.query_data_ids("detector", order_by="-detector", limit=2)
261 self.assertCountEqual(data_ids, expected_coordinates[2:])
263 with self.assertLogs("lsst.daf.butler", level="WARNING") as wcm:
264 data_ids = butler.query_data_ids("detector", order_by="-detector", limit=-2)
265 self.assertCountEqual(data_ids, expected_coordinates[2:])
266 self.assertIn("More data IDs are available", wcm.output[0])
268 data_ids = butler.query_data_ids("detector", limit=0)
269 self.assertEqual(len(data_ids), 0)
271 def test_simple_dataset_query(self) -> None:
272 butler = self.make_butler("base.yaml", "datasets.yaml")
273 with butler.query() as query:
274 refs_q = list(query.datasets("bias", "imported_g").order_by("detector"))
275 refs_simple = butler.query_datasets("bias", "imported_g", order_by="detector")
276 self.assertCountEqual(refs_q, refs_simple)
278 for refs in (refs_q, refs_simple):
279 self.assertEqual(len(refs), 3)
280 self.assertEqual(refs[0].id, UUID("e15ab039-bc8b-4135-87c5-90902a7c0b22"))
281 self.assertEqual(refs[1].id, UUID("51352db4-a47a-447c-b12d-a50b206b17cd"))
282 for detector, ref in enumerate(refs, 1):
283 self.assertEqual(ref.datasetType.name, "bias")
284 self.assertEqual(ref.dataId["instrument"], "Cam1")
285 self.assertEqual(ref.dataId["detector"], detector)
286 self.assertEqual(ref.run, "imported_g")
288 # Try again with limit.
289 with butler.query() as query:
290 refs_q = list(query.datasets("bias", "imported_g").order_by("detector").limit(2))
291 refs_simple = butler.query_datasets("bias", "imported_g", order_by="detector", limit=2)
292 self.assertCountEqual(refs_q, refs_simple)
293 self.assertEqual(len(refs_q), 2)
294 self.assertEqual(refs_q[0].id, UUID("e15ab039-bc8b-4135-87c5-90902a7c0b22"))
295 self.assertEqual(refs_q[1].id, UUID("51352db4-a47a-447c-b12d-a50b206b17cd"))
297 # limit=0 means test the query but don't return anything and
298 # don't complain.
299 refs_simple = butler.query_datasets("bias", "imported_g", limit=0, explain=True)
300 self.assertEqual(len(refs_simple), 0)
302 # Explicitly run with no restrictions.
303 refs_simple = butler.query_datasets("bias", collections="*", find_first=False, limit=None)
304 self.assertEqual(len(refs_simple), 6)
306 # Now limit the number of results and look for a warning.
307 with self.assertLogs("lsst.daf.butler", level="WARNING") as lcm:
308 refs_simple = butler.query_datasets("bias", collections="*", find_first=False, limit=-4)
309 self.assertEqual(len(refs_simple), 4)
310 self.assertIn("More datasets are available", lcm.output[0])
312 with self.assertRaises(InvalidQueryError) as cm:
313 butler.query_datasets("bias", "*", detector=100, instrument="Unknown", find_first=True)
314 self.assertIn("Can not use wildcards", str(cm.exception))
315 with self.assertRaises(EmptyQueryResultError) as cm2:
316 butler.query_datasets("bias", "*", detector=100, instrument="Unknown", find_first=False)
317 self.assertIn("doomed", str(cm2.exception))
319 # Test for a regression of an issue where "band" was not being included
320 # in the data ID, despite being one of the dimensions in the "flat"
321 # dataset type.
322 #
323 # "band" is implied by "physical_filter", so it's technically not a
324 # 'required' dimension. However, the contract of query_datasets is
325 # that hasFull() should be true, so implied dimensions must be
326 # included.
327 refs = butler.query_datasets("flat", "imported_r", where="detector = 2", instrument="Cam1")
328 self.assertEqual(len(refs), 1)
329 flat = refs[0]
330 self.assertTrue(flat.dataId.hasFull())
331 self.assertEqual(flat.datasetType.name, "flat")
332 self.assertEqual(flat.dataId["instrument"], "Cam1")
333 self.assertEqual(flat.dataId["detector"], 2)
334 self.assertEqual(flat.dataId["physical_filter"], "Cam1-R1")
335 self.assertEqual(flat.dataId["band"], "r")
337 def test_general_query(self) -> None:
338 """Test Query.general and its result."""
339 butler = self.make_butler("base.yaml", "datasets.yaml")
340 dimensions = butler.dimensions["detector"].minimal_group
342 # Do simple dimension queries.
343 with butler.query() as query:
344 query = query.join_dimensions(dimensions)
345 rows = list(query.general(dimensions).order_by("detector"))
346 self.assertEqual(
347 rows,
348 [
349 {"instrument": "Cam1", "detector": 1},
350 {"instrument": "Cam1", "detector": 2},
351 {"instrument": "Cam1", "detector": 3},
352 {"instrument": "Cam1", "detector": 4},
353 ],
354 )
355 rows = list(
356 query.general(dimensions, "detector.full_name", "purpose").order_by(
357 "-detector.purpose", "full_name"
358 )
359 )
360 self.assertEqual(
361 rows,
362 [
363 {
364 "instrument": "Cam1",
365 "detector": 4,
366 "detector.full_name": "Bb",
367 "detector.purpose": "WAVEFRONT",
368 },
369 {
370 "instrument": "Cam1",
371 "detector": 1,
372 "detector.full_name": "Aa",
373 "detector.purpose": "SCIENCE",
374 },
375 {
376 "instrument": "Cam1",
377 "detector": 2,
378 "detector.full_name": "Ab",
379 "detector.purpose": "SCIENCE",
380 },
381 {
382 "instrument": "Cam1",
383 "detector": 3,
384 "detector.full_name": "Ba",
385 "detector.purpose": "SCIENCE",
386 },
387 ],
388 )
389 rows = list(
390 query.general(dimensions, "detector.full_name", "purpose").where(
391 "instrument = 'Cam1' AND purpose = 'WAVEFRONT'"
392 )
393 )
394 self.assertEqual(
395 rows,
396 [
397 {
398 "instrument": "Cam1",
399 "detector": 4,
400 "detector.full_name": "Bb",
401 "detector.purpose": "WAVEFRONT",
402 },
403 ],
404 )
405 result = query.general(dimensions, dimension_fields={"detector": {"full_name"}})
406 self.assertEqual(set(row["detector.full_name"] for row in result), {"Aa", "Ab", "Ba", "Bb"})
408 # Use "flat" whose dimension group includes implied dimension.
409 flat = butler.get_dataset_type("flat")
410 dimensions = butler.dimensions.conform(["detector", "physical_filter"])
412 # Do simple dataset queries in RUN collection.
413 with butler.query() as query:
414 query = query.join_dataset_search("flat", "imported_g")
415 # This just returns data IDs.
416 rows = list(query.general(dimensions).order_by("detector"))
417 self.assertEqual(
418 rows,
419 [
420 {"instrument": "Cam1", "detector": 2, "physical_filter": "Cam1-G", "band": "g"},
421 {"instrument": "Cam1", "detector": 3, "physical_filter": "Cam1-G", "band": "g"},
422 {"instrument": "Cam1", "detector": 4, "physical_filter": "Cam1-G", "band": "g"},
423 ],
424 )
426 result = query.general(dimensions, dataset_fields={"flat": ...}, find_first=True).order_by(
427 "detector"
428 )
429 ids = {row["flat.dataset_id"] for row in result}
430 self.assertEqual(
431 ids,
432 {
433 UUID("60c8a65c-7290-4c38-b1de-e3b1cdcf872d"),
434 UUID("84239e7f-c41f-46d5-97b9-a27976b98ceb"),
435 UUID("fd51bce1-2848-49d6-a378-f8a122f5139a"),
436 },
437 )
439 # Check what iter_tuples() returns
440 row_tuples = list(result.iter_tuples(flat))
441 self.assertEqual(len(row_tuples), 3)
442 for row_tuple in row_tuples:
443 self.assertEqual(len(row_tuple.refs), 1)
444 self.assertEqual(row_tuple.refs[0].datasetType, flat)
445 self.assertTrue(row_tuple.refs[0].dataId.hasFull())
446 self.assertFalse(row_tuple.refs[0].dataId.hasRecords())
447 self.assertTrue(row_tuple.data_id.hasFull())
448 self.assertFalse(row_tuple.data_id.hasRecords())
449 self.assertEqual(row_tuple.data_id.dimensions, dimensions)
450 self.assertEqual(row_tuple.raw_row["flat.run"], "imported_g")
452 flat1, flat2, flat3 = (row_tuple.refs[0] for row_tuple in row_tuples)
454 # Query datasets CALIBRATION/TAGGED collections.
455 butler.registry.registerCollection("tagged", CollectionType.TAGGED)
456 butler.registry.registerCollection("calib", CollectionType.CALIBRATION)
458 # Add two refs to tagged collection.
459 butler.registry.associate("tagged", [flat1, flat2])
461 # Certify some calibs.
462 t1 = astropy.time.Time("2020-01-01T01:00:00", format="isot", scale="tai")
463 t2 = astropy.time.Time("2020-01-01T02:00:00", format="isot", scale="tai")
464 t3 = astropy.time.Time("2020-01-01T03:00:00", format="isot", scale="tai")
465 butler.registry.certify("calib", [flat1], Timespan(t1, t2))
466 butler.registry.certify("calib", [flat3], Timespan(t2, t3))
467 butler.registry.certify("calib", [flat1], Timespan(t3, None))
468 butler.registry.certify("calib", [flat2], Timespan.makeEmpty())
470 # Query tagged collection.
471 with butler.query() as query:
472 query = query.join_dataset_search("flat", ["tagged"])
474 result = query.general(
475 dimensions, "flat.dataset_id", "flat.run", "flat.collection", find_first=False
476 )
477 row_tuples = list(result.iter_tuples(flat))
478 self.assertEqual(len(row_tuples), 2)
479 self.assertEqual({row_tuple.refs[0] for row_tuple in row_tuples}, {flat1, flat2})
480 self.assertEqual({row_tuple.raw_row["flat.collection"] for row_tuple in row_tuples}, {"tagged"})
482 # Query calib collection.
483 with butler.query() as query:
484 query = query.join_dataset_search("flat", ["calib"])
485 result = query.general(
486 dimensions,
487 "flat.dataset_id",
488 "flat.run",
489 "flat.collection",
490 "flat.timespan",
491 find_first=False,
492 )
493 row_tuples = list(result.iter_tuples(flat))
494 self.assertEqual(len(row_tuples), 4)
495 self.assertEqual({row_tuple.refs[0] for row_tuple in row_tuples}, {flat1, flat2, flat3})
496 self.assertEqual({row_tuple.raw_row["flat.collection"] for row_tuple in row_tuples}, {"calib"})
497 self.assertEqual(
498 {row_tuple.raw_row["flat.timespan"] for row_tuple in row_tuples},
499 {Timespan(t1, t2), Timespan(t2, t3), Timespan(t3, None), Timespan.makeEmpty()},
500 )
502 # Query both tagged and calib collection.
503 with butler.query() as query:
504 query = query.join_dataset_search("flat", ["tagged", "calib"])
505 result = query.general(
506 dimensions,
507 "flat.dataset_id",
508 "flat.run",
509 "flat.collection",
510 "flat.timespan",
511 find_first=False,
512 )
513 row_tuples = list(result.iter_tuples(flat))
514 self.assertEqual(len(row_tuples), 6)
515 self.assertEqual(
516 {row_tuple.raw_row["flat.collection"] for row_tuple in row_tuples}, {"calib", "tagged"}
517 )
518 self.assertEqual(
519 {row_tuple.raw_row["flat.timespan"] for row_tuple in row_tuples},
520 {Timespan(t1, t2), Timespan(t2, t3), Timespan(t3, None), Timespan.makeEmpty(), None},
521 )
523 dimensions = butler.dimensions["detector"].minimal_group
525 # Include dimension records into query.
526 with butler.query() as query:
527 query = query.join_dimensions(dimensions)
528 result = query.general(dimensions).order_by("detector")
529 rows = list(result.with_dimension_records())
530 self.assertEqual(
531 rows[0],
532 {
533 "instrument": "Cam1",
534 "detector": 1,
535 "instrument.visit_max": 1024,
536 "instrument.visit_system": 1,
537 "instrument.exposure_max": 512,
538 "instrument.detector_max": 4,
539 "instrument.class_name": "lsst.pipe.base.Instrument",
540 "detector.full_name": "Aa",
541 "detector.name_in_raft": "a",
542 "detector.raft": "A",
543 "detector.purpose": "SCIENCE",
544 },
545 )
547 dimensions = butler.dimensions.conform(["detector", "physical_filter"])
549 # DataIds should come with records.
550 with butler.query() as query:
551 query = query.join_dataset_search("flat", "imported_g")
552 result = query.general(dimensions, dataset_fields={"flat": ...}, find_first=True).order_by(
553 "detector"
554 )
555 result = result.with_dimension_records()
556 row_tuples = list(result.iter_tuples(flat))
557 self.assertEqual(len(row_tuples), 3)
558 for row_tuple in row_tuples:
559 self.assertTrue(row_tuple.data_id.hasRecords())
560 self.assertEqual(len(row_tuple.refs), 1)
561 self.assertTrue(row_tuple.refs[0].dataId.hasRecords())
563 def test_query_ingest_date(self) -> None:
564 """Test general query returning ingest_date field."""
565 before_ingest = astropy.time.Time.now()
566 butler = self.make_butler("base.yaml", "datasets.yaml")
567 dimensions = butler.dimensions.conform(["detector", "physical_filter"])
569 # Check that returned type of ingest_date is astropy Time, must work
570 # for schema versions 1 and 2 of datasets manager.
571 with butler.query() as query:
572 query = query.join_dataset_search("flat", "imported_g")
573 rows = list(query.general(dimensions, dataset_fields={"flat": {"ingest_date"}}, find_first=False))
574 self.assertEqual(len(rows), 3)
575 for row in rows:
576 self.assertIsInstance(row["flat.ingest_date"], astropy.time.Time)
578 # Check that WHERE accepts astropy time
579 with butler.query() as query:
580 query = query.join_dataset_search("flat", "imported_g")
581 query1 = query.where("flat.ingest_date < :before_ingest", bind={"before_ingest": before_ingest})
582 rows = list(query1.general(dimensions))
583 self.assertEqual(len(rows), 0)
584 query1 = query.where("flat.ingest_date >= :before_ingest", bind={"before_ingest": before_ingest})
585 rows = list(query1.general(dimensions))
586 self.assertEqual(len(rows), 3)
587 # Same with a time in string literal.
588 query1 = query.where(f"flat.ingest_date < T'mjd/{before_ingest.tai.mjd}'")
589 rows = list(query1.general(dimensions))
590 self.assertEqual(len(rows), 0)
592 def test_implied_union_record_query(self) -> None:
593 """Test queries for a dimension ('band') that uses "implied union"
594 storage, in which its values are the union of the values for it in a
595 another dimension (physical_filter) that implies it.
596 """
597 butler = self.make_butler("base.yaml")
598 band = butler.dimensions["band"]
599 self.assertEqual(band.implied_union_target, butler.dimensions["physical_filter"])
600 with butler.query() as query:
601 self.assertCountEqual(
602 list(query.dimension_records("band")),
603 [band.RecordClass(name="g"), band.RecordClass(name="r")],
604 )
605 self.assertCountEqual(
606 list(query.where(physical_filter="Cam1-R1", instrument="Cam1").dimension_records("band")),
607 [band.RecordClass(name="r")],
608 )
609 self.assertCountEqual(
610 butler.query_dimension_records("band"),
611 [band.RecordClass(name="g"), band.RecordClass(name="r")],
612 )
613 self.assertCountEqual(
614 butler.query_dimension_records("band", physical_filter="Cam1-R1", instrument="Cam1"),
615 [band.RecordClass(name="r")],
616 )
618 def test_dataset_constrained_record_query(self) -> None:
619 """Test a query for dimension records constrained by the existence of
620 datasets of a particular type.
621 """
622 butler = self.make_butler("base.yaml", "datasets.yaml")
623 butler.registry.insertDimensionData("instrument", {"name": "Cam2"})
624 butler.collections.register("empty", CollectionType.RUN)
625 butler.collections.register("chain", CollectionType.CHAINED)
626 butler.collections.redefine_chain("chain", ["imported_g", "empty", "imported_r"])
627 with butler.query() as query:
628 # No collections here or in defaults is an error.
629 with self.assertRaises(NoDefaultCollectionError):
630 query.join_dataset_search("bias").dimension_records("detector").any()
631 butler.registry.defaults = RegistryDefaults(collections=["chain"])
632 with butler.query() as query:
633 _x = query.expression_factory
634 # Simplest case: this collection only has the first 3 detectors.
635 self.check_detector_records(
636 query.join_dataset_search("bias", collections=["imported_g"]).dimension_records("detector"),
637 [1, 2, 3],
638 )
639 # Together these collections have two biases for two of the
640 # detectors, but this shouldn't cause duplicate results.
641 self.check_detector_records(
642 query.join_dataset_search("bias", collections=["imported_g", "imported_r"]).dimension_records(
643 "detector"
644 ),
645 )
646 # Again we've got the potential for duplicates due to multiple
647 # datasets with the same data ID, and this time we force the
648 # deduplication to happen outside the dataset-search subquery by
649 # adding a WHERE filter on a dataset column. We also use the
650 # defaulted collection ('chain') to supply the collection.
651 self.check_detector_records(
652 query.join_dataset_search("bias")
653 .where(
654 _x.any(
655 _x.all(_x["bias"].collection == "imported_g", _x.detector.raft == "B"),
656 _x.all(_x["bias"].collection == "imported_r", _x.detector.raft == "A"),
657 ),
658 instrument="Cam1",
659 )
660 .dimension_records("detector"),
661 [2, 3],
662 )
663 # Flats have dimensions (physical_filter and band) we would
664 # normally include in query for detector records. This also should
665 # not cause duplicates.
666 self.check_detector_records(
667 query.join_dataset_search("flat", collections=["imported_g"]).dimension_records("detector"),
668 [2, 3, 4],
669 )
670 # No results, but for reasons we can't determine before we run the
671 # query.
672 self.check_detector_records(
673 query.join_dataset_search("flat", collections=["imported_g"])
674 .where(_x.band == "r")
675 .dimension_records("detector"),
676 [],
677 )
678 # No results, and we can diagnose why before we run the query.
679 self.check_detector_records(
680 query.join_dataset_search("bias", collections=["empty"]).dimension_records("detector"),
681 [],
682 messages=[
683 "Search for dataset type 'bias' in ['empty'] is doomed to fail.",
684 "No datasets of type 'bias' in collection 'empty'.",
685 ],
686 doomed=True,
687 )
688 self.check_detector_records(
689 query.join_dataset_search("bias", collections=["imported_g"])
690 .where(instrument="Cam2")
691 .dimension_records("detector"),
692 [],
693 messages=[
694 "Search for dataset type 'bias' in ['imported_g'] is doomed to fail.",
695 "No datasets with instrument='Cam2' in collection 'imported_g'.",
696 ],
697 doomed=True,
698 )
700 def test_duplicate_overlaps(self) -> None:
701 """Test for query option that enables duplicate rows in queries that
702 use skypix overalps.
703 """
704 butler = self.make_butler("base.yaml", "spatial.yaml")
705 butler.registry.defaults = RegistryDefaults(instrument="Cam1", skymap="SkyMap1")
706 with butler.query() as query:
707 data_ids = list(query.data_ids(["visit", "detector", "patch"]).where(visit=1, detector=1))
708 self.assertCountEqual(
709 [(data_id["tract"], data_id["patch"]) for data_id in data_ids], [(0, 0), (0, 2)]
710 )
712 query._allow_duplicate_overlaps = True
713 data_ids = list(query.data_ids(["visit", "detector", "patch"]).where(visit=1, detector=1))
714 self.assertCountEqual(
715 [(data_id["tract"], data_id["patch"]) for data_id in data_ids],
716 [(0, 0), (0, 0), (0, 2), (0, 2)],
717 )
719 def test_spatial_overlaps(self) -> None:
720 """Test queries for dimension records with spatial overlaps.
722 Run tests/data/registry/spatial.py to plot the various regions used in
723 this test.
724 """
725 butler = self.make_butler("base.yaml", "spatial.yaml")
726 # Set default governor data ID values both to test that code path and
727 # to keep us from having to repeat them in every 'where' call below.
728 butler.registry.defaults = RegistryDefaults(instrument="Cam1", skymap="SkyMap1")
729 htm7 = butler.dimensions.skypix_dimensions["htm7"]
730 with butler.query() as query:
731 _x = query.expression_factory
732 # Query for detectors from a particular visit that overlap an
733 # explicit region.
734 self.check_detector_records(
735 query.where(
736 _x.visit_detector_region.region.overlaps(htm7.pixelization.pixel(253954)),
737 visit=1,
738 ).dimension_records("detector"),
739 [1, 3, 4],
740 has_postprocessing=True,
741 )
742 self.check_detector_records_returned(
743 butler.query_dimension_records(
744 "detector",
745 where="visit_detector_region.region OVERLAPS :region",
746 bind={"region": htm7.pixelization.pixel(253954)},
747 visit=1,
748 ),
749 ids=[1, 3, 4],
750 )
751 # Query for detectors from a particular visit that overlap an htm7
752 # ID. This is basically the same query as the last one, but
753 # expressed as a spatial join, and we can recognize that
754 # postprocessing is not needed (while in the last case it did
755 # nothing, but we couldn't tell that in advance because the query
756 # didn't know the region came from htm7).
757 self.check_detector_records(
758 query.where(
759 _x.visit_detector_region.region.overlaps(_x.htm7.region),
760 visit=1,
761 htm7=253954,
762 ).dimension_records("detector"),
763 [1, 3, 4],
764 has_postprocessing=False,
765 )
766 # Repeat the last query but with the spatial join implicit rather
767 # than explicit.
768 self.check_detector_records(
769 query.where(
770 visit=1,
771 htm7=253954,
772 ).dimension_records("detector"),
773 [1, 3, 4],
774 has_postprocessing=False,
775 )
776 self.check_detector_records_returned(
777 butler.query_dimension_records(
778 "detector",
779 visit=1,
780 htm7=253954,
781 ),
782 ids=[1, 3, 4],
783 )
784 # Query for the detectors from any visit that overlap a region:
785 # this gets contributions from multiple visits, and would have
786 # duplicates if we didn't get rid of them via GROUP BY.
787 self.check_detector_records(
788 query.where(
789 _x.visit_detector_region.region.overlaps(htm7.pixelization.pixel(253954)),
790 ).dimension_records("detector"),
791 [1, 2, 3, 4],
792 has_postprocessing=True,
793 )
794 self.check_detector_records_returned(
795 butler.query_dimension_records(
796 "detector",
797 where="visit_detector_region.region OVERLAPS :region",
798 bind={"region": htm7.pixelization.pixel(253954)},
799 ),
800 ids=[1, 2, 3, 4],
801 )
802 # Once again we rewrite the region-constraint query as a spatial
803 # join, which drops the postprocessing. This join has to be
804 # explicit because `visit` no longer gets into the query dimensions
805 # some other way, and without it `detector` is not spatial.
806 self.check_detector_records(
807 query.where(
808 _x.visit_detector_region.region.overlaps(_x.htm7.region),
809 htm7=253954,
810 ).dimension_records("detector"),
811 [1, 2, 3, 4],
812 has_postprocessing=False,
813 )
814 # Query for detectors from any visit that overlap a patch. This
815 # requires joining visit_detector_region to htm7 and htm7 to patch,
816 # and then some postprocessing. We want to make sure there are no
817 # duplicates from a detector and patch both overlapping multiple
818 # htm7 tiles (which affects detectors 1 and 2) and that
819 # postprocessing filters out detector 4, which has one htm7 tile in
820 # common with the patch but does not actually overlap it.
821 self.check_detector_records(
822 query.where(
823 _x.visit_detector_region.region.overlaps(_x.patch.region),
824 tract=0,
825 patch=4,
826 ).dimension_records("detector"),
827 [1, 2, 3],
828 has_postprocessing=True,
829 )
830 # Same as above, but with a materialization.
831 self.check_detector_records(
832 query.where(
833 _x.visit_detector_region.region.overlaps(_x.patch.region),
834 tract=0,
835 patch=4,
836 )
837 .materialize(dimensions=["detector"])
838 .dimension_records("detector"),
839 [1, 2, 3],
840 has_postprocessing=True,
841 )
842 # Query for that patch's region and express the previous query as
843 # a region-constraint instead of a spatial join.
844 (patch_record,) = query.where(tract=0, patch=4).dimension_records("patch")
845 self.check_detector_records(
846 query.where(
847 _x.visit_detector_region.region.overlaps(patch_record.region),
848 ).dimension_records("detector"),
849 [1, 2, 3],
850 has_postprocessing=True,
851 )
852 self.check_detector_records_returned(
853 butler.query_dimension_records(
854 "detector",
855 where="visit_detector_region.region OVERLAPS :region",
856 bind={"region": patch_record.region},
857 ),
858 ids=[1, 2, 3],
859 )
860 # Query for detectors where a patch/visit+detector overlap is
861 # satisfied, in the case where there are no rows with an overlap,
862 # but the union of the patch regions overlaps the union of the
863 # visit+detector regions.
864 self.check_detector_records(
865 query.where(
866 _x.visit_detector_region.region.overlaps(_x.patch.region),
867 _x.any(
868 _x.all(_x.tract == 1, _x.visit == 1),
869 _x.all(_x.tract == 0, _x.patch == 0, _x.visit == 2),
870 ),
871 ).dimension_records("detector"),
872 [],
873 has_postprocessing=True,
874 )
875 # Combine postprocessing with order_by and limit.
876 self.check_detector_records(
877 query.where(
878 _x.visit_detector_region.region.overlaps(patch_record.region),
879 )
880 .dimension_records("detector")
881 .order_by(_x.detector.desc)
882 .limit(2),
883 [3, 2],
884 has_postprocessing=True,
885 )
886 self.check_detector_records_returned(
887 butler.query_dimension_records(
888 "detector",
889 where="visit_detector_region.region OVERLAPS :region",
890 bind={"region": patch_record.region},
891 order_by="-detector",
892 limit=2,
893 ),
894 ids=[3, 2],
895 )
896 # Try a case where there are some records before postprocessing but
897 # none afterwards.
898 self.check_detector_records(
899 query.where(
900 _x.visit_detector_region.region.overlaps(patch_record.region),
901 detector=4,
902 ).dimension_records("detector"),
903 [],
904 has_postprocessing=True,
905 )
906 self.check_detector_records_returned(
907 butler.query_dimension_records(
908 "detector",
909 where="visit_detector_region.region OVERLAPS :region",
910 bind={"region": patch_record.region},
911 detector=4,
912 explain=False,
913 ),
914 ids=[],
915 )
916 # Check spatial queries using points instead of regions.
917 # This (ra, dec) is a point in the center of the region for visit
918 # 1, detector 3.
919 ra = 0.25209391431545386 # degrees
920 dec = 0.9269112711026793 # degrees
922 def _check_visit_id(query: Query) -> None:
923 result = list(query.data_ids(["visit", "detector"]))
924 self.assertEqual(len(result), 1)
925 id = result[0]
926 self.assertEqual(id["visit"], 1)
927 self.assertEqual(id["detector"], 3)
929 # Basic POINT() syntax.
930 _check_visit_id(query.where(f"visit_detector_region.region OVERLAPS POINT({ra}, {dec})"))
931 _check_visit_id(query.where(f"POINT({ra}, {dec}) OVERLAPS visit_detector_region.region"))
933 # dec of 1 is close enough to still be in the region, and tests
934 # conversion of integer to float.
935 _check_visit_id(query.where(f"visit_detector_region.region OVERLAPS POINT({ra}, 1)"))
937 # Negative values are allowed for dec, since it's defined as -90 to
938 # 90. Tract 1, patch 4 slightly overlaps some negative dec values.
939 result = list(query.where("patch.region OVERLAPS POINT(0.335, -0.000000001)").data_ids(["patch"]))
940 self.assertEqual(len(result), 1)
941 id = result[0]
942 self.assertEqual(id["patch"], 4)
943 self.assertEqual(id["tract"], 1)
944 # Out of bounds dec values are not allowed.
945 with self.assertRaisesRegex(ValueError, "invalid latitude angle"):
946 list(query.where("patch.region OVERLAPS POINT(0.335, -91)").data_ids(["patch"]))
948 # Negative ra values are allowed.
949 _check_visit_id(query.where(f"POINT({ra - 360}, {dec}) OVERLAPS visit_detector_region.region"))
951 # Substitute ra and dec values via bind instead of literals in the
952 # string.
953 _check_visit_id(
954 query.where(
955 "visit_detector_region.region OVERLAPS POINT(:ra, :dec)", bind={"ra": ra, "dec": dec}
956 )
957 )
959 # Bind in a point object instead of specifying ra/dec separately.
960 _check_visit_id(
961 query.where(
962 "visit_detector_region.region OVERLAPS :my_point",
963 bind={"my_point": LonLat.fromDegrees(ra, dec)},
964 )
965 )
966 _check_visit_id(
967 query.where(
968 "visit_detector_region.region OVERLAPS :my_point",
969 bind={"my_point": astropy.coordinates.SkyCoord(ra, dec, frame="icrs", unit="deg")},
970 )
971 )
972 # Make sure alternative coordinate frames in astropy SkyCoord are
973 # handled.
974 _check_visit_id(
975 query.where(
976 "visit_detector_region.region OVERLAPS :my_point",
977 bind={
978 "my_point": astropy.coordinates.SkyCoord(
979 ra, dec, frame="icrs", unit="deg"
980 ).transform_to("galactic")
981 },
982 )
983 )
985 # Compare against literal values using ExpressionFactory.
986 _check_visit_id(
987 query.where(_x.visit_detector_region.region.overlaps(LonLat.fromDegrees(ra, dec)))
988 )
989 _check_visit_id(
990 query.where(
991 _x.visit_detector_region.region.overlaps(
992 astropy.coordinates.SkyCoord(ra, dec, frame="icrs", unit="deg")
993 )
994 )
995 )
997 # Check errors for invalid syntax.
998 with self.assertRaisesRegex(
999 InvalidQueryError, r"Expression 'visit.id' in POINT\(\) is not a literal number."
1000 ):
1001 query.where(f"visit_detector_region.region OVERLAPS POINT(visit.id, {dec})")
1002 with self.assertRaisesRegex(
1003 InvalidQueryError, r"Expression ''not-a-number'' in POINT\(\) is not a literal number."
1004 ):
1005 query.where(f"visit_detector_region.region OVERLAPS POINT({ra}, 'not-a-number')")
1007 # astropy's SkyCoord can be array-valued, but we expect only a
1008 # single point.
1009 array_point = astropy.coordinates.SkyCoord(
1010 ra=[10, 11, 12, 13], dec=[41, -5, 42, 0], unit="deg", frame="icrs"
1011 )
1012 with self.assertRaisesRegex(ValueError, "Astropy SkyCoord contained an array of points"):
1013 query.where(
1014 "visit_detector_region.region OVERLAPS :my_point",
1015 bind={"my_point": array_point},
1016 )
1018 def test_auto_spatial_joins(self) -> None:
1019 """Test the addition of automatic spatial joins in the presence and
1020 absence of datasets with dimensions that cross spatial families.
1021 """
1022 butler = self.make_butler("base.yaml", "spatial.yaml")
1023 # Set default governor data ID values both to test that code path and
1024 # to keep us from having to repeat them in every 'where' call below.
1025 butler.registry.defaults = RegistryDefaults(instrument="Cam1", skymap="SkyMap1")
1026 # Add some datasets with {tract, visit, detector} dimensions.
1027 # These will cover all {visit, detector}s that overlap tract=0.
1028 cat = DatasetType(
1029 "cat",
1030 dimensions=butler.dimensions.conform(["visit", "detector", "tract"]),
1031 storageClass="ArrowTable",
1032 )
1033 butler.registry.registerDatasetType(cat)
1034 butler.collections.register("run1")
1035 butler.registry.insertDatasets(
1036 cat,
1037 [
1038 {"instrument": "Cam1", "skymap": "SkyMap1", "visit": 1, "detector": 1, "tract": 0},
1039 {"instrument": "Cam1", "skymap": "SkyMap1", "visit": 1, "detector": 2, "tract": 0},
1040 {"instrument": "Cam1", "skymap": "SkyMap1", "visit": 1, "detector": 3, "tract": 0},
1041 {"instrument": "Cam1", "skymap": "SkyMap1", "visit": 1, "detector": 4, "tract": 0},
1042 {"instrument": "Cam1", "skymap": "SkyMap1", "visit": 2, "detector": 1, "tract": 0},
1043 {"instrument": "Cam1", "skymap": "SkyMap1", "visit": 2, "detector": 2, "tract": 0},
1044 ],
1045 run="run1",
1046 )
1047 with butler.query() as query:
1048 # When we query for just these dimensions with that dataset type
1049 # included, we shouldn't need a spatial join, because we assume
1050 # it's embedded in the rows of that dataset search.
1051 q1 = query.join_dataset_search(cat, "run1").data_ids(["visit", "detector", "tract"])
1052 # If there's no explicit spatial join, there's no postprocessing,
1053 # and hence we can do an exact count with discard=False.
1054 self.assertEqual(q1.count(exact=True, discard=False), 6)
1055 self.assertCountEqual(
1056 [(row["visit"], row["detector"], row["tract"]) for row in q1],
1057 [
1058 (1, 1, 0),
1059 (1, 2, 0),
1060 (1, 3, 0),
1061 (1, 4, 0),
1062 (2, 1, 0),
1063 (2, 2, 0),
1064 ],
1065 )
1066 # When we query for these dimensions and another one that wants
1067 # a more precise spatial join, we should get that more precise
1068 # spatial join.
1069 q1 = query.join_dataset_search(cat, "run1").data_ids(["visit", "detector", "tract", "patch"])
1070 self.assertCountEqual(
1071 [(row["visit"], row["detector"], row["tract"], row["patch"]) for row in q1],
1072 [
1073 (1, 1, 0, 0),
1074 (1, 1, 0, 2),
1075 (1, 2, 0, 0),
1076 (1, 2, 0, 1),
1077 (1, 2, 0, 2),
1078 (1, 2, 0, 3),
1079 (1, 3, 0, 2),
1080 (1, 3, 0, 4),
1081 (1, 4, 0, 2),
1082 (1, 4, 0, 3),
1083 (2, 1, 0, 4),
1084 (2, 2, 0, 4),
1085 (2, 2, 0, 5),
1086 ],
1087 )
1089 def test_common_skypix_overlaps(self) -> None:
1090 """Test spatial overlap queries that return htm7 records."""
1091 butler = self.make_butler("base.yaml", "spatial.yaml")
1092 # Insert some datasets that use a skypix dimension, since some queries
1093 # are only possible if a superset of the skypix IDs are in the query
1094 # already.
1095 cat1 = DatasetType("cat1", dimensions=butler.dimensions.conform(["htm7"]), storageClass="ArrowTable")
1096 butler.registry.registerDatasetType(cat1)
1097 butler.registry.registerCollection("refcats", CollectionType.RUN)
1098 butler.registry.insertDatasets(cat1, [{"htm7": i} for i in range(253952, 253968)], run="refcats")
1099 with butler.query() as query:
1100 _x = query.expression_factory
1101 # Explicit join to patch.
1102 self.assertCountEqual(
1103 [
1104 record.id
1105 for record in query.where(
1106 _x.htm7.region.overlaps(_x.patch.region), skymap="SkyMap1", tract=0, patch=4
1107 ).dimension_records("htm7")
1108 ],
1109 [253954, 253955],
1110 )
1111 # Implicit join to patch.
1112 self.assertCountEqual(
1113 [
1114 record.id
1115 for record in query.where(skymap="SkyMap1", tract=0, patch=4).dimension_records("htm7")
1116 ],
1117 [253954, 253955],
1118 )
1119 self.assertCountEqual(
1120 [
1121 record.id
1122 for record in butler.query_dimension_records("htm7", skymap="SkyMap1", tract=0, patch=4)
1123 ],
1124 [253954, 253955],
1125 )
1126 # Constraint on the patch region (with the query not knowing it
1127 # corresponds to that patch).
1128 (patch,) = query.where(skymap="SkyMap1", tract=0, patch=4).dimension_records("patch")
1129 self.assertCountEqual(
1130 [
1131 record.id
1132 for record in query.join_dataset_search("cat1", collections=["refcats"])
1133 .where(_x.htm7.region.overlaps(patch.region))
1134 .dimension_records("htm7")
1135 ],
1136 [253954, 253955],
1137 )
1139 def test_spatial_constraint_queries(self) -> None:
1140 """Test queries in which one spatial dimension in the constraint (data
1141 ID or ``where`` string) constrains a different spatial dimension in the
1142 query result columns.
1143 """
1144 butler = self.make_butler("base.yaml", "spatial.yaml")
1145 with butler.query() as query:
1146 # This tests the case where the 'tract' region is needed for
1147 # postprocessing, to compare against the visit region, but is not
1148 # needed in the resulting data ID.
1149 self.assertCountEqual(
1150 [0],
1151 [
1152 data_id["tract"]
1153 for data_id in query.data_ids(["tract"]).where({"instrument": "Cam1", "visit": 1})
1154 ],
1155 )
1156 self.assertCountEqual(
1157 [0],
1158 [
1159 data_id["tract"]
1160 for data_id in butler.query_data_ids(["tract"], instrument="Cam1", visit=1)
1161 ],
1162 )
1164 # This tests the case where the 'tract' region is needed in
1165 # postprocessing AND is also returned in the result rows.
1166 region_hex = (
1167 "7022408b0df0feef3f20378b0df0fe6f3fe23d8b0df0fe8f3ff1d8af0460ffef3f"
1168 "efcfaf0460ff6f3f75e0830388ff873f31aaeb0730ffef3fb0a5eb0730ff7f3f65"
1169 "bdf00564ff873f31aaeb0730ffef3fb1aeeb0730ff7f3f65bdf00564ff873f3e1c"
1170 "2f0fe0feef3f6e57630b28ff873fef52630b28ff873f911ade5e30fdef3f2d9626"
1171 "47e4fd873f0d952647e4fd973f553df64a80fdef3fd438f64a80fd7f3f20af3838"
1172 "20fe973f58462440b0fdef3f573d2440b0fd6f3fe2351b3044fe973f22408b0df0"
1173 "feef3f20378b0df0fe6f3f61428b0df0fe8f3f"
1174 )
1175 self.assertEqual(
1176 [(0, region_hex)],
1177 [
1178 (record.id, record.region.encode().hex())
1179 for record in query.dimension_records("tract").where({"instrument": "Cam1", "visit": 1})
1180 ],
1181 )
1182 self.assertEqual(
1183 [(0, region_hex)],
1184 [
1185 (record.id, record.region.encode().hex())
1186 for record in butler.query_dimension_records("tract", instrument="Cam1", visit=1)
1187 ],
1188 )
1190 def test_data_coordinate_upload(self) -> None:
1191 """Test queries for dimension records with a data coordinate upload."""
1192 butler = self.make_butler("base.yaml", "spatial.yaml")
1193 with butler.query() as query:
1194 # Query with a data ID upload that has an irrelevant row (there's
1195 # no data with "Cam2").
1196 self.check_detector_records(
1197 query.join_data_coordinates(
1198 [
1199 DataCoordinate.standardize(instrument="Cam1", detector=1, universe=butler.dimensions),
1200 DataCoordinate.standardize(instrument="Cam1", detector=3, universe=butler.dimensions),
1201 DataCoordinate.standardize(instrument="Cam2", detector=4, universe=butler.dimensions),
1202 ]
1203 ).dimension_records("detector"),
1204 [1, 3],
1205 )
1206 # Query with a data ID upload that directly contains duplicates,
1207 # which should not appear in the results.
1208 self.check_detector_records(
1209 query.join_data_coordinates(
1210 [
1211 DataCoordinate.standardize(instrument="Cam1", detector=1, universe=butler.dimensions),
1212 DataCoordinate.standardize(instrument="Cam1", detector=3, universe=butler.dimensions),
1213 DataCoordinate.standardize(instrument="Cam1", detector=3, universe=butler.dimensions),
1214 ]
1215 ).dimension_records("detector"),
1216 [1, 3],
1217 )
1218 # Query with a data ID upload that has extra dimensions that could
1219 # also introduce duplicates if we're not careful.
1220 self.check_detector_records(
1221 query.join_data_coordinates(
1222 [
1223 DataCoordinate.standardize(
1224 instrument="Cam1", visit=1, detector=1, universe=butler.dimensions
1225 ),
1226 DataCoordinate.standardize(
1227 instrument="Cam1", visit=2, detector=3, universe=butler.dimensions
1228 ),
1229 DataCoordinate.standardize(
1230 instrument="Cam1", visit=1, detector=3, universe=butler.dimensions
1231 ),
1232 ]
1233 ).dimension_records("detector"),
1234 [1, 3],
1235 )
1236 # Query with a data ID upload that has extra dimensions that are
1237 # used in a constraint.
1238 self.check_detector_records(
1239 query.join_data_coordinates(
1240 [
1241 DataCoordinate.standardize(
1242 instrument="Cam1", visit=1, detector=1, universe=butler.dimensions
1243 ),
1244 DataCoordinate.standardize(
1245 instrument="Cam1", visit=2, detector=3, universe=butler.dimensions
1246 ),
1247 DataCoordinate.standardize(
1248 instrument="Cam1", visit=1, detector=3, universe=butler.dimensions
1249 ),
1250 ]
1251 )
1252 .where(instrument="Cam1", visit=2)
1253 .dimension_records("detector"),
1254 [3],
1255 )
1256 # Query with a data ID upload that must be spatially joined to
1257 # the other dimensions. This join is added automatically.
1258 self.check_detector_records(
1259 query.join_data_coordinates(
1260 [
1261 DataCoordinate.standardize(
1262 skymap="SkyMap1", tract=1, patch=1, universe=butler.dimensions
1263 ),
1264 DataCoordinate.standardize(
1265 skymap="SkyMap1", tract=1, patch=2, universe=butler.dimensions
1266 ),
1267 DataCoordinate.standardize(
1268 skymap="SkyMap1", tract=1, patch=3, universe=butler.dimensions
1269 ),
1270 ]
1271 )
1272 .where(instrument="Cam1", visit=2)
1273 .dimension_records("detector"),
1274 [2, 3, 4],
1275 has_postprocessing=True,
1276 )
1277 # Query with a data ID upload that embeds a spatial relationship.
1278 # This prevents automatic creation of a spatial join. To make the
1279 # test more interesting, the spatial relationship embedded in these
1280 # data IDs is nonsense: it includes combinations that do not
1281 # overlap, while leaving out combinations that do overlap.
1282 self.check_detector_records(
1283 query.join_data_coordinates(
1284 [
1285 DataCoordinate.standardize(
1286 skymap="SkyMap1",
1287 tract=1,
1288 patch=1,
1289 instrument="Cam1",
1290 visit=1,
1291 detector=1,
1292 universe=butler.dimensions,
1293 ),
1294 DataCoordinate.standardize(
1295 skymap="SkyMap1",
1296 tract=1,
1297 patch=1,
1298 instrument="Cam1",
1299 visit=1,
1300 detector=2,
1301 universe=butler.dimensions,
1302 ),
1303 DataCoordinate.standardize(
1304 skymap="SkyMap1",
1305 tract=1,
1306 patch=3,
1307 instrument="Cam1",
1308 visit=1,
1309 detector=3,
1310 universe=butler.dimensions,
1311 ),
1312 ]
1313 )
1314 .where(skymap="SkyMap1", tract=1, patch=1)
1315 .dimension_records("detector"),
1316 [1, 2],
1317 )
1318 # Query with an empty data ID upload (not a useful thing to do,
1319 # but a way to probe edge-case behavior).
1320 self.check_detector_records(
1321 query.join_data_coordinates(
1322 [
1323 DataCoordinate.make_empty(universe=butler.dimensions),
1324 ]
1325 ).dimension_records("detector"),
1326 [1, 2, 3, 4],
1327 )
1329 def test_data_coordinate_upload_force_temp_table(self) -> None:
1330 """Test queries for dimension records with a data coordinate upload
1331 that is so big it has to go into a temporary table rather than be
1332 included directly into the query via bind params (by making the
1333 threshold for making a a temporary table tiny).
1335 This test assumes a DirectQueryDriver and is automatically skipped when
1336 some other driver is found.
1337 """
1338 butler = self.make_butler("base.yaml", "spatial.yaml")
1339 with butler.query() as query:
1340 if not isinstance(query._driver, DirectQueryDriver):
1341 raise unittest.SkipTest("Test requires meddling with DirectQueryDriver internals.")
1342 query._driver._constant_rows_limit = 2
1343 data_coordinates = [
1344 DataCoordinate.standardize(instrument="Cam1", detector=1, universe=butler.dimensions),
1345 DataCoordinate.standardize(instrument="Cam1", detector=3, universe=butler.dimensions),
1346 DataCoordinate.standardize(instrument="Cam1", detector=4, universe=butler.dimensions),
1347 ]
1348 self.check_detector_records(
1349 query.join_data_coordinates(data_coordinates).dimension_records("detector"),
1350 [1, 3, 4],
1351 )
1353 # Make sure it can fall back to a VALUES clause if temporary tables
1354 # are not supported by the DB.
1355 with unittest.mock.patch.object(query._driver.db, "_allow_temporary_tables", False):
1356 self.check_detector_records(
1357 query.join_data_coordinates(data_coordinates).dimension_records("detector"),
1358 [1, 3, 4],
1359 )
1361 def test_materialization(self) -> None:
1362 """Test querying for dimension records against a materialized previous
1363 query.
1364 """
1365 butler = self.make_butler("base.yaml", "datasets.yaml", "spatial.yaml")
1366 with butler.query() as query:
1367 _x = query.expression_factory
1368 # Simple case where the materialization has just the dimensions
1369 # we need for the rest of the query.
1370 self.check_detector_records(
1371 query.where(_x.detector.raft == "A", instrument="Cam1")
1372 .materialize()
1373 .dimension_records("detector"),
1374 [1, 2],
1375 )
1376 # This materialization has extra dimensions that could cause
1377 # duplicates if we don't SELECT DISTINCT them away.
1378 self.check_detector_records(
1379 query.join_dimensions(["visit", "detector"])
1380 .where(_x.detector.raft == "A", instrument="Cam1")
1381 .materialize()
1382 .dimension_records("detector"),
1383 [1, 2],
1384 )
1385 # Materialize a spatial-join, which should prevent the creation
1386 # of a spatial join in the downstream query.
1387 self.check_detector_records(
1388 query.join_dimensions(["visit", "detector", "tract"])
1389 .materialize()
1390 .where(skymap="SkyMap1", tract=0, instrument="Cam1", visit=2)
1391 .dimension_records("detector"),
1392 [1, 2],
1393 has_postprocessing=True,
1394 )
1395 # Materialize with a dataset join.
1396 self.check_detector_records(
1397 query.join_dataset_search("bias", collections=["imported_g"])
1398 .materialize(datasets=["bias"])
1399 .dimension_records("detector"),
1400 [1, 2, 3],
1401 )
1403 def test_materialization_find_first(self) -> None:
1404 """Test querying for datasets with find_first against a materialized
1405 query.
1406 """
1407 butler = self.make_butler("ci_hsc-subset.yaml", "ci_hsc-subset-skymap.yaml")
1409 run = "HSC/runs/ci_hsc/20240806T180642Z"
1410 extra_run = "HSC/runs/ci_hsc/20240806T180642Z-extra"
1412 # Find few datasets to duplicate.
1413 refs = butler.query_datasets("calexp", run, limit=3)
1414 data_ids = [ref.dataId for ref in refs]
1416 butler.collections.register(extra_run)
1417 butler.registry.insertDatasets("calexp", data_ids, extra_run)
1419 collections = [run, extra_run, "skymaps"]
1420 with butler.query() as query:
1421 query = query.join_dimensions(
1422 [
1423 "instrument",
1424 "physical_filter",
1425 "band",
1426 "visit",
1427 "detector",
1428 "day_obs",
1429 "skymap",
1430 "tract",
1431 ]
1432 )
1433 query = query.join_dataset_search("skyMap", collections)
1434 query = query.join_dataset_search("calexp", collections)
1435 query = query.where({}, "instrument='HSC' AND skymap='discrete/ci_hsc'", bind=None)
1436 m_query = query.materialize()
1437 _ = list(m_query.datasets("skyMap", collections))
1438 _ = list(m_query.datasets("calexp", collections))
1440 def test_materialization_no_results(self) -> None:
1441 """Test querying for datasets when materialized table is empty."""
1442 butler = self.make_butler("ci_hsc-subset.yaml", "ci_hsc-subset-skymap.yaml")
1444 run = "HSC/runs/ci_hsc/20240806T180642Z"
1446 # Register a dataset type but do not add any datasets.
1447 butler.registry.registerDatasetType(
1448 DatasetType("nothing", ["visit", "detector"], "int", universe=butler.dimensions)
1449 )
1451 collections = [run]
1452 with butler.query() as query:
1453 query = query.join_dimensions(
1454 [
1455 "instrument",
1456 "physical_filter",
1457 "band",
1458 "visit",
1459 "detector",
1460 "day_obs",
1461 "skymap",
1462 "tract",
1463 ]
1464 )
1465 query = query.join_dataset_search("calexp", collections)
1466 query = query.join_dataset_search("nothing", collections)
1467 query = query.where({}, "instrument='HSC' AND skymap='discrete/ci_hsc'", bind=None)
1468 no_results = "\n".join(query.explain_no_results())
1469 self.assertIn("No datasets of type 'nothing'", no_results)
1471 m_query = query.materialize()
1472 result = m_query.datasets("nothing")
1473 self.assertFalse(result.any())
1474 no_results = "\n".join(result.explain_no_results())
1475 self.assertIn("No datasets of type 'nothing'", no_results)
1477 def test_timespan_results(self) -> None:
1478 """Test returning dimension records that include timespans."""
1479 butler = self.make_butler("base.yaml", "spatial.yaml")
1480 with butler.query() as query:
1481 query_results = list(query.dimension_records("visit"))
1482 simple_results = butler.query_dimension_records("visit")
1483 for results in (query_results, simple_results):
1484 self.assertCountEqual(
1485 [(record.id, record.timespan.begin, record.timespan.end) for record in results],
1486 [
1487 (
1488 1,
1489 astropy.time.Time("2021-09-09T03:00:00", format="isot", scale="tai"),
1490 astropy.time.Time("2021-09-09T03:01:00", format="isot", scale="tai"),
1491 ),
1492 (
1493 2,
1494 astropy.time.Time("2021-09-09T03:02:00", format="isot", scale="tai"),
1495 astropy.time.Time("2021-09-09T03:03:00", format="isot", scale="tai"),
1496 ),
1497 ],
1498 )
1500 def test_direct_driver_paging(self) -> None:
1501 """Test queries for dimension records that require multiple pages (by
1502 making the page size tiny for DirectQueryDriver).
1504 For RemoteQueryDriver, we can't manipulate the page size so this just
1505 checks that the driver context manager logic is executing.
1506 """
1507 butler = self.make_butler("base.yaml")
1508 # Basic test where pages should be transparent.
1509 with butler.query() as query:
1510 if isinstance(query._driver, DirectQueryDriver):
1511 query._driver._raw_page_size = 2
1512 self.check_detector_records(
1513 query.dimension_records("detector"),
1514 [1, 2, 3, 4],
1515 )
1516 # Test that it's an error to continue query iteration after closing the
1517 # context manager.
1518 with butler.query() as query:
1519 if isinstance(query._driver, DirectQueryDriver):
1520 query._driver._raw_page_size = 2
1521 iterator = iter(query.dimension_records("detector"))
1522 next(iterator)
1523 with self.assertRaisesRegex(RuntimeError, "Cannot continue query result iteration"):
1524 list(iterator)
1526 def test_column_expressions(self) -> None:
1527 """Test queries with a wide variant of column expressions."""
1528 butler = self.make_butler("base.yaml", "spatial.yaml")
1529 butler.registry.defaults = RegistryDefaults(instrument="Cam1")
1530 with butler.query() as query:
1531 _x = query.expression_factory
1532 self.check_detector_records(
1533 query.where(_x.not_(_x.detector != 2)).dimension_records("detector"),
1534 [2],
1535 )
1536 self.check_detector_records_returned(
1537 butler.query_dimension_records("detector", where="NOT (detector != 2)"),
1538 [2],
1539 )
1540 self.check_detector_records(
1541 # Empty string expression should evaluate to True.
1542 query.where(_x.detector == 2, "").dimension_records("detector"),
1543 [2],
1544 )
1545 self.check_detector_records(
1546 query.where(_x.literal(2) == _x.detector).dimension_records("detector"),
1547 [2],
1548 )
1549 self.check_detector_records(
1550 query.where(_x.literal(2) == _x.detector + 1).dimension_records("detector"),
1551 [1],
1552 )
1553 self.check_detector_records(
1554 query.where(-_x.detector == -3).dimension_records("detector"),
1555 [3],
1556 )
1557 self.check_detector_records(
1558 query.where(_x.detector == 1, _x.detector == 2).dimension_records("detector"),
1559 [],
1560 messages=["'where' expression requires both detector=2 and detector=1."],
1561 )
1562 self.assertCountEqual(
1563 [
1564 record.id
1565 for record in query.where(
1566 # Datetime equal to the "begin" of the timespan.
1567 _x.visit.timespan.overlaps(
1568 astropy.time.Time("2021-09-09T03:00:00", format="isot", scale="tai")
1569 )
1570 ).dimension_records("visit")
1571 ],
1572 # Timespan begin bound is inclusive, so the record should
1573 # match.
1574 [1],
1575 )
1576 self.assertCountEqual(
1577 [
1578 record.id
1579 for record in query.where(
1580 # Datetime equal to the "end" of the timespan.
1581 _x.visit.timespan.overlaps(
1582 astropy.time.Time("2021-09-09T03:01:00", format="isot", scale="tai")
1583 )
1584 ).dimension_records("visit")
1585 ],
1586 # Timespan end bound is exclusive, so we should get no records.
1587 [],
1588 )
1589 self.assertCountEqual(
1590 [
1591 record.id
1592 for record in query.where(
1593 # In the middle of the timespan.
1594 _x.visit.timespan.overlaps(
1595 astropy.time.Time("2021-09-09T03:02:30", format="isot", scale="tai")
1596 )
1597 ).dimension_records("visit")
1598 ],
1599 [2],
1600 )
1601 self.assertCountEqual(
1602 [
1603 record.id
1604 for record in butler.query_dimension_records(
1605 # In the middle of the timespan.
1606 "visit",
1607 where="visit.timespan OVERLAPS(:ts)",
1608 bind={"ts": astropy.time.Time("2021-09-09T03:02:30", format="isot", scale="tai")},
1609 )
1610 ],
1611 [2],
1612 )
1613 self.assertCountEqual(
1614 [
1615 record.id
1616 for record in query.where(
1617 _x.visit.timespan.overlaps(
1618 Timespan(
1619 begin=astropy.time.Time("2021-09-09T03:02:30", format="isot", scale="tai"),
1620 end=None,
1621 )
1622 )
1623 ).dimension_records("visit")
1624 ],
1625 [2],
1626 )
1627 self.assertCountEqual(
1628 [
1629 record.id
1630 for record in query.where(
1631 _x.not_(
1632 _x.visit.timespan.end
1633 < astropy.time.Time("2021-09-09T03:02:30", format="isot", scale="tai"),
1634 )
1635 ).dimension_records("visit")
1636 ],
1637 [2],
1638 )
1639 self.assertCountEqual(
1640 [
1641 record.id
1642 for record in query.where(
1643 _x.visit.timespan.begin
1644 > astropy.time.Time("2021-09-09T03:01:30", format="isot", scale="tai")
1645 ).dimension_records("visit")
1646 ],
1647 [2],
1648 )
1649 self.assertCountEqual(
1650 [
1651 record.id
1652 for record in query.where(
1653 (_x.visit.exposure_time + -(5.0 * _x.visit.zenith_angle)) > 0.0
1654 ).dimension_records("visit")
1655 ],
1656 [1],
1657 )
1658 self.assertCountEqual(
1659 [
1660 record.id
1661 for record in query.where(_x.visit.exposure_time - 5.0 >= 50.0).dimension_records("visit")
1662 ],
1663 [1],
1664 )
1665 self.assertCountEqual(
1666 [record.id for record in query.where(_x.visit.id % 2 != 0).dimension_records("visit")],
1667 [1],
1668 )
1669 self.assertCountEqual(
1670 [
1671 record.id
1672 for record in query.where(_x.visit.zenith_angle / 5.0 <= 1.0).dimension_records("visit")
1673 ],
1674 [1],
1675 )
1676 self.assertCountEqual(
1677 [record.id for record in query.where(_x.visit.timespan.is_null).dimension_records("visit")],
1678 [],
1679 )
1680 self.assertCountEqual(
1681 [
1682 record.id
1683 for record in query.where(_x.visit.exposure_time.is_null).dimension_records("visit")
1684 ],
1685 [],
1686 )
1688 # Allow comparison of float columns with int literals
1689 self.assertCountEqual(
1690 [record.id for record in query.where("visit.exposure_time > 50").dimension_records("visit")],
1691 [1],
1692 )
1693 self.assertCountEqual(
1694 [record.id for record in query.where(_x.visit.exposure_time > 50).dimension_records("visit")],
1695 [1],
1696 )
1698 self.check_detector_records(
1699 query.where(_x.detector.in_iterable([1, 3, 4])).dimension_records("detector"),
1700 [1, 3, 4],
1701 )
1702 self.check_detector_records_returned(
1703 butler.query_dimension_records(
1704 "detector", where="detector IN (:det)", bind={"det": [1, 3, 4]}
1705 ),
1706 [1, 3, 4],
1707 )
1708 self.check_detector_records(
1709 query.where(_x.detector.in_range(start=2, stop=None)).dimension_records("detector"),
1710 [2, 3, 4],
1711 )
1712 self.check_detector_records(
1713 query.where(_x.detector.in_range(start=1, stop=3)).dimension_records("detector"),
1714 [1, 2],
1715 )
1716 self.check_detector_records(
1717 query.where(_x.detector.in_range(start=1, stop=None, step=2)).dimension_records("detector"),
1718 [1, 3],
1719 )
1720 self.check_detector_records(
1721 query.where(_x.detector.in_range(start=1, stop=2)).dimension_records("detector"),
1722 [1],
1723 )
1724 # This is a complex way to write a much simpler query ("where
1725 # detector.raft == 'A'"), but it tests code paths that would
1726 # otherwise require a lot more test setup.
1727 self.check_detector_records(
1728 query.where(
1729 _x.detector.in_query(_x.detector, query.where(_x.detector.raft == "A"))
1730 ).dimension_records("detector"),
1731 [1, 2],
1732 )
1733 # Error to reference tract without skymap in a WHERE clause.
1734 with self.assertRaises(InvalidQueryError):
1735 list(query.where(_x.tract == 4).dimension_records("patch"))
1737 def test_boolean_columns(self) -> None:
1738 """Test that boolean columns work as expected when specifying
1739 expressions.
1740 """
1741 # Exposure is the only dimension that has boolean columns, and this set
1742 # of data has most of the pre-requisites for exposure set up.
1743 butler = self.make_butler("base.yaml", "spatial.yaml")
1744 butler.registry.insertDimensionData("group", {"instrument": "Cam1", "name": "1"})
1746 base_data = {"instrument": "Cam1", "physical_filter": "Cam1-R1", "group": "1", "day_obs": 20210909}
1748 TRUE_ID = 1000
1749 FALSE_ID_1 = 2001
1750 FALSE_ID_2 = 2002
1751 NULL_ID_1 = 3000
1752 records = [
1753 {"id": TRUE_ID, "obs_id": "true-1", "can_see_sky": True},
1754 {"id": FALSE_ID_1, "obs_id": "false-1", "can_see_sky": False, "observation_type": "science"},
1755 {"id": FALSE_ID_2, "obs_id": "false-2", "can_see_sky": False, "observation_type": None},
1756 {"id": NULL_ID_1, "obs_id": "null-1", "can_see_sky": None},
1757 ]
1758 for record in records:
1759 butler.registry.insertDimensionData("exposure", base_data | record)
1761 # Go through the registry interface to cover the old query system, too.
1762 # This can be removed once the old query system is removed.
1763 def _run_registry_query(where: str) -> list[int]:
1764 return _get_exposure_ids_from_dimension_records(
1765 butler.registry.queryDimensionRecords("exposure", where=where, instrument="Cam1")
1766 )
1768 def _run_simple_query(where: str) -> list[int]:
1769 return _get_exposure_ids_from_dimension_records(
1770 butler.query_dimension_records("exposure", where=where, instrument="Cam1")
1771 )
1773 def _run_query(where: str) -> list[int]:
1774 with butler.query() as query:
1775 return _get_exposure_ids_from_dimension_records(
1776 query.dimension_records("exposure").where(where, instrument="Cam1")
1777 )
1779 # Test boolean columns in the `where` string syntax.
1780 for test, query_func in [
1781 ("registry", _run_registry_query),
1782 ("new-query", _run_query),
1783 ("simple", _run_simple_query),
1784 ]:
1785 with self.subTest(test):
1786 # Boolean columns should be usable standalone as an expression.
1787 self.assertCountEqual(query_func("exposure.can_see_sky"), [TRUE_ID])
1789 # You can find false values in the column with NOT. The NOT of
1790 # NULL is NULL, consistent with SQL semantics -- so records
1791 # with NULL can_see_sky are not included here.
1792 self.assertCountEqual(query_func("NOT exposure.can_see_sky"), [FALSE_ID_1, FALSE_ID_2])
1794 # Make sure the bare column composes with other expressions
1795 # correctly.
1796 self.assertCountEqual(
1797 query_func("exposure.can_see_sky OR exposure = 2001"), [TRUE_ID, FALSE_ID_1]
1798 )
1800 # Find nulls and non-nulls.
1801 #
1802 # This is run only against the new query system. It appears that the
1803 # `= NULL` syntax never had test coverage in the old query system and
1804 # doesn't work for any column types. Not worth fixing since we are
1805 # dropping that code soon.
1806 nulls = [NULL_ID_1]
1807 non_nulls = [TRUE_ID, FALSE_ID_1, FALSE_ID_2]
1808 self.assertCountEqual(_run_query("exposure.can_see_sky = NULL"), nulls)
1809 self.assertCountEqual(_run_query("exposure.can_see_sky != NULL"), non_nulls)
1810 self.assertCountEqual(_run_query("NULL = exposure.can_see_sky"), nulls)
1811 self.assertCountEqual(_run_query("NULL != exposure.can_see_sky"), non_nulls)
1813 # You can't do a NULL check on an arbitrary boolean predicate.
1814 with self.assertRaises(InvalidQueryError):
1815 _run_query("NULL = (exposure.can_see_sky AND exposure = 2001)")
1817 # Check null finding for non-boolean columns, too.
1818 self.assertEqual(
1819 _run_query("exposure.observation_type = NULL AND NOT exposure.can_see_sky"), [FALSE_ID_2]
1820 )
1821 self.assertEqual(
1822 _run_query("exposure.observation_type != NULL AND NOT exposure.can_see_sky"), [FALSE_ID_1]
1823 )
1824 self.assertEqual(
1825 _run_query("NULL = exposure.observation_type AND NOT exposure.can_see_sky"), [FALSE_ID_2]
1826 )
1827 self.assertEqual(
1828 _run_query("NULL != exposure.observation_type AND NOT exposure.can_see_sky"), [FALSE_ID_1]
1829 )
1831 # Test boolean columns in ExpressionFactory.
1832 with butler.query() as query:
1833 x = query.expression_factory
1835 def do_query(constraint: Predicate) -> list[int]:
1836 return _get_exposure_ids_from_dimension_records(
1837 query.dimension_records("exposure").where(constraint, instrument="Cam1")
1838 )
1840 # Boolean columns should be usable standalone as a Predicate.
1841 self.assertCountEqual(do_query(x.exposure.can_see_sky.as_boolean()), [TRUE_ID])
1843 # You can find false values in the column with NOT. The NOT of
1844 # NULL is NULL, consistent with SQL semantics -- so records
1845 # with NULL can_see_sky are not included here.
1846 self.assertCountEqual(
1847 do_query(x.exposure.can_see_sky.as_boolean().logical_not()), [FALSE_ID_1, FALSE_ID_2]
1848 )
1850 # Searching for nulls works.
1851 self.assertCountEqual(do_query(x.exposure.can_see_sky.is_null), [NULL_ID_1])
1853 # Attempting to use operators that only apply to non-boolean types
1854 # is an error.
1855 with self.assertRaisesRegex(
1856 InvalidQueryError,
1857 r"Boolean expression 'exposure.can_see_sky' can't be used directly in other expressions."
1858 r" Call the 'as_boolean\(\)' method to convert it to a Predicate instead.",
1859 ):
1860 x.exposure.can_see_sky == 1
1862 # Non-boolean types can't be converted directly to Predicate.
1863 with self.assertRaisesRegex(
1864 InvalidQueryError,
1865 r"Expression 'exposure.observation_type' with type 'string' can't be used directly"
1866 r" as a boolean value.",
1867 ):
1868 x.exposure.observation_type.as_boolean()
1870 def test_dataset_region_queries(self) -> None:
1871 """Test region queries for datasets."""
1872 # Import data to play with.
1873 butler = self.make_butler("base.yaml", "ci_hsc-subset.yaml")
1875 run = "HSC/runs/ci_hsc/20240806T180642Z"
1876 with butler.query() as query:
1877 # Return everything.
1878 results = query.datasets("calexp", collections=run)
1879 # Sort by data coordinate.
1880 refs = sorted(results.with_dimension_records(), key=attrgetter("dataId"))
1881 self.assertEqual(len(refs), 33)
1883 # Use a region from the first visit.
1884 first_visit_region = refs[0].dataId.visit.region # type: ignore
1886 # Get a visit detector region from the first ref.
1887 with butler.query() as query:
1888 data_id = refs[0].dataId.mapping
1889 records = list(query.dimension_records("visit_detector_region").where(**data_id)) # type: ignore
1890 self.assertEqual(len(records), 1)
1892 for pos, use_bind, count in (
1893 ("CIRCLE 320. -0.25 10.", True, 33), # Match everything.
1894 ("CIRCLE 321.0 -0.4 0.01", True, 1), # Should be small region on 1 detector.
1895 ("CIRCLE 321.1 -0.35 0.02", True, 2),
1896 ("CIRCLE 321.1 -0.48 0.05", True, 1), # Center off the region.
1897 ("CIRCLE 321.0 -0.5 0.01", True, 0), # No overlap.
1898 (first_visit_region.to_ivoa_pos(), True, 33), # Visit region overlaps everything.
1899 (records[0].region.to_ivoa_pos(), True, 17), # Some overlap.
1900 ("CIRCLE(320., -0.25, 10.)", False, 33), # Match everything.
1901 ("CIRCLE(321.0, -0.4, 0.01)", False, 1), # Should be small region on 1 detector.
1902 ("CIRCLE(321.0, -0.5, 0.01)", False, 0), # No overlap.
1903 ("BOX(320, -0.25, 5, 5)", False, 33), # Match everything.
1904 ("BOX(321.0, -0.4, 0.01, 0.01)", False, 1), # Should be small region on 1 detector.
1905 ("BOX(321.0, -0.5, 0.01, 0.01)", False, 0), # No overlap.
1906 ("POLYGON(320, -10, 320, 10, 340, 10, 340, -10)", False, 33), # Match everything.
1907 ("POLYGON(320.99, -0.401, 320.99, -0.399, 321.01, -0.399, 321.01, -0.401)", False, 1),
1908 ("POLYGON(320.99, -0.501, 320.99, -0.499, 321.01, -0.499, 321.01, -0.501)", False, 0),
1909 ("REGION('CIRCLE 320. -0.25 10.')", False, 33), # Match everything.
1910 ("REGION('RANGE 310 330 -10 10')", False, 33), # Match everything.
1911 ("REGION('RANGE 320.99 321.01 -0.401 -0.399')", False, 1), # Small region on 1 detector.
1912 ("REGION('POLYGON 320.99 -0.501 320.99 -0.499 321.01 -0.499 321.01 -0.501')", False, 0),
1913 ):
1914 if use_bind:
1915 overlap_where = "visit_detector_region.region OVERLAPS :POS"
1916 bind = {"POS": Region.from_ivoa_pos(pos)}
1917 else:
1918 overlap_where = f"visit_detector_region.region OVERLAPS {pos}"
1919 bind = {}
1920 with butler.query() as query:
1921 results = query.datasets("calexp", collections=run)
1922 results = results.where(f"instrument = 'HSC' AND {overlap_where}", bind=bind)
1923 refs = list(results)
1924 self.assertEqual(len(refs), count, f"POS={pos} REFS={refs}")
1926 simple_refs = butler.query_datasets(
1927 "calexp",
1928 collections=run,
1929 instrument="HSC",
1930 where=overlap_where,
1931 bind=bind,
1932 explain=False,
1933 )
1934 self.assertCountEqual(refs, simple_refs)
1936 def test_dataset_time_queries(self) -> None:
1937 """Test temporal queries for datasets."""
1938 # Import data to play with.
1939 butler = self.make_butler("base.yaml", "ci_hsc-subset.yaml")
1941 # Some times from the test data.
1942 v_903334_pre = astropy.time.Time("2013-01-01T12:00:00", scale="tai", format="isot")
1943 v_903334_mid = astropy.time.Time("2013-06-17T13:29:20", scale="tai", format="isot")
1944 v_904014_pre = astropy.time.Time("2013-11-01T12:00:00", scale="tai", format="isot")
1945 v_904014_post = astropy.time.Time("2013-12-21T12:00:00", scale="tai", format="isot")
1947 with butler.query() as query:
1948 run = "HSC/runs/ci_hsc/20240806T180642Z"
1949 results = query.datasets("calexp", collections=run)
1951 # Use a time during the middle of a visit.
1952 v_903334 = results.where(
1953 "instrument = 'HSC' and visit.timespan OVERLAPS(:ts)", bind={"ts": v_903334_mid}
1954 )
1955 self.assertEqual(len(list(v_903334)), 4)
1957 # Timespan covering first half of the data.
1958 first_half = results.where(
1959 "instrument = 'HSC' and visit.timespan OVERLAPS(:t1, :t2)",
1960 bind={"t1": v_903334_pre, "t2": v_904014_pre},
1961 )
1962 self.assertEqual(len(list(first_half)), 17)
1964 # Query using a timespan object.
1965 with_ts = results.where(
1966 "instrument = 'HSC' and visit.timespan OVERLAPS(:ts)",
1967 bind={"ts": Timespan(v_904014_pre, v_904014_post)},
1968 )
1969 self.assertEqual(len(list(with_ts)), 16)
1971 def test_calibration_join_queries(self) -> None:
1972 """Test using the 'general' query result type to join observations to
1973 calibration datasets temporally.
1975 We have to use general results because we want calibration DatasetRefs
1976 and data IDs that include the observation identifiers (which are not
1977 part of the calibration dataset dimensions).
1978 """
1979 butler = self.make_butler("base.yaml", "datasets.yaml")
1980 # Set up some timestamps.
1981 t1 = astropy.time.Time("2020-01-01T01:00:00", format="isot", scale="tai")
1982 t2 = astropy.time.Time("2020-01-01T02:00:00", format="isot", scale="tai")
1983 t3 = astropy.time.Time("2020-01-01T03:00:00", format="isot", scale="tai")
1984 t4 = astropy.time.Time("2020-01-01T04:00:00", format="isot", scale="tai")
1985 t5 = astropy.time.Time("2020-01-01T05:00:00", format="isot", scale="tai")
1986 # Insert some exposure records with timespans between each sequential
1987 # pair of those.
1988 butler.registry.insertDimensionData(
1989 "day_obs", {"instrument": "Cam1", "id": 20200101, "timespan": Timespan(t1, t5)}
1990 )
1991 butler.registry.insertDimensionData(
1992 "group",
1993 {"instrument": "Cam1", "name": "group0"},
1994 {"instrument": "Cam1", "name": "group1"},
1995 {"instrument": "Cam1", "name": "group2"},
1996 {"instrument": "Cam1", "name": "group3"},
1997 )
1998 butler.registry.insertDimensionData(
1999 "exposure",
2000 {
2001 "instrument": "Cam1",
2002 "id": 0,
2003 "group": "group0",
2004 "obs_id": "zero",
2005 "physical_filter": "Cam1-G",
2006 "day_obs": 20200101,
2007 "timespan": Timespan(t1, t2),
2008 },
2009 {
2010 "instrument": "Cam1",
2011 "id": 1,
2012 "group": "group1",
2013 "obs_id": "one",
2014 "physical_filter": "Cam1-G",
2015 "day_obs": 20200101,
2016 "timespan": Timespan(t2, t3),
2017 },
2018 {
2019 "instrument": "Cam1",
2020 "id": 2,
2021 "group": "group2",
2022 "obs_id": "two",
2023 "physical_filter": "Cam1-G",
2024 "day_obs": 20200101,
2025 "timespan": Timespan(t3, t4),
2026 },
2027 {
2028 "instrument": "Cam1",
2029 "id": 3,
2030 "group": "group3",
2031 "obs_id": "three",
2032 "physical_filter": "Cam1-G",
2033 "day_obs": 20200101,
2034 "timespan": Timespan(t4, t5),
2035 },
2036 )
2037 # Get references to the datasets we imported.
2038 bias = butler.get_dataset_type("bias")
2039 bias2a = butler.find_dataset("bias", instrument="Cam1", detector=2, collections="imported_g")
2040 assert bias2a is not None
2041 bias3a = butler.find_dataset("bias", instrument="Cam1", detector=3, collections="imported_g")
2042 assert bias3a is not None
2043 bias2b = butler.find_dataset("bias", instrument="Cam1", detector=2, collections="imported_r")
2044 assert bias2b is not None
2045 bias3b = butler.find_dataset("bias", instrument="Cam1", detector=3, collections="imported_r")
2046 assert bias3b is not None
2047 # Register the main calibration collection we'll be working with.
2048 collection = "Cam1/calibs"
2049 butler.collections.register(collection, type=CollectionType.CALIBRATION)
2050 # Certify 2a dataset with [t2, t4) validity.
2051 butler.registry.certify(collection, [bias2a], Timespan(begin=t2, end=t4))
2052 # Certify 3a over [t1, t3).
2053 butler.registry.certify(collection, [bias3a], Timespan(begin=t1, end=t3))
2054 # Certify 2b and 3b together over [t4, ∞).
2055 butler.registry.certify(collection, [bias2b, bias3b], Timespan(begin=t4, end=None))
2056 # Query for (bias, exposure, detector) combinations.
2057 base_data_id = DataCoordinate.standardize(instrument="Cam1", universe=butler.dimensions)
2058 with butler.query() as q:
2059 x = q.expression_factory
2060 q = q.join_dimensions(["exposure"])
2061 q = q.join_dataset_search("bias", [collection])
2062 # Query for all calibs with an explicit temporal join.
2063 self.assertCountEqual(
2064 [
2065 (data_id, refs[0])
2066 for data_id, refs, _ in q.where(
2067 x["bias"].timespan.overlaps(x.exposure.timespan), base_data_id
2068 )
2069 .general(
2070 butler.dimensions.conform(["exposure", "detector"]),
2071 dataset_fields={"bias": ...},
2072 find_first=True,
2073 )
2074 .iter_tuples(bias)
2075 ],
2076 [
2077 (DataCoordinate.standardize(base_data_id, detector=2, exposure=1), bias2a),
2078 (DataCoordinate.standardize(base_data_id, detector=2, exposure=2), bias2a),
2079 (DataCoordinate.standardize(base_data_id, detector=3, exposure=0), bias3a),
2080 (DataCoordinate.standardize(base_data_id, detector=3, exposure=1), bias3a),
2081 (DataCoordinate.standardize(base_data_id, detector=2, exposure=3), bias2b),
2082 (DataCoordinate.standardize(base_data_id, detector=3, exposure=3), bias3b),
2083 ],
2084 )
2085 # Query for all calibs with the temporal join implicit and the
2086 # dimensions given as an incomplete list (detector is added by
2087 # the dataset results).
2088 self.assertCountEqual(
2089 [
2090 (data_id, refs[0])
2091 for data_id, refs, _ in q.where(base_data_id)
2092 .general(["exposure"], dataset_fields={"bias": ...}, find_first=True)
2093 .iter_tuples(bias)
2094 ],
2095 [
2096 (DataCoordinate.standardize(base_data_id, detector=2, exposure=1), bias2a),
2097 (DataCoordinate.standardize(base_data_id, detector=2, exposure=2), bias2a),
2098 (DataCoordinate.standardize(base_data_id, detector=3, exposure=0), bias3a),
2099 (DataCoordinate.standardize(base_data_id, detector=3, exposure=1), bias3a),
2100 (DataCoordinate.standardize(base_data_id, detector=2, exposure=3), bias2b),
2101 (DataCoordinate.standardize(base_data_id, detector=3, exposure=3), bias3b),
2102 ],
2103 )
2105 # Query in multiple collections, with one of the collections being a
2106 # calibration collection. This triggers special cases related to
2107 # timespan columns in the query code.
2108 refs = butler.query_datasets(
2109 "bias",
2110 collections=[collection, "imported_g"],
2111 where="instrument = 'Cam1' and detector = 2",
2112 find_first=False,
2113 )
2114 self.assertEqual(
2115 sorted([str(ref.id) for ref in refs]),
2116 ["51352db4-a47a-447c-b12d-a50b206b17cd", "87f3e68d-258d-41b7-8ea5-edf3557ccb30"],
2117 )
2118 with butler.query() as query:
2119 query = query.join_dataset_search("bias", [collection, "imported_g"])
2120 query = query.where("instrument = 'Cam1' and detector = 2")
2122 results = list(
2123 query.general(
2124 ["detector"],
2125 dataset_fields={"bias": {"dataset_id"}},
2126 find_first=False,
2127 ).iter_tuples()
2128 )
2129 # Dataset ID should be de-duplicated with no collection/timespan
2130 # column present.
2131 self.assertEqual(
2132 sorted([str(x.raw_row["bias.dataset_id"]) for x in results]),
2133 ["51352db4-a47a-447c-b12d-a50b206b17cd", "87f3e68d-258d-41b7-8ea5-edf3557ccb30"],
2134 )
2136 results = list(
2137 query.general(
2138 ["detector"],
2139 dataset_fields={"bias": {"dataset_id", "timespan"}},
2140 find_first=False,
2141 ).iter_tuples()
2142 )
2143 # We should have one row for each timespan associated with each
2144 # dataset ID. The extra copy of "51352..." comes from the run
2145 # collection, with a timespan of None.
2146 self.assertEqual(
2147 sorted([str(x.raw_row["bias.dataset_id"]) for x in results]),
2148 [
2149 "51352db4-a47a-447c-b12d-a50b206b17cd",
2150 "51352db4-a47a-447c-b12d-a50b206b17cd",
2151 "87f3e68d-258d-41b7-8ea5-edf3557ccb30",
2152 ],
2153 )
2155 results = list(
2156 query.general(
2157 ["detector"],
2158 dataset_fields={"bias": {"dataset_id", "collection"}},
2159 find_first=False,
2160 ).iter_tuples()
2161 )
2162 # We should have one row for each collection associated with each
2163 # dataset ID.
2164 self.assertEqual(
2165 sorted([(str(x.raw_row["bias.dataset_id"]), x.raw_row["bias.collection"]) for x in results]),
2166 [
2167 ("51352db4-a47a-447c-b12d-a50b206b17cd", "Cam1/calibs"),
2168 ("51352db4-a47a-447c-b12d-a50b206b17cd", "imported_g"),
2169 ("87f3e68d-258d-41b7-8ea5-edf3557ccb30", "Cam1/calibs"),
2170 ],
2171 )
2173 def test_collection_query_info(self) -> None:
2174 butler = self.make_butler("base.yaml", "datasets.yaml")
2176 info = butler.collections.query_info("imported_g", include_summary=True)
2177 self.assertEqual(len(info), 1)
2178 dataset_types = info[0].dataset_types
2179 assert dataset_types is not None
2180 self.assertCountEqual(dataset_types, ["flat", "bias"])
2182 info = butler.collections.query_info("imported_g", include_summary=True, summary_datasets=["flat"])
2183 self.assertEqual(len(info), 1)
2184 dataset_types = info[0].dataset_types
2185 assert dataset_types is not None
2186 self.assertCountEqual(dataset_types, ["flat"])
2188 def test_dataset_queries(self) -> None:
2189 butler = self.make_butler("base.yaml", "spatial.yaml")
2191 # Need a dataset with some spatial information to trigger aggregate
2192 # value logic in queries.
2193 butler.registry.registerDatasetType(
2194 DatasetType("dt", ["visit", "detector"], "int", universe=butler.dimensions)
2195 )
2196 butler.collections.register("run1")
2197 butler.registry.insertDatasets("dt", [{"instrument": "Cam1", "visit": 1, "detector": 1}], "run1")
2199 # Tests for a regression of DM-46340, where invalid SQL would be
2200 # generated when the list of collections is a single run collection and
2201 # there is region-postprocessing logic involved. This was due to
2202 # missing type information associated with the "run" dataset field.
2203 result = butler.query_datasets(
2204 "dt",
2205 "run1",
2206 where="instrument='Cam1' and skymap='SkyMap1' and visit=1 and tract=0",
2207 with_dimension_records=True,
2208 )
2209 self.assertEqual(result[0].dataId, {"instrument": "Cam1", "visit": 1, "detector": 1})
2210 self.assertEqual(result[0].run, "run1")
2212 # A similar issue to the "run" issue above was occuring with the
2213 # 'collection' dataset field.
2214 with butler.query() as query:
2215 rows = list(
2216 query.join_dataset_search("dt", "run1")
2217 .where("instrument='Cam1' and skymap='SkyMap1' and visit=1 and tract=0")
2218 .general(
2219 dimensions=["visit", "detector"],
2220 dataset_fields={"dt": set(["collection"])},
2221 find_first=True,
2222 )
2223 )
2224 self.assertEqual(len(rows), 1)
2225 self.assertEqual(rows[0]["visit"], 1)
2226 self.assertEqual(rows[0]["dt.collection"], "run1")
2228 # Test that dataset fields like ingest_date can be used in the 'where'
2229 # clause.
2230 result = butler.query_datasets("dt", "run1", where="ingest_date > T'2000-01-01'")
2231 self.assertEqual(len(result), 1)
2232 result = butler.query_datasets("dt", "run1", where="ingest_date < T'2000-01-01'", explain=False)
2233 self.assertEqual(len(result), 0)
2234 result = butler.query_datasets(
2235 "dt", "run1", where="ingest_date OVERLAPS (T'2000-01-01', T'2099-01-01')"
2236 )
2237 self.assertEqual(len(result), 1)
2238 result = butler.query_datasets(
2239 "dt", "run1", where="(T'2000-01-01', T'2099-01-01') OVERLAPS ingest_date"
2240 )
2241 self.assertEqual(len(result), 1)
2242 result = butler.query_datasets(
2243 "dt", "run1", where="(T'2000-01-01', T'2001-01-01') OVERLAPS ingest_date", explain=False
2244 )
2245 self.assertEqual(len(result), 0)
2247 def test_multiple_instrument_queries(self) -> None:
2248 """Test that multiple-instrument queries are not rejected as having
2249 governor dimension ambiguities.
2250 """
2251 butler = self.make_butler("base.yaml")
2252 butler.registry.insertDimensionData("instrument", {"name": "Cam2"})
2253 self.assertCountEqual(
2254 butler.query_data_ids(["detector"], where="instrument='Cam1' OR instrument='Cam2'"),
2255 [
2256 DataCoordinate.standardize(instrument="Cam1", detector=n, universe=butler.dimensions)
2257 for n in range(1, 5)
2258 ],
2259 )
2260 self.assertCountEqual(
2261 butler.query_data_ids(
2262 ["detector"],
2263 where="(instrument='Cam1' OR instrument='Cam2') AND visit.region OVERLAPS :region",
2264 bind={"region": Region.from_ivoa_pos("CIRCLE 320. -0.25 10.")},
2265 explain=False,
2266 ),
2267 # No visits in this test dataset means no result, but the point of
2268 # the test is just that the query can be constructed at all.
2269 [],
2270 )
2271 self.assertCountEqual(
2272 butler.query_data_ids(
2273 ["instrument"],
2274 where="(instrument='Cam1' AND detector=2) OR (instrument='Cam2' AND detector=500)",
2275 explain=False,
2276 ),
2277 [DataCoordinate.standardize(instrument="Cam1", universe=butler.dimensions)],
2278 )
2280 def test_default_data_id(self) -> None:
2281 butler = self.make_butler("base.yaml")
2282 butler.registry.insertDimensionData("instrument", {"name": "Cam2"})
2283 butler.registry.insertDimensionData(
2284 "physical_filter", {"instrument": "Cam2", "name": "Cam2-G", "band": "g"}
2285 )
2287 # With no default data ID, queries should return results for all
2288 # instruments.
2289 result = butler.query_dimension_records("physical_filter")
2290 names = [x.name for x in result]
2291 self.assertCountEqual(names, ["Cam1-G", "Cam1-R1", "Cam1-R2", "Cam2-G"])
2293 result = butler.query_dimension_records("physical_filter", where="band='g'")
2294 names = [x.name for x in result]
2295 self.assertCountEqual(names, ["Cam1-G", "Cam2-G"])
2297 # When there is no default data ID and a where clause references
2298 # something depending on instrument, it throws an error as a
2299 # sanity check.
2300 # In this case, 'instrument' is not part of the dimensions returned by
2301 # the query, so there is extra logic needed to detect the need for the
2302 # default data ID.
2303 with self.assertRaisesRegex(
2304 InvalidQueryError,
2305 "Query 'where' expression references a dimension dependent on instrument"
2306 " without constraining it directly.",
2307 ):
2308 butler.query_data_ids(["band"], where="physical_filter='Cam1-G'")
2310 # Override the default data ID to specify a default instrument for
2311 # subsequent tests.
2312 butler.registry.defaults = RegistryDefaults(instrument="Cam1")
2314 # When a where clause references something depending on instrument, use
2315 # the default data ID to constrain the instrument.
2316 # In this case, 'instrument' is not part of the dimensions returned by
2317 # the query, so there is extra logic needed to detect the need for the
2318 # default data ID.
2319 data_ids = butler.query_data_ids(["band"], where="physical_filter='Cam1-G'")
2320 self.assertEqual([x["band"] for x in data_ids], ["g"])
2321 # Default data ID instrument=Cam1 does not match Cam2, so there are no
2322 # results.
2323 data_ids = butler.query_data_ids(["band"], where="physical_filter='Cam2-G'", explain=False)
2324 self.assertEqual(data_ids, [])
2325 # Overriding the default lets us get the results.
2326 data_ids = butler.query_data_ids(["band"], where="instrument='Cam2' and physical_filter='Cam2-G'")
2327 self.assertEqual([x["band"] for x in data_ids], ["g"])
2329 # Query for a dimension that depends on instrument should pull in the
2330 # default data ID instrument="Cam1" to constrain results.
2331 result = butler.query_dimension_records("physical_filter")
2332 names = [x.name for x in result]
2333 self.assertCountEqual(names, ["Cam1-G", "Cam1-R1", "Cam1-R2"])
2335 # Query for a dimension that depends on instrument should pull in the
2336 # default data ID instrument="Cam1" to constrain results, if the where
2337 # clause does not explicitly specify instrument.
2338 result = butler.query_dimension_records("physical_filter", where="band='g'")
2339 names = [x.name for x in result]
2340 self.assertEqual(names, ["Cam1-G"])
2342 # Queries that specify instrument explicitly in the where clause
2343 # should ignore the default data ID.
2344 result = butler.query_dimension_records("physical_filter", where="instrument='Cam2'")
2345 names = [x.name for x in result]
2346 self.assertCountEqual(names, ["Cam2-G"])
2348 result = butler.query_dimension_records("physical_filter", where="instrument IN ('Cam2')")
2349 names = [x.name for x in result]
2350 self.assertCountEqual(names, ["Cam2-G"])
2352 def test_unusual_column_literals(self) -> None:
2353 butler = self.make_butler("base.yaml")
2355 # Users frequently use numpy integer types as literals in queries.
2356 result = butler.query_dimension_records(
2357 "detector", data_id={"instrument": "Cam1", "detector": int64(1)}
2358 )
2359 names = [x.full_name for x in result]
2360 self.assertEqual(names, ["Aa"])
2362 result = butler.query_dimension_records(
2363 "detector", where="instrument='Cam1' and detector=:an_integer", bind={"an_integer": int64(2)}
2364 )
2365 names = [x.full_name for x in result]
2366 self.assertEqual(names, ["Ab"])
2368 with butler.query() as query:
2369 x = query.expression_factory
2370 result = list(
2371 query.dimension_records("detector").where(x.instrument == "Cam1", x.detector == int64(3))
2372 )
2373 names = [x.full_name for x in result]
2374 self.assertEqual(names, ["Ba"])
2376 def test_query_all_datasets(self) -> None:
2377 butler = self.make_butler("base.yaml", "datasets.yaml")
2379 # Make sure that refs are coming out well-formed.
2380 datasets = butler.query_all_datasets("imported_r", where="detector = 2", instrument="Cam1")
2381 datasets.sort(key=lambda ref: ref.datasetType.name)
2382 self.assertEqual(len(datasets), 2)
2383 bias = datasets[0]
2384 self.assertEqual(bias.datasetType.name, "bias")
2385 self.assertEqual(bias.dataId["instrument"], "Cam1")
2386 self.assertEqual(bias.dataId["detector"], 2)
2387 self.assertEqual(bias.run, "imported_r")
2388 self.assertEqual(bias.id, UUID("87f3e68d-258d-41b7-8ea5-edf3557ccb30"))
2389 flat = datasets[1]
2390 self.assertEqual(flat.datasetType.name, "flat")
2391 self.assertEqual(flat.dataId["instrument"], "Cam1")
2392 self.assertEqual(flat.dataId["detector"], 2)
2393 self.assertEqual(flat.dataId["physical_filter"], "Cam1-R1")
2394 self.assertEqual(flat.dataId["band"], "r")
2395 self.assertEqual(flat.run, "imported_r")
2396 self.assertEqual(flat.id, UUID("c1296796-56c5-4acf-9b49-40d920c6f840"))
2398 # Querying for everything finds everything.
2399 results = butler.query_all_datasets("*", find_first=False)
2400 self.assertEqual(len(results), 13)
2402 # constraining by data ID works
2403 detector_1_ids = ("d0bb04cd-d697-4a83-ba53-cdfcd58e3a0c", "e15ab039-bc8b-4135-87c5-90902a7c0b22")
2404 results = butler.query_all_datasets(
2405 "*", data_id={"detector": 1, "instrument": "Cam1"}, find_first=False
2406 )
2407 self.assertCountEqual(detector_1_ids, _ref_uuids(results))
2409 # bind values work.
2410 results = butler.query_all_datasets(
2411 "*", where="detector=:my_bind and instrument='Cam1'", bind={"my_bind": 1}, find_first=False
2412 )
2413 self.assertCountEqual(detector_1_ids, _ref_uuids(results))
2415 # find_first requires ordered collections.
2416 with self.assertRaisesRegex(InvalidQueryError, "Can not use wildcards"):
2417 results = butler.query_all_datasets("*")
2419 butler.collections.register("chain", CollectionType.CHAINED)
2420 butler.collections.redefine_chain("chain", ["imported_g", "imported_r"])
2421 results = butler.query_all_datasets(
2422 "chain", where="detector=2 and instrument = 'Cam1'", find_first=True
2423 )
2424 # find_first searches the collection chain in order.
2425 self.assertCountEqual(
2426 _ref_uuids(results),
2427 [
2428 "51352db4-a47a-447c-b12d-a50b206b17cd", # imported_g bias
2429 "60c8a65c-7290-4c38-b1de-e3b1cdcf872d", # imported_g flat
2430 "c1296796-56c5-4acf-9b49-40d920c6f840", # imported_r flat
2431 # There is also a bias dataset with detector=2 in imported_r,
2432 # but it is masked by the presence of the same data ID in
2433 # imported_g.
2434 ],
2435 )
2437 # collection searches work.
2438 results = butler.query_all_datasets(
2439 "*g", where="detector=1 and instrument = 'Cam1'", find_first=False
2440 )
2441 self.assertEqual(_ref_uuids(results), ["e15ab039-bc8b-4135-87c5-90902a7c0b22"])
2443 # we raise for missing collections with explicit names.
2444 with self.assertRaises(MissingCollectionError):
2445 results = butler.query_all_datasets("nonexistent")
2446 # we don't raise for collection wildcard searches that find nothing.
2447 results = butler.query_all_datasets("nonexistent*", find_first=False)
2448 self.assertEqual(results, [])
2450 # dataset type searches work.
2451 results = butler.query_all_datasets(
2452 "*", name="b*", where="detector=1 and instrument = 'Cam1'", find_first=False
2453 )
2454 self.assertEqual(_ref_uuids(results), ["e15ab039-bc8b-4135-87c5-90902a7c0b22"])
2456 # Missing dataset types raise.
2457 with self.assertRaises(MissingDatasetTypeError):
2458 results = butler.query_all_datasets("chain", name=["notfound", "flat"])
2459 with self.assertRaises(MissingDatasetTypeError):
2460 results = butler.query_all_datasets("chain", name="notfound*")
2462 # Limit of 3 lands at the boundary of a dataset type.
2463 # Limit of 4 is in the middle of a dataset type.
2464 for limit in [3, 4]:
2465 with self.subTest(limit=limit):
2466 results = butler.query_all_datasets("imported_g", limit=limit)
2467 self.assertEqual(len(results), limit)
2468 with self.assertLogs(level="WARNING") as log:
2469 results = butler.query_all_datasets("imported_g", limit=-limit)
2470 self.assertEqual(len(results), limit)
2471 self.assertIn("requested limit", log.output[0])
2473 results = butler.query_all_datasets("imported_g", limit=0)
2474 self.assertEqual(len(results), 0)
2476 # 'where' constraints that don't apply to all dataset types follow the
2477 # same rules as query_datasets.
2478 results = butler.query_all_datasets(
2479 "*", where="detector = 2 and band = 'g' and instrument = 'Cam1'", find_first=False
2480 )
2481 self.assertCountEqual(
2482 _ref_uuids(results),
2483 [
2484 # bias does not have 'band'
2485 "51352db4-a47a-447c-b12d-a50b206b17cd",
2486 "87f3e68d-258d-41b7-8ea5-edf3557ccb30",
2487 # flat does have 'band', and we filter based on it
2488 "60c8a65c-7290-4c38-b1de-e3b1cdcf872d",
2489 ],
2490 )
2492 # Default collections and data ID apply.
2493 butler.registry.defaults = RegistryDefaults(collections="imported_g")
2494 results = butler.query_all_datasets(where="detector = 2")
2495 self.assertCountEqual(
2496 _ref_uuids(results),
2497 ["51352db4-a47a-447c-b12d-a50b206b17cd", "60c8a65c-7290-4c38-b1de-e3b1cdcf872d"],
2498 )
2500 def test_irrelevant_governor_constraints(self) -> None:
2501 """Test that constraining an irrelevant governor dimension doesn't
2502 break dataset queries.
2503 """
2504 butler = self.make_butler("base.yaml", "spatial.yaml")
2505 butler.registry.insertDimensionData("instrument", {"name": "Cam2"})
2506 a = DatasetType("a", {"detector"}, "StructuredDataDict", universe=butler.dimensions)
2507 b = DatasetType("b", {"tract"}, "StructuredDataDict", universe=butler.dimensions)
2508 butler.registry.registerDatasetType(a)
2509 butler.registry.registerDatasetType(b)
2510 collection = "run1"
2511 butler.collections.register(collection)
2512 (ref_a,) = butler.registry.insertDatasets(a, [{"instrument": "Cam1", "detector": 2}], run=collection)
2513 (ref_b,) = butler.registry.insertDatasets(b, [{"skymap": "SkyMap1", "tract": 1}], run=collection)
2514 # First, some sanity-check query that's mostly to check the test setup.
2515 self.assertEqual(butler.query_datasets("a", collections=collection), [ref_a])
2516 self.assertEqual(butler.query_datasets("b", collections=collection), [ref_b])
2517 # Now check that we can get both with an irrelevant constraint.
2518 # In the first case, there is a dataset of a different type that is
2519 # consistent with the constraint in the collection:
2520 self.assertEqual(butler.query_datasets("a", collections=collection, skymap="SkyMap1"), [ref_a])
2521 # In the second case there is a dataset of a different type that is
2522 # inconsistent with the constraint in the collection:
2523 self.assertEqual(butler.query_datasets("b", collections=collection, instrument="Cam2"), [ref_b])
2525 def test_inferred_primary_key(self) -> None:
2526 """Test expressions that have an unqualified reference to a primary key
2527 field whose dimension must be inferred from context.
2528 """
2529 butler = self.make_butler("base.yaml")
2530 self.assertEqual(
2531 butler.query_dimension_records("detector", instrument="Cam1", where="id=2"),
2532 butler.query_dimension_records("detector", instrument="Cam1", detector=2),
2533 )
2535 def test_glob_expression(self) -> None:
2536 """Test GLOB() function in user expressions."""
2537 butler = self.make_butler("base.yaml")
2539 tests = (
2540 ("full_name", "*", 4),
2541 ("full_name", "\\*", 0),
2542 ("full_name", "A*", 2),
2543 ("full_name", "A?", 2),
2544 ("full_name", "??", 4),
2545 ("full_name", "*a", 2),
2546 ("full_name", "A[ab]", 0),
2547 ("purpose", "*EN?E", 3),
2548 ("purpose", "\\*CIENC\\*", 0),
2549 ("full_name", "%", 0),
2550 ("full_name", "__", 0),
2551 ("full_name", "a", 0),
2552 ("full_name", "", 0),
2553 )
2555 for column, pattern, count in tests:
2556 # Pattern as a literal string.
2557 records = butler.query_dimension_records(
2558 "detector", instrument="Cam1", where=f"GLOB({column}, '{pattern}')", explain=False
2559 )
2560 self.assertEqual(len(records), count)
2562 # Check that bind works with pattern.
2563 records = butler.query_dimension_records(
2564 "detector",
2565 instrument="Cam1",
2566 where=f"GLOB({column}, :pattern)",
2567 explain=False,
2568 bind={"pattern": pattern},
2569 )
2570 self.assertEqual(len(records), count)
2572 # Check that glob works on dimension itself, not just metadata.
2573 records = butler.query_dimension_records(
2574 "detector", where="GLOB(instrument, '?a*1') AND GLOB(full_name, '*')"
2575 )
2576 self.assertEqual(len(records), 4)
2578 # Check exceptions.
2579 with self.assertRaisesRegex(InvalidQueryError, "first argument must be a string column"):
2580 butler.query_dimension_records("detector", instrument="Cam1", where="GLOB(detector, '*')")
2582 # This ofails at parser level because parser expects string literal.
2583 with self.assertRaisesRegex(InvalidQueryError, "Failed to parse expression"):
2584 butler.query_dimension_records("detector", instrument="Cam1", where="GLOB(full_name, full_name)")
2586 def test_dataset_id_queries(self) -> None:
2587 """Test queries on dataset_id."""
2588 butler = self.make_butler("base.yaml", "datasets.yaml")
2590 dataset_id = UUID("e15ab039-bc8b-4135-87c5-90902a7c0b22")
2592 refs = butler.query_datasets(
2593 "bias",
2594 "imported_g",
2595 instrument="Cam1",
2596 where="dataset_id = :ID",
2597 bind={"ID": dataset_id},
2598 )
2599 self.assertEqual({ref.id for ref in refs}, {dataset_id})
2601 dataset_ids = {
2602 UUID("87f3e68d-258d-41b7-8ea5-edf3557ccb30"),
2603 UUID("dc0ef017-dc94-4118-b431-d65b1ef89a5f"),
2604 UUID("e255067d-dcc5-4f39-9824-0baa5817d3e5"),
2605 }
2606 refs = butler.query_datasets(
2607 "bias",
2608 "imported_r",
2609 instrument="Cam1",
2610 where="bias.dataset_id IN (:IDS)",
2611 bind={"IDS": dataset_ids},
2612 )
2613 self.assertEqual({ref.id for ref in refs}, dataset_ids)
2615 refs = butler.query_datasets(
2616 "bias",
2617 "imported_g",
2618 instrument="Cam1",
2619 where="dataset_id = UUID('e15ab039-bc8b-4135-87c5-90902a7c0b22')",
2620 )
2621 self.assertEqual({ref.id for ref in refs}, {dataset_id})
2623 refs = butler.query_datasets(
2624 "bias",
2625 "imported_r",
2626 instrument="Cam1",
2627 where=(
2628 "bias.dataset_id IN ("
2629 "UUID('87f3e68d-258d-41b7-8ea5-edf3557ccb30'), "
2630 "UUID('dc0ef017-dc94-4118-b431-d65b1ef89a5f'), "
2631 "UUID('e255067d-dcc5-4f39-9824-0baa5817d3e5')"
2632 ")"
2633 ),
2634 )
2635 self.assertEqual({ref.id for ref in refs}, dataset_ids)
2638def _get_exposure_ids_from_dimension_records(dimension_records: Iterable[DimensionRecord]) -> list[int]:
2639 output = []
2640 for rec in dimension_records:
2641 id = rec.dataId["exposure"]
2642 assert isinstance(id, int)
2643 output.append(id)
2645 return output
2648def _ref_uuids(refs: list[DatasetRef]) -> list[str]:
2649 return [str(ref.id) for ref in refs]