Remove reading partition info pixels from Norder/Npix

Author: delucchi-cmu

src/hats/catalog/partition_info.py

@@ -11,8 +11,7 @@
 
 import hats.pixel_math.healpix_shim as hp
 from hats.io import file_io, paths
-from hats.io.parquet_metadata import read_row_group_fragments, row_group_stat_single_value
-from hats.pixel_math import HealpixPixel
+from hats.pixel_math.healpix_pixel import INVALID_PIXEL, HealpixPixel
 
 
 class PartitionInfo:
@@ -102,71 +101,39 @@ def read_from_dir(cls, catalog_base_dir: str | Path | UPath | None) -> Partition
         return cls(pixel_list, catalog_base_dir)
 
     @classmethod
-    def read_from_file(cls, metadata_file: str | Path | UPath, strict: bool = False) -> PartitionInfo:
+    def read_from_file(cls, metadata_file: str | Path | UPath) -> PartitionInfo:
         """Read partition info from a `_metadata` file to create an object
 
         Args:
             metadata_file (UPath): path to the `_metadata` file
-            strict (bool): use strict parsing of _metadata file. this is slower, but
-                gives more helpful error messages in the case of invalid data.
 
         Returns:
             A `PartitionInfo` object with the data from the file
         """
-        return cls(cls._read_from_metadata_file(metadata_file, strict))
+        return cls(cls._read_from_metadata_file(metadata_file))
 
     @classmethod
-    def _read_from_metadata_file(
-        cls, metadata_file: str | Path | UPath, strict: bool = False
-    ) -> list[HealpixPixel]:
+    def _read_from_metadata_file(cls, metadata_file: str | Path | UPath) -> list[HealpixPixel]:
         """Read partition info list from a `_metadata` file.
 
         Args:
             metadata_file (UPath): path to the `_metadata` file
-            strict (bool): use strict parsing of _metadata file. this is slower, but
-                gives more helpful error messages in the case of invalid data.
 
         Returns:
             A `PartitionInfo` object with the data from the file
         """
-        if strict:
-            pixel_list = [
-                HealpixPixel(
-                    row_group_stat_single_value(row_group, cls.METADATA_ORDER_COLUMN_NAME),
-                    row_group_stat_single_value(row_group, cls.METADATA_PIXEL_COLUMN_NAME),
-                )
-                for row_group in read_row_group_fragments(metadata_file)
-            ]
-        else:
-            total_metadata = file_io.read_parquet_metadata(metadata_file)
-            num_row_groups = total_metadata.num_row_groups
-
-            first_row_group = total_metadata.row_group(0)
-            norder_column = -1
-            npix_column = -1
-
-            for i in range(0, first_row_group.num_columns):
-                column = first_row_group.column(i)
-                if column.path_in_schema == cls.METADATA_ORDER_COLUMN_NAME:
-                    norder_column = i
-                elif column.path_in_schema == cls.METADATA_PIXEL_COLUMN_NAME:
-                    npix_column = i
-
-            if norder_column == -1 or npix_column == -1:
-                raise ValueError("Metadata missing Norder or Npix column")
-
-            row_group_index = np.arange(0, num_row_groups)
-
-            pixel_list = [
-                HealpixPixel(
-                    total_metadata.row_group(index).column(norder_column).statistics.min,
-                    total_metadata.row_group(index).column(npix_column).statistics.min,
-                )
-                for index in row_group_index
-            ]
+        total_metadata = file_io.read_parquet_metadata(metadata_file)
+        if total_metadata.num_row_groups == 0 or total_metadata.row_group(0).num_columns == 0:
+            raise ValueError(f"Insufficient metadata in file {metadata_file}")
+
+        pixel_list = [
+            paths.get_healpix_from_path(total_metadata.row_group(index).column(0).file_path)
+            for index in range(0, total_metadata.num_row_groups)
+        ]
+        pixel_list = [p for p in pixel_list if p != INVALID_PIXEL]
+        if len(pixel_list) == 0:
+            raise ValueError(f"Insufficient metadata in file {metadata_file}")
         ## Remove duplicates, preserving order.
-        ## In the case of association partition join info, we may have multiple entries
-        ## for the primary order/pixels.
         return list(dict.fromkeys(pixel_list))
 
     @classmethod

src/hats/io/__init__.py

@@ -1,6 +1,6 @@
 """Utilities for reading and writing catalog files"""
 
-from .parquet_metadata import read_row_group_fragments, row_group_stat_single_value, write_parquet_metadata
+from .parquet_metadata import write_parquet_metadata
 from .paths import (
     get_common_metadata_pointer,
     get_parquet_metadata_pointer,

src/hats/io/parquet_metadata.py

@@ -7,73 +7,13 @@
 import numpy as np
 import pandas as pd
 import pyarrow.dataset as pds
-import pyarrow.parquet as pq
 from upath import UPath
 
 from hats.io import file_io, paths
 from hats.io.file_io.file_pointer import get_upath
-from hats.pixel_math.healpix_pixel import INVALID_PIXEL, HealpixPixel
 from hats.pixel_math.healpix_pixel_function import get_pixel_argsort
 
 
-def row_group_stat_single_value(row_group, stat_key: str):
-    """Convenience method to find the min and max inside a statistics dictionary,
-    and raise an error if they're unequal.
-
-    Args:
-        row_group: dataset fragment row group
-        stat_key (str): column name of interest.
-    Returns:
-        The value of the specified row group statistic
-    """
-    if stat_key not in row_group.statistics:
-        raise ValueError(f"row group doesn't have expected key {stat_key}")
-    stat_dict = row_group.statistics[stat_key]
-    min_val = stat_dict["min"]
-    max_val = stat_dict["max"]
-    if min_val != max_val:
-        raise ValueError(f"stat min != max ({min_val} != {max_val})")
-    return min_val
-
-
-def get_healpix_pixel_from_metadata(
-    metadata: pq.FileMetaData, norder_column: str = "Norder", npix_column: str = "Npix"
-) -> HealpixPixel:
-    """Get the healpix pixel according to a parquet file's metadata.
-
-    This is determined by the value of Norder and Npix in the table's data
-
-    Args:
-        metadata (pyarrow.parquet.FileMetaData): full metadata for a single file.
-
-    Returns:
-        Healpix pixel representing the Norder and Npix from the first row group.
-    """
-    if metadata.num_row_groups <= 0 or metadata.num_columns <= 0:
-        raise ValueError("metadata is for empty table")
-    order = -1
-    pixel = -1
-    first_row_group = metadata.row_group(0)
-    for i in range(0, first_row_group.num_columns):
-        column = first_row_group.column(i)
-        if column.path_in_schema == norder_column:
-            if column.statistics.min != column.statistics.max:
-                raise ValueError(
-                    f"{norder_column} stat min != max ({column.statistics.min} != {column.statistics.max})"
-                )
-            order = column.statistics.min
-        elif column.path_in_schema == npix_column:
-            if column.statistics.min != column.statistics.max:
-                raise ValueError(
-                    f"{npix_column} stat min != max ({column.statistics.min} != {column.statistics.max})"
-                )
-            pixel = column.statistics.min
-
-    if order == -1 or pixel == -1:
-        raise ValueError(f"Metadata missing Norder ({norder_column}) or Npix ({npix_column}) column")
-    return HealpixPixel(order, pixel)
-
-
 def write_parquet_metadata(
     catalog_path: str | Path | UPath,
     order_by_healpix=True,
@@ -122,8 +62,6 @@ def write_parquet_metadata(
 
         if order_by_healpix:
             healpix_pixel = paths.get_healpix_from_path(relative_path)
-            if healpix_pixel == INVALID_PIXEL:
-                healpix_pixel = get_healpix_pixel_from_metadata(single_metadata)
 
             healpix_pixels.append(healpix_pixel)
         metadata_collector.append(single_metadata)

src/hats/io/validation.py

@@ -102,18 +102,9 @@ def handle_error(msg):
             print(f"Found {len(expected_pixels)} partitions.")
 
         ## Compare the pixels in _metadata with partition_info.csv
-        # Use both strategies of reading the partition info: strict and !strict.
-        metadata_pixels = sort_pixels(
-            PartitionInfo.read_from_file(metadata_file, strict=True).get_healpix_pixels()
-        )
-        if not np.array_equal(expected_pixels, metadata_pixels):
-            handle_error("Partition pixels differ between catalog and _metadata file (strict)")
-
-        metadata_pixels = sort_pixels(
-            PartitionInfo.read_from_file(metadata_file, strict=False).get_healpix_pixels()
-        )
+        metadata_pixels = sort_pixels(PartitionInfo.read_from_file(metadata_file).get_healpix_pixels())
         if not np.array_equal(expected_pixels, metadata_pixels):
-            handle_error("Partition pixels differ between catalog and _metadata file (non-strict)")
+            handle_error("Partition pixels differ between catalog and _metadata file")
 
         partition_info_file = get_partition_info_pointer(pointer)
         partition_info = PartitionInfo.read_from_csv(partition_info_file)

tests/hats/catalog/test_partition_info.py

@@ -29,29 +29,20 @@ def test_load_partition_info_from_metadata(small_sky_dir, small_sky_source_dir,
     partitions = PartitionInfo.read_from_file(metadata_file)
     assert partitions.get_healpix_pixels() == small_sky_source_pixels
 
-    partitions = PartitionInfo.read_from_file(metadata_file, strict=True)
-    assert partitions.get_healpix_pixels() == small_sky_source_pixels
-
 
 def test_load_partition_info_from_metadata_fail(tmp_path):
     empty_dataframe = pd.DataFrame()
     metadata_filename = tmp_path / "empty_metadata.parquet"
     empty_dataframe.to_parquet(metadata_filename)
-    with pytest.raises(ValueError, match="missing Norder"):
+    with pytest.raises(ValueError, match="Insufficient metadata"):
         PartitionInfo.read_from_file(metadata_filename)
 
-    with pytest.raises(ValueError, match="at least one column"):
-        PartitionInfo.read_from_file(metadata_filename, strict=True)
-
     non_healpix_dataframe = pd.DataFrame({"data": [0], "Npix": [45]})
     metadata_filename = tmp_path / "non_healpix_metadata.parquet"
     non_healpix_dataframe.to_parquet(metadata_filename)
-    with pytest.raises(ValueError, match="missing Norder"):
+    with pytest.raises(ValueError, match="Insufficient metadata"):
         PartitionInfo.read_from_file(metadata_filename)
 
-    with pytest.raises(ValueError, match="empty file path"):
-        PartitionInfo.read_from_file(metadata_filename, strict=True)
-
 
 def test_load_partition_info_from_dir_fail(tmp_path):
     empty_dataframe = pd.DataFrame()
@@ -65,7 +56,7 @@ def test_load_partition_info_from_dir_fail(tmp_path):
     metadata_filename = tmp_path / "dataset" / "_metadata"
     empty_dataframe.to_parquet(metadata_filename)
     with pytest.warns(UserWarning, match="slow"):
-        with pytest.raises(ValueError, match="missing Norder"):
+        with pytest.raises(ValueError, match="Insufficient metadata"):
             PartitionInfo.read_from_dir(tmp_path)
 
 

tests/hats/io/test_parquet_metadata.py

@@ -5,15 +5,9 @@
 import pandas as pd
 import pyarrow as pa
 import pyarrow.parquet as pq
-import pytest
 
 from hats.io import file_io, paths
-from hats.io.parquet_metadata import (
-    aggregate_column_statistics,
-    read_row_group_fragments,
-    row_group_stat_single_value,
-    write_parquet_metadata,
-)
+from hats.io.parquet_metadata import aggregate_column_statistics, write_parquet_metadata
 
 
 def test_write_parquet_metadata(tmp_path, small_sky_dir, small_sky_schema, check_parquet_schema):
@@ -126,24 +120,6 @@ def test_write_index_parquet_metadata(tmp_path, check_parquet_schema):
     )
 
 
-def test_row_group_fragments(small_sky_order1_dir):
-    partition_info_file = paths.get_parquet_metadata_pointer(small_sky_order1_dir)
-
-    num_row_groups = 0
-    for _ in read_row_group_fragments(partition_info_file):
-        num_row_groups += 1
-
-    assert num_row_groups == 4
-
-
-def test_row_group_fragments_with_dir(small_sky_order1_dir):
-    num_row_groups = 0
-    for _ in read_row_group_fragments(small_sky_order1_dir):
-        num_row_groups += 1
-
-    assert num_row_groups == 4
-
-
 def test_aggregate_column_statistics(small_sky_order1_dir):
     partition_info_file = paths.get_parquet_metadata_pointer(small_sky_order1_dir)
 
@@ -193,83 +169,3 @@ def test_aggregate_column_statistics_with_nulls(tmp_path):
     assert data_stats["min_value"] == 1
     assert data_stats["max_value"] == 6
     assert data_stats["null_count"] == 4
-
-
-def test_row_group_stats(small_sky_dir):
-    partition_info_file = paths.get_parquet_metadata_pointer(small_sky_dir)
-    first_row_group = next(read_row_group_fragments(partition_info_file))
-
-    assert row_group_stat_single_value(first_row_group, "Norder") == 0
-    assert row_group_stat_single_value(first_row_group, "Npix") == 11
-
-    with pytest.raises(ValueError, match="doesn't have expected key"):
-        row_group_stat_single_value(first_row_group, "NOT HERE")
-
-    with pytest.raises(ValueError, match="stat min != max"):
-        row_group_stat_single_value(first_row_group, "ra")
-
-
-# def test_get_healpix_pixel_from_metadata(small_sky_dir):
-#     partition_info_file = paths.get_parquet_metadata_pointer(small_sky_dir)
-#     single_metadata = file_io.read_parquet_metadata(partition_info_file)
-#     pixel = get_healpix_pixel_from_metadata(single_metadata)
-#     assert pixel == HealpixPixel(0, 11)
-
-
-# def test_get_healpix_pixel_from_metadata_min_max(tmp_path):
-#     good_healpix_dataframe = pd.DataFrame({"data": [0, 1], "Norder": [1, 1], "Npix": [44, 44]})
-#     metadata_filename = tmp_path / "non_healpix_metadata.parquet"
-#     good_healpix_dataframe.to_parquet(metadata_filename)
-#     single_metadata = file_io.read_parquet_metadata(metadata_filename)
-#     pixel = get_healpix_pixel_from_metadata(single_metadata)
-#     assert pixel == HealpixPixel(1, 44)
-
-#     non_healpix_dataframe = pd.DataFrame({"data": [0, 1], "Npix": [45, 44]})
-#     non_healpix_dataframe.to_parquet(metadata_filename)
-#     single_metadata = file_io.read_parquet_metadata(metadata_filename)
-#     with pytest.raises(ValueError, match="Npix stat min != max"):
-#         get_healpix_pixel_from_metadata(single_metadata)
-
-#     non_healpix_dataframe = pd.DataFrame({"data": [0, 1], "Norder": [5, 6]})
-#     non_healpix_dataframe.to_parquet(metadata_filename)
-#     single_metadata = file_io.read_parquet_metadata(metadata_filename)
-#     with pytest.raises(ValueError, match="Norder stat min != max"):
-#         get_healpix_pixel_from_metadata(single_metadata)
-
-
-# def test_get_healpix_pixel_from_metadata_fail(tmp_path):
-#     empty_dataframe = pd.DataFrame()
-#     metadata_filename = tmp_path / "empty_metadata.parquet"
-#     empty_dataframe.to_parquet(metadata_filename)
-#     single_metadata = file_io.read_parquet_metadata(metadata_filename)
-#     with pytest.raises(ValueError, match="empty table"):
-#         get_healpix_pixel_from_metadata(single_metadata)
-
-#     non_healpix_dataframe = pd.DataFrame({"data": [0], "Npix": [45]})
-#     metadata_filename = tmp_path / "non_healpix_metadata.parquet"
-#     non_healpix_dataframe.to_parquet(metadata_filename)
-#     single_metadata = file_io.read_parquet_metadata(metadata_filename)
-#     with pytest.raises(ValueError, match="missing Norder"):
-#         get_healpix_pixel_from_metadata(single_metadata)
-
-
-# def test_get_healpix_pixel_from_metadata_columns(tmp_path):
-#     """Test fetching the healpix pixel from columns with non-default names."""
-#     non_healpix_dataframe = pd.DataFrame({"data": [1], "Npix": [45], "join_Norder": [2], "join_Npix": [3]})
-#     metadata_filename = tmp_path / "non_healpix_metadata.parquet"
-#     non_healpix_dataframe.to_parquet(metadata_filename)
-#     single_metadata = file_io.read_parquet_metadata(metadata_filename)
-#     with pytest.raises(ValueError, match="missing Norder"):
-#         get_healpix_pixel_from_metadata(single_metadata)
-
-#     pixel = get_healpix_pixel_from_metadata(single_metadata, norder_column="data")
-#     assert pixel == HealpixPixel(1, 45)
-
-#     pixel = get_healpix_pixel_from_metadata(
-#         single_metadata, norder_column="join_Norder", npix_column="join_Npix"
-#     )
-#     assert pixel == HealpixPixel(2, 3)
-
-#     ## People can do silly things!
-#     pixel = get_healpix_pixel_from_metadata(single_metadata, norder_column="data", npix_column="join_Npix")
-#     assert pixel == HealpixPixel(1, 3)