Skip to content

Commit

Permalink
Set _hipscat_index as pandas index, where possible. (#415)
Browse files Browse the repository at this point in the history
* Set _hipscat_index as pandas index, where possible.

* Fix bad merge.
  • Loading branch information
delucchi-cmu authored Sep 16, 2024
1 parent cb30209 commit c478401
Show file tree
Hide file tree
Showing 4 changed files with 10 additions and 1 deletion.
11 changes: 10 additions & 1 deletion src/lsdb/loaders/hipscat/abstract_catalog_loader.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand All @@ -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
Binary file modified tests/data/small_sky/Norder=0/Dir=0/Npix=11.parquet
Binary file not shown.
Binary file modified tests/data/small_sky/_common_metadata
Binary file not shown.
Binary file modified tests/data/small_sky/_metadata
Binary file not shown.

0 comments on commit c478401

Please sign in to comment.