Added convenience accessor array.get (jorgecarleitao#1416)

ritchie46 · Apr 5, 2023 · f87e9d2 · f87e9d2
1 parent 9f48859
commit f87e9d2
Show file tree

Hide file tree

Showing 6 changed files with 78 additions and 0 deletions.
diff --git a/src/array/binary/mod.rs b/src/array/binary/mod.rs
@@ -151,6 +151,19 @@ impl<O: Offset> BinaryArray<O> {
         self.values.get_unchecked(start..end)
     }
 
+    /// Returns the element at index `i` or `None` if it is null
+    /// # Panics
+    /// iff `i >= self.len()`
+    #[inline]
+    pub fn get(&self, i: usize) -> Option<&[u8]> {
+        if !self.is_null(i) {
+            // soundness: Array::is_null panics if i >= self.len
+            unsafe { Some(self.value_unchecked(i)) }
+        } else {
+            None
+        }
+    }
+
     /// Returns the [`DataType`] of this array.
     #[inline]
     pub fn data_type(&self) -> &DataType {

diff --git a/src/array/boolean/mod.rs b/src/array/boolean/mod.rs
@@ -143,6 +143,19 @@ impl BooleanArray {
         self.values.get_bit_unchecked(i)
     }
 
+    /// Returns the element at index `i` or `None` if it is null
+    /// # Panics
+    /// iff `i >= self.len()`
+    #[inline]
+    pub fn get(&self, i: usize) -> Option<bool> {
+        if !self.is_null(i) {
+            // soundness: Array::is_null panics if i >= self.len
+            unsafe { Some(self.value_unchecked(i)) }
+        } else {
+            None
+        }
+    }
+
     /// Slices this [`BooleanArray`].
     /// # Implementation
     /// This operation is `O(1)` as it amounts to increase up to two ref counts.

diff --git a/src/array/fixed_size_binary/mod.rs b/src/array/fixed_size_binary/mod.rs
@@ -157,6 +157,19 @@ impl FixedSizeBinaryArray {
             .get_unchecked(i * self.size..(i + 1) * self.size)
     }
 
+    /// Returns the element at index `i` or `None` if it is null
+    /// # Panics
+    /// iff `i >= self.len()`
+    #[inline]
+    pub fn get(&self, i: usize) -> Option<&[u8]> {
+        if !self.is_null(i) {
+            // soundness: Array::is_null panics if i >= self.len
+            unsafe { Some(self.value_unchecked(i)) }
+        } else {
+            None
+        }
+    }
+
     /// Returns a new [`FixedSizeBinaryArray`] with a different logical type.
     /// This is `O(1)`.
     /// # Panics

diff --git a/src/array/fixed_size_list/mod.rs b/src/array/fixed_size_list/mod.rs
@@ -166,6 +166,19 @@ impl FixedSizeListArray {
     pub unsafe fn value_unchecked(&self, i: usize) -> Box<dyn Array> {
         self.values.sliced_unchecked(i * self.size, self.size)
     }
+
+    /// Returns the element at index `i` or `None` if it is null
+    /// # Panics
+    /// iff `i >= self.len()`
+    #[inline]
+    pub fn get(&self, i: usize) -> Option<Box<dyn Array>> {
+        if !self.is_null(i) {
+            // soundness: Array::is_null panics if i >= self.len
+            unsafe { Some(self.value_unchecked(i)) }
+        } else {
+            None
+        }
+    }
 }
 
 impl FixedSizeListArray {

diff --git a/src/array/primitive/mod.rs b/src/array/primitive/mod.rs
@@ -196,6 +196,19 @@ impl<T: NativeType> PrimitiveArray<T> {
         *self.values.get_unchecked(i)
     }
 
+    /// Returns the element at index `i` or `None` if it is null
+    /// # Panics
+    /// iff `i >= self.len()`
+    #[inline]
+    pub fn get(&self, i: usize) -> Option<T> {
+        if !self.is_null(i) {
+            // soundness: Array::is_null panics if i >= self.len
+            unsafe { Some(self.value_unchecked(i)) }
+        } else {
+            None
+        }
+    }
+
     /// Slices this [`PrimitiveArray`] by an offset and length.
     /// # Implementation
     /// This operation is `O(1)`.

diff --git a/src/array/utf8/mod.rs b/src/array/utf8/mod.rs
@@ -171,6 +171,19 @@ impl<O: Offset> Utf8Array<O> {
         std::str::from_utf8_unchecked(slice)
     }
 
+    /// Returns the element at index `i` or `None` if it is null
+    /// # Panics
+    /// iff `i >= self.len()`
+    #[inline]
+    pub fn get(&self, i: usize) -> Option<&str> {
+        if !self.is_null(i) {
+            // soundness: Array::is_null panics if i >= self.len
+            unsafe { Some(self.value_unchecked(i)) }
+        } else {
+            None
+        }
+    }
+
     /// Returns the [`DataType`] of this array.
     #[inline]
     pub fn data_type(&self) -> &DataType {