Coverage for tests / test_dimension_record_containers.py: 11%
301 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-04-30 08:41 +0000
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/>.
28import copy
29import json
30import unittest
32import pyarrow as pa
33import pyarrow.parquet as pq
34import pydantic
36from lsst.daf.butler import (
37 DataCoordinate,
38 DimensionDataAttacher,
39 DimensionDataExtractor,
40 DimensionRecordSet,
41 DimensionRecordSetDeserializer,
42 DimensionRecordTable,
43 SerializableDimensionData,
44 SerializedKeyValueDimensionRecord,
45)
46from lsst.daf.butler.arrow_utils import TimespanArrowType
47from lsst.daf.butler.tests.utils import create_populated_sqlite_registry
49DIMENSION_DATA_FILES = [
50 "resource://lsst.daf.butler/tests/registry_data/base.yaml",
51 "resource://lsst.daf.butler/tests/registry_data/spatial.yaml",
52]
55class DimensionRecordContainersTestCase(unittest.TestCase):
56 """Tests for the DimensionRecordTable class."""
58 @classmethod
59 def setUpClass(cls):
60 # Create an in-memory SQLite database and Registry just to import the
61 # YAML data.
62 cls.butler = create_populated_sqlite_registry(*DIMENSION_DATA_FILES)
63 cls.enterClassContext(cls.butler)
64 cls.records = {
65 element: tuple(list(cls.butler.registry.queryDimensionRecords(element)))
66 for element in ("visit", "skymap", "patch")
67 }
68 cls.universe = cls.butler.dimensions
69 cls.data_ids = list(cls.butler.registry.queryDataIds(["visit", "patch"]).expanded())
71 def assert_record_sets_equal(self, a: DimensionRecordSet, b: DimensionRecordSet) -> None:
72 self.assertEqual(a, b) # only checks data ID equality
73 self.assertCountEqual(list(a), list(b))
75 def test_record_table_schema_visit(self):
76 """Test that the Arrow schema for 'visit' has the right types,
77 including dictionary encoding.
78 """
79 schema = DimensionRecordTable.make_arrow_schema(self.universe["visit"])
80 self.assertEqual(schema.field("instrument").type, pa.dictionary(pa.int32(), pa.string()))
81 self.assertEqual(schema.field("id").type, pa.uint64())
82 self.assertEqual(schema.field("physical_filter").type, pa.dictionary(pa.int32(), pa.string()))
83 self.assertEqual(schema.field("name").type, pa.string())
84 self.assertEqual(schema.field("observation_reason").type, pa.dictionary(pa.int32(), pa.string()))
86 def test_record_table_schema_skymap(self):
87 """Test that the Arrow schema for 'skymap' has the right types,
88 including dictionary encoding.
89 """
90 schema = DimensionRecordTable.make_arrow_schema(self.universe["skymap"])
91 self.assertEqual(schema.field("name").type, pa.string())
92 self.assertEqual(schema.field("hash").type, pa.binary())
94 def test_empty_record_table_visit(self):
95 """Test methods on a table that was initialized with no records.
97 We use 'visit' records for this test because they have both timespans
98 and regions, and those are the tricky column types for interoperability
99 with Arrow.
100 """
101 table = DimensionRecordTable(self.universe["visit"])
102 self.assertEqual(len(table), 0)
103 self.assertEqual(list(table), [])
104 self.assertEqual(table.element, self.universe["visit"])
105 with self.assertRaises(IndexError):
106 table[0]
107 self.assertEqual(len(table.column("instrument")), 0)
108 self.assertEqual(len(table.column("id")), 0)
109 self.assertEqual(len(table.column("physical_filter")), 0)
110 self.assertEqual(len(table.column("name")), 0)
111 self.assertEqual(len(table.column("day_obs")), 0)
112 table.extend(self.records["visit"])
113 self.assertCountEqual(table, self.records["visit"])
114 self.assertEqual(
115 table.to_arrow().schema, DimensionRecordTable.make_arrow_schema(self.universe["visit"])
116 )
118 def test_empty_record_table_skymap(self):
119 """Test methods on a table that was initialized with no records.
121 We use 'skymap' records for this test because that's the only one with
122 a "hash" column.
123 """
124 table = DimensionRecordTable(self.universe["skymap"])
125 self.assertEqual(len(table), 0)
126 self.assertEqual(list(table), [])
127 self.assertEqual(table.element, self.universe["skymap"])
128 with self.assertRaises(IndexError):
129 table[0]
130 self.assertEqual(len(table.column("name")), 0)
131 self.assertEqual(len(table.column("hash")), 0)
132 table.extend(self.records["skymap"])
133 self.assertCountEqual(table, self.records["skymap"])
134 self.assertEqual(
135 table.to_arrow().schema, DimensionRecordTable.make_arrow_schema(self.universe["skymap"])
136 )
138 def test_full_record_table_visit(self):
139 """Test methods on a table that was initialized with an iterable.
141 We use 'visit' records for this test because they have both timespans
142 and regions, and those are the tricky column types for interoperability
143 with Arrow.
144 """
145 table = DimensionRecordTable(self.universe["visit"], self.records["visit"])
146 self.assertEqual(len(table), 2)
147 self.assertEqual(table[0], self.records["visit"][0])
148 self.assertEqual(table[1], self.records["visit"][1])
149 self.assertEqual(list(table), list(self.records["visit"]))
150 self.assertListEqual([x.region for x in table], [x.region for x in self.records["visit"]])
151 self.assertListEqual([x.timespan for x in table], [x.timespan for x in self.records["visit"]])
152 self.assertEqual(table.element, self.universe["visit"])
153 self.assertEqual(table.column("instrument")[0].as_py(), "Cam1")
154 self.assertEqual(table.column("instrument")[1].as_py(), "Cam1")
155 self.assertEqual(table.column("id")[0].as_py(), 1)
156 self.assertEqual(table.column("id")[1].as_py(), 2)
157 self.assertEqual(table.column("name")[0].as_py(), "1")
158 self.assertEqual(table.column("name")[1].as_py(), "2")
159 self.assertEqual(table.column("physical_filter")[0].as_py(), "Cam1-G")
160 self.assertEqual(table.column("physical_filter")[1].as_py(), "Cam1-R1")
161 self.assertEqual(table.column("day_obs")[0].as_py(), 20210909)
162 self.assertEqual(table.column("day_obs")[1].as_py(), 20210909)
163 self.assertEqual(list(table[:1]), list(self.records["visit"][:1]))
164 self.assertEqual(list(table[1:]), list(self.records["visit"][1:]))
165 table.extend(self.records["visit"])
166 self.assertCountEqual(table, self.records["visit"] + self.records["visit"])
168 def test_full_record_table_skymap(self):
169 """Test methods on a table that was initialized with an iterable.
171 We use 'skymap' records for this test because that's the only one with
172 a "hash" column.
173 """
174 table = DimensionRecordTable(self.universe["skymap"], self.records["skymap"])
175 self.assertEqual(len(table), 1)
176 self.assertEqual(table[0], self.records["skymap"][0])
177 self.assertEqual(list(table), list(self.records["skymap"]))
178 self.assertEqual(table.element, self.universe["skymap"])
179 self.assertEqual(table.column("name")[0].as_py(), "SkyMap1")
180 self.assertEqual(table.column("hash")[0].as_py(), b"notreallyahashofanything!")
181 table.extend(self.records["skymap"])
182 self.assertCountEqual(table, self.records["skymap"] + self.records["skymap"])
184 def test_record_table_parquet_visit(self):
185 """Test round-tripping a dimension record table through Parquet.
187 We use 'visit' records for this test because they have both timespans
188 and regions, and those are the tricky column types for interoperability
189 with Arrow.
190 """
191 original_records = list(self.records["visit"])
192 table1 = DimensionRecordTable(self.universe["visit"], self.records["visit"])
193 stream = pa.BufferOutputStream()
194 pq.write_table(table1.to_arrow(), stream)
195 table2 = DimensionRecordTable(
196 universe=self.universe, table=pq.read_table(pa.BufferReader(stream.getvalue()))
197 )
198 self.assertEqual(list(table1), list(table2))
199 self.assertEqual(original_records, list(table2))
200 self.assertListEqual([x.region for x in table1], [x.region for x in table2])
201 self.assertListEqual([x.region for x in original_records], [x.region for x in table2])
202 self.assertListEqual([x.timespan for x in table1], [x.timespan for x in table2])
203 self.assertListEqual([x.timespan for x in original_records], [x.timespan for x in table2])
205 def test_record_table_parquet_skymap(self):
206 """Test round-tripping a dimension record table through Parquet.
208 We use 'skymap' records for this test because that's the only one with
209 a "hash" column.
210 """
211 table1 = DimensionRecordTable(self.universe["skymap"], self.records["skymap"])
212 stream = pa.BufferOutputStream()
213 pq.write_table(table1.to_arrow(), stream)
214 table2 = DimensionRecordTable(
215 universe=self.universe, table=pq.read_table(pa.BufferReader(stream.getvalue()))
216 )
217 self.assertEqual(list(table1), list(table2))
219 def test_record_chunk_init(self):
220 """Test constructing a DimensionRecordTable from an iterable in chunks.
222 We use 'patch' records for this test because there are enough of them
223 to have multiple chunks.
224 """
225 table1 = DimensionRecordTable(self.universe["patch"], self.records["patch"], batch_size=5)
226 self.assertEqual(len(table1), 12)
227 self.assertEqual([len(batch) for batch in table1.to_arrow().to_batches()], [5, 5, 2])
228 self.assertEqual(list(table1), list(self.records["patch"]))
230 def test_arrow_timespan_nulls(self):
231 """Test that null values for Timespan can be round-tripped through
232 arrow.
233 """
234 schema = pa.schema([pa.field("timespan", TimespanArrowType(), nullable=True)])
235 table = pa.Table.from_pylist([{"timespan": None}], schema=schema)
236 output_value = table.to_pylist()[0]["timespan"]
237 self.assertIsNone(output_value)
239 def test_record_set_const(self):
240 """Test attributes and methods of `DimensionRecordSet` that do not
241 modify the set.
243 We use 'patch' records for this test because there are enough of them
244 to do nontrivial set-operation tests.
245 """
246 element = self.universe["patch"]
247 records = self.records["patch"]
248 set1 = DimensionRecordSet(element, records[:7])
249 self.assertEqual(set1, DimensionRecordSet("patch", records[:7], universe=self.universe))
250 # DimensionRecordSets do not compare as equal with other set types,
251 # even with the same content.
252 self.assertNotEqual(set1, set(records[:7]))
253 with self.assertRaises(TypeError):
254 DimensionRecordSet("patch", records[:7])
255 self.assertEqual(set1.element, self.universe["patch"])
256 self.assertEqual(len(set1), 7)
257 self.assertEqual(list(set1), list(records[:7]))
258 self.assertIn(records[4], set1)
259 self.assertIn(records[5].dataId, set1)
260 self.assertNotIn(self.records["visit"][0], set1)
261 self.assertTrue(set1.issubset(DimensionRecordSet(element, records[:8])))
262 self.assertFalse(set1.issubset(DimensionRecordSet(element, records[1:6])))
263 with self.assertRaises(ValueError):
264 set1.issubset(DimensionRecordSet(self.universe["tract"]))
265 self.assertTrue(set1.issuperset(DimensionRecordSet(element, records[1:6])))
266 self.assertFalse(set1.issuperset(DimensionRecordSet(element, records[:8])))
267 with self.assertRaises(ValueError):
268 set1.issuperset(DimensionRecordSet(self.universe["tract"]))
269 self.assertTrue(set1.isdisjoint(DimensionRecordSet(element, records[7:])))
270 self.assertFalse(set1.isdisjoint(DimensionRecordSet(element, records[5:8])))
271 with self.assertRaises(ValueError):
272 set1.isdisjoint(DimensionRecordSet(self.universe["tract"]))
273 self.assertEqual(
274 set1.intersection(DimensionRecordSet(element, records[5:])),
275 DimensionRecordSet(element, records[5:7]),
276 )
277 self.assertEqual(
278 set1.intersection(DimensionRecordSet(element, records[5:])),
279 DimensionRecordSet(element, records[5:7]),
280 )
281 with self.assertRaises(ValueError):
282 set1.intersection(DimensionRecordSet(self.universe["tract"]))
283 self.assertEqual(
284 set1.difference(DimensionRecordSet(element, records[5:])),
285 DimensionRecordSet(element, records[:5]),
286 )
287 with self.assertRaises(ValueError):
288 set1.difference(DimensionRecordSet(self.universe["tract"]))
289 self.assertEqual(
290 set1.union(DimensionRecordSet(element, records[5:9])),
291 DimensionRecordSet(element, records[:9]),
292 )
293 with self.assertRaises(ValueError):
294 set1.union(DimensionRecordSet(self.universe["tract"]))
295 self.assertEqual(set1.find(records[0].dataId), records[0])
296 with self.assertRaises(LookupError):
297 set1.find(self.records["patch"][8].dataId)
298 with self.assertRaises(ValueError):
299 set1.find(self.records["visit"][0].dataId)
300 self.assertEqual(set1.find_with_required_values(records[0].dataId.required_values), records[0])
302 def test_record_set_add(self):
303 """Test DimensionRecordSet.add."""
304 set1 = DimensionRecordSet("patch", self.records["patch"][:2], universe=self.universe)
305 set1.add(self.records["patch"][2])
306 with self.assertRaises(ValueError):
307 set1.add(self.records["visit"][0])
308 self.assertEqual(set1, DimensionRecordSet("patch", self.records["patch"][:3], universe=self.universe))
309 set1.add(self.records["patch"][2])
310 self.assertEqual(list(set1), list(self.records["patch"][:3]))
312 def test_record_set_find_or_add(self):
313 """Test DimensionRecordSet.find and find_with_required_values with
314 a 'or_add' callback.
315 """
316 set1 = DimensionRecordSet("patch", self.records["patch"][:2], universe=self.universe)
317 set1.find(self.records["patch"][2].dataId, or_add=lambda _c, _r: self.records["patch"][2])
318 with self.assertRaises(ValueError):
319 set1.find(self.records["visit"][0].dataId, or_add=lambda _c, _r: self.records["visit"][0])
320 self.assertEqual(set1, DimensionRecordSet("patch", self.records["patch"][:3], universe=self.universe))
322 set1.find_with_required_values(
323 self.records["patch"][3].dataId.required_values, or_add=lambda _c, _r: self.records["patch"][3]
324 )
325 self.assertEqual(set1, DimensionRecordSet("patch", self.records["patch"][:4], universe=self.universe))
327 def test_record_set_update_from_data_coordinates(self):
328 """Test DimensionRecordSet.update_from_data_coordinates."""
329 set1 = DimensionRecordSet("patch", self.records["patch"][:2], universe=self.universe)
330 set1.update_from_data_coordinates(self.data_ids)
331 for data_id in self.data_ids:
332 self.assertIn(data_id.records["patch"], set1)
334 def test_record_set_discard(self):
335 """Test DimensionRecordSet.discard."""
336 set1 = DimensionRecordSet("patch", self.records["patch"][:2], universe=self.universe)
337 set2 = copy.deepcopy(set1)
338 # These discards should do nothing.
339 set1.discard(self.records["patch"][2])
340 self.assertEqual(set1, set2)
341 set1.discard(self.records["patch"][2].dataId)
342 self.assertEqual(set1, set2)
343 with self.assertRaises(ValueError):
344 set1.discard(self.records["visit"][0])
345 self.assertEqual(set1, set2)
346 with self.assertRaises(ValueError):
347 set1.discard(self.records["visit"][0].dataId)
348 self.assertEqual(set1, set2)
349 # These ones should remove a record from each set.
350 set1.discard(self.records["patch"][1])
351 set2.discard(self.records["patch"][1].dataId)
352 self.assertEqual(set1, set2)
353 self.assertNotIn(self.records["patch"][1], set1)
354 self.assertNotIn(self.records["patch"][1], set2)
356 def test_record_set_remove(self):
357 """Test DimensionRecordSet.remove."""
358 set1 = DimensionRecordSet("patch", self.records["patch"][:2], universe=self.universe)
359 set2 = copy.deepcopy(set1)
360 # These removes should raise with strong exception safety.
361 with self.assertRaises(KeyError):
362 set1.remove(self.records["patch"][2])
363 self.assertEqual(set1, set2)
364 with self.assertRaises(KeyError):
365 set1.remove(self.records["patch"][2].dataId)
366 self.assertEqual(set1, set2)
367 with self.assertRaises(ValueError):
368 set1.remove(self.records["visit"][0])
369 self.assertEqual(set1, set2)
370 with self.assertRaises(ValueError):
371 set1.remove(self.records["visit"][0].dataId)
372 self.assertEqual(set1, set2)
373 # These ones should remove a record from each set.
374 set1.remove(self.records["patch"][1])
375 set2.remove(self.records["patch"][1].dataId)
376 self.assertEqual(set1, set2)
377 self.assertNotIn(self.records["patch"][1], set1)
378 self.assertNotIn(self.records["patch"][1], set2)
380 def test_record_set_pop(self):
381 """Test DimensionRecordSet.pop."""
382 set1 = DimensionRecordSet("patch", self.records["patch"][:2], universe=self.universe)
383 set2 = copy.deepcopy(set1)
384 record1 = set1.pop()
385 set2.remove(record1)
386 self.assertNotIn(record1, set1)
387 self.assertEqual(set1, set2)
388 record2 = set1.pop()
389 set2.remove(record2)
390 self.assertNotIn(record2, set1)
391 self.assertEqual(set1, set2)
392 self.assertFalse(set1)
394 def test_record_key_value_serialization(self):
395 """Test DimensionRecord.serialize_key_value, deserialize_key, and
396 deserialize_value.
397 """
398 for element, records in self.records.items():
399 with self.subTest(element=element):
400 for record1 in records:
401 raw1 = record1.serialize_key_value()
402 self.assertIsInstance(raw1, list)
403 raw2 = json.loads(json.dumps(raw1))
404 self.assertEqual(raw1, raw2)
405 key, raw_value = record1.definition.RecordClass.deserialize_key(raw2)
406 self.assertEqual(key, record1.dataId.required_values)
407 record2 = record1.definition.RecordClass.deserialize_value(key, raw_value)
408 self.assertEqual(record1, record2)
410 def test_record_set_serialization(self):
411 """Test DimensionRecordSet.serialize_records and deserialize_records,
412 as well as DimensionRecordSetDeserializer.
413 """
414 adapter = pydantic.TypeAdapter(list[SerializedKeyValueDimensionRecord])
415 for element, records in self.records.items():
416 with self.subTest(element=element):
417 rs1 = DimensionRecordSet(element, records, universe=self.universe)
418 raw = adapter.validate_json(adapter.dump_json(rs1.serialize_records()))
419 rs2 = DimensionRecordSet(element, universe=self.universe)
420 rs2.deserialize_records(raw)
421 self.assert_record_sets_equal(rs1, rs2)
422 deserializer = DimensionRecordSetDeserializer.from_raw(self.universe[element], raw)
423 self.assertEqual(len(deserializer), len(records))
424 for record1 in rs1:
425 record2 = deserializer[record1.dataId.required_values]
426 self.assertEqual(record1, record2)
428 def test_extract_serialize_attach(self):
429 """Test serializing expanded data IDs by first extracting their
430 dimension records, serializing them in normalized form, and then
431 attaching the records to deserialized data IDs.
432 """
433 dimensions = self.universe.conform(["visit", "patch", "htm5"])
434 expanded_data_ids_1: list[DataCoordinate] = []
435 pixelization = dimensions.universe.skypix_dimensions["htm5"].pixelization
436 for data_id in self.data_ids:
437 for htm_begin, htm_end in pixelization.envelope(data_id.records["patch"].region):
438 for htm_id in range(htm_begin, htm_end):
439 new_data_id = self.butler.registry.expandDataId(data_id, htm5=htm_id)
440 self.assertEqual(new_data_id.dimensions, dimensions)
441 expanded_data_ids_1.append(new_data_id)
442 extractor = DimensionDataExtractor.from_dimension_group(
443 dimensions,
444 ignore=["tract"],
445 ignore_cached=True,
446 include_skypix=False,
447 )
448 extractor.update(expanded_data_ids_1)
449 thin_data_ids = [
450 DataCoordinate.from_full_values(dimensions, data_id.full_values)
451 for data_id in expanded_data_ids_1
452 ]
453 model1 = SerializableDimensionData.from_record_sets(extractor.records.values())
454 self.assertEqual(model1.root.keys(), {"visit", "patch", "day_obs"})
455 # Delete one patch record from the model to probe more logic branches
456 # in the deserialization; like the tract records, we'll recover it by
457 # querying the butler.
458 del model1.root["patch"][-1]
459 model2 = SerializableDimensionData.model_validate_json(model1.model_dump_json())
460 with self.butler.query() as query:
461 attacher1 = DimensionDataAttacher(
462 deserializers=model2.make_deserializers(self.universe),
463 cache=query._driver._dimension_record_cache,
464 dimensions=dimensions,
465 )
466 expanded_data_ids_2 = attacher1.attach(dimensions, thin_data_ids, query=query)
467 # This equality check is necessary but not sufficient for this test,
468 # because DataCoordinate equality doesn't look at records.
469 self.assertEqual(expanded_data_ids_1, expanded_data_ids_2)
470 for a, b in zip(expanded_data_ids_1, expanded_data_ids_2):
471 self.assertEqual(dict(a.records), dict(b.records))
472 # Caching in the attacher should ensure that there is only one copy
473 # of each dimension record in the new set.
474 all_values = {(data_id["tract"], data_id["patch"]) for data_id in expanded_data_ids_2}
475 all_identities = {
476 id(data_id.records["patch"]): data_id.records["patch"] for data_id in expanded_data_ids_2
477 }
478 self.assertGreater(len(all_values), 1)
479 self.assertEqual(len(all_identities), len(all_values), msg=str(all_identities))
480 # Serialize and round-trip again, via attacher.serialized(), and attach
481 # all records again.
482 model3 = attacher1.serialized()
483 attacher2 = DimensionDataAttacher(
484 deserializers=model3.make_deserializers(self.universe), dimensions=dimensions
485 )
486 expanded_data_ids_4 = attacher2.attach(dimensions, thin_data_ids)
487 self.assertEqual(expanded_data_ids_1, expanded_data_ids_4)
488 for a, b in zip(expanded_data_ids_1, expanded_data_ids_4):
489 self.assertEqual(dict(a.records), dict(b.records))
490 # Attach data IDs with the original attacher again, after clearing out
491 # its deserializers. The cached records should all be present, and it
492 # should work without any butler queries.
493 attacher1.deserializers.clear()
494 expanded_data_ids_3 = attacher1.attach(dimensions, thin_data_ids)
495 self.assertEqual(expanded_data_ids_1, expanded_data_ids_3)
496 for a, b in zip(expanded_data_ids_1, expanded_data_ids_3):
497 self.assertEqual(dict(a.records), dict(b.records))
498 # Delete a record and try again to check error-handling.
499 removed_patch_data_id = next(iter(attacher1.records["patch"])).dataId
500 attacher1.records["patch"].remove(removed_patch_data_id)
501 with self.assertRaisesRegex(LookupError, r"No dimension record.*patch.*"):
502 attacher1.attach(dimensions, thin_data_ids)
505if __name__ == "__main__":
506 unittest.main()