polars_plan/plans/ir/
format.rs

1use std::fmt::{self, Display, Formatter};
2
3use polars_core::schema::Schema;
4use polars_io::RowIndex;
5use polars_utils::format_list_truncated;
6use polars_utils::slice_enum::Slice;
7use recursive::recursive;
8
9use self::ir::dot::ScanSourcesDisplay;
10use crate::prelude::*;
11
12const INDENT_INCREMENT: usize = 2;
13
14pub struct IRDisplay<'a> {
15    is_streaming: bool,
16    lp: IRPlanRef<'a>,
17}
18
19#[derive(Clone, Copy)]
20pub struct ExprIRDisplay<'a> {
21    pub(crate) node: Node,
22    pub(crate) output_name: &'a OutputName,
23    pub(crate) expr_arena: &'a Arena<AExpr>,
24}
25
26impl<'a> ExprIRDisplay<'a> {
27    pub fn display_node(node: Node, expr_arena: &'a Arena<AExpr>) -> Self {
28        Self {
29            node,
30            output_name: &OutputName::None,
31            expr_arena,
32        }
33    }
34}
35
36/// Utility structure to display several [`ExprIR`]'s in a nice way
37pub(crate) struct ExprIRSliceDisplay<'a, T: AsExpr> {
38    pub(crate) exprs: &'a [T],
39    pub(crate) expr_arena: &'a Arena<AExpr>,
40}
41
42pub(crate) trait AsExpr {
43    fn node(&self) -> Node;
44    fn output_name(&self) -> &OutputName;
45}
46
47impl AsExpr for Node {
48    fn node(&self) -> Node {
49        *self
50    }
51    fn output_name(&self) -> &OutputName {
52        &OutputName::None
53    }
54}
55
56impl AsExpr for ExprIR {
57    fn node(&self) -> Node {
58        self.node()
59    }
60    fn output_name(&self) -> &OutputName {
61        self.output_name_inner()
62    }
63}
64
65#[allow(clippy::too_many_arguments)]
66fn write_scan(
67    f: &mut dyn fmt::Write,
68    name: &str,
69    sources: &ScanSources,
70    indent: usize,
71    n_columns: i64,
72    total_columns: usize,
73    predicate: &Option<ExprIRDisplay<'_>>,
74    pre_slice: Option<Slice>,
75    row_index: Option<&RowIndex>,
76) -> fmt::Result {
77    write!(
78        f,
79        "{:indent$}{name} SCAN {}",
80        "",
81        ScanSourcesDisplay(sources)
82    )?;
83
84    let total_columns = total_columns - usize::from(row_index.is_some());
85    if n_columns > 0 {
86        write!(
87            f,
88            "\n{:indent$}PROJECT {n_columns}/{total_columns} COLUMNS",
89            "",
90        )?;
91    } else {
92        write!(f, "\n{:indent$}PROJECT */{total_columns} COLUMNS", "")?;
93    }
94    if let Some(predicate) = predicate {
95        write!(f, "\n{:indent$}SELECTION: {predicate}", "")?;
96    }
97    if let Some(pre_slice) = pre_slice {
98        write!(f, "\n{:indent$}SLICE: {pre_slice:?}", "")?;
99    }
100    if let Some(row_index) = row_index {
101        write!(f, "\n{:indent$}ROW_INDEX: {}", "", row_index.name)?;
102        if row_index.offset != 0 {
103            write!(f, " (offset: {})", row_index.offset)?;
104        }
105    }
106    Ok(())
107}
108
109impl<'a> IRDisplay<'a> {
110    pub fn new(lp: IRPlanRef<'a>) -> Self {
111        if let Some(streaming_lp) = lp.extract_streaming_plan() {
112            return Self::new_streaming(streaming_lp);
113        }
114
115        Self {
116            is_streaming: false,
117            lp,
118        }
119    }
120
121    fn new_streaming(lp: IRPlanRef<'a>) -> Self {
122        Self {
123            is_streaming: true,
124            lp,
125        }
126    }
127
128    fn root(&self) -> &IR {
129        self.lp.root()
130    }
131
132    fn with_root(&self, root: Node) -> Self {
133        Self {
134            is_streaming: false,
135            lp: self.lp.with_root(root),
136        }
137    }
138
139    fn display_expr(&self, root: &'a ExprIR) -> ExprIRDisplay<'a> {
140        ExprIRDisplay {
141            node: root.node(),
142            output_name: root.output_name_inner(),
143            expr_arena: self.lp.expr_arena,
144        }
145    }
146
147    fn display_expr_slice(&self, exprs: &'a [ExprIR]) -> ExprIRSliceDisplay<'a, ExprIR> {
148        ExprIRSliceDisplay {
149            exprs,
150            expr_arena: self.lp.expr_arena,
151        }
152    }
153
154    #[recursive]
155    fn _format(&self, f: &mut Formatter, indent: usize) -> fmt::Result {
156        let indent = if self.is_streaming {
157            writeln!(f, "{:indent$}STREAMING:", "")?;
158            indent + INDENT_INCREMENT
159        } else {
160            if indent != 0 {
161                writeln!(f)?;
162            }
163            indent
164        };
165
166        let sub_indent = indent + INDENT_INCREMENT;
167        use IR::*;
168
169        let ir_node = self.root();
170        let schema = ir_node.schema(self.lp.lp_arena);
171        let schema = schema.as_ref();
172        match ir_node {
173            Union { inputs, options } => {
174                write_ir_non_recursive(f, ir_node, self.lp.expr_arena, schema, indent)?;
175                let name = if let Some(slice) = options.slice {
176                    format!("SLICED UNION: {slice:?}")
177                } else {
178                    "UNION".to_string()
179                };
180
181                // 3 levels of indentation
182                // - 0 => UNION ... END UNION
183                // - 1 => PLAN 0, PLAN 1, ... PLAN N
184                // - 2 => actual formatting of plans
185                let sub_sub_indent = sub_indent + INDENT_INCREMENT;
186                for (i, plan) in inputs.iter().enumerate() {
187                    write!(f, "\n{:sub_indent$}PLAN {i}:", "")?;
188                    self.with_root(*plan)._format(f, sub_sub_indent)?;
189                }
190                write!(f, "\n{:indent$}END {name}", "")
191            },
192            HConcat { inputs, .. } => {
193                let sub_sub_indent = sub_indent + INDENT_INCREMENT;
194                write_ir_non_recursive(f, ir_node, self.lp.expr_arena, schema, indent)?;
195                for (i, plan) in inputs.iter().enumerate() {
196                    write!(f, "\n{:sub_indent$}PLAN {i}:", "")?;
197                    self.with_root(*plan)._format(f, sub_sub_indent)?;
198                }
199                write!(f, "\n{:indent$}END HCONCAT", "")
200            },
201            GroupBy { input, .. } => {
202                write_ir_non_recursive(f, ir_node, self.lp.expr_arena, schema, indent)?;
203                write!(f, "\n{:sub_indent$}FROM", "")?;
204                self.with_root(*input)._format(f, sub_indent)?;
205                Ok(())
206            },
207            Join {
208                input_left,
209                input_right,
210                left_on,
211                right_on,
212                options,
213                ..
214            } => {
215                let left_on = self.display_expr_slice(left_on);
216                let right_on = self.display_expr_slice(right_on);
217
218                // Fused cross + filter (show as nested loop join)
219                if let Some(JoinTypeOptionsIR::Cross { predicate }) = &options.options {
220                    let predicate = self.display_expr(predicate);
221                    let name = "NESTED LOOP";
222                    write!(f, "{:indent$}{name} JOIN ON {predicate}:", "")?;
223                    write!(f, "\n{:indent$}LEFT PLAN:", "")?;
224                    self.with_root(*input_left)._format(f, sub_indent)?;
225                    write!(f, "\n{:indent$}RIGHT PLAN:", "")?;
226                    self.with_root(*input_right)._format(f, sub_indent)?;
227                    write!(f, "\n{:indent$}END {name} JOIN", "")
228                } else {
229                    let how = &options.args.how;
230                    write!(f, "{:indent$}{how} JOIN:", "")?;
231                    write!(f, "\n{:indent$}LEFT PLAN ON: {left_on}", "")?;
232                    self.with_root(*input_left)._format(f, sub_indent)?;
233                    write!(f, "\n{:indent$}RIGHT PLAN ON: {right_on}", "")?;
234                    self.with_root(*input_right)._format(f, sub_indent)?;
235                    write!(f, "\n{:indent$}END {how} JOIN", "")
236                }
237            },
238            MapFunction {
239                input, function, ..
240            } => {
241                if let Some(streaming_lp) = function.to_streaming_lp() {
242                    IRDisplay::new_streaming(streaming_lp)._format(f, indent)
243                } else {
244                    write_ir_non_recursive(f, ir_node, self.lp.expr_arena, schema, indent)?;
245                    self.with_root(*input)._format(f, sub_indent)
246                }
247            },
248            SinkMultiple { inputs } => {
249                write_ir_non_recursive(f, ir_node, self.lp.expr_arena, schema, indent)?;
250
251                // 3 levels of indentation
252                // - 0 => SINK_MULTIPLE ... END SINK_MULTIPLE
253                // - 1 => PLAN 0, PLAN 1, ... PLAN N
254                // - 2 => actual formatting of plans
255                let sub_sub_indent = sub_indent + 2;
256                for (i, plan) in inputs.iter().enumerate() {
257                    write!(f, "\n{:sub_indent$}PLAN {i}:", "")?;
258                    self.with_root(*plan)._format(f, sub_sub_indent)?;
259                }
260                write!(f, "\n{:indent$}END SINK_MULTIPLE", "")
261            },
262            #[cfg(feature = "merge_sorted")]
263            MergeSorted {
264                input_left,
265                input_right,
266                key: _,
267            } => {
268                write_ir_non_recursive(f, ir_node, self.lp.expr_arena, schema, indent)?;
269                write!(f, ":")?;
270
271                write!(f, "\n{:indent$}LEFT PLAN:", "")?;
272                self.with_root(*input_left)._format(f, sub_indent)?;
273                write!(f, "\n{:indent$}RIGHT PLAN:", "")?;
274                self.with_root(*input_right)._format(f, sub_indent)?;
275                write!(f, "\n{:indent$}END MERGE_SORTED", "")
276            },
277            ir_node => {
278                write_ir_non_recursive(f, ir_node, self.lp.expr_arena, schema, indent)?;
279                for input in ir_node.get_inputs().iter() {
280                    self.with_root(*input)._format(f, sub_indent)?;
281                }
282                Ok(())
283            },
284        }
285    }
286}
287
288impl<'a> ExprIRDisplay<'a> {
289    fn with_slice<T: AsExpr>(&self, exprs: &'a [T]) -> ExprIRSliceDisplay<'a, T> {
290        ExprIRSliceDisplay {
291            exprs,
292            expr_arena: self.expr_arena,
293        }
294    }
295
296    fn with_root<T: AsExpr>(&self, root: &'a T) -> Self {
297        Self {
298            node: root.node(),
299            output_name: root.output_name(),
300            expr_arena: self.expr_arena,
301        }
302    }
303}
304
305impl Display for IRDisplay<'_> {
306    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
307        self._format(f, 0)
308    }
309}
310
311impl fmt::Debug for IRDisplay<'_> {
312    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
313        Display::fmt(&self, f)
314    }
315}
316
317impl<T: AsExpr> Display for ExprIRSliceDisplay<'_, T> {
318    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
319        // Display items in slice delimited by a comma
320
321        use std::fmt::Write;
322
323        let mut iter = self.exprs.iter();
324
325        f.write_char('[')?;
326        if let Some(fst) = iter.next() {
327            let fst = ExprIRDisplay {
328                node: fst.node(),
329                output_name: fst.output_name(),
330                expr_arena: self.expr_arena,
331            };
332            write!(f, "{fst}")?;
333        }
334
335        for expr in iter {
336            let expr = ExprIRDisplay {
337                node: expr.node(),
338                output_name: expr.output_name(),
339                expr_arena: self.expr_arena,
340            };
341            write!(f, ", {expr}")?;
342        }
343
344        f.write_char(']')?;
345
346        Ok(())
347    }
348}
349
350impl<T: AsExpr> fmt::Debug for ExprIRSliceDisplay<'_, T> {
351    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
352        Display::fmt(self, f)
353    }
354}
355
356impl Display for ExprIRDisplay<'_> {
357    #[recursive]
358    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
359        let root = self.expr_arena.get(self.node);
360
361        use AExpr::*;
362        match root {
363            Window {
364                function,
365                partition_by,
366                order_by,
367                options,
368            } => {
369                let function = self.with_root(function);
370                let partition_by = self.with_slice(partition_by);
371                match options {
372                    #[cfg(feature = "dynamic_group_by")]
373                    WindowType::Rolling(options) => {
374                        write!(
375                            f,
376                            "{function}.rolling(by='{}', offset={}, period={})",
377                            options.index_column, options.offset, options.period
378                        )
379                    },
380                    _ => {
381                        if let Some((order_by, _)) = order_by {
382                            let order_by = self.with_root(order_by);
383                            write!(
384                                f,
385                                "{function}.over(partition_by: {partition_by}, order_by: {order_by})"
386                            )
387                        } else {
388                            write!(f, "{function}.over({partition_by})")
389                        }
390                    },
391                }
392            },
393            Len => write!(f, "len()"),
394            Explode { expr, skip_empty } => {
395                let expr = self.with_root(expr);
396                if *skip_empty {
397                    write!(f, "{expr}.explode(skip_empty)")
398                } else {
399                    write!(f, "{expr}.explode()")
400                }
401            },
402            Alias(expr, name) => {
403                let expr = self.with_root(expr);
404                write!(f, "{expr}.alias(\"{name}\")")
405            },
406            Column(name) => write!(f, "col(\"{name}\")"),
407            Literal(v) => write!(f, "{v:?}"),
408            BinaryExpr { left, op, right } => {
409                let left = self.with_root(left);
410                let right = self.with_root(right);
411                write!(f, "[({left}) {op:?} ({right})]")
412            },
413            Sort { expr, options } => {
414                let expr = self.with_root(expr);
415                if options.descending {
416                    write!(f, "{expr}.sort(desc)")
417                } else {
418                    write!(f, "{expr}.sort(asc)")
419                }
420            },
421            SortBy {
422                expr,
423                by,
424                sort_options,
425            } => {
426                let expr = self.with_root(expr);
427                let by = self.with_slice(by);
428                write!(f, "{expr}.sort_by(by={by}, sort_option={sort_options:?})",)
429            },
430            Filter { input, by } => {
431                let input = self.with_root(input);
432                let by = self.with_root(by);
433
434                write!(f, "{input}.filter({by})")
435            },
436            Gather {
437                expr,
438                idx,
439                returns_scalar,
440            } => {
441                let expr = self.with_root(expr);
442                let idx = self.with_root(idx);
443                expr.fmt(f)?;
444
445                if *returns_scalar {
446                    write!(f, ".get({idx})")
447                } else {
448                    write!(f, ".gather({idx})")
449                }
450            },
451            Agg(agg) => {
452                use IRAggExpr::*;
453                match agg {
454                    Min {
455                        input,
456                        propagate_nans,
457                    } => {
458                        self.with_root(input).fmt(f)?;
459                        if *propagate_nans {
460                            write!(f, ".nan_min()")
461                        } else {
462                            write!(f, ".min()")
463                        }
464                    },
465                    Max {
466                        input,
467                        propagate_nans,
468                    } => {
469                        self.with_root(input).fmt(f)?;
470                        if *propagate_nans {
471                            write!(f, ".nan_max()")
472                        } else {
473                            write!(f, ".max()")
474                        }
475                    },
476                    Median(expr) => write!(f, "{}.median()", self.with_root(expr)),
477                    Mean(expr) => write!(f, "{}.mean()", self.with_root(expr)),
478                    First(expr) => write!(f, "{}.first()", self.with_root(expr)),
479                    Last(expr) => write!(f, "{}.last()", self.with_root(expr)),
480                    Implode(expr) => write!(f, "{}.implode()", self.with_root(expr)),
481                    NUnique(expr) => write!(f, "{}.n_unique()", self.with_root(expr)),
482                    Sum(expr) => write!(f, "{}.sum()", self.with_root(expr)),
483                    AggGroups(expr) => write!(f, "{}.groups()", self.with_root(expr)),
484                    Count(expr, _) => write!(f, "{}.count()", self.with_root(expr)),
485                    Var(expr, _) => write!(f, "{}.var()", self.with_root(expr)),
486                    Std(expr, _) => write!(f, "{}.std()", self.with_root(expr)),
487                    Quantile { expr, .. } => write!(f, "{}.quantile()", self.with_root(expr)),
488                }
489            },
490            Cast {
491                expr,
492                dtype,
493                options,
494            } => {
495                self.with_root(expr).fmt(f)?;
496                if options.is_strict() {
497                    write!(f, ".strict_cast({dtype:?})")
498                } else {
499                    write!(f, ".cast({dtype:?})")
500                }
501            },
502            Ternary {
503                predicate,
504                truthy,
505                falsy,
506            } => {
507                let predicate = self.with_root(predicate);
508                let truthy = self.with_root(truthy);
509                let falsy = self.with_root(falsy);
510                write!(f, "when({predicate}).then({truthy}).otherwise({falsy})",)
511            },
512            Function {
513                input, function, ..
514            } => {
515                let fst = self.with_root(&input[0]);
516                fst.fmt(f)?;
517                if input.len() >= 2 {
518                    write!(f, ".{function}({})", self.with_slice(&input[1..]))
519                } else {
520                    write!(f, ".{function}()")
521                }
522            },
523            AnonymousFunction { input, options, .. } => {
524                let fst = self.with_root(&input[0]);
525                fst.fmt(f)?;
526                if input.len() >= 2 {
527                    write!(f, ".{}({})", options.fmt_str, self.with_slice(&input[1..]))
528                } else {
529                    write!(f, ".{}()", options.fmt_str)
530                }
531            },
532            Slice {
533                input,
534                offset,
535                length,
536            } => {
537                let input = self.with_root(input);
538                let offset = self.with_root(offset);
539                let length = self.with_root(length);
540
541                write!(f, "{input}.slice(offset={offset}, length={length})")
542            },
543        }?;
544
545        match self.output_name {
546            OutputName::None => {},
547            OutputName::LiteralLhs(_) => {},
548            OutputName::ColumnLhs(_) => {},
549            #[cfg(feature = "dtype-struct")]
550            OutputName::Field(_) => {},
551            OutputName::Alias(name) => write!(f, r#".alias("{name}")"#)?,
552        }
553
554        Ok(())
555    }
556}
557
558impl fmt::Debug for ExprIRDisplay<'_> {
559    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
560        Display::fmt(self, f)
561    }
562}
563
564pub(crate) struct ColumnsDisplay<'a>(pub(crate) &'a Schema);
565
566impl fmt::Display for ColumnsDisplay<'_> {
567    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
568        let len = self.0.len();
569        let mut iter_names = self.0.iter_names().enumerate();
570
571        const MAX_LEN: usize = 32;
572        const ADD_PER_ITEM: usize = 4;
573
574        let mut current_len = 0;
575
576        if let Some((_, fst)) = iter_names.next() {
577            write!(f, "\"{fst}\"")?;
578
579            current_len += fst.len() + ADD_PER_ITEM;
580        }
581
582        for (i, col) in iter_names {
583            current_len += col.len() + ADD_PER_ITEM;
584
585            if current_len > MAX_LEN {
586                write!(f, ", ... {} other ", len - i)?;
587                if len - i == 1 {
588                    f.write_str("column")?;
589                } else {
590                    f.write_str("columns")?;
591                }
592
593                break;
594            }
595
596            write!(f, ", \"{col}\"")?;
597        }
598
599        Ok(())
600    }
601}
602
603impl fmt::Debug for Operator {
604    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
605        Display::fmt(self, f)
606    }
607}
608
609impl fmt::Debug for LiteralValue {
610    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
611        use LiteralValue::*;
612
613        match self {
614            Self::Scalar(sc) => write!(f, "{}", sc.value()),
615            Self::Series(s) => {
616                let name = s.name();
617                if name.is_empty() {
618                    write!(f, "Series")
619                } else {
620                    write!(f, "Series[{name}]")
621                }
622            },
623            Range(range) => fmt::Debug::fmt(range, f),
624            Dyn(d) => fmt::Debug::fmt(d, f),
625        }
626    }
627}
628
629impl fmt::Debug for DynLiteralValue {
630    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
631        match self {
632            Self::Int(v) => write!(f, "dyn int: {v}"),
633            Self::Float(v) => write!(f, "dyn float: {}", v),
634            Self::Str(v) => write!(f, "dyn str: {v}"),
635            Self::List(_) => todo!(),
636        }
637    }
638}
639
640impl fmt::Debug for RangeLiteralValue {
641    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
642        write!(f, "range({}, {})", self.low, self.high)
643    }
644}
645
646pub fn write_ir_non_recursive(
647    f: &mut dyn fmt::Write,
648    ir: &IR,
649    expr_arena: &Arena<AExpr>,
650    schema: &Schema,
651    indent: usize,
652) -> fmt::Result {
653    match ir {
654        #[cfg(feature = "python")]
655        IR::PythonScan { options } => {
656            let total_columns = options.schema.len();
657            let n_columns = options
658                .with_columns
659                .as_ref()
660                .map(|s| s.len() as i64)
661                .unwrap_or(-1);
662
663            let predicate = match &options.predicate {
664                PythonPredicate::Polars(e) => Some(e.display(expr_arena)),
665                PythonPredicate::PyArrow(_) => None,
666                PythonPredicate::None => None,
667            };
668
669            write_scan(
670                f,
671                "PYTHON",
672                &ScanSources::default(),
673                indent,
674                n_columns,
675                total_columns,
676                &predicate,
677                options
678                    .n_rows
679                    .map(|len| polars_utils::slice_enum::Slice::Positive { offset: 0, len }),
680                None,
681            )
682        },
683        IR::Slice {
684            input: _,
685            offset,
686            len,
687        } => {
688            write!(f, "{:indent$}SLICE[offset: {offset}, len: {len}]", "")
689        },
690        IR::Filter {
691            input: _,
692            predicate,
693        } => {
694            let predicate = predicate.display(expr_arena);
695            // this one is writeln because we don't increase indent (which inserts a line)
696            write!(f, "{:indent$}FILTER {predicate}", "")?;
697            write!(f, "\n{:indent$}FROM", "")
698        },
699        IR::Scan {
700            sources,
701            file_info,
702            predicate,
703            scan_type,
704            unified_scan_args,
705            hive_parts: _,
706            output_schema: _,
707        } => {
708            let n_columns = unified_scan_args
709                .projection
710                .as_ref()
711                .map(|columns| columns.len() as i64)
712                .unwrap_or(-1);
713
714            let predicate = predicate.as_ref().map(|p| p.display(expr_arena));
715
716            write_scan(
717                f,
718                (&**scan_type).into(),
719                sources,
720                indent,
721                n_columns,
722                file_info.schema.len(),
723                &predicate,
724                unified_scan_args.pre_slice.clone(),
725                unified_scan_args.row_index.as_ref(),
726            )
727        },
728        IR::DataFrameScan {
729            df: _,
730            schema,
731            output_schema,
732        } => {
733            let total_columns = schema.len();
734            let (n_columns, projected) = if let Some(schema) = output_schema {
735                (
736                    format!("{}", schema.len()),
737                    format_list_truncated!(schema.iter_names(), 4, '"'),
738                )
739            } else {
740                ("*".to_string(), "".to_string())
741            };
742            write!(
743                f,
744                "{:indent$}DF {}; PROJECT{} {}/{} COLUMNS",
745                "",
746                format_list_truncated!(schema.iter_names(), 4, '"'),
747                projected,
748                n_columns,
749                total_columns,
750            )
751        },
752        IR::SimpleProjection { input: _, columns } => {
753            let num_columns = columns.as_ref().len();
754            let total_columns = schema.len();
755
756            let columns = ColumnsDisplay(columns.as_ref());
757            write!(
758                f,
759                "{:indent$}simple π {num_columns}/{total_columns} [{columns}]",
760                ""
761            )
762        },
763        IR::Select {
764            input: _,
765            expr,
766            schema: _,
767            options: _,
768        } => {
769            // @NOTE: Maybe there should be a clear delimiter here?
770            let exprs = ExprIRSliceDisplay {
771                exprs: expr,
772                expr_arena,
773            };
774            write!(f, "{:indent$}SELECT {exprs}", "")?;
775            Ok(())
776        },
777        IR::Sort {
778            input: _,
779            by_column,
780            slice: _,
781            sort_options: _,
782        } => {
783            let by_column = ExprIRSliceDisplay {
784                exprs: by_column,
785                expr_arena,
786            };
787            write!(f, "{:indent$}SORT BY {by_column}", "")
788        },
789        IR::Cache {
790            input: _,
791            id,
792            cache_hits,
793        } => write!(
794            f,
795            "{:indent$}CACHE[id: {:x}, cache_hits: {}]",
796            "", *id, *cache_hits
797        ),
798        IR::GroupBy {
799            input: _,
800            keys,
801            aggs,
802            schema: _,
803            maintain_order: _,
804            options: _,
805            apply,
806        } => write_group_by(f, indent, expr_arena, keys, aggs, apply.as_deref()),
807        IR::Join {
808            input_left: _,
809            input_right: _,
810            schema: _,
811            left_on,
812            right_on,
813            options,
814        } => {
815            let left_on = ExprIRSliceDisplay {
816                exprs: left_on,
817                expr_arena,
818            };
819            let right_on = ExprIRSliceDisplay {
820                exprs: right_on,
821                expr_arena,
822            };
823
824            // Fused cross + filter (show as nested loop join)
825            if let Some(JoinTypeOptionsIR::Cross { predicate }) = &options.options {
826                let predicate = predicate.display(expr_arena);
827                write!(f, "{:indent$}NESTED_LOOP JOIN ON {predicate}", "")?;
828            } else {
829                let how = &options.args.how;
830                write!(f, "{:indent$}{how} JOIN", "")?;
831                write!(f, "\n{:indent$}LEFT PLAN ON: {left_on}", "")?;
832                write!(f, "\n{:indent$}RIGHT PLAN ON: {right_on}", "")?;
833            }
834
835            Ok(())
836        },
837        IR::HStack {
838            input: _,
839            exprs,
840            schema: _,
841            options: _,
842        } => {
843            // @NOTE: Maybe there should be a clear delimiter here?
844            let exprs = ExprIRSliceDisplay { exprs, expr_arena };
845
846            write!(f, "{:indent$} WITH_COLUMNS:", "",)?;
847            write!(f, "\n{:indent$} {exprs} ", "")
848        },
849        IR::Distinct { input: _, options } => {
850            write!(
851                f,
852                "{:indent$}UNIQUE[maintain_order: {:?}, keep_strategy: {:?}] BY {:?}",
853                "", options.maintain_order, options.keep_strategy, options.subset
854            )
855        },
856        IR::MapFunction { input: _, function } => write!(f, "{:indent$}{function}", ""),
857        IR::Union { inputs: _, options } => {
858            let name = if let Some(slice) = options.slice {
859                format!("SLICED UNION: {slice:?}")
860            } else {
861                "UNION".to_string()
862            };
863            write!(f, "{:indent$}{name}", "")
864        },
865        IR::HConcat {
866            inputs: _,
867            schema: _,
868            options: _,
869        } => write!(f, "{:indent$}HCONCAT", ""),
870        IR::ExtContext {
871            input: _,
872            contexts: _,
873            schema: _,
874        } => write!(f, "{:indent$}EXTERNAL_CONTEXT", ""),
875        IR::Sink { input: _, payload } => {
876            let name = match payload {
877                SinkTypeIR::Memory => "SINK (memory)",
878                SinkTypeIR::File { .. } => "SINK (file)",
879                SinkTypeIR::Partition { .. } => "SINK (partition)",
880            };
881            write!(f, "{:indent$}{name}", "")
882        },
883        IR::SinkMultiple { inputs: _ } => write!(f, "{:indent$}SINK_MULTIPLE", ""),
884        #[cfg(feature = "merge_sorted")]
885        IR::MergeSorted {
886            input_left: _,
887            input_right: _,
888            key,
889        } => write!(f, "{:indent$}MERGE SORTED ON '{key}'", ""),
890        IR::Invalid => write!(f, "{:indent$}INVALID", ""),
891    }
892}
893
894pub fn write_group_by(
895    f: &mut dyn fmt::Write,
896    indent: usize,
897    expr_arena: &Arena<AExpr>,
898    keys: &[ExprIR],
899    aggs: &[ExprIR],
900    apply: Option<&dyn DataFrameUdf>,
901) -> fmt::Result {
902    let sub_indent = indent + INDENT_INCREMENT;
903    let keys = ExprIRSliceDisplay {
904        exprs: keys,
905        expr_arena,
906    };
907    write!(f, "{:indent$}AGGREGATE", "")?;
908    if apply.is_some() {
909        write!(f, "\n{:sub_indent$}MAP_GROUPS BY {keys}", "")?;
910        write!(f, "\n{:sub_indent$}FROM", "")?;
911    } else {
912        let aggs = ExprIRSliceDisplay {
913            exprs: aggs,
914            expr_arena,
915        };
916        write!(f, "\n{:sub_indent$}{aggs} BY {keys}", "")?;
917    }
918
919    Ok(())
920}
polars_plan/plans/ir/format.rs

polars_plan/plans/ir/
format.rs
