diff --git a/arrow-row/src/lib.rs b/arrow-row/src/lib.rs index 83ebe79f726e..8ed2debb6a4b 100644 --- a/arrow-row/src/lib.rs +++ b/arrow-row/src/lib.rs @@ -5446,12 +5446,12 @@ mod tests { } #[test] - fn test_values_buffer_smaller_when_utf8_validation_disabled() { - fn get_values_buffer_len(col: ArrayRef) -> (usize, usize) { + fn test_utf8_validation_doesnt_affect_values_buffer_size() { + fn assert_values_buffer_lens(col: ArrayRef) -> usize { // 1. Convert cols into rows let converter = RowConverter::new(vec![SortField::new(DataType::Utf8View)]).unwrap(); - // 2a. Convert rows into colsa (validate_utf8 = false) + // 2a. Convert rows into cols (validate_utf8 = false) let rows = converter.convert_columns(&[col]).unwrap(); let converted = converter.convert_rows(&rows).unwrap(); let unchecked_values_len = converted[0].as_string_view().data_buffers()[0].len(); @@ -5463,7 +5463,9 @@ mod tests { .convert_rows(rows.iter().map(|b| parser.parse(b.expect("valid bytes")))) .unwrap(); let checked_values_len = converted[0].as_string_view().data_buffers()[0].len(); - (unchecked_values_len, checked_values_len) + // Regardless of utf8 validation flag, we should always have minimal data in buffers + assert_eq!(unchecked_values_len, checked_values_len); + checked_values_len } // Case1. StringViewArray with inline strings @@ -5474,22 +5476,18 @@ mod tests { Some("tiny"), // short(4) ])) as ArrayRef; - let (unchecked_values_len, checked_values_len) = get_values_buffer_len(col); + let values_len = assert_values_buffer_lens(col); // Since there are no long (>12) strings, len of values buffer is 0 - assert_eq!(unchecked_values_len, 0); - // When utf8 validation enabled, values buffer includes inline strings (5+5+4) - assert_eq!(checked_values_len, 14); + assert_eq!(values_len, 0); // Case2. StringViewArray with long(>12) strings let col = Arc::new(StringViewArray::from_iter([ - Some("this is a very long string over 12 bytes"), - Some("another long string to test the buffer"), + Some("1234567890123"), // 13 + Some("12345678901234"), // 14 ])) as ArrayRef; - let (unchecked_values_len, checked_values_len) = get_values_buffer_len(col); - // Since there are no inline strings, expected length of values buffer is the same - assert!(unchecked_values_len > 0); - assert_eq!(unchecked_values_len, checked_values_len); + let values_len = assert_values_buffer_lens(col); + assert_eq!(values_len, 13 + 14); // Case3. StringViewArray with both short and long strings let col = Arc::new(StringViewArray::from_iter([ @@ -5499,10 +5497,9 @@ mod tests { Some("short"), // 5 (short) ])) as ArrayRef; - let (unchecked_values_len, checked_values_len) = get_values_buffer_len(col); + let values_len = assert_values_buffer_lens(col); // Since there is single long string, len of values buffer is 13 - assert_eq!(unchecked_values_len, 13); - assert!(checked_values_len > unchecked_values_len); + assert_eq!(values_len, 13); } #[test] diff --git a/arrow-row/src/variable.rs b/arrow-row/src/variable.rs index b421b5e27ad1..ca9700b90fc1 100644 --- a/arrow-row/src/variable.rs +++ b/arrow-row/src/variable.rs @@ -309,47 +309,44 @@ pub fn decode_binary( } } -fn decode_binary_view_inner( +fn decode_binary_view_inner( rows: &mut [&[u8]], options: SortOptions, - validate_utf8: bool, ) -> BinaryViewArray { let len = rows.len(); let inline_str_max_len = MAX_INLINE_VIEW_LEN as usize; let nulls = decode_nulls_sentinel(rows, options); - // If we are validating UTF-8, decode all string values (including short strings) - // into the values buffer and validate UTF-8 once. If not validating, - // we save memory by only copying long strings to the values buffer, as short strings - // will be inlined into the view and do not need to be stored redundantly. - let values_capacity = if validate_utf8 { - // Capacity for all long and short strings - rows.iter().map(|row| decoded_len(row, options)).sum() + // Capacity for all long strings plus room for one short string + let mut values_capacity = inline_str_max_len; + let mut inline_capacity = 0; + for row in rows.iter() { + let len = decoded_len(row, options); + if len > inline_str_max_len { + values_capacity += len; + } else if VALIDATE_UTF8 { + inline_capacity += len; + } + } + let mut values = MutableBuffer::new(values_capacity); + let mut view_utf8_validation_buffer = if VALIDATE_UTF8 { + Vec::with_capacity(inline_capacity) } else { - // Capacity for all long strings plus room for one short string - rows.iter().fold(0, |acc, row| { - let len = decoded_len(row, options); - if len > inline_str_max_len { - acc + len - } else { - acc - } - }) + inline_str_max_len + Vec::new() }; - let mut values = MutableBuffer::new(values_capacity); - let mut views = BufferBuilder::::new(len); - for row in rows { + let mut views = vec![0_u128; len]; + for (i, row) in rows.iter_mut().enumerate() { let start_offset = values.len(); let offset = decode_blocks(row, options, |b| values.extend_from_slice(b)); - // Measure string length via change in values buffer. - // Used to check if decoded value should be truncated (short string) when validate_utf8 is false + // Measure string length via change in values buffer. This way we can + // overwrite short strings in the values buffer as we inline those to + // views. let decoded_len = values.len() - start_offset; if row[0] == null_sentinel(options) { debug_assert_eq!(offset, 1); debug_assert_eq!(start_offset, values.len()); - views.append(0); } else { // Safety: we just appended the data to the end of the buffer let val = unsafe { values.get_unchecked_mut(start_offset..) }; @@ -358,21 +355,21 @@ fn decode_binary_view_inner( val.iter_mut().for_each(|o| *o = !*o); } - let view = make_view(val, 0, start_offset as u32); - views.append(view); + views[i] = make_view(val, 0, start_offset as u32); - // truncate inline string in values buffer if validate_utf8 is false - if !validate_utf8 && decoded_len <= inline_str_max_len { + if decoded_len <= inline_str_max_len { + if VALIDATE_UTF8 { + view_utf8_validation_buffer.extend_from_slice(val); + } values.truncate(start_offset); } } *row = &row[offset..]; } - if validate_utf8 { - // the values contains all data, no matter if it is short or long - // we can validate utf8 in one go. - std::str::from_utf8(values.as_slice()).unwrap(); + if VALIDATE_UTF8 { + std::str::from_utf8(&values).unwrap(); + std::str::from_utf8(&view_utf8_validation_buffer).unwrap(); } // SAFETY: @@ -382,7 +379,7 @@ fn decode_binary_view_inner( /// Decodes a binary view array from `rows` with the provided `options` pub fn decode_binary_view(rows: &mut [&[u8]], options: SortOptions) -> BinaryViewArray { - decode_binary_view_inner(rows, options, false) + decode_binary_view_inner::(rows, options) } /// Decodes a string array from `rows` with the provided `options` @@ -418,7 +415,11 @@ pub unsafe fn decode_string_view( options: SortOptions, validate_utf8: bool, ) -> StringViewArray { - let view = decode_binary_view_inner(rows, options, validate_utf8); + let view = if validate_utf8 { + decode_binary_view_inner::(rows, options) + } else { + decode_binary_view_inner::(rows, options) + }; unsafe { view.to_string_view_unchecked() } }