perf(promql): fuse direct rate projection into range manipulation

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

src/promql/src/extension_plan/range_manipulate.rs

@@ -18,7 +18,9 @@ use std::sync::Arc;
 use std::task::{Context, Poll};
 
 use common_telemetry::{debug, warn};
-use datafusion::arrow::array::{Array, ArrayRef, Int64Array, TimestampMillisecondArray};
+use datafusion::arrow::array::{
+    Array, ArrayRef, Float64Array, Int64Array, TimestampMillisecondArray,
+};
 use datafusion::arrow::compute;
 use datafusion::arrow::datatypes::{DataType, Field, SchemaRef, TimeUnit};
 use datafusion::arrow::error::ArrowError;
@@ -30,13 +32,16 @@ use datafusion::error::{DataFusionError, Result as DataFusionResult};
 use datafusion::execution::context::TaskContext;
 use datafusion::logical_expr::{EmptyRelation, Expr, LogicalPlan, UserDefinedLogicalNodeCore};
 use datafusion::physical_expr::EquivalenceProperties;
+use datafusion::physical_expr::ScalarFunctionExpr;
+use datafusion::physical_expr::expressions::{Column, Literal};
 use datafusion::physical_plan::metrics::{
     BaselineMetrics, Count, ExecutionPlanMetricsSet, MetricBuilder, MetricValue, MetricsSet,
 };
+use datafusion::physical_plan::projection::{ProjectionExec, ProjectionExpr};
 use datafusion::physical_plan::{
-    ChildStats, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan,
+    ChildStats, ChildrenPropertiesMode, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan,
     InputDistributionRequirements, PhysicalExpr, PlanProperties, RecordBatchStream,
-    SendableRecordBatchStream, Statistics, StatisticsArgs,
+    ReplaceChildrenOptions, SendableRecordBatchStream, Statistics, StatisticsArgs,
 };
 use datafusion_expr::ident;
 use datatypes::timestamp::timestamp_array_to_primitive;
@@ -207,6 +212,7 @@ impl RangeManipulate {
             output_schema,
             metric: ExecutionPlanMetricsSet::new(),
             properties,
+            fused_rate: None,
         })
     }
 
@@ -434,7 +440,7 @@ impl UserDefinedLogicalNodeCore for RangeManipulate {
     }
 }
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct RangeManipulateExec {
     offset: Millisecond,
     start: Millisecond,
@@ -449,6 +455,161 @@ pub struct RangeManipulateExec {
     output_schema: SchemaRef,
     metric: ExecutionPlanMetricsSet,
     properties: Arc<PlanProperties>,
+    fused_rate: Option<FusedRate>,
+}
+
+#[derive(Debug, Clone)]
+struct FusedRate {
+    rate_source: usize,
+    rate_position: usize,
+    timestamp_position: usize,
+    projections: Vec<ProjectionExpr>,
+    folded_schema: SchemaRef,
+}
+
+impl RangeManipulateExec {
+    /// Fuse a projection containing one supported PromQL rate expression into this range operator.
+    ///
+    /// Returns `None` when the raw columns or projection shape do not match the rate kernel ABI.
+    pub fn try_fuse_rate_projection(
+        &self,
+        projection: &ProjectionExec,
+    ) -> DataFusionResult<Option<Arc<dyn ExecutionPlan>>> {
+        if self.offset != 0 {
+            return Ok(None);
+        }
+        let raw_input_schema = self.input.schema();
+        let Some((time_index, _)) = raw_input_schema.column_with_name(&self.time_index_column)
+        else {
+            return Ok(None);
+        };
+        if raw_input_schema
+            .fields()
+            .get(time_index)
+            .map(|f| f.data_type())
+            != Some(&DataType::Timestamp(TimeUnit::Millisecond, None))
+        {
+            return Ok(None);
+        }
+        let Some((time_range_index, _)) =
+            self.output_schema.column_with_name(&self.time_range_column)
+        else {
+            return Ok(None);
+        };
+        let Some((rate_field_name, rate_field_index)) =
+            self.field_columns.iter().find_map(|name| {
+                self.output_schema
+                    .column_with_name(name)
+                    .map(|(index, _)| (name, index))
+            })
+        else {
+            return Ok(None);
+        };
+        let Some((raw_rate_index, raw_rate_field)) =
+            raw_input_schema.column_with_name(rate_field_name)
+        else {
+            return Ok(None);
+        };
+        if raw_rate_field.data_type() != &DataType::Float64 {
+            return Ok(None);
+        }
+        if raw_input_schema.field(time_index).is_nullable() {
+            return Ok(None);
+        }
+        let projections = projection.expr().to_vec();
+        let mut rate_position = None;
+        let mut timestamp_position = None;
+        for (position, item) in projections.iter().enumerate() {
+            if let Some(function) = item.expr.downcast_ref::<ScalarFunctionExpr>() {
+                if rate_position.is_some()
+                    || crate::functions::Rate::scalar_udf() != *function.fun()
+                    || function.args().len() != 4
+                {
+                    return Ok(None);
+                }
+                let Some(ts_arg) = function.args()[0].downcast_ref::<Column>() else {
+                    return Ok(None);
+                };
+                let Some(value_arg) = function.args()[1].downcast_ref::<Column>() else {
+                    return Ok(None);
+                };
+                let Some(eval_arg) = function.args()[2].downcast_ref::<Column>() else {
+                    return Ok(None);
+                };
+                let Some(range_arg) = function.args()[3].downcast_ref::<Literal>() else {
+                    return Ok(None);
+                };
+                if ts_arg.index() != time_range_index
+                    || eval_arg.index() != time_index
+                    || value_arg.index() != rate_field_index
+                {
+                    return Ok(None);
+                }
+                let datafusion::common::ScalarValue::Int64(Some(duration)) = range_arg.value()
+                else {
+                    return Ok(None);
+                };
+                if *duration != self.range {
+                    return Ok(None);
+                }
+                rate_position = Some(position);
+            } else {
+                let Some(column) = item.expr.downcast_ref::<Column>() else {
+                    return Ok(None);
+                };
+                if column.index() == time_range_index
+                    || self.field_columns.iter().any(|name| {
+                        self.output_schema
+                            .column_with_name(name)
+                            .is_some_and(|(i, _)| i == column.index())
+                    })
+                {
+                    return Ok(None);
+                }
+                if column.index() == time_index && timestamp_position.replace(position).is_some() {
+                    return Ok(None);
+                }
+            }
+        }
+        let Some(rate_position) = rate_position else {
+            return Ok(None);
+        };
+        let Some(timestamp_position) = timestamp_position else {
+            return Ok(None);
+        };
+        if self.output_schema.fields().get(time_range_index).is_none()
+            || raw_input_schema.fields().get(raw_rate_index).is_none()
+            || projections.iter().any(|item| {
+                item.expr
+                    .downcast_ref::<Column>()
+                    .is_some_and(|column| self.output_schema.fields().get(column.index()).is_none())
+            })
+        {
+            return Ok(None);
+        }
+        let base = Arc::new(Self {
+            fused_rate: None,
+            ..self.clone()
+        }) as Arc<dyn ExecutionPlan>;
+        let rebuilt = ProjectionExec::try_new_with_schema_metadata(
+            projections.clone(),
+            base,
+            projection.schema().as_ref(),
+        )?;
+        let fused = Self {
+            fused_rate: Some(FusedRate {
+                rate_source: raw_rate_index,
+                rate_position,
+                timestamp_position,
+                projections,
+                folded_schema: self.output_schema.clone(),
+            }),
+            output_schema: projection.schema(),
+            properties: rebuilt.properties().clone(),
+            ..self.clone()
+        };
+        Ok(Some(Arc::new(fused)))
+    }
 }
 
 impl ExecutionPlan for RangeManipulateExec {
@@ -495,13 +656,32 @@ impl ExecutionPlan for RangeManipulateExec {
     ) -> DataFusionResult<Arc<dyn ExecutionPlan>> {
         assert!(!children.is_empty());
         let exec_input = children[0].clone();
-        let properties = exec_input.properties();
-        let properties = Arc::new(PlanProperties::new(
-            EquivalenceProperties::new(self.output_schema.clone()),
-            properties.partitioning.clone(),
-            properties.emission_type,
-            properties.boundedness,
-        ));
+        let properties = if let Some(fused) = &self.fused_rate {
+            let ordinary = Arc::new(Self {
+                fused_rate: None,
+                output_schema: fused.folded_schema.clone(),
+                ..self.as_ref().clone()
+            })
+            .replace_children(
+                vec![exec_input.clone()],
+                ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
+            )?;
+            ProjectionExec::try_new_with_schema_metadata(
+                fused.projections.clone(),
+                ordinary,
+                self.output_schema.as_ref(),
+            )?
+            .properties()
+            .clone()
+        } else {
+            let properties = exec_input.properties();
+            Arc::new(PlanProperties::new(
+                EquivalenceProperties::new(self.output_schema.clone()),
+                properties.partitioning.clone(),
+                properties.emission_type,
+                properties.boundedness,
+            ))
+        };
         Ok(Arc::new(Self {
             offset: self.offset,
             start: self.start,
@@ -515,6 +695,7 @@ impl ExecutionPlan for RangeManipulateExec {
             input: children[0].clone(),
             metric: self.metric.clone(),
             properties,
+            fused_rate: self.fused_rate.clone(),
         }))
     }
 
@@ -566,6 +747,7 @@ impl ExecutionPlan for RangeManipulateExec {
             input,
             metric: baseline_metric,
             num_series,
+            fused_rate: self.fused_rate.clone(),
         }))
     }
 
@@ -603,7 +785,11 @@ impl ExecutionPlan for RangeManipulateExec {
     }
 
     fn name(&self) -> &str {
-        "RangeManipulateExec"
+        if self.fused_rate.is_some() {
+            "PromRangeRateExec"
+        } else {
+            "RangeManipulateExec"
+        }
     }
 }
 
@@ -615,8 +801,17 @@ impl DisplayAs for RangeManipulateExec {
             | DisplayFormatType::TreeRender => {
                 write!(
                     f,
-                    "PromRangeManipulateExec: req range=[{}..{}], interval=[{}], eval range=[{}], time index=[{}]",
-                    self.start, self.end, self.interval, self.range, self.time_index_column
+                    "{}: req range=[{}..{}], interval=[{}], eval range=[{}], time index=[{}]",
+                    if self.fused_rate.is_some() {
+                        "PromRangeRateExec"
+                    } else {
+                        "PromRangeManipulateExec"
+                    },
+                    self.start,
+                    self.end,
+                    self.interval,
+                    self.range,
+                    self.time_index_column
                 )
             }
         }
@@ -639,6 +834,7 @@ pub struct RangeManipulateStream {
     metric: BaselineMetrics,
     /// Number of series processed.
     num_series: Count,
+    fused_rate: Option<FusedRate>,
 }
 
 impl RecordBatchStream for RangeManipulateStream {
@@ -681,7 +877,6 @@ impl RangeManipulateStream {
     // But they are not exactly the same, because we don't eager-evaluate on the data in this plan.
     // And the generated timestamp is not aligned to the step. It's expected to do later.
     pub fn manipulate(&self, input: RecordBatch) -> DataFusionResult<Option<RecordBatch>> {
-        let mut other_columns = (0..input.columns().len()).collect::<HashSet<_>>();
         // calculate the range
         let (ranges, (start, end)) = self.calculate_range(&input)?;
         // ignore this if all ranges are empty
@@ -689,7 +884,64 @@ impl RangeManipulateStream {
             return Ok(None);
         }
 
+        if let Some(fused) = &self.fused_rate {
+            let timestamps = input
+                .column(self.time_index)
+                .as_any()
+                .downcast_ref::<TimestampMillisecondArray>()
+                .ok_or_else(|| {
+                    DataFusionError::Execution("Fused rate requires millisecond timestamps".into())
+                })?;
+            let values = input
+                .column(fused.rate_source)
+                .as_any()
+                .downcast_ref::<Float64Array>()
+                .ok_or_else(|| {
+                    DataFusionError::Execution("Fused rate requires Float64 values".into())
+                })?;
+            let aligned = if start != self.start || end != self.end {
+                Self::build_aligned_ts_array(start, end, self.interval)
+            } else {
+                self.aligned_ts_array.clone()
+            };
+            let eval_ts = aligned
+                .as_any()
+                .downcast_ref::<TimestampMillisecondArray>()
+                .ok_or_else(|| {
+                    DataFusionError::Execution(
+                        "Fused rate requires millisecond evaluation timestamps".into(),
+                    )
+                })?;
+            let rates = crate::functions::Rate::calculate_ranges(
+                timestamps, values, &ranges, eval_ts, self.range,
+            )?;
+            let take_indices = Int64Array::from(vec![0; ranges.len()]);
+            let mut columns = Vec::with_capacity(fused.projections.len());
+            for (position, projection) in fused.projections.iter().enumerate() {
+                if position == fused.rate_position {
+                    columns.push(Arc::new(rates.clone()) as ArrayRef);
+                } else if position == fused.timestamp_position {
+                    columns.push(aligned.clone());
+                } else {
+                    let Some(column) = projection.expr.downcast_ref::<Column>() else {
+                        return Err(DataFusionError::Execution(
+                            "Invalid fused rate projection mapping".into(),
+                        ));
+                    };
+                    columns.push(compute::take(
+                        input.column(column.index()),
+                        &take_indices,
+                        None,
+                    )?);
+                }
+            }
+            return RecordBatch::try_new(self.output_schema.clone(), columns)
+                .map(Some)
+                .map_err(|e| DataFusionError::ArrowError(Box::new(e), None));
+        }
+
         // transform columns
+        let mut other_columns = (0..input.columns().len()).collect::<HashSet<_>>();
         let mut new_columns = input.columns().to_vec();
         for index in self.field_columns.iter() {
             let _ = other_columns.remove(index);
@@ -841,15 +1093,20 @@ mod test {
     use datafusion::arrow::datatypes::{
         ArrowPrimitiveType, DataType, Field, Int64Type, Schema, TimestampMillisecondType,
     };
+    use datafusion::common::ScalarValue;
     use datafusion::common::ToDFSchema;
+    use datafusion::config::ConfigOptions;
     use datafusion::datasource::memory::MemorySourceConfig;
     use datafusion::datasource::source::DataSourceExec;
     use datafusion::logical_expr::{
         EmptyRelation, Extension, LogicalPlan, UserDefinedLogicalNodeCore,
     };
     use datafusion::physical_expr::Partitioning;
+    use datafusion::physical_expr::ScalarFunctionExpr;
+    use datafusion::physical_expr::expressions::Literal as PhysicalLiteral;
     use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
     use datafusion::physical_plan::memory::MemoryStream;
+    use datafusion::physical_plan::projection::ProjectionExec;
     use datafusion::physical_plan::{ChildrenPropertiesMode, ReplaceChildrenOptions};
     use datafusion::prelude::SessionContext;
     use datatypes::arrow::array::TimestampMillisecondArray;
@@ -857,6 +1114,281 @@ mod test {
 
     use super::*;
 
+    fn rate_projection(schema: &SchemaRef, duration: i64) -> ProjectionExec {
+        rate_projection_with_udf(schema, duration, crate::functions::Rate::scalar_udf())
+    }
+
+    fn rate_projection_with_udf(
+        schema: &SchemaRef,
+        duration: i64,
+        udf: datafusion::logical_expr::ScalarUDF,
+    ) -> ProjectionExec {
+        let config = Arc::new(ConfigOptions::default());
+        let args = vec![
+            Arc::new(Column::new("timestamp_range", 5)) as Arc<dyn PhysicalExpr>,
+            Arc::new(Column::new("value", 3)),
+            Arc::new(Column::new("timestamp", 4)),
+            Arc::new(PhysicalLiteral::new(ScalarValue::Int64(Some(duration)))),
+        ];
+        let rate_field = Arc::new(Field::new("rate", DataType::Float64, true));
+        let rate = Arc::new(ScalarFunctionExpr::new(
+            "prom_rate",
+            Arc::new(udf),
+            args,
+            rate_field,
+            config.clone(),
+        ));
+        ProjectionExec::try_new(
+            vec![
+                (
+                    Arc::new(Column::new("timestamp", 4)) as Arc<dyn PhysicalExpr>,
+                    "timestamp".to_string(),
+                ),
+                (Arc::new(Column::new("dc", 1)), "dc".to_string()),
+                (rate as Arc<dyn PhysicalExpr>, "rate".to_string()),
+            ],
+            Arc::new(datafusion::physical_plan::empty::EmptyExec::new(
+                schema.clone(),
+            )),
+        )
+        .unwrap()
+    }
+
+    async fn fusion_fixture(
+        start: i64,
+        end: i64,
+        interval: i64,
+        range: i64,
+    ) -> (
+        RangeManipulateExec,
+        Arc<dyn ExecutionPlan>,
+        ProjectionExec,
+        RecordBatch,
+    ) {
+        let raw_schema = Arc::new(Schema::new(vec![
+            Field::new("host", DataType::Utf8, true),
+            Field::new("dc", DataType::Utf8, true),
+            Field::new("unused", DataType::Int64, true),
+            Field::new("value", DataType::Float64, true),
+            Field::new(
+                "timestamp",
+                DataType::Timestamp(TimeUnit::Millisecond, None),
+                false,
+            ),
+        ]));
+        let batch = RecordBatch::try_new(
+            raw_schema.clone(),
+            vec![
+                Arc::new(StringArray::from(vec![Some("h"); 7])),
+                Arc::new(StringArray::from(vec![Some("west"); 7])),
+                Arc::new(Int64Array::from(vec![Some(1); 7])),
+                Arc::new(Float64Array::from(vec![
+                    Some(1.0),
+                    Some(3.0),
+                    Some(0.5),
+                    None,
+                    Some(2.0),
+                    Some(4.0),
+                    Some(7.0),
+                ])),
+                Arc::new(TimestampMillisecondArray::from(vec![
+                    0, 10, 20, 30, 40, 50, 60,
+                ])),
+            ],
+        )
+        .unwrap();
+        let source = Arc::new(DataSourceExec::new(Arc::new(
+            MemorySourceConfig::try_new(&[vec![batch.clone()]], raw_schema.clone(), None).unwrap(),
+        ))) as Arc<dyn ExecutionPlan>;
+        let logical = LogicalPlan::EmptyRelation(EmptyRelation {
+            produce_one_row: false,
+            schema: raw_schema.to_dfschema_ref().unwrap(),
+        });
+        let logical = RangeManipulate::new(
+            start,
+            end,
+            interval,
+            0,
+            range,
+            "timestamp".into(),
+            vec!["value".into()],
+            logical,
+        )
+        .unwrap();
+        let range_exec = logical.to_execution_plan(source.clone());
+        let range_exec = range_exec
+            .downcast_ref::<RangeManipulateExec>()
+            .unwrap()
+            .clone();
+        let folded_schema = range_exec.schema();
+        let projection = rate_projection(&folded_schema, range);
+        (range_exec, source, projection, batch)
+    }
+
+    #[tokio::test]
+    async fn fused_rate_projection_matches_range_projection_for_clipped_windows() {
+        for (start, end, interval, duration) in
+            [(15, 45, 10, 30), (20, 20, 10, 20), (-15, 105, 10, 30)]
+        {
+            let (range, source, projection, _) =
+                fusion_fixture(start, end, interval, duration).await;
+            let fused = range
+                .try_fuse_rate_projection(&projection)
+                .unwrap()
+                .expect("supported rate projection");
+            assert_eq!(fused.name(), "PromRangeRateExec");
+            let fused = fused.downcast_ref::<RangeManipulateExec>().unwrap().clone();
+            assert_eq!(fused.schema(), projection.schema());
+            let replacement = range.input.clone();
+            let fused = Arc::new(fused)
+                .replace_children(
+                    vec![replacement],
+                    ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
+                )
+                .unwrap();
+            assert_eq!(fused.schema(), projection.schema());
+            let ordinary_range = range.input.clone();
+            let ordinary = ProjectionExec::try_new_with_schema_metadata(
+                projection.expr().to_vec(),
+                range_to_exec(&range, ordinary_range),
+                projection.schema().as_ref(),
+            )
+            .unwrap();
+            let expected = datafusion::physical_plan::collect(
+                Arc::new(ordinary),
+                SessionContext::default().task_ctx(),
+            )
+            .await
+            .unwrap();
+            let actual =
+                datafusion::physical_plan::collect(fused, SessionContext::default().task_ctx())
+                    .await
+                    .unwrap();
+            assert_eq!(actual.len(), expected.len());
+            for (actual, expected) in actual.iter().zip(expected.iter()) {
+                assert_eq!(actual.schema(), expected.schema());
+                assert_eq!(actual.column(0), expected.column(0));
+                assert_eq!(actual.column(1), expected.column(1));
+                let a = actual
+                    .column(2)
+                    .as_any()
+                    .downcast_ref::<Float64Array>()
+                    .unwrap();
+                let e = expected
+                    .column(2)
+                    .as_any()
+                    .downcast_ref::<Float64Array>()
+                    .unwrap();
+                assert_eq!(a.len(), e.len());
+                for (a, e) in a.iter().zip(e.iter()) {
+                    match (a, e) {
+                        (None, None) => {}
+                        (Some(a), Some(e)) => assert_eq!(a.to_bits(), e.to_bits()),
+                        _ => panic!("rate null mismatch"),
+                    }
+                }
+            }
+            let _ = source;
+        }
+    }
+
+    #[tokio::test]
+    async fn fused_replacement_recomputes_projection_properties() {
+        let (range, _, projection, _) = fusion_fixture(15, 45, 10, 30).await;
+        let fused = range
+            .try_fuse_rate_projection(&projection)
+            .unwrap()
+            .unwrap();
+        let replacement = Arc::new(DataSourceExec::new(Arc::new(
+            MemorySourceConfig::try_new(&[vec![], vec![]], range.input.schema(), None).unwrap(),
+        ))) as Arc<dyn ExecutionPlan>;
+        let fused = fused
+            .replace_children(
+                vec![replacement.clone()],
+                ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
+            )
+            .unwrap();
+        let ordinary = Arc::new(range.clone())
+            .replace_children(
+                vec![replacement],
+                ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
+            )
+            .unwrap();
+        let ordinary = ProjectionExec::try_new_with_schema_metadata(
+            projection.expr().to_vec(),
+            ordinary,
+            projection.schema().as_ref(),
+        )
+        .unwrap();
+        assert!(matches!(
+            fused.properties().partitioning,
+            Partitioning::UnknownPartitioning(2)
+        ));
+        assert_eq!(
+            fused.properties().partitioning.partition_count(),
+            ordinary.properties().partitioning.partition_count()
+        );
+        assert_eq!(
+            fused.properties().emission_type,
+            ordinary.properties().emission_type
+        );
+        assert_eq!(
+            fused.properties().boundedness,
+            ordinary.properties().boundedness
+        );
+    }
+
+    #[tokio::test]
+    async fn same_name_custom_rate_udf_is_not_fused() {
+        let (range, _, projection, _) = fusion_fixture(15, 45, 10, 30).await;
+        let custom_udf = datafusion_expr::create_udf(
+            "prom_rate",
+            vec![
+                RangeArray::convert_data_type(DataType::Timestamp(TimeUnit::Millisecond, None)),
+                RangeArray::convert_data_type(DataType::Float64),
+                DataType::Timestamp(TimeUnit::Millisecond, None),
+                DataType::Int64,
+            ],
+            DataType::Float64,
+            datafusion::logical_expr::Volatility::Volatile,
+            Arc::new(|_| {
+                Ok(datafusion::physical_plan::ColumnarValue::Scalar(
+                    ScalarValue::Float64(Some(1.0)),
+                ))
+            }),
+        );
+        let custom_projection = rate_projection_with_udf(&range.schema(), 30, custom_udf);
+        assert_eq!(
+            custom_projection.expr()[2]
+                .expr
+                .downcast_ref::<ScalarFunctionExpr>()
+                .unwrap()
+                .name(),
+            projection.expr()[2]
+                .expr
+                .downcast_ref::<ScalarFunctionExpr>()
+                .unwrap()
+                .name()
+        );
+        assert!(
+            range
+                .try_fuse_rate_projection(&custom_projection)
+                .unwrap()
+                .is_none()
+        );
+    }
+
+    fn range_to_exec(
+        range: &RangeManipulateExec,
+        source: Arc<dyn ExecutionPlan>,
+    ) -> Arc<dyn ExecutionPlan> {
+        Arc::new(RangeManipulateExec {
+            input: source,
+            fused_rate: None,
+            ..range.clone()
+        })
+    }
+
     const TIME_INDEX_COLUMN: &str = "timestamp";
 
     fn project_batch(batch: &RecordBatch, indices: &[usize]) -> RecordBatch {
@@ -941,6 +1473,7 @@ mod test {
             input: memory_exec,
             metric: ExecutionPlanMetricsSet::new(),
             properties,
+            fused_rate: None,
         });
         let session_context = SessionContext::default();
         let result = datafusion::physical_plan::collect(normalize_exec, session_context.task_ctx())
@@ -1570,6 +2103,7 @@ mod test {
             input: Box::pin(empty_stream),
             metric: BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0),
             num_series: Count::new(),
+            fused_rate: None,
         };
 
         // Create test data with timestamps not aligned to query pattern
@@ -1699,6 +2233,7 @@ mod test {
             input: Box::pin(empty_stream),
             metric: BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0),
             num_series: Count::new(),
+            fused_rate: None,
         };
         let batch = RecordBatch::try_new(
             schema,

src/promql/src/functions/extrapolate_rate.rs

@@ -31,7 +31,7 @@
 
 use std::fmt::Display;
 use std::ops::Range;
-use std::sync::Arc;
+use std::sync::{Arc, LazyLock};
 
 use datafusion::arrow::array::{Float64Array, Float64Builder, TimestampMillisecondArray};
 use datafusion::arrow::datatypes::TimeUnit;
@@ -179,13 +179,35 @@ impl<const IS_COUNTER: bool, const IS_RATE: bool> ExtrapolatedRate<IS_COUNTER, I
             .values()
             .as_any()
             .downcast_ref::<TimestampMillisecondArray>()
-            .expect("validated by extract_range_dict")
-            .values();
+            .expect("validated by extract_range_dict");
         let value_array = value_dict
             .values()
             .as_any()
             .downcast_ref::<Float64Array>()
             .expect("validated by extract_range_dict");
+        let result = Self::calc_ranges(
+            all_timestamps,
+            value_array,
+            keys.iter().copied().map(unpack),
+            &eval_ts_array,
+            self.range_length,
+        )?;
+        Ok(ColumnarValue::Array(Arc::new(result)))
+    }
+
+    /// Shared numeric kernel for the dictionary range caller and raw range callers.
+    fn calc_ranges<I>(
+        timestamps: &TimestampMillisecondArray,
+        value_array: &Float64Array,
+        ranges: I,
+        eval_ts_array: &TimestampMillisecondArray,
+        range_length: i64,
+    ) -> DfResult<Float64Array>
+    where
+        I: Iterator<Item = (u32, u32)> + ExactSizeIterator + Clone,
+    {
+        let num_windows = ranges.len();
+        let all_timestamps = timestamps.values();
         // A NULL field value means the series has no sample at that timestamp, so the padding
         // under a null slot must not be read. Skip the per-window null scan when the whole
         // backing array is null-free, which is the common case.
@@ -194,7 +216,6 @@ impl<const IS_COUNTER: bool, const IS_RATE: bool> ExtrapolatedRate<IS_COUNTER, I
         let eval_ts = eval_ts_array.values();
 
         let mut result_builder = Float64Builder::with_capacity(num_windows);
-        let range_length = self.range_length;
         let range_length_secs = range_length as f64 / 1000.0;
 
         // Range windows normally overlap heavily, so scanning every one for resets costs far
@@ -205,10 +226,10 @@ impl<const IS_COUNTER: bool, const IS_RATE: bool> ExtrapolatedRate<IS_COUNTER, I
         let mut reset_index = if IS_COUNTER && !has_nulls {
             let budget = all_values.len().saturating_sub(1);
             let mut scanned_pairs = 0usize;
-            keys.iter()
-                .any(|&key| {
-                    scanned_pairs =
-                        scanned_pairs.saturating_add(unpack(key).1.saturating_sub(1) as usize);
+            ranges
+                .clone()
+                .any(|(_, length)| {
+                    scanned_pairs = scanned_pairs.saturating_add(length.saturating_sub(1) as usize);
                     scanned_pairs > budget
                 })
                 .then(|| CounterResetIndex::new(all_values))
@@ -216,8 +237,7 @@ impl<const IS_COUNTER: bool, const IS_RATE: bool> ExtrapolatedRate<IS_COUNTER, I
             None
         };
 
-        for index in 0..num_windows {
-            let (raw_offset, raw_length) = unpack(keys[index]);
+        for (index, (raw_offset, raw_length)) in ranges.enumerate() {
             let offset = raw_offset as usize;
             let length = raw_length as usize;
 
@@ -303,8 +323,43 @@ impl<const IS_COUNTER: bool, const IS_RATE: bool> ExtrapolatedRate<IS_COUNTER, I
             result_builder.append_value(result_value * factor);
         }
 
-        let result = ColumnarValue::Array(Arc::new(result_builder.finish()));
-        Ok(result)
+        Ok(result_builder.finish())
+    }
+}
+
+impl ExtrapolatedRate<true, true> {
+    /// Evaluates raw typed range bounds using the same kernel as the dictionary UDF.
+    pub(crate) fn calculate_ranges(
+        timestamps: &TimestampMillisecondArray,
+        values: &Float64Array,
+        ranges: &[(u32, u32)],
+        evaluation_timestamps: &TimestampMillisecondArray,
+        range_length: i64,
+    ) -> DfResult<Float64Array> {
+        if evaluation_timestamps.len() != ranges.len() {
+            return Err(DataFusionError::Execution(format!(
+                "{}: evaluation timestamp vector should have the same number of rows as range inputs, found {} and {}",
+                Self::func_name(),
+                evaluation_timestamps.len(),
+                ranges.len()
+            )));
+        }
+        for &(offset, length) in ranges {
+            let end = (offset as usize).checked_add(length as usize);
+            if end.is_none_or(|end| end > timestamps.len() || end > values.len()) {
+                return Err(DataFusionError::Execution(format!(
+                    "{}: range ({offset}, {length}) exceeds timestamp or value array bounds",
+                    Self::func_name()
+                )));
+            }
+        }
+        ExtrapolatedRate::<true, true>::calc_ranges(
+            timestamps,
+            values,
+            ranges.iter().copied(),
+            evaluation_timestamps,
+            range_length,
+        )
     }
 }
 
@@ -488,7 +543,9 @@ impl ExtrapolatedRate<true, true> {
     }
 
     pub fn scalar_udf() -> ScalarUDF {
-        Self::scalar_udf_with_name(Self::name())
+        static RATE_UDF: LazyLock<ScalarUDF> =
+            LazyLock::new(|| Rate::scalar_udf_with_name(Rate::name()));
+        RATE_UDF.clone()
     }
 }
 
@@ -627,6 +684,72 @@ mod test {
         }
     }
 
+    #[test]
+    fn raw_rate_ranges_match_dictionary_calc() {
+        let timestamps = Arc::new(TimestampMillisecondArray::from_iter_values(
+            (0..13).map(|i| i * 1_000),
+        ));
+        let ranges = [(0, 0), (0, 1), (0, 3), (0, 4), (1, 4), (2, 4)];
+        let eval_ts = Arc::new(TimestampMillisecondArray::from_iter_values([
+            5_000, 5_000, 5_000, 5_000, 6_000, 7_000,
+        ]));
+        for values in [
+            Float64Array::from(vec![1e16, 1.0, 0.0, 1.0, 2.0, 0.0, 3.0, 2.0, 4.0]),
+            Float64Array::from_iter([
+                Some(1.0),
+                Some(2.0),
+                Some(3.0),
+                None,
+                Some(5.0),
+                Some(1.0),
+                Some(2.0),
+                Some(f64::NAN),
+                Some(4.0),
+            ]),
+        ] {
+            let values = Arc::new(values);
+            let ts_dict = RangeArray::from_ranges(timestamps.clone(), ranges)
+                .unwrap()
+                .into_dict();
+            let values_dict = RangeArray::from_ranges(values.clone(), ranges)
+                .unwrap()
+                .into_dict();
+            let input = [
+                ColumnarValue::Array(Arc::new(ts_dict)),
+                ColumnarValue::Array(Arc::new(values_dict)),
+                ColumnarValue::Array(eval_ts.clone()),
+                ColumnarValue::Scalar(ScalarValue::Int64(Some(5_000))),
+            ];
+            let dictionary_array = extract_array(&Rate::new(5_000).calc(&input).unwrap()).unwrap();
+            let dictionary = dictionary_array
+                .as_any()
+                .downcast_ref::<Float64Array>()
+                .unwrap();
+            let raw =
+                Rate::calculate_ranges(&timestamps, &values, &ranges, &eval_ts, 5_000).unwrap();
+            assert_eq!(dictionary.len(), raw.len());
+            for (expected, actual) in dictionary.iter().zip(raw.iter()) {
+                match (expected, actual) {
+                    (None, None) => {}
+                    (Some(expected), Some(actual)) => {
+                        assert_eq!(expected.to_bits(), actual.to_bits())
+                    }
+                    _ => panic!("dictionary and raw range nulls differ"),
+                }
+            }
+            assert!(
+                Rate::calculate_ranges(
+                    &timestamps,
+                    &values,
+                    &[(8, 2)],
+                    &TimestampMillisecondArray::from(vec![Some(10_000)]),
+                    5_000,
+                )
+                .is_err()
+            );
+        }
+    }
+
     #[test]
     fn counter_resets_accumulate_into_the_running_result() {
         // Summed on their own, 1e16 and 1.0 round to 1e16, which then cancels against the

src/query/src/optimizer.rs

@@ -23,6 +23,7 @@ pub(crate) mod json_schema_concretize;
 pub(crate) mod json_type_concretize;
 pub mod parallelize_scan;
 pub mod pass_distribution;
+pub mod promql_rate_fusion;
 pub mod promql_tsid_narrow_join;
 pub mod remove_duplicate;
 pub mod scan_hint;

src/query/src/optimizer/promql_rate_fusion.rs

@@ -0,0 +1,158 @@
+// Copyright 2023 Greptime Team
+//
+// Licensed 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.
+
+#[cfg(test)]
+mod tests;
+
+use std::any::TypeId;
+use std::sync::Arc;
+
+use datafusion::config::ConfigOptions;
+use datafusion::physical_optimizer::PhysicalOptimizerRule;
+use datafusion::physical_plan::aggregates::{AggregateExec, AggregateMode};
+use datafusion::physical_plan::filter::FilterExec;
+use datafusion::physical_plan::projection::ProjectionExec;
+use datafusion::physical_plan::{ChildrenPropertiesMode, ExecutionPlan, ReplaceChildrenOptions};
+use datafusion_common::Result as DfResult;
+use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
+use datafusion_physical_expr::ScalarFunctionExpr;
+use datafusion_physical_expr::expressions::{Column, IsNotNullExpr};
+use promql::extension_plan::RangeManipulateExec;
+use promql::functions::Rate;
+
+/// Fuses the physical PromQL rate projection into its range manipulation node.
+#[derive(Debug)]
+pub struct PromqlRateFusion;
+
+impl PhysicalOptimizerRule for PromqlRateFusion {
+    fn optimize(
+        &self,
+        plan: Arc<dyn ExecutionPlan>,
+        _config: &ConfigOptions,
+    ) -> DfResult<Arc<dyn ExecutionPlan>> {
+        plan.transform_up(Self::rewrite).data()
+    }
+
+    fn name(&self) -> &str {
+        "PromqlRateFusion"
+    }
+
+    fn schema_check(&self) -> bool {
+        true
+    }
+}
+
+impl PromqlRateFusion {
+    fn rewrite(plan: Arc<dyn ExecutionPlan>) -> DfResult<Transformed<Arc<dyn ExecutionPlan>>> {
+        let Some(aggregate) = plan.downcast_ref::<AggregateExec>() else {
+            return Ok(Transformed::no(plan));
+        };
+        if aggregate.mode() != &AggregateMode::Partial
+            || aggregate.aggr_expr().len() != 1
+            || aggregate.group_expr().has_grouping_set()
+            || aggregate.limit_options().is_some()
+            || aggregate.filter_expr().iter().any(Option::is_some)
+        {
+            return Ok(Transformed::no(plan));
+        }
+        let sum = &aggregate.aggr_expr()[0];
+        if sum.is_distinct()
+            || !sum.order_bys().is_empty()
+            || sum.fun().name() != "sum"
+            || sum.fun().inner().as_ref().type_id()
+                != TypeId::of::<datafusion::functions_aggregate::sum::Sum>()
+            || sum.field().data_type() != &arrow_schema::DataType::Float64
+        {
+            return Ok(Transformed::no(plan));
+        }
+        let sum_args = sum.expressions();
+        let [sum_arg] = sum_args.as_slice() else {
+            return Ok(Transformed::no(plan));
+        };
+        let Some(sum_column) = sum_arg.downcast_ref::<Column>() else {
+            return Ok(Transformed::no(plan));
+        };
+
+        let Some(filter) = aggregate.input().downcast_ref::<FilterExec>() else {
+            return Ok(Transformed::no(plan));
+        };
+        if filter.fetch().is_some() {
+            return Ok(Transformed::no(plan));
+        }
+        let Some(not_null) = filter.predicate().downcast_ref::<IsNotNullExpr>() else {
+            return Ok(Transformed::no(plan));
+        };
+        let Some(rate_column) = not_null.arg().downcast_ref::<Column>() else {
+            return Ok(Transformed::no(plan));
+        };
+        if rate_column.index() != sum_column.index() {
+            return Ok(Transformed::no(plan));
+        }
+        let Some(projection) = filter.input().downcast_ref::<ProjectionExec>() else {
+            return Ok(Transformed::no(plan));
+        };
+        let rate_exprs = projection
+            .expr()
+            .iter()
+            .enumerate()
+            .filter_map(|(index, expr)| {
+                expr.expr
+                    .downcast_ref::<ScalarFunctionExpr>()
+                    .filter(|func| func.name() == Rate::name() && func.args().len() == 4)
+                    .map(|func| (index, func))
+            })
+            .collect::<Vec<_>>();
+        if rate_exprs.len() != 1
+            || projection.expr().len() != aggregate.group_expr().expr().len() + 1
+        {
+            return Ok(Transformed::no(plan));
+        }
+        let (rate_output_index, _) = rate_exprs[0];
+        if rate_output_index != rate_column.index()
+            || projection.expr().iter().enumerate().any(|(index, expr)| {
+                index != rate_output_index && expr.expr.downcast_ref::<Column>().is_none()
+            })
+        {
+            return Ok(Transformed::no(plan));
+        }
+        for group_expr in aggregate.group_expr().expr() {
+            let Some(group_column) = group_expr.0.downcast_ref::<Column>() else {
+                return Ok(Transformed::no(plan));
+            };
+            if group_column.index() == rate_column.index()
+                || projection
+                    .expr()
+                    .get(group_column.index())
+                    .is_none_or(|expr| expr.expr.downcast_ref::<Column>().is_none())
+            {
+                return Ok(Transformed::no(plan));
+            }
+        }
+        let Some(range) = projection.input().downcast_ref::<RangeManipulateExec>() else {
+            return Ok(Transformed::no(plan));
+        };
+        let Some(fused) = range.try_fuse_rate_projection(projection)? else {
+            return Ok(Transformed::no(plan));
+        };
+        let new_filter = aggregate.input().clone().replace_children(
+            vec![fused],
+            ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
+        )?;
+        let new_aggregate = plan.replace_children(
+            vec![new_filter],
+            ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
+        )?;
+        Ok(Transformed::yes(new_aggregate))
+    }
+}

src/query/src/optimizer/promql_rate_fusion/tests.rs

@@ -0,0 +1,491 @@
+// Copyright 2023 Greptime Team
+//
+// Licensed 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.
+
+use std::collections::BTreeMap;
+use std::sync::Arc;
+
+use datafusion::arrow::array::{Array, Float64Array, StringArray, TimestampMillisecondArray};
+use datafusion::arrow::datatypes::{DataType, Field, Schema};
+use datafusion::arrow::record_batch::RecordBatch;
+use datafusion::config::ConfigOptions;
+use datafusion::datasource::memory::MemorySourceConfig;
+use datafusion::datasource::source::DataSourceExec;
+use datafusion::execution::context::{SessionConfig, TaskContext};
+use datafusion::execution::disk_manager::{DiskManagerBuilder, DiskManagerMode};
+use datafusion::execution::memory_pool::{FairSpillPool, MemoryPool};
+use datafusion::execution::runtime_env::{RuntimeEnv, RuntimeEnvBuilder};
+use datafusion::physical_expr::PhysicalSortExpr;
+use datafusion::physical_expr::aggregate::AggregateExprBuilder;
+use datafusion::physical_expr::expressions::{Column, IsNotNullExpr, Literal};
+use datafusion::physical_expr::{PhysicalExpr, ScalarFunctionExpr};
+use datafusion::physical_optimizer::PhysicalOptimizerRule;
+use datafusion::physical_plan::aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy};
+use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec;
+use datafusion::physical_plan::filter::FilterExec;
+use datafusion::physical_plan::projection::ProjectionExec;
+use datafusion::physical_plan::repartition::RepartitionExec;
+use datafusion::physical_plan::sorts::sort::SortExec;
+use datafusion::physical_plan::{ExecutionPlan, displayable};
+use datafusion_common::ScalarValue;
+use datafusion_physical_expr::Partitioning;
+use futures::StreamExt;
+use promql::extension_plan::RangeManipulate;
+use promql::functions::Rate;
+
+use super::PromqlRateFusion;
+
+const PARTITIONS: usize = 4;
+const GROUPS: usize = 64;
+const EVALS: usize = 256;
+const RAW_POINTS: usize = 512;
+const STEP_MS: i64 = 15_000;
+const RANGE_MS: i64 = 300_000;
+// Calibrated by the owner of the physical-rule integration; both plans intentionally
+// receive the same finite, spill-capable runtime, not a deny-all spill pool.
+const MEMORY_LIMIT: usize = 1024 * 1024;
+
+struct Fixture {
+    root: Arc<dyn ExecutionPlan>,
+}
+
+struct Run {
+    runtime: Arc<RuntimeEnv>,
+    pool: Arc<FairSpillPool>,
+}
+
+fn fixture(eval_end: usize) -> Fixture {
+    let schema = Arc::new(Schema::new(vec![
+        Field::new(
+            "ts",
+            DataType::Timestamp(datafusion::arrow::datatypes::TimeUnit::Millisecond, None),
+            false,
+        ),
+        Field::new("val", DataType::Float64, true),
+        Field::new("dc", DataType::Utf8, false),
+    ]));
+    let mut partitions = vec![Vec::new(); PARTITIONS];
+    for group in 0..GROUPS {
+        // Two source series per group exercise partial aggregation across partitions.
+        for series in 0..2 {
+            let partition = (group + series) % PARTITIONS;
+            let times = (0..RAW_POINTS)
+                .map(|i| i as i64 * STEP_MS)
+                .collect::<Vec<_>>();
+            let values = (0..RAW_POINTS)
+                .map(|i| {
+                    // Integer slopes and one counter reset produce exactly representable
+                    // rates; nulls exercise the projection filter.
+                    if i % 97 == 0 {
+                        None
+                    } else if (300..304).contains(&i) {
+                        Some((i - 300) as f64)
+                    } else {
+                        Some((i as i64 * (group as i64 % 7 + 2) + series as i64 * 20) as f64)
+                    }
+                })
+                .collect::<Vec<_>>();
+            let labels = vec![format!("datacenter-{group:03}"); RAW_POINTS];
+            partitions[partition].push(
+                RecordBatch::try_new(
+                    schema.clone(),
+                    vec![
+                        Arc::new(TimestampMillisecondArray::from(times)),
+                        Arc::new(Float64Array::from(values)),
+                        Arc::new(StringArray::from(labels)),
+                    ],
+                )
+                .unwrap(),
+            );
+        }
+    }
+    let source = Arc::new(DataSourceExec::new(Arc::new(
+        MemorySourceConfig::try_new(&partitions, schema.clone(), None).unwrap(),
+    ))) as Arc<dyn ExecutionPlan>;
+    let empty = datafusion::logical_expr::LogicalPlan::EmptyRelation(
+        datafusion::logical_expr::EmptyRelation {
+            produce_one_row: false,
+            schema: Arc::new(
+                datafusion_common::DFSchema::try_from(schema.as_ref().clone()).unwrap(),
+            ),
+        },
+    );
+    let range = RangeManipulate::new(
+        (RAW_POINTS as i64 - eval_end as i64) * STEP_MS,
+        (RAW_POINTS as i64 - 1) * STEP_MS,
+        STEP_MS,
+        0,
+        RANGE_MS,
+        "ts".to_string(),
+        vec!["val".to_string()],
+        empty,
+    )
+    .unwrap()
+    .to_execution_plan(Arc::clone(&source));
+
+    let ts = Arc::new(Column::new("ts", 0)) as Arc<dyn PhysicalExpr>;
+    let values = Arc::new(Column::new("val", 1)) as Arc<dyn PhysicalExpr>;
+    let eval_time = Arc::new(Column::new("ts", 0)) as Arc<dyn PhysicalExpr>;
+    let range_literal =
+        Arc::new(Literal::new(ScalarValue::Int64(Some(RANGE_MS)))) as Arc<dyn PhysicalExpr>;
+    let rate = Arc::new(ScalarFunctionExpr::new(
+        "prom_rate",
+        Arc::new(Rate::scalar_udf()),
+        vec![
+            Arc::new(Column::new("ts_range", 3)),
+            values,
+            eval_time,
+            range_literal,
+        ],
+        Arc::new(Field::new("rate", DataType::Float64, true)),
+        Arc::new(ConfigOptions::default()),
+    )) as Arc<dyn PhysicalExpr>;
+    let dc = Arc::new(Column::new("dc", 2)) as Arc<dyn PhysicalExpr>;
+    let projection = Arc::new(
+        ProjectionExec::try_new(
+            vec![
+                (ts, "ts".to_string()),
+                (rate, "rate".to_string()),
+                (dc, "dc".to_string()),
+            ],
+            Arc::clone(&range),
+        )
+        .unwrap(),
+    );
+    let filter = Arc::new(
+        FilterExec::try_new(
+            Arc::new(IsNotNullExpr::new(Arc::new(Column::new("rate", 1)))),
+            projection.clone(),
+        )
+        .unwrap(),
+    ) as Arc<dyn ExecutionPlan>;
+    let group = PhysicalGroupBy::new_single(vec![
+        (
+            Arc::new(Column::new("dc", 2)) as Arc<dyn PhysicalExpr>,
+            "dc".to_string(),
+        ),
+        (
+            Arc::new(Column::new("ts", 0)) as Arc<dyn PhysicalExpr>,
+            "ts".to_string(),
+        ),
+    ]);
+    let sum = Arc::new(
+        AggregateExprBuilder::new(
+            datafusion::functions_aggregate::sum::sum_udaf(),
+            vec![Arc::new(Column::new("rate", 1))],
+        )
+        .schema(filter.schema())
+        .alias("sum_rate")
+        .build()
+        .unwrap(),
+    );
+    let partial = Arc::new(
+        AggregateExec::try_new(
+            AggregateMode::Partial,
+            group.clone(),
+            vec![sum.clone()],
+            vec![None],
+            filter,
+            projection.schema(),
+        )
+        .unwrap(),
+    ) as Arc<dyn ExecutionPlan>;
+    let repartition = Arc::new(
+        RepartitionExec::try_new(
+            partial,
+            Partitioning::Hash(
+                vec![
+                    Arc::new(Column::new("dc", 0)),
+                    Arc::new(Column::new("ts", 1)),
+                ],
+                PARTITIONS,
+            ),
+        )
+        .unwrap(),
+    ) as Arc<dyn ExecutionPlan>;
+    let final_sum = Arc::new(
+        AggregateExprBuilder::new(
+            datafusion::functions_aggregate::sum::sum_udaf(),
+            vec![Arc::new(Column::new("sum_rate", 2))],
+        )
+        .schema(repartition.schema())
+        .alias("sum_rate")
+        .build()
+        .unwrap(),
+    );
+    let final_aggregate = Arc::new(
+        AggregateExec::try_new(
+            AggregateMode::FinalPartitioned,
+            group.as_final(),
+            vec![final_sum],
+            vec![None],
+            repartition.clone(),
+            repartition.schema(),
+        )
+        .unwrap(),
+    ) as Arc<dyn ExecutionPlan>;
+    let coalesced = Arc::new(CoalescePartitionsExec::new(final_aggregate));
+    let sort = SortExec::new(
+        [
+            PhysicalSortExpr::new(Arc::new(Column::new("dc", 0)), Default::default()),
+            PhysicalSortExpr::new(Arc::new(Column::new("ts", 1)), Default::default()),
+        ]
+        .into(),
+        coalesced,
+    );
+    Fixture {
+        root: Arc::new(sort),
+    }
+}
+
+fn run_env() -> Run {
+    let pool = Arc::new(FairSpillPool::new(MEMORY_LIMIT));
+    let runtime = Arc::new(
+        RuntimeEnvBuilder::new()
+            .with_memory_pool(pool.clone())
+            .with_disk_manager_builder(
+                DiskManagerBuilder::default().with_mode(DiskManagerMode::OsTmpDirectory),
+            )
+            .build()
+            .unwrap(),
+    );
+    Run { runtime, pool }
+}
+
+fn collect_spill_stats(plan: &Arc<dyn ExecutionPlan>) -> (usize, usize, usize) {
+    let mut stats = (0, 0, 0);
+    if let Some(aggregate) = plan.downcast_ref::<AggregateExec>()
+        && matches!(
+            aggregate.mode(),
+            AggregateMode::Final | AggregateMode::FinalPartitioned
+        )
+        && let Some(metrics) = aggregate.metrics()
+    {
+        stats.0 += metrics.spill_count().unwrap_or(0);
+        stats.1 += metrics.spilled_bytes().unwrap_or(0);
+        stats.2 += metrics.spilled_rows().unwrap_or(0);
+    }
+    for child in plan.children() {
+        let child_stats = collect_spill_stats(child);
+        stats.0 += child_stats.0;
+        stats.1 += child_stats.1;
+        stats.2 += child_stats.2;
+    }
+    stats
+}
+
+fn task_context(runtime: &Arc<RuntimeEnv>, migration: bool) -> Arc<TaskContext> {
+    let mut config = SessionConfig::new()
+        .with_target_partitions(PARTITIONS)
+        .with_batch_size(128)
+        .with_sort_spill_reservation_bytes(32 * 1024);
+    config.options_mut().execution.enable_migration_aggregate = migration;
+    let state = datafusion::execution::SessionStateBuilder::new()
+        .with_config(config)
+        .with_runtime_env(Arc::clone(runtime))
+        .build();
+    Arc::new(TaskContext::from(&state))
+}
+
+async fn execute_full_with_held_batch(
+    plan: Arc<dyn ExecutionPlan>,
+    runtime: &Arc<RuntimeEnv>,
+    migration: bool,
+) -> Vec<RecordBatch> {
+    let mut stream = plan.execute(0, task_context(runtime, migration)).unwrap();
+    let held = stream.next().await.unwrap().unwrap();
+    let held_snapshot = rows(std::slice::from_ref(&held));
+    let mut batches = vec![held.clone()];
+    while let Some(batch) = stream.next().await {
+        batches.push(batch.unwrap());
+    }
+    if batches.iter().map(RecordBatch::num_rows).sum::<usize>() > 128 {
+        assert!(batches.len() >= 2, "expected multiple output batches");
+    }
+    assert_eq!(rows(std::slice::from_ref(&held)), held_snapshot);
+    batches
+}
+
+fn rows(batches: &[RecordBatch]) -> BTreeMap<(String, i64), f64> {
+    let mut output = BTreeMap::new();
+    for batch in batches {
+        let dcs = batch
+            .column(0)
+            .as_any()
+            .downcast_ref::<StringArray>()
+            .unwrap();
+        let times = batch
+            .column(1)
+            .as_any()
+            .downcast_ref::<TimestampMillisecondArray>()
+            .unwrap();
+        let sums = batch
+            .column(2)
+            .as_any()
+            .downcast_ref::<Float64Array>()
+            .unwrap();
+        for row in 0..batch.num_rows() {
+            let key = (dcs.value(row).to_string(), times.value(row));
+            assert!(
+                !output.contains_key(&key),
+                "duplicate output key {key:?}; partitioned sort lost rows or duplicated groups"
+            );
+            output.insert(
+                key,
+                if sums.is_null(row) {
+                    f64::NAN
+                } else {
+                    sums.value(row)
+                },
+            );
+        }
+    }
+    output
+}
+
+fn assert_results_equal(
+    actual: &BTreeMap<(String, i64), f64>,
+    expected: &BTreeMap<(String, i64), f64>,
+    migration: bool,
+) {
+    assert_eq!(
+        actual.keys().collect::<Vec<_>>(),
+        expected.keys().collect::<Vec<_>>()
+    );
+    for (key, expected_value) in expected {
+        let actual_value = actual.get(key).unwrap();
+        assert!(
+            actual_value.to_bits() == expected_value.to_bits()
+                || (actual_value - expected_value).abs() <= 1e-12,
+            "{key:?}: actual={actual_value}, expected={expected_value}, migration={migration}"
+        );
+    }
+}
+
+#[tokio::test]
+async fn bounded_full_pipeline_rate_fusion_matches_spilling_baseline_and_releases_resources() {
+    for migration in [true, false] {
+        for eval_end in [1, EVALS] {
+            let fixture = fixture(eval_end);
+            let baseline_run = run_env();
+            let fused_run = run_env();
+            let mut config = ConfigOptions::default();
+            config.execution.enable_migration_aggregate = migration;
+            let optimized = PromqlRateFusion
+                .optimize(Arc::clone(&fixture.root), &config)
+                .unwrap();
+            let before = displayable(fixture.root.as_ref()).indent(true).to_string();
+            let after = displayable(optimized.as_ref()).indent(true).to_string();
+            assert_ne!(after, before, "migration={migration}, eval_end={eval_end}");
+            assert!(after.contains("PromRangeRateExec"), "{after}");
+            for retained in [
+                "FilterExec",
+                "Partial",
+                "FinalPartitioned",
+                "RepartitionExec",
+                "SortExec",
+            ] {
+                assert!(after.contains(retained), "missing {retained} in {after}");
+            }
+
+            let expected = execute_full_with_held_batch(
+                Arc::clone(&fixture.root),
+                &baseline_run.runtime,
+                migration,
+            )
+            .await;
+            let actual =
+                execute_full_with_held_batch(Arc::clone(&optimized), &fused_run.runtime, migration)
+                    .await;
+            let expected_rows = rows(&expected);
+            let actual_rows = rows(&actual);
+            assert_results_equal(&actual_rows, &expected_rows, migration);
+            assert_eq!(expected_rows.len(), GROUPS * eval_end);
+            assert_eq!(actual_rows.len(), GROUPS * eval_end);
+            for group in 0..GROUPS {
+                for eval in 0..eval_end {
+                    let key = (
+                        format!("datacenter-{group:03}"),
+                        (RAW_POINTS as i64 - eval_end as i64 + eval as i64) * STEP_MS,
+                    );
+                    assert!(
+                        expected_rows.contains_key(&key),
+                        "missing baseline key {key:?}"
+                    );
+                    assert!(actual_rows.contains_key(&key), "missing fused key {key:?}");
+                }
+            }
+            let spill = collect_spill_stats(&optimized);
+            let baseline_spill = collect_spill_stats(&fixture.root);
+            if eval_end == EVALS {
+                assert!(
+                    spill.0 > 0 && spill.1 > 0 && spill.2 > 0,
+                    "fused final aggregate spill={spill:?}"
+                );
+                assert!(
+                    baseline_spill.0 > 0 && baseline_spill.1 > 0 && baseline_spill.2 > 0,
+                    "baseline final aggregate spill={baseline_spill:?}"
+                );
+            }
+            assert!(baseline_run.pool.reserved() <= MEMORY_LIMIT);
+            assert!(fused_run.pool.reserved() <= MEMORY_LIMIT);
+            assert_eq!(
+                baseline_run.pool.reserved(),
+                0,
+                "baseline reservations remain"
+            );
+            assert_eq!(fused_run.pool.reserved(), 0, "fused reservations remain");
+            for run in [&baseline_run, &fused_run] {
+                for dir in run.runtime.disk_manager.temp_dir_paths() {
+                    if dir.exists() {
+                        assert_eq!(
+                            std::fs::read_dir(dir).unwrap().count(),
+                            0,
+                            "spill files leaked"
+                        );
+                    }
+                }
+            }
+        }
+    }
+}
+
+#[tokio::test]
+async fn early_drop_preserves_held_output_and_releases_memory_reservations() {
+    for migration in [true, false] {
+        let fixture = fixture(EVALS);
+        let run = run_env();
+        let mut config = ConfigOptions::default();
+        config.execution.enable_migration_aggregate = migration;
+        let optimized = PromqlRateFusion
+            .optimize(Arc::clone(&fixture.root), &config)
+            .unwrap();
+        let mut stream = optimized
+            .execute(0, task_context(&run.runtime, migration))
+            .unwrap();
+        let held = stream.next().await.unwrap().unwrap();
+        let held_snapshot = rows(std::slice::from_ref(&held));
+        assert!(!held_snapshot.is_empty());
+        drop(stream);
+        tokio::task::yield_now().await;
+        assert_eq!(rows(std::slice::from_ref(&held)), held_snapshot);
+        drop((optimized, fixture));
+        assert!(held.num_rows() > 0);
+        assert_eq!(run.pool.reserved(), 0);
+        for dir in run.runtime.disk_manager.temp_dir_paths() {
+            if dir.exists() {
+                assert_eq!(std::fs::read_dir(dir).unwrap().count(), 0);
+            }
+        }
+    }
+}

src/query/src/query_engine/state.rs

@@ -71,6 +71,7 @@ use crate::optimizer::json_schema_concretize::JsonSchemaConcretizeRule;
 use crate::optimizer::json_type_concretize::JsonTypeConcretizeRule;
 use crate::optimizer::parallelize_scan::ParallelizeScan;
 use crate::optimizer::pass_distribution::PassDistribution;
+use crate::optimizer::promql_rate_fusion::PromqlRateFusion;
 use crate::optimizer::promql_tsid_narrow_join::PromqlTsidNarrowJoin;
 use crate::optimizer::remove_duplicate::RemoveDuplicate;
 use crate::optimizer::scan_hint::ScanHintRule;
@@ -257,8 +258,9 @@ impl QueryEngineState {
         physical_optimizer
             .rules
             .push(Arc::new(EnsureGlobalLimitForFetch));
-        // Add rule to remove duplicate nodes generated by other rules. Run this in the last.
+        // Add rule to remove duplicate nodes generated by preceding rules.
         physical_optimizer.rules.push(Arc::new(RemoveDuplicate));
+        physical_optimizer.rules.push(Arc::new(PromqlRateFusion));
         // Place SanityCheckPlan at the end of the list to ensure that it runs after all other rules.
         Self::remove_physical_optimizer_rule(
             &mut physical_optimizer.rules,