diff --git a/vortex-array/benches/varbinview_try_new.rs b/vortex-array/benches/varbinview_try_new.rs index fd87bc70a83..a50f0b588f3 100644 --- a/vortex-array/benches/varbinview_try_new.rs +++ b/vortex-array/benches/varbinview_try_new.rs @@ -96,3 +96,32 @@ fn all_valid_exclusive(bencher: Bencher, len: usize) { .vortex_expect("try_new must succeed") }); } + +/// Outlined ASCII strings in one buffer, with every other slot null. +#[divan::bench(args = SIZES)] +fn outlined_nullable_exclusive(bencher: Bencher, len: usize) { + let data = b"a string that is outlined"; + let buffer = ByteBuffer::from(data.repeat(len)); + let views: Vec = (0..len) + .map(|i| { + let offset = u32::try_from(i * data.len()).vortex_expect("offset fits in u32"); + BinaryView::make_view(data, 0, offset) + }) + .collect(); + let buffers: Arc<[ByteBuffer]> = Arc::new([buffer]); + let dtype = DType::Utf8(Nullability::Nullable); + bencher + .with_inputs(|| { + ( + Buffer::copy_from(&views), + Arc::clone(&buffers), + dtype.clone(), + nullable_validity(len), + SESSION.create_execution_ctx(), + ) + }) + .bench_values(|(views, buffers, dtype, validity, mut ctx)| { + VarBinViewArray::try_new(views, buffers, dtype, validity, &mut ctx) + .vortex_expect("try_new must succeed") + }); +} diff --git a/vortex-array/src/arrays/validation_tests.rs b/vortex-array/src/arrays/validation_tests.rs index fc7a6ee0a30..30ebdf40322 100644 --- a/vortex-array/src/arrays/validation_tests.rs +++ b/vortex-array/src/arrays/validation_tests.rs @@ -223,4 +223,36 @@ mod tests { assert!(matches!(result, Err(VortexError::InvalidArgument(_, _)))); assert!(result.is_err()); } + + #[rstest::rstest] + #[case::whole(vec![0, 2, 3], "éa".as_bytes().to_vec(), Validity::NonNullable, true)] + #[case::inside_a_char(vec![0, 1, 3], "éa".as_bytes().to_vec(), Validity::NonNullable, false)] + #[case::invalid_bytes(vec![0, 1, 2], vec![b'a', 0xFF], Validity::NonNullable, false)] + #[case::invalid_bytes_at_a_null( + vec![0, 1, 2, 3], + vec![b'a', 0xFF, b'b'], + Validity::from_iter([true, false, true]), + true + )] + #[case::decreasing_offsets(vec![0, 2, 1], "abc".as_bytes().to_vec(), Validity::NonNullable, false)] + #[case::offset_past_the_end(vec![0, 5, 1], "abc".as_bytes().to_vec(), Validity::NonNullable, false)] + fn test_varbin_utf8_validation( + #[case] offsets: Vec, + #[case] bytes: Vec, + #[case] validity: Validity, + #[case] ok: bool, + ) { + let nullability = if matches!(validity, Validity::NonNullable) { + Nullability::NonNullable + } else { + Nullability::Nullable + }; + let result = VarBinArray::try_new( + Buffer::from(offsets).into_array(), + ByteBuffer::from(bytes), + DType::Utf8(nullability), + validity, + ); + assert_eq!(result.is_ok(), ok); + } } diff --git a/vortex-array/src/arrays/varbin/array.rs b/vortex-array/src/arrays/varbin/array.rs index a763fb4f1de..82aa3251253 100644 --- a/vortex-array/src/arrays/varbin/array.rs +++ b/vortex-array/src/arrays/varbin/array.rs @@ -32,6 +32,12 @@ use crate::legacy_session; use crate::match_each_integer_ptype; use crate::validity::Validity; +/// Returns `true` if `index` is the end of `bytes` or the start of a UTF-8 char in it. +fn is_char_boundary_at(bytes: &[u8], index: usize) -> bool { + // Continuation bytes have the form `0b10xx_xxxx`. + index == bytes.len() || bytes[index] & 0b1100_0000 != 0b1000_0000 +} + #[array_slots(VarBin)] pub struct VarBinSlots { /// The offsets array defining the start/end of each variable-length binary element. @@ -241,7 +247,12 @@ impl VarBinData { #[allow(clippy::disallowed_methods)] fn validate_utf8(offsets: &ArrayRef, bytes: &[u8], validity: &Validity) -> VortexResult<()> { let validate_at = |i: usize, start: usize, end: usize| -> VortexResult<()> { - let string_bytes = &bytes[start..end]; + let string_bytes = bytes.get(start..end).ok_or_else(|| { + vortex_err!( + InvalidArgument: "offsets {start}..{end} at index {i} are out of order or out of bounds for bytes of length {}", + bytes.len() + ) + })?; simdutf8::basic::from_utf8(string_bytes).map_err(|_| { #[expect(clippy::unwrap_used)] // run validation using `compat` package to get more detailed error message @@ -277,6 +288,21 @@ impl VarBinData { bytes.len() ); + // When the offsets never decrease, the strings tile `bytes[first..last]`. If that range + // is valid UTF-8 as a whole, every string is valid UTF-8 if and only if every offset + // falls on a char boundary. Otherwise, for example for invalid bytes at a null, check + // the strings one by one. + let first_offset: usize = offsets_slice[0].as_(); + if offsets_slice.windows(2).all(|o| o[0] <= o[1]) + && first_offset <= last_offset + && simdutf8::basic::from_utf8(&bytes[first_offset..last_offset]).is_ok() + && offsets_slice + .iter() + .all(|&o| is_char_boundary_at(bytes, o.as_())) + { + return Ok(()); + } + for (i, (start, end)) in offsets_slice .windows(2) .map(|o| (o[0].as_(), o[1].as_())) diff --git a/vortex-array/src/arrays/varbinview/array.rs b/vortex-array/src/arrays/varbinview/array.rs index cfccfa2a00a..b2db8f79b65 100644 --- a/vortex-array/src/arrays/varbinview/array.rs +++ b/vortex-array/src/arrays/varbinview/array.rs @@ -147,6 +147,33 @@ fn for_each_invalid_range(mask: &Mask, len: usize, mut f: impl FnMut(usize, usiz } } +/// Calls `f` for each valid index of `mask`, and stops at the first error. +fn try_for_each_valid( + mask: &Mask, + mut f: impl FnMut(usize) -> VortexResult<()>, +) -> VortexResult<()> { + match mask.bit_buffer() { + AllOr::All => (0..mask.len()).try_for_each(f), + AllOr::None => Ok(()), + AllOr::Some(buffer) => buffer.set_indices().try_for_each(&mut f), + } +} + +/// Returns `true` if all the inlined bytes of `view` are ASCII. +fn is_inlined_ascii(view: &BinaryView) -> bool { + const HIGH_BITS: u128 = u128::from_le_bytes([0x80; 16]); + + // The first 4 bytes hold the length, and the inlined bytes follow. + let data_bits = ((1u128 << (8 * view.len())) - 1) << 32; + view.as_u128() & data_bits & HIGH_BITS == 0 +} + +/// Returns `true` if `byte` starts a UTF-8 char, that is, it is not a continuation byte. +fn is_char_boundary(byte: u8) -> bool { + // Continuation bytes have the form `0b10xx_xxxx`. + byte & 0b1100_0000 != 0b1000_0000 +} + impl VarBinViewData { fn dtype_parts(dtype: &DType) -> VortexResult<(bool, Nullability)> { match dtype { @@ -333,28 +360,8 @@ impl VarBinViewData { ctx: &mut ExecutionCtx, ) -> VortexResult<()> { let check_utf8 = Self::check_nullability(dtype, validity)?; - match validity { - // Array-backed validity is the only variant that needs an execution context: execute it - // into a mask once and zip it with the views, validating only the valid (non-null) - // entries. - Validity::Array(_) => { - let mask = validity.execute_mask(views.len(), ctx)?; - for ((idx, view), valid) in views.iter().enumerate().zip(mask.iter()) { - if valid { - Self::validate_view(idx, view, buffers, check_utf8)?; - } - } - } - // Every entry is null, so there is nothing to validate. - Validity::AllInvalid => {} - // No nulls: validate every view. - Validity::NonNullable | Validity::AllValid => { - for (idx, view) in views.iter().enumerate() { - Self::validate_view(idx, view, buffers, check_utf8)?; - } - } - } - Ok(()) + let mask = validity.execute_mask(views.len(), ctx)?; + Self::validate_views(views, &mask, buffers, check_utf8) } /// Validates components like validate() and replaces views at null slots to empty views @@ -366,66 +373,31 @@ impl VarBinViewData { ctx: &mut ExecutionCtx, ) -> VortexResult> { let check_utf8 = Self::check_nullability(dtype, validity)?; - let empty = BinaryView::empty_view(); - let len = views.len(); - - match validity { - Validity::Array(_) => { - let mask = validity.execute_mask(len, ctx)?; - match views.try_into_mut() { - Ok(mut views) => { - let slice = views.as_mut_slice(); - for (idx, valid) in mask.iter().enumerate() { - if valid { - Self::validate_view(idx, &slice[idx], buffers, check_utf8)?; - } else { - slice[idx] = empty; - } - } - Ok(views.freeze()) - } - Err(views) => { - let mut needs_replace = false; - for ((idx, view), valid) in views.iter().enumerate().zip(mask.iter()) { - if valid { - Self::validate_view(idx, view, buffers, check_utf8)?; - } else if *view != empty { - needs_replace = true; - } - } - if !needs_replace { - return Ok(views); - } - let mut views = views.into_mut(); - let slice = views.as_mut_slice(); - for_each_invalid_range(&mask, len, |start, end| { - slice[start..end].fill(empty) - }); - Ok(views.freeze()) + let mask = validity.execute_mask(views.len(), ctx)?; + + // With nulls and views we own, validate and replace in a single pass: a second pass over + // many short null runs costs more than a branch per view. + let views = match (mask.bit_buffer(), views.try_into_mut()) { + (AllOr::Some(bits), Ok(mut views)) => { + let buffers_utf8 = Self::buffers_utf8(&mask, buffers, check_utf8); + let buffers_utf8 = buffers_utf8.as_deref(); + let empty = BinaryView::empty_view(); + let slice = views.as_mut_slice(); + for (idx, (view, valid)) in slice.iter_mut().zip(bits.iter()).enumerate() { + if valid { + Self::validate_view(idx, view, buffers, buffers_utf8)?; + } else { + *view = empty; } } + return Ok(views.freeze()); } - // Every entry is null, so there is nothing to validate: replace all views with empty. - Validity::AllInvalid => match views.try_into_mut() { - Ok(mut views) => { - views.as_mut_slice().fill(empty); - Ok(views.freeze()) - } - Err(views) if views.iter().all(|view| *view == empty) => Ok(views), - Err(views) => { - let mut views = views.into_mut(); - views.as_mut_slice().fill(empty); - Ok(views.freeze()) - } - }, - // No nulls: validate every view, nothing to replace. - Validity::NonNullable | Validity::AllValid => { - for (idx, view) in views.iter().enumerate() { - Self::validate_view(idx, view, buffers, check_utf8)?; - } - Ok(views) - } - } + (_, Ok(views)) => views.freeze(), + (_, Err(views)) => views, + }; + + Self::validate_views(&views, &mask, buffers, check_utf8)?; + Ok(Self::replace_null_views(views, &mask)) } fn check_nullability(dtype: &DType, validity: &Validity) -> VortexResult { @@ -439,56 +411,154 @@ impl VarBinViewData { Ok(is_utf8) } + /// Replaces the views at null slots with empty views. Copies the views only if a null slot + /// holds a non-empty view and the buffer is shared. + fn replace_null_views(views: Buffer, mask: &Mask) -> Buffer { + if mask.all_true() { + return views; + } + + let empty = BinaryView::empty_view(); + let len = views.len(); + + let mut views = match views.try_into_mut() { + Ok(views) => views, + Err(views) => { + let mut needs_replace = false; + for_each_invalid_range(mask, len, |start, end| { + needs_replace |= views[start..end].iter().any(|view| *view != empty); + }); + if !needs_replace { + return views; + } + views.into_mut() + } + }; + + let slice = views.as_mut_slice(); + for_each_invalid_range(mask, len, |start, end| slice[start..end].fill(empty)); + views.freeze() + } + + /// Validates the views at the valid slots of `mask`. + /// + /// A per-view UTF-8 check costs a function call per view, which dominates for short strings. + /// To avoid it, each data buffer is first checked once as a whole. If a buffer is valid UTF-8, + /// a view into it is valid UTF-8 if and only if it starts and ends on a char boundary. + /// + /// A buffer can hold bytes that no valid view references, for example after a slice, after a + /// filter, or at null views. So the whole-buffer check is only done while the total size of + /// the checked buffers stays within a fixed cost per valid view. The views into a buffer that + /// is not checked, or that is not valid UTF-8, are checked one by one. + fn validate_views( + views: &[BinaryView], + mask: &Mask, + buffers: &[ByteBuffer], + check_utf8: bool, + ) -> VortexResult<()> { + let buffers_utf8 = Self::buffers_utf8(mask, buffers, check_utf8); + let buffers_utf8 = buffers_utf8.as_deref(); + + try_for_each_valid(mask, |idx| { + Self::validate_view(idx, &views[idx], buffers, buffers_utf8) + }) + } + + /// The approximate cost of a per-view UTF-8 check, in bytes of a whole-buffer check. + const PER_VIEW_UTF8_COST: usize = 64; + + /// Returns `None` if the UTF-8 check is off. Otherwise, returns for each buffer `true` if the + /// whole buffer was checked and is valid UTF-8. + /// + /// The budget only counts a fixed cost per valid view, not the bytes the views reference: for + /// long strings, the call per view is cheap next to the string, and checking a buffer larger + /// than the cache and then reading it again for the char boundaries is slower. + fn buffers_utf8(mask: &Mask, buffers: &[ByteBuffer], check_utf8: bool) -> Option> { + if !check_utf8 { + return None; + } + + let mut budget = mask.true_count().saturating_mul(Self::PER_VIEW_UTF8_COST); + + Some( + buffers + .iter() + .map(|buf| { + if buf.len() > budget { + return false; + } + + budget -= buf.len(); + simdutf8::basic::from_utf8(buf).is_ok() + }) + .collect(), + ) + } + + /// Checks the bounds, the prefix, and the UTF-8 of a view. + /// + /// `buffers_utf8` is `None` if the UTF-8 check is off. Otherwise, it tells which buffers are + /// valid UTF-8 as a whole. fn validate_view( idx: usize, view: &BinaryView, - buffers: &Arc<[ByteBuffer]>, - check_utf8: bool, + buffers: &[ByteBuffer], + buffers_utf8: Option<&[bool]>, ) -> VortexResult<()> { - let valid_utf8 = |bytes: &[u8]| !check_utf8 || simdutf8::basic::from_utf8(bytes).is_ok(); if view.is_inlined() { - // Validate the inline bytestring - let bytes = &view.as_inlined().data[..view.len() as usize]; - vortex_ensure!( - valid_utf8(bytes), - InvalidArgument: "view at index {idx}: inlined bytes failed utf-8 validation" - ); - } else { - // Validate the view pointer - let view = view.as_view(); - let buf_index = view.buffer_index as usize; - let start_offset = view.offset as usize; - let end_offset = start_offset.saturating_add(view.size as usize); + if buffers_utf8.is_some() && !is_inlined_ascii(view) { + let bytes = view.as_inlined().value(); + vortex_ensure!( + simdutf8::basic::from_utf8(bytes).is_ok(), + InvalidArgument: "view at index {idx}: inlined bytes failed utf-8 validation" + ); + } + return Ok(()); + } - let buf = buffers.get(buf_index).ok_or_else(|| - vortex_err!(InvalidArgument: "view at index {idx} references invalid buffer: {buf_index} out of bounds for VarBinViewData with {} buffers", - buffers.len()))?; + let view = view.as_view(); + let buf_index = view.buffer_index as usize; + let start_offset = view.offset as usize; + let end_offset = start_offset.saturating_add(view.size as usize); - vortex_ensure!( - start_offset < buf.len(), - InvalidArgument: "start offset {start_offset} out of bounds for buffer {buf_index} with size {}", - buf.len(), - ); + let buf = buffers.get(buf_index).ok_or_else(|| + vortex_err!(InvalidArgument: "view at index {idx} references invalid buffer: {buf_index} out of bounds for VarBinViewData with {} buffers", + buffers.len()))?; - vortex_ensure!( - end_offset <= buf.len(), - InvalidArgument: "end offset {end_offset} out of bounds for buffer {buf_index} with size {}", - buf.len(), - ); + vortex_ensure!( + start_offset < buf.len(), + InvalidArgument: "start offset {start_offset} out of bounds for buffer {buf_index} with size {}", + buf.len(), + ); - // Make sure the prefix data matches the buffer data. - let bytes = &buf[start_offset..end_offset]; - vortex_ensure!( - view.prefix == bytes[..4], - InvalidArgument: "VarBinView prefix does not match full string" - ); + vortex_ensure!( + end_offset <= buf.len(), + InvalidArgument: "end offset {end_offset} out of bounds for buffer {buf_index} with size {}", + buf.len(), + ); + + // Make sure the prefix data matches the buffer data. + vortex_ensure!( + view.prefix == buf[start_offset..start_offset + 4], + InvalidArgument: "VarBinView prefix does not match full string" + ); + + let Some(buffers_utf8) = buffers_utf8 else { + return Ok(()); + }; + + let valid_utf8 = if buffers_utf8[buf_index] { + is_char_boundary(buf[start_offset]) + && (end_offset == buf.len() || is_char_boundary(buf[end_offset])) + } else { + simdutf8::basic::from_utf8(&buf[start_offset..end_offset]).is_ok() + }; + + vortex_ensure!( + valid_utf8, + InvalidArgument: "view at index {idx}: outlined bytes fails utf-8 validation" + ); - // Validate the full string - vortex_ensure!( - valid_utf8(bytes), - InvalidArgument: "view at index {idx}: outlined bytes fails utf-8 validation" - ); - } Ok(()) } diff --git a/vortex-array/src/arrays/varbinview/tests.rs b/vortex-array/src/arrays/varbinview/tests.rs index 80238403536..1350e1eca96 100644 --- a/vortex-array/src/arrays/varbinview/tests.rs +++ b/vortex-array/src/arrays/varbinview/tests.rs @@ -3,6 +3,7 @@ use std::sync::Arc; +use rstest::rstest; use vortex_buffer::BitBuffer; use vortex_buffer::Buffer; use vortex_buffer::ByteBuffer; @@ -141,3 +142,88 @@ pub fn deserialize_null_views() -> VortexResult<()> { assert_eq!(decoded.views()[1], BinaryView::empty_view()); Ok(()) } + +/// Validates `views` into a single data buffer as nullable UTF-8. +fn validate_utf8_views( + buffer: &[u8], + views: Vec, + validity: Validity, +) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let buffers: Arc<[ByteBuffer]> = Arc::new([ByteBuffer::from(buffer.to_vec())]); + let views = Buffer::copy_from(views); + let dtype = DType::Utf8(Nullability::Nullable); + + VarBinViewData::validate(&views, &buffers, &dtype, &validity, &mut ctx) +} + +/// Makes a view of `buffer[start..end]` at buffer index 0. +fn view_of(buffer: &[u8], start: usize, end: usize) -> BinaryView { + BinaryView::make_view(&buffer[start..end], 0, u32::try_from(start).unwrap()) +} + +#[test] +pub fn validate_multibyte_utf8_views() -> VortexResult<()> { + let first = "zażółć gęślą jaźń"; + let buffer = format!("{first}łódź pod mostem"); + let buffer = buffer.as_bytes(); + + let views = vec![ + view_of(buffer, 0, first.len()), + view_of(buffer, first.len(), buffer.len()), + BinaryView::new_inlined("żółw".as_bytes()), + ]; + + validate_utf8_views(buffer, views, Validity::AllValid) +} + +#[rstest] +#[case::whole(0, 18, true)] +#[case::ascii_middle(2, 16, true)] +#[case::starts_inside_char(1, 18, false)] +#[case::ends_inside_char(0, 17, false)] +pub fn validate_view_char_boundaries( + #[case] start: usize, + #[case] end: usize, + #[case] valid: bool, +) { + // "ż" is two bytes, so the buffer is 18 bytes. The first view covers the whole buffer, so the + // second view is checked against a buffer range that is valid UTF-8 as a whole. + let buffer = "żaaaaaaaaaaaaaaż".as_bytes(); + let views = vec![ + view_of(buffer, 0, buffer.len()), + view_of(buffer, start, end), + ]; + + assert_eq!( + validate_utf8_views(buffer, views, Validity::AllValid).is_ok(), + valid + ); +} + +#[test] +pub fn validate_ignores_invalid_utf8_outside_valid_views() -> VortexResult<()> { + let buffer = b"valid string one\xffvalid string two\xff\xfe\xfd garbage bytes"; + + let views = vec![ + view_of(buffer, 0, 16), + view_of(buffer, 17, 33), + view_of(buffer, 33, buffer.len()), + ]; + let validity = Validity::from_bit_buffer( + BitBuffer::from_iter([true, true, false]), + Nullability::Nullable, + ); + + validate_utf8_views(buffer, views, validity) +} + +#[test] +pub fn validate_rejects_invalid_utf8_in_valid_views() { + let buffer = b"valid string\xffmore bytes"; + let views = vec![view_of(buffer, 0, buffer.len())]; + assert!(validate_utf8_views(buffer, views, Validity::AllValid).is_err()); + + let views = vec![BinaryView::new_inlined(b"ab\xff")]; + assert!(validate_utf8_views(buffer, views, Validity::AllValid).is_err()); +}