uni_plugin/traits/
operator.rs

1//! Physical operator + optimizer-rule plugins.
2
3use std::sync::Arc;
4
5use datafusion::arrow::datatypes::SchemaRef;
6use datafusion::execution::context::SessionContext;
7use datafusion::optimizer::OptimizerRule;
8use datafusion::physical_optimizer::PhysicalOptimizerRule;
9use datafusion::physical_plan::ExecutionPlan;
10
11use crate::errors::FnError;
12
13/// Per-planner-invocation context for [`OperatorProvider::plan`].
14#[non_exhaustive]
15pub struct PlannerArgs<'a> {
16    /// Reference to the executing DataFusion `SessionContext`.
17    pub session_ctx: &'a SessionContext,
18    /// Input physical plans the operator should consume.
19    pub input_plans: &'a [Arc<dyn ExecutionPlan>],
20    /// Free-form JSON configuration.
21    pub config_json: &'a str,
22    /// Optional schema hint for the operator's output.
23    pub schema_hint: Option<SchemaRef>,
24}
25
26impl std::fmt::Debug for PlannerArgs<'_> {
27    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
28        f.debug_struct("PlannerArgs")
29            .field("session_ctx", &"<SessionContext>")
30            .field("input_plans.len", &self.input_plans.len())
31            .field("config_json", &self.config_json)
32            .field("schema_hint", &self.schema_hint)
33            .finish()
34    }
35}
36
37/// A custom physical operator factory.
38pub trait OperatorProvider: Send + Sync {
39    /// The logical name of this operator (`"hash_join_geo"`, …).
40    fn logical_name(&self) -> &str;
41
42    /// Construct an `ExecutionPlan` for an instance of this operator.
43    ///
44    /// # Errors
45    ///
46    /// Returns [`FnError`] on planning failure (incompatible inputs, bad
47    /// configuration).
48    fn plan(&self, args: PlannerArgs<'_>) -> Result<Arc<dyn ExecutionPlan>, FnError>;
49}
50
51/// Phase at which an `OptimizerRule` runs.
52#[derive(Clone, Copy, Debug, PartialEq, Eq)]
53#[non_exhaustive]
54pub enum OptimizerPhase {
55    /// Logical optimizer.
56    Logical,
57    /// Physical optimizer.
58    Physical,
59    /// Both — the rule is applied at logical and physical phases.
60    Both,
61}
62
63/// A registered optimizer-rule provider.
64///
65/// A provider that runs at the logical phase returns a logical
66/// [`OptimizerRule`] from [`rule`](Self::rule); a provider that runs at
67/// the physical phase returns a [`PhysicalOptimizerRule`] from
68/// [`physical_rule`](Self::physical_rule). A `Both` provider must
69/// supply both. The host iterates the registered providers, inspects
70/// `phase`, and installs each rule into the matching DataFusion
71/// optimizer chain.
72///
73/// The default `physical_rule` returns `None`, so existing
74/// logical-only providers compile unchanged across the 1.6 → 1.7
75/// minor bump.
76pub trait OptimizerRuleProvider: Send + Sync {
77    /// The DataFusion logical `OptimizerRule` to apply.
78    ///
79    /// Logical-phase and `Both`-phase providers must return a real
80    /// rule. Physical-only providers may return any rule (the host
81    /// ignores it when `phase()` is [`OptimizerPhase::Physical`]);
82    /// returning a sentinel/no-op is conventional. The default impl
83    /// returns a no-op rule that never rewrites.
84    fn rule(&self) -> Arc<dyn OptimizerRule + Send + Sync> {
85        Arc::new(NoopOptimizerRule)
86    }
87
88    /// The DataFusion physical [`PhysicalOptimizerRule`] to apply.
89    ///
90    /// Physical-phase and `Both`-phase providers should return
91    /// `Some(...)`. The default `None` keeps existing logical-only
92    /// providers source-compatible.
93    fn physical_rule(&self) -> Option<Arc<dyn PhysicalOptimizerRule + Send + Sync>> {
94        None
95    }
96
97    /// Phase the rule runs at.
98    fn phase(&self) -> OptimizerPhase;
99
100    /// Ordering hint — lower precedence rules run first.
101    fn precedence(&self) -> i32 {
102        0
103    }
104}
105
106/// No-op logical `OptimizerRule` used as the default for
107/// [`OptimizerRuleProvider::rule`].
108///
109/// Returned by the trait's default `rule()` implementation so that
110/// physical-only providers do not have to construct a sentinel
111/// themselves. The rule is `Bottom-Up` and never transforms the plan.
112#[derive(Debug, Default)]
113pub struct NoopOptimizerRule;
114
115impl OptimizerRule for NoopOptimizerRule {
116    fn name(&self) -> &str {
117        "uni_noop_optimizer_rule"
118    }
119
120    fn apply_order(&self) -> Option<datafusion::optimizer::ApplyOrder> {
121        Some(datafusion::optimizer::ApplyOrder::BottomUp)
122    }
123
124    fn rewrite(
125        &self,
126        plan: datafusion::logical_expr::LogicalPlan,
127        _config: &dyn datafusion::optimizer::OptimizerConfig,
128    ) -> Result<
129        datafusion::common::tree_node::Transformed<datafusion::logical_expr::LogicalPlan>,
130        datafusion::error::DataFusionError,
131    > {
132        Ok(datafusion::common::tree_node::Transformed::no(plan))
133    }
134}
uni_plugin/traits/operator.rs

uni_plugin/traits/
operator.rs
