diff --git a/src/frontend/src/optimizer/plan_node/batch.rs b/src/frontend/src/optimizer/plan_node/batch.rs index bb3cbb96d6c5d..5eeafab64f1cd 100644 --- a/src/frontend/src/optimizer/plan_node/batch.rs +++ b/src/frontend/src/optimizer/plan_node/batch.rs @@ -28,6 +28,7 @@ pub trait BatchPlanRef: PhysicalPlanRef { fn order(&self) -> &Order; } +/// Prelude for batch plan nodes. pub mod prelude { pub use super::super::generic::{GenericPlanRef, PhysicalPlanRef}; pub use super::super::Batch; diff --git a/src/frontend/src/optimizer/plan_node/mod.rs b/src/frontend/src/optimizer/plan_node/mod.rs index fd72155cec3b3..beaaa18726dfe 100644 --- a/src/frontend/src/optimizer/plan_node/mod.rs +++ b/src/frontend/src/optimizer/plan_node/mod.rs @@ -51,18 +51,18 @@ use self::stream::StreamPlanRef; use self::utils::Distill; use super::property::{Distribution, FunctionalDependencySet, Order}; -pub trait PlanNodeMeta { - fn node_type(&self) -> PlanNodeType; - fn plan_base(&self) -> PlanBaseRef<'_>; - fn convention(&self) -> Convention; -} - +/// A marker trait for different conventions, used for enforcing type safety. +/// +/// Implementors are [`Logical`], [`Batch`], and [`Stream`]. pub trait ConventionMarker: 'static + Sized { + /// The extra fields in the [`PlanBase`] of this convention. type Extra: 'static + Eq + Hash + Clone + Debug; + /// Get the [`Convention`] enum value. fn value() -> Convention; } +/// The marker for logical convention. pub struct Logical; impl ConventionMarker for Logical { type Extra = plan_base::NoExtra; @@ -72,6 +72,7 @@ impl ConventionMarker for Logical { } } +/// The marker for batch convention. pub struct Batch; impl ConventionMarker for Batch { type Extra = plan_base::BatchExtra; @@ -81,6 +82,7 @@ impl ConventionMarker for Batch { } } +/// The marker for stream convention. pub struct Stream; impl ConventionMarker for Stream { type Extra = plan_base::StreamExtra; @@ -90,31 +92,54 @@ impl ConventionMarker for Stream { } } -pub trait StaticPlanNodeMeta { +/// The trait for accessing the meta data and [`PlanBase`] for plan nodes. +pub trait PlanNodeMeta { type Convention: ConventionMarker; + const NODE_TYPE: PlanNodeType; + /// Get the reference to the [`PlanBase`] with corresponding convention. fn plan_base(&self) -> &PlanBase; + /// Get the reference to the [`PlanBase`] with erased convention. + /// + /// This is mainly used for implementing [`AnyPlanNodeMeta`]. Callers should prefer + /// [`PlanNodeMeta::plan_base`] instead as it is more type-safe. fn plan_base_ref(&self) -> PlanBaseRef<'_>; } -impl

PlanNodeMeta for P -where - P: StaticPlanNodeMeta, -{ - fn node_type(&self) -> PlanNodeType { - P::NODE_TYPE - } +// Intentionally made private. +mod plan_node_meta { + use super::*; - fn plan_base(&self) -> PlanBaseRef<'_> { - StaticPlanNodeMeta::plan_base_ref(self) + /// The object-safe version of [`PlanNodeMeta`], used as a super trait of [`PlanNode`]. + /// + /// Check [`PlanNodeMeta`] for more details. + pub trait AnyPlanNodeMeta { + fn node_type(&self) -> PlanNodeType; + fn plan_base(&self) -> PlanBaseRef<'_>; + fn convention(&self) -> Convention; } - fn convention(&self) -> Convention { - P::Convention::value() + /// Implement [`AnyPlanNodeMeta`] for all [`PlanNodeMeta`]. + impl

AnyPlanNodeMeta for P + where + P: PlanNodeMeta, + { + fn node_type(&self) -> PlanNodeType { + P::NODE_TYPE + } + + fn plan_base(&self) -> PlanBaseRef<'_> { + PlanNodeMeta::plan_base_ref(self) + } + + fn convention(&self) -> Convention { + P::Convention::value() + } } } +use plan_node_meta::AnyPlanNodeMeta; /// The common trait over all plan nodes. Used by optimizer framework which will treat all node as /// `dyn PlanNode` @@ -136,7 +161,7 @@ pub trait PlanNode: + ToPb + ToLocalBatch + PredicatePushdown - + PlanNodeMeta + + AnyPlanNodeMeta { } @@ -477,7 +502,8 @@ impl PlanTreeNode for PlanRef { } } -impl PlanNodeMeta for PlanRef { +/// Implement again for the `dyn` newtype wrapper. +impl AnyPlanNodeMeta for PlanRef { fn node_type(&self) -> PlanNodeType { self.0.node_type() } @@ -491,6 +517,8 @@ impl PlanNodeMeta for PlanRef { } } +/// Allow access to all fields defined in [`GenericPlanRef`] for the type-erased plan node. +// TODO: may also implement on `dyn PlanNode` directly. impl GenericPlanRef for PlanRef { fn id(&self) -> PlanNodeId { self.plan_base().id() @@ -513,12 +541,16 @@ impl GenericPlanRef for PlanRef { } } +/// Allow access to all fields defined in [`PhysicalPlanRef`] for the type-erased plan node. +// TODO: may also implement on `dyn PlanNode` directly. impl PhysicalPlanRef for PlanRef { fn distribution(&self) -> &Distribution { self.plan_base().distribution() } } +/// Allow access to all fields defined in [`StreamPlanRef`] for the type-erased plan node. +// TODO: may also implement on `dyn PlanNode` directly. impl StreamPlanRef for PlanRef { fn append_only(&self) -> bool { self.plan_base().append_only() @@ -533,6 +565,8 @@ impl StreamPlanRef for PlanRef { } } +/// Allow access to all fields defined in [`BatchPlanRef`] for the type-erased plan node. +// TODO: may also implement on `dyn PlanNode` directly. impl BatchPlanRef for PlanRef { fn order(&self) -> &Order { self.plan_base().order() @@ -592,7 +626,8 @@ pub(crate) fn pretty_config() -> PrettyConfig { } } -// TODO: remove this direct implementation always require `GenericPlanRef`. +/// Directly implement methods for [`PlanNode`] to access the fields defined in [`GenericPlanRef`]. +// TODO: always require `GenericPlanRef` to make it more consistent. impl dyn PlanNode { pub fn id(&self) -> PlanNodeId { self.plan_base().id() @@ -1101,7 +1136,7 @@ macro_rules! impl_plan_node_meta { $( [<$convention $name>] ),* } - $(impl StaticPlanNodeMeta for [<$convention $name>] { + $(impl PlanNodeMeta for [<$convention $name>] { type Convention = $convention; const NODE_TYPE: PlanNodeType = PlanNodeType::[<$convention $name>]; diff --git a/src/frontend/src/optimizer/plan_node/plan_base.rs b/src/frontend/src/optimizer/plan_node/plan_base.rs index 241226417405f..0d2e649379112 100644 --- a/src/frontend/src/optimizer/plan_node/plan_base.rs +++ b/src/frontend/src/optimizer/plan_node/plan_base.rs @@ -21,6 +21,7 @@ use super::*; use crate::optimizer::optimizer_context::OptimizerContextRef; use crate::optimizer::property::{Distribution, FunctionalDependencySet, Order}; +/// No extra fields for logical plan nodes. #[derive(Clone, Debug, PartialEq, Eq, Hash)] pub struct NoExtra; @@ -37,6 +38,8 @@ mod physical_common { pub dist: Distribution, } + /// A helper trait to reuse code for accessing the common physical fields of batch and stream + /// plan bases. pub trait GetPhysicalCommon { fn physical(&self) -> &PhysicalCommonExtra; fn physical_mut(&mut self) -> &mut PhysicalCommonExtra; @@ -92,14 +95,17 @@ impl GetPhysicalCommon for BatchExtra { } } -/// the common fields of all nodes, please make a field named `base` in -/// every planNode and correctly value it when construct the planNode. +/// The common fields of all plan nodes with different conventions. +/// +/// Please make a field named `base` in every planNode and correctly value +/// it when construct the planNode. /// /// All fields are intentionally made private and immutable, as they should /// normally be the same as the given [`GenericPlanNode`] when constructing. /// /// - To access them, use traits including [`GenericPlanRef`], -/// [`PhysicalPlanRef`], [`StreamPlanRef`] and [`BatchPlanRef`]. +/// [`PhysicalPlanRef`], [`StreamPlanRef`] and [`BatchPlanRef`] with +/// compile-time checks. /// - To mutate them, use methods like `new_*` or `clone_with_*`. #[derive(Educe)] #[educe(PartialEq, Eq, Hash, Clone, Debug)] @@ -118,6 +124,7 @@ pub struct PlanBase { stream_key: Option>, functional_dependency: FunctionalDependencySet, + /// Extra fields for different conventions. extra: C::Extra, } @@ -309,6 +316,11 @@ impl PlanBase { } } +/// Reference to [`PlanBase`] with erased conventions. +/// +/// Used for accessing fields on a type-erased plan node. All traits of [`GenericPlanRef`], +/// [`PhysicalPlanRef`], [`StreamPlanRef`] and [`BatchPlanRef`] are implemented for this type, +/// so runtime checks are required when calling methods on it. #[derive(PartialEq, Eq, Hash, Clone, Copy, Debug, enum_as_inner::EnumAsInner)] pub enum PlanBaseRef<'a> { Logical(&'a PlanBase), @@ -317,6 +329,7 @@ pub enum PlanBaseRef<'a> { } impl PlanBaseRef<'_> { + /// Get the convention of this plan base. pub fn convention(self) -> Convention { match self { PlanBaseRef::Logical(_) => Convention::Logical, @@ -345,9 +358,9 @@ macro_rules! dispatch_plan_base { /// For example, callers writing `GenericPlanRef::schema(&foo.plan_base())` will lead to a /// borrow checker error, as it borrows [`PlanBaseRef`] again, which is already a reference. /// -/// As a workaround, we directly let the getters below take the ownership of [`PlanBaseRef`], when -/// callers write `foo.plan_base().schema()`, the compiler will prefer these ones over the ones -/// defined in traits like [`GenericPlanRef`]. +/// As a workaround, we directly let the getters below take the ownership of [`PlanBaseRef`], +/// which is `Copy`. When callers write `foo.plan_base().schema()`, the compiler will prefer +/// these ones over the ones defined in traits like [`GenericPlanRef`]. impl<'a> PlanBaseRef<'a> { pub(super) fn schema(self) -> &'a Schema { dispatch_plan_base!(self, [Logical, Stream, Batch], GenericPlanRef::schema) diff --git a/src/frontend/src/optimizer/plan_node/stream.rs b/src/frontend/src/optimizer/plan_node/stream.rs index f33a15d2570b3..394a64b656ad3 100644 --- a/src/frontend/src/optimizer/plan_node/stream.rs +++ b/src/frontend/src/optimizer/plan_node/stream.rs @@ -31,6 +31,7 @@ pub trait StreamPlanRef: PhysicalPlanRef { fn watermark_columns(&self) -> &FixedBitSet; } +/// Prelude for stream plan nodes. pub mod prelude { pub use super::super::generic::{GenericPlanRef, PhysicalPlanRef}; pub use super::super::Stream;