list take: don't create null buffer if there are no nulls

pola-rs · Aug 5, 2021 · d29d7b6 · d29d7b6
1 parent 943157e
commit d29d7b6
Showing 1 changed file with 21 additions and 22 deletions.
diff --git a/polars/polars-core/src/chunked_array/kernels/take.rs b/polars/polars-core/src/chunked_array/kernels/take.rs
@@ -563,28 +563,27 @@ pub(crate) unsafe fn take_list_unchecked(
         .unwrap();
     let taken = taken.chunks()[0].clone();
 
-    // determine null buffer, which are a function of `values` and `indices`
-    let mut validity = MutableBitmap::with_capacity(indices.len());
-    validity.extend_constant(indices.len(), true);
-
-    {
-        offsets
-            .as_slice()
-            .windows(2)
-            .enumerate()
-            .for_each(|(i, window): (usize, &[i64])| {
-                if window[0] == window[1] {
-                    // offsets are equal, slot is null
-                    validity.set(i, false);
-                }
-            });
-    }
-    ListArray::from_data(
-        values.data_type().clone(),
-        offsets.into(),
-        taken,
-        Some(validity.into()),
-    )
+    let validity =
+        if values.null_count() > 0 || indices.null_count() > 0 {
+            // determine null buffer, which are a function of `values` and `indices`
+            let mut validity = MutableBitmap::with_capacity(indices.len());
+            validity.extend_constant(indices.len(), true);
+
+            {
+                offsets.as_slice().windows(2).enumerate().for_each(
+                    |(i, window): (usize, &[i64])| {
+                        if window[0] == window[1] {
+                            // offsets are equal, slot is null
+                            validity.set(i, false);
+                        }
+                    },
+                );
+            }
+            Some(validity.into())
+        } else {
+            None
+        };
+    ListArray::from_data(values.data_type().clone(), offsets.into(), taken, validity)
 }
 
 #[cfg(test)]