Flesh out testing for string primary key probes

[peterdukelarsen]

Motivation

Split out some testing from #38022 to keep that PR manageable.

Part of: 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.

Verification

[ublubu]

Thanks for all these tests. Walking through the examples gave me some ideas.

[ublubu]

Does this comment refer to L352 above? "utf8mb4_0900_ai_ci"

[peterdukelarsen]

Indirectly I guess. It's more about the immediate following lines being mixed case.

Line 352 is relevant because it is what sets it to case-insensitive sorting explicitly (although that's mysql's default). It is more specifically about why I have a specific mix of upper and lower case "A" "b" "C" in the following lines.

[ublubu]

I know what's happening here, but it'd be nice to have some comments during the traversal to say things like:

• Segment the key space with 1-character prefixes as fenceposts.
• Subdivide the [A, b) segment with 2-character prefixes as fenceposts.
• " [b, C) "
• " [C, ...) ". No further subdivision is possible, so we find no new fenceposts.

It may be worth noting that we must keep all the too-short fenceposts because of the a, aa, aaa situation.
That is, row a lives in its own partition [a, aa). If we forgot about fencepost (prefix) a, we'd start at aa (and row a would be lost).

---

Re: too-short prefixes

When we call prefix_of_first_row_not_matching_prefix, there are two ways we might get len(prefix) < max_prefix_length:

1. We got prefix from a previous iteration with a lower max_prefix_length.
2. Our last call to prefix_of_first_row_not_matching_prefix found a too-short row.

If we start with max_prefix_length=1 and only increment the length by 1 between rounds, scenario 2 can't happen...
unless there's a row with length 0.
(But if we start at max_prefix_length=2 or increment by 2, then scenario 2 can happen.)

Regardless of whether it's scenario 1 or 2, it's possible that a too-short prefix means we have a short row that we must account for. (See suggested note in previous section.)

---

I had a think, and I believe there's a clean way to explain everything:

Suppose we're segmenting the key range [a, aa, aaa, b, ba, c].
Suppose we want to use 2-character prefixes as our fenceposts.

problem: a isn't long enough to match any 2-character prefix.
solution: Let 0 be a NULL character. Treat each key as having a bunch of NULLs after it.
For example, a, a0, and a00 are equivalent. Now we can use the prefix a0 to match a.

Now, the algorithm (which we're already following) can be described like this:

1. You have a range of keys. For example, the range encompassing all possible keys is ['', inf), from the empty prefix up to infinity.

2. You subdivide the range of keys by increasing the prefix length.

   • '' (length 0) becomes '<null>' (length 1), which only matches the empty string. Then prefix_of_first_row_not_matching_prefix gives you the prefix 'a'. Then 'b', etc.
   • 'a' becomes 'a<null>', so the next 2-character prefix is 'aa'.
   • If you go straight from 0 to 2, '' becomes '<null><null>'. The next prefix is 'a<null>'. Then 'aa', etc.

3. Return to step 1 for any sub-range that is still too big.

At the end, you should have a contiguous sequence of ranges, each fenceposted by a single prefix.

This way we don't have to think too hard about any special cases. And we can easily explain why we use different logic to match against too-short prefixes.

[peterdukelarsen]

This is really well-said and gets at the core of what I was discussing with Marty yesterday. What I landed on doing was actually intentionally throwing out the short keys if the buckets are too big and we need to step down into them.

The two key contextual pieces there are:

1. I'm planning to have a lower bound of 10-20k for a single prefix that we try to split. So losing 1 row out of 10k doesn't meaningfully impact the results
2. I'm planning to limit the budget of specific calls to the DB to exit early -- i.e. ~3k for a big table. This means we're already sunk in the case of a jagged key.

All-together for the actual algorithm I think that leaves us safe to represent it as a list of Prefixes, where a prefix is defined by the prefix itself and an upper bound that we carry through.

[ublubu]

Suggested change

	async fn test_wild_card_char_in_data() -> Result<(), anyhow::Error> {
	async fn test_wildcard_char_in_data() -> Result<(), anyhow::Error> {

[ublubu]

I've reviewed up to here 😅

[def-]

This test doesn't seem useful at the moment. Something like this:

diff --git a/src/mysql-util/src/probe.rs b/src/mysql-util/src/probe.rs
index c8dde4121c..a79f7c7fa4 100644
--- a/src/mysql-util/src/probe.rs
+++ b/src/mysql-util/src/probe.rs
@@ -676,14 +676,22 @@ mod tests {
         // 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.
+        // one interval, under the prefix bounding it. 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?;
+                let (n, under_lo) = count_range(&mut conn, DB, lo, walked.get(i + 1)).await?;
+                // The total below telescopes to `COUNT(id >= walked[0])` for
+                // any ascending walk, so these two carry the test: an empty
+                // interval means a spurious boundary, a key `lo` is not a
+                // prefix of means a skipped one.
+                let ctx = format!("len={len} lo={lo:?} walked={walked:?}");
+                assert!(n > 0, "empty interval: {ctx}");
+                assert_eq!(under_lo, n, "keys not under the prefix: {ctx}");
+                total += n;
             }
             assert_eq!(
                 total,
@@ -1053,27 +1061,34 @@ 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.
+    /// Keys in `[lo, hi)` of `db`'s table: the total, and how many have `lo`
+    /// as a prefix. Counted by the server so comparisons happen under the
+    /// column's collation, via `LEFT` rather than `LIKE` so the oracle does
+    /// not reuse the escaping under test.
     async fn count_range(
         conn: &mut mysql_async::Conn,
         db: &str,
         lo: &str,
         hi: Option<&String>,
-    ) -> Result<u64, anyhow::Error> {
+    ) -> Result<(u64, u64), anyhow::Error> {
         let mut clause = "id >= ?".to_string();
-        let mut params: Vec<Value> = vec![lo.into()];
+        // The select list's two placeholders bind before the clause's.
+        let mut params: Vec<Value> = vec![lo.into(), lo.into(), lo.into()];
         if let Some(hi) = hi {
             clause.push_str(" AND id < ?");
             params.push(hi.as_str().into());
         }
-        let count: Option<u64> = conn
+        let row: Option<(u64, Option<u64>)> = conn
             .exec_first(
-                format!("SELECT COUNT(*) FROM {db}.t WHERE {clause}"),
+                format!(
+                    "SELECT COUNT(*), SUM(LEFT(id, CHAR_LENGTH(?)) = ?) FROM {db}.t WHERE {clause}"
+                ),
                 Params::Positional(params),
             )
             .await?;
-        Ok(count.expect("COUNT returns a row"))
+        // The SUM is NULL over an empty range.
+        let (total, prefixed) = row.expect("COUNT returns a row");
+        Ok((total, prefixed.unwrap_or(0)))
     }

     /// Sum of this session's `Handler_read_*` counters: how many index or row

Via the QA LLM review: https://github.com/MaterializeInc/qa-llm-review/blob/master/commit-bugs/done/analysis-pr-38030.md