From d416dd6c43348245cd02b5ddc6917cd132aa947d Mon Sep 17 00:00:00 2001 From: Luigi Lucas de Carvalho Silva Date: Sat, 4 Jul 2026 10:37:06 -0300 Subject: [PATCH] perf: distribute crossmatch pair processing and extend SLURM walltime - stage projected crossmatch pairs directly from workers - remove global adjacency dictionaries from Dask task graphs - aggregate and merge compared_to relationships with distributed operations - reuse neighbor tables across both sides of crossmatches - compute saturation diagnostics through scalar Dask reductions - preserve optional geometric diagnostics behind explicit configuration - apply the distributed path to auto and catalog crossmatches - extend driver and worker walltimes to 12 hours - add tests for staged projections, distributed adjacency, and saturation --- config.py | 2 +- config.template.yaml | 2 +- packages/crossmatch_auto.py | 140 +++++-- packages/crossmatch_cross.py | 387 ++++++++++++++---- packages/crossmatch_diagnostics.py | 18 + sbatch.template | 2 +- test_configs/config.test.concatenate.yaml | 2 +- ...test.homogenized_columns_out_of_range.yaml | 2 +- test_configs/config.test.mark.yaml | 2 +- test_configs/config.test.remove.yaml | 2 +- .../config.test.survey_col_missing.yaml | 2 +- .../config.test.survey_not_supported.yaml | 2 +- ...test.survey_supported_but_wrong_flags.yaml | 2 +- tests/test_crossmatch_adjacency.py | 63 +++ tests/test_crossmatch_diagnostics.py | 24 ++ tests/test_crossmatch_saturation.py | 25 +- 16 files changed, 557 insertions(+), 120 deletions(-) diff --git a/config.py b/config.py index d99603a..21a9013 100644 --- a/config.py +++ b/config.py @@ -16,7 +16,7 @@ class Instance(BaseModel): memory: str = "40GB" queue: str = "cpu_pipelines" account: str = "hpc-pipelines" - job_extra_directives: list[str] = ["--propagate", "--time=04:00:00"] + job_extra_directives: list[str] = ["--propagate", "--time=12:00:00"] class Scale(BaseModel): minimum_jobs: int = 11 diff --git a/config.template.yaml b/config.template.yaml index 80a8755..3dc14eb 100644 --- a/config.template.yaml +++ b/config.template.yaml @@ -21,7 +21,7 @@ executor: # account: "hpc-public" # job_extra_directives: # - "--propagate" -# - "--time=2:00:00" +# - "--time=12:00:00" # scale: # minimum_jobs: 11 # maximum_jobs: 11 diff --git a/packages/crossmatch_auto.py b/packages/crossmatch_auto.py index c69b265..1e5092b 100644 --- a/packages/crossmatch_auto.py +++ b/packages/crossmatch_auto.py @@ -15,6 +15,7 @@ # ----------------------- import logging import os +import shutil import time from typing import TYPE_CHECKING, Dict, Iterable, List, Set @@ -26,17 +27,26 @@ # ----------------------- import numpy as np import pandas as pd +import dask +import dask.dataframe as dd +from dask.distributed import get_client # ----------------------- # Project # ----------------------- -from specz import DTYPE_STR # Arrow-backed string dtype +from specz import DTYPE_STR, _build_collection_with_retry +from crossmatch_cross import ( + _attach_distributed_neighbors, + _log_neighbor_saturation_distributed, + _normalize_ddf_expected_types, + _safe_to_parquet, +) from crossmatch_diagnostics import ( - compute_projected_pairs, log_component_size_diagnostics, log_neighbor_count_diagnostics, log_pair_separation_diagnostics, project_catalog_for_pair_crossmatch, + stage_projected_pairs, ) from utils import get_phase_logger @@ -261,8 +271,9 @@ def _self_xmatch_pairs( warn_fraction: float, fail_fraction: float | None, geometry_diagnostics_enabled: bool = False, -) -> Dict[str, Set[str]]: - """Run a self-crossmatch and return an adjacency (CRD_ID -> neighbor ids). + staging_path: str = "self_crossmatch_pairs", +) -> tuple[dd.DataFrame, int, int]: + """Run a self-crossmatch and return narrow, unique ID pairs. Args: catalog: LSDB catalog to crossmatch with itself. @@ -271,7 +282,7 @@ def _self_xmatch_pairs( logger: Logger. Returns: - Dict[str, Set[str]]: Mapping of CRD_ID to neighbor ids. + tuple: Staged pair dataframe, undirected link count, and node count. """ logger.info( 'Running self-crossmatch: radius=%.3f" n_neighbors=%d', @@ -294,11 +305,38 @@ def _self_xmatch_pairs( pair_cols.append("_dist_arcsec") if total_by_source is not None and "sourceleft" in xmatched.columns: pair_cols.append("sourceleft") - pairs_df = compute_projected_pairs(xmatched._ddf, pair_cols) - if len(pairs_df) == 0: + raw_path = os.path.join(staging_path, "raw") + pairs_path = os.path.join(staging_path, "pairs") + raw_pairs = stage_projected_pairs(xmatched._ddf, pair_cols, raw_path) + pairs = raw_pairs.astype( + {"CRD_IDleft": DTYPE_STR, "CRD_IDright": DTYPE_STR} + ) + pairs = pairs[pairs["CRD_IDleft"] != pairs["CRD_IDright"]] + pairs = pairs.drop_duplicates(subset=["CRD_IDleft", "CRD_IDright"]) + _safe_to_parquet(pairs, pairs_path, write_index=False) + pairs = dd.read_parquet(pairs_path, engine="pyarrow") + pair_count, node_count = dask.compute( + pairs.map_partitions(len).sum(), + dd.concat([pairs["CRD_IDleft"], pairs["CRD_IDright"]]).nunique(), + ) + pair_count = int(pair_count) + node_count = int(node_count) + if not pair_count: logger.info("Self-crossmatch: no pairs found; `compared_to` remains unchanged.") - return {} - if geometry_diagnostics_enabled: + if total_by_source is not None and pair_count: + _log_neighbor_saturation_distributed( + pairs, + id_col="CRD_IDleft", + source_col="sourceleft" if "sourceleft" in pairs.columns else None, + limit=n_neighbors, + logger=logger, + context="Self-crossmatch", + total_by_source=total_by_source, + warn_fraction=warn_fraction, + fail_fraction=fail_fraction, + ) + if geometry_diagnostics_enabled and pair_count: + pairs_df = pairs.compute() log_pair_separation_diagnostics( pairs_df, left_col="CRD_IDleft", @@ -307,32 +345,21 @@ def _self_xmatch_pairs( logger=logger, context="Self-crossmatch", ) - if total_by_source is not None: - _log_neighbor_saturation( - pairs_df, - limit=n_neighbors, - logger=logger, - total_by_source=total_by_source, - warn_fraction=warn_fraction, - fail_fraction=fail_fraction, + adj = _adjacency_from_pairs( + pairs_df["CRD_IDleft"], pairs_df["CRD_IDright"] ) - pairs_df = pairs_df.astype({"CRD_IDleft": "string", "CRD_IDright": "string"}) - pairs_df = pairs_df[ - pairs_df["CRD_IDleft"] != pairs_df["CRD_IDright"] - ].drop_duplicates() - - adj = _adjacency_from_pairs(pairs_df["CRD_IDleft"], pairs_df["CRD_IDright"]) - if geometry_diagnostics_enabled: log_component_size_diagnostics( adj, logger=logger, context="Self-crossmatch", ) - total_links = sum(len(v) for v in adj.values()) + total_links = 2 * pair_count logger.info( - "Self-crossmatch: %d unique pairs across %d nodes", total_links, len(adj) + "Self-crossmatch: %d links across %d nodes without global adjacency", + total_links, + node_count, ) - return adj + return pairs, total_links, node_count def _update_compared_to( @@ -450,8 +477,12 @@ def crossmatch_auto( collection_path_auto, ) - # 1) Self-crossmatch -> adjacency (CRD_ID -> neighbor ids) - pairs_adj = _self_xmatch_pairs( + staging_root = os.path.join(parent_dir, f".{artifact_auto}_distributed_pairs") + shutil.rmtree(staging_root, ignore_errors=True) + os.makedirs(staging_root, exist_ok=True) + + # 1) Self-crossmatch -> distributed narrow pair table + pairs_ddf, total_links, n_nodes = _self_xmatch_pairs( catalog, radius, k, @@ -460,22 +491,51 @@ def crossmatch_auto( warn_fraction, fail_fraction, geometry_diagnostics_enabled, + staging_root, ) - # 2) Update `compared_to` - updated = _update_compared_to(catalog, pairs_adj) - total_links = sum(len(v) for v in pairs_adj.values()) - n_nodes = len(pairs_adj) - - # 3) Persist as a NEW collection in the parent directory + # 2) Persist as a NEW collection in the parent directory. For non-empty + # pairs, aggregate and join through Dask instead of embedding one global + # adjacency mapping into every partition task. logger.info( "Writing collection: base_dir=%s catalog_name=%s", parent_dir, artifact_auto ) - updated.write_catalog( - collection_path_auto, - as_collection=True, - overwrite=True, - ) + if total_links == 0: + catalog.write_catalog( + collection_path_auto, + as_collection=True, + overwrite=True, + ) + shutil.rmtree(staging_root, ignore_errors=True) + else: + neighbors_path = os.path.join(staging_root, "neighbors") + merged_path = os.path.join(staging_root, "merged") + updated_ddf = _attach_distributed_neighbors( + catalog._ddf, + pairs_ddf, + id_column="CRD_IDleft", + neighbor_column="CRD_IDright", + staging_path=neighbors_path, + symmetric=True, + ) + updated_ddf = _normalize_ddf_expected_types( + updated_ddf, translation_config + ) + _safe_to_parquet(updated_ddf, merged_path, write_index=False) + _build_collection_with_retry( + parquet_path=merged_path, + logs_dir=logs_dir, + logger=logger, + client=get_client(), + try_margin=True, + margin_threshold=float( + (translation_config or {}).get("margin_threshold_arcsec", 5.0) + ), + output_path=collection_path_auto, + output_artifact_name=artifact_auto, + catalog_artifact_name=artifact_auto, + ) + shutil.rmtree(staging_root, ignore_errors=True) logger.info("Write complete: path=%s", collection_path_auto) # END (per-catalog) diff --git a/packages/crossmatch_cross.py b/packages/crossmatch_cross.py index e7cb74d..f2de300 100644 --- a/packages/crossmatch_cross.py +++ b/packages/crossmatch_cross.py @@ -27,12 +27,14 @@ USE_LSDB_CONCAT: bool = True # Toggle behavior as described above. BACKEND_LSDB_LABEL = "LSDB+write_catalog" BACKEND_LEGACY_LABEL = "Dask+Parquet+import" +BACKEND_DISTRIBUTED_PAIRS_LABEL = "distributed-pairs+Parquet+import" # ----------------------- # Standard library # ----------------------- import logging import os +import shutil import time from typing import Dict, Iterable, List, Set @@ -41,6 +43,7 @@ # ----------------------- import numpy as np import pandas as pd +import dask import dask.dataframe as dd # ----------------------- @@ -48,11 +51,11 @@ # ----------------------- from utils import get_phase_logger from crossmatch_diagnostics import ( - compute_projected_pairs, log_component_size_diagnostics, log_neighbor_count_diagnostics, log_pair_separation_diagnostics, project_catalog_for_pair_crossmatch, + stage_projected_pairs, ) from specz import ( _build_collection_with_retry, @@ -91,7 +94,7 @@ def _get_backend_label() -> str: Returns: str: Active backend label. """ - return BACKEND_LSDB_LABEL if USE_LSDB_CONCAT else BACKEND_LEGACY_LABEL + return BACKEND_DISTRIBUTED_PAIRS_LABEL # ----------------------- @@ -211,6 +214,82 @@ def _catalog_source_totals(catalog) -> dict: return {"": int(catalog._ddf.map_partitions(len).sum().compute())} +def _log_neighbor_saturation_distributed( + pairs, + *, + id_col: str, + source_col: str | None, + limit: int, + logger: logging.LoggerAdapter, + context: str, + total_by_source: dict, + warn_fraction: float, + fail_fraction: float | None, +) -> None: + """Log saturation summaries without collecting pair rows on the driver.""" + group_cols = [id_col] + sources: list[object] = [None] + if source_col and source_col in pairs.columns: + group_cols.insert(0, source_col) + sources = pairs[source_col].dropna().drop_duplicates().compute().tolist() + + counts = pairs.groupby(group_cols).size().rename("count").to_frame().reset_index() + for source in sources: + source_counts = counts["count"] + label = context + total_key = "" + if source is not None and source_col is not None: + source_counts = counts.loc[counts[source_col] == source, "count"] + label = f"{context} source={source}" + total_key = str(source) + + count, maximum, p50, p90, p99, ge2, ge5, saturated = dask.compute( + source_counts.count(), + source_counts.max(), + source_counts.quantile(0.50), + source_counts.quantile(0.90), + source_counts.quantile(0.99), + source_counts.ge(2).sum(), + source_counts.ge(5).sum(), + source_counts.ge(limit).sum(), + ) + count = int(count) + saturated = int(saturated) + total = int(total_by_source.get(total_key, count)) + fraction = saturated / total if total else 0.0 + logger.info( + "%s returned-match diagnostics: objects_with_matches=%d " + "p50=%.1f p90=%.1f p99=%.1f max=%d fraction_ge_2=%.6f " + "fraction_ge_5=%.6f fraction_at_limit=%.6f limit=%d", + label, + count, + float(p50) if count else 0.0, + float(p90) if count else 0.0, + float(p99) if count else 0.0, + int(maximum) if count else 0, + int(ge2) / count if count else 0.0, + int(ge5) / count if count else 0.0, + fraction, + limit, + ) + log_method = logger.warning if fraction >= warn_fraction else logger.info + log_method( + "%s neighbor saturation: at_limit=%d total_objects=%d " + "fraction=%.6f limit=%d", + label, + saturated, + total, + fraction, + limit, + ) + if fail_fraction is not None and fraction >= fail_fraction: + raise RuntimeError( + f"{label}: neighbor saturation is {fraction:.6f}, reaching " + f"failure threshold {fail_fraction:.6f}; increase " + "crossmatch_n_neighbors" + ) + + def _merge_compared_to_partition( part: pd.DataFrame, pairs_adj: Dict[str, Iterable[str]], @@ -276,6 +355,119 @@ def _parse_existing(val) -> Set[str]: return p +def _merge_compared_to_column_partition( + part: pd.DataFrame, + new_column: str = "_new_compared_to", +) -> pd.DataFrame: + """Union an aggregated neighbor column into ``compared_to``.""" + p = part.copy() + if "compared_to" not in p.columns: + p["compared_to"] = pd.Series( + pd.array([pd.NA] * len(p), dtype=DTYPE_STR), index=p.index + ) + + def _tokens(value) -> set[str]: + if isinstance(value, (list, set, tuple)): + values = value + else: + if pd.isna(value): + return set() + values = str(value).split(",") + return { + token + for raw in values + if (token := str(raw).strip()) and token != "" + } + + merged_values: list[object] = [] + for crd_id, old_value, new_value in zip( + p["CRD_ID"].astype(str), p["compared_to"], p[new_column] + ): + neighbors = _tokens(old_value) | _tokens(new_value) + neighbors.discard(crd_id) + merged_values.append(", ".join(sorted(neighbors)) if neighbors else pd.NA) + + p["compared_to"] = pd.Series( + pd.array(merged_values, dtype=DTYPE_STR), index=p.index + ) + return p.drop(columns=[new_column]) + + +def _join_unique_neighbors(values: pd.Series) -> str: + """Return a deterministic comma-separated set for one grouped CRD_ID.""" + return ", ".join(sorted(set(values.dropna().astype(str)))) + + +def _aggregate_distributed_neighbors( + pairs_ddf, + *, + id_column: str, + neighbor_column: str, + staging_path: str, + symmetric: bool = False, +): + """Aggregate and stage a reusable distributed neighbor table.""" + directional = pairs_ddf[[id_column, neighbor_column]].rename( + columns={id_column: "CRD_ID", neighbor_column: "_neighbor"} + ) + if symmetric: + reverse = pairs_ddf[[id_column, neighbor_column]].rename( + columns={neighbor_column: "CRD_ID", id_column: "_neighbor"} + ) + directional = dd.concat([directional, reverse]) + directional = directional.astype( + {"CRD_ID": DTYPE_STR, "_neighbor": DTYPE_STR} + ) + directional = directional[ + directional["CRD_ID"] != directional["_neighbor"] + ] + neighbor_meta = pd.Series( + name="_new_compared_to", + dtype="string", + index=pd.Index([], name="CRD_ID", dtype="string"), + ) + neighbors = ( + directional.groupby("CRD_ID")["_neighbor"] + .apply(_join_unique_neighbors, meta=neighbor_meta) + .to_frame() + .reset_index() + ) + _safe_to_parquet(neighbors, staging_path, write_index=False) + return dd.read_parquet(staging_path, engine="pyarrow") + + +def _merge_distributed_neighbors(ddf, neighbors): + """Merge one dataframe with a pre-aggregated neighbor table.""" + normalized_ddf = ddf.assign(CRD_ID=ddf["CRD_ID"].astype(DTYPE_STR)) + neighbors = neighbors.assign(CRD_ID=neighbors["CRD_ID"].astype(DTYPE_STR)) + merged = normalized_ddf.merge(neighbors, how="left", on="CRD_ID") + meta = _merge_compared_to_column_partition(merged._meta) + return merged.map_partitions( + _merge_compared_to_column_partition, + meta=meta, + ) + + +def _attach_distributed_neighbors( + ddf, + pairs_ddf, + *, + id_column: str, + neighbor_column: str, + staging_path: str, + symmetric: bool = False, +): + """Aggregate pairs once and attach their neighbors to one dataframe.""" + neighbors = _aggregate_distributed_neighbors( + pairs_ddf, + id_column=id_column, + neighbor_column=neighbor_column, + staging_path=staging_path, + symmetric=symmetric, + ) + return _merge_distributed_neighbors(ddf, neighbors) + + def _ensure_compared_to_meta(meta_df: pd.DataFrame) -> pd.DataFrame: """Return a meta dataframe with a typed `compared_to` column. @@ -761,6 +953,60 @@ def _concat_parquet_import( return collection_path +def _distributed_pairs_update_and_import( + left_cat, + right_cat, + pairs_ddf, + temp_dir: str, + logs_dir: str, + step, + client, + translation_config: dict | None, + logger: logging.LoggerAdapter, +) -> str: + """Update both catalogs without embedding global adjacency in task graphs.""" + staging_root = os.path.join(temp_dir, f"distributed_pairs_step{step}") + merged_path = os.path.join(staging_root, "merged") + collection_path_target = os.path.join(temp_dir, f"merged_step{step}_hats") + logger.info( + "Distributed compared_to update: step=%s staging=%s", + step, + staging_root, + ) + + neighbors = _aggregate_distributed_neighbors( + pairs_ddf, + id_column="CRD_IDleft", + neighbor_column="CRD_IDright", + staging_path=os.path.join(staging_root, "neighbors"), + symmetric=True, + ) + left_ddf = _merge_distributed_neighbors(left_cat._ddf, neighbors) + right_ddf = _merge_distributed_neighbors(right_cat._ddf, neighbors) + merged = dd.concat([left_ddf, right_ddf]) + merged = _normalize_ddf_expected_types(merged, translation_config) + _safe_to_parquet(merged, merged_path, write_index=False) + logger.info("Distributed updated Parquet written: step=%s path=%s", step, merged_path) + + schema_hints_local = _get_expr_schema_hints(translation_config) + collection_path = _build_collection_with_retry( + parquet_path=merged_path, + logs_dir=logs_dir, + logger=logger, + client=client, + try_margin=True, + schema_hints=schema_hints_local or None, + margin_threshold=float( + (translation_config or {}).get("margin_threshold_arcsec", 5.0) + ), + output_path=collection_path_target, + output_artifact_name=f"merged_step{step}_hats", + catalog_artifact_name=f"merged_step{step}", + ) + shutil.rmtree(staging_root, ignore_errors=True) + return collection_path + + # ----------------------- # Main logic # ----------------------- @@ -851,79 +1097,79 @@ def crossmatch_tiebreak( ) logger.info("Crossmatch done (%.2fs)", time.time() - t0) - # 2) Build adjacency from CRD_ID pairs + # 2) Project and stage pair columns inside workers. No pair rows are + # gathered on the driver in the production path. t0 = time.time() pair_cols = ["CRD_IDleft", "CRD_IDright"] if geometry_diagnostics_enabled and "_dist_arcsec" in xmatched._ddf.columns: pair_cols.append("_dist_arcsec") if saturation_enabled and "sourceleft" in xmatched._ddf.columns: pair_cols.append("sourceleft") - pairs_df = compute_projected_pairs(xmatched._ddf, pair_cols) - if len(pairs_df) == 0: - pairs_adj: Dict[str, Set[str]] = {} + staging_root = os.path.join(temp_dir, f"distributed_pairs_step{step}") + raw_pairs_path = os.path.join(staging_root, "raw_pairs") + pairs_path = os.path.join(staging_root, "pairs") + shutil.rmtree(staging_root, ignore_errors=True) + os.makedirs(staging_root, exist_ok=True) + raw_pairs = stage_projected_pairs(xmatched._ddf, pair_cols, raw_pairs_path) + pairs = raw_pairs.astype( + {"CRD_IDleft": DTYPE_STR, "CRD_IDright": DTYPE_STR} + ) + pairs = pairs[pairs["CRD_IDleft"] != pairs["CRD_IDright"]] + pairs = pairs.drop_duplicates(subset=["CRD_IDleft", "CRD_IDright"]) + _safe_to_parquet(pairs, pairs_path, write_index=False) + pairs = dd.read_parquet(pairs_path, engine="pyarrow") + + pair_count, node_count = dask.compute( + pairs.map_partitions(len).sum(), + dd.concat([pairs["CRD_IDleft"], pairs["CRD_IDright"]]).nunique(), + ) + pair_count = int(pair_count) + node_count = int(node_count) + total_links = 2 * pair_count + if not pair_count: logger.info("No pairs found; `compared_to` remains unchanged.") - else: - if geometry_diagnostics_enabled: - log_pair_separation_diagnostics( - pairs_df, - left_col="CRD_IDleft", - right_col="CRD_IDright", - radius_arcsec=radius, - logger=logger, - context=f"crossmatch step={step}", - ) - if saturation_enabled: - _log_neighbor_saturation( - pairs_df, - id_col="CRD_IDleft", - source_col="sourceleft" if "sourceleft" in pairs_df else None, - limit=k, - logger=logger, - context=f"crossmatch step={step}", - total_by_source=total_by_source, - warn_fraction=warn_fraction, - fail_fraction=fail_fraction, - ) - pairs_df = pairs_df.astype({"CRD_IDleft": "string", "CRD_IDright": "string"}) - pairs_df = pairs_df[ - pairs_df["CRD_IDleft"] != pairs_df["CRD_IDright"] - ].drop_duplicates() - pairs_adj = _adjacency_from_pairs( - pairs_df["CRD_IDleft"], pairs_df["CRD_IDright"] + if saturation_enabled and pair_count: + _log_neighbor_saturation_distributed( + pairs, + id_col="CRD_IDleft", + source_col="sourceleft" if "sourceleft" in pairs.columns else None, + limit=k, + logger=logger, + context=f"crossmatch step={step}", + total_by_source=total_by_source, + warn_fraction=warn_fraction, + fail_fraction=fail_fraction, + ) + if geometry_diagnostics_enabled and pair_count: + diagnostic_pairs = pairs.compute() + log_pair_separation_diagnostics( + diagnostic_pairs, + left_col="CRD_IDleft", + right_col="CRD_IDright", + radius_arcsec=radius, + logger=logger, + context=f"crossmatch step={step}", + ) + diagnostic_adj = _adjacency_from_pairs( + diagnostic_pairs["CRD_IDleft"], diagnostic_pairs["CRD_IDright"] + ) + log_component_size_diagnostics( + diagnostic_adj, + logger=logger, + context=f"crossmatch step={step}", ) - if geometry_diagnostics_enabled: - log_component_size_diagnostics( - pairs_adj, - logger=logger, - context=f"crossmatch step={step}", - ) - total_links = sum(len(v) for v in pairs_adj.values()) logger.info( - "Adjacency built: links=%d nodes=%d (%.2fs)", + "Pair summary built without global adjacency: links=%d nodes=%d (%.2fs)", total_links, - len(pairs_adj), + node_count, time.time() - t0, ) - # 3) Update `compared_to` on both catalogs, partition-wise - t0 = time.time() - - left_meta = _ensure_compared_to_meta(left_cat._ddf._meta) - right_meta = _ensure_compared_to_meta(right_cat._ddf._meta) - - left_updated = left_cat.map_partitions( - _merge_compared_to_partition, pairs_adj, meta=left_meta - ) - right_updated = right_cat.map_partitions( - _merge_compared_to_partition, pairs_adj, meta=right_meta - ) - logger.info("Compared_to updated on partitions (%.2fs)", time.time() - t0) - - # 4) Export path A: LSDB concat + write_catalog - if USE_LSDB_CONCAT: + # With no new edges, retain the cheaper spatial concat path. + if total_links == 0: collection_path = _concat_and_write_hats( - left_updated, - right_updated, + left_cat, + right_cat, temp_dir, step, translation_config, @@ -934,29 +1180,32 @@ def crossmatch_tiebreak( "END crossmatch_update_compared_to: step=%s links=%d nodes=%d output=%s (%.2fs)", step, total_links, - len(pairs_adj), + node_count, collection_path, time.time() - t0_all, ) + shutil.rmtree(staging_root, ignore_errors=True) return collection_path - # 4b) Export path B (legacy): Dask concat + Parquet + import - collection_path = _concat_parquet_import( - left_updated, - right_updated, + # Aggregate and join neighbor strings in the distributed dataframe. The + # resulting Parquet is reimported so spatial HATS metadata and margins are + # rebuilt from the updated rows. + collection_path = _distributed_pairs_update_and_import( + left_cat, + right_cat, + pairs, temp_dir, logs_dir, step, client, translation_config, - logger=logger, - log_steps=True, + logger, ) logger.info( "END crossmatch_update_compared_to: step=%s links=%d nodes=%d output=%s (%.2fs)", step, total_links, - len(pairs_adj), + node_count, collection_path, time.time() - t0_all, ) diff --git a/packages/crossmatch_diagnostics.py b/packages/crossmatch_diagnostics.py index 77c7269..dc67735 100644 --- a/packages/crossmatch_diagnostics.py +++ b/packages/crossmatch_diagnostics.py @@ -5,6 +5,7 @@ from collections.abc import Mapping, Set as AbstractSet import logging +import dask.dataframe as dd import numpy as np import pandas as pd @@ -64,6 +65,23 @@ def compute_projected_pairs(ddf, columns: list[str]) -> pd.DataFrame: return result +def stage_projected_pairs(ddf, columns: list[str], path: str): + """Project pair columns inside workers and stage them as Parquet. + + Unlike :func:`compute_projected_pairs`, this never gathers all pair rows on + the driver. The returned Dask dataframe starts from the staged dataset, so + downstream shuffles also avoid retaining the original LSDB graph. + """ + meta = _project_pairs_partition(ddf._meta, columns).iloc[:0] + projected = ddf.map_partitions( + _project_pairs_partition, + columns, + meta=meta, + ) + projected.to_parquet(path, engine="pyarrow", write_index=False) + return dd.read_parquet(path, engine="pyarrow") + + def log_pair_separation_diagnostics( pairs: pd.DataFrame, *, diff --git a/sbatch.template b/sbatch.template index 1b5f872..d860b1a 100644 --- a/sbatch.template +++ b/sbatch.template @@ -3,7 +3,7 @@ #SBATCH --job-name={jobname} #SBATCH --chdir={cwd} #SBATCH --account=hpc-pipelines -#SBATCH --time=04:00:00 +#SBATCH --time=12:00:00 #SBATCH --ntasks=1 #SBATCH --cpus-per-task=4 #SBATCH --mem=40G diff --git a/test_configs/config.test.concatenate.yaml b/test_configs/config.test.concatenate.yaml index acad2e7..8aa7da9 100644 --- a/test_configs/config.test.concatenate.yaml +++ b/test_configs/config.test.concatenate.yaml @@ -21,7 +21,7 @@ executor: # account: "hpc-public" # job_extra_directives: # - "--propagate" -# - "--time=2:00:00" +# - "--time=12:00:00" # scale: # minimum_jobs: 4 # maximum_jobs: 12 diff --git a/test_configs/config.test.homogenized_columns_out_of_range.yaml b/test_configs/config.test.homogenized_columns_out_of_range.yaml index 839624f..2922cd2 100644 --- a/test_configs/config.test.homogenized_columns_out_of_range.yaml +++ b/test_configs/config.test.homogenized_columns_out_of_range.yaml @@ -21,7 +21,7 @@ executor: # account: "hpc-public" # job_extra_directives: # - "--propagate" -# - "--time=2:00:00" +# - "--time=12:00:00" # scale: # minimum_jobs: 4 # maximum_jobs: 12 diff --git a/test_configs/config.test.mark.yaml b/test_configs/config.test.mark.yaml index d3fb288..e2846fe 100644 --- a/test_configs/config.test.mark.yaml +++ b/test_configs/config.test.mark.yaml @@ -21,7 +21,7 @@ executor: # account: "hpc-public" # job_extra_directives: # - "--propagate" -# - "--time=2:00:00" +# - "--time=12:00:00" # scale: # minimum_jobs: 4 # maximum_jobs: 12 diff --git a/test_configs/config.test.remove.yaml b/test_configs/config.test.remove.yaml index 666b9a3..4964120 100644 --- a/test_configs/config.test.remove.yaml +++ b/test_configs/config.test.remove.yaml @@ -21,7 +21,7 @@ executor: # account: "hpc-public" # job_extra_directives: # - "--propagate" -# - "--time=2:00:00" +# - "--time=12:00:00" # scale: # minimum_jobs: 4 # maximum_jobs: 12 diff --git a/test_configs/config.test.survey_col_missing.yaml b/test_configs/config.test.survey_col_missing.yaml index 27162f9..931618a 100644 --- a/test_configs/config.test.survey_col_missing.yaml +++ b/test_configs/config.test.survey_col_missing.yaml @@ -21,7 +21,7 @@ executor: # account: "hpc-public" # job_extra_directives: # - "--propagate" -# - "--time=2:00:00" +# - "--time=12:00:00" # scale: # minimum_jobs: 4 # maximum_jobs: 12 diff --git a/test_configs/config.test.survey_not_supported.yaml b/test_configs/config.test.survey_not_supported.yaml index 951146e..c6e93e4 100644 --- a/test_configs/config.test.survey_not_supported.yaml +++ b/test_configs/config.test.survey_not_supported.yaml @@ -21,7 +21,7 @@ executor: # account: "hpc-public" # job_extra_directives: # - "--propagate" -# - "--time=2:00:00" +# - "--time=12:00:00" # scale: # minimum_jobs: 4 # maximum_jobs: 12 diff --git a/test_configs/config.test.survey_supported_but_wrong_flags.yaml b/test_configs/config.test.survey_supported_but_wrong_flags.yaml index 4cf7a45..881f891 100644 --- a/test_configs/config.test.survey_supported_but_wrong_flags.yaml +++ b/test_configs/config.test.survey_supported_but_wrong_flags.yaml @@ -21,7 +21,7 @@ executor: # account: "hpc-public" # job_extra_directives: # - "--propagate" -# - "--time=2:00:00" +# - "--time=12:00:00" # scale: # minimum_jobs: 4 # maximum_jobs: 12 diff --git a/tests/test_crossmatch_adjacency.py b/tests/test_crossmatch_adjacency.py index fa62598..1bb6451 100644 --- a/tests/test_crossmatch_adjacency.py +++ b/tests/test_crossmatch_adjacency.py @@ -17,6 +17,7 @@ ) from crossmatch_cross import ( # noqa: E402 _adjacency_from_pairs as cross_adjacency, + _attach_distributed_neighbors, _merge_compared_to_partition as cross_merge_compared_to, ) @@ -46,3 +47,65 @@ def test_adjacency_is_symmetric_deduplicated_and_excludes_self(adjacency): result = adjacency(left, right) assert result == {"A": {"B"}, "B": {"A", "C"}, "C": {"B"}} + + +def test_distributed_neighbor_join_matches_adjacency_semantics(tmp_path): + import dask.dataframe as dd + + frame = pd.DataFrame( + { + "CRD_ID": pd.Series(["A", "B", "D"], dtype="string"), + "compared_to": pd.Series(["OLD", pd.NA, pd.NA], dtype="string"), + } + ) + pairs = pd.DataFrame( + { + "CRD_IDleft": pd.Series(["A", "A", "A", "B"], dtype="string"), + "CRD_IDright": pd.Series(["B", "C", "C", "B"], dtype="string"), + } + ) + + result = _attach_distributed_neighbors( + dd.from_pandas(frame, npartitions=2), + dd.from_pandas(pairs, npartitions=2), + id_column="CRD_IDleft", + neighbor_column="CRD_IDright", + staging_path=str(tmp_path / "neighbors"), + ).compute(scheduler="synchronous") + compared = result.set_index("CRD_ID")["compared_to"].to_dict() + + assert compared["A"] == "B, C, OLD" + assert pd.isna(compared["B"]) + assert pd.isna(compared["D"]) + + +def test_distributed_self_neighbors_are_symmetric(tmp_path): + import dask.dataframe as dd + + frame = pd.DataFrame( + { + "CRD_ID": pd.Series(["A", "B", "C"], dtype="string"), + "compared_to": pd.Series([pd.NA, pd.NA, pd.NA], dtype="string"), + } + ) + pairs = pd.DataFrame( + { + "CRD_IDleft": pd.Series(["A", "B"], dtype="string"), + "CRD_IDright": pd.Series(["B", "C"], dtype="string"), + } + ) + + result = _attach_distributed_neighbors( + dd.from_pandas(frame, npartitions=2), + dd.from_pandas(pairs, npartitions=2), + id_column="CRD_IDleft", + neighbor_column="CRD_IDright", + staging_path=str(tmp_path / "neighbors"), + symmetric=True, + ).compute(scheduler="synchronous") + + assert result.set_index("CRD_ID")["compared_to"].to_dict() == { + "A": "B", + "B": "A, C", + "C": "B", + } diff --git a/tests/test_crossmatch_diagnostics.py b/tests/test_crossmatch_diagnostics.py index 3000055..65cdbbc 100644 --- a/tests/test_crossmatch_diagnostics.py +++ b/tests/test_crossmatch_diagnostics.py @@ -12,6 +12,7 @@ log_neighbor_count_diagnostics, log_pair_separation_diagnostics, project_catalog_for_pair_crossmatch, + stage_projected_pairs, ) @@ -74,6 +75,29 @@ def test_compute_projected_pairs_transfers_only_requested_plain_columns(): assert result["CRD_IDleft"].tolist() == ["A", "B"] +def test_stage_projected_pairs_roundtrips_without_unused_columns(tmp_path): + source = dd.from_pandas( + pd.DataFrame( + { + "CRD_IDleft": ["A", "B"], + "CRD_IDright": ["C", "D"], + "large_unused_column": ["x" * 100, "y" * 100], + } + ), + npartitions=2, + ) + + staged = stage_projected_pairs( + source, + ["CRD_IDleft", "CRD_IDright"], + str(tmp_path / "pairs"), + ) + result = staged.compute(scheduler="synchronous") + + assert result.columns.tolist() == ["CRD_IDleft", "CRD_IDright"] + assert result["CRD_IDleft"].tolist() == ["A", "B"] + + def test_pair_diagnostics_exclude_self_matches_and_log_radius_fractions(): logger = RecordingLogger() pairs = pd.DataFrame( diff --git a/tests/test_crossmatch_saturation.py b/tests/test_crossmatch_saturation.py index 86952bb..3203a7b 100644 --- a/tests/test_crossmatch_saturation.py +++ b/tests/test_crossmatch_saturation.py @@ -3,6 +3,7 @@ from pathlib import Path import pandas as pd +import dask.dataframe as dd import pytest sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "packages")) @@ -11,7 +12,10 @@ tables_io.types = types.SimpleNamespace(PD_DATAFRAME="PD_DATAFRAME") sys.modules["tables_io"] = tables_io -from crossmatch_cross import _log_neighbor_saturation # noqa: E402 +from crossmatch_cross import ( # noqa: E402 + _log_neighbor_saturation, + _log_neighbor_saturation_distributed, +) class RecordingLogger: @@ -75,3 +79,22 @@ def test_saturation_can_fail_or_be_disabled(): _log_neighbor_saturation( logger=RecordingLogger(), fail_fraction=None, **kwargs ) + + +def test_distributed_saturation_returns_only_scalar_summaries(): + logger = RecordingLogger() + + _log_neighbor_saturation_distributed( + dd.from_pandas(_saturated_pairs(), npartitions=2), + id_col="CRD_IDleft", + source_col="sourceleft", + limit=2, + logger=logger, + context="test", + total_by_source={"survey": 100}, + warn_fraction=0.01, + fail_fraction=None, + ) + + assert len(logger.warning_calls) == 1 + assert any("returned-match diagnostics" in call[0] for call in logger.info_calls)