From ca81704fb7dc592e869dc593f0c6b41bc0e904b7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Mon, 28 Apr 2025 11:32:47 +0200 Subject: [PATCH 1/9] Speedup arith --- arrow-arith/src/arity.rs | 48 ++++++++++++++++++++-------------------- 1 file changed, 24 insertions(+), 24 deletions(-) diff --git a/arrow-arith/src/arity.rs b/arrow-arith/src/arity.rs index 9b3272abb617..544908081c5e 100644 --- a/arrow-arith/src/arity.rs +++ b/arrow-arith/src/arity.rs @@ -21,7 +21,7 @@ use arrow_array::builder::BufferBuilder; use arrow_array::*; use arrow_buffer::buffer::NullBuffer; use arrow_buffer::ArrowNativeType; -use arrow_buffer::{Buffer, MutableBuffer}; +use arrow_buffer::MutableBuffer; use arrow_data::ArrayData; use arrow_schema::ArrowError; @@ -125,12 +125,8 @@ where let nulls = NullBuffer::union(a.logical_nulls().as_ref(), b.logical_nulls().as_ref()); let values = a.values().iter().zip(b.values()).map(|(l, r)| op(*l, *r)); - // JUSTIFICATION - // Benefit - // ~60% speedup - // Soundness - // `values` is an iterator with a known size from a PrimitiveArray - let buffer = unsafe { Buffer::from_trusted_len_iter(values) }; + + let buffer: Vec<_> = values.collect(); Ok(PrimitiveArray::new(buffer.into(), nulls)) } @@ -251,14 +247,16 @@ where /// /// Return an error if the arrays have different lengths or /// the operation is under erroneous -pub fn try_binary( - a: A, - b: B, +pub fn try_binary( + a: &PrimitiveArray, + b: &PrimitiveArray, op: F, ) -> Result, ArrowError> where + A: ArrowPrimitiveType, + B: ArrowPrimitiveType, O: ArrowPrimitiveType, - F: Fn(A::Item, B::Item) -> Result, + F: Fn(A::Native, B::Native) -> Result, { if a.len() != b.len() { return Err(ArrowError::ComputeError( @@ -271,7 +269,7 @@ where let len = a.len(); if a.null_count() == 0 && b.null_count() == 0 { - try_binary_no_nulls(len, a, b, op) + try_binary_no_nulls(a, b, op) } else { let nulls = NullBuffer::union(a.logical_nulls().as_ref(), b.logical_nulls().as_ref()).unwrap(); @@ -369,23 +367,25 @@ fn create_union_null_buffer( /// This intentional inline(never) attribute helps LLVM optimize the loop. #[inline(never)] -fn try_binary_no_nulls( - len: usize, - a: A, - b: B, +fn try_binary_no_nulls( + a: &PrimitiveArray, + b: &PrimitiveArray, op: F, ) -> Result, ArrowError> where + A: ArrowPrimitiveType, + B: ArrowPrimitiveType, O: ArrowPrimitiveType, - F: Fn(A::Item, B::Item) -> Result, + B: ArrowPrimitiveType, + F: Fn(A::Native, B::Native) -> Result, { - let mut buffer = MutableBuffer::new(len * O::Native::get_byte_width()); - for idx in 0..len { - unsafe { - buffer.push_unchecked(op(a.value_unchecked(idx), b.value_unchecked(idx))?); - }; - } - Ok(PrimitiveArray::new(buffer.into(), None)) + let new_values = a + .values() + .into_iter() + .zip(b.values().into_iter()) + .map(|(l, r)| op(*l, *r)) + .collect::, ArrowError>>()?; + Ok(PrimitiveArray::new(new_values.into(), None)) } /// This intentional inline(never) attribute helps LLVM optimize the loop. From c31011caddd1820584017a5818d95e6058d48e78 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Tue, 29 Apr 2025 20:17:40 +0200 Subject: [PATCH 2/9] Speedup arith scalar --- arrow-arith/src/arity.rs | 2 -- arrow-array/src/array/primitive_array.rs | 7 +------ 2 files changed, 1 insertion(+), 8 deletions(-) diff --git a/arrow-arith/src/arity.rs b/arrow-arith/src/arity.rs index 544908081c5e..a4d74269f335 100644 --- a/arrow-arith/src/arity.rs +++ b/arrow-arith/src/arity.rs @@ -20,8 +20,6 @@ use arrow_array::builder::BufferBuilder; use arrow_array::*; use arrow_buffer::buffer::NullBuffer; -use arrow_buffer::ArrowNativeType; -use arrow_buffer::MutableBuffer; use arrow_data::ArrayData; use arrow_schema::ArrowError; diff --git a/arrow-array/src/array/primitive_array.rs b/arrow-array/src/array/primitive_array.rs index 57aa23bf9040..e6b272d685cf 100644 --- a/arrow-array/src/array/primitive_array.rs +++ b/arrow-array/src/array/primitive_array.rs @@ -828,12 +828,7 @@ impl PrimitiveArray { { let nulls = self.nulls().cloned(); let values = self.values().iter().map(|v| op(*v)); - // JUSTIFICATION - // Benefit - // ~60% speedup - // Soundness - // `values` is an iterator with a known size because arrays are sized. - let buffer = unsafe { Buffer::from_trusted_len_iter(values) }; + let buffer: Vec<_> = values.collect(); PrimitiveArray::new(buffer.into(), nulls) } From 9afc9408bb5759563f0e26ac8c0db106b36ed43e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Tue, 29 Apr 2025 20:23:31 +0200 Subject: [PATCH 3/9] Speedup arith --- arrow-array/src/array/primitive_array.rs | 15 +++++---------- 1 file changed, 5 insertions(+), 10 deletions(-) diff --git a/arrow-array/src/array/primitive_array.rs b/arrow-array/src/array/primitive_array.rs index e6b272d685cf..d0e5e563106c 100644 --- a/arrow-array/src/array/primitive_array.rs +++ b/arrow-array/src/array/primitive_array.rs @@ -729,10 +729,8 @@ impl PrimitiveArray { /// Creates a PrimitiveArray based on a constant value with `count` elements pub fn from_value(value: T::Native, count: usize) -> Self { - unsafe { - let val_buf = Buffer::from_trusted_len_iter((0..count).map(|_| value)); - Self::new(val_buf.into(), None) - } + let val_buf: Vec<_> = vec![value; count]; + Self::new(val_buf.into(), None) } /// Returns an iterator that returns the values of `array.value(i)` for an iterator with each element `i` @@ -1030,12 +1028,9 @@ impl PrimitiveArray { F: FnMut(U::Item) -> T::Native, { let nulls = left.logical_nulls(); - let buffer = unsafe { - // SAFETY: i in range 0..left.len() - let iter = (0..left.len()).map(|i| op(left.value_unchecked(i))); - // SAFETY: upper bound is trusted because `iter` is over a range - Buffer::from_trusted_len_iter(iter) - }; + let buffer: Vec<_> = (0..left.len()) + .map(|i| op(unsafe { left.value_unchecked(i) })) + .collect(); PrimitiveArray::new(buffer.into(), nulls) } From 675239eab68330e117c9b44130d52bc024ac9950 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Tue, 29 Apr 2025 21:20:36 +0200 Subject: [PATCH 4/9] Speedup arith --- arrow-arith/src/arity.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/arrow-arith/src/arity.rs b/arrow-arith/src/arity.rs index a4d74269f335..c842216b974d 100644 --- a/arrow-arith/src/arity.rs +++ b/arrow-arith/src/arity.rs @@ -122,7 +122,7 @@ where let nulls = NullBuffer::union(a.logical_nulls().as_ref(), b.logical_nulls().as_ref()); - let values = a.values().iter().zip(b.values()).map(|(l, r)| op(*l, *r)); + let values = a.values().into_iter().zip(b.values()).map(|(l, r)| op(*l, *r)); let buffer: Vec<_> = values.collect(); Ok(PrimitiveArray::new(buffer.into(), nulls)) From 5cc033e4abcf1297665bc2e263111835373d9471 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Tue, 29 Apr 2025 21:26:51 +0200 Subject: [PATCH 5/9] Speedup arith --- arrow-array/src/array/primitive_array.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/arrow-array/src/array/primitive_array.rs b/arrow-array/src/array/primitive_array.rs index d0e5e563106c..07f0532ccc9d 100644 --- a/arrow-array/src/array/primitive_array.rs +++ b/arrow-array/src/array/primitive_array.rs @@ -825,7 +825,7 @@ impl PrimitiveArray { F: Fn(T::Native) -> O::Native, { let nulls = self.nulls().cloned(); - let values = self.values().iter().map(|v| op(*v)); + let values = self.values().into_iter().map(|v| op(*v)); let buffer: Vec<_> = values.collect(); PrimitiveArray::new(buffer.into(), nulls) } From f556fa6d90b28904de51c95a04f8e30c83b24592 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Tue, 29 Apr 2025 21:44:55 +0200 Subject: [PATCH 6/9] Speedup arith --- arrow-arith/src/arity.rs | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/arrow-arith/src/arity.rs b/arrow-arith/src/arity.rs index c842216b974d..56dfb07e7d2f 100644 --- a/arrow-arith/src/arity.rs +++ b/arrow-arith/src/arity.rs @@ -122,7 +122,11 @@ where let nulls = NullBuffer::union(a.logical_nulls().as_ref(), b.logical_nulls().as_ref()); - let values = a.values().into_iter().zip(b.values()).map(|(l, r)| op(*l, *r)); + let values = a + .values() + .into_iter() + .zip(b.values()) + .map(|(l, r)| op(*l, *r)); let buffer: Vec<_> = values.collect(); Ok(PrimitiveArray::new(buffer.into(), nulls)) From 3f5f977762c73901fb59831aa44cef234aba1dc9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Wed, 30 Apr 2025 12:49:41 +0200 Subject: [PATCH 7/9] Speedup bitwise_unary_op_helper --- arrow-buffer/src/buffer/ops.rs | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/arrow-buffer/src/buffer/ops.rs b/arrow-buffer/src/buffer/ops.rs index c69e5c6deb10..a299f163caa4 100644 --- a/arrow-buffer/src/buffer/ops.rs +++ b/arrow-buffer/src/buffer/ops.rs @@ -103,11 +103,13 @@ where F: FnMut(u64) -> u64, { // reserve capacity and set length so we can get a typed view of u64 chunks - let mut result = - MutableBuffer::new(ceil(len_in_bits, 8)).with_bitset(len_in_bits / 64 * 8, false); + let mut result = MutableBuffer::new(ceil(len_in_bits, 8)); let left_chunks = left.bit_chunks(offset_in_bits, len_in_bits); + // SAFETY: `MutableBuffer::set_len` is sound because it is initalized right after + unsafe { result.set_len(left_chunks.iter().len() * 8) }; + let result_chunks = result.typed_data_mut::().iter_mut(); result_chunks From 7be4bc064451e8051b35fccc9786e982cbb480c0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Thu, 1 May 2025 10:10:52 +0200 Subject: [PATCH 8/9] Add back safety note --- arrow-array/src/array/primitive_array.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/arrow-array/src/array/primitive_array.rs b/arrow-array/src/array/primitive_array.rs index 07f0532ccc9d..c4f1228915de 100644 --- a/arrow-array/src/array/primitive_array.rs +++ b/arrow-array/src/array/primitive_array.rs @@ -1029,6 +1029,7 @@ impl PrimitiveArray { { let nulls = left.logical_nulls(); let buffer: Vec<_> = (0..left.len()) + // SAFETY: i is within bounds of the array .map(|i| op(unsafe { left.value_unchecked(i) })) .collect(); From 55cb9bb4833f75eef78e8a6b05012e2b0f60e81f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Thu, 1 May 2025 10:11:42 +0200 Subject: [PATCH 9/9] Add back safety note --- arrow-array/src/array/primitive_array.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/arrow-array/src/array/primitive_array.rs b/arrow-array/src/array/primitive_array.rs index c4f1228915de..2793f12babdf 100644 --- a/arrow-array/src/array/primitive_array.rs +++ b/arrow-array/src/array/primitive_array.rs @@ -1029,7 +1029,7 @@ impl PrimitiveArray { { let nulls = left.logical_nulls(); let buffer: Vec<_> = (0..left.len()) - // SAFETY: i is within bounds of the array + // SAFETY: i in range 0..left.len() .map(|i| op(unsafe { left.value_unchecked(i) })) .collect();