execution.rs

/// Implementation of PlanVisitor for DataFusion ExecutionPlan nodes.
use crate::visitor::PlanVisitor;
use crate::{JoinAlgorithm, LeafNode, PlanNodeMetadata};

use datafusion::arrow::datatypes::Schema;
use datafusion::execution::context::SessionContext;
use datafusion::physical_plan::ExecutionPlan;
use datafusion::physical_plan::joins::{HashJoinExec, NestedLoopJoinExec};
use datafusion::physical_plan::test::exec::MockExec;
use datafusion::physical_plan::{ExecutionPlanVisitor, accept};
use std::collections::HashSet;

use std::sync::Arc;

use tracing::{debug, trace, warn};

pub struct DFExecutionPlanVisitor {
    metadata: Vec<PlanNodeMetadata>,
    ctx: Arc<SessionContext>,
}

impl DFExecutionPlanVisitor {
    pub fn new(ctx: Arc<SessionContext>) -> Self {
        DFExecutionPlanVisitor {
            metadata: Vec::new(),
            ctx,
        }
    }

    // Doesn't work if tables have same name
    async fn _find_table_name_by_schema(
        &self,
        target_schema: &Schema,
    ) -> Option<String> {
        for catalog_name in self.ctx.catalog_names() {
            let catalog = self.ctx.catalog(&catalog_name)?;
            for schema_name in catalog.schema_names() {
                let schema_provider = catalog.schema(&schema_name)?;
                for table_name in schema_provider.table_names() {
                    if let Some(table) =
                        schema_provider.table(&table_name).await.ok().unwrap()
                    {
                        // Create a set of tables columns
                        let schema = table.schema();
                        let a_cols = schema
                            .fields()
                            .iter()
                            .map(|f| f.name())
                            .collect::<HashSet<_>>()
                            .clone();
                        let b_cols = target_schema
                            .fields()
                            .iter()
                            .map(|f| f.name())
                            .collect::<Vec<_>>();
                        // Check if the table's schema matches the target schema
                        let mut is_match = true;
                        for col in b_cols {
                            if !a_cols.contains(col) {
                                is_match = false;
                                break;
                            }
                        }
                        if is_match {
                            return Some(table_name.clone());
                        }
                    }
                }
            }
        }
        None
    }
}

impl PlanVisitor<dyn ExecutionPlan> for DFExecutionPlanVisitor {
    fn visit(&mut self, plan: &dyn ExecutionPlan) -> Vec<PlanNodeMetadata> {
        // TODO: Making another temp visitor isn't the best way to do this, try something else later
        let mut visitor = DFExecutionPlanVisitor::new(self.ctx.clone());

        debug!("stats: {:?}", plan.partition_statistics(None));

        // Visit the plan and collect metadata
        accept(plan, &mut visitor).unwrap();
        // print!("{:?}", visitor.metadata);
        visitor.metadata
    }
}

impl ExecutionPlanVisitor for DFExecutionPlanVisitor {
    type Error = datafusion::error::DataFusionError;

    fn pre_visit<'a>(
        &mut self,
        _plan: &'a dyn ExecutionPlan,
    ) -> Result<bool, Self::Error> {
        // We don't use previsit, but must implement

        Ok(true)
    }

    /// Visit all nodes of a plan in post order, saving information related to joins.
    ///
    /// Consider the following join tree:
    /// ```text
    ///
    ///           HashJoin
    ///          /        \
    ///     HashJoin     NestedLoopJoin
    ///     /     \       /          \
    ///    A       B     C            D
    ///```
    /// When we visit the nodes in post order, we will visit A, B, HashJoin(AB), C, D, NestedLoopJoin(CD), HashJoin(AB, CD)
    ///
    /// We implement the post_visit to collect the required metadata from each node and save it to output
    ///
    fn post_visit(
        &mut self,
        plan: &dyn ExecutionPlan,
    ) -> Result<bool, Self::Error> {
        // TODO: Replace this with a trait that maps nodes to NodeMetadatas
        match plan.name() {
            "NestedLoopJoinExec" => {
                trace!("Found NestedLoopJoinExec");
                if let Some(_exec) =
                    plan.as_any().downcast_ref::<NestedLoopJoinExec>()
                {
                    self.metadata.push(PlanNodeMetadata::Join {
                        join_algorithm: JoinAlgorithm::NestedLoopJoin,
                        join_card: None,
                        memoize: None,
                    });
                }
            }
            "HashJoinExec" => {
                trace!("Found HashJoinExec");
                if let Some(_exec) =
                    plan.as_any().downcast_ref::<HashJoinExec>()
                {
                    self.metadata.push(PlanNodeMetadata::Join {
                        join_algorithm: JoinAlgorithm::HashJoin,
                        join_card: None,
                        memoize: None,
                    });
                }
            }
            "DataSourceExec" => {
                trace!("Found DataSourceExec");
                // println!("{:?}", plan);
                // let target_schema = plan.schema();
                // let table_name_future = self.find_table_name_by_schema(target_schema.as_ref());
                // let table_name = block_on(table_name_future);

                // Extract table name from csv
                let file_name = plan
                    .as_any()
                    .downcast_ref::<datafusion::datasource::memory::DataSourceExec>()
                    .unwrap()
                    .data_source()
                    .as_any()
                    .downcast_ref::<datafusion::datasource::physical_plan::FileScanConfig>()
                    .unwrap()
                    .file_groups
                    .get(0)
                    .unwrap()
                    .files()
                    .get(0)
                    .unwrap()
                    .path()
                    .filename()
                    .unwrap();

                trace!("File name: {:?}", file_name);
                let mut table_name: Option<String> = None;
                if file_name.ends_with(".csv") {
                    table_name =
                        Some(file_name.replace(".csv", "").to_string());
                }
                if let Some(name) = table_name {
                    self.metadata
                        .push(PlanNodeMetadata::Leaf(LeafNode::new(name)));
                } else {
                    warn!(
                        "Could not find table name for schema: {:?}",
                        plan.schema()
                    );
                }
            }
            "MockExec" => {
                trace!("Found MockExec");
                let table_name = plan
                    .as_any()
                    .downcast_ref::<MockExec>()
                    .unwrap()
                    .schema()
                    .metadata()
                    .get("table_name")
                    .unwrap()
                    .clone();
                self.metadata
                    .push(PlanNodeMetadata::Leaf(LeafNode::new(table_name)));
            }
            _ => {
                /* No hint for other plan types */
                trace!("Skipping plan node: {}", plan.name());
            }
        }

        Ok(true)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use datafusion::arrow::datatypes::{DataType, Field, Schema};
    use datafusion::common::NullEquality;
    use datafusion::physical_expr::expressions::Column;
    use datafusion::physical_plan::joins::{HashJoinExec, PartitionMode};
    use datafusion::prelude::JoinType;
    use std::collections::HashMap;
    use std::sync::Arc;
    use std::vec;

    #[test]
    fn basic_post_visit_test() {
        let ctx = Arc::new(SessionContext::new());
        let mut visitor = DFExecutionPlanVisitor::new(ctx.clone());

        // Init a very simple DF execution plan
        // ```text
        //  HashJoin
        //  /      \
        // A        B
        // ```

        let a = MockExec::new(
            vec![],
            Arc::new(Schema::new_with_metadata(
                vec![Field::new("A", DataType::Utf8, false)],
                HashMap::from([("table_name".to_string(), "A".to_string())]),
            )),
        );

        let b = MockExec::new(
            vec![],
            Arc::new(Schema::new_with_metadata(
                vec![Field::new("B", DataType::Utf8, false)],
                HashMap::from([("table_name".to_string(), "B".to_string())]),
            )),
        );

        let hash_join = HashJoinExec::try_new(
            Arc::new(a),
            Arc::new(b),
            vec![(
                Arc::new(Column::new("A", 0)),
                Arc::new(Column::new("B", 0)),
            )],
            None,
            &JoinType::Inner,
            None,
            PartitionMode::Auto,
            NullEquality::NullEqualsNull,
        )
        .unwrap();

        let nodes = visitor.visit(&hash_join);

        let expected_nodes = vec![
            PlanNodeMetadata::Leaf(LeafNode::new("A".to_string())),
            PlanNodeMetadata::Leaf(LeafNode::new("B".to_string())),
            PlanNodeMetadata::Join {
                join_algorithm: JoinAlgorithm::HashJoin,
                join_card: None,
                memoize: None,
            },
        ];

        assert_eq!(nodes, expected_nodes);
    }
}