datafusion_materialized_views/rewrite/

exploitation.rs

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

/*!

This module implements a query rewriting optimizer, also known as "view exploitation"
in some academic sources. The "view matching" subproblem is implemented in the [`SpjNormalForm`] code,
which is used by the [`ViewMatcher`] logical optimizer to compare queries with materialized views.

The query rewriting process spans both the logical and physical planning phases and can be described as follows:

1.  During logical optimization, the [`ViewMatcher`] rule scans all available materialized views
    and attempts to match them against each sub-expression of the query plan by comparing their SPJ normal forms.
    If a match is found, the sub-expression is replaced with a [`OneOf`] node, which contains the original sub-expression
    and one or more candidate rewrites using materialized views.
2.  During physical planning, the [`ViewExploitationPlanner`] identifies [`OneOf`] nodes and generates a [`OneOfExec`]
    physical plan node, which contains all candidate physical plans corresponding to the logical plans in the original [`OneOf`] node.
3.  DataFusion is allowed to run its usual physical optimization rules, which may add additional operators such as sorts or repartitions
    to the candidate plans. Filter, sort, and projection pushdown into the `OneOfExec` nodes are important as these can affect cost
    estimations in the next phase.
4.  Finally, a user-defined cost function is used to choose the "best" candidate within each `OneOfExec` node.
    The [`PruneCandidates`] physical optimizer rule is used to finalize the choice by replacing each `OneOfExec` node
    with its selected best candidate plan.

In the [reference paper](https://dsg.uwaterloo.ca/seminars/notes/larson-paper.pdf) for this implementation, the authors mention
that the database's builtin cost optimizer takes care of selecting the best rewrite. However, DataFusion lacks cost-based optimization.
While we do use a user-defined cost function to select the best candidate at each `OneOfExec`, this requires cooperation from the planner
to push down relevant information such as projections, sorts, and filters into the `OneOfExec` nodes.

*/

use std::collections::HashMap;
use std::{collections::HashSet, sync::Arc};

use async_trait::async_trait;
use datafusion::catalog::TableProvider;
use datafusion::datasource::provider_as_source;
use datafusion::execution::context::SessionState;
use datafusion::execution::{SendableRecordBatchStream, TaskContext};
use datafusion::physical_expr::{PhysicalSortExpr, PhysicalSortRequirement};
use datafusion::physical_expr_common::sort_expr::format_physical_sort_requirement_list;
use datafusion::physical_optimizer::PhysicalOptimizerRule;
use datafusion::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties};
use datafusion::physical_planner::{ExtensionPlanner, PhysicalPlanner};
use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion, TreeNodeRewriter};
use datafusion_common::{DataFusionError, Result, TableReference};
use datafusion_expr::{Extension, LogicalPlan, UserDefinedLogicalNode, UserDefinedLogicalNodeCore};
use datafusion_optimizer::OptimizerRule;
use datafusion_physical_expr::OrderingRequirements;
use itertools::Itertools;
use ordered_float::OrderedFloat;

use crate::materialized::cast_to_materialized;

use super::normal_form::SpjNormalForm;
use super::QueryRewriteOptions;

/// A cost function. Used to evaluate the best physical plan among multiple equivalent choices.
pub type CostFn = Arc<
    dyn for<'a> Fn(Box<dyn Iterator<Item = &'a dyn ExecutionPlan> + 'a>) -> Vec<f64> + Send + Sync,
>;

/// A logical optimizer that generates candidate logical plans in the form of [`OneOf`] nodes.
#[derive(Debug)]
pub struct ViewMatcher {
    mv_plans: HashMap<TableReference, (Arc<dyn TableProvider>, SpjNormalForm)>,
}

impl ViewMatcher {
    /// Loads table information and processes all views, returning a new optimizer rule.
    pub async fn try_new_from_state(session_state: &SessionState) -> Result<Self> {
        let mut mv_plans: HashMap<TableReference, _> = HashMap::default();
        for (resolved_table_ref, table) in
            super::util::list_tables(session_state.catalog_list().as_ref()).await?
        {
            let Some(mv) = cast_to_materialized(table.as_ref())? else {
                continue;
            };

            if !mv.config().use_in_query_rewrite {
                continue;
            }

            // Analyze the plan to normalize things such as wildcard expressions
            let analyzed_plan = session_state.analyzer().execute_and_check(
                mv.query(),
                session_state.config_options(),
                |_, _| {},
            )?;

            match SpjNormalForm::new(&analyzed_plan) {
                Err(e) => {
                    log::trace!("can't support view matching for {resolved_table_ref}: {e}")
                }
                Ok(normal_form) => {
                    mv_plans.insert(resolved_table_ref.clone().into(), (table, normal_form));
                }
            }
        }

        Ok(ViewMatcher { mv_plans })
    }
}

impl OptimizerRule for ViewMatcher {
    fn rewrite(
        &self,
        plan: LogicalPlan,
        config: &dyn datafusion_optimizer::OptimizerConfig,
    ) -> Result<Transformed<LogicalPlan>> {
        if !config
            .options()
            .extensions
            .get::<QueryRewriteOptions>()
            .cloned()
            .unwrap_or_default()
            .enabled
        {
            return Ok(Transformed::no(plan));
        }

        plan.rewrite(&mut ViewMatchingRewriter { parent: self })
    }

    fn supports_rewrite(&self) -> bool {
        true
    }

    fn name(&self) -> &str {
        "view_matcher"
    }
}

/// A logical plan rewriter that looks for opportunities for substitution.
struct ViewMatchingRewriter<'a> {
    parent: &'a ViewMatcher,
}

impl TreeNodeRewriter for ViewMatchingRewriter<'_> {
    type Node = LogicalPlan;

    fn f_down(&mut self, node: Self::Node) -> Result<Transformed<Self::Node>> {
        if matches!(&node, LogicalPlan::Extension(ext) if ext.node.as_any().is::<OneOf>()) {
            return Ok(Transformed::new(node, false, TreeNodeRecursion::Jump));
        }

        let table_reference = match locate_single_table_scan(&node)? {
            None => return Ok(Transformed::no(node)),
            Some(table_reference) => table_reference,
        };

        log::trace!("rewriting node: {}", node.display());

        let form = match SpjNormalForm::new(&node) {
            Err(e) => {
                log::trace!(
                    "couldn't generate rewrites: for {table_reference}, recursing deeper: {e}"
                );
                return Ok(Transformed::no(node));
            }
            Ok(form) => form,
        };

        // Generate candidate substitutions
        let candidates = self
            .parent
            .mv_plans
            .iter()
            .filter_map(|(table_ref, (table, plan))| {
                // Only attempt rewrite if the view references our table in the first place
                plan.referenced_tables()
                    .contains(&table_reference)
                    .then(|| {
                        form.rewrite_from(
                            plan,
                            table_ref.clone(),
                            provider_as_source(Arc::clone(table)),
                        )
                        .transpose()
                    })
                    .flatten()
            })
            .flat_map(|res| match res {
                Err(e) => {
                    log::trace!("error rewriting: {e}");
                    None
                }
                Ok(plan) => Some(plan),
            })
            .collect::<Vec<_>>();

        if candidates.is_empty() {
            log::trace!("no candidates");
            Ok(Transformed::no(node))
        } else {
            Ok(Transformed::new(
                LogicalPlan::Extension(Extension {
                    node: Arc::new(OneOf {
                        branches: Some(node).into_iter().chain(candidates).collect_vec(),
                    }),
                }),
                true,
                TreeNodeRecursion::Jump,
            ))
        }
    }
}

fn locate_single_table_scan(node: &LogicalPlan) -> Result<Option<TableReference>> {
    let mut table_reference = None;
    node.apply(|plan| {
        if let LogicalPlan::TableScan(scan) = plan {
            match table_reference {
                Some(_) => {
                    // A table reference was already found, but there is another.
                    // Erase the original and stop, return None.
                    table_reference = None;
                    return Ok(TreeNodeRecursion::Stop);
                }
                None => table_reference = Some(scan.table_name.clone()),
            }
        }
        Ok(TreeNodeRecursion::Continue)
    })?;

    // Either there are multiple or there are no table scans.
    Ok(table_reference)
}

/// [`ExtensionPlanner`]` that chooses the best plan from a `OneOf` node.
pub struct ViewExploitationPlanner {
    cost: CostFn,
}

impl ViewExploitationPlanner {
    /// Initialize this ViewExploitationPlanner with a given cost function.
    pub fn new(cost: CostFn) -> Self {
        Self { cost }
    }
}

#[async_trait]
impl ExtensionPlanner for ViewExploitationPlanner {
    /// Choose the best candidate and use it for the physical plan.
    async fn plan_extension(
        &self,
        _planner: &dyn PhysicalPlanner,
        node: &dyn UserDefinedLogicalNode,
        logical_inputs: &[&LogicalPlan],
        physical_inputs: &[Arc<dyn ExecutionPlan>],
        _session_state: &SessionState,
    ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
        if node.as_any().downcast_ref::<OneOf>().is_none() {
            return Ok(None);
        }

        if logical_inputs
            .iter()
            .map(|plan| plan.schema())
            .any(|schema| schema != logical_inputs[0].schema())
        {
            return Err(DataFusionError::Plan(
                "candidate logical plans should have the same schema".to_string(),
            ));
        }

        if physical_inputs
            .iter()
            .map(|plan| plan.schema())
            .any(|schema| schema != physical_inputs[0].schema())
        {
            return Err(DataFusionError::Plan(
                "candidate physical plans should have the same schema".to_string(),
            ));
        }

        Ok(Some(Arc::new(OneOfExec::try_new(
            physical_inputs.to_vec(),
            None,
            Arc::clone(&self.cost),
        )?)))
    }
}

/// A custom logical plan node that denotes multiple equivalent sub-expressions in a logical plan.
/// Used for rewriting queries to use available materialized views, where beneficial.
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Hash)]
pub struct OneOf {
    branches: Vec<LogicalPlan>,
}

impl UserDefinedLogicalNodeCore for OneOf {
    fn name(&self) -> &str {
        "OneOf"
    }

    fn inputs(&self) -> Vec<&LogicalPlan> {
        self.branches
            .iter()
            .sorted_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
            .collect_vec()
    }

    fn schema(&self) -> &datafusion_common::DFSchemaRef {
        self.branches[0].schema()
    }

    fn expressions(&self) -> Vec<datafusion::prelude::Expr> {
        vec![]
    }

    fn prevent_predicate_push_down_columns(&self) -> std::collections::HashSet<String> {
        HashSet::new() // all predicates can be pushed down
    }

    fn fmt_for_explain(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "OneOf")
    }

    fn with_exprs_and_inputs(
        &self,
        _exprs: Vec<datafusion::prelude::Expr>,
        inputs: Vec<LogicalPlan>,
    ) -> Result<Self> {
        Ok(Self { branches: inputs })
    }
}

/// A physical plan that represents a unique selection between multiple logically equivalent candidate physical plans.
#[derive(Clone)]
pub struct OneOfExec {
    candidates: Vec<Arc<dyn ExecutionPlan>>,
    // Optionally declare a required input ordering
    // This will inform DataFusion to add sorts to children,
    // which will improve cost estimation of candidates
    required_input_ordering: Option<OrderingRequirements>,
    // Index of the candidate with the best cost
    best: usize,
    // Cost function to use in optimization
    cost: CostFn,
}

impl std::fmt::Debug for OneOfExec {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("OneOfExec")
            .field("candidates", &self.candidates)
            .field("required_input_ordering", &self.required_input_ordering)
            .field("best", &self.best)
            .finish_non_exhaustive()
    }
}

impl OneOfExec {
    /// Create a new `OneOfExec`
    pub fn try_new(
        candidates: Vec<Arc<dyn ExecutionPlan>>,
        required_input_ordering: Option<OrderingRequirements>,
        cost: CostFn,
    ) -> Result<Self> {
        if candidates.is_empty() {
            return Err(DataFusionError::Plan(
                "can't create OneOfExec with empty children".to_string(),
            ));
        }

        let best = cost(Box::new(candidates.iter().map(|c| c.as_ref())))
            .iter()
            .position_min_by_key(|&cost| OrderedFloat(*cost))
            .unwrap();

        Ok(Self {
            candidates,
            required_input_ordering,
            best,
            cost,
        })
    }

    /// Return the best of this `OneOfExec`'s children, using the cost function provided to
    /// this plan at initialization timee
    pub fn best(&self) -> Arc<dyn ExecutionPlan> {
        Arc::clone(&self.candidates[self.best])
    }

    /// Modify this plan's required input ordering.
    /// Used for sort pushdown
    pub fn with_required_input_ordering(self, requirement: Option<OrderingRequirements>) -> Self {
        Self {
            required_input_ordering: requirement,
            ..self
        }
    }
}

impl ExecutionPlan for OneOfExec {
    fn name(&self) -> &str {
        "OneOfExec"
    }

    fn as_any(&self) -> &dyn std::any::Any {
        self
    }

    fn properties(&self) -> &PlanProperties {
        self.candidates[self.best].properties()
    }

    fn required_input_ordering(&self) -> Vec<Option<OrderingRequirements>> {
        vec![self.required_input_ordering.clone(); self.children().len()]
    }

    fn maintains_input_order(&self) -> Vec<bool> {
        vec![true; self.candidates.len()]
    }

    fn benefits_from_input_partitioning(&self) -> Vec<bool> {
        vec![false; self.candidates.len()]
    }

    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
        self.candidates.iter().collect_vec()
    }

    fn with_new_children(
        self: Arc<Self>,
        children: Vec<Arc<dyn ExecutionPlan>>,
    ) -> Result<Arc<dyn ExecutionPlan>> {
        if children.len() == 1 {
            return Ok(Arc::clone(&children[0]));
        }

        Ok(Arc::new(Self::try_new(
            children,
            self.required_input_ordering.clone(),
            Arc::clone(&self.cost),
        )?))
    }

    fn supports_limit_pushdown(&self) -> bool {
        true
    }

    fn execute(
        &self,
        partition: usize,
        context: Arc<TaskContext>,
    ) -> Result<SendableRecordBatchStream> {
        self.candidates[self.best].execute(partition, context)
    }

    fn statistics(&self) -> Result<datafusion_common::Statistics> {
        self.candidates[self.best].partition_statistics(None)
    }

    fn partition_statistics(
        &self,
        partition: Option<usize>,
    ) -> Result<datafusion_common::Statistics> {
        self.candidates[self.best].partition_statistics(partition)
    }
}

impl DisplayAs for OneOfExec {
    fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        let costs = (self.cost)(Box::new(self.children().iter().map(|arc| arc.as_ref())));
        match t {
            DisplayFormatType::Default | DisplayFormatType::Verbose => {
                write!(
                    f,
                    "OneOfExec(best={}), costs=[{}], required_input_ordering=[{}]",
                    self.best,
                    costs.into_iter().join(","),
                    format_physical_sort_requirement_list(
                        &self
                            .required_input_ordering
                            .as_ref()
                            .map(|req| {
                                req.clone()
                                    .into_single()
                                    .into_iter()
                                    .map(PhysicalSortExpr::from)
                                    .map(PhysicalSortRequirement::from)
                                    .collect_vec()
                            })
                            .unwrap_or_default(),
                    )
                )
            }
            DisplayFormatType::TreeRender => Ok(()),
        }
    }
}

/// Finalize selection of best candidate plan from a OneOfExec.
#[derive(Debug, Clone, Default)]
pub struct PruneCandidates;

impl PhysicalOptimizerRule for PruneCandidates {
    fn optimize(
        &self,
        plan: Arc<dyn ExecutionPlan>,
        _config: &datafusion::config::ConfigOptions,
    ) -> Result<Arc<dyn ExecutionPlan>> {
        // Search for any OneOfExec nodes.
        plan.transform(&|plan: Arc<dyn ExecutionPlan>| {
            if let Some(one_of_exec) = plan.as_any().downcast_ref::<OneOfExec>() {
                Ok(Transformed::new(
                    one_of_exec.best(),
                    true,
                    TreeNodeRecursion::Jump,
                ))
            } else {
                Ok(Transformed::no(plan))
            }
        })
        .map(|t| t.data)
    }

    fn name(&self) -> &str {
        "PruneCandidates"
    }

    fn schema_check(&self) -> bool {
        true
    }
}