diff --git a/src/promql/src/extension_plan/range_manipulate.rs b/src/promql/src/extension_plan/range_manipulate.rs index 8ea7ba6b191..bb73376464e 100644 --- a/src/promql/src/extension_plan/range_manipulate.rs +++ b/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, + fused_rate: Option, +} + +#[derive(Debug, Clone)] +struct FusedRate { + rate_source: usize, + rate_position: usize, + timestamp_position: usize, + projections: Vec, + 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>> { + 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::() { + 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::() else { + return Ok(None); + }; + let Some(value_arg) = function.args()[1].downcast_ref::() else { + return Ok(None); + }; + let Some(eval_arg) = function.args()[2].downcast_ref::() else { + return Ok(None); + }; + let Some(range_arg) = function.args()[3].downcast_ref::() 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::() 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::() + .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; + 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> { 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, } 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> { - let mut other_columns = (0..input.columns().len()).collect::>(); // 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::() + .ok_or_else(|| { + DataFusionError::Execution("Fused rate requires millisecond timestamps".into()) + })?; + let values = input + .column(fused.rate_source) + .as_any() + .downcast_ref::() + .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::() + .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::() 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::>(); 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, + 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, + "timestamp".to_string(), + ), + (Arc::new(Column::new("dc", 1)), "dc".to_string()), + (rate as Arc, "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, + 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; + 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::() + .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::().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::() + .unwrap(); + let e = expected + .column(2) + .as_any() + .downcast_ref::() + .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; + 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::() + .unwrap() + .name(), + projection.expr()[2] + .expr + .downcast_ref::() + .unwrap() + .name() + ); + assert!( + range + .try_fuse_rate_projection(&custom_projection) + .unwrap() + .is_none() + ); + } + + fn range_to_exec( + range: &RangeManipulateExec, + source: Arc, + ) -> Arc { + 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, diff --git a/src/promql/src/functions/extrapolate_rate.rs b/src/promql/src/functions/extrapolate_rate.rs index 263ae6337ba..83e8351e700 100644 --- a/src/promql/src/functions/extrapolate_rate.rs +++ b/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 ExtrapolatedRate() - .expect("validated by extract_range_dict") - .values(); + .expect("validated by extract_range_dict"); let value_array = value_dict .values() .as_any() .downcast_ref::() .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( + timestamps: &TimestampMillisecondArray, + value_array: &Float64Array, + ranges: I, + eval_ts_array: &TimestampMillisecondArray, + range_length: i64, + ) -> DfResult + where + I: Iterator + 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 ExtrapolatedRate ExtrapolatedRate budget }) .then(|| CounterResetIndex::new(all_values)) @@ -216,8 +237,7 @@ impl ExtrapolatedRate ExtrapolatedRate { + /// 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 { + 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::::calc_ranges( + timestamps, + values, + ranges.iter().copied(), + evaluation_timestamps, + range_length, + ) } } @@ -488,7 +543,9 @@ impl ExtrapolatedRate { } pub fn scalar_udf() -> ScalarUDF { - Self::scalar_udf_with_name(Self::name()) + static RATE_UDF: LazyLock = + 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::() + .unwrap(); + let raw = + Rate::calculate_ranges(×tamps, &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( + ×tamps, + &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 diff --git a/src/query/src/optimizer.rs b/src/query/src/optimizer.rs index fec001e7506..bb589a64890 100644 --- a/src/query/src/optimizer.rs +++ b/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; diff --git a/src/query/src/optimizer/promql_rate_fusion.rs b/src/query/src/optimizer/promql_rate_fusion.rs new file mode 100644 index 00000000000..f380b24580a --- /dev/null +++ b/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, + _config: &ConfigOptions, + ) -> DfResult> { + plan.transform_up(Self::rewrite).data() + } + + fn name(&self) -> &str { + "PromqlRateFusion" + } + + fn schema_check(&self) -> bool { + true + } +} + +impl PromqlRateFusion { + fn rewrite(plan: Arc) -> DfResult>> { + let Some(aggregate) = plan.downcast_ref::() 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::() + || 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::() else { + return Ok(Transformed::no(plan)); + }; + + let Some(filter) = aggregate.input().downcast_ref::() else { + return Ok(Transformed::no(plan)); + }; + if filter.fetch().is_some() { + return Ok(Transformed::no(plan)); + } + let Some(not_null) = filter.predicate().downcast_ref::() else { + return Ok(Transformed::no(plan)); + }; + let Some(rate_column) = not_null.arg().downcast_ref::() 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::() else { + return Ok(Transformed::no(plan)); + }; + let rate_exprs = projection + .expr() + .iter() + .enumerate() + .filter_map(|(index, expr)| { + expr.expr + .downcast_ref::() + .filter(|func| func.name() == Rate::name() && func.args().len() == 4) + .map(|func| (index, func)) + }) + .collect::>(); + 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::().is_none() + }) + { + return Ok(Transformed::no(plan)); + } + for group_expr in aggregate.group_expr().expr() { + let Some(group_column) = group_expr.0.downcast_ref::() 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::().is_none()) + { + return Ok(Transformed::no(plan)); + } + } + let Some(range) = projection.input().downcast_ref::() 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)) + } +} diff --git a/src/query/src/optimizer/promql_rate_fusion/tests.rs b/src/query/src/optimizer/promql_rate_fusion/tests.rs new file mode 100644 index 00000000000..716a2f8a2c8 --- /dev/null +++ b/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, +} + +struct Run { + runtime: Arc, + pool: Arc, +} + +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::>(); + 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::>(); + 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; + 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; + let values = Arc::new(Column::new("val", 1)) as Arc; + let eval_time = Arc::new(Column::new("ts", 0)) as Arc; + let range_literal = + Arc::new(Literal::new(ScalarValue::Int64(Some(RANGE_MS)))) as Arc; + 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; + let dc = Arc::new(Column::new("dc", 2)) as Arc; + 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; + let group = PhysicalGroupBy::new_single(vec![ + ( + Arc::new(Column::new("dc", 2)) as Arc, + "dc".to_string(), + ), + ( + Arc::new(Column::new("ts", 0)) as Arc, + "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; + 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; + 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; + 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) -> (usize, usize, usize) { + let mut stats = (0, 0, 0); + if let Some(aggregate) = plan.downcast_ref::() + && 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, migration: bool) -> Arc { + 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, + runtime: &Arc, + migration: bool, +) -> Vec { + 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::() > 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::() + .unwrap(); + let times = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + let sums = batch + .column(2) + .as_any() + .downcast_ref::() + .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::>(), + expected.keys().collect::>() + ); + 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); + } + } + } +} diff --git a/src/query/src/query_engine/state.rs b/src/query/src/query_engine/state.rs index aa3b9417b14..059f030a568 100644 --- a/src/query/src/query_engine/state.rs +++ b/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,