use std::sync::Arc;
use datafusion::{
common::tree_node::{Transformed, TreeNode},
config::ConfigOptions,
error::Result as DFResult,
physical_optimizer::PhysicalOptimizerRule,
physical_plan::ExecutionPlan,
};
use super::TakeExec;
pub struct CoalesceTake;
impl PhysicalOptimizerRule for CoalesceTake {
fn optimize(
&self,
plan: Arc<dyn ExecutionPlan>,
_config: &ConfigOptions,
) -> DFResult<Arc<dyn ExecutionPlan>> {
plan.transform_down(&|plan| {
if let Some(take) = plan.as_any().downcast_ref::<TakeExec>() {
let child = &take.children()[0];
if let Some(exec_child) = child.as_any().downcast_ref::<TakeExec>() {
let upstream_plan = exec_child.children();
return Ok(Transformed::Yes(plan.with_new_children(upstream_plan)?));
}
}
Ok(Transformed::No(plan))
})
}
fn name(&self) -> &str {
"coalesce_take"
}
fn schema_check(&self) -> bool {
true
}
}