This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/paimon-mosaic.git
The following commit(s) were added to refs/heads/main by this push:
new 57bff28 perf(core): speed up schema BPE construction (#89)
57bff28 is described below
commit 57bff28e5ca33d212baeb144f052001141b2c929
Author: jianguotian <[email protected]>
AuthorDate: Sun Sep 20 20:03:30 2026 +0800
perf(core): speed up schema BPE construction (#89)
---
core/src/bpe.rs | 262 +++++++++++++++++++++++++++++++++++++++++++++++------
core/src/schema.rs | 11 +--
2 files changed, 237 insertions(+), 36 deletions(-)
diff --git a/core/src/bpe.rs b/core/src/bpe.rs
index d2d8be8..f334e1e 100644
--- a/core/src/bpe.rs
+++ b/core/src/bpe.rs
@@ -17,37 +17,41 @@
use std::collections::HashMap;
-const TOKEN_BASE: usize = 0x80;
+const TOKEN_BASE: u8 = 0x80;
const MAX_RULES: usize = 128;
+const PAIR_TABLE_SIZE: usize = 1 << 16;
+const DENSE_PAIR_THRESHOLD: usize = 4_096;
pub fn is_ascii_only(names: &[&[u8]]) -> bool {
names.iter().all(|name| name.iter().all(|&b| b & 0x80 == 0))
}
pub fn build_vocabulary(names: &[&[u8]]) -> Vec<[u8; 2]> {
- let mut tokens: Vec<Vec<u16>> = names
- .iter()
- .map(|name| name.iter().map(|&b| b as u16).collect())
- .collect();
+ build_vocabulary_and_encode(names).0
+}
+
+pub(crate) fn build_vocabulary_and_encode(names: &[&[u8]]) -> (Vec<[u8; 2]>,
Vec<Vec<u8>>) {
+ let mut tokens: Vec<Vec<u8>> = names.iter().map(|name|
name.to_vec()).collect();
let mut rules = Vec::new();
+ let use_dense_counter = tokens
+ .iter()
+ .map(|seq| seq.len().saturating_sub(1))
+ .sum::<usize>()
+ >= DENSE_PAIR_THRESHOLD;
+ let mut pair_counts = use_dense_counter.then(|| vec![0u32;
PAIR_TABLE_SIZE]);
+ let mut touched_pairs = Vec::new();
for _ in 0..MAX_RULES {
- let mut pair_counts: HashMap<u32, u32> = HashMap::new();
- for seq in &tokens {
- for w in seq.windows(2) {
- let pair = (w[0] as u32) << 16 | w[1] as u32;
- *pair_counts.entry(pair).or_default() += 1;
- }
- }
-
- let best = pair_counts.iter().max_by_key(|&(&k, &v)| (v, k));
+ let best = if let Some(pair_counts) = &mut pair_counts {
+ most_frequent_pair_dense(&tokens, pair_counts, &mut touched_pairs)
+ } else {
+ most_frequent_pair_hashmap(&tokens)
+ };
match best {
- Some((&pair, &count)) if count > 1 => {
- let left = (pair >> 16) as u16;
- let right = (pair & 0xFFFF) as u16;
- let new_token = (TOKEN_BASE + rules.len()) as u16;
- rules.push([left as u8, right as u8]);
+ Some((left, right, count)) if count > 1 => {
+ let new_token = TOKEN_BASE + rules.len() as u8;
+ rules.push([left, right]);
for seq in &mut tokens {
replace_pair(seq, left, right, new_token);
@@ -57,10 +61,55 @@ pub fn build_vocabulary(names: &[&[u8]]) -> Vec<[u8; 2]> {
}
}
- rules
+ (rules, tokens)
+}
+
+fn most_frequent_pair_dense(
+ tokens: &[Vec<u8>],
+ pair_counts: &mut [u32],
+ touched_pairs: &mut Vec<u16>,
+) -> Option<(u8, u8, u32)> {
+ debug_assert_eq!(pair_counts.len(), PAIR_TABLE_SIZE);
+ for pair in touched_pairs.drain(..) {
+ pair_counts[pair as usize] = 0;
+ }
+
+ let mut best_pair = 0usize;
+ let mut best_count = 0u32;
+ for seq in tokens {
+ for pair in seq.windows(2) {
+ let pair = ((pair[0] as usize) << 8) | pair[1] as usize;
+ if pair_counts[pair] == 0 {
+ touched_pairs.push(pair as u16);
+ }
+ pair_counts[pair] += 1;
+ let count = pair_counts[pair];
+ if count > best_count || (count == best_count && pair > best_pair)
{
+ best_pair = pair;
+ best_count = count;
+ }
+ }
+ }
+
+ (best_count != 0).then_some(((best_pair >> 8) as u8, (best_pair & 0xff) as
u8, best_count))
}
-fn replace_pair(seq: &mut Vec<u16>, left: u16, right: u16, new_token: u16) {
+fn most_frequent_pair_hashmap(tokens: &[Vec<u8>]) -> Option<(u8, u8, u32)> {
+ let mut pair_counts: HashMap<u16, u32> = HashMap::new();
+ for seq in tokens {
+ for pair in seq.windows(2) {
+ let pair = (pair[0] as u16) << 8 | pair[1] as u16;
+ *pair_counts.entry(pair).or_default() += 1;
+ }
+ }
+
+ pair_counts
+ .into_iter()
+ .max_by_key(|&(pair, count)| (count, pair))
+ .map(|(pair, count)| ((pair >> 8) as u8, pair as u8, count))
+}
+
+fn replace_pair(seq: &mut Vec<u8>, left: u8, right: u8, new_token: u8) {
let mut i = 0;
let mut out = 0;
while i < seq.len() {
@@ -77,16 +126,16 @@ fn replace_pair(seq: &mut Vec<u16>, left: u16, right: u16,
new_token: u16) {
}
pub fn encode(name: &[u8], rules: &[[u8; 2]]) -> Vec<u8> {
- let mut tokens: Vec<u16> = name.iter().map(|&b| b as u16).collect();
+ let mut tokens = name.to_vec();
for (r, rule) in rules.iter().enumerate() {
- let left = rule[0] as u16;
- let right = rule[1] as u16;
- let new_token = (TOKEN_BASE + r) as u16;
+ let left = rule[0];
+ let right = rule[1];
+ let new_token = TOKEN_BASE + r as u8;
replace_pair(&mut tokens, left, right, new_token);
}
- tokens.iter().map(|&t| t as u8).collect()
+ tokens
}
pub fn decode(encoded: &[u8], rules: &[[u8; 2]]) -> Vec<u8> {
@@ -98,10 +147,10 @@ pub fn decode(encoded: &[u8], rules: &[[u8; 2]]) ->
Vec<u8> {
}
fn expand(token: usize, rules: &[[u8; 2]], out: &mut Vec<u8>) {
- if token < TOKEN_BASE {
+ if token < TOKEN_BASE as usize {
out.push(token as u8);
} else {
- let idx = token - TOKEN_BASE;
+ let idx = token - TOKEN_BASE as usize;
expand(rules[idx][0] as usize, rules, out);
expand(rules[idx][1] as usize, rules, out);
}
@@ -111,6 +160,81 @@ fn expand(token: usize, rules: &[[u8; 2]], out: &mut
Vec<u8>) {
mod tests {
use super::*;
+ const LEGACY_TOKEN_BASE: u8 = 0x80;
+ const LEGACY_MAX_RULES: usize = 128;
+
+ fn reference_replace_pair(seq: &mut Vec<u16>, left: u16, right: u16,
new_token: u16) {
+ let mut i = 0;
+ let mut out = 0;
+ while i < seq.len() {
+ if i + 1 < seq.len() && seq[i] == left && seq[i + 1] == right {
+ seq[out] = new_token;
+ i += 2;
+ } else {
+ seq[out] = seq[i];
+ i += 1;
+ }
+ out += 1;
+ }
+ seq.truncate(out);
+ }
+
+ fn reference_build_vocabulary(names: &[&[u8]]) -> Vec<[u8; 2]> {
+ let mut tokens: Vec<Vec<u16>> = names
+ .iter()
+ .map(|name| name.iter().map(|&b| b as u16).collect())
+ .collect();
+ let mut rules = Vec::new();
+
+ for _ in 0..LEGACY_MAX_RULES {
+ let mut pair_counts: HashMap<u32, u32> = HashMap::new();
+ for seq in &tokens {
+ for pair in seq.windows(2) {
+ let pair = (pair[0] as u32) << 16 | pair[1] as u32;
+ *pair_counts.entry(pair).or_default() += 1;
+ }
+ }
+
+ match pair_counts
+ .iter()
+ .max_by_key(|&(&pair, &count)| (count, pair))
+ {
+ Some((&pair, &count)) if count > 1 => {
+ let left = (pair >> 16) as u16;
+ let right = pair as u16;
+ let new_token = LEGACY_TOKEN_BASE as u16 + rules.len() as
u16;
+ rules.push([left as u8, right as u8]);
+ for seq in &mut tokens {
+ reference_replace_pair(seq, left, right, new_token);
+ }
+ }
+ _ => break,
+ }
+ }
+
+ rules
+ }
+
+ fn legacy_decode(encoded: &[u8], rules: &[[u8; 2]]) -> Vec<u8> {
+ let mut out = Vec::with_capacity(encoded.len() * 2);
+ for &token in encoded {
+ legacy_expand(token as usize, rules, &mut out);
+ }
+ out
+ }
+
+ fn legacy_expand(token: usize, rules: &[[u8; 2]], out: &mut Vec<u8>) {
+ if token < LEGACY_TOKEN_BASE as usize {
+ out.push(token as u8);
+ } else {
+ let rule = rules
+ .get(token - LEGACY_TOKEN_BASE as usize)
+ .expect("encoded token must reference a frozen legacy rule");
+ legacy_expand(rule[0] as usize, rules, out);
+ legacy_expand(rule[1] as usize, rules, out);
+ }
+ }
+
#[test]
fn test_bpe_encode() {
let names: Vec<&[u8]> = vec![
@@ -128,9 +252,91 @@ mod tests {
}
}
+ #[test]
+ fn test_build_vocabulary_and_encode_matches_separate_encoding() {
+ let mut owned_names = vec![b"aaaaa".to_vec()];
+ owned_names.extend((0..128).map(|i| {
+ format!(
+
"vehicle_data_collection_powertrain_group_{:03}_signal_measurement_{i:05}_value",
+ i % 17
+ )
+ .into_bytes()
+ }));
+ let names: Vec<&[u8]> =
owned_names.iter().map(Vec::as_slice).collect();
+
+ let (rules, encoded) = build_vocabulary_and_encode(&names);
+ assert_eq!(rules, reference_build_vocabulary(&names));
+ assert_eq!(
+ encoded,
+ names
+ .iter()
+ .map(|name| encode(name, &rules))
+ .collect::<Vec<_>>()
+ );
+ }
+
+ #[test]
+ fn test_full_legacy_token_space_round_trips() {
+ let owned_names: Vec<Vec<u8>> = (0..4096)
+ .map(|i| {
+ format!("prefix_{:04}_{}_{}_{}_suffix", i % 257, i, i % 97, i
% 31).into_bytes()
+ })
+ .collect();
+ let names: Vec<&[u8]> =
owned_names.iter().map(Vec::as_slice).collect();
+
+ let (rules, encoded) = build_vocabulary_and_encode(&names);
+
+ assert_eq!(rules.len(), LEGACY_MAX_RULES);
+ assert_eq!(rules, reference_build_vocabulary(&names));
+ assert!(encoded.iter().flatten().any(|&token| token == 0xff));
+ for (original, encoded_name) in names.iter().zip(&encoded) {
+ assert_eq!(legacy_decode(encoded_name, &rules), *original);
+ assert_eq!(decode(encoded_name, &rules), *original);
+ }
+ }
+
#[test]
fn test_ascii_only() {
assert!(is_ascii_only(&[b"hello", b"world"]));
assert!(!is_ascii_only(&[b"hello", &[0x80, 0x81]]));
}
+
+ #[test]
+ fn test_dense_pair_counter_reuses_storage_and_preserves_tie_breaking() {
+ let mut counts = vec![0; 1 << 16];
+ let mut touched = Vec::new();
+
+ let first = vec![vec![1, 2, 1, 3]];
+ assert_eq!(
+ most_frequent_pair_dense(&first, &mut counts, &mut touched),
+ Some((2, 1, 1))
+ );
+
+ let second = vec![vec![255, 254, 255, 254]];
+ assert_eq!(
+ most_frequent_pair_dense(&second, &mut counts, &mut touched),
+ Some((255, 254, 2))
+ );
+ }
+
+ #[test]
+ fn test_dense_pair_counter_matches_hashmap_reference() {
+ let mut state = 1u64;
+ for _ in 0..32 {
+ let mut names = Vec::new();
+ for _ in 0..256 {
+ let len = 4 + (state as usize % 28);
+ let mut name = Vec::with_capacity(len);
+ for _ in 0..len {
+ state = state
+ .wrapping_mul(6_364_136_223_846_793_005)
+ .wrapping_add(1);
+ name.push(b'a' + ((state >> 32) % 26) as u8);
+ }
+ names.push(name);
+ }
+ let refs: Vec<&[u8]> = names.iter().map(Vec::as_slice).collect();
+ assert_eq!(build_vocabulary(&refs),
reference_build_vocabulary(&refs));
+ }
+ }
}
diff --git a/core/src/schema.rs b/core/src/schema.rs
index e5c6d78..ecfc6d1 100644
--- a/core/src/schema.rs
+++ b/core/src/schema.rs
@@ -240,12 +240,11 @@ impl MosaicSchema {
pub fn serialize(&self) -> Vec<u8> {
let num_columns = self.columns.len();
- let raw_names: Vec<Vec<u8>> = self
+ let raw_refs: Vec<&[u8]> = self
.columns
.iter()
- .map(|c| c.name.as_bytes().to_vec())
+ .map(|column| column.name.as_bytes())
.collect();
- let raw_refs: Vec<&[u8]> = raw_names.iter().map(|v|
v.as_slice()).collect();
let plain_size = front_coded_size(&raw_refs);
@@ -254,12 +253,8 @@ impl MosaicSchema {
let mut bpe_size = usize::MAX;
if bpe::is_ascii_only(&raw_refs) {
- let rules = bpe::build_vocabulary(&raw_refs);
+ let (rules, names) = bpe::build_vocabulary_and_encode(&raw_refs);
if !rules.is_empty() {
- let names: Vec<Vec<u8>> = raw_refs
- .iter()
- .map(|name| bpe::encode(name, &rules))
- .collect();
let name_refs: Vec<&[u8]> = names.iter().map(|v|
v.as_slice()).collect();
bpe_size = 1 + rules.len() * 2 + front_coded_size(&name_refs);
bpe_rules = rules;
