Improve error message for pack_flat with NaN values in index (#450)

Author: hombit

src/nested_pandas/series/packer.py

@@ -93,12 +93,28 @@ def pack_flat(
     nested_pandas.series.dtype.NestedDtype : The dtype of the output series.
     nested_pandas.series.packer.pack_lists : Pack a dataframe of nested arrays.
     """
-
     if on is not None:
         df = df.set_index(on)
     # pandas knows when index is pre-sorted, so it would do nothing if it is already sorted
     sorted_flat = df.sort_index(kind="stable")
-    return pack_sorted_df_into_struct(sorted_flat, name=name)
+    try:
+        return pack_sorted_df_into_struct(sorted_flat, name=name)
+    except ValueError:
+        # Check if the error is due to NaN values and raise a more informative message
+        if any(sorted_flat.index.get_level_values(i).hasnans for i in range(sorted_flat.index.nlevels)):
+            if on is None:
+                raise ValueError(
+                    "The index contains NaN values. "
+                    "NaN values are not supported because they cannot be used for grouping rows. "
+                    "Please remove or fill NaN values before packing."
+                ) from None
+            cols = [on] if isinstance(on, str) else list(on)
+            raise ValueError(
+                f"Column(s) {cols} contain NaN values. "
+                "NaN values are not supported because they cannot be used for grouping rows. "
+                "Please remove or fill NaN values before packing."
+            ) from None
+        raise
 
 
 def pack_seq(

tests/nested_pandas/series/test_packer.py

@@ -598,3 +598,49 @@ def test_calculate_sorted_index_offsets_raises_when_not_sorted():
     index = pd.Index([1, 2, 1, 2, 3, 3, 4, 4, 4])
     with pytest.raises(ValueError):
         packer.calculate_sorted_index_offsets(index)
+
+
+@pytest.mark.parametrize(
+    "index",
+    [
+        pd.Index([1.0, 1.0, 2.0, 2.0, np.nan]),
+        pd.MultiIndex.from_arrays(([1, 1, 2, 2, 2], [1.0, 2.0, np.nan, 1.0, 2.0])),
+    ],
+)
+def test_pack_flat_raises_with_nan_in_index(index):
+    """Test pack_flat() raises informative error when index contains NaN values.
+
+    This is a regression test for https://github.com/lincc-frameworks/nested-pandas/issues/440
+    """
+    df = pd.DataFrame(
+        data={
+            "a": [1, 2, 3, 4, 5],
+            "b": [0, 1, 0, 1, 0],
+        },
+        index=index,
+    )
+    with pytest.raises(ValueError, match="The index contains NaN values"):
+        packer.pack_flat(df)
+
+
+@pytest.mark.parametrize(
+    "col_data,col_dtype",
+    [
+        ([1.0, 1.0, 2.0, 2.0, np.nan], None),
+        ([1.0, 1.0, 2.0, 2.0, None], pd.ArrowDtype(pa.float64())),
+    ],
+)
+def test_pack_flat_raises_with_nan_in_on_column(col_data, col_dtype):
+    """Test pack_flat() raises informative error when 'on' column contains NaN values.
+
+    This is a regression test for https://github.com/lincc-frameworks/nested-pandas/issues/440
+    """
+    df = pd.DataFrame(
+        data={
+            "a": [1, 2, 3, 4, 5],
+            "b": [0, 1, 0, 1, 0],
+            "c": pd.array(col_data, dtype=col_dtype),
+        },
+    )
+    with pytest.raises(ValueError, match=r"Column\(s\) \['c'\] contain NaN values"):
+        packer.pack_flat(df, on="c")