perf(rust, python): ensure set_at_idx is O(1) (#5977)

pola-rs · Jan 1, 2023 · 9c3a659 · 9c3a659
1 parent dc67fd0
commit 9c3a659
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 16 deletions.
diff --git a/polars/polars-ops/src/chunked_array/set.rs b/polars/polars-ops/src/chunked_array/set.rs
@@ -90,25 +90,22 @@ unsafe fn set_at_idx_impl<V, T: NativeType>(
             mut_validity.into()
         })
     } else {
-        let mut validity = MutableBitmap::default();
+        let mut null_idx = vec![];
         for (idx, val) in idx.iter().zip(values_iter) {
             match val {
-                Some(value) => {
-                    if validity.is_empty() {
-                        validity.extend_constant(len, true);
-                    }
-                    validity.set_unchecked(*idx as usize, true);
-                    *new_values_slice.get_unchecked_mut(*idx as usize) = value
-                }
+                Some(value) => *new_values_slice.get_unchecked_mut(*idx as usize) = value,
                 None => {
-                    if validity.is_empty() {
-                        validity.extend_constant(len, true);
-                    }
-                    validity.set_unchecked(*idx as usize, false)
+                    null_idx.push(*idx);
                 }
             }
         }
-        if !validity.is_empty() {
+        // only make a validity bitmap when null values are set
+        if !null_idx.is_empty() {
+            let mut validity = MutableBitmap::with_capacity(len);
+            validity.extend_constant(len, true);
+            for idx in null_idx {
+                validity.set_unchecked(idx as usize, false)
+            }
             arr.set_validity(Some(validity.into()))
         }
     }

diff --git a/py-polars/polars/internals/series/series.py b/py-polars/polars/internals/series/series.py
@@ -862,7 +862,11 @@ def __setitem__(
         key: int | Series | np.ndarray[Any, Any] | Sequence[object] | tuple[object],
         value: Any,
     ) -> None:
-        if isinstance(value, Sequence) and not isinstance(value, str):
+        # do the single idx as first branch as those are likely in a tight loop
+        if isinstance(key, int) and not isinstance(key, bool):
+            self.set_at_idx(key, value)
+            return None
+        elif isinstance(value, Sequence) and not isinstance(value, str):
             if self.is_numeric() or self.is_datelike():
                 self.set_at_idx(key, value)  # type: ignore[arg-type]
                 return None
@@ -889,8 +893,6 @@ def __setitem__(
         elif isinstance(key, (list, tuple)):
             s = wrap_s(sequence_to_pyseries("", key, dtype=UInt32))
             self.__setitem__(s, value)
-        elif isinstance(key, int) and not isinstance(key, bool):
-            self.__setitem__([key], value)
         else:
             raise ValueError(f'cannot use "{key}" for indexing')