Merge 28959e3 into 2782705

Author: delucchi-cmu

src/hats/catalog/healpix_dataset/healpix_dataset.py

@@ -418,6 +418,27 @@ def has_healpix_column(self):
                 return True
         return False
 
+    def get_pixel_paths(self):
+        """Generate paths to all pixel files.
+
+        Pixels will be traversed in "breadth-first" healpix order. If any spatial filters
+        have been applied to this catalog, only those pixels that remain will be included.
+
+        Yields
+        ------
+        UPath
+            Universal Pathlib pointing to either an npix directory, or to a single
+            pixel partition data file.
+        """
+        if not self.on_disk:
+            warnings.warn("Calling read_pixel_to_pandas on an in-memory catalog. No results.")
+            return
+
+        for pixel in self.get_healpix_pixels():
+            yield paths.pixel_catalog_file(
+                self.catalog_base_dir, pixel, npix_suffix=self.catalog_info.npix_suffix
+            )
+
     def read_pixel_to_pandas(self, pixel: HealpixPixel, **kwargs) -> npd.NestedFrame:
         """Read the parquet file(s) for this pixel into a pandas dataframe.
 

tests/hats/catalog/loaders/test_read_hats.py

@@ -125,10 +125,16 @@ def test_read_hats_nonstandard_npix_suffix(
     cat = read_hats(small_sky_npix_alt_suffix_dir)
     result = cat.read_pixel_to_pandas(cat.get_healpix_pixels()[0])
     assert len(result) == 131
+    for path in cat.get_pixel_paths():
+        assert path.exists()
+        assert str(path).endswith(".parq")
 
     cat = read_hats(small_sky_npix_as_dir_dir)
     result = cat.read_pixel_to_pandas(cat.get_healpix_pixels()[0])
     assert len(result) == 131
+    for path in cat.get_pixel_paths():
+        assert path.exists()
+        assert path.is_dir()
 
 
 def test_read_hats_original_schema(small_sky_order1_dir):

tests/hats/catalog/test_catalog.py

@@ -235,6 +235,35 @@ def test_cone_filter(small_sky_order1_catalog):
     assert filtered_catalog.original_schema is not None
 
 
+def test_get_pixel_paths(small_sky_order1_catalog):
+    generator_length = 0
+    for path in small_sky_order1_catalog.get_pixel_paths():
+        assert path.exists()
+        generator_length += 1
+    assert generator_length == 4
+
+    ra = 315
+    dec = -66.443
+    radius = 0.1
+
+    filtered_catalog = small_sky_order1_catalog.filter_by_cone(ra, dec, radius)
+    generator_length = 0
+    for path in filtered_catalog.get_pixel_paths():
+        assert path.exists()
+        generator_length += 1
+    assert generator_length == 1
+
+
+def test_get_pixel_paths_in_memory(in_memory_catalog):
+    with pytest.warns(UserWarning, match="in-memory"):
+        with pytest.raises(StopIteration):
+            next(in_memory_catalog.get_pixel_paths())
+
+    with pytest.warns(UserWarning, match="in-memory"):
+        for _ in in_memory_catalog.get_pixel_paths():
+            assert False, "Iterator should be empty."
+
+
 def test_cone_filter_big(small_sky_order1_catalog):
     filtered_catalog = small_sky_order1_catalog.filter_by_cone(315, -66.443, 30 * 3600)
     assert len(filtered_catalog.get_healpix_pixels()) == 4