diff --git a/src/mysql-util/src/partition.rs b/src/mysql-util/src/partition.rs index 48d9bf0a2ada8..39f0358a4cbae 100644 --- a/src/mysql-util/src/partition.rs +++ b/src/mysql-util/src/partition.rs @@ -14,8 +14,17 @@ use crate::{KeyProber, MySqlError, QualifiedTableRef}; /// Computes up to `num_workers - 1` partition boundaries that divide the primary key space /// into `num_workers` roughly even partitions. -/// This should be run in a repeatable read transaction against a primary key varchar/char column -/// with the `utf8mb4_bin` collation. +/// +/// Nothing here validates the setup: the caller must abide by these +/// constraints or undefined/untested behavior could occur, e.g. boundaries +/// that fail to partition the key space or a walk that does not converge. +/// * `pk_col` is the table's single-column primary key. +/// * The column type is CHAR or VARCHAR with a declared length of at most +/// [`crate::probe::MAX_KEY_LENGTH`] characters. +/// * The column collation is `utf8mb4_bin`. +/// * The connection is inside a REPEATABLE READ transaction, so the probes +/// (several queries each) all see one snapshot of the table. +/// /// `min_split_threshold` is the smallest estimated row count granularity partitioning will /// target, which means if the algorithm processes a prefix estimated to cover less /// than min_split_threshold rows it won't bother splitting it up further. This is useful to @@ -59,8 +68,8 @@ struct Prefix { depth: usize, } -async fn partition( - db: &mut KeyProber<'_>, +async fn partition( + db: &mut D, workers: usize, estimated_row_count: u64, min_split_threshold: u64, @@ -90,8 +99,8 @@ async fn partition( compute_boundaries(db, workers, estimated_row_count, target_max_rows_per_prefix).await } -async fn compute_boundaries( - db: &mut KeyProber<'_>, +async fn compute_boundaries( + db: &mut D, workers: usize, estimated_row_count: u64, target_rows_per_prefix: u64, @@ -156,8 +165,8 @@ async fn compute_boundaries( /// /// Note: This will drop the key "a" on the floor, along with any keys /// sorting below their own prefix (below-space characters at this depth). -async fn children_prefixes( - db: &mut KeyProber<'_>, +async fn children_prefixes( + db: &mut D, parent: &Prefix, ) -> Result, MySqlError> { let depth = parent.depth + 1; @@ -188,3 +197,364 @@ async fn children_prefixes( } } } + +/// Probing operations of [`KeyProber`], as a trait so tests can substitute +/// an in-memory implementation. See [`KeyProber`]'s methods for each +/// operation's contract. +trait PrimaryKeyProber { + async fn estimate_range_rows( + &mut self, + start: &str, + end: Option<&str>, + ) -> Result; + + async fn prefix_of_first_key_in_range( + &mut self, + start: &str, + end: Option<&str>, + len: usize, + ) -> Result, MySqlError>; + + async fn prefix_of_first_row_not_matching_prefix( + &mut self, + cur: &str, + end: Option<&str>, + len: usize, + ) -> Result, MySqlError>; +} + +impl<'a> PrimaryKeyProber for KeyProber<'a> { + async fn estimate_range_rows( + &mut self, + start: &str, + end: Option<&str>, + ) -> Result { + KeyProber::estimate_range_rows(self, start, end).await + } + + async fn prefix_of_first_key_in_range( + &mut self, + start: &str, + end: Option<&str>, + len: usize, + ) -> Result, MySqlError> { + KeyProber::prefix_of_first_key_in_range(self, start, end, len).await + } + + async fn prefix_of_first_row_not_matching_prefix( + &mut self, + cur: &str, + end: Option<&str>, + len: usize, + ) -> Result, MySqlError> { + KeyProber::prefix_of_first_row_not_matching_prefix(self, cur, end, len).await + } +} + +#[cfg(test)] +mod tests { + use mysql_async::prelude::Queryable; + use mz_ore::cast::CastFrom; + + use super::*; + use crate::probe::tests::{connect, drop_db, setup_table}; + + /// In-memory [`PrimaryKeyProber`] over a sorted key list with exact + /// "estimates". Byte order stands in for the collation, so the PAD SPACE + /// below-space cases are deliberately out of scope here, the live tests + /// cover them. + struct MockDb { + keys: Vec, + } + + impl MockDb { + fn bounds(&self, start: &str, end: Option<&str>) -> (usize, usize) { + // The lower bound is exclusive, a key equal to `start` is skipped. + let lo = self.keys.partition_point(|k| k.as_str() <= start); + let hi = match end { + Some(e) => self.keys.partition_point(|k| k.as_str() < e), + None => self.keys.len(), + }; + (lo, hi.max(lo)) + } + } + + impl PrimaryKeyProber for MockDb { + async fn estimate_range_rows( + &mut self, + start: &str, + end: Option<&str>, + ) -> Result { + let (lo, hi) = self.bounds(start, end); + Ok(u64::cast_from(hi - lo)) + } + + async fn prefix_of_first_key_in_range( + &mut self, + start: &str, + end: Option<&str>, + len: usize, + ) -> Result, MySqlError> { + let (lo, hi) = self.bounds(start, end); + if lo >= hi { + return Ok(None); + } + Ok(Some(self.keys[lo].chars().take(len).collect())) + } + + async fn prefix_of_first_row_not_matching_prefix( + &mut self, + cur: &str, + end: Option<&str>, + len: usize, + ) -> Result, MySqlError> { + let (_, hi) = self.bounds("", end); + // Find the last key matching `cur`, byte prefixes stand in for + // the collation's LIKE matching. + let Some(last_match) = self.keys[..hi].iter().rposition(|k| k.starts_with(cur)) else { + return Ok(None); + }; + Ok(self.keys[last_match + 1..hi] + .first() + .map(|k| k.chars().take(len).collect())) + } + } + + fn keys(n: usize) -> Vec { + (0..n).map(|i| format!("{i:06}")).collect() + } + + const MIN_ROWS_PER_WORKER: u64 = 50_000; + + #[mz_ore::test(tokio::test)] + async fn single_worker_gets_no_boundaries() -> Result<(), MySqlError> { + let mut db = MockDb { keys: keys(1000) }; + let count = u64::cast_from(db.keys.len()); + let boundaries = partition(&mut db, 1, count, MIN_ROWS_PER_WORKER).await?; + assert!(boundaries.is_empty()); + Ok(()) + } + + #[mz_ore::test(tokio::test)] + async fn small_table_gets_no_boundaries() -> Result<(), MySqlError> { + // All keys share one depth-1 prefix and fit under `min_rows_per_worker`, + // so the single open-ended range yields no boundary. + let mut db = MockDb { keys: keys(10_000) }; + let count = u64::cast_from(db.keys.len()); + let boundaries = partition(&mut db, 4, count, MIN_ROWS_PER_WORKER).await?; + assert!(boundaries.is_empty()); + Ok(()) + } + + #[mz_ore::test(tokio::test)] + async fn empty_table_gets_no_boundaries() -> Result<(), MySqlError> { + let mut db = MockDb { keys: vec![] }; + let boundaries = partition(&mut db, 4, 0, MIN_ROWS_PER_WORKER).await?; + assert!(boundaries.is_empty()); + Ok(()) + } + + #[mz_ore::test(tokio::test)] + async fn splits_evenly_across_workers() -> Result<(), MySqlError> { + let mut db = MockDb { + keys: keys(200_000), + }; + let count = u64::cast_from(db.keys.len()); + let boundaries = partition(&mut db, 4, count, MIN_ROWS_PER_WORKER).await?; + assert_eq!(boundaries.len(), 3); + // Boundaries must be sorted and split the keys into ~50k chunks. + let mut prev = 0; + for b in &boundaries { + let idx = db.keys.partition_point(|k| k.as_str() < b.as_str()); + let share = idx - prev; + assert!( + (40_000..=60_000).contains(&share), + "uneven share {share} at boundary {b:?} (all: {boundaries:?})", + ); + prev = idx; + } + assert!((40_000..=60_000).contains(&(db.keys.len() - prev))); + Ok(()) + } + + #[mz_ore::test(tokio::test)] + async fn low_min_rows_per_worker_splits_small_tables() -> Result<(), MySqlError> { + let mut db = MockDb { keys: keys(1000) }; + let count = u64::cast_from(db.keys.len()); + let boundaries = partition(&mut db, 4, count, 10).await?; + assert_eq!(boundaries.len(), 3); + let mut prev = 0; + for b in &boundaries { + let idx = db.keys.partition_point(|k| k.as_str() < b.as_str()); + let share = idx - prev; + assert!( + (150..=350).contains(&share), + "uneven share {share} at boundary {b:?} (all: {boundaries:?})", + ); + prev = idx; + } + Ok(()) + } + + #[mz_ore::test(tokio::test)] + async fn short_key_does_not_block_splitting() -> Result<(), MySqlError> { + // One key is a bare "U" and every other key extends it. The walk + // skips the exact key (exclusive lower bounds) and must keep + // splitting inside the extensions at greater depths instead of + // stalling on the all-encompassing "U" prefix. + let mut all_keys = vec!["U".to_string()]; + all_keys.extend((0..1000).map(|i| format!("U{i:06}"))); + let mut db = MockDb { keys: all_keys }; + let count = u64::cast_from(db.keys.len()); + let boundaries = partition(&mut db, 4, count, 10).await?; + assert_eq!(boundaries.len(), 3); + for b in &boundaries { + assert!( + b.starts_with('U') && b.len() > 1, + "boundary {b:?} does not subdivide the extensions (all: {boundaries:?})" + ); + } + Ok(()) + } + + #[mz_ore::test(tokio::test)] + async fn fractional_target_still_terminates() -> Result<(), MySqlError> { + // count / (workers * 4) is fractional and the minimum is zero, so + // the target floors at one row instead of splitting forever. + let mut db = MockDb { keys: keys(3) }; + let boundaries = partition(&mut db, 4, 3, 0).await?; + assert_eq!(boundaries, vec!["000001", "000002"]); + Ok(()) + } + + // Live tests against MySQL (when available) for more realistic results. + + /// Splitting must reach inside the extensions of the bare key 'a' and + /// yield boundaries MySQL agrees are strictly increasing. + #[mz_ore::test(tokio::test)] + #[cfg_attr(miri, ignore)] + async fn skewed_partitions_with_wildcards_and_short_keys() -> Result<(), anyhow::Error> { + let Some(mut conn) = connect().await? else { + return Ok(()); + }; + + // A bare key 'a' that 900 keys extend, 100 keys under 'b', and LIKE + // metacharacters. + let mut all_keys = vec![ + "a".to_string(), + "c_1".to_string(), + "c%2".to_string(), + "c\\3".to_string(), + "c|4".to_string(), + ]; + all_keys.extend((0..900).map(|i| format!("a{i:05}"))); + all_keys.extend((0..100).map(|i| format!("b{i:05}"))); + + const DB: &str = "mz_partition_test"; + let table = setup_table(&mut conn, DB, "utf8mb4_bin", &all_keys).await?; + let total = u64::cast_from(all_keys.len()); + + // A minimum above the table size yields no boundaries at all. + let bounds = partition_table(&mut conn, table.clone(), "id", 4, total, 50_000).await?; + assert!(bounds.is_empty(), "{bounds:?}"); + + // A low minimum splits inside the 'a' extensions rather than stopping + // at the exact key. + let bounds = partition_table(&mut conn, table, "id", 4, total, 10).await?; + assert_eq!(bounds.len(), 3, "{bounds:?}"); + + // MySQL agrees the boundaries are strictly increasing. + for pair in bounds.windows(2) { + let increasing: Option = conn + .exec_first("SELECT ? < ?", (&pair[0], &pair[1])) + .await?; + assert_eq!(increasing, Some(1), "{bounds:?}"); + } + let counts = partition_counts(&mut conn, DB, &bounds, total).await?; + // ~1/4 of the table is 250 so 100 leaves lots of room for error. + assert!(counts.iter().all(|&c| c > 100), "{counts:?}"); + + drop_db(&mut conn, DB).await?; + conn.disconnect().await?; + Ok(()) + } + + #[mz_ore::test(tokio::test)] + #[cfg_attr(miri, ignore)] + async fn skew_empty_string_and_below_space_characters_inaccuracy() -> Result<(), anyhow::Error> + { + let Some(mut conn) = connect().await? else { + return Ok(()); + }; + + // ~10k keys in total, about as many under "c" as the rest combined. + // Tabs and empty strings will be dropped, which will show up in the resulting skew. + let mut all_keys = vec![String::new()]; + add_1k_keys(&mut all_keys, "\t"); + add_1k_keys(&mut all_keys, "a"); + add_1k_keys(&mut all_keys, "b"); + add_1k_keys(&mut all_keys, "b\t"); + add_1k_keys(&mut all_keys, "c"); + add_1k_keys(&mut all_keys, "ca"); + add_1k_keys(&mut all_keys, "cb"); + add_1k_keys(&mut all_keys, "cc"); + add_1k_keys(&mut all_keys, "cd"); + add_1k_keys(&mut all_keys, "d"); + + const DB: &str = "mz_partition_live_mixed_test"; + let table = setup_table(&mut conn, DB, "utf8mb4_bin", &all_keys).await?; + let total = u64::cast_from(all_keys.len()); + + // Partition for 4 workers with a minimum split size around 250. + let bounds = partition_table(&mut conn, table, "id", 4, total, 250).await?; + assert_eq!(bounds.len(), 3); + let counts = partition_counts(&mut conn, DB, &bounds, total).await?; + // ~8k keys are visible, so each count gets at least 2k under perfect + // partitioning, and the ranges partition cleanly except for the + // hidden tab prefixes. Asserting each count above 1600 makes room + // for single partitions being misallocated (~250) and some + // inaccuracy on top of that (~150). + assert!(counts.iter().all(|&c| c > 1600), "{counts:?}"); + + // Each hidden group piles into the partition left of the next visible + // boundary, here all of them ('', tabs, b-tabs) land in the first. Keep + // the assertion low to ensure there's room for estimate variability. + // This is a performance degradation edge case, not a correctness + // issue. + assert!(counts[0] > 2600, "{counts:?}"); + + drop_db(&mut conn, DB).await?; + conn.disconnect().await?; + Ok(()) + } + + fn add_1k_keys(all_keys: &mut Vec, prefix: &str) { + all_keys.extend((0..1000).map(|i| format!("{prefix}{i:03}"))); + } + + /// Rows per snapshot partition of `bounds`, i.e. the half-open ranges + /// `[..b0), [b0, b1), .., [bn, ..)`. The server counts, so the + /// comparisons happen under the column's collation. + async fn partition_counts( + conn: &mut mysql_async::Conn, + db: &str, + bounds: &[String], + total: u64, + ) -> Result, anyhow::Error> { + let mut counts = Vec::with_capacity(bounds.len() + 1); + let mut below = 0; + for bound in bounds { + let cumulative: Option = conn + .exec_first( + format!("SELECT COUNT(*) FROM {db}.t WHERE id < ?"), + (bound.as_str(),), + ) + .await?; + let cumulative = cumulative.expect("COUNT returns a row"); + counts.push(cumulative - below); + below = cumulative; + } + counts.push(total - below); + Ok(counts) + } +} diff --git a/src/mysql-util/src/probe.rs b/src/mysql-util/src/probe.rs index c9d1405105880..66df0c157a811 100644 --- a/src/mysql-util/src/probe.rs +++ b/src/mysql-util/src/probe.rs @@ -259,8 +259,9 @@ where Ok(estimate) } +/// The live MySQL harness here is shared with [`crate::partition`]'s tests. #[cfg(test)] -mod tests { +pub(crate) mod tests { use std::collections::BTreeSet; use mz_ore::cast::CastFrom; @@ -969,7 +970,7 @@ mod tests { /// Connects to the server named by `MZ_TEST_MYSQL_URL`, or `None` to skip /// the test when it is unset. Skipping is a local-only convenience, CI /// must always provide the URL. - async fn connect() -> Result, anyhow::Error> { + pub(crate) async fn connect() -> Result, anyhow::Error> { let Ok(url) = std::env::var("MZ_TEST_MYSQL_URL") else { if mz_ore::env::is_var_truthy("CI") { panic!("CI is supposed to run this test but something has gone wrong!"); @@ -997,7 +998,7 @@ mod tests { /// Recreates scratch database `db` holding one table `t` whose string /// primary key `id` is pinned to the given `collation`, containing /// `keys`, with fresh statistics. Returns a ref for [`KeyProber::new`]. - async fn setup_table<'a>( + pub(crate) async fn setup_table<'a>( conn: &mut mysql_async::Conn, db: &'a str, collation: &str, @@ -1013,11 +1014,20 @@ mod tests { COLLATE {collation} PRIMARY KEY NOT NULL)" )) .await?; - conn.exec_batch( - format!("INSERT INTO {db}.t VALUES (?)"), - keys.iter().map(|id| (id.as_ref(),)), - ) - .await?; + + for chunk in keys.chunks(1000) { + conn.exec_drop( + format!( + "INSERT INTO {db}.t VALUES {}", + vec!["(?)"; chunk.len()].join(",") + ), + chunk + .iter() + .map(|id| id.as_ref().into()) + .collect::>(), + ) + .await?; + } #[allow(clippy::disallowed_methods)] conn.query_drop(format!("ANALYZE TABLE {db}.t")).await?; Ok(QualifiedTableRef { @@ -1027,7 +1037,10 @@ mod tests { } /// Drops the scratch database `db`. - async fn drop_db(conn: &mut mysql_async::Conn, db: &str) -> Result<(), anyhow::Error> { + pub(crate) async fn drop_db( + conn: &mut mysql_async::Conn, + db: &str, + ) -> Result<(), anyhow::Error> { #[allow(clippy::disallowed_methods)] conn.query_drop(format!("DROP DATABASE {db}")).await?; Ok(())