1use 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#[derive(Debug)]
67pub(crate) struct PendingPeek {
68 pub(crate) conn_id: ConnectionId,
70 pub(crate) cluster_id: ClusterId,
72 pub(crate) depends_on: BTreeSet<GlobalId>,
74 pub(crate) ctx_extra: ExecuteContextGuard,
77 pub(crate) is_fast_path: bool,
79}
80
81#[derive(Debug)]
86pub enum PeekResponseUnary {
87 Rows(Box<dyn RowIterator + Send + Sync>),
88 Error(String),
89 Canceled,
90 DependencyDropped(DroppedDependency),
96}
97
98#[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 pub fn query_terminated_error(&self) -> String {
122 format!("query could not complete because {self} was dropped")
123 }
124
125 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 Constant(Result<Vec<(Row, Diff)>, EvalError>, SqlRelationType),
177 PeekExisting(GlobalId, GlobalId, Option<Vec<Row>>, mz_expr::SafeMfpPlan),
180 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 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 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 pub intermediate_result_type: SqlRelationType,
436 pub source_arity: usize,
437 pub source_ids: BTreeSet<GlobalId>,
438}
439
440#[derive(Clone, Debug)]
442pub enum PeekPlan {
443 FastPath(FastPathPlan),
444 SlowPath(PeekDataflowPlan),
446}
447
448fn 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
461fn 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
473pub 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 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 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 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 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 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 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 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 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 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 _ => {}
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 #[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 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(&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 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 id: index_id,
861 key: index_key,
862 permutation: index_permutation,
863 thinned_arity: index_thinned_arity,
864 }) => {
865 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 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 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 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 let (rows_tx, rows_rx) = tokio::sync::oneshot::channel();
934
935 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 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 let Some(index_id) = drop_dataflow {
971 self.drop_compute_collections(vec![(compute_instance, index_id)]);
972 }
973 return Err(e);
974 }
975
976 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 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 #[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 let (tx, mut rx) = tokio::sync::mpsc::channel(1);
1122 mz_ore::task::spawn(|| "read_peek_batches", async move {
1123 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 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 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 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 #[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 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 pub(crate) fn handle_peek_notification(
1287 &mut self,
1288 uuid: Uuid,
1289 notification: PeekNotification,
1290 otel_ctx: OpenTelemetryContext,
1291 ) {
1292 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 }
1327
1328 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 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 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 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 if result.is_err() {
1406 let _ = ctx_guard.defuse();
1407 }
1408 result
1409 }
1410
1411 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 let send_err = |tx: oneshot::Sender<Result<ExecuteResponse, AdapterError>>,
1435 e: AdapterError| {
1436 let _ = tx.send(Err(e));
1437 };
1438
1439 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 let sink_id = df_desc.sink_id();
1454
1455 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 drop(self.add_active_compute_sink(sink_id, ActiveComputeSink::CopyTo(active_copy_to)));
1468
1469 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 self.remove_active_compute_sink(sink_id).await;
1480 send_err(tx, e);
1481 return;
1482 }
1483
1484 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 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)] 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}