Merge branch 'main' into limit_pushdown

Author: xudong963

Cargo.toml

@@ -28,6 +28,7 @@ keywords = ["arrow", "arrow-rs", "datafusion"]
 rust-version = "1.85.1"
 
 [dependencies]
+aquamarine = "0.6.0"
 arrow = "56.0.0"
 arrow-schema = "56.0.0"
 async-trait = "0.1.89"

src/lib.rs

@@ -42,6 +42,9 @@
 pub mod materialized;
 
 /// An implementation of Query Rewriting, an optimization that rewrites queries to make use of materialized views.
+///
+/// The implementation is based heavily on [this paper](https://dsg.uwaterloo.ca/seminars/notes/larson-paper.pdf),
+/// *Optimizing Queries Using Materialized Views: A Practical, Scalable Solution*.
 pub mod rewrite;
 
 /// Configuration options for materialized view related features.

src/materialized.rs

@@ -15,7 +15,6 @@
 // specific language governing permissions and limitations
 // under the License.
 
-/// Track dependencies of materialized data in object storage
 pub mod dependencies;
 
 /// Pluggable metadata sources for incremental view maintenance
@@ -41,6 +40,7 @@ use datafusion::{
     catalog::TableProvider,
     datasource::listing::{ListingTable, ListingTableUrl},
 };
+use datafusion_common::DataFusionError;
 use datafusion_expr::LogicalPlan;
 use itertools::Itertools;
 
@@ -110,6 +110,14 @@ pub trait Materialized: ListingTableLike {
     fn config(&self) -> MaterializedConfig {
         MaterializedConfig::default()
     }
+
+    /// Which partition columns are 'static'.
+    /// Static partition columns are those that are used in incremental view maintenance.
+    /// These should be a prefix of the full set of partition columns returned by [`ListingTableLike::partition_columns`].
+    /// The rest of the partition columns are 'dynamic' and their values will be generated at runtime during incremental refresh.
+    fn static_partition_columns(&self) -> Vec<String> {
+        <Self as ListingTableLike>::partition_columns(self)
+    }
 }
 
 /// Register a [`Materialized`] implementation in this registry.
@@ -122,13 +130,38 @@ pub fn register_materialized<T: Materialized>() {
 }
 
 /// Attempt to cast the given TableProvider into a [`Materialized`].
-/// If the table's type has not been registered using [`register_materialized`], will return `None`.
-pub fn cast_to_materialized(table: &dyn TableProvider) -> Option<&dyn Materialized> {
-    TABLE_TYPE_REGISTRY.cast_to_materialized(table).or_else(|| {
-        TABLE_TYPE_REGISTRY
-            .cast_to_decorator(table)
-            .and_then(|decorator| cast_to_materialized(decorator.base()))
-    })
+/// If the table's type has not been registered using [`register_materialized`], will return `Ok(None)`.
+///
+/// Does a runtime check on some invariants of `Materialized` and returns an error if they are violated.
+/// In particular, checks that the static partition columns are a prefix of all partition columns.
+pub fn cast_to_materialized(
+    table: &dyn TableProvider,
+) -> Result<Option<&dyn Materialized>, DataFusionError> {
+    let materialized = match TABLE_TYPE_REGISTRY
+        .cast_to_materialized(table)
+        .map(Ok)
+        .or_else(|| {
+            TABLE_TYPE_REGISTRY
+                .cast_to_decorator(table)
+                .and_then(|decorator| cast_to_materialized(decorator.base()).transpose())
+        })
+        .transpose()?
+    {
+        None => return Ok(None),
+        Some(m) => m,
+    };
+
+    let static_partition_cols = materialized.static_partition_columns();
+    let all_partition_cols = materialized.partition_columns();
+
+    if materialized.partition_columns()[..static_partition_cols.len()] != static_partition_cols[..]
+    {
+        return Err(DataFusionError::Plan(format!(
+            "Materialized view's static partition columns ({static_partition_cols:?}) must be a prefix of all partition columns ({all_partition_cols:?})"
+        )));
+    }
+
+    Ok(Some(materialized))
 }
 
 /// A `TableProvider` that decorates other `TableProvider`s.

src/materialized/dependencies.rs

@@ -722,7 +722,7 @@ impl FileMetadataBuilder {
     }
 }
 
-/// Provides [`ObjectMetadata`] data to the [`FileMetadata`] table provider.
+/// Provides [`ObjectMeta`] data to the [`FileMetadata`] table provider.
 #[async_trait]
 pub trait FileMetadataProvider: std::fmt::Debug + Send + Sync {
     /// List all files in the store for the given `url` prefix.

src/materialized/file_metadata.rs

@@ -17,8 +17,6 @@
 
 use datafusion::{common::extensions_options, config::ConfigExtension};
 
-/// Implements a query rewriting optimizer, also known as "view exploitation"
-/// in some academic sources.
 pub mod exploitation;
 
 pub mod normal_form;

src/rewrite.rs

@@ -15,6 +15,34 @@
 // 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};
 
@@ -59,7 +87,7 @@ impl ViewMatcher {
         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 {
+            let Some(mv) = cast_to_materialized(table.as_ref())? else {
                 continue;
             };
 

src/rewrite/exploitation.rs

@@ -17,7 +17,7 @@
 
 /*!
 
-This module contains code primarily used for view matching. We implement the view matching algorithm from [this paper](https://courses.cs.washington.edu/courses/cse591d/01sp/opt_views.pdf),
+This module contains code primarily used for view matching. We implement the view matching algorithm from [this paper](https://dsg.uwaterloo.ca/seminars/notes/larson-paper.pdf),
 which provides a method for determining when one Select-Project-Join query can be rewritten in terms of another Select-Project-Join query.
 
 The implementation is contained in [`SpjNormalForm::rewrite_from`]. The method can be summarized as follows: