perf: Only use first element to sniff types in sequences in __getitem__ and filter (#2384)

* perf: Only use first element to sniff types in sequences in `__getitem__` and `filter`

* use is_list_of

Author: MarcoGorelli

narwhals/_arrow/dataframe.py

@@ -355,11 +355,7 @@ def __getitem__(
             return self._with_native(self.native.slice(start, stop - start))
 
         elif isinstance(item, Sequence) or is_numpy_array_1d(item):
-            if (
-                isinstance(item, Sequence)
-                and all(isinstance(x, str) for x in item)
-                and len(item) > 0
-            ):
+            if isinstance(item, Sequence) and len(item) > 0 and isinstance(item[0], str):
                 return self._with_native(self.native.select(cast("Indices", item)))
             if isinstance(item, Sequence) and len(item) == 0:
                 return self._with_native(self.native.slice(0, 0))

narwhals/_arrow/series.py

@@ -33,6 +33,7 @@
 from narwhals.utils import Implementation
 from narwhals.utils import generate_temporary_column_name
 from narwhals.utils import import_dtypes_module
+from narwhals.utils import is_list_of
 from narwhals.utils import not_implemented
 from narwhals.utils import validate_backend_version
 
@@ -306,9 +307,7 @@ def len(self: Self, *, _return_py_scalar: bool = True) -> int:
         return maybe_extract_py_scalar(len(self.native), _return_py_scalar)
 
     def filter(self: Self, predicate: ArrowSeries | list[bool | None]) -> Self:
-        if not (
-            isinstance(predicate, list) and all(isinstance(x, bool) for x in predicate)
-        ):
+        if not is_list_of(predicate, bool):
             _, other_native = extract_native(self, predicate)
         else:
             other_native = predicate

narwhals/_pandas_like/dataframe.py

@@ -331,11 +331,11 @@ def __getitem__(
                 return self._with_native(
                     self.native.__class__(), validate_column_names=False
                 )
-            if all(isinstance(x, int) for x in item[1]):  # type: ignore[var-annotated]
+            if isinstance(item[1][0], int):
                 return self._with_native(
                     self.native.iloc[item], validate_column_names=False
                 )
-            if all(isinstance(x, str) for x in item[1]):  # type: ignore[var-annotated]
+            if isinstance(item[1][0], str):
                 indexer = (
                     item[0],
                     self.native.columns.get_indexer(item[1]),
@@ -383,7 +383,7 @@ def __getitem__(
             return PandasLikeSeries.from_native(native_series, context=self)
 
         elif is_sequence_but_not_str(item) or is_numpy_array_1d(item):
-            if all(isinstance(x, str) for x in item) and len(item) > 0:
+            if len(item) > 0 and isinstance(item[0], str):
                 return self._with_native(
                     select_columns_by_name(
                         self.native,

narwhals/_pandas_like/series.py

@@ -31,6 +31,7 @@
 from narwhals.exceptions import InvalidOperationError
 from narwhals.utils import Implementation
 from narwhals.utils import import_dtypes_module
+from narwhals.utils import is_list_of
 from narwhals.utils import parse_version
 from narwhals.utils import validate_backend_version
 
@@ -376,9 +377,7 @@ def arg_max(self: Self) -> int:
     # Binary comparisons
 
     def filter(self: Self, predicate: Any) -> PandasLikeSeries:
-        if not (
-            isinstance(predicate, list) and all(isinstance(x, bool) for x in predicate)
-        ):
+        if not is_list_of(predicate, bool):
             _, other_native = align_and_extract_native(self, predicate)
         else:
             other_native = predicate

narwhals/dataframe.py

@@ -34,6 +34,7 @@
 from narwhals.utils import flatten
 from narwhals.utils import generate_repr
 from narwhals.utils import is_compliant_lazyframe
+from narwhals.utils import is_list_of
 from narwhals.utils import is_sequence_but_not_str
 from narwhals.utils import issue_deprecation_warning
 from narwhals.utils import parse_version
@@ -192,11 +193,9 @@ def filter(
         *predicates: IntoExpr | Iterable[IntoExpr] | list[bool],
         **constraints: Any,
     ) -> Self:
-        if not (
-            len(predicates) == 1
-            and isinstance(predicates[0], list)
-            and all(isinstance(x, bool) for x in predicates[0])
-        ):
+        if len(predicates) == 1 and is_list_of(predicates[0], bool):
+            predicate = predicates[0]
+        else:
             from narwhals.functions import col
 
             flat_predicates = flatten(predicates)
@@ -210,8 +209,6 @@ def filter(
             predicate = plx.all_horizontal(
                 *chain(compliant_predicates, compliant_constraints)
             )
-        else:
-            predicate = predicates[0]
         return self._with_compliant(self._compliant_frame.filter(predicate))
 
     def sort(
@@ -2786,10 +2783,7 @@ def filter(
             <BLANKLINE>
         """
         if (
-            len(predicates) == 1
-            and isinstance(predicates[0], list)
-            and all(isinstance(x, bool) for x in predicates[0])
-            and not constraints
+            len(predicates) == 1 and is_list_of(predicates[0], bool) and not constraints
         ):  # pragma: no cover
             msg = "`LazyFrame.filter` is not supported with Python boolean masks - use expressions instead."
             raise TypeError(msg)

narwhals/utils.py

@@ -1271,6 +1271,11 @@ def is_sequence_but_not_str(sequence: Any | Sequence[_T]) -> TypeIs[Sequence[_T]
     return isinstance(sequence, Sequence) and not isinstance(sequence, str)
 
 
+def is_list_of(obj: Any, tp: type[_T]) -> TypeIs[list[type[_T]]]:
+    # Check if an object is a list of `tp`, only sniffing the first element.
+    return bool(isinstance(obj, list) and obj and isinstance(obj[0], tp))
+
+
 def find_stacklevel() -> int:
     """Find the first place in the stack that is not inside narwhals.