From 257fa0baa8e405d1db988480c6f8b0a683814f56 Mon Sep 17 00:00:00 2001 From: Rajiv Ranganath Date: Tue, 31 Oct 2023 10:28:08 +0530 Subject: [PATCH 1/3] partiql-visualize-dot: add Extract code from `partiql-rust-cli/src/visualize` Signed-off-by: Rajiv Ranganath --- Cargo.toml | 2 +- partiql-visualize-dot/Cargo.toml | 28 + partiql-visualize-dot/src/ast_to_dot.rs | 660 +++++++++++++++++++++++ partiql-visualize-dot/src/common.rs | 5 + partiql-visualize-dot/src/lib.rs | 8 + partiql-visualize-dot/src/plan_to_dot.rs | 211 ++++++++ 6 files changed, 913 insertions(+), 1 deletion(-) create mode 100644 partiql-visualize-dot/Cargo.toml create mode 100644 partiql-visualize-dot/src/ast_to_dot.rs create mode 100644 partiql-visualize-dot/src/common.rs create mode 100644 partiql-visualize-dot/src/lib.rs create mode 100644 partiql-visualize-dot/src/plan_to_dot.rs diff --git a/Cargo.toml b/Cargo.toml index 77d41bf1..d8291978 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -25,7 +25,7 @@ members = [ "partiql-rewriter", "partiql-types", "partiql-value", - + "partiql-visualize-dot", "extension/partiql-extension-ion", "extension/partiql-extension-ion-functions", diff --git a/partiql-visualize-dot/Cargo.toml b/partiql-visualize-dot/Cargo.toml new file mode 100644 index 00000000..6a0e6776 --- /dev/null +++ b/partiql-visualize-dot/Cargo.toml @@ -0,0 +1,28 @@ +[package] +name = "partiql-visualize-dot" +description = "Visulize PartiQL AST and Logical Plan in DOT format" +authors.workspace = true +homepage.workspace = true +repository.workspace = true +license = "Apache-2.0" +readme = "../README.md" +keywords = ["sql", "parser", "query", "compilers", "interpreters"] +categories = ["database", "compilers"] +exclude = [ + "**/.git/**", + "**/.github/**", + "**/.travis.yml", + "**/.appveyor.yml", +] +version.workspace = true +edition.workspace = true + +[lib] +bench = false + +[dependencies] +partiql-ast = { path = "../partiql-ast", version = "0.5.*" } +partiql-logical = { path = "../partiql-logical", version = "0.5.*" } + +dot-writer = "0.1.*" +itertools = "0.10.*" \ No newline at end of file diff --git a/partiql-visualize-dot/src/ast_to_dot.rs b/partiql-visualize-dot/src/ast_to_dot.rs new file mode 100644 index 00000000..63b30590 --- /dev/null +++ b/partiql-visualize-dot/src/ast_to_dot.rs @@ -0,0 +1,660 @@ +use crate::common::{ToDotGraph, FG_COLOR}; +use dot_writer::{Attributes, DotWriter, Node, NodeId, Scope, Shape}; +use partiql_ast::ast; + +/* +subgraph cluster_legend { + rank = same; + variable[shape=Mdiamond] + literal[shape=rect] + "node"[shape=ellipse] +} + */ + +trait ScopeExt<'d, 'w> { + fn node_auto_labelled(&mut self, lbl: &str) -> Node<'_, 'w>; + fn cluster_auto_labelled(&mut self, lbl: &str) -> Scope<'_, 'w>; + fn with_cluster(&mut self, lbl: &str, func: F) -> R + where + F: FnMut(Scope<'_, 'w>) -> R; +} + +impl<'d, 'w> ScopeExt<'d, 'w> for Scope<'d, 'w> { + #[inline] + fn node_auto_labelled(&mut self, lbl: &str) -> Node<'_, 'w> { + let mut node = self.node_auto(); + node.set_label(lbl); + node + } + + fn cluster_auto_labelled(&mut self, lbl: &str) -> Scope<'_, 'w> { + let mut cluster = self.cluster(); + cluster.set("label", lbl, lbl.contains(" ")); + cluster + } + + fn with_cluster(&mut self, lbl: &str, mut func: F) -> R + where + F: FnMut(Scope<'_, 'w>) -> R, + { + let cluster = self.cluster_auto_labelled(lbl); + func(cluster) + } +} + +trait ChildEdgeExt { + fn edges(self, out: &mut Scope, from: &NodeId, lbl: &str) -> Targets; +} + +impl ChildEdgeExt for Targets { + fn edges(self, out: &mut Scope, from: &NodeId, lbl: &str) -> Targets { + for target in &self { + out.edge(&from, &target).attributes().set_label(lbl); + } + self + } +} + +type Targets = Vec; + +pub struct AstToDot {} + +impl Default for AstToDot { + fn default() -> Self { + AstToDot {} + } +} + +impl ToDotGraph for AstToDot +where + AstToDot: ToDot, +{ + fn to_graph(mut self, ast: &T) -> String { + let mut output_bytes = Vec::new(); + + { + let mut writer = DotWriter::from(&mut output_bytes); + writer.set_pretty_print(true); + let mut digraph = writer.digraph(); + digraph + .graph_attributes() + .set_rank_direction(dot_writer::RankDirection::TopBottom) + .set("rankdir", "0.05", false) + .set("fontcolor", FG_COLOR, false) + .set("pencolor", FG_COLOR, false); + digraph.node_attributes().set("color", FG_COLOR, false).set( + "fontcolor", + FG_COLOR, + false, + ); + digraph.edge_attributes().set("color", FG_COLOR, false).set( + "fontcolor", + FG_COLOR, + false, + ); + + self.to_dot(&mut digraph, ast); + } + + return String::from_utf8(output_bytes).expect("invalid utf8"); + } +} + +trait ToDot { + fn to_dot(&mut self, out: &mut Scope, ast: &T) -> Targets; +} + +impl ToDot> for AstToDot +where + AstToDot: ToDot, +{ + fn to_dot(&mut self, out: &mut Scope, ast: &Box) -> Targets { + self.to_dot(out, &**ast) + } +} + +impl ToDot> for AstToDot +where + AstToDot: ToDot, +{ + fn to_dot(&mut self, out: &mut Scope, asts: &Vec) -> Targets { + let mut res = Vec::with_capacity(asts.len()); + for ast in asts { + res.extend(self.to_dot(out, &ast)); + } + res + } +} + +impl ToDot> for AstToDot +where + AstToDot: ToDot, +{ + fn to_dot(&mut self, out: &mut Scope, ast: &Option) -> Targets { + match ast { + None => vec![], + Some(ast) => self.to_dot(out, &ast), + } + } +} + +impl ToDot> for AstToDot +where + AstToDot: ToDot, +{ + fn to_dot(&mut self, out: &mut Scope, ast: &ast::AstNode) -> Targets { + self.to_dot(out, &ast.node) + } +} + +impl ToDot for AstToDot { + fn to_dot(&mut self, out: &mut Scope, ast: &ast::Expr) -> Targets { + let mut expr_subgraph = out.subgraph(); + + use ast::Expr; + match &ast { + Expr::Lit(l) => self.to_dot(&mut expr_subgraph, l), + Expr::VarRef(v) => self.to_dot(&mut expr_subgraph, v), + Expr::BinOp(bop) => self.to_dot(&mut expr_subgraph, bop), + Expr::UniOp(unop) => self.to_dot(&mut expr_subgraph, unop), + Expr::Like(like) => self.to_dot(&mut expr_subgraph, like), + Expr::Between(btwn) => self.to_dot(&mut expr_subgraph, btwn), + Expr::In(in_expr) => self.to_dot(&mut expr_subgraph, in_expr), + Expr::Case(_) => todo!(), + Expr::Struct(_) => todo!(), + Expr::Bag(_) => todo!(), + Expr::List(_) => todo!(), + Expr::Sexp(_) => todo!(), + Expr::Path(p) => self.to_dot(&mut expr_subgraph, p), + Expr::Call(c) => self.to_dot(&mut expr_subgraph, c), + Expr::CallAgg(c) => self.to_dot(&mut expr_subgraph, c), + Expr::Query(q) => self.to_dot(&mut expr_subgraph, q), + Expr::Error => todo!(), + } + } +} + +#[inline] +fn lit_to_str(ast: &ast::Lit) -> String { + use ast::Lit; + match ast { + Lit::Null => "NULL".to_string(), + Lit::Missing => "MISSING".to_string(), + Lit::Int8Lit(l) => l.to_string(), + Lit::Int16Lit(l) => l.to_string(), + Lit::Int32Lit(l) => l.to_string(), + Lit::Int64Lit(l) => l.to_string(), + Lit::DecimalLit(l) => l.to_string(), + Lit::NumericLit(l) => l.to_string(), + Lit::RealLit(l) => l.to_string(), + Lit::FloatLit(l) => l.to_string(), + Lit::DoubleLit(l) => l.to_string(), + Lit::BoolLit(l) => (if *l { "TRUE" } else { "FALSE" }).to_string(), + Lit::IonStringLit(l) => format!("`{}`", l), + Lit::CharStringLit(l) => format!("'{}'", l), + Lit::NationalCharStringLit(l) => format!("'{}'", l), + Lit::BitStringLit(l) => format!("b'{}'", l), + Lit::HexStringLit(l) => format!("x'{}'", l), + Lit::BagLit(_b) => todo!("bag literals"), + Lit::ListLit(_b) => todo!("list literals"), + Lit::StructLit(_b) => todo!("struct literals"), + Lit::TypedLit(val_str, ty) => { + format!("{} '{}'", type_to_str(ty), val_str) + } + } +} + +#[inline] +fn custom_type_param_to_str(param: &ast::CustomTypeParam) -> String { + use ast::CustomTypeParam; + match param { + CustomTypeParam::Lit(lit) => lit_to_str(lit), + CustomTypeParam::Type(ty) => type_to_str(ty), + } +} + +#[inline] +fn custom_type_part_to_str(part: &ast::CustomTypePart) -> String { + use ast::CustomTypePart; + match part { + CustomTypePart::Name(name) => symbol_primitive_to_label(name), + CustomTypePart::Parameterized(name, args) => { + let name = symbol_primitive_to_label(name); + let args = args + .iter() + .map(custom_type_param_to_str) + .collect::>() + .join(","); + format!("{}({})", name, args) + } + } +} + +#[inline] +fn type_to_str(ty: &ast::Type) -> String { + use ast::Type; + match ty { + Type::CustomType(cty) => cty + .parts + .iter() + .map(custom_type_part_to_str) + .collect::>() + .join(" "), + _ => format!("{:?}", ty), + } +} + +impl ToDot for AstToDot { + fn to_dot(&mut self, out: &mut Scope, ast: &ast::Lit) -> Targets { + let lbl = lit_to_str(ast); + + let mut node = out.node_auto(); + node.set_label(&lbl).set_shape(Shape::Rectangle); + + vec![node.id()] + } +} + +impl ToDot for AstToDot { + fn to_dot(&mut self, out: &mut Scope, ast: &ast::BinOp) -> Targets { + use ast::BinOpKind; + let lbl = match ast.kind { + BinOpKind::Add => "+", + BinOpKind::Div => "/", + BinOpKind::Exp => "^", + BinOpKind::Mod => "%", + BinOpKind::Mul => "*", + BinOpKind::Sub => "-", + BinOpKind::And => "AND", + BinOpKind::Or => "OR", + BinOpKind::Concat => "||", + BinOpKind::Eq => "=", + BinOpKind::Gt => ">", + BinOpKind::Gte => ">=", + BinOpKind::Lt => "<", + BinOpKind::Lte => "<=", + BinOpKind::Ne => "<>", + BinOpKind::Is => "IS", + }; + let id = out.node_auto_labelled(lbl).id(); + + self.to_dot(out, &ast.lhs).edges(out, &id, ""); + self.to_dot(out, &ast.rhs).edges(out, &id, ""); + + vec![id] + } +} + +impl ToDot for AstToDot { + fn to_dot(&mut self, out: &mut Scope, ast: &ast::UniOp) -> Targets { + use ast::UniOpKind; + let lbl = match ast.kind { + UniOpKind::Pos => "+", + UniOpKind::Neg => "-", + UniOpKind::Not => "NOT", + }; + let id = out.node_auto_labelled(lbl).id(); + + self.to_dot(out, &ast.expr).edges(out, &id, ""); + + vec![id] + } +} + +impl ToDot for AstToDot { + fn to_dot(&mut self, out: &mut Scope, ast: &ast::Like) -> Targets { + let id = out.node_auto_labelled("LIKE").id(); + + self.to_dot(out, &ast.value).edges(out, &id, "value"); + self.to_dot(out, &ast.pattern).edges(out, &id, "pattern"); + self.to_dot(out, &ast.escape).edges(out, &id, "escape"); + + vec![id] + } +} + +impl ToDot for AstToDot { + fn to_dot(&mut self, out: &mut Scope, ast: &ast::Between) -> Targets { + let id = out.node_auto_labelled("BETWEEN").id(); + + self.to_dot(out, &ast.value).edges(out, &id, "value"); + self.to_dot(out, &ast.from).edges(out, &id, "from"); + self.to_dot(out, &ast.to).edges(out, &id, "to"); + + vec![id] + } +} + +impl ToDot for AstToDot { + fn to_dot(&mut self, out: &mut Scope, ast: &ast::In) -> Targets { + let id = out.node_auto_labelled("IN").id(); + + self.to_dot(out, &ast.lhs).edges(out, &id, ""); + self.to_dot(out, &ast.rhs).edges(out, &id, ""); + + vec![id] + } +} + +impl ToDot for AstToDot { + fn to_dot(&mut self, out: &mut Scope, ast: &ast::Query) -> Targets { + let id = out.node_auto_labelled("Query").id(); + + self.to_dot(out, &ast.set).edges(out, &id, ""); + self.to_dot(out, &ast.order_by).edges(out, &id, "order_by"); + self.to_dot(out, &ast.limit_offset) + .edges(out, &id, "limit_offset"); + vec![id] + } +} +impl ToDot for AstToDot { + fn to_dot(&mut self, out: &mut Scope, ast: &ast::QuerySet) -> Targets { + use ast::QuerySet; + match &ast { + QuerySet::BagOp(_) => todo!(), + QuerySet::Select(select) => self.to_dot(out, select), + QuerySet::Expr(e) => self.to_dot(out, e), + QuerySet::Values(_) => todo!(), + QuerySet::Table(_) => todo!(), + } + } +} + +impl ToDot for AstToDot { + fn to_dot(&mut self, out: &mut Scope, ast: &ast::Select) -> Targets { + let id = out.node_auto_labelled("Select").id(); + + out.with_cluster("PROJECT", |mut cl| self.to_dot(&mut cl, &ast.project)) + .edges(out, &id, ""); + out.with_cluster("FROM", |mut cl| self.to_dot(&mut cl, &ast.from)) + .edges(out, &id, ""); + out.with_cluster("FROM LET", |mut cl| self.to_dot(&mut cl, &ast.from_let)) + .edges(out, &id, ""); + out.with_cluster("WHERE", |mut cl| self.to_dot(&mut cl, &ast.where_clause)) + .edges(out, &id, ""); + out.with_cluster("GROUP BY", |mut cl| self.to_dot(&mut cl, &ast.group_by)) + .edges(out, &id, ""); + out.with_cluster("HAVING", |mut cl| self.to_dot(&mut cl, &ast.having)) + .edges(out, &id, ""); + + vec![id] + } +} + +impl ToDot for AstToDot { + fn to_dot(&mut self, out: &mut Scope, ast: &ast::WhereClause) -> Targets { + self.to_dot(out, &ast.expr) + } +} + +impl ToDot for AstToDot { + fn to_dot(&mut self, out: &mut Scope, ast: &ast::HavingClause) -> Targets { + self.to_dot(out, &ast.expr) + } +} + +impl ToDot for AstToDot { + fn to_dot(&mut self, out: &mut Scope, ast: &ast::LimitOffsetClause) -> Targets { + let mut list = self.to_dot(out, &ast.limit); + list.extend(self.to_dot(out, &ast.offset)); + list + } +} + +impl ToDot for AstToDot { + fn to_dot(&mut self, out: &mut Scope, ast: &ast::Projection) -> Targets { + let lbl = match &ast.setq { + Some(ast::SetQuantifier::Distinct) => "Projection | Distinct", + _ => "Projection | All", + }; + let id = out.node_auto_labelled(lbl).id(); + + use ast::ProjectionKind; + let children = { + let mut expr_subgraph = out.subgraph(); + + match &ast.kind { + ProjectionKind::ProjectStar => vec![expr_subgraph.node_auto_labelled("*").id()], + ProjectionKind::ProjectList(items) => { + let mut list = vec![]; + for item in items { + list.extend(self.to_dot(&mut expr_subgraph, item)); + } + list + } + ProjectionKind::ProjectPivot { .. } => todo!(), + ProjectionKind::ProjectValue(_) => todo!(), + } + }; + + children.edges(out, &id, ""); + + vec![id] + } +} + +impl ToDot for AstToDot { + fn to_dot(&mut self, out: &mut Scope, ast: &ast::ProjectItem) -> Targets { + match ast { + ast::ProjectItem::ProjectAll(all) => { + let id = out.node_auto_labelled("ProjectAll").id(); + self.to_dot(out, &all.expr).edges(out, &id, ""); + vec![id] + } + ast::ProjectItem::ProjectExpr(expr) => { + let id = out.node_auto_labelled("ProjectExpr").id(); + self.to_dot(out, &expr.expr).edges(out, &id, ""); + self.to_dot(out, &expr.as_alias).edges(out, &id, "as"); + vec![id] + } + } + } +} + +fn symbol_primitive_to_label(sym: &ast::SymbolPrimitive) -> String { + use ast::CaseSensitivity; + match &sym.case { + CaseSensitivity::CaseSensitive => format!("'{}'", sym.value), + CaseSensitivity::CaseInsensitive => format!("{}", sym.value), + } +} + +impl ToDot for AstToDot { + fn to_dot(&mut self, out: &mut Scope, ast: &ast::SymbolPrimitive) -> Targets { + let lbl = symbol_primitive_to_label(ast); + let id = out.node_auto_labelled(&lbl).id(); + vec![id] + } +} + +impl ToDot for AstToDot { + fn to_dot(&mut self, out: &mut Scope, ast: &ast::VarRef) -> Targets { + let lbl = symbol_primitive_to_label(&ast.name); + let lbl = match &ast.qualifier { + ast::ScopeQualifier::Unqualified => lbl, + ast::ScopeQualifier::Qualified => format!("@{}", lbl), + }; + let id = out.node_auto_labelled(&lbl).id(); + + vec![id] + } +} + +impl ToDot for AstToDot { + fn to_dot(&mut self, _out: &mut Scope, _ast: &ast::OrderByExpr) -> Targets { + todo!("OrderByExpr"); + } +} + +impl ToDot for AstToDot { + fn to_dot(&mut self, _out: &mut Scope, _ast: &ast::GroupByExpr) -> Targets { + todo!("GroupByExpr"); + } +} + +impl ToDot for AstToDot { + fn to_dot(&mut self, out: &mut Scope, ast: &ast::FromClause) -> Targets { + self.to_dot(out, &ast.source) + } +} + +impl ToDot for AstToDot { + fn to_dot(&mut self, out: &mut Scope, ast: &ast::FromSource) -> Targets { + match &ast { + ast::FromSource::FromLet(fl) => self.to_dot(out, fl), + ast::FromSource::Join(j) => self.to_dot(out, j), + } + } +} + +impl ToDot for AstToDot { + fn to_dot(&mut self, out: &mut Scope, ast: &ast::FromLet) -> Targets { + let lbl = match &ast.kind { + ast::FromLetKind::Scan => "Scan", + ast::FromLetKind::Unpivot => "Unpivot", + }; + let id = out.node_auto_labelled(lbl).id(); + + self.to_dot(out, &ast.expr).edges(out, &id, ""); + self.to_dot(out, &ast.as_alias).edges(out, &id, "as"); + self.to_dot(out, &ast.at_alias).edges(out, &id, "at"); + self.to_dot(out, &ast.by_alias).edges(out, &id, "by"); + + vec![id] + } +} + +impl ToDot for AstToDot { + fn to_dot(&mut self, out: &mut Scope, ast: &ast::Join) -> Targets { + let lbl = match &ast.kind { + ast::JoinKind::Inner => "Inner Join", + ast::JoinKind::Left => "Left Join", + ast::JoinKind::Right => "Right Join", + ast::JoinKind::Full => "Full Join", + ast::JoinKind::Cross => "Cross Join", + }; + let id = out.node_auto_labelled(lbl).id(); + + self.to_dot(out, &ast.left).edges(out, &id, "left"); + self.to_dot(out, &ast.right).edges(out, &id, "right"); + self.to_dot(out, &ast.predicate) + .edges(out, &id, "predicate"); + + vec![id] + } +} + +impl ToDot for AstToDot { + fn to_dot(&mut self, out: &mut Scope, ast: &ast::JoinSpec) -> Targets { + match &ast { + ast::JoinSpec::On(fl) => { + let id = out.node_auto_labelled("On").id(); + self.to_dot(out, fl).edges(out, &id, ""); + vec![id] + } + ast::JoinSpec::Using(j) => { + let id = out.node_auto_labelled("Using").id(); + self.to_dot(out, j).edges(out, &id, ""); + vec![id] + } + ast::JoinSpec::Natural => vec![out.node_auto_labelled("Natural").id()], + } + } +} + +impl ToDot for AstToDot { + fn to_dot(&mut self, out: &mut Scope, ast: &ast::Call) -> Targets { + let id = out.node_auto_labelled("Call").id(); + + self.to_dot(out, &ast.func_name).edges(out, &id, "name"); + self.to_dot(out, &ast.args).edges(out, &id, "args"); + + vec![id] + } +} + +impl ToDot for AstToDot { + fn to_dot(&mut self, out: &mut Scope, ast: &ast::CallArg) -> Targets { + use ast::CallArg; + match ast { + ast::CallArg::Star() => vec![out.node_auto_labelled("*").id()], + ast::CallArg::Positional(e) => self.to_dot(out, e), + ast::CallArg::Named(call_arg_named) => { + let id = out.node_auto_labelled("Named").id(); + self.to_dot(out, &call_arg_named.name) + .edges(out, &id, "name"); + self.to_dot(out, &call_arg_named.value) + .edges(out, &id, "value"); + vec![id] + } + CallArg::PositionalType(ty) => { + let mut node = out.node_auto_labelled(&type_to_str(ty)); + node.set("shape", "parallelogram", false); + vec![node.id()] + } + CallArg::NamedType(call_arg_named_type) => { + let id = out.node_auto_labelled("Named").id(); + self.to_dot(out, &call_arg_named_type.name) + .edges(out, &id, "name"); + + let ty_target = { + let mut ty_node = out.node_auto_labelled(&type_to_str(&call_arg_named_type.ty)); + ty_node.set("shape", "parallelogram", false); + vec![ty_node.id()] + }; + ty_target.edges(out, &id, "type"); + + vec![id] + } + } + } +} + +impl ToDot for AstToDot { + fn to_dot(&mut self, out: &mut Scope, ast: &ast::CallAgg) -> Targets { + // Set quantifier is defined in `CallAgg.args` + let id = out.node_auto_labelled("CallAgg").id(); + + self.to_dot(out, &ast.func_name).edges(out, &id, "name"); + self.to_dot(out, &ast.args).edges(out, &id, "args"); + + vec![id] + } +} + +impl ToDot for AstToDot { + fn to_dot(&mut self, out: &mut Scope, ast: &ast::Path) -> Targets { + let id = out.node_auto_labelled("Path").id(); + + self.to_dot(out, &ast.root).edges(out, &id, "root"); + self.to_dot(out, &ast.steps).edges(out, &id, "steps"); + + vec![id] + } +} + +impl ToDot for AstToDot { + fn to_dot(&mut self, out: &mut Scope, ast: &ast::PathStep) -> Targets { + match &ast { + ast::PathStep::PathExpr(e) => self.to_dot(out, e), + ast::PathStep::PathWildCard => vec![out.node_auto_labelled("*").id()], + ast::PathStep::PathUnpivot => vec![out.node_auto_labelled("Unpivot").id()], + } + } +} + +impl ToDot for AstToDot { + fn to_dot(&mut self, out: &mut Scope, ast: &ast::PathExpr) -> Targets { + let id = out.node_auto_labelled("PathExpr").id(); + + self.to_dot(out, &ast.index).edges(out, &id, "index"); + + vec![id] + } +} + +impl ToDot for AstToDot { + fn to_dot(&mut self, _out: &mut Scope, _ast: &ast::Let) -> Targets { + todo!("Let"); + } +} diff --git a/partiql-visualize-dot/src/common.rs b/partiql-visualize-dot/src/common.rs new file mode 100644 index 00000000..1eb41aba --- /dev/null +++ b/partiql-visualize-dot/src/common.rs @@ -0,0 +1,5 @@ +pub trait ToDotGraph { + fn to_graph(self, data: &T) -> String; +} + +pub(crate) const FG_COLOR: &'static str = "\"#839496\""; diff --git a/partiql-visualize-dot/src/lib.rs b/partiql-visualize-dot/src/lib.rs new file mode 100644 index 00000000..7ecc7654 --- /dev/null +++ b/partiql-visualize-dot/src/lib.rs @@ -0,0 +1,8 @@ +mod ast_to_dot; +mod plan_to_dot; + +pub(crate) mod common; + +pub use ast_to_dot::AstToDot; +pub use common::ToDotGraph; +pub use plan_to_dot::PlanToDot; diff --git a/partiql-visualize-dot/src/plan_to_dot.rs b/partiql-visualize-dot/src/plan_to_dot.rs new file mode 100644 index 00000000..4090acb1 --- /dev/null +++ b/partiql-visualize-dot/src/plan_to_dot.rs @@ -0,0 +1,211 @@ +use dot_writer::{Attributes, DotWriter, NodeId, Scope, Shape}; +use itertools::Itertools; +use partiql_logical::{ + AggregateExpression, BinaryOp, BindingsOp, JoinKind, LogicalPlan, ValueExpr, +}; + +use std::collections::HashMap; + +use crate::common::{ToDotGraph, FG_COLOR}; + +pub struct PlanToDot {} + +impl Default for PlanToDot { + fn default() -> Self { + PlanToDot {} + } +} + +impl PlanToDot { + pub(crate) fn to_dot(&self, scope: &mut Scope, plan: &LogicalPlan) { + let mut graph_nodes = HashMap::new(); + for (opid, op) in plan.operators_by_id() { + graph_nodes.insert(opid, self.op_to_dot(scope, op)); + } + + for (src, dst, branch) in plan.flows() { + let src = graph_nodes.get(src).expect("src op"); + let dst = graph_nodes.get(dst).expect("dst op"); + + scope + .edge(src, dst) + .attributes() + .set_label(&branch.to_string()); + } + } + + fn op_to_dot(&self, scope: &mut Scope, op: &BindingsOp) -> NodeId { + let mut node = scope.node_auto(); + let label = match op { + BindingsOp::Scan(s) => { + format!("{{scan | {} | as {} }}", expr_to_str(&s.expr), s.as_key) + } + BindingsOp::Pivot(p) => format!( + "{{pivot | {} | at {} }}", + expr_to_str(&p.value), + expr_to_str(&p.key) + ), + BindingsOp::Unpivot(u) => format!( + "{{unpivot | {} | as {} | at {} }}", + expr_to_str(&u.expr), + &u.as_key, + &u.at_key.as_deref().unwrap_or("") + ), + BindingsOp::Filter(f) => format!("{{filter | {} }}", expr_to_str(&f.expr)), + BindingsOp::OrderBy(o) => { + let specs = o + .specs + .iter() + .map(|s| { + format!( + "{} {:?} NULLS {:?}", + expr_to_str(&s.expr), + s.order, + s.null_order + ) + }) + .join(" | "); + format!("{{order by | {} }}", specs) + } + BindingsOp::LimitOffset(lo) => { + let clauses = [ + lo.limit + .as_ref() + .map(|e| format!("limit {}", expr_to_str(&e))), + lo.offset + .as_ref() + .map(|e| format!("offset {}", expr_to_str(&e))), + ] + .iter() + .filter_map(|o| o.as_ref()) + .join(" | "); + format!("{{ {clauses} }}") + } + BindingsOp::Join(join) => { + let kind = match join.kind { + JoinKind::Inner => "inner", + JoinKind::Left => "left", + JoinKind::Right => "right", + JoinKind::Full => "full", + JoinKind::Cross => "cross", + }; + format!( + "{{ {} join | {} }}", + kind, + join.on + .as_ref() + .map(|e| expr_to_str(e)) + .unwrap_or("".to_string()) + ) + } + BindingsOp::BagOp(_) => "bag op (TODO)".to_string(), + BindingsOp::Project(p) => { + format!( + "{{project | {} }}", + p.exprs + .iter() + .map(|(k, e)| format!("{}:{}", k, expr_to_str(e))) + .join(" | "), + ) + } + BindingsOp::ProjectAll => { + format!("{{project * }}") + } + BindingsOp::ProjectValue(pv) => { + format!("{{project value | {} }}", expr_to_str(&pv.expr)) + } + BindingsOp::ExprQuery(eq) => { + format!("{{ {} }}", expr_to_str(&eq.expr)) + } + BindingsOp::Distinct => "distinct".to_string(), + BindingsOp::GroupBy(g) => { + format!( + "{{group by | {:?} | {{ keys | {{ {} }} }} | {{ aggs | {{ {} }} }} | as {} }}", + g.strategy, + g.exprs + .iter() + .map(|(k, e)| format!("{}:{}", k, expr_to_str(e))) + .join(" | "), + g.aggregate_exprs.iter().map(agg_expr_to_str).join(" | "), + g.group_as_alias.as_deref().unwrap_or(""), + ) + } + BindingsOp::Having(h) => { + format!("{{ having | {} }}", expr_to_str(&h.expr)) + } + BindingsOp::Sink => "sink".to_string(), + }; + node.set_shape(Shape::Mrecord) + .set_label(&format!("{label}")); + + node.id() + } +} + +fn expr_to_str(expr: &ValueExpr) -> String { + match expr { + ValueExpr::BinaryExpr(BinaryOp::And, lhs, rhs) => { + format!( + "{{ AND | {{ {} | {} }} }}", + expr_to_str(lhs), + expr_to_str(rhs) + ) + } + ValueExpr::BinaryExpr(BinaryOp::Or, lhs, rhs) => { + format!( + "{{ OR | {{ {} | {} }} }}", + expr_to_str(lhs), + expr_to_str(rhs) + ) + } + expr => { + let expr: String = format!("{:?}", expr).escape_default().collect(); + let expr = expr.replace('{', "\\{"); + let expr = expr.replace('}', "\\}"); + let expr = expr.replace('<', "\\<"); + let expr = expr.replace('>', "\\>"); + expr + } + } +} + +fn agg_expr_to_str(agg_expr: &AggregateExpression) -> String { + let expr: String = format!("{:?}", agg_expr.expr).escape_default().collect(); + let expr = expr.replace('{', "\\{"); + let expr = expr.replace('}', "\\}"); + format!( + "{}:{}:{:?}:{:?}", + agg_expr.name, expr, agg_expr.func, agg_expr.setq + ) +} + +impl ToDotGraph> for PlanToDot { + fn to_graph(self, plan: &LogicalPlan) -> String { + let mut output_bytes = Vec::new(); + + { + let mut writer = DotWriter::from(&mut output_bytes); + writer.set_pretty_print(true); + let mut digraph = writer.digraph(); + digraph + .graph_attributes() + .set_rank_direction(dot_writer::RankDirection::TopBottom) + .set("fontcolor", FG_COLOR, false) + .set("pencolor", FG_COLOR, false); + digraph.node_attributes().set("color", FG_COLOR, false).set( + "fontcolor", + FG_COLOR, + false, + ); + digraph.edge_attributes().set("color", FG_COLOR, false).set( + "fontcolor", + FG_COLOR, + false, + ); + + self.to_dot(&mut digraph, plan); + } + + return String::from_utf8(output_bytes).expect("invalid utf8"); + } +} From 8cb79659990d82aee360e28719875154302e50fa Mon Sep 17 00:00:00 2001 From: Rajiv Ranganath Date: Mon, 6 Nov 2023 08:53:46 +0530 Subject: [PATCH 2/3] address PR comments Signed-off-by: Rajiv Ranganath --- CHANGELOG.md | 1 + Cargo.toml | 2 +- .../partiql-extension-visualize}/Cargo.toml | 10 +++++----- extension/partiql-extension-visualize/README.md | 5 +++++ .../partiql-extension-visualize}/src/ast_to_dot.rs | 0 .../partiql-extension-visualize}/src/common.rs | 0 .../partiql-extension-visualize}/src/lib.rs | 0 .../partiql-extension-visualize}/src/plan_to_dot.rs | 0 8 files changed, 12 insertions(+), 6 deletions(-) rename {partiql-visualize-dot => extension/partiql-extension-visualize}/Cargo.toml (63%) create mode 100644 extension/partiql-extension-visualize/README.md rename {partiql-visualize-dot => extension/partiql-extension-visualize}/src/ast_to_dot.rs (100%) rename {partiql-visualize-dot => extension/partiql-extension-visualize}/src/common.rs (100%) rename {partiql-visualize-dot => extension/partiql-extension-visualize}/src/lib.rs (100%) rename {partiql-visualize-dot => extension/partiql-extension-visualize}/src/plan_to_dot.rs (100%) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9795d8aa..0188426e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Changed ### Added +- Add `partiql-extension-visualize` for visualizing AST and logical plan ### Fixed diff --git a/Cargo.toml b/Cargo.toml index dc06d712..60a53595 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -25,10 +25,10 @@ members = [ "partiql-rewriter", "partiql-types", "partiql-value", - "partiql-visualize-dot", "extension/partiql-extension-ion", "extension/partiql-extension-ion-functions", + "extension/partiql-extension-visualize", ] [profile.dev.build-override] diff --git a/partiql-visualize-dot/Cargo.toml b/extension/partiql-extension-visualize/Cargo.toml similarity index 63% rename from partiql-visualize-dot/Cargo.toml rename to extension/partiql-extension-visualize/Cargo.toml index 6a0e6776..8acc120c 100644 --- a/partiql-visualize-dot/Cargo.toml +++ b/extension/partiql-extension-visualize/Cargo.toml @@ -1,11 +1,11 @@ [package] -name = "partiql-visualize-dot" -description = "Visulize PartiQL AST and Logical Plan in DOT format" +name = "partiql-extension-visualize" +description = "Visualize PartiQL AST and Logical Plan" authors.workspace = true homepage.workspace = true repository.workspace = true license = "Apache-2.0" -readme = "../README.md" +readme = "../../README.md" keywords = ["sql", "parser", "query", "compilers", "interpreters"] categories = ["database", "compilers"] exclude = [ @@ -21,8 +21,8 @@ edition.workspace = true bench = false [dependencies] -partiql-ast = { path = "../partiql-ast", version = "0.5.*" } -partiql-logical = { path = "../partiql-logical", version = "0.5.*" } +partiql-ast = { path = "../../partiql-ast", version = "0.6.*" } +partiql-logical = { path = "../../partiql-logical", version = "0.6.*" } dot-writer = "0.1.*" itertools = "0.10.*" \ No newline at end of file diff --git a/extension/partiql-extension-visualize/README.md b/extension/partiql-extension-visualize/README.md new file mode 100644 index 00000000..21d11608 --- /dev/null +++ b/extension/partiql-extension-visualize/README.md @@ -0,0 +1,5 @@ +# PartiQL Extension Visualize + +PoC for visualizing AST and Logical Plan. At the moment, it **only** supports visualizing in DOT format. + +It should be considered experimental, subject to change etc. diff --git a/partiql-visualize-dot/src/ast_to_dot.rs b/extension/partiql-extension-visualize/src/ast_to_dot.rs similarity index 100% rename from partiql-visualize-dot/src/ast_to_dot.rs rename to extension/partiql-extension-visualize/src/ast_to_dot.rs diff --git a/partiql-visualize-dot/src/common.rs b/extension/partiql-extension-visualize/src/common.rs similarity index 100% rename from partiql-visualize-dot/src/common.rs rename to extension/partiql-extension-visualize/src/common.rs diff --git a/partiql-visualize-dot/src/lib.rs b/extension/partiql-extension-visualize/src/lib.rs similarity index 100% rename from partiql-visualize-dot/src/lib.rs rename to extension/partiql-extension-visualize/src/lib.rs diff --git a/partiql-visualize-dot/src/plan_to_dot.rs b/extension/partiql-extension-visualize/src/plan_to_dot.rs similarity index 100% rename from partiql-visualize-dot/src/plan_to_dot.rs rename to extension/partiql-extension-visualize/src/plan_to_dot.rs From 2e527b8cc6d4a34697970eb052afa95600518d1f Mon Sep 17 00:00:00 2001 From: Rajiv Ranganath Date: Wed, 8 Nov 2023 10:56:16 +0530 Subject: [PATCH 3/3] address PR comments Signed-off-by: Rajiv Ranganath --- extension/partiql-extension-visualize/Cargo.toml | 11 +++++++++-- extension/partiql-extension-visualize/src/common.rs | 1 + extension/partiql-extension-visualize/src/lib.rs | 9 ++++++++- 3 files changed, 18 insertions(+), 3 deletions(-) diff --git a/extension/partiql-extension-visualize/Cargo.toml b/extension/partiql-extension-visualize/Cargo.toml index 8acc120c..5e755af0 100644 --- a/extension/partiql-extension-visualize/Cargo.toml +++ b/extension/partiql-extension-visualize/Cargo.toml @@ -24,5 +24,12 @@ bench = false partiql-ast = { path = "../../partiql-ast", version = "0.6.*" } partiql-logical = { path = "../../partiql-logical", version = "0.6.*" } -dot-writer = "0.1.*" -itertools = "0.10.*" \ No newline at end of file +dot-writer = { version = "0.1.*", optional = true } +itertools = { version = "0.10.*", optional = true } + +[features] +default = [] +visualize-dot = [ + "dep:dot-writer", + "dep:itertools", +] diff --git a/extension/partiql-extension-visualize/src/common.rs b/extension/partiql-extension-visualize/src/common.rs index 1eb41aba..1a79bc27 100644 --- a/extension/partiql-extension-visualize/src/common.rs +++ b/extension/partiql-extension-visualize/src/common.rs @@ -2,4 +2,5 @@ pub trait ToDotGraph { fn to_graph(self, data: &T) -> String; } +#[cfg(feature = "visualize-dot")] pub(crate) const FG_COLOR: &'static str = "\"#839496\""; diff --git a/extension/partiql-extension-visualize/src/lib.rs b/extension/partiql-extension-visualize/src/lib.rs index 7ecc7654..d22e7d51 100644 --- a/extension/partiql-extension-visualize/src/lib.rs +++ b/extension/partiql-extension-visualize/src/lib.rs @@ -1,8 +1,15 @@ +#[cfg(feature = "visualize-dot")] mod ast_to_dot; + +#[cfg(feature = "visualize-dot")] mod plan_to_dot; pub(crate) mod common; +#[cfg(feature = "visualize-dot")] pub use ast_to_dot::AstToDot; -pub use common::ToDotGraph; + +#[cfg(feature = "visualize-dot")] pub use plan_to_dot::PlanToDot; + +pub use common::ToDotGraph;