diff --git a/crates/polars-plan/src/dsl/meta.rs b/crates/polars-plan/src/dsl/meta.rs index ec608e901bdf..b23c5181cc13 100644 --- a/crates/polars-plan/src/dsl/meta.rs +++ b/crates/polars-plan/src/dsl/meta.rs @@ -1,10 +1,11 @@ -use std::fmt::{Binary, Display}; +use std::fmt::Display; use std::ops::BitAnd; use super::*; use crate::plans::conversion::is_regex_projection; use crate::plans::ir::tree_format::TreeFmtVisitor; use crate::plans::visitor::{AexprNode, TreeWalker}; +use crate::prelude::tree_format::TreeFmtVisitorDisplay; /// Specialized expressions for Categorical dtypes. pub struct MetaNameSpace(pub(crate) Expr); @@ -159,10 +160,13 @@ impl MetaNameSpace { /// Get a hold to an implementor of the `Display` trait that will format as /// the expression as a tree - pub fn into_tree_formatter(self) -> PolarsResult { + pub fn into_tree_formatter(self, display_as_dot: bool) -> PolarsResult { let mut arena = Default::default(); let node = to_aexpr(self.0, &mut arena)?; let mut visitor = TreeFmtVisitor::default(); + if display_as_dot { + visitor.display = TreeFmtVisitorDisplay::DisplayDot; + } AexprNode::new(node).visit(&mut visitor, &arena)?; diff --git a/crates/polars-plan/src/plans/ir/tree_format.rs b/crates/polars-plan/src/plans/ir/tree_format.rs index 6d8d93270c39..d7b44756e662 100644 --- a/crates/polars-plan/src/plans/ir/tree_format.rs +++ b/crates/polars-plan/src/plans/ir/tree_format.rs @@ -382,12 +382,20 @@ impl<'a> TreeFmtNode<'a> { } } +#[derive(Default)] +pub enum TreeFmtVisitorDisplay { + #[default] + DisplayText, + DisplayDot, +} + #[derive(Default)] pub(crate) struct TreeFmtVisitor { levels: Vec>, prev_depth: usize, depth: usize, width: usize, + pub(crate) display: TreeFmtVisitorDisplay, } impl Visitor for TreeFmtVisitor { @@ -868,59 +876,64 @@ impl fmt::Display for Canvas { } } -impl fmt::Display for TreeFmtVisitor { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> std::fmt::Result { - fmt::Debug::fmt(self, f) - } -} +fn tree_fmt_text(tree: &TreeFmtVisitor, f: &mut fmt::Formatter<'_>) -> std::fmt::Result { + let tree_view: TreeView<'_> = tree.levels.as_slice().into(); + let canvas: Canvas = tree_view.into(); + write!(f, "{canvas}")?; -impl fmt::Debug for TreeFmtVisitor { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> std::fmt::Result { - let tree_view: TreeView<'_> = self.levels.as_slice().into(); - let canvas: Canvas = tree_view.into(); - write!(f, "{canvas}")?; - - Ok(()) - } + Ok(()) } // GraphViz Output // Create a simple DOT graph String from TreeFmtVisitor - -// Adding a non-standard display trait is not the right way to do this. -// Not sure how to cleanly communicate this function to the higher level which -// only wants traits. -impl fmt::Binary for TreeFmtVisitor { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> std::fmt::Result { - // Build a dot graph as a string - let tree_view: TreeView<'_> = self.levels.as_slice().into(); - let mut relations: Vec = Vec::new(); - - // Non-empty cells (nodes) and their connections (edges) - for (i, row) in tree_view.matrix.iter().enumerate() { - for (j, cell) in row.iter().enumerate() { - if !cell.text.is_empty() { - // Add node - let node_label = &cell.text.join("\n"); - let node_desc = format!("n{i}{j} [label=\"{node_label}\", ordering=\"out\"]"); - relations.push(node_desc); - - // Add child edges - if i < tree_view.rows.len() - 1 { - // Iter in reversed order to undo the reversed child order when iterating expressions - for child_col in cell.children_columns.iter().rev() { - let next_row = i + 1; - let edge = format!("n{i}{j} -- n{next_row}{child_col}"); - relations.push(edge); - } +fn tree_fmt_dot(tree: &TreeFmtVisitor, f: &mut fmt::Formatter<'_>) -> std::fmt::Result { + // Build a dot graph as a string + let tree_view: TreeView<'_> = tree.levels.as_slice().into(); + let mut relations: Vec = Vec::new(); + + // Non-empty cells (nodes) and their connections (edges) + for (i, row) in tree_view.matrix.iter().enumerate() { + for (j, cell) in row.iter().enumerate() { + if !cell.text.is_empty() { + // Add node + let node_label = &cell.text.join("\n"); + let node_desc = format!("n{i}{j} [label=\"{node_label}\", ordering=\"out\"]"); + relations.push(node_desc); + + // Add child edges + if i < tree_view.rows.len() - 1 { + // Iter in reversed order to undo the reversed child order when iterating expressions + for child_col in cell.children_columns.iter().rev() { + let next_row = i + 1; + let edge = format!("n{i}{j} -- n{next_row}{child_col}"); + relations.push(edge); } } } } + } - let graph_str = relations.join("\n "); - let s = format!("graph {{\n {graph_str}\n}}"); - write!(f, "{s}")?; - Ok(()) + let graph_str = relations.join("\n "); + let s = format!("graph {{\n {graph_str}\n}}"); + write!(f, "{s}")?; + Ok(()) +} + +fn tree_fmt(tree: &TreeFmtVisitor, f: &mut fmt::Formatter<'_>) -> std::fmt::Result { + match tree.display { + TreeFmtVisitorDisplay::DisplayText => tree_fmt_text(tree, f), + TreeFmtVisitorDisplay::DisplayDot => tree_fmt_dot(tree, f), + } +} + +impl fmt::Display for TreeFmtVisitor { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> std::fmt::Result { + tree_fmt(self, f) + } +} + +impl fmt::Debug for TreeFmtVisitor { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> std::fmt::Result { + tree_fmt(self, f) } } diff --git a/crates/polars-python/src/expr/meta.rs b/crates/polars-python/src/expr/meta.rs index 1850319180d7..d0e3a8b3e1df 100644 --- a/crates/polars-python/src/expr/meta.rs +++ b/crates/polars-python/src/expr/meta.rs @@ -102,23 +102,21 @@ impl PyExpr { self.inner.clone().meta()._into_selector().into() } - fn meta_tree_format(&self) -> PyResult { + fn compute_tree_format(&self, display_as_dot: bool) -> Result { let e = self .inner .clone() .meta() - .into_tree_formatter() + .into_tree_formatter(display_as_dot) .map_err(PyPolarsErr::from)?; Ok(format!("{e}")) } + fn meta_tree_format(&self) -> PyResult { + self.compute_tree_format(false) + } + fn meta_show_graph(&self) -> PyResult { - let e = self - .inner - .clone() - .meta() - .into_tree_formatter() - .map_err(PyPolarsErr::from)?; - Ok(format!("{e:b}")) + self.compute_tree_format(true) } }