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 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/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..46cd237 100644 --- a/src/materialized/file_metadata.rs +++ b/src/materialized/file_metadata.rs @@ -277,17 +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 if let Some(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; - }; + 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; 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)); + } +}