diff --git a/src/lsdb/loaders/hipscat/abstract_catalog_loader.py b/src/lsdb/loaders/hipscat/abstract_catalog_loader.py index 3e655848..ed10f8f0 100644 --- a/src/lsdb/loaders/hipscat/abstract_catalog_loader.py +++ b/src/lsdb/loaders/hipscat/abstract_catalog_loader.py @@ -13,6 +13,7 @@ from hipscat.io.file_io import file_io from hipscat.pixel_math import HealpixPixel from hipscat.pixel_math.healpix_pixel_function import get_pixel_argsort +from hipscat.pixel_math.hipscat_id import HIPSCAT_ID_COLUMN from upath import UPath from lsdb.catalog.catalog import DaskDFPixelMap @@ -86,6 +87,9 @@ def _create_dask_meta_schema(self, schema: pa.Schema) -> npd.NestedFrame: dask_meta_schema = schema.empty_table().to_pandas(types_mapper=self.config.get_dtype_mapper()) if self.config.columns is not None: dask_meta_schema = dask_meta_schema[self.config.columns] + + if dask_meta_schema.index.name != HIPSCAT_ID_COLUMN and HIPSCAT_ID_COLUMN in dask_meta_schema.columns: + dask_meta_schema = dask_meta_schema.set_index(HIPSCAT_ID_COLUMN) return npd.NestedFrame(dask_meta_schema) def _get_kwargs(self) -> dict: @@ -104,6 +108,11 @@ def read_pixel( **kwargs, ): """Utility method to read a single pixel's parquet file from disk.""" - return file_io.read_parquet_file_to_pandas( + dataframe = file_io.read_parquet_file_to_pandas( hc.io.pixel_catalog_file(catalog.catalog_base_dir, pixel, query_url_params), columns=columns, **kwargs ) + + if dataframe.index.name != HIPSCAT_ID_COLUMN and HIPSCAT_ID_COLUMN in dataframe.columns: + dataframe = dataframe.set_index(HIPSCAT_ID_COLUMN) + + return dataframe diff --git a/tests/data/small_sky/Norder=0/Dir=0/Npix=11.parquet b/tests/data/small_sky/Norder=0/Dir=0/Npix=11.parquet index e0cb8d94..e4aa453a 100644 Binary files a/tests/data/small_sky/Norder=0/Dir=0/Npix=11.parquet and b/tests/data/small_sky/Norder=0/Dir=0/Npix=11.parquet differ diff --git a/tests/data/small_sky/_common_metadata b/tests/data/small_sky/_common_metadata index 4cf7a744..8b9d1c75 100644 Binary files a/tests/data/small_sky/_common_metadata and b/tests/data/small_sky/_common_metadata differ diff --git a/tests/data/small_sky/_metadata b/tests/data/small_sky/_metadata index 26df207b..7134c142 100644 Binary files a/tests/data/small_sky/_metadata and b/tests/data/small_sky/_metadata differ