From 29f909629c0759fdc1b129447d7e1325120c9d5c Mon Sep 17 00:00:00 2001 From: Peter Larsen Date: Mon, 3 Aug 2026 20:34:35 -0400 Subject: [PATCH 01/10] mysql-util: add a PK-prefix partitioner for parallel snapshots Discovers boundaries that split a table string primary key space into per-worker ranges of roughly equal estimated row counts. Ranges the optimizer estimates too large are recursively subdivided at each distinct key prefix one character longer, probing through KeyProber, then accumulated into per-worker buckets, so discovery costs EXPLAIN index dives instead of an O(rows) index pass. Inaccurate estimates skew bucket sizes but never correctness: any ordered boundary list partitions the key space. All key ordering happens server-side under the column collation. The walk guards against non-advancing prefixes and caps children per split so a misbehaving server cannot hang it. KeyProber steps past exact keys shorter than the prefix length, so a lone short key among keys extending it cannot leave a range unsplittable. Also documents the caller contracts on like_prefix_pattern and explain_row_estimate. --- src/mysql-util/src/lib.rs | 3 + src/mysql-util/src/partition.rs | 316 ++++++++++++++++++++++++++++++++ src/mysql-util/src/probe.rs | 9 + 3 files changed, 328 insertions(+) create mode 100644 src/mysql-util/src/partition.rs diff --git a/src/mysql-util/src/lib.rs b/src/mysql-util/src/lib.rs index 6537ee2ffa876..61f4ecfeaa54c 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_by_pk_prefix; + mod aws_rds; #[derive(Debug, Clone)] diff --git a/src/mysql-util/src/partition.rs b/src/mysql-util/src/partition.rs new file mode 100644 index 0000000000000..b3da0b66459ac --- /dev/null +++ b/src/mysql-util/src/partition.rs @@ -0,0 +1,316 @@ +// 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. + +//! Discovers boundaries that split a table's string primary key space into +//! ranges of roughly equal row counts, for parallel snapshot reads. +//! +//! Ranges are split on character prefixes: a range the optimizer estimates +//! too large is subdivided at each distinct key prefix one character longer +//! than the prefix it was last split at, recursively, until every range is +//! small. Adjacent small ranges are then accumulated back into one bucket per +//! worker. Row estimates come from `EXPLAIN`, so probing a range costs an +//! index dive instead of a scan, and inaccurate estimates skew bucket sizes +//! but never correctness: any ordered boundary list partitions the key space. +//! +//! All key matching and ordering happens server-side under the key column's +//! own collation. Rust never orders the returned prefixes, it only checks +//! them for byte equality as a termination guard. + +use mz_ore::cast::CastLossy; + +use crate::{KeyProber, MySqlError, QualifiedTableRef}; + +/// Longest key prefix to split on, bounds the refinement loop. +const MAX_DEPTH: usize = 16; +/// Hard cap on children from a single split, bounds the walk even if the +/// server's ordering misbehaves. +const MAX_CHILDREN_PER_SPLIT: usize = 4096; + +/// Computes up to `workers - 1` exclusive upper bounds, in key order, that +/// split `table` into per-worker key ranges of roughly equal estimated row +/// counts. Worker `i` reads keys in `[boundaries[i - 1], boundaries[i])`, +/// with the first and last range open ended. Returns an empty list when the +/// table is too small to be worth splitting. +/// +/// `pk_col` is the raw (unquoted) name of the key column. It must be a string +/// column, prefixes of a numeric column do not order consistently with its +/// values, and it should be the leading column of an index or every probe +/// becomes a full table scan. `estimated_row_count` seeds the bucket sizing +/// and may be approximate. No range is split below `min_bucket_rows` rows. +pub async fn partition_table_by_pk_prefix( + conn: &mut mysql_async::Conn, + table: QualifiedTableRef<'_>, + pk_col: &str, + workers: usize, + estimated_row_count: u64, + min_bucket_rows: 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, workers, estimated_row_count, min_bucket_rows).await?; + tracing::trace!( + schema = schema_name, + table = table_name, + ?boundaries, + "partitioned table by pk prefix" + ); + Ok(boundaries) +} + +/// What the partitioner needs from the database. Keys are treated as strings +/// and split by prefix, so implementations must order prefixes consistently +/// with the full keys they abbreviate. +/// +/// Bounds are exclusive on both sides and optional: a `start` of `None` +/// means the beginning of the key space, an `end` of `None` means unbounded. +/// A key exactly equal to a bound or prefix argument is skipped, its +/// extensions still surface through the exclusive bound. +trait PartitionDb { + /// Row count estimate for keys in `(start, end)`. May be arbitrarily + /// inaccurate. + async fn estimate_range_rows( + &mut self, + start: Option<&str>, + end: Option<&str>, + ) -> Result; + + /// The prefix of up to `len` characters of the first key in + /// `(start, end)`, or `None` if the range holds no rows. + async fn first_prefix( + &mut self, + start: Option<&str>, + end: Option<&str>, + len: usize, + ) -> Result, MySqlError>; + + /// The prefix of up to `len` characters of the first key past the last + /// key matching `cur`, taken from keys below `end`. `None` if no such + /// key exists. Every key extending `cur` is covered by the match, + /// including `cur` itself as an exact key. + async fn next_prefix( + &mut self, + cur: &str, + end: Option<&str>, + len: usize, + ) -> Result, MySqlError>; +} + +/// A half open key range `[start, end)` tracked by the splitting loop. +#[derive(Debug)] +struct Range { + /// Exclusive start, `None` at the beginning of the key space. + start: Option, + /// Exclusive end, `None` for the final open range. + end: Option, + /// Row estimate for the range, at least 1. + estimated_rows: f64, + /// Prefix length this range was split at. + depth: usize, +} + +async fn partition( + db: &mut D, + workers: usize, + estimated_row_count: u64, + min_bucket_rows: u64, +) -> Result, MySqlError> { + if workers <= 1 { + return Ok(Vec::new()); + } + let estimated_row_count = f64::cast_lossy(estimated_row_count.max(1)); + let min_bucket_rows = f64::cast_lossy(min_bucket_rows.max(1)); + + // Aim for more buckets than workers. With exactly one bucket per worker + // an under-estimated table size could make the whole table look too small + // to split. Small buckets are recombined when boundaries are assigned. + let mut buckets = workers; + while buckets < 8 { + buckets *= 2; + } + let target_bucket_rows = (estimated_row_count / f64::cast_lossy(buckets)).max(min_bucket_rows); + // Split ranges down to well below the bucket size so that accumulated + // bucket boundaries can land close to their targets. + let target_split_rows = (target_bucket_rows / 8.0).max(min_bucket_rows); + tracing::debug!( + estimated_row_count, + buckets, + target_bucket_rows, + target_split_rows, + "partitioning key space by prefix" + ); + + let mut ranges = vec![Range { + start: None, + end: None, + estimated_rows: estimated_row_count, + depth: 0, + }]; + loop { + let mut split_any = false; + let mut next: Vec = Vec::with_capacity(ranges.len()); + for range in ranges { + if range.estimated_rows > target_split_rows && range.depth < MAX_DEPTH { + split_any = true; + // An empty child list means the range holds no rows (a + // phantom estimate) and is dropped. A single child spanning + // the whole parent comes back with a greater depth, so + // retrying it splits on a longer prefix and MAX_DEPTH bounds + // the loop. + next.extend(split_range(db, &range, target_split_rows).await?); + } else { + next.push(range); + } + } + ranges = next; + if !split_any { + break; + } + } + + // Ranges are in key order. Emit a boundary each time the cumulative + // estimated rows pass the next worker's share. + let total: f64 = ranges.iter().map(|r| r.estimated_rows).sum(); + let per_worker = total / f64::cast_lossy(workers); + tracing::debug!( + ranges = ranges.len(), + total_estimated_rows = total, + per_worker, + "assigning prefix ranges to workers" + ); + let mut boundaries: Vec = Vec::with_capacity(workers - 1); + let mut rows_seen = 0.0; + for range in &ranges { + if boundaries.len() == workers - 1 { + break; + } + rows_seen += range.estimated_rows; + if rows_seen >= f64::cast_lossy(boundaries.len() + 1) * per_worker { + // The final range's end is None (open), it can never be a boundary. + if let Some(end) = &range.end { + // Only a misbehaving server can repeat an end (the walk stops + // on non-advancing prefixes). Skip it, a duplicate boundary + // would fail the strict monotonicity check downstream. + if boundaries.last() != Some(end) { + boundaries.push(end.clone()); + } + } + } + } + Ok(boundaries) +} + +/// Splits `parent` at every distinct key prefix one character longer than the +/// prefix `parent` was split at, except that a tail whose estimate already +/// fits `target_rows` stays a single child. +async fn split_range( + db: &mut D, + parent: &Range, + target_rows: f64, +) -> Result, MySqlError> { + let len = parent.depth + 1; + let mut children = Vec::new(); + + let Some(mut cur) = db + .first_prefix(parent.start.as_deref(), parent.end.as_deref(), len) + .await? + else { + return Ok(children); + }; + // The first child inherits the parent's start. + let mut start = parent.start.clone(); + + loop { + // If the remaining rows fit the target, emit them as one child and + // stop. The children cap likewise closes out the split with whatever + // remains. + let remaining = db + .estimate_range_rows(start.as_deref(), parent.end.as_deref()) + .await?; + let remaining = f64::cast_lossy(remaining).max(1.0); + if remaining <= target_rows || children.len() + 1 >= MAX_CHILDREN_PER_SPLIT { + children.push(Range { + start, + end: parent.end.clone(), + estimated_rows: remaining, + depth: len, + }); + return Ok(children); + } + // When `cur` is an exact key shorter than `len`, every key extending + // it matches `cur`, so next_prefix exhausts even though those + // extensions still need visiting. The first key past `cur` exposes + // them, so the walk can keep splitting instead of retrying the whole + // range at greater depths forever. + let next = match db.next_prefix(&cur, parent.end.as_deref(), len).await? { + Some(next) => Some(next), + None => { + db.first_prefix(Some(&cur), parent.end.as_deref(), len) + .await? + } + }; + // A prefix equal to `cur` cannot advance the walk (only a misbehaving + // server produces one), stop splitting here. + match next.filter(|next| next != &cur) { + Some(next) => { + let estimated_rows = db + .estimate_range_rows(start.as_deref(), Some(&next)) + .await?; + children.push(Range { + start: start.clone(), + end: Some(next.clone()), + estimated_rows: f64::cast_lossy(estimated_rows).max(1.0), + depth: len, + }); + start = Some(next.clone()); + cur = next; + } + None => { + children.push(Range { + start, + end: parent.end.clone(), + estimated_rows: remaining, + depth: len, + }); + return Ok(children); + } + } + } +} + +impl<'a> PartitionDb for KeyProber<'a> { + async fn estimate_range_rows( + &mut self, + start: Option<&str>, + end: Option<&str>, + ) -> Result { + // A missing optimizer estimate reads as an empty range, which the + // splitting loop drops or leaves unsplit. + Ok(KeyProber::estimate_range_rows(self, start, end) + .await? + .unwrap_or(0)) + } + + async fn first_prefix( + &mut self, + start: Option<&str>, + end: Option<&str>, + len: usize, + ) -> Result, MySqlError> { + KeyProber::prefix_of_first_key_in_range(self, start, end, len).await + } + + async fn next_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 + } +} diff --git a/src/mysql-util/src/probe.rs b/src/mysql-util/src/probe.rs index 439c979c73711..b957bf488e1bd 100644 --- a/src/mysql-util/src/probe.rs +++ b/src/mysql-util/src/probe.rs @@ -231,6 +231,13 @@ fn like_prefix_pattern(prefix: &str) -> String { pattern } +/// Runs `EXPLAIN` on `select` and returns the optimizer's estimate of rows +/// examined, from the `rows` column of the plan. +/// +/// The caller must pass a single-table `SELECT`. Joins and subqueries produce +/// multiple plan rows whose estimates do not combine additively, only the +/// first row is read. Returns `None` when the optimizer reports no estimate. +/// The estimate can be arbitrarily stale, callers must tolerate inaccuracy. async fn explain_row_estimate

( conn: &mut mysql_async::Conn, select: &str, @@ -239,6 +246,8 @@ async fn explain_row_estimate

( where P: Into + Send, { + // NOTE: The format must be pinned because newer MySQL versions default + // `explain_format` to `TREE`, which has no `rows` column. let plan: Option = conn .exec_first(format!("EXPLAIN FORMAT=TRADITIONAL {select}"), params) .await?; From d6b454e35e067f6bdfe6ec75a2fe266e2c564f3c Mon Sep 17 00:00:00 2001 From: Peter Larsen Date: Tue, 4 Aug 2026 15:23:26 -0400 Subject: [PATCH 02/10] mysql-util: rely on the request budget as the partition walk bound Drop MAX_DEPTH, MAX_CHILDREN_PER_SPLIT, and the non-advancing prefix guard. On healthy data the walk terminates because child ranges shrink and fresh estimates track them. The pathological cases (phantom estimates, misbehaving servers) will be bounded by the per-table request budget once it lands, rather than by per-mechanism caps. Reformulate bucket sizing as a per-worker share divided by BUCKETS_PER_WORKER, dropping the double-to-8 bucket floor for small worker counts. --- src/mysql-util/src/partition.rs | 128 +++++++++++++------------------- 1 file changed, 53 insertions(+), 75 deletions(-) diff --git a/src/mysql-util/src/partition.rs b/src/mysql-util/src/partition.rs index b3da0b66459ac..bdac980c02a0d 100644 --- a/src/mysql-util/src/partition.rs +++ b/src/mysql-util/src/partition.rs @@ -26,11 +26,9 @@ use mz_ore::cast::CastLossy; use crate::{KeyProber, MySqlError, QualifiedTableRef}; -/// Longest key prefix to split on, bounds the refinement loop. -const MAX_DEPTH: usize = 16; -/// Hard cap on children from a single split, bounds the walk even if the -/// server's ordering misbehaves. -const MAX_CHILDREN_PER_SPLIT: usize = 4096; +/// When partitioning the data, how many buckets per worker to target to +/// try to get more even splits. +const BUCKETS_PER_WORKER: f64 = 8.0; /// Computes up to `workers - 1` exclusive upper bounds, in key order, that /// split `table` into per-worker key ranges of roughly equal estimated row @@ -63,45 +61,6 @@ pub async fn partition_table_by_pk_prefix( Ok(boundaries) } -/// What the partitioner needs from the database. Keys are treated as strings -/// and split by prefix, so implementations must order prefixes consistently -/// with the full keys they abbreviate. -/// -/// Bounds are exclusive on both sides and optional: a `start` of `None` -/// means the beginning of the key space, an `end` of `None` means unbounded. -/// A key exactly equal to a bound or prefix argument is skipped, its -/// extensions still surface through the exclusive bound. -trait PartitionDb { - /// Row count estimate for keys in `(start, end)`. May be arbitrarily - /// inaccurate. - async fn estimate_range_rows( - &mut self, - start: Option<&str>, - end: Option<&str>, - ) -> Result; - - /// The prefix of up to `len` characters of the first key in - /// `(start, end)`, or `None` if the range holds no rows. - async fn first_prefix( - &mut self, - start: Option<&str>, - end: Option<&str>, - len: usize, - ) -> Result, MySqlError>; - - /// The prefix of up to `len` characters of the first key past the last - /// key matching `cur`, taken from keys below `end`. `None` if no such - /// key exists. Every key extending `cur` is covered by the match, - /// including `cur` itself as an exact key. - async fn next_prefix( - &mut self, - cur: &str, - end: Option<&str>, - len: usize, - ) -> Result, MySqlError>; -} - -/// A half open key range `[start, end)` tracked by the splitting loop. #[derive(Debug)] struct Range { /// Exclusive start, `None` at the beginning of the key space. @@ -129,19 +88,16 @@ async fn partition( // Aim for more buckets than workers. With exactly one bucket per worker // an under-estimated table size could make the whole table look too small // to split. Small buckets are recombined when boundaries are assigned. - let mut buckets = workers; - while buckets < 8 { - buckets *= 2; - } - let target_bucket_rows = (estimated_row_count / f64::cast_lossy(buckets)).max(min_bucket_rows); - // Split ranges down to well below the bucket size so that accumulated - // bucket boundaries can land close to their targets. - let target_split_rows = (target_bucket_rows / 8.0).max(min_bucket_rows); + let estimated_rows_per_worker = + (estimated_row_count / f64::cast_lossy(workers)).max(min_bucket_rows); + let target_rows_per_bucket = + (estimated_rows_per_worker / BUCKETS_PER_WORKER).max(min_bucket_rows); + tracing::debug!( estimated_row_count, - buckets, - target_bucket_rows, - target_split_rows, + workers, + estimated_rows_per_worker, + target_rows_per_bucket, "partitioning key space by prefix" ); @@ -155,14 +111,11 @@ async fn partition( let mut split_any = false; let mut next: Vec = Vec::with_capacity(ranges.len()); for range in ranges { - if range.estimated_rows > target_split_rows && range.depth < MAX_DEPTH { + if range.estimated_rows > target_rows_per_bucket { split_any = true; // An empty child list means the range holds no rows (a - // phantom estimate) and is dropped. A single child spanning - // the whole parent comes back with a greater depth, so - // retrying it splits on a longer prefix and MAX_DEPTH bounds - // the loop. - next.extend(split_range(db, &range, target_split_rows).await?); + // phantom estimate) and is dropped. + next.extend(split_range(db, &range, target_rows_per_bucket).await?); } else { next.push(range); } @@ -206,18 +159,18 @@ async fn partition( } /// Splits `parent` at every distinct key prefix one character longer than the -/// prefix `parent` was split at, except that a tail whose estimate already -/// fits `target_rows` stays a single child. +/// prefix `parent`, i.e. "a" depth: 1 for a table with keys "a", "aa", "aaa", "ab" would be split to +/// "a" depth: 2 estimate 1, "aa" depth: 2 estimate 2, "ab" depth 2, estimate 1. async fn split_range( db: &mut D, parent: &Range, target_rows: f64, ) -> Result, MySqlError> { - let len = parent.depth + 1; + let depth = parent.depth + 1; let mut children = Vec::new(); let Some(mut cur) = db - .first_prefix(parent.start.as_deref(), parent.end.as_deref(), len) + .first_prefix(parent.start.as_deref(), parent.end.as_deref(), depth) .await? else { return Ok(children); @@ -226,31 +179,30 @@ async fn split_range( let mut start = parent.start.clone(); loop { - // If the remaining rows fit the target, emit them as one child and - // stop. The children cap likewise closes out the split with whatever - // remains. + // Small optimization to return early if the remaining rows past the current start prefix + // fits within the threshold we're looking for. let remaining = db .estimate_range_rows(start.as_deref(), parent.end.as_deref()) .await?; let remaining = f64::cast_lossy(remaining).max(1.0); - if remaining <= target_rows || children.len() + 1 >= MAX_CHILDREN_PER_SPLIT { + if remaining <= target_rows { children.push(Range { start, end: parent.end.clone(), estimated_rows: remaining, - depth: len, + depth, }); return Ok(children); } - // When `cur` is an exact key shorter than `len`, every key extending + // When `cur` is an exact key shorter than `depth`, every key extending // it matches `cur`, so next_prefix exhausts even though those // extensions still need visiting. The first key past `cur` exposes // them, so the walk can keep splitting instead of retrying the whole // range at greater depths forever. - let next = match db.next_prefix(&cur, parent.end.as_deref(), len).await? { + let next = match db.next_prefix(&cur, parent.end.as_deref(), depth).await? { Some(next) => Some(next), None => { - db.first_prefix(Some(&cur), parent.end.as_deref(), len) + db.first_prefix(Some(&cur), parent.end.as_deref(), depth) .await? } }; @@ -265,7 +217,7 @@ async fn split_range( start: start.clone(), end: Some(next.clone()), estimated_rows: f64::cast_lossy(estimated_rows).max(1.0), - depth: len, + depth, }); start = Some(next.clone()); cur = next; @@ -275,7 +227,7 @@ async fn split_range( start, end: parent.end.clone(), estimated_rows: remaining, - depth: len, + depth, }); return Ok(children); } @@ -283,6 +235,32 @@ async fn split_range( } } +/// Wrapper around KeyProber for testing purposes. +trait PartitionDb { + async fn estimate_range_rows( + &mut self, + start: Option<&str>, + end: Option<&str>, + ) -> Result; + + /// A `start` of `None` means the beginning of the key space. + async fn first_prefix( + &mut self, + start: Option<&str>, + end: Option<&str>, + len: usize, + ) -> Result, MySqlError>; + + /// Every key extending `cur` is covered by the match, including `cur` + /// itself as an exact key. + async fn next_prefix( + &mut self, + cur: &str, + end: Option<&str>, + len: usize, + ) -> Result, MySqlError>; +} + impl<'a> PartitionDb for KeyProber<'a> { async fn estimate_range_rows( &mut self, From 5101bb2fb2f5d1c079d2bc0360a1000dde061128 Mon Sep 17 00:00:00 2001 From: Peter Larsen Date: Tue, 4 Aug 2026 15:26:38 -0400 Subject: [PATCH 03/10] mysql-util: split partition into target, split, and assignment stages partition() becomes a composition of three stages: bucket_target_rows (pure sizing math), split_into_ranges (the only stage touching PartitionDb), and assign_boundaries (pure bucket accumulation). The pure stages are now directly unit-testable and the signatures document the data flow. --- src/mysql-util/src/partition.rs | 35 ++++++++++++++++++++++++++------- 1 file changed, 28 insertions(+), 7 deletions(-) diff --git a/src/mysql-util/src/partition.rs b/src/mysql-util/src/partition.rs index bdac980c02a0d..845c7394bb41e 100644 --- a/src/mysql-util/src/partition.rs +++ b/src/mysql-util/src/partition.rs @@ -84,7 +84,14 @@ async fn partition( } let estimated_row_count = f64::cast_lossy(estimated_row_count.max(1)); let min_bucket_rows = f64::cast_lossy(min_bucket_rows.max(1)); + let target_rows_per_bucket = bucket_target_rows(workers, estimated_row_count, min_bucket_rows); + let ranges = split_into_ranges(db, estimated_row_count, target_rows_per_bucket).await?; + Ok(assign_boundaries(&ranges, workers)) +} +/// The estimated row count a range should be split down to before it stops +/// being subdivided. +fn bucket_target_rows(workers: usize, estimated_row_count: f64, min_bucket_rows: f64) -> f64 { // Aim for more buckets than workers. With exactly one bucket per worker // an under-estimated table size could make the whole table look too small // to split. Small buckets are recombined when boundaries are assigned. @@ -100,7 +107,16 @@ async fn partition( target_rows_per_bucket, "partitioning key space by prefix" ); + target_rows_per_bucket +} +/// Splits the whole key space into ranges of at most roughly +/// `target_rows_per_bucket` estimated rows each, returned in key order. +async fn split_into_ranges( + db: &mut D, + estimated_row_count: f64, + target_rows_per_bucket: f64, +) -> Result, MySqlError> { let mut ranges = vec![Range { start: None, end: None, @@ -125,9 +141,14 @@ async fn partition( break; } } + Ok(ranges) +} - // Ranges are in key order. Emit a boundary each time the cumulative - // estimated rows pass the next worker's share. +/// Accumulates `ranges` (in key order) into `workers` buckets of roughly +/// equal estimated rows and returns the bucket edges as boundaries. +fn assign_boundaries(ranges: &[Range], workers: usize) -> Vec { + // Emit a boundary each time the cumulative estimated rows pass the next + // worker's share. let total: f64 = ranges.iter().map(|r| r.estimated_rows).sum(); let per_worker = total / f64::cast_lossy(workers); tracing::debug!( @@ -138,7 +159,7 @@ async fn partition( ); let mut boundaries: Vec = Vec::with_capacity(workers - 1); let mut rows_seen = 0.0; - for range in &ranges { + for range in ranges { if boundaries.len() == workers - 1 { break; } @@ -146,16 +167,16 @@ async fn partition( if rows_seen >= f64::cast_lossy(boundaries.len() + 1) * per_worker { // The final range's end is None (open), it can never be a boundary. if let Some(end) = &range.end { - // Only a misbehaving server can repeat an end (the walk stops - // on non-advancing prefixes). Skip it, a duplicate boundary - // would fail the strict monotonicity check downstream. + // A server returning non-advancing prefixes can repeat an + // end. Skip it, a duplicate boundary would fail the strict + // monotonicity check downstream. if boundaries.last() != Some(end) { boundaries.push(end.clone()); } } } } - Ok(boundaries) + boundaries } /// Splits `parent` at every distinct key prefix one character longer than the From bd4eb5ab0f72ee26a877fc27c106c417179bc198 Mon Sep 17 00:00:00 2001 From: Peter Larsen Date: Tue, 4 Aug 2026 18:37:58 -0400 Subject: [PATCH 04/10] mysql-util: simplify partitioner naming and drop the dedup guard Rename the sizing knobs to say what they mean (TARGET_RANGES_PER_WORKER, min_rows_per_worker, target_max_rows_per_range), trim module and function docs to the essentials, and stop deduplicating repeated boundary ends. Duplicate ends only arise from non-advancing servers, and the snapshot layer validates boundary monotonicity server-side before using boundaries. --- src/mysql-util/src/partition.rs | 94 +++++++++++++++------------------ 1 file changed, 42 insertions(+), 52 deletions(-) diff --git a/src/mysql-util/src/partition.rs b/src/mysql-util/src/partition.rs index 845c7394bb41e..6e665e6661a15 100644 --- a/src/mysql-util/src/partition.rs +++ b/src/mysql-util/src/partition.rs @@ -10,48 +10,43 @@ //! Discovers boundaries that split a table's string primary key space into //! ranges of roughly equal row counts, for parallel snapshot reads. //! -//! Ranges are split on character prefixes: a range the optimizer estimates -//! too large is subdivided at each distinct key prefix one character longer -//! than the prefix it was last split at, recursively, until every range is -//! small. Adjacent small ranges are then accumulated back into one bucket per -//! worker. Row estimates come from `EXPLAIN`, so probing a range costs an -//! index dive instead of a scan, and inaccurate estimates skew bucket sizes -//! but never correctness: any ordered boundary list partitions the key space. -//! //! All key matching and ordering happens server-side under the key column's -//! own collation. Rust never orders the returned prefixes, it only checks -//! them for byte equality as a termination guard. +//! own collation. Rust never orders the returned prefixes. use mz_ore::cast::CastLossy; use crate::{KeyProber, MySqlError, QualifiedTableRef}; -/// When partitioning the data, how many buckets per worker to target to -/// try to get more even splits. -const BUCKETS_PER_WORKER: f64 = 8.0; +/// When partitioning the data, how many ranges per worker should we break the +/// keyspace into. This helps avoid underestimates for large row counts resulting +/// in severe skew. +const TARGET_RANGES_PER_WORKER: f64 = 8.0; -/// Computes up to `workers - 1` exclusive upper bounds, in key order, that -/// split `table` into per-worker key ranges of roughly equal estimated row -/// counts. Worker `i` reads keys in `[boundaries[i - 1], boundaries[i])`, -/// with the first and last range open ended. Returns an empty list when the -/// table is too small to be worth splitting. +/// Computes up to `num_workers - 1` partition boundaries that divide the primary key space +/// into `num_workers` roughly even partitions. The boundaries are ordered according to MySQL's +/// internal collation rules. This should only be used against a single CHAR(N) or VARCHAR +/// primary key column. /// -/// `pk_col` is the raw (unquoted) name of the key column. It must be a string -/// column, prefixes of a numeric column do not order consistently with its -/// values, and it should be the leading column of an index or every probe -/// becomes a full table scan. `estimated_row_count` seeds the bucket sizing -/// and may be approximate. No range is split below `min_bucket_rows` rows. +/// NOTE: Run this inside a REPEATABLE READ transaction. The underlying probes +/// read separate snapshots otherwise, and concurrent inserts can make the +/// prefix walk see the same prefix repeatedly instead of advancing. pub async fn partition_table_by_pk_prefix( conn: &mut mysql_async::Conn, table: QualifiedTableRef<'_>, pk_col: &str, - workers: usize, + num_workers: usize, estimated_row_count: u64, - min_bucket_rows: u64, + min_rows_per_worker: 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, workers, estimated_row_count, min_bucket_rows).await?; + let boundaries = partition( + &mut db, + num_workers, + estimated_row_count, + min_rows_per_worker, + ) + .await?; tracing::trace!( schema = schema_name, table = table_name, @@ -77,41 +72,43 @@ async fn partition( db: &mut D, workers: usize, estimated_row_count: u64, - min_bucket_rows: u64, + min_rows_per_worker: u64, ) -> Result, MySqlError> { if workers <= 1 { return Ok(Vec::new()); } let estimated_row_count = f64::cast_lossy(estimated_row_count.max(1)); - let min_bucket_rows = f64::cast_lossy(min_bucket_rows.max(1)); - let target_rows_per_bucket = bucket_target_rows(workers, estimated_row_count, min_bucket_rows); - let ranges = split_into_ranges(db, estimated_row_count, target_rows_per_bucket).await?; + let min_rows_per_worker = f64::cast_lossy(min_rows_per_worker.max(1)); + let target_max_rows_per_range = + get_target_max_rows_per_range(workers, estimated_row_count, min_rows_per_worker); + // Should be many more ranges than workers unless the overall row count of the table is quite small. + let ranges = split_into_ranges(db, estimated_row_count, target_max_rows_per_range).await?; Ok(assign_boundaries(&ranges, workers)) } -/// The estimated row count a range should be split down to before it stops -/// being subdivided. -fn bucket_target_rows(workers: usize, estimated_row_count: f64, min_bucket_rows: f64) -> f64 { - // Aim for more buckets than workers. With exactly one bucket per worker - // an under-estimated table size could make the whole table look too small - // to split. Small buckets are recombined when boundaries are assigned. +fn get_target_max_rows_per_range( + workers: usize, + estimated_row_count: f64, + min_rows_per_worker: f64, +) -> f64 { + // Break up the key space into smaller ranges to more accurately rebuild the per-worker ranges later with less skew. let estimated_rows_per_worker = - (estimated_row_count / f64::cast_lossy(workers)).max(min_bucket_rows); - let target_rows_per_bucket = - (estimated_rows_per_worker / BUCKETS_PER_WORKER).max(min_bucket_rows); + (estimated_row_count / f64::cast_lossy(workers)).max(min_rows_per_worker); + // Respect min_rows_per_worker as a lower bound for the granularity with which we attempt to break up the table. + // No need to add this overhead for small tables. Estimate accuracy isn't super clear for small numbers. + let target_rows_per_range = + (estimated_rows_per_worker / TARGET_RANGES_PER_WORKER).max(min_rows_per_worker); tracing::debug!( estimated_row_count, workers, estimated_rows_per_worker, - target_rows_per_bucket, - "partitioning key space by prefix" + target_rows_per_range, + "partitioning key space" ); - target_rows_per_bucket + target_rows_per_range } -/// Splits the whole key space into ranges of at most roughly -/// `target_rows_per_bucket` estimated rows each, returned in key order. async fn split_into_ranges( db: &mut D, estimated_row_count: f64, @@ -129,8 +126,6 @@ async fn split_into_ranges( for range in ranges { if range.estimated_rows > target_rows_per_bucket { split_any = true; - // An empty child list means the range holds no rows (a - // phantom estimate) and is dropped. next.extend(split_range(db, &range, target_rows_per_bucket).await?); } else { next.push(range); @@ -167,12 +162,7 @@ fn assign_boundaries(ranges: &[Range], workers: usize) -> Vec { if rows_seen >= f64::cast_lossy(boundaries.len() + 1) * per_worker { // The final range's end is None (open), it can never be a boundary. if let Some(end) = &range.end { - // A server returning non-advancing prefixes can repeat an - // end. Skip it, a duplicate boundary would fail the strict - // monotonicity check downstream. - if boundaries.last() != Some(end) { - boundaries.push(end.clone()); - } + boundaries.push(end.clone()); } } } From 83cdc0bbd84737d44db3a1f9c8191e48c5c87910 Mon Sep 17 00:00:00 2001 From: Peter Larsen Date: Wed, 5 Aug 2026 13:13:47 -0400 Subject: [PATCH 05/10] mysql-util: name the partitioner probe trait PrimaryKeyProber The trait is the partitioner-owned seam over the concrete KeyProber, a distinct name keeps it from shadowing the prober it wraps. --- src/mysql-util/src/partition.rs | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/src/mysql-util/src/partition.rs b/src/mysql-util/src/partition.rs index 6e665e6661a15..57b3f4a00259a 100644 --- a/src/mysql-util/src/partition.rs +++ b/src/mysql-util/src/partition.rs @@ -68,7 +68,7 @@ struct Range { depth: usize, } -async fn partition( +async fn partition( db: &mut D, workers: usize, estimated_row_count: u64, @@ -109,7 +109,7 @@ fn get_target_max_rows_per_range( target_rows_per_range } -async fn split_into_ranges( +async fn split_into_ranges( db: &mut D, estimated_row_count: f64, target_rows_per_bucket: f64, @@ -172,7 +172,7 @@ fn assign_boundaries(ranges: &[Range], workers: usize) -> Vec { /// Splits `parent` at every distinct key prefix one character longer than the /// prefix `parent`, i.e. "a" depth: 1 for a table with keys "a", "aa", "aaa", "ab" would be split to /// "a" depth: 2 estimate 1, "aa" depth: 2 estimate 2, "ab" depth 2, estimate 1. -async fn split_range( +async fn split_range( db: &mut D, parent: &Range, target_rows: f64, @@ -247,7 +247,7 @@ async fn split_range( } /// Wrapper around KeyProber for testing purposes. -trait PartitionDb { +trait PrimaryKeyProber { async fn estimate_range_rows( &mut self, start: Option<&str>, @@ -272,7 +272,7 @@ trait PartitionDb { ) -> Result, MySqlError>; } -impl<'a> PartitionDb for KeyProber<'a> { +impl<'a> PrimaryKeyProber for KeyProber<'a> { async fn estimate_range_rows( &mut self, start: Option<&str>, From c827ecd0cb27036c643531af595fc3bd9dafee29 Mon Sep 17 00:00:00 2001 From: Peter Larsen Date: Thu, 6 Aug 2026 13:04:13 -0400 Subject: [PATCH 06/10] mysql-util: shorten partitioner naming partition_table_by_pk_prefix becomes partition_table, the trait methods take the probe names they delegate to, and Range names its exclusive lower bound prefix. --- src/mysql-util/src/lib.rs | 2 +- src/mysql-util/src/partition.rs | 117 ++++++++------------------------ src/mysql-util/src/probe.rs | 9 --- 3 files changed, 28 insertions(+), 100 deletions(-) diff --git a/src/mysql-util/src/lib.rs b/src/mysql-util/src/lib.rs index 61f4ecfeaa54c..37e1445dac2b9 100644 --- a/src/mysql-util/src/lib.rs +++ b/src/mysql-util/src/lib.rs @@ -46,7 +46,7 @@ pub mod probe; pub use probe::{KeyProber, MAX_KEY_LENGTH}; pub mod partition; -pub use partition::partition_table_by_pk_prefix; +pub use partition::partition_table; mod aws_rds; diff --git a/src/mysql-util/src/partition.rs b/src/mysql-util/src/partition.rs index 57b3f4a00259a..a87544866af58 100644 --- a/src/mysql-util/src/partition.rs +++ b/src/mysql-util/src/partition.rs @@ -7,12 +7,6 @@ // the Business Source License, use of this software will be governed // by the Apache License, Version 2.0. -//! Discovers boundaries that split a table's string primary key space into -//! ranges of roughly equal row counts, for parallel snapshot reads. -//! -//! All key matching and ordering happens server-side under the key column's -//! own collation. Rust never orders the returned prefixes. - use mz_ore::cast::CastLossy; use crate::{KeyProber, MySqlError, QualifiedTableRef}; @@ -23,14 +17,8 @@ use crate::{KeyProber, MySqlError, QualifiedTableRef}; const TARGET_RANGES_PER_WORKER: f64 = 8.0; /// Computes up to `num_workers - 1` partition boundaries that divide the primary key space -/// into `num_workers` roughly even partitions. The boundaries are ordered according to MySQL's -/// internal collation rules. This should only be used against a single CHAR(N) or VARCHAR -/// primary key column. -/// -/// NOTE: Run this inside a REPEATABLE READ transaction. The underlying probes -/// read separate snapshots otherwise, and concurrent inserts can make the -/// prefix walk see the same prefix repeatedly instead of advancing. -pub async fn partition_table_by_pk_prefix( +/// into `num_workers` roughly even partitions. This should be run in a transaction. +pub async fn partition_table( conn: &mut mysql_async::Conn, table: QualifiedTableRef<'_>, pk_col: &str, @@ -58,8 +46,8 @@ pub async fn partition_table_by_pk_prefix( #[derive(Debug)] struct Range { - /// Exclusive start, `None` at the beginning of the key space. - start: Option, + /// `None` for the beginning of the key space. + prefix: Option, /// Exclusive end, `None` for the final open range. end: Option, /// Row estimate for the range, at least 1. @@ -68,8 +56,8 @@ struct Range { depth: usize, } -async fn partition( - db: &mut D, +async fn partition( + db: &mut KeyProber<'_>, workers: usize, estimated_row_count: u64, min_rows_per_worker: u64, @@ -109,13 +97,13 @@ fn get_target_max_rows_per_range( target_rows_per_range } -async fn split_into_ranges( - db: &mut D, +async fn split_into_ranges( + db: &mut KeyProber<'_>, estimated_row_count: f64, target_rows_per_bucket: f64, ) -> Result, MySqlError> { let mut ranges = vec![Range { - start: None, + prefix: None, end: None, estimated_rows: estimated_row_count, depth: 0, @@ -172,8 +160,8 @@ fn assign_boundaries(ranges: &[Range], workers: usize) -> Vec { /// Splits `parent` at every distinct key prefix one character longer than the /// prefix `parent`, i.e. "a" depth: 1 for a table with keys "a", "aa", "aaa", "ab" would be split to /// "a" depth: 2 estimate 1, "aa" depth: 2 estimate 2, "ab" depth 2, estimate 1. -async fn split_range( - db: &mut D, +async fn split_range( + db: &mut KeyProber<'_>, parent: &Range, target_rows: f64, ) -> Result, MySqlError> { @@ -181,24 +169,27 @@ async fn split_range( let mut children = Vec::new(); let Some(mut cur) = db - .first_prefix(parent.start.as_deref(), parent.end.as_deref(), depth) + .prefix_of_first_key_in_range(parent.prefix.as_deref(), parent.end.as_deref(), depth) .await? else { return Ok(children); }; // The first child inherits the parent's start. - let mut start = parent.start.clone(); + let mut start = parent.prefix.clone(); loop { // Small optimization to return early if the remaining rows past the current start prefix // fits within the threshold we're looking for. + // A missing optimizer estimate reads as an empty range, which the + // splitting loop drops or leaves unsplit. let remaining = db .estimate_range_rows(start.as_deref(), parent.end.as_deref()) - .await?; + .await? + .unwrap_or(0); let remaining = f64::cast_lossy(remaining).max(1.0); if remaining <= target_rows { children.push(Range { - start, + prefix: start, end: parent.end.clone(), estimated_rows: remaining, depth, @@ -210,10 +201,13 @@ async fn split_range( // extensions still need visiting. The first key past `cur` exposes // them, so the walk can keep splitting instead of retrying the whole // range at greater depths forever. - let next = match db.next_prefix(&cur, parent.end.as_deref(), depth).await? { + let next = match db + .prefix_of_first_row_not_matching_prefix(&cur, parent.end.as_deref(), depth) + .await? + { Some(next) => Some(next), None => { - db.first_prefix(Some(&cur), parent.end.as_deref(), depth) + db.prefix_of_first_key_in_range(Some(&cur), parent.end.as_deref(), depth) .await? } }; @@ -223,9 +217,10 @@ async fn split_range( Some(next) => { let estimated_rows = db .estimate_range_rows(start.as_deref(), Some(&next)) - .await?; + .await? + .unwrap_or(0); children.push(Range { - start: start.clone(), + prefix: start.clone(), end: Some(next.clone()), estimated_rows: f64::cast_lossy(estimated_rows).max(1.0), depth, @@ -235,7 +230,7 @@ async fn split_range( } None => { children.push(Range { - start, + prefix: start, end: parent.end.clone(), estimated_rows: remaining, depth, @@ -245,61 +240,3 @@ async fn split_range( } } } - -/// Wrapper around KeyProber for testing purposes. -trait PrimaryKeyProber { - async fn estimate_range_rows( - &mut self, - start: Option<&str>, - end: Option<&str>, - ) -> Result; - - /// A `start` of `None` means the beginning of the key space. - async fn first_prefix( - &mut self, - start: Option<&str>, - end: Option<&str>, - len: usize, - ) -> Result, MySqlError>; - - /// Every key extending `cur` is covered by the match, including `cur` - /// itself as an exact key. - async fn next_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: Option<&str>, - end: Option<&str>, - ) -> Result { - // A missing optimizer estimate reads as an empty range, which the - // splitting loop drops or leaves unsplit. - Ok(KeyProber::estimate_range_rows(self, start, end) - .await? - .unwrap_or(0)) - } - - async fn first_prefix( - &mut self, - start: Option<&str>, - end: Option<&str>, - len: usize, - ) -> Result, MySqlError> { - KeyProber::prefix_of_first_key_in_range(self, start, end, len).await - } - - async fn next_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 - } -} diff --git a/src/mysql-util/src/probe.rs b/src/mysql-util/src/probe.rs index b957bf488e1bd..439c979c73711 100644 --- a/src/mysql-util/src/probe.rs +++ b/src/mysql-util/src/probe.rs @@ -231,13 +231,6 @@ fn like_prefix_pattern(prefix: &str) -> String { pattern } -/// Runs `EXPLAIN` on `select` and returns the optimizer's estimate of rows -/// examined, from the `rows` column of the plan. -/// -/// The caller must pass a single-table `SELECT`. Joins and subqueries produce -/// multiple plan rows whose estimates do not combine additively, only the -/// first row is read. Returns `None` when the optimizer reports no estimate. -/// The estimate can be arbitrarily stale, callers must tolerate inaccuracy. async fn explain_row_estimate

( conn: &mut mysql_async::Conn, select: &str, @@ -246,8 +239,6 @@ async fn explain_row_estimate

( where P: Into + Send, { - // NOTE: The format must be pinned because newer MySQL versions default - // `explain_format` to `TREE`, which has no `rows` column. let plan: Option = conn .exec_first(format!("EXPLAIN FORMAT=TRADITIONAL {select}"), params) .await?; From 374af215cb5ede7593b5950f3d37abaa9d564ca9 Mon Sep 17 00:00:00 2001 From: Peter Larsen Date: Thu, 6 Aug 2026 16:48:55 -0400 Subject: [PATCH 07/10] mysql-util: rework partitioner walk into a BFS with surrogate sort keys Simplify children_prefixes into a clean prefix walk that accepts skipping exact keys shorter than the probe depth. Split breadth first with a coarse target of 1/max(workers, 8) of the table, restore key order via per-parent ordinal sort keys instead of client-side key comparison, and fold boundary assignment into compute_boundaries. Estimates are u64 end to end and a missing optimizer estimate is now a named MissingRowEstimate error raised inside estimate_range_rows. --- src/mysql-util/src/lib.rs | 10 ++ src/mysql-util/src/partition.rs | 244 +++++++++++++------------------- src/mysql-util/src/probe.rs | 32 ++--- 3 files changed, 126 insertions(+), 160 deletions(-) diff --git a/src/mysql-util/src/lib.rs b/src/mysql-util/src/lib.rs index 37e1445dac2b9..ba1620d02e46f 100644 --- a/src/mysql-util/src/lib.rs +++ b/src/mysql-util/src/lib.rs @@ -111,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 index a87544866af58..f417551f979ae 100644 --- a/src/mysql-util/src/partition.rs +++ b/src/mysql-util/src/partition.rs @@ -8,16 +8,14 @@ // by the Apache License, Version 2.0. use mz_ore::cast::CastLossy; +use mz_ore::str::redact; use crate::{KeyProber, MySqlError, QualifiedTableRef}; -/// When partitioning the data, how many ranges per worker should we break the -/// keyspace into. This helps avoid underestimates for large row counts resulting -/// in severe skew. -const TARGET_RANGES_PER_WORKER: f64 = 8.0; - /// 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 transaction. +/// 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. pub async fn partition_table( conn: &mut mysql_async::Conn, table: QualifiedTableRef<'_>, @@ -38,22 +36,26 @@ pub async fn partition_table( tracing::trace!( schema = schema_name, table = table_name, - ?boundaries, + // The boundaries are user data, redacted outside of CI. + boundaries = ?redact(&boundaries), "partitioned table by pk prefix" ); Ok(boundaries) } #[derive(Debug)] -struct Range { - /// `None` for the beginning of the key space. - prefix: Option, - /// Exclusive end, `None` for the final open range. +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 range, at least 1. - estimated_rows: f64, - /// Prefix length this range was split at. + /// Row estimate for the prefix, at least 1. + estimated_rows: u64, + /// Length this prefix was split at. depth: usize, + /// Use the position within each parent as a surrogate sort key to maintain the sort ordering + /// specified by MySQL. + surrogate_sort_key: Vec, } async fn partition( @@ -65,178 +67,132 @@ async fn partition( if workers <= 1 { return Ok(Vec::new()); } - let estimated_row_count = f64::cast_lossy(estimated_row_count.max(1)); - let min_rows_per_worker = f64::cast_lossy(min_rows_per_worker.max(1)); - let target_max_rows_per_range = - get_target_max_rows_per_range(workers, estimated_row_count, min_rows_per_worker); - // Should be many more ranges than workers unless the overall row count of the table is quite small. - let ranges = split_into_ranges(db, estimated_row_count, target_max_rows_per_range).await?; - Ok(assign_boundaries(&ranges, workers)) -} + let estimated_row_count = estimated_row_count.max(1); -fn get_target_max_rows_per_range( - workers: usize, - estimated_row_count: f64, - min_rows_per_worker: f64, -) -> f64 { - // Break up the key space into smaller ranges to more accurately rebuild the per-worker ranges later with less skew. - let estimated_rows_per_worker = - (estimated_row_count / f64::cast_lossy(workers)).max(min_rows_per_worker); - // Respect min_rows_per_worker as a lower bound for the granularity with which we attempt to break up the table. - // No need to add this overhead for small tables. Estimate accuracy isn't super clear for small numbers. - let target_rows_per_range = - (estimated_rows_per_worker / TARGET_RANGES_PER_WORKER).max(min_rows_per_worker); + // 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 before selecting partitions. + // Prefixes estimate at least one row, so a target below one never + // converges. + let target_max_rows_per_prefix = (f64::cast_lossy(estimated_row_count) + / f64::cast_lossy(workers.max(8))) + .max(f64::cast_lossy(min_rows_per_worker)) + .max(1.0); - tracing::debug!( - estimated_row_count, - workers, - estimated_rows_per_worker, - target_rows_per_range, - "partitioning key space" - ); - target_rows_per_range + compute_boundaries(db, workers, estimated_row_count, target_max_rows_per_prefix).await } -async fn split_into_ranges( +async fn compute_boundaries( db: &mut KeyProber<'_>, - estimated_row_count: f64, - target_rows_per_bucket: f64, -) -> Result, MySqlError> { - let mut ranges = vec![Range { - prefix: None, + workers: usize, + estimated_row_count: u64, + target_rows_per_prefix: f64, +) -> Result, MySqlError> { + // BFS of prefixes, splitting until estimates fall under the target. + let mut final_prefixes: Vec = vec![]; + let mut pending_prefixes = vec![Prefix { + prefix: String::new(), end: None, estimated_rows: estimated_row_count, depth: 0, + surrogate_sort_key: Vec::new(), }]; - loop { - let mut split_any = false; - let mut next: Vec = Vec::with_capacity(ranges.len()); - for range in ranges { - if range.estimated_rows > target_rows_per_bucket { - split_any = true; - next.extend(split_range(db, &range, target_rows_per_bucket).await?); - } else { - next.push(range); + + while !pending_prefixes.is_empty() { + let mut next: Vec = Vec::with_capacity(pending_prefixes.len()); + for prefix in pending_prefixes { + let children = children_prefixes(db, &prefix).await?; + // No child is visible under this prefix, its keys are the bare + // prefix itself or sort below their own deeper prefixes. Keep the + // parent as a leaf so its key space stays accounted for. + if children.is_empty() { + final_prefixes.push(prefix); + continue; + } + for child in children { + if f64::cast_lossy(child.estimated_rows) > target_rows_per_prefix { + next.push(child); + } else { + final_prefixes.push(child); + } } } - ranges = next; - if !split_any { - break; - } + pending_prefixes = next; } - Ok(ranges) -} + final_prefixes.sort_unstable_by(|a, b| a.surrogate_sort_key.cmp(&b.surrogate_sort_key)); -/// Accumulates `ranges` (in key order) into `workers` buckets of roughly -/// equal estimated rows and returns the bucket edges as boundaries. -fn assign_boundaries(ranges: &[Range], workers: usize) -> Vec { - // Emit a boundary each time the cumulative estimated rows pass the next - // worker's share. - let total: f64 = ranges.iter().map(|r| r.estimated_rows).sum(); + // 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: f64 = final_prefixes + .iter() + .map(|r| f64::cast_lossy(r.estimated_rows)) + .sum(); let per_worker = total / f64::cast_lossy(workers); tracing::debug!( - ranges = ranges.len(), + prefixes = final_prefixes.len(), total_estimated_rows = total, per_worker, - "assigning prefix ranges to workers" + "assigning prefixes to workers" ); let mut boundaries: Vec = Vec::with_capacity(workers - 1); let mut rows_seen = 0.0; - for range in ranges { + for prefix in &final_prefixes { if boundaries.len() == workers - 1 { break; } - rows_seen += range.estimated_rows; + rows_seen += f64::cast_lossy(prefix.estimated_rows); if rows_seen >= f64::cast_lossy(boundaries.len() + 1) * per_worker { - // The final range's end is None (open), it can never be a boundary. - if let Some(end) = &range.end { + // The final prefix's end is None (open), it can never be a boundary. + if let Some(end) = &prefix.end { boundaries.push(end.clone()); } } } - boundaries + Ok(boundaries) } -/// Splits `parent` at every distinct key prefix one character longer than the -/// prefix `parent`, i.e. "a" depth: 1 for a table with keys "a", "aa", "aaa", "ab" would be split to -/// "a" depth: 2 estimate 1, "aa" depth: 2 estimate 2, "ab" depth 2, estimate 1. -async fn split_range( +/// 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). +/// They are only invisible to probing, the snapshot ranges built from the +/// boundaries still cover them. +async fn children_prefixes( db: &mut KeyProber<'_>, - parent: &Range, - target_rows: f64, -) -> Result, MySqlError> { + parent: &Prefix, +) -> Result, MySqlError> { let depth = parent.depth + 1; let mut children = Vec::new(); + // Guaranteed to return None or a key longer than the current prefix assuming the upper + // bound correctly caps keys to the current prefix and we're in a transaction where + // new keys with a shorter length can't be inserted. Note that this only holds for + // collations that sort character-by-character. let Some(mut cur) = db - .prefix_of_first_key_in_range(parent.prefix.as_deref(), parent.end.as_deref(), depth) + .prefix_of_first_key_in_range(&parent.prefix, parent.end.as_deref(), depth) .await? else { return Ok(children); }; - // The first child inherits the parent's start. - let mut start = parent.prefix.clone(); loop { - // Small optimization to return early if the remaining rows past the current start prefix - // fits within the threshold we're looking for. - // A missing optimizer estimate reads as an empty range, which the - // splitting loop drops or leaves unsplit. - let remaining = db - .estimate_range_rows(start.as_deref(), parent.end.as_deref()) - .await? - .unwrap_or(0); - let remaining = f64::cast_lossy(remaining).max(1.0); - if remaining <= target_rows { - children.push(Range { - prefix: start, - end: parent.end.clone(), - estimated_rows: remaining, - depth, - }); - return Ok(children); - } - // When `cur` is an exact key shorter than `depth`, every key extending - // it matches `cur`, so next_prefix exhausts even though those - // extensions still need visiting. The first key past `cur` exposes - // them, so the walk can keep splitting instead of retrying the whole - // range at greater depths forever. - let next = match db + let next = db .prefix_of_first_row_not_matching_prefix(&cur, parent.end.as_deref(), depth) - .await? - { - Some(next) => Some(next), - None => { - db.prefix_of_first_key_in_range(Some(&cur), parent.end.as_deref(), depth) - .await? - } - }; - // A prefix equal to `cur` cannot advance the walk (only a misbehaving - // server produces one), stop splitting here. - match next.filter(|next| next != &cur) { - Some(next) => { - let estimated_rows = db - .estimate_range_rows(start.as_deref(), Some(&next)) - .await? - .unwrap_or(0); - children.push(Range { - prefix: start.clone(), - end: Some(next.clone()), - estimated_rows: f64::cast_lossy(estimated_rows).max(1.0), - depth, - }); - start = Some(next.clone()); - cur = next; - } - None => { - children.push(Range { - prefix: start, - end: parent.end.clone(), - estimated_rows: remaining, - depth, - }); - return Ok(children); - } + .await?; + let end = next.clone().or_else(|| parent.end.clone()); + let estimated_rows = db.estimate_range_rows(&cur, end.as_deref()).await?; + let mut surrogate_sort_key = parent.surrogate_sort_key.clone(); + surrogate_sort_key.push(children.len()); + children.push(Prefix { + prefix: cur, + end: end.clone(), + estimated_rows: estimated_rows.max(1), + depth, + surrogate_sort_key, + }); + 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?; From ab64a4a78982c09311c27de29ea92f8bfbb23891 Mon Sep 17 00:00:00 2001 From: Peter Larsen Date: Tue, 11 Aug 2026 08:37:37 -0400 Subject: [PATCH 08/10] mysql-util: remove surrogate key, revive O(depth*prefixes) algorithm and smaller partitions for more accuracy --- src/mysql-util/src/partition.rs | 65 ++++++++++++--------------------- 1 file changed, 24 insertions(+), 41 deletions(-) diff --git a/src/mysql-util/src/partition.rs b/src/mysql-util/src/partition.rs index f417551f979ae..d6f04263a83f6 100644 --- a/src/mysql-util/src/partition.rs +++ b/src/mysql-util/src/partition.rs @@ -53,9 +53,6 @@ struct Prefix { estimated_rows: u64, /// Length this prefix was split at. depth: usize, - /// Use the position within each parent as a surrogate sort key to maintain the sort ordering - /// specified by MySQL. - surrogate_sort_key: Vec, } async fn partition( @@ -70,11 +67,11 @@ async fn partition( 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 before selecting partitions. - // Prefixes estimate at least one row, so a target below one never - // converges. + // 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. Breaking down to smaller partitions results in more accurate splits, so we keep the + // 4x multiple of the worker count for > 2 workers. let target_max_rows_per_prefix = (f64::cast_lossy(estimated_row_count) - / f64::cast_lossy(workers.max(8))) + / f64::cast_lossy(workers * 4)) .max(f64::cast_lossy(min_rows_per_worker)) .max(1.0); @@ -88,54 +85,49 @@ async fn compute_boundaries( target_rows_per_prefix: f64, ) -> Result, MySqlError> { // BFS of prefixes, splitting until estimates fall under the target. - let mut final_prefixes: Vec = vec![]; - let mut pending_prefixes = vec![Prefix { + let mut ordered_prefixes = vec![Prefix { prefix: String::new(), end: None, estimated_rows: estimated_row_count, depth: 0, - surrogate_sort_key: Vec::new(), }]; - while !pending_prefixes.is_empty() { - let mut next: Vec = Vec::with_capacity(pending_prefixes.len()); - for prefix in pending_prefixes { - let children = children_prefixes(db, &prefix).await?; - // No child is visible under this prefix, its keys are the bare - // prefix itself or sort below their own deeper prefixes. Keep the - // parent as a leaf so its key space stays accounted for. - if children.is_empty() { - final_prefixes.push(prefix); - continue; - } - for child in children { - if f64::cast_lossy(child.estimated_rows) > target_rows_per_prefix { - next.push(child); - } else { - final_prefixes.push(child); - } + loop { + let mut next_ordered_prefixes: Vec = vec![]; + let mut split_any = false; + for prefix in ordered_prefixes { + if f64::cast_lossy(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); } } - pending_prefixes = next; + ordered_prefixes = next_ordered_prefixes; + if !split_any { + break; + } } - final_prefixes.sort_unstable_by(|a, b| a.surrogate_sort_key.cmp(&b.surrogate_sort_key)); // 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: f64 = final_prefixes + let total: f64 = ordered_prefixes .iter() .map(|r| f64::cast_lossy(r.estimated_rows)) .sum(); let per_worker = total / f64::cast_lossy(workers); tracing::debug!( - prefixes = final_prefixes.len(), + 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.0; - for prefix in &final_prefixes { + for prefix in &ordered_prefixes { if boundaries.len() == workers - 1 { break; } @@ -155,8 +147,6 @@ 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). -/// They are only invisible to probing, the snapshot ranges built from the -/// boundaries still cover them. async fn children_prefixes( db: &mut KeyProber<'_>, parent: &Prefix, @@ -164,10 +154,6 @@ async fn children_prefixes( let depth = parent.depth + 1; let mut children = Vec::new(); - // Guaranteed to return None or a key longer than the current prefix assuming the upper - // bound correctly caps keys to the current prefix and we're in a transaction where - // new keys with a shorter length can't be inserted. Note that this only holds for - // collations that sort character-by-character. let Some(mut cur) = db .prefix_of_first_key_in_range(&parent.prefix, parent.end.as_deref(), depth) .await? @@ -181,14 +167,11 @@ async fn children_prefixes( .await?; let end = next.clone().or_else(|| parent.end.clone()); let estimated_rows = db.estimate_range_rows(&cur, end.as_deref()).await?; - let mut surrogate_sort_key = parent.surrogate_sort_key.clone(); - surrogate_sort_key.push(children.len()); children.push(Prefix { prefix: cur, end: end.clone(), estimated_rows: estimated_rows.max(1), depth, - surrogate_sort_key, }); match next { Some(next) => cur = next, From 1774d2d3b2ea7678f50a5f368f3b6a7bbd010cce Mon Sep 17 00:00:00 2001 From: Peter Larsen Date: Tue, 11 Aug 2026 18:51:25 -0400 Subject: [PATCH 09/10] mysql-util: use integer row counts and rename min_split_threshold --- src/mysql-util/src/partition.rs | 32 ++++++++++++++------------------ 1 file changed, 14 insertions(+), 18 deletions(-) diff --git a/src/mysql-util/src/partition.rs b/src/mysql-util/src/partition.rs index d6f04263a83f6..1f6f3f6429b96 100644 --- a/src/mysql-util/src/partition.rs +++ b/src/mysql-util/src/partition.rs @@ -7,7 +7,7 @@ // the Business Source License, use of this software will be governed // by the Apache License, Version 2.0. -use mz_ore::cast::CastLossy; +use mz_ore::cast::CastFrom; use mz_ore::str::redact; use crate::{KeyProber, MySqlError, QualifiedTableRef}; @@ -22,7 +22,7 @@ pub async fn partition_table( pk_col: &str, num_workers: usize, estimated_row_count: u64, - min_rows_per_worker: 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); @@ -30,7 +30,7 @@ pub async fn partition_table( &mut db, num_workers, estimated_row_count, - min_rows_per_worker, + min_split_threshold, ) .await?; tracing::trace!( @@ -59,7 +59,7 @@ async fn partition( db: &mut KeyProber<'_>, workers: usize, estimated_row_count: u64, - min_rows_per_worker: u64, + min_split_threshold: u64, ) -> Result, MySqlError> { if workers <= 1 { return Ok(Vec::new()); @@ -70,10 +70,9 @@ async fn partition( // 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. Breaking down to smaller partitions results in more accurate splits, so we keep the // 4x multiple of the worker count for > 2 workers. - let target_max_rows_per_prefix = (f64::cast_lossy(estimated_row_count) - / f64::cast_lossy(workers * 4)) - .max(f64::cast_lossy(min_rows_per_worker)) - .max(1.0); + 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 } @@ -82,7 +81,7 @@ async fn compute_boundaries( db: &mut KeyProber<'_>, workers: usize, estimated_row_count: u64, - target_rows_per_prefix: f64, + target_rows_per_prefix: u64, ) -> Result, MySqlError> { // BFS of prefixes, splitting until estimates fall under the target. let mut ordered_prefixes = vec![Prefix { @@ -96,7 +95,7 @@ async fn compute_boundaries( let mut next_ordered_prefixes: Vec = vec![]; let mut split_any = false; for prefix in ordered_prefixes { - if f64::cast_lossy(prefix.estimated_rows) > target_rows_per_prefix { + 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. @@ -114,11 +113,8 @@ async fn compute_boundaries( // 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: f64 = ordered_prefixes - .iter() - .map(|r| f64::cast_lossy(r.estimated_rows)) - .sum(); - let per_worker = total / f64::cast_lossy(workers); + 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, @@ -126,13 +122,13 @@ async fn compute_boundaries( "assigning prefixes to workers" ); let mut boundaries: Vec = Vec::with_capacity(workers - 1); - let mut rows_seen = 0.0; + let mut rows_seen = 0; for prefix in &ordered_prefixes { if boundaries.len() == workers - 1 { break; } - rows_seen += f64::cast_lossy(prefix.estimated_rows); - if rows_seen >= f64::cast_lossy(boundaries.len() + 1) * per_worker { + 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()); From a3c03ee2dc11f34431dbcfbf701f70a67d1721d5 Mon Sep 17 00:00:00 2001 From: Peter Larsen Date: Tue, 11 Aug 2026 19:34:06 -0400 Subject: [PATCH 10/10] mysql-util: explain the partition target selection and split threshold --- src/mysql-util/src/partition.rs | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/src/mysql-util/src/partition.rs b/src/mysql-util/src/partition.rs index 1f6f3f6429b96..48d9bf0a2ada8 100644 --- a/src/mysql-util/src/partition.rs +++ b/src/mysql-util/src/partition.rs @@ -16,6 +16,10 @@ use crate::{KeyProber, MySqlError, QualifiedTableRef}; /// 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<'_>, @@ -68,8 +72,17 @@ async fn partition( // 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. Breaking down to smaller partitions results in more accurate splits, so we keep the - // 4x multiple of the worker count for > 2 workers. + // 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); @@ -165,7 +178,7 @@ async fn children_prefixes( let estimated_rows = db.estimate_range_rows(&cur, end.as_deref()).await?; children.push(Prefix { prefix: cur, - end: end.clone(), + end, estimated_rows: estimated_rows.max(1), depth, });