diff --git a/src/mysql-util/src/lib.rs b/src/mysql-util/src/lib.rs index 6537ee2ffa876..ba1620d02e46f 100644 --- a/src/mysql-util/src/lib.rs +++ b/src/mysql-util/src/lib.rs @@ -45,6 +45,9 @@ pub use decoding::pack_mysql_row; pub mod probe; pub use probe::{KeyProber, MAX_KEY_LENGTH}; +pub mod partition; +pub use partition::partition_table; + mod aws_rds; #[derive(Debug, Clone)] @@ -108,6 +111,16 @@ pub enum MySqlError { column_name: String, error: String, }, + #[error( + "missing row estimate in '{qualified_table_name}' for key range ({lower_bound}, {upper_bound})" + )] + MissingRowEstimate { + qualified_table_name: String, + /// Redacted at construction, safe to log. + lower_bound: String, + /// Redacted at construction, safe to log. + upper_bound: String, + }, #[error("unsupported data types: {columns:?}")] UnsupportedDataTypes { columns: Vec }, #[error("duplicated column names in table '{qualified_table_name}': {columns:?}")] diff --git a/src/mysql-util/src/partition.rs b/src/mysql-util/src/partition.rs new file mode 100644 index 0000000000000..48d9bf0a2ada8 --- /dev/null +++ b/src/mysql-util/src/partition.rs @@ -0,0 +1,190 @@ +// Copyright Materialize, Inc. and contributors. All rights reserved. +// +// Use of this software is governed by the Business Source License +// included in the LICENSE file. +// +// As of the Change Date specified in that file, in accordance with +// the Business Source License, use of this software will be governed +// by the Apache License, Version 2.0. + +use mz_ore::cast::CastFrom; +use mz_ore::str::redact; + +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. +/// `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 +/// avoid unnecessary work for smaller tables limiting the overhead of partitioning. +pub async fn partition_table( + conn: &mut mysql_async::Conn, + table: QualifiedTableRef<'_>, + pk_col: &str, + num_workers: usize, + estimated_row_count: u64, + min_split_threshold: u64, +) -> Result, MySqlError> { + let (schema_name, table_name) = (table.schema_name, table.table_name); + let mut db = KeyProber::new(conn, table, pk_col); + let boundaries = partition( + &mut db, + num_workers, + estimated_row_count, + min_split_threshold, + ) + .await?; + tracing::trace!( + schema = schema_name, + table = table_name, + // The boundaries are user data, redacted outside of CI. + boundaries = ?redact(&boundaries), + "partitioned table by pk prefix" + ); + Ok(boundaries) +} + +#[derive(Debug)] +struct Prefix { + /// Empty for the beginning of the key space. + prefix: String, + /// Exclusive end, `None` for the final open prefix. + end: Option, + /// Row estimate for the prefix, at least 1. + estimated_rows: u64, + /// Length this prefix was split at. + depth: usize, +} + +async fn partition( + db: &mut KeyProber<'_>, + workers: usize, + estimated_row_count: u64, + min_split_threshold: u64, +) -> Result, MySqlError> { + if workers <= 1 { + return Ok(Vec::new()); + } + let estimated_row_count = estimated_row_count.max(1); + + // Estimates vary wildly especially near the full table size (see `KeyProber::estimate_range_rows` for more details). + // Estimates tend to get more useful as smaller chunks, so break up the table into at least 1/8ths (2 workers * 4) + // before selecting partitions. 1/8th was selected by feel due to a couple of observed inaccuracies: + // 1. Large estimates were observed as capped at 1/2 the estimated table size when the estimates were big. + // 2. Medium or approaching 1/2 estimated table size estimates were observed as large overestimates (~2x) + // So, that's a potential 4x swing and then a 2x safety factor to not push too close to the edge. + // + // Breaking down to smaller partitions results in more accurate splits, so we keep the + // 4x multiple of the worker count for > 2 workers. Initial testing was with an 8x multiplier, selected + // arbitrarily. From first principles, you can expect that if a prefix containing ~target_max_rows_per_prefix rows + // lands right on a boundary (i.e. the worker was 99% full for its range) the worker will get a slot worth + // 99% + 1/multiplier (in this case 25%) of the normal worker share resulting in skew with ~124% of the rows it + // should own. + let target_max_rows_per_prefix = (estimated_row_count / u64::cast_from(workers * 4)) + .max(min_split_threshold) + .max(1); + + compute_boundaries(db, workers, estimated_row_count, target_max_rows_per_prefix).await +} + +async fn compute_boundaries( + db: &mut KeyProber<'_>, + workers: usize, + estimated_row_count: u64, + target_rows_per_prefix: u64, +) -> Result, MySqlError> { + // BFS of prefixes, splitting until estimates fall under the target. + let mut ordered_prefixes = vec![Prefix { + prefix: String::new(), + end: None, + estimated_rows: estimated_row_count, + depth: 0, + }]; + + loop { + let mut next_ordered_prefixes: Vec = vec![]; + let mut split_any = false; + for prefix in ordered_prefixes { + if prefix.estimated_rows > target_rows_per_prefix { + split_any = true; + // Partitioning children can drop some rows from the parent prefix range. This + // is acceptable given the approximate nature of the algorithm. + let children = children_prefixes(db, &prefix).await?; + next_ordered_prefixes.extend(children); + } else { + next_ordered_prefixes.push(prefix); + } + } + ordered_prefixes = next_ordered_prefixes; + if !split_any { + break; + } + } + + // Recompute the total after partitioning the table to get more even splits because the actual row count and the + // granularly estimated row count can diverge from the original top level estimate. + let total: u64 = ordered_prefixes.iter().map(|r| r.estimated_rows).sum(); + let per_worker = total / u64::cast_from(workers); + tracing::debug!( + prefixes = ordered_prefixes.len(), + total_estimated_rows = total, + per_worker, + "assigning prefixes to workers" + ); + let mut boundaries: Vec = Vec::with_capacity(workers - 1); + let mut rows_seen = 0; + for prefix in &ordered_prefixes { + if boundaries.len() == workers - 1 { + break; + } + rows_seen += prefix.estimated_rows; + if rows_seen >= u64::cast_from(boundaries.len() + 1) * per_worker { + // The final prefix's end is None (open), it can never be a boundary. + if let Some(end) = &prefix.end { + boundaries.push(end.clone()); + } + } + } + Ok(boundaries) +} + +/// Splits `parent` into prefixes one character longer. i.e. prefix "a", upper bound "b" in table +/// with pks: ["a", "ab", "abc", "abd", "af", "bb"] will return: ["ab", "af"]. +/// +/// 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<'_>, + parent: &Prefix, +) -> Result, MySqlError> { + let depth = parent.depth + 1; + let mut children = Vec::new(); + + let Some(mut cur) = db + .prefix_of_first_key_in_range(&parent.prefix, parent.end.as_deref(), depth) + .await? + else { + return Ok(children); + }; + + loop { + let next = db + .prefix_of_first_row_not_matching_prefix(&cur, parent.end.as_deref(), depth) + .await?; + let end = next.clone().or_else(|| parent.end.clone()); + let estimated_rows = db.estimate_range_rows(&cur, end.as_deref()).await?; + children.push(Prefix { + prefix: cur, + end, + estimated_rows: estimated_rows.max(1), + depth, + }); + match next { + Some(next) => cur = next, + None => return Ok(children), + } + } +} diff --git a/src/mysql-util/src/probe.rs b/src/mysql-util/src/probe.rs index 439c979c73711..c9d1405105880 100644 --- a/src/mysql-util/src/probe.rs +++ b/src/mysql-util/src/probe.rs @@ -9,6 +9,7 @@ use mysql_async::prelude::Queryable; use mysql_async::{Params, Value}; +use mz_ore::str::redact; use crate::{MySqlError, QualifiedTableRef, quote_identifier}; @@ -70,7 +71,7 @@ impl<'a> KeyProber<'a> { &mut self, lower_bound_exclusive: &str, upper_bound_exclusive: Option<&str>, - ) -> Result, MySqlError> { + ) -> Result { let (clause, params) = self.range_filter(Some(lower_bound_exclusive), upper_bound_exclusive); let select = format!( @@ -78,7 +79,15 @@ impl<'a> KeyProber<'a> { col = self.col, table = self.table, ); - explain_row_estimate(&mut *self.conn, &select, Params::Positional(params)).await + explain_row_estimate(&mut *self.conn, &select, Params::Positional(params)) + .await? + .ok_or_else(|| MySqlError::MissingRowEstimate { + qualified_table_name: self.table_name.clone(), + // The bounds are column values, redact them so the error + // stays loggable outside of CI. + lower_bound: format!("{:?}", redact(&lower_bound_exclusive)), + upper_bound: format!("{:?}", redact(&upper_bound_exclusive)), + }) } /// Grabs a prefix of up to `max_prefix_length` characters for the first @@ -341,14 +350,11 @@ mod tests { // Estimates are index dives, near reality but never exact by // contract, so the bounds are deliberately loose. - let all = p.estimate_range_rows("", None).await?.expect("estimate"); + let all = p.estimate_range_rows("", None).await?; assert!((500..=2000).contains(&all), "all={all}"); - let half = p - .estimate_range_rows("a00500", None) - .await? - .expect("estimate"); + let half = p.estimate_range_rows("a00500", None).await?; assert!((250..=1000).contains(&half), "half={half}"); - let none = p.estimate_range_rows("zzz", None).await?.expect("estimate"); + let none = p.estimate_range_rows("zzz", None).await?; assert!(none <= 5, "none={none}"); drop_db(&mut conn, DB).await?; @@ -817,15 +823,9 @@ mod tests { // Range estimates come from index dives on the real B-tree, not the // stale table statistics, so they still reflect the actual data. - let all = prober - .estimate_range_rows("", None) - .await? - .expect("estimate"); + let all = prober.estimate_range_rows("", None).await?; assert!((500..=2000).contains(&all), "all={all}"); - let range = prober - .estimate_range_rows("a00100", Some("a00200")) - .await? - .expect("estimate"); + let range = prober.estimate_range_rows("a00100", Some("a00200")).await?; assert!((50..=200).contains(&range), "range={range}"); drop_db(&mut conn, DB).await?;