Add configuration for a healpix column name. (#601)

* Add data/test for non-order-29 indexed dataset.

* More testing and allowing margin cache with no healpix 29

* Set the properties on new catalogs.

* Refresh unit test data.

Author: delucchi-cmu

src/hats_import/catalog/arguments.py

@@ -12,6 +12,7 @@
 from hats.io.paths import DATASET_DIR, PARTITION_ORDER
 from hats.io.validation import is_valid_catalog
 from hats.pixel_math import spatial_index
+from hats.pixel_math.spatial_index import SPATIAL_INDEX_COLUMN, SPATIAL_INDEX_ORDER
 from typing_extensions import Self
 from upath import UPath
 
@@ -178,6 +179,13 @@ def to_table_properties(
                     "hats_skymap_alt_orders": self.skymap_alt_orders,
                 }
             )
+        if self.add_healpix_29:
+            info.update(
+                {
+                    "hats_col_healpix": SPATIAL_INDEX_COLUMN,
+                    "hats_col_healpix_order": SPATIAL_INDEX_ORDER,
+                }
+            )
         properties = TableProperties(**info)
 
         if properties.default_columns and column_names:

src/hats_import/margin_cache/margin_cache.py

@@ -32,6 +32,8 @@ def generate_margin_cache(args, client):
                     margin_pair_file=resume_plan.margin_pair_file,
                     output_path=args.tmp_path,
                     margin_order=args.margin_order,
+                    healpix_column=args.catalog.catalog_info.healpix_column,
+                    healpix_order=args.catalog.catalog_info.healpix_order,
                 )
             )
         resume_plan.wait_for_mapping(futures)

src/hats_import/margin_cache/margin_cache_arguments.py

@@ -51,6 +51,8 @@ def _check_arguments(self):
             self.catalog = self.catalog.filter_from_pixel_list(self.debug_filter_pixel_list)
             if len(self.catalog.get_healpix_pixels()) == 0:
                 raise ValueError("debug_filter_pixel_list has created empty catalog")
+        if not self.catalog.has_healpix_column():
+            raise ValueError("Only catalogs with some healpix column (e.g. _healpix_29) can have a margin")
 
         if self.fine_filtering:
             raise NotImplementedError("Fine filtering temporarily removed.")
@@ -86,5 +88,7 @@ def to_table_properties(
             "margin_threshold": self.margin_threshold,
             "hats_order": highest_order,
             "moc_sky_fraction": f"{moc_sky_fraction:0.5f}",
+            "hats_col_healpix": self.catalog.catalog_info.healpix_column,
+            "hats_col_healpix_order": self.catalog.catalog_info.healpix_order,
         } | self.extra_property_dict()
         return TableProperties(**info)

src/hats_import/margin_cache/margin_cache_map_reduce.py

@@ -4,7 +4,7 @@
 import pyarrow.parquet as pq
 from hats.io import file_io, paths
 from hats.pixel_math.healpix_pixel import HealpixPixel
-from hats.pixel_math.spatial_index import SPATIAL_INDEX_COLUMN, spatial_index_to_healpix
+from hats.pixel_math.spatial_index import spatial_index_to_healpix
 
 from hats_import.margin_cache.margin_cache_resume_plan import MarginCachePlan
 from hats_import.pipeline_resume_plan import get_pixel_cache_directory, print_task_failure
@@ -19,6 +19,8 @@ def map_pixel_shards(
     margin_pair_file,
     output_path,
     margin_order,
+    healpix_column,
+    healpix_order,
 ):
     """Creates margin cache shards from a source partition file."""
     try:
@@ -39,8 +41,9 @@ def map_pixel_shards(
         )
 
         margin_pixel_list = spatial_index_to_healpix(
-            data[SPATIAL_INDEX_COLUMN].to_numpy(),
+            data[healpix_column].to_numpy(),
             target_order=margin_order,
+            spatial_index_order=healpix_order,
         )
         margin_pixel_filter = pd.DataFrame(
             {"margin_pixel": margin_pixel_list, "filter_value": np.arange(0, len(margin_pixel_list))}
@@ -60,6 +63,7 @@ def map_pixel_shards(
                 pixel=pixel,
                 output_path=output_path,
                 source_pixel=source_pixel,
+                healpix_column=healpix_column,
             )
 
         MarginCachePlan.mapping_key_done(output_path, mapping_key, num_rows)
@@ -74,6 +78,7 @@ def _to_pixel_shard(
     pixel,
     output_path,
     source_pixel,
+    healpix_column,
 ):
     """Do boundary checking for the cached partition and then output remaining data."""
     margin_data = filtered_data
@@ -88,7 +93,7 @@ def _to_pixel_shard(
 
         shard_path = paths.pixel_catalog_file(partition_dir, source_pixel)
 
-        margin_data = margin_data.sort_by(SPATIAL_INDEX_COLUMN)
+        margin_data = margin_data.sort_by(healpix_column)
 
         pq.write_table(margin_data, shard_path.path, filesystem=shard_path.fs)
     return num_rows