1use std::sync::Arc;
22use std::time::{Duration, Instant};
23
24use mz_compute_types::plan::LirRelationExpr;
25use mz_compute_types::sinks::{ComputeSinkConnection, ComputeSinkDesc, MetricSinkConnection};
26use mz_expr::func::variadic::Coalesce;
27use mz_expr::{MirRelationExpr, MirScalarExpr, func};
28use mz_repr::explain::trace_plan;
29use mz_repr::{
30 ColumnName, Datum, GlobalId, RelationDesc, ReprRelationType, ReprScalarType, Row, SqlScalarType,
31};
32use mz_sql::names::QualifiedItemName;
33use mz_sql::optimizer_metrics::OptimizerMetrics;
34use mz_transform::TransformCtx;
35use mz_transform::dataflow::DataflowMetainfo;
36use mz_transform::normalize_lets::normalize_lets;
37use mz_transform::typecheck::{SharedTypecheckingContext, empty_typechecking_context};
38use timely::progress::Antichain;
39
40use crate::optimize::dataflows::{
41 ComputeInstanceSnapshot, DataflowBuilder, ExprPrep, ExprPrepMaintained,
42};
43use crate::optimize::{
44 LirDataflowDescription, MirDataflowDescription, Optimize, OptimizerCatalog, OptimizerConfig,
45 OptimizerError, optimize_mir_local,
46};
47
48const METRIC_NAME_PATTERN: &str = "^[a-zA-Z_:][a-zA-Z0-9_:]*$";
53
54pub struct Optimizer {
60 typecheck_ctx: SharedTypecheckingContext,
62 catalog: Arc<dyn OptimizerCatalog>,
64 compute_instance: ComputeInstanceSnapshot,
66 view_id: GlobalId,
69 sink_id: GlobalId,
71 config: OptimizerConfig,
73 metrics: OptimizerMetrics,
75 duration: Duration,
77}
78
79impl Optimizer {
80 pub fn new(
81 catalog: Arc<dyn OptimizerCatalog>,
82 compute_instance: ComputeInstanceSnapshot,
83 view_id: GlobalId,
84 sink_id: GlobalId,
85 config: OptimizerConfig,
86 metrics: OptimizerMetrics,
87 ) -> Self {
88 Self {
89 typecheck_ctx: empty_typechecking_context(),
90 catalog,
91 compute_instance,
92 view_id,
93 sink_id,
94 config,
95 metrics,
96 duration: Default::default(),
97 }
98 }
99}
100
101pub struct MetricSink {
103 name: QualifiedItemName,
104 from: GlobalId,
105 prefix: String,
108}
109
110impl MetricSink {
111 pub fn new(name: QualifiedItemName, from: GlobalId, prefix: String) -> Self {
113 Self { name, from, prefix }
114 }
115}
116
117#[derive(Clone, Debug)]
120pub struct GlobalMirPlan {
121 df_desc: MirDataflowDescription,
122 df_meta: DataflowMetainfo,
123}
124
125impl GlobalMirPlan {
126 pub fn df_desc(&self) -> &MirDataflowDescription {
127 &self.df_desc
128 }
129}
130
131#[derive(Clone, Debug)]
134pub struct GlobalLirPlan {
135 df_desc: LirDataflowDescription,
136 df_meta: DataflowMetainfo,
137}
138
139impl GlobalLirPlan {
140 pub fn df_desc(&self) -> &LirDataflowDescription {
141 &self.df_desc
142 }
143}
144
145impl Optimize<MetricSink> for Optimizer {
146 type To = GlobalMirPlan;
147
148 fn optimize(&mut self, metric_sink: MetricSink) -> Result<Self::To, OptimizerError> {
149 let time = Instant::now();
150
151 let from_entry = self.catalog.get_entry(&metric_sink.from);
152 let full_name = self
153 .catalog
154 .resolve_full_name(&metric_sink.name, from_entry.conn_id());
155 let from_desc = from_entry
156 .relation_desc()
157 .expect("can only create a metric sink on items with a valid description")
158 .into_owned();
159
160 let mut df_builder = {
161 let compute = self.compute_instance.clone();
162 DataflowBuilder::new(&*self.catalog, compute).with_config(&self.config)
163 };
164 let mut df_desc = MirDataflowDescription::new(full_name.to_string());
165 let mut df_meta = DataflowMetainfo::default();
166
167 df_builder.import_into_dataflow(&metric_sink.from, &mut df_desc, &self.config.features)?;
168 df_builder.maybe_reoptimize_imported_views(&mut df_desc, &self.config)?;
169
170 let (shaped_expr, shaped_desc) =
174 shape_metric_sink_source(metric_sink.from, &from_desc, &metric_sink.prefix);
175 let mut local_ctx = TransformCtx::local(
176 &self.config.features,
177 &self.typecheck_ctx,
178 &mut df_meta,
179 Some(&mut self.metrics),
180 Some(self.view_id),
181 );
182 let shaped_expr = optimize_mir_local(shaped_expr, &mut local_ctx)?;
183
184 df_builder.import_view_into_dataflow(
185 &self.view_id,
186 &shaped_expr,
187 &mut df_desc,
188 &self.config.features,
189 )?;
190 df_builder.maybe_reoptimize_imported_views(&mut df_desc, &self.config)?;
191
192 let sink_description = ComputeSinkDesc {
193 from: self.view_id,
194 from_desc: shaped_desc,
195 connection: ComputeSinkConnection::MetricSink(MetricSinkConnection {}),
196 with_snapshot: true,
197 up_to: Antichain::new(),
198 non_null_assertions: Vec::new(),
199 refresh_schedule: None,
200 };
201 df_desc.export_sink(self.sink_id, sink_description);
202
203 let style = ExprPrepMaintained;
205 df_desc.visit_children(
206 |r| style.prep_relation_expr(r),
207 |s| style.prep_scalar_expr(s),
208 )?;
209
210 let mut transform_ctx = TransformCtx::global(
212 &df_builder,
213 &mz_transform::EmptyStatisticsOracle,
214 &self.config.features,
215 &self.typecheck_ctx,
216 &mut df_meta,
217 Some(&mut self.metrics),
218 );
219 mz_transform::optimize_dataflow(&mut df_desc, &mut transform_ctx, false)?;
221
222 self.duration += time.elapsed();
223
224 Ok(GlobalMirPlan { df_desc, df_meta })
225 }
226}
227
228impl Optimize<GlobalMirPlan> for Optimizer {
229 type To = GlobalLirPlan;
230
231 fn optimize(&mut self, plan: GlobalMirPlan) -> Result<Self::To, OptimizerError> {
232 let time = Instant::now();
233
234 let GlobalMirPlan {
235 mut df_desc,
236 df_meta,
237 } = plan;
238
239 for build in df_desc.objects_to_build.iter_mut() {
241 normalize_lets(&mut build.plan.0, &self.config.features)?
242 }
243
244 let df_desc = LirRelationExpr::finalize_dataflow(
246 df_desc,
247 &self.config.features,
248 Some(self.metrics.lowering()),
249 )?;
250
251 trace_plan(&df_desc);
253
254 self.duration += time.elapsed();
255 self.metrics
256 .observe_e2e_optimization_time("metric_sink", self.duration);
257
258 Ok(GlobalLirPlan { df_desc, df_meta })
259 }
260}
261
262impl GlobalLirPlan {
263 pub fn unapply(self) -> (LirDataflowDescription, DataflowMetainfo) {
265 (self.df_desc, self.df_meta)
266 }
267}
268
269fn shape_metric_sink_source(
291 from_id: GlobalId,
292 from_desc: &RelationDesc,
293 prefix: &str,
294) -> (MirRelationExpr, RelationDesc) {
295 let get_idx = |name: &str| {
300 from_desc
301 .get_by_name(&ColumnName::from(name))
302 .expect("metric-sink source relation must expose the canonical columns")
303 };
304 let (metric_name_idx, metric_name_ct) = get_idx("metric_name");
305 let (metric_type_idx, metric_type_ct) = get_idx("metric_type");
306 let (labels_idx, labels_ct) = get_idx("labels");
307 let (value_idx, value_ct) = get_idx("value");
308 let (help_idx, help_ct) = get_idx("help");
309
310 let repr_typ = ReprRelationType::from(from_desc.typ());
311 let arity = repr_typ.column_types.len();
312 let labels_repr_type = ReprScalarType::from(&labels_ct.scalar_type);
313
314 let empty_map_row = {
315 let mut row = Row::default();
316 row.packer().push_dict_with(|_| {});
317 row
318 };
319 let labels_coalesced = MirScalarExpr::call_variadic(
320 Coalesce,
321 vec![
322 MirScalarExpr::column(labels_idx),
323 MirScalarExpr::literal_from_single_element_row(empty_map_row, labels_repr_type),
324 ],
325 );
326 let help_coalesced = MirScalarExpr::call_variadic(
327 Coalesce,
328 vec![
329 MirScalarExpr::column(help_idx),
330 MirScalarExpr::literal_ok(Datum::String(""), ReprScalarType::String),
331 ],
332 );
333
334 let metric_type_literal = |s: &'static str| {
335 MirScalarExpr::column(metric_type_idx).call_binary(
336 MirScalarExpr::literal_ok(Datum::String(s), ReprScalarType::String),
337 func::Eq,
338 )
339 };
340 let metric_kind = metric_type_literal("gauge").if_then_else(
341 MirScalarExpr::literal_ok(Datum::Int32(0), ReprScalarType::Int32),
342 metric_type_literal("counter").if_then_else(
343 MirScalarExpr::literal_ok(Datum::Int32(1), ReprScalarType::Int32),
344 MirScalarExpr::literal_null(ReprScalarType::Int32),
345 ),
346 );
347
348 let prefixed_name = MirScalarExpr::literal_ok(Datum::String(prefix), ReprScalarType::String)
354 .call_binary(
355 MirScalarExpr::column(metric_name_idx),
356 func::TextConcatBinary,
357 );
358
359 let published_name = MirScalarExpr::column(arity + 2);
365 let name_valid = published_name
366 .clone()
367 .call_is_null()
368 .not()
369 .and(published_name.call_binary(
370 MirScalarExpr::literal_ok(Datum::String(METRIC_NAME_PATTERN), ReprScalarType::String),
371 func::IsRegexpMatchCaseSensitive,
372 ));
373
374 let shaped_expr = MirRelationExpr::global_get(from_id, repr_typ)
375 .map(vec![
376 labels_coalesced,
377 help_coalesced,
378 prefixed_name,
379 metric_kind,
380 name_valid,
381 ])
382 .project(vec![
383 arity + 2, metric_type_idx,
385 arity, value_idx,
387 arity + 1, arity + 3, arity + 4, ]);
391
392 let mut labels_shaped_ct = labels_ct.clone();
393 labels_shaped_ct.nullable = false;
394 let mut help_shaped_ct = help_ct.clone();
395 help_shaped_ct.nullable = false;
396 let shaped_desc = RelationDesc::from_names_and_types([
397 ("metric_name", metric_name_ct.clone()),
398 ("metric_type", metric_type_ct.clone()),
399 ("labels", labels_shaped_ct),
400 ("value", value_ct.clone()),
401 ("help", help_shaped_ct),
402 ("metric_kind", SqlScalarType::Int32.nullable(true)),
403 ("name_valid", SqlScalarType::Bool.nullable(true)),
404 ]);
405
406 (shaped_expr, shaped_desc)
407}
408
409#[cfg(test)]
410mod tests {
411 use std::collections::BTreeMap;
412
413 use mz_catalog::memory::objects::{CatalogEntry, CatalogItem, Table, TableDataSource};
414 use mz_controller_types::ClusterId;
415 use mz_expr::Eval;
416 use mz_ore::metrics::MetricsRegistry;
417 use mz_repr::adt::mz_acl_item::PrivilegeMap;
418 use mz_repr::role_id::RoleId;
419 use mz_repr::{
420 CatalogItemId, RelationVersion, RelationVersionSelector, RowArena, SqlColumnType,
421 VersionedRelationDesc,
422 };
423 use mz_sql::names::{
424 FullItemName, ItemQualifiers, RawDatabaseSpecifier, ResolvedDatabaseSpecifier, ResolvedIds,
425 SchemaId, SchemaSpecifier,
426 };
427 use mz_sql::session::vars::SystemVars;
428
429 use super::*;
430
431 fn source_desc() -> RelationDesc {
435 RelationDesc::builder()
436 .with_column("metric_name", SqlScalarType::String.nullable(true))
437 .with_column("metric_type", SqlScalarType::String.nullable(false))
438 .with_column(
439 "labels",
440 SqlScalarType::Map {
441 value_type: Box::new(SqlScalarType::String),
442 custom_id: None,
443 }
444 .nullable(true),
445 )
446 .with_column("value", SqlScalarType::Float64.nullable(true))
447 .with_column("help", SqlScalarType::String.nullable(true))
448 .with_column("extra", SqlScalarType::String.nullable(true))
449 .finish()
450 }
451
452 #[mz_ore::test]
453 fn shaped_desc_column_contract() {
454 let (_expr, desc) =
455 shape_metric_sink_source(GlobalId::Transient(0), &source_desc(), "app_");
456
457 let cols: Vec<(String, SqlColumnType)> = desc
458 .iter()
459 .map(|(name, ty)| (name.as_str().to_string(), ty.clone()))
460 .collect();
461
462 let names: Vec<&str> = cols.iter().map(|(n, _)| n.as_str()).collect();
465 assert_eq!(
466 names,
467 vec![
468 "metric_name",
469 "metric_type",
470 "labels",
471 "value",
472 "help",
473 "metric_kind",
474 "name_valid",
475 ]
476 );
477
478 let by_name = |name: &str| {
479 cols.iter()
480 .find(|(n, _)| n == name)
481 .expect("column present in shaped desc")
482 .1
483 .clone()
484 };
485
486 assert!(!by_name("labels").nullable);
488 assert!(!by_name("help").nullable);
489
490 assert!(by_name("metric_name").nullable);
492 assert!(by_name("value").nullable);
493
494 assert_eq!(by_name("metric_kind"), SqlScalarType::Int32.nullable(true));
496 assert_eq!(by_name("name_valid"), SqlScalarType::Bool.nullable(true));
497 }
498
499 #[mz_ore::test]
500 fn shaped_expr_projects_seven_columns() {
501 let (expr, _desc) =
502 shape_metric_sink_source(GlobalId::Transient(0), &source_desc(), "app_");
503
504 match &expr {
507 MirRelationExpr::Project { outputs, .. } => {
508 assert_eq!(outputs.len(), 7);
509 }
510 other => panic!("expected a Project at the root of the shaped expr, got {other:?}"),
511 }
512 }
513
514 fn shaped_map_scalars(desc: &RelationDesc, prefix: &str) -> Vec<MirScalarExpr> {
517 let (expr, _desc) = shape_metric_sink_source(GlobalId::Transient(0), desc, prefix);
518 match expr {
519 MirRelationExpr::Project { input, .. } => match *input {
520 MirRelationExpr::Map { scalars, .. } => scalars,
521 other => panic!("expected a Map under the Project, got {other:?}"),
522 },
523 other => panic!("expected a Project at the root, got {other:?}"),
524 }
525 }
526
527 fn eval_shaped_row<'a>(
531 scalars: &'a [MirScalarExpr],
532 input: &[Datum<'a>],
533 arena: &'a RowArena,
534 ) -> Vec<Datum<'a>> {
535 let mut row = input.to_vec();
536 for scalar in scalars {
537 let datum = scalar.eval(&row, arena).expect("scalar eval succeeds");
538 row.push(datum);
539 }
540 row
541 }
542
543 #[mz_ore::test]
544 fn metric_kind_classifies_type() {
545 let scalars = shaped_map_scalars(&source_desc(), "app_");
546 let metric_kind = &scalars[3];
547 let arena = RowArena::new();
548 for (metric_type, expected) in [
550 ("gauge", Datum::Int32(0)),
551 ("counter", Datum::Int32(1)),
552 ("histogram", Datum::Null),
553 ("summary", Datum::Null),
554 ] {
555 let row = [
556 Datum::Null,
557 Datum::String(metric_type),
558 Datum::Null,
559 Datum::Null,
560 Datum::Null,
561 Datum::Null,
562 ];
563 assert_eq!(
564 metric_kind
565 .eval(&row, &arena)
566 .expect("metric_kind eval succeeds"),
567 expected,
568 "metric_type = {metric_type}",
569 );
570 }
571 }
572
573 #[mz_ore::test]
574 fn name_valid_matches_prometheus_grammar() {
575 let scalars = shaped_map_scalars(&source_desc(), "app_");
580 let arena = RowArena::new();
581 for (metric_name, expected) in [
582 (Datum::String("http_requests_total"), Datum::True),
583 (Datum::String("with:colons_and_1_digit"), Datum::True),
584 (Datum::String("1_leading_digit"), Datum::True),
585 (Datum::String("has-a-dash"), Datum::False),
586 (Datum::String(""), Datum::True),
587 (Datum::Null, Datum::False),
588 ] {
589 let input = [
590 metric_name,
591 Datum::Null,
592 Datum::Null,
593 Datum::Null,
594 Datum::Null,
595 Datum::Null,
596 ];
597 let row = eval_shaped_row(&scalars, &input, &arena);
600 assert_eq!(
601 *row.last().expect("row has appended columns"),
602 expected,
603 "metric_name = {metric_name:?}",
604 );
605 }
606 }
607
608 #[derive(Debug)]
611 struct SingleTableCatalog {
612 entry: CatalogEntry,
613 }
614
615 const TABLE_ITEM_ID: CatalogItemId = CatalogItemId::User(1);
616 const TABLE_GID: GlobalId = GlobalId::User(1);
617 const SINK_GID: GlobalId = GlobalId::User(2);
618
619 impl SingleTableCatalog {
620 fn new() -> Self {
621 let table = Table {
622 create_sql: None,
623 desc: VersionedRelationDesc::new(source_desc()),
624 collections: BTreeMap::from([(RelationVersion::root(), TABLE_GID)]),
625 conn_id: None,
626 resolved_ids: ResolvedIds::empty(),
627 custom_logical_compaction_window: None,
628 is_retained_metrics_object: false,
629 data_source: TableDataSource::TableWrites {
630 defaults: Vec::new(),
631 },
632 };
633 let entry = CatalogEntry {
634 item: CatalogItem::Table(table),
635 referenced_by: Vec::new(),
636 used_by: Vec::new(),
637 id: TABLE_ITEM_ID,
638 oid: 20_000,
639 name: QualifiedItemName {
640 qualifiers: ItemQualifiers {
641 database_spec: ResolvedDatabaseSpecifier::Ambient,
642 schema_spec: SchemaSpecifier::Id(SchemaId::User(1)),
643 },
644 item: "t".to_string(),
645 },
646 owner_id: RoleId::User(1),
647 privileges: PrivilegeMap::default(),
648 };
649 Self { entry }
650 }
651 }
652
653 impl OptimizerCatalog for SingleTableCatalog {
654 fn get_entry(&self, _id: &GlobalId) -> mz_catalog::memory::objects::CatalogCollectionEntry {
655 mz_catalog::memory::objects::CatalogCollectionEntry {
656 entry: self.entry.clone(),
657 version: RelationVersionSelector::Latest,
658 }
659 }
660
661 fn get_entry_by_item_id(&self, _id: &CatalogItemId) -> &CatalogEntry {
662 &self.entry
663 }
664
665 fn resolve_full_name(
666 &self,
667 name: &QualifiedItemName,
668 _conn_id: Option<&mz_adapter_types::connection::ConnectionId>,
669 ) -> FullItemName {
670 FullItemName {
671 database: RawDatabaseSpecifier::Ambient,
672 schema: "public".to_string(),
673 item: name.item.clone(),
674 }
675 }
676
677 fn get_indexes_on(
678 &self,
679 _id: GlobalId,
680 _cluster: ClusterId,
681 ) -> Box<dyn Iterator<Item = (GlobalId, &mz_catalog::memory::objects::Index)> + '_>
682 {
683 Box::new(std::iter::empty())
684 }
685 }
686
687 #[mz_ore::test]
690 fn optimizer_exports_one_metric_sink() {
691 let catalog = Arc::new(SingleTableCatalog::new());
692 let cluster_id = ClusterId::user(1).expect("valid cluster id");
693 let compute_instance = ComputeInstanceSnapshot::new_without_collections(cluster_id);
694 let view_id = GlobalId::Transient(1);
695 let config = OptimizerConfig::from(&SystemVars::default());
696 let metrics = OptimizerMetrics::register_into(&MetricsRegistry::new(), Duration::MAX);
697
698 let mut optimizer = Optimizer::new(
699 catalog,
700 compute_instance,
701 view_id,
702 SINK_GID,
703 config,
704 metrics,
705 );
706
707 let name = QualifiedItemName {
708 qualifiers: ItemQualifiers {
709 database_spec: ResolvedDatabaseSpecifier::Ambient,
710 schema_spec: SchemaSpecifier::Id(SchemaId::User(1)),
711 },
712 item: "s".to_string(),
713 };
714 let global_mir_plan = optimizer
715 .optimize(MetricSink::new(name, TABLE_GID, "app_".to_string()))
716 .expect("MIR optimization succeeds");
717 let global_lir_plan = optimizer
718 .optimize(global_mir_plan)
719 .expect("LIR optimization succeeds");
720
721 let df_desc = global_lir_plan.df_desc();
722 assert!(df_desc.index_exports.is_empty());
723 let sink_exports: Vec<_> = df_desc.sink_exports.iter().collect();
724 assert_eq!(sink_exports.len(), 1);
725 let (sink_id, sink_desc) = sink_exports[0];
726 assert_eq!(*sink_id, SINK_GID);
727 assert!(matches!(
728 sink_desc.connection,
729 ComputeSinkConnection::MetricSink(_)
730 ));
731 assert_eq!(sink_desc.from, view_id);
733 let shaped_names: Vec<&str> = sink_desc
734 .from_desc
735 .iter_names()
736 .map(|n| n.as_str())
737 .collect();
738 assert_eq!(
739 shaped_names,
740 vec![
741 "metric_name",
742 "metric_type",
743 "labels",
744 "value",
745 "help",
746 "metric_kind",
747 "name_valid",
748 ]
749 );
750 }
751}