rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-11-15 15:26:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

chore: merge narthex into feat/narthex-cross-table

Browse Source 
Signed-off-by: Dennis Zhuang <killme2008@gmail.com>
This commit is contained in:
Dennis Zhuang
2026-10-27 18:20:21 +08:00
parent ba8156be47 b29a471575
commit 02eb3aa60c
54 changed files with 2269 additions and 9689 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/query/src/datafusion.rs
+133 -13
View File
@@ -30,7 +30,9 @@ use common_error::ext::BoxedError;
use common_function::function::FunctionContext;
use common_function::function_factory::ScalarFunctionFactory;
use common_query::{Output, OutputData, OutputMeta};
use common_recordbatch::adapter::{RecordBatchStreamAdapter, RegionQueryStatCounters};
use common_recordbatch::adapter::{
RecordBatchStreamAdapter, RegionQueryStatCounters, SERIES_SCAN_EXEC_NAME,
};
use common_recordbatch::{EmptyRecordBatchStream, SendableRecordBatchStream};
use common_telemetry::tracing;
use datafusion::catalog::TableFunction;
@@ -50,7 +52,7 @@ use snafu::{OptionExt, ResultExt, ensure};
use sqlparser::ast::AnalyzeFormat;
use table::TableRef;
use table::requests::{DeleteRequest, InsertRequest};
use table::table::scan::{REGION_SCAN_EXEC_NAME, RegionScanExec};
use table::table::scan::{REGION_SCAN_EXEC_NAME, RegionScanExec, SeriesScanExec};
use tracing::Span;
use crate::analyze::DistAnalyzeExec;
@@ -87,10 +89,16 @@ fn query_load_region_id(plan: &Arc<dyn ExecutionPlan>) -> Option<u64> {
let mut stack = vec![plan.clone()];
while let Some(plan) = stack.pop() {
if plan.name() == REGION_SCAN_EXEC_NAME
&& let Some(scan) = plan.downcast_ref::<RegionScanExec>()
&& let Some(scan_region_id) = scan.query_load_region_id()
{
let scan_region_id = if plan.name() == REGION_SCAN_EXEC_NAME {
plan.downcast_ref::<RegionScanExec>()
.and_then(RegionScanExec::query_load_region_id)
} else if plan.name() == SERIES_SCAN_EXEC_NAME {
plan.downcast_ref::<SeriesScanExec>()
.and_then(SeriesScanExec::query_load_region_id)
} else {
None
};
if let Some(scan_region_id) = scan_region_id {
match region_id {
Some(region_id) if region_id != scan_region_id => return None,
Some(_) => {}
@@ -108,7 +116,7 @@ fn query_load_region_id(plan: &Arc<dyn ExecutionPlan>) -> Option<u64> {
// Unlike the Prometheus read-load reporting in `MergeScanExec`, the heartbeat
// counters must be updated before metrics leave the datanode process. The
// `RecordBatchStreamAdapter` that resolves `RecordBatchMetrics` does not know
// the owning `MitoRegion`, so we extract the counters from `RegionScanExec` and
// the owning `MitoRegion`, so we extract the counters from the scan execs and
// pass them to the adapter. If a plan contains scans from different regions,
// return `None` to avoid charging the whole plan metrics to one region.
fn query_stat_counters(plan: &Arc<dyn ExecutionPlan>) -> Option<RegionQueryStatCounters> {
@@ -116,10 +124,16 @@ fn query_stat_counters(plan: &Arc<dyn ExecutionPlan>) -> Option<RegionQueryStatC
let mut stack = vec![plan.clone()];
while let Some(plan) = stack.pop() {
if plan.name() == REGION_SCAN_EXEC_NAME
&& let Some(scan) = plan.downcast_ref::<RegionScanExec>()
&& let Some(scan_counters) = scan.query_stat_counters()
{
let scan_counters = if plan.name() == REGION_SCAN_EXEC_NAME {
plan.downcast_ref::<RegionScanExec>()
.and_then(RegionScanExec::query_stat_counters)
} else if plan.name() == SERIES_SCAN_EXEC_NAME {
plan.downcast_ref::<SeriesScanExec>()
.and_then(SeriesScanExec::query_stat_counters)
} else {
None
};
if let Some(scan_counters) = scan_counters {
match &counters {
Some(counters)
if !Arc::ptr_eq(&counters.query_cpu_time, &scan_counters.query_cpu_time)
@@ -758,7 +772,7 @@ mod tests {
};
use datafusion::physical_plan::display::{DisplayAs, DisplayFormatType};
use datafusion::physical_plan::expressions::PhysicalSortExpr;
use datafusion::physical_plan::joins::{HashJoinExec, JoinOn, PartitionMode};
use datafusion::physical_plan::joins::{CrossJoinExec, HashJoinExec, JoinOn, PartitionMode};
use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet;
use datafusion::physical_plan::{ExecutionPlan, PhysicalExpr};
use datafusion::prelude::{col, lit};
@@ -770,7 +784,8 @@ mod tests {
use session::context::{QueryContext, QueryContextBuilder};
use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder, RegionMetadataRef};
use store_api::region_engine::{
PartitionRange, PrepareRequest, QueryScanContext, RegionScanner, ScannerProperties,
PartitionRange, PrepareRequest, QueryScanContext, RawLabelSource, RegionScanner,
ScannerProperties,
};
use store_api::storage::{RegionId, ScanRequest};
use table::metadata::{TableInfoBuilder, TableMetaBuilder};
@@ -844,6 +859,17 @@ mod tests {
true
}
fn raw_series_schema(
&self,
_label_source: &RawLabelSource,
) -> std::result::Result<Option<Arc<datatypes::series::samples::SeriesSchema>>, BoxedError>
{
Ok(Some(Arc::new(datatypes::series::samples::SeriesSchema {
fields: Arc::from([datatypes::series::samples::FieldSpec { name: "ts".into() }]),
eval_grid: None,
})))
}
fn add_dyn_filter_to_predicate(
&mut self,
filter_exprs: Vec<Arc<dyn PhysicalExpr>>,
@@ -885,6 +911,13 @@ mod tests {
query_load_region_id: Option<RegionId>,
query_stat_counters: Option<RegionQueryStatCounters>,
) -> Arc<dyn ExecutionPlan> {
build_region_scan_exec(query_load_region_id, query_stat_counters)
}
fn build_region_scan_exec(
query_load_region_id: Option<RegionId>,
query_stat_counters: Option<RegionQueryStatCounters>,
) -> Arc<RegionScanExec> {
let schema = Arc::new(datatypes::schema::Schema::new(vec![ColumnSchema::new(
"ts",
ConcreteDataType::timestamp_millisecond_datatype(),
@@ -921,6 +954,23 @@ mod tests {
Arc::new(RegionScanExec::new(Box::new(scanner), ScanRequest::default(), None).unwrap())
}
fn build_series_scan(
query_load_region_id: Option<RegionId>,
query_stat_counters: Option<RegionQueryStatCounters>,
) -> Arc<dyn ExecutionPlan> {
build_region_scan_exec(query_load_region_id, query_stat_counters)
.try_to_series_scan(RawLabelSource::Explicit)
.unwrap()
.expect("RecordingScanner supports a native series scan")
}
fn build_mixed_scan_plan(
left: Arc<dyn ExecutionPlan>,
right: Arc<dyn ExecutionPlan>,
) -> Arc<dyn ExecutionPlan> {
Arc::new(CrossJoinExec::new(left, right))
}
fn query_stat_counters_for_test() -> RegionQueryStatCounters {
RegionQueryStatCounters {
query_cpu_time: Arc::new(AtomicU64::new(0)),
@@ -928,6 +978,33 @@ mod tests {
}
}
#[test]
fn query_load_region_id_reads_native_series_scan() {
let region_id = RegionId::new(1024, 42);
let scan = build_series_scan(Some(region_id), None);
assert_eq!(super::query_load_region_id(&scan), Some(region_id.as_u64()));
}
#[test]
fn query_load_region_id_accepts_mixed_scans_for_same_region() {
let region_id = RegionId::new(1024, 42);
let left = build_query_load_region_scan(Some(region_id));
let right = build_series_scan(Some(region_id), None);
let plan = build_mixed_scan_plan(left, right);
assert_eq!(super::query_load_region_id(&plan), Some(region_id.as_u64()));
}
#[test]
fn query_load_region_id_rejects_mixed_region_scans() {
let left = build_query_load_region_scan(Some(RegionId::new(1024, 42)));
let right = build_series_scan(Some(RegionId::new(1024, 43)), None);
let plan = build_mixed_scan_plan(left, right);
assert!(super::query_load_region_id(&plan).is_none());
}
#[test]
fn query_load_region_id_ignores_scans_without_region_id() {
let query_load_region_id = RegionId::new(1024, 42);
@@ -958,6 +1035,49 @@ mod tests {
);
}
#[test]
fn query_stat_counters_reads_native_series_scan() {
let counters = query_stat_counters_for_test();
let scan = build_series_scan(None, Some(counters.clone()));
let actual = super::query_stat_counters(&scan).unwrap();
assert!(Arc::ptr_eq(
&actual.query_cpu_time,
&counters.query_cpu_time
));
assert!(Arc::ptr_eq(
&actual.query_scanned_bytes,
&counters.query_scanned_bytes
));
}
#[test]
fn query_stat_counters_returns_shared_counter_for_mixed_scan_plan() {
let counters = query_stat_counters_for_test();
let left = build_query_stat_counter_region_scan(counters.clone());
let right = build_series_scan(None, Some(counters.clone()));
let plan = build_mixed_scan_plan(left, right);
let actual = super::query_stat_counters(&plan).unwrap();
assert!(Arc::ptr_eq(
&actual.query_cpu_time,
&counters.query_cpu_time
));
assert!(Arc::ptr_eq(
&actual.query_scanned_bytes,
&counters.query_scanned_bytes
));
}
#[test]
fn query_stat_counters_rejects_mixed_counter_plan_with_series_scan() {
let left = build_query_stat_counter_region_scan(query_stat_counters_for_test());
let right = build_series_scan(None, Some(query_stat_counters_for_test()));
let plan = build_mixed_scan_plan(left, right);
assert!(super::query_stat_counters(&plan).is_none());
}
#[test]
fn query_stat_counters_returns_shared_counter_for_multi_scan_plan() {
let counters = query_stat_counters_for_test();
src/query/src/dist_plan/merge_scan.rs
+2 -289
View File
@@ -69,7 +69,6 @@ use crate::dist_plan::dyn_filter_bridge::{
query_context_with_refreshed_initial_dyn_filter_regs,
register_dyn_filter_subscribers_for_region, register_remote_dyn_filters,
};
use crate::dist_plan::planner::CaptureNativeRawRoute;
use crate::dist_plan::{
FilterId, RemoteDynFilterProducerId, RemoteDynFilterRegistryLease, Subscriber,
};
@@ -495,32 +494,6 @@ pub struct MergeScanExec {
captured_remote_dyn_filters: Arc<Mutex<Vec<CapturedDynFilter>>>,
target_partition: usize,
enable_per_region_metrics: bool,
/// Optional route provenance captured by the distributed planner from the
/// same physical route that selected `regions`. Ordinary execution never
/// reads it, so `None` keeps every existing plan unchanged.
native_raw_route: Option<CaptureNativeRawRoute>,
}
/// Borrowed handoff view over the ordinary plan state a native raw consumer
/// needs to read the captured route.
///
/// It borrows the existing [`MergeScanExec`] fields instead of duplicating them;
/// the output schema is available from the public `ExecutionPlan::schema()` of
/// the plan this source came from.
// TODO(narthex): remove this allowance once the remote native PromQL planner
// (#9670) consumes the captured route; this PR only captures it.
#[allow(dead_code)]
pub(crate) struct NativeRawSource<'a> {
/// Frozen route provenance in ordinary region order.
pub(crate) route: &'a CaptureNativeRawRoute,
/// The exact ordinary input logical plan sent to every region.
pub(crate) plan: &'a LogicalPlan,
/// The table the ordinary plan scans.
pub(crate) table: &'a TableName,
/// The handler that performs the region reads.
pub(crate) region_query_handler: &'a RegionQueryHandlerRef,
/// The query context of the ordinary execution.
pub(crate) query_ctx: &'a QueryContextRef,
}
impl std::fmt::Debug for MergeScanExec {
@@ -604,43 +577,6 @@ impl MergeScanExec {
captured_remote_dyn_filters: Arc::default(),
target_partition,
enable_per_region_metrics,
native_raw_route: None,
})
}
/// Attaches the native raw route captured from the same physical route that
/// selected this plan's regions.
///
/// The artifact is provenance for a native raw consumer; ordinary execution
/// never reads it.
pub(crate) fn with_native_raw_route(mut self, native_raw_route: CaptureNativeRawRoute) -> Self {
self.native_raw_route = Some(native_raw_route);
self
}
/// Returns the frozen region targets captured for this plan together with
/// the ordinary execution state a native raw consumer needs to read them.
///
/// Returns `None` unless the planner proved the route and ordinary region
/// dispatch stays unchanged: a scan that already captured remote dynamic
/// filters, or a flow sink scan, keeps its ordinary path, so a consumer must
/// fall back as well. An assigned producer id alone is not active intent,
/// because the analyzer assigns one to every merge scan.
// TODO(narthex): remove this allowance once the remote native PromQL planner
// (#9670) consumes the captured route; this PR only captures it.
#[allow(dead_code)]
pub(crate) fn native_raw_source(&self) -> Option<NativeRawSource<'_>> {
if !self.captured_remote_dyn_filters.lock().unwrap().is_empty() || self.is_flow_sink_scan()
{
return None;
}
let route = self.native_raw_route.as_ref()?;
Some(NativeRawSource {
route,
plan: &self.plan,
table: &self.table,
region_query_handler: &self.region_query_handler,
query_ctx: &self.query_ctx,
})
}
@@ -1431,7 +1367,7 @@ mod tests {
use datatypes::vectors::{Int64Vector, StringVector, TimestampMillisecondVector};
use futures_util::{Stream, TryStreamExt};
use session::ReadPreference;
use session::context::{QueryContext, QueryContextBuilder};
use session::context::QueryContext;
use session::query_id::QueryId;
use snafu::IntoError;
use table::table::scan::REGION_SCAN_EXEC_NAME;
@@ -1441,7 +1377,7 @@ mod tests {
use super::*;
use crate::dist_plan::DynFilterRegistryManager;
use crate::options::{FLOW_SINK_TABLE_ID, QueryOptions};
use crate::options::QueryOptions;
use crate::query_engine::{QueryEngineContext, QueryEngineState};
use crate::region_query::RegionQueryHandler;
@@ -1610,217 +1546,6 @@ mod tests {
);
}
#[test]
fn merge_scan_ordinary_construction_has_no_native_raw_route() {
// The artifact is provenance only: an ordinary plan carries none and
// executes exactly as before.
let exec = merge_scan_exec_with_sorted_input(1, 1);
assert!(exec.native_raw_source().is_none());
}
#[test]
fn merge_scan_keeps_captured_native_raw_route() {
let region_id = RegionId::new(1024, 1);
let exec =
merge_scan_exec_with_sorted_input(1, 1).with_native_raw_route(CaptureNativeRawRoute {
targets: vec![(region_id, test_target(1))],
});
let source = exec.native_raw_source().unwrap();
assert_eq!(vec![(region_id, test_target(1))], source.route.targets);
// The bundle borrows the ordinary plan state, not a copy of it.
assert_eq!(1, source.plan.schema().fields().len());
assert_eq!("table", source.table.table_name);
assert_eq!(1, exec.partition_count());
}
#[test]
fn merge_scan_declines_native_raw_route_for_captured_dyn_filters() {
// An assigned remote dynamic-filter producer id is identity, not intent:
// the analyzer assigns one to every merge scan, so a plan that captured no
// filters keeps the handoff view.
let region_id = RegionId::new(1024, 1);
let exec = remote_dyn_filter_test_exec(
Arc::new(TestRegionQueryHandler::default()),
QueryContext::arc(),
)
.with_native_raw_route(CaptureNativeRawRoute {
targets: vec![(region_id, test_target(1))],
});
assert!(exec.captured_remote_dyn_filters().is_empty());
assert!(exec.native_raw_source().is_some());
// A filter this producer accepted is really captured, and only then does the
// plan keep its ordinary dispatch.
install_remote_dyn_filter(&exec);
assert_eq!(1, exec.captured_remote_dyn_filters().len());
assert!(exec.native_raw_source().is_none());
// A rewrite of the same scan keeps the frozen route and the captured filters,
// so the rewritten plan still declines the handoff view.
let plan: Arc<dyn ExecutionPlan> = Arc::new(exec);
let replaced = plan
.replace_children(
vec![],
datafusion::physical_plan::ReplaceChildrenOptions::new(
datafusion::physical_plan::ChildrenPropertiesMode::Recompute,
),
)
.unwrap();
let exec = replaced.downcast_ref::<MergeScanExec>().unwrap();
assert_eq!(
vec![(region_id, test_target(1))],
exec.native_raw_route.as_ref().unwrap().targets
);
assert_eq!(1, exec.captured_remote_dyn_filters().len());
assert!(exec.native_raw_source().is_none());
}
#[tokio::test]
async fn merge_scan_execution_ignores_native_raw_route() {
// Attaching the artifact must not change what an ordinary execution
// reads: every region is still dispatched through the handler's own
// target selection, with the ordinary read preference, and the same
// batches come back in the same order.
let regions = vec![
RegionId::new(1024, 2),
RegionId::new(1024, 1),
RegionId::new(1024, 3),
];
let batch_for = |region_number: u32| {
let base = i64::from(region_number) * 10;
record_batch(
int64_schema(&["a", "b"]),
vec![
Arc::new(Int64Vector::from_slice([base + 1])) as _,
Arc::new(Int64Vector::from_slice([base + 2])) as _,
],
)
};
let responses = || -> Vec<(RegionId, RecordBatch)> {
regions
.iter()
.map(|region_id| (*region_id, batch_for(region_id.region_number())))
.collect()
};
// Region striping keeps the selection order, so this is the whole output.
let expected = [(21, 22), (11, 12), (31, 32)];
let baseline_handler = Arc::new(TestRegionQueryHandler::new(responses()));
let baseline = collect_merge_scan(merge_scan_exec_with_handler(
regions.clone(),
expected_int64_schema(),
baseline_handler.clone(),
1,
))
.await
.unwrap();
assert_eq!(expected.len(), baseline.len());
for (batch, values) in baseline.iter().zip(expected) {
assert_int64_batch(batch, values);
}
// Deliberately different peers, and in a deliberately different order:
// an execution that read the artifact would dispatch to these.
let handler = Arc::new(TestRegionQueryHandler::new(responses()));
let with_route = collect_merge_scan(
merge_scan_exec_with_handler(
regions.clone(),
expected_int64_schema(),
handler.clone(),
1,
)
.with_native_raw_route(CaptureNativeRawRoute {
targets: vec![
(regions[2], test_target(7)),
(regions[0], test_target(8)),
(regions[1], test_target(9)),
],
}),
)
.await
.unwrap();
assert_eq!(baseline, with_route);
assert_eq!(expected.len(), with_route.len());
for (batch, values) in with_route.iter().zip(expected) {
assert_int64_batch(batch, values);
}
for exec_handler in [&baseline_handler, &handler] {
assert_eq!(regions.len(), exec_handler.select_target_calls());
assert_eq!(
vec![ReadPreference::Leader; regions.len()],
exec_handler.read_preferences()
);
// The ordinary plan dispatches to the target its handler selected,
// never to a peer from the artifact.
assert_eq!(vec![1; regions.len()], exec_handler.do_get_targets());
}
}
#[test]
fn merge_scan_with_new_children_preserves_native_raw_route() {
let region_id = RegionId::new(1024, 1);
let plan: Arc<dyn ExecutionPlan> = Arc::new(
merge_scan_exec_with_sorted_input(1, 1).with_native_raw_route(CaptureNativeRawRoute {
targets: vec![(region_id, test_target(3))],
}),
);
let replaced = plan
.replace_children(
vec![],
datafusion::physical_plan::ReplaceChildrenOptions::new(
datafusion::physical_plan::ChildrenPropertiesMode::Recompute,
),
)
.unwrap();
let exec = replaced.downcast_ref::<MergeScanExec>().unwrap();
assert_eq!(
vec![(region_id, test_target(3))],
exec.native_raw_source().unwrap().route.targets
);
}
#[test]
fn merge_scan_declines_native_raw_route_for_flow_sink_scan() {
// A flow sink read keeps its ordinary dispatch, so the handoff view is
// not offered even though the artifact exists for these regions.
let region_id = RegionId::new(1024, 1);
let query_ctx = Arc::new(
QueryContextBuilder::default()
.set_extension(
FLOW_SINK_TABLE_ID.to_string(),
region_id.table_id().to_string(),
)
.build(),
);
let plan = LogicalPlanBuilder::empty(true).build().unwrap();
let schema = plan.schema().as_arrow().clone();
let exec = MergeScanExec::new(
&SessionStateBuilder::new().build(),
TableName::new("catalog", "schema", "table"),
vec![region_id],
plan,
&schema,
Arc::new(TestRegionQueryHandler::default()),
query_ctx,
1,
None,
false,
)
.unwrap()
.with_native_raw_route(CaptureNativeRawRoute {
targets: vec![(region_id, test_target(1))],
});
assert!(exec.is_flow_sink_scan());
assert!(exec.native_raw_source().is_none());
}
#[test]
fn merge_scan_advertises_ordering_when_each_partition_reads_at_most_one_region() {
let exec = merge_scan_exec_with_sorted_input(3, 3);
@@ -2398,18 +2123,6 @@ mod tests {
..Default::default()
}
}
fn select_target_calls(&self) -> usize {
self.select_target_calls.load(Ordering::SeqCst)
}
fn read_preferences(&self) -> Vec<ReadPreference> {
self.read_preferences.lock().unwrap().clone()
}
fn do_get_targets(&self) -> Vec<u64> {
self.do_get_targets.lock().unwrap().clone()
}
}
struct TestRecordBatchStream {
src/query/src/dist_plan/planner.rs
+18 -1027
View File
File diff suppressed because it is too large Load Diff
src/query/src/options.rs
+4 -2
View File
@@ -110,8 +110,10 @@ pub struct QueryOptions {
/// - `fair`: Fair divides memory available after unspillable reservations
/// evenly among spillable reservations and may trigger earlier spills.
pub experimental_memory_pool_policy: QueryMemoryPoolPolicy,
/// Enables the experimental native PromQL execution path for compatible
/// instant queries over shared Metric raw series. Default: `false`.
/// Enables the experimental native PromQL execution path for compatible queries. This
/// path currently applies only to standalone queries using in-process region transport,
/// and also requires the per-request `query.experimental_native_promql=true` extension.
/// Default: `false`.
pub experimental_native_promql: bool,
}
src/query/src/promql/native_compile.rs
+5 -7
View File
@@ -15,8 +15,8 @@
//! Compiles a PromQL AST into a [`NativeTree`].
//!
//! This is the front half of the native evaluation model: it reads the parsed statement and writes
//! down what it means as bound nodes. It runs nothing, and nothing calls it yet, so the subset it
//! accepts is one recursive rule:
//! down what it means as bound nodes. The native planner uses this compiler to build a tree for
//! standalone execution. Its accepted expressions are described by the recursive rule below:
//!
//! ```text
//! <expression> one of, recursively:
@@ -718,11 +718,9 @@ fn binary_op(op: &TokenType) -> Option<BinaryOp> {
/// Compiles an aggregation into an aggregate node over the expression it groups.
///
/// Every aggregator that folds a group's values into one value per grid point binds, whether or not
/// its kernel exists yet: `sum`, `avg`, `count`, `min`, `max`, `group`, `stddev`, `stdvar` and
/// `quantile`. Binding a shape the compiler can describe and letting the runtime refuse a kind it
/// cannot execute keeps this the one place that says what an aggregate means, so a kernel landing
/// later changes no binding.
/// The aggregators that fold a group's values into one value per grid point bind to an aggregate
/// node: `sum`, `avg`, `count`, `min`, `max`, `group`, `stddev`, `stdvar` and `quantile`. The
/// runtime validates whether the resulting tree can execute.
///
/// The rest of the parser's aggregates are a different shape and bind as a vector operation
/// ([`vector_aggregate`]) rather than as a kind that would compute something else: `topk`,
src/query/src/promql/native_compile/tests.rs
+39
View File
@@ -4645,6 +4645,45 @@ fn timestamp_multi_field_input_is_not_supported_natively() {
.unwrap_or_else(|error| panic!("`{query}` over one field was refused: {error}"));
}
}
#[test]
fn timestamp_anchored_selector_retains_offset_and_ordinary_selection_keeps_it() {
for (query, expected_offset, expected_anchor, expected_bounds) in [
(
"timestamp(x @ 50 offset 60s)",
60_000,
Some(50_000),
(-15_000, -10_000),
),
("timestamp(x offset 2s)", 2_000, None, (-7_000, 98_000)),
(
"x @ 50 offset 60s",
60_000,
Some(50_000),
(-15_000, -10_000),
),
] {
let compiled = compile_selectors(
&eval_stmt(query),
&|_| Ok(raw_schema()),
SeriesOrdering::KeyThenTime,
)
.unwrap_or_else(|error| panic!("{query}: {error}"));
let requirement = &compiled.inputs[0].requirement;
assert_eq!(requirement.offset_ms, expected_offset, "{query}");
assert_eq!(requirement.at_ms, expected_anchor, "{query}");
let (start, end, _) = requirement.raw_scan_bounds().unwrap();
assert_eq!((start, end), expected_bounds, "{query}");
if query.starts_with("timestamp(") && expected_anchor.is_some() {
let NativeExpr::Instant { op, .. } = &compiled.tree.root().expr else {
panic!("{query} did not compile as instant timestamp");
};
assert_eq!(op.offset_ms, expected_offset, "{query}");
assert_eq!(op.at_ms, expected_anchor, "{query}");
}
}
}
#[test]
fn selector_requirements_capture_bounds_and_source_occurrence() {
let compiled = compile_selectors(
src/query/src/promql/planner/native.rs
+36 -4
View File
@@ -969,11 +969,13 @@ fn collect_selectors<'a>(expr: &'a PromExpr, out: &mut Vec<&'a VectorSelector>)
match expr {
PromExpr::VectorSelector(selector) => out.push(selector),
PromExpr::MatrixSelector(MatrixSelector { vs, .. }) => out.push(vs),
PromExpr::Paren(ParenExpr { expr }) => return collect_selectors(expr, out),
PromExpr::Paren(ParenExpr { expr }) => {
return collect_selectors(expr, out);
}
PromExpr::Call(call)
if matches!(
call.func.name,
"rate" | "increase" | "sum_over_time" | "abs"
"rate" | "increase" | "sum_over_time" | "abs" | "timestamp"
) && call.args.args.len() == 1 =>
{
return collect_selectors(&call.args.args[0], out);
@@ -983,7 +985,9 @@ fn collect_selectors<'a>(expr: &'a PromExpr, out: &mut Vec<&'a VectorSelector>)
}
// Literals are permitted only as compiler-recognized broadcasts in a bounded binary tree.
PromExpr::NumberLiteral(_) => return true,
PromExpr::Unary(unary) => return collect_selectors(&unary.expr, out),
PromExpr::Unary(unary) => {
return collect_selectors(&unary.expr, out);
}
PromExpr::Call(call) => {
if matches!(call.func.name, "label_replace" | "label_join") {
return false;
@@ -994,7 +998,9 @@ fn collect_selectors<'a>(expr: &'a PromExpr, out: &mut Vec<&'a VectorSelector>)
}
}
}
PromExpr::Aggregate(aggregate) => return collect_selectors(&aggregate.expr, out),
PromExpr::Aggregate(aggregate) => {
return collect_selectors(&aggregate.expr, out);
}
PromExpr::Subquery(_) => return false,
_ => {}
}
@@ -1133,6 +1139,32 @@ mod tests {
.collect()
}
#[test]
fn selector_walk_keeps_plain_selector_identity() {
for (query, expected) in [
("timestamp((x @ 2 offset -1s))", vec!["x"]),
("abs(timestamp(x @ 2 offset -1s))", vec!["x"]),
("timestamp(x @ start() offset 1s)", vec!["x"]),
("timestamp(x @ end() offset -1s)", vec!["x"]),
(
"timestamp(x @ 2 offset 1s) + y @ 2 offset 1s",
vec!["x", "y"],
),
] {
let expr = promql_parser::parser::parse(query).unwrap();
let mut selectors = Vec::new();
assert!(collect_selectors(&expr, &mut selectors));
assert_eq!(
selectors
.iter()
.map(|s| s.name.as_deref().unwrap())
.collect::<Vec<_>>(),
expected,
"{query}"
);
}
}
#[test]
fn collect_selectors_descends_into_aggregate_inputs() {
for query in [
src/query/src/query_engine/state.rs
-3
View File
@@ -313,7 +313,6 @@ impl QueryEngineState {
partition_rule_manager,
region_query_handler.clone(),
options.enable_per_region_metrics,
options.experimental_native_promql,
)))
.with_optimizer_rules(optimizer.rules)
.with_physical_optimizer_rules(physical_optimizer.rules)
@@ -626,7 +625,6 @@ impl DfQueryPlanner {
partition_rule_manager: Option<PartitionRuleManagerRef>,
region_query_handler: Option<RegionQueryHandlerRef>,
enable_per_region_metrics: bool,
capture_native_raw_route: bool,
) -> Self {
let mut planners: Vec<Arc<dyn ExtensionPlanner + Send + Sync>> = vec![
Arc::new(PromExtensionPlanner),
@@ -641,7 +639,6 @@ impl DfQueryPlanner {
partition_rule_manager,
region_query_handler,
enable_per_region_metrics,
capture_native_raw_route,
)));
planners.push(Arc::new(MergeSortExtensionPlanner {}));
}
src/query/tests/native_instant_at.rs
+129
View File
@@ -179,6 +179,33 @@ fn native_instant_at_timestamp_values_and_labels() {
run("timestamp(abs(a @ 1))", 0, 4000, 3000, &values, &times, 1)[0].0,
[0.0, 1.0, 2.0, 3.0, 4.0]
);
// The raw source axis is deliberately unpruned: timestamp selection must enforce both
// windows itself. At @11s offset -2s with 5s lookback, 10.5s is eligible but 11.5s is not.
assert_eq!(
run(
"timestamp(a @ 11 offset -2s)",
11_000,
11_000,
5_000,
&[Some(10_500.0), Some(11_500.0)],
&[10_500, 11_500],
1,
)[0]
.0,
[10.5]
);
assert!(
run(
"timestamp(a @ 50 offset 60s)",
50_000,
50_000,
30_000,
&[Some(10_000.0)],
&[10_000],
1,
)
.is_empty()
);
assert_eq!(
run("a @ -1.5", 2000, 6000, 1000, &values, &times, 1)[0].0,
[7.0; 5]
@@ -204,6 +231,108 @@ fn native_instant_at_timestamp_values_and_labels() {
);
}
#[test]
fn native_anchored_timestamp_selection_boundaries() {
let stale = f64::from_bits(common_query::prometheus::PROMETHEUS_STALE_NAN_BITS);
let all_times = [-1500, 500, 1500, 3500];
let all_values = [Some(7.0), Some(1.0), Some(2.0), Some(3.0)];
type Case<'a> = (&'a str, u64, &'a [i64], &'a [Option<f64>], &'a [f64]);
let cases: &[Case<'_>] = &[
(
"timestamp(a @ 2 offset 1s)",
5_000,
&all_times,
&all_values,
&[0.5],
),
(
"timestamp(a @ 2 offset -1s)",
5_000,
&all_times,
&all_values,
&[1.5],
),
(
"timestamp(a @ 2 offset 2s)",
750,
&all_times,
&all_values,
&[],
),
("timestamp(a @ 0.5)", 0, &all_times, &all_values, &[0.5]),
("timestamp(a @ 1)", 0, &all_times, &all_values, &[]),
(
"timestamp(a @ 2 offset 1s)",
1_000,
&all_times,
&all_values,
&[],
),
// The lower endpoint is exclusive for either sign of offset.
(
"timestamp(a @ 2 offset 1s)",
1_500,
&[500],
&[Some(1.0)],
&[],
),
(
"timestamp(a @ 2 offset -1s)",
1_500,
&[1500],
&[Some(2.0)],
&[],
),
// The upper endpoint is inclusive.
(
"timestamp(a @ 2 offset 1s)",
5_000,
&[1000],
&[Some(4.0)],
&[1.0],
),
(
"timestamp(a @ 2 offset -1s)",
5_000,
&[2000, 3000],
&[Some(2.0), Some(3.0)],
&[2.0],
),
// A stale marker at the eligible upper endpoint suppresses the preceding value.
(
"timestamp(a @ 2 offset 1s)",
5_000,
&[500, 1000],
&[Some(1.0), Some(stale)],
&[],
),
// A stale marker above the intersection's upper endpoint is ignored.
(
"timestamp(a @ 2 offset 1s)",
5_000,
&[1000, 1500],
&[Some(4.0), Some(stale)],
&[1.0],
),
// A stale marker inside the eligible window prevents fallback to an older finite sample.
(
"timestamp(a @ 2 offset -1s)",
5_000,
&[500, 1500],
&[Some(1.0), Some(stale)],
&[],
),
];
for (query, lookback, times, values, expected) in cases {
let actual = run(query, 2_000, 2_000, *lookback, values, times, 1)
.into_iter()
.flat_map(|series| series.0)
.collect::<Vec<_>>();
assert_eq!(actual, *expected, "{query} with lookback {lookback}");
}
}
/// Two source series that differ only by a metadata label `timestamp()` drops must not be emitted
/// as two rows of one label set: the second is refused at the position it collides on, which is the
/// same rule a map follows when it drops the metric name.