Minor fixes (#171)

Co-authored-by: Derek Jones <[email protected]>

src/nested_pandas/nestedframe/core.py

@@ -10,7 +10,6 @@
 from pandas._libs import lib
 from pandas._typing import Any, AnyAll, Axis, IndexLabel
 from pandas.api.extensions import no_default
-from pandas.api.types import is_bool_dtype
 from pandas.core.computation.expr import PARSERS, PandasExprVisitor
 
 from nested_pandas.nestedframe.utils import extract_nest_names
@@ -522,14 +521,11 @@ def query(self, expr: str, *, inplace: bool = False, **kwargs) -> NestedFrame |
         # to the nest and repack.  Otherwise, apply it to this instance as usual,
         # since it operated on the base attributes.
         if isinstance(result, _SeriesFromNest):
-            if not is_bool_dtype(result.dtype):
-                raise ValueError("Query condition must evaluate to a boolean Series")
-
             nest_name, flat_nest = result.nest_name, result.flat_nest
-
             # Reset index to "ordinal" like [0, 0, 0, 1, 1, 2, 2, 2]
-            flat_nest = flat_nest.set_index(self[nest_name].array.list_index)
-            query_result = result.set_axis(self[nest_name].array.list_index)
+            list_index = self[nest_name].array.get_list_index()
+            flat_nest = flat_nest.set_index(list_index)
+            query_result = result.set_axis(list_index)
             # Selecting flat values matching the query result
             new_flat_nest = flat_nest[query_result]
             new_df = self._set_filtered_flat_df(nest_name, new_flat_nest)
@@ -682,7 +678,7 @@ def dropna(
         if subset is not None:
             subset = [col.split(".")[-1] for col in subset]
         target_flat = self[target].nest.to_flat()
-        target_flat = target_flat.set_index(self[target].array.list_index)
+        target_flat = target_flat.set_index(self[target].array.get_list_index())
         if inplace:
             target_flat.dropna(
                 axis=axis,

src/nested_pandas/series/ext_array.py

@@ -648,8 +648,7 @@ def num_chunks(self) -> int:
         """Number of chunks in underlying pyarrow.ChunkedArray"""
         return self._chunked_array.num_chunks
 
-    @property
-    def list_index(self) -> np.ndarray:
+    def get_list_index(self) -> np.ndarray:
         """Keys mapping values to lists"""
         list_index = np.arange(len(self))
         return np.repeat(list_index, np.diff(self.list_offsets))