Cleanup uses of Array::data_ref (#3880) (#3918)

* Cleanup uses of Array::data_ref (#3880) * Further cleanup and fixes
apache · Mar 23, 2023 · de9a90b · de9a90b
1 parent 6af2bf6
commit de9a90b
Show file tree

Hide file tree

Showing 12 changed files with 170 additions and 299 deletions.
diff --git a/arrow-arith/src/arithmetic.rs b/arrow-arith/src/arithmetic.rs
@@ -103,14 +103,13 @@ fn math_checked_divide_op_on_iters<T, F>(
     left: impl Iterator<Item = Option<T::Native>>,
     right: impl Iterator<Item = Option<T::Native>>,
     op: F,
-    len: usize,
-    null_bit_buffer: Option<arrow_buffer::Buffer>,
+    nulls: Option<arrow_buffer::NullBuffer>,
 ) -> Result<PrimitiveArray<T>, ArrowError>
 where
     T: ArrowNumericType,
     F: Fn(T::Native, T::Native) -> Result<T::Native, ArrowError>,
 {
-    let buffer = if null_bit_buffer.is_some() {
+    let buffer = if nulls.is_some() {
         let values = left.zip(right).map(|(left, right)| {
             if let (Some(l), Some(r)) = (left, right) {
                 op(l, r)
@@ -130,18 +129,7 @@ where
         unsafe { arrow_buffer::Buffer::try_from_trusted_len_iter(values) }
     }?;
 
-    let data = unsafe {
-        arrow_data::ArrayData::new_unchecked(
-            T::DATA_TYPE,
-            len,
-            None,
-            null_bit_buffer,
-            0,
-            vec![buffer],
-            vec![],
-        )
-    };
-    Ok(PrimitiveArray::<T>::from(data))
+    Ok(PrimitiveArray::new(T::DATA_TYPE, buffer.into(), nulls))
 }
 
 /// Calculates the modulus operation `left % right` on two SIMD inputs.
@@ -284,20 +272,16 @@ where
     }
 
     // Create the combined `Bitmap`
-    let null_bit_buffer = arrow_data::bit_mask::combine_option_bitmap(
-        &[left.data_ref(), right.data_ref()],
-        left.len(),
-    );
+    let nulls = arrow_buffer::NullBuffer::union(left.nulls(), right.nulls());
 
     let lanes = T::lanes();
     let buffer_size = left.len() * std::mem::size_of::<T::Native>();
     let mut result =
         arrow_buffer::MutableBuffer::new(buffer_size).with_bitset(buffer_size, false);
 
-    match &null_bit_buffer {
+    match &nulls {
         Some(b) => {
-            // combine_option_bitmap returns a slice or new buffer starting at 0
-            let valid_chunks = b.bit_chunks(0, left.len());
+            let valid_chunks = b.inner().bit_chunks();
 
             // process data in chunks of 64 elements since we also get 64 bits of validity information at a time
 
@@ -372,18 +356,7 @@ where
         }
     }
 
-    let data = unsafe {
-        arrow_data::ArrayData::new_unchecked(
-            T::DATA_TYPE,
-            left.len(),
-            None,
-            null_bit_buffer,
-            0,
-            vec![result.into()],
-            vec![],
-        )
-    };
-    Ok(PrimitiveArray::<T>::from(data))
+    Ok(PrimitiveArray::new(T::DATA_TYPE, result.into(), nulls))
 }
 
 /// Applies $OP to $LEFT and $RIGHT which are two dictionaries which have (the same) key type $KT
@@ -628,10 +601,7 @@ where
         )));
     }
 
-    let null_bit_buffer = arrow_data::bit_mask::combine_option_bitmap(
-        &[left.data_ref(), right.data_ref()],
-        left.len(),
-    );
+    let nulls = arrow_buffer::NullBuffer::union(left.nulls(), right.nulls());
 
     // Safety justification: Since the inputs are valid Arrow arrays, all values are
     // valid indexes into the dictionary (which is verified during construction)
@@ -653,13 +623,7 @@ where
             .take_iter_unchecked(right.keys_iter())
     };
 
-    math_checked_divide_op_on_iters(
-        left_iter,
-        right_iter,
-        op,
-        left.len(),
-        null_bit_buffer,
-    )
+    math_checked_divide_op_on_iters(left_iter, right_iter, op, nulls)
 }
 
 #[cfg(feature = "dyn_arith_dict")]