Enhance catalog preparation and processing for HATS inputs

docs/index.rst

@@ -4,6 +4,9 @@ science_catalogs
 ``science_catalogs`` is a reusable toolkit for preparing science-ready catalogs
 with Dask, LSDB, and HATS.
 
+``prepare_catalog`` accepts a single input path and detects whether it points
+to a regular file collection or an existing HATS catalog.
+
 Beta public API
 ----------------------------------------------------------------------------------------
 
@@ -55,6 +58,7 @@ A runnable example with fictitious parquet inputs lives in
 
 - sample parquet files
 - YAML configs for parquet and HATS output
+- documentation for using an existing HATS catalog as input
 - an interactive walkthrough
 - a smoke test script
 

docs/quickstart.rst

@@ -8,6 +8,40 @@ This package is designed around two layers:
 - ``materialize_catalog`` and ``materialize_lsdb_catalog`` return ``data`` plus
   the written output ``path``.
 
+Use ``materialize_catalog`` only when the final dataframe is expected to fit in
+memory. For larger catalogs, prefer ``build_catalog``, ``write_catalog``, or
+``materialize_lsdb_catalog`` so the workflow stays distributed and disk-backed.
+
+Input modes
+----------------------------------------------------------------------------------------
+
+``prepare_catalog`` can start from two kinds of inputs:
+
+- a regular file or directory
+- an existing HATS catalog opened through LSDB
+
+For file-based inputs, use a single path:
+
+.. code-block:: yaml
+
+   input:
+     catalog_path: /path/to/files_or_file
+     catalog_pattern: "*.parquet"
+
+If ``catalog_path`` is a single file, it is used directly. If it is a regular
+directory, ``catalog_pattern`` is used to find the files to process.
+
+For HATS-based inputs, point ``catalog_path`` to the existing catalog
+directory:
+
+.. code-block:: yaml
+
+   input:
+     catalog_path: /path/to/existing_hats_catalog
+
+In both cases, ``user_selected_cols`` still defines the columns projected into
+the processing step.
+
 Minimal interactive flow
 ----------------------------------------------------------------------------------------
 

examples/notebooks/demo_notebook.ipynb

@@ -15,12 +15,14 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "from glob import glob\n",
+    "\n",
     "import dask.dataframe as dd\n",
     "import lsdb\n",
     "import pandas as pd\n",
-    "from glob import glob\n",
     "from dask.distributed import Client, LocalCluster\n",
-    "from science_catalogs import prepare_catalog, write_catalog, open_lsdb_catalog"
+    "\n",
+    "from science_catalogs import open_lsdb_catalog, prepare_catalog, write_catalog"
    ]
   },
   {

pyproject.toml

@@ -40,9 +40,9 @@ dependencies = [
     "jinja2==3.1.*",
     "h5py",
     "cdshealpix>=0.8.1,<0.9.0",
-    "lsdb>=0.9.0,<0.11.0",
-    "hats>=0.9.0,<0.11.0",
-    "hats-import>=0.9.0,<0.11.0",
+    "lsdb>=0.9.2,<0.11.0",
+    "hats>=0.9.2,<0.11.0",
+    "hats-import>=0.9.2,<0.11.0",
 ]
 
 [project.optional-dependencies]
@@ -121,5 +121,28 @@ ignore = [
     "UP028",
 ]
 
+[tool.ruff.lint.per-file-ignores]
+"*.ipynb" = [
+    "E",
+    "W",
+    "F",
+    "N",
+    "UP",
+    "B",
+    "SIM",
+    "I",
+    "D101",
+    "D102",
+    "D103",
+    "D106",
+    "D206",
+    "D207",
+    "D208",
+    "D300",
+    "D417",
+    "D419",
+    "NPY201",
+]
+
 [tool.coverage.run]
 omit = ["src/science_catalogs/_version.py"]

src/science_catalogs/catalog.py

@@ -13,7 +13,7 @@
 from dask.distributed import Client, wait
 
 from science_catalogs.executor import get_executor
-from science_catalogs.processing import process_file_df
+from science_catalogs.processing import process_dataframe, process_file_df
 from science_catalogs.utils.config import decide_suffix_and_flags
 from science_catalogs.utils.dust import configure_dustmaps_path
 from science_catalogs.utils.partitioning import reorder_and_rechunk
@@ -41,8 +41,80 @@ def load_catalog_config(config_path: str) -> dict[str, Any]:
         return yaml.safe_load(_file) or {}
 
 
-def prepare_catalog(config_path: str, config: dict[str, Any] | None = None) -> PreparedCatalog:
-    """Build the lazy processed catalog without creating workflow-specific artifacts."""
+def _is_hats_catalog_path(path: Path) -> bool:
+    """Check whether a directory already contains a valid HATS catalog."""
+    if not path.is_dir():
+        return False
+
+    from hats.io.validation import is_valid_catalog
+
+    return bool(is_valid_catalog(path))
+
+
+def _glob_input_files(base_path: Path, pattern: str) -> list[str]:
+    """Collect matching input files below a directory."""
+    return sorted(glob.glob((base_path / pattern).as_posix()))
+
+
+def _resolve_input_source(inputs: dict[str, Any]) -> dict[str, Any]:
+    """Resolve a single input path as HATS, a single file, or a directory of files."""
+    raw_catalog_path = inputs.get("catalog_path")
+    raw_catalog_folder = inputs.get("catalog_folder")
+    pattern = inputs.get("catalog_pattern", "*.parquet")
+
+    if raw_catalog_path:
+        catalog_path = Path(raw_catalog_path).expanduser()
+    elif raw_catalog_folder:
+        catalog_path = Path(raw_catalog_folder).expanduser()
+    else:
+        raise ValueError("input.catalog_path is required")
+
+    if not catalog_path.exists():
+        raise FileNotFoundError(f"Input path does not exist: {catalog_path}")
+
+    if _is_hats_catalog_path(catalog_path):
+        return {"source": "hats", "catalog_path": str(catalog_path)}
+
+    if catalog_path.is_file():
+        return {"source": "files", "input_files": [str(catalog_path)]}
+
+    if catalog_path.is_dir():
+        input_files = _glob_input_files(catalog_path, pattern)
+        if not input_files:
+            raise FileNotFoundError(f"No input files found under {catalog_path} matching pattern {pattern!r}")
+        return {"source": "files", "input_files": input_files}
+
+    raise ValueError(f"Unsupported input path type: {catalog_path}")
+
+
+def _build_processed_meta(
+    ddf: dd.DataFrame,
+    cfg: dict[str, Any],
+    *,
+    will_mag: bool,
+    will_dered_flux: bool,
+    will_dered_mag: bool,
+):
+    """Infer the processed partition schema for Dask map_partitions."""
+    meta_input = ddf._meta.copy()
+    meta_output = process_dataframe(
+        meta_input,
+        cfg,
+        will_mag=will_mag,
+        will_dered_flux=will_dered_flux,
+        will_dered_mag=will_dered_mag,
+        source_name="<meta>",
+    )
+    return meta_output.iloc[:0]
+
+
+def prepare_catalog(
+    config_path: str,
+    config: dict[str, Any] | None = None,
+    *,
+    client=None,
+) -> PreparedCatalog:
+    """Build the lazy processed catalog from file inputs or an existing HATS catalog."""
     cfg = config if config is not None else load_catalog_config(config_path)
 
     inputs = cfg.get("input", {})
@@ -57,28 +129,45 @@ def prepare_catalog(config_path: str, config: dict[str, Any] | None = None) -> P
         inputs.get("compute_dereddening", True),
     )
 
-    input_files = [
-        f
-        for f in glob.glob(
-            Path(inputs.get("catalog_folder", "")).expanduser().as_posix()
-            + "/"
-            + inputs.get("catalog_pattern", "*.parquet")
+    input_source = _resolve_input_source(inputs)
+
+    if input_source["source"] == "files":
+        input_files = input_source["input_files"]
+        delayed_dfs = [
+            delayed(process_file_df)(
+                p,
+                cfg_path=config_path,
+                will_mag=will_mag,
+                will_dered_flux=will_dered_flux,
+                will_dered_mag=will_dered_mag,
+            )
+            for p in input_files
+        ]
+        ddf = dd.from_delayed(delayed_dfs)
+    else:
+        input_files = [input_source["catalog_path"]]
+        selected_columns = list(inputs.get("user_selected_cols", []) or []) or None
+        hats_catalog = open_lsdb_catalog(
+            input_source["catalog_path"],
+            client=client,
+            columns=selected_columns,
         )
-    ]
-    if not input_files:
-        raise FileNotFoundError("No input files found for catalog_pattern")
-
-    delayed_dfs = [
-        delayed(process_file_df)(
-            p,
-            cfg_path=config_path,
+        ddf = hats_catalog.to_dask_dataframe()
+        ddf = ddf.map_partitions(
+            process_dataframe,
+            cfg,
             will_mag=will_mag,
             will_dered_flux=will_dered_flux,
             will_dered_mag=will_dered_mag,
+            source_name=input_source["catalog_path"],
+            meta=_build_processed_meta(
+                ddf,
+                cfg,
+                will_mag=will_mag,
+                will_dered_flux=will_dered_flux,
+                will_dered_mag=will_dered_mag,
+            ),
         )
-        for p in input_files
-    ]
-    ddf = dd.from_delayed(delayed_dfs)
     ddf_out = reorder_and_rechunk(ddf, output_cfg)
 
     return PreparedCatalog(
@@ -237,7 +326,7 @@ def build_catalog(
             wait(cluster_comm)
         client.run(lambda: gc.collect())
 
-        prepared = prepare_catalog(config_path, config=cfg)
+        prepared = prepare_catalog(config_path, config=cfg, client=client)
         resolved_output_dir = _resolve_output_dir(prepared, output_dir)
         written_paths = write_catalog(
             prepared,

src/science_catalogs/processing.py

@@ -1,4 +1,4 @@
-"""Per-file processing logic."""
+"""Catalog processing logic shared by file and HATS inputs."""
 
 import astropy.units as u
 import numpy as np
@@ -20,25 +20,30 @@ def _keep_input_columns(input_cfg):
     return bool(input_cfg.get(LEGACY_KEEP_INPUT_COLUMNS_KEY, False))
 
 
-def process_file_df(
-    path: str,
-    cfg_path: str,
+def _load_processing_config(cfg_path: str) -> dict:
+    """Load the processing configuration from disk without importing catalog helpers."""
+    import yaml
+
+    with open(cfg_path, "r", encoding="utf-8") as _file:
+        return yaml.safe_load(_file) or {}
+
+
+def process_dataframe(
+    df,
+    cfg: dict,
+    *,
     will_mag: bool,
     will_dered_flux: bool,
     will_dered_mag: bool,
+    source_name: str = "<dataframe>",
 ):
-    """Read, filter, transform, and return a single catalog file as a dataframe."""
-    import yaml
+    """Filter and transform a dataframe according to the catalog configuration."""
+    df = df.copy()
+    input_cfg = cfg.get("input", {})
+    dust_cfg = cfg.get("dust", {})
+    output_cfg = cfg.get("output", {})
+    invalid = cfg.get("invalid_handling", {})
 
-    with open(cfg_path, "r", encoding="utf-8") as _f:
-        cfgw = yaml.safe_load(_f) or {}
-
-    input_cfg = cfgw.get("input", {})
-    dust_cfg = cfgw.get("dust", {})
-    output_cfg = cfgw.get("output", {})
-    invalid = cfgw.get("invalid_handling", {})
-
-    input_user_selected_cols = list(input_cfg.get("user_selected_cols", []) or [])
     is_id_index = bool(input_cfg.get("is_id_in_index", False))
 
     col_pattern = input_cfg.get("col_pattern")
@@ -50,14 +55,12 @@ def process_file_df(
     mag_offset = as_float_or_none(output_cfg.get("mag_offset"))
     a_ebv = dict(output_cfg.get("A_EBV", {}))
 
-    df = detect_and_read(path, input_user_selected_cols)
-
     filt = input_cfg.get("filter", {})
     if filt.get("enabled"):
         col = filt.get("column")
         val = filt.get("value")
         if col not in df.columns:
-            raise ValueError(f"Boolean column '{col}' not found in {path}")
+            raise ValueError(f"Boolean column '{col}' not found in {source_name}")
         df = df[df[col] == val]
         if filt.get("drop_column_after_filter") and col in df.columns:
             df = df.drop(columns=[col])
@@ -66,7 +69,7 @@ def process_file_df(
     if init.get("enabled"):
         cut_col = init.get("column")
         if cut_col not in df.columns:
-            raise ValueError(f"Initial-cut column '{cut_col}' not found in {path}")
+            raise ValueError(f"Initial-cut column '{cut_col}' not found in {source_name}")
         col_type = str(init.get("column_type", "flux")).strip().lower()
         mag_val = as_float_or_none(init.get("mag_value"))
         flux_val = as_float_or_none(init.get("flux_value"))
@@ -179,7 +182,7 @@ def apply_replacement(arr, mask, replacement_value):
         final_output_cols.update([final_col, final_err_col])
 
         if col_in not in df.columns or err_in not in df.columns:
-            raise ValueError(f"Missing column(s) {[col_in, err_in]} in file {path}")
+            raise ValueError(f"Missing column(s) {[col_in, err_in]} in source {source_name}")
 
         values = df[col_in].astype(float, copy=False).values
         errors = df[err_in].astype(float, copy=False).values
@@ -239,4 +242,26 @@ def apply_replacement(arr, mask, replacement_value):
     return df
 
 
-__all__ = ["process_file_df", "MAG_CONV"]
+def process_file_df(
+    path: str,
+    cfg_path: str,
+    will_mag: bool,
+    will_dered_flux: bool,
+    will_dered_mag: bool,
+):
+    """Read, filter, transform, and return a single catalog file as a dataframe."""
+    cfgw = _load_processing_config(cfg_path)
+    input_cfg = cfgw.get("input", {})
+    input_user_selected_cols = list(input_cfg.get("user_selected_cols", []) or [])
+    df = detect_and_read(path, input_user_selected_cols)
+    return process_dataframe(
+        df,
+        cfgw,
+        will_mag=will_mag,
+        will_dered_flux=will_dered_flux,
+        will_dered_mag=will_dered_mag,
+        source_name=path,
+    )
+
+
+__all__ = ["process_dataframe", "process_file_df", "MAG_CONV"]