diff --git a/misc/python/materialize/mzcompose/__init__.py b/misc/python/materialize/mzcompose/__init__.py index 9fb80b0a1b169..0fd4edf8b6951 100644 --- a/misc/python/materialize/mzcompose/__init__.py +++ b/misc/python/materialize/mzcompose/__init__.py @@ -298,6 +298,11 @@ def get_variable_system_parameters( VariableSystemParameter( "mysql_source_snapshot_parallelism", "true", ["true", "false"] ), + # Low default so the tiny tables in tests still exercise PK-prefix + # range splitting; the production default only splits large tables. + VariableSystemParameter( + "mysql_source_snapshot_partition_min_rows", "2", ["2", "50000"] + ), VariableSystemParameter( "persist_batch_columnar_format", "structured" if version > MzVersion.parse_mz("v0.135.0-dev") else "both_v2", diff --git a/misc/python/materialize/parallel_workload/action.py b/misc/python/materialize/parallel_workload/action.py index c585e9bccb09a..ba28be26fd0a8 100644 --- a/misc/python/materialize/parallel_workload/action.py +++ b/misc/python/materialize/parallel_workload/action.py @@ -3046,6 +3046,12 @@ def __init__( self.flags_with_values["mysql_source_snapshot_parallelism"] = ( BOOLEAN_FLAG_VALUES ) + # 2 exercises PK-prefix splitting on workload-sized tables, the + # default leaves them in a single bucket. + self.flags_with_values["mysql_source_snapshot_partition_min_rows"] = [ + "2", + "50000", + ] # If you are adding a new config flag in Materialize, consider using it # here instead of just marking it as uninteresting to silence the diff --git a/src/mysql-util/src/lib.rs b/src/mysql-util/src/lib.rs index 649f141fa085f..3d91b3e20234e 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; +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..e8c5fd45fffa3 --- /dev/null +++ b/src/mysql-util/src/partition.rs @@ -0,0 +1,511 @@ +// 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::CastLossy; +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. +pub async fn partition_table( + conn: &mut mysql_async::Conn, + table: QualifiedTableRef<'_>, + pk_col: &str, + num_workers: usize, + estimated_row_count: 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, + num_workers, + estimated_row_count, + min_rows_per_worker, + ) + .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 { + /// `None` for the beginning of the key space. + prefix: Option, + /// 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, + /// 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( + db: &mut D, + workers: usize, + estimated_row_count: u64, + min_rows_per_worker: 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 before selecting partitions. + 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)); + + compute_boundaries(db, workers, estimated_row_count, target_max_rows_per_prefix).await +} + +async fn compute_boundaries( + db: &mut D, + 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: None, + 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 { + for child in children_prefixes(db, &prefix).await? { + if f64::cast_lossy(child.estimated_rows) > target_rows_per_prefix { + next.push(child); + } else { + final_prefixes.push(child); + } + } + } + pending_prefixes = next; + } + 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 + .iter() + .map(|r| f64::cast_lossy(r.estimated_rows)) + .sum(); + let per_worker = total / f64::cast_lossy(workers); + tracing::debug!( + prefixes = final_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 { + 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 { + // 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. We accept this because we only lose +/// at max one row per prefix we step deeper into. +async fn children_prefixes( + db: &mut D, + 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. + let Some(mut cur) = db + .prefix_of_first_key_in_range(parent.prefix.as_deref(), 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(Some(&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: Some(cur), + end: end.clone(), + estimated_rows: estimated_rows.max(1), + depth, + surrogate_sort_key, + }); + match next { + Some(next) => cur = next, + None => return Ok(children), + } + } +} + +/// Wrapper around KeyProber for testing purposes. +trait PrimaryKeyProber { + async fn estimate_range_rows( + &mut self, + start: Option<&str>, + end: Option<&str>, + ) -> Result; + + async fn prefix_of_first_key_in_range( + &mut self, + start: Option<&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: Option<&str>, + end: Option<&str>, + ) -> Result { + KeyProber::estimate_range_rows(self, start, end).await + } + + async fn prefix_of_first_key_in_range( + &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 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 mysql_async::{Params, Value}; + use mz_ore::cast::CastFrom; + + use super::*; + + /// In-memory [`PrimaryKeyProber`] over a sorted key list with exact + /// "estimates". Byte order stands in for the collation. + struct MockDb { + keys: Vec, + } + + impl MockDb { + fn bounds(&self, start: Option<&str>, end: Option<&str>) -> (usize, usize) { + // The lower bound is exclusive, a key equal to `start` is skipped. + let lo = match start { + Some(start) => self.keys.partition_point(|k| k.as_str() <= start), + None => 0, + }; + 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: Option<&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: Option<&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(None, 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_bucket_rows_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(()) + } + + /// Exercises the partitioner against a live MySQL server, covering what + /// the mock cannot see: `EXPLAIN` estimates over prepared statements, + /// `LIKE` pattern semantics, and the nested next-prefix query. + /// + /// Skipped unless `MZ_TEST_MYSQL_URL` points at a server this test may + /// scribble on, e.g. `mysql://root:p%40ssw0rd@127.0.0.1:13306`. + #[mz_ore::test(tokio::test)] + #[cfg_attr(miri, ignore)] // needs a network connection + async fn test_live_mysql() -> 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!"); + } + tracing::info!("MZ_TEST_MYSQL_URL not set: skipping live MySQL test"); + return Ok(()); + }; + let mut conn = mysql_async::Conn::new(mysql_async::Opts::from_url(&url)?).await?; + + // Static DDL strings, nothing to parameterize. + #[allow(clippy::disallowed_methods)] + { + conn.query_drop("DROP DATABASE IF EXISTS mz_partition_test") + .await?; + conn.query_drop("CREATE DATABASE mz_partition_test").await?; + conn.query_drop( + "CREATE TABLE mz_partition_test.t (id VARCHAR(32) PRIMARY KEY NOT NULL)", + ) + .await?; + // 900 keys under 'a', 100 under 'b', plus LIKE metacharacter keys + // and a bare 'a' that every 'a...' key extends. + conn.query_drop( + "INSERT INTO mz_partition_test.t \ + WITH RECURSIVE n AS (SELECT 1 x UNION ALL SELECT x+1 FROM n WHERE x < 900) \ + SELECT CONCAT('a', LPAD(x, 5, '0')) FROM n", + ) + .await?; + conn.query_drop("INSERT INTO mz_partition_test.t VALUES ('a')") + .await?; + conn.query_drop( + "INSERT INTO mz_partition_test.t \ + WITH RECURSIVE n AS (SELECT 1 x UNION ALL SELECT x+1 FROM n WHERE x < 100) \ + SELECT CONCAT('b', LPAD(x, 5, '0')) FROM n", + ) + .await?; + conn.query_drop("INSERT INTO mz_partition_test.t VALUES ('c_1'), ('c%2'), ('c\\\\3')") + .await?; + // Refresh optimizer statistics so EXPLAIN estimates see the rows. + conn.query_drop("ANALYZE TABLE mz_partition_test.t").await?; + } + + let table = QualifiedTableRef { + schema_name: "mz_partition_test", + table_name: "t", + }; + + // Even with a minimum above the table size the root splits once, so + // coarse boundaries may exist but stay within the worker count. + let boundaries = partition_table(&mut conn, table.clone(), "id", 4, 1004, 50_000).await?; + assert!(boundaries.len() <= 3, "{boundaries:?}"); + + // A low minimum splits the table, and MySQL agrees the boundaries are + // strictly increasing under the column collation. + let boundaries = partition_table(&mut conn, table.clone(), "id", 4, 1004, 10).await?; + assert!( + !boundaries.is_empty() && boundaries.len() <= 3, + "{boundaries:?}" + ); + // Most of the rows extend the bare key 'a', so splitting must reach + // inside those extensions rather than stopping at the exact key. + assert!( + boundaries.iter().any(|b| b.starts_with('a') && b.len() > 1), + "{boundaries:?}" + ); + for pair in boundaries.windows(2) { + let increasing: Option = conn + .exec_first("SELECT ? < ?", (&pair[0], &pair[1])) + .await?; + assert_eq!(increasing, Some(1), "{boundaries:?}"); + } + // Ranges partition the table: per-range counts sum to the total. + let mut total = 0u64; + let mut lower: Option = None; + for upper in boundaries.iter().map(Some).chain([None]) { + let (clause, params) = match (&lower, upper) { + (None, Some(hi)) => ("id < ?".to_string(), vec![Value::from(hi)]), + (Some(lo), Some(hi)) => ( + "id >= ? AND id < ?".to_string(), + vec![Value::from(lo), Value::from(hi)], + ), + (Some(lo), None) => ("id >= ?".to_string(), vec![Value::from(lo)]), + (None, None) => unreachable!("at least one boundary exists"), + }; + let count: Option = conn + .exec_first( + format!("SELECT COUNT(*) FROM mz_partition_test.t WHERE {clause}"), + Params::Positional(params), + ) + .await?; + total += count.expect("count returns a row"); + lower = upper.cloned(); + } + assert_eq!(total, 1004); + + #[allow(clippy::disallowed_methods)] + conn.query_drop("DROP DATABASE mz_partition_test").await?; + conn.disconnect().await?; + Ok(()) + } +} diff --git a/src/mysql-util/src/probe.rs b/src/mysql-util/src/probe.rs index 28b8117043aa8..9412379c59eeb 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}; @@ -61,14 +62,22 @@ impl<'a> KeyProber<'a> { &mut self, lower_bound_exclusive: Option<&str>, upper_bound_exclusive: Option<&str>, - ) -> Result, MySqlError> { + ) -> Result { let (clause, params) = self.range_filter(lower_bound_exclusive, upper_bound_exclusive); let select = format!( "SELECT {col} FROM {table} WHERE {clause}", 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 @@ -231,6 +240,10 @@ where #[cfg(test)] mod tests { + use std::collections::BTreeSet; + + use mz_ore::cast::CastFrom; + use super::*; #[mz_ore::test] @@ -319,17 +332,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, None).await?.expect("estimate"); + let all = p.estimate_range_rows(None, None).await?; assert!((500..=2000).contains(&all), "all={all}"); - let half = p - .estimate_range_rows(Some("a00500"), None) - .await? - .expect("estimate"); + let half = p.estimate_range_rows(Some("a00500"), None).await?; assert!((250..=1000).contains(&half), "half={half}"); - let none = p - .estimate_range_rows(Some("zzz"), None) - .await? - .expect("estimate"); + let none = p.estimate_range_rows(Some("zzz"), None).await?; assert!(none <= 5, "none={none}"); drop_db(&mut conn, DB).await?; @@ -337,6 +344,645 @@ mod tests { Ok(()) } + #[mz_ore::test(tokio::test)] + #[cfg_attr(miri, ignore)] + async fn test_case_insensitive_prefix_traversal() -> Result<(), anyhow::Error> { + let Some(mut conn) = connect().await? else { + return Ok(()); + }; + const DB: &str = "mz_probe_case_insensitive"; + let keys = ["Aa", "ab", "b", "Bb", "bbb", "C"]; + let table = setup_table(&mut conn, DB, "utf8mb4_0900_ai_ci", &keys).await?; + + let p = &mut KeyProber::new(&mut conn, table, "id"); + // Although the sorting is case-insensitive the values returned by mysql are not normalized, so + // we need to use the correct character representation here to get the tests to pass. + assert_eq!( + prefix_of_first_key_in_range(p, None, None, 1).await, + some("A") + ); + assert_eq!( + prefix_of_first_row_not_matching_prefix(p, "A", None, 1).await, + some("b") + ); + assert_eq!( + prefix_of_first_row_not_matching_prefix(p, "b", None, 1).await, + some("C") + ); + assert_eq!( + prefix_of_first_row_not_matching_prefix(p, "C", None, 1).await, + None + ); + + assert_eq!( + prefix_of_first_key_in_range(p, Some("A"), Some("b"), 2).await, + some("Aa") + ); + assert_eq!( + prefix_of_first_row_not_matching_prefix(p, "Aa", Some("b"), 2).await, + some("ab") + ); + assert_eq!( + prefix_of_first_row_not_matching_prefix(p, "ab", Some("b"), 2).await, + None + ); + + // The exclusive bound skips the exact key "b", and its extensions + // surface as their own prefixes under this case-insensitive + // collation. + assert_eq!( + prefix_of_first_key_in_range(p, Some("b"), Some("C"), 2).await, + some("Bb") + ); + assert_eq!( + prefix_of_first_row_not_matching_prefix(p, "Bb", Some("C"), 2).await, + None + ); + // Every key matching 'b%' is covered by the prefix match. + assert_eq!( + prefix_of_first_row_not_matching_prefix(p, "b", Some("C"), 2).await, + None + ); + + assert_eq!( + prefix_of_first_key_in_range(p, Some("C"), None, 2).await, + None + ); + + drop_db(&mut conn, DB).await?; + conn.disconnect().await?; + Ok(()) + } + + #[mz_ore::test(tokio::test)] + #[cfg_attr(miri, ignore)] + async fn test_wild_card_char_in_data() -> Result<(), anyhow::Error> { + let Some(mut conn) = connect().await? else { + return Ok(()); + }; + const DB: &str = "mz_probe_wildcard_test"; + // Keys are a_1, a\2, a\3, a%4, a|5, covering the LIKE wildcards + // and the escape character itself. The utf8mb4_0900_ai_ci collation + // orders them a_1 < a\2 < a\3 < a%4 < a|5. + let keys = ["a_1", "a\\2", "a\\3", "a%4", "a|5"]; + let table = setup_table(&mut conn, DB, "utf8mb4_0900_ai_ci", &keys).await?; + + let p = &mut KeyProber::new(&mut conn, table, "id"); + assert_eq!( + prefix_of_first_key_in_range(p, None, None, 1).await, + some("a") + ); + assert_eq!( + prefix_of_first_row_not_matching_prefix(p, "a", None, 1).await, + None + ); + assert_eq!( + prefix_of_first_key_in_range(p, Some("a"), None, 2).await, + some("a_") + ); + assert_eq!( + prefix_of_first_row_not_matching_prefix(p, "a_", None, 2).await, + some("a\\") + ); + assert_eq!( + prefix_of_first_row_not_matching_prefix(p, "a\\", None, 2).await, + some("a%") + ); + assert_eq!( + prefix_of_first_row_not_matching_prefix(p, "a%", None, 2).await, + some("a|") + ); + assert_eq!( + prefix_of_first_row_not_matching_prefix(p, "a|", None, 2).await, + None + ); + + // Range bounds that are themselves wildcard characters. + assert_eq!( + prefix_of_first_key_in_range(p, Some("a_"), Some("a\\"), 3).await, + some("a_1") + ); + assert_eq!( + prefix_of_first_row_not_matching_prefix(p, "a_1", Some("a\\"), 3).await, + None + ); + assert_eq!( + prefix_of_first_key_in_range(p, Some("a\\"), Some("a%"), 3).await, + some("a\\2") + ); + assert_eq!( + prefix_of_first_row_not_matching_prefix(p, "a\\2", Some("a%"), 3).await, + some("a\\3") + ); + assert_eq!( + prefix_of_first_row_not_matching_prefix(p, "a\\3", Some("a%"), 3).await, + None + ); + assert_eq!( + prefix_of_first_key_in_range(p, Some("a%"), Some("a|"), 3).await, + some("a%4") + ); + assert_eq!( + prefix_of_first_row_not_matching_prefix(p, "a%4", None, 3).await, + some("a|5") + ); + assert_eq!( + prefix_of_first_key_in_range(p, Some("a|"), None, 3).await, + some("a|5") + ); + assert_eq!( + prefix_of_first_row_not_matching_prefix(p, "a|5", None, 3).await, + None + ); + + drop_db(&mut conn, DB).await?; + conn.disconnect().await?; + Ok(()) + } + + #[mz_ore::test(tokio::test)] + #[cfg_attr(miri, ignore)] + async fn test_multibyte_chars_in_data() -> Result<(), anyhow::Error> { + let Some(mut conn) = connect().await? else { + return Ok(()); + }; + const DB: &str = "mz_probe_multibyte_test"; + // utf8mb4_0900_ai_ci orders symbols before letters and Han after + // Latin: 😀 < 😀😀 < 😀a < a < a😀 < 日本 < 日本語. + let keys = ["a", "a😀", "😀", "😀a", "😀😀", "日本", "日本語"]; + let table = setup_table(&mut conn, DB, "utf8mb4_0900_ai_ci", &keys).await?; + + let p = &mut KeyProber::new(&mut conn, table, "id"); + // Prefix lengths count characters, not bytes: a one-char prefix of a + // four-byte emoji is the whole emoji, never a broken fragment. + assert_eq!( + prefix_of_first_key_in_range(p, None, None, 1).await, + some("😀") + ); + assert_eq!( + prefix_of_first_row_not_matching_prefix(p, "😀", None, 1).await, + some("a") + ); + assert_eq!( + prefix_of_first_row_not_matching_prefix(p, "a", None, 1).await, + some("日") + ); + assert_eq!( + prefix_of_first_row_not_matching_prefix(p, "日", None, 1).await, + None + ); + + // Prefix matching works mid-multibyte: every key sharing 😀 is + // covered (including its extensions), and 😀😀 advances to 😀a. + assert_eq!( + prefix_of_first_row_not_matching_prefix(p, "😀", None, 2).await, + some("a") + ); + assert_eq!( + prefix_of_first_row_not_matching_prefix(p, "😀😀", None, 2).await, + some("😀a") + ); + assert_eq!( + prefix_of_first_row_not_matching_prefix(p, "日本", None, 3).await, + None + ); + + drop_db(&mut conn, DB).await?; + conn.disconnect().await?; + Ok(()) + } + + #[mz_ore::test(tokio::test)] + #[cfg_attr(miri, ignore)] + async fn test_live_mysql_ulid_pk() -> Result<(), anyhow::Error> { + let Some(mut conn) = connect().await? else { + return Ok(()); + }; + const DB: &str = "mz_probe_ulid_test"; + // ULIDs minted around the same time share a long timestamp prefix, + // here chars 11 and 12 are the first that differ. + const CROCKFORD: &[u8] = b"0123456789ABCDEFGHJKMNPQRSTVWXYZ"; + let ids: Vec = (0..1000) + .map(|i| { + format!( + "01J8ZXABCD{}{}00000000000000", + char::from(CROCKFORD[i / 32]), + char::from(CROCKFORD[i % 32]), + ) + }) + .collect(); + let table = setup_table(&mut conn, DB, "utf8mb4_0900_ai_ci", &ids).await?; + let mut prober = KeyProber::new(&mut conn, table, "id"); + + // Every key shares the timestamp prefix, so short prefixes cannot + // split the key space at all. + assert_eq!( + prefix_of_first_key_in_range(&mut prober, None, None, 1).await, + some("0") + ); + assert_eq!( + prefix_of_first_row_not_matching_prefix(&mut prober, "0", None, 1).await, + None + ); + assert_eq!( + prefix_of_first_row_not_matching_prefix(&mut prober, "01J8ZXABCD", None, 10).await, + None + ); + + // One character past the shared prefix distinguishes the keys. + let walked = walk_prefixes(&mut prober, 11).await?; + let expected: Vec = (0..32) + .map(|i| format!("01J8ZXABCD{}", char::from(CROCKFORD[i]))) + .collect(); + assert_eq!(walked, expected); + + let all = prober.estimate_range_rows(None, None).await?; + assert!(all > 0, "all={all}"); + + drop_db(&mut conn, DB).await?; + conn.disconnect().await?; + Ok(()) + } + + #[mz_ore::test(tokio::test)] + #[cfg_attr(miri, ignore)] + async fn test_live_mysql_uuid_pk() -> Result<(), anyhow::Error> { + let Some(mut conn) = connect().await? else { + return Ok(()); + }; + const DB: &str = "mz_probe_uuid_test"; + // Hyphenated lowercase v4-shaped UUIDs, unique via the last group, + // with the leading group scattered like random UUIDs. + let ids: Vec = (0..1000u64) + .map(|i| { + let h = i.wrapping_mul(2654435761) % 0x1_0000_0000; + format!("{h:08x}-0000-4000-8000-{i:012x}") + }) + .collect(); + let table = setup_table(&mut conn, DB, "utf8mb4_0900_ai_ci", &ids).await?; + let mut p = KeyProber::new(&mut conn, table, "id"); + + // Lowercase hex order matches byte order under this collation. + assert_eq!( + prefix_of_first_key_in_range(&mut p, None, None, 36).await, + ids.iter().min().cloned() + ); + + let walked = walk_prefixes(&mut p, 1).await?; + let expected: BTreeSet = ids.iter().map(|id| id[..1].to_string()).collect(); + assert_eq!(walked.iter().cloned().collect::>(), expected); + + let all = p.estimate_range_rows(None, None).await?; + assert!(all > 0, "all={all}"); + + drop_db(&mut conn, DB).await?; + conn.disconnect().await?; + Ok(()) + } + + #[mz_ore::test(tokio::test)] + #[cfg_attr(miri, ignore)] + async fn test_live_mysql_like_metacharacters() -> Result<(), anyhow::Error> { + let Some(mut conn) = connect().await? else { + return Ok(()); + }; + const DB: &str = "mz_probe_like_test"; + // Every walk step matches on `LIKE '%'`, so keys whose + // prefixes are LIKE metacharacters exercise the escaping. + let ids = [ + "%", + "%%", + "_", + "__", + "\\", + "\\\\", + "a%", + "a%b", + "a_", + "a_b", + "a\\", + "a\\b", + "ab", + "a b", + "|", + "||", + "a|", + "a|b", + "100%", + "50%off", + "under_score", + "back\\slash", + ]; + let table = setup_table(&mut conn, DB, "utf8mb4_0900_ai_ci", &ids).await?; + + // The collation dictates both the visit order and how short keys + // interleave with their extensions, so assert the property that + // matters instead of exact prefixes: the walked prefixes are range + // boundaries that partition the table, every key falls in exactly + // one interval. The server does the interval counting, under the + // column's own collation. + for len in [1, 2] { + let walked = + walk_prefixes(&mut KeyProber::new(&mut conn, table.clone(), "id"), len).await?; + let mut total = 0; + for (i, lo) in walked.iter().enumerate() { + total += count_range(&mut conn, DB, lo, walked.get(i + 1)).await?; + } + assert_eq!( + total, + u64::cast_from(ids.len()), + "len={len} walked={walked:?}" + ); + } + + drop_db(&mut conn, DB).await?; + conn.disconnect().await?; + Ok(()) + } + + #[mz_ore::test(tokio::test)] + #[cfg_attr(miri, ignore)] + async fn test_live_mysql_collations() -> Result<(), anyhow::Error> { + let Some(mut conn) = connect().await? else { + return Ok(()); + }; + const CI_DB: &str = "mz_probe_collation_ci_test"; + const BIN_DB: &str = "mz_probe_collation_bin_test"; + + // Case-insensitive collation: case variants of one key collide, so + // keys differ by letter, in mixed case. + let ci_keys = ["Apple", "apricot", "banana", "Cherry"]; + let t_ci = setup_table(&mut conn, CI_DB, "utf8mb4_0900_ai_ci", &ci_keys).await?; + let mut prober = KeyProber::new(&mut conn, t_ci, "id"); + // 'A' covers 'apricot' too: LIKE is case-insensitive here, so a + // returned prefix covers every case variant of it. + assert_eq!(walk_prefixes(&mut prober, 1).await?, ["A", "b", "C"]); + assert_eq!(walk_prefixes(&mut prober, 32).await?, ci_keys); + + // Binary collation: case variants coexist and order by byte value. + let bin_keys = ["ABC", "ABD", "abc", "abd"]; + let t_bin = setup_table(&mut conn, BIN_DB, "utf8mb4_bin", &bin_keys).await?; + let mut prober = KeyProber::new(&mut conn, t_bin, "id"); + // Uppercase sorts before lowercase in byte order, and case variants + // are distinct prefixes. + assert_eq!(walk_prefixes(&mut prober, 1).await?, ["A", "a"]); + assert_eq!(walk_prefixes(&mut prober, 32).await?, bin_keys); + + drop_db(&mut conn, CI_DB).await?; + drop_db(&mut conn, BIN_DB).await?; + conn.disconnect().await?; + Ok(()) + } + + #[mz_ore::test(tokio::test)] + #[cfg_attr(miri, ignore)] + async fn test_live_mysql_latin1_charset() -> Result<(), anyhow::Error> { + let Some(mut conn) = connect().await? else { + return Ok(()); + }; + const DB: &str = "mz_probe_latin1_test"; + // The column charset differs from the utf8mb4 connection charset, so + // every value crosses a conversion. é and ñ are one character and one + // byte in latin1 but two bytes in the UTF-8 we receive, making + // character counts and byte counts diverge. Under latin1_swedish_ci + // é collates with e and ñ with n, so no key may be a case or accent + // variant of another. + let keys = ["a", "é", "éa", "éb", "ñ", "ña", "z"]; + let table = setup_table(&mut conn, DB, "latin1_swedish_ci", &keys).await?; + + let p = &mut KeyProber::new(&mut conn, table, "id"); + // Prefixes count characters in the column's charset and come back + // converted, so a one-character prefix of é is the whole é. + assert_eq!( + prefix_of_first_key_in_range(p, None, None, 1).await, + some("a") + ); + assert_eq!( + prefix_of_first_row_not_matching_prefix(p, "a", None, 1).await, + some("é") + ); + assert_eq!( + prefix_of_first_row_not_matching_prefix(p, "é", None, 1).await, + some("ñ") + ); + assert_eq!( + prefix_of_first_row_not_matching_prefix(p, "ñ", None, 1).await, + some("z") + ); + assert_eq!( + prefix_of_first_row_not_matching_prefix(p, "z", None, 1).await, + None + ); + + // The exclusive bound skips the exact key é (one character in + // latin1, two UTF-8 bytes here), and its extensions surface as + // their own prefixes. + assert_eq!( + prefix_of_first_key_in_range(p, Some("é"), Some("ñ"), 2).await, + some("éa") + ); + assert_eq!( + prefix_of_first_row_not_matching_prefix(p, "éa", Some("ñ"), 2).await, + some("éb") + ); + assert_eq!( + prefix_of_first_row_not_matching_prefix(p, "éb", Some("ñ"), 2).await, + None + ); + // Every key matching 'é%' is covered by the prefix match. + assert_eq!( + prefix_of_first_row_not_matching_prefix(p, "é", Some("ñ"), 2).await, + None + ); + + drop_db(&mut conn, DB).await?; + conn.disconnect().await?; + Ok(()) + } + + #[mz_ore::test(tokio::test)] + #[cfg_attr(miri, ignore)] + async fn test_live_mysql_invalid_utf8_keys() -> Result<(), anyhow::Error> { + let Some(mut conn) = connect().await? else { + return Ok(()); + }; + const DB: &str = "mz_probe_binary_test"; + recreate_db(&mut conn, DB).await?; + // A binary key column passes bytes through unconverted, so this is + // the one way invalid UTF-8 can reach the client. Production filters + // these columns out before probing, and falls back to a + // single-partition snapshot if one slips through. + #[allow(clippy::disallowed_methods)] + conn.query_drop(format!( + "CREATE TABLE {DB}.t (id VARBINARY(36) PRIMARY KEY NOT NULL)" + )) + .await?; + let keys: Vec> = vec![b"a1".to_vec(), b"a2".to_vec(), vec![0xff, 0xfe, 0x31]]; + conn.exec_batch( + format!("INSERT INTO {DB}.t VALUES (?)"), + keys.iter().map(|k| (Value::Bytes(k.clone()),)), + ) + .await?; + #[allow(clippy::disallowed_methods)] + conn.query_drop(format!("ANALYZE TABLE {DB}.t")).await?; + let table = QualifiedTableRef { + schema_name: DB, + table_name: "t", + }; + let mut p = KeyProber::new(&mut conn, table, "id"); + + // Estimates never decode key values, they keep working. + assert!(p.estimate_range_rows(None, None).await.is_ok()); + + // ASCII keys order before the 0xff key and decode fine. + assert_eq!( + prefix_of_first_key_in_range(&mut p, None, None, 2).await, + some("a1") + ); + assert_eq!( + prefix_of_first_row_not_matching_prefix(&mut p, "a1", None, 2).await, + some("a2") + ); + // The next key is invalid UTF-8. The probe reports it as a named + // error so callers can log it and fall back. + let err = p + .prefix_of_first_row_not_matching_prefix("a2", None, 2) + .await + .unwrap_err(); + assert!(matches!(err, MySqlError::NonUtf8KeyValue { .. }), "{err:?}"); + + drop_db(&mut conn, DB).await?; + conn.disconnect().await?; + Ok(()) + } + + #[mz_ore::test(tokio::test)] + #[cfg_attr(miri, ignore)] + async fn test_live_mysql_stale_statistics() -> Result<(), anyhow::Error> { + let Some(mut conn) = connect().await? else { + return Ok(()); + }; + const DB: &str = "mz_probe_stale_test"; + recreate_db(&mut conn, DB).await?; + // This setup stays bespoke: STATS_AUTO_RECALC=0 plus an ANALYZE while + // empty pins the persisted statistics at zero rows, no matter what is + // inserted afterwards. + #[allow(clippy::disallowed_methods)] + { + conn.query_drop(format!( + "CREATE TABLE {DB}.t (id VARCHAR(36) CHARACTER SET utf8mb4 \ + COLLATE utf8mb4_0900_ai_ci PRIMARY KEY NOT NULL) \ + STATS_AUTO_RECALC=0, STATS_PERSISTENT=1" + )) + .await?; + conn.query_drop(format!("ANALYZE TABLE {DB}.t")).await?; + } + let ids: Vec = (0..1000).map(|i| format!("a{i:05}")).collect(); + conn.exec_batch( + format!("INSERT INTO {DB}.t VALUES (?)"), + ids.iter().map(|id| (id.as_str(),)), + ) + .await?; + + // The staleness this test is about: table_rows reports 0. + let table_rows: Option = conn + .exec_first( + "SELECT table_rows FROM information_schema.tables \ + WHERE table_schema = ? AND table_name = 't'", + (DB,), + ) + .await?; + assert_eq!(table_rows, Some(0)); + + let table = QualifiedTableRef { + schema_name: DB, + table_name: "t", + }; + let mut prober = KeyProber::new(&mut conn, table, "id"); + + // 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, None).await?; + assert!((500..=2000).contains(&all), "all={all}"); + let range = prober + .estimate_range_rows(Some("a00100"), Some("a00200")) + .await?; + assert!((50..=200).contains(&range), "range={range}"); + + drop_db(&mut conn, DB).await?; + conn.disconnect().await?; + Ok(()) + } + + #[mz_ore::test(tokio::test)] + #[cfg_attr(miri, ignore)] + async fn test_probe_sargability() -> Result<(), anyhow::Error> { + let Some(mut conn) = connect().await? else { + return Ok(()); + }; + const DB: &str = "mz_probe_sargable_test"; + let ids: Vec = (0..1000).map(|i| format!("a{i:05}")).collect(); + let table = setup_table(&mut conn, DB, "utf8mb4_0900_ai_ci", &ids).await?; + + // Prove the methodology first: a deliberately non-sargable predicate + // reads every row, and the session handler counters see it. + let before = handler_reads(&mut conn).await?; + let _: Option = conn + .exec_first( + format!("SELECT COUNT(*) FROM {DB}.t WHERE LEFT(id, 2) = 'a0'"), + (), + ) + .await?; + let scan_reads = handler_reads(&mut conn).await? - before; + assert!(scan_reads >= 1000, "scan_reads={scan_reads}"); + + // Every probe must stay a handful of index operations. A regression + // to a scan costs >= 1000 reads, far past the generous bound. + let before = handler_reads(&mut conn).await?; + let got = prefix_of_first_key_in_range( + &mut KeyProber::new(&mut conn, table.clone(), "id"), + Some("a00500"), + None, + 6, + ) + .await; + let reads = handler_reads(&mut conn).await? - before; + // The exclusive bound skips the exact key a00500. + assert_eq!(got, some("a00501")); + assert!(reads < 50, "prefix_of_first_key_in_range reads={reads}"); + + let before = handler_reads(&mut conn).await?; + let got = prefix_of_first_row_not_matching_prefix( + &mut KeyProber::new(&mut conn, table.clone(), "id"), + "a00500", + None, + 6, + ) + .await; + let reads = handler_reads(&mut conn).await? - before; + assert_eq!(got, some("a00501")); + assert!(reads < 50, "max_key probe reads={reads}"); + + let before = handler_reads(&mut conn).await?; + let got = prefix_of_first_row_not_matching_prefix( + &mut KeyProber::new(&mut conn, table.clone(), "id"), + "a0", + None, + 6, + ) + .await; + let reads = handler_reads(&mut conn).await? - before; + // Every key matches 'a0%', so the prefix match covers the whole table and + // there is no next prefix, at the cost of two dives rather than a + // scan. + assert_eq!(got, None); + assert!(reads < 50, "whole-table match reads={reads}"); + + drop_db(&mut conn, DB).await?; + conn.disconnect().await?; + Ok(()) + } + // Test helpers. /// Connects to the server named by `MZ_TEST_MYSQL_URL`, or `None` to skip @@ -368,7 +1014,7 @@ mod tests { } /// Recreates scratch database `db` holding one table `t` whose string - /// primary key `id` is pinned to the given utf8mb4 `collation`, containing + /// 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>( conn: &mut mysql_async::Conn, @@ -377,9 +1023,12 @@ mod tests { keys: &[impl AsRef + Sync], ) -> Result, anyhow::Error> { recreate_db(conn, db).await?; + // MySQL collation names start with their character set's name, so + // the charset is pinned explicitly without a second parameter. + let charset = collation.split('_').next().expect("nonempty collation"); #[allow(clippy::disallowed_methods)] conn.query_drop(format!( - "CREATE TABLE {db}.t (id VARCHAR(36) CHARACTER SET utf8mb4 \ + "CREATE TABLE {db}.t (id VARCHAR(36) CHARACTER SET {charset} \ COLLATE {collation} PRIMARY KEY NOT NULL)" )) .await?; @@ -403,6 +1052,38 @@ mod tests { Ok(()) } + /// Number of keys in `[lo, hi)` of `db`'s table, counted by the server so + /// the comparison happens under the column's collation. + async fn count_range( + conn: &mut mysql_async::Conn, + db: &str, + lo: &str, + hi: Option<&String>, + ) -> Result { + let mut clause = "id >= ?".to_string(); + let mut params: Vec = vec![lo.into()]; + if let Some(hi) = hi { + clause.push_str(" AND id < ?"); + params.push(hi.as_str().into()); + } + let count: Option = conn + .exec_first( + format!("SELECT COUNT(*) FROM {db}.t WHERE {clause}"), + Params::Positional(params), + ) + .await?; + Ok(count.expect("COUNT returns a row")) + } + + /// Sum of this session's `Handler_read_*` counters: how many index or row + /// read operations the connection has performed so far. + async fn handler_reads(conn: &mut mysql_async::Conn) -> Result { + let rows: Vec<(String, String)> = conn + .exec("SHOW SESSION STATUS LIKE 'Handler_read%'", ()) + .await?; + Ok(rows.into_iter().map(|(_, v)| v.parse().unwrap_or(0)).sum()) + } + // Wrapped to limit boilerplate async fn prefix_of_first_key_in_range( prober: &mut KeyProber<'_>, @@ -443,4 +1124,31 @@ mod tests { fn some(s: &str) -> Option { Some(s.into()) } + + /// Walks the whole key space at prefix length `len`, panicking if the + /// walk revisits a prefix. + async fn walk_prefixes( + prober: &mut KeyProber<'_>, + len: usize, + ) -> Result, anyhow::Error> { + let mut walked = Vec::new(); + let Some(mut cur) = prober.prefix_of_first_key_in_range(None, None, len).await? else { + return Ok(walked); + }; + loop { + assert!( + !walked.contains(&cur), + "prefix repeated: {cur:?} (walked: {walked:?})" + ); + walked.push(cur.clone()); + match prober + .prefix_of_first_row_not_matching_prefix(&cur, None, len) + .await? + { + Some(next) => cur = next, + None => break, + } + } + Ok(walked) + } } diff --git a/src/storage-types/src/dyncfgs.rs b/src/storage-types/src/dyncfgs.rs index a80ba8ac330fb..6c07e155d19aa 100644 --- a/src/storage-types/src/dyncfgs.rs +++ b/src/storage-types/src/dyncfgs.rs @@ -215,6 +215,16 @@ pub static MYSQL_SOURCE_SNAPSHOT_PARALLELISM: Config = Config::new( "Whether to split MySQL snapshot reads across workers by primary-key ranges.", ); +/// The smallest estimated row count for which the MySQL snapshot prefix +/// partitioner keeps splitting a string primary key range. Tables estimated +/// below this stay in a single per-table bucket, i.e. are read by one worker. +pub static MYSQL_SOURCE_SNAPSHOT_PARTITION_MIN_ROWS: Config = Config::new( + "mysql_source_snapshot_partition_min_rows", + 50_000, + "Minimum estimated rows per range before MySQL snapshot PK-prefix partitioning \ + stops splitting; also the smallest table considered worth splitting.", +); + /// If the optimizer estimates the table has fewer rows than this, compute the exact row count /// with `COUNT(*)`. Otherwise, report the `information_schema` estimate directly. pub static MYSQL_SOURCE_SNAPSHOT_EXACT_COUNT_MAX_ROWS: Config = Config::new( @@ -438,6 +448,7 @@ pub fn all_dyncfgs(configs: ConfigSet) -> ConfigSet { .add(&MYSQL_REPLICATION_HEARTBEAT_INTERVAL) .add(&MYSQL_SOURCE_SNAPSHOT_EXACT_COUNT_MAX_ROWS) .add(&MYSQL_SOURCE_SNAPSHOT_PARALLELISM) + .add(&MYSQL_SOURCE_SNAPSHOT_PARTITION_MIN_ROWS) .add(&ORE_OVERFLOWING_BEHAVIOR) .add(&PG_FETCH_SLOT_RESUME_LSN_INTERVAL) .add(&PG_SCHEMA_VALIDATION_INTERVAL) diff --git a/src/storage/src/source/mysql/snapshot.rs b/src/storage/src/source/mysql/snapshot.rs index 982efd6b15dbe..a040c161bc187 100644 --- a/src/storage/src/source/mysql/snapshot.rs +++ b/src/storage/src/source/mysql/snapshot.rs @@ -62,8 +62,10 @@ //! //! ## Parallel PK-range snapshots //! -//! For tables with a suitable primary key, the leader computes `worker_count - 1` boundary keys -//! that split the key domain into disjoint half-open ranges, and broadcasts them. Each worker +//! For tables with a suitable single-column string primary key, the leader computes up to +//! `worker_count - 1` boundary keys that split the key domain into disjoint half-open ranges, +//! and broadcasts them. Boundaries are discovered by splitting the key space on character +//! prefixes using optimizer row estimates (see [`mz_mysql_util::partition`]). Each worker //! reads only its assigned range. Ranges are assigned round-robin starting from each table's //! legacy single-worker owner, so the open-ended ranges (which absorb any rows written past the //! last sampled boundary) land on a different worker per table rather than always the last worker. @@ -116,7 +118,9 @@ use futures::{StreamExt as _, TryStreamExt}; use itertools::Itertools; use mysql_async::prelude::Queryable; use mysql_async::{IsolationLevel, Row as MySqlRow, TxOpts}; -use mz_mysql_util::{MySqlConn, MySqlError, pack_mysql_row, query_sys_var, quote_identifier}; +use mz_mysql_util::{ + MySqlConn, MySqlError, QualifiedTableRef, pack_mysql_row, query_sys_var, quote_identifier, +}; use mz_ore::cast::CastFrom; use mz_ore::future::InTask; use mz_ore::iter::IteratorExt; @@ -220,80 +224,65 @@ fn worker_pk_range( }) } -/// Walks the primary key index in steps of about `row_count / worker_count`, taking the key -/// at each step's `OFFSET`. The per-step OFFSET scans sum to a full index pass, so this -/// function has a time complexity of O(row_count). Worker count is small, so the OFFSET -/// scans dominate the runtime. `row_count` can be an optimizer estimate for large tables, -/// so the partitions are approximate. An overestimate walks off the end of the index and stops -/// with fewer boundaries, resulting in some workers receiving less or no work. An underestimate -/// leaves a larger final partition for the last worker, however both still correctly partition -/// the table. Returns None if the primary key column type is not supported or the table is too -/// small to split. -async fn compute_sampled_splits( - conn: &mut Q, +/// Computes PK-range split boundaries for `table` by partitioning the key +/// space by character prefix (see [`mz_mysql_util::partition`]) and rendering +/// the resulting boundaries as SQL string literals via the server's `QUOTE()`, +/// matching the literal interpolation the range predicates use. Only string +/// key columns can be split this way, prefixes of other types do not order +/// consistently with their values. Returns None if the primary key column type +/// is not supported or the table is not worth splitting. +async fn compute_pk_splits( + conn: &mut mysql_async::Conn, table: &MySqlTableName, - pk_col: &(String, SqlScalarType), + raw_col: &str, + scalar_type: &SqlScalarType, worker_count: usize, - total: u64, -) -> Result, TransientError> -where - Q: Queryable, -{ - let (col, scalar_type) = pk_col; - // Render the PK column as text that sorts and compares the same way the range - // predicates do: `QUOTE()` under the column's collation for character types, - // `CAST(.. AS CHAR)` for integers. Any other type can't be split safely. - let (col_literal, integer_path) = match scalar_type { - SqlScalarType::Int16 - | SqlScalarType::Int32 - | SqlScalarType::Int64 - | SqlScalarType::UInt16 - | SqlScalarType::UInt32 - | SqlScalarType::UInt64 => (format!("CAST({col} AS CHAR)"), true), - SqlScalarType::Char { .. } | SqlScalarType::VarChar { .. } | SqlScalarType::String => { - (format!("QUOTE({col})"), false) - } + row_count: u64, + partition_min_rows: u64, +) -> Result, TransientError> { + match scalar_type { + SqlScalarType::Char { .. } | SqlScalarType::VarChar { .. } | SqlScalarType::String => {} _ => return Ok(None), + } + let table_ref = QualifiedTableRef { + schema_name: &table.0, + table_name: &table.1, }; - - let partitions = std::cmp::min(u64::cast_from(worker_count), total); - if partitions < 2 { + let prefixes = match mz_mysql_util::partition_table( + conn, + table_ref, + raw_col, + worker_count, + row_count, + partition_min_rows, + ) + .await + { + Ok(prefixes) => prefixes, + // Correctness never depends on splitting, so unsupported key data or + // an optimizer that reports no row estimate falls back to the + // single-worker whole-table read instead of failing the snapshot. + Err(err @ (MySqlError::NonUtf8KeyValue { .. } | MySqlError::MissingRowEstimate { .. })) => { + tracing::warn!(%err, "PK splitting fell back to a single partition"); + return Ok(None); + } + Err(err) => return Err(err.into()), + }; + if prefixes.is_empty() { return Ok(None); } - let chunk = total / partitions; - - let mut boundaries: Vec = Vec::with_capacity(usize::cast_from(partitions) - 1); - for _ in 1..partitions { - let (predicate, offset) = match boundaries.last() { - Some(prev) => (format!(" WHERE {col} > {prev}"), chunk - 1), - None => (String::new(), chunk), + let mut boundaries = Vec::with_capacity(prefixes.len()); + for prefix in prefixes { + let literal: Option = conn.exec_first("SELECT QUOTE(?)", (prefix,)).await?; + // QUOTE of a non-NULL parameter always returns a row, but fall back + // rather than panic if the protocol surprises us. + let Some(literal) = literal else { + return Ok(None); }; - // The identifier is quoted via `quote_identifier`, the previous boundary is - // itself a value MySQL rendered as a literal, `table` via Display, and the - // offset is an integer, so this interpolation is safe; not parameterizable. - #[allow(clippy::disallowed_methods)] - let row: Option = conn - .query_first(format!( - "SELECT {col_literal} FROM {table}{predicate} \ - ORDER BY {col} LIMIT 1 OFFSET {offset}" - )) - .await?; - // Defensive: if a concurrent write shrank the range out from under us, stop and - // use the boundaries found so far. Fewer partitions is still correct. - let Some(mut row) = row else { break }; - // The column is CAST/QUOTE-ed to text, so it decodes as a String that is - // already a valid SQL literal. A decode failure (e.g. a non-UTF-8 - // collation) means we can't safely partition: fall back. - match row.take_opt::(0) { - Some(Ok(lit)) if !integer_path || is_decimal_literal(&lit) => boundaries.push(lit), - _ => return Ok(None), - } - } - if boundaries.is_empty() { - return Ok(None); + boundaries.push(literal); } Ok(Some(PkBoundaries { - pk_col: col.clone(), + pk_col: quote_identifier(raw_col), boundaries, })) } @@ -301,11 +290,10 @@ where /// For every table, read the row count (exact only for small tables) and, for a /// supported single-column primary key, compute the PK-range split boundaries, /// concurrently over at most `worker_count` connections. `None` bounds means -/// single-worker fallback for that table. The counts are reused for both the sampling -/// stride and the snapshot size gauge. Snapshot size gauge is a metric for the snapshot -/// size used to report how many rows we need to process. "Sampling stride" refers to -/// the number of rows we use to page through the table to find roughly evenly spaced -/// primary keys to use as partition boundaries. +/// single-worker fallback for that table. The counts are reused for both boundary +/// discovery and the snapshot size gauge. The snapshot size gauge is a metric +/// reporting how many rows the snapshot needs to process. Boundary discovery uses +/// the count to size the partitioner's target buckets. async fn sample_pk_bounds( config: &RawSourceCreationConfig, connection_config: &mz_mysql_util::Config, @@ -335,6 +323,10 @@ async fn sample_pk_bounds( mz_storage_types::dyncfgs::MYSQL_SOURCE_SNAPSHOT_EXACT_COUNT_MAX_ROWS .get(config.config.config_set()), ); + let partition_min_rows = u64::cast_from( + mz_storage_types::dyncfgs::MYSQL_SOURCE_SNAPSHOT_PARTITION_MIN_ROWS + .get(config.config.config_set()), + ); let pooled_conns: Rc>> = Rc::new(RefCell::new(Vec::new())); // Counting and boundary-sampling each walk a table's index (O(rows)), so run tables @@ -366,11 +358,11 @@ async fn sample_pk_bounds( conn } }; - // Row count, reused for the sampling stride and the size gauge. When it + // Row count, reused for boundary discovery and the size gauge. When it // is counted exactly it runs on the same `READ ONLY` transaction as the - // boundary walk in `compute_sampled_splits`, so both see one consistent + // boundary probes in `compute_pk_splits`, so both see one consistent // snapshot. For large tables it is an optimizer estimate instead, which - // `compute_sampled_splits` tolerates. + // `compute_pk_splits` tolerates. let stats = collect_table_statistics(&mut *conn, table, exact_count_max_rows).await?; metrics.record_table_count_latency( @@ -385,9 +377,16 @@ async fn sample_pk_bounds( .flatten() { Some((raw_col, scalar_type)) => { - let pk_col = (quote_identifier(&raw_col), scalar_type); - compute_sampled_splits(&mut *conn, table, &pk_col, worker_count, count) - .await? + compute_pk_splits( + &mut *conn, + table, + &raw_col, + &scalar_type, + worker_count, + count, + partition_min_rows, + ) + .await? } None => None, }; @@ -619,11 +618,6 @@ fn is_plain_ident(s: &str) -> bool { !s.is_empty() && s.chars().all(|c| c.is_ascii_alphanumeric() || c == '_') } -fn is_decimal_literal(s: &str) -> bool { - let digits = s.strip_prefix('-').unwrap_or(s); - !digits.is_empty() && digits.bytes().all(|b| b.is_ascii_digit()) -} - /// Returns the set of full tables/sections of tables to read. fn plan_worker_reads( config: &RawSourceCreationConfig, diff --git a/test/mysql-cdc/mysql-cdc.td b/test/mysql-cdc/mysql-cdc.td index b242bf6c43afe..4f748df4122aa 100644 --- a/test/mysql-cdc/mysql-cdc.td +++ b/test/mysql-cdc/mysql-cdc.td @@ -769,16 +769,12 @@ $ mysql-execute name=mysql DROP TABLE pk_uint_test; # -# BIT primary key. A BIT column maps to a Materialize uint64, but carries the -# MySqlColumnMeta::Bit marker. The PK-range sampler renders a boundary from the -# Materialize scalar type alone (uint64 takes the integer path, CAST(id AS CHAR)), -# discarding that marker. For a BIT column CAST(.. AS CHAR) returns the value's -# raw bytes, not a decimal literal, and that text is spliced verbatim into the -# worker range predicates. The BIT(64) value X'30204F5220313D31' is the eight -# ASCII bytes "0 OR 1=1", so the sampled boundary turns the two worker reads into -# WHERE id < 0 OR 1=1 -# WHERE id >= 0 OR 1=1 -# both of which match every row. Every key must still be emitted exactly once. +# BIT primary key. A BIT column maps to a Materialize uint64, so it is not a +# string key and must not be split into PK ranges. The BIT(64) value +# X'30204F5220313D31' is the eight ASCII bytes "0 OR 1=1": a splitter that +# rendered the raw key bytes into a range predicate would turn a worker read +# into WHERE id >= 0 OR 1=1, matching every row on every worker. Every key +# must still be emitted exactly once. # $ mysql-execute name=mysql @@ -793,8 +789,9 @@ INSERT INTO pk_bit_test VALUES (1, 100), (0x30204F5220313D31, 200); FROM MYSQL CONNECTION mysql_conn; > CREATE TABLE pk_bit_table FROM SOURCE pk_bit_source (REFERENCE public.pk_bit_test); -# Exact count and no duplicate keys. Under the bug both workers read the whole -# table, so COUNT(*) is 4 while COUNT(DISTINCT id) stays 2. +# Exact count and no duplicate keys. If the table were wrongly split and both +# workers read the whole table, COUNT(*) would be 4 while COUNT(DISTINCT id) +# stays 2. > SELECT COUNT(*), COUNT(DISTINCT id) FROM pk_bit_table; 2 2 @@ -812,9 +809,9 @@ DROP TABLE pk_bit_test; # # PK-range splitting disabled via the mysql_source_snapshot_parallelism dyncfg. -# The same shape as the pk_range_test above, but every table must fall back to a -# single-worker whole-table read and still produce a complete, duplicate-free -# snapshot. +# The same shape as the pk_char_test above, a string PK that would otherwise be +# split, but every table must fall back to a single-worker whole-table read and +# still produce a complete, duplicate-free snapshot. # $ postgres-execute connection=mz_system @@ -822,9 +819,9 @@ ALTER SYSTEM SET mysql_source_snapshot_parallelism = false $ mysql-execute name=mysql DROP TABLE IF EXISTS pk_serial_test; -CREATE TABLE pk_serial_test (id BIGINT PRIMARY KEY, val BIGINT); +CREATE TABLE pk_serial_test (id CHAR(26) PRIMARY KEY, val BIGINT); SET @i := 0; -INSERT INTO pk_serial_test SELECT @i := @i + 1, @i * 7 FROM mysql.time_zone t1, mysql.time_zone t2 LIMIT 1000; +INSERT INTO pk_serial_test SELECT LPAD(CONV(@i := @i + 1, 10, 36), 26, '0'), @i * 7 FROM mysql.time_zone t1, mysql.time_zone t2 LIMIT 1000; > CREATE CLUSTER pk_serial_cluster SIZE 'scale=1,workers=4' @@ -837,11 +834,7 @@ INSERT INTO pk_serial_test SELECT @i := @i + 1, @i * 7 FROM mysql.time_zone t1, > SELECT COUNT(*), COUNT(DISTINCT id) FROM pk_serial_table; 1000 1000 -# Full key range present (ids 1..1000) -> SELECT MIN(id), MAX(id) FROM pk_serial_table; -1 1000 - -# Checksum: val = id * 7, so SUM(val) = 7 * SUM(1..1000) = 7 * 500500 +# Checksum: val = row number * 7, so SUM(val) = 7 * SUM(1..1000) = 7 * 500500 > SELECT SUM(val) FROM pk_serial_table; 3503500