Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Improvements to UTF-8 statistics truncation #6870

Open
wants to merge 9 commits into
base: main
Choose a base branch
from
Open

Improvements to UTF-8 statistics truncation #6870

Changes from 5 commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
0f7af02
fix a few edge cases with utf-8 incrementing
etseidl Dec 10, 2024
52706f9
add todo
etseidl Dec 11, 2024
80fa0dd
simplify truncation
etseidl Dec 11, 2024
f1726ab
add another test
etseidl Dec 11, 2024
7b88e91
Merge branch 'apache:main' into increment_utf8
etseidl Dec 11, 2024
c4d9474
note case where string should render right to left
etseidl Dec 12, 2024
7a7fd0e
rework entirely, also avoid UTF8 processing if not required by the sc…
etseidl Dec 13, 2024
400f5f8
more consistent naming
etseidl Dec 14, 2024
006a388
modify some tests to truncate in the middle of a multibyte char
etseidl Dec 14, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
138 changes: 110 additions & 28 deletions parquet/src/column/writer/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1444,15 +1444,31 @@ fn increment(mut data: Vec<u8>) -> Option<Vec<u8>> {
/// Try and increment the the string's bytes from right to left, returning when the result
/// is a valid UTF8 string. Returns `None` when it can't increment any byte.
fn increment_utf8(mut data: Vec<u8>) -> Option<Vec<u8>> {
const UTF8_CONTINUATION: u8 = 0x80;
const UTF8_CONTINUATION_MASK: u8 = 0xc0;

let mut len = data.len();
for idx in (0..data.len()).rev() {
let original = data[idx];
let (byte, overflow) = original.overflowing_add(1);
if !overflow {
data[idx] = byte;
if str::from_utf8(&data).is_ok() {
etseidl marked this conversation as resolved.
Show resolved Hide resolved
if len != data.len() {
data.truncate(len);
}
return Some(data);
}
data[idx] = original;
// Incrementing "original" did not yield a valid unicode character, so it overflowed
// its available bits. If it was a continuation byte (b10xxxxxx) then set to min
// continuation (b10000000). Otherwise it was the first byte so set reset the first
// byte back to its original value (so data remains a valid string) and reduce "len".
if original & UTF8_CONTINUATION_MASK == UTF8_CONTINUATION {
etseidl marked this conversation as resolved.
Show resolved Hide resolved
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If this isn't super perf critical, can we switch over to operating on codepoints?
(i assume this is for stats only, so not a hot path?)

Copy link
Contributor

@alamb alamb Dec 12, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If this isn't super perf critical, can we switch over to operating on codepoints? (i assume this is for stats only, so not a hot path?)

I agree switching to arithmetic on codepoints would be easier to reason about.

I double checked and this function is only called while writing stats (at most once per page, and once per column chunk ):
https://github.com/search?q=repo%3Aapache%2Farrow-rs%20increment_utf8&type=code

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ok, here's what I came up with (stealing liberally from the code in apache/datafusion#12978 😄) that avoids as much translation/allocation as I could manage.

/// caller guarantees that data.len() > length
fn truncate_and_inc_utf8(data: &str, length: usize) -> Option<Vec<u8>> {
    // UTF-8 is max 4 bytes, so start search 3 back from desired length
    let lower_bound = length.saturating_sub(3);
    let split = (lower_bound..=length).rfind(|x| data.is_char_boundary(*x))?;
    increment_utf8_str(data.get(..split)?)
}

fn increment_utf8_str(data: &str) -> Option<Vec<u8>> {
    for (idx, code_point) in data.char_indices().rev() {
        let curr_len = code_point.len_utf8();
        let original = code_point as u32;
        if let Some(next_char) = char::from_u32(original + 1) {
            // do not allow increasing byte width of incremented char
            if next_char.len_utf8() == curr_len {
                let mut result = data.as_bytes()[..idx + curr_len].to_vec();
                next_char.encode_utf8(&mut result[idx..]);
                return Some(result);
            }
        }
    }
    None
}

This winds up being a little faster than the byte-by-byte implementation in this PR in the latter's best case, and as much as 30X faster for the worst case.

The check for increasing byte count can be skipped if we don't care about overshooting a little. This also doesn't include the check for non-characters as all we care about is a valid UTF-8 bound...the bounds shouldn't be interpreted by readers anyway as they are not exact. The only case this doesn't handle is incrementing U+D7FF. The range U+D800..U+DFFF overlaps with UTF-16, so those code points are not valid characters. The byte-by-byte implementation will increment to the next valid code point of U+E000, whereas this code will punt and say U+D7FF cannot be incremented. This isn't really a practical concern, though, since the nearest assigned code point is U+D7FB.

@alamb @findepi if this looks reasonable I'll clean it up and push it.

data[idx] = UTF8_CONTINUATION;
} else {
data[idx] = original;
len = idx;
}
}
}

Expand All @@ -1462,6 +1478,7 @@ fn increment_utf8(mut data: Vec<u8>) -> Option<Vec<u8>> {
#[cfg(test)]
mod tests {
use crate::file::properties::DEFAULT_COLUMN_INDEX_TRUNCATE_LENGTH;
use core::str;
use rand::distributions::uniform::SampleUniform;
use std::sync::Arc;

Expand Down Expand Up @@ -3140,39 +3157,67 @@ mod tests {

#[test]
fn test_increment_utf8() {
let test_inc = |o: &str, expected: &str| {
if let Ok(v) = String::from_utf8(increment_utf8(o.as_bytes().to_vec()).unwrap()) {
// Got the expected result...
assert_eq!(v, expected);
// and it's greater than the original string
assert!(*v > *o);
// Also show that BinaryArray level comparison works here
let mut greater = ByteArray::new();
greater.set_data(Bytes::from(v));
let mut original = ByteArray::new();
original.set_data(Bytes::from(o.as_bytes().to_vec()));
assert!(greater > original);
etseidl marked this conversation as resolved.
Show resolved Hide resolved
} else {
panic!("Expected incremented UTF8 string to also be valid.");
}
};

// Basic ASCII case
let v = increment_utf8("hello".as_bytes().to_vec()).unwrap();
assert_eq!(&v, "hellp".as_bytes());
test_inc("hello", "hellp");

// 1-byte ending in max 1-byte
test_inc("a\u{7f}", "b");

// Also show that BinaryArray level comparison works here
let mut greater = ByteArray::new();
greater.set_data(Bytes::from(v));
let mut original = ByteArray::new();
original.set_data(Bytes::from("hello".as_bytes().to_vec()));
assert!(greater > original);
// 1-byte max should not truncate as it would need 2-byte code points
assert!(increment_utf8("\u{7f}\u{7f}".as_bytes().to_vec()).is_none());

// UTF8 string
let s = "❤️🧡💛💚💙💜";
let v = increment_utf8(s.as_bytes().to_vec()).unwrap();
test_inc("❤️🧡💛💚💙💜", "❤️🧡💛💚💙💝");

if let Ok(new) = String::from_utf8(v) {
assert_ne!(&new, s);
assert_eq!(new, "❤️🧡💛💚💙💝");
assert!(new.as_bytes().last().unwrap() > s.as_bytes().last().unwrap());
} else {
panic!("Expected incremented UTF8 string to also be valid.")
}
// 2-byte without overflow
test_inc("éééé", "éééê");

// Max UTF8 character - should be a No-Op
let s = char::MAX.to_string();
assert_eq!(s.len(), 4);
let v = increment_utf8(s.as_bytes().to_vec());
assert!(v.is_none());
// 2-byte that overflows lowest byte
test_inc("\u{ff}\u{ff}", "\u{ff}\u{100}");
etseidl marked this conversation as resolved.
Show resolved Hide resolved

// 2-byte ending in max 2-byte
test_inc("a\u{7ff}", "b");

// Max 2-byte should not truncate as it would need 3-byte code points
assert!(increment_utf8("\u{7ff}\u{7ff}".as_bytes().to_vec()).is_none());

// 3-byte without overflow
test_inc("ࠀࠀ", "ࠀࠁ");
alamb marked this conversation as resolved.
Show resolved Hide resolved

// 3-byte ending in max 3-byte
test_inc("a\u{ffff}", "b");

// Max 3-byte should not truncate as it would need 4-byte code points
assert!(increment_utf8("\u{ffff}\u{ffff}".as_bytes().to_vec()).is_none());

// Handle multi-byte UTF8 characters
let s = "a\u{10ffff}";
let v = increment_utf8(s.as_bytes().to_vec());
assert_eq!(&v.unwrap(), "b\u{10ffff}".as_bytes());
// 4-byte without overflow
test_inc("𐀀𐀀", "𐀀𐀁");

// 4-byte ending in max unicode
test_inc("a\u{10ffff}", "b");

// Max 4-byte should not truncate
assert!(increment_utf8("\u{10ffff}\u{10ffff}".as_bytes().to_vec()).is_none());

// Skip over surrogate pair range (0xD800..=0xDFFF)
test_inc("a\u{D7FF}", "a\u{e000}");
}

#[test]
Expand All @@ -3182,7 +3227,6 @@ mod tests {
let r = truncate_utf8(data, data.as_bytes().len()).unwrap();
assert_eq!(r.len(), data.as_bytes().len());
assert_eq!(&r, data.as_bytes());
println!("len is {}", data.len());

// We slice it away from the UTF8 boundary
let r = truncate_utf8(data, 13).unwrap();
Expand All @@ -3192,6 +3236,44 @@ mod tests {
// One multi-byte code point, and a length shorter than it, so we can't slice it
let r = truncate_utf8("\u{0836}", 1);
assert!(r.is_none());

// test truncate and increment for max bounds on utf-8 stats
// 7-bit (i.e. ASCII)
let r = truncate_utf8("yyyyyyyyy", 8)
.and_then(increment_utf8)
.unwrap();
assert_eq!(&r, "yyyyyyyz".as_bytes());

// 2-byte without overflow
let r = truncate_utf8("ééééé", 8).and_then(increment_utf8).unwrap();
assert_eq!(&r, "éééê".as_bytes());

// 2-byte that overflows lowest byte
let r = truncate_utf8("\u{ff}\u{ff}\u{ff}\u{ff}\u{ff}", 8)
.and_then(increment_utf8)
.unwrap();
assert_eq!(&r, "\u{ff}\u{ff}\u{ff}\u{100}".as_bytes());

// max 2-byte should not truncate as it would need 3-byte code points
let r = truncate_utf8("߿߿߿߿߿", 8).and_then(increment_utf8);
assert!(r.is_none());

// 3-byte without overflow
let r = truncate_utf8("ࠀࠀࠀ", 8).and_then(increment_utf8).unwrap();
assert_eq!(&r, "ࠀࠁ".as_bytes());
assert_eq!(r.len(), 6);

// max 3-byte should not truncate as it would need 4-byte code points
let r = truncate_utf8("\u{ffff}\u{ffff}\u{ffff}", 8).and_then(increment_utf8);
assert!(r.is_none());

// 4-byte without overflow
let r = truncate_utf8("𐀀𐀀𐀀", 8).and_then(increment_utf8).unwrap();
assert_eq!(&r, "𐀀𐀁".as_bytes());

// max 4-byte should not truncate
let r = truncate_utf8("\u{10ffff}\u{10ffff}", 8).and_then(increment_utf8);
assert!(r.is_none());
}

#[test]
Expand Down
Loading