diff --git a/src/hats_import/margin_cache/margin_cache_map_reduce.py b/src/hats_import/margin_cache/margin_cache_map_reduce.py index f4cb5c21..fc12998c 100644 --- a/src/hats_import/margin_cache/margin_cache_map_reduce.py +++ b/src/hats_import/margin_cache/margin_cache_map_reduce.py @@ -25,8 +25,10 @@ def map_pixel_shards( """Creates margin cache shards from a source partition file.""" try: schema = file_io.read_parquet_metadata(original_catalog_metadata).schema.to_arrow_schema() + + # first pass to compute margin filter data = pq.read_table( - partition_file.path, filesystem=partition_file.fs, schema=schema + partition_file.path, filesystem=partition_file.fs, schema=schema, columns=[healpix_column] ).combine_chunks() # Constrain the possible margin pairs, first by only those `margin_order` pixels @@ -49,6 +51,16 @@ def map_pixel_shards( {"margin_pixel": margin_pixel_list, "filter_value": np.arange(0, len(margin_pixel_list))} ).merge(margin_pairs, on="margin_pixel") + # If no row falls into any margin pixel, skip the full-schema read entirely. + if len(margin_pixel_filter) == 0: + MarginCachePlan.mapping_key_done(output_path, mapping_key, 0) + return + + # Second pass: read the full table once, now that we know rows are needed. + data = pq.read_table( + partition_file.path, filesystem=partition_file.fs, schema=schema + ).combine_chunks() + # For every possible output pixel, find the full margin_order pixel filter list, # perform the filter, and pass along to helper method to compute fine filter # and write out shard file. diff --git a/tests/hats_import/margin_cache/test_margin_cache_map_reduce.py b/tests/hats_import/margin_cache/test_margin_cache_map_reduce.py index b0d79494..6f9a72b8 100644 --- a/tests/hats_import/margin_cache/test_margin_cache_map_reduce.py +++ b/tests/hats_import/margin_cache/test_margin_cache_map_reduce.py @@ -1,8 +1,10 @@ import os +from unittest.mock import patch import numpy as np import pandas as pd import pyarrow as pa +import pyarrow.parquet as pq import pytest from hats import pixel_math from hats.io import paths @@ -197,3 +199,34 @@ def test_reduce_margin_shards(tmp_path): validate_result_dataframe(result_path, 720) assert not os.path.exists(shard_dir) + + +def test_map_pixel_shards_no_match(tmp_path, test_data_dir, small_sky_source_catalog): + """When no row falls into any margin pixel, the full-schema read is skipped. + + Regression guard for issue #685: the early-exit branch must short-circuit + before the second pq.read_table call, not just produce the same output. + """ + intermediate_dir = tmp_path / "intermediate" + os.makedirs(intermediate_dir / "mapping") + + with patch.object(pq, "read_table", wraps=pq.read_table) as read_table_spy: + margin_cache_map_reduce.map_pixel_shards( + paths.pixel_catalog_file(small_sky_source_catalog, HealpixPixel(1, 47)), + source_pixel=HealpixPixel(1, 47), + mapping_key="1_47", + original_catalog_metadata=small_sky_source_catalog / "dataset" / "_common_metadata", + margin_pair_file=test_data_dir / "margin_pairs" / "small_sky_source_pairs.csv", + output_path=intermediate_dir, + margin_order=1, + healpix_column="_healpix_29", + healpix_order=29, + ) + + # Exactly one read happened, and it asked for only the healpix column. + assert read_table_spy.call_count == 1 + assert read_table_spy.call_args.kwargs.get("columns") == ["_healpix_29"] + + # And the observable behavior: marker written, no shards. + assert os.path.exists(intermediate_dir / "mapping" / "1_47_done") + assert not any(p.name.startswith("order_") for p in intermediate_dir.iterdir())