diff --git a/src/hipscat/catalog/catalog.py b/src/hipscat/catalog/catalog.py index edfb0ca2..88460308 100644 --- a/src/hipscat/catalog/catalog.py +++ b/src/hipscat/catalog/catalog.py @@ -2,7 +2,6 @@ from __future__ import annotations -import dataclasses from typing import Any, Dict, List, Tuple, Union import healpy as hp @@ -121,19 +120,6 @@ def filter_by_polygon(self, vertices: List[SphericalCoordinates] | List[Cartesia validate_polygon(vertices) return self.filter_from_pixel_list(filter_pixels_by_polygon(self.pixel_tree, vertices)) - def filter_from_pixel_list(self, pixels: List[HealpixPixel]) -> Catalog: - """Filter the pixels in the catalog to only include the requested pixels. - - Args: - pixels (List[HealpixPixels]): the pixels to include - - Returns: - A new catalog with only those pixels. Note that we reset the total_rows - to None, instead of performing a scan over the new pixel sizes. - """ - filtered_catalog_info = dataclasses.replace(self.catalog_info, total_rows=None) - return Catalog(filtered_catalog_info, pixels) - def generate_negative_tree_pixels(self) -> List[HealpixPixel]: """Get the leaf nodes at each healpix order that have zero catalog data. diff --git a/src/hipscat/catalog/healpix_dataset/healpix_dataset.py b/src/hipscat/catalog/healpix_dataset/healpix_dataset.py index dffd4eed..3608529d 100644 --- a/src/hipscat/catalog/healpix_dataset/healpix_dataset.py +++ b/src/hipscat/catalog/healpix_dataset/healpix_dataset.py @@ -1,7 +1,8 @@ +import dataclasses from typing import Any, Dict, List, Tuple, Union import pandas as pd -from typing_extensions import TypeAlias +from typing_extensions import Self, TypeAlias from hipscat.catalog.dataset import BaseCatalogInfo, Dataset from hipscat.catalog.partition_info import PartitionInfo @@ -102,3 +103,16 @@ def _check_files_exist(cls, catalog_base_dir: FilePointer, storage_options: dict raise FileNotFoundError( f"_metadata or partition info file is required in catalog directory {catalog_base_dir}" ) + + def filter_from_pixel_list(self, pixels: List[HealpixPixel]) -> Self: + """Filter the pixels in the catalog to only include the requested pixels. + + Args: + pixels (List[HealpixPixels]): the pixels to include + + Returns: + A new catalog with only those pixels. Note that we reset the total_rows + to None, instead of performing a scan over the new pixel sizes. + """ + filtered_catalog_info = dataclasses.replace(self.catalog_info, total_rows=None) + return self.__class__(filtered_catalog_info, pixels)