storage: Fast approximate snapshot partitioning

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",

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

src/storage-types/src/dyncfgs.rs

@@ -215,6 +215,16 @@ pub static MYSQL_SOURCE_SNAPSHOT_PARALLELISM: Config<bool> = 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<usize> = 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<usize> = 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)

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,92 +224,76 @@ 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<Q>(
-    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<Option<PkBoundaries>, 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<Option<PkBoundaries>, 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<String> = 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<String> = 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<MySqlRow> = 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::<String, usize>(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,
     }))
 }
 
 /// 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<RefCell<Vec<MySqlConn>>> = 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,

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,19 +809,19 @@ 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
 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