diff --git a/src/nested_pandas/series/accessor.py b/src/nested_pandas/series/accessor.py index f6253bd..b7936bb 100644 --- a/src/nested_pandas/series/accessor.py +++ b/src/nested_pandas/series/accessor.py @@ -223,8 +223,11 @@ def query_flat(self, query: str) -> pd.Series: The filtered series. """ flat = self.to_flat().query(query) + if len(flat) == 0: - return pd.Series([], dtype=self._series.dtype) + return pd.Series( + [], dtype=self._series.dtype, index=pd.Index([], dtype=flat.index.dtype, name=flat.index.name) + ) return pack_sorted_df_into_struct(flat) def get_flat_index(self) -> pd.Index: diff --git a/tests/nested_pandas/series/test_accessor.py b/tests/nested_pandas/series/test_accessor.py index a4ac4db..b63c879 100644 --- a/tests/nested_pandas/series/test_accessor.py +++ b/tests/nested_pandas/series/test_accessor.py @@ -1,3 +1,4 @@ +import nested_pandas as npd import numpy as np import pandas as pd import pyarrow as pa @@ -487,6 +488,17 @@ def test_query_flat_empty_rows(): assert_series_equal(filtered, desired) +def test_query_flat_with_empty_result(): + """Make sure the index is properly set for empty result cases""" + base = npd.NestedFrame({"a": []}, index=pd.Index([], dtype=np.float64)) + nested = npd.NestedFrame({"b": []}, index=pd.Index([], dtype=np.float64)) + + ndf = base.add_nested(nested, "nested") + + res = ndf.nested.nest.query_flat("b > 2") + assert res.index.dtype == np.float64 + + @pytest.mark.parametrize( "df", [