milenkovicm commented on PR #18739:
URL: https://github.com/apache/datafusion/pull/18739#issuecomment-3622438362
to make it easier to understand, please have a look, it illustrates my point
(this is datafusion 50),
```rust
use datafusion::config::{ConfigEntry, ConfigExtension, ConfigOptions,
ExtensionOptions};
use datafusion::execution::SessionStateBuilder;
use datafusion::physical_optimizer::PhysicalOptimizerRule;
use datafusion::physical_plan::ExecutionPlan;
use datafusion::prelude::{SessionConfig, SessionContext};
use std::any::Any;
use std::fmt::Debug;
use std::sync::Arc;
#[tokio::main]
async fn main() -> datafusion::error::Result<()> {
// you may want to keep your rule_context
// accessible or not, im not sure what you want
let rule_context = Arc::new(MyRuleContext::default());
let my_rule = Arc::new(MyRule::new(rule_context.clone()));
// if for any reason you need to reconfigure your rule
// use can use SQL SET or directly on session config
let my_rule_config = MyRuleConfig::default();
// assign your rule context to a session config
// so it gets dropped when your session gets dropped
// use of session extension
let session_config = SessionConfig::new()
.with_extension(rule_context.clone())
.with_option_extension(my_rule_config);
let session_state = SessionStateBuilder::new_with_default_features()
.with_config(session_config)
.with_physical_optimizer_rule(my_rule)
.build();
let ctx = SessionContext::new_with_state(session_state);
// this should trigger rule
ctx.sql("select 1").await?.show().await?;
Ok(())
}
#[derive(Debug, Default)]
struct MyRuleContext {
// probably your stage
// need some kind of week hash
// map or similar
// if you want to track physical expressions
}
#[derive(Debug)]
struct MyRule {
context: Arc<MyRuleContext>,
}
impl MyRule {
fn new(context: Arc<MyRuleContext>) -> Self {
Self { context }
}
}
impl PhysicalOptimizerRule for MyRule {
fn optimize(
&self,
plan: Arc<dyn ExecutionPlan>,
config: &ConfigOptions,
) -> datafusion::common::Result<Arc<dyn ExecutionPlan>> {
// configure / reconfigure your rule
// any way you need
let my_config: Option<&MyRuleConfig> = config.extensions.get();
println!("MyRule config: {:?}", my_config);
println!("MyRule context: {:?}", self.context);
Ok(plan)
}
fn name(&self) -> &str {
"MyRule"
}
fn schema_check(&self) -> bool {
false
}
}
#[derive(Debug, Default, Clone)]
struct MyRuleConfig {
// whatever you need to configure your
// rule goes here
}
impl ExtensionOptions for MyRuleConfig {
fn as_any(&self) -> &dyn Any {
self
}
fn as_any_mut(&mut self) -> &mut dyn Any {
self
}
fn cloned(&self) -> Box<dyn ExtensionOptions> {
Box::new(self.clone())
}
fn set(&mut self, _key: &str, _value: &str) ->
datafusion::common::Result<()> {
todo!("to be implemented later")
}
fn entries(&self) -> Vec<ConfigEntry> {
todo!("to be implemented later")
}
}
impl ConfigExtension for MyRuleConfig {
const PREFIX: &'static str = "my_rule";
}
```
it should print:
```text
MyRule config: Some(MyRuleConfig)
MyRule context: MyRuleContext
+----------+
| Int64(1) |
+----------+
| 1 |
+----------+
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
