Simplify NestedPandasDataset: use nested-pandas bracket operator, add docstrings

Agent-Logs-Url: https://github.com/lincc-frameworks/hyrax/sessions/637f35ea-20ee-4019-ba3e-8d6e2c3627c8

Co-authored-by: mtauraso <31012+mtauraso@users.noreply.github.com>

Author: Copilot

src/hyrax/datasets/nested_pandas_dataset.py

@@ -1,6 +1,5 @@
 from pathlib import Path
 from types import MethodType
-from typing import Any
 
 from hyrax.datasets.dataset_registry import HyraxDataset
 
@@ -21,7 +20,7 @@ def __init__(self, config: dict, data_location: Path | str | None = None):
         self._register_getters()
         super().__init__(config)
 
-    def _load_nested_frame(self, read_kwargs: dict[str, Any]):
+    def _load_nested_frame(self, read_kwargs: dict):
         try:
             import nested_pandas as npd
         except ImportError as err:
@@ -38,23 +37,10 @@ def _all_available_fields(self) -> list[str]:
             fields.extend(self.nested_frame.get_subcolumns())
         return fields
 
-    def _value_for_field(self, idx: int, field: str):
-        if "." not in field:
-            return self.nested_frame.iloc[idx][field]
-
-        nested_col, sub_col = field.split(".", maxsplit=1)
-        if nested_col not in self.nested_frame.columns:
-            return self.nested_frame.iloc[idx][field]
-
-        nested_value = self.nested_frame.iloc[idx][nested_col]
-        if nested_value is None:
-            return None
-        return nested_value[sub_col]
-
     def _register_getters(self) -> None:
         def _make_getter(field_name: str):
             def getter(self, idx, _field_name=field_name):
-                return self._value_for_field(idx, _field_name)
+                return self.nested_frame[_field_name].loc[self.nested_frame.index[idx]]
 
             return getter
 

tests/hyrax/test_nested_pandas_dataset.py

@@ -19,6 +19,7 @@ def _write_nested_lightcurve_parquet(path: Path) -> None:
 
 
 def test_nested_pandas_dataset_reads_parquet_and_exposes_nested_fields(tmp_path):
+    """Test that NestedPandasDataset reads a parquet file and exposes nested field getters."""
     data_path = tmp_path / "random_lightcurve.parquet"
     _write_nested_lightcurve_parquet(data_path)
 
@@ -34,6 +35,7 @@ def test_nested_pandas_dataset_reads_parquet_and_exposes_nested_fields(tmp_path)
 
 
 def test_nested_pandas_dataset_passes_read_kwargs(tmp_path):
+    """Test that read_kwargs are forwarded to nested_pandas.read_parquet."""
     data_path = tmp_path / "random_lightcurve.parquet"
     _write_nested_lightcurve_parquet(data_path)
 
@@ -54,5 +56,6 @@ def test_nested_pandas_dataset_passes_read_kwargs(tmp_path):
 
 
 def test_nested_pandas_dataset_requires_data_location():
+    """Test that omitting data_location raises a ValueError."""
     with pytest.raises(ValueError):
         NestedPandasDataset(config={"data_set": {}})