add reserve calls

pola-rs · Feb 27, 2025 · 398e25a · 398e25a
1 parent 521eb17
commit 398e25a
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 0 deletions.
diff --git a/crates/polars-compute/src/filter/mod.rs b/crates/polars-compute/src/filter/mod.rs
@@ -100,6 +100,7 @@ pub fn filter_with_bitmap(array: &dyn Array, mask: &Bitmap) -> Box<dyn Array> {
         _ => {
             let iter = SlicesIterator::new(mask);
             let mut mutable = make_builder(array.dtype());
+            mutable.reserve(iter.slots());
             iter.for_each(|(start, len)| {
                 mutable.subslice_extend(array, start, len, ShareStrategy::Always)
             });

diff --git a/crates/polars-compute/src/if_then_else/array.rs b/crates/polars-compute/src/if_then_else/array.rs
@@ -11,6 +11,7 @@ impl IfThenElseKernel for FixedSizeListArray {
         let inner_dt = if_true.dtype().inner_dtype().unwrap();
         let mut builder =
             FixedSizeListArrayBuilder::new(if_true.dtype().clone(), make_builder(inner_dt));
+        builder.reserve(mask.len());
         if_then_else_extend(
             &mut builder,
             mask,
@@ -30,6 +31,7 @@ impl IfThenElseKernel for FixedSizeListArray {
         let inner_dt = if_false.dtype().inner_dtype().unwrap();
         let mut builder =
             FixedSizeListArrayBuilder::new(if_false.dtype().clone(), make_builder(inner_dt));
+        builder.reserve(mask.len());
         if_then_else_extend(
             &mut builder,
             mask,
@@ -49,6 +51,7 @@ impl IfThenElseKernel for FixedSizeListArray {
         let inner_dt = if_true.dtype().inner_dtype().unwrap();
         let mut builder =
             FixedSizeListArrayBuilder::new(if_true.dtype().clone(), make_builder(inner_dt));
+        builder.reserve(mask.len());
         if_then_else_extend(
             &mut builder,
             mask,
@@ -72,6 +75,7 @@ impl IfThenElseKernel for FixedSizeListArray {
             std::iter::once(if_false).collect_arr_trusted_with_dtype(dtype.clone());
         let inner_dt = dtype.inner_dtype().unwrap();
         let mut builder = FixedSizeListArrayBuilder::new(dtype.clone(), make_builder(inner_dt));
+        builder.reserve(mask.len());
         if_then_else_extend(
             &mut builder,
             mask,

diff --git a/crates/polars-compute/src/if_then_else/list.rs b/crates/polars-compute/src/if_then_else/list.rs
@@ -10,6 +10,7 @@ impl IfThenElseKernel for ListArray<i64> {
     fn if_then_else(mask: &Bitmap, if_true: &Self, if_false: &Self) -> Self {
         let inner_dt = if_true.dtype().inner_dtype().unwrap();
         let mut builder = ListArrayBuilder::new(if_true.dtype().clone(), make_builder(inner_dt));
+        builder.reserve(mask.len());
         if_then_else_extend(
             &mut builder,
             mask,
@@ -28,6 +29,7 @@ impl IfThenElseKernel for ListArray<i64> {
             std::iter::once(if_true).collect_arr_trusted_with_dtype(if_false.dtype().clone());
         let inner_dt = if_false.dtype().inner_dtype().unwrap();
         let mut builder = ListArrayBuilder::new(if_false.dtype().clone(), make_builder(inner_dt));
+        builder.reserve(mask.len());
         if_then_else_extend(
             &mut builder,
             mask,
@@ -46,6 +48,7 @@ impl IfThenElseKernel for ListArray<i64> {
             std::iter::once(if_false).collect_arr_trusted_with_dtype(if_true.dtype().clone());
         let inner_dt = if_true.dtype().inner_dtype().unwrap();
         let mut builder = ListArrayBuilder::new(if_true.dtype().clone(), make_builder(inner_dt));
+        builder.reserve(mask.len());
         if_then_else_extend(
             &mut builder,
             mask,
@@ -69,6 +72,7 @@ impl IfThenElseKernel for ListArray<i64> {
             std::iter::once(if_false).collect_arr_trusted_with_dtype(dtype.clone());
         let inner_dt = dtype.inner_dtype().unwrap();
         let mut builder = ListArrayBuilder::new(dtype.clone(), make_builder(inner_dt));
+        builder.reserve(mask.len());
         if_then_else_extend(
             &mut builder,
             mask,