From 3846bfcddfebf0361c0b947cc2dda71c56af4e8c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Dani=C3=ABl=20Heres?= <danielheres@gmail.com>
Date: Thu, 24 Oct 2024 14:46:57 +0200
Subject: [PATCH] Wip

---
 arrow-select/src/take.rs | 26 +++++++++++++++++---------
 1 file changed, 17 insertions(+), 9 deletions(-)

diff --git a/arrow-select/src/take.rs b/arrow-select/src/take.rs
index e2d99e86b170..d08f3edf61fd 100644
--- a/arrow-select/src/take.rs
+++ b/arrow-select/src/take.rs
@@ -473,15 +473,23 @@ fn take_bytes<T: ByteArrayType, IndexType: ArrowPrimitiveType>(
         }));
         nulls = None
     } else if indices.null_count() == 0 {
-        offsets.extend(indices.values().iter().map(|index| {
-            let index = index.as_usize();
-            if array.is_valid(index) {
-                let s: &[u8] = array.value(index).as_ref();
-                values.extend_from_slice(s.as_ref());
-            }
-            T::Offset::usize_as(values.len())
-        }));
-        nulls = Some(take_bits(array.nulls().unwrap().inner(), indices).into_inner());
+        let nulls_buf = take_bits(array.nulls().unwrap().inner(), indices);
+
+        offsets.extend(
+            indices
+                .values()
+                .iter()
+                .zip(nulls_buf.iter())
+                .map(|(index, valid)| {
+                    let index = index.as_usize();
+                    if valid {
+                        let s: &[u8] = array.value(index).as_ref();
+                        values.extend_from_slice(s.as_ref());
+                    }
+                    T::Offset::usize_as(values.len())
+                }),
+        );
+        nulls = Some(nulls_buf.into_inner());
     } else if array.null_count() == 0 {
         offsets.extend(indices.values().iter().enumerate().map(|(i, index)| {
             if indices.is_valid(i) {