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

polars_plan/plans/ir/
format.rs
