From 485e977e56dcc1f7d207e5aebdd87e5f4c9dd67f Mon Sep 17 00:00:00 2001 From: Qi Zhu Date: Fri, 24 Jul 2026 11:54:13 +0800 Subject: [PATCH 1/4] Readiness-aware view exploitation: distinguish predicate-pruned empty MVs from unpopulated ones A materialized view candidate whose scan yields an empty plan is ambiguous: either the predicate genuinely pruned everything (the empty result is correct and extremely cheap), or the MV simply has not been populated yet (trusting the empty plan silently returns wrong results). Cost functions previously could not tell these apart. This adds: - RewriteReadiness (Ready / NotReady / Unknown) and a Materialized::rewrite_readiness() method (default: Unknown) so table providers can report their lifecycle state. - CandidateMetadata (Base | Materialized { table_ref, readiness }) carried per branch via a RewriteContext attached to OneOf/OneOfExec, aligned with branch order. To make that alignment hold, OneOf::inputs() now returns branches in stored order (previously it re-sorted with LogicalPlan::partial_cmp on every call, which could move the base branch off branches[0] after a rebuild and expose a candidate's schema from OneOf::schema()); candidates are sorted by table reference at construction, and the base plan is pinned at branches[0]. - Two-gate NotReady filtering: NotReady candidates are dropped at logical rewrite time, and again at physical planning time after a refresh. - Physical-time readiness refresh in plan_extension: by the time the physical planner runs, every candidate's scan() has executed, so providers whose readiness depends on scan-time side effects (lazy index loading) are re-consulted. Providers can also opt into ReadinessAnnotatedExec, a transparent wrapper that carries readiness captured atomically from the scan's own snapshot, which the refresh prefers over re-sampling (avoids a TOCTOU race with concurrent index publication); the wrapper is stripped before the OneOfExec is built. - CostFn now receives a CostContext (candidate plans + rewrite metadata) instead of a bare plan iterator, so cost functions can implement policies like 'trust Ready empty candidates, penalize Unknown empty ones'. 43 new unit tests cover the readiness types, annotation and stripping, the two filtering gates, the refresh, and branch-order/metadata alignment across plan rebuilds. --- src/materialized.rs | 19 + src/rewrite.rs | 2 + src/rewrite/exploitation.rs | 1919 ++++++++++++++++++++++++++++++++++- src/rewrite/readiness.rs | 570 +++++++++++ 4 files changed, 2465 insertions(+), 45 deletions(-) create mode 100644 src/rewrite/readiness.rs diff --git a/src/materialized.rs b/src/materialized.rs index c3cf30e..e25e7ff 100644 --- a/src/materialized.rs +++ b/src/materialized.rs @@ -100,6 +100,12 @@ pub fn cast_to_listing_table(table: &dyn TableProvider) -> Option<&dyn ListingTa }) } +// Re-exported for backward compatibility with downstream code that used +// to import `RewriteReadiness` from `crate::materialized`. The enum +// itself lives in `crate::rewrite::readiness` now (see PR #55 review), +// so all readiness-related types are grouped in one file. +pub use crate::rewrite::readiness::RewriteReadiness; + /// A hive-partitioned table in object storage that is defined by a user-provided query. pub trait Materialized: ListingTableLike { /// The query that defines this materialized view. @@ -118,6 +124,19 @@ pub trait Materialized: ListingTableLike { fn static_partition_columns(&self) -> Vec { ::partition_columns(self) } + + /// Report whether this MV is currently safe to route queries to. See + /// [`RewriteReadiness`] for the semantics of each variant. Consulted by + /// [`ViewMatcher`](crate::rewrite::exploitation::ViewMatcher) during LP + /// rewrite; `NotReady` MVs are dropped from the candidate set upstream + /// of the cost function, so they never win a rewrite. + /// + /// Default is `Unknown`, which means "include as candidate but the + /// cost function decides" — backward-compatible for providers that + /// don't distinguish lifecycle states. + fn rewrite_readiness(&self) -> RewriteReadiness { + RewriteReadiness::Unknown + } } /// Register a [`Materialized`] implementation in this registry. diff --git a/src/rewrite.rs b/src/rewrite.rs index da24cc5..2ff20ad 100644 --- a/src/rewrite.rs +++ b/src/rewrite.rs @@ -21,6 +21,8 @@ pub mod exploitation; pub mod normal_form; +pub mod readiness; + mod util; extensions_options! { diff --git a/src/rewrite/exploitation.rs b/src/rewrite/exploitation.rs index 997d28a..2529b55 100644 --- a/src/rewrite/exploitation.rs +++ b/src/rewrite/exploitation.rs @@ -67,12 +67,92 @@ use ordered_float::OrderedFloat; use crate::materialized::cast_to_materialized; use super::normal_form::SpjNormalForm; +use super::readiness::{ + drop_not_ready_after_refresh, refresh_candidate_readiness, strip_readiness_annotation, +}; use super::QueryRewriteOptions; +// Re-export readiness types for backward compatibility. Callers that +// used to import from `crate::rewrite::exploitation` (the pre-move +// location) keep working; the canonical home is now +// `crate::rewrite::readiness`. +pub use super::readiness::{ + readiness_from_plan, CandidateMetadata, ReadinessAnnotatedExec, RewriteReadiness, +}; + +/// Logical rewrite metadata propagated alongside equivalent candidate plans. +#[derive(Debug, Clone, Default, PartialEq, PartialOrd, Eq, Hash)] +pub struct RewriteContext { + root_table_refs: Vec, + /// Per-branch metadata aligned with the containing `OneOf` / `OneOfExec` + /// branch order. See [`CandidateMetadata`] for the alignment invariant. + /// Empty when no callsite has attached candidate metadata (backward + /// compatible with callers that pre-date the per-branch extensions). + candidates: Vec, +} + +impl RewriteContext { + /// Create a new rewrite context from the root table refs visible during rewrite. + /// The candidate list starts empty; use [`Self::with_candidates`] to attach it. + pub fn new(root_table_refs: Vec) -> Self { + Self { + root_table_refs, + candidates: Vec::new(), + } + } + + /// Attach the per-branch candidate metadata list. Consumes and returns + /// `self` so the builder chain stays readable at the OneOf construction + /// site. The caller must respect the alignment invariant documented on + /// [`CandidateMetadata`]. + pub fn with_candidates(mut self, candidates: Vec) -> Self { + self.candidates = candidates; + self + } + + /// Returns the root table refs that produced this rewrite opportunity. + pub fn root_table_refs(&self) -> &[String] { + &self.root_table_refs + } + + /// Returns the per-branch candidate metadata, aligned with the branch + /// order in the containing `OneOf` / `OneOfExec`. + pub fn candidates(&self) -> &[CandidateMetadata] { + &self.candidates + } +} + +/// Inputs provided to a cost function when selecting the best candidate plan. +pub struct CostContext<'a> { + candidate_plans: Box + 'a>, + rewrite_context: &'a RewriteContext, +} + +impl<'a> CostContext<'a> { + /// Create a new cost context. + pub fn new( + candidate_plans: Box + 'a>, + rewrite_context: &'a RewriteContext, + ) -> Self { + Self { + candidate_plans, + rewrite_context, + } + } + + /// Consume the context and return the candidate plans iterator. + pub fn into_candidate_plans(self) -> Box + 'a> { + self.candidate_plans + } + + /// Returns rewrite metadata for the current candidate set. + pub fn rewrite_context(&self) -> &RewriteContext { + self.rewrite_context + } +} + /// A cost function. Used to evaluate the best physical plan among multiple equivalent choices. -pub type CostFn = Arc< - dyn for<'a> Fn(Box + 'a>) -> Vec + Send + Sync, ->; +pub type CostFn = Arc Fn(CostContext<'a>) -> Vec + Send + Sync>; /// A logical optimizer that generates candidate logical plans in the form of [`OneOf`] nodes. #[derive(Debug)] @@ -120,6 +200,19 @@ impl ViewMatcher { &self.mv_plans } + /// Test-only constructor that skips the full `try_new_from_state` catalog + /// walk and takes a pre-built `mv_plans` HashMap directly. Used by unit + /// tests that need to drive `ViewMatchingRewriter` without setting up a + /// SessionContext / SchemaProvider — e.g. tests that inject mock + /// `Materialized` providers with a specific `rewrite_readiness()` value + /// to assert the upstream filter. + #[cfg(test)] + pub(crate) fn from_mv_plans_for_test( + mv_plans: HashMap, SpjNormalForm)>, + ) -> Self { + Self { mv_plans } + } + /// Returns materialized views that potentially reference the given table. /// /// This is a preliminary filter - it only checks if the MV references the table @@ -217,43 +310,111 @@ impl TreeNodeRewriter for ViewMatchingRewriter<'_> { Ok(form) => form, }; - // Generate candidate substitutions - let candidates = self + // Generate candidate substitutions, filtering out MVs that report + // `NotReady` upstream so unpopulated / in-flight MVs never reach the + // cost function. Lifecycle correctness lives in the provider (see + // `Materialized::rewrite_readiness`); the cost function then only + // sees candidates it can safely route queries to. + // + // `Ready` and `Unknown` both survive here, but the readiness value is + // preserved and carried through into `CandidateMetadata::Materialized` + // so cost functions can distinguish the two: `Ready` == the provider + // guarantees the MV is safe to answer the query; `Unknown` == the + // provider doesn't know and the cost function should apply its own + // conservative policy (e.g. avoid trusting an EmptyExec candidate as + // predicate-pruned). Erasing the readiness at this layer would silently + // promote every default (`Unknown`) provider to `Ready`. + let candidates: Vec<(LogicalPlan, TableReference, RewriteReadiness)> = 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() + if !plan.referenced_tables().contains(&table_reference) { + return None; + } + // Drop `NotReady` MVs upfront. `Ready` and `Unknown` both pass + // through — `Unknown` is the trait default, so we don't break + // pre-existing providers that don't distinguish lifecycle states. + // + // A `cast_to_materialized` error means the provider violates a + // `Materialized` invariant (e.g. static partition columns not a + // prefix). Surface it in the log rather than silently defaulting + // to `Unknown` so catalog issues stay debuggable, then treat the + // MV as Unknown (the pre-existing "include as candidate, cost + // function decides" default) so the rewrite path keeps functioning. + let readiness = match cast_to_materialized(table.as_ref()) { + Ok(Some(mv)) => mv.rewrite_readiness(), + Ok(None) => RewriteReadiness::Unknown, + Err(e) => { + log::warn!( + "cast_to_materialized failed for {table_ref}: {e}; \ + treating as Unknown readiness" + ); + RewriteReadiness::Unknown + } + }; + if matches!(readiness, RewriteReadiness::NotReady) { + log::trace!("skipping NotReady MV: {table_ref}"); + return None; + } + form.rewrite_from( + plan, + table_ref.clone(), + provider_as_source(Arc::clone(table)), + ) + .transpose() + .map(|res| res.map(|lp| (lp, table_ref.clone(), readiness))) }) .flat_map(|res| match res { Err(e) => { log::trace!("error rewriting: {e}"); None } - Ok(plan) => Some(plan), + Ok(triple) => Some(triple), }) - .collect::>(); + .collect(); if candidates.is_empty() { log::trace!("no candidates"); Ok(Transformed::no(node)) } else { + // Establish and preserve the alignment invariant documented on + // `CandidateMetadata`: + // * `branches[0]` is always the base LP (so `OneOf::schema()`, + // which reads `branches[0].schema()`, always exposes the + // original query's schema — never an MV's schema, which may + // differ in nullability or column qualification); + // * `branches[1..]` are the MV candidates in a deterministic + // order derived from each MV's `TableReference`, which is + // stable under downstream LP transformations (identity + // projection elimination, always-true filter folding, etc. + // do not change an MV's registered name). This keeps the + // metadata aligned with the physical candidate the cost + // function receives even after several optimizer passes + // rebuild the `OneOf` via `with_exprs_and_inputs`. + let mut candidates = candidates; + candidates.sort_by_key(|a| a.1.to_string()); + + let branches: Vec = std::iter::once(node) + .chain(candidates.iter().map(|(lp, _, _)| lp.clone())) + .collect(); + let metadata: Vec = std::iter::once(CandidateMetadata::Base) + .chain(candidates.into_iter().map(|(_, table_ref, readiness)| { + CandidateMetadata::Materialized { + table_ref, + readiness, + } + })) + .collect(); + Ok(Transformed::new( LogicalPlan::Extension(Extension { - node: Arc::new(OneOf { - branches: Some(node).into_iter().chain(candidates).collect_vec(), - }), + node: Arc::new(OneOf::with_rewrite_context( + branches, + RewriteContext::new(vec![table_reference.to_string()]) + .with_candidates(metadata), + )), }), true, TreeNodeRecursion::Jump, @@ -304,16 +465,23 @@ impl ExtensionPlanner for ViewExploitationPlanner { node: &dyn UserDefinedLogicalNode, logical_inputs: &[&LogicalPlan], physical_inputs: &[Arc], - _session_state: &SessionState, + session_state: &SessionState, ) -> Result>> { - if node.as_any().downcast_ref::().is_none() { + let Some(one_of) = node.as_any().downcast_ref::() else { return Ok(None); - } + }; + // Different table providers may expose equivalent fields with different + // nullability. Names and data types still have to match. if logical_inputs .iter() .map(|plan| plan.schema()) - .any(|schema| schema != logical_inputs[0].schema()) + .any(|schema| { + !schemas_equal_ignoring_nullability( + schema.as_arrow(), + logical_inputs[0].schema().as_arrow(), + ) + }) { return Err(DataFusionError::Plan( "candidate logical plans should have the same schema".to_string(), @@ -323,17 +491,68 @@ impl ExtensionPlanner for ViewExploitationPlanner { if physical_inputs .iter() .map(|plan| plan.schema()) - .any(|schema| schema != physical_inputs[0].schema()) + .any(|schema| { + !schemas_equal_ignoring_nullability(&schema, &physical_inputs[0].schema()) + }) { return Err(DataFusionError::Plan( "candidate physical plans should have the same schema".to_string(), )); } + // Physical-time readiness resolution. By the time this planner runs, + // DataFusion has already planned each branch into `physical_inputs`, + // which means it has already invoked `TableProvider::scan()` on every + // candidate MV. For providers whose readiness depends on scan-time + // side effects (e.g. a lazy-index provider that only loads its index + // on first `scan()`), the readiness reported at LP-rewrite time is + // stale by now — re-consulting the provider here gives the definitive + // value the cost function should see (the lazy-index cold-start + // scenario). + let refreshed_context = refresh_candidate_readiness( + one_of.rewrite_context().clone(), + physical_inputs, + session_state, + ) + .await; + + // Strip `ReadinessAnnotatedExec` from every physical input now that + // we've read its readiness. The wrapper is a signaling device local + // to this planner; keeping it in the plan handed to the cost + // function / `OneOfExec` / downstream optimizers would block passes + // that depend on `ExecutionPlan` trait methods the wrapper doesn't + // proxy (limit pushdown, projection pushdown, `with_preserve_order`, + // etc.). + let stripped_inputs: Vec> = physical_inputs + .iter() + .map(|plan| strip_readiness_annotation(Arc::clone(plan))) + .collect::>()?; + + // Enforce the invariant documented on `CandidateMetadata`: NotReady + // never reaches the cost function. A provider can transition to + // NotReady between LP rewrite and physical planning (rare, but + // possible if scan detected an unpopulated state), so this filter is + // the post-scan analogue of `f_down`'s upstream drop. Metadata and + // `physical_inputs` are filtered in lockstep to preserve the + // by-index alignment cost functions rely on. + let (kept_inputs, kept_candidates) = + drop_not_ready_after_refresh(&stripped_inputs, refreshed_context.candidates()); + + // If the filter left only the Base branch we can skip the OneOfExec + // wrapper entirely and just return the base physical plan — nothing + // for the cost function to choose between. + if kept_candidates.len() == 1 && matches!(kept_candidates[0], CandidateMetadata::Base) { + return Ok(Some(Arc::clone(&kept_inputs[0]))); + } + + let filtered_context = RewriteContext::new(refreshed_context.root_table_refs().to_vec()) + .with_candidates(kept_candidates); + Ok(Some(Arc::new(OneOfExec::try_new( - physical_inputs.to_vec(), + kept_inputs, None, Arc::clone(&self.cost), + filtered_context, )?))) } } @@ -343,6 +562,30 @@ impl ExtensionPlanner for ViewExploitationPlanner { #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Hash)] pub struct OneOf { branches: Vec, + rewrite_context: RewriteContext, +} + +impl OneOf { + /// Create a new OneOf node with the given branches. + pub fn new(branches: Vec) -> Self { + Self::with_rewrite_context(branches, RewriteContext::default()) + } + + /// Create a new OneOf node with the given branches and rewrite context. + pub fn with_rewrite_context( + branches: Vec, + rewrite_context: RewriteContext, + ) -> Self { + Self { + branches, + rewrite_context, + } + } + + /// Returns logical rewrite metadata for this candidate set. + pub fn rewrite_context(&self) -> &RewriteContext { + &self.rewrite_context + } } impl UserDefinedLogicalNodeCore for OneOf { @@ -351,10 +594,20 @@ impl UserDefinedLogicalNodeCore for 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() + // Return branches in stored order so `RewriteContext::candidates` stays + // aligned across `with_exprs_and_inputs` rebuilds. Determinism is + // supplied by `ViewMatchingRewriter::f_down`, which puts the base at + // index 0 and sorts MV candidates by `TableReference` (an identifier + // that survives downstream LP transformations). + // + // The previous implementation re-sorted on every call using + // `LogicalPlan::partial_cmp`, which flipped the order whenever an + // optimizer pass changed a branch's LP shape (identity projection + // elimination, always-true filter folding, etc.). That silently + // misattributed per-branch metadata after any tree rebuild — and + // moved the base branch off `branches[0]`, so `OneOf::schema()` + // could expose an MV's schema instead of the query's. + self.branches.iter().collect_vec() } fn schema(&self) -> &datafusion_common::DFSchemaRef { @@ -378,7 +631,10 @@ impl UserDefinedLogicalNodeCore for OneOf { _exprs: Vec, inputs: Vec, ) -> Result { - Ok(Self { branches: inputs }) + Ok(Self { + branches: inputs, + rewrite_context: self.rewrite_context.clone(), + }) } } @@ -394,6 +650,7 @@ pub struct OneOfExec { best: usize, // Cost function to use in optimization cost: CostFn, + rewrite_context: RewriteContext, } impl std::fmt::Debug for OneOfExec { @@ -402,6 +659,7 @@ impl std::fmt::Debug for OneOfExec { .field("candidates", &self.candidates) .field("required_input_ordering", &self.required_input_ordering) .field("best", &self.best) + .field("rewrite_context", &self.rewrite_context) .finish_non_exhaustive() } } @@ -412,6 +670,7 @@ impl OneOfExec { candidates: Vec>, required_input_ordering: Option, cost: CostFn, + rewrite_context: RewriteContext, ) -> Result { if candidates.is_empty() { return Err(DataFusionError::Plan( @@ -419,16 +678,20 @@ impl OneOfExec { )); } - let best = cost(Box::new(candidates.iter().map(|c| c.as_ref()))) - .iter() - .position_min_by_key(|&cost| OrderedFloat(*cost)) - .unwrap(); + let best = cost(CostContext::new( + Box::new(candidates.iter().map(|c| c.as_ref())), + &rewrite_context, + )) + .iter() + .position_min_by_key(|&cost| OrderedFloat(*cost)) + .unwrap(); Ok(Self { candidates, required_input_ordering, best, cost, + rewrite_context, }) } @@ -438,6 +701,11 @@ impl OneOfExec { Arc::clone(&self.candidates[self.best]) } + /// Returns rewrite metadata for this candidate set. + pub fn rewrite_context(&self) -> &RewriteContext { + &self.rewrite_context + } + /// Modify this plan's required input ordering. /// Used for sort pushdown pub fn with_required_input_ordering(self, requirement: Option) -> Self { @@ -449,14 +717,14 @@ impl OneOfExec { } impl ExecutionPlan for OneOfExec { - fn name(&self) -> &str { - "OneOfExec" - } - fn as_any(&self) -> &dyn std::any::Any { self } + fn name(&self) -> &str { + "OneOfExec" + } + fn properties(&self) -> &PlanProperties { self.candidates[self.best].properties() } @@ -489,6 +757,7 @@ impl ExecutionPlan for OneOfExec { children, self.required_input_ordering.clone(), Arc::clone(&self.cost), + self.rewrite_context.clone(), )?)) } @@ -504,10 +773,6 @@ impl ExecutionPlan for OneOfExec { self.candidates[self.best].execute(partition, context) } - fn statistics(&self) -> Result { - self.candidates[self.best].partition_statistics(None) - } - fn partition_statistics( &self, partition: Option, @@ -518,7 +783,10 @@ impl ExecutionPlan for OneOfExec { 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()))); + let costs = (self.cost)(CostContext::new( + Box::new(self.children().iter().map(|arc| arc.as_ref())), + &self.rewrite_context, + )); match t { DisplayFormatType::Default | DisplayFormatType::Verbose => { write!( @@ -580,3 +848,1564 @@ impl PhysicalOptimizerRule for PruneCandidates { true } } + +/// Compare two Arrow schemas ignoring field nullability. +fn schemas_equal_ignoring_nullability(a: &arrow_schema::Schema, b: &arrow_schema::Schema) -> bool { + a.fields().len() == b.fields().len() + && a.fields() + .iter() + .zip(b.fields().iter()) + .all(|(f1, f2)| f1.name() == f2.name() && f1.data_type() == f2.data_type()) +} + +#[cfg(test)] +mod tests_nullability { + use super::*; + use arrow_schema::{DataType, Field, Schema}; + + #[test] + fn schemas_equal_when_only_nullability_differs() { + let a = Schema::new(vec![ + Field::new("ticker", DataType::Utf8, false), + Field::new("date", DataType::Utf8, false), + Field::new("price", DataType::Float64, false), + ]); + let b = Schema::new(vec![ + Field::new("ticker", DataType::Utf8, true), + Field::new("date", DataType::Utf8, true), + Field::new("price", DataType::Float64, true), + ]); + assert!(schemas_equal_ignoring_nullability(&a, &b)); + } + + #[test] + fn schemas_not_equal_when_types_differ() { + let a = Schema::new(vec![Field::new("x", DataType::Int32, false)]); + let b = Schema::new(vec![Field::new("x", DataType::Int64, false)]); + assert!(!schemas_equal_ignoring_nullability(&a, &b)); + } + + #[test] + fn schemas_not_equal_when_names_differ() { + let a = Schema::new(vec![Field::new("ticker", DataType::Utf8, true)]); + let b = Schema::new(vec![Field::new("symbol", DataType::Utf8, true)]); + assert!(!schemas_equal_ignoring_nullability(&a, &b)); + } + + #[test] + fn schemas_not_equal_when_field_count_differs() { + let a = Schema::new(vec![ + Field::new("x", DataType::Int32, false), + Field::new("y", DataType::Int32, false), + ]); + let b = Schema::new(vec![Field::new("x", DataType::Int32, false)]); + assert!(!schemas_equal_ignoring_nullability(&a, &b)); + } +} + +#[cfg(test)] +mod tests_rewrite_context { + use super::*; + use arrow_schema::Schema; + use datafusion::physical_plan::empty::EmptyExec; + use datafusion_expr::LogicalPlanBuilder; + use std::sync::Mutex; + + #[test] + fn one_of_preserves_rewrite_context_when_rebuilt() { + let plan = LogicalPlanBuilder::empty(false) + .build() + .expect("empty plan"); + let one_of = OneOf::with_rewrite_context( + vec![plan.clone()], + RewriteContext::new(vec!["catalog.schema.root_table".to_string()]), + ); + + let rebuilt = + UserDefinedLogicalNodeCore::with_exprs_and_inputs(&one_of, vec![], vec![plan]) + .expect("rebuild one_of"); + + assert_eq!( + rebuilt.rewrite_context().root_table_refs(), + ["catalog.schema.root_table".to_string()] + ); + } + + #[test] + fn one_of_exec_passes_rewrite_context_to_cost_function() { + let seen = Arc::new(Mutex::new(Vec::::new())); + let seen_clone = Arc::clone(&seen); + let cost: CostFn = Arc::new(move |ctx| { + *seen_clone.lock().expect("lock seen") = + ctx.rewrite_context().root_table_refs().to_vec(); + ctx.into_candidate_plans().map(|_| 1.0).collect() + }); + let context = RewriteContext::new(vec!["catalog.schema.root_table".to_string()]); + let schema = Arc::new(Schema::empty()); + let candidates = vec![ + Arc::new(EmptyExec::new(Arc::clone(&schema))) as Arc, + Arc::new(EmptyExec::new(schema)) as Arc, + ]; + + let exec = + OneOfExec::try_new(candidates, None, cost, context.clone()).expect("one_of exec"); + + assert_eq!(exec.rewrite_context(), &context); + assert_eq!(*seen.lock().expect("lock seen"), context.root_table_refs()); + } + + #[test] + fn one_of_exec_with_new_children_preserves_rewrite_context() { + let cost: CostFn = Arc::new(|ctx| ctx.into_candidate_plans().map(|_| 1.0).collect()); + let context = RewriteContext::new(vec!["catalog.schema.root_table".to_string()]); + let schema = Arc::new(Schema::empty()); + let exec = Arc::new( + OneOfExec::try_new( + vec![ + Arc::new(EmptyExec::new(Arc::clone(&schema))) as Arc, + Arc::new(EmptyExec::new(Arc::clone(&schema))) as Arc, + ], + None, + cost, + context.clone(), + ) + .expect("one_of exec"), + ); + + let rebuilt = exec + .with_new_children(vec![ + Arc::new(EmptyExec::new(Arc::clone(&schema))) as Arc, + Arc::new(EmptyExec::new(schema)) as Arc, + ]) + .expect("rebuild exec"); + let rebuilt = rebuilt + .as_any() + .downcast_ref::() + .expect("expected OneOfExec"); + + assert_eq!(rebuilt.rewrite_context(), &context); + } + + #[test] + fn rewrite_context_defaults_candidates_to_empty() { + // `RewriteContext::new` must leave the candidate metadata empty so any + // caller that pre-dates the per-branch metadata extension keeps its + // original behaviour (downstream discriminators fall back to their + // existing signals when the slice is empty). + let ctx = RewriteContext::new(vec!["catalog.schema.root".to_string()]); + assert!(ctx.candidates().is_empty()); + } + + #[test] + fn rewrite_context_builder_attaches_candidate_metadata() { + let mv_a_ref = TableReference::bare("catalog.schema.mv_a"); + let mv_b_ref = TableReference::bare("catalog.schema.mv_b"); + let ctx = + RewriteContext::new(vec!["catalog.schema.root".to_string()]).with_candidates(vec![ + CandidateMetadata::Base, + CandidateMetadata::Materialized { + table_ref: mv_a_ref.clone(), + readiness: RewriteReadiness::Ready, + }, + CandidateMetadata::Materialized { + table_ref: mv_b_ref.clone(), + readiness: RewriteReadiness::Unknown, + }, + ]); + + let cands = ctx.candidates(); + assert_eq!(cands.len(), 3); + assert!(matches!(cands[0], CandidateMetadata::Base)); + assert_eq!( + cands[1], + CandidateMetadata::Materialized { + table_ref: mv_a_ref, + readiness: RewriteReadiness::Ready, + } + ); + assert_eq!( + cands[2], + CandidateMetadata::Materialized { + table_ref: mv_b_ref, + readiness: RewriteReadiness::Unknown, + } + ); + } + + #[test] + fn one_of_inputs_returns_stored_order_without_resort() { + // Regression test: `OneOf::inputs()` used to re-sort by + // `LogicalPlan::partial_cmp`, which flipped branch order whenever a + // downstream optimizer pass mutated a branch (identity Projection + // elimination, always-true Filter folding, etc.) and silently + // misattributed per-branch metadata. `inputs()` must return the + // stored order verbatim so alignment with + // `RewriteContext::candidates()` survives every rebuild. + let base = LogicalPlanBuilder::empty(false) + .build() + .expect("empty base plan"); + let candidate_a = LogicalPlanBuilder::empty(true) + .build() + .expect("candidate a plan"); + // Store in a deliberately non-sorted order to prove `inputs()` is a + // no-op — if it re-sorted, `candidate_a` would come before `base` + // (nullable schema sorts differently) and this assertion would fail. + let branches = vec![base.clone(), candidate_a.clone()]; + let one_of = OneOf::with_rewrite_context( + branches.clone(), + RewriteContext::new(vec!["catalog.schema.root".to_string()]).with_candidates(vec![ + CandidateMetadata::Base, + CandidateMetadata::Materialized { + table_ref: TableReference::bare("catalog.schema.mv_a"), + readiness: RewriteReadiness::Ready, + }, + ]), + ); + + let exposed = UserDefinedLogicalNodeCore::inputs(&one_of); + assert_eq!(exposed.len(), 2); + assert_eq!(*exposed[0], base, "index 0 must remain the base branch"); + assert_eq!(*exposed[1], candidate_a); + } + + #[test] + fn one_of_schema_always_reports_base_branch_schema() { + // `OneOf::schema()` reads `branches[0].schema()`. The ordering fix must + // never move the base branch off index 0 — if an MV's LP sorted + // first, the parent `OneOf` would expose an MV's schema (possibly + // with different nullability / column qualification) instead of the + // query's schema, silently changing DataFusion's downstream planning. + // + // We construct the OneOf directly here (bypassing `f_down`) to check + // that whichever order the caller stores, `schema()` picks up + // `branches[0]`. `ViewMatchingRewriter::f_down` guarantees the caller + // always puts base at index 0, so together this locks in the + // "index 0 is base, and its schema is the query schema" invariant. + let base_nonnull = LogicalPlanBuilder::empty(false) + .build() + .expect("empty base plan"); + let candidate_nullable = LogicalPlanBuilder::empty(true) + .build() + .expect("nullable candidate plan"); + + let one_of_base_first = OneOf::with_rewrite_context( + vec![base_nonnull.clone(), candidate_nullable.clone()], + RewriteContext::default(), + ); + assert_eq!( + UserDefinedLogicalNodeCore::schema(&one_of_base_first), + base_nonnull.schema() + ); + } + + #[test] + fn one_of_alignment_survives_with_exprs_and_inputs_rebuild() { + // Once branches are stored in a stable order, DataFusion's + // `TreeNodeRewriter` walks children by index and rebuilds the parent + // via `with_exprs_and_inputs`, keeping the positional contract. + // Even after a transformation swaps one branch for a differently- + // shaped LP, the rebuilt OneOf must keep every candidate's metadata + // aligned with the (possibly transformed) branch at the same index. + let base = LogicalPlanBuilder::empty(false) + .build() + .expect("empty base plan"); + let candidate = LogicalPlanBuilder::empty(true) + .build() + .expect("candidate plan"); + + let mv_meta = CandidateMetadata::Materialized { + table_ref: TableReference::bare("catalog.schema.mv_a"), + readiness: RewriteReadiness::Ready, + }; + let one_of = OneOf::with_rewrite_context( + vec![base.clone(), candidate.clone()], + RewriteContext::new(vec!["catalog.schema.root".to_string()]) + .with_candidates(vec![CandidateMetadata::Base, mv_meta.clone()]), + ); + + // Simulate an optimizer pass rewriting the MV branch (e.g. injecting + // a Filter). `with_exprs_and_inputs` returns a new OneOf with the + // rewritten branch at the SAME index — never re-sorted. + let rewritten_candidate = LogicalPlanBuilder::empty(false) + .build() + .expect("rewritten candidate plan"); + let rebuilt = UserDefinedLogicalNodeCore::with_exprs_and_inputs( + &one_of, + vec![], + vec![base.clone(), rewritten_candidate.clone()], + ) + .expect("rebuild OneOf"); + + // Base still at index 0. + assert_eq!( + UserDefinedLogicalNodeCore::schema(&rebuilt), + base.schema(), + "OneOf::schema() must still resolve to the base branch schema" + ); + // Metadata still aligned: index 0 = Base, index 1 = the same MV + // identity even though its LP changed. + assert_eq!(rebuilt.rewrite_context().candidates().len(), 2); + assert!(matches!( + rebuilt.rewrite_context().candidates()[0], + CandidateMetadata::Base + )); + assert_eq!(rebuilt.rewrite_context().candidates()[1], mv_meta); + // Exposed inputs match stored order, MV at index 1 is the transformed LP. + let exposed = UserDefinedLogicalNodeCore::inputs(&rebuilt); + assert_eq!(*exposed[0], base); + assert_eq!(*exposed[1], rewritten_candidate); + } + + #[test] + fn multi_mv_candidates_sort_lexicographically_by_table_ref() { + // `f_down` sorts MVs by `TableReference` (transformation-invariant). + // We reproduce the ordering step here with three candidates whose + // table_refs are deliberately in reverse-lex order to prove the sort + // brings them into canonical `mv_a < mv_b < mv_c` sequence, with + // the base still pinned at index 0. + let base = LogicalPlanBuilder::empty(false) + .build() + .expect("empty base plan"); + let mv_a_lp = LogicalPlanBuilder::empty(true).build().expect("mv_a plan"); + let mv_b_lp = LogicalPlanBuilder::empty(true).build().expect("mv_b plan"); + let mv_c_lp = LogicalPlanBuilder::empty(true).build().expect("mv_c plan"); + + // Deliberately non-sorted insertion order (as would happen with + // `mv_plans` HashMap iteration). + let mut candidates: Vec<(LogicalPlan, TableReference, RewriteReadiness)> = vec![ + ( + mv_c_lp.clone(), + TableReference::bare("catalog.schema.mv_c"), + RewriteReadiness::Ready, + ), + ( + mv_a_lp.clone(), + TableReference::bare("catalog.schema.mv_a"), + RewriteReadiness::Ready, + ), + ( + mv_b_lp.clone(), + TableReference::bare("catalog.schema.mv_b"), + RewriteReadiness::Ready, + ), + ]; + candidates.sort_by_key(|a| a.1.to_string()); + + // Base first, then sorted MVs. + let branches: Vec = std::iter::once(base.clone()) + .chain(candidates.iter().map(|(lp, _, _)| lp.clone())) + .collect(); + let metadata: Vec = std::iter::once(CandidateMetadata::Base) + .chain(candidates.into_iter().map(|(_, r, readiness)| { + CandidateMetadata::Materialized { + table_ref: r, + readiness, + } + })) + .collect(); + + assert_eq!(branches[0], base, "index 0 must be base"); + // MVs land in lex order regardless of insertion order. + assert_eq!( + metadata[1], + CandidateMetadata::Materialized { + table_ref: TableReference::bare("catalog.schema.mv_a"), + readiness: RewriteReadiness::Ready, + } + ); + assert_eq!( + metadata[2], + CandidateMetadata::Materialized { + table_ref: TableReference::bare("catalog.schema.mv_b"), + readiness: RewriteReadiness::Ready, + } + ); + assert_eq!( + metadata[3], + CandidateMetadata::Materialized { + table_ref: TableReference::bare("catalog.schema.mv_c"), + readiness: RewriteReadiness::Ready, + } + ); + } + + #[test] + fn multi_mv_sort_key_is_table_ref_not_readiness() { + // Regression guard: sorting the candidate list by + // `TableReference::to_string()` must be stable regardless of the + // readiness field. Otherwise the alignment invariant between the + // branch order and the metadata order would depend on a field the + // cost function is *also* reading — a subtle coupling. Build a + // mixed-readiness set and check that the lexicographic table_ref + // order is preserved: mv_a (Unknown), mv_b (Ready), mv_c (Unknown). + let mv_a_lp = LogicalPlanBuilder::empty(true).build().expect("mv_a plan"); + let mv_b_lp = LogicalPlanBuilder::empty(true).build().expect("mv_b plan"); + let mv_c_lp = LogicalPlanBuilder::empty(true).build().expect("mv_c plan"); + let mut candidates: Vec<(LogicalPlan, TableReference, RewriteReadiness)> = vec![ + ( + mv_c_lp, + TableReference::bare("catalog.schema.mv_c"), + RewriteReadiness::Unknown, + ), + ( + mv_a_lp, + TableReference::bare("catalog.schema.mv_a"), + RewriteReadiness::Unknown, + ), + ( + mv_b_lp, + TableReference::bare("catalog.schema.mv_b"), + RewriteReadiness::Ready, + ), + ]; + candidates.sort_by_key(|a| a.1.to_string()); + + let metadata: Vec = std::iter::once(CandidateMetadata::Base) + .chain(candidates.into_iter().map(|(_, r, readiness)| { + CandidateMetadata::Materialized { + table_ref: r, + readiness, + } + })) + .collect(); + + // Lex order by table_ref, readiness untouched. + assert_eq!( + metadata[1], + CandidateMetadata::Materialized { + table_ref: TableReference::bare("catalog.schema.mv_a"), + readiness: RewriteReadiness::Unknown, + } + ); + assert_eq!( + metadata[2], + CandidateMetadata::Materialized { + table_ref: TableReference::bare("catalog.schema.mv_b"), + readiness: RewriteReadiness::Ready, + } + ); + assert_eq!( + metadata[3], + CandidateMetadata::Materialized { + table_ref: TableReference::bare("catalog.schema.mv_c"), + readiness: RewriteReadiness::Unknown, + } + ); + } + + #[test] + fn candidate_metadata_partial_eq_discriminates_readiness() { + // `PartialEq` on `CandidateMetadata::Materialized` must consider + // the readiness field. Otherwise a cost function relying on + // `matches!(..., readiness: Ready)` or a `HashSet` + // dedup would silently conflate a `Ready` variant with an `Unknown` + // variant on the same table_ref — reintroducing the exact regression + // this PR fixes. + let ready = CandidateMetadata::Materialized { + table_ref: TableReference::bare("catalog.schema.mv_a"), + readiness: RewriteReadiness::Ready, + }; + let unknown = CandidateMetadata::Materialized { + table_ref: TableReference::bare("catalog.schema.mv_a"), + readiness: RewriteReadiness::Unknown, + }; + assert_ne!(ready, unknown); + + // Same readiness + same table_ref must still compare equal so + // dedup / cache-key uses still work. + let ready_again = CandidateMetadata::Materialized { + table_ref: TableReference::bare("catalog.schema.mv_a"), + readiness: RewriteReadiness::Ready, + }; + assert_eq!(ready, ready_again); + } + + #[test] + fn one_of_alignment_survives_rebuild_with_unknown_readiness() { + // Regression for the `Unknown` half of the readiness contract. + // `one_of_alignment_survives_with_exprs_and_inputs_rebuild` only + // covers `Ready`; this variant proves that an `Unknown` readiness + // label is preserved verbatim across a `with_exprs_and_inputs` + // rebuild too, i.e. no downstream code path erases it once a + // subsequent optimizer pass rebuilds the parent OneOf. + let base = LogicalPlanBuilder::empty(false) + .build() + .expect("empty base plan"); + let candidate = LogicalPlanBuilder::empty(true) + .build() + .expect("candidate plan"); + + let mv_meta = CandidateMetadata::Materialized { + table_ref: TableReference::bare("catalog.schema.mv_cold_lazy"), + readiness: RewriteReadiness::Unknown, + }; + let one_of = OneOf::with_rewrite_context( + vec![base.clone(), candidate.clone()], + RewriteContext::new(vec!["catalog.schema.root".to_string()]) + .with_candidates(vec![CandidateMetadata::Base, mv_meta.clone()]), + ); + + let rewritten_candidate = LogicalPlanBuilder::empty(false) + .build() + .expect("rewritten candidate plan"); + let rebuilt = UserDefinedLogicalNodeCore::with_exprs_and_inputs( + &one_of, + vec![], + vec![base.clone(), rewritten_candidate], + ) + .expect("rebuild OneOf"); + + // The Unknown-readiness metadata survived verbatim. + assert_eq!(rebuilt.rewrite_context().candidates()[1], mv_meta); + } +} + +/// End-to-end coverage for `ViewMatchingRewriter::f_down`'s `NotReady` +/// filter. Uses a `MockMv` provider whose `rewrite_readiness()` return +/// value is set per-instance so a single test can register three MVs — +/// one Ready, one NotReady, one Unknown — and assert the rewriter drops +/// only the NotReady one. +#[cfg(test)] +mod tests_view_matcher_readiness_filter { + use super::*; + use crate::materialized::{ListingTableLike, Materialized, RewriteReadiness}; + use crate::rewrite::normal_form::SpjNormalForm; + use arrow_schema::{DataType, Field, Schema, SchemaRef}; + use datafusion::catalog::{Session, TableProvider}; + use datafusion::datasource::listing::ListingTableUrl; + use datafusion::datasource::TableType; + use datafusion_common::Result; + use datafusion_expr::builder::LogicalTableSource; + use datafusion_expr::{Expr, LogicalPlan, LogicalPlanBuilder}; + use std::collections::HashMap; + use std::sync::{Arc, Mutex}; + + /// Mock Materialized provider with a caller-selectable readiness value. + /// The `query` is a trivial `TableScan(source_table)` so the rewriter's + /// `SpjNormalForm` handling can trivially match a base LP against it. + /// + /// `readiness` uses interior mutability so tests can simulate lifecycle + /// transitions (e.g. a cold lazy index reporting `Unknown` during the + /// first rewrite pass and `Ready` after the scan initialises it) without + /// rebuilding the whole `ViewMatcher`. + #[derive(Debug)] + pub(super) struct MockMv { + source_table: String, + query: LogicalPlan, + readiness: Mutex, + } + + impl MockMv { + pub(super) fn new(source_table: &str, readiness: RewriteReadiness) -> Arc { + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); + let source = Arc::new(LogicalTableSource::new(schema)); + let query = LogicalPlanBuilder::scan(source_table, source, None) + .expect("scan builder") + .build() + .expect("build MV query"); + Arc::new(Self { + source_table: source_table.to_string(), + query, + readiness: Mutex::new(readiness), + }) + } + + pub(super) fn set_readiness(&self, readiness: RewriteReadiness) { + *self.readiness.lock().expect("readiness mutex") = readiness; + } + } + + #[async_trait] + impl TableProvider for MockMv { + fn as_any(&self) -> &dyn std::any::Any { + self + } + fn schema(&self) -> SchemaRef { + Arc::new(self.query.schema().as_arrow().clone()) + } + fn table_type(&self) -> TableType { + TableType::Base + } + async fn scan( + &self, + _state: &dyn Session, + _projection: Option<&Vec>, + _filters: &[Expr], + _limit: Option, + ) -> Result> { + // Mimic a lazy-index provider whose `scan()` transitions the + // provider from cold (`Unknown`) to `Ready` because it has just + // initialised the index. Only promotes; does not overwrite an + // existing `Ready` and never toggles `NotReady` (which stays + // `NotReady` — a provider that reported unpopulated shouldn't + // silently become populated by being scanned). + if self.rewrite_readiness() == RewriteReadiness::Unknown { + self.set_readiness(RewriteReadiness::Ready); + } + Ok(Arc::new(datafusion::physical_plan::empty::EmptyExec::new( + self.schema(), + ))) + } + } + + impl ListingTableLike for MockMv { + fn table_paths(&self) -> Vec { + vec![ + ListingTableUrl::parse(format!("mem:///{}", self.source_table)).expect("parse url"), + ] + } + fn partition_columns(&self) -> Vec { + Vec::new() + } + fn file_ext(&self) -> String { + String::new() + } + } + + impl Materialized for MockMv { + fn query(&self) -> LogicalPlan { + self.query.clone() + } + fn rewrite_readiness(&self) -> RewriteReadiness { + *self.readiness.lock().expect("readiness mutex") + } + } + + /// Build a `ViewMatcher` whose `mv_plans` HashMap contains three MVs, all + /// covering the same source table, one per readiness variant. The MV + /// registered keys use the readiness name so tests can identify which + /// candidates survived by their `TableReference`. + fn build_matcher_with_three_readiness_states(source_table: &str) -> ViewMatcher { + crate::materialized::register_materialized::(); + let mut plans: HashMap, SpjNormalForm)> = + HashMap::new(); + for (name, readiness) in [ + ("mv_ready", RewriteReadiness::Ready), + ("mv_not_ready", RewriteReadiness::NotReady), + ("mv_unknown", RewriteReadiness::Unknown), + ] { + let mv = MockMv::new(source_table, readiness); + let form = SpjNormalForm::new(&mv.query).expect("MV LP is a valid SPJ"); + plans.insert( + TableReference::bare(name), + (mv.clone() as Arc, form), + ); + } + ViewMatcher::from_mv_plans_for_test(plans) + } + + /// Drive `ViewMatchingRewriter` over a `TableScan(source_table)` LP and + /// return the rewritten plan. + fn rewrite_scan_lp(matcher: &ViewMatcher, source_table: &str) -> LogicalPlan { + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); + let source = Arc::new(LogicalTableSource::new(schema)); + let query = LogicalPlanBuilder::scan(source_table, source, None) + .expect("scan builder") + .build() + .expect("build query LP"); + + query + .rewrite(&mut ViewMatchingRewriter { parent: matcher }) + .expect("rewrite scan") + .data + } + + fn candidate_table_refs(plan: &LogicalPlan) -> Vec { + let LogicalPlan::Extension(ext) = plan else { + panic!("expected OneOf extension, got {plan:?}"); + }; + let one_of = ext + .node + .as_any() + .downcast_ref::() + .expect("OneOf node"); + one_of + .rewrite_context() + .candidates() + .iter() + .filter_map(|c| match c { + CandidateMetadata::Materialized { table_ref, .. } => { + Some(table_ref.table().to_string()) + } + CandidateMetadata::Base => None, + }) + .collect() + } + + #[test] + fn f_down_excludes_not_ready_and_admits_ready_and_unknown() { + let matcher = build_matcher_with_three_readiness_states("target_t"); + let rewritten = rewrite_scan_lp(&matcher, "target_t"); + + let mv_refs = candidate_table_refs(&rewritten); + // Ready + Unknown survive, NotReady is dropped. + assert!( + mv_refs.iter().any(|r| r == "mv_ready"), + "Ready MV must be a candidate. got={mv_refs:?}" + ); + assert!( + mv_refs.iter().any(|r| r == "mv_unknown"), + "Unknown MV must be a candidate. got={mv_refs:?}" + ); + assert!( + mv_refs.iter().all(|r| r != "mv_not_ready"), + "NotReady MV must be filtered out. got={mv_refs:?}" + ); + + // Base is always index 0. + let LogicalPlan::Extension(ext) = &rewritten else { + panic!("expected Extension"); + }; + let one_of = ext + .node + .as_any() + .downcast_ref::() + .expect("OneOf node"); + assert!(matches!( + one_of.rewrite_context().candidates()[0], + CandidateMetadata::Base + )); + } + + #[test] + fn f_down_returns_original_lp_when_every_mv_is_not_ready() { + // Register only a NotReady MV — the rewriter must fall through and + // return the original scan unchanged, not wrap it in a candidate-less + // OneOf. + crate::materialized::register_materialized::(); + let mut plans: HashMap, SpjNormalForm)> = + HashMap::new(); + let mv = MockMv::new("target_t", RewriteReadiness::NotReady); + let form = SpjNormalForm::new(&mv.query).expect("valid SPJ"); + plans.insert( + TableReference::bare("mv_not_ready_only"), + (mv as Arc, form), + ); + let matcher = ViewMatcher::from_mv_plans_for_test(plans); + + let rewritten = rewrite_scan_lp(&matcher, "target_t"); + + // No OneOf wrapper — plan is the plain TableScan. + assert!( + !matches!(&rewritten, LogicalPlan::Extension(ext) if ext.node.as_any().is::()), + "no candidate MV survived, so the rewriter must return the base LP untouched. got={rewritten:?}" + ); + } + + /// Return the full metadata slice from the OneOf wrapping `plan`. + fn candidate_metadata(plan: &LogicalPlan) -> Vec { + let LogicalPlan::Extension(ext) = plan else { + panic!("expected OneOf extension, got {plan:?}"); + }; + let one_of = ext + .node + .as_any() + .downcast_ref::() + .expect("OneOf node"); + one_of.rewrite_context().candidates().to_vec() + } + + #[test] + fn f_down_encodes_ready_and_unknown_readiness_verbatim() { + // Regression test for the "readiness must reach the cost function" + // contract. Prior to this fix `f_down` collapsed `Ready` and `Unknown` + // into a single `CandidateMetadata::Materialized { table_ref }` shape, + // which meant a cost function that trusted `Materialized` candidates + // as predicate-pruned would silently promote every default-`Unknown` + // provider to `Ready`. `CandidateMetadata::Materialized` now carries + // the raw readiness so downstream policy can branch on it. + let matcher = build_matcher_with_three_readiness_states("target_t"); + let rewritten = rewrite_scan_lp(&matcher, "target_t"); + + let metadata = candidate_metadata(&rewritten); + // Base + Ready + Unknown = 3 entries; NotReady was dropped upstream. + assert_eq!(metadata.len(), 3, "expected Base + Ready + Unknown"); + assert!(matches!(metadata[0], CandidateMetadata::Base)); + + let ready = metadata + .iter() + .find_map(|c| match c { + CandidateMetadata::Materialized { + table_ref, + readiness, + } if table_ref.table() == "mv_ready" => Some(*readiness), + _ => None, + }) + .expect("Ready MV should be a candidate"); + assert_eq!(ready, RewriteReadiness::Ready); + + let unknown = metadata + .iter() + .find_map(|c| match c { + CandidateMetadata::Materialized { + table_ref, + readiness, + } if table_ref.table() == "mv_unknown" => Some(*readiness), + _ => None, + }) + .expect("Unknown MV should be a candidate"); + assert_eq!(unknown, RewriteReadiness::Unknown); + } + + #[test] + fn f_down_metadata_length_matches_branch_count() { + // Alignment invariant: `RewriteContext::candidates()` must have one + // entry per branch in the OneOf so the cost function can index into + // the two slices in lockstep. This is easy to break when a new field + // is added to `CandidateMetadata::Materialized`. + let matcher = build_matcher_with_three_readiness_states("target_t"); + let rewritten = rewrite_scan_lp(&matcher, "target_t"); + + let LogicalPlan::Extension(ext) = &rewritten else { + panic!("expected Extension"); + }; + let one_of = ext + .node + .as_any() + .downcast_ref::() + .expect("OneOf node"); + let branches = UserDefinedLogicalNodeCore::inputs(one_of); + let metadata = one_of.rewrite_context().candidates(); + assert_eq!( + branches.len(), + metadata.len(), + "branches ({}) must have the same length as metadata ({}): {:?}", + branches.len(), + metadata.len(), + metadata + ); + } + + #[test] + fn f_down_reflects_provider_readiness_on_each_rewrite() { + // What this test proves (a metadata-channel property): + // 1. `f_down` re-consults the provider's `rewrite_readiness()` on + // every rewrite pass rather than caching the first result, and + // 2. propagates the current value verbatim into + // `CandidateMetadata::Materialized::readiness`. + // + // Motivation (see PR #55 review by @xudong963): a provider that + // uses a lazy-loaded index can report `Unknown` while the index is + // cold and `Ready` once something warms it. If the rewriter cached + // the cold `Unknown` reading, a provider that later transitions to + // `Ready` would never surface that fact to the cost function, and + // policies keyed on `Ready` would never activate. + // + // What this test does *not* prove: + // * That the transition is triggered by a real scan. Here it is + // flipped explicitly via `set_readiness`. In production the + // trigger would be a lazy-index provider's `scan()` loading its + // index on first use. Full + // end-to-end coverage of that path — including the + // `lazy_index = true` cold-start behaviour @xudong963 asked + // about — belongs in the downstream provider's own tests, not + // here. + // * That a provider reporting `Unknown` in the cold phase actually + // avoids the "scan never plans, index stays cold" deadlock. That + // depends on the cost function's policy for `Unknown` candidates + // (this crate does not ship one). + crate::materialized::register_materialized::(); + let mv = MockMv::new("target_t", RewriteReadiness::Unknown); + let form = SpjNormalForm::new(&mv.query).expect("valid SPJ"); + let mut plans: HashMap, SpjNormalForm)> = + HashMap::new(); + plans.insert( + TableReference::bare("mv_cold_lazy"), + (Arc::clone(&mv) as Arc, form), + ); + let matcher = ViewMatcher::from_mv_plans_for_test(plans); + + // First rewrite: provider is reporting `Unknown` and must be + // admitted with that value in the metadata. + let cold = rewrite_scan_lp(&matcher, "target_t"); + let cold_readiness = candidate_metadata(&cold) + .iter() + .find_map(|c| match c { + CandidateMetadata::Materialized { + table_ref, + readiness, + } if table_ref.table() == "mv_cold_lazy" => Some(*readiness), + _ => None, + }) + .expect("Unknown provider must be admitted as a candidate"); + assert_eq!(cold_readiness, RewriteReadiness::Unknown); + + // Simulate the transition a lazy-index provider would perform when + // its scan warms the index. We flip the value directly here; the + // property under test is that the *next* rewrite observes the new + // value, not how the transition itself is triggered. + mv.set_readiness(RewriteReadiness::Ready); + + // Second rewrite: `f_down` must have re-consulted the provider + // (no caching) and the new readiness must be visible verbatim to + // downstream cost functions. + let warm = rewrite_scan_lp(&matcher, "target_t"); + let warm_readiness = candidate_metadata(&warm) + .iter() + .find_map(|c| match c { + CandidateMetadata::Materialized { + table_ref, + readiness, + } if table_ref.table() == "mv_cold_lazy" => Some(*readiness), + _ => None, + }) + .expect("provider that transitions to Ready must remain a candidate"); + assert_eq!(warm_readiness, RewriteReadiness::Ready); + } +} + +/// Coverage for the physical-time readiness refresh that runs from +/// `ViewExploitationPlanner::plan_extension` once DataFusion has planned +/// each candidate branch into a physical tree (which invokes each +/// provider's `scan()`, giving lazy indexes a chance to warm). +#[cfg(test)] +mod tests_physical_time_readiness_refresh { + use super::tests_view_matcher_readiness_filter::*; + use super::*; + use crate::materialized::register_materialized; + use crate::rewrite::readiness::ReadinessAnnotatedExec; + use datafusion::execution::context::SessionContext; + use datafusion_expr::LogicalPlanBuilder; + use std::sync::Arc; + + /// Pull the readiness carried by the `Materialized { table_ref: name, .. }` + /// entry out of a `RewriteContext`, panicking if it isn't present. + fn readiness_for(context: &RewriteContext, name: &str) -> RewriteReadiness { + context + .candidates() + .iter() + .find_map(|c| match c { + CandidateMetadata::Materialized { + table_ref, + readiness, + } if table_ref.table() == name => Some(*readiness), + _ => None, + }) + .unwrap_or_else(|| { + panic!( + "expected Materialized candidate for '{name}' in {:?}", + context.candidates() + ) + }) + } + + /// Build a `RewriteContext` shaped like what `f_down` would produce: a + /// Base branch plus one `Materialized` branch pointing at the + /// fully-qualified `table_ref`. + fn one_mv_context( + table_ref: TableReference, + lp_time_readiness: RewriteReadiness, + ) -> RewriteContext { + RewriteContext::new(vec!["target_t".to_string()]).with_candidates(vec![ + CandidateMetadata::Base, + CandidateMetadata::Materialized { + table_ref, + readiness: lp_time_readiness, + }, + ]) + } + + /// Build the `TableReference` that the default catalog would register a + /// mock provider under, given the caller's simple table name. + fn default_qualified_table_ref( + state: &datafusion::execution::context::SessionState, + table: &str, + ) -> TableReference { + let opts = &state.config().options().catalog; + TableReference::full( + opts.default_catalog.clone(), + opts.default_schema.clone(), + table.to_string(), + ) + } + + #[tokio::test] + async fn refresh_reflects_post_scan_transition_to_ready() { + // Simulate the lazy-index cold-start scenario end-to-end at the + // `RewriteContext` layer. Provider starts Unknown, its `scan()` is + // invoked (mimicking what physical planning would do), the scan + // impl promotes readiness to Ready, and the refresh helper picks + // up the new value from the same catalog handle. + register_materialized::(); + let ctx = SessionContext::new(); + let mv = MockMv::new("target_t", RewriteReadiness::Unknown); + ctx.register_table("mv_lazy", Arc::clone(&mv) as Arc) + .expect("register mv_lazy"); + let state = ctx.state(); + let mv_ref = default_qualified_table_ref(&state, "mv_lazy"); + + // LP-time snapshot: cold. + let context = one_mv_context(mv_ref, RewriteReadiness::Unknown); + + // Physical planning would have called scan() on the MV branch by + // now; call it directly to simulate that side effect. + let session_for_scan = state.clone(); + mv.scan(&session_for_scan, None, &[], None) + .await + .expect("scan flips readiness"); + + // Now refresh — should observe Ready. + let refreshed = refresh_candidate_readiness(context, &[], &state).await; + assert_eq!( + readiness_for(&refreshed, "mv_lazy"), + RewriteReadiness::Ready + ); + } + + #[tokio::test] + async fn refresh_falls_back_to_lp_time_value_when_provider_missing() { + // Provider is not registered under the recorded `table_ref`. The + // helper must return the caller's LP-time value verbatim rather + // than defaulting to Unknown, so a transient catalog hiccup can't + // silently downgrade a previously-observed Ready to Unknown. + let ctx = SessionContext::new(); + let context = one_mv_context( + TableReference::full("datafusion", "public", "does_not_exist"), + RewriteReadiness::Ready, + ); + + let refreshed = refresh_candidate_readiness(context, &[], &ctx.state()).await; + assert_eq!( + readiness_for(&refreshed, "does_not_exist"), + RewriteReadiness::Ready, + "missing catalog entry must not overwrite LP-time readiness" + ); + } + + #[tokio::test] + async fn refresh_promotes_unknown_candidate_to_not_ready_when_scan_uncovers_unpopulated() { + // Complements `refresh_reflects_post_scan_transition_to_ready` + // (Unknown -> Ready). Here we cover Unknown -> NotReady, which is + // the shape a lazy-index provider would use to say "scan tried + // and confirmed there's no data". `drop_not_ready_after_refresh` + // then removes the candidate before `OneOfExec::try_new` invokes + // the cost function, preserving the invariant that NotReady never + // reaches cost policy. + register_materialized::(); + let ctx = SessionContext::new(); + let mv = MockMv::new("target_t", RewriteReadiness::Unknown); + ctx.register_table("mv_regressed", Arc::clone(&mv) as Arc) + .expect("register mv_regressed"); + let state = ctx.state(); + let mv_ref = default_qualified_table_ref(&state, "mv_regressed"); + + let context = one_mv_context(mv_ref, RewriteReadiness::Unknown); + + // Simulated scan concluded that the MV is unpopulated. + mv.set_readiness(RewriteReadiness::NotReady); + + let refreshed = refresh_candidate_readiness(context, &[], &state).await; + assert_eq!( + readiness_for(&refreshed, "mv_regressed"), + RewriteReadiness::NotReady + ); + } + + #[tokio::test] + async fn refresh_leaves_base_variant_untouched() { + // `CandidateMetadata::Base` has no provider; refresh must skip it. + let ctx = SessionContext::new(); + let context = RewriteContext::new(vec!["target_t".to_string()]) + .with_candidates(vec![CandidateMetadata::Base]); + + let refreshed = refresh_candidate_readiness(context, &[], &ctx.state()).await; + assert_eq!(refreshed.candidates().len(), 1); + assert!(matches!(refreshed.candidates()[0], CandidateMetadata::Base)); + } + + #[test] + fn drop_not_ready_filter_keeps_ready_and_unknown_and_base() { + use arrow_schema::Schema; + use datafusion::physical_plan::empty::EmptyExec; + let schema = Arc::new(Schema::empty()); + let make = || Arc::new(EmptyExec::new(Arc::clone(&schema))) as Arc; + let physical_inputs = vec![make(), make(), make(), make()]; + let candidates = vec![ + CandidateMetadata::Base, + CandidateMetadata::Materialized { + table_ref: TableReference::bare("a"), + readiness: RewriteReadiness::Ready, + }, + CandidateMetadata::Materialized { + table_ref: TableReference::bare("b"), + readiness: RewriteReadiness::NotReady, + }, + CandidateMetadata::Materialized { + table_ref: TableReference::bare("c"), + readiness: RewriteReadiness::Unknown, + }, + ]; + + let (kept_inputs, kept_candidates) = + drop_not_ready_after_refresh(&physical_inputs, &candidates); + assert_eq!(kept_inputs.len(), 3, "NotReady branch must be dropped"); + assert_eq!(kept_candidates.len(), 3); + assert!(matches!(kept_candidates[0], CandidateMetadata::Base)); + assert!(matches!( + &kept_candidates[1], + CandidateMetadata::Materialized { table_ref, readiness: RewriteReadiness::Ready } + if table_ref.table() == "a" + )); + assert!(matches!( + &kept_candidates[2], + CandidateMetadata::Materialized { table_ref, readiness: RewriteReadiness::Unknown } + if table_ref.table() == "c" + )); + } + + #[test] + fn drop_not_ready_filter_passes_through_when_metadata_missing() { + // Older callsites construct `OneOf` without per-branch metadata. + // The filter must not corrupt those: with a length mismatch, hand + // the inputs back unchanged so downstream keeps working. + use arrow_schema::Schema; + use datafusion::physical_plan::empty::EmptyExec; + let schema = Arc::new(Schema::empty()); + let make = || Arc::new(EmptyExec::new(Arc::clone(&schema))) as Arc; + let physical_inputs = vec![make(), make()]; + let candidates: Vec = Vec::new(); + + let (kept_inputs, kept_candidates) = + drop_not_ready_after_refresh(&physical_inputs, &candidates); + assert_eq!(kept_inputs.len(), 2); + assert!(kept_candidates.is_empty()); + } + + #[test] + fn drop_not_ready_filter_drops_misaligned_metadata_instead_of_propagating_it() { + // Regression for the "misalignment-poisons-downstream" concern: if a + // caller supplies a non-empty but wrong-length metadata slice, the + // filter must not carry that slice through to `OneOfExec`, or the + // cost function would end up indexing candidates by the branch + // position of an unrelated candidate. We keep the physical inputs + // intact so the OneOfExec still builds, but the metadata is dropped + // so downstream sees the same shape as the pre-readiness code path. + use arrow_schema::Schema; + use datafusion::physical_plan::empty::EmptyExec; + let schema = Arc::new(Schema::empty()); + let make = || Arc::new(EmptyExec::new(Arc::clone(&schema))) as Arc; + // 3 physical inputs but only 2 metadata entries — mismatched. + let physical_inputs = vec![make(), make(), make()]; + let candidates = vec![ + CandidateMetadata::Base, + CandidateMetadata::Materialized { + table_ref: TableReference::bare("a"), + readiness: RewriteReadiness::Ready, + }, + ]; + + let (kept_inputs, kept_candidates) = + drop_not_ready_after_refresh(&physical_inputs, &candidates); + assert_eq!( + kept_inputs.len(), + 3, + "physical inputs must survive so OneOfExec still builds" + ); + assert!( + kept_candidates.is_empty(), + "misaligned metadata must not be propagated to downstream cost functions" + ); + } + + #[test] + fn drop_not_ready_filter_collapses_all_mvs_to_base_only() { + // If every MV branch transitions to `NotReady`, only Base survives. + // The `plan_extension` short-circuit uses this exact shape to skip + // the OneOfExec wrapper and return the base plan directly. + use arrow_schema::Schema; + use datafusion::physical_plan::empty::EmptyExec; + let schema = Arc::new(Schema::empty()); + let make = || Arc::new(EmptyExec::new(Arc::clone(&schema))) as Arc; + let physical_inputs = vec![make(), make(), make()]; + let candidates = vec![ + CandidateMetadata::Base, + CandidateMetadata::Materialized { + table_ref: TableReference::bare("a"), + readiness: RewriteReadiness::NotReady, + }, + CandidateMetadata::Materialized { + table_ref: TableReference::bare("b"), + readiness: RewriteReadiness::NotReady, + }, + ]; + + let (kept_inputs, kept_candidates) = + drop_not_ready_after_refresh(&physical_inputs, &candidates); + assert_eq!(kept_inputs.len(), 1); + assert_eq!(kept_candidates.len(), 1); + assert!(matches!(kept_candidates[0], CandidateMetadata::Base)); + } + + /// Bare-bones `PhysicalPlanner` for `plan_extension` end-to-end tests. + /// `plan_extension` never invokes the planner argument (the field is + /// `_planner`), so this stub can panic in every method. + struct StubPlanner; + + #[async_trait::async_trait] + impl datafusion::physical_planner::PhysicalPlanner for StubPlanner { + async fn create_physical_plan( + &self, + _logical_plan: &LogicalPlan, + _session_state: &datafusion::execution::context::SessionState, + ) -> datafusion_common::Result> { + unimplemented!("StubPlanner::create_physical_plan not exercised") + } + fn create_physical_expr( + &self, + _expr: &datafusion_expr::Expr, + _input_dfschema: &datafusion_common::DFSchema, + _session_state: &datafusion::execution::context::SessionState, + ) -> datafusion_common::Result> { + unimplemented!("StubPlanner::create_physical_expr not exercised") + } + } + + fn make_empty_exec() -> Arc { + Arc::new(datafusion::physical_plan::empty::EmptyExec::new(Arc::new( + arrow_schema::Schema::empty(), + ))) as Arc + } + + #[tokio::test] + async fn plan_extension_updates_metadata_after_scan_transition() { + // End-to-end path through `ViewExploitationPlanner::plan_extension`: + // provider is Unknown at LP-time, the (mock) scan promotes it to + // Ready, and the returned OneOfExec's metadata reflects the + // post-scan value. This is the cold-lazy scenario Xudong called + // out, wired through the actual `plan_extension` composition of + // refresh + filter + OneOfExec construction. + register_materialized::(); + let ctx = SessionContext::new(); + let mv = MockMv::new("target_t", RewriteReadiness::Unknown); + ctx.register_table("mv_lazy", Arc::clone(&mv) as Arc) + .expect("register mv_lazy"); + let state = ctx.state(); + + // Simulate the scan that physical planning would have performed: + // MockMv::scan promotes readiness Unknown -> Ready. + mv.scan(&state, None, &[], None) + .await + .expect("scan flips readiness"); + + let mv_ref = default_qualified_table_ref(&state, "mv_lazy"); + let base_lp = LogicalPlanBuilder::empty(false).build().expect("base lp"); + let mv_lp = LogicalPlanBuilder::empty(false).build().expect("mv lp"); + let one_of_node = OneOf::with_rewrite_context( + vec![base_lp.clone(), mv_lp.clone()], + RewriteContext::new(vec!["target_t".to_string()]).with_candidates(vec![ + CandidateMetadata::Base, + CandidateMetadata::Materialized { + table_ref: mv_ref.clone(), + readiness: RewriteReadiness::Unknown, + }, + ]), + ); + + let base_phys = make_empty_exec(); + let mv_phys = make_empty_exec(); + + let cost_fn: CostFn = Arc::new(|ctx| ctx.into_candidate_plans().map(|_| 1.0_f64).collect()); + let planner = ViewExploitationPlanner::new(cost_fn); + let stub = StubPlanner; + + let result = planner + .plan_extension( + &stub, + &one_of_node, + &[&base_lp, &mv_lp], + &[Arc::clone(&base_phys), Arc::clone(&mv_phys)], + &state, + ) + .await + .expect("plan_extension") + .expect("plan_extension returned Some"); + + let one_of_exec = result + .as_any() + .downcast_ref::() + .expect("result must be OneOfExec after refresh"); + let candidates = one_of_exec.rewrite_context().candidates(); + assert_eq!(candidates.len(), 2, "Base + refreshed Materialized"); + assert!(matches!(candidates[0], CandidateMetadata::Base)); + let (fetched_ref, fetched_readiness) = candidates + .iter() + .find_map(|c| match c { + CandidateMetadata::Materialized { + table_ref, + readiness, + } => Some((table_ref.clone(), *readiness)), + _ => None, + }) + .expect("expected Materialized entry"); + assert_eq!(fetched_ref, mv_ref); + assert_eq!( + fetched_readiness, + RewriteReadiness::Ready, + "post-scan refresh must have promoted Unknown -> Ready" + ); + } + + #[tokio::test] + async fn plan_extension_short_circuits_to_base_when_all_mvs_go_not_ready() { + // Provider was `Unknown` at LP-time (cold lazy). `scan()` runs + // during physical planning and discovers the MV is unpopulated, + // so the provider now reports `NotReady`. `plan_extension` must + // refresh (Unknown triggers the refresh path), pick up the new + // `NotReady`, filter out the MV, and — because only Base survives + // — return the base physical plan directly instead of wrapping + // it in an otherwise-useless single-branch OneOfExec. + register_materialized::(); + let ctx = SessionContext::new(); + let mv = MockMv::new("target_t", RewriteReadiness::Unknown); + ctx.register_table("mv_regressed", Arc::clone(&mv) as Arc) + .expect("register mv_regressed"); + let state = ctx.state(); + + // Simulate scan discovering the MV is unusable. + mv.set_readiness(RewriteReadiness::NotReady); + + let mv_ref = default_qualified_table_ref(&state, "mv_regressed"); + let base_lp = LogicalPlanBuilder::empty(false).build().expect("base lp"); + let mv_lp = LogicalPlanBuilder::empty(false).build().expect("mv lp"); + let one_of_node = OneOf::with_rewrite_context( + vec![base_lp.clone(), mv_lp.clone()], + RewriteContext::new(vec!["target_t".to_string()]).with_candidates(vec![ + CandidateMetadata::Base, + CandidateMetadata::Materialized { + table_ref: mv_ref, + readiness: RewriteReadiness::Unknown, + }, + ]), + ); + + let base_phys = make_empty_exec(); + let mv_phys = make_empty_exec(); + + let cost_fn: CostFn = Arc::new(|ctx| ctx.into_candidate_plans().map(|_| 1.0_f64).collect()); + let planner = ViewExploitationPlanner::new(cost_fn); + let stub = StubPlanner; + + let result = planner + .plan_extension( + &stub, + &one_of_node, + &[&base_lp, &mv_lp], + &[Arc::clone(&base_phys), Arc::clone(&mv_phys)], + &state, + ) + .await + .expect("plan_extension") + .expect("plan_extension returned Some"); + + assert!( + result.as_any().downcast_ref::().is_none(), + "should short-circuit to base plan, not wrap in OneOfExec" + ); + assert!( + Arc::ptr_eq(&result, &base_phys), + "short-circuit must return the base physical input verbatim" + ); + } + + #[tokio::test] + async fn refresh_does_not_move_base_off_index_zero() { + // The alignment invariant on `CandidateMetadata` says Base sits at + // index 0. Refresh iterates in order and rebuilds the vector, so + // it should preserve position. + register_materialized::(); + let ctx = SessionContext::new(); + let mv_a = MockMv::new("target_t", RewriteReadiness::Ready); + let mv_b = MockMv::new("target_t", RewriteReadiness::Unknown); + ctx.register_table("mv_a", Arc::clone(&mv_a) as Arc) + .expect("register mv_a"); + ctx.register_table("mv_b", Arc::clone(&mv_b) as Arc) + .expect("register mv_b"); + let state = ctx.state(); + let mv_a_ref = default_qualified_table_ref(&state, "mv_a"); + let mv_b_ref = default_qualified_table_ref(&state, "mv_b"); + + let context = RewriteContext::new(vec!["target_t".to_string()]).with_candidates(vec![ + CandidateMetadata::Base, + CandidateMetadata::Materialized { + table_ref: mv_a_ref, + readiness: RewriteReadiness::Ready, + }, + CandidateMetadata::Materialized { + table_ref: mv_b_ref, + readiness: RewriteReadiness::Unknown, + }, + ]); + + let refreshed = refresh_candidate_readiness(context, &[], &state).await; + assert!(matches!(refreshed.candidates()[0], CandidateMetadata::Base)); + assert_eq!(refreshed.candidates().len(), 3); + } + + #[tokio::test] + async fn refresh_prefers_annotated_readiness_over_catalog_sampling() { + // A provider that wraps its scan output in `ReadinessAnnotatedExec` + // captures readiness atomically at scan time. Refresh must read + // that captured value verbatim and NOT re-sample the provider — + // sampling is racy against concurrent snapshot swaps. + // + // Set up the scenario where the two values disagree: + // * physical_input carries an annotated `Ready` (what scan saw) + // * the live provider now reports `Unknown` (as if a concurrent + // swap happened after scan returned) + // The annotated value must win. + use arrow_schema::Schema; + use datafusion::physical_plan::empty::EmptyExec; + register_materialized::(); + let ctx = SessionContext::new(); + let mv = MockMv::new("target_t", RewriteReadiness::Unknown); + ctx.register_table("mv_annotated", Arc::clone(&mv) as Arc) + .expect("register mv_annotated"); + let state = ctx.state(); + let mv_ref = default_qualified_table_ref(&state, "mv_annotated"); + + let context = one_mv_context(mv_ref, RewriteReadiness::Unknown); + let schema = Arc::new(Schema::empty()); + let base_input: Arc = Arc::new(EmptyExec::new(Arc::clone(&schema))); + let inner: Arc = Arc::new(EmptyExec::new(schema)); + let annotated_input: Arc = + Arc::new(ReadinessAnnotatedExec::new(inner, RewriteReadiness::Ready)); + + // Provider says Unknown, annotated input says Ready — annotated wins. + let refreshed = + refresh_candidate_readiness(context, &[base_input, annotated_input], &state).await; + assert_eq!( + readiness_for(&refreshed, "mv_annotated"), + RewriteReadiness::Ready, + "annotated readiness must override sampled value" + ); + } + + #[tokio::test] + async fn refresh_falls_back_to_sampling_when_input_not_annotated() { + // Backward-compat: providers that haven't opted into + // `ReadinessAnnotatedExec` still get the old sampling behaviour. + // Here the provider is unwrapped and the scan (mock) promotes + // readiness from Unknown to Ready — refresh must observe it via + // catalog sampling. + register_materialized::(); + let ctx = SessionContext::new(); + let mv = MockMv::new("target_t", RewriteReadiness::Unknown); + ctx.register_table("mv_unwrapped", Arc::clone(&mv) as Arc) + .expect("register mv_unwrapped"); + let state = ctx.state(); + let mv_ref = default_qualified_table_ref(&state, "mv_unwrapped"); + + let context = one_mv_context(mv_ref, RewriteReadiness::Unknown); + // Simulated warmup — provider now reports Ready. + mv.set_readiness(RewriteReadiness::Ready); + + // Plain (unwrapped) physical_input from the "provider" side. + use arrow_schema::Schema; + use datafusion::physical_plan::empty::EmptyExec; + let schema = Arc::new(Schema::empty()); + let base_input: Arc = Arc::new(EmptyExec::new(Arc::clone(&schema))); + let mv_input: Arc = Arc::new(EmptyExec::new(schema)); + + let refreshed = refresh_candidate_readiness(context, &[base_input, mv_input], &state).await; + assert_eq!( + readiness_for(&refreshed, "mv_unwrapped"), + RewriteReadiness::Ready, + "sampling fallback must find the warmed provider" + ); + } + + #[test] + fn readiness_from_plan_walks_nested_children() { + // `ReadinessAnnotatedExec` doesn't have to be at the top of the + // physical tree — DataFusion may wrap the scan's output in + // e.g. a `RepartitionExec` before it reaches `plan_extension`. + // The walker must reach an annotation nested any depth down. + use arrow_schema::Schema; + use datafusion::physical_plan::empty::EmptyExec; + use datafusion::physical_plan::repartition::RepartitionExec; + use datafusion_physical_expr::Partitioning; + let schema = Arc::new(Schema::empty()); + let leaf: Arc = Arc::new(EmptyExec::new(schema)); + let annotated: Arc = + Arc::new(ReadinessAnnotatedExec::new(leaf, RewriteReadiness::Ready)); + let wrapped: Arc = Arc::new( + RepartitionExec::try_new(annotated, Partitioning::RoundRobinBatch(1)) + .expect("repartition"), + ); + + let found = readiness_from_plan(&wrapped); + assert_eq!(found, Some(RewriteReadiness::Ready)); + } + + #[test] + fn readiness_from_plan_returns_none_when_no_annotation() { + use arrow_schema::Schema; + use datafusion::physical_plan::empty::EmptyExec; + let schema = Arc::new(Schema::empty()); + let plan: Arc = Arc::new(EmptyExec::new(schema)); + assert_eq!(readiness_from_plan(&plan), None); + } + + #[tokio::test] + async fn plan_extension_strips_readiness_annotation_before_returning() { + // Regression test for the concern that a leaked + // `ReadinessAnnotatedExec` would appear in the plan handed to the + // cost function / OneOfExec / downstream optimizers and block + // trait-method-driven optimizations (limit pushdown, projection + // pushdown, with_preserve_order, ...). After plan_extension runs, + // no ReadinessAnnotatedExec must remain anywhere in the resulting + // OneOfExec's child subtrees. + use arrow_schema::Schema; + use datafusion::physical_plan::empty::EmptyExec; + register_materialized::(); + let ctx = SessionContext::new(); + let mv = MockMv::new("target_t", RewriteReadiness::Ready); + ctx.register_table("mv_wrapped", Arc::clone(&mv) as Arc) + .expect("register mv_wrapped"); + let state = ctx.state(); + let mv_ref = default_qualified_table_ref(&state, "mv_wrapped"); + let base_lp = LogicalPlanBuilder::empty(false).build().expect("base lp"); + let mv_lp = LogicalPlanBuilder::empty(false).build().expect("mv lp"); + let one_of_node = OneOf::with_rewrite_context( + vec![base_lp.clone(), mv_lp.clone()], + RewriteContext::new(vec!["target_t".to_string()]).with_candidates(vec![ + CandidateMetadata::Base, + CandidateMetadata::Materialized { + table_ref: mv_ref, + readiness: RewriteReadiness::Ready, + }, + ]), + ); + + let schema = Arc::new(Schema::empty()); + let base_phys: Arc = Arc::new(EmptyExec::new(Arc::clone(&schema))); + // Provider wraps its scan output — plan_extension must strip this + // before building the OneOfExec. + let mv_phys: Arc = Arc::new(ReadinessAnnotatedExec::new( + Arc::new(EmptyExec::new(schema)), + RewriteReadiness::Ready, + )); + + let cost_fn: CostFn = Arc::new(|ctx| ctx.into_candidate_plans().map(|_| 1.0_f64).collect()); + let planner = ViewExploitationPlanner::new(cost_fn); + let stub = StubPlanner; + + let result = planner + .plan_extension( + &stub, + &one_of_node, + &[&base_lp, &mv_lp], + &[Arc::clone(&base_phys), Arc::clone(&mv_phys)], + &state, + ) + .await + .expect("plan_extension") + .expect("plan_extension returned Some"); + + // No ReadinessAnnotatedExec anywhere in the returned plan tree. + assert!( + readiness_from_plan(&result).is_none(), + "ReadinessAnnotatedExec must be stripped before plan_extension returns" + ); + } +} diff --git a/src/rewrite/readiness.rs b/src/rewrite/readiness.rs new file mode 100644 index 0000000..7d3ba41 --- /dev/null +++ b/src/rewrite/readiness.rs @@ -0,0 +1,570 @@ +// 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. + +//! Readiness-related types and helpers used by the query-rewrite optimizer. +//! +//! Two-stage readiness contract: +//! +//! 1. At LP-rewrite time, `ViewMatcher` consults +//! [`Materialized::rewrite_readiness`](crate::materialized::Materialized::rewrite_readiness) +//! and drops candidates whose readiness is [`RewriteReadiness::NotReady`]. +//! `Ready` and `Unknown` both enter the candidate set and their raw +//! readiness is recorded on [`CandidateMetadata::Materialized::readiness`]. +//! 2. At physical-planning time (once `TableProvider::scan()` has run on every +//! branch), `ViewExploitationPlanner` refreshes the readiness on each +//! candidate. Providers that opt into [`ReadinessAnnotatedExec`] have +//! their readiness read straight off the plan tree (atomically captured +//! at scan time); others fall back to sampling the current provider +//! state, which is best-effort but racy under concurrent snapshot swaps. +//! +//! The cost function only ever sees `Base` or `Materialized { Ready | Unknown }` +//! entries — `NotReady` is filtered by both gates before it can reach cost +//! policy. + +use std::sync::Arc; + +use datafusion::catalog::CatalogProviderList; +use datafusion::execution::context::SessionState; +use datafusion::execution::TaskContext; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, SendableRecordBatchStream, +}; +use datafusion_common::tree_node::{Transformed, TreeNode}; +use datafusion_common::{DataFusionError, Result, TableReference}; + +use crate::materialized::cast_to_materialized; + +use super::exploitation::RewriteContext; + +/// Whether a materialized view is currently safe to route queries to. Reported +/// by [`Materialized::rewrite_readiness`](crate::materialized::Materialized::rewrite_readiness) +/// and consulted by +/// [`ViewMatcher`](crate::rewrite::exploitation::ViewMatcher) during LP rewrite +/// so that unpopulated / in-flight MVs are excluded from the candidate set +/// upstream of the cost function. +/// +/// Keeping this a lifecycle abstraction (rather than a proxy such as file +/// count) means the trait doesn't couple to any specific storage layout — +/// providers describe their own readiness however they want (index loaded, +/// snapshot published, migration complete, staleness threshold satisfied, +/// etc.) and only report the answer. +// `PartialOrd`/`Ord` are derived so types that embed `RewriteReadiness` (e.g. +// candidate metadata) can derive ordering traits when needed. The variant +// ordering has no lifecycle meaning — callers must not depend on +// `Ready < NotReady < Unknown`. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub enum RewriteReadiness { + /// The MV is populated and can safely answer the query. The + /// `ViewMatchingRewriter` will include it as a rewrite candidate. + Ready, + /// The MV should not be used yet (index not loaded, ingest task never + /// ran after a version bump, snapshot rebuild in progress, etc.). + /// The `ViewMatchingRewriter` will drop the MV from the candidate set + /// so a query never gets routed to it and silently returns empty. + NotReady, + /// The provider cannot cheaply determine readiness. The + /// `ViewMatchingRewriter` treats this as "include as candidate" and + /// propagates the `Unknown` value to the cost function via + /// `CandidateMetadata::Materialized { readiness, .. }`, so the caller + /// can pick whatever policy fits — e.g. fall back to the base scan + /// cost rather than trust an EmptyExec candidate as predicate-pruned. + /// Default value returned by the trait's blanket impl so + /// backward-compatible providers keep the pre-existing "always a + /// candidate" behaviour. + Unknown, +} + +/// Per-branch metadata inside a `OneOf` / `OneOfExec`. One variant per branch, +/// aligned by index with the containing `branches` / `candidates` vector. +/// +/// Lifecycle judgement (populated / unpopulated / stale / ...) is made in two +/// stages via [`Materialized::rewrite_readiness`](crate::materialized::Materialized::rewrite_readiness): +/// +/// 1. At LP-rewrite time, +/// [`ViewMatcher`](crate::rewrite::exploitation::ViewMatcher) drops `NotReady` +/// providers and admits `Ready` + `Unknown` as candidates. +/// 2. At physical-planning time, once DataFusion has invoked each provider's +/// `scan()` (giving lazy indexes a chance to warm), +/// [`ViewExploitationPlanner`](crate::rewrite::exploitation::ViewExploitationPlanner) +/// re-consults `rewrite_readiness()` and updates the metadata to the current +/// value; providers that transitioned to `NotReady` between the two stages are +/// dropped here. +/// +/// So any `Materialized` variant that reaches the cost function reflects the +/// **post-scan** readiness, and `NotReady` is guaranteed to have been filtered +/// (upstream or downstream of scan) before the cost function runs. `Ready` +/// and `Unknown` remain distinguishable so cost policy for the two can diverge +/// (typically: trust `Ready` as safe to route, and fall back to a conservative +/// estimate for `Unknown`). +/// +/// The invariant enforced at construction inside +/// [`ViewMatcher`](crate::rewrite::exploitation::ViewMatcher): +/// +/// * Index 0 is always [`CandidateMetadata::Base`] — the query's original LP, +/// with no MV rewrite applied. `OneOf::schema()` reads `branches[0].schema()` +/// and therefore always exposes the query's schema (not an MV's). +/// * Indices 1..N are [`CandidateMetadata::Materialized`] entries, sorted +/// deterministically by `table_ref`. Sorting on the MV's registered +/// `TableReference` (rather than on the branch LP) keeps the order stable +/// under downstream LP transformations that happen between LP rewrite and +/// physical planning — identity projection elimination, always-true filter +/// folding, etc. change the LP but never the MV's registered name, so the +/// alignment between this metadata and the physical candidate the cost +/// function receives survives every optimizer pass. +#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Hash)] +pub enum CandidateMetadata { + /// The query's original branch (index 0 in every OneOf). + Base, + /// A materialized-view rewrite of the query. Only MVs that reported + /// `Ready` or `Unknown` readiness reach this state; `NotReady` MVs + /// are filtered upstream by + /// [`ViewMatcher`](crate::rewrite::exploitation::ViewMatcher). + /// Cost functions consult the `readiness` field to distinguish the + /// two cases: `Unknown` (the trait default returned by + /// [`RewriteReadiness`]) must not be silently treated as `Ready`, + /// since that would regress providers that haven't declared a + /// lifecycle. + Materialized { + /// Registered `TableReference` of the source MV. Stored directly + /// rather than as `.to_string()` so identifiers with dots or + /// quoting round-trip losslessly through the physical-time + /// catalog resolution. Stable across LP transformations; safe to + /// use as a sort key. + table_ref: TableReference, + /// Provider readiness as observed at physical-planning time — i.e. + /// after DataFusion has invoked `TableProvider::scan()` on this + /// branch. Set once at LP rewrite from the provider's initial + /// report and refreshed from the same provider inside + /// [`ViewExploitationPlanner::plan_extension`](crate::rewrite::exploitation::ViewExploitationPlanner) + /// so lazy-index providers surface their warmed-up value. + /// Downstream cost functions branch on this to implement their + /// `Unknown` policy (e.g. fall back to the base scan cost rather + /// than trust an EmptyExec as predicate-pruned). + readiness: RewriteReadiness, + }, +} + +/// Wraps the `ExecutionPlan` returned by a `Materialized` provider's +/// `scan()` together with the readiness value captured at scan time. +/// +/// Providers use this to close the race between "which snapshot did +/// scan read from" and "what does `rewrite_readiness()` return now". +/// Sampling `rewrite_readiness()` again after `scan()` has returned is +/// racy — a concurrent snapshot swap can publish new state between the +/// two calls, letting an `EmptyExec` from the old snapshot be labelled +/// with readiness from the new one. If the provider computes both +/// atomically (from the same state Arc, under the same read lock, etc.) +/// and wraps the returned plan with `ReadinessAnnotatedExec`, the +/// physical-time refresh in +/// [`ViewExploitationPlanner::plan_extension`](crate::rewrite::exploitation::ViewExploitationPlanner) +/// reads the annotated value verbatim instead of re-sampling. +/// +/// The wrapper is transparent: all `ExecutionPlan` methods delegate to +/// the inner plan, so downstream operators see the same shape as if the +/// provider returned `inner` directly. +/// +/// Providers that don't wrap fall back to the pre-existing (racy) +/// sampling path in the refresh, which remains supported for backward +/// compatibility. +#[derive(Debug)] +pub struct ReadinessAnnotatedExec { + inner: Arc, + readiness: RewriteReadiness, + properties: PlanProperties, +} + +impl ReadinessAnnotatedExec { + /// Wrap `inner` with the readiness captured atomically at scan time. + pub fn new(inner: Arc, readiness: RewriteReadiness) -> Self { + let properties = inner.properties().clone(); + Self { + inner, + readiness, + properties, + } + } + + /// Readiness captured at the moment the provider produced `inner`. + pub fn readiness(&self) -> RewriteReadiness { + self.readiness + } + + /// The wrapped plan. + pub fn inner(&self) -> &Arc { + &self.inner + } +} + +impl DisplayAs for ReadinessAnnotatedExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!(f, "ReadinessAnnotatedExec: readiness={:?}", self.readiness) + } + DisplayFormatType::TreeRender => Ok(()), + } + } +} + +impl ExecutionPlan for ReadinessAnnotatedExec { + fn name(&self) -> &str { + "ReadinessAnnotatedExec" + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn properties(&self) -> &PlanProperties { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.inner] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + if children.len() != 1 { + return Err(DataFusionError::Plan(format!( + "ReadinessAnnotatedExec expects exactly one child, got {}", + children.len() + ))); + } + Ok(Arc::new(ReadinessAnnotatedExec::new( + children.into_iter().next().unwrap(), + self.readiness, + ))) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + self.inner.execute(partition, context) + } + + fn partition_statistics( + &self, + partition: Option, + ) -> Result { + self.inner.partition_statistics(partition) + } +} + +/// Depth-first search for the nearest `ReadinessAnnotatedExec` in a +/// physical plan tree. Returns its readiness value, or `None` if the +/// tree doesn't contain one (e.g. the provider hasn't opted in). Used +/// by the physical-time refresh to prefer atomically-captured readiness +/// over racy re-sampling. +pub fn readiness_from_plan(plan: &Arc) -> Option { + if let Some(annotated) = plan.as_any().downcast_ref::() { + return Some(annotated.readiness()); + } + for child in plan.children() { + if let Some(readiness) = readiness_from_plan(child) { + return Some(readiness); + } + } + None +} + +/// Remove every `ReadinessAnnotatedExec` from a physical plan tree and +/// return the plan with each wrapper collapsed to its inner. Used by +/// `ViewExploitationPlanner::plan_extension` after `readiness_from_plan` +/// has extracted the annotation — the wrapper's job is done at that +/// point and it shouldn't leak into the plan handed to downstream +/// optimizers, cost functions, `OneOfExec`, or physical-plan codecs. +/// +/// The wrapper is deliberately minimal (it only delegates `properties`, +/// `execute`, and `partition_statistics`); leaving it in place would +/// silently block optimizer passes such as limit pushdown, projection +/// pushdown, or `with_preserve_order`, which look at +/// `ExecutionPlan` trait methods we don't proxy. +pub fn strip_readiness_annotation(plan: Arc) -> Result> { + plan.transform(&|node: Arc| { + if let Some(annotated) = node.as_any().downcast_ref::() { + Ok(Transformed::yes(Arc::clone(annotated.inner()))) + } else { + Ok(Transformed::no(node)) + } + }) + .map(|t| t.data) +} + +/// Pair-wise filter that drops every `Materialized` candidate whose +/// refreshed readiness is `NotReady`, along with its aligned entry in +/// `physical_inputs`. Base is never filtered. +/// +/// If `candidates` is empty or is not aligned with `physical_inputs` +/// (older callsites that constructed the `OneOf` without per-branch +/// metadata, or a caller that supplied a mismatched vector), the inputs +/// are passed through unchanged and the metadata is dropped entirely — +/// preserving a misaligned slice would let it flow into `OneOfExec` and +/// attribute readiness to the wrong branch inside the cost function. +/// Callsites without metadata keep their original behaviour (the empty +/// slice is what the pre-readiness code path also carried). +pub(super) fn drop_not_ready_after_refresh( + physical_inputs: &[Arc], + candidates: &[CandidateMetadata], +) -> (Vec>, Vec) { + if candidates.len() != physical_inputs.len() { + return (physical_inputs.to_vec(), Vec::new()); + } + let mut kept_inputs = Vec::with_capacity(physical_inputs.len()); + let mut kept_candidates = Vec::with_capacity(candidates.len()); + for (input, cand) in physical_inputs.iter().zip(candidates.iter()) { + let drop = matches!( + cand, + CandidateMetadata::Materialized { + readiness: RewriteReadiness::NotReady, + .. + } + ); + if drop { + continue; + } + kept_inputs.push(Arc::clone(input)); + kept_candidates.push(cand.clone()); + } + (kept_inputs, kept_candidates) +} + +/// Re-consult every `Materialized` candidate's `rewrite_readiness()` and +/// return a `RewriteContext` whose metadata reflects the current values. +/// +/// Called from `ViewExploitationPlanner::plan_extension` once +/// physical inputs have been planned (i.e. after `TableProvider::scan()` +/// has run on every candidate branch). Providers whose readiness only +/// becomes definite after scan initialization (lazy indexes, warmup jobs, +/// snapshot swaps that happen inside `scan`) surface their new value here +/// so the cost function sees the definitive readiness rather than the +/// stale LP-time snapshot. +/// +/// For each `Materialized` candidate the refresh prefers the atomic +/// readiness captured at scan time (via `ReadinessAnnotatedExec` in the +/// candidate's `physical_input`) over racy re-sampling. Providers that +/// haven't opted into the wrapper fall back to sampling the current +/// `rewrite_readiness()` through the catalog — this is best-effort and +/// can be stale under concurrent snapshot swaps (see `ReadinessAnnotatedExec`). +/// +/// Lookup failures on the fallback path (table not found in the +/// catalog, provider is no longer a `Materialized`, `cast_to_materialized` +/// error) preserve the LP-time value — never downgrade what we already +/// observed. +pub(super) async fn refresh_candidate_readiness( + context: RewriteContext, + physical_inputs: &[Arc], + session_state: &SessionState, +) -> RewriteContext { + let catalog_list = session_state.catalog_list(); + let default_catalog = &session_state.config().options().catalog.default_catalog; + let default_schema = &session_state.config().options().catalog.default_schema; + + let candidates = context.candidates(); + // Alignment-tolerant zipping: if a caller supplied metadata whose + // length differs from `physical_inputs`, we can't safely pair them, + // so we skip the annotated lookup for the misaligned indices. The + // downstream `drop_not_ready_after_refresh` filter handles the same + // case defensively. + let aligned = candidates.len() == physical_inputs.len(); + + let refreshed: Vec = + futures::future::join_all(candidates.iter().enumerate().map(|(idx, c)| async move { + match c { + CandidateMetadata::Base => CandidateMetadata::Base, + CandidateMetadata::Materialized { + table_ref, + readiness, + } => { + // Prefer atomically-captured readiness from the physical + // input over sampling — closes the race between the + // snapshot scan read from and the current provider state. + let annotated = if aligned { + readiness_from_plan(&physical_inputs[idx]) + } else { + None + }; + + let new_readiness = match annotated { + Some(r) => r, + None => resolve_current_readiness( + catalog_list.as_ref(), + default_catalog, + default_schema, + table_ref, + ) + .await + .unwrap_or(*readiness), + }; + + CandidateMetadata::Materialized { + table_ref: table_ref.clone(), + readiness: new_readiness, + } + } + } + })) + .await; + + context.with_candidates(refreshed) +} + +/// Resolve `table_ref` through the catalog list and re-consult the +/// provider's `rewrite_readiness()`. Returns `None` on any lookup failure +/// so the caller can fall back to the LP-time value. +async fn resolve_current_readiness( + catalog_list: &dyn CatalogProviderList, + default_catalog: &str, + default_schema: &str, + table_ref: &TableReference, +) -> Option { + // Resolve the (possibly bare or partial) reference against the + // session's default catalog/schema without ever going through + // `Display`/`parse_str`, which is not lossless for quoted or dotted + // identifiers. + let resolved = table_ref.clone().resolve(default_catalog, default_schema); + let catalog = catalog_list.catalog(resolved.catalog.as_ref())?; + let schema = catalog.schema(resolved.schema.as_ref())?; + let table = match schema.table(resolved.table.as_ref()).await { + Ok(Some(t)) => t, + Ok(None) => { + log::trace!("refresh_candidate_readiness: no table {table_ref} in catalog"); + return None; + } + Err(e) => { + log::warn!("refresh_candidate_readiness: catalog lookup failed for {table_ref}: {e}"); + return None; + } + }; + match cast_to_materialized(table.as_ref()) { + Ok(Some(mv)) => Some(mv.rewrite_readiness()), + Ok(None) => None, + Err(e) => { + log::warn!( + "refresh_candidate_readiness: cast_to_materialized failed for {table_ref}: {e}" + ); + None + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn readiness_variants_are_distinct_and_comparable() { + assert_ne!(RewriteReadiness::Ready, RewriteReadiness::NotReady); + assert_ne!(RewriteReadiness::Ready, RewriteReadiness::Unknown); + assert_ne!(RewriteReadiness::NotReady, RewriteReadiness::Unknown); + assert_eq!(RewriteReadiness::Ready, RewriteReadiness::Ready); + } + + #[test] + fn readiness_is_copy_and_hashable() { + // Ensure the variants can be stored / matched cheaply from the + // rewrite path without cloning. + fn requires_copy() {} + fn requires_hash() {} + requires_copy::(); + requires_hash::(); + } + + #[test] + fn readiness_annotated_exec_delegates_properties_and_children() { + use arrow_schema::Schema; + use datafusion::physical_plan::empty::EmptyExec; + let schema = Arc::new(Schema::empty()); + let inner: Arc = Arc::new(EmptyExec::new(Arc::clone(&schema))); + let annotated = ReadinessAnnotatedExec::new(Arc::clone(&inner), RewriteReadiness::Ready); + + assert_eq!(annotated.readiness(), RewriteReadiness::Ready); + // Same schema/partitioning as inner — the wrapper is transparent. + assert_eq!( + format!("{:?}", annotated.properties()), + format!("{:?}", inner.properties()) + ); + // Inner exposed as the sole child. + let children = annotated.children(); + assert_eq!(children.len(), 1); + assert!(Arc::ptr_eq(children[0], &inner)); + } + + #[test] + fn strip_returns_inner_for_bare_annotated_exec() { + use arrow_schema::Schema; + use datafusion::physical_plan::empty::EmptyExec; + let schema = Arc::new(Schema::empty()); + let inner: Arc = Arc::new(EmptyExec::new(schema)); + let annotated: Arc = Arc::new(ReadinessAnnotatedExec::new( + Arc::clone(&inner), + RewriteReadiness::Ready, + )); + + let stripped = strip_readiness_annotation(annotated).expect("strip"); + assert!( + stripped + .as_any() + .downcast_ref::() + .is_none(), + "top-level wrapper must be gone after strip" + ); + assert!(Arc::ptr_eq(&stripped, &inner)); + } + + #[test] + fn strip_replaces_annotated_exec_inside_nested_plan() { + // Providers wrap their scan output, but DataFusion may put a + // Filter/Projection/Repartition on top before `plan_extension` sees + // the plan. Strip must find the wrapper wherever it sits. + use arrow_schema::Schema; + use datafusion::physical_plan::empty::EmptyExec; + use datafusion::physical_plan::repartition::RepartitionExec; + use datafusion_physical_expr::Partitioning; + let schema = Arc::new(Schema::empty()); + let leaf: Arc = Arc::new(EmptyExec::new(schema)); + let annotated: Arc = Arc::new(ReadinessAnnotatedExec::new( + Arc::clone(&leaf), + RewriteReadiness::Ready, + )); + let wrapped: Arc = Arc::new( + RepartitionExec::try_new(annotated, Partitioning::RoundRobinBatch(1)) + .expect("repartition"), + ); + + let stripped = strip_readiness_annotation(wrapped).expect("strip"); + assert!( + readiness_from_plan(&stripped).is_none(), + "no ReadinessAnnotatedExec should remain anywhere in the tree" + ); + } + + #[test] + fn strip_is_noop_when_no_annotation_present() { + use arrow_schema::Schema; + use datafusion::physical_plan::empty::EmptyExec; + let schema = Arc::new(Schema::empty()); + let plan: Arc = Arc::new(EmptyExec::new(schema)); + + let stripped = strip_readiness_annotation(Arc::clone(&plan)).expect("strip"); + assert!(Arc::ptr_eq(&stripped, &plan)); + } +} From 6e7359d385a943e4f4ef8526c12a6f658c37bd34 Mon Sep 17 00:00:00 2001 From: Qi Zhu <821684824@qq.com> Date: Tue, 28 Jul 2026 14:33:37 +0800 Subject: [PATCH 2/4] fix(clippy): silence 1.97 unneeded_wildcard_pattern + question_mark - dependencies.rs: LogicalPlan::Window pattern already covers window_expr via .. - file_metadata.rs: use ? on the else branch instead of explicit return None --- src/materialized/dependencies.rs | 6 +----- src/materialized/file_metadata.rs | 5 ++--- 2 files changed, 3 insertions(+), 8 deletions(-) diff --git a/src/materialized/dependencies.rs b/src/materialized/dependencies.rs index 7295b36..f707184 100644 --- a/src/materialized/dependencies.rs +++ b/src/materialized/dependencies.rs @@ -492,11 +492,7 @@ fn pushdown_projection_inexact(plan: LogicalPlan, indices: &HashSet) -> R ) .map(LogicalPlan::Filter) } - LogicalPlan::Window(Window { - input, - window_expr: _, - .. - }) => { + LogicalPlan::Window(Window { input, .. }) => { // Window nodes take their input and append window expressions to the end. // If our projection doesn't include window expressions, we can just turn // the window into a regular projection. diff --git a/src/materialized/file_metadata.rs b/src/materialized/file_metadata.rs index bca69ea..f7c1b95 100644 --- a/src/materialized/file_metadata.rs +++ b/src/materialized/file_metadata.rs @@ -282,11 +282,10 @@ impl FileMetadataExec { { let right_literal = binary_expr.right().as_any().downcast_ref::()?; (left_column, right_literal) - } else if let Some(right_column) = binary_expr.right().as_any().downcast_ref::() { + } else { + let right_column = binary_expr.right().as_any().downcast_ref::()?; let left_literal = binary_expr.left().as_any().downcast_ref::()?; (right_column, left_literal) - } else { - return None; }; if column.index() != column_idx { From 80044543dcbb4cb01f9b7912c411efabf8672acf Mon Sep 17 00:00:00 2001 From: Qi Zhu <821684824@qq.com> Date: Tue, 28 Jul 2026 14:36:40 +0800 Subject: [PATCH 3/4] ci: use `check licenses` (plural) for cargo-deny The current cargo-deny CLI rejects the singular `license` subcommand: error: invalid value 'license' for '[WHICH]...' Use the plural form so the license check runs. --- .github/workflows/ci.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 43ec6b0..c906fab 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -134,7 +134,7 @@ jobs: - uses: actions/checkout@v3 - uses: EmbarkStudios/cargo-deny-action@v2 with: - command: check license + command: check licenses coverage: if: github.event.pull_request.draft == false From 6316aebaf9439a4452ca75c1e4f3f66782de9b36 Mon Sep 17 00:00:00 2001 From: Qi Zhu <821684824@qq.com> Date: Tue, 28 Jul 2026 14:43:28 +0800 Subject: [PATCH 4/4] chore: cargo fmt after clippy fix --- src/materialized/file_metadata.rs | 19 +++++++++---------- 1 file changed, 9 insertions(+), 10 deletions(-) diff --git a/src/materialized/file_metadata.rs b/src/materialized/file_metadata.rs index f7c1b95..46cd237 100644 --- a/src/materialized/file_metadata.rs +++ b/src/materialized/file_metadata.rs @@ -277,16 +277,15 @@ impl FileMetadataExec { return None; } - let (column, literal) = if let Some(left_column) = - binary_expr.left().as_any().downcast_ref::() - { - let right_literal = binary_expr.right().as_any().downcast_ref::()?; - (left_column, right_literal) - } else { - let right_column = binary_expr.right().as_any().downcast_ref::()?; - let left_literal = binary_expr.left().as_any().downcast_ref::()?; - (right_column, left_literal) - }; + let (column, literal) = + if let Some(left_column) = binary_expr.left().as_any().downcast_ref::() { + let right_literal = binary_expr.right().as_any().downcast_ref::()?; + (left_column, right_literal) + } else { + let right_column = binary_expr.right().as_any().downcast_ref::()?; + let left_literal = binary_expr.left().as_any().downcast_ref::()?; + (right_column, left_literal) + }; if column.index() != column_idx { return None;