Merge pull request #128 from lincc-frameworks/empty_query_fix

fix for empty query results

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:

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",
     [