Mysql string primary key partitioner by peterdukelarsen · Pull Request #38093 · MaterializeInc/materialize

peterdukelarsen · 2026-08-06T19:46:13Z

Motivation

Part of SS-97.

Working on a faster way to identify partition boundaries for string primary keys.

Description

Adds code for partitioning the primary keys of a table with a single-column string primary key.

The core idea for partitioning without needing to understand MySQL's string sort order rules is to walk through string prefixes, in order, estimating the sizes and refining as you go based on table size estimates.

We're doing a BFS here so in the future if we want to exit early (which is planned for a subsequent PR) we have a better shot of useful partitioning.

linear-code · 2026-08-06T19:46:17Z

SS-97

def-

Small extra tests:

diff --git a/src/mysql-util/src/partition.rs b/src/mysql-util/src/partition.rs
index 19fc9e1be2..12e87f759e 100644
--- a/src/mysql-util/src/partition.rs
+++ b/src/mysql-util/src/partition.rs
@@ -187,3 +187,81 @@ async fn children_prefixes(
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use mysql_async::prelude::Queryable;
+    use mz_ore::cast::CastFrom;
+
+    use super::*;
+    use crate::probe::tests::{connect, drop_db, setup_table};
+
+    /// Keys starting below space collate below `''` under a PAD SPACE collation,
+    /// and `utf8mb4_bin` is one that still compares character by character as
+    /// [`partition_table`] requires. Treating `''` as the start of the key space
+    /// hides that region from the walk, so it yields no boundaries and is never
+    /// counted towards sizing the ones it does yield. The snapshot range below
+    /// the first boundary is open-started, so it absorbs the whole region.
+    #[mz_ore::test(tokio::test)]
+    #[cfg_attr(miri, ignore)]
+    async fn test_live_mysql_partition_keys_below_empty_string() -> Result<(), anyhow::Error> {
+        let Some(mut conn) = connect().await? else {
+            return Ok(());
+        };
+        // One table entirely below `''`, one four fifths below it.
+        for (db, below) in [
+            ("mz_partition_all_below_test", 200),
+            ("mz_partition_part_below_test", 160),
+        ] {
+            let mut keys: Vec<String> = (0..below).map(|i| format!("\t{i:03}")).collect();
+            keys.extend((below..200).map(|i| format!("a{i:03}")));
+            let table = setup_table(&mut conn, db, "utf8mb4_bin", &keys).await?;
+            // Table statistics count every row, including those below `''`, so
+            // that is the estimate a caller has to work with.
+            let total = u64::cast_from(keys.len());
+
+            for workers in [2, 4, 8] {
+                let bounds =
+                    partition_table(&mut conn, table.clone(), "id", workers, total, 1).await?;
+                let counts = partition_counts(&mut conn, db, &bounds, total).await?;
+                // An even split gives each worker `1 / workers` of the table, so
+                // capping at three fifths is loose. It fails only once the
+                // partitioning has effectively collapsed onto one worker.
+                let largest = counts.iter().copied().max().expect("nonempty");
+                assert!(
+                    largest * 5 <= total * 3,
+                    "{db} workers={workers} largest={largest} of {total}, counts={counts:?}"
+                );
+            }
+            drop_db(&mut conn, db).await?;
+        }
+        conn.disconnect().await?;
+        Ok(())
+    }
+
+    /// Rows per snapshot partition of `bounds`, i.e. the half-open ranges
+    /// `[..b0), [b0, b1), .., [bn, ..)`. The server counts, so the comparisons
+    /// happen under the column's own collation.
+    async fn partition_counts(
+        conn: &mut mysql_async::Conn,
+        db: &str,
+        bounds: &[String],
+        total: u64,
+    ) -> Result<Vec<u64>, anyhow::Error> {
+        let mut counts = Vec::with_capacity(bounds.len() + 1);
+        let mut below = 0;
+        for bound in bounds {
+            let cumulative: Option<u64> = conn
+                .exec_first(
+                    format!("SELECT COUNT(*) FROM {db}.t WHERE id < ?"),
+                    (bound.as_str(),),
+                )
+                .await?;
+            let cumulative = cumulative.expect("COUNT returns a row");
+            counts.push(cumulative - below);
+            below = cumulative;
+        }
+        counts.push(total - below);
+        Ok(counts)
+    }
+}
diff --git a/src/mysql-util/src/probe.rs b/src/mysql-util/src/probe.rs
index 676a50368d..c24e16e887 100644
--- a/src/mysql-util/src/probe.rs
+++ b/src/mysql-util/src/probe.rs
@@ -229,8 +229,9 @@ where
     Ok(estimate)
 }

+/// The live MySQL harness here is shared with [`crate::partition`]'s tests.
 #[cfg(test)]
-mod tests {
+pub(crate) mod tests {
     use std::collections::BTreeSet;

     use mz_ore::cast::CastFrom;
@@ -722,6 +723,34 @@ mod tests {
         Ok(())
     }

+    #[mz_ore::test(tokio::test)]
+    #[cfg_attr(miri, ignore)]
+    async fn test_live_mysql_keys_below_empty_string() -> Result<(), anyhow::Error> {
+        let Some(mut conn) = connect().await? else {
+            return Ok(());
+        };
+        const DB: &str = "mz_probe_below_empty_test";
+        // `utf8mb4_bin` compares character by character but is PAD SPACE, so
+        // comparing a key against `''` pads `''` out with spaces and every key
+        // starting below space collates below `''`. An empty lower bound has to
+        // mean no lower bound, or that whole region of the key space is
+        // invisible to both `prefix_of_first_key_in_range` and the
+        // `max_key_with_prefix` step that advances off a prefix.
+        let keys = ["\u{1}a", "\u{9}b", "a1", "b1"];
+        let table = setup_table(&mut conn, DB, "utf8mb4_bin", &keys).await?;
+
+        let p = &mut KeyProber::new(&mut conn, table, "id");
+        assert_eq!(
+            prefix_of_first_key_in_range(p, "", None, 1).await,
+            some("\u{1}")
+        );
+        assert_eq!(walk_prefixes(p, 2).await?, keys);
+
+        drop_db(&mut conn, 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> {
@@ -974,7 +1003,7 @@ mod tests {
     /// Connects to the server named by `MZ_TEST_MYSQL_URL`, or `None` to skip
     /// the test when it is unset. Skipping is a local-only convenience, CI
     /// must always provide the URL.
-    async fn connect() -> Result<Option<mysql_async::Conn>, anyhow::Error> {
+    pub(crate) async fn connect() -> Result<Option<mysql_async::Conn>, 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!");
@@ -1002,7 +1031,7 @@ mod tests {
     /// Recreates scratch database `db` holding one table `t` whose string
     /// primary key `id` is pinned to the given `collation`, containing
     /// `keys`, with fresh statistics. Returns a ref for [`KeyProber::new`].
-    async fn setup_table<'a>(
+    pub(crate) async fn setup_table<'a>(
         conn: &mut mysql_async::Conn,
         db: &'a str,
         collation: &str,
@@ -1032,7 +1061,10 @@ mod tests {
     }

     /// Drops the scratch database `db`.
-    async fn drop_db(conn: &mut mysql_async::Conn, db: &str) -> Result<(), anyhow::Error> {
+    pub(crate) async fn drop_db(
+        conn: &mut mysql_async::Conn,
+        db: &str,
+    ) -> Result<(), anyhow::Error> {
         #[allow(clippy::disallowed_methods)]
         conn.query_drop(format!("DROP DATABASE {db}")).await?;
         Ok(())

Currently both fail:

Details

    thread 'probe::tests::test_live_mysql_keys_below_empty_string' (867008) panicked at src/mysql-util/src/probe.rs:743:9:
    assertion `left == right` failed
      left: Some("a")
     right: Some("\u{1}")
    stack backtrace:
       0: __rustc::rust_begin_unwind
       1: core::panicking::panic_fmt
       2: core::panicking::assert_failed_inner
       3: core::panicking::assert_failed::<core::option::Option<alloc::string::String>, core::option::Option<alloc::string::String>>
       4: mz_mysql_util::probe::tests::test_live_mysql_keys_below_empty_string::{closure#0}::test_impl::{closure#0}
       5: mz_mysql_util::probe::tests::test_live_mysql_keys_below_empty_string::{closure#0}
       6: <core::pin::Pin<&mut dyn core::future::future::Future<Output = core::result::Result<(), anyhow::Error>>> as core::future::future::Future>::poll
       7: <tracing::instrument::Instrumented<core::pin::Pin<&mut dyn core::future::future::Future<Output = core::result::Result<(), anyhow::Error>>>> as core::future::future::Future>::poll
       8: <tokio::runtime::park::CachedParkThread>::block_on::<tracing::instrument::Instrumented<core::pin::Pin<&mut dyn core::future::future::Future<Output = core::result::Result<(), anyhow::Error>>>>>::{closure#0}
       9: <tokio::runtime::scheduler::current_thread::CoreGuard>::block_on::<core::pin::Pin<&mut tracing::instrument::Instrumented<core::pin::Pin<&mut dyn core::future::future::Future<Output = core::result::Result<(), anyhow::Error>>>>>>::{closure#0}::{closure#0}::{closure#0}
      10: <tokio::runtime::scheduler::current_thread::CoreGuard>::block_on::<core::pin::Pin<&mut tracing::instrument::Instrumented<core::pin::Pin<&mut dyn core::future::future::Future<Output = core::result::Result<(), anyhow::Error>>>>>>::{closure#0}::{closure#0}
      11: <tokio::runtime::scheduler::current_thread::Context>::enter::<core::task::poll::Poll<core::result::Result<(), anyhow::Error>>, <tokio::runtime::scheduler::current_thread::CoreGuard>::block_on<core::pin::Pin<&mut tracing::instrument::Instrumented<core::pin::Pin<&mut dyn core::future::future::Future<Output = core::result::Result<(), anyhow::Error>>>>>>::{closure#0}::{closure#0}>
      12: <tokio::runtime::scheduler::current_thread::CoreGuard>::block_on::<core::pin::Pin<&mut tracing::instrument::Instrumented<core::pin::Pin<&mut dyn core::future::future::Future<Output = core::result::Result<(), anyhow::Error>>>>>>::{closure#0}
      13: <tokio::runtime::scheduler::current_thread::CoreGuard>::enter::<<tokio::runtime::scheduler::current_thread::CoreGuard>::block_on<core::pin::Pin<&mut tracing::instrument::Instrumented<core::pin::Pin<&mut dyn core::future::future::Future<Output = core::result::Result<(), anyhow::Error>>>>>>::{closure#0}, core::option::Option<core::result::Result<(), anyhow::Error>>>::{closure#0}
      14: <tokio::runtime::context::scoped::Scoped<tokio::runtime::scheduler::Context>>::set::<<tokio::runtime::scheduler::current_thread::CoreGuard>::enter<<tokio::runtime::scheduler::current_thread::CoreGuard>::block_on<core::pin::Pin<&mut tracing::instrument::Instrumented<core::pin::Pin<&mut dyn core::future::future::Future<Output = core::result::Result<(), anyhow::Error>>>>>>::{closure#0}, core::option::Option<core::result::Result<(), anyhow::Error>>>::{closure#0}, (alloc::boxed::Box<tokio::runtime::scheduler::current_thread::Core>, core::option::Option<core::result::Result<(), anyhow::Error>>)>
      15: tokio::runtime::context::set_scheduler::<(alloc::boxed::Box<tokio::runtime::scheduler::current_thread::Core>, core::option::Option<core::result::Result<(), anyhow::Error>>), <tokio::runtime::scheduler::current_thread::CoreGuard>::enter<<tokio::runtime::scheduler::current_thread::CoreGuard>::block_on<core::pin::Pin<&mut tracing::instrument::Instrumented<core::pin::Pin<&mut dyn core::future::future::Future<Output = core::result::Result<(), anyhow::Error>>>>>>::{closure#0}, core::option::Option<core::result::Result<(), anyhow::Error>>>::{closure#0}>::{closure#0}
      16: <std::thread::local::LocalKey<tokio::runtime::context::Context>>::try_with::<tokio::runtime::context::set_scheduler<(alloc::boxed::Box<tokio::runtime::scheduler::current_thread::Core>, core::option::Option<core::result::Result<(), anyhow::Error>>), <tokio::runtime::scheduler::current_thread::CoreGuard>::enter<<tokio::runtime::scheduler::current_thread::CoreGuard>::block_on<core::pin::Pin<&mut tracing::instrument::Instrumented<core::pin::Pin<&mut dyn core::future::future::Future<Output = core::result::Result<(), anyhow::Error>>>>>>::{closure#0}, core::option::Option<core::result::Result<(), anyhow::Error>>>::{closure#0}>::{closure#0}, (alloc::boxed::Box<tokio::runtime::scheduler::current_thread::Core>, core::option::Option<core::result::Result<(), anyhow::Error>>)>
      17: <std::thread::local::LocalKey<tokio::runtime::context::Context>>::with::<tokio::runtime::context::set_scheduler<(alloc::boxed::Box<tokio::runtime::scheduler::current_thread::Core>, core::option::Option<core::result::Result<(), anyhow::Error>>), <tokio::runtime::scheduler::current_thread::CoreGuard>::enter<<tokio::runtime::scheduler::current_thread::CoreGuard>::block_on<core::pin::Pin<&mut tracing::instrument::Instrumented<core::pin::Pin<&mut dyn core::future::future::Future<Output = core::result::Result<(), anyhow::Error>>>>>>::{closure#0}, core::option::Option<core::result::Result<(), anyhow::Error>>>::{closure#0}>::{closure#0}, (alloc::boxed::Box<tokio::runtime::scheduler::current_thread::Core>, core::option::Option<core::result::Result<(), anyhow::Error>>)>
      18: tokio::runtime::context::set_scheduler::<(alloc::boxed::Box<tokio::runtime::scheduler::current_thread::Core>, core::option::Option<core::result::Result<(), anyhow::Error>>), <tokio::runtime::scheduler::current_thread::CoreGuard>::enter<<tokio::runtime::scheduler::current_thread::CoreGuard>::block_on<core::pin::Pin<&mut tracing::instrument::Instrumented<core::pin::Pin<&mut dyn core::future::future::Future<Output = core::result::Result<(), anyhow::Error>>>>>>::{closure#0}, core::option::Option<core::result::Result<(), anyhow::Error>>>::{closure#0}>
      19: <tokio::runtime::scheduler::current_thread::CoreGuard>::enter::<<tokio::runtime::scheduler::current_thread::CoreGuard>::block_on<core::pin::Pin<&mut tracing::instrument::Instrumented<core::pin::Pin<&mut dyn core::future::future::Future<Output = core::result::Result<(), anyhow::Error>>>>>>::{closure#0}, core::option::Option<core::result::Result<(), anyhow::Error>>>
      20: <tokio::runtime::scheduler::current_thread::CoreGuard>::block_on::<core::pin::Pin<&mut tracing::instrument::Instrumented<core::pin::Pin<&mut dyn core::future::future::Future<Output = core::result::Result<(), anyhow::Error>>>>>>
      21: <tokio::runtime::scheduler::current_thread::CurrentThread>::block_on::<tracing::instrument::Instrumented<core::pin::Pin<&mut dyn core::future::future::Future<Output = core::result::Result<(), anyhow::Error>>>>>::{closure#0}
      22: tokio::runtime::context::runtime::enter_runtime::<<tokio::runtime::scheduler::current_thread::CurrentThread>::block_on<tracing::instrument::Instrumented<core::pin::Pin<&mut dyn core::future::future::Future<Output = core::result::Result<(), anyhow::Error>>>>>::{closure#0}, core::result::Result<(), anyhow::Error>>
      23: <tokio::runtime::scheduler::current_thread::CurrentThread>::block_on::<tracing::instrument::Instrumented<core::pin::Pin<&mut dyn core::future::future::Future<Output = core::result::Result<(), anyhow::Error>>>>>
      24: <tokio::runtime::runtime::Runtime>::block_on_inner::<core::pin::Pin<&mut dyn core::future::future::Future<Output = core::result::Result<(), anyhow::Error>>>>
      25: <tokio::runtime::runtime::Runtime>::block_on::<core::pin::Pin<&mut dyn core::future::future::Future<Output = core::result::Result<(), anyhow::Error>>>>
      26: mz_mysql_util::probe::tests::test_live_mysql_keys_below_empty_string
      27: mz_mysql_util::probe::tests::test_live_mysql_keys_below_empty_string::{closure#0}
      28: <mz_mysql_util::probe::tests::test_live_mysql_keys_below_empty_string::{closure#0} as core::ops::function::FnOnce<()>>::call_once
    note: Some details are omitted, run with `RUST_BACKTRACE=full` for a verbose backtrace.

  Cancelling due to test failure: 1 test still running
        FAIL [   3.262s] mz-mysql-util partition::tests::test_live_mysql_partition_keys_below_empty_string
  stdout ───

    running 1 test
    test partition::tests::test_live_mysql_partition_keys_below_empty_string ... FAILED

    failures:

    failures:
        partition::tests::test_live_mysql_partition_keys_below_empty_string

    test result: FAILED. 0 passed; 1 failed; 0 ignored; 0 measured; 32 filtered out; finished in 3.26s

  stderr ───

    thread 'partition::tests::test_live_mysql_partition_keys_below_empty_string' (867009) panicked at src/mysql-util/src/partition.rs:231:17:
    mz_partition_all_below_test workers=2 largest=200 of 200, counts=[200]
    stack backtrace:
       0: __rustc::rust_begin_unwind
       1: core::panicking::panic_fmt
       2: mz_mysql_util::partition::tests::test_live_mysql_partition_keys_below_empty_string::{closure#0}::test_impl::{closure#0}
       3: mz_mysql_util::partition::tests::test_live_mysql_partition_keys_below_empty_string::{closure#0}
       4: <core::pin::Pin<&mut dyn core::future::future::Future<Output = core::result::Result<(), anyhow::Error>>> as core::future::future::Future>::poll
       5: <tracing::instrument::Instrumented<core::pin::Pin<&mut dyn core::future::future::Future<Output = core::result::Result<(), anyhow::Error>>>> as core::future::future::Future>::poll
       6: <tokio::runtime::park::CachedParkThread>::block_on::<tracing::instrument::Instrumented<core::pin::Pin<&mut dyn core::future::future::Future<Output = core::result::Result<(), anyhow::Error>>>>>::{closure#0}
       7: <tokio::runtime::scheduler::current_thread::CoreGuard>::block_on::<core::pin::Pin<&mut tracing::instrument::Instrumented<core::pin::Pin<&mut dyn core::future::future::Future<Output = core::result::Result<(), anyhow::Error>>>>>>::{closure#0}::{closure#0}::{closure#0}
       8: <tokio::runtime::scheduler::current_thread::CoreGuard>::block_on::<core::pin::Pin<&mut tracing::instrument::Instrumented<core::pin::Pin<&mut dyn core::future::future::Future<Output = core::result::Result<(), anyhow::Error>>>>>>::{closure#0}::{closure#0}
       9: <tokio::runtime::scheduler::current_thread::Context>::enter::<core::task::poll::Poll<core::result::Result<(), anyhow::Error>>, <tokio::runtime::scheduler::current_thread::CoreGuard>::block_on<core::pin::Pin<&mut tracing::instrument::Instrumented<core::pin::Pin<&mut dyn core::future::future::Future<Output = core::result::Result<(), anyhow::Error>>>>>>::{closure#0}::{closure#0}>
      10: <tokio::runtime::scheduler::current_thread::CoreGuard>::block_on::<core::pin::Pin<&mut tracing::instrument::Instrumented<core::pin::Pin<&mut dyn core::future::future::Future<Output = core::result::Result<(), anyhow::Error>>>>>>::{closure#0}
      11: <tokio::runtime::scheduler::current_thread::CoreGuard>::enter::<<tokio::runtime::scheduler::current_thread::CoreGuard>::block_on<core::pin::Pin<&mut tracing::instrument::Instrumented<core::pin::Pin<&mut dyn core::future::future::Future<Output = core::result::Result<(), anyhow::Error>>>>>>::{closure#0}, core::option::Option<core::result::Result<(), anyhow::Error>>>::{closure#0}
      12: <tokio::runtime::context::scoped::Scoped<tokio::runtime::scheduler::Context>>::set::<<tokio::runtime::scheduler::current_thread::CoreGuard>::enter<<tokio::runtime::scheduler::current_thread::CoreGuard>::block_on<core::pin::Pin<&mut tracing::instrument::Instrumented<core::pin::Pin<&mut dyn core::future::future::Future<Output = core::result::Result<(), anyhow::Error>>>>>>::{closure#0}, core::option::Option<core::result::Result<(), anyhow::Error>>>::{closure#0}, (alloc::boxed::Box<tokio::runtime::scheduler::current_thread::Core>, core::option::Option<core::result::Result<(), anyhow::Error>>)>
      13: tokio::runtime::context::set_scheduler::<(alloc::boxed::Box<tokio::runtime::scheduler::current_thread::Core>, core::option::Option<core::result::Result<(), anyhow::Error>>), <tokio::runtime::scheduler::current_thread::CoreGuard>::enter<<tokio::runtime::scheduler::current_thread::CoreGuard>::block_on<core::pin::Pin<&mut tracing::instrument::Instrumented<core::pin::Pin<&mut dyn core::future::future::Future<Output = core::result::Result<(), anyhow::Error>>>>>>::{closure#0}, core::option::Option<core::result::Result<(), anyhow::Error>>>::{closure#0}>::{closure#0}
      14: <std::thread::local::LocalKey<tokio::runtime::context::Context>>::try_with::<tokio::runtime::context::set_scheduler<(alloc::boxed::Box<tokio::runtime::scheduler::current_thread::Core>, core::option::Option<core::result::Result<(), anyhow::Error>>), <tokio::runtime::scheduler::current_thread::CoreGuard>::enter<<tokio::runtime::scheduler::current_thread::CoreGuard>::block_on<core::pin::Pin<&mut tracing::instrument::Instrumented<core::pin::Pin<&mut dyn core::future::future::Future<Output = core::result::Result<(), anyhow::Error>>>>>>::{closure#0}, core::option::Option<core::result::Result<(), anyhow::Error>>>::{closure#0}>::{closure#0}, (alloc::boxed::Box<tokio::runtime::scheduler::current_thread::Core>, core::option::Option<core::result::Result<(), anyhow::Error>>)>
      15: <std::thread::local::LocalKey<tokio::runtime::context::Context>>::with::<tokio::runtime::context::set_scheduler<(alloc::boxed::Box<tokio::runtime::scheduler::current_thread::Core>, core::option::Option<core::result::Result<(), anyhow::Error>>), <tokio::runtime::scheduler::current_thread::CoreGuard>::enter<<tokio::runtime::scheduler::current_thread::CoreGuard>::block_on<core::pin::Pin<&mut tracing::instrument::Instrumented<core::pin::Pin<&mut dyn core::future::future::Future<Output = core::result::Result<(), anyhow::Error>>>>>>::{closure#0}, core::option::Option<core::result::Result<(), anyhow::Error>>>::{closure#0}>::{closure#0}, (alloc::boxed::Box<tokio::runtime::scheduler::current_thread::Core>, core::option::Option<core::result::Result<(), anyhow::Error>>)>
      16: tokio::runtime::context::set_scheduler::<(alloc::boxed::Box<tokio::runtime::scheduler::current_thread::Core>, core::option::Option<core::result::Result<(), anyhow::Error>>), <tokio::runtime::scheduler::current_thread::CoreGuard>::enter<<tokio::runtime::scheduler::current_thread::CoreGuard>::block_on<core::pin::Pin<&mut tracing::instrument::Instrumented<core::pin::Pin<&mut dyn core::future::future::Future<Output = core::result::Result<(), anyhow::Error>>>>>>::{closure#0}, core::option::Option<core::result::Result<(), anyhow::Error>>>::{closure#0}>
      17: <tokio::runtime::scheduler::current_thread::CoreGuard>::enter::<<tokio::runtime::scheduler::current_thread::CoreGuard>::block_on<core::pin::Pin<&mut tracing::instrument::Instrumented<core::pin::Pin<&mut dyn core::future::future::Future<Output = core::result::Result<(), anyhow::Error>>>>>>::{closure#0}, core::option::Option<core::result::Result<(), anyhow::Error>>>
      18: <tokio::runtime::scheduler::current_thread::CoreGuard>::block_on::<core::pin::Pin<&mut tracing::instrument::Instrumented<core::pin::Pin<&mut dyn core::future::future::Future<Output = core::result::Result<(), anyhow::Error>>>>>>
      19: <tokio::runtime::scheduler::current_thread::CurrentThread>::block_on::<tracing::instrument::Instrumented<core::pin::Pin<&mut dyn core::future::future::Future<Output = core::result::Result<(), anyhow::Error>>>>>::{closure#0}
      20: tokio::runtime::context::runtime::enter_runtime::<<tokio::runtime::scheduler::current_thread::CurrentThread>::block_on<tracing::instrument::Instrumented<core::pin::Pin<&mut dyn core::future::future::Future<Output = core::result::Result<(), anyhow::Error>>>>>::{closure#0}, core::result::Result<(), anyhow::Error>>
      21: <tokio::runtime::scheduler::current_thread::CurrentThread>::block_on::<tracing::instrument::Instrumented<core::pin::Pin<&mut dyn core::future::future::Future<Output = core::result::Result<(), anyhow::Error>>>>>
      22: <tokio::runtime::runtime::Runtime>::block_on_inner::<core::pin::Pin<&mut dyn core::future::future::Future<Output = core::result::Result<(), anyhow::Error>>>>
      23: <tokio::runtime::runtime::Runtime>::block_on::<core::pin::Pin<&mut dyn core::future::future::Future<Output = core::result::Result<(), anyhow::Error>>>>
      24: mz_mysql_util::partition::tests::test_live_mysql_partition_keys_below_empty_string
      25: mz_mysql_util::partition::tests::test_live_mysql_partition_keys_below_empty_string::{closure#0}
      26: <mz_mysql_util::partition::tests::test_live_mysql_partition_keys_below_empty_string::{closure#0} as core::ops::function::FnOnce<()>>::call_once
    note: Some details are omitted, run with `RUST_BACKTRACE=full` for a verbose backtrace.

Currently only seems to be a performance issue where we wouldn't parallelize the ingestion, so not too bad.

martykulma · 2026-08-11T19:38:54Z

+    let estimated_row_count = estimated_row_count.max(1);
+
+    // Estimates vary wildly especially near the full table size (see `KeyProber::estimate_range_rows` for more details).
+    // Estimates tend to get more useful as smaller chunks, so break up the table into at least 1/8ths (2 workers * 4)


The comment should explain why 1/8th (e.g. due to some testing, arbitrary, etc.).

Added some explanation.

martykulma · 2026-08-11T19:51:13Z

+    db: &mut KeyProber<'_>,
+    workers: usize,
+    estimated_row_count: u64,
+    min_rows_per_worker: u64,


min_rows_per_worker seems to only guide the splitting, it doesn't enforce any bounds on the rows per worker. If the intention is to guide splitting, maybe a better name would be min_split_threshold with some explanation in a doc comment.

yeah, in my mind it's "the minimum number of rows to bother trying to split down to". Will update and document.

martykulma · 2026-08-11T20:06:29Z

+/// Computes up to `num_workers - 1` partition boundaries that divide the primary key space
+/// into `num_workers` roughly even partitions.
+/// This should be run in a repeatable read transaction against a primary key varchar/char column
+/// with the `utf8mb4_bin` collation.


The doc needs to make it very clear that these are invariants expected of the caller. Since this function doesn't assert any of those things, a heads up on what goes wrong if the invariants aren't upheld.

KeyProber docs should call out similar invariants.

### Motivation Split out some testing from #38022 to keep that PR manageable. Part of: [SS-97](https://linear.app/materializeinc/issue/SS-97) ### Description Adds testing for more diverse character types, collations, and for asserting searches use the B-tree rather than scanning the full table. ### For Context I've caught a couple of bugs in the later PR from discussion with Marty and a review with Dennis: #38093. 1. For PAD SPACE collations certain characters will sort before a shorter string. Our algorithm still works cleanly for that case at a high level, just throwing out the lower bounds which for utf8mb4_bin are all things like null, carriage return, and tab. The one tricky spot we needed to handle correctly was the upper bound, for which we're able to append NUL characters to get behavior like what we would have expected initially. I tried fixing this directly initially, but thought the fix was more complex than the perf issue was worth (see [here](#38147)). 2. Certain collations can't be supported with this approach because multiple characters in a row can in fact sort differently from it's prefix. There are tests here highlighting that and then most of the tests have been moved to utf8mb4_bin to match the first collation we intend to support.

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.

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.

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.

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.

The trait is the partitioner-owned seam over the concrete KeyProber, a distinct name keeps it from shadowing the prober it wraps.

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.

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.

…and smaller partitions for more accuracy

martykulma

🐟

peterdukelarsen force-pushed the pl/mysql-pk-partitioner branch 2 times, most recently from 8753eb8 to 50b66ec Compare August 6, 2026 23:36

peterdukelarsen marked this pull request as ready for review August 7, 2026 00:03

peterdukelarsen requested a review from a team as a code owner August 7, 2026 00:03

peterdukelarsen requested a review from a team August 7, 2026 00:03

peterdukelarsen force-pushed the pl/mysql-pk-partitioner branch 5 times, most recently from 6cdaec2 to fe76d04 Compare August 8, 2026 01:44

peterdukelarsen force-pushed the pl/mysql-pk-partitioner branch from fe76d04 to fcfd0ea Compare August 9, 2026 17:35

def- reviewed Aug 10, 2026

View reviewed changes

peterdukelarsen force-pushed the pl/mysql-pk-partitioner branch from 04e3e3f to 5fee710 Compare August 11, 2026 20:14

martykulma reviewed Aug 11, 2026

View reviewed changes

peterdukelarsen force-pushed the pl/mysql-pk-partitioner branch 2 times, most recently from 47030b9 to d765772 Compare August 11, 2026 23:02

peterdukelarsen force-pushed the pl/mysql-pk-partitioner branch from 7b3cfb1 to f993f7b Compare August 11, 2026 23:42

peterdukelarsen requested review from a team and ggevay as code owners August 11, 2026 23:42

peterdukelarsen requested a review from jdonelson August 11, 2026 23:42

peterdukelarsen changed the base branch from pl/mysql-key-probes-tests to main August 11, 2026 23:42

peterdukelarsen added 10 commits August 11, 2026 20:11

mysql-util: name the partitioner probe trait PrimaryKeyProber

83cdc0b

The trait is the partitioner-owned seam over the concrete KeyProber, a distinct name keeps it from shadowing the prober it wraps.

mysql-util: shorten partitioner naming

c827ecd

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.

mysql-util: remove surrogate key, revive O(depth*prefixes) algorithm …

ab64a4a

…and smaller partitions for more accuracy

mysql-util: use integer row counts and rename min_split_threshold

1774d2d

mysql-util: explain the partition target selection and split threshold

a3c03ee

peterdukelarsen force-pushed the pl/mysql-pk-partitioner branch from f993f7b to a3c03ee Compare August 12, 2026 00:12

martykulma approved these changes Aug 12, 2026

View reviewed changes

peterdukelarsen merged commit a266159 into main Aug 12, 2026
79 checks passed

peterdukelarsen deleted the pl/mysql-pk-partitioner branch August 12, 2026 01:07

peterdukelarsen mentioned this pull request Aug 12, 2026

Add tests for MySQL primary key prefix partitioner #38089

Open

Navigation Menu

Uh oh!

Mysql string primary key partitioner - #38093

Mysql string primary key partitioner#38093
peterdukelarsen merged 10 commits into
mainfrom
pl/mysql-pk-partitioner

peterdukelarsen commented Aug 6, 2026 •

edited

Loading

Uh oh!

linear-code Bot commented Aug 6, 2026

Uh oh!

def- left a comment

Uh oh!

Uh oh!

martykulma Aug 11, 2026

Uh oh!

peterdukelarsen Aug 11, 2026

Uh oh!

martykulma Aug 11, 2026

Uh oh!

peterdukelarsen Aug 11, 2026 •

edited

Loading

Uh oh!

martykulma Aug 11, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

martykulma left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

peterdukelarsen commented Aug 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Description

Uh oh!

linear-code Bot commented Aug 6, 2026

Uh oh!

def- left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

martykulma Aug 11, 2026

Choose a reason for hiding this comment

Uh oh!

peterdukelarsen Aug 11, 2026

Choose a reason for hiding this comment

Uh oh!

martykulma Aug 11, 2026

Choose a reason for hiding this comment

Uh oh!

peterdukelarsen Aug 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

martykulma Aug 11, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

martykulma left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

peterdukelarsen commented Aug 6, 2026 •

edited

Loading

peterdukelarsen Aug 11, 2026 •

edited

Loading