Skip to main content

mz_adapter/coord/
peek.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! Logic and types for creating, executing, and tracking peeks.
11//!
12//! This module determines if a dataflow can be short-cut, by returning constant values
13//! or by reading out of existing arrangements, and implements the appropriate plan.
14
15use std::collections::{BTreeMap, BTreeSet};
16use std::fmt;
17use std::num::NonZeroUsize;
18use std::ops::Deref;
19use std::sync::Arc;
20
21use differential_dataflow::consolidation::consolidate;
22use itertools::Itertools;
23use mz_adapter_types::connection::ConnectionId;
24use mz_cluster_client::ReplicaId;
25use mz_compute_client::controller::PeekNotification;
26use mz_compute_client::protocol::command::PeekTarget;
27use mz_compute_client::protocol::response::PeekResponse;
28use mz_compute_types::ComputeInstanceId;
29use mz_compute_types::dataflows::{DataflowDescription, IndexImport};
30use mz_controller_types::ClusterId;
31use mz_expr::explain::{HumanizedExplain, HumanizerMode, fmt_text_constant_rows};
32use mz_expr::row::RowCollection;
33use mz_expr::{
34    EvalError, Id, MirRelationExpr, MirScalarExpr, OptimizedMirRelationExpr, RowSetFinishing,
35    RowSetFinishingIncremental, permutation_for_arrangement,
36};
37use mz_ore::cast::CastFrom;
38use mz_ore::collections::CollectionExt;
39use mz_ore::soft_assert_eq_or_log;
40use mz_ore::str::{StrExt, separated};
41use mz_ore::task;
42use mz_ore::tracing::OpenTelemetryContext;
43use mz_persist_client::Schemas;
44use mz_persist_types::codec_impls::UnitSchema;
45use mz_repr::explain::text::DisplayText;
46use mz_repr::explain::{CompactScalars, IndexUsageType, PlanRenderingContext, UsedIndexes};
47use mz_repr::{
48    Diff, GlobalId, IntoRowIterator, RelationDesc, Row, RowIterator, SqlRelationType,
49    preserves_order,
50};
51use mz_storage_types::sources::SourceData;
52use serde::{Deserialize, Serialize};
53use timely::progress::Antichain;
54use tokio::sync::oneshot;
55use tracing::{Instrument, Span};
56use uuid::Uuid;
57
58use crate::active_compute_sink::{ActiveComputeSink, ActiveCopyTo};
59use crate::coord::timestamp_selection::TimestampDetermination;
60use crate::optimize::OptimizerError;
61use crate::statement_logging::WatchSetCreation;
62use crate::statement_logging::{StatementEndedExecutionReason, StatementExecutionStrategy};
63use crate::{AdapterError, ExecuteContextGuard, ExecuteResponse};
64
65/// A peek is a request to read data from a maintained arrangement.
66#[derive(Debug)]
67pub(crate) struct PendingPeek {
68    /// The connection that initiated the peek.
69    pub(crate) conn_id: ConnectionId,
70    /// The cluster that the peek is being executed on.
71    pub(crate) cluster_id: ClusterId,
72    /// All `GlobalId`s that the peek depend on.
73    pub(crate) depends_on: BTreeSet<GlobalId>,
74    /// Context about the execute that produced this peek,
75    /// needed by the coordinator for retiring it.
76    pub(crate) ctx_extra: ExecuteContextGuard,
77    /// Is this a fast-path peek, i.e. one that doesn't require a dataflow?
78    pub(crate) is_fast_path: bool,
79}
80
81/// The response from a `Peek`, with row multiplicities represented in unary.
82///
83/// Note that each `Peek` expects to generate exactly one `PeekResponse`, i.e.
84/// we expect a 1:1 contract between `Peek` and `PeekResponseUnary`.
85#[derive(Debug)]
86pub enum PeekResponseUnary {
87    Rows(Box<dyn RowIterator + Send + Sync>),
88    Error(String),
89    Canceled,
90    /// A dependency was dropped during execution.
91    ///
92    /// N.B. This is a bit of a workaround for the fact that our Error variant
93    /// is unstructured and right now we specifically care about this error and
94    /// need to render differently based on context.
95    DependencyDropped(DroppedDependency),
96}
97
98/// A dependency that was dropped while a peek or subscribe was in flight.
99///
100/// The `name` fields hold the bare name (e.g. `db.schema.t` or `c`); `Display`
101/// applies SQL identifier quoting to produce `relation "db.schema.t"` or
102/// `cluster "c"` for direct use in error wording.
103#[derive(Clone, Debug)]
104pub enum DroppedDependency {
105    Relation { name: String },
106    Cluster { name: String },
107}
108
109impl fmt::Display for DroppedDependency {
110    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
111        match self {
112            Self::Relation { name } => write!(f, "relation {}", name.quoted()),
113            Self::Cluster { name } => write!(f, "cluster {}", name.quoted()),
114        }
115    }
116}
117
118impl DroppedDependency {
119    /// User-facing error for a query (peek or subscribe) that could not finish
120    /// because this dependency was dropped mid-flight.
121    pub fn query_terminated_error(&self) -> String {
122        format!("query could not complete because {self} was dropped")
123    }
124
125    /// Convert this dropped dependency into an [`AdapterError::ConcurrentDependencyDrop`].
126    pub fn to_concurrent_dependency_drop(&self) -> AdapterError {
127        let (kind, name) = match self {
128            Self::Relation { name } => ("relation", name.clone()),
129            Self::Cluster { name } => ("cluster", name.clone()),
130        };
131        AdapterError::ConcurrentDependencyDrop {
132            dependency_kind: kind,
133            dependency_id: name,
134        }
135    }
136}
137
138#[derive(Clone, Debug)]
139pub struct PeekDataflowPlan {
140    pub(crate) desc: DataflowDescription<mz_compute_types::plan::LirRelationExpr, ()>,
141    pub(crate) id: GlobalId,
142    key: Vec<MirScalarExpr>,
143    permutation: Vec<usize>,
144    thinned_arity: usize,
145}
146
147impl PeekDataflowPlan {
148    pub fn new(
149        desc: DataflowDescription<mz_compute_types::plan::LirRelationExpr, ()>,
150        id: GlobalId,
151        typ: &SqlRelationType,
152    ) -> Self {
153        let arity = typ.arity();
154        let key = typ
155            .default_key()
156            .into_iter()
157            .map(MirScalarExpr::column)
158            .collect::<Vec<_>>();
159        let (permutation, thinning) = permutation_for_arrangement(&key, arity);
160        Self {
161            desc,
162            id,
163            key,
164            permutation,
165            thinned_arity: thinning.len(),
166        }
167    }
168}
169
170#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq, Ord, PartialOrd)]
171pub enum FastPathPlan {
172    /// The view evaluates to a constant result that can be returned.
173    ///
174    /// The [SqlRelationType] is unnecessary for evaluating the constant result but
175    /// may be helpful when printing out an explanation.
176    Constant(Result<Vec<(Row, Diff)>, EvalError>, SqlRelationType),
177    /// The view can be read out of an existing arrangement.
178    /// (coll_id, idx_id, values to look up, mfp to apply)
179    PeekExisting(GlobalId, GlobalId, Option<Vec<Row>>, mz_expr::SafeMfpPlan),
180    /// The view can be read directly out of Persist.
181    PeekPersist(GlobalId, Option<Row>, mz_expr::SafeMfpPlan),
182}
183
184impl<'a, T: 'a> DisplayText<PlanRenderingContext<'a, T>> for FastPathPlan {
185    fn fmt_text(
186        &self,
187        f: &mut fmt::Formatter<'_>,
188        ctx: &mut PlanRenderingContext<'a, T>,
189    ) -> fmt::Result {
190        if ctx.config.verbose_syntax {
191            self.fmt_verbose_text(f, ctx)
192        } else {
193            self.fmt_default_text(f, ctx)
194        }
195    }
196}
197
198impl FastPathPlan {
199    pub fn fmt_default_text<'a, T>(
200        &self,
201        f: &mut fmt::Formatter<'_>,
202        ctx: &mut PlanRenderingContext<'a, T>,
203    ) -> fmt::Result {
204        let mode = HumanizedExplain::new(ctx.config.redacted);
205
206        match self {
207            FastPathPlan::Constant(rows, _) => {
208                write!(f, "{}→Constant ", ctx.indent)?;
209
210                match rows {
211                    Ok(rows) => writeln!(f, "({} rows)", rows.len())?,
212                    Err(err) => {
213                        if mode.redacted() {
214                            writeln!(f, "(error: █)")?;
215                        } else {
216                            writeln!(f, "(error: {})", err.to_string().quoted(),)?;
217                        }
218                    }
219                }
220            }
221            FastPathPlan::PeekExisting(coll_id, idx_id, literal_constraints, mfp) => {
222                let coll = ctx
223                    .humanizer
224                    .humanize_id(*coll_id)
225                    .unwrap_or_else(|| coll_id.to_string());
226                let idx = ctx
227                    .humanizer
228                    .humanize_id(*idx_id)
229                    .unwrap_or_else(|| idx_id.to_string());
230                writeln!(f, "{}→Map/Filter/Project", ctx.indent)?;
231                ctx.indent.set();
232
233                ctx.indent += 1;
234
235                mode.expr(mfp.deref(), None).fmt_default_text(f, ctx)?;
236                let printed = !mfp.expressions.is_empty() || !mfp.predicates.is_empty();
237
238                if printed {
239                    ctx.indent += 1;
240                }
241                if let Some(literal_constraints) = literal_constraints {
242                    writeln!(f, "{}→Index Lookup on {coll} (using {idx})", ctx.indent)?;
243                    ctx.indent += 1;
244                    let values = separated("; ", mode.seq(literal_constraints, None));
245                    writeln!(f, "{}Lookup values: {values}", ctx.indent)?;
246                } else {
247                    writeln!(f, "{}→Indexed {coll} (using {idx})", ctx.indent)?;
248                }
249
250                ctx.indent.reset();
251            }
252            FastPathPlan::PeekPersist(global_id, literal_constraint, mfp) => {
253                let coll = ctx
254                    .humanizer
255                    .humanize_id(*global_id)
256                    .unwrap_or_else(|| global_id.to_string());
257                writeln!(f, "{}→Map/Filter/Project", ctx.indent)?;
258                ctx.indent.set();
259
260                ctx.indent += 1;
261
262                mode.expr(mfp.deref(), None).fmt_default_text(f, ctx)?;
263                let printed = !mfp.expressions.is_empty() || !mfp.predicates.is_empty();
264
265                if printed {
266                    ctx.indent += 1;
267                }
268                if let Some(literal_constraint) = literal_constraint {
269                    writeln!(f, "{}→ReadStorage Lookup on {coll}", ctx.indent)?;
270                    ctx.indent += 1;
271                    let value = mode.expr(literal_constraint, None);
272                    writeln!(f, "{}Lookup value: {value}", ctx.indent)?;
273                } else {
274                    writeln!(f, "{}→ReadStorage {coll}", ctx.indent)?;
275                }
276
277                ctx.indent.reset();
278            }
279        }
280
281        Ok(())
282    }
283
284    pub fn fmt_verbose_text<'a, T>(
285        &self,
286        f: &mut fmt::Formatter<'_>,
287        ctx: &mut PlanRenderingContext<'a, T>,
288    ) -> fmt::Result {
289        let redacted = ctx.config.redacted;
290        let mode = HumanizedExplain::new(redacted);
291
292        // TODO(aalexandrov): factor out common PeekExisting and PeekPersist
293        // code.
294        match self {
295            FastPathPlan::Constant(Ok(rows), _) => {
296                if !rows.is_empty() {
297                    writeln!(f, "{}Constant", ctx.indent)?;
298                    *ctx.as_mut() += 1;
299                    fmt_text_constant_rows(
300                        f,
301                        rows.iter().map(|(row, diff)| (row, diff)),
302                        ctx.as_mut(),
303                        redacted,
304                    )?;
305                    *ctx.as_mut() -= 1;
306                } else {
307                    writeln!(f, "{}Constant <empty>", ctx.as_mut())?;
308                }
309                Ok(())
310            }
311            FastPathPlan::Constant(Err(err), _) => {
312                if redacted {
313                    writeln!(f, "{}Error █", ctx.as_mut())
314                } else {
315                    writeln!(f, "{}Error {}", ctx.as_mut(), err.to_string().escaped())
316                }
317            }
318            FastPathPlan::PeekExisting(coll_id, idx_id, literal_constraints, mfp) => {
319                ctx.as_mut().set();
320                let (map, filter, project) = mfp.as_map_filter_project();
321
322                let cols = if !ctx.config.humanized_exprs {
323                    None
324                } else if let Some(cols) = ctx.humanizer.column_names_for_id(*idx_id) {
325                    // FIXME: account for thinning and permutation
326                    // See mz_expr::permutation_for_arrangement
327                    // See permute_oneshot_mfp_around_index
328                    let cols = itertools::chain(
329                        cols.iter().cloned(),
330                        std::iter::repeat(String::new()).take(map.len()),
331                    )
332                    .collect();
333                    Some(cols)
334                } else {
335                    None
336                };
337
338                if project.len() != mfp.input_arity + map.len()
339                    || !project.iter().enumerate().all(|(i, o)| i == *o)
340                {
341                    let outputs = mode.seq(&project, cols.as_ref());
342                    let outputs = CompactScalars(outputs);
343                    writeln!(f, "{}Project ({})", ctx.as_mut(), outputs)?;
344                    *ctx.as_mut() += 1;
345                }
346                if !filter.is_empty() {
347                    let predicates = separated(" AND ", mode.seq(&filter, cols.as_ref()));
348                    writeln!(f, "{}Filter {}", ctx.as_mut(), predicates)?;
349                    *ctx.as_mut() += 1;
350                }
351                if !map.is_empty() {
352                    let scalars = mode.seq(&map, cols.as_ref());
353                    let scalars = CompactScalars(scalars);
354                    writeln!(f, "{}Map ({})", ctx.as_mut(), scalars)?;
355                    *ctx.as_mut() += 1;
356                }
357                MirRelationExpr::fmt_indexed_filter(
358                    f,
359                    ctx,
360                    coll_id,
361                    idx_id,
362                    literal_constraints.clone(),
363                    None,
364                )?;
365                writeln!(f)?;
366                ctx.as_mut().reset();
367                Ok(())
368            }
369            FastPathPlan::PeekPersist(gid, literal_constraint, mfp) => {
370                ctx.as_mut().set();
371                let (map, filter, project) = mfp.as_map_filter_project();
372
373                let cols = if !ctx.config.humanized_exprs {
374                    None
375                } else if let Some(cols) = ctx.humanizer.column_names_for_id(*gid) {
376                    let cols = itertools::chain(
377                        cols.iter().cloned(),
378                        std::iter::repeat(String::new()).take(map.len()),
379                    )
380                    .collect::<Vec<_>>();
381                    Some(cols)
382                } else {
383                    None
384                };
385
386                if project.len() != mfp.input_arity + map.len()
387                    || !project.iter().enumerate().all(|(i, o)| i == *o)
388                {
389                    let outputs = mode.seq(&project, cols.as_ref());
390                    let outputs = CompactScalars(outputs);
391                    writeln!(f, "{}Project ({})", ctx.as_mut(), outputs)?;
392                    *ctx.as_mut() += 1;
393                }
394                if !filter.is_empty() {
395                    let predicates = separated(" AND ", mode.seq(&filter, cols.as_ref()));
396                    writeln!(f, "{}Filter {}", ctx.as_mut(), predicates)?;
397                    *ctx.as_mut() += 1;
398                }
399                if !map.is_empty() {
400                    let scalars = mode.seq(&map, cols.as_ref());
401                    let scalars = CompactScalars(scalars);
402                    writeln!(f, "{}Map ({})", ctx.as_mut(), scalars)?;
403                    *ctx.as_mut() += 1;
404                }
405                let human_id = ctx
406                    .humanizer
407                    .humanize_id(*gid)
408                    .unwrap_or_else(|| gid.to_string());
409                write!(f, "{}PeekPersist {human_id}", ctx.as_mut())?;
410                if let Some(literal) = literal_constraint {
411                    let value = mode.expr(literal, None);
412                    writeln!(f, " [value={}]", value)?;
413                } else {
414                    writeln!(f, "")?;
415                }
416                ctx.as_mut().reset();
417                Ok(())
418            }
419        }?;
420        Ok(())
421    }
422}
423
424#[derive(Debug)]
425pub struct PlannedPeek {
426    pub plan: PeekPlan,
427    pub determination: TimestampDetermination,
428    pub conn_id: ConnectionId,
429    /// The result type _after_ reading out of the "source" and applying any
430    /// [MapFilterProject](mz_expr::MapFilterProject), but _before_ applying a
431    /// [RowSetFinishing].
432    ///
433    /// This is _the_ `result_type` as far as compute is concerned and further
434    /// changes through projections happen purely in the adapter.
435    pub intermediate_result_type: SqlRelationType,
436    pub source_arity: usize,
437    pub source_ids: BTreeSet<GlobalId>,
438}
439
440/// Possible ways in which the coordinator could produce the result for a goal view.
441#[derive(Clone, Debug)]
442pub enum PeekPlan {
443    FastPath(FastPathPlan),
444    /// The view must be installed as a dataflow and then read.
445    SlowPath(PeekDataflowPlan),
446}
447
448/// Convert `mfp` to an executable, non-temporal plan.
449/// It should be non-temporal, as OneShot preparation populates `mz_now`.
450///
451/// If the `mfp` can't be converted into a non-temporal plan, this returns an _internal_ error.
452fn mfp_to_safe_plan(
453    mfp: mz_expr::MapFilterProject,
454) -> Result<mz_expr::SafeMfpPlan, OptimizerError> {
455    mfp.into_plan()
456        .map_err(OptimizerError::InternalUnsafeMfpPlan)?
457        .into_nontemporal()
458        .map_err(|e| OptimizerError::InternalUnsafeMfpPlan(format!("{:?}", e)))
459}
460
461/// If it can't convert `mfp` into a `SafeMfpPlan`, this returns an _internal_ error.
462fn permute_oneshot_mfp_around_index(
463    mfp: mz_expr::MapFilterProject,
464    key: &[MirScalarExpr],
465) -> Result<mz_expr::SafeMfpPlan, OptimizerError> {
466    let input_arity = mfp.input_arity;
467    let mut safe_mfp = mfp_to_safe_plan(mfp)?;
468    let (permute, thinning) = permutation_for_arrangement(key, input_arity);
469    safe_mfp.permute_fn(|c| permute[c], key.len() + thinning.len());
470    Ok(safe_mfp)
471}
472
473/// Determine if the dataflow plan can be implemented without an actual dataflow.
474///
475/// If the optimized plan is a `Constant` or a `Get` of a maintained arrangement,
476/// we can avoid building a dataflow (and either just return the results, or peek
477/// out of the arrangement, respectively).
478pub fn create_fast_path_plan(
479    dataflow_plan: &mut DataflowDescription<OptimizedMirRelationExpr>,
480    view_id: GlobalId,
481    finishing: Option<&RowSetFinishing>,
482    persist_fast_path_limit: usize,
483    persist_fast_path_order: bool,
484) -> Result<Option<FastPathPlan>, OptimizerError> {
485    // At this point, `dataflow_plan` contains our best optimized dataflow.
486    // We will check the plan to see if there is a fast path to escape full dataflow construction.
487
488    // We need to restrict ourselves to settings where the inserted transient view is the first thing
489    // to build (no dependent views). There is likely an index to build as well, but we may not be sure.
490    if dataflow_plan.objects_to_build.len() >= 1 && dataflow_plan.objects_to_build[0].id == view_id
491    {
492        let mut mir = &*dataflow_plan.objects_to_build[0].plan.as_inner_mut();
493        if let Some((rows, found_typ)) = mir.as_const() {
494            // In the case of a constant, we can return the result now.
495            let plan = FastPathPlan::Constant(
496                rows.clone(),
497                mz_repr::SqlRelationType::from_repr(found_typ),
498            );
499            return Ok(Some(plan));
500        } else {
501            // If there is a TopK that would be completely covered by the finishing, then jump
502            // through the TopK.
503            if let MirRelationExpr::TopK {
504                input,
505                group_key,
506                order_key,
507                limit,
508                offset,
509                monotonic: _,
510                expected_group_size: _,
511            } = mir
512            {
513                if let Some(finishing) = finishing {
514                    if group_key.is_empty() && *order_key == finishing.order_by && *offset == 0 {
515                        // The following is roughly `limit >= finishing.limit + finishing.offset`,
516                        // but with Options.
517                        let finishing_limits_at_least_as_topk = match (limit, finishing.limit) {
518                            (None, _) => true,
519                            (Some(..), None) => false,
520                            (Some(topk_limit), Some(finishing_limit)) => {
521                                if let Some(l) = topk_limit.as_literal_int64() {
522                                    i128::cast_from(l)
523                                        >= i128::cast_from(*finishing_limit)
524                                            + i128::cast_from(finishing.offset)
525                                } else {
526                                    false
527                                }
528                            }
529                        };
530                        if finishing_limits_at_least_as_topk {
531                            mir = input;
532                        }
533                    }
534                }
535            }
536            // In the case of a linear operator around an indexed view, we
537            // can skip creating a dataflow and instead pull all the rows in
538            // index and apply the linear operator against them.
539            let (mfp, mir) = mz_expr::MapFilterProject::extract_from_expression(mir);
540            match mir {
541                MirRelationExpr::Get {
542                    id: Id::Global(get_id),
543                    typ: repr_typ,
544                    ..
545                } => {
546                    // Just grab any arrangement if an arrangement exists
547                    for (index_id, IndexImport { desc, .. }) in dataflow_plan.index_imports.iter() {
548                        if desc.on_id == *get_id {
549                            return Ok(Some(FastPathPlan::PeekExisting(
550                                *get_id,
551                                *index_id,
552                                None,
553                                permute_oneshot_mfp_around_index(mfp, &desc.key)?,
554                            )));
555                        }
556                    }
557
558                    // If there is no arrangement, consider peeking the persist shard directly.
559                    // Generally, we consider a persist peek when the query can definitely be satisfied
560                    // by scanning through a small, constant number of Persist key-values.
561                    let safe_mfp = mfp_to_safe_plan(mfp)?;
562                    let (_maps, filters, projection) = safe_mfp.as_map_filter_project();
563
564                    let persist_fast_path_order_relation_typ = if persist_fast_path_order {
565                        Some(
566                            dataflow_plan
567                                .source_imports
568                                .get(get_id)
569                                .expect("Get's ID is also imported")
570                                .desc
571                                .typ
572                                .clone(),
573                        )
574                    } else {
575                        None
576                    };
577
578                    let literal_constraint =
579                        if let Some(relation_typ) = &persist_fast_path_order_relation_typ {
580                            let mut row = Row::default();
581                            let mut packer = row.packer();
582                            for (idx, col) in relation_typ.column_types.iter().enumerate() {
583                                if !preserves_order(&col.scalar_type) {
584                                    break;
585                                }
586                                let col_expr = MirScalarExpr::column(idx);
587
588                                let Some((literal, _)) = filters
589                                    .iter()
590                                    .filter_map(|f| f.expr_eq_literal(&col_expr))
591                                    .next()
592                                else {
593                                    break;
594                                };
595                                packer.extend_by_row(&literal);
596                            }
597                            if row.is_empty() { None } else { Some(row) }
598                        } else {
599                            None
600                        };
601
602                    let finish_ok = match &finishing {
603                        None => false,
604                        Some(RowSetFinishing {
605                            order_by,
606                            limit,
607                            offset,
608                            ..
609                        }) => {
610                            let order_ok =
611                                if let Some(relation_typ) = &persist_fast_path_order_relation_typ {
612                                    order_by.iter().enumerate().all(|(idx, order)| {
613                                        // Map the ordering column back to the column in the source data.
614                                        // (If it's not one of the input columns, we can't make any guarantees.)
615                                        let column_idx = projection[order.column];
616                                        if column_idx >= safe_mfp.input_arity {
617                                            return false;
618                                        }
619                                        let column_type = &relation_typ.column_types[column_idx];
620                                        let index_ok = idx == column_idx;
621                                        let nulls_ok = !column_type.nullable || order.nulls_last;
622                                        let asc_ok = !order.desc;
623                                        let type_ok = preserves_order(&column_type.scalar_type);
624                                        index_ok && nulls_ok && asc_ok && type_ok
625                                    })
626                                } else {
627                                    order_by.is_empty()
628                                };
629                            let limit_ok = limit.map_or(false, |l| {
630                                usize::cast_from(l) + *offset < persist_fast_path_limit
631                            });
632                            order_ok && limit_ok
633                        }
634                    };
635
636                    let key_constraint = if let Some(literal) = &literal_constraint {
637                        let prefix_len = literal.iter().count();
638                        repr_typ
639                            .keys
640                            .iter()
641                            .any(|k| k.iter().all(|idx| *idx < prefix_len))
642                    } else {
643                        false
644                    };
645
646                    // We can generate a persist peek when:
647                    // - We have a literal constraint that includes an entire key (so we'll return at most one value)
648                    // - We can return the first N key values (no filters, small limit, consistent order)
649                    if key_constraint || (filters.is_empty() && finish_ok) {
650                        return Ok(Some(FastPathPlan::PeekPersist(
651                            *get_id,
652                            literal_constraint,
653                            safe_mfp,
654                        )));
655                    }
656                }
657                MirRelationExpr::Join { implementation, .. } => {
658                    if let mz_expr::JoinImplementation::IndexedFilter(coll_id, idx_id, key, vals) =
659                        implementation
660                    {
661                        return Ok(Some(FastPathPlan::PeekExisting(
662                            *coll_id,
663                            *idx_id,
664                            Some(vals.clone()),
665                            permute_oneshot_mfp_around_index(mfp, key)?,
666                        )));
667                    }
668                }
669                // nothing can be done for non-trivial expressions.
670                _ => {}
671            }
672        }
673    }
674    Ok(None)
675}
676
677impl FastPathPlan {
678    pub fn used_indexes(&self, finishing: Option<&RowSetFinishing>) -> UsedIndexes {
679        match self {
680            FastPathPlan::Constant(..) => UsedIndexes::default(),
681            FastPathPlan::PeekExisting(_coll_id, idx_id, literal_constraints, _mfp) => {
682                if literal_constraints.is_some() {
683                    UsedIndexes::new([(*idx_id, vec![IndexUsageType::Lookup(*idx_id)])].into())
684                } else if finishing.map_or(false, |f| f.limit.is_some() && f.order_by.is_empty()) {
685                    UsedIndexes::new([(*idx_id, vec![IndexUsageType::FastPathLimit])].into())
686                } else {
687                    UsedIndexes::new([(*idx_id, vec![IndexUsageType::FullScan])].into())
688                }
689            }
690            FastPathPlan::PeekPersist(..) => UsedIndexes::default(),
691        }
692    }
693}
694
695impl crate::coord::Coordinator {
696    /// Implements a peek plan produced by `create_plan` above.
697    ///
698    /// On success this takes the contents of `ctx_extra`, the
699    /// statement-logging guard: a constant peek is retired immediately, and a
700    /// streaming peek moves them into `pending_peeks` for
701    /// `handle_peek_notification` to retire. On an error return the contents
702    /// are left intact and the caller must take ownership of them: `retire`
703    /// if the caller is the sole end-logger, `defuse` if something else logs
704    /// the error end. Dropping the guard armed emits a spurious `Aborted`,
705    /// double-ending the statement.
706    #[mz_ore::instrument(level = "debug")]
707    pub async fn implement_peek_plan(
708        &mut self,
709        ctx_extra: &mut ExecuteContextGuard,
710        plan: PlannedPeek,
711        finishing: RowSetFinishing,
712        compute_instance: ComputeInstanceId,
713        target_replica: Option<ReplicaId>,
714        max_result_size: u64,
715        max_returned_query_size: Option<u64>,
716    ) -> Result<ExecuteResponse, AdapterError> {
717        let PlannedPeek {
718            plan: fast_path,
719            determination,
720            conn_id,
721            intermediate_result_type,
722            source_arity,
723            source_ids,
724        } = plan;
725
726        // If the dataflow optimizes to a constant expression, we can immediately return the result.
727        if let PeekPlan::FastPath(FastPathPlan::Constant(rows, _)) = fast_path {
728            let mut rows = match rows {
729                Ok(rows) => rows,
730                Err(e) => return Err(e.into()),
731            };
732            // Consolidate down the results to get correct totals.
733            consolidate(&mut rows);
734
735            let mut results = Vec::new();
736            for (row, count) in rows {
737                if count.is_negative() {
738                    Err(EvalError::InvalidParameterValue(
739                        format!("Negative multiplicity in constant result: {}", count).into(),
740                    ))?
741                };
742                if count.is_positive() {
743                    let count = usize::cast_from(
744                        u64::try_from(count.into_inner())
745                            .expect("known to be positive from check above"),
746                    );
747                    results.push((
748                        row,
749                        NonZeroUsize::new(count).expect("known to be non-zero from check above"),
750                    ));
751                }
752            }
753            let row_collection = RowCollection::new(results, &finishing.order_by);
754            let duration_histogram = self.metrics.row_set_finishing_seconds();
755
756            let (ret, reason) = match finishing.finish(
757                row_collection,
758                max_result_size,
759                max_returned_query_size,
760                &duration_histogram,
761            ) {
762                Ok((rows, row_size_bytes)) => {
763                    let result_size = u64::cast_from(row_size_bytes);
764                    let rows_returned = u64::cast_from(rows.count());
765                    (
766                        Ok(Self::send_immediate_rows(rows)),
767                        StatementEndedExecutionReason::Success {
768                            result_size: Some(result_size),
769                            rows_returned: Some(rows_returned),
770                            execution_strategy: Some(StatementExecutionStrategy::Constant),
771                        },
772                    )
773                }
774                Err(error) => (
775                    Err(AdapterError::ResultSize(error.clone())),
776                    StatementEndedExecutionReason::Errored { error },
777                ),
778            };
779            self.retire_execution(reason, std::mem::take(ctx_extra).defuse());
780            return ret;
781        }
782
783        let timestamp = determination.timestamp_context.timestamp_or_default();
784        if let Some(id) = ctx_extra.contents() {
785            self.set_statement_execution_timestamp(id, timestamp)
786        }
787
788        // The remaining cases are a peek into a maintained arrangement, or building a dataflow.
789        // In both cases we will want to peek, and the main difference is that we might want to
790        // build a dataflow and drop it once the peek is issued. The peeks are also constructed
791        // differently.
792
793        // Acquire a read hold for the peek target so its `since` cannot advance past
794        // `timestamp` before `compute.peek()` runs. On the slow path we ship the dataflow
795        // first: the implied hold from `create_dataflow` pins the new collection's `since`
796        // at `as_of`, so the subsequent `acquire_read_hold` lands at `as_of <= timestamp`.
797        let (peek_command, drop_dataflow, is_fast_path, peek_target, strategy, read_hold) =
798            match fast_path {
799                PeekPlan::FastPath(FastPathPlan::PeekExisting(
800                    _coll_id,
801                    idx_id,
802                    literal_constraints,
803                    map_filter_project,
804                )) => {
805                    let read_hold = self
806                        .controller
807                        .compute
808                        .acquire_read_hold(compute_instance, idx_id)
809                        .map_err(
810                            AdapterError::concurrent_dependency_drop_from_collection_update_error,
811                        )?;
812                    (
813                        (literal_constraints, timestamp, map_filter_project),
814                        None,
815                        true,
816                        PeekTarget::Index { id: idx_id },
817                        StatementExecutionStrategy::FastPath,
818                        read_hold,
819                    )
820                }
821                PeekPlan::FastPath(FastPathPlan::PeekPersist(
822                    coll_id,
823                    literal_constraint,
824                    map_filter_project,
825                )) => {
826                    let peek_command = (
827                        literal_constraint.map(|r| vec![r]),
828                        timestamp,
829                        map_filter_project,
830                    );
831                    let metadata = self
832                        .controller
833                        .storage
834                        .collection_metadata(coll_id)
835                        .expect("storage collection for fast-path peek")
836                        .clone();
837                    let read_hold = self
838                        .controller
839                        .storage_collections
840                        .acquire_read_holds(vec![coll_id])
841                        .map_err(AdapterError::concurrent_dependency_drop_from_collection_missing)?
842                        .into_element();
843                    (
844                        peek_command,
845                        None,
846                        true,
847                        PeekTarget::Persist {
848                            id: coll_id,
849                            metadata,
850                        },
851                        StatementExecutionStrategy::PersistFastPath,
852                        read_hold,
853                    )
854                }
855                PeekPlan::SlowPath(PeekDataflowPlan {
856                    desc: dataflow,
857                    // n.b. this index_id identifies a transient index the
858                    // caller created, so it is guaranteed to be on
859                    // `compute_instance`.
860                    id: index_id,
861                    key: index_key,
862                    permutation: index_permutation,
863                    thinned_arity: index_thinned_arity,
864                }) => {
865                    // The slow-path peek read-hold strategy below acquires a hold for
866                    // `index_id` only. That is sufficient today because slow-path peek
867                    // dataflows have a single export equal to `index_id`. If we ever
868                    // ship multi-output dataflows on this path, the hold acquisition
869                    // needs to be revisited.
870                    let exports: Vec<GlobalId> = dataflow.export_ids().collect();
871                    soft_assert_eq_or_log!(
872                        exports.as_slice(),
873                        &[index_id],
874                        "slow-path peek dataflow must export exactly [index_id]",
875                    );
876                    if exports.as_slice() != [index_id] {
877                        return Err(AdapterError::internal(
878                            "peek error",
879                            format!(
880                                "slow-path peek dataflow exports {exports:?}, expected [{index_id}]",
881                            ),
882                        ));
883                    }
884
885                    // Very important: actually create the dataflow (here, so we can destructure).
886                    self.controller
887                        .compute
888                        .create_dataflow(compute_instance, dataflow, None)
889                        .map_err(
890                            AdapterError::concurrent_dependency_drop_from_dataflow_creation_error,
891                        )?;
892
893                    // Acquire a bare hold on the freshly-shipped index. On failure we must
894                    // drop the dataflow ourselves, otherwise it leaks.
895                    let acquire_result = self
896                        .controller
897                        .compute
898                        .acquire_read_hold(compute_instance, index_id)
899                        .map_err(
900                            AdapterError::concurrent_dependency_drop_from_collection_update_error,
901                        );
902                    let read_hold = match acquire_result {
903                        Ok(hold) => hold,
904                        Err(e) => {
905                            self.drop_compute_collections(vec![(compute_instance, index_id)]);
906                            return Err(e);
907                        }
908                    };
909
910                    // Create an identity MFP operator.
911                    let mut map_filter_project = mz_expr::MapFilterProject::new(source_arity);
912                    map_filter_project.permute_fn(
913                        |c| index_permutation[c],
914                        index_key.len() + index_thinned_arity,
915                    );
916                    let map_filter_project = mfp_to_safe_plan(map_filter_project)?;
917
918                    (
919                        (None, timestamp, map_filter_project),
920                        Some(index_id),
921                        false,
922                        PeekTarget::Index { id: index_id },
923                        StatementExecutionStrategy::Standard,
924                        read_hold,
925                    )
926                }
927                PeekPlan::FastPath(_) => {
928                    unreachable!()
929                }
930            };
931
932        // Endpoints for sending and receiving peek responses.
933        let (rows_tx, rows_rx) = tokio::sync::oneshot::channel();
934
935        // Generate unique UUID. Guaranteed to be unique to all pending peeks, there's an very
936        // small but unlikely chance that it's not unique to completed peeks.
937        let mut uuid = Uuid::new_v4();
938        while self.pending_peeks.contains_key(&uuid) {
939            uuid = Uuid::new_v4();
940        }
941
942        let (literal_constraints, timestamp, map_filter_project) = peek_command;
943
944        // At this stage we don't know column names for the result because we
945        // only know the peek's result type as a bare SqlRelationType.
946        let peek_result_column_names =
947            (0..intermediate_result_type.arity()).map(|i| format!("peek_{i}"));
948        let peek_result_desc =
949            RelationDesc::new(intermediate_result_type, peek_result_column_names);
950
951        let peek_result = self
952            .controller
953            .compute
954            .peek(
955                compute_instance,
956                peek_target,
957                literal_constraints,
958                uuid,
959                timestamp,
960                peek_result_desc,
961                finishing.clone(),
962                map_filter_project,
963                read_hold,
964                target_replica,
965                rows_tx,
966            )
967            .map_err(AdapterError::concurrent_dependency_drop_from_peek_error);
968        if let Err(e) = peek_result {
969            // If we shipped a transient dataflow above, drop it now to avoid leaking it.
970            if let Some(index_id) = drop_dataflow {
971                self.drop_compute_collections(vec![(compute_instance, index_id)]);
972            }
973            return Err(e);
974        }
975
976        // Register the pending peek only after compute.peek() succeeds. If it
977        // fails (e.g. concurrent replica/cluster drop), inserting first would
978        // leak entries in these maps and misattribute statement execution reasons.
979        self.pending_peeks.insert(
980            uuid,
981            PendingPeek {
982                conn_id: conn_id.clone(),
983                cluster_id: compute_instance,
984                depends_on: source_ids,
985                ctx_extra: std::mem::take(ctx_extra),
986                is_fast_path,
987            },
988        );
989        self.client_pending_peeks
990            .entry(conn_id)
991            .or_default()
992            .insert(uuid, compute_instance);
993
994        let duration_histogram = self.metrics.row_set_finishing_seconds();
995
996        // If a dataflow was created, drop it now that the peek is queued. This is
997        // required: `add_collection` installs implied/warmup holds owned by the
998        // controller and only released via `drop_collections`, so without this call
999        // the transient dataflow's `since` would stay pinned at `as_of` forever. The
1000        // peek's own read hold keeps the collection alive on the cluster until the
1001        // response arrives.
1002        if let Some(index_id) = drop_dataflow {
1003            self.drop_compute_collections(vec![(compute_instance, index_id)]);
1004        }
1005
1006        let persist_client = self.persist_client.clone();
1007        let peek_stash_read_batch_size_bytes =
1008            mz_compute_types::dyncfgs::PEEK_RESPONSE_STASH_READ_BATCH_SIZE_BYTES
1009                .get(self.catalog().system_config().dyncfgs());
1010        let peek_stash_read_memory_budget_bytes =
1011            mz_compute_types::dyncfgs::PEEK_RESPONSE_STASH_READ_MEMORY_BUDGET_BYTES
1012                .get(self.catalog().system_config().dyncfgs());
1013
1014        let peek_response_stream = Self::create_peek_response_stream(
1015            rows_rx,
1016            finishing,
1017            max_result_size,
1018            max_returned_query_size,
1019            duration_histogram,
1020            persist_client,
1021            peek_stash_read_batch_size_bytes,
1022            peek_stash_read_memory_budget_bytes,
1023        );
1024
1025        Ok(crate::ExecuteResponse::SendingRowsStreaming {
1026            rows: Box::pin(peek_response_stream),
1027            instance_id: compute_instance,
1028            strategy,
1029        })
1030    }
1031
1032    /// Creates an async stream that processes peek responses and yields rows.
1033    ///
1034    /// TODO(peek-seq): Move this out of `coord` once we delete the old peek sequencing.
1035    #[mz_ore::instrument(level = "debug")]
1036    pub(crate) fn create_peek_response_stream(
1037        rows_rx: tokio::sync::oneshot::Receiver<PeekResponse>,
1038        finishing: RowSetFinishing,
1039        max_result_size: u64,
1040        max_returned_query_size: Option<u64>,
1041        duration_histogram: prometheus::Histogram,
1042        mut persist_client: mz_persist_client::PersistClient,
1043        peek_stash_read_batch_size_bytes: usize,
1044        peek_stash_read_memory_budget_bytes: usize,
1045    ) -> impl futures::Stream<Item = PeekResponseUnary> {
1046        async_stream::stream!({
1047            let result = rows_rx.await;
1048
1049            let rows = match result {
1050                Ok(rows) => rows,
1051                Err(e) => {
1052                    yield PeekResponseUnary::Error(e.to_string());
1053                    return;
1054                }
1055            };
1056
1057            match rows {
1058                PeekResponse::Rows(rows) => {
1059                    let rows = RowCollection::merge_sorted(&rows, &finishing.order_by);
1060                    match finishing.finish(
1061                        rows,
1062                        max_result_size,
1063                        max_returned_query_size,
1064                        &duration_histogram,
1065                    ) {
1066                        Ok((rows, _size_bytes)) => yield PeekResponseUnary::Rows(Box::new(rows)),
1067                        Err(e) => yield PeekResponseUnary::Error(e),
1068                    }
1069                }
1070                PeekResponse::Stashed(response) => {
1071                    let response = *response;
1072
1073                    let shard_id = response.shard_id;
1074
1075                    let mut batches = Vec::new();
1076                    for proto_batch in response.batches.into_iter() {
1077                        let batch =
1078                            persist_client.batch_from_transmittable_batch(&shard_id, proto_batch);
1079
1080                        batches.push(batch);
1081                    }
1082                    tracing::trace!(?batches, "stashed peek response");
1083
1084                    let as_of = Antichain::from_elem(mz_repr::Timestamp::default());
1085                    let read_schemas: Schemas<SourceData, ()> = Schemas {
1086                        id: None,
1087                        key: Arc::new(response.relation_desc.clone()),
1088                        val: Arc::new(UnitSchema),
1089                    };
1090
1091                    let mut row_cursor = persist_client
1092                        .read_batches_consolidated::<_, _, _, i64>(
1093                            response.shard_id,
1094                            as_of,
1095                            read_schemas,
1096                            batches,
1097                            |_stats| true,
1098                            peek_stash_read_memory_budget_bytes,
1099                        )
1100                        .await
1101                        .expect("invalid usage");
1102
1103                    // NOTE: Using the cursor creates Futures that are not Sync,
1104                    // so we can't drive them on the main Coordinator loop.
1105                    // Spawning a task has the additional benefit that we get to
1106                    // delete batches once we're done.
1107                    //
1108                    // Batch deletion is best-effort, though, and there are
1109                    // multiple known ways in which they can leak, among them:
1110                    //
1111                    // - ProtoBatch is lost in flight
1112                    // - ProtoBatch is lost because when combining PeekResponse
1113                    // from workers a cancellation or error "overrides" other
1114                    // results, meaning we drop them
1115                    // - This task here is not run to completion before it can
1116                    // delete all batches
1117                    //
1118                    // This is semi-ok, because persist needs a reaper of leaked
1119                    // batches already, and so we piggy-back on that, even if it
1120                    // might not exist as of today.
1121                    let (tx, mut rx) = tokio::sync::mpsc::channel(1);
1122                    mz_ore::task::spawn(|| "read_peek_batches", async move {
1123                        // We always send our inline rows first. Ordering
1124                        // doesn't matter because we can only be in this case
1125                        // when there is no ORDER BY.
1126                        //
1127                        // We _could_ write these out as a Batch, and include it
1128                        // in the batches we read via the Consolidator. If we
1129                        // wanted to get a consistent ordering. That's not
1130                        // needed for correctness! But might be nice for more
1131                        // aesthetic reasons.
1132                        for rows in response.inline_rows {
1133                            let result = tx.send(rows).await;
1134                            if result.is_err() {
1135                                tracing::debug!("receiver went away");
1136                            }
1137                        }
1138
1139                        let mut current_batch = Vec::new();
1140                        let mut current_batch_size: usize = 0;
1141
1142                        'outer: while let Some(rows) = row_cursor.next().await {
1143                            for ((source_data, _val), _ts, diff) in rows {
1144                                let row = source_data
1145                                    .0
1146                                    .expect("we are not sending errors on this code path");
1147
1148                                let diff = usize::try_from(diff)
1149                                    .expect("peek responses cannot have negative diffs");
1150
1151                                if diff > 0 {
1152                                    let diff =
1153                                        NonZeroUsize::new(diff).expect("checked to be non-zero");
1154                                    current_batch_size =
1155                                        current_batch_size.saturating_add(row.byte_len());
1156                                    current_batch.push((row, diff));
1157                                }
1158
1159                                if current_batch_size > peek_stash_read_batch_size_bytes {
1160                                    // We're re-encoding the rows as a RowCollection
1161                                    // here, for which we pay in CPU time. We're in a
1162                                    // slow path already, since we're returning a big
1163                                    // stashed result so this is worth the convenience
1164                                    // of that for now.
1165                                    let result = tx
1166                                        .send(RowCollection::new(
1167                                            current_batch.drain(..).collect_vec(),
1168                                            &[],
1169                                        ))
1170                                        .await;
1171                                    if result.is_err() {
1172                                        tracing::debug!("receiver went away");
1173                                        // Don't return but break so we fall out to the
1174                                        // batch delete logic below.
1175                                        break 'outer;
1176                                    }
1177
1178                                    current_batch_size = 0;
1179                                }
1180                            }
1181                        }
1182
1183                        if current_batch.len() > 0 {
1184                            let result = tx.send(RowCollection::new(current_batch, &[])).await;
1185                            if result.is_err() {
1186                                tracing::debug!("receiver went away");
1187                            }
1188                        }
1189
1190                        let batches = row_cursor.into_lease();
1191                        tracing::trace!(?response.shard_id, "cleaning up batches of peek result");
1192                        for batch in batches {
1193                            batch.delete().await;
1194                        }
1195                    });
1196
1197                    assert!(
1198                        finishing.is_streamable(response.relation_desc.arity()),
1199                        "can only get stashed responses when the finishing is streamable"
1200                    );
1201
1202                    tracing::trace!("query result is streamable!");
1203
1204                    assert!(finishing.is_streamable(response.relation_desc.arity()));
1205                    let mut incremental_finishing = RowSetFinishingIncremental::new(
1206                        finishing.offset,
1207                        finishing.limit,
1208                        finishing.project,
1209                        max_returned_query_size,
1210                    );
1211
1212                    let mut got_zero_rows = true;
1213                    while let Some(rows) = rx.recv().await {
1214                        got_zero_rows = false;
1215
1216                        let result_rows = incremental_finishing.finish_incremental(
1217                            rows,
1218                            max_result_size,
1219                            &duration_histogram,
1220                        );
1221
1222                        match result_rows {
1223                            Ok(result_rows) => yield PeekResponseUnary::Rows(Box::new(result_rows)),
1224                            Err(e) => yield PeekResponseUnary::Error(e),
1225                        }
1226                    }
1227
1228                    // Even when there's zero rows, clients still expect an
1229                    // empty PeekResponse.
1230                    if got_zero_rows {
1231                        let row_iter = vec![].into_row_iter();
1232                        yield PeekResponseUnary::Rows(Box::new(row_iter));
1233                    }
1234                }
1235                PeekResponse::Canceled => {
1236                    yield PeekResponseUnary::Canceled;
1237                }
1238                PeekResponse::Error(e) => {
1239                    yield PeekResponseUnary::Error(e);
1240                }
1241            }
1242        })
1243    }
1244
1245    /// Cancel and remove all pending peeks that were initiated by the client with `conn_id`.
1246    #[mz_ore::instrument(level = "debug")]
1247    pub(crate) fn cancel_pending_peeks(&mut self, conn_id: &ConnectionId) {
1248        if let Some(uuids) = self.client_pending_peeks.remove(conn_id) {
1249            self.metrics
1250                .canceled_peeks
1251                .inc_by(u64::cast_from(uuids.len()));
1252
1253            let mut inverse: BTreeMap<ComputeInstanceId, BTreeSet<Uuid>> = Default::default();
1254            for (uuid, compute_instance) in &uuids {
1255                inverse.entry(*compute_instance).or_default().insert(*uuid);
1256            }
1257            for (compute_instance, uuids) in inverse {
1258                // It's possible that this compute instance no longer exists because it was dropped
1259                // while the peek was in progress. In this case we ignore the error and move on
1260                // because the dataflow no longer exists.
1261                // TODO(jkosh44) Dropping a cluster should actively cancel all pending queries.
1262                for uuid in uuids {
1263                    let _ = self.controller.compute.cancel_peek(
1264                        compute_instance,
1265                        uuid,
1266                        PeekResponse::Canceled,
1267                    );
1268                }
1269            }
1270
1271            let peeks = uuids
1272                .iter()
1273                .filter_map(|(uuid, _)| self.pending_peeks.remove(uuid))
1274                .collect::<Vec<_>>();
1275            for peek in peeks {
1276                self.retire_execution(
1277                    StatementEndedExecutionReason::Canceled,
1278                    peek.ctx_extra.defuse(),
1279                );
1280            }
1281        }
1282    }
1283
1284    /// Handle a peek notification and retire the corresponding execution. Does nothing for
1285    /// already-removed peeks.
1286    pub(crate) fn handle_peek_notification(
1287        &mut self,
1288        uuid: Uuid,
1289        notification: PeekNotification,
1290        otel_ctx: OpenTelemetryContext,
1291    ) {
1292        // We expect exactly one peek response, which we forward. Then we clean up the
1293        // peek's state in the coordinator.
1294        if let Some(PendingPeek {
1295            conn_id: _,
1296            cluster_id: _,
1297            depends_on: _,
1298            ctx_extra,
1299            is_fast_path,
1300        }) = self.remove_pending_peek(&uuid)
1301        {
1302            let reason = match notification {
1303                PeekNotification::Success {
1304                    rows: num_rows,
1305                    result_size,
1306                } => {
1307                    let strategy = if is_fast_path {
1308                        StatementExecutionStrategy::FastPath
1309                    } else {
1310                        StatementExecutionStrategy::Standard
1311                    };
1312                    StatementEndedExecutionReason::Success {
1313                        result_size: Some(result_size),
1314                        rows_returned: Some(num_rows),
1315                        execution_strategy: Some(strategy),
1316                    }
1317                }
1318                PeekNotification::Error(error) => StatementEndedExecutionReason::Errored { error },
1319                PeekNotification::Canceled => StatementEndedExecutionReason::Canceled,
1320            };
1321            otel_ctx.attach_as_parent();
1322            self.retire_execution(reason, ctx_extra.defuse());
1323        }
1324        // Cancellation may cause us to receive responses for peeks no
1325        // longer in `self.pending_peeks`, so we quietly ignore them.
1326    }
1327
1328    /// Clean up a peek's state.
1329    pub(crate) fn remove_pending_peek(&mut self, uuid: &Uuid) -> Option<PendingPeek> {
1330        let pending_peek = self.pending_peeks.remove(uuid);
1331        if let Some(pending_peek) = &pending_peek {
1332            let uuids = self
1333                .client_pending_peeks
1334                .get_mut(&pending_peek.conn_id)
1335                .expect("coord peek state is inconsistent");
1336            uuids.remove(uuid);
1337            if uuids.is_empty() {
1338                self.client_pending_peeks.remove(&pending_peek.conn_id);
1339            }
1340        }
1341        pending_peek
1342    }
1343
1344    /// Implements a slow-path peek by creating a transient dataflow.
1345    /// This is called from the command handler for ExecuteSlowPathPeek.
1346    ///
1347    /// (For now, this method simply delegates to implement_peek_plan by constructing
1348    /// the necessary PlannedPeek structure.)
1349    pub(crate) async fn implement_slow_path_peek(
1350        &mut self,
1351        dataflow_plan: PeekDataflowPlan,
1352        determination: TimestampDetermination,
1353        finishing: RowSetFinishing,
1354        compute_instance: ComputeInstanceId,
1355        target_replica: Option<ReplicaId>,
1356        intermediate_result_type: SqlRelationType,
1357        source_ids: BTreeSet<GlobalId>,
1358        conn_id: ConnectionId,
1359        max_result_size: u64,
1360        max_query_result_size: Option<u64>,
1361        watch_set: Option<WatchSetCreation>,
1362    ) -> Result<ExecuteResponse, AdapterError> {
1363        // Install watch sets for statement lifecycle logging if enabled.
1364        // This must happen _before_ creating ExecuteContextExtra, so that if it fails,
1365        // we don't have an ExecuteContextExtra that needs to be retired (the frontend
1366        // will handle logging for the error case).
1367        let statement_logging_id = watch_set.as_ref().map(|ws| ws.logging_id);
1368        if let Some(ws) = watch_set {
1369            self.install_peek_watch_sets(conn_id.clone(), ws)
1370                .map_err(|e| {
1371                    AdapterError::concurrent_dependency_drop_from_watch_set_install_error(e)
1372                })?;
1373        }
1374
1375        let source_arity = intermediate_result_type.arity();
1376
1377        let planned_peek = PlannedPeek {
1378            plan: PeekPlan::SlowPath(dataflow_plan),
1379            determination,
1380            conn_id,
1381            intermediate_result_type,
1382            source_arity,
1383            source_ids,
1384        };
1385
1386        // TODO(peek-seq): After the old peek sequencing is completely removed, we should merge the
1387        // relevant parts of the old `implement_peek_plan` into this method, and remove the old
1388        // `implement_peek_plan`.
1389        let mut ctx_guard =
1390            ExecuteContextGuard::new(statement_logging_id, self.internal_cmd_tx.clone());
1391        let result = self
1392            .implement_peek_plan(
1393                &mut ctx_guard,
1394                planned_peek,
1395                finishing,
1396                compute_instance,
1397                target_replica,
1398                max_result_size,
1399                max_query_result_size,
1400            )
1401            .await;
1402        // On error `implement_peek_plan` left the guard's contents intact (see
1403        // its doc comment) and the frontend logs the error end, so we defuse
1404        // rather than let the guard's `Drop` emit a spurious `Aborted`.
1405        if result.is_err() {
1406            let _ = ctx_guard.defuse();
1407        }
1408        result
1409    }
1410
1411    /// Implements a `COPY TO` command by installing peek watch sets,
1412    /// shipping the dataflow, and spawning a background task to wait for completion.
1413    /// This is called from the command handler for ExecuteCopyTo.
1414    ///
1415    /// (The S3 preflight check must be completed successfully via the
1416    /// `CopyToPreflight` command _before_ calling this method. The preflight is
1417    /// handled separately to avoid blocking the coordinator's main task with
1418    /// slow S3 network operations.)
1419    ///
1420    /// This method does NOT block waiting for completion. Instead, it spawns a background task that
1421    /// will send the response through the provided tx channel when the COPY TO completes.
1422    /// All errors (setup or execution) are sent through tx.
1423    pub(crate) async fn implement_copy_to(
1424        &mut self,
1425        df_desc: DataflowDescription<mz_compute_types::plan::LirRelationExpr>,
1426        compute_instance: ComputeInstanceId,
1427        target_replica: Option<ReplicaId>,
1428        source_ids: BTreeSet<GlobalId>,
1429        conn_id: ConnectionId,
1430        watch_set: Option<WatchSetCreation>,
1431        tx: oneshot::Sender<Result<ExecuteResponse, AdapterError>>,
1432    ) {
1433        // Helper to send error and return early
1434        let send_err = |tx: oneshot::Sender<Result<ExecuteResponse, AdapterError>>,
1435                        e: AdapterError| {
1436            let _ = tx.send(Err(e));
1437        };
1438
1439        // Install watch sets for statement lifecycle logging if enabled.
1440        // If this fails, we just send the error back. The frontend will handle logging
1441        // for the error case (no ExecuteContextExtra is created here).
1442        if let Some(ws) = watch_set {
1443            if let Err(e) = self.install_peek_watch_sets(conn_id.clone(), ws) {
1444                let err = AdapterError::concurrent_dependency_drop_from_watch_set_install_error(e);
1445                send_err(tx, err);
1446                return;
1447            }
1448        }
1449
1450        // Note: We don't create an ExecuteContextExtra here because the frontend handles
1451        // all statement logging for COPY TO operations.
1452
1453        let sink_id = df_desc.sink_id();
1454
1455        // Create and register ActiveCopyTo.
1456        // Note: sink_tx/sink_rx is the channel for the compute sink to notify completion
1457        // This is different from the command's tx which sends the response to the client
1458        let (sink_tx, sink_rx) = oneshot::channel();
1459        let active_copy_to = ActiveCopyTo {
1460            conn_id: conn_id.clone(),
1461            tx: sink_tx,
1462            cluster_id: compute_instance,
1463            depends_on: source_ids,
1464        };
1465
1466        // Add metadata for the new COPY TO. CopyTo returns a `ready` future, so it is safe to drop.
1467        drop(self.add_active_compute_sink(sink_id, ActiveComputeSink::CopyTo(active_copy_to)));
1468
1469        // Try to ship the dataflow. We handle errors gracefully because dependencies might have
1470        // disappeared during sequencing.
1471        if let Err(e) = self
1472            .try_ship_dataflow(df_desc, compute_instance, target_replica)
1473            .await
1474            .map_err(AdapterError::concurrent_dependency_drop_from_dataflow_creation_error)
1475        {
1476            // Clean up the active compute sink that was added above, since the dataflow was never
1477            // created. If we don't do this, the sink_id remains in drop_sinks but no collection
1478            // exists in the compute controller, causing a panic when the connection terminates.
1479            self.remove_active_compute_sink(sink_id).await;
1480            send_err(tx, e);
1481            return;
1482        }
1483
1484        // Spawn background task to wait for completion
1485        // We must NOT await sink_rx here directly, as that would block the coordinator's main task
1486        // from processing the completion message. Instead, we spawn a background task that will
1487        // send the result through tx when the COPY TO completes.
1488        let span = Span::current();
1489        task::spawn(
1490            || "copy to completion",
1491            async move {
1492                let res = sink_rx.await;
1493                let result = match res {
1494                    Ok(res) => res,
1495                    Err(_) => Err(AdapterError::Internal("copy to sender dropped".into())),
1496                };
1497
1498                let _ = tx.send(result);
1499            }
1500            .instrument(span),
1501        );
1502    }
1503
1504    /// Constructs an [`ExecuteResponse`] that that will send some rows to the
1505    /// client immediately, as opposed to asking the dataflow layer to send along
1506    /// the rows after some computation.
1507    pub(crate) fn send_immediate_rows<I>(rows: I) -> ExecuteResponse
1508    where
1509        I: IntoRowIterator,
1510        I::Iter: Send + Sync + 'static,
1511    {
1512        let rows = Box::new(rows.into_row_iter());
1513        ExecuteResponse::SendingRowsImmediate { rows }
1514    }
1515}
1516
1517#[cfg(test)]
1518mod tests {
1519    use mz_expr::func::IsNull;
1520    use mz_expr::{MapFilterProject, UnaryFunc};
1521    use mz_ore::str::Indent;
1522    use mz_repr::explain::text::text_string_at;
1523    use mz_repr::explain::{DummyHumanizer, ExplainConfig, PlanRenderingContext};
1524    use mz_repr::{Datum, SqlColumnType, SqlScalarType};
1525
1526    use super::*;
1527
1528    #[mz_ore::test]
1529    #[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function `rust_psm_stack_pointer` on OS `linux`
1530    fn test_fast_path_plan_as_text() {
1531        let typ = SqlRelationType::new(vec![SqlColumnType {
1532            scalar_type: SqlScalarType::String,
1533            nullable: false,
1534        }]);
1535        let constant_err = FastPathPlan::Constant(Err(EvalError::DivisionByZero), typ.clone());
1536        let no_lookup = FastPathPlan::PeekExisting(
1537            GlobalId::User(8),
1538            GlobalId::User(10),
1539            None,
1540            MapFilterProject::new(4)
1541                .map(Some(MirScalarExpr::column(0).or(MirScalarExpr::column(2))))
1542                .project([1, 4])
1543                .into_plan()
1544                .expect("invalid plan")
1545                .into_nontemporal()
1546                .expect("invalid nontemporal"),
1547        );
1548        let lookup = FastPathPlan::PeekExisting(
1549            GlobalId::User(9),
1550            GlobalId::User(11),
1551            Some(vec![Row::pack(Some(Datum::Int32(5)))]),
1552            MapFilterProject::new(3)
1553                .filter(Some(
1554                    MirScalarExpr::column(0).call_unary(UnaryFunc::IsNull(IsNull)),
1555                ))
1556                .into_plan()
1557                .expect("invalid plan")
1558                .into_nontemporal()
1559                .expect("invalid nontemporal"),
1560        );
1561
1562        let humanizer = DummyHumanizer;
1563        let config = ExplainConfig {
1564            redacted: false,
1565            verbose_syntax: true,
1566            ..Default::default()
1567        };
1568        let ctx_gen = || {
1569            let indent = Indent::default();
1570            let annotations = BTreeMap::new();
1571            PlanRenderingContext::<FastPathPlan>::new(
1572                indent,
1573                &humanizer,
1574                annotations,
1575                &config,
1576                BTreeSet::default(),
1577            )
1578        };
1579
1580        let constant_err_exp = "Error \"division by zero\"\n";
1581        let no_lookup_exp = "Project (#1, #4)\n  Map ((#0 OR #2))\n    ReadIndex on=u8 [DELETED INDEX]=[*** full scan ***]\n";
1582        let lookup_exp =
1583            "Filter (#0) IS NULL\n  ReadIndex on=u9 [DELETED INDEX]=[lookup value=(5)]\n";
1584
1585        assert_eq!(text_string_at(&constant_err, ctx_gen), constant_err_exp);
1586        assert_eq!(text_string_at(&no_lookup, ctx_gen), no_lookup_exp);
1587        assert_eq!(text_string_at(&lookup, ctx_gen), lookup_exp);
1588
1589        let mut constant_rows = vec![
1590            (Row::pack(Some(Datum::String("hello"))), Diff::ONE),
1591            (Row::pack(Some(Datum::String("world"))), 2.into()),
1592            (Row::pack(Some(Datum::String("star"))), 500.into()),
1593        ];
1594        let constant_exp1 =
1595            "Constant\n  - (\"hello\")\n  - ((\"world\") x 2)\n  - ((\"star\") x 500)\n";
1596        assert_eq!(
1597            text_string_at(
1598                &FastPathPlan::Constant(Ok(constant_rows.clone()), typ.clone()),
1599                ctx_gen
1600            ),
1601            constant_exp1
1602        );
1603        constant_rows
1604            .extend((0..20).map(|i| (Row::pack(Some(Datum::String(&i.to_string()))), Diff::ONE)));
1605        let constant_exp2 = "Constant\n  total_rows (diffs absed): 523\n  first_rows:\n    - (\"hello\")\
1606        \n    - ((\"world\") x 2)\n    - ((\"star\") x 500)\n    - (\"0\")\n    - (\"1\")\
1607        \n    - (\"2\")\n    - (\"3\")\n    - (\"4\")\n    - (\"5\")\n    - (\"6\")\
1608        \n    - (\"7\")\n    - (\"8\")\n    - (\"9\")\n    - (\"10\")\n    - (\"11\")\
1609        \n    - (\"12\")\n    - (\"13\")\n    - (\"14\")\n    - (\"15\")\n    - (\"16\")\n";
1610        assert_eq!(
1611            text_string_at(&FastPathPlan::Constant(Ok(constant_rows), typ), ctx_gen),
1612            constant_exp2
1613        );
1614    }
1615}