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::IsSorted { input, .. } => ("is_sorted".to_owned(), &[*input][..]),
495 PhysNodeKind::OrderedUnion { inputs } => ("ordered-union".to_string(), inputs.as_slice()),
496 PhysNodeKind::UnorderedUnion { inputs } => {
497 ("unordered-union".to_string(), inputs.as_slice())
498 },
499 PhysNodeKind::Zip {
500 inputs,
501 zip_behavior,
502 } => {
503 let label = match zip_behavior {
504 ZipBehavior::NullExtend => "zip-null-extend",
505 ZipBehavior::Broadcast => "zip-broadcast",
506 ZipBehavior::Strict => "zip-strict",
507 };
508 (label.to_string(), inputs.as_slice())
509 },
510 PhysNodeKind::Multiplexer { input } => ("multiplexer".to_string(), from_ref(input)),
511 PhysNodeKind::MultiScan {
512 scan_sources,
513 file_reader_builder,
514 cloud_options: _,
515 file_projection_builder,
516 output_schema,
517 row_index,
518 pre_slice,
519 predicate,
520 predicate_file_skip_applied: _,
521 hive_parts,
522 include_file_paths,
523 cast_columns_policy: _,
524 missing_columns_policy: _,
525 forbid_extra_columns: _,
526 deletion_files,
527 table_statistics: _,
528 file_schema: _,
529 disable_morsel_split: _,
530 } => {
531 let mut out = format!("multi-scan[{}]", file_reader_builder.reader_name());
532 let mut f = EscapeLabel(&mut out);
533
534 write!(f, "\n{} source", scan_sources.len()).unwrap();
535
536 if scan_sources.len() != 1 {
537 write!(f, "s").unwrap();
538 }
539
540 write!(
541 f,
542 "\nproject: {} total, {} from file",
543 output_schema.len(),
544 file_projection_builder.num_projections(),
545 )
546 .unwrap();
547
548 if let Some(ri) = row_index {
549 write!(f, "\nrow index: name: {}, offset: {:?}", ri.name, ri.offset).unwrap();
550 }
551
552 if let Some(col_name) = include_file_paths {
553 write!(f, "\nfile path column: {col_name}").unwrap();
554 }
555
556 if let Some(pre_slice) = pre_slice {
557 write!(f, "\nslice: offset: ").unwrap();
558
559 match pre_slice {
560 Slice::Positive { offset, len: _ } => write!(f, "{}", *offset),
561 Slice::Negative {
562 offset_from_end,
563 len: _,
564 } => write!(f, "-{}", *offset_from_end),
565 }
566 .unwrap();
567
568 write!(f, ", len: {}", pre_slice.len()).unwrap()
569 }
570
571 if let Some(predicate) = predicate {
572 write!(f, "\nfilter: {}", predicate.display(expr_arena)).unwrap();
573 }
574
575 if let Some(v) = hive_parts.as_ref().map(|h| h.df().width()) {
576 write!(f, "\nhive: {v} column").unwrap();
577
578 if v != 1 {
579 write!(f, "s").unwrap();
580 }
581 }
582
583 if let Some(deletion_files) = deletion_files {
584 write!(f, "\n{deletion_files}").unwrap();
585 }
586
587 (out, &[][..])
588 },
589 PhysNodeKind::GroupBy {
590 inputs,
591 key_per_input,
592 aggs_per_input,
593 } => {
594 let mut out = String::from("group-by");
595 for (key, aggs) in key_per_input.iter().zip(aggs_per_input) {
596 write!(
597 &mut out,
598 "\\nkey:\\n{}\\naggs:\\n{}",
599 fmt_exprs_to_label(key, expr_arena, FormatExprStyle::Select),
600 fmt_exprs_to_label(aggs, expr_arena, FormatExprStyle::Select)
601 )
602 .ok();
603 }
604 (out, inputs.as_slice())
605 },
606 #[cfg(feature = "dynamic_group_by")]
607 PhysNodeKind::DynamicGroupBy {
608 input,
609 options,
610 aggs,
611 slice,
612 } => {
613 use polars_time::prelude::{Label, StartBy};
614
615 let DynamicGroupOptions {
616 index_column,
617 every,
618 period,
619 offset,
620 label,
621 include_boundaries,
622 closed_window,
623 start_by,
624 } = options;
625 let mut s = String::new();
626 let f = &mut s;
627 f.write_str("dynamic-group-by\\n").unwrap();
628 write!(f, "index column: {index_column}\\n").unwrap();
629 write!(f, "every: {every}").unwrap();
630 if every != period {
631 write!(f, ", period: {period}").unwrap();
632 }
633 if !offset.is_zero() {
634 write!(f, ", offset: {offset}").unwrap();
635 }
636 f.write_str("\\n").unwrap();
637 if *label != Label::Left {
638 write!(f, "label: {}\\n", <&'static str>::from(label)).unwrap();
639 }
640 if *include_boundaries {
641 write!(f, "include_boundaries: true\\n").unwrap();
642 }
643 if *start_by != StartBy::WindowBound {
644 write!(f, "start_by: {}\\n", <&'static str>::from(start_by)).unwrap();
645 }
646 if *closed_window != ClosedWindow::Left {
647 write!(
648 f,
649 "closed_window: {}\\n",
650 <&'static str>::from(closed_window)
651 )
652 .unwrap();
653 }
654 if let Some((offset, length)) = slice {
655 write!(f, "slice: {offset}, {length}\\n").unwrap();
656 }
657 write!(
658 f,
659 "aggs:\\n{}",
660 fmt_exprs_to_label(aggs, expr_arena, FormatExprStyle::Select)
661 )
662 .unwrap();
663
664 (s, from_ref(input))
665 },
666 #[cfg(feature = "dynamic_group_by")]
667 PhysNodeKind::RollingGroupBy {
668 input,
669 index_column,
670 period,
671 offset,
672 closed,
673 slice,
674 aggs,
675 } => {
676 let mut s = String::new();
677 let f = &mut s;
678 f.write_str("rolling-group-by\\n").unwrap();
679 write!(f, "index column: {index_column}\\n").unwrap();
680 write!(f, "period: {period}, offset: {offset}\\n").unwrap();
681 write!(f, "closed: {}\\n", <&'static str>::from(*closed)).unwrap();
682 if let Some((offset, length)) = slice {
683 write!(f, "slice: {offset}, {length}\\n").unwrap();
684 }
685 write!(
686 f,
687 "aggs:\\n{}",
688 fmt_exprs_to_label(aggs, expr_arena, FormatExprStyle::Select)
689 )
690 .unwrap();
691
692 (s, from_ref(input))
693 },
694
695 #[cfg(feature = "is_first_distinct")]
696 PhysNodeKind::IsFirstDistinct {
697 input,
698 out_name,
699 columns,
700 } => {
701 let mut s = String::new();
702 let mut f = EscapeLabel(&mut s);
703 writeln!(f, "is-first-distinct").unwrap();
704 writeln!(f, "key: {}", columns.join(", ")).unwrap();
705 write!(f, "out: {out_name}").unwrap();
706 (s, from_ref(input))
707 },
708 PhysNodeKind::MergeJoin {
709 input_left,
710 input_right,
711 left_on,
712 right_on,
713 args,
714 ..
715 } => {
716 let mut tmp_arena: Arena<AExpr> = Arena::with_capacity(2);
717 let left_on_exprs = left_on
718 .iter()
719 .map(|on| ExprIR::from_column_name(on.clone(), &mut tmp_arena))
720 .collect_vec();
721 let right_on_exprs = right_on
722 .iter()
723 .map(|on| ExprIR::from_column_name(on.clone(), &mut tmp_arena))
724 .collect_vec();
725 let label = fmt_join_label(
726 "merge-join",
727 &fmt_exprs_to_label(&left_on_exprs, &tmp_arena, FormatExprStyle::NoAliases),
728 &fmt_exprs_to_label(&right_on_exprs, &tmp_arena, FormatExprStyle::NoAliases),
729 args,
730 );
731 (label, &[*input_left, *input_right][..])
732 },
733 PhysNodeKind::InMemoryJoin {
734 input_left,
735 input_right,
736 left_on,
737 right_on,
738 args,
739 ..
740 }
741 | PhysNodeKind::EquiJoin {
742 input_left,
743 input_right,
744 left_on,
745 right_on,
746 args,
747 }
748 | PhysNodeKind::SemiAntiJoin {
749 input_left,
750 input_right,
751 left_on,
752 right_on,
753 args,
754 output_bool: _,
755 } => {
756 let base_label = match phys_sm[node_key].kind {
757 PhysNodeKind::MergeJoin { .. } => "merge-join",
758 PhysNodeKind::EquiJoin { .. } => "equi-join",
759 PhysNodeKind::InMemoryJoin { .. } => "in-memory-join",
760 PhysNodeKind::SemiAntiJoin {
761 output_bool: false, ..
762 } if args.how.is_semi() => "semi-join",
763 PhysNodeKind::SemiAntiJoin {
764 output_bool: false, ..
765 } if args.how.is_anti() => "anti-join",
766 PhysNodeKind::SemiAntiJoin {
767 output_bool: true, ..
768 } if args.how.is_semi() => "is-in",
769 PhysNodeKind::SemiAntiJoin {
770 output_bool: true, ..
771 } if args.how.is_anti() => "is-not-in",
772 _ => unreachable!(),
773 };
774 let label = fmt_join_label(
775 base_label,
776 &fmt_exprs_to_label(left_on, expr_arena, FormatExprStyle::NoAliases),
777 &fmt_exprs_to_label(right_on, expr_arena, FormatExprStyle::NoAliases),
778 args,
779 );
780 (label, &[*input_left, *input_right][..])
781 },
782 #[cfg(feature = "iejoin")]
783 PhysNodeKind::RangeJoin {
784 input_left,
785 input_right,
786 left_on,
787 right_on,
788 args,
789 ..
790 } => {
791 let mut tmp_arena: Arena<AExpr> = Arena::with_capacity(3);
792 let left_on_exprs = left_on
793 .iter()
794 .map(|on| ExprIR::from_column_name(on.clone(), &mut tmp_arena))
795 .collect_vec();
796 let right_on_exprs = right_on
797 .iter()
798 .map(|on| ExprIR::from_column_name(on.clone(), &mut tmp_arena))
799 .collect_vec();
800 let label = fmt_join_label(
801 "range-join",
802 &fmt_exprs_to_label(&left_on_exprs, &tmp_arena, FormatExprStyle::NoAliases),
803 &fmt_exprs_to_label(&right_on_exprs, &tmp_arena, FormatExprStyle::NoAliases),
804 args,
805 );
806 (label, &[*input_left, *input_right][..])
807 },
808 PhysNodeKind::CrossJoin {
809 input_left,
810 input_right,
811 args: _,
812 } => ("cross-join".to_string(), &[*input_left, *input_right][..]),
813 PhysNodeKind::AsOfJoin {
814 input_left,
815 input_right,
816 left_on,
817 right_on,
818 args,
819 ..
820 } => {
821 let label = fmt_join_label(
822 "asof-join",
823 &escape_graphviz(&left_on[..]),
824 &escape_graphviz(&right_on[..]),
825 args,
826 );
827 (label, &[*input_left, *input_right][..])
828 },
829 #[cfg(feature = "merge_sorted")]
830 PhysNodeKind::MergeSorted {
831 input_left,
832 input_right,
833 ..
834 } => ("merge-sorted".to_string(), &[*input_left, *input_right][..]),
835 PhysNodeKind::Gather { input, idxs, .. } => ("gather".to_string(), &[*input, *idxs][..]),
836 #[cfg(feature = "ewma")]
837 PhysNodeKind::EwmMean { input, options: _ } => ("ewm-mean".to_string(), &[*input][..]),
838 #[cfg(feature = "ewma")]
839 PhysNodeKind::EwmSum { input, options: _ } => ("ewm-sum".to_string(), &[*input][..]),
840 #[cfg(feature = "ewma")]
841 PhysNodeKind::EwmVar { input, options: _ } => ("ewm-var".to_string(), &[*input][..]),
842 #[cfg(feature = "ewma")]
843 PhysNodeKind::EwmStd { input, options: _ } => ("ewm-std".to_string(), &[*input][..]),
844 #[cfg(any(
845 feature = "dtype-date",
846 feature = "dtype-datetime",
847 feature = "dtype-time"
848 ))]
849 PhysNodeKind::StrptimeInfer {
850 input,
851 dtype,
852 ambiguous_is_raise,
853 ..
854 } => {
855 let mut s = String::new();
856 let mut f = EscapeLabel(&mut s);
857 writeln!(f, "strptime-infer").unwrap();
858 writeln!(f, "dtype: {dtype}").unwrap();
859 let ambiguous = if *ambiguous_is_raise { "raise" } else { "null" };
860 write!(f, "ambiguous: {ambiguous}").unwrap();
861 (s, &[*input][..])
862 },
863 };
864
865 let node_id = node_key.data().as_ffi();
866 let style = NodeStyle::for_node_kind(kind);
867
868 if let Some(attrs) = style.node_attrs() {
869 out.push(format!("{node_id} [label=\"{label}\",{attrs}];"));
870 } else {
871 out.push(format!("{node_id} [label=\"{label}\"];"));
872 }
873 for input in inputs {
874 visualize_plan_rec(input.node, phys_sm, expr_arena, visited, out);
875 out.push(format!(
876 "{} -> {};",
877 input.node.data().as_ffi(),
878 node_key.data().as_ffi()
879 ));
880 }
881}
882
883pub fn visualize_plan(
884 root: PhysNodeKey,
885 phys_sm: &SlotMap<PhysNodeKey, PhysNode>,
886 expr_arena: &Arena<AExpr>,
887) -> String {
888 let mut visited: SecondaryMap<PhysNodeKey, ()> = SecondaryMap::new();
889 let mut out = Vec::with_capacity(phys_sm.len() + 3);
890 out.push("digraph polars {\nrankdir=\"BT\"\nnode [fontname=\"Monospace\"]".to_string());
891 out.push(NodeStyle::legend());
892 visualize_plan_rec(root, phys_sm, expr_arena, &mut visited, &mut out);
893 out.push("}".to_string());
894 out.join("\n")
895}