use super::display::{GraphvizVisitor, IndentVisitor};
use super::expr::{Column, Expr};
use super::extension::UserDefinedLogicalNode;
use crate::datasource::TableProvider;
use crate::error::DataFusionError;
use crate::logical_plan::dfschema::DFSchemaRef;
use crate::sql::parser::FileType;
use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
use std::fmt::Formatter;
use std::{
collections::HashSet,
fmt::{self, Display},
sync::Arc,
};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum JoinType {
Inner,
Left,
Right,
Full,
Semi,
Anti,
}
impl Display for JoinType {
fn fmt(&self, f: &mut Formatter) -> fmt::Result {
let join_type = match self {
JoinType::Inner => "Inner",
JoinType::Left => "Left",
JoinType::Right => "Right",
JoinType::Full => "Full",
JoinType::Semi => "Semi",
JoinType::Anti => "Anti",
};
write!(f, "{}", join_type)
}
}
#[derive(Debug, Clone, Copy)]
pub enum JoinConstraint {
On,
Using,
}
#[derive(Clone)]
pub struct Projection {
pub expr: Vec<Expr>,
pub input: Arc<LogicalPlan>,
pub schema: DFSchemaRef,
pub alias: Option<String>,
}
#[derive(Clone)]
pub struct Filter {
pub predicate: Expr,
pub input: Arc<LogicalPlan>,
}
#[derive(Clone)]
pub struct Window {
pub input: Arc<LogicalPlan>,
pub window_expr: Vec<Expr>,
pub schema: DFSchemaRef,
}
#[derive(Clone)]
pub struct TableScan {
pub table_name: String,
pub source: Arc<dyn TableProvider>,
pub projection: Option<Vec<usize>>,
pub projected_schema: DFSchemaRef,
pub filters: Vec<Expr>,
pub limit: Option<usize>,
}
#[derive(Clone)]
pub struct CrossJoin {
pub left: Arc<LogicalPlan>,
pub right: Arc<LogicalPlan>,
pub schema: DFSchemaRef,
}
#[derive(Clone)]
pub struct Repartition {
pub input: Arc<LogicalPlan>,
pub partitioning_scheme: Partitioning,
}
#[derive(Clone)]
pub struct Union {
pub inputs: Vec<LogicalPlan>,
pub schema: DFSchemaRef,
pub alias: Option<String>,
}
#[derive(Clone)]
pub struct CreateMemoryTable {
pub name: String,
pub input: Arc<LogicalPlan>,
}
#[derive(Clone)]
pub struct CreateExternalTable {
pub schema: DFSchemaRef,
pub name: String,
pub location: String,
pub file_type: FileType,
pub has_header: bool,
}
#[derive(Clone)]
pub struct DropTable {
pub name: String,
pub if_exist: bool,
pub schema: DFSchemaRef,
}
#[derive(Clone)]
pub struct Explain {
pub verbose: bool,
pub plan: Arc<LogicalPlan>,
pub stringified_plans: Vec<StringifiedPlan>,
pub schema: DFSchemaRef,
}
#[derive(Clone)]
pub struct Analyze {
pub verbose: bool,
pub input: Arc<LogicalPlan>,
pub schema: DFSchemaRef,
}
#[derive(Clone)]
pub struct Extension {
pub node: Arc<dyn UserDefinedLogicalNode + Send + Sync>,
}
#[derive(Clone)]
pub struct EmptyRelation {
pub produce_one_row: bool,
pub schema: DFSchemaRef,
}
#[derive(Clone)]
pub struct Limit {
pub n: usize,
pub input: Arc<LogicalPlan>,
}
#[derive(Clone)]
pub struct Values {
pub schema: DFSchemaRef,
pub values: Vec<Vec<Expr>>,
}
#[derive(Clone)]
pub struct Aggregate {
pub input: Arc<LogicalPlan>,
pub group_expr: Vec<Expr>,
pub aggr_expr: Vec<Expr>,
pub schema: DFSchemaRef,
}
#[derive(Clone)]
pub struct Sort {
pub expr: Vec<Expr>,
pub input: Arc<LogicalPlan>,
}
#[derive(Clone)]
pub struct Join {
pub left: Arc<LogicalPlan>,
pub right: Arc<LogicalPlan>,
pub on: Vec<(Column, Column)>,
pub join_type: JoinType,
pub join_constraint: JoinConstraint,
pub schema: DFSchemaRef,
pub null_equals_null: bool,
}
#[derive(Clone)]
pub enum LogicalPlan {
Projection(Projection),
Filter(Filter),
Window(Window),
Aggregate(Aggregate),
Sort(Sort),
Join(Join),
CrossJoin(CrossJoin),
Repartition(Repartition),
Union(Union),
TableScan(TableScan),
EmptyRelation(EmptyRelation),
Limit(Limit),
CreateExternalTable(CreateExternalTable),
CreateMemoryTable(CreateMemoryTable),
DropTable(DropTable),
Values(Values),
Explain(Explain),
Analyze(Analyze),
Extension(Extension),
}
impl LogicalPlan {
pub fn schema(&self) -> &DFSchemaRef {
match self {
LogicalPlan::EmptyRelation(EmptyRelation { schema, .. }) => schema,
LogicalPlan::Values(Values { schema, .. }) => schema,
LogicalPlan::TableScan(TableScan {
projected_schema, ..
}) => projected_schema,
LogicalPlan::Projection(Projection { schema, .. }) => schema,
LogicalPlan::Filter(Filter { input, .. }) => input.schema(),
LogicalPlan::Window(Window { schema, .. }) => schema,
LogicalPlan::Aggregate(Aggregate { schema, .. }) => schema,
LogicalPlan::Sort(Sort { input, .. }) => input.schema(),
LogicalPlan::Join(Join { schema, .. }) => schema,
LogicalPlan::CrossJoin(CrossJoin { schema, .. }) => schema,
LogicalPlan::Repartition(Repartition { input, .. }) => input.schema(),
LogicalPlan::Limit(Limit { input, .. }) => input.schema(),
LogicalPlan::CreateExternalTable(CreateExternalTable { schema, .. }) => {
schema
}
LogicalPlan::Explain(explain) => &explain.schema,
LogicalPlan::Analyze(analyze) => &analyze.schema,
LogicalPlan::Extension(extension) => extension.node.schema(),
LogicalPlan::Union(Union { schema, .. }) => schema,
LogicalPlan::CreateMemoryTable(CreateMemoryTable { input, .. }) => {
input.schema()
}
LogicalPlan::DropTable(DropTable { schema, .. }) => schema,
}
}
pub fn all_schemas(&self) -> Vec<&DFSchemaRef> {
match self {
LogicalPlan::TableScan(TableScan {
projected_schema, ..
}) => vec![projected_schema],
LogicalPlan::Values(Values { schema, .. }) => vec![schema],
LogicalPlan::Window(Window { input, schema, .. })
| LogicalPlan::Projection(Projection { input, schema, .. })
| LogicalPlan::Aggregate(Aggregate { input, schema, .. }) => {
let mut schemas = input.all_schemas();
schemas.insert(0, schema);
schemas
}
LogicalPlan::Join(Join {
left,
right,
schema,
..
})
| LogicalPlan::CrossJoin(CrossJoin {
left,
right,
schema,
}) => {
let mut schemas = left.all_schemas();
schemas.extend(right.all_schemas());
schemas.insert(0, schema);
schemas
}
LogicalPlan::Union(Union { schema, .. }) => {
vec![schema]
}
LogicalPlan::Extension(extension) => vec![extension.node.schema()],
LogicalPlan::Explain(Explain { schema, .. })
| LogicalPlan::Analyze(Analyze { schema, .. })
| LogicalPlan::EmptyRelation(EmptyRelation { schema, .. })
| LogicalPlan::CreateExternalTable(CreateExternalTable { schema, .. }) => {
vec![schema]
}
LogicalPlan::Limit(Limit { input, .. })
| LogicalPlan::Repartition(Repartition { input, .. })
| LogicalPlan::Sort(Sort { input, .. })
| LogicalPlan::CreateMemoryTable(CreateMemoryTable { input, .. })
| LogicalPlan::Filter(Filter { input, .. }) => input.all_schemas(),
LogicalPlan::DropTable(_) => vec![],
}
}
pub fn explain_schema() -> SchemaRef {
SchemaRef::new(Schema::new(vec![
Field::new("plan_type", DataType::Utf8, false),
Field::new("plan", DataType::Utf8, false),
]))
}
pub fn expressions(self: &LogicalPlan) -> Vec<Expr> {
match self {
LogicalPlan::Projection(Projection { expr, .. }) => expr.clone(),
LogicalPlan::Values(Values { values, .. }) => {
values.iter().flatten().cloned().collect()
}
LogicalPlan::Filter(Filter { predicate, .. }) => vec![predicate.clone()],
LogicalPlan::Repartition(Repartition {
partitioning_scheme,
..
}) => match partitioning_scheme {
Partitioning::Hash(expr, _) => expr.clone(),
_ => vec![],
},
LogicalPlan::Window(Window { window_expr, .. }) => window_expr.clone(),
LogicalPlan::Aggregate(Aggregate {
group_expr,
aggr_expr,
..
}) => group_expr.iter().chain(aggr_expr.iter()).cloned().collect(),
LogicalPlan::Join(Join { on, .. }) => on
.iter()
.flat_map(|(l, r)| vec![Expr::Column(l.clone()), Expr::Column(r.clone())])
.collect(),
LogicalPlan::Sort(Sort { expr, .. }) => expr.clone(),
LogicalPlan::Extension(extension) => extension.node.expressions(),
LogicalPlan::TableScan { .. }
| LogicalPlan::EmptyRelation(_)
| LogicalPlan::Limit(_)
| LogicalPlan::CreateExternalTable(_)
| LogicalPlan::CreateMemoryTable(_)
| LogicalPlan::DropTable(_)
| LogicalPlan::CrossJoin(_)
| LogicalPlan::Analyze { .. }
| LogicalPlan::Explain { .. }
| LogicalPlan::Union(_) => {
vec![]
}
}
}
pub fn inputs(self: &LogicalPlan) -> Vec<&LogicalPlan> {
match self {
LogicalPlan::Projection(Projection { input, .. }) => vec![input],
LogicalPlan::Filter(Filter { input, .. }) => vec![input],
LogicalPlan::Repartition(Repartition { input, .. }) => vec![input],
LogicalPlan::Window(Window { input, .. }) => vec![input],
LogicalPlan::Aggregate(Aggregate { input, .. }) => vec![input],
LogicalPlan::Sort(Sort { input, .. }) => vec![input],
LogicalPlan::Join(Join { left, right, .. }) => vec![left, right],
LogicalPlan::CrossJoin(CrossJoin { left, right, .. }) => vec![left, right],
LogicalPlan::Limit(Limit { input, .. }) => vec![input],
LogicalPlan::Extension(extension) => extension.node.inputs(),
LogicalPlan::Union(Union { inputs, .. }) => inputs.iter().collect(),
LogicalPlan::Explain(explain) => vec![&explain.plan],
LogicalPlan::Analyze(analyze) => vec![&analyze.input],
LogicalPlan::CreateMemoryTable(CreateMemoryTable { input, .. }) => {
vec![input]
}
LogicalPlan::TableScan { .. }
| LogicalPlan::EmptyRelation { .. }
| LogicalPlan::Values { .. }
| LogicalPlan::CreateExternalTable(_)
| LogicalPlan::DropTable(_) => vec![],
}
}
pub fn using_columns(&self) -> Result<Vec<HashSet<Column>>, DataFusionError> {
struct UsingJoinColumnVisitor {
using_columns: Vec<HashSet<Column>>,
}
impl PlanVisitor for UsingJoinColumnVisitor {
type Error = DataFusionError;
fn pre_visit(&mut self, plan: &LogicalPlan) -> Result<bool, Self::Error> {
if let LogicalPlan::Join(Join {
join_constraint: JoinConstraint::Using,
on,
..
}) = plan
{
self.using_columns.push(
on.iter()
.map(|entry| [&entry.0, &entry.1])
.flatten()
.cloned()
.collect::<HashSet<Column>>(),
);
}
Ok(true)
}
}
let mut visitor = UsingJoinColumnVisitor {
using_columns: vec![],
};
self.accept(&mut visitor)?;
Ok(visitor.using_columns)
}
}
#[derive(Debug, Clone)]
pub enum Partitioning {
RoundRobinBatch(usize),
Hash(Vec<Expr>, usize),
}
pub trait PlanVisitor {
type Error;
fn pre_visit(&mut self, plan: &LogicalPlan)
-> std::result::Result<bool, Self::Error>;
fn post_visit(
&mut self,
_plan: &LogicalPlan,
) -> std::result::Result<bool, Self::Error> {
Ok(true)
}
}
impl LogicalPlan {
pub fn accept<V>(&self, visitor: &mut V) -> std::result::Result<bool, V::Error>
where
V: PlanVisitor,
{
if !visitor.pre_visit(self)? {
return Ok(false);
}
let recurse = match self {
LogicalPlan::Projection(Projection { input, .. }) => input.accept(visitor)?,
LogicalPlan::Filter(Filter { input, .. }) => input.accept(visitor)?,
LogicalPlan::Repartition(Repartition { input, .. }) => {
input.accept(visitor)?
}
LogicalPlan::Window(Window { input, .. }) => input.accept(visitor)?,
LogicalPlan::Aggregate(Aggregate { input, .. }) => input.accept(visitor)?,
LogicalPlan::Sort(Sort { input, .. }) => input.accept(visitor)?,
LogicalPlan::Join(Join { left, right, .. })
| LogicalPlan::CrossJoin(CrossJoin { left, right, .. }) => {
left.accept(visitor)? && right.accept(visitor)?
}
LogicalPlan::Union(Union { inputs, .. }) => {
for input in inputs {
if !input.accept(visitor)? {
return Ok(false);
}
}
true
}
LogicalPlan::Limit(Limit { input, .. }) => input.accept(visitor)?,
LogicalPlan::CreateMemoryTable(CreateMemoryTable { input, .. }) => {
input.accept(visitor)?
}
LogicalPlan::Extension(extension) => {
for input in extension.node.inputs() {
if !input.accept(visitor)? {
return Ok(false);
}
}
true
}
LogicalPlan::Explain(explain) => explain.plan.accept(visitor)?,
LogicalPlan::Analyze(analyze) => analyze.input.accept(visitor)?,
LogicalPlan::TableScan { .. }
| LogicalPlan::EmptyRelation(_)
| LogicalPlan::Values(_)
| LogicalPlan::CreateExternalTable(_)
| LogicalPlan::DropTable(_) => true,
};
if !recurse {
return Ok(false);
}
if !visitor.post_visit(self)? {
return Ok(false);
}
Ok(true)
}
}
impl LogicalPlan {
pub fn display_indent(&self) -> impl fmt::Display + '_ {
struct Wrapper<'a>(&'a LogicalPlan);
impl<'a> fmt::Display for Wrapper<'a> {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
let with_schema = false;
let mut visitor = IndentVisitor::new(f, with_schema);
self.0.accept(&mut visitor).unwrap();
Ok(())
}
}
Wrapper(self)
}
pub fn display_indent_schema(&self) -> impl fmt::Display + '_ {
struct Wrapper<'a>(&'a LogicalPlan);
impl<'a> fmt::Display for Wrapper<'a> {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
let with_schema = true;
let mut visitor = IndentVisitor::new(f, with_schema);
self.0.accept(&mut visitor).unwrap();
Ok(())
}
}
Wrapper(self)
}
pub fn display_graphviz(&self) -> impl fmt::Display + '_ {
struct Wrapper<'a>(&'a LogicalPlan);
impl<'a> fmt::Display for Wrapper<'a> {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
writeln!(
f,
"// Begin DataFusion GraphViz Plan (see https://graphviz.org)"
)?;
writeln!(f, "digraph {{")?;
let mut visitor = GraphvizVisitor::new(f);
visitor.pre_visit_plan("LogicalPlan")?;
self.0.accept(&mut visitor).unwrap();
visitor.post_visit_plan()?;
visitor.set_with_schema(true);
visitor.pre_visit_plan("Detailed LogicalPlan")?;
self.0.accept(&mut visitor).unwrap();
visitor.post_visit_plan()?;
writeln!(f, "}}")?;
writeln!(f, "// End DataFusion GraphViz Plan")?;
Ok(())
}
}
Wrapper(self)
}
pub fn display(&self) -> impl fmt::Display + '_ {
struct Wrapper<'a>(&'a LogicalPlan);
impl<'a> fmt::Display for Wrapper<'a> {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
match &*self.0 {
LogicalPlan::EmptyRelation(_) => write!(f, "EmptyRelation"),
LogicalPlan::Values(Values { ref values, .. }) => {
let str_values: Vec<_> = values
.iter()
.take(5)
.map(|row| {
let item = row
.iter()
.map(|expr| expr.to_string())
.collect::<Vec<_>>()
.join(", ");
format!("({})", item)
})
.collect();
let elipse = if values.len() > 5 { "..." } else { "" };
write!(f, "Values: {}{}", str_values.join(", "), elipse)
}
LogicalPlan::TableScan(TableScan {
ref table_name,
ref projection,
ref filters,
ref limit,
..
}) => {
write!(
f,
"TableScan: {} projection={:?}",
table_name, projection
)?;
if !filters.is_empty() {
write!(f, ", filters={:?}", filters)?;
}
if let Some(n) = limit {
write!(f, ", limit={}", n)?;
}
Ok(())
}
LogicalPlan::Projection(Projection {
ref expr, alias, ..
}) => {
write!(f, "Projection: ")?;
for (i, expr_item) in expr.iter().enumerate() {
if i > 0 {
write!(f, ", ")?;
}
write!(f, "{:?}", expr_item)?;
}
if let Some(a) = alias {
write!(f, ", alias={}", a)?;
}
Ok(())
}
LogicalPlan::Filter(Filter {
predicate: ref expr,
..
}) => write!(f, "Filter: {:?}", expr),
LogicalPlan::Window(Window {
ref window_expr, ..
}) => {
write!(f, "WindowAggr: windowExpr=[{:?}]", window_expr)
}
LogicalPlan::Aggregate(Aggregate {
ref group_expr,
ref aggr_expr,
..
}) => write!(
f,
"Aggregate: groupBy=[{:?}], aggr=[{:?}]",
group_expr, aggr_expr
),
LogicalPlan::Sort(Sort { expr, .. }) => {
write!(f, "Sort: ")?;
for (i, expr_item) in expr.iter().enumerate() {
if i > 0 {
write!(f, ", ")?;
}
write!(f, "{:?}", expr_item)?;
}
Ok(())
}
LogicalPlan::Join(Join {
on: ref keys,
join_constraint,
join_type,
..
}) => {
let join_expr: Vec<String> =
keys.iter().map(|(l, r)| format!("{} = {}", l, r)).collect();
match join_constraint {
JoinConstraint::On => {
write!(f, "{} Join: {}", join_type, join_expr.join(", "))
}
JoinConstraint::Using => {
write!(
f,
"{} Join: Using {}",
join_type,
join_expr.join(", ")
)
}
}
}
LogicalPlan::CrossJoin(_) => {
write!(f, "CrossJoin:")
}
LogicalPlan::Repartition(Repartition {
partitioning_scheme,
..
}) => match partitioning_scheme {
Partitioning::RoundRobinBatch(n) => write!(
f,
"Repartition: RoundRobinBatch partition_count={}",
n
),
Partitioning::Hash(expr, n) => {
let hash_expr: Vec<String> =
expr.iter().map(|e| format!("{:?}", e)).collect();
write!(
f,
"Repartition: Hash({}) partition_count={}",
hash_expr.join(", "),
n
)
}
},
LogicalPlan::Limit(Limit { ref n, .. }) => write!(f, "Limit: {}", n),
LogicalPlan::CreateExternalTable(CreateExternalTable {
ref name,
..
}) => {
write!(f, "CreateExternalTable: {:?}", name)
}
LogicalPlan::CreateMemoryTable(CreateMemoryTable {
name, ..
}) => {
write!(f, "CreateMemoryTable: {:?}", name)
}
LogicalPlan::DropTable(DropTable { name, if_exist, .. }) => {
write!(f, "DropTable: {:?} if not exist:={}", name, if_exist)
}
LogicalPlan::Explain { .. } => write!(f, "Explain"),
LogicalPlan::Analyze { .. } => write!(f, "Analyze"),
LogicalPlan::Union(_) => write!(f, "Union"),
LogicalPlan::Extension(e) => e.node.fmt_for_explain(f),
}
}
}
Wrapper(self)
}
}
impl fmt::Debug for LogicalPlan {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
self.display_indent().fmt(f)
}
}
#[derive(Debug, Clone, PartialEq)]
pub enum PlanType {
InitialLogicalPlan,
OptimizedLogicalPlan {
optimizer_name: String,
},
FinalLogicalPlan,
InitialPhysicalPlan,
OptimizedPhysicalPlan {
optimizer_name: String,
},
FinalPhysicalPlan,
}
impl fmt::Display for PlanType {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
match self {
PlanType::InitialLogicalPlan => write!(f, "initial_logical_plan"),
PlanType::OptimizedLogicalPlan { optimizer_name } => {
write!(f, "logical_plan after {}", optimizer_name)
}
PlanType::FinalLogicalPlan => write!(f, "logical_plan"),
PlanType::InitialPhysicalPlan => write!(f, "initial_physical_plan"),
PlanType::OptimizedPhysicalPlan { optimizer_name } => {
write!(f, "physical_plan after {}", optimizer_name)
}
PlanType::FinalPhysicalPlan => write!(f, "physical_plan"),
}
}
}
#[derive(Debug, Clone, PartialEq)]
#[allow(clippy::rc_buffer)]
pub struct StringifiedPlan {
pub plan_type: PlanType,
pub plan: Arc<String>,
}
impl StringifiedPlan {
pub fn new(plan_type: PlanType, plan: impl Into<String>) -> Self {
StringifiedPlan {
plan_type,
plan: Arc::new(plan.into()),
}
}
pub fn should_display(&self, verbose_mode: bool) -> bool {
match self.plan_type {
PlanType::FinalLogicalPlan | PlanType::FinalPhysicalPlan => true,
_ => verbose_mode,
}
}
}
pub trait ToStringifiedPlan {
fn to_stringified(&self, plan_type: PlanType) -> StringifiedPlan;
}
impl ToStringifiedPlan for LogicalPlan {
fn to_stringified(&self, plan_type: PlanType) -> StringifiedPlan {
StringifiedPlan::new(plan_type, self.display_indent().to_string())
}
}
#[cfg(test)]
mod tests {
use super::super::{col, lit, LogicalPlanBuilder};
use super::*;
fn employee_schema() -> Schema {
Schema::new(vec![
Field::new("id", DataType::Int32, false),
Field::new("first_name", DataType::Utf8, false),
Field::new("last_name", DataType::Utf8, false),
Field::new("state", DataType::Utf8, false),
Field::new("salary", DataType::Int32, false),
])
}
fn display_plan() -> LogicalPlan {
LogicalPlanBuilder::scan_empty(
Some("employee_csv"),
&employee_schema(),
Some(vec![0, 3]),
)
.unwrap()
.filter(col("state").eq(lit("CO")))
.unwrap()
.project(vec![col("id")])
.unwrap()
.build()
.unwrap()
}
#[test]
fn test_display_indent() {
let plan = display_plan();
let expected = "Projection: #employee_csv.id\
\n Filter: #employee_csv.state = Utf8(\"CO\")\
\n TableScan: employee_csv projection=Some([0, 3])";
assert_eq!(expected, format!("{}", plan.display_indent()));
}
#[test]
fn test_display_indent_schema() {
let plan = display_plan();
let expected = "Projection: #employee_csv.id [id:Int32]\
\n Filter: #employee_csv.state = Utf8(\"CO\") [id:Int32, state:Utf8]\
\n TableScan: employee_csv projection=Some([0, 3]) [id:Int32, state:Utf8]";
assert_eq!(expected, format!("{}", plan.display_indent_schema()));
}
#[test]
fn test_display_graphviz() {
let plan = display_plan();
let graphviz = format!("{}", plan.display_graphviz());
assert!(
graphviz.contains(
r#"// Begin DataFusion GraphViz Plan (see https://graphviz.org)"#
),
"\n{}",
plan.display_graphviz()
);
assert!(
graphviz.contains(
r#"[shape=box label="TableScan: employee_csv projection=Some([0, 3])"]"#
),
"\n{}",
plan.display_graphviz()
);
assert!(graphviz.contains(r#"[shape=box label="TableScan: employee_csv projection=Some([0, 3])\nSchema: [id:Int32, state:Utf8]"]"#),
"\n{}", plan.display_graphviz());
assert!(
graphviz.contains(r#"// End DataFusion GraphViz Plan"#),
"\n{}",
plan.display_graphviz()
);
}
#[derive(Debug, Default)]
struct OkVisitor {
strings: Vec<String>,
}
impl PlanVisitor for OkVisitor {
type Error = String;
fn pre_visit(
&mut self,
plan: &LogicalPlan,
) -> std::result::Result<bool, Self::Error> {
let s = match plan {
LogicalPlan::Projection { .. } => "pre_visit Projection",
LogicalPlan::Filter { .. } => "pre_visit Filter",
LogicalPlan::TableScan { .. } => "pre_visit TableScan",
_ => unimplemented!("unknown plan type"),
};
self.strings.push(s.into());
Ok(true)
}
fn post_visit(
&mut self,
plan: &LogicalPlan,
) -> std::result::Result<bool, Self::Error> {
let s = match plan {
LogicalPlan::Projection { .. } => "post_visit Projection",
LogicalPlan::Filter { .. } => "post_visit Filter",
LogicalPlan::TableScan { .. } => "post_visit TableScan",
_ => unimplemented!("unknown plan type"),
};
self.strings.push(s.into());
Ok(true)
}
}
#[test]
fn visit_order() {
let mut visitor = OkVisitor::default();
let plan = test_plan();
let res = plan.accept(&mut visitor);
assert!(res.is_ok());
assert_eq!(
visitor.strings,
vec![
"pre_visit Projection",
"pre_visit Filter",
"pre_visit TableScan",
"post_visit TableScan",
"post_visit Filter",
"post_visit Projection"
]
);
}
#[derive(Debug, Default)]
struct OptionalCounter {
val: Option<usize>,
}
impl OptionalCounter {
fn new(val: usize) -> Self {
Self { val: Some(val) }
}
fn dec(&mut self) -> bool {
if Some(0) == self.val {
true
} else {
self.val = self.val.take().map(|i| i - 1);
false
}
}
}
#[derive(Debug, Default)]
struct StoppingVisitor {
inner: OkVisitor,
return_false_from_pre_in: OptionalCounter,
return_false_from_post_in: OptionalCounter,
}
impl PlanVisitor for StoppingVisitor {
type Error = String;
fn pre_visit(
&mut self,
plan: &LogicalPlan,
) -> std::result::Result<bool, Self::Error> {
if self.return_false_from_pre_in.dec() {
return Ok(false);
}
self.inner.pre_visit(plan)
}
fn post_visit(
&mut self,
plan: &LogicalPlan,
) -> std::result::Result<bool, Self::Error> {
if self.return_false_from_post_in.dec() {
return Ok(false);
}
self.inner.post_visit(plan)
}
}
#[test]
fn early_stopping_pre_visit() {
let mut visitor = StoppingVisitor {
return_false_from_pre_in: OptionalCounter::new(2),
..Default::default()
};
let plan = test_plan();
let res = plan.accept(&mut visitor);
assert!(res.is_ok());
assert_eq!(
visitor.inner.strings,
vec!["pre_visit Projection", "pre_visit Filter",]
);
}
#[test]
fn early_stopping_post_visit() {
let mut visitor = StoppingVisitor {
return_false_from_post_in: OptionalCounter::new(1),
..Default::default()
};
let plan = test_plan();
let res = plan.accept(&mut visitor);
assert!(res.is_ok());
assert_eq!(
visitor.inner.strings,
vec![
"pre_visit Projection",
"pre_visit Filter",
"pre_visit TableScan",
"post_visit TableScan",
]
);
}
#[derive(Debug, Default)]
struct ErrorVisitor {
inner: OkVisitor,
return_error_from_pre_in: OptionalCounter,
return_error_from_post_in: OptionalCounter,
}
impl PlanVisitor for ErrorVisitor {
type Error = String;
fn pre_visit(
&mut self,
plan: &LogicalPlan,
) -> std::result::Result<bool, Self::Error> {
if self.return_error_from_pre_in.dec() {
return Err("Error in pre_visit".into());
}
self.inner.pre_visit(plan)
}
fn post_visit(
&mut self,
plan: &LogicalPlan,
) -> std::result::Result<bool, Self::Error> {
if self.return_error_from_post_in.dec() {
return Err("Error in post_visit".into());
}
self.inner.post_visit(plan)
}
}
#[test]
fn error_pre_visit() {
let mut visitor = ErrorVisitor {
return_error_from_pre_in: OptionalCounter::new(2),
..Default::default()
};
let plan = test_plan();
let res = plan.accept(&mut visitor);
if let Err(e) = res {
assert_eq!("Error in pre_visit", e);
} else {
panic!("Expected an error");
}
assert_eq!(
visitor.inner.strings,
vec!["pre_visit Projection", "pre_visit Filter",]
);
}
#[test]
fn error_post_visit() {
let mut visitor = ErrorVisitor {
return_error_from_post_in: OptionalCounter::new(1),
..Default::default()
};
let plan = test_plan();
let res = plan.accept(&mut visitor);
if let Err(e) = res {
assert_eq!("Error in post_visit", e);
} else {
panic!("Expected an error");
}
assert_eq!(
visitor.inner.strings,
vec![
"pre_visit Projection",
"pre_visit Filter",
"pre_visit TableScan",
"post_visit TableScan",
]
);
}
fn test_plan() -> LogicalPlan {
let schema = Schema::new(vec![
Field::new("id", DataType::Int32, false),
Field::new("state", DataType::Utf8, false),
]);
LogicalPlanBuilder::scan_empty(None, &schema, Some(vec![0, 1]))
.unwrap()
.filter(col("state").eq(lit("CO")))
.unwrap()
.project(vec![col("id")])
.unwrap()
.build()
.unwrap()
}
}