1use std::fmt::Write;
2
3use polars_ops::frame::JoinArgs;
4use polars_plan::dsl::PartitionStrategyIR;
5use polars_plan::plans::expr_ir::ExprIR;
6use polars_plan::plans::{AExpr, EscapeLabel};
7use polars_plan::prelude::FileWriteFormat;
8use polars_time::ClosedWindow;
9#[cfg(feature = "dynamic_group_by")]
10use polars_time::DynamicGroupOptions;
11use polars_utils::arena::Arena;
12use polars_utils::itertools::Itertools;
13use polars_utils::slice_enum::Slice;
14use slotmap::{Key, SecondaryMap, SlotMap};
15
16use super::{PhysNode, PhysNodeKey, PhysNodeKind};
17use crate::physical_plan::ZipBehavior;
18
19pub enum NodeStyle {
21 InMemoryFallback,
22 MemoryIntensive,
23 Generic,
24}
25
26impl NodeStyle {
27 const COLOR_IN_MEM_FALLBACK: &str = "0.0 0.3 1.0"; const COLOR_MEM_INTENSIVE: &str = "0.16 0.3 1.0"; pub fn for_node_kind(kind: &PhysNodeKind) -> Self {
32 use PhysNodeKind as K;
33 match kind {
34 K::InMemoryMap { .. } | K::InMemoryJoin { .. } | K::ColumnarFunction { .. } => {
35 Self::InMemoryFallback
36 },
37 K::InMemorySource { .. }
38 | K::InputIndependentSelect { .. }
39 | K::NegativeSlice { .. }
40 | K::InMemorySink { .. }
41 | K::Sort { .. }
42 | K::GroupBy { .. }
43 | K::EquiJoin { .. }
44 | K::SemiAntiJoin { .. }
45 | K::CrossJoin { .. }
46 | K::Multiplexer { .. }
47 | K::Gather { .. } => Self::MemoryIntensive,
48 #[cfg(feature = "iejoin")]
49 K::RangeJoin { .. } => Self::MemoryIntensive,
50 #[cfg(feature = "merge_sorted")]
51 K::MergeSorted { .. } => Self::MemoryIntensive,
52 _ => Self::Generic,
53 }
54 }
55
56 pub fn node_attrs(&self) -> Option<String> {
58 match self {
59 Self::InMemoryFallback => Some(format!(
60 "style=filled,fillcolor=\"{}\"",
61 Self::COLOR_IN_MEM_FALLBACK
62 )),
63 Self::MemoryIntensive => Some(format!(
64 "style=filled,fillcolor=\"{}\"",
65 Self::COLOR_MEM_INTENSIVE
66 )),
67 Self::Generic => None,
68 }
69 }
70
71 pub fn legend() -> String {
73 format!(
74 "fontname=\"Helvetica\"\nfontsize=\"10\"\nlabelloc=\"b\"\nlabel=<<BR/><BR/><B>Legend</B><BR/><BR/>◯ streaming engine node <FONT COLOR=\"{}\">⬤</FONT> potentially memory-intensive node <FONT COLOR=\"{}\">⬤</FONT> in-memory engine fallback>",
75 Self::COLOR_MEM_INTENSIVE,
76 Self::COLOR_IN_MEM_FALLBACK,
77 )
78 }
79}
80
81fn escape_graphviz(s: &str) -> String {
82 s.replace('\\', "\\\\")
83 .replace('\n', "\\n")
84 .replace('"', "\\\"")
85}
86
87fn fmt_expr(f: &mut dyn Write, expr: &ExprIR, expr_arena: &Arena<AExpr>) -> std::fmt::Result {
88 let without_alias = ExprIR::from_node(expr.node(), expr_arena);
90 write!(
91 f,
92 "{} = {}",
93 expr.output_name(),
94 without_alias.display(expr_arena)
95 )
96}
97
98pub enum FormatExprStyle {
99 Select,
100 NoAliases,
101}
102
103pub fn fmt_exprs_to_label(
104 exprs: &[ExprIR],
105 expr_arena: &Arena<AExpr>,
106 style: FormatExprStyle,
107) -> String {
108 let mut buffer = String::new();
109 let mut f = EscapeLabel(&mut buffer);
110 fmt_exprs(&mut f, exprs, expr_arena, style);
111 buffer
112}
113
114pub fn fmt_exprs(
115 f: &mut dyn Write,
116 exprs: &[ExprIR],
117 expr_arena: &Arena<AExpr>,
118 style: FormatExprStyle,
119) {
120 if matches!(style, FormatExprStyle::Select) {
121 let mut formatted = Vec::new();
122
123 let mut max_name_width = 0;
124 let mut max_expr_width = 0;
125
126 for e in exprs {
127 let mut name = String::new();
128 let mut expr = String::new();
129
130 let without_alias = ExprIR::from_node(e.node(), expr_arena);
132
133 write!(name, "{}", e.output_name()).unwrap();
134 write!(expr, "{}", without_alias.display(expr_arena)).unwrap();
135
136 max_name_width = max_name_width.max(name.chars().count());
137 max_expr_width = max_expr_width.max(expr.chars().count());
138
139 formatted.push((name, expr));
140 }
141
142 for (name, expr) in formatted {
143 writeln!(f, "{name:>max_name_width$} = {expr:<max_expr_width$}").unwrap();
144 }
145 } else {
146 let Some(e) = exprs.first() else {
147 return;
148 };
149
150 fmt_expr(f, e, expr_arena).unwrap();
151
152 for e in &exprs[1..] {
153 f.write_str("\n").unwrap();
154 fmt_expr(f, e, expr_arena).unwrap();
155 }
156 }
157}
158
159fn fmt_join_label(base_label: &str, left_on: &str, right_on: &str, args: &JoinArgs) -> String {
160 let mut label = base_label.to_string();
161 write!(label, r"\nleft_on:\n{}", left_on).unwrap();
162 write!(label, r"\nright_on:\n{}", right_on).unwrap();
163 if args.how.is_equi() {
164 write!(
165 label,
166 r"\nhow: {}",
167 escape_graphviz(&format!("{:?}", args.how))
168 )
169 .unwrap();
170 }
171 if args.nulls_equal {
172 write!(label, r"\njoin-nulls").unwrap();
173 }
174 label
175}
176
177#[recursive::recursive]
178fn visualize_plan_rec(
179 node_key: PhysNodeKey,
180 phys_sm: &SlotMap<PhysNodeKey, PhysNode>,
181 expr_arena: &Arena<AExpr>,
182 visited: &mut SecondaryMap<PhysNodeKey, ()>,
183 out: &mut Vec<String>,
184) {
185 if visited.contains_key(node_key) {
186 return;
187 }
188 visited.insert(node_key, ());
189
190 let kind = &phys_sm[node_key].kind;
191
192 use std::slice::from_ref;
193 let (label, inputs) = match kind {
194 PhysNodeKind::InMemorySource {
195 df,
196 disable_morsel_split: _,
197 } => (
198 format!(
199 "in-memory-source\\ncols: {}",
200 df.get_column_names_owned().join(", ")
201 ),
202 &[][..],
203 ),
204 #[cfg(feature = "python")]
205 PhysNodeKind::PythonScan { .. } => ("streaming-python-scan".to_string(), &[][..]),
206 PhysNodeKind::SinkMultiple { sinks } => {
207 for sink in sinks {
208 visualize_plan_rec(*sink, phys_sm, expr_arena, visited, out);
209 }
210 return;
211 },
212 PhysNodeKind::Select {
213 input,
214 selectors,
215 extend_original,
216 } => {
217 let label = if *extend_original {
218 "with-columns"
219 } else {
220 "select"
221 };
222 (
223 format!(
224 "{label}\\n{}",
225 fmt_exprs_to_label(selectors, expr_arena, FormatExprStyle::Select)
226 ),
227 from_ref(input),
228 )
229 },
230 PhysNodeKind::WithRowIndex {
231 input,
232 name,
233 offset,
234 } => (
235 format!("with-row-index\\nname: {name}\\noffset: {offset:?}"),
236 from_ref(input),
237 ),
238 PhysNodeKind::InputIndependentSelect { selectors } => (
239 format!(
240 "input-independent-select\\n{}",
241 fmt_exprs_to_label(selectors, expr_arena, FormatExprStyle::Select)
242 ),
243 &[][..],
244 ),
245 PhysNodeKind::Reduce { input, exprs } => (
246 format!(
247 "reduce\\n{}",
248 fmt_exprs_to_label(exprs, expr_arena, FormatExprStyle::Select)
249 ),
250 from_ref(input),
251 ),
252 PhysNodeKind::StreamingSlice {
253 input,
254 offset,
255 length,
256 } => (
257 format!("slice\\noffset: {offset}, length: {length}"),
258 from_ref(input),
259 ),
260 PhysNodeKind::NegativeSlice {
261 input,
262 offset,
263 length,
264 } => (
265 format!("slice\\noffset: {offset}, length: {length}"),
266 from_ref(input),
267 ),
268 PhysNodeKind::DynamicSlice {
269 input,
270 offset,
271 length,
272 } => ("slice".to_owned(), &[*input, *offset, *length][..]),
273 PhysNodeKind::Shift {
274 input,
275 offset,
276 fill: Some(fill),
277 } => ("shift".to_owned(), &[*input, *offset, *fill][..]),
278 PhysNodeKind::Shift {
279 input,
280 offset,
281 fill: None,
282 } => ("shift".to_owned(), &[*input, *offset][..]),
283 PhysNodeKind::Filter { input, predicate } => (
284 format!(
285 "filter\\n{}",
286 fmt_exprs_to_label(from_ref(predicate), expr_arena, FormatExprStyle::Select)
287 ),
288 from_ref(input),
289 ),
290 PhysNodeKind::SimpleProjection { input, columns } => {
291 let mut label = "select".to_string();
292 let mut f = EscapeLabel(&mut label);
293 if columns.iter().all(|(out, col)| out == col) {
294 write!(f, "\n{}", &columns.values().join(", ")).unwrap();
295 } else {
296 for (out, col) in columns {
297 if out == col {
298 write!(f, "\n{col}").unwrap();
299 } else {
300 write!(f, "\n{out} = {col}").unwrap();
301 }
302 }
303 };
304 (label, from_ref(input))
305 },
306 PhysNodeKind::InMemorySink { input } => ("in-memory-sink".to_string(), from_ref(input)),
307 PhysNodeKind::CallbackSink { input, .. } => ("callback-sink".to_string(), from_ref(input)),
308 PhysNodeKind::FileSink { input, options } => match options.file_format {
309 #[cfg(feature = "parquet")]
310 FileWriteFormat::Parquet(_) => ("parquet-sink".to_string(), from_ref(input)),
311 #[cfg(feature = "ipc")]
312 FileWriteFormat::Ipc(_) => ("ipc-sink".to_string(), from_ref(input)),
313 #[cfg(feature = "csv")]
314 FileWriteFormat::Csv(_) => ("csv-sink".to_string(), from_ref(input)),
315 #[cfg(feature = "json")]
316 FileWriteFormat::NDJson(_) => ("ndjson-sink".to_string(), from_ref(input)),
317 },
318 PhysNodeKind::PartitionedSink { input, options } => {
319 let variant = match options.partition_strategy {
320 PartitionStrategyIR::Keyed { .. } => "partition-keyed",
321 PartitionStrategyIR::FileSize => "partition-file-size",
322 };
323
324 match options.file_format {
325 #[cfg(feature = "parquet")]
326 FileWriteFormat::Parquet(_) => (format!("{variant}[parquet]"), from_ref(input)),
327 #[cfg(feature = "ipc")]
328 FileWriteFormat::Ipc(_) => (format!("{variant}[ipc]"), from_ref(input)),
329 #[cfg(feature = "csv")]
330 FileWriteFormat::Csv(_) => (format!("{variant}[csv]"), from_ref(input)),
331 #[cfg(feature = "json")]
332 FileWriteFormat::NDJson(_) => (format!("{variant}[ndjson]"), from_ref(input)),
333 }
334 },
335 PhysNodeKind::InMemoryMap {
336 input,
337 map: _,
338 format_str,
339 } => {
340 let mut label = String::new();
341 label.push_str("in-memory-map");
342 if let Some(format_str) = format_str {
343 label.write_str("\\n").unwrap();
344
345 let mut f = EscapeLabel(&mut label);
346 f.write_str(format_str).unwrap();
347 }
348 (label, from_ref(input))
349 },
350 PhysNodeKind::Map {
351 input,
352 map: _,
353 format_str,
354 } => {
355 let mut label = String::new();
356 label.push_str("map");
357 if let Some(format_str) = format_str {
358 label.push_str("\\n");
359
360 let mut f = EscapeLabel(&mut label);
361 f.write_str(format_str).unwrap();
362 }
363 (label, from_ref(input))
364 },
365 PhysNodeKind::ColumnarFunction {
366 inputs,
367 func: _,
368 arg_map: _,
369 output_name,
370 format_str,
371 } => {
372 let mut label = String::new();
373 label.push_str("columnar-function");
374 if let Some(format_str) = format_str {
375 label.push_str("\\n");
376
377 let mut f = EscapeLabel(&mut label);
378 write!(f, "{output_name} = {format_str}(...)").unwrap();
379 }
380 (label, &inputs[..])
381 },
382 PhysNodeKind::SortedGroupBy {
383 input,
384 key,
385 aggs,
386 slice,
387 } => {
388 let mut s = String::new();
389 s.push_str("sorted-group-by\\n");
390 let f = &mut s;
391 write!(f, "key: {key}\\n").unwrap();
392 if let Some((offset, length)) = slice {
393 write!(f, "slice: {offset}, {length}\\n").unwrap();
394 }
395 write!(
396 f,
397 "aggs:\\n{}",
398 fmt_exprs_to_label(aggs, expr_arena, FormatExprStyle::Select)
399 )
400 .unwrap();
401
402 (s, from_ref(input))
403 },
404 PhysNodeKind::Sort {
405 input,
406 by_column,
407 slice: _,
408 sort_options: _,
409 } => (
410 format!(
411 "sort\\n{}",
412 fmt_exprs_to_label(by_column, expr_arena, FormatExprStyle::NoAliases)
413 ),
414 from_ref(input),
415 ),
416 PhysNodeKind::TopK {
417 input,
418 k,
419 by_column,
420 reverse,
421 nulls_last: _,
422 dyn_pred: _,
423 } => {
424 let name = if reverse.iter().all(|r| *r) {
425 "bottom-k"
426 } else {
427 "top-k"
428 };
429 (
430 format!(
431 "{name}\\n{}",
432 fmt_exprs_to_label(by_column, expr_arena, FormatExprStyle::NoAliases)
433 ),
434 &[*input, *k][..],
435 )
436 },
437 PhysNodeKind::Repeat { value, repeats } => ("repeat".to_owned(), &[*value, *repeats][..]),
438 #[cfg(feature = "cum_agg")]
439 PhysNodeKind::CumAgg { input, kind } => {
440 use crate::nodes::cum_agg::CumAggKind;
441
442 (
443 format!(
444 "cum_{}",
445 match kind {
446 CumAggKind::Min => "min",
447 CumAggKind::Max => "max",
448 CumAggKind::Sum => "sum",
449 CumAggKind::Count => "count",
450 CumAggKind::Prod => "prod",
451 }
452 ),
453 &[*input][..],
454 )
455 },
456 PhysNodeKind::GatherEvery { input, n, offset } => (
457 format!("gather_every\\nn: {n}, offset: {offset}"),
458 &[*input][..],
459 ),
460 PhysNodeKind::ForwardFill { input, limit }
461 | PhysNodeKind::BackwardFill { input, limit } => (
462 {
463 let mut out = if matches!(kind, PhysNodeKind::ForwardFill { .. }) {
464 String::from("forward_fill")
465 } else {
466 String::from("backward_fill")
467 };
468 if let Some(limit) = limit {
469 use std::fmt::Write;
470 writeln!(&mut out).unwrap();
471 write!(&mut out, "limit: {limit}").unwrap();
472 }
473 out
474 },
475 &[*input][..],
476 ),
477 #[cfg(feature = "interpolate")]
478 PhysNodeKind::Interpolate { input, method } => {
479 (format!("interpolate\\nmethod: {method:?}"), &[*input][..])
480 },
481 PhysNodeKind::Rle(input) => ("rle".to_owned(), &[*input][..]),
482 PhysNodeKind::RleId(input) => ("rle_id".to_owned(), &[*input][..]),
483 PhysNodeKind::SortedUnique { input, keys } => {
484 let mut out = String::from("sorted-unique\n");
485 for key in keys.iter() {
486 writeln!(&mut out, "{key}",).unwrap();
487 }
488 (out, &[*input][..])
489 },
490 PhysNodeKind::PeakMinMax { input, is_peak_max } => (
491 if *is_peak_max { "peak_max" } else { "peak_min" }.to_owned(),
492 &[*input][..],
493 ),
494 PhysNodeKind::OrderedUnion { inputs } => ("ordered-union".to_string(), inputs.as_slice()),
495 PhysNodeKind::UnorderedUnion { inputs } => {
496 ("unordered-union".to_string(), inputs.as_slice())
497 },
498 PhysNodeKind::Zip {
499 inputs,
500 zip_behavior,
501 } => {
502 let label = match zip_behavior {
503 ZipBehavior::NullExtend => "zip-null-extend",
504 ZipBehavior::Broadcast => "zip-broadcast",
505 ZipBehavior::Strict => "zip-strict",
506 };
507 (label.to_string(), inputs.as_slice())
508 },
509 PhysNodeKind::Multiplexer { input } => ("multiplexer".to_string(), from_ref(input)),
510 PhysNodeKind::MultiScan {
511 scan_sources,
512 file_reader_builder,
513 cloud_options: _,
514 file_projection_builder,
515 output_schema,
516 row_index,
517 pre_slice,
518 predicate,
519 predicate_file_skip_applied: _,
520 hive_parts,
521 include_file_paths,
522 cast_columns_policy: _,
523 missing_columns_policy: _,
524 forbid_extra_columns: _,
525 deletion_files,
526 table_statistics: _,
527 file_schema: _,
528 disable_morsel_split: _,
529 } => {
530 let mut out = format!("multi-scan[{}]", file_reader_builder.reader_name());
531 let mut f = EscapeLabel(&mut out);
532
533 write!(f, "\n{} source", scan_sources.len()).unwrap();
534
535 if scan_sources.len() != 1 {
536 write!(f, "s").unwrap();
537 }
538
539 write!(
540 f,
541 "\nproject: {} total, {} from file",
542 output_schema.len(),
543 file_projection_builder.num_projections(),
544 )
545 .unwrap();
546
547 if let Some(ri) = row_index {
548 write!(f, "\nrow index: name: {}, offset: {:?}", ri.name, ri.offset).unwrap();
549 }
550
551 if let Some(col_name) = include_file_paths {
552 write!(f, "\nfile path column: {col_name}").unwrap();
553 }
554
555 if let Some(pre_slice) = pre_slice {
556 write!(f, "\nslice: offset: ").unwrap();
557
558 match pre_slice {
559 Slice::Positive { offset, len: _ } => write!(f, "{}", *offset),
560 Slice::Negative {
561 offset_from_end,
562 len: _,
563 } => write!(f, "-{}", *offset_from_end),
564 }
565 .unwrap();
566
567 write!(f, ", len: {}", pre_slice.len()).unwrap()
568 }
569
570 if let Some(predicate) = predicate {
571 write!(f, "\nfilter: {}", predicate.display(expr_arena)).unwrap();
572 }
573
574 if let Some(v) = hive_parts.as_ref().map(|h| h.df().width()) {
575 write!(f, "\nhive: {v} column").unwrap();
576
577 if v != 1 {
578 write!(f, "s").unwrap();
579 }
580 }
581
582 if let Some(deletion_files) = deletion_files {
583 write!(f, "\n{deletion_files}").unwrap();
584 }
585
586 (out, &[][..])
587 },
588 PhysNodeKind::GroupBy {
589 inputs,
590 key_per_input,
591 aggs_per_input,
592 } => {
593 let mut out = String::from("group-by");
594 for (key, aggs) in key_per_input.iter().zip(aggs_per_input) {
595 write!(
596 &mut out,
597 "\\nkey:\\n{}\\naggs:\\n{}",
598 fmt_exprs_to_label(key, expr_arena, FormatExprStyle::Select),
599 fmt_exprs_to_label(aggs, expr_arena, FormatExprStyle::Select)
600 )
601 .ok();
602 }
603 (out, inputs.as_slice())
604 },
605 #[cfg(feature = "dynamic_group_by")]
606 PhysNodeKind::DynamicGroupBy {
607 input,
608 options,
609 aggs,
610 slice,
611 } => {
612 use polars_time::prelude::{Label, StartBy};
613
614 let DynamicGroupOptions {
615 index_column,
616 every,
617 period,
618 offset,
619 label,
620 include_boundaries,
621 closed_window,
622 start_by,
623 } = options;
624 let mut s = String::new();
625 let f = &mut s;
626 f.write_str("dynamic-group-by\\n").unwrap();
627 write!(f, "index column: {index_column}\\n").unwrap();
628 write!(f, "every: {every}").unwrap();
629 if every != period {
630 write!(f, ", period: {period}").unwrap();
631 }
632 if !offset.is_zero() {
633 write!(f, ", offset: {offset}").unwrap();
634 }
635 f.write_str("\\n").unwrap();
636 if *label != Label::Left {
637 write!(f, "label: {}\\n", <&'static str>::from(label)).unwrap();
638 }
639 if *include_boundaries {
640 write!(f, "include_boundaries: true\\n").unwrap();
641 }
642 if *start_by != StartBy::WindowBound {
643 write!(f, "start_by: {}\\n", <&'static str>::from(start_by)).unwrap();
644 }
645 if *closed_window != ClosedWindow::Left {
646 write!(
647 f,
648 "closed_window: {}\\n",
649 <&'static str>::from(closed_window)
650 )
651 .unwrap();
652 }
653 if let Some((offset, length)) = slice {
654 write!(f, "slice: {offset}, {length}\\n").unwrap();
655 }
656 write!(
657 f,
658 "aggs:\\n{}",
659 fmt_exprs_to_label(aggs, expr_arena, FormatExprStyle::Select)
660 )
661 .unwrap();
662
663 (s, from_ref(input))
664 },
665 #[cfg(feature = "dynamic_group_by")]
666 PhysNodeKind::RollingGroupBy {
667 input,
668 index_column,
669 period,
670 offset,
671 closed,
672 slice,
673 aggs,
674 } => {
675 let mut s = String::new();
676 let f = &mut s;
677 f.write_str("rolling-group-by\\n").unwrap();
678 write!(f, "index column: {index_column}\\n").unwrap();
679 write!(f, "period: {period}, offset: {offset}\\n").unwrap();
680 write!(f, "closed: {}\\n", <&'static str>::from(*closed)).unwrap();
681 if let Some((offset, length)) = slice {
682 write!(f, "slice: {offset}, {length}\\n").unwrap();
683 }
684 write!(
685 f,
686 "aggs:\\n{}",
687 fmt_exprs_to_label(aggs, expr_arena, FormatExprStyle::Select)
688 )
689 .unwrap();
690
691 (s, from_ref(input))
692 },
693
694 #[cfg(feature = "is_first_distinct")]
695 PhysNodeKind::IsFirstDistinct {
696 input,
697 out_name,
698 columns,
699 } => {
700 let mut s = String::new();
701 let mut f = EscapeLabel(&mut s);
702 writeln!(f, "is-first-distinct").unwrap();
703 writeln!(f, "key: {}", columns.join(", ")).unwrap();
704 write!(f, "out: {out_name}").unwrap();
705 (s, from_ref(input))
706 },
707 PhysNodeKind::MergeJoin {
708 input_left,
709 input_right,
710 left_on,
711 right_on,
712 args,
713 ..
714 } => {
715 let mut tmp_arena: Arena<AExpr> = Arena::with_capacity(2);
716 let left_on_exprs = left_on
717 .iter()
718 .map(|on| ExprIR::from_column_name(on.clone(), &mut tmp_arena))
719 .collect_vec();
720 let right_on_exprs = right_on
721 .iter()
722 .map(|on| ExprIR::from_column_name(on.clone(), &mut tmp_arena))
723 .collect_vec();
724 let label = fmt_join_label(
725 "merge-join",
726 &fmt_exprs_to_label(&left_on_exprs, &tmp_arena, FormatExprStyle::NoAliases),
727 &fmt_exprs_to_label(&right_on_exprs, &tmp_arena, FormatExprStyle::NoAliases),
728 args,
729 );
730 (label, &[*input_left, *input_right][..])
731 },
732 PhysNodeKind::InMemoryJoin {
733 input_left,
734 input_right,
735 left_on,
736 right_on,
737 args,
738 ..
739 }
740 | PhysNodeKind::EquiJoin {
741 input_left,
742 input_right,
743 left_on,
744 right_on,
745 args,
746 }
747 | PhysNodeKind::SemiAntiJoin {
748 input_left,
749 input_right,
750 left_on,
751 right_on,
752 args,
753 output_bool: _,
754 } => {
755 let base_label = match phys_sm[node_key].kind {
756 PhysNodeKind::MergeJoin { .. } => "merge-join",
757 PhysNodeKind::EquiJoin { .. } => "equi-join",
758 PhysNodeKind::InMemoryJoin { .. } => "in-memory-join",
759 PhysNodeKind::SemiAntiJoin {
760 output_bool: false, ..
761 } if args.how.is_semi() => "semi-join",
762 PhysNodeKind::SemiAntiJoin {
763 output_bool: false, ..
764 } if args.how.is_anti() => "anti-join",
765 PhysNodeKind::SemiAntiJoin {
766 output_bool: true, ..
767 } if args.how.is_semi() => "is-in",
768 PhysNodeKind::SemiAntiJoin {
769 output_bool: true, ..
770 } if args.how.is_anti() => "is-not-in",
771 _ => unreachable!(),
772 };
773 let label = fmt_join_label(
774 base_label,
775 &fmt_exprs_to_label(left_on, expr_arena, FormatExprStyle::NoAliases),
776 &fmt_exprs_to_label(right_on, expr_arena, FormatExprStyle::NoAliases),
777 args,
778 );
779 (label, &[*input_left, *input_right][..])
780 },
781 #[cfg(feature = "iejoin")]
782 PhysNodeKind::RangeJoin {
783 input_left,
784 input_right,
785 left_on,
786 right_on,
787 args,
788 ..
789 } => {
790 let mut tmp_arena: Arena<AExpr> = Arena::with_capacity(3);
791 let left_on_exprs = left_on
792 .iter()
793 .map(|on| ExprIR::from_column_name(on.clone(), &mut tmp_arena))
794 .collect_vec();
795 let right_on_exprs = right_on
796 .iter()
797 .map(|on| ExprIR::from_column_name(on.clone(), &mut tmp_arena))
798 .collect_vec();
799 let label = fmt_join_label(
800 "range-join",
801 &fmt_exprs_to_label(&left_on_exprs, &tmp_arena, FormatExprStyle::NoAliases),
802 &fmt_exprs_to_label(&right_on_exprs, &tmp_arena, FormatExprStyle::NoAliases),
803 args,
804 );
805 (label, &[*input_left, *input_right][..])
806 },
807 PhysNodeKind::CrossJoin {
808 input_left,
809 input_right,
810 args: _,
811 } => ("cross-join".to_string(), &[*input_left, *input_right][..]),
812 PhysNodeKind::AsOfJoin {
813 input_left,
814 input_right,
815 left_on,
816 right_on,
817 args,
818 ..
819 } => {
820 let label = fmt_join_label(
821 "asof-join",
822 &escape_graphviz(&left_on[..]),
823 &escape_graphviz(&right_on[..]),
824 args,
825 );
826 (label, &[*input_left, *input_right][..])
827 },
828 #[cfg(feature = "merge_sorted")]
829 PhysNodeKind::MergeSorted {
830 input_left,
831 input_right,
832 ..
833 } => ("merge-sorted".to_string(), &[*input_left, *input_right][..]),
834 PhysNodeKind::Gather { input, idxs, .. } => ("gather".to_string(), &[*input, *idxs][..]),
835 #[cfg(feature = "ewma")]
836 PhysNodeKind::EwmMean { input, options: _ } => ("ewm-mean".to_string(), &[*input][..]),
837 #[cfg(feature = "ewma")]
838 PhysNodeKind::EwmVar { input, options: _ } => ("ewm-var".to_string(), &[*input][..]),
839 #[cfg(feature = "ewma")]
840 PhysNodeKind::EwmStd { input, options: _ } => ("ewm-std".to_string(), &[*input][..]),
841 #[cfg(any(
842 feature = "dtype-date",
843 feature = "dtype-datetime",
844 feature = "dtype-time"
845 ))]
846 PhysNodeKind::StrptimeInfer {
847 input,
848 dtype,
849 ambiguous_is_raise,
850 ..
851 } => {
852 let mut s = String::new();
853 let mut f = EscapeLabel(&mut s);
854 writeln!(f, "strptime-infer").unwrap();
855 writeln!(f, "dtype: {dtype}").unwrap();
856 let ambiguous = if *ambiguous_is_raise { "raise" } else { "null" };
857 write!(f, "ambiguous: {ambiguous}").unwrap();
858 (s, &[*input][..])
859 },
860 };
861
862 let node_id = node_key.data().as_ffi();
863 let style = NodeStyle::for_node_kind(kind);
864
865 if let Some(attrs) = style.node_attrs() {
866 out.push(format!("{node_id} [label=\"{label}\",{attrs}];"));
867 } else {
868 out.push(format!("{node_id} [label=\"{label}\"];"));
869 }
870 for input in inputs {
871 visualize_plan_rec(input.node, phys_sm, expr_arena, visited, out);
872 out.push(format!(
873 "{} -> {};",
874 input.node.data().as_ffi(),
875 node_key.data().as_ffi()
876 ));
877 }
878}
879
880pub fn visualize_plan(
881 root: PhysNodeKey,
882 phys_sm: &SlotMap<PhysNodeKey, PhysNode>,
883 expr_arena: &Arena<AExpr>,
884) -> String {
885 let mut visited: SecondaryMap<PhysNodeKey, ()> = SecondaryMap::new();
886 let mut out = Vec::with_capacity(phys_sm.len() + 3);
887 out.push("digraph polars {\nrankdir=\"BT\"\nnode [fontname=\"Monospace\"]".to_string());
888 out.push(NodeStyle::legend());
889 visualize_plan_rec(root, phys_sm, expr_arena, &mut visited, &mut out);
890 out.push("}".to_string());
891 out.join("\n")
892}