HATS Phase‑2: lazy partition-aware dataset with bundle cache, docs, and tests

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
+}

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_<field>(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
+}