diff --git a/docs/notebooks/hats_dataset.ipynb b/docs/notebooks/hats_dataset.ipynb new file mode 100644 index 000000000..471e38e84 --- /dev/null +++ b/docs/notebooks/hats_dataset.ipynb @@ -0,0 +1,86 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "58f22b5f", + "metadata": {}, + "source": [ + "# Hyrax HATS Dataset Example\n", + "\n", + "Minimal Phase-1 example using `HyraxHATSDataset` with an LSDB-written HATS catalog." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a0e1ea1d", + "metadata": {}, + "outputs": [], + "source": [ + "from pathlib import Path\n", + "import tempfile\n", + "\n", + "import pandas as pd\n", + "import lsdb\n", + "import hyrax" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "dab296ae", + "metadata": {}, + "outputs": [], + "source": [ + "tmp_dir = Path(tempfile.mkdtemp())\n", + "catalog_path = tmp_dir / 'sample_hats'\n", + "df = pd.DataFrame({\n", + " 'object_id': [1001, 1002, 1003],\n", + " 'coord_ra': [150.1, 150.2, 150.3],\n", + " 'coord_dec': [2.1, 2.2, 2.3],\n", + " 'mag-r': [19.1, 19.4, 18.9],\n", + "})\n", + "lsdb.from_dataframe(df, ra_column='coord_ra', dec_column='coord_dec').write_catalog(catalog_path)\n", + "catalog_path" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8abb4866", + "metadata": {}, + "outputs": [], + "source": [ + "h = hyrax.Hyrax()\n", + "h.config['data_request'] = {\n", + " 'train': {\n", + " 'data': {\n", + " 'dataset_class': 'HyraxHATSDataset',\n", + " 'data_location': str(catalog_path),\n", + " 'fields': ['coord_ra', 'mag-r'],\n", + " 'primary_id_field': 'object_id',\n", + " }\n", + " }\n", + "}\n", + "prepared = h.prepare()['train']\n", + "dataset = prepared._primary_or_first_dataset()\n", + "len(dataset), dataset.get_object_id(0), dataset.get_coord_ra(1), getattr(dataset, 'get_mag-r')(2)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2b5568db", + "metadata": {}, + "outputs": [], + "source": [ + "sample = dataset.sample_data()\n", + "assert 'data' in sample\n", + "sample['data']" + ] + } + ], + "metadata": {}, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/notebooks/hats_dataset_phase2_loopback.ipynb b/docs/notebooks/hats_dataset_phase2_loopback.ipynb new file mode 100644 index 000000000..b01780e2b --- /dev/null +++ b/docs/notebooks/hats_dataset_phase2_loopback.ipynb @@ -0,0 +1,186 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Hyrax HATS Phase 2 Example with `HyraxLoopback`\n", + "\n", + "This notebook demonstrates how to configure and exercise `HyraxHATSDataset` (phase-2 lazy, partition-aware path) using the built-in `HyraxLoopback` model.\n", + "\n", + "It is designed so you can replace one path (`HATS_PATH`) with your own catalog (for example your ~7GB catalog) and iterate quickly.\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## What this notebook demonstrates\n", + "\n", + "1. Setting up a `data_request` that uses `HyraxHATSDataset`.\n", + "2. Phase-2 HATS options (`bundle_size`, `max_cached_bundles`, `project_columns`, `strict_metadata`).\n", + "3. Accessing rows/fields through `get_(idx)`.\n", + "4. Observing bundle-cache behavior.\n", + "5. Running end-to-end `infer()` with `HyraxLoopback`.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from pathlib import Path\n", + "import numpy as np\n", + "\n", + "import hyrax\n", + "\n", + "print('Hyrax version:', hyrax.__version__)\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# TODO: replace with your local HATS catalog path\n", + "HATS_PATH = Path('/path/to/your/hats_catalog')\n", + "\n", + "if not HATS_PATH.exists():\n", + " print(f'WARNING: {HATS_PATH} does not exist yet. Update HATS_PATH before running full notebook.')\n", + "\n", + "# Adjust these to columns in your catalog\n", + "PRIMARY_ID_FIELD = 'object_id'\n", + "FEATURE_FIELDS = ['coord_ra', 'coord_dec']\n", + "\n", + "print('Primary ID:', PRIMARY_ID_FIELD)\n", + "print('Feature fields:', FEATURE_FIELDS)\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "h = hyrax.Hyrax()\n", + "h.set_config('model.name', 'HyraxLoopback')\n", + "\n", + "# Optional, keeps data loading deterministic while exploring\n", + "h.set_config('data_loader.shuffle', False)\n", + "\n", + "data_request = {\n", + " 'infer': {\n", + " 'catalog': {\n", + " 'dataset_class': 'HyraxHATSDataset',\n", + " 'data_location': str(HATS_PATH),\n", + " 'fields': FEATURE_FIELDS,\n", + " 'primary_id_field': PRIMARY_ID_FIELD,\n", + " 'split_fraction': 1.0,\n", + " # Phase-2 options\n", + " 'hats': {\n", + " 'bundle_size': 256,\n", + " 'max_cached_bundles': 64,\n", + " 'project_columns': 'auto', # 'auto' | 'all' | explicit list\n", + " 'strict_metadata': True,\n", + " },\n", + " }\n", + " }\n", + "}\n", + "\n", + "h.set_config('data_request', data_request)\n", + "data_request\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "prepared = h.prepare()\n", + "dataset = prepared['infer']._primary_or_first_dataset()\n", + "\n", + "print('Dataset class:', dataset.__class__.__name__)\n", + "print('Length:', len(dataset))\n", + "print('Columns (first 15):', dataset.column_names[:15])\n", + "print('Sample data keys:', list(dataset.sample_data()['data'].keys())[:10])\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Demonstrate dynamic getters\n", + "first_idx = 0\n", + "primary_getter = getattr(dataset, f'get_{PRIMARY_ID_FIELD}')\n", + "\n", + "print('ID at row 0:', primary_getter(first_idx))\n", + "for field in FEATURE_FIELDS:\n", + " getter = getattr(dataset, f'get_{field}')\n", + " print(f'{field} at row 0:', getter(first_idx))\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Demonstrate bundle-cache effects via repeated nearby access\n", + "# (uses internal stats for exploration while phase-2 is maturing)\n", + "cache = dataset._accessor.cache\n", + "print('Initial cache hits/misses:', cache.hits, cache.misses)\n", + "\n", + "for idx in [0, 1, 2, 2, 1, 0]:\n", + " _ = getattr(dataset, f'get_{FEATURE_FIELDS[0]}')(idx)\n", + "\n", + "print('Post-access cache hits/misses:', cache.hits, cache.misses)\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Run end-to-end inference. HyraxLoopback returns model input values.\n", + "inference_results = h.infer()\n", + "\n", + "print('Number of inference rows:', len(inference_results))\n", + "print('First 3 result ids:', inference_results.ids()[:3])\n", + "\n", + "for i in range(3):\n", + " print(f'Result {i}:', inference_results[i])\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Suggested experiments with your 7GB catalog\n", + "\n", + "- Increase/decrease `bundle_size` and observe cache hit/miss behavior and wall time.\n", + "- Compare `project_columns='auto'` versus `'all'`.\n", + "- Try different `fields` lists in `data_request` to validate projected reads.\n", + "- Enable debug logging to see cache and catalog-open behavior in detail.\n" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} \ No newline at end of file diff --git a/specs/hasts_dataset_phase2.md b/specs/hasts_dataset_phase2.md new file mode 100644 index 000000000..38b92f6ab --- /dev/null +++ b/specs/hasts_dataset_phase2.md @@ -0,0 +1,284 @@ +# HATS Dataset Phase 2 Plan (Lazy, Partition-Aware Access) + +## Scope and intent + +This document defines **Phase 2** for `HyraxHATSDataset`: move from full eager materialization (`catalog.compute()`) to a **lazy, partition-aware row access model** that remains compatible with Hyrax's `get_(idx)` interface. + +Phase 2 is intentionally implementation-oriented so we can review architecture and sequencing before coding. + +--- + +## Why Phase 2 is needed + +The current Phase 1 dataset loads an entire LSDB catalog into a pandas DataFrame. This is simple, but does not scale for large HATS catalogs and defeats LSDB/HATS advantages: + +- HATS is partitioned by sky region (HEALPix), and LSDB is designed to read only needed partitions. +- Random row reads can be expensive if each read triggers broad scans. +- A Hyrax access pattern that reads many nearby indices should exploit locality. + +Phase 2 will preserve Hyrax ergonomics while reducing memory pressure and improving large-catalog behavior. + +--- + +## Phase 2 goals + +1. **Do not materialize full catalog by default** during dataset initialization. +2. Support `__len__` and `get_(idx)` with deterministic global row indexing. +3. Introduce a **row-bundle cache** keyed by partition (and optional row-group window) to avoid repeated remote reads. +4. Allow optional **column projection** and basic **filter pushdown**. +5. Keep existing Hyrax config usage stable where possible. +6. Add tests/benchmarks/docs that make behavior and tradeoffs explicit. + +--- + +## Explicit non-goals for this phase + +- Rewriting Hyrax's data request execution model. +- Full spatial query DSL in Hyrax config. +- Distributed scheduler tuning (Dask cluster orchestration). +- Margin catalog joins and crossmatch orchestration. + +--- + +## Proposed architecture + +## 1) New internal components in `hats_dataset.py` + +### A. `HATSPartitionIndex` + +A lightweight mapping between global row indices and HATS partition-local offsets. + +Responsibilities: + +- Read partition metadata once. +- Build cumulative row counts per partition. +- Resolve global index `idx` -> `(partition_id, local_offset)` via binary search. +- Expose total row count for `__len__`. + +Notes: + +- Source of truth should be LSDB/HATS metadata artifacts when available (avoid scanning data leaves). +- If row-count metadata is missing/inconsistent, fail with clear message and fallback options. + +### B. `HATSRowBundleCache` + +A bounded in-memory cache that stores recently accessed row bundles. + +Responsibilities: + +- Cache unit: `(partition_id, selected_columns, bundle_start, bundle_size)`. +- Configurable `max_bundles` and `bundle_size`. +- LRU eviction. +- Optional metrics counters (hit/miss, bytes, load time). + +### C. `HATSLazyAccessor` + +Execution layer that fetches rows from LSDB/HATS using partition-aware reads. + +Responsibilities: + +- Use `HATSPartitionIndex` to locate partition/local offset. +- Fetch only required columns where possible. +- Load a small contiguous bundle around requested local row to increase cache hit probability. +- Return scalar/array values normalized to current dataset semantics. + +## 2) `HyraxHATSDataset` responsibilities after refactor + +- Initialize LSDB catalog lazily (`lsdb.read_hats` / equivalent lazy-open path). +- Build normalized column map for getter generation. +- Instantiate `HATSPartitionIndex` and `HATSRowBundleCache`. +- Generate dynamic getters that route to `HATSLazyAccessor.get_value(idx, field)`. +- Keep `sample_data` behavior (read first row lazily). + +--- + +## Data access strategy + +## A. Global index resolution + +Use cumulative row counts per partition: + +- `partition_row_counts = [n0, n1, ..., nk]` +- `prefix = cumulative_sum(partition_row_counts)` +- Resolve via `bisect_right(prefix, idx)`. + +Complexities: + +- Build: `O(P)` for `P` partitions. +- Lookup: `O(log P)` per access. + +## B. Bundle reads + +Given `(partition, local_offset)`: + +- Compute `bundle_start = floor(local_offset / bundle_size) * bundle_size`. +- Read rows `[bundle_start : bundle_start + bundle_size)` from that partition. +- Cache bundle. +- Extract requested row. + +Rationale: + +- Maintains bounded memory. +- Improves locality if training/evaluation iterates near-adjacent indices. +- Avoids one-file-per-row penalties. + +## C. Column projection + +Column set for bundle reads: + +- Always include `primary_id_field`. +- Include fields requested in data request if present. +- Permit override: `hats_columns = "all" | [..]` (new config option). + +Default for phase 2: projected columns = union of required columns for active getters in request context, with safe fallback to all columns if projection path is unavailable. + +## D. Optional filter pushdown + +Add optional config section: + +```toml +[data_request.train.catalog.hats] +filters = [ + ["coord_ra", ">=", 120.0], + ["coord_ra", "<", 130.0], +] +``` + +Apply at catalog-open/read stage using LSDB-compatible filtering primitives where available. + +If a filter is unsupported, fail fast with precise error and examples. + +--- + +## Configuration additions (proposed) + +All options optional with conservative defaults. + +```toml +[data_request...hats] +lazy = true # default true in phase 2 +bundle_size = 256 # rows per partition bundle +max_cached_bundles = 64 # LRU capacity +project_columns = "auto" # auto | all | explicit list +filters = [] # optional predicate list +strict_metadata = true # fail if partition row counts unavailable +``` + +Backward compatibility: + +- Existing phase-1 configs should continue to work. +- We may keep an escape hatch `lazy=false` for troubleshooting. + +--- + +## Error handling and edge cases + +1. **Missing metadata for row counts** + - If `strict_metadata=true`: raise `ValueError` with actionable guidance. + - Else: fallback strategy (documented as potentially expensive). + +2. **Index out of range** + - Raise `IndexError` consistently from central accessor. + +3. **Field missing from projected columns** + - Raise `KeyError` mentioning normalized/original column names. + +4. **Column name normalization collisions** + - Keep deterministic suffix policy (`_2`, `_3`, ...), and retain reverse map in object state. + +5. **Remote I/O failures / transient read errors** + - Surface original exception class with contextual message (partition id, columns). + +--- + +## Testing plan + +Add or update: `tests/hyrax/test_hats_dataset.py` and dedicated lazy-access tests. + +## Unit tests + +1. **Partition index construction** from synthetic multi-partition catalogs. +2. **Global index mapping** correctness across partition boundaries. +3. **Getter correctness** under lazy reads for scalar and array-like fields. +4. **Cache behavior**: deterministic hit/miss with controlled access sequence. +5. **Column projection**: ensure only required columns requested (via mock/spies). +6. **Normalization collisions**: generated getter names are unique and stable. + +## Integration tests + +1. End-to-end Hyrax prepare/train split with lazy HATS dataset. +2. `fields` in data_request interacts correctly with projected reads. +3. Optional filter path returns reduced row domain with correct indexing semantics. + +## Performance regression checks + +Small benchmark script (non-CI optional): + +- Compare phase-1 eager vs phase-2 lazy memory and wall-clock for: + - sequential access of N rows, + - random access of N rows, + - clustered random access. + +Record median metrics and add to docs. + +--- + +## Implementation sequence + +### Milestone 1: Internal scaffolding + +- Add internal index/cache/accessor classes. +- Wire dataset getters through accessor. +- Preserve old behavior behind compatibility switch. + +### Milestone 2: Metadata-driven indexing + +- Implement robust partition row-count extraction. +- Add strict/fallback modes. +- Validate length and index mapping. + +### Milestone 3: Bundle cache + projection + +- Implement bundle fetch/cache. +- Add projection policy (`auto`, `all`, explicit list). +- Add cache metrics hooks. + +### Milestone 4: Filter pushdown (optional, guarded) + +- Parse filter config schema. +- Apply LSDB-supported filters and validation. + +### Milestone 5: Hardening + +- Extend tests (unit + integration). +- Add benchmark notes and developer docs. +- Finalize error messages and migration notes. + +--- + +## Acceptance criteria (Phase 2 done) + +1. No full-catalog `.compute()` on initialization in default lazy path. +2. `__len__` and dynamic `get_(idx)` pass full test suite. +3. Repeated local accesses show measurable cache benefit in benchmark. +4. Existing phase-1 configs run unchanged (or with clearly documented warnings). +5. Docs clearly describe lazy semantics, config knobs, and limitations. + +--- + +## Resolutions before implementation + +1. Prefer `lsdb.open_catalog` where available, and fall back to `lsdb.read_hats` when required by the installed LSDB version or catalog behavior. +2. Start with the simplest viable row-count metadata path (derive counts via partition-aware length operations), then iterate if catalog variants require richer metadata parsing. +3. Preserve globally stable row ordering within a fixed underlying HATS representation; if HATS format evolution changes physical ordering, index drift across versions is acceptable and should be documented. +4. Keep cache stats in debug logging only for Phase 2 (no public user-facing stats API yet). + +--- + +## References + +- Existing phase-1 spec: `specs/hats_dataset.md`. +- LSDB docs on HATS structure and lazy data access patterns (for implementation alignment): + - https://docs.lsdb.io/en/latest/data-access/hats.html + - https://docs.lsdb.io/en/latest/tutorials/catalog_object.html + - https://github.com/astronomy-commons/lsdb diff --git a/specs/hats_dataset.md b/specs/hats_dataset.md new file mode 100644 index 000000000..7abbd0058 --- /dev/null +++ b/specs/hats_dataset.md @@ -0,0 +1,149 @@ +# HATS Dataset Interoperability Plan (LSDB ↔ Hyrax) + +## Context + +Hyrax dataset classes must provide: + +1. `__init__(self, config, data_location=None)` +2. `__len__(self)` +3. `get_(self, idx)` for each requested field +4. `get_(self, idx)` for the configured primary ID field + +The current LSST-specific HATS path computes an entire catalog eagerly. We want a generic HATS dataset class that enables interoperability between LSDB and Hyrax while preserving Hyrax's `get_*` access model. + +--- + +## Goals + +1. Add a generic built-in dataset class: `HyraxHATSDataset`. +2. Support reading HATS catalogs through LSDB. +3. Provide dynamic `get_*` getters for HATS columns. +4. Keep API and behavior aligned with Hyrax dataset conventions. +5. Provide tests and docs for the MVP so work is durable if interrupted. + +--- + +## Non-goals for Phase 1 + +1. Full lazy partition-aware row locator index. +2. Advanced cone/polygon searches at dataset initialization time. +3. Margin catalog join behavior. +4. Remote-object-store performance tuning. + +These remain planned for subsequent phases. + +--- + +## Target user experience + +Users can configure: + +```toml +[data_request.train.catalog] +dataset_class = "HyraxHATSDataset" +data_location = "/path/or/url/to/hats_catalog" +fields = ["object_id", "coord_ra", "coord_dec"] +primary_id_field = "object_id" +``` + +And Hyrax will call generated methods such as: + +- `get_object_id(idx)` +- `get_coord_ra(idx)` +- `get_coord_dec(idx)` + +--- + +## Architecture + +### Class: `HyraxHATSDataset` + +Location: + +- `src/hyrax/datasets/hats_dataset.py` + +Responsibilities: + +1. Read HATS catalog via LSDB (`lsdb.read_hats`). +2. Materialize a pandas DataFrame (MVP implementation). +3. Normalize column names to valid Python getter suffixes. +4. Dynamically add `get_(idx)` methods for each normalized column. +5. Expose `__len__`, `__getitem__`, and `sample_data` helpers consistent with existing built-in datasets. + +### Column naming + +HATS column names may contain characters invalid in method names. + +- Normalize using `re.sub(r"\W", "_", column_name)`. +- Ensure uniqueness by suffixing (`_2`, `_3`, ...`) when collisions occur. +- Rename DataFrame columns to normalized names so user-requested fields can map directly to generated getters. + +### Primary ID behavior + +For phase 1, users should set `primary_id_field` to the exact column name present in the HATS catalog. + +Column names are preserved as-is in this phase. + +--- + +## Config surface (Phase 1) + +### Required + +- `data_location` (HATS root path/URL) + +### Optional + +- No phase-1 column-subsetting option is implemented in the dataset class itself. + +Hyrax still controls which fields are *requested at runtime* via `data_request...fields`. + +--- + +## Error handling + +1. Missing `data_location` => `ValueError` +2. LSDB import/runtime issues bubble up naturally (explicit dependency) + +--- + +## Testing strategy (Phase 1) + +Create dedicated unit tests: + +- `tests/hyrax/test_hats_dataset.py` + +Tests: + +1. Initialization and length with synthetic HATS catalog. +2. Dynamic getter generation and retrieval correctness. +3. `sample_data` shape and values. +4. Behavior when `fields` is set in data_request (dataset still exposes all column getters; DataProvider requests selected fields). + +Use `lsdb.from_dataframe(...).write_catalog(tmp_path)` to create fixture data. + +--- + +## Incremental phases after MVP + +### Phase 2 + +1. Partition-aware lazy row access without full DataFrame compute. +2. Row bundle cache to avoid repeated row fetch. +3. Optional filter pushdown and spatial prefilter. + +### Phase 3 + +1. Shared HATS backend utility for LSSTDataset + HyraxHATSDataset. +2. Interop helpers (`to_lsdb_catalog`). +3. Benchmarks and docs notebook. + +--- + +## Deliverables checklist + +- [x] Spec document in `specs/hats_dataset.md` +- [x] `HyraxHATSDataset` implementation +- [x] Export from `hyrax.datasets` +- [x] Unit tests +- [ ] Optional docs touch-ups diff --git a/src/hyrax/datasets/__init__.py b/src/hyrax/datasets/__init__.py index 8ed68fc78..2bf67b6d5 100644 --- a/src/hyrax/datasets/__init__.py +++ b/src/hyrax/datasets/__init__.py @@ -70,6 +70,7 @@ from .result_factories import create_results_writer, load_results_dataset from .dataset_registry import HyraxDataset from .hyrax_csv_dataset import HyraxCSVDataset +from .hats_dataset import HyraxHATSDataset from .mmu_dataset import MultimodalUniverseDataset from .data_cache import DataCache @@ -88,6 +89,7 @@ "HyraxRandomDataset", "HyraxRandomDatasetBase", "HyraxCSVDataset", + "HyraxHATSDataset", "MultimodalUniverseDataset", "DataCache", ] diff --git a/src/hyrax/datasets/hats_dataset.py b/src/hyrax/datasets/hats_dataset.py new file mode 100644 index 000000000..dcc003bf3 --- /dev/null +++ b/src/hyrax/datasets/hats_dataset.py @@ -0,0 +1,317 @@ +from __future__ import annotations + +import bisect +import logging +from collections import OrderedDict +from pathlib import Path +from types import MethodType +from typing import Any + +from hyrax.datasets.dataset_registry import HyraxDataset + +logger = logging.getLogger(__name__) + + +class HATSPartitionIndex: + """Partition row-count index used to resolve global row indices lazily.""" + + def __init__(self, catalog: Any, strict_metadata: bool = True): + self.partition_row_counts = self._build_partition_row_counts(catalog, strict_metadata=strict_metadata) + self._prefix_counts = [] + + running = 0 + for count in self.partition_row_counts: + running += count + self._prefix_counts.append(running) + + @property + def total_rows(self) -> int: + return 0 if len(self._prefix_counts) == 0 else self._prefix_counts[-1] + + def resolve(self, idx: int) -> tuple[int, int]: + if idx < 0 or idx >= self.total_rows: + raise IndexError(f"Index {idx} is out of range for dataset of length {self.total_rows}.") + + partition_id = bisect.bisect_right(self._prefix_counts, idx) + previous_total = 0 if partition_id == 0 else self._prefix_counts[partition_id - 1] + local_offset = idx - previous_total + return partition_id, local_offset + + @staticmethod + def _build_partition_row_counts(catalog: Any, strict_metadata: bool = True) -> list[int]: + ddf = getattr(catalog, "_ddf", None) + if ddf is None: + ddf = getattr(catalog, "ddf", None) + + if ddf is not None and hasattr(ddf, "map_partitions"): + try: + counts = ddf.map_partitions(len).compute() + return [int(x) for x in counts] + except Exception: + if strict_metadata: + msg = "Could not derive per-partition row counts from catalog metadata." + raise ValueError(msg) from None + logger.debug("Falling back to eager row-count compute for HATS catalog.", exc_info=True) + + try: + fallback_length = len(catalog.compute()) + logger.debug("Using eager fallback row-count metadata for HATS catalog.") + return [fallback_length] + except Exception as exc: + msg = "Could not derive HATS partition row counts; unable to build lazy row index." + raise ValueError(msg) from exc + + +class HATSRowBundleCache: + """Simple LRU cache for partition row bundles.""" + + def __init__(self, max_bundles: int = 64): + self.max_bundles = max(1, int(max_bundles)) + self._cache: OrderedDict[tuple[Any, ...], Any] = OrderedDict() + self.hits = 0 + self.misses = 0 + + def get(self, key: tuple[Any, ...]): + if key in self._cache: + self._cache.move_to_end(key) + self.hits += 1 + logger.debug("HATS row-bundle cache hit for key=%s", key) + return self._cache[key] + + self.misses += 1 + logger.debug("HATS row-bundle cache miss for key=%s", key) + return None + + def set(self, key: tuple[Any, ...], value: Any): + self._cache[key] = value + self._cache.move_to_end(key) + + if len(self._cache) > self.max_bundles: + evicted_key, _ = self._cache.popitem(last=False) + logger.debug("HATS row-bundle cache evicted key=%s", evicted_key) + + +class HATSLazyAccessor: + """Lazy row fetch/access helper for Hyrax HATS datasets.""" + + def __init__( + self, + catalog: Any, + partition_index: HATSPartitionIndex, + cache: HATSRowBundleCache, + *, + bundle_size: int = 256, + project_columns: str | list[str] = "auto", + required_columns: list[str] | None = None, + ): + self.catalog = catalog + self.partition_index = partition_index + self.cache = cache + self.bundle_size = max(1, int(bundle_size)) + self.project_columns = project_columns + self.required_columns = required_columns or [] + + def _bundle_key(self, partition_id: int, bundle_start: int, columns: list[str]) -> tuple[Any, ...]: + return (partition_id, bundle_start, self.bundle_size, tuple(columns)) + + def _projected_columns(self, requested_column: str | None) -> list[str]: + all_columns = self._catalog_columns() + if self.project_columns == "all": + return list(all_columns) + + requested = [requested_column] if requested_column else [] + + if isinstance(self.project_columns, list): + columns = list(dict.fromkeys(self.project_columns + requested)) + return columns + + columns = list(dict.fromkeys(self.required_columns + requested)) + return columns if len(columns) > 0 else requested + + def _catalog_columns(self) -> list[str]: + columns = getattr(self.catalog, "columns", None) + if columns is not None: + return list(columns) + + ddf = getattr(self.catalog, "_ddf", None) or getattr(self.catalog, "ddf", None) + if ddf is not None and hasattr(ddf, "columns"): + return list(ddf.columns) + + return [] + + def _read_partition_dataframe(self, partition_id: int, columns: list[str]): + ddf = getattr(self.catalog, "_ddf", None) or getattr(self.catalog, "ddf", None) + if ddf is not None and hasattr(ddf, "get_partition"): + partition = ddf.get_partition(partition_id) + if len(columns) > 0: + partition = partition[columns] + return partition.compute() + + # Conservative fallback path: compute full catalog and treat as one partition. + df = self.catalog.compute() + if len(columns) > 0: + return df[columns] + return df + + def get_row(self, idx: int): + partition_id, local_offset = self.partition_index.resolve(idx) + columns = self._projected_columns(requested_column=None) + + bundle_start = (local_offset // self.bundle_size) * self.bundle_size + cache_key = self._bundle_key(partition_id, bundle_start, columns) + bundle = self.cache.get(cache_key) + + if bundle is None: + partition_df = self._read_partition_dataframe(partition_id, columns) + bundle = partition_df.iloc[bundle_start : bundle_start + self.bundle_size] + self.cache.set(cache_key, bundle) + + return bundle.iloc[local_offset - bundle_start] + + def get_value(self, idx: int, column: str): + partition_id, local_offset = self.partition_index.resolve(idx) + columns = self._projected_columns(requested_column=column) + + bundle_start = (local_offset // self.bundle_size) * self.bundle_size + cache_key = self._bundle_key(partition_id, bundle_start, columns) + bundle = self.cache.get(cache_key) + + if bundle is None: + partition_df = self._read_partition_dataframe(partition_id, columns) + bundle = partition_df.iloc[bundle_start : bundle_start + self.bundle_size] + self.cache.set(cache_key, bundle) + + import numpy as np + import pandas as pd + + ret_val = bundle.iloc[local_offset - bundle_start][column] + if isinstance(ret_val, pd.Series): + ret_val = ret_val.to_numpy() + elif isinstance(ret_val, (list, tuple)): + ret_val = np.asarray(ret_val) + return ret_val + + +class HyraxHATSDataset(HyraxDataset): + """Generic Hyrax dataset for HATS catalogs loaded through LSDB. + + Phase-2 default behavior is lazy and partition-aware, with optional fallback + to eager compute if metadata access is unavailable. + """ + + def __init__(self, config: dict, data_location: Path = None): + if data_location is None: + raise ValueError("A `data_location` to a HATS catalog must be provided.") + + self.data_location = data_location + self._config = config + + hats_config = self._extract_hats_config(config) + + self.catalog = self._open_hats_catalog(data_location) + self.column_names = self._catalog_columns() + + strict_metadata = hats_config.get("strict_metadata", True) + bundle_size = hats_config.get("bundle_size", 256) + max_cached_bundles = hats_config.get("max_cached_bundles", 64) + project_columns = hats_config.get("project_columns", "auto") + + required_columns = self._required_columns(config, self.column_names) + + partition_index = HATSPartitionIndex(self.catalog, strict_metadata=strict_metadata) + cache = HATSRowBundleCache(max_bundles=max_cached_bundles) + self._accessor = HATSLazyAccessor( + self.catalog, + partition_index, + cache, + bundle_size=bundle_size, + project_columns=project_columns, + required_columns=required_columns, + ) + + def _make_getter(column: str): + def getter(self, idx: int, _col: str = column): + return self._accessor.get_value(idx, _col) + + return getter + + for col in self.column_names: + method_name = f"get_{col}" + if not hasattr(self, method_name): + setattr(self, method_name, MethodType(_make_getter(col), self)) + + super().__init__(config) + + @staticmethod + def _extract_hats_config(config: dict) -> dict: + request = config.get("data_request", {}) + for split_cfg in request.values(): + if not isinstance(split_cfg, dict): + continue + for dataset_cfg in split_cfg.values(): + if isinstance(dataset_cfg, dict) and dataset_cfg.get("dataset_class") == "HyraxHATSDataset": + return dataset_cfg.get("hats", {}) + return {} + + @staticmethod + def _required_columns(config: dict, all_columns: list[str]) -> list[str]: + request = config.get("data_request", {}) + required = [] + + for split_cfg in request.values(): + if not isinstance(split_cfg, dict): + continue + for dataset_cfg in split_cfg.values(): + if not isinstance(dataset_cfg, dict): + continue + fields = dataset_cfg.get("fields", []) + if isinstance(fields, list): + required.extend(fields) + primary_id_field = dataset_cfg.get("primary_id_field") + if primary_id_field: + required.append(primary_id_field) + + required = [col for col in dict.fromkeys(required) if col in all_columns] + return required + + @staticmethod + def _open_hats_catalog(data_location: Path): + import lsdb + + if hasattr(lsdb, "open_catalog"): + try: + return lsdb.open_catalog(data_location) + except Exception: + logger.debug("Failed to open HATS catalog via lsdb.open_catalog; falling back.", exc_info=True) + + if hasattr(lsdb, "read_hats"): + return lsdb.read_hats(data_location) + + raise RuntimeError("Installed LSDB does not expose open_catalog or read_hats.") + + def _catalog_columns(self) -> list[str]: + columns = getattr(self.catalog, "columns", None) + if columns is not None: + return list(columns) + + ddf = getattr(self.catalog, "_ddf", None) or getattr(self.catalog, "ddf", None) + if ddf is not None and hasattr(ddf, "columns"): + return list(ddf.columns) + + # Fallback for unusual catalog wrappers. + return list(self.catalog.compute().columns) + + def __len__(self) -> int: + return self._accessor.partition_index.total_rows + + def __getitem__(self, idx): + """Currently required by Hyrax machinery, but likely to be phased out.""" + row = self._accessor.get_row(idx) + return {"data": {col: row[col] for col in self.column_names}} + + def sample_data(self): + """Return the first record in dictionary form as a sample.""" + if len(self) == 0: + return {"data": {col: None for col in self.column_names}} + row = self._accessor.get_row(0) + return {"data": {col: row[col] for col in self.column_names}} diff --git a/tests/hyrax/test_hats_dataset.py b/tests/hyrax/test_hats_dataset.py new file mode 100644 index 000000000..28e133b79 --- /dev/null +++ b/tests/hyrax/test_hats_dataset.py @@ -0,0 +1,95 @@ +from pathlib import Path + +import pandas as pd +import pytest + +import hyrax + +lsdb = pytest.importorskip("lsdb") + + +@pytest.fixture(scope="function") +def hats_catalog_path(tmp_path: Path) -> Path: + df = pd.DataFrame( + { + "object_id": [1001, 1002, 1003], + "coord_ra": [150.1, 150.2, 150.3], + "coord_dec": [2.1, 2.2, 2.3], + "mag-r": [19.1, 19.4, 18.9], + } + ) + catalog = lsdb.from_dataframe(df, ra_column="coord_ra", dec_column="coord_dec") + + out_path = tmp_path / "sample_hats" + catalog.write_catalog(out_path) + + return out_path + + +@pytest.fixture(scope="function") +def test_hyrax_hats_dataset(hats_catalog_path: Path): + h = hyrax.Hyrax() + h.config["data_request"] = { + "train": { + "data": { + "dataset_class": "HyraxHATSDataset", + "data_location": str(hats_catalog_path), + "primary_id_field": "object_id", + "split_fraction": 1.0, + } + } + } + return h + + +def test_hats_dataset_initialization(test_hyrax_hats_dataset): + dataset = test_hyrax_hats_dataset.prepare() + assert len(dataset["train"]) == 3 + + +def test_hats_dataset_dynamic_getters(test_hyrax_hats_dataset): + dataset = test_hyrax_hats_dataset.prepare() + hats_dataset = dataset["train"]._primary_or_first_dataset() + + assert hasattr(hats_dataset, "get_object_id") + assert hasattr(hats_dataset, "get_coord_ra") + assert hasattr(hats_dataset, "get_coord_dec") + assert hasattr(hats_dataset, "get_mag-r") + + assert hats_dataset.get_object_id(0) == 1001 + assert hats_dataset.get_coord_ra(1) == pytest.approx(150.2) + assert getattr(hats_dataset, "get_mag-r")(2) == pytest.approx(18.9) + + +def test_hats_dataset_sample_data(test_hyrax_hats_dataset): + dataset = test_hyrax_hats_dataset.prepare() + hats_dataset = dataset["train"]._primary_or_first_dataset() + sample = hats_dataset.sample_data() + + assert "data" in sample + assert sample["data"]["object_id"] == 1001 + assert sample["data"]["coord_ra"] == pytest.approx(150.1) + assert sample["data"]["mag-r"] == pytest.approx(19.1) + + +def test_hats_dataset_with_data_request_fields_still_builds_all_getters(hats_catalog_path: Path): + h = hyrax.Hyrax() + h.config["data_request"] = { + "train": { + "data": { + "dataset_class": "HyraxHATSDataset", + "data_location": str(hats_catalog_path), + "fields": ["coord_ra"], + "primary_id_field": "object_id", + "split_fraction": 1.0, + } + } + } + + dataset = h.prepare() + hats_dataset = dataset["train"]._primary_or_first_dataset() + + assert hasattr(hats_dataset, "get_object_id") + assert hasattr(hats_dataset, "get_coord_ra") + assert hasattr(hats_dataset, "get_coord_dec") + assert hasattr(hats_dataset, "get_mag-r") diff --git a/tests/hyrax/test_hats_dataset_phase2.py b/tests/hyrax/test_hats_dataset_phase2.py new file mode 100644 index 000000000..733263d6f --- /dev/null +++ b/tests/hyrax/test_hats_dataset_phase2.py @@ -0,0 +1,132 @@ +import sys +import types + +import pandas as pd +import pytest + +from hyrax.datasets.hats_dataset import HATSPartitionIndex, HyraxHATSDataset + + +class _FakePartition: + def __init__(self, df): + self._df = df + + def __getitem__(self, columns): + return _FakePartition(self._df[columns]) + + def compute(self): + return self._df + + +class _FakeCountResult: + def __init__(self, counts): + self._counts = counts + + def compute(self): + return self._counts + + +class _FakeDDF: + def __init__(self, partitions): + self._partitions = partitions + self.columns = partitions[0].columns if partitions else [] + + def map_partitions(self, fn): + return _FakeCountResult([len(partition) for partition in self._partitions]) + + def get_partition(self, partition_id): + return _FakePartition(self._partitions[partition_id]) + + +class _FakeCatalog: + def __init__(self, partitions): + self._ddf = _FakeDDF(partitions) + self.columns = partitions[0].columns if partitions else [] + self.compute_calls = 0 + + def compute(self): + self.compute_calls += 1 + return pd.concat(self._ddf._partitions, ignore_index=True) + + +@pytest.fixture +def fake_catalog(): + p0 = pd.DataFrame({"object_id": [1, 2], "coord_ra": [11.0, 12.0], "mag-r": [20.1, 20.2]}) + p1 = pd.DataFrame({"object_id": [3, 4, 5], "coord_ra": [13.0, 14.0, 15.0], "mag-r": [20.3, 20.4, 20.5]}) + return _FakeCatalog([p0, p1]) + + +def _config_with_hats_options(): + return { + "data_request": { + "train": { + "data": { + "dataset_class": "HyraxHATSDataset", + "data_location": "/fake/path", + "fields": ["coord_ra"], + "primary_id_field": "object_id", + "hats": { + "bundle_size": 2, + "max_cached_bundles": 2, + "project_columns": "auto", + "strict_metadata": True, + }, + } + } + } + } + + +def test_partition_index_resolve(fake_catalog): + index = HATSPartitionIndex(fake_catalog) + + assert index.total_rows == 5 + assert index.resolve(0) == (0, 0) + assert index.resolve(1) == (0, 1) + assert index.resolve(2) == (1, 0) + assert index.resolve(4) == (1, 2) + + with pytest.raises(IndexError): + index.resolve(5) + + +def test_hats_dataset_prefers_open_catalog(monkeypatch, fake_catalog): + fake_lsdb = types.SimpleNamespace( + open_catalog=lambda _loc: fake_catalog, + read_hats=lambda _loc: pytest.fail("read_hats should not be called when open_catalog succeeds"), + ) + monkeypatch.setitem(sys.modules, "lsdb", fake_lsdb) + + dataset = HyraxHATSDataset(_config_with_hats_options(), data_location="/fake/path") + + assert len(dataset) == 5 + assert dataset.get_object_id(3) == 4 + assert getattr(dataset, "get_mag-r")(4) == pytest.approx(20.5) + assert fake_catalog.compute_calls == 0 + + +def test_hats_dataset_falls_back_to_read_hats(monkeypatch, fake_catalog): + def _boom(_loc): + raise RuntimeError("open failed") + + fake_lsdb = types.SimpleNamespace(open_catalog=_boom, read_hats=lambda _loc: fake_catalog) + monkeypatch.setitem(sys.modules, "lsdb", fake_lsdb) + + dataset = HyraxHATSDataset(_config_with_hats_options(), data_location="/fake/path") + assert len(dataset) == 5 + assert dataset.get_coord_ra(2) == pytest.approx(13.0) + + +def test_hats_dataset_uses_bundle_cache(monkeypatch, fake_catalog): + fake_lsdb = types.SimpleNamespace(open_catalog=lambda _loc: fake_catalog, read_hats=lambda _loc: fake_catalog) + monkeypatch.setitem(sys.modules, "lsdb", fake_lsdb) + + dataset = HyraxHATSDataset(_config_with_hats_options(), data_location="/fake/path") + + _ = dataset.get_coord_ra(0) + misses_after_first = dataset._accessor.cache.misses + _ = dataset.get_object_id(1) + hits_after_second = dataset._accessor.cache.hits + + assert misses_after_first == 1 + assert hits_after_second >= 1