Expand column statistics to limit by pixels (#472)

* Expand column statistics to limit by pixels, and be more friendly for nested parquet.

* Ooops - black formatting.

* Plumb more include_pixlels.

* Bad merge - missing import.

* Another missing import?

Author: delucchi-cmu

src/hats/catalog/healpix_dataset/healpix_dataset.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import warnings
 from pathlib import Path
 
 import astropy.units as u
@@ -16,6 +17,7 @@
 from hats.catalog.partition_info import PartitionInfo
 from hats.inspection import plot_pixels
 from hats.inspection.visualize_catalog import plot_moc
+from hats.io.parquet_metadata import aggregate_column_statistics
 from hats.pixel_math import HealpixPixel
 from hats.pixel_math.box_filter import generate_box_moc, wrap_ra_angles
 from hats.pixel_math.validators import (
@@ -256,3 +258,33 @@ def plot_moc(self, **kwargs):
         plot_args = {"title": default_title}
         plot_args.update(kwargs)
         return plot_moc(self.moc, **plot_args)
+
+    def aggregate_column_statistics(
+        self,
+        exclude_hats_columns: bool = True,
+        exclude_columns: list[str] = None,
+        include_columns: list[str] = None,
+        include_pixels: list[HealpixPixel] = None,
+    ):
+        """Read footer statistics in parquet metadata, and report on global min/max values.
+
+        Args:
+            exclude_hats_columns (bool): exclude HATS spatial and partitioning fields
+                from the statistics. Defaults to True.
+            exclude_columns (List[str]): additional columns to exclude from the statistics.
+            include_columns (List[str]): if specified, only return statistics for the column
+                names provided. Defaults to None, and returns all non-hats columns.
+        """
+        if not self.on_disk:
+            warnings.warn("Calling aggregate_column_statistics on an in-memory catalog. No results.")
+            return pd.DataFrame()
+
+        if include_pixels is None:
+            include_pixels = self.get_healpix_pixels()
+        return aggregate_column_statistics(
+            self.catalog_base_dir / "dataset" / "_metadata",
+            exclude_hats_columns=exclude_hats_columns,
+            exclude_columns=exclude_columns,
+            include_columns=include_columns,
+            include_pixels=include_pixels,
+        )

src/hats/io/parquet_metadata.py

@@ -11,6 +11,7 @@
 
 from hats.io import file_io, paths
 from hats.io.file_io.file_pointer import get_upath
+from hats.pixel_math.healpix_pixel import HealpixPixel
 from hats.pixel_math.healpix_pixel_function import get_pixel_argsort
 
 
@@ -131,6 +132,7 @@ def aggregate_column_statistics(
     exclude_hats_columns: bool = True,
     exclude_columns: list[str] = None,
     include_columns: list[str] = None,
+    include_pixels: list[HealpixPixel] = None,
 ):
     """Read footer statistics in parquet metadata, and report on global min/max values.
 
@@ -157,6 +159,7 @@ def aggregate_column_statistics(
     column_names = [
         first_row_group.column(col).path_in_schema for col in range(0, first_row_group.num_columns)
     ]
+    column_names = [name.removesuffix(".list.element") for name in column_names]
     good_column_indexes = [
         index
         for index, name in enumerate(column_names)
@@ -166,42 +169,43 @@ def aggregate_column_statistics(
     if not good_column_indexes:
         return pd.DataFrame()
     column_names = [column_names[i] for i in good_column_indexes]
-    extrema = [
-        (
-            (None, None, 0)
-            if first_row_group.column(col).statistics is None
-            else (
-                first_row_group.column(col).statistics.min,
-                first_row_group.column(col).statistics.max,
-                first_row_group.column(col).statistics.null_count,
-            )
-        )
-        for col in good_column_indexes
-    ]
+    extrema = None
 
-    for row_group_index in range(1, num_row_groups):
+    for row_group_index in range(0, num_row_groups):
         row_group = total_metadata.row_group(row_group_index)
+        if include_pixels is not None:
+            pixel = paths.get_healpix_from_path(row_group.column(0).file_path)
+            if pixel not in include_pixels:
+                continue
         row_stats = [
             (
-                (None, None, 0)
+                (None, None, 0, 0)
                 if row_group.column(col).statistics is None
                 else (
                     row_group.column(col).statistics.min,
                     row_group.column(col).statistics.max,
                     row_group.column(col).statistics.null_count,
+                    row_group.column(col).num_values,
                 )
             )
             for col in good_column_indexes
         ]
+        if extrema is None:
+            extrema = row_stats
         ## This is annoying, but avoids extra copies, or none comparison.
-        extrema = [
-            (
-                (_nonemin(extrema[col][0], row_stats[col][0])),
-                (_nonemax(extrema[col][1], row_stats[col][1])),
-                extrema[col][2] + row_stats[col][2],
-            )
-            for col in range(0, len(good_column_indexes))
-        ]
+        else:
+            extrema = [
+                (
+                    (_nonemin(extrema[col][0], row_stats[col][0])),
+                    (_nonemax(extrema[col][1], row_stats[col][1])),
+                    extrema[col][2] + row_stats[col][2],
+                    extrema[col][3] + row_stats[col][3],
+                )
+                for col in range(0, len(good_column_indexes))
+            ]
+
+    if extrema is None:
+        return pd.DataFrame()
 
     stats_lists = np.array(extrema).T
 
@@ -211,6 +215,7 @@ def aggregate_column_statistics(
             "min_value": stats_lists[0],
             "max_value": stats_lists[1],
             "null_count": stats_lists[2],
+            "row_count": stats_lists[3],
         }
     ).set_index("column_names")
     return frame

tests/conftest.py

@@ -17,6 +17,7 @@
 SMALL_SKY_NPIX_AS_DIR_NAME = "small_sky_npix_as_dir"
 SMALL_SKY_ORDER1_DIR_NAME = "small_sky_order1"
 SMALL_SKY_SOURCE_OBJECT_INDEX_DIR_NAME = "small_sky_source_object_index"
+SMALL_SKY_NESTED_DIR_NAME = "small_sky_nested"
 
 TEST_DIR = os.path.dirname(__file__)
 
@@ -53,6 +54,11 @@ def small_sky_order1_dir(test_data_dir):
     return test_data_dir / SMALL_SKY_ORDER1_DIR_NAME
 
 
+@pytest.fixture
+def small_sky_nested_dir(test_data_dir):
+    return test_data_dir / SMALL_SKY_NESTED_DIR_NAME
+
+
 @pytest.fixture
 def small_sky_source_object_index_dir(test_data_dir):
     return test_data_dir / SMALL_SKY_SOURCE_OBJECT_INDEX_DIR_NAME

tests/data/generate_data.ipynb

@@ -41,6 +41,7 @@
     "from hats.catalog.association_catalog.partition_join_info import PartitionJoinInfo\n",
     "from hats.catalog.dataset.table_properties import TableProperties\n",
     "from hats.io.file_io import remove_directory\n",
+    "import lsdb\n",
     "from hats.pixel_math.spatial_index import healpix_to_spatial_index\n",
     "\n",
     "tmp_path = tempfile.TemporaryDirectory()\n",
@@ -446,6 +447,40 @@
     "    runner.pipeline_with_client(args, client)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Nested catalog: small_sky_nested\n",
+    "\n",
+    "Nests light curves from `small_sky_source` into `small_sky_order1` object catalog."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "remove_directory(\"./small_sky_nested\")\n",
+    "\n",
+    "small_sky_object = lsdb.read_hats(\"small_sky_order1\")\n",
+    "small_sky_source = lsdb.read_hats(\"small_sky_source\")\n",
+    "small_sky_nested = small_sky_object.join_nested(\n",
+    "    small_sky_source, left_on=\"id\", right_on=\"object_id\", nested_column_name=\"lc\"\n",
+    ")\n",
+    "small_sky_nested = small_sky_nested.map_partitions(\n",
+    "    lambda df, p: df.assign(Norder=p.order, Npix=p.pixel, Dir=p.pixel // 10000), include_pixel=True\n",
+    ")\n",
+    "lsdb.io.to_hats(\n",
+    "    small_sky_nested,\n",
+    "    base_catalog_path=\"small_sky_nested\",\n",
+    "    catalog_name=\"small_sky_nested\",\n",
+    "    histogram_order=5,\n",
+    "    overwrite=True,\n",
+    ")"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,

tests/data/small_sky_nested/dataset/Norder=1/Dir=0/Npix=47.parquet

@@ -0,0 +1,14 @@
+Norder,Npix
+2,176
+2,177
+2,178
+2,179
+2,180
+2,181
+2,182
+2,183
+2,184
+2,185
+2,186
+2,187
+1,47

tests/data/small_sky_nested/dataset/Norder=2/Dir=0/Npix=176.parquet

@@ -0,0 +1,15 @@
+#HATS catalog
+obs_collection=small_sky_nested
+dataproduct_type=object
+hats_nrows=131
+hats_col_ra=ra
+hats_col_dec=dec
+hats_npix_suffix=.parquet
+hats_max_rows=1000000
+hats_order=1
+moc_sky_fraction=0.08333
+hats_builder=hats-import v0.4.6.dev1+gf00cd7a
+hats_creation_date=2025-03-05T16\:16UTC
+hats_estsize=27
+hats_release_date=2024-09-18
+hats_version=v0.1

tests/data/small_sky_nested/dataset/Norder=2/Dir=0/Npix=177.parquet

@@ -93,17 +93,34 @@ def test_load_catalog_small_sky_order1(small_sky_order1_dir):
 
 
 def test_aggregate_column_statistics(small_sky_order1_dir):
+    def assert_column_stat_as_floats(
+        result_frame, column_name, min_value=None, max_value=None, row_count=None
+    ):
+        assert column_name in result_frame.index
+        data_stats = result_frame.loc[column_name]
+        assert float(data_stats["min_value"]) >= min_value
+        assert float(data_stats["max_value"]) <= max_value
+        assert int(data_stats["null_count"]) == 0
+        assert int(data_stats["row_count"]) == row_count
+
     cat = read_hats(small_sky_order1_dir)
 
     result_frame = cat.aggregate_column_statistics()
     assert len(result_frame) == 5
+    assert_column_stat_as_floats(result_frame, "dec", min_value=-69.5, max_value=-25.5, row_count=131)
 
     result_frame = cat.aggregate_column_statistics(exclude_hats_columns=False)
+    assert_column_stat_as_floats(result_frame, "Norder", min_value=1, max_value=1, row_count=131)
     assert len(result_frame) == 9
 
     result_frame = cat.aggregate_column_statistics(include_columns=["ra", "dec"])
     assert len(result_frame) == 2
 
+    filtered_catalog = cat.filter_by_cone(315, -66.443, 0.1)
+    result_frame = filtered_catalog.aggregate_column_statistics()
+    assert len(result_frame) == 5
+    assert_column_stat_as_floats(result_frame, "dec", min_value=-69.5, max_value=-47.5, row_count=42)
+
 
 def test_aggregate_column_statistics_inmemory(catalog_info, catalog_pixels):
     catalog = Catalog(catalog_info, catalog_pixels)

tests/data/small_sky_nested/dataset/Norder=2/Dir=0/Npix=178.parquet

@@ -8,6 +8,7 @@
 
 from hats.io import file_io, paths
 from hats.io.parquet_metadata import aggregate_column_statistics, write_parquet_metadata
+from hats.pixel_math.healpix_pixel import HealpixPixel
 
 
 def test_write_parquet_metadata(tmp_path, small_sky_dir, small_sky_schema, check_parquet_schema):
@@ -136,6 +137,75 @@ def test_aggregate_column_statistics(small_sky_order1_dir):
     assert len(result_frame) == 0
 
 
+def assert_column_stat_as_floats(
+    result_frame, column_name, min_value=None, max_value=None, null_count=0, row_count=None
+):
+    assert column_name in result_frame.index
+    data_stats = result_frame.loc[column_name]
+    assert float(data_stats["min_value"]) >= min_value
+    assert float(data_stats["max_value"]) <= max_value
+    assert int(data_stats["null_count"]) == null_count
+    assert int(data_stats["row_count"]) == row_count
+
+
+def test_aggregate_column_statistics_with_pixel(small_sky_order1_dir):
+    partition_info_file = paths.get_parquet_metadata_pointer(small_sky_order1_dir)
+
+    result_frame = aggregate_column_statistics(partition_info_file)
+    assert len(result_frame) == 5
+    assert_column_stat_as_floats(result_frame, "dec", min_value=-69.5, max_value=-25.5, row_count=131)
+
+    result_frame = aggregate_column_statistics(partition_info_file, include_pixels=[HealpixPixel(1, 45)])
+    assert len(result_frame) == 5
+    assert_column_stat_as_floats(result_frame, "dec", min_value=-60.5, max_value=-25.5, row_count=29)
+
+    result_frame = aggregate_column_statistics(partition_info_file, include_pixels=[HealpixPixel(1, 47)])
+    assert len(result_frame) == 5
+    assert_column_stat_as_floats(result_frame, "dec", min_value=-36.5, max_value=-25.5, row_count=18)
+
+    result_frame = aggregate_column_statistics(
+        partition_info_file, include_pixels=[HealpixPixel(1, 45), HealpixPixel(1, 47)]
+    )
+    assert len(result_frame) == 5
+    assert_column_stat_as_floats(result_frame, "dec", min_value=-60.5, max_value=-25.5, row_count=47)
+
+    result_frame = aggregate_column_statistics(partition_info_file, include_pixels=[HealpixPixel(1, 4)])
+    assert len(result_frame) == 0
+
+
+def test_aggregate_column_statistics_with_nested(small_sky_nested_dir):
+    partition_info_file = paths.get_parquet_metadata_pointer(small_sky_nested_dir)
+
+    ## Will have 13 returned columns (5 object and 8 light curve)
+    ## Since object_dec is copied from object.dec, the min/max are the same,
+    ## but there are MANY more rows of light curve dec values.
+    result_frame = aggregate_column_statistics(partition_info_file)
+    assert len(result_frame) == 13
+    assert_column_stat_as_floats(result_frame, "dec", min_value=-69.5, max_value=-25.5, row_count=131)
+    assert_column_stat_as_floats(
+        result_frame, "lc.object_dec", min_value=-69.5, max_value=-25.5, row_count=16135
+    )
+
+    ## Only peeking at a single pixel, we should see the same dec min/max as
+    ## we see above for the flat object table.
+    result_frame = aggregate_column_statistics(partition_info_file, include_pixels=[HealpixPixel(1, 47)])
+    assert len(result_frame) == 13
+    assert_column_stat_as_floats(result_frame, "dec", min_value=-36.5, max_value=-25.5, row_count=18)
+    assert_column_stat_as_floats(
+        result_frame, "lc.source_id", min_value=70008, max_value=87148, row_count=2358
+    )
+    assert_column_stat_as_floats(result_frame, "lc.mag", min_value=15, max_value=21, row_count=2358)
+
+    ## Test that we can request light curve columns, using the shorter name
+    ## e.g. full path in the file is "lc.source_id.list.element"
+    result_frame = aggregate_column_statistics(
+        partition_info_file, include_columns=["ra", "dec", "lc.source_ra", "lc.source_dec", "lc.mag"]
+    )
+    assert len(result_frame) == 5
+    assert_column_stat_as_floats(result_frame, "dec", min_value=-69.5, max_value=-25.5, row_count=131)
+    assert_column_stat_as_floats(result_frame, "lc.mag", min_value=15, max_value=21, row_count=16135)
+
+
 def test_aggregate_column_statistics_with_nulls(tmp_path):
     file_io.make_directory(tmp_path / "dataset")
 
@@ -160,12 +230,5 @@ def test_aggregate_column_statistics_with_nulls(tmp_path):
     result_frame = aggregate_column_statistics(tmp_path / "dataset" / "_metadata", exclude_hats_columns=False)
     assert len(result_frame) == 2
 
-    data_stats = result_frame.loc["data"]
-    assert data_stats["min_value"] == -1
-    assert data_stats["max_value"] == 2
-    assert data_stats["null_count"] == 4
-
-    data_stats = result_frame.loc["Npix"]
-    assert data_stats["min_value"] == 1
-    assert data_stats["max_value"] == 6
-    assert data_stats["null_count"] == 4
+    assert_column_stat_as_floats(result_frame, "data", min_value=-1, max_value=2, null_count=4, row_count=6)
+    assert_column_stat_as_floats(result_frame, "Npix", min_value=1, max_value=6, null_count=4, row_count=6)