From bfe278e30d83ca19b78a707ba216cb47605903de Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Sun, 11 Dec 2022 11:15:29 -0800 Subject: [PATCH 1/4] Provide into_builder for bytearray --- arrow-array/src/array/byte_array.rs | 86 +++++++++++++++++++ arrow-array/src/array/string_array.rs | 37 ++++++++ .../src/builder/generic_bytes_builder.rs | 55 +++++++++++- arrow-array/src/builder/primitive_builder.rs | 4 +- 4 files changed, 179 insertions(+), 3 deletions(-) diff --git a/arrow-array/src/array/byte_array.rs b/arrow-array/src/array/byte_array.rs index eb528384eace..fb5eeb9fac4e 100644 --- a/arrow-array/src/array/byte_array.rs +++ b/arrow-array/src/array/byte_array.rs @@ -16,6 +16,7 @@ // under the License. use crate::array::{empty_offsets, print_long_array}; +use crate::builder::GenericByteBuilder; use crate::iterator::ArrayIter; use crate::raw_pointer::RawPtrBox; use crate::types::bytes::ByteArrayNativeType; @@ -139,6 +140,91 @@ impl GenericByteArray { pub fn iter(&self) -> ArrayIter<&Self> { ArrayIter::new(self) } + + /// Returns `GenericByteBuilder` of this byte array for mutating its values if the underlying + /// offset and data buffers are not shared by others. + pub fn into_builder(self) -> Result, Self> { + let len = self.len(); + let null_bit_buffer = self + .data + .null_buffer() + .map(|b| b.bit_slice(self.data.offset(), len)); + + let element_len = std::mem::size_of::(); + let offset_buffer = self.data.buffers()[0] + .slice_with_length(self.data.offset() * element_len, (len + 1) * element_len); + + let element_len = std::mem::size_of::(); + let value_len = + T::Offset::as_usize(self.value_offsets()[len] - self.value_offsets()[0]); + let value_buffer = self.data.buffers()[1] + .slice_with_length(self.data.offset() * element_len, value_len * element_len); + + drop(self.data); + + let try_mutable_null_buffer = match null_bit_buffer { + None => Ok(None), + Some(null_buffer) => { + // Null buffer exists, tries to make it mutable + null_buffer.into_mutable().map(Some) + } + }; + + let try_mutable_buffers = match try_mutable_null_buffer { + Ok(mutable_null_buffer) => { + // Got mutable null buffer, tries to get mutable value buffer + let try_mutable_offset_buffer = offset_buffer.into_mutable(); + let try_mutable_value_buffer = value_buffer.into_mutable(); + + // try_mutable_offset_buffer.map(...).map_err(...) doesn't work as the compiler complains + // mutable_null_buffer is moved into map closure. + match (try_mutable_offset_buffer, try_mutable_value_buffer) { + (Ok(mutable_offset_buffer), Ok(mutable_value_buffer)) => { + Ok(GenericByteBuilder::::new_from_buffer( + mutable_offset_buffer, + mutable_value_buffer, + mutable_null_buffer, + )) + } + (Ok(mutable_offset_buffer), Err(value_buffer)) => Err(( + mutable_offset_buffer.into(), + value_buffer, + mutable_null_buffer.map(|b| b.into()), + )), + (Err(offset_buffer), Ok(mutable_value_buffer)) => Err(( + offset_buffer, + mutable_value_buffer.into(), + mutable_null_buffer.map(|b| b.into()), + )), + (Err(offset_buffer), Err(value_buffer)) => Err(( + offset_buffer, + value_buffer, + mutable_null_buffer.map(|b| b.into()), + )), + } + } + Err(mutable_null_buffer) => { + // Unable to get mutable null buffer + Err((offset_buffer, value_buffer, Some(mutable_null_buffer))) + } + }; + + match try_mutable_buffers { + Ok(builder) => Ok(builder), + Err((offset_buffer, value_buffer, null_bit_buffer)) => { + let builder = ArrayData::builder(T::DATA_TYPE) + .len(len) + .add_buffer(offset_buffer) + .add_buffer(value_buffer) + .null_bit_buffer(null_bit_buffer); + + let array_data = unsafe { builder.build_unchecked() }; + let array = GenericByteArray::::from(array_data); + + Err(array) + } + } + } } impl std::fmt::Debug for GenericByteArray { diff --git a/arrow-array/src/array/string_array.rs b/arrow-array/src/array/string_array.rs index c8db589e3c28..0ea5368f9499 100644 --- a/arrow-array/src/array/string_array.rs +++ b/arrow-array/src/array/string_array.rs @@ -285,6 +285,7 @@ mod tests { use crate::builder::{ListBuilder, StringBuilder}; use arrow_buffer::Buffer; use arrow_schema::Field; + use num::cast::AsPrimitive; #[test] fn test_string_array_from_u8_slice() { @@ -697,4 +698,40 @@ mod tests { assert_eq!(string.len(), 0); assert_eq!(string.value_offsets(), &[0]); } + + #[test] + fn test_into_builder() { + let array: StringArray = vec!["hello", "arrow"].into(); + + // Modify values + let mut builder = array.into_builder().unwrap(); + + let (offset_slice, value_slice, _) = builder.slices_mut(); + + assert_eq!(offset_slice, &[0, 5, 10]); + + let expected_slice = "helloarrow".as_bytes(); + assert_eq!(value_slice, expected_slice); + + value_slice[0] = 'H'.as_(); + value_slice[1] = 'E'.as_(); + value_slice[2] = 'L'.as_(); + value_slice[3] = 'L'.as_(); + value_slice[4] = 'O'.as_(); + + let expected: StringArray = vec!["HELLO", "arrow"].into(); + let array = builder.finish(); + assert_eq!(expected, array); + + // Modify offsets + let mut builder = array.into_builder().unwrap(); + + let (offset_slice, _, _) = builder.slices_mut(); + + offset_slice[1] = 7; + + let expected: StringArray = vec!["HELLOar", "row"].into(); + let array = builder.finish(); + assert_eq!(expected, array); + } } diff --git a/arrow-array/src/builder/generic_bytes_builder.rs b/arrow-array/src/builder/generic_bytes_builder.rs index 9f9078c708c8..a7fb51795cdd 100644 --- a/arrow-array/src/builder/generic_bytes_builder.rs +++ b/arrow-array/src/builder/generic_bytes_builder.rs @@ -19,7 +19,7 @@ use crate::builder::null_buffer_builder::NullBufferBuilder; use crate::builder::{ArrayBuilder, BufferBuilder, UInt8BufferBuilder}; use crate::types::{ByteArrayType, GenericBinaryType, GenericStringType}; use crate::{ArrayRef, GenericByteArray, OffsetSizeTrait}; -use arrow_buffer::{ArrowNativeType, Buffer}; +use arrow_buffer::{ArrowNativeType, Buffer, MutableBuffer}; use arrow_data::ArrayDataBuilder; use std::any::Any; use std::sync::Arc; @@ -53,6 +53,30 @@ impl GenericByteBuilder { } } + /// Creates a new primitive array builder from buffers + pub fn new_from_buffer( + offsets_buffer: MutableBuffer, + value_buffer: MutableBuffer, + null_buffer: Option, + ) -> Self { + let offsets_builder = BufferBuilder::::new_from_buffer(offsets_buffer); + let value_builder = BufferBuilder::::new_from_buffer(value_buffer); + + let null_buffer_builder = null_buffer + .map(|buffer| { + NullBufferBuilder::new_from_buffer(buffer, offsets_builder.len() - 1) + }) + .unwrap_or_else(|| { + NullBufferBuilder::new_with_len(offsets_builder.len() - 1) + }); + + Self { + offsets_builder, + value_builder, + null_buffer_builder, + } + } + /// Appends a value into the builder. #[inline] pub fn append_value(&mut self, value: impl AsRef) { @@ -118,10 +142,39 @@ impl GenericByteBuilder { self.value_builder.as_slice() } + /// Returns the current values buffer as a mutable slice + pub fn values_slice_mut(&mut self) -> &mut [u8] { + self.value_builder.as_slice_mut() + } + /// Returns the current offsets buffer as a slice pub fn offsets_slice(&self) -> &[T::Offset] { self.offsets_builder.as_slice() } + + /// Returns the current offsets buffer as a mutable slice + pub fn offsets_slice_mut(&mut self) -> &mut [T::Offset] { + self.offsets_builder.as_slice_mut() + } + + /// Returns the current null buffer as a slice + pub fn validity_slice(&self) -> Option<&[u8]> { + self.null_buffer_builder.as_slice() + } + + /// Returns the current null buffer as a mutable slice + pub fn validity_slice_mut(&mut self) -> Option<&mut [u8]> { + self.null_buffer_builder.as_slice_mut() + } + + /// Returns the current offsets, values buffer and null buffer as a slice + pub fn slices_mut(&mut self) -> (&mut [T::Offset], &mut [u8], Option<&mut [u8]>) { + ( + self.offsets_builder.as_slice_mut(), + self.value_builder.as_slice_mut(), + self.null_buffer_builder.as_slice_mut(), + ) + } } impl std::fmt::Debug for GenericByteBuilder { diff --git a/arrow-array/src/builder/primitive_builder.rs b/arrow-array/src/builder/primitive_builder.rs index fa1dc3ad1264..f3f3f3728db9 100644 --- a/arrow-array/src/builder/primitive_builder.rs +++ b/arrow-array/src/builder/primitive_builder.rs @@ -286,12 +286,12 @@ impl PrimitiveBuilder { self.values_builder.as_slice_mut() } - /// Returns the current values buffer as a slice + /// Returns the current null buffer as a slice pub fn validity_slice(&self) -> Option<&[u8]> { self.null_buffer_builder.as_slice() } - /// Returns the current values buffer as a mutable slice + /// Returns the current null buffer as a mutable slice pub fn validity_slice_mut(&mut self) -> Option<&mut [u8]> { self.null_buffer_builder.as_slice_mut() } From 453bd3508c051649543b3f4fa856245fba54d233 Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Tue, 20 Dec 2022 17:03:39 -0800 Subject: [PATCH 2/4] For review --- arrow-array/src/array/byte_array.rs | 4 ++-- arrow-array/src/array/string_array.rs | 11 +++++++++++ arrow-array/src/builder/generic_bytes_builder.rs | 13 ++++++------- 3 files changed, 19 insertions(+), 9 deletions(-) diff --git a/arrow-array/src/array/byte_array.rs b/arrow-array/src/array/byte_array.rs index fb5eeb9fac4e..2cb04efb8e89 100644 --- a/arrow-array/src/array/byte_array.rs +++ b/arrow-array/src/array/byte_array.rs @@ -179,13 +179,13 @@ impl GenericByteArray { // try_mutable_offset_buffer.map(...).map_err(...) doesn't work as the compiler complains // mutable_null_buffer is moved into map closure. match (try_mutable_offset_buffer, try_mutable_value_buffer) { - (Ok(mutable_offset_buffer), Ok(mutable_value_buffer)) => { + (Ok(mutable_offset_buffer), Ok(mutable_value_buffer)) => unsafe { Ok(GenericByteBuilder::::new_from_buffer( mutable_offset_buffer, mutable_value_buffer, mutable_null_buffer, )) - } + }, (Ok(mutable_offset_buffer), Err(value_buffer)) => Err(( mutable_offset_buffer.into(), value_buffer, diff --git a/arrow-array/src/array/string_array.rs b/arrow-array/src/array/string_array.rs index 0ea5368f9499..a39dc57c00c4 100644 --- a/arrow-array/src/array/string_array.rs +++ b/arrow-array/src/array/string_array.rs @@ -734,4 +734,15 @@ mod tests { let array = builder.finish(); assert_eq!(expected, array); } + + #[test] + fn test_into_builder_err() { + let array: StringArray = vec!["hello", "arrow"].into(); + + // Clone it, so we cannot get a mutable builder back + let shared_array = array.clone(); + + let err_return = array.into_builder().unwrap_err(); + assert_eq!(&err_return, &shared_array); + } } diff --git a/arrow-array/src/builder/generic_bytes_builder.rs b/arrow-array/src/builder/generic_bytes_builder.rs index a7fb51795cdd..6c005db72f40 100644 --- a/arrow-array/src/builder/generic_bytes_builder.rs +++ b/arrow-array/src/builder/generic_bytes_builder.rs @@ -53,8 +53,12 @@ impl GenericByteBuilder { } } - /// Creates a new primitive array builder from buffers - pub fn new_from_buffer( + /// Creates a new [`GenericByteBuilder`] from buffers. + /// + /// # Safety + /// This doesn't verify buffer contents as it assumes the buffers are from existing and + /// valid [`GenericByteArray`]. + pub unsafe fn new_from_buffer( offsets_buffer: MutableBuffer, value_buffer: MutableBuffer, null_buffer: Option, @@ -152,11 +156,6 @@ impl GenericByteBuilder { self.offsets_builder.as_slice() } - /// Returns the current offsets buffer as a mutable slice - pub fn offsets_slice_mut(&mut self) -> &mut [T::Offset] { - self.offsets_builder.as_slice_mut() - } - /// Returns the current null buffer as a slice pub fn validity_slice(&self) -> Option<&[u8]> { self.null_buffer_builder.as_slice() From 5f09047195779ee2fb038ff0b93a6ecad678b278 Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Wed, 28 Dec 2022 15:31:55 -0800 Subject: [PATCH 3/4] Remove slices_mut --- arrow-array/src/array/string_array.rs | 15 +-------------- arrow-array/src/builder/generic_bytes_builder.rs | 9 --------- 2 files changed, 1 insertion(+), 23 deletions(-) diff --git a/arrow-array/src/array/string_array.rs b/arrow-array/src/array/string_array.rs index a39dc57c00c4..12d48950adda 100644 --- a/arrow-array/src/array/string_array.rs +++ b/arrow-array/src/array/string_array.rs @@ -706,9 +706,7 @@ mod tests { // Modify values let mut builder = array.into_builder().unwrap(); - let (offset_slice, value_slice, _) = builder.slices_mut(); - - assert_eq!(offset_slice, &[0, 5, 10]); + let value_slice = builder.values_slice_mut(); let expected_slice = "helloarrow".as_bytes(); assert_eq!(value_slice, expected_slice); @@ -722,17 +720,6 @@ mod tests { let expected: StringArray = vec!["HELLO", "arrow"].into(); let array = builder.finish(); assert_eq!(expected, array); - - // Modify offsets - let mut builder = array.into_builder().unwrap(); - - let (offset_slice, _, _) = builder.slices_mut(); - - offset_slice[1] = 7; - - let expected: StringArray = vec!["HELLOar", "row"].into(); - let array = builder.finish(); - assert_eq!(expected, array); } #[test] diff --git a/arrow-array/src/builder/generic_bytes_builder.rs b/arrow-array/src/builder/generic_bytes_builder.rs index 6c005db72f40..3dc48b2b1d1b 100644 --- a/arrow-array/src/builder/generic_bytes_builder.rs +++ b/arrow-array/src/builder/generic_bytes_builder.rs @@ -165,15 +165,6 @@ impl GenericByteBuilder { pub fn validity_slice_mut(&mut self) -> Option<&mut [u8]> { self.null_buffer_builder.as_slice_mut() } - - /// Returns the current offsets, values buffer and null buffer as a slice - pub fn slices_mut(&mut self) -> (&mut [T::Offset], &mut [u8], Option<&mut [u8]>) { - ( - self.offsets_builder.as_slice_mut(), - self.value_builder.as_slice_mut(), - self.null_buffer_builder.as_slice_mut(), - ) - } } impl std::fmt::Debug for GenericByteBuilder { From d564b4ce29a79d6639313a0545bf4c3361cdc90d Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Thu, 29 Dec 2022 10:32:33 -0800 Subject: [PATCH 4/4] Modify test and remove values_slice_mut --- arrow-array/src/array/string_array.rs | 16 +++------------- arrow-array/src/builder/generic_bytes_builder.rs | 5 ----- 2 files changed, 3 insertions(+), 18 deletions(-) diff --git a/arrow-array/src/array/string_array.rs b/arrow-array/src/array/string_array.rs index 12d48950adda..4a4152adc678 100644 --- a/arrow-array/src/array/string_array.rs +++ b/arrow-array/src/array/string_array.rs @@ -285,7 +285,6 @@ mod tests { use crate::builder::{ListBuilder, StringBuilder}; use arrow_buffer::Buffer; use arrow_schema::Field; - use num::cast::AsPrimitive; #[test] fn test_string_array_from_u8_slice() { @@ -703,21 +702,12 @@ mod tests { fn test_into_builder() { let array: StringArray = vec!["hello", "arrow"].into(); - // Modify values + // Append values let mut builder = array.into_builder().unwrap(); - let value_slice = builder.values_slice_mut(); + builder.append_value("rust"); - let expected_slice = "helloarrow".as_bytes(); - assert_eq!(value_slice, expected_slice); - - value_slice[0] = 'H'.as_(); - value_slice[1] = 'E'.as_(); - value_slice[2] = 'L'.as_(); - value_slice[3] = 'L'.as_(); - value_slice[4] = 'O'.as_(); - - let expected: StringArray = vec!["HELLO", "arrow"].into(); + let expected: StringArray = vec!["hello", "arrow", "rust"].into(); let array = builder.finish(); assert_eq!(expected, array); } diff --git a/arrow-array/src/builder/generic_bytes_builder.rs b/arrow-array/src/builder/generic_bytes_builder.rs index 3dc48b2b1d1b..195628f4712f 100644 --- a/arrow-array/src/builder/generic_bytes_builder.rs +++ b/arrow-array/src/builder/generic_bytes_builder.rs @@ -146,11 +146,6 @@ impl GenericByteBuilder { self.value_builder.as_slice() } - /// Returns the current values buffer as a mutable slice - pub fn values_slice_mut(&mut self) -> &mut [u8] { - self.value_builder.as_slice_mut() - } - /// Returns the current offsets buffer as a slice pub fn offsets_slice(&self) -> &[T::Offset] { self.offsets_builder.as_slice()