stay packed for nested[['sub_column']]

Author: dougbrn

src/nested_pandas/series/accessor.py

@@ -508,9 +508,9 @@ def __getitem__(self, key: str | list[str]) -> NestedSeries:
         # on the number of fields requested and their dtypes
         if isinstance(key, list):
             new_array = self._series.array.view_fields(key)
-            if len(key) == 1 and not isinstance(new_array.dtype.field_dtype(key[0]), NestedDtype):
-                # If only one field is requested, return it as a pd.Series
-                return self._series[key[0]]
+            # if len(key) == 1 and not isinstance(new_array.dtype.field_dtype(key[0]), NestedDtype):
+            #    # If only one field is requested, return it as a pd.Series
+            #    return self._series[key[0]]
             return NestedSeries(new_array, index=self._series.index, name=self._series.name)
 
         # If the key is a single string, return the flat series for that field

tests/nested_pandas/nestedframe/test_nestedframe.py

@@ -182,6 +182,7 @@ def test_get_nested_columns():
 
     base = base.add_nested(nested, "nested")
 
+    # import pdb;pdb.set_trace()
     df = base[["a", "b", "nested.c"]]
     assert np.all(df.columns == ["a", "b", "nested"])
     assert df.dtypes["nested"].field_names == ["c"]

tests/nested_pandas/series/test_accessor.py

@@ -769,27 +769,6 @@ def test___getitem___series_masking():
         _ = nf["nested"].nest[mask[0:23]]
 
 
-def test___getitem___single_list_field():
-    """Test that the .nest["field"] works for a single list field."""
-    struct_array = pa.StructArray.from_arrays(
-        arrays=[
-            pa.array([np.array([1, 2, 3]), np.array([4, 5, 6])]),
-            pa.array([np.array([6, 4, 2]), np.array([1, 2, 3])]),
-        ],
-        names=["a", "b"],
-    )
-    series = NestedSeries(struct_array, dtype=NestedDtype(struct_array.type), index=[5, 7])
-    assert_series_equal(
-        series.nest[["a"]],
-        pd.Series(
-            data=[1, 2, 3, 4, 5, 6],
-            index=[5, 5, 5, 7, 7, 7],
-            name="a",
-            dtype=pd.ArrowDtype(pa.int64()),
-        ),
-    )
-
-
 def test___setitem__():
     """Test that the .nest["field"] = ... works for a single field."""
     struct_array = pa.StructArray.from_arrays(