compute: thin peek results by partitioning, not sorting

doc/user/data/metrics.yml

@@ -1032,18 +1032,18 @@ metrics:
   help: Time checking trace frontiers.
   source: src/compute/src/metrics.rs
   visibility: internal
-- name: mz_index_peek_result_sort_seconds_bucket
-  help: Time sorting intermediate results during peek collection.
+- name: mz_index_peek_result_thinning_seconds_bucket
+  help: Time thinning intermediate peek results down to the rows the finishing can need, summed over the activations an index peek took and reported once it is done. Peeks that are cancelled or dropped before finishing report nothing.
   labels:
   - le
   source: src/compute/src/metrics.rs
   visibility: internal
-- name: mz_index_peek_result_sort_seconds_count
-  help: Time sorting intermediate results during peek collection.
+- name: mz_index_peek_result_thinning_seconds_count
+  help: Time thinning intermediate peek results down to the rows the finishing can need, summed over the activations an index peek took and reported once it is done. Peeks that are cancelled or dropped before finishing report nothing.
   source: src/compute/src/metrics.rs
   visibility: internal
-- name: mz_index_peek_result_sort_seconds_sum
-  help: Time sorting intermediate results during peek collection.
+- name: mz_index_peek_result_thinning_seconds_sum
+  help: Time thinning intermediate peek results down to the rows the finishing can need, summed over the activations an index peek took and reported once it is done. Peeks that are cancelled or dropped before finishing report nothing.
   source: src/compute/src/metrics.rs
   visibility: internal
 - name: mz_index_peek_row_collection_seconds_bucket
@@ -1061,45 +1061,45 @@ metrics:
   source: src/compute/src/metrics.rs
   visibility: internal
 - name: mz_index_peek_row_iteration_seconds_bucket
-  help: Time iterating rows and evaluating MFP.
+  help: Time iterating rows and evaluating MFP, summed over the activations an index peek took and reported once it is done. Peeks that are cancelled or dropped before finishing report nothing.
   labels:
   - le
   source: src/compute/src/metrics.rs
   visibility: internal
 - name: mz_index_peek_row_iteration_seconds_count
-  help: Time iterating rows and evaluating MFP.
+  help: Time iterating rows and evaluating MFP, summed over the activations an index peek took and reported once it is done. Peeks that are cancelled or dropped before finishing report nothing.
   source: src/compute/src/metrics.rs
   visibility: internal
 - name: mz_index_peek_row_iteration_seconds_sum
-  help: Time iterating rows and evaluating MFP.
+  help: Time iterating rows and evaluating MFP, summed over the activations an index peek took and reported once it is done. Peeks that are cancelled or dropped before finishing report nothing.
   source: src/compute/src/metrics.rs
   visibility: internal
 - name: mz_index_peek_seek_fulfillment_seconds_bucket
-  help: Time in seek_fulfillment method including frontier checks and data collection.
+  help: Worker time spent in seek_fulfillment for an index peek, including every not-yet-ready frontier check, summed over activations and reported once the peek is done.
   labels:
   - le
   source: src/compute/src/metrics.rs
   visibility: internal
 - name: mz_index_peek_seek_fulfillment_seconds_count
-  help: Time in seek_fulfillment method including frontier checks and data collection.
+  help: Worker time spent in seek_fulfillment for an index peek, including every not-yet-ready frontier check, summed over activations and reported once the peek is done.
   source: src/compute/src/metrics.rs
   visibility: internal
 - name: mz_index_peek_seek_fulfillment_seconds_sum
-  help: Time in seek_fulfillment method including frontier checks and data collection.
+  help: Worker time spent in seek_fulfillment for an index peek, including every not-yet-ready frontier check, summed over activations and reported once the peek is done.
   source: src/compute/src/metrics.rs
   visibility: internal
 - name: mz_index_peek_total_seconds_bucket
-  help: Total time processing index peeks, from process_peek entry to response. Excluding peeks that use the peek response stash.
+  help: Worker time spent serving an index peek, summed over the activations it took and reported once it is done. Peeks that are cancelled or dropped before finishing report nothing.
   labels:
   - le
   source: src/compute/src/metrics.rs
   visibility: internal
 - name: mz_index_peek_total_seconds_count
-  help: Total time processing index peeks, from process_peek entry to response. Excluding peeks that use the peek response stash.
+  help: Worker time spent serving an index peek, summed over the activations it took and reported once it is done. Peeks that are cancelled or dropped before finishing report nothing.
   source: src/compute/src/metrics.rs
   visibility: internal
 - name: mz_index_peek_total_seconds_sum
-  help: Total time processing index peeks, from process_peek entry to response. Excluding peeks that use the peek response stash.
+  help: Worker time spent serving an index peek, summed over the activations it took and reported once it is done. Peeks that are cancelled or dropped before finishing report nothing.
   source: src/compute/src/metrics.rs
   visibility: internal
 - name: mz_kafka_partition_offset_max
@@ -1527,6 +1527,20 @@ metrics:
   help: The time it takes to parse a SQL statement. (Works for both Simple Queries and the Extended Query protocol.)
   source: src/adapter/src/metrics.rs
   visibility: internal
+- name: mz_peek_processing_seconds_bucket
+  help: Time a worker spent serving pending peeks in one activation, before returning to scheduling dataflows and handling commands. This is the latency the peek yielding budgets exist to bound. Only recorded for activations that had at least one pending peek.
+  labels:
+  - le
+  source: src/compute/src/metrics.rs
+  visibility: internal
+- name: mz_peek_processing_seconds_count
+  help: Time a worker spent serving pending peeks in one activation, before returning to scheduling dataflows and handling commands. This is the latency the peek yielding budgets exist to bound. Only recorded for activations that had at least one pending peek.
+  source: src/compute/src/metrics.rs
+  visibility: internal
+- name: mz_peek_processing_seconds_sum
+  help: Time a worker spent serving pending peeks in one activation, before returning to scheduling dataflows and handling commands. This is the latency the peek yielding budgets exist to bound. Only recorded for activations that had at least one pending peek.
+  source: src/compute/src/metrics.rs
+  visibility: internal
 - name: mz_persist_*_bytes
   help: total encoded size of * batches written
   source: src/persist-client/src/internal/metrics.rs

misc/python/materialize/mzcompose/__init__.py

@@ -48,7 +48,15 @@
 # a new feature causes benchmarks to become flaky, consider that this can also
 # impact customers' experience and try to find a solution other than disabling
 # the feature here!
-ADDITIONAL_BENCHMARKING_SYSTEM_PARAMETERS = {}
+ADDITIONAL_BENCHMARKING_SYSTEM_PARAMETERS = {
+    # Benchmarks measure against an older Materialize image that does not know
+    # these parameters and therefore ignores them, so a non-production value
+    # here would only slow down one side of the comparison. Pin both to their
+    # production defaults, so peek-heavy scenarios measure the same thing on
+    # both sides.
+    "peek_yielding": "work:100000,time:10",
+    "peek_yielding_total": "work:1000000,time:100",
+}
 
 
 def get_minimal_system_parameters(
@@ -345,6 +353,27 @@ def get_variable_system_parameters(
         VariableSystemParameter(
             "persist_encoding_enable_dictionary", "true", ["true", "false"]
         ),
+        # A work budget well below production, so that any index peek over more
+        # than ~1000 rows resumes at least once and CI covers the resumable scan
+        # path rather than only the completes-in-one-slice path. Peeks over tiny
+        # relations, which is most of sqllogictest, still finish in one slice;
+        # that case is covered directly by
+        # test/clusterd-test-driver/scripts/peek_yielding.spec.
+        #
+        # Kept within an order of magnitude of production on purpose. The value
+        # reaches every mzcompose composition, and a very small budget buys
+        # little extra coverage while multiplying the timely steps a large peek
+        # needs. The randomized runs go lower.
+        VariableSystemParameter(
+            "peek_yielding",
+            "work:1024,time:10",
+            ["work:16,time:10", "work:1024,time:10", "work:100000,time:10"],
+        ),
+        VariableSystemParameter(
+            "peek_yielding_total",
+            "work:8192,time:100",
+            ["work:128,time:100", "work:8192,time:100", "work:1000000,time:100"],
+        ),
         VariableSystemParameter(
             "persist_fast_path_limit",
             "1000",

misc/python/materialize/parallel_workload/action.py

@@ -1853,6 +1853,18 @@ def __init__(
             BOOLEAN_FLAG_VALUES
         )
         self.flags_with_values["enable_eager_delta_joins"] = BOOLEAN_FLAG_VALUES
+        # Small work budgets force index peeks to yield and resume, which is
+        # where the interesting state lives.
+        self.flags_with_values["peek_yielding"] = [
+            "'work:16,time:10'",
+            "'work:64,time:10'",
+            "'work:100000,time:10'",
+        ]
+        self.flags_with_values["peek_yielding_total"] = [
+            "'work:128,time:100'",
+            "'work:1024,time:100'",
+            "'work:1000000,time:100'",
+        ]
         self.flags_with_values["enable_public_metrics_endpoint"] = BOOLEAN_FLAG_VALUES
         self.flags_with_values["enable_scoped_system_parameters"] = BOOLEAN_FLAG_VALUES
         self.flags_with_values["persist_batch_structured_key_lower_len"] = [

src/clusterd-test-driver/src/driver.rs

@@ -183,11 +183,18 @@ impl Driver {
     /// materialized-view sink's output shard, which is how `SELECT * FROM mv` reads.
     /// A persist peek blocks (async-friendly) until the shard seals through `ts`, so
     /// it doubles as a wait for the writing sink to catch up.
+    /// Peeks `target` at `ts`, optionally restricted to `literal_constraints`.
+    ///
+    /// Each literal is a row of the index's key columns, the same shape the
+    /// optimizer produces for `WHERE key IN (..)`. Only valid against an index
+    /// peek whose key the literals match. The replica sorts them itself, so they
+    /// may be passed in any order. `None` scans the whole arrangement.
     pub async fn peek(
         &self,
         target: PeekTarget,
         result_desc: RelationDesc,
         ts: Timestamp,
+        literal_constraints: Option<Vec<Row>>,
     ) -> anyhow::Result<Vec<Row>> {
         let uuid = uuid::Uuid::new_v4();
         let rx = self.responses.register_peek(uuid);
@@ -201,7 +208,7 @@ impl Driver {
         let peek = Peek {
             target,
             result_desc: result_desc.clone(),
-            literal_constraints: None,
+            literal_constraints,
             uuid,
             timestamp: ts,
             finishing: RowSetFinishing::trivial(arity),
@@ -232,8 +239,12 @@ impl Driver {
         target: PeekTarget,
         result_desc: RelationDesc,
         ts: Timestamp,
+        literal_constraints: Option<Vec<Row>>,
     ) -> anyhow::Result<usize> {
-        Ok(self.peek(target, result_desc, ts).await?.len())
+        Ok(self
+            .peek(target, result_desc, ts, literal_constraints)
+            .await?
+            .len())
     }
 
     /// Registers a subscribe-sink buffer for `id`, so the response pump accumulates

src/clusterd-test-driver/src/script.rs

@@ -45,8 +45,8 @@ use mz_expr_parser::{TestCatalog, try_parse_mir};
 use mz_persist_client::PersistClient;
 use mz_persist_types::{PersistLocation, ShardId};
 use mz_repr::{
-    GlobalId, RelationDesc, ReprRelationType, Row, SqlColumnType, SqlRelationType, SqlScalarType,
-    Timestamp, strconv,
+    Datum, GlobalId, RelationDesc, ReprRelationType, Row, SqlColumnType, SqlRelationType,
+    SqlScalarType, Timestamp, strconv,
 };
 use mz_storage_types::controller::CollectionMetadata;
 use serde::{Deserialize, Serialize};
@@ -462,6 +462,24 @@ pub enum Command {
         /// The timestamp to count at.
         ts: u64,
     },
+    /// Peek `id` directly and emit only how many rows came back.
+    ///
+    /// Unlike [`Command::Count`], which tallies through an ephemeral reduce
+    /// dataflow and then peeks its single-row result, this peeks `id` itself. So
+    /// the peek walks `id`'s whole arrangement, which is what makes it useful
+    /// for exercising the peek scan over many rows without a golden that has to
+    /// spell every one of them out.
+    PeekCount {
+        /// The index's global id.
+        id: u64,
+        /// The timestamp to peek at.
+        ts: u64,
+        /// Restrict the peek to these key values, as `WHERE key IN (..)` would.
+        ///
+        /// Only meaningful for a single-column `Int64` key, which is what the
+        /// sample schema's first column is. `None` scans the whole arrangement.
+        literals: Option<Vec<i64>>,
+    },
     /// Submit (without scheduling) a dataflow built from generic MIR — the
     /// abstraction behind index / materialized-view / subscribe / copy-to.
     ///
@@ -664,6 +682,7 @@ impl ScriptState {
                 PeekTarget::Index { id: out_index_id },
                 count_desc,
                 Timestamp::from(ts),
+                None,
             )
             .await?;
         match rows.as_slice() {
@@ -1049,9 +1068,29 @@ impl ScriptState {
                         id: GlobalId::User(id),
                     },
                 };
-                let rows = self.driver.peek(target, desc, Timestamp::from(ts)).await?;
+                let rows = self
+                    .driver
+                    .peek(target, desc, Timestamp::from(ts), None)
+                    .await?;
                 Ok(render_rows(&rows))
             }
+            Command::PeekCount { id, ts, literals } => {
+                let desc = self.resolve_schema(&None)?;
+                let target = PeekTarget::Index {
+                    id: GlobalId::User(id),
+                };
+                let literal_constraints = literals.map(|literals| {
+                    literals
+                        .into_iter()
+                        .map(|literal| Row::pack_slice(&[Datum::Int64(literal)]))
+                        .collect()
+                });
+                let count = self
+                    .driver
+                    .peek_count(target, desc, Timestamp::from(ts), literal_constraints)
+                    .await?;
+                Ok(count.to_string())
+            }
             Command::AwaitSubscribe {
                 id,
                 up_to,

src/clusterd-test-driver/src/text.rs

@@ -279,6 +279,24 @@ fn opt_string(args: &BTreeMap<String, String>, key: &str) -> Option<String> {
 }
 
 /// Parse a `[a,b,c]` list of `usize`s (`[]` is empty).
+fn parse_i64_list(s: &str) -> anyhow::Result<Vec<i64>> {
+    let inner = s
+        .strip_prefix('[')
+        .and_then(|s| s.strip_suffix(']'))
+        .ok_or_else(|| anyhow!("expected a list like `[0,1]`, got `{s}`"))?;
+    if inner.trim().is_empty() {
+        return Ok(vec![]);
+    }
+    inner
+        .split(',')
+        .map(|part| {
+            part.trim()
+                .parse()
+                .with_context(|| format!("bad list element `{part}`"))
+        })
+        .collect()
+}
+
 fn parse_usize_list(s: &str) -> anyhow::Result<Vec<usize>> {
     let inner = s
         .strip_prefix('[')
@@ -490,6 +508,14 @@ fn parse_command(input: &str) -> anyhow::Result<Command> {
             schema: opt_string(&args, "schema"),
             ts: req_u64(&args, "ts")?,
         },
+        "peek-count" => Command::PeekCount {
+            id: req_u64(&args, "id")?,
+            ts: req_u64(&args, "ts")?,
+            literals: opt_string(&args, "literals")
+                .as_deref()
+                .map(parse_i64_list)
+                .transpose()?,
+        },
         "await-subscribe" => Command::AwaitSubscribe {
             id: req_u64(&args, "id")?,
             up_to: req_u64(&args, "up-to")?,

src/clusterd-test-driver/tests/index_smoke.rs

@@ -94,7 +94,12 @@ async fn index_over_small_shard() {
         .await
         .expect("frontier");
     let n = driver
-        .peek_count(PeekTarget::Index { id: index_id }, desc, Timestamp::from(0))
+        .peek_count(
+            PeekTarget::Index { id: index_id },
+            desc,
+            Timestamp::from(0),
+            None,
+        )
         .await
         .expect("peek");
     assert_eq!(n, 10_000);

src/compute-types/src/dyncfgs.rs

@@ -236,6 +236,28 @@ pub const LINEAR_JOIN_YIELDING: Config<&str> = Config::new(
      work, respectively, rather than falling back to some default.",
 );
 
+/// The yielding behavior with which a single index peek should be processed.
+pub const PEEK_YIELDING: Config<&str> = Config::new(
+    "peek_yielding",
+    "work:100000,time:10",
+    "How much work a compute worker may spend on one index peek before moving on to the \
+     next pending peek. Either 'work:<cursor steps>' or 'time:<milliseconds>' or \
+     'work:<cursor steps>,time:<milliseconds>'. Note that omitting one of 'work' or 'time' \
+     will entirely disable peek yielding by time or work, respectively, rather than falling \
+     back to some default.",
+);
+
+/// The yielding behavior with which index peeks as a whole should be processed.
+pub const PEEK_YIELDING_TOTAL: Config<&str> = Config::new(
+    "peek_yielding_total",
+    "work:1000000,time:100",
+    "How much work a compute worker may spend on index peeks in one worker activation, \
+     across all of them. Same format as 'peek_yielding'. Each pending peek gets at most one \
+     'peek_yielding' turn per activation, so this only binds once more than \
+     'peek_yielding_total' / 'peek_yielding' peeks are scanning at once. Peeks that don't get \
+     a turn are served first on the next activation.",
+);
+
 /// Enable lgalloc.
 pub const ENABLE_LGALLOC: Config<bool> =
     Config::new("enable_lgalloc", true, "Enable lgalloc.").scoped(ParameterScope::Replica);
@@ -556,6 +578,8 @@ pub fn all_dyncfgs(configs: ConfigSet) -> ConfigSet {
         .add(&ENABLE_COMPUTE_TEMPORAL_BUCKETING)
         .add(&TEMPORAL_BUCKETING_SUMMARY)
         .add(&LINEAR_JOIN_YIELDING)
+        .add(&PEEK_YIELDING)
+        .add(&PEEK_YIELDING_TOTAL)
         .add(&ENABLE_LGALLOC)
         .add(&LGALLOC_BACKGROUND_INTERVAL)
         .add(&LGALLOC_FILE_GROWTH_DAMPENER)