mirror of
https://github.com/GreptimeTeam/greptimedb.git
synced 2026-11-14 06:46:56 +00:00
feat(query): run a PromQL statement through the fused native node (#9868)
Wires the native PromQL route into the production planner so a real PromQL statement reaches `FusedSeriesExec` and runs. Scope is a standalone, single-region point of contact: `rate(x[5m])` and `abs(rate(x[5m]))` execute through the node, and every shape the compiler does not accept keeps the plan the existing planner builds. Depends on the node `FusedSeries` / `FusedSeriesExec`, which landed in `narthex` as #9820. ## The route A statement is offered to `native_compile::compile` at its root inside `PromPlanner`, before the expression is split into the nodes a range function and the math functions over it would have shared. The route is offered only when **both** hold: the engine-wide `query.experimental_native_promql` option (retained in `QueryEngineState`) and a per-request query-context extension of the same name. Either alone keeps the existing plan, so a client cannot turn the route on against the server's configuration. A refused shape is the route declining the query, not an approximation: planning continues exactly as it would with the flag off. ## Getting the node to the region **Placement.** `check_extension_plan` answered `Unsupported` for the node, which the analyzer treats as a boundary. It now answers commutative when the input is exactly the table scan (the region builds the node from the scan's own entrypoint and cannot look through a filter or projection), the node binds a metric name, and no node of its tree folds across series. **Encodability.** `should_expand` asks whether a plan encodes to Substrait before it asks whether a node commutes. That is right for a region reached over the wire and wrong for standalone, whose region server is in the same process. The engine now knows whether its regions are reached in process, and the exception is scoped to a plan whose **root** is `FusedSeries`; every other plan keeps the ordinary whole-plan fast path and the per-node check. Both transports are covered by the analyzer tests. **Scan hint.** A region only prepares its series scan for a per-series request. The existing `ScanHintRule` now attaches `PerSeries` to the scan that is the immediate input of a fused node; the duplicate helper that used to do it from the planner is deleted. **Time bounds.** The native scan is the selector's own scan, so it gets the same bounds the legacy selector path computes: `build_time_index_filter` with the selector range, `offset` and `@` anchor, pushed into `TableScan.filters` so the plan stays a bare `TableScan`. Without this the scan read the table's whole retention. ## Cross-region correctness A region evaluates the whole tree over the series its own scan returns, so a series split across storage regions loses window contributions instead of failing — `rate(cpu[2m])` on a time-partitioned table returns nothing where the complete series answers `1.0`. `single_partition_scan()` only covers scan partitions inside one region. The route therefore declines unless **every partition column of the table is one of the selector's own tag columns**: no partition columns is a single region, and a tag-determined partition key cannot split one series. Anything else, including internal columns, falls back to the existing plan. No exchange and no region-count probing. ## Verification Real standalone, `experimental_native_promql = true`, physical plan from `datafusion::physical_planner=debug`: ``` FusedSeriesExec: rate(cpu[5m]) NativeTree root Window rate [300000ms] ... SeriesScanExec ... ``` - `rate(cpu[5m])` and `abs(rate(cpu[5m]))` return the same series, timestamps and values as the existing path on a dense series, a sparse one and a two-point series. **Multi-partition, with a flush in the middle.** `[query] parallelism = 4`, 8 series with a `shard` tag written in two rounds with `ADMIN FLUSH_TABLE` between them, so each series' samples span a memtable and an SST. The physical plan reports `SeriesScanExec: ... partitions=4` and `FusedSeriesExec`, and for both `rate(cpu[5m])` and `abs(rate(cpu[5m]))` all 40 emitted rows match the legacy plan position by position (timestamps, values, `host`/`shard` labels). A series split across partitions would have made `rate` lose the earlier sample, so the row-for-row match is the check that each series stayed in one scan partition. **Scan parallelism is regression-tested in the workspace, not just at the seam.** Moving the range-assignment rule into `table::table::scan` initially broke two callers that no single-crate test covered (`src/cmd/src/datanode/scanbench.rs` and `parallelize_scan`'s own test module); `cargo check --workspace --all-targets` is what caught them, and it now passes along with `query --lib` 897 passed. - `/v1/prometheus/api/v1/query` and `/query_range` responses are identical between the two routes. - `/v1/promql` now emits the same column order as the existing path (time index first). Sqlness, standalone and distributed: `distinct`, `latest_per_series`, `neumann` pass. Focused unit tests, 13 passed / 0 failed (`native_fused_*` and the request-opt-in parser after the admission fix; the analyzer and scan-hint checks below are separate invocations): `native_fused_*` (including the time-partitioned fallback and the tag-partitioned still-native case), the request-opt-in parser, both analyzer encodability checks under **both** transports, the ordinary-plan fast-path assertion, and the three `ScanHintRule` cases. Clippy `-p query --all-targets -D warnings` and fmt clean. ## Admission: what the route refuses The region's raw-series preflight refuses an append-mode table and any time index that is not millisecond, and the physical layer turns that refusal into a planning error rather than a fallback. The logical route therefore declines those tables itself, so such a query keeps the existing plan instead of failing after the logical plan is already fixed. Both cases have fallback tests asserting the native and flag-off plans are identical. ## What the route refuses, and why that is a fallback rather than an error Every refusal below happens while the logical plan is still this function's to choose, so the query keeps the plan the existing planner builds. That matters because once a `FusedSeries` exists the logical plan is fixed, and the physical layer can only report an error - it cannot undo the choice. - **A shape the compiler refuses**, a table it cannot bind, a remaining label matcher, `__field__`, a subquery (the selector walk does not enter one), or more than one selector. - **A tree that reads the whole vector at a point** (`Aggregate`, vector-vector `Binary`, and `Vector`: `sort`, `absent`, `topk`, `histogram_quantile`, ...). One executor runs per scan partition, and a partition holds only part of a metric's series. - **A table whose own settings the region cannot frame**: a non-millisecond time index, append mode, a partition column that is not a tag, a field that is not numeric (a decimal is not: the region asks `ConcreteDataType::is_numeric`, which excludes them), or an engine without a raw-series scanner. Mito's series scanner is the one that implements the raw path; a file-engine table is read by a `SinglePartitionScanner`, which cannot frame a series. A metric-engine logical table is refused too: it reads through a mito scanner, but the planner only reaches this node with a bare scan when it could not resolve the physical table, and that table's `TableInfo` carries no `append_mode` - so the region's own eligibility could not be checked here. That is a false negative that costs the query its native route, not an error. - **A plan that cannot reach a region**: regions reached over the wire (a distributed frontend) and a PromQL statement embedded in a hybrid CTE, where the node would sit under a SQL root and the in-process encodability exception - which is scoped to a plan whose root is the node - would not apply. Two shapes fall back for a simpler reason, and are recorded as limits rather than fixed here because supporting them would widen this POC to scalar broadcast: a negation or a binary operation mentioning a selector (`-m`, `rate(m[5m]) * 100`) is not found by the selector walk, and a statement whose bounds precede the Unix epoch cannot be represented as a `SystemTime`. Both answer correctly on the existing path. - **`TQL EXPLAIN` / `TQL ANALYZE` still print the existing plan.** Those statements wrap the expression in an explain extension before this route sees it, and the selector walk does not enter a wrapper, so the route declines. The query itself answers the same way; what differs is that the displayed plan does not show the node the same statement executes through `TQL EVAL`. Closing that needs the wrapper recognised and re-applied around a native plan, which is deliberately not part of this POC. ## Known limits, stated plainly - **The value column is named `val`, not the expression-derived name** the legacy projection carries (`prom_rate(ts_range,val,ts,Int64(300000))`). Prometheus API output is unaffected because it reads column roles, not names; the greptimedb-shaped `/v1/promql` response shows the difference. - **Scan parallelism is preserved.** The fused leaf needs a *prepared* scan, because `ParallelizeScan` runs after physical planning and cannot see a `SeriesScanExec`. Preparation now uses the same range-assignment rule the parallelizer uses — `table::table::scan::assign_partition_ranges`, moved there and shared by both callers — with the session's `target_partitions`, and it leaves an already-prepared scan untouched. The inherited single-partition restriction in `resolve_named_series_binding` is gone, so a multi-partition scan resolves its binding; the series producer still keeps one series inside one scan partition, so each partition feeds `FusedSeriesExec` with whole series. - **Each series must stay inside one region.** The guard declines the route unless every partition column is one of the selector's own tag columns, so a table partitioned by time, or by anything that is not a series label, falls back to the existing plan instead of answering wrongly. A tag-partitioned table is accepted and does use the native route, because a tag-determined partition key cannot split one series. The guard is a sufficient condition, not a proof for every partitioning scheme: it has not been exercised against a concurrent repartition lifecycle. - **Anything with a label matcher falls back**, because a matcher that survives preprocessing selects fewer rows than the scan below the node reads. - **Distributed is unsupported**: in a frontend the metric scan is rewritten to `MergeScanLogicalPlan`, the node is not placed there, and the flag is off by default. - **No aggregates, scalars or subquery execution.**
This commit is contained in:
Generated
+1
@@ -11378,6 +11378,7 @@ dependencies = [
|
||||
"regex",
|
||||
"snafu 0.8.6",
|
||||
"store-api",
|
||||
"table",
|
||||
"tokio",
|
||||
"twox-hash",
|
||||
]
|
||||
|
||||
@@ -47,7 +47,7 @@ use mito2::sst::file_ref::FileReferenceManager;
|
||||
use moka::future::CacheBuilder;
|
||||
use object_store::manager::ObjectStoreManager;
|
||||
use object_store::util::normalize_dir;
|
||||
use query::optimizer::parallelize_scan::ParallelizeScan;
|
||||
use query::optimizer::parallelize_scan::assign_partition_ranges;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use snafu::{OptionExt, ResultExt};
|
||||
use sqlparser::ast::ExprWithAlias as SqlExprWithAlias;
|
||||
@@ -995,8 +995,7 @@ impl ScanbenchCommand {
|
||||
let all_ranges: Vec<_> = original_partitions.into_iter().flatten().collect();
|
||||
|
||||
// Distribute ranges across partitions
|
||||
let mut partitions =
|
||||
ParallelizeScan::assign_partition_range(all_ranges, self.parallelism);
|
||||
let mut partitions = assign_partition_ranges(all_ranges, self.parallelism);
|
||||
|
||||
// Sort ranges within each partition by start time ascending
|
||||
for partition in &mut partitions {
|
||||
|
||||
@@ -734,6 +734,9 @@ impl StartCommand {
|
||||
procedure_executor.clone(),
|
||||
process_manager,
|
||||
)
|
||||
// The region servers this frontend reads are the ones built above, in this process, so a
|
||||
// region input can travel as a logical plan instead of Substrait bytes.
|
||||
.with_in_process_region_transport()
|
||||
.with_local_file_access(local_file_access);
|
||||
|
||||
plugins::setup_frontend_plugins_post_build(&mut plugins, &plugin_opts, &fe_instance)
|
||||
|
||||
@@ -78,6 +78,9 @@ pub struct FrontendBuilder {
|
||||
procedure_executor: ProcedureExecutorRef,
|
||||
process_manager: ProcessManagerRef,
|
||||
local_file_access: LocalFileAccess,
|
||||
/// Whether this frontend reaches the regions of its datanodes in process. Declared by whoever
|
||||
/// assembles the frontend, because only it knows what its node manager talks to.
|
||||
in_process_region_transport: bool,
|
||||
}
|
||||
|
||||
impl FrontendBuilder {
|
||||
@@ -105,6 +108,7 @@ impl FrontendBuilder {
|
||||
procedure_executor,
|
||||
process_manager,
|
||||
local_file_access: LocalFileAccess::Disabled,
|
||||
in_process_region_transport: false,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -157,6 +161,20 @@ impl FrontendBuilder {
|
||||
self
|
||||
}
|
||||
|
||||
/// Declares that every region this frontend queries is read in this very process.
|
||||
///
|
||||
/// The caller must own the node manager this frontend was built with, and know that its
|
||||
/// `DatanodeManager` hands region reads to a `RegionServer` in this process — as standalone
|
||||
/// startup does, since it passes the datanode's own region server to the node manager it
|
||||
/// installs. The distributed planner then hands a region a `LogicalPlan` instead of Substrait
|
||||
/// bytes, which `RegionServer::handle_read` accepts. A frontend whose node manager reaches
|
||||
/// remote datanodes must leave it off: those datanodes only accept Substrait bytes
|
||||
/// (`RegionServer::handle_remote_read`), so a plan that does not encode would fail there.
|
||||
pub fn with_in_process_region_transport(mut self) -> Self {
|
||||
self.in_process_region_transport = true;
|
||||
self
|
||||
}
|
||||
|
||||
pub fn options(&self) -> &FrontendOptions {
|
||||
&self.options
|
||||
}
|
||||
@@ -315,7 +333,7 @@ impl FrontendBuilder {
|
||||
|
||||
let mut query_options = self.options.query.clone();
|
||||
query_options.enable_per_region_metrics = self.options.logging.enable_per_region_metrics;
|
||||
let query_engine = QueryEngineFactory::try_new_with_plugins(
|
||||
let query_engine = QueryEngineFactory::try_new_with_plugins_and_region_transport(
|
||||
self.catalog_manager.clone(),
|
||||
Some(partition_manager.clone()),
|
||||
Some(region_query_handler.clone()),
|
||||
@@ -323,6 +341,7 @@ impl FrontendBuilder {
|
||||
Some(procedure_service_handler),
|
||||
Some(Arc::new(flow_service)),
|
||||
true,
|
||||
self.in_process_region_transport,
|
||||
plugins.clone(),
|
||||
query_options,
|
||||
)
|
||||
|
||||
@@ -256,10 +256,41 @@ impl SeriesScan {
|
||||
.is_some_and(|context| context.format == crate::config::SeriesIndexFormat::Columnar)
|
||||
}
|
||||
|
||||
/// Returns whether raw series come from the columnar prepared source.
|
||||
///
|
||||
/// Derived only from the pinned series index version and every configured partition range, so
|
||||
/// preflight, the ordering guarantee, the partition reader and scanner-wide preparation agree.
|
||||
/// `false` means the stored sources are read by the ordinary scan.
|
||||
fn uses_prepared_columnar(&self) -> Result<bool, BoxedError> {
|
||||
if !self.uses_columnar() {
|
||||
return Ok(false);
|
||||
}
|
||||
let ranges = self
|
||||
.properties
|
||||
.partitions
|
||||
.iter()
|
||||
.flatten()
|
||||
.copied()
|
||||
.collect::<Vec<_>>();
|
||||
prepared::validate_inputs(&self.stream_ctx, &ranges).map_err(BoxedError::new)
|
||||
}
|
||||
|
||||
/// Returns the raw ordering guarantee for a scan whose columnar readiness is already known.
|
||||
fn raw_ordering(&self, prepared_columnar: bool) -> SeriesOrdering {
|
||||
if prepared_columnar {
|
||||
return SeriesOrdering::KeyThenTime;
|
||||
}
|
||||
match self.mode {
|
||||
SeriesScanMode::Legacy => SeriesOrdering::KeyThenTime,
|
||||
SeriesScanMode::TwoPhase => SeriesOrdering::Contiguous,
|
||||
}
|
||||
}
|
||||
|
||||
/// Prepares the raw scan and reports whether its data comes from the columnar prepared source.
|
||||
fn prepare_raw_scan(
|
||||
&self,
|
||||
label_source: &RawLabelSource,
|
||||
) -> Result<Option<RawScanPreparation>, BoxedError> {
|
||||
) -> Result<Option<(RawScanPreparation, bool)>, BoxedError> {
|
||||
let input = &self.stream_ctx.input;
|
||||
let metadata = input.mapper.metadata();
|
||||
match label_source {
|
||||
@@ -386,25 +417,19 @@ impl SeriesScan {
|
||||
);
|
||||
return Ok(None);
|
||||
}
|
||||
if self.uses_columnar() {
|
||||
let ranges = self
|
||||
.properties
|
||||
.partitions
|
||||
.iter()
|
||||
.flatten()
|
||||
.copied()
|
||||
.collect::<Vec<_>>();
|
||||
prepared::validate_inputs(&self.stream_ctx, &ranges).map_err(BoxedError::new)?;
|
||||
}
|
||||
// Structural eligibility errors still reject the raw path; ordinary scans cover sources
|
||||
// without columnar artifacts, so missing coverage is not an error here.
|
||||
let prepared_columnar = self.uses_prepared_columnar()?;
|
||||
let schema = Arc::new(SeriesSchema {
|
||||
fields: specs.into(),
|
||||
eval_grid: None,
|
||||
});
|
||||
let typed_schema = input.mapper.input_arrow_schema(false);
|
||||
let mut predicates = Vec::new();
|
||||
let predicate = if self.uses_columnar() {
|
||||
let predicate = if prepared_columnar {
|
||||
input.predicate_group().predicate()
|
||||
} else {
|
||||
// The ordinary scan applies region predicates before this residual filtering.
|
||||
input.predicate_group().predicate_without_region()
|
||||
};
|
||||
if let Some(predicate) = predicate {
|
||||
@@ -452,11 +477,14 @@ impl SeriesScan {
|
||||
predicates.push(physical);
|
||||
}
|
||||
}
|
||||
Ok(Some(RawScanPreparation {
|
||||
schema,
|
||||
fields,
|
||||
predicates,
|
||||
}))
|
||||
Ok(Some((
|
||||
RawScanPreparation {
|
||||
schema,
|
||||
fields,
|
||||
predicates,
|
||||
},
|
||||
prepared_columnar,
|
||||
)))
|
||||
}
|
||||
|
||||
fn supports_two_phase(input: &ScanInput) -> bool {
|
||||
@@ -844,12 +872,11 @@ impl RegionScanner for SeriesScan {
|
||||
}
|
||||
|
||||
fn raw_series_ordering(&self) -> SeriesOrdering {
|
||||
if self.uses_columnar() {
|
||||
return SeriesOrdering::KeyThenTime;
|
||||
}
|
||||
match self.mode {
|
||||
SeriesScanMode::Legacy => SeriesOrdering::KeyThenTime,
|
||||
SeriesScanMode::TwoPhase => SeriesOrdering::Contiguous,
|
||||
// Readiness is fallible here, but every fallible entry point rejects the same structural
|
||||
// errors before execution, so the `false` fallback below is never observed in a query.
|
||||
match self.uses_prepared_columnar() {
|
||||
Ok(prepared_columnar) => self.raw_ordering(prepared_columnar),
|
||||
Err(_) => self.raw_ordering(false),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -859,7 +886,7 @@ impl RegionScanner for SeriesScan {
|
||||
) -> Result<Option<Arc<SeriesSchema>>, BoxedError> {
|
||||
Ok(self
|
||||
.prepare_raw_scan(label_source)?
|
||||
.map(|preparation| preparation.schema))
|
||||
.map(|(preparation, _)| preparation.schema))
|
||||
}
|
||||
|
||||
fn scan_series_partition(
|
||||
@@ -869,10 +896,10 @@ impl RegionScanner for SeriesScan {
|
||||
partition: usize,
|
||||
label_source: &RawLabelSource,
|
||||
) -> Result<Option<SendableRawSeriesStream>, BoxedError> {
|
||||
let Some(preparation) = self.prepare_raw_scan(label_source)? else {
|
||||
let Some((preparation, prepared_columnar)) = self.prepare_raw_scan(label_source)? else {
|
||||
return Ok(None);
|
||||
};
|
||||
if self.uses_columnar() {
|
||||
if prepared_columnar {
|
||||
let partitions = self.properties.num_partitions().max(1);
|
||||
if partition >= partitions {
|
||||
return Err(BoxedError::new(
|
||||
@@ -897,10 +924,12 @@ impl RegionScanner for SeriesScan {
|
||||
return Ok(Some(raw::from_stream(
|
||||
schema,
|
||||
labels,
|
||||
SeriesOrdering::KeyThenTime,
|
||||
self.raw_ordering(prepared_columnar),
|
||||
Box::pin(stream),
|
||||
)));
|
||||
}
|
||||
// Stored sources without columnar artifacts (memtables, unindexed SSTs) are read by the
|
||||
// ordinary scan, so the raw path stays available.
|
||||
let metrics = new_partition_metrics(
|
||||
&self.stream_ctx,
|
||||
ctx.explain_verbose,
|
||||
@@ -911,7 +940,7 @@ impl RegionScanner for SeriesScan {
|
||||
let batches = self
|
||||
.scan_batch_in_partition(ctx, partition, metrics.clone(), metrics_set)
|
||||
.map_err(BoxedError::new)?;
|
||||
let ordering = self.raw_series_ordering();
|
||||
let ordering = self.raw_ordering(prepared_columnar);
|
||||
Ok(Some(raw::make_stream(
|
||||
batches,
|
||||
preparation.schema,
|
||||
@@ -957,9 +986,13 @@ impl RegionScanner for SeriesScan {
|
||||
&self.metrics_list,
|
||||
);
|
||||
Box::pin(async move {
|
||||
let Some(preparation) = preparation? else {
|
||||
let Some((preparation, prepared_columnar)) = preparation? else {
|
||||
return Ok(None);
|
||||
};
|
||||
if !prepared_columnar {
|
||||
// Known unsupported coverage: the caller keeps the ordinary scan instead.
|
||||
return Ok(None);
|
||||
}
|
||||
state
|
||||
.prepare(stream_ctx, ranges, labels, preparation, metrics)
|
||||
.await
|
||||
|
||||
@@ -137,20 +137,30 @@ fn same_labels(left: &RawLabelSource, right: &RawLabelSource) -> bool {
|
||||
}
|
||||
|
||||
/// Validates the pinned assignments and coverage without opening artifacts.
|
||||
pub(crate) fn validate_inputs(context: &StreamContext, ranges: &[PartitionRange]) -> Result<()> {
|
||||
collect_inputs(context, ranges).map(|_| ())
|
||||
///
|
||||
/// Returns whether every assigned source can be read from columnar artifacts. `false` is not an
|
||||
/// error: the caller scans those sources with the ordinary reader instead.
|
||||
pub(crate) fn validate_inputs(context: &StreamContext, ranges: &[PartitionRange]) -> Result<bool> {
|
||||
collect_inputs(context, ranges).map(|inputs| inputs.is_some())
|
||||
}
|
||||
|
||||
/// Collects the assigned SSTs and their covering columnar series indexes.
|
||||
///
|
||||
/// `None` means the assignments are structurally valid but not ready for columnar reads: an
|
||||
/// assigned memtable, an SST without columnar series coverage, or an SST without columnar
|
||||
/// samples. Structural corruption is still an error, and every assignment is validated before
|
||||
/// either answer is returned, so a later invalid assignment is never masked by missing artifacts.
|
||||
fn collect_inputs(
|
||||
context: &StreamContext,
|
||||
ranges: &[PartitionRange],
|
||||
) -> Result<(Vec<usize>, HashSet<FileId>)> {
|
||||
) -> Result<Option<(Vec<usize>, HashSet<FileId>)>> {
|
||||
let input = &context.input;
|
||||
let metadata = input.region_metadata();
|
||||
let index = input
|
||||
.series_index
|
||||
.as_ref()
|
||||
.ok_or_else(|| invalid("missing columnar index context"))?;
|
||||
let mut ready = true;
|
||||
let mut files = BTreeMap::<usize, HashSet<i64>>::new();
|
||||
for range in ranges {
|
||||
let range = context
|
||||
@@ -159,7 +169,9 @@ fn collect_inputs(
|
||||
.ok_or_else(|| invalid("invalid prepared scan range"))?;
|
||||
for row_group in &range.row_group_indices {
|
||||
if context.is_mem_range_index(*row_group) {
|
||||
return Err(invalid("columnar prepared reads require flushed memtables"));
|
||||
// Artifacts exist for flushed SSTs only.
|
||||
ready = false;
|
||||
continue;
|
||||
}
|
||||
if !context.is_file_range_index(*row_group) {
|
||||
return Err(invalid(
|
||||
@@ -203,24 +215,18 @@ fn collect_inputs(
|
||||
}
|
||||
}
|
||||
if !covered {
|
||||
return Err(invalid(format!(
|
||||
"SST {} has no columnar series coverage; run build_series_index()",
|
||||
file.file_id()
|
||||
)));
|
||||
ready = false;
|
||||
}
|
||||
index
|
||||
if index
|
||||
.version
|
||||
.samples
|
||||
.get(&file.file_id().file_id())
|
||||
.filter(|handle| handle.entry().artifact.format == ArtifactFormat::ColumnarSamples)
|
||||
.ok_or_else(|| {
|
||||
invalid(format!(
|
||||
"SST {} has no columnar samples; run build_series_index()",
|
||||
file.file_id()
|
||||
))
|
||||
})?;
|
||||
.is_none_or(|handle| handle.entry().artifact.format != ArtifactFormat::ColumnarSamples)
|
||||
{
|
||||
ready = false;
|
||||
}
|
||||
}
|
||||
Ok((file_indices, indexes))
|
||||
Ok(ready.then_some((file_indices, indexes)))
|
||||
}
|
||||
|
||||
struct Candidate {
|
||||
@@ -254,7 +260,12 @@ impl ColumnarPreparedSource {
|
||||
.series_index
|
||||
.as_ref()
|
||||
.ok_or_else(|| invalid("missing columnar index context"))?;
|
||||
let (files, indexes) = collect_inputs(&context, &ranges)?;
|
||||
let Some((files, indexes)) = collect_inputs(&context, &ranges)? else {
|
||||
// Readiness was checked before preparation, so this is a real capability change.
|
||||
return Err(invalid(
|
||||
"columnar series coverage changed after series preparation",
|
||||
));
|
||||
};
|
||||
let mut samples = Vec::<Arc<dyn SamplesFormatReader>>::new();
|
||||
for file_idx in files {
|
||||
let file = &input.files[file_idx];
|
||||
|
||||
@@ -25,13 +25,14 @@ use datafusion::execution::TaskContext;
|
||||
use datafusion::physical_plan::{ExecutionPlan, metrics::ExecutionPlanMetricsSet};
|
||||
use datafusion_expr::{col, lit};
|
||||
use datatypes::prelude::ConcreteDataType;
|
||||
use datatypes::series::SeriesOrdering;
|
||||
use datatypes::series::raw::RawSeriesChunk;
|
||||
use datatypes::series::samples::SampleRef;
|
||||
use object_store::ObjectStore;
|
||||
use object_store::layers::mock::{self, MockLayerBuilder, oio};
|
||||
use store_api::codec::PrimaryKeyEncoding;
|
||||
use store_api::region_engine::{
|
||||
BoundTag, PrepareRequest, QueryScanContext, RegionEngine, RegionScanner,
|
||||
BoundTag, PartitionRange, PrepareRequest, QueryScanContext, RegionEngine, RegionScanner,
|
||||
};
|
||||
use store_api::region_request::{
|
||||
RegionBuildIndexRequest, RegionDeleteRequest, RegionPutRequest, RegionRequest,
|
||||
@@ -51,6 +52,10 @@ use crate::test_util::sst_util::{new_sparse_primary_key, sst_region_metadata_wit
|
||||
use crate::test_util::{TestEnv, flush_region};
|
||||
|
||||
async fn fixture() -> (TestEnv, MitoEngine, MitoRegionRef) {
|
||||
fixture_with(false).await
|
||||
}
|
||||
|
||||
async fn fixture_with(experimental_series_scan_v2: bool) -> (TestEnv, MitoEngine, MitoRegionRef) {
|
||||
let mut env = TestEnv::with_prefix("prepared-columnar").await;
|
||||
let mut metadata = sst_region_metadata_with_encoding(PrimaryKeyEncoding::Sparse);
|
||||
metadata
|
||||
@@ -66,6 +71,7 @@ async fn fixture() -> (TestEnv, MitoEngine, MitoRegionRef) {
|
||||
MitoConfig {
|
||||
experimental_enable_series_index: true,
|
||||
experimental_series_index_format: SelectedFormat::Columnar,
|
||||
experimental_series_scan_v2,
|
||||
experimental_series_index_maintenance_interval: Duration::from_secs(3600),
|
||||
min_compaction_interval: Duration::from_secs(3600),
|
||||
..Default::default()
|
||||
@@ -192,6 +198,41 @@ impl mock::Read for ReadRecorder {
|
||||
}
|
||||
}
|
||||
|
||||
/// Fails reads of the selected artifact paths so tests can prove that object-store failures
|
||||
/// surface as errors instead of being classified as not-ready fallbacks.
|
||||
struct FailingReader {
|
||||
inner: oio::Reader,
|
||||
fail: bool,
|
||||
}
|
||||
impl mock::Read for FailingReader {
|
||||
async fn read(
|
||||
&self,
|
||||
range: mock::BytesRange,
|
||||
) -> object_store::Result<(mock::RpRead, object_store::Buffer)> {
|
||||
if self.fail {
|
||||
return Err(injected_failure());
|
||||
}
|
||||
self.inner.read(range).await
|
||||
}
|
||||
|
||||
async fn open(
|
||||
&self,
|
||||
range: mock::BytesRange,
|
||||
) -> object_store::Result<(mock::RpRead, Box<dyn oio::ReadStreamDyn>)> {
|
||||
if self.fail {
|
||||
return Err(injected_failure());
|
||||
}
|
||||
self.inner.open(range).await
|
||||
}
|
||||
}
|
||||
|
||||
fn injected_failure() -> mock::Error {
|
||||
mock::Error::new(
|
||||
mock::ErrorKind::Unexpected,
|
||||
"injected series manifest read failure",
|
||||
)
|
||||
}
|
||||
|
||||
fn scanner(region: &MitoRegionRef, store: ObjectStore) -> SeriesScan {
|
||||
scanner_with_format(region, store, SelectedFormat::Columnar)
|
||||
}
|
||||
@@ -269,6 +310,35 @@ fn rows(chunks: &[RawSeriesChunk]) -> Vec<(Vec<u8>, i64, f64)> {
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn samples(chunks: &[RawSeriesChunk]) -> Vec<(i64, f64)> {
|
||||
rows(chunks)
|
||||
.into_iter()
|
||||
.map(|(_, time, value)| (time, value))
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Reads the raw path of every partition, checks each stream declares `ordering`, and returns the
|
||||
/// samples ordered by time so a check does not depend on how range grouping spread a memtable.
|
||||
async fn all_raw_samples(
|
||||
scanner: &dyn RegionScanner,
|
||||
ctx: &QueryScanContext,
|
||||
metrics: &ExecutionPlanMetricsSet,
|
||||
binding: &RawLabelSource,
|
||||
ordering: SeriesOrdering,
|
||||
) -> Vec<(i64, f64)> {
|
||||
let mut collected = Vec::new();
|
||||
for partition in 0..scanner.properties().num_partitions().max(1) {
|
||||
let stream = scanner
|
||||
.scan_series_partition(ctx, metrics, partition, binding)
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(ordering, stream.ordering());
|
||||
collected.extend(samples(&stream.try_collect::<Vec<_>>().await.unwrap()));
|
||||
}
|
||||
collected.sort_unstable_by_key(|(time, _)| *time);
|
||||
collected
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn prepared_sources_share_discovery_and_read_independently() {
|
||||
let (_env, engine, region) = fixture().await;
|
||||
@@ -386,89 +456,133 @@ async fn prepared_sources_share_discovery_and_read_independently() {
|
||||
engine.stop().await.unwrap();
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::legacy(false)]
|
||||
#[case::two_phase(true)]
|
||||
#[tokio::test]
|
||||
async fn native_scan_reports_incomplete_coverage_and_uses_prepared_source() {
|
||||
let (_env, engine, region) = fixture().await;
|
||||
async fn native_scan_falls_back_to_raw_sources_until_columnar_coverage(#[case] v2: bool) {
|
||||
let (_env, engine, region) = fixture_with(v2).await;
|
||||
let request = ScanRequest {
|
||||
distribution: Some(TimeSeriesDistribution::PerSeries),
|
||||
..Default::default()
|
||||
};
|
||||
// Three files are below the automatic index threshold.
|
||||
let ctx = QueryScanContext::default();
|
||||
let metrics = ExecutionPlanMetricsSet::new();
|
||||
// The three flushed SSTs only have columnar artifacts after `build`.
|
||||
let fallback_ordering = if v2 {
|
||||
SeriesOrdering::Contiguous
|
||||
} else {
|
||||
SeriesOrdering::KeyThenTime
|
||||
};
|
||||
let raw_scanner = engine
|
||||
.handle_query(region.region_id, request.clone())
|
||||
.await
|
||||
.unwrap();
|
||||
let binding = labels(raw_scanner.metadata());
|
||||
assert!(raw_scanner.raw_series_schema(&binding).is_err());
|
||||
assert!(
|
||||
raw_scanner
|
||||
.scan_series_partition(
|
||||
&QueryScanContext::default(),
|
||||
&ExecutionPlanMetricsSet::new(),
|
||||
0,
|
||||
&binding,
|
||||
)
|
||||
.is_err()
|
||||
assert_eq!(fallback_ordering, raw_scanner.raw_series_ordering());
|
||||
assert!(raw_scanner.raw_series_schema(&binding).unwrap().is_some());
|
||||
assert_eq!(
|
||||
vec![(1000, 1.0), (2000, 1.0)],
|
||||
all_raw_samples(
|
||||
raw_scanner.as_ref(),
|
||||
&ctx,
|
||||
&metrics,
|
||||
&binding,
|
||||
fallback_ordering
|
||||
)
|
||||
.await
|
||||
);
|
||||
// Known unready coverage selects the ordinary scan instead of a prepared source.
|
||||
assert!(
|
||||
raw_scanner
|
||||
.prepare_series(
|
||||
&QueryScanContext::default(),
|
||||
&ExecutionPlanMetricsSet::new(),
|
||||
&binding
|
||||
)
|
||||
.prepare_series(&ctx, &metrics, &binding)
|
||||
.await
|
||||
.is_err()
|
||||
.unwrap()
|
||||
.is_none()
|
||||
);
|
||||
|
||||
build(&engine, ®ion).await;
|
||||
let raw_scanner = engine
|
||||
let mut raw_scanner = engine
|
||||
.handle_query(region.region_id, request.clone())
|
||||
.await
|
||||
.unwrap();
|
||||
// Move every assigned range into a new last partition: readiness, the ordering guarantee and
|
||||
// the reader must all follow the new assignment, not the pre-prepare partition list.
|
||||
let assigned = raw_scanner
|
||||
.properties()
|
||||
.partitions
|
||||
.iter()
|
||||
.flatten()
|
||||
.copied()
|
||||
.collect::<Vec<_>>();
|
||||
raw_scanner
|
||||
.prepare(PrepareRequest::default().with_ranges(vec![vec![], assigned]))
|
||||
.unwrap();
|
||||
assert_eq!(2, raw_scanner.properties().num_partitions());
|
||||
let binding = labels(raw_scanner.metadata());
|
||||
assert_eq!(
|
||||
SeriesOrdering::KeyThenTime,
|
||||
raw_scanner.raw_series_ordering()
|
||||
);
|
||||
let source = raw_scanner
|
||||
.prepare_series(&ctx, &metrics, &binding)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(1, source.num_series());
|
||||
// The first partition owns no range and the second owns all of them.
|
||||
assert_eq!(
|
||||
vec![(1000, 1.0), (2000, 1.0)],
|
||||
all_raw_samples(
|
||||
raw_scanner.as_ref(),
|
||||
&ctx,
|
||||
&metrics,
|
||||
&binding,
|
||||
SeriesOrdering::KeyThenTime
|
||||
)
|
||||
.await
|
||||
);
|
||||
let plan = Arc::new(RegionScanExec::new(raw_scanner, request.clone(), None).unwrap())
|
||||
.try_to_series_scan(binding.clone())
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
let frame = plan.execute(0, Arc::new(TaskContext::default())).unwrap();
|
||||
let chunks = unpack_raw_series(frame, Some(plan.series_schema().clone()))
|
||||
.unwrap()
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(SeriesOrdering::KeyThenTime, plan.raw_series_ordering());
|
||||
let frame = plan.execute(1, Arc::new(TaskContext::default())).unwrap();
|
||||
let unpacked = unpack_raw_series(frame, Some(plan.series_schema().clone())).unwrap();
|
||||
// The packed frame declares the same guarantee the plan advertises.
|
||||
assert_eq!(plan.raw_series_ordering(), unpacked.ordering());
|
||||
assert_eq!(
|
||||
vec![1000, 2000],
|
||||
rows(&chunks).iter().map(|(_, t, _)| *t).collect::<Vec<_>>()
|
||||
vec![(1000, 1.0), (2000, 1.0)],
|
||||
samples(&unpacked.try_collect::<Vec<_>>().await.unwrap())
|
||||
);
|
||||
|
||||
// An assigned memtable is known unready, so the raw path keeps the ordinary scan and must
|
||||
// still return the new row.
|
||||
write(&engine, ®ion, 0, 3000, Some(4.0)).await;
|
||||
let raw_scanner = engine
|
||||
.handle_query(region.region_id, request)
|
||||
.await
|
||||
.unwrap();
|
||||
let binding = labels(raw_scanner.metadata());
|
||||
assert!(raw_scanner.raw_series_schema(&binding).is_err());
|
||||
assert_eq!(fallback_ordering, raw_scanner.raw_series_ordering());
|
||||
assert!(raw_scanner.raw_series_schema(&binding).unwrap().is_some());
|
||||
assert!(
|
||||
raw_scanner
|
||||
.scan_series_partition(
|
||||
&QueryScanContext::default(),
|
||||
&ExecutionPlanMetricsSet::new(),
|
||||
0,
|
||||
&binding,
|
||||
)
|
||||
.is_err()
|
||||
);
|
||||
assert!(
|
||||
raw_scanner
|
||||
.prepare_series(
|
||||
&QueryScanContext::default(),
|
||||
&ExecutionPlanMetricsSet::new(),
|
||||
&binding
|
||||
)
|
||||
.prepare_series(&ctx, &metrics, &binding)
|
||||
.await
|
||||
.err()
|
||||
.unwrap()
|
||||
.output_msg()
|
||||
.contains("memtables")
|
||||
.is_none()
|
||||
);
|
||||
assert_eq!(
|
||||
vec![(1000, 1.0), (2000, 1.0), (3000, 4.0)],
|
||||
all_raw_samples(
|
||||
raw_scanner.as_ref(),
|
||||
&ctx,
|
||||
&metrics,
|
||||
&binding,
|
||||
fallback_ordering
|
||||
)
|
||||
.await
|
||||
);
|
||||
engine.stop().await.unwrap();
|
||||
}
|
||||
@@ -1046,7 +1160,8 @@ async fn preflight_rejects_unsupported_assignments_without_artifact_io() {
|
||||
};
|
||||
let mut context = make_context();
|
||||
let ranges = context.partition_ranges();
|
||||
validate_inputs(&context, &ranges).unwrap();
|
||||
// Complete columnar coverage is the only ready answer.
|
||||
assert!(validate_inputs(&context, &ranges).unwrap());
|
||||
context.ranges[ranges[0].identifier].row_group_indices = vec![RowGroupIndex {
|
||||
index: context.input.num_files(),
|
||||
row_group_index: -1,
|
||||
@@ -1083,30 +1198,30 @@ async fn preflight_rejects_unsupported_assignments_without_artifact_io() {
|
||||
.contains("foreign SST")
|
||||
);
|
||||
|
||||
// Missing columnar samples are a not-ready answer, not an error.
|
||||
let mut context = make_context();
|
||||
let index = context.input.series_index.as_mut().unwrap();
|
||||
index.version = Arc::new(SeriesIndexVersion {
|
||||
series_indexes: index.version.series_indexes.clone(),
|
||||
..Default::default()
|
||||
});
|
||||
assert!(!validate_inputs(&context, &ranges).unwrap());
|
||||
// Invalid assignments remain errors even when sample coverage is incomplete.
|
||||
let mut invalid = ranges.clone();
|
||||
invalid.push(PartitionRange {
|
||||
identifier: ranges[0].identifier + ranges.len(),
|
||||
..ranges[0]
|
||||
});
|
||||
assert!(
|
||||
validate_inputs(&context, &ranges)
|
||||
validate_inputs(&context, &invalid)
|
||||
.unwrap_err()
|
||||
.to_string()
|
||||
.contains("no columnar samples")
|
||||
.contains("invalid prepared scan range")
|
||||
);
|
||||
let scan = SeriesScan::new(context.input, false);
|
||||
let binding = labels(scan.metadata());
|
||||
assert!(scan.raw_series_schema(&binding).is_err());
|
||||
assert!(
|
||||
scan.scan_series_partition(
|
||||
&QueryScanContext::default(),
|
||||
&ExecutionPlanMetricsSet::new(),
|
||||
0,
|
||||
&binding
|
||||
)
|
||||
.is_err()
|
||||
);
|
||||
// The raw path stays available for the ordinary scan and preparation is skipped.
|
||||
assert!(scan.raw_series_schema(&binding).unwrap().is_some());
|
||||
assert!(
|
||||
scan.prepare_series(
|
||||
&QueryScanContext::default(),
|
||||
@@ -1114,8 +1229,49 @@ async fn preflight_rejects_unsupported_assignments_without_artifact_io() {
|
||||
&binding
|
||||
)
|
||||
.await
|
||||
.is_err()
|
||||
.unwrap()
|
||||
.is_none()
|
||||
);
|
||||
assert!(reads.lock().unwrap().is_empty());
|
||||
engine.stop().await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn series_artifact_read_failure_propagates() {
|
||||
let (_env, engine, region) = fixture().await;
|
||||
build(&engine, ®ion).await;
|
||||
// The published catalog is complete; the injected manifest read failure must still propagate.
|
||||
let store = region.series_index_store.clone().unwrap().layer(
|
||||
MockLayerBuilder::default()
|
||||
.reader_factory(Arc::new(|path, _, inner| {
|
||||
Box::new(FailingReader {
|
||||
inner,
|
||||
fail: path.contains("/series/") && path.ends_with("/manifest.json"),
|
||||
})
|
||||
}))
|
||||
.build()
|
||||
.unwrap(),
|
||||
);
|
||||
let context = StreamContext::seq_scan_ctx(scan_input(
|
||||
®ion,
|
||||
store,
|
||||
SelectedFormat::Columnar,
|
||||
(0..region.version().metadata.column_metadatas.len()).collect(),
|
||||
));
|
||||
let ranges = context.partition_ranges();
|
||||
assert!(validate_inputs(&context, &ranges).unwrap());
|
||||
let scan = SeriesScan::new(context.input, false);
|
||||
let binding = labels(scan.metadata());
|
||||
assert!(scan.raw_series_schema(&binding).unwrap().is_some());
|
||||
let error = scan
|
||||
.prepare_series(
|
||||
&QueryScanContext::default(),
|
||||
&ExecutionPlanMetricsSet::new(),
|
||||
&binding,
|
||||
)
|
||||
.await
|
||||
.err()
|
||||
.expect("a failed artifact read must not report a not-ready source");
|
||||
assert!(error.output_msg().contains("injected"), "{error}");
|
||||
engine.stop().await.unwrap();
|
||||
}
|
||||
|
||||
@@ -75,6 +75,7 @@ use partition::expr::{Operand, PartitionExpr, RestrictedOp};
|
||||
use partition::multi_dim::MultiDimPartitionRule;
|
||||
use query::parser::QueryStatement;
|
||||
use query::plan::extract_and_rewrite_full_table_names;
|
||||
use query::promql::planner::NATIVE_PROMQL_EXTENSION_KEY;
|
||||
use query::query_engine::DefaultSerializer;
|
||||
use query::sql::create_table_stmt;
|
||||
use session::context::QueryContextRef;
|
||||
@@ -1164,7 +1165,12 @@ impl StatementExecutor {
|
||||
)
|
||||
.await?
|
||||
}
|
||||
Statement::Tql(query) => self.plan_tql(query.clone(), &ctx).await?,
|
||||
Statement::Tql(query) => {
|
||||
// The serialized view plan must not depend on the native execution extension.
|
||||
let mut view_ctx = ctx.as_ref().clone();
|
||||
view_ctx.set_extension(NATIVE_PROMQL_EXTENSION_KEY, "false");
|
||||
self.plan_tql(query.clone(), &Arc::new(view_ctx)).await?
|
||||
}
|
||||
_ => {
|
||||
return InvalidViewStmtSnafu {}.fail();
|
||||
}
|
||||
|
||||
+47
-49
@@ -16,38 +16,57 @@ operators. Statement-to-plan translation lives in `src/query/src/promql/`.
|
||||
| Kernels | `src/promql/src/functions.rs`, `src/promql/src/functions/` | PromQL functions |
|
||||
| Matching | `src/promql/src/vector_match.rs` | Binary vector matching |
|
||||
| Extension operators | `src/promql/src/extension_plan.rs`, `src/promql/src/extension_plan/` | Logical nodes and execution plans |
|
||||
| Native model | `src/promql/src/native/` | Compiled tree, chunks, display, metrics |
|
||||
| Native model | `src/promql/src/native/` | Compiled tree, chunks, display, metrics and execution |
|
||||
|
||||
## Two evaluation paths, and their status
|
||||
## Evaluation paths
|
||||
|
||||
The easiest thing to get wrong here. Both exist; only one runs.
|
||||
|
||||
- **Extension operators are live.** `PromPlanner` builds `SeriesDivide`,
|
||||
`SeriesNormalize`, `RangeManipulate`, `InstantManipulate`, `EmptyMetric` and
|
||||
`UnionDistinctOn`; `PromExtensionPlanner`
|
||||
(`src/promql/src/extension_plan/planner.rs`) maps each logical node to its
|
||||
execution plan, one branch per node, `None` for anything else.
|
||||
- **Nothing plans into the native model yet.** `src/query/src/promql/native_compile.rs`
|
||||
compiles a statement into a `NativeTree`, but no planner calls it, so a tree
|
||||
built here reaches no execution. `src/promql/src/native/` holds the tree, the
|
||||
chunk protocol, display and metrics, and — once a runtime lands — the driver
|
||||
that runs a tree node by node. Compiling and running are separate concerns
|
||||
owned by separate files; do not assume a tree is executable because it
|
||||
compiles.
|
||||
- **Extension operators:** the existing PromQL path plans operators such as
|
||||
`SeriesDivide`, `SeriesNormalize`, `RangeManipulate`, `InstantManipulate`,
|
||||
`EmptyMetric` and `UnionDistinctOn`. `PromExtensionPlanner`
|
||||
(`src/promql/src/extension_plan/planner.rs`) maps the logical nodes to their
|
||||
execution plans.
|
||||
- **Experimental fused native-series POC (default off):** the query planner offers
|
||||
a PromQL statement only when the engine's `in_process_region_transport` and
|
||||
`QueryOptions::experimental_native_promql` options are enabled and the per-request
|
||||
`query.experimental_native_promql` extension is true. Admission is limited to a
|
||||
bare vector selector (parentheses allowed) and unary `rate`, `increase`,
|
||||
`sum_over_time`, and `abs` calls over a selector; range selectors occur only
|
||||
inside those calls. It is a standalone route: unsupported expressions and all
|
||||
other cases use the existing planner. Accepted plans use process-local
|
||||
`FusedSeries`, with ordinary `RecordBatch` output collected through DataFusion.
|
||||
Scan N→Fused N retains region/partition parallelism, but this POC does not merge
|
||||
one logical series across regions or scan partitions, guarantee canonical final
|
||||
labels across them, or use global label-set occupancy; `NativeExecutor` retains
|
||||
its partition-local `Emitted` state. Repartitioning redistributes rows but does not perform
|
||||
semantic series merging. A future gather/repartition plus merge/finalize is
|
||||
outside this POC. Metric scans retain the static `__table_id` predicate.
|
||||
- **Admission and SQL output:** the guards live in
|
||||
`src/query/src/promql/planner/native.rs`. Non-`Float64` fields, alternative
|
||||
(`or`) matchers and `rate`/`increase` on temporality-tagged tables keep the
|
||||
existing plan; the latter preserves legacy per-series delta summation.
|
||||
Allowing query fallback through the engine option, session variable or
|
||||
`query_fallback` hint also disables native admission. Native SQL/TQL output
|
||||
has column order `time index, fields, tags`, raw field names with nullable
|
||||
`Float64` values, and nullable `Utf8` tags (including integer table tags).
|
||||
Unlike a legacy raw selector, it does not preserve tag types; legacy function
|
||||
expressions can also use different field names and column order. The native
|
||||
time index is a non-nullable millisecond timestamp without a time zone.
|
||||
|
||||
## Change coupling
|
||||
|
||||
- **An extension node reaches a datanode only if three things line up**: a payload
|
||||
branch in `src/common/substrait/src/extension_serializer.rs`, a physical branch
|
||||
in `PromExtensionPlanner`, and a `Commutativity::Commutative` arm in
|
||||
`check_extension_plan` (`src/query/src/dist_plan/commutativity.rs`). Missing the
|
||||
third is silent: the node falls through to `Unsupported`, the distributed
|
||||
boundary is placed above it, and it never reaches a region — with no error.
|
||||
- Distributed extension operators need matching handling in the Substrait
|
||||
serializer (`src/common/substrait/src/extension_serializer.rs`), the physical
|
||||
`PromExtensionPlanner`, and `check_extension_plan`
|
||||
(`src/query/src/dist_plan/commutativity.rs`). Process-local `FusedSeries` is
|
||||
the exception: it is not serialized to a datanode. Its distributed placement
|
||||
is restricted to a direct table scan, a bound metric name and trees that read
|
||||
one series at a time. Native admission requires partition columns to be tags,
|
||||
but region/partition-local execution does not merge a logical series across
|
||||
regions or scan partitions; do not treat that placement guard as a global
|
||||
series-merge or final-label guarantee.
|
||||
- A new logical node also needs its physical planner reachable from
|
||||
`DfQueryPlanner::create_physical_plan`, assembled in
|
||||
`src/query/src/query_engine/state.rs`.
|
||||
- User-visible changes need planner coverage plus a sqlness case under
|
||||
`tests/cases/standalone/common/promql/`.
|
||||
|
||||
## Gotchas
|
||||
|
||||
@@ -72,9 +91,9 @@ The easiest thing to get wrong here. Both exist; only one runs.
|
||||
the producer ordering the adapter adds. Nothing in that metadata binds the key
|
||||
to a decode — the frame decoder compares those format keys and the field list,
|
||||
and never a label identity. `unpack_raw_series` therefore requires a
|
||||
caller-supplied `RawLabelSource`: decode with the source that was actually bound
|
||||
when the stream was produced, because the frame schema does not replace that
|
||||
binding. The metadata describes format, not label identity.
|
||||
caller-supplied `RawLabelSource`: decode with the source that was actually
|
||||
bound when the stream was produced, because the frame schema does not replace
|
||||
that binding. The metadata describes format, not label identity.
|
||||
- **A wrong label binding is silently wrong, not an error.** Storage validates the
|
||||
binding it is given (metadata identity, tag-ness, name matches, no duplicates)
|
||||
but cannot tell whether that set is the one the query meant, and the metric name
|
||||
@@ -85,26 +104,5 @@ The easiest thing to get wrong here. Both exist; only one runs.
|
||||
```bash
|
||||
cargo nextest run -p promql
|
||||
cargo nextest run -p query # planner and native-compile changes
|
||||
cargo sqlness bare -t <case> # PromQL cases live under tests/cases/standalone/common/promql/
|
||||
cargo sqlness bare -t <case> # PromQL cases under tests/cases/standalone/common/promql/
|
||||
```
|
||||
|
||||
## Not wired up yet
|
||||
|
||||
The storage side of the raw path exists; what is missing is planner wiring. Both
|
||||
gaps below are structural, not bugs.
|
||||
|
||||
- **The native tree has no planner entry point.** `native_compile::compile` has no
|
||||
caller, and a `NativeTree` becomes a physical plan only once something builds it
|
||||
while it still holds the AST and the selector binding — the point where the
|
||||
matchers, `__field__` and the lookback are still available. Compiling later, from
|
||||
a plan, loses them.
|
||||
- **No planner reaches a raw scan.** mito2 implements
|
||||
`RegionScanner::raw_series_schema` (`mito2/src/read/series_scan.rs`), so
|
||||
`RegionScanExec::try_to_series_scan` can build `SeriesScanExec` from a prepared
|
||||
scanner; the trait's `Ok(None)` default only refuses engines that have not
|
||||
implemented it. What is missing is the planner side: nothing selects that
|
||||
conversion, and `SeriesScanExec::children()` is empty, so wrapping a
|
||||
`RegionScanExec` before partition preparation leaves `ParallelizeScan` unable to
|
||||
assign partitions. Either give the partitioner a child it can see, or convert to
|
||||
the native leaf only after partition preparation — the second is the contract
|
||||
`RegionScanExec::try_to_series_scan` documents.
|
||||
|
||||
@@ -28,6 +28,7 @@ prost.workspace = true
|
||||
regex.workspace = true
|
||||
snafu.workspace = true
|
||||
store-api.workspace = true
|
||||
table.workspace = true
|
||||
twox-hash.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
|
||||
@@ -37,18 +37,15 @@
|
||||
//! was not compiled against, or an output schema the conversion cannot classify.
|
||||
|
||||
use std::cmp::Ordering;
|
||||
use std::collections::hash_map::Entry;
|
||||
use std::collections::{HashMap, VecDeque};
|
||||
use std::collections::VecDeque;
|
||||
use std::fmt::{self, Formatter};
|
||||
use std::hash::{Hash, Hasher};
|
||||
use std::pin::Pin;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::sync::Arc;
|
||||
use std::task::{Context, Poll};
|
||||
|
||||
use common_query::series_frame::{ensure_frame_matches_label_source, unpack_raw_series};
|
||||
use datafusion::arrow::array::{
|
||||
Array, ArrayRef, Float64Array, StringArray, TimestampMillisecondArray,
|
||||
};
|
||||
use datafusion::arrow::array::{ArrayRef, Float64Array, StringArray, TimestampMillisecondArray};
|
||||
use datafusion::arrow::buffer::ScalarBuffer;
|
||||
use datafusion::arrow::datatypes::{DataType, SchemaRef, TimeUnit};
|
||||
use datafusion::arrow::record_batch::RecordBatch;
|
||||
@@ -56,9 +53,7 @@ use datafusion::common::DFSchemaRef;
|
||||
use datafusion::common::tree_node::TreeNodeRecursion;
|
||||
use datafusion::error::{DataFusionError, Result as DataFusionResult};
|
||||
use datafusion::execution::context::TaskContext;
|
||||
use datafusion::execution::memory_pool::{
|
||||
MemoryConsumer, MemoryPool, MemoryReservation, UnboundedMemoryPool,
|
||||
};
|
||||
use datafusion::execution::memory_pool::UnboundedMemoryPool;
|
||||
use datafusion::logical_expr::{Expr, LogicalPlan, UserDefinedLogicalNodeCore};
|
||||
use datafusion::physical_expr::{EquivalenceProperties, PhysicalExpr};
|
||||
use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricsSet};
|
||||
@@ -69,17 +64,15 @@ use datafusion::physical_plan::{
|
||||
};
|
||||
use datatypes::series::SeriesOrdering;
|
||||
use datatypes::series::raw::RawSeriesChunk;
|
||||
use datatypes::series::samples::{DenseSamplesRef, SampleRef, Samples, Series, SeriesSchema};
|
||||
use datatypes::series::samples::{DenseSamplesRef, SampleRef, Samples, SeriesSchema};
|
||||
use datatypes::series::time::EvalGrid;
|
||||
use futures::Stream;
|
||||
use store_api::region_engine::{RawLabelSource, SendableRawSeriesStream};
|
||||
|
||||
use crate::evaluated::EvaluatedChunk;
|
||||
use crate::grid_series::{LabelSet, label_set};
|
||||
use crate::native::execute::{NativeExecutionState, NativeExecutor};
|
||||
use crate::native::metrics::{self, NodePath};
|
||||
use crate::native::plan::{NativeExpr, NativeTree};
|
||||
use crate::native::runs::{Run, absorb};
|
||||
use crate::source::materialize_raw_labels;
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -456,89 +449,6 @@ fn tree_input(tree: &NativeTree) -> DataFusionResult<TreeInput> {
|
||||
Ok(TreeInput { path, schema })
|
||||
}
|
||||
|
||||
/// Refuses a tree whose logical series can be split across the partitions this node runs.
|
||||
///
|
||||
/// Keeping the metric name is not what keeps two raw series apart. An empty tag value and an absent
|
||||
/// tag are the same canonical label set, so `{host="h", zone=""}` and `{host="h"}` are one series to
|
||||
/// Prometheus while the scan encodes them as two raw keys. A window that evaluated its input without
|
||||
/// merging it therefore passes such a pair through as two series under one label set, exactly as the
|
||||
/// functions that drop `__name__` do by merging them, and a map drops `__name__` outright.
|
||||
///
|
||||
/// A stateful node above one of those needs both halves together, and the runtime does provide that
|
||||
/// within one execution: a window whose input can hold two series under one label set collects the
|
||||
/// fragments, merges each label set once, and only then evaluates. That merge is per execution, and
|
||||
/// this node runs one execution per partition. Two raw series that become one label set can land in
|
||||
/// different partitions - the scan keeps a *series* in one partition, which says nothing about two
|
||||
/// different series that later become indistinguishable - so each execution merges only the half it
|
||||
/// was given. The consumer above then answers from one half and reports nothing for the other, which
|
||||
/// is an empty result rather than an error.
|
||||
///
|
||||
/// Whether a window merged its input is runtime state this check does not have: it runs over a bare
|
||||
/// tree, where a window's input is only something that produces samples. A window below another
|
||||
/// window is therefore treated as one that did not merge, whatever function it runs, which refuses
|
||||
/// `rate` under `rate` (which drops the name) and `last_over_time` under `rate` (which keeps it)
|
||||
/// alike.
|
||||
///
|
||||
/// Reading such a tree needs an exchange above this node that merges a logical series across
|
||||
/// partitions, which does not exist yet. Until it does the tree is refused where it can hurt: over
|
||||
/// one partition there is nothing to split, so only a tree run over more than one is checked.
|
||||
///
|
||||
/// The compiler cannot build a nested window today - [`compile`](crate::promql::native_compile::compile)
|
||||
/// requires a range function to apply directly to a selector, `<range function>(<selector>[<range>])`,
|
||||
/// and a selector is raw samples - so no tree it produces has a window above a window, and this
|
||||
/// check costs no executable query. It is checked because [`FusedSeriesExec::try_new`] accepts any
|
||||
/// well-formed tree, and a tree it accepts is a tree it must answer correctly.
|
||||
fn reject_a_split_logical_series(
|
||||
tree: &NativeTree,
|
||||
partition_count: usize,
|
||||
) -> DataFusionResult<()> {
|
||||
if partition_count == 1 {
|
||||
return Ok(());
|
||||
}
|
||||
// The walk is outermost first, so it meets nodes in the reverse of the order data flows: a
|
||||
// consumer met before a rewriter is one that the rewriter feeds.
|
||||
let mut consumer: Option<String> = None;
|
||||
let mut rejection: Option<(String, String)> = None;
|
||||
metrics::walk(tree.root(), &mut |node, path| {
|
||||
if rejection.is_some() {
|
||||
return;
|
||||
}
|
||||
// A node can leave two raw series on one canonical label set. A map always can: it drops
|
||||
// `__name__` and every empty-valued tag. An instant selection preserves its input labels,
|
||||
// rather than dropping them, but can still leave two raw series on one canonical label set:
|
||||
// `{zone=""}` and the absent `zone` are one label set and two raw keys. A window can
|
||||
// whenever it did not merge its input, which a bare tree says nothing about: a check on the
|
||||
// function alone accepts `last_over_time`, whose output keeps the name and can still hold
|
||||
// both raw series. Every map, instant selection, and window is therefore taken to be a
|
||||
// possible rewriter, and the refusal below is what narrows that to the shape that costs
|
||||
// correctness - a window with a rewriter below it, a window being the only consumer here.
|
||||
let rewrites = matches!(
|
||||
node.expr,
|
||||
NativeExpr::Map { .. } | NativeExpr::Instant { .. } | NativeExpr::Window { .. }
|
||||
);
|
||||
if rewrites && let Some(above) = &consumer {
|
||||
rejection = Some((path.as_str().to_string(), above.clone()));
|
||||
return;
|
||||
}
|
||||
let consumes = matches!(node.expr, NativeExpr::Window { .. });
|
||||
if consumes && consumer.is_none() {
|
||||
consumer = Some(path.as_str().to_string());
|
||||
}
|
||||
});
|
||||
let Some((rewriter, above)) = rejection else {
|
||||
return Ok(());
|
||||
};
|
||||
Err(DataFusionError::Plan(format!(
|
||||
"FusedSeries runs a tree where `{rewriter}` can leave two raw series on one canonical label \
|
||||
set and `{above}` above it reads a whole series, over {partition_count} partitions: an \
|
||||
empty tag value and an absent tag are the same label set, the runtime merges such \
|
||||
fragments once per execution, and this node runs one execution per partition, so `{above}` \
|
||||
would read one half and answer an empty result. Reading this tree needs an exchange between \
|
||||
`{rewriter}` and `{above}` that merges a logical series across partitions, which is not \
|
||||
implemented"
|
||||
)))
|
||||
}
|
||||
|
||||
/// Checks that `tree` reads the schema `binding` declares, on the grid it declares.
|
||||
///
|
||||
/// The output grid is compared by [`EvalGrid::spec`], not by allocation: that is the rule
|
||||
@@ -762,11 +672,6 @@ pub struct FusedSeriesExec {
|
||||
input: Arc<dyn ExecutionPlan>,
|
||||
/// The columns this node produces.
|
||||
output_schema: SchemaRef,
|
||||
/// The grid positions every label set in this plan's output already occupies, and what they
|
||||
/// keep charged on the memory pool.
|
||||
///
|
||||
/// Plan-level rather than per-partition: see [`LabelSets`].
|
||||
labels: SharedLabelSets,
|
||||
/// Per-node metrics, registered by each execution of the tree.
|
||||
metric: ExecutionPlanMetricsSet,
|
||||
/// This node's partitioning, emission and boundedness.
|
||||
@@ -807,8 +712,6 @@ impl FusedSeriesExec {
|
||||
)?;
|
||||
let classification =
|
||||
classify_output(&output_schema, &tree_input.schema, &tags, &time_column)?;
|
||||
// Read here rather than below, because whether the split below can happen decides whether
|
||||
// the tree shape is one this node may run at all.
|
||||
let input_properties = input.properties();
|
||||
let partition_count = input_properties.partitioning.partition_count();
|
||||
if partition_count == 0 {
|
||||
@@ -816,7 +719,6 @@ impl FusedSeriesExec {
|
||||
"FusedSeriesExec input has no partitions".into(),
|
||||
));
|
||||
}
|
||||
reject_a_split_logical_series(&tree, partition_count)?;
|
||||
// The frame states its own label layout and series ordering; this plan's binding is an
|
||||
// expectation checked against that declaration, not an input to decoding.
|
||||
ensure_frame_matches_label_source(input.schema().as_ref(), &label_source)?;
|
||||
@@ -857,10 +759,6 @@ impl FusedSeriesExec {
|
||||
input_properties.emission_type,
|
||||
input_properties.boundedness,
|
||||
));
|
||||
// Shared by every partition's stream: a label set colliding is a property of the whole
|
||||
// output vector, not of one partition. Building it needs no pool; the owner's reservation is
|
||||
// attached in `execute`, which is the first point that has one.
|
||||
let labels: SharedLabelSets = Arc::new(Mutex::new(LabelSets::default()));
|
||||
Ok(Self {
|
||||
promql,
|
||||
tree,
|
||||
@@ -873,7 +771,6 @@ impl FusedSeriesExec {
|
||||
tags: tags.into(),
|
||||
input,
|
||||
output_schema,
|
||||
labels,
|
||||
metric: ExecutionPlanMetricsSet::new(),
|
||||
properties,
|
||||
})
|
||||
@@ -967,14 +864,6 @@ impl ExecutionPlan for FusedSeriesExec {
|
||||
partition: usize,
|
||||
context: Arc<TaskContext>,
|
||||
) -> DataFusionResult<SendableRecordBatchStream> {
|
||||
// The label-set state is the plan's, so the reservation that pays for it is attached exactly
|
||||
// once, to the pool of the first execution. Every partition calls this and only the first
|
||||
// takes effect: re-registering per stream would move a charge the other partitions keep
|
||||
// growing, and would give it back when that stream ends while the entries stay.
|
||||
self.labels
|
||||
.lock()
|
||||
.unwrap_or_else(|e| e.into_inner())
|
||||
.attach(&context.runtime_env().memory_pool);
|
||||
let input = self.input.execute(partition, Arc::clone(&context))?;
|
||||
// Batches become raw series chunks through the shared frame decoder, which is the inverse of
|
||||
// what the producer packed: one chunk per frame row, each series closed exactly once by the
|
||||
@@ -990,8 +879,6 @@ impl ExecutionPlan for FusedSeriesExec {
|
||||
&self.metric,
|
||||
partition,
|
||||
)?;
|
||||
// What the label sets hold is charged to the reservation attached above, so nothing is
|
||||
// reserved per stream: the entries outlive every stream and are released with the plan.
|
||||
Ok(Box::pin(FusedSeriesStream {
|
||||
input: Some(input),
|
||||
executor: Some(NativeExecutor::new(state)),
|
||||
@@ -999,7 +886,6 @@ impl ExecutionPlan for FusedSeriesExec {
|
||||
classification: Arc::clone(&self.classification),
|
||||
tags: Arc::clone(&self.tags),
|
||||
output_schema: self.output_schema.clone(),
|
||||
labels: Arc::clone(&self.labels),
|
||||
ready: VecDeque::new(),
|
||||
done: false,
|
||||
}))
|
||||
@@ -1100,263 +986,6 @@ fn field_column(samples: &Samples) -> DataFusionResult<ArrayRef> {
|
||||
}
|
||||
}
|
||||
|
||||
/// The bytes one label set holds, for the memory reservation.
|
||||
///
|
||||
/// It counts each label's name and value with the pair that refers to them, which is the same
|
||||
/// convention [`crate::native::execute`]'s own label-set table keeps.
|
||||
fn labels_bytes(labels: &LabelSet) -> usize {
|
||||
labels
|
||||
.iter()
|
||||
.map(|(name, value)| size_of::<(Arc<str>, Arc<str>)>() + name.len() + value.len())
|
||||
.sum()
|
||||
}
|
||||
|
||||
/// Renders the label set Prometheus compares the way it prints one, for the duplicate-label-set
|
||||
/// error.
|
||||
fn render_labels(labels: &LabelSet) -> String {
|
||||
let parts = labels
|
||||
.iter()
|
||||
.map(|(name, value)| format!("{name}={value:?}"))
|
||||
.collect::<Vec<_>>();
|
||||
format!("{{{}}}", parts.join(", "))
|
||||
}
|
||||
|
||||
/// The grid positions each final label set already occupies, per field and as merged, sorted runs.
|
||||
///
|
||||
/// Shared by every partition's stream, because a label-set collision is a property of the whole
|
||||
/// output vector and not of one partition. Two series that carry the same labels after the range
|
||||
/// function and its maps drop a label land in different partitions, and the runtime's own per-node
|
||||
/// `Emitted` table never sees both: it belongs to one execution. Checking only within a partition
|
||||
/// would miss exactly those collisions.
|
||||
///
|
||||
/// The key is [`crate::grid_series::label_set`]'s normalization of the final labels rather than
|
||||
/// [`datatypes::series::MetricName`] itself: a label with an empty value is no label to
|
||||
/// Prometheus, so `{host="h", zone=""}` and `{host="h"}` are one label set and have to be compared
|
||||
/// as one. Encoding the tags as they are would file them apart and let one point hold two values.
|
||||
///
|
||||
/// Occupancy is per field and per position, not per chunk: two series whose label sets became equal
|
||||
/// while holding values in *different* fields merge cleanly in Prometheus, so only the same field
|
||||
/// at the same position collides. That is the rule
|
||||
/// [`crate::grid_series::merge_same_labelset`] applies.
|
||||
///
|
||||
/// The map belongs to the plan, not to one stream: a partition that ends leaves what it claimed
|
||||
/// for the partitions still running, so the entries are only released with the plan itself. The
|
||||
/// charge has the same lifetime, which is what makes the bytes the map holds - one entry per label
|
||||
/// set of the output, kept until the plan dies - visible to the memory pool. Charging each stream
|
||||
/// for what it inserted instead returned those bytes when that stream ended, while the entries the
|
||||
/// later partitions are checked against stayed: everything the map held after the first partition
|
||||
/// was unaccounted for.
|
||||
#[derive(Default)]
|
||||
struct LabelSets {
|
||||
/// The grid positions each final label set already occupies, one run list per field of the
|
||||
/// series schema. Each list is sorted, disjoint and has the runs that touch joined, which is
|
||||
/// what lets a claim look a position up instead of walking the history.
|
||||
seen: HashMap<LabelSet, Vec<Vec<Run>>>,
|
||||
/// What `seen` keeps charged on the task's memory pool. `None` until the first execution,
|
||||
/// since construction has no `TaskContext` and therefore no pool to register against.
|
||||
reservation: Option<MemoryReservation>,
|
||||
}
|
||||
|
||||
impl LabelSets {
|
||||
/// Attaches the reservation to the pool of the first execution.
|
||||
///
|
||||
/// Later executions keep the first reservation: it already pays for what `seen` holds, and
|
||||
/// moving it to another pool would drop the bytes other partitions are still growing.
|
||||
fn attach(&mut self, pool: &Arc<dyn MemoryPool>) {
|
||||
if self.reservation.is_none() {
|
||||
self.reservation =
|
||||
Some(MemoryConsumer::new("FusedSeriesExec label sets").register(pool));
|
||||
}
|
||||
}
|
||||
|
||||
/// Claims the occupied grid positions of one chunk of a final series for its label set.
|
||||
///
|
||||
/// A collision is an error only when both series hold a sample of *one field* at the same
|
||||
/// evaluation timestamp, which is Prometheus' range-query rule. Series with disjoint points are
|
||||
/// legal, and so are series that fill the same points field by field: Prometheus merges both
|
||||
/// into one output series, while this node emits them as separate batches and defers that merge
|
||||
/// to matrix assembly, so every sample survives.
|
||||
fn claim(&mut self, series: &Series) -> DataFusionResult<()> {
|
||||
let runs = occupied_runs(series)?;
|
||||
// A chunk holding no value occupies nothing, so there is no collision to remember and no
|
||||
// entry to write. Such a chunk is not rare: a selection over a series that staleness or the
|
||||
// range left empty still closes it, and hands over a closing chunk with no rows. Recording
|
||||
// one entry per such series would grow the table - and its charge, which the plan holds for
|
||||
// its whole life - for a query that never answers a value, until an empty result exhausted
|
||||
// the pool. The executor's own `Emitted::occupy` takes the same early return.
|
||||
if runs.iter().all(Vec::is_empty) {
|
||||
return Ok(());
|
||||
}
|
||||
let key = label_set(&series.metric);
|
||||
let Self { seen, reservation } = self;
|
||||
// Nothing reaches a claim without an `execute` first, and an execution that ran attached the
|
||||
// reservation. Claiming without one would leave the map growing uncharged, so it is refused
|
||||
// rather than silently skipped.
|
||||
let Some(reservation) = reservation.as_mut() else {
|
||||
return Err(DataFusionError::Execution(
|
||||
"FusedSeriesExec label-set state is not attached to a memory pool".to_string(),
|
||||
));
|
||||
};
|
||||
match seen.entry(key) {
|
||||
Entry::Occupied(entry) => {
|
||||
if let Some(field) = colliding_field(entry.get(), &runs) {
|
||||
return Err(DataFusionError::Execution(format!(
|
||||
"vector cannot contain metrics with the same labelset: {} (field \
|
||||
{field})",
|
||||
render_labels(entry.key())
|
||||
)));
|
||||
}
|
||||
let stored = entry.into_mut();
|
||||
// An entry recorded from a chunk of fewer fields is padded, so every field of this
|
||||
// chunk has a slot of its own: a slot is a field index of the schema, so both sides
|
||||
// count the same field the same way. A chunk of *fewer* fields than the entry
|
||||
// leaves the remaining slots to the chunks that fill them.
|
||||
// Stage only the affected run lists, and reserve before committing. A forward
|
||||
// append whose existing capacity suffices needs no copy of its stored history.
|
||||
let mut growth = 0usize;
|
||||
let mut shrink = 0usize;
|
||||
let mut merged_runs = Vec::with_capacity(runs.len());
|
||||
for (index, incoming) in runs.iter().enumerate() {
|
||||
let Some(old) = stored.get(index) else {
|
||||
let copy = incoming.clone();
|
||||
growth += copy.capacity();
|
||||
merged_runs.push(Some(copy));
|
||||
continue;
|
||||
};
|
||||
let append_fits = incoming.is_empty()
|
||||
|| (old.last().is_none_or(|last| incoming[0].start >= last.end)
|
||||
&& old.capacity()
|
||||
>= old.len() + incoming.len()
|
||||
- usize::from(old.last().is_some_and(|last| {
|
||||
incoming
|
||||
.first()
|
||||
.is_some_and(|first| first.start == last.end)
|
||||
})));
|
||||
if append_fits {
|
||||
merged_runs.push(None);
|
||||
} else {
|
||||
let mut merged = old.clone();
|
||||
absorb(&mut merged, incoming);
|
||||
if merged.capacity() >= old.capacity() {
|
||||
growth += merged.capacity() - old.capacity();
|
||||
} else {
|
||||
shrink += old.capacity() - merged.capacity();
|
||||
}
|
||||
merged_runs.push(Some(merged));
|
||||
}
|
||||
}
|
||||
reservation.try_grow(growth * size_of::<Run>())?;
|
||||
for (index, (incoming, merged)) in runs.iter().zip(merged_runs).enumerate() {
|
||||
if let Some(merged) = merged {
|
||||
if index < stored.len() {
|
||||
stored[index] = merged;
|
||||
} else {
|
||||
stored.push(merged);
|
||||
}
|
||||
} else if index < stored.len() {
|
||||
absorb(&mut stored[index], incoming);
|
||||
}
|
||||
}
|
||||
reservation.try_shrink(shrink * size_of::<Run>())?;
|
||||
}
|
||||
Entry::Vacant(entry) => {
|
||||
// Only the label-set state is this node's own cross-poll state that grows with
|
||||
// the output, so it is the only thing reserved here. The runs are charged by the
|
||||
// capacity they were gathered with rather than by their length, so that a later
|
||||
// claim appending into that slack is charged nothing for it.
|
||||
let bytes = labels_bytes(entry.key())
|
||||
+ runs.iter().map(|runs| runs.capacity()).sum::<usize>() * size_of::<Run>();
|
||||
reservation.try_grow(bytes)?;
|
||||
entry.insert(runs);
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
type SharedLabelSets = Arc<Mutex<LabelSets>>;
|
||||
|
||||
/// The grid positions the rows of one chunk hold a value at, one run list per field.
|
||||
///
|
||||
/// Each [`Run`] is a half-open range of absolute grid indices, and the runs of one field are sorted
|
||||
/// and disjoint. Occupancy is kept per field because it is what a collision is: two series whose
|
||||
/// label sets became equal merge field by field in Prometheus, so only two values of the same field
|
||||
/// at one position collide. Marking the whole position as taken as soon as *any* field held a value
|
||||
/// would refuse the complementary pair that merges cleanly.
|
||||
fn occupied_runs(series: &Series) -> DataFusionResult<Vec<Vec<Run>>> {
|
||||
// Rows are addressed on the statement's grid, so a row sits at the absolute position two
|
||||
// series can collide on. An explicit axis has no such position, and a chunk never carries
|
||||
// one: this is the refusal `EvaluatedChunk` already makes when the chunk is built.
|
||||
if series.time.grid().is_none() {
|
||||
return Err(DataFusionError::Execution(
|
||||
"FusedSeriesExec requires evaluated series on a time grid".to_string(),
|
||||
));
|
||||
}
|
||||
let rows = series.time.len();
|
||||
let range = series.time.as_grid_range();
|
||||
let mut runs: Vec<Vec<Run>> = vec![Vec::new(); series.fields.len()];
|
||||
// A field that is one dense, null-free float run holds a value on every row the chunk carries,
|
||||
// so it occupies exactly the range the chunk reports and no row of it has to be inspected. This
|
||||
// is the shape the fused evaluator produces for a range function over ordinary float samples.
|
||||
// Every other field is walked below, and a field filled here is skipped there.
|
||||
let mut walked: Vec<(usize, &Samples)> = Vec::new();
|
||||
for (index, samples) in series.fields.iter().enumerate() {
|
||||
match (range, samples.dense()) {
|
||||
(Some(range), Some(DenseSamplesRef::Float(values)))
|
||||
if rows > 0 && values.len() == rows && values.null_count() == 0 =>
|
||||
{
|
||||
runs[index] = vec![Run {
|
||||
start: range.start(),
|
||||
end: range.start() + rows,
|
||||
}];
|
||||
}
|
||||
_ => walked.push((index, samples)),
|
||||
}
|
||||
}
|
||||
if walked.is_empty() {
|
||||
return Ok(runs);
|
||||
}
|
||||
// `positions` walks a chunk that forms one run and one whose rows form several the same way,
|
||||
// in position order, so the holes between two runs are never read as rows.
|
||||
let mut cursors: Vec<_> = walked.iter().map(|(_, samples)| samples.cursor()).collect();
|
||||
for (position, row) in series.time.positions() {
|
||||
for ((index, _), cursor) in walked.iter().zip(&mut cursors) {
|
||||
if !matches!(cursor.at(row), Some(sample) if !matches!(sample, SampleRef::Missing)) {
|
||||
continue;
|
||||
}
|
||||
match runs[*index].last_mut() {
|
||||
Some(run) if run.end == position => run.end += 1,
|
||||
_ => runs[*index].push(Run {
|
||||
start: position,
|
||||
end: position + 1,
|
||||
}),
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(runs)
|
||||
}
|
||||
|
||||
/// The field of `runs` that holds a value where `stored` already holds one, if any.
|
||||
///
|
||||
/// A collision is two values of one field at one position, so each field is checked on its own.
|
||||
/// Each incoming run is looked up rather than the stored history walked: the stored runs are sorted
|
||||
/// and disjoint, so the only run that can share a position with `run` is the first one that does not
|
||||
/// end before it, which is the look-up the runtime's own `Emitted::occupy` makes. Walking the
|
||||
/// history instead costs a sparse series one pass over everything it stored per chunk - a query
|
||||
/// whose chunks move forward one position at a time, a hole apart, is legal and would cost the
|
||||
/// square of its chunk count under the lock every partition shares.
|
||||
fn colliding_field(stored: &[Vec<Run>], runs: &[Vec<Run>]) -> Option<usize> {
|
||||
for (field, (stored, new)) in stored.iter().zip(runs).enumerate() {
|
||||
for run in new {
|
||||
let at = stored.partition_point(|stored| stored.end <= run.start);
|
||||
if stored.get(at).is_some_and(|stored| stored.overlaps(*run)) {
|
||||
return Some(field);
|
||||
}
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// The batches one execution of [`FusedSeriesExec`] yields.
|
||||
pub struct FusedSeriesStream {
|
||||
/// The input frames, read as raw series chunks, or `None` once the stream ended in an error.
|
||||
@@ -1375,10 +1004,6 @@ pub struct FusedSeriesStream {
|
||||
tags: Arc<[String]>,
|
||||
/// The schema of every batch this stream yields.
|
||||
output_schema: SchemaRef,
|
||||
/// The label-set occupancy shared with the other partitions' streams.
|
||||
///
|
||||
/// What it holds is charged to the plan's own reservation, attached by the first execution.
|
||||
labels: SharedLabelSets,
|
||||
/// The batches the executor settled since the last poll of the input.
|
||||
ready: VecDeque<RecordBatch>,
|
||||
/// Whether the input ended and everything it settled has been handed out.
|
||||
@@ -1447,9 +1072,6 @@ impl FusedSeriesStream {
|
||||
self.done = true;
|
||||
self.executor = None;
|
||||
self.input = None;
|
||||
// The label-set map and its charge are the plan's, so a stream stopping early gives back
|
||||
// nothing: the entries another partition is still checking against have to stay charged, and
|
||||
// are released when the plan itself is dropped.
|
||||
// A batch queued before the error was settled by a run that has now stopped.
|
||||
self.ready.clear();
|
||||
}
|
||||
@@ -1457,12 +1079,6 @@ impl FusedSeriesStream {
|
||||
/// Converts the chunks the executor settled into batches, in the order it produced them.
|
||||
fn convert(&mut self, chunks: Vec<EvaluatedChunk>) -> DataFusionResult<()> {
|
||||
for chunk in &chunks {
|
||||
// The label set is claimed before its batch is queued, so a collision surfaces at the
|
||||
// chunk that reached it and nothing this plan would refuse is handed out after it.
|
||||
self.labels
|
||||
.lock()
|
||||
.unwrap_or_else(|e| e.into_inner())
|
||||
.claim(&chunk.series)?;
|
||||
let batch = series_to_record_batch(
|
||||
&self.classification,
|
||||
&self.tags,
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -16,15 +16,16 @@ use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use datafusion::catalog::Session;
|
||||
use datafusion::error::Result as DfResult;
|
||||
use datafusion::error::{DataFusionError, Result as DfResult};
|
||||
use datafusion::logical_expr::physical_planning_context::PhysicalPlanningContext;
|
||||
use datafusion::logical_expr::{LogicalPlan, UserDefinedLogicalNode};
|
||||
use datafusion::physical_plan::ExecutionPlan;
|
||||
use datafusion::physical_planner::{ExtensionPlanner, PhysicalPlanner};
|
||||
use table::table::scan::RegionScanExec;
|
||||
|
||||
use crate::extension_plan::{
|
||||
Absent, EmptyMetric, HistogramFold, InstantManipulate, RangeManipulate, ScalarCalculate,
|
||||
SeriesDivide, SeriesNormalize, UnionDistinctOn,
|
||||
Absent, EmptyMetric, FusedSeries, HistogramFold, InstantManipulate, RangeManipulate,
|
||||
ScalarCalculate, SeriesDivide, SeriesNormalize, UnionDistinctOn,
|
||||
};
|
||||
|
||||
pub struct PromExtensionPlanner;
|
||||
@@ -65,8 +66,102 @@ impl ExtensionPlanner for PromExtensionPlanner {
|
||||
)))
|
||||
} else if let Some(node) = node.as_any().downcast_ref::<Absent>() {
|
||||
Ok(Some(node.to_execution_plan(physical_inputs[0].clone())))
|
||||
} else if let Some(node) = node.as_any().downcast_ref::<FusedSeries>() {
|
||||
plan_fused_series(
|
||||
node,
|
||||
physical_inputs,
|
||||
session.config_options().execution.target_partitions,
|
||||
)
|
||||
} else {
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The label `__name__` is reported under, which is a column of this node's output but not a tag
|
||||
/// column of any region.
|
||||
const METRIC_NAME: &str = "__name__";
|
||||
|
||||
/// Builds the physical node of a [`FusedSeries`] over the raw series leaf its input can provide.
|
||||
///
|
||||
/// The node reads a raw series frame, so what it needs below it is the region scan itself: the
|
||||
/// scan's own raw-series entrypoint is what packs a series' chunks contiguously and declares their
|
||||
/// ordering. Anything else in between - a filter, a projection, an exchange - produces a different
|
||||
/// frame, or no frame at all, and this layer cannot look under it.
|
||||
///
|
||||
/// `Ok(None)` is therefore the answer for a child that is not a region scan: the logical plan is
|
||||
/// already fixed and DataFusion reports an unplanned extension node normally. A region scan that is
|
||||
/// found but refuses the raw path is a different case — the logical route chose this node for this
|
||||
/// scan and nothing here can undo that choice, because the logical plan is not available — so it is
|
||||
/// reported as an error that names the reason instead of silently producing no plan.
|
||||
/// The fused leaf needs a prepared scan: range assignment runs after physical planning and cannot
|
||||
/// see a `SeriesScanExec`. Preparation keeps the scan's parallel layout, and the series-producer
|
||||
/// contract guarantees that one series stays in one scan partition.
|
||||
fn plan_fused_series(
|
||||
fused: &FusedSeries,
|
||||
physical_inputs: &[Arc<dyn ExecutionPlan>],
|
||||
target_partitions: usize,
|
||||
) -> DfResult<Option<Arc<dyn ExecutionPlan>>> {
|
||||
let Some(input) = physical_inputs.first() else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(scan) = input.downcast_ref::<RegionScanExec>() else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let binding = fused.binding();
|
||||
let Some(metric_name) = binding.metric_name.as_deref() else {
|
||||
return Err(DataFusionError::Plan(
|
||||
"the FusedSeries node binds no metric name, so the region scan below it cannot be \
|
||||
asked for the series of one metric"
|
||||
.to_string(),
|
||||
));
|
||||
};
|
||||
// `__name__` is the metric name rather than a column of the region, so it is not one of the
|
||||
// logical tags the scan is asked to decode; the frame names each series' own metric.
|
||||
let logical_tags = binding
|
||||
.tags
|
||||
.iter()
|
||||
.filter(|tag| tag.as_str() != METRIC_NAME)
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
let field_names = binding
|
||||
.series_schema
|
||||
.fields
|
||||
.iter()
|
||||
.map(|field| field.name.to_string())
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
let scan = if scan.is_partition_set() {
|
||||
Arc::new(scan.clone())
|
||||
} else {
|
||||
let partitions = table::table::scan::assign_partition_ranges(
|
||||
scan.get_partition_ranges(),
|
||||
target_partitions.max(1),
|
||||
);
|
||||
Arc::new(
|
||||
scan.with_new_partitions(partitions, target_partitions.max(1))
|
||||
.map_err(|error| DataFusionError::External(error.into_inner()))?,
|
||||
)
|
||||
};
|
||||
// The scan resolves the label source and the ordering it will declare itself, against its own
|
||||
// metadata snapshot, and hands both back through the leaf it builds. Naming them here instead
|
||||
// would be a second copy of a binding only the scanner can make: the physical metadata and
|
||||
// column ids a `RawLabelSource::PrimaryKey` carries are not reachable from a region scan's
|
||||
// public surface, so the layer that resolves them is the scanner.
|
||||
let Some(series_scan) =
|
||||
scan.try_to_named_series_scan(metric_name, &logical_tags, &field_names)?
|
||||
else {
|
||||
return Err(DataFusionError::Plan(format!(
|
||||
"the region scan under FusedSeries could not serve a raw series frame for metric \
|
||||
{metric_name:?}: the scanner refused the raw series path for this scan, so this node \
|
||||
has no input to read"
|
||||
)));
|
||||
};
|
||||
let label_source = series_scan.label_source().clone();
|
||||
let ordering = series_scan.raw_series_ordering();
|
||||
|
||||
fused
|
||||
.to_execution_plan(series_scan, label_source, ordering)
|
||||
.map(Some)
|
||||
}
|
||||
|
||||
@@ -349,6 +349,28 @@ pub struct NativeNode {
|
||||
pub output: NativeType,
|
||||
}
|
||||
|
||||
/// Whether `node` reads the whole input vector at a point rather than one series at a time.
|
||||
///
|
||||
/// Three kinds do: an `Aggregate` folds every series of its group, a `Binary` pairs two vectors'
|
||||
/// series by label, and a `Vector` operation (`sort`, `absent`, `topk`, `limitk`,
|
||||
/// `histogram_quantile`, ...) reorders or reads the whole vector at a grid point. Everything else -
|
||||
/// an input, a scalar, a selection, a window, a map - answers one series from one series' rows.
|
||||
///
|
||||
/// Both a region-side placement and the native route's admission ask this, and they must agree: a
|
||||
/// tree admitted here but refused there becomes a node the physical layer cannot build, and one
|
||||
/// refused here but admitted there would be pushed to a region that holds part of the metric.
|
||||
pub fn reads_across_series(node: &NativeNode) -> bool {
|
||||
match &node.expr {
|
||||
NativeExpr::Input(_) | NativeExpr::Scalar(_) => false,
|
||||
NativeExpr::Instant { input, .. }
|
||||
| NativeExpr::Window { input, .. }
|
||||
| NativeExpr::Map { input, .. } => reads_across_series(input),
|
||||
NativeExpr::Aggregate { .. } | NativeExpr::Binary { .. } | NativeExpr::Vector { .. } => {
|
||||
true
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A bound evaluation tree.
|
||||
///
|
||||
/// The tree itself carries no run-time state: every execution builds its own, so one compiled plan
|
||||
|
||||
@@ -21,12 +21,9 @@
|
||||
//! merged away, while a chunk that fills positions the set left empty merges cleanly, as
|
||||
//! [`crate::grid_series::merge_same_labelset`] does.
|
||||
//!
|
||||
//! The state is kept per label set, per field, as runs, so a chunk of many positions costs one run
|
||||
//! rather than one entry per position. Keeping it in one place is what keeps the two layers that
|
||||
//! check it — the native runtime's own [`crate::native::execute::Emitted`] table and the
|
||||
//! cross-partition [`crate::extension_plan::fused_series::LabelSets`] — from drifting apart: both
|
||||
//! ask the same functions for the same answers, so the fast path one of them earned is the fast path
|
||||
//! the other gets.
|
||||
//! The native runtime's partition-local [`crate::native::execute::Emitted`] table keeps this state
|
||||
//! per label set and field as runs, so a chunk of many positions costs one run rather than one
|
||||
//! entry per position.
|
||||
|
||||
/// One run of grid positions, half-open, that hold a value in one field.
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
|
||||
@@ -37,7 +37,7 @@ use datafusion_optimizer::push_down_filter::PushDownFilter;
|
||||
use datafusion_optimizer::rewrite_set_comparison::RewriteSetComparison;
|
||||
use datafusion_optimizer::scalar_subquery_to_join::ScalarSubqueryToJoin;
|
||||
use datafusion_optimizer::simplify_expressions::SimplifyExpressions;
|
||||
use promql::extension_plan::SeriesDivide;
|
||||
use promql::extension_plan::{FusedSeries, SeriesDivide};
|
||||
use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan};
|
||||
use table::metadata::TableType;
|
||||
use table::table::adapter::DfTableProviderAdapter;
|
||||
@@ -106,8 +106,35 @@ impl ExtensionOptions for DistPlannerOptions {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct DistPlannerAnalyzer;
|
||||
/// Rewrites a logical plan into a distributed plan: the part above a merge scan runs on the
|
||||
/// frontend, and every `MergeScan` carries the sub-plan that runs in a region.
|
||||
#[derive(Debug, Default)]
|
||||
pub struct DistPlannerAnalyzer {
|
||||
/// Whether every region a rewritten plan reads is reached in process.
|
||||
///
|
||||
/// A region input normally has to be encoded to Substrait before it can leave the frontend, and
|
||||
/// [`PlanRewriter::should_expand`] cuts the plan at the first node that does not encode. In
|
||||
/// process, a root `FusedSeries` plan travels as a `LogicalPlan` instead and skips that gate;
|
||||
/// other plans retain the ordinary encoding checks.
|
||||
///
|
||||
/// Whoever builds the query engine sets this, once, from the transport it wired. It is neither a
|
||||
/// query-time decision nor a property of a region.
|
||||
in_process_region_transport: bool,
|
||||
}
|
||||
|
||||
impl DistPlannerAnalyzer {
|
||||
/// Creates the analyzer.
|
||||
///
|
||||
/// `in_process_region_transport` must only be `true` for an engine whose region reads are served
|
||||
/// by region servers in the same process. Its plans are then handed to
|
||||
/// `RegionServer::handle_read` as `LogicalPlan`s; the remote read path
|
||||
/// (`RegionServer::handle_remote_read`) decodes Substrait bytes and cannot do that.
|
||||
pub fn new(in_process_region_transport: bool) -> Self {
|
||||
Self {
|
||||
in_process_region_transport,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl AnalyzerRule for DistPlannerAnalyzer {
|
||||
fn name(&self) -> &str {
|
||||
@@ -250,12 +277,15 @@ fn unwrap_dictionary_literals(plan: LogicalPlan) -> DfResult<Transformed<Logical
|
||||
impl DistPlannerAnalyzer {
|
||||
/// Try push down as many nodes as possible
|
||||
fn try_push_down(&self, plan: LogicalPlan) -> DfResult<LogicalPlan> {
|
||||
let in_process_region_transport = self.in_process_region_transport;
|
||||
// Use the subquery-aware transform so expression subqueries (including
|
||||
// nested scalar subqueries left in place by DataFusion 55's
|
||||
// `enable_physical_uncorrelated_scalar_subquery`) are visited at every
|
||||
// depth and their table scans get wrapped in `MergeScan`.
|
||||
let plan = plan.transform_with_subqueries(&Self::inspect_plan_with_subquery)?;
|
||||
let mut rewriter = PlanRewriter::new(&plan.data);
|
||||
let plan = plan.transform_with_subqueries(&|plan| {
|
||||
Self::inspect_plan_with_subquery(plan, in_process_region_transport)
|
||||
})?;
|
||||
let mut rewriter = PlanRewriter::new(&plan.data, in_process_region_transport);
|
||||
let result = plan.data.rewrite(&mut rewriter)?.data;
|
||||
Self::assign_merge_scan_remote_dyn_filter_producer_ids(result)
|
||||
}
|
||||
@@ -267,7 +297,10 @@ impl DistPlannerAnalyzer {
|
||||
Self::assign_merge_scan_remote_dyn_filter_producer_ids(result)
|
||||
}
|
||||
|
||||
fn inspect_plan_with_subquery(plan: LogicalPlan) -> DfResult<Transformed<LogicalPlan>> {
|
||||
fn inspect_plan_with_subquery(
|
||||
plan: LogicalPlan,
|
||||
in_process_region_transport: bool,
|
||||
) -> DfResult<Transformed<LogicalPlan>> {
|
||||
// Workaround for https://github.com/GreptimeTeam/greptimedb/issues/5469 and https://github.com/GreptimeTeam/greptimedb/issues/5799
|
||||
// FIXME(yingwen): Remove the `Limit` plan once we update DataFusion.
|
||||
if let LogicalPlan::Limit(_) | LogicalPlan::Distinct(_) = &plan {
|
||||
@@ -277,7 +310,12 @@ impl DistPlannerAnalyzer {
|
||||
let exprs = plan
|
||||
.expressions_consider_join()
|
||||
.into_iter()
|
||||
.map(|e| e.transform(&Self::transform_subquery).map(|x| x.data))
|
||||
.map(|e| {
|
||||
e.transform(|subquery_expr| {
|
||||
Self::transform_subquery(subquery_expr, in_process_region_transport)
|
||||
})
|
||||
.map(|x| x.data)
|
||||
})
|
||||
.collect::<DfResult<Vec<_>>>()?;
|
||||
|
||||
// Some plans that are special treated (should not call `with_new_exprs` on them)
|
||||
@@ -289,27 +327,33 @@ impl DistPlannerAnalyzer {
|
||||
}
|
||||
}
|
||||
|
||||
fn transform_subquery(expr: Expr) -> DfResult<Transformed<Expr>> {
|
||||
fn transform_subquery(
|
||||
expr: Expr,
|
||||
in_process_region_transport: bool,
|
||||
) -> DfResult<Transformed<Expr>> {
|
||||
match expr {
|
||||
Expr::Exists(exists) => Ok(Transformed::yes(Expr::Exists(Exists {
|
||||
subquery: Self::handle_subquery(exists.subquery)?,
|
||||
subquery: Self::handle_subquery(exists.subquery, in_process_region_transport)?,
|
||||
negated: exists.negated,
|
||||
}))),
|
||||
Expr::InSubquery(in_subquery) => Ok(Transformed::yes(Expr::InSubquery(InSubquery {
|
||||
expr: in_subquery.expr,
|
||||
subquery: Self::handle_subquery(in_subquery.subquery)?,
|
||||
subquery: Self::handle_subquery(in_subquery.subquery, in_process_region_transport)?,
|
||||
negated: in_subquery.negated,
|
||||
}))),
|
||||
Expr::ScalarSubquery(scalar_subquery) => Ok(Transformed::yes(Expr::ScalarSubquery(
|
||||
Self::handle_subquery(scalar_subquery)?,
|
||||
Self::handle_subquery(scalar_subquery, in_process_region_transport)?,
|
||||
))),
|
||||
|
||||
_ => Ok(Transformed::no(expr)),
|
||||
}
|
||||
}
|
||||
|
||||
fn handle_subquery(subquery: Subquery) -> DfResult<Subquery> {
|
||||
let mut rewriter = PlanRewriter::new(&subquery.subquery);
|
||||
fn handle_subquery(
|
||||
subquery: Subquery,
|
||||
in_process_region_transport: bool,
|
||||
) -> DfResult<Subquery> {
|
||||
let mut rewriter = PlanRewriter::new(&subquery.subquery, in_process_region_transport);
|
||||
let mut rewrote_subquery = subquery
|
||||
.subquery
|
||||
.as_ref()
|
||||
@@ -400,7 +444,12 @@ enum RewriterStatus {
|
||||
struct PlanRewriter {
|
||||
/// Whether the whole plan this rewriter walks can be encoded to Substrait.
|
||||
/// Used by [`PlanRewriter::should_expand`] to skip the per-node encoding check.
|
||||
///
|
||||
/// Used for the ordinary standalone fast path when the root plan is encodable. It is not
|
||||
/// needed when the root itself is the in-process `FusedSeries` exception.
|
||||
whole_plan_encodable: bool,
|
||||
/// In-process transport skips the encodability gate only when the rewritten root itself is FusedSeries.
|
||||
skip_encodability_gate: bool,
|
||||
/// Current level in the tree
|
||||
level: usize,
|
||||
/// Simulated stack for the `rewrite` recursion
|
||||
@@ -456,14 +505,28 @@ struct PlanRewriter {
|
||||
new_child_plan: Option<LogicalPlan>,
|
||||
}
|
||||
|
||||
fn is_fused_series_plan(plan: &LogicalPlan) -> bool {
|
||||
matches!(plan, LogicalPlan::Extension(extension) if extension.node.as_any().is::<FusedSeries>())
|
||||
}
|
||||
|
||||
impl PlanRewriter {
|
||||
/// `plan` must be the root of the tree that is about to be rewritten, otherwise
|
||||
/// [`PlanRewriter::should_expand`] may skip a check it has to perform.
|
||||
fn new(plan: &LogicalPlan) -> Self {
|
||||
fn new(plan: &LogicalPlan, in_process_region_transport: bool) -> Self {
|
||||
let native_root = match plan {
|
||||
LogicalPlan::Analyze(analyze) => analyze.input.as_ref(),
|
||||
_ => plan,
|
||||
};
|
||||
let skip_encodability_gate =
|
||||
in_process_region_transport && is_fused_series_plan(native_root);
|
||||
Self {
|
||||
whole_plan_encodable: DFLogicalSubstraitConvertor
|
||||
.encode(plan, DefaultSerializer)
|
||||
.is_ok(),
|
||||
// Preserve the ordinary whole-plan fast path unless this is the in-process root
|
||||
// exception: no node in that tree needs an encoding check.
|
||||
skip_encodability_gate,
|
||||
whole_plan_encodable: !skip_encodability_gate
|
||||
&& DFLogicalSubstraitConvertor
|
||||
.encode(plan, DefaultSerializer)
|
||||
.is_ok(),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
@@ -495,7 +558,15 @@ impl PlanRewriter {
|
||||
// is the dominant cost on deep PromQL plans.
|
||||
// When the root does not encode, the plan holds at least one node that has to stay
|
||||
// on the frontend and only the per-node check locates it.
|
||||
if !self.whole_plan_encodable
|
||||
//
|
||||
// For the in-process root `FusedSeries` exception, encodability says nothing about where
|
||||
// a node may run, and asking it would cut the plan at the first node with no serializer
|
||||
// branch before `check_plan` could decide. Skipping it only removes a reason to expand: every
|
||||
// check below still runs, and a node they call `NonCommutative`/`Unimplemented`/
|
||||
// `Unsupported` still expands, so a node no serializer lists only rides to a region when
|
||||
// the categorizer places it there.
|
||||
if !self.skip_encodability_gate
|
||||
&& !self.whole_plan_encodable
|
||||
&& let Err(e) = DFLogicalSubstraitConvertor.encode(plan, DefaultSerializer)
|
||||
{
|
||||
debug!(
|
||||
|
||||
@@ -68,6 +68,38 @@ fn collect_merge_scan_remote_dyn_filter_producer_ids(
|
||||
producer_ids.extend(producer_id_list);
|
||||
}
|
||||
|
||||
/// The analyzer as an engine whose regions are reached over the wire builds it: every plan a region
|
||||
/// receives from it has to be encodable.
|
||||
fn analyzer() -> DistPlannerAnalyzer {
|
||||
DistPlannerAnalyzer::default()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ordinary_encodable_plan_keeps_in_process_encodability_gate() {
|
||||
let table = TestTable::table_with_name(0, "t".to_string());
|
||||
let table_source = Arc::new(DefaultTableSource::new(Arc::new(
|
||||
DfTableProviderAdapter::new(table),
|
||||
)));
|
||||
let plan = LogicalPlanBuilder::scan_with_filters("t", table_source, None, vec![])
|
||||
.unwrap()
|
||||
.build()
|
||||
.unwrap();
|
||||
|
||||
assert!(
|
||||
DFLogicalSubstraitConvertor
|
||||
.encode(&plan, DefaultSerializer)
|
||||
.is_ok()
|
||||
);
|
||||
|
||||
let in_process_rewriter = PlanRewriter::new(&plan, true);
|
||||
assert!(in_process_rewriter.whole_plan_encodable);
|
||||
assert!(!in_process_rewriter.skip_encodability_gate);
|
||||
|
||||
let wire_rewriter = PlanRewriter::new(&plan, false);
|
||||
assert!(wire_rewriter.whole_plan_encodable);
|
||||
assert!(!wire_rewriter.skip_encodability_gate);
|
||||
}
|
||||
|
||||
struct MergeScanRemoteDynFilterProducerIdCollector<'a> {
|
||||
producer_ids: &'a mut Vec<RemoteDynFilterProducerId>,
|
||||
}
|
||||
@@ -395,24 +427,29 @@ fn frontend_only_histogram_folds_stay_above_merge_scan() {
|
||||
"{name}"
|
||||
);
|
||||
|
||||
let result = DistPlannerAnalyzer {}
|
||||
.analyze(plan, &ConfigOptions::default())
|
||||
.unwrap();
|
||||
let result_text = result.to_string();
|
||||
assert!(
|
||||
result_text.contains("HistogramFold:"),
|
||||
"{name}\n{result_text}"
|
||||
);
|
||||
for analyzer in [
|
||||
DistPlannerAnalyzer::default(),
|
||||
DistPlannerAnalyzer::new(true),
|
||||
] {
|
||||
let result = analyzer
|
||||
.analyze(plan.clone(), &ConfigOptions::default())
|
||||
.unwrap();
|
||||
let result_text = result.to_string();
|
||||
assert!(
|
||||
result_text.contains("HistogramFold:"),
|
||||
"{name}\n{result_text}"
|
||||
);
|
||||
|
||||
let remote_input = find_merge_scan(&result).unwrap().input();
|
||||
let remote_text = remote_input.to_string();
|
||||
assert!(
|
||||
!remote_text.contains("HistogramFold:"),
|
||||
"{name}\n{remote_text}"
|
||||
);
|
||||
DFLogicalSubstraitConvertor
|
||||
.encode(remote_input, DefaultSerializer)
|
||||
.unwrap();
|
||||
let remote_input = find_merge_scan(&result).unwrap().input();
|
||||
let remote_text = remote_input.to_string();
|
||||
assert!(
|
||||
!remote_text.contains("HistogramFold:"),
|
||||
"{name}\n{remote_text}"
|
||||
);
|
||||
DFLogicalSubstraitConvertor
|
||||
.encode(remote_input, DefaultSerializer)
|
||||
.unwrap();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -442,19 +479,24 @@ fn unencodable_unnest_stays_above_merge_scan() {
|
||||
.is_err()
|
||||
);
|
||||
|
||||
let result = DistPlannerAnalyzer {}
|
||||
.analyze(plan, &ConfigOptions::default())
|
||||
.unwrap();
|
||||
let result_text = result.to_string();
|
||||
assert!(result_text.contains("Unnest:"), "{result_text}");
|
||||
for analyzer in [
|
||||
DistPlannerAnalyzer::default(),
|
||||
DistPlannerAnalyzer::new(true),
|
||||
] {
|
||||
let result = analyzer
|
||||
.analyze(plan.clone(), &ConfigOptions::default())
|
||||
.unwrap();
|
||||
let result_text = result.to_string();
|
||||
assert!(result_text.contains("Unnest:"), "{result_text}");
|
||||
|
||||
let remote_input = find_merge_scan(&result).unwrap().input();
|
||||
let remote_text = remote_input.to_string();
|
||||
assert!(!remote_text.contains("Unnest:"), "{remote_text}");
|
||||
assert!(remote_text.contains("make_array"), "{remote_text}");
|
||||
DFLogicalSubstraitConvertor
|
||||
.encode(remote_input, DefaultSerializer)
|
||||
.unwrap();
|
||||
let remote_input = find_merge_scan(&result).unwrap().input();
|
||||
let remote_text = remote_input.to_string();
|
||||
assert!(!remote_text.contains("Unnest:"), "{remote_text}");
|
||||
assert!(remote_text.contains("make_array"), "{remote_text}");
|
||||
DFLogicalSubstraitConvertor
|
||||
.encode(remote_input, DefaultSerializer)
|
||||
.unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -481,9 +523,7 @@ fn dictionary_literals_are_unwrapped_before_pushdown() {
|
||||
.build()
|
||||
.unwrap();
|
||||
|
||||
let result = DistPlannerAnalyzer {}
|
||||
.analyze(plan, &ConfigOptions::default())
|
||||
.unwrap();
|
||||
let result = analyzer().analyze(plan, &ConfigOptions::default()).unwrap();
|
||||
assert!(
|
||||
result
|
||||
.to_string()
|
||||
@@ -578,7 +618,7 @@ fn expand_proj_sort_proj() {
|
||||
.unwrap();
|
||||
|
||||
let config = ConfigOptions::default();
|
||||
let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap();
|
||||
let result = analyzer().analyze(plan, &config).unwrap();
|
||||
|
||||
let expected = [
|
||||
"Projection: t.number",
|
||||
@@ -625,7 +665,7 @@ fn expand_proj_sort_partial_proj() {
|
||||
.unwrap();
|
||||
|
||||
let config = ConfigOptions::default();
|
||||
let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap();
|
||||
let result = analyzer().analyze(plan, &config).unwrap();
|
||||
|
||||
let expected = [
|
||||
"Projection: t.number",
|
||||
@@ -661,7 +701,7 @@ fn expand_sort_limit() {
|
||||
.unwrap();
|
||||
|
||||
let config = ConfigOptions::default();
|
||||
let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap();
|
||||
let result = analyzer().analyze(plan, &config).unwrap();
|
||||
|
||||
let expected = [
|
||||
"Projection: t.pk1, t.pk2, t.pk3, t.ts, t.number",
|
||||
@@ -699,7 +739,7 @@ fn expand_sort_alias_limit() {
|
||||
.unwrap();
|
||||
|
||||
let config = ConfigOptions::default();
|
||||
let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap();
|
||||
let result = analyzer().analyze(plan, &config).unwrap();
|
||||
|
||||
let expected = [
|
||||
"Projection: something",
|
||||
@@ -738,7 +778,7 @@ fn expand_sort_alias_conflict_limit() {
|
||||
.unwrap();
|
||||
|
||||
let config = ConfigOptions::default();
|
||||
let result = DistPlannerAnalyzer {}.analyze(plan.clone(), &config);
|
||||
let result = analyzer().analyze(plan.clone(), &config);
|
||||
assert!(result.is_err(), "Expected error for ambiguous alias");
|
||||
assert!(format!("{result:?}").contains("AmbiguousReference"));
|
||||
|
||||
@@ -746,7 +786,7 @@ fn expand_sort_alias_conflict_limit() {
|
||||
config.extensions.insert(DistPlannerOptions {
|
||||
allow_query_fallback: true,
|
||||
});
|
||||
let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap();
|
||||
let result = analyzer().analyze(plan, &config).unwrap();
|
||||
|
||||
let expected = [
|
||||
"Limit: skip=0, fetch=10",
|
||||
@@ -780,7 +820,7 @@ fn expand_sort_alias_conflict_but_not_really_limit() {
|
||||
.unwrap();
|
||||
|
||||
let config = ConfigOptions::default();
|
||||
let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap();
|
||||
let result = analyzer().analyze(plan, &config).unwrap();
|
||||
|
||||
let expected = [
|
||||
"Projection: t.pk1",
|
||||
@@ -818,7 +858,7 @@ fn expand_limit_sort() {
|
||||
.unwrap();
|
||||
|
||||
let config = ConfigOptions::default();
|
||||
let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap();
|
||||
let result = analyzer().analyze(plan, &config).unwrap();
|
||||
|
||||
let expected = [
|
||||
"Sort: t.pk1 ASC NULLS LAST",
|
||||
@@ -853,7 +893,7 @@ fn expand_sort_limit_sort() {
|
||||
.unwrap();
|
||||
|
||||
let config = ConfigOptions::default();
|
||||
let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap();
|
||||
let result = analyzer().analyze(plan, &config).unwrap();
|
||||
|
||||
let expected = [
|
||||
"Sort: t.pk1 ASC NULLS LAST",
|
||||
@@ -898,7 +938,7 @@ fn expand_proj_step_aggr() {
|
||||
.unwrap();
|
||||
|
||||
let config = ConfigOptions::default();
|
||||
let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap();
|
||||
let result = analyzer().analyze(plan, &config).unwrap();
|
||||
|
||||
let expected = [
|
||||
"Projection: min(t.number)",
|
||||
@@ -943,7 +983,7 @@ fn expand_proj_alias_fake_part_col_aggr() {
|
||||
.unwrap();
|
||||
|
||||
let config = ConfigOptions::default();
|
||||
let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap();
|
||||
let result = analyzer().analyze(plan, &config).unwrap();
|
||||
|
||||
let expected = [
|
||||
"Projection: pk1, pk2, min(t.number)",
|
||||
@@ -987,7 +1027,7 @@ fn expand_proj_alias_aliased_part_col_aggr() {
|
||||
.unwrap();
|
||||
|
||||
let config = ConfigOptions::default();
|
||||
let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap();
|
||||
let result = analyzer().analyze(plan, &config).unwrap();
|
||||
|
||||
let expected = [
|
||||
"Projection: pk42, pk43, min(t.number)",
|
||||
@@ -1022,7 +1062,7 @@ fn expand_expr_over_part_col_aggr() {
|
||||
.unwrap();
|
||||
|
||||
let config = ConfigOptions::default();
|
||||
let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap();
|
||||
let result = analyzer().analyze(plan, &config).unwrap();
|
||||
|
||||
let expected = [
|
||||
"Projection: substr(t.pk1,Int64(1),Int64(1)), t.pk2, min(t.number)",
|
||||
@@ -1057,7 +1097,7 @@ fn expand_part_col_aggr_step_aggr() {
|
||||
.unwrap();
|
||||
|
||||
let config = ConfigOptions::default();
|
||||
let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap();
|
||||
let result = analyzer().analyze(plan, &config).unwrap();
|
||||
|
||||
let expected = [
|
||||
"Projection: min(max(t.number))",
|
||||
@@ -1090,7 +1130,7 @@ fn expand_step_aggr_step_aggr() {
|
||||
.unwrap();
|
||||
|
||||
let config = ConfigOptions::default();
|
||||
let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap();
|
||||
let result = analyzer().analyze(plan, &config).unwrap();
|
||||
|
||||
let expected = [
|
||||
"Aggregate: groupBy=[[]], aggr=[[min(max(t.number))]]",
|
||||
@@ -1135,7 +1175,7 @@ fn expand_part_col_aggr_part_col_aggr() {
|
||||
assert_eq!(expected_original, plan.to_string());
|
||||
|
||||
let config = ConfigOptions::default();
|
||||
let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap();
|
||||
let result = analyzer().analyze(plan, &config).unwrap();
|
||||
|
||||
let expected = [
|
||||
"Projection: t.pk1, t.pk2, min(max(t.number))",
|
||||
@@ -1167,7 +1207,7 @@ fn expand_step_aggr_proj() {
|
||||
.unwrap();
|
||||
|
||||
let config = ConfigOptions::default();
|
||||
let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap();
|
||||
let result = analyzer().analyze(plan, &config).unwrap();
|
||||
|
||||
let expected = [
|
||||
"Projection: min(t.number)",
|
||||
@@ -1225,7 +1265,7 @@ fn expand_complex_col_req_sort_pql() {
|
||||
.unwrap();
|
||||
|
||||
let config = ConfigOptions::default();
|
||||
let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap();
|
||||
let result = analyzer().analyze(plan, &config).unwrap();
|
||||
|
||||
let expected = [
|
||||
"Projection: t.pk1, t.pk2",
|
||||
@@ -1271,7 +1311,7 @@ fn expand_proj_sort_step_aggr_limit() {
|
||||
.unwrap();
|
||||
|
||||
let config = ConfigOptions::default();
|
||||
let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap();
|
||||
let result = analyzer().analyze(plan, &config).unwrap();
|
||||
let expected = [
|
||||
"Limit: skip=0, fetch=10",
|
||||
" Aggregate: groupBy=[[]], aggr=[[min(t.number)]]",
|
||||
@@ -1312,7 +1352,7 @@ fn expand_proj_sort_limit_step_aggr() {
|
||||
.unwrap();
|
||||
|
||||
let config = ConfigOptions::default();
|
||||
let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap();
|
||||
let result = analyzer().analyze(plan, &config).unwrap();
|
||||
let expected = [
|
||||
"Aggregate: groupBy=[[]], aggr=[[min(t.number)]]",
|
||||
" Projection: t.number",
|
||||
@@ -1354,7 +1394,7 @@ fn expand_proj_limit_step_aggr_sort() {
|
||||
.unwrap();
|
||||
|
||||
let config = ConfigOptions::default();
|
||||
let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap();
|
||||
let result = analyzer().analyze(plan, &config).unwrap();
|
||||
let expected = [
|
||||
"Sort: min(t.number) ASC NULLS LAST",
|
||||
" Aggregate: groupBy=[[]], aggr=[[min(t.number)]]",
|
||||
@@ -1396,7 +1436,7 @@ fn expand_proj_sort_part_col_aggr_limit() {
|
||||
.unwrap();
|
||||
|
||||
let config = ConfigOptions::default();
|
||||
let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap();
|
||||
let result = analyzer().analyze(plan, &config).unwrap();
|
||||
let expected = [
|
||||
"Limit: skip=0, fetch=10",
|
||||
" Aggregate: groupBy=[[t.pk1, t.pk2]], aggr=[[min(t.number)]]",
|
||||
@@ -1438,7 +1478,7 @@ fn expand_proj_sort_limit_part_col_aggr() {
|
||||
.unwrap();
|
||||
|
||||
let config = ConfigOptions::default();
|
||||
let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap();
|
||||
let result = analyzer().analyze(plan, &config).unwrap();
|
||||
let expected = [
|
||||
"Aggregate: groupBy=[[t.pk1, t.pk2]], aggr=[[min(t.number)]]",
|
||||
" Projection: t.number, t.pk1, t.pk2",
|
||||
@@ -1480,7 +1520,7 @@ fn expand_proj_part_col_aggr_limit_sort() {
|
||||
.unwrap();
|
||||
|
||||
let config = ConfigOptions::default();
|
||||
let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap();
|
||||
let result = analyzer().analyze(plan, &config).unwrap();
|
||||
let expected = [
|
||||
"Sort: t.pk2 ASC NULLS LAST",
|
||||
" Projection: t.pk1, t.pk2, min(t.number)",
|
||||
@@ -1522,7 +1562,7 @@ fn expand_proj_part_col_aggr_sort_limit() {
|
||||
.unwrap();
|
||||
|
||||
let config = ConfigOptions::default();
|
||||
let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap();
|
||||
let result = analyzer().analyze(plan, &config).unwrap();
|
||||
let expected = [
|
||||
"Projection: t.pk1, t.pk2, min(t.number)",
|
||||
" Limit: skip=0, fetch=10",
|
||||
@@ -1557,9 +1597,7 @@ fn test_simplify_select_now_expression() {
|
||||
.unwrap();
|
||||
|
||||
let config = ConfigOptions::default();
|
||||
let result = DistPlannerAnalyzer {}
|
||||
.analyze(plan.clone(), &config)
|
||||
.unwrap();
|
||||
let result = analyzer().analyze(plan.clone(), &config).unwrap();
|
||||
|
||||
common_telemetry::info!("Analyzed plan: {}", result);
|
||||
|
||||
@@ -1615,7 +1653,7 @@ fn sibling_merge_scans_have_unique_remote_dyn_filter_producer_ids() {
|
||||
.unwrap();
|
||||
|
||||
let config = ConfigOptions::default();
|
||||
let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap();
|
||||
let result = analyzer().analyze(plan, &config).unwrap();
|
||||
|
||||
let mut producer_ids = Vec::new();
|
||||
collect_merge_scan_remote_dyn_filter_producer_id_list(&result, &mut producer_ids);
|
||||
@@ -1670,7 +1708,7 @@ fn pre_merge_scan_optimizer_eliminates_projected_false_filter() {
|
||||
.unwrap();
|
||||
|
||||
let config = ConfigOptions::default();
|
||||
let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap();
|
||||
let result = analyzer().analyze(plan, &config).unwrap();
|
||||
|
||||
assert_eq!("EmptyRelation: rows=0", result.to_string());
|
||||
}
|
||||
@@ -1718,7 +1756,7 @@ fn test_simplify_now_expression() {
|
||||
.unwrap();
|
||||
|
||||
let config = ConfigOptions::default();
|
||||
let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap();
|
||||
let result = analyzer().analyze(plan, &config).unwrap();
|
||||
|
||||
let plan_str = result.to_string();
|
||||
common_telemetry::info!("Analyzed plan: {}", plan_str);
|
||||
@@ -1757,7 +1795,7 @@ fn expand_proj_limit_part_col_aggr_sort() {
|
||||
.unwrap();
|
||||
|
||||
let config = ConfigOptions::default();
|
||||
let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap();
|
||||
let result = analyzer().analyze(plan, &config).unwrap();
|
||||
// Pre-MergeScan optimizer intentionally excludes PushDownLimit, so the
|
||||
// remote plan shows an explicit Limit node instead of `fetch=10` on
|
||||
// TableScan.
|
||||
@@ -1802,7 +1840,7 @@ fn expand_proj_limit_sort_part_col_aggr() {
|
||||
.unwrap();
|
||||
|
||||
let config = ConfigOptions::default();
|
||||
let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap();
|
||||
let result = analyzer().analyze(plan, &config).unwrap();
|
||||
// Pre-MergeScan optimizer intentionally excludes PushDownLimit, so the
|
||||
// remote plan shows an explicit Limit node instead of `fetch=10` on
|
||||
// TableScan.
|
||||
@@ -1840,7 +1878,7 @@ fn expand_step_aggr_limit() {
|
||||
.unwrap();
|
||||
|
||||
let config = ConfigOptions::default();
|
||||
let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap();
|
||||
let result = analyzer().analyze(plan, &config).unwrap();
|
||||
let expected = [
|
||||
"Limit: skip=0, fetch=10",
|
||||
" Projection: t.pk1, min(t.number)",
|
||||
@@ -1873,7 +1911,7 @@ fn expand_step_aggr_avg_limit() {
|
||||
.unwrap();
|
||||
|
||||
let config = ConfigOptions::default();
|
||||
let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap();
|
||||
let result = analyzer().analyze(plan, &config).unwrap();
|
||||
let expected = [
|
||||
"Limit: skip=0, fetch=10",
|
||||
" Projection: t.pk1, avg(t.number)",
|
||||
@@ -1906,7 +1944,7 @@ fn expand_part_col_aggr_limit() {
|
||||
.unwrap();
|
||||
|
||||
let config = ConfigOptions::default();
|
||||
let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap();
|
||||
let result = analyzer().analyze(plan, &config).unwrap();
|
||||
let expected = [
|
||||
"Projection: t.pk1, t.pk2, min(t.number)",
|
||||
" Limit: skip=0, fetch=10",
|
||||
@@ -1940,7 +1978,7 @@ fn transform_simple_projection_filter() {
|
||||
.unwrap();
|
||||
|
||||
let config = ConfigOptions::default();
|
||||
let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap();
|
||||
let result = analyzer().analyze(plan, &config).unwrap();
|
||||
let expected = [
|
||||
"Distinct:",
|
||||
" MergeScan [is_placeholder=false]",
|
||||
@@ -1968,7 +2006,7 @@ fn transform_aggregator() {
|
||||
.unwrap();
|
||||
|
||||
let config = ConfigOptions::default();
|
||||
let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap();
|
||||
let result = analyzer().analyze(plan, &config).unwrap();
|
||||
let expected = "Projection: avg(t.number)\
|
||||
\n MergeScan [is_placeholder=false, remote_input=[\
|
||||
\nAggregate: groupBy=[[]], aggr=[[avg(t.number)]]\
|
||||
@@ -1994,7 +2032,7 @@ fn transform_distinct_order() {
|
||||
.unwrap();
|
||||
|
||||
let config = ConfigOptions::default();
|
||||
let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap();
|
||||
let result = analyzer().analyze(plan, &config).unwrap();
|
||||
let expected = [
|
||||
"Projection: t.number",
|
||||
" MergeScan [is_placeholder=false, remote_input=[
|
||||
@@ -2022,7 +2060,7 @@ fn transform_single_limit() {
|
||||
.unwrap();
|
||||
|
||||
let config = ConfigOptions::default();
|
||||
let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap();
|
||||
let result = analyzer().analyze(plan, &config).unwrap();
|
||||
let expected = "Projection: t.number\
|
||||
\n MergeScan [is_placeholder=false, remote_input=[
|
||||
Limit: skip=0, fetch=1
|
||||
@@ -2063,7 +2101,7 @@ fn transform_unalighed_join_with_alias() {
|
||||
.unwrap();
|
||||
|
||||
let config = ConfigOptions::default();
|
||||
let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap();
|
||||
let result = analyzer().analyze(plan, &config).unwrap();
|
||||
let expected = [
|
||||
"Limit: skip=0, fetch=1",
|
||||
" LeftSemi Join: t.number = right.number",
|
||||
@@ -2099,7 +2137,7 @@ fn transform_subquery_sort_alias() {
|
||||
.build()
|
||||
.unwrap();
|
||||
let config = ConfigOptions::default();
|
||||
let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap();
|
||||
let result = analyzer().analyze(plan, &config).unwrap();
|
||||
let expected = [
|
||||
"Projection: a.pk1, a.pk2, a.pk3, a.ts, a.number",
|
||||
" MergeSort: a.number ASC NULLS LAST",
|
||||
@@ -2130,7 +2168,7 @@ fn transform_sort_subquery_alias() {
|
||||
.build()
|
||||
.unwrap();
|
||||
let config = ConfigOptions::default();
|
||||
let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap();
|
||||
let result = analyzer().analyze(plan, &config).unwrap();
|
||||
let expected = [
|
||||
"Projection: a.pk1, a.pk2, a.pk3, a.ts, a.number",
|
||||
" MergeSort: a.number ASC NULLS LAST",
|
||||
@@ -2169,7 +2207,7 @@ fn remote_dyn_filter_producer_ids_do_not_collide_between_subquery_and_outer_plan
|
||||
.unwrap()
|
||||
.build()
|
||||
.unwrap();
|
||||
let rewritten = DistPlannerAnalyzer {}.try_push_down(outer_plan).unwrap();
|
||||
let rewritten = analyzer().try_push_down(outer_plan).unwrap();
|
||||
|
||||
let mut producer_ids = BTreeSet::new();
|
||||
collect_merge_scan_remote_dyn_filter_producer_ids(&rewritten, &mut producer_ids);
|
||||
@@ -2201,7 +2239,7 @@ fn date_bin_ts_group_by() {
|
||||
.unwrap();
|
||||
|
||||
let config = ConfigOptions::default();
|
||||
let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap();
|
||||
let result = analyzer().analyze(plan, &config).unwrap();
|
||||
|
||||
let expected = [
|
||||
r#"Projection: date_bin(IntervalDayTime("IntervalDayTime { days: 0, milliseconds: 60000 }"),t.ts), min(t.number)"#,
|
||||
@@ -2249,9 +2287,7 @@ fn test_last_value_order_by() {
|
||||
try_encode_decode_substrait(&plan, ctx.state());
|
||||
|
||||
let config = ConfigOptions::default();
|
||||
let result = DistPlannerAnalyzer {}
|
||||
.analyze(plan.clone(), &config)
|
||||
.unwrap();
|
||||
let result = analyzer().analyze(plan.clone(), &config).unwrap();
|
||||
|
||||
let expected = [
|
||||
"Projection: last_value(t.ts) ORDER BY [t.ts ASC NULLS FIRST]",
|
||||
@@ -2312,9 +2348,7 @@ fn test_last_value_no_order_by() {
|
||||
try_encode_decode_substrait(&split.lower_state, ctx.state());
|
||||
|
||||
let config = ConfigOptions::default();
|
||||
let result = DistPlannerAnalyzer {}
|
||||
.analyze(plan.clone(), &config)
|
||||
.unwrap();
|
||||
let result = analyzer().analyze(plan.clone(), &config).unwrap();
|
||||
|
||||
let expected = [
|
||||
"Projection: last_value(t.ts)",
|
||||
@@ -2344,9 +2378,7 @@ fn test_table_scan_projection() {
|
||||
.unwrap();
|
||||
|
||||
let config = ConfigOptions::default();
|
||||
let result = DistPlannerAnalyzer {}
|
||||
.analyze(plan.clone(), &config)
|
||||
.unwrap();
|
||||
let result = analyzer().analyze(plan.clone(), &config).unwrap();
|
||||
let expected = [
|
||||
"Projection: t.ts",
|
||||
" MergeScan [is_placeholder=false, remote_input=[",
|
||||
@@ -2380,9 +2412,7 @@ fn test_table_scan_cast_projection_pushdown() {
|
||||
.unwrap();
|
||||
|
||||
let config = ConfigOptions::default();
|
||||
let result = DistPlannerAnalyzer {}
|
||||
.analyze(plan.clone(), &config)
|
||||
.unwrap();
|
||||
let result = analyzer().analyze(plan.clone(), &config).unwrap();
|
||||
let expected = [
|
||||
"Projection: ts",
|
||||
" MergeScan [is_placeholder=false, remote_input=[",
|
||||
@@ -2426,9 +2456,7 @@ fn test_cast_filter_simplified_after_pushdown() {
|
||||
.unwrap();
|
||||
|
||||
let config = ConfigOptions::default();
|
||||
let result = DistPlannerAnalyzer {}
|
||||
.analyze(plan.clone(), &config)
|
||||
.unwrap();
|
||||
let result = analyzer().analyze(plan.clone(), &config).unwrap();
|
||||
let expected = [
|
||||
"Projection: ts",
|
||||
" MergeScan [is_placeholder=false, remote_input=[",
|
||||
@@ -2483,7 +2511,7 @@ fn test_join_side_local_filter_pushdown_into_merge_scan() {
|
||||
.unwrap();
|
||||
|
||||
let config = ConfigOptions::default();
|
||||
let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap();
|
||||
let result = analyzer().analyze(plan, &config).unwrap();
|
||||
assert_remote_table_scan_filters_are_safe(&result);
|
||||
|
||||
let predicate = col("t1.pk1").eq(lit("v"));
|
||||
@@ -2555,7 +2583,7 @@ fn test_left_join_left_side_filter_pushdown_into_merge_scan() {
|
||||
.unwrap();
|
||||
|
||||
let config = ConfigOptions::default();
|
||||
let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap();
|
||||
let result = analyzer().analyze(plan, &config).unwrap();
|
||||
assert_remote_table_scan_filters_are_safe(&result);
|
||||
|
||||
let plan_str = result.to_string();
|
||||
@@ -2611,7 +2639,7 @@ fn test_join_cross_table_predicate_not_pushed_to_single_side() {
|
||||
.unwrap();
|
||||
|
||||
let config = ConfigOptions::default();
|
||||
let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap();
|
||||
let result = analyzer().analyze(plan, &config).unwrap();
|
||||
assert_remote_table_scan_filters_are_safe(&result);
|
||||
|
||||
let plan_str = result.to_string();
|
||||
@@ -2668,9 +2696,7 @@ fn scheduled_now_yields_stable_literal_in_remote_plan() {
|
||||
scheduled_time: Some(scheduled_time),
|
||||
});
|
||||
|
||||
let result = DistPlannerAnalyzer {}
|
||||
.analyze(plan.clone(), &config)
|
||||
.unwrap();
|
||||
let result = analyzer().analyze(plan.clone(), &config).unwrap();
|
||||
let result_str = result.to_string();
|
||||
common_telemetry::info!("Analyzed plan with scheduled time: {}", result_str);
|
||||
|
||||
@@ -2706,9 +2732,7 @@ fn scheduled_now_yields_stable_literal_in_remote_plan() {
|
||||
// Absent the extension (default config), the same plan simplifies to a
|
||||
// wall-clock literal — the remote_input contains a variable timestamp.
|
||||
let default_config = ConfigOptions::default();
|
||||
let wall_result = DistPlannerAnalyzer {}
|
||||
.analyze(plan, &default_config)
|
||||
.unwrap();
|
||||
let wall_result = analyzer().analyze(plan, &default_config).unwrap();
|
||||
let wall_str = wall_result.to_string();
|
||||
// The wall-clock result must NOT contain the scheduled literal.
|
||||
assert!(
|
||||
@@ -2769,7 +2793,7 @@ fn scheduled_current_timestamp_filter_folds_before_remote_pushdown() {
|
||||
scheduled_time: Some(scheduled_time),
|
||||
});
|
||||
|
||||
let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap();
|
||||
let result = analyzer().analyze(plan, &config).unwrap();
|
||||
assert_remote_table_scan_filters_are_safe(&result);
|
||||
|
||||
let result_str = result.to_string();
|
||||
@@ -2816,7 +2840,7 @@ fn scheduled_none_falls_back_to_wall_clock() {
|
||||
scheduled_time: None,
|
||||
});
|
||||
|
||||
let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap();
|
||||
let result = analyzer().analyze(plan, &config).unwrap();
|
||||
let result_str = result.to_string();
|
||||
|
||||
// Must still simplify now() — remotely a literal, top is now().
|
||||
@@ -2884,7 +2908,7 @@ fn expand_nested_scalar_subquery() {
|
||||
.unwrap();
|
||||
|
||||
let config = ConfigOptions::default();
|
||||
let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap();
|
||||
let result = analyzer().analyze(plan, &config).unwrap();
|
||||
let result_str = result.to_string();
|
||||
|
||||
// The main scan and both subquery levels must each be wrapped in a MergeScan.
|
||||
|
||||
@@ -21,8 +21,9 @@ use datafusion::error::Result as DfResult;
|
||||
use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion};
|
||||
use datafusion_expr::{Expr, LogicalPlan, UserDefinedLogicalNode};
|
||||
use promql::extension_plan::{
|
||||
EmptyMetric, InstantManipulate, RangeManipulate, SeriesDivide, SeriesNormalize,
|
||||
EmptyMetric, FusedSeries, InstantManipulate, RangeManipulate, SeriesDivide, SeriesNormalize,
|
||||
};
|
||||
use promql::native::plan::reads_across_series;
|
||||
use store_api::metric_engine_consts::DATA_SCHEMA_TSID_COLUMN_NAME;
|
||||
|
||||
use crate::dist_plan::MergeScanLogicalPlan;
|
||||
@@ -240,6 +241,38 @@ impl Categorizer {
|
||||
|
||||
Commutativity::Commutative
|
||||
}
|
||||
name if name == FusedSeries::name() => {
|
||||
let fused = plan.as_any().downcast_ref::<FusedSeries>().unwrap();
|
||||
// A region builds this node from the region scan itself: the conversion takes that
|
||||
// scan's raw-series entrypoint and cannot look through anything in between, so a
|
||||
// node whose input is not that scan has nothing to read where it would be placed
|
||||
// and stays where it was planned.
|
||||
if !matches!(fused.inputs().as_slice(), [LogicalPlan::TableScan(_)]) {
|
||||
return Commutativity::NonCommutative;
|
||||
}
|
||||
// The region asks the scan for the series of one named metric, so a node that binds
|
||||
// no name cannot be built there at all.
|
||||
if fused.binding().metric_name.is_none() {
|
||||
return Commutativity::NonCommutative;
|
||||
}
|
||||
// A region evaluates the tree over whatever series its own scan returns, so this
|
||||
// is only answerable when every node of the tree reads one series at a time. An
|
||||
// aggregate or a binary operation would fold or pair across series, and one region
|
||||
// holds only part of them.
|
||||
if reads_across_series(fused.tree().root()) {
|
||||
return Commutativity::NonCommutative;
|
||||
}
|
||||
|
||||
// A region evaluates the whole tree over the series its scan returns, so this is
|
||||
// only correct while each series reaches exactly one region: a series split
|
||||
// across regions would be ranged per fragment, and a range over fragments is not
|
||||
// the range over the series. Native admission
|
||||
// (`promql::planner::native::try_plan_native_fused`) only builds the node when
|
||||
// every partition column of the table is one of the binding's tag columns, so
|
||||
// the partition key is a function of the label set and a series cannot straddle
|
||||
// regions.
|
||||
Commutativity::Commutative
|
||||
}
|
||||
name if name == SeriesNormalize::name()
|
||||
|| name == InstantManipulate::name()
|
||||
|| name == RangeManipulate::name() =>
|
||||
@@ -349,16 +382,39 @@ impl Categorizer {
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
use std::time::{Duration, UNIX_EPOCH};
|
||||
|
||||
use datafusion::arrow::datatypes::{
|
||||
DataType as ArrowDataType, Field, TimeUnit as ArrowTimeUnit,
|
||||
};
|
||||
use datafusion::common::{DFSchema, DFSchemaRef};
|
||||
use datafusion::datasource::DefaultTableSource;
|
||||
use datafusion_common::Column;
|
||||
use datafusion_common::config::ConfigOptions;
|
||||
use datafusion_expr::LogicalPlanBuilder;
|
||||
use datafusion_expr::expr::ScalarFunction;
|
||||
use datafusion_expr::{Extension, LogicalPlan, UserDefinedLogicalNodeCore};
|
||||
use datafusion_functions::core::coalesce;
|
||||
use datafusion_optimizer::analyzer::AnalyzerRule;
|
||||
use datatypes::series::SeriesOrdering;
|
||||
use datatypes::series::samples::{FieldSpec, SeriesSchema};
|
||||
use promql::extension_plan::{FusedSeries, FusedSeriesBinding};
|
||||
use promql::functions::{
|
||||
NativeHistogramDrop, NativeHistogramFraction, NativeHistogramQuantile,
|
||||
};
|
||||
use promql::native::plan::{
|
||||
AggregateKind, AggregateStage, AggregateStateSchema, BoundAggregate, BoundInput, Grouping,
|
||||
NativeExpr, NativeNode, NativeType,
|
||||
};
|
||||
use promql_parser::parser;
|
||||
use store_api::storage::RegionId;
|
||||
use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan};
|
||||
|
||||
use super::*;
|
||||
use crate::dist_plan::DistPlannerAnalyzer;
|
||||
use crate::optimizer::test_util::mock_table_provider;
|
||||
use crate::promql::native_compile::{compile, statement_grid};
|
||||
use crate::query_engine::DefaultSerializer;
|
||||
|
||||
#[test]
|
||||
fn series_divide_by_tsid_is_commutative() {
|
||||
@@ -393,6 +449,388 @@ mod tests {
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
/// The statement one `FusedSeries` test node is bound to.
|
||||
fn native_statement() -> promql_parser::parser::EvalStmt {
|
||||
promql_parser::parser::EvalStmt {
|
||||
expr: parser::parse("some_metric").unwrap(),
|
||||
start: UNIX_EPOCH,
|
||||
end: UNIX_EPOCH
|
||||
.checked_add(Duration::from_secs(100_000))
|
||||
.unwrap(),
|
||||
interval: Duration::from_secs(5),
|
||||
lookback_delta: Duration::from_secs(1),
|
||||
}
|
||||
}
|
||||
|
||||
/// A `FusedSeries` node over `input`, bound to the tree a bare selector compiles to.
|
||||
///
|
||||
/// The tree is the simplest shape the compiler accepts, because what these tests change is
|
||||
/// where the node may be placed, not what it computes.
|
||||
fn fused_series_plan(input: LogicalPlan) -> LogicalPlan {
|
||||
fused_series_plan_named(input, Some("some_metric"))
|
||||
}
|
||||
|
||||
/// The same node, with the metric name it binds chosen by the caller.
|
||||
fn fused_series_plan_named(input: LogicalPlan, metric_name: Option<&str>) -> LogicalPlan {
|
||||
let stmt = native_statement();
|
||||
let series_schema = Arc::new(SeriesSchema {
|
||||
fields: Arc::from([FieldSpec { name: "v0".into() }]),
|
||||
eval_grid: None,
|
||||
});
|
||||
let tree = compile(
|
||||
&stmt,
|
||||
Arc::clone(&series_schema),
|
||||
SeriesOrdering::KeyThenTime,
|
||||
)
|
||||
.unwrap();
|
||||
let binding = FusedSeriesBinding {
|
||||
series_schema,
|
||||
grid: statement_grid(&stmt).unwrap(),
|
||||
metric_name: metric_name.map(Arc::from),
|
||||
time_index: "ts".to_string(),
|
||||
tags: Vec::new(),
|
||||
};
|
||||
let output_schema: DFSchemaRef = Arc::new(
|
||||
DFSchema::new_with_metadata(
|
||||
vec![
|
||||
(
|
||||
None,
|
||||
Arc::new(Field::new("v0", ArrowDataType::Float64, true)),
|
||||
),
|
||||
(
|
||||
None,
|
||||
Arc::new(Field::new(
|
||||
"ts",
|
||||
ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None),
|
||||
false,
|
||||
)),
|
||||
),
|
||||
],
|
||||
Default::default(),
|
||||
)
|
||||
.unwrap(),
|
||||
);
|
||||
|
||||
LogicalPlan::Extension(Extension {
|
||||
node: Arc::new(
|
||||
FusedSeries::try_new(
|
||||
Arc::from("some_metric"),
|
||||
tree,
|
||||
binding,
|
||||
input,
|
||||
output_schema,
|
||||
)
|
||||
.unwrap(),
|
||||
),
|
||||
})
|
||||
}
|
||||
|
||||
/// The region's own view of the metric's table: a scan whose provider is the region provider.
|
||||
fn table_scan_plan() -> LogicalPlan {
|
||||
let provider = mock_table_provider(RegionId::new(1, 1));
|
||||
LogicalPlanBuilder::scan(
|
||||
"some_metric",
|
||||
Arc::new(DefaultTableSource::new(Arc::new(provider))),
|
||||
None,
|
||||
)
|
||||
.unwrap()
|
||||
.build()
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
fn holds_fused_series(plan: &LogicalPlan) -> bool {
|
||||
if let LogicalPlan::Extension(extension) = plan
|
||||
&& extension.node.as_any().is::<FusedSeries>()
|
||||
{
|
||||
return true;
|
||||
}
|
||||
plan.inputs().into_iter().any(holds_fused_series)
|
||||
}
|
||||
|
||||
/// An extension node with no `ExtensionSerializer` branch: neither one under its own name nor
|
||||
/// one it shares, since the serializer lists names and not shapes. It stands in for the nodes
|
||||
/// that cannot reach a region as bytes for their own reasons, which the encoding gate must keep
|
||||
/// honouring. The categorizer refuses it, as it refuses every name it does not know.
|
||||
#[derive(Debug, PartialEq, Eq, Hash)]
|
||||
struct UnserializableNode {
|
||||
input: LogicalPlan,
|
||||
schema: DFSchemaRef,
|
||||
}
|
||||
|
||||
impl PartialOrd for UnserializableNode {
|
||||
fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
|
||||
let ordering = self.input.partial_cmp(&other.input);
|
||||
match ordering {
|
||||
Some(std::cmp::Ordering::Equal) if self.schema != other.schema => None,
|
||||
ordering => ordering,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl UnserializableNode {
|
||||
fn new(input: LogicalPlan) -> Self {
|
||||
let schema = input.schema().clone();
|
||||
Self { input, schema }
|
||||
}
|
||||
}
|
||||
|
||||
impl UserDefinedLogicalNodeCore for UnserializableNode {
|
||||
fn name(&self) -> &str {
|
||||
"UnserializableNode"
|
||||
}
|
||||
|
||||
fn inputs(&self) -> Vec<&LogicalPlan> {
|
||||
vec![&self.input]
|
||||
}
|
||||
|
||||
fn schema(&self) -> &DFSchemaRef {
|
||||
&self.schema
|
||||
}
|
||||
|
||||
fn expressions(&self) -> Vec<Expr> {
|
||||
vec![]
|
||||
}
|
||||
|
||||
fn fmt_for_explain(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
|
||||
write!(f, "UnserializableNode")
|
||||
}
|
||||
|
||||
fn with_exprs_and_inputs(
|
||||
&self,
|
||||
_exprs: Vec<Expr>,
|
||||
mut inputs: Vec<LogicalPlan>,
|
||||
) -> DfResult<Self> {
|
||||
Ok(Self::new(inputs.swap_remove(0)))
|
||||
}
|
||||
}
|
||||
|
||||
/// The plan that node builds over `input`.
|
||||
fn unserializable_plan(input: LogicalPlan) -> LogicalPlan {
|
||||
LogicalPlan::Extension(Extension {
|
||||
node: Arc::new(UnserializableNode::new(input)),
|
||||
})
|
||||
}
|
||||
|
||||
fn holds_unserializable_node(plan: &LogicalPlan) -> bool {
|
||||
if let LogicalPlan::Extension(extension) = plan
|
||||
&& extension.node.as_any().is::<UnserializableNode>()
|
||||
{
|
||||
return true;
|
||||
}
|
||||
plan.inputs().into_iter().any(holds_unserializable_node)
|
||||
}
|
||||
|
||||
fn holds_table_scan(plan: &LogicalPlan) -> bool {
|
||||
if matches!(plan, LogicalPlan::TableScan(_)) {
|
||||
return true;
|
||||
}
|
||||
plan.inputs().into_iter().any(holds_table_scan)
|
||||
}
|
||||
|
||||
/// The merge scan `plan` is, if it is one.
|
||||
fn as_merge_scan(plan: &LogicalPlan) -> Option<&MergeScanLogicalPlan> {
|
||||
if let LogicalPlan::Extension(extension) = plan {
|
||||
extension
|
||||
.node
|
||||
.as_any()
|
||||
.downcast_ref::<MergeScanLogicalPlan>()
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
/// The first merge scan in `plan`'s visible tree. A merge scan's remote input is not one of its
|
||||
/// children, so this does not look into one.
|
||||
fn first_merge_scan(plan: &LogicalPlan) -> Option<&MergeScanLogicalPlan> {
|
||||
if let Some(merge_scan) = as_merge_scan(plan) {
|
||||
return Some(merge_scan);
|
||||
}
|
||||
plan.inputs().into_iter().find_map(first_merge_scan)
|
||||
}
|
||||
|
||||
fn analyze(analyzer: DistPlannerAnalyzer, plan: LogicalPlan) -> LogicalPlan {
|
||||
analyzer.analyze(plan, &ConfigOptions::default()).unwrap()
|
||||
}
|
||||
|
||||
/// Whether the remote input of any merge scan in `plan` holds a node `wanted` accepts.
|
||||
fn any_remote_input_holds(plan: &LogicalPlan, wanted: fn(&LogicalPlan) -> bool) -> bool {
|
||||
let mut holds = false;
|
||||
plan.apply(|node| {
|
||||
if let Some(merge_scan) = as_merge_scan(node) {
|
||||
holds |= wanted(merge_scan.input());
|
||||
}
|
||||
Ok(TreeNodeRecursion::Continue)
|
||||
})
|
||||
.unwrap();
|
||||
holds
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn fused_series_over_its_scan_is_commutative() {
|
||||
let plan = fused_series_plan(table_scan_plan());
|
||||
|
||||
assert!(matches!(
|
||||
Categorizer::check_plan(&plan, None).unwrap(),
|
||||
Commutativity::Commutative
|
||||
));
|
||||
}
|
||||
|
||||
/// The region builds the node from the scan's own raw-series entrypoint, so an input that is
|
||||
/// not that scan leaves the node with nothing to read where it would be placed.
|
||||
#[test]
|
||||
fn fused_series_over_anything_but_a_scan_is_not_commutative() {
|
||||
let scan = table_scan_plan();
|
||||
let input = LogicalPlanBuilder::from(scan)
|
||||
.filter(datafusion_expr::col("k0").eq(datafusion_expr::lit("a")))
|
||||
.unwrap()
|
||||
.build()
|
||||
.unwrap();
|
||||
let plan = fused_series_plan(input);
|
||||
|
||||
assert!(matches!(
|
||||
Categorizer::check_plan(&plan, None).unwrap(),
|
||||
Commutativity::NonCommutative
|
||||
));
|
||||
}
|
||||
|
||||
/// The region asks the scan for one named metric's series, so a node binding no name cannot
|
||||
/// be built there.
|
||||
#[test]
|
||||
fn fused_series_that_binds_no_metric_name_is_not_commutative() {
|
||||
let plan = fused_series_plan_named(table_scan_plan(), None);
|
||||
|
||||
assert!(matches!(
|
||||
Categorizer::check_plan(&plan, None).unwrap(),
|
||||
Commutativity::NonCommutative
|
||||
));
|
||||
}
|
||||
|
||||
/// A tree built by something other than the compiler can hold a node that folds several
|
||||
/// series. The check is a property of the tree, so it is exercised on the node itself: the
|
||||
/// compiler emits no such node today, and nothing else would build a tree containing one.
|
||||
#[test]
|
||||
fn a_tree_that_folds_across_series_is_not_per_series() {
|
||||
let raw = NativeType::Samples(Arc::new(SeriesSchema {
|
||||
fields: Arc::from([FieldSpec { name: "v0".into() }]),
|
||||
eval_grid: None,
|
||||
}));
|
||||
let input = NativeNode {
|
||||
expr: NativeExpr::Input(BoundInput {
|
||||
output: raw.clone(),
|
||||
ordering: SeriesOrdering::KeyThenTime,
|
||||
}),
|
||||
output: raw,
|
||||
};
|
||||
let group = Grouping::By(Arc::from(["k0".into()]));
|
||||
let aggregate = NativeNode {
|
||||
expr: NativeExpr::Aggregate {
|
||||
input: Box::new(input),
|
||||
op: BoundAggregate {
|
||||
kind: AggregateKind::Sum,
|
||||
group: group.clone(),
|
||||
stage: AggregateStage::Partial,
|
||||
},
|
||||
},
|
||||
output: NativeType::AggregateState(Arc::new(AggregateStateSchema {
|
||||
group,
|
||||
accumulators: Arc::from([]),
|
||||
grid: statement_grid(&native_statement()).unwrap(),
|
||||
})),
|
||||
};
|
||||
assert!(reads_across_series(&aggregate));
|
||||
|
||||
// The same walk over what the compiler emits stays per-series.
|
||||
let single_selector = compile(
|
||||
&native_statement(),
|
||||
Arc::new(SeriesSchema {
|
||||
fields: Arc::from([FieldSpec { name: "v0".into() }]),
|
||||
eval_grid: None,
|
||||
}),
|
||||
SeriesOrdering::KeyThenTime,
|
||||
)
|
||||
.unwrap();
|
||||
assert!(!reads_across_series(single_selector.root()));
|
||||
}
|
||||
|
||||
/// The encoding gate is what kept the node on the frontend, and it only exists for a plan that
|
||||
/// has to be encoded.
|
||||
///
|
||||
/// `PlanRewriter::should_expand` asked two questions about a node before it let the descent
|
||||
/// continue, and the Substrait one came first: unless the *whole* root plan encoded, every
|
||||
/// subtree that failed to encode was expanded on, whichever way `check_plan` answered. A
|
||||
/// `FusedSeries` node has no serializer branch — `ExtensionSerializer` answers `NotImplemented`
|
||||
/// for every name it does not list — so over the wire the boundary is still drawn at the node
|
||||
/// and the region receives the bare scan.
|
||||
///
|
||||
/// A plan handed to `RegionServer::handle_read` as a `LogicalPlan` is never encoded, so there
|
||||
/// the gate has nothing to say and the commutativity arm decides: the node rides inside the
|
||||
/// merge scan.
|
||||
#[test]
|
||||
fn fused_series_rides_to_the_region_only_when_regions_are_reached_in_process() {
|
||||
let plan = fused_series_plan(table_scan_plan());
|
||||
|
||||
assert!(
|
||||
DFLogicalSubstraitConvertor
|
||||
.encode(&plan, DefaultSerializer)
|
||||
.is_err(),
|
||||
"the fused node encodes now, so the encoding gate decides nothing"
|
||||
);
|
||||
assert!(
|
||||
DFLogicalSubstraitConvertor
|
||||
.encode(&table_scan_plan(), DefaultSerializer)
|
||||
.is_ok(),
|
||||
"the scan under the fused node encodes, so the fused node is the only reason this plan does not"
|
||||
);
|
||||
|
||||
// Over the wire the plan has to encode, so the boundary is drawn at the node.
|
||||
let encoded = analyze(DistPlannerAnalyzer::new(false), plan.clone());
|
||||
assert!(
|
||||
!any_remote_input_holds(&encoded, holds_fused_series),
|
||||
"the analyzer placed the node although its plan has to encode: {encoded}"
|
||||
);
|
||||
|
||||
// In process nothing is encoded, so the categorizer places the node.
|
||||
let in_process = analyze(DistPlannerAnalyzer::new(true), plan);
|
||||
assert!(
|
||||
any_remote_input_holds(&in_process, holds_fused_series),
|
||||
"the analyzer kept the node on the frontend although regions are reached in process: {in_process}"
|
||||
);
|
||||
}
|
||||
|
||||
/// Skipping the encoding gate removes one reason to expand and adds none: a node that does not
|
||||
/// encode for its own reasons, and that the categorizer refuses, still draws the boundary at
|
||||
/// itself when regions are reached in process.
|
||||
///
|
||||
/// Without this, the skip would amount to placing every unencodable node region-side.
|
||||
#[test]
|
||||
fn a_plan_that_does_not_encode_for_another_reason_still_expands_at_that_node() {
|
||||
let plan = unserializable_plan(table_scan_plan());
|
||||
assert!(
|
||||
DFLogicalSubstraitConvertor
|
||||
.encode(&plan, DefaultSerializer)
|
||||
.is_err(),
|
||||
"the stand-in node has a serializer branch now, so it no longer proves anything"
|
||||
);
|
||||
|
||||
let analyzed = analyze(DistPlannerAnalyzer::new(true), plan);
|
||||
|
||||
assert!(
|
||||
!any_remote_input_holds(&analyzed, holds_unserializable_node),
|
||||
"the analyzer placed a node the categorizer refuses: {analyzed}"
|
||||
);
|
||||
let inputs = analyzed.inputs();
|
||||
let [input] = inputs.as_slice() else {
|
||||
panic!("expected the node to keep one input in: {analyzed}");
|
||||
};
|
||||
// `expand` keeps a projection over the merge scan to restore the schema of what it wrapped.
|
||||
let merge_scan = first_merge_scan(input)
|
||||
.unwrap_or_else(|| panic!("expected the boundary right below the node in: {analyzed}"));
|
||||
assert!(
|
||||
holds_table_scan(merge_scan.input()),
|
||||
"expected the scan under the node to be the region input in: {analyzed}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
pub type Transformer = Arc<dyn Fn(&LogicalPlan) -> Option<LogicalPlan>>;
|
||||
|
||||
@@ -12,7 +12,6 @@
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
use std::collections::BinaryHeap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use common_telemetry::debug;
|
||||
@@ -22,9 +21,10 @@ use datafusion::physical_plan::ExecutionPlan;
|
||||
use datafusion::physical_plan::sorts::sort::SortExec;
|
||||
use datafusion_common::tree_node::{Transformed, TreeNode};
|
||||
use datafusion_common::{DataFusionError, Result};
|
||||
use store_api::region_engine::PartitionRange;
|
||||
use table::table::scan::RegionScanExec;
|
||||
|
||||
pub use table::table::scan::assign_partition_ranges;
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct ParallelizeScan;
|
||||
|
||||
@@ -71,7 +71,7 @@ impl ParallelizeScan {
|
||||
|
||||
// assign ranges to each partition
|
||||
let mut partition_ranges =
|
||||
Self::assign_partition_range(ranges, expected_partition_num);
|
||||
assign_partition_ranges(ranges, expected_partition_num);
|
||||
debug!(
|
||||
"Assign {total_range_num} ranges to {expected_partition_num} partitions"
|
||||
);
|
||||
@@ -116,176 +116,15 @@ impl ParallelizeScan {
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
/// Distribute [`PartitionRange`]s to each partition.
|
||||
///
|
||||
/// Currently we assign ranges to partitions according to their rows so each partition
|
||||
/// has similar number of rows. This method always return `expected_partition_num` partitions.
|
||||
pub fn assign_partition_range(
|
||||
mut ranges: Vec<PartitionRange>,
|
||||
expected_partition_num: usize,
|
||||
) -> Vec<Vec<PartitionRange>> {
|
||||
if ranges.is_empty() {
|
||||
// Returns a single partition with no range.
|
||||
return vec![vec![]; expected_partition_num];
|
||||
}
|
||||
|
||||
if ranges.len() == 1 {
|
||||
let mut vec = vec![vec![]; expected_partition_num];
|
||||
vec[0] = ranges;
|
||||
return vec;
|
||||
}
|
||||
|
||||
// Sort ranges by number of rows in descending order.
|
||||
ranges.sort_by_key(|b| std::cmp::Reverse(b.num_rows));
|
||||
let mut partition_ranges = vec![vec![]; expected_partition_num];
|
||||
|
||||
#[derive(Eq, PartialEq)]
|
||||
struct HeapNode {
|
||||
num_rows: usize,
|
||||
partition_idx: usize,
|
||||
}
|
||||
|
||||
impl Ord for HeapNode {
|
||||
fn cmp(&self, other: &Self) -> std::cmp::Ordering {
|
||||
// Reverse for min-heap.
|
||||
self.num_rows.cmp(&other.num_rows).reverse()
|
||||
}
|
||||
}
|
||||
|
||||
impl PartialOrd for HeapNode {
|
||||
fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
|
||||
Some(self.cmp(other))
|
||||
}
|
||||
}
|
||||
|
||||
let mut part_heap =
|
||||
BinaryHeap::from_iter((0..expected_partition_num).map(|partition_idx| HeapNode {
|
||||
num_rows: 0,
|
||||
partition_idx,
|
||||
}));
|
||||
|
||||
// Assigns the range to the partition with the smallest number of rows.
|
||||
for range in ranges {
|
||||
let mut node = part_heap.pop().unwrap();
|
||||
let partition_idx = node.partition_idx;
|
||||
node.num_rows += range.num_rows;
|
||||
partition_ranges[partition_idx].push(range);
|
||||
part_heap.push(node);
|
||||
}
|
||||
|
||||
partition_ranges
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod test {
|
||||
use common_time::Timestamp;
|
||||
use common_time::timestamp::TimeUnit;
|
||||
use store_api::region_engine::PartitionRange;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_assign_partition_range() {
|
||||
let ranges = vec![
|
||||
PartitionRange {
|
||||
start: Timestamp::new(0, TimeUnit::Second),
|
||||
end: Timestamp::new(10, TimeUnit::Second),
|
||||
num_rows: 100,
|
||||
identifier: 1,
|
||||
},
|
||||
PartitionRange {
|
||||
start: Timestamp::new(10, TimeUnit::Second),
|
||||
end: Timestamp::new(20, TimeUnit::Second),
|
||||
num_rows: 200,
|
||||
identifier: 2,
|
||||
},
|
||||
PartitionRange {
|
||||
start: Timestamp::new(20, TimeUnit::Second),
|
||||
end: Timestamp::new(30, TimeUnit::Second),
|
||||
num_rows: 150,
|
||||
identifier: 3,
|
||||
},
|
||||
PartitionRange {
|
||||
start: Timestamp::new(30, TimeUnit::Second),
|
||||
end: Timestamp::new(40, TimeUnit::Second),
|
||||
num_rows: 250,
|
||||
identifier: 4,
|
||||
},
|
||||
];
|
||||
|
||||
// assign to 2 partitions
|
||||
let expected_partition_num = 2;
|
||||
let result =
|
||||
ParallelizeScan::assign_partition_range(ranges.clone(), expected_partition_num);
|
||||
let expected = vec![
|
||||
vec![
|
||||
PartitionRange {
|
||||
start: Timestamp::new(30, TimeUnit::Second),
|
||||
end: Timestamp::new(40, TimeUnit::Second),
|
||||
num_rows: 250,
|
||||
identifier: 4,
|
||||
},
|
||||
PartitionRange {
|
||||
start: Timestamp::new(0, TimeUnit::Second),
|
||||
end: Timestamp::new(10, TimeUnit::Second),
|
||||
num_rows: 100,
|
||||
identifier: 1,
|
||||
},
|
||||
],
|
||||
vec![
|
||||
PartitionRange {
|
||||
start: Timestamp::new(10, TimeUnit::Second),
|
||||
end: Timestamp::new(20, TimeUnit::Second),
|
||||
num_rows: 200,
|
||||
identifier: 2,
|
||||
},
|
||||
PartitionRange {
|
||||
start: Timestamp::new(20, TimeUnit::Second),
|
||||
end: Timestamp::new(30, TimeUnit::Second),
|
||||
num_rows: 150,
|
||||
identifier: 3,
|
||||
},
|
||||
],
|
||||
];
|
||||
assert_eq!(result, expected);
|
||||
|
||||
// assign 4 ranges to 5 partitions.
|
||||
let expected_partition_num = 5;
|
||||
let result = ParallelizeScan::assign_partition_range(ranges, expected_partition_num);
|
||||
let expected = vec![
|
||||
vec![PartitionRange {
|
||||
start: Timestamp::new(30, TimeUnit::Second),
|
||||
end: Timestamp::new(40, TimeUnit::Second),
|
||||
num_rows: 250,
|
||||
identifier: 4,
|
||||
}],
|
||||
vec![PartitionRange {
|
||||
start: Timestamp::new(0, TimeUnit::Second),
|
||||
end: Timestamp::new(10, TimeUnit::Second),
|
||||
num_rows: 100,
|
||||
identifier: 1,
|
||||
}],
|
||||
vec![PartitionRange {
|
||||
start: Timestamp::new(10, TimeUnit::Second),
|
||||
end: Timestamp::new(20, TimeUnit::Second),
|
||||
num_rows: 200,
|
||||
identifier: 2,
|
||||
}],
|
||||
vec![],
|
||||
vec![PartitionRange {
|
||||
start: Timestamp::new(20, TimeUnit::Second),
|
||||
end: Timestamp::new(30, TimeUnit::Second),
|
||||
num_rows: 150,
|
||||
identifier: 3,
|
||||
}],
|
||||
];
|
||||
assert_eq!(result, expected);
|
||||
|
||||
// assign 0 ranges to 5 partitions. Should return 5 empty ranges.
|
||||
let result = ParallelizeScan::assign_partition_range(vec![], 5);
|
||||
assert_eq!(result.len(), 5);
|
||||
}
|
||||
use super::assign_partition_ranges;
|
||||
|
||||
#[test]
|
||||
fn test_assign_unbalance_partition_range() {
|
||||
@@ -318,8 +157,7 @@ mod test {
|
||||
|
||||
// assign to 2 partitions
|
||||
let expected_partition_num = 2;
|
||||
let result =
|
||||
ParallelizeScan::assign_partition_range(ranges.clone(), expected_partition_num);
|
||||
let result = assign_partition_ranges(ranges.clone(), expected_partition_num);
|
||||
let expected = vec![
|
||||
vec![PartitionRange {
|
||||
start: Timestamp::new(30, TimeUnit::Second),
|
||||
|
||||
@@ -26,7 +26,7 @@ use datafusion_expr::expr::Sort;
|
||||
use datafusion_expr::{Expr, LogicalPlan, utils};
|
||||
use datafusion_optimizer::{OptimizerConfig, OptimizerRule};
|
||||
use datatypes::arrow::datatypes::{DataType, TimeUnit as ArrowTimeUnit};
|
||||
use promql::extension_plan::{InstantManipulate, SeriesDivide, SeriesNormalize};
|
||||
use promql::extension_plan::{FusedSeries, InstantManipulate, SeriesDivide, SeriesNormalize};
|
||||
use store_api::metric_engine_consts::DATA_SCHEMA_TSID_COLUMN_NAME;
|
||||
use store_api::storage::{TimeSeriesDistribution, TimeSeriesRowSelector};
|
||||
|
||||
@@ -37,6 +37,9 @@ use crate::dummy_catalog::DummyTableProvider;
|
||||
/// - the nearest order requirement to the leaf table scan node as ordering hint.
|
||||
/// - the group by columns when all aggregate functions are `last_value` as
|
||||
/// time series row selector hint.
|
||||
/// - the per-series distribution when the leaf table scan node is the input of a
|
||||
/// [`FusedSeries`] node, which reads the raw series frames a region only packs
|
||||
/// for that distribution.
|
||||
///
|
||||
/// [`ScanRequest`]: store_api::storage::ScanRequest
|
||||
#[derive(Debug)]
|
||||
@@ -68,6 +71,7 @@ impl ScanHintRule {
|
||||
fn set_hints(
|
||||
plan: LogicalPlan,
|
||||
rewriter: &mut ScanHintRewriter,
|
||||
fused_series_input: bool,
|
||||
) -> Result<Transformed<LogicalPlan>> {
|
||||
let LogicalPlan::TableScan(mut table_scan) = plan else {
|
||||
return Ok(Transformed::no(plan));
|
||||
@@ -89,8 +93,10 @@ impl ScanHintRule {
|
||||
};
|
||||
let use_last_row = rewriter.inside_single_evaluation && filters_preserve_last_row;
|
||||
|
||||
let has_hint =
|
||||
rewriter.order_expr.is_some() || rewriter.ts_row_selector.is_some() || use_last_row;
|
||||
let has_hint = rewriter.order_expr.is_some()
|
||||
|| rewriter.ts_row_selector.is_some()
|
||||
|| use_last_row
|
||||
|| fused_series_input;
|
||||
if !has_hint {
|
||||
return Ok(Transformed::no(LogicalPlan::TableScan(table_scan)));
|
||||
}
|
||||
@@ -101,6 +107,14 @@ impl ScanHintRule {
|
||||
// local as well as the new LastRow hint.
|
||||
let adapter = original.clone_for_scan();
|
||||
Self::apply_hints(&adapter, rewriter);
|
||||
if fused_series_input {
|
||||
// The scan below a `FusedSeries` node reads packs of raw series, which a region only
|
||||
// produces for a per-series request; the existing PromQL pipeline asks for the same
|
||||
// distribution through a sort on `(series key, time)`, which the native node does not
|
||||
// have below it. Applying it after the order hints makes the plan shape the deciding
|
||||
// one: the node's own input is what reads the frames.
|
||||
adapter.with_distribution(TimeSeriesDistribution::PerSeries);
|
||||
}
|
||||
if use_last_row {
|
||||
// Apply the instant-derived hint after the aggregate hint. Both
|
||||
// select LastRow today, and this ordering preserves the existing
|
||||
@@ -275,6 +289,12 @@ struct ScanHintRewriter {
|
||||
ts_stack: Vec<Option<(HashSet<Column>, Column)>>,
|
||||
inside_single_evaluation: bool,
|
||||
single_evaluation_stack: Vec<bool>,
|
||||
/// Whether the node being visited is the input of a [`FusedSeries`] node.
|
||||
///
|
||||
/// The flag only describes that one edge: it is set while a `FusedSeries` node is visited, so
|
||||
/// the node visited next — that node's own input — sees it, and a node further down does not.
|
||||
fused_series_input: bool,
|
||||
fused_series_input_stack: Vec<bool>,
|
||||
}
|
||||
|
||||
impl TreeNodeRewriter for ScanHintRewriter {
|
||||
@@ -285,6 +305,13 @@ impl TreeNodeRewriter for ScanHintRewriter {
|
||||
self.ts_stack.push(self.ts_row_selector.clone());
|
||||
self.single_evaluation_stack
|
||||
.push(self.inside_single_evaluation);
|
||||
self.fused_series_input_stack.push(self.fused_series_input);
|
||||
// Read the edge this node sits on before recording what this node is, so a `FusedSeries`
|
||||
// node hands the flag to its own input and to nothing below it.
|
||||
let fused_series_input = std::mem::replace(
|
||||
&mut self.fused_series_input,
|
||||
matches!(&node, LogicalPlan::Extension(extension) if extension.node.as_any().is::<FusedSeries>()),
|
||||
);
|
||||
|
||||
if let LogicalPlan::Sort(sort) = &node {
|
||||
self.order_expr = Some(sort.expr.clone());
|
||||
@@ -321,7 +348,7 @@ impl TreeNodeRewriter for ScanHintRewriter {
|
||||
}
|
||||
}
|
||||
|
||||
ScanHintRule::set_hints(node, self)
|
||||
ScanHintRule::set_hints(node, self, fused_series_input)
|
||||
}
|
||||
|
||||
fn f_up(&mut self, node: LogicalPlan) -> Result<Transformed<LogicalPlan>> {
|
||||
@@ -334,6 +361,9 @@ impl TreeNodeRewriter for ScanHintRewriter {
|
||||
if let Some(previous) = self.single_evaluation_stack.pop() {
|
||||
self.inside_single_evaluation = previous;
|
||||
}
|
||||
if let Some(previous) = self.fused_series_input_stack.pop() {
|
||||
self.fused_series_input = previous;
|
||||
}
|
||||
Ok(Transformed::no(node))
|
||||
}
|
||||
}
|
||||
@@ -446,6 +476,7 @@ impl ScanHintRewriter {
|
||||
mod test {
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, UNIX_EPOCH};
|
||||
|
||||
use datafusion::functions_aggregate::first_last::last_value_udaf;
|
||||
use datafusion::functions_aggregate::min_max::max_udaf;
|
||||
@@ -454,21 +485,26 @@ mod test {
|
||||
use datafusion::logical_expr::{WindowFrame, WindowFunctionDefinition};
|
||||
use datafusion::prelude::JoinType;
|
||||
use datafusion_common::tree_node::TreeNodeRecursion;
|
||||
use datafusion_common::{DFSchema, DFSchemaRef};
|
||||
use datafusion_expr::expr::{
|
||||
AggregateFunction, AggregateFunctionParams, Cast, WindowFunctionParams,
|
||||
};
|
||||
use datafusion_expr::expr_fn::scalar_subquery;
|
||||
use datafusion_expr::{Extension, LogicalPlan, LogicalPlanBuilder, col, lit};
|
||||
use datafusion_optimizer::OptimizerContext;
|
||||
use datatypes::arrow::datatypes::DataType;
|
||||
use datatypes::arrow::datatypes::{DataType, Field, TimeUnit};
|
||||
use datatypes::data_type::ConcreteDataType;
|
||||
use datatypes::schema::ColumnSchema;
|
||||
use promql::extension_plan::RangeManipulate;
|
||||
use datatypes::series::SeriesOrdering;
|
||||
use datatypes::series::samples::{FieldSpec, SeriesSchema};
|
||||
use promql::extension_plan::{FusedSeries, FusedSeriesBinding, RangeManipulate};
|
||||
use promql_parser::parser;
|
||||
use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder};
|
||||
use store_api::metric_engine_consts::DATA_SCHEMA_TSID_COLUMN_NAME;
|
||||
use store_api::storage::{RegionId, TimeSeriesRowSelector};
|
||||
|
||||
use super::*;
|
||||
use crate::promql::native_compile::{compile, statement_grid};
|
||||
|
||||
fn scan_requests(plan: &LogicalPlan) -> Vec<store_api::storage::ScanRequest> {
|
||||
scan_requests_with_names(plan)
|
||||
@@ -1520,4 +1556,123 @@ mod test {
|
||||
Some(TimeSeriesDistribution::PerSeries)
|
||||
);
|
||||
}
|
||||
|
||||
/// A `FusedSeries` node over `input`, bound to the tree a bare selector compiles to.
|
||||
///
|
||||
/// The rule reads the plan shape, not what the node computes, so the simplest shape the
|
||||
/// compiler accepts is enough for these tests.
|
||||
fn fused_series(input: LogicalPlan) -> LogicalPlan {
|
||||
let stmt = promql_parser::parser::EvalStmt {
|
||||
expr: parser::parse("some_metric").unwrap(),
|
||||
start: UNIX_EPOCH,
|
||||
end: UNIX_EPOCH
|
||||
.checked_add(Duration::from_secs(100_000))
|
||||
.unwrap(),
|
||||
interval: Duration::from_secs(5),
|
||||
lookback_delta: Duration::from_secs(1),
|
||||
};
|
||||
let series_schema = Arc::new(SeriesSchema {
|
||||
fields: Arc::from([FieldSpec { name: "v0".into() }]),
|
||||
eval_grid: None,
|
||||
});
|
||||
let tree = compile(
|
||||
&stmt,
|
||||
Arc::clone(&series_schema),
|
||||
SeriesOrdering::KeyThenTime,
|
||||
)
|
||||
.unwrap();
|
||||
let binding = FusedSeriesBinding {
|
||||
series_schema,
|
||||
grid: statement_grid(&stmt).unwrap(),
|
||||
metric_name: Some("some_metric".into()),
|
||||
time_index: "ts".to_string(),
|
||||
tags: Vec::new(),
|
||||
};
|
||||
let output_schema: DFSchemaRef = Arc::new(
|
||||
DFSchema::new_with_metadata(
|
||||
vec![
|
||||
(None, Arc::new(Field::new("v0", DataType::Float64, true))),
|
||||
(
|
||||
None,
|
||||
Arc::new(Field::new(
|
||||
"ts",
|
||||
DataType::Timestamp(TimeUnit::Millisecond, None),
|
||||
false,
|
||||
)),
|
||||
),
|
||||
],
|
||||
Default::default(),
|
||||
)
|
||||
.unwrap(),
|
||||
);
|
||||
|
||||
LogicalPlan::Extension(Extension {
|
||||
node: Arc::new(
|
||||
FusedSeries::try_new(
|
||||
Arc::from("some_metric"),
|
||||
tree,
|
||||
binding,
|
||||
input,
|
||||
output_schema,
|
||||
)
|
||||
.unwrap(),
|
||||
),
|
||||
})
|
||||
}
|
||||
|
||||
/// The region only packs raw series frames for a per-series request, and the node below is
|
||||
/// what reads them, so the node's own input asks for that distribution.
|
||||
#[test]
|
||||
fn a_scan_under_fused_series_asks_for_the_per_series_distribution() {
|
||||
let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
|
||||
let plan = fused_series(scan_plan(provider.clone(), "t"));
|
||||
|
||||
let rewritten = ScanHintRule
|
||||
.rewrite(plan, &OptimizerContext::default())
|
||||
.unwrap()
|
||||
.data;
|
||||
|
||||
assert_eq!(
|
||||
scan_requests(&rewritten)[0].distribution,
|
||||
Some(TimeSeriesDistribution::PerSeries)
|
||||
);
|
||||
// The hint belongs to this use-site only, never to the shared catalog provider.
|
||||
assert_eq!(provider.scan_request().distribution, None);
|
||||
}
|
||||
|
||||
/// The ordinary path: a scan with no fused node above it is left exactly as it was.
|
||||
#[test]
|
||||
fn a_scan_not_under_fused_series_keeps_its_request() {
|
||||
let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
|
||||
let plan = scan_plan(provider.clone(), "t");
|
||||
|
||||
let rewritten = ScanHintRule
|
||||
.rewrite(plan, &OptimizerContext::default())
|
||||
.unwrap()
|
||||
.data;
|
||||
|
||||
assert_eq!(scan_requests(&rewritten)[0].distribution, None);
|
||||
assert_eq!(provider.scan_request().distribution, None);
|
||||
}
|
||||
|
||||
/// Only the node's own input is a raw series reader; a scan further down feeds something else
|
||||
/// and keeps the request it had.
|
||||
#[test]
|
||||
fn a_scan_below_fused_series_input_keeps_its_request() {
|
||||
let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
|
||||
let scan = scan_plan(provider.clone(), "t");
|
||||
let input = LogicalPlanBuilder::from(scan)
|
||||
.project(vec![col("ts"), col("v0")])
|
||||
.unwrap()
|
||||
.build()
|
||||
.unwrap();
|
||||
let plan = fused_series(input);
|
||||
|
||||
let rewritten = ScanHintRule
|
||||
.rewrite(plan, &OptimizerContext::default())
|
||||
.unwrap()
|
||||
.data;
|
||||
|
||||
assert_eq!(scan_requests(&rewritten)[0].distribution, None);
|
||||
}
|
||||
}
|
||||
|
||||
+141
-6
@@ -50,6 +50,8 @@ use sql::statements::statement::Statement;
|
||||
use sql::statements::tql::Tql;
|
||||
use sqlparser::ast::{AccessExpr, Value, visit_expressions_mut};
|
||||
|
||||
use crate::datafusion::QUERY_FALLBACK_HINT;
|
||||
use crate::dist_plan::DistPlannerOptions;
|
||||
use crate::error::{
|
||||
CteColumnSchemaMismatchSnafu, PlanSqlSnafu, QueryPlanSnafu, Result, SqlSnafu,
|
||||
UnimplementedSnafu,
|
||||
@@ -316,8 +318,27 @@ impl DfLogicalPlanner {
|
||||
}
|
||||
|
||||
#[tracing::instrument(skip_all)]
|
||||
async fn plan_pql(&self, stmt: &EvalStmt, query_ctx: QueryContextRef) -> Result<LogicalPlan> {
|
||||
async fn plan_pql(
|
||||
&self,
|
||||
stmt: &EvalStmt,
|
||||
query_ctx: QueryContextRef,
|
||||
native_route_allowed: bool,
|
||||
) -> Result<LogicalPlan> {
|
||||
let mut scheduled_state = self.derive_session_state_with_scheduled_time(&query_ctx)?;
|
||||
// Whether the distributed planner may fall back to the fallback plan rewriter. The
|
||||
// execution engine resolves the same three sources the same way: the engine option, which
|
||||
// it copies into the session state it plans with, and then, per query, the session
|
||||
// variable or the `query_fallback` hint. Each source only ever turns the fallback on, so
|
||||
// the value they end up with is the disjunction of them.
|
||||
let allow_query_fallback = scheduled_state
|
||||
.config_options()
|
||||
.extensions
|
||||
.get::<DistPlannerOptions>()
|
||||
.is_some_and(|options| options.allow_query_fallback)
|
||||
|| query_ctx.configuration_parameter().allow_query_fallback()
|
||||
|| query_ctx
|
||||
.extension(QUERY_FALLBACK_HINT)
|
||||
.is_some_and(|fallback| fallback.to_lowercase().parse::<bool>().unwrap_or(false));
|
||||
let promql_annotations = query_ctx.remote_query_id().map(promql_annotation_collector);
|
||||
if let Some(collector) = &promql_annotations {
|
||||
scheduled_state
|
||||
@@ -345,6 +366,16 @@ impl DfLogicalPlanner {
|
||||
stmt,
|
||||
&self.engine_state,
|
||||
promql_annotations,
|
||||
// A query allowed to fall back keeps the existing plan. The native node's physical
|
||||
// planner needs the bare region scan as its input; the fallback rewriter instead wraps
|
||||
// every table scan in `MergeScan`, so the node above it would have no physical plan and
|
||||
// the query would fail after the fallback was already taken. Declining here keeps the
|
||||
// whole query on the legacy path.
|
||||
native_route_allowed
|
||||
&& !allow_query_fallback
|
||||
&& self.engine_state.experimental_native_promql()
|
||||
&& self.engine_state.in_process_region_transport()
|
||||
&& crate::promql::planner::native::native_promql_requested(&query_ctx),
|
||||
)
|
||||
.await
|
||||
.map_err(BoxedError::new)
|
||||
@@ -461,8 +492,16 @@ impl DfLogicalPlanner {
|
||||
alias: eval.alias.clone(),
|
||||
};
|
||||
let stmt = QueryLanguageParser::parse_promql(&prom_query, &query_ctx)?;
|
||||
|
||||
self.plan(&stmt, query_ctx).await
|
||||
// A hybrid CTE registers this plan as a table the outer SQL query reads, so the
|
||||
// `FusedSeries` node would sit under a SQL root. The in-process encodability
|
||||
// exception is scoped to a plan whose root is the node, so here the node would be
|
||||
// cut off and left in the frontend, where the physical planner cannot build it -
|
||||
// a query that used to answer would fail instead. The route is therefore not
|
||||
// offered for an embedded statement.
|
||||
let QueryStatement::Promql(eval, _alias) = &stmt else {
|
||||
unreachable!("a TQL eval parses to a PromQL statement");
|
||||
};
|
||||
self.plan_pql(eval, query_ctx, false).await
|
||||
}
|
||||
Tql::Explain(_) => UnimplementedSnafu {
|
||||
operation: "TQL EXPLAIN in CTEs",
|
||||
@@ -699,7 +738,8 @@ impl LogicalPlanner for DfLogicalPlanner {
|
||||
async fn plan(&self, stmt: &QueryStatement, query_ctx: QueryContextRef) -> Result<LogicalPlan> {
|
||||
match stmt {
|
||||
QueryStatement::Sql(stmt) => self.plan_sql(stmt, query_ctx).await,
|
||||
QueryStatement::Promql(stmt, _alias) => self.plan_pql(stmt, query_ctx).await,
|
||||
// The statement's own plan is what a region receives, so the native route is offered.
|
||||
QueryStatement::Promql(stmt, _alias) => self.plan_pql(stmt, query_ctx, true).await,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -760,7 +800,9 @@ mod tests {
|
||||
use table::test_util::EmptyTable;
|
||||
|
||||
use super::*;
|
||||
use crate::datafusion::QUERY_FALLBACK_HINT;
|
||||
use crate::parser::{PromQuery, QueryLanguageParser};
|
||||
use crate::promql::planner::NATIVE_PROMQL_EXTENSION_KEY;
|
||||
use crate::{QueryEngineFactory, QueryEngineRef};
|
||||
|
||||
async fn create_test_engine() -> QueryEngineRef {
|
||||
@@ -819,6 +861,14 @@ mod tests {
|
||||
}
|
||||
|
||||
fn create_promql_test_engine() -> QueryEngineRef {
|
||||
create_promql_test_engine_with(crate::options::QueryOptions::default(), false)
|
||||
}
|
||||
|
||||
/// The PromQL fixture, additionally stating the engine options and how its regions are reached.
|
||||
fn create_promql_test_engine_with(
|
||||
query_options: crate::options::QueryOptions,
|
||||
in_process_region_transport: bool,
|
||||
) -> QueryEngineRef {
|
||||
let catalog_manager = MemoryCatalogManager::with_default_setup();
|
||||
let physical_table_name = "phy";
|
||||
let physical_table_id = 999u32;
|
||||
@@ -909,15 +959,19 @@ mod tests {
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
QueryEngineFactory::new(
|
||||
QueryEngineFactory::try_new_with_plugins_and_region_transport(
|
||||
catalog_manager,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
false,
|
||||
crate::options::QueryOptions::default(),
|
||||
in_process_region_transport,
|
||||
Default::default(),
|
||||
query_options,
|
||||
)
|
||||
.unwrap()
|
||||
.query_engine()
|
||||
}
|
||||
|
||||
@@ -950,6 +1004,87 @@ mod tests {
|
||||
engine.planner().plan(&stmt, query_ctx).await.unwrap()
|
||||
}
|
||||
|
||||
/// Plans `rate(some_metric[5m])` on the fixture above with the fallback sources the caller
|
||||
/// states — the engine option, the session variable and the `query_fallback` hint — while the
|
||||
/// native route is enabled and requested, so only the fallback guard decides.
|
||||
async fn plan_native_promql_with_fallback(
|
||||
engine_fallback: bool,
|
||||
session_fallback: bool,
|
||||
hint: Option<&str>,
|
||||
) -> String {
|
||||
let engine = create_promql_test_engine_with(
|
||||
crate::options::QueryOptions {
|
||||
allow_query_fallback: engine_fallback,
|
||||
experimental_native_promql: true,
|
||||
..Default::default()
|
||||
},
|
||||
// The native route plans a region input the region server in this process accepts.
|
||||
true,
|
||||
);
|
||||
// The native route request and the fallback hint both reach the planner as query context
|
||||
// extensions, the way a client sends them.
|
||||
let mut builder = QueryContextBuilder::default()
|
||||
.set_extension(NATIVE_PROMQL_EXTENSION_KEY.to_string(), "true".to_string());
|
||||
if let Some(hint) = hint {
|
||||
builder = builder.set_extension(QUERY_FALLBACK_HINT.to_string(), hint.to_string());
|
||||
}
|
||||
let query_ctx = builder.build();
|
||||
query_ctx
|
||||
.configuration_parameter()
|
||||
.set_allow_query_fallback(session_fallback);
|
||||
let query_ctx = Arc::new(query_ctx);
|
||||
let stmt = QueryLanguageParser::parse_promql(
|
||||
&PromQuery {
|
||||
query: "rate(some_metric[5m])".to_string(),
|
||||
start: "0".to_string(),
|
||||
end: "10".to_string(),
|
||||
step: "5s".to_string(),
|
||||
lookback: "300s".to_string(),
|
||||
alias: None,
|
||||
},
|
||||
&query_ctx,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
engine
|
||||
.planner()
|
||||
.plan(&stmt, query_ctx)
|
||||
.await
|
||||
.unwrap()
|
||||
.display_indent()
|
||||
.to_string()
|
||||
}
|
||||
|
||||
/// The native route is admitted only when the query cannot fall back, and every source the
|
||||
/// execution engine resolves the fallback from decides that at plan time too.
|
||||
#[tokio::test]
|
||||
async fn test_plan_pql_declines_native_route_when_query_may_fall_back() {
|
||||
// (engine option, session variable, hint, native route admitted)
|
||||
for (engine, session, hint, native) in [
|
||||
// Nothing asks for the fallback, so the route is planned as before.
|
||||
(false, false, None, true),
|
||||
// A hint that does not parse as `true` leaves the fallback off, as it does at execution.
|
||||
(false, false, Some("false"), true),
|
||||
(false, false, Some("yes"), true),
|
||||
// Either per-query source turns the fallback on, so the route is declined.
|
||||
(false, true, None, false),
|
||||
(false, false, Some("true"), false),
|
||||
(false, false, Some("TRUE"), false),
|
||||
// Each source only turns the fallback on: neither a `false` hint nor a `false` session
|
||||
// variable clears the engine option, and a `false` hint does not clear the variable.
|
||||
(true, false, None, false),
|
||||
(true, false, Some("false"), false),
|
||||
(true, true, Some("false"), false),
|
||||
] {
|
||||
let plan = plan_native_promql_with_fallback(engine, session, hint).await;
|
||||
assert_eq!(
|
||||
plan.contains("FusedSeries"),
|
||||
native,
|
||||
"engine={engine}, session={session}, hint={hint:?}: {plan}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// Plans `sql` and runs the DataFusion analyzer, which is where
|
||||
/// `InsertAssignmentRule` sits. Planning alone stops short of it, so these
|
||||
/// assertions would not see the assignment rewrite at all.
|
||||
|
||||
@@ -1116,7 +1116,11 @@ fn instant_selection(context: &Context, selector: &VectorSelector) -> Result<Nat
|
||||
|
||||
/// The grid the statement is evaluated on: `(end - start) / step + 1` points, the last of which
|
||||
/// does not pass `end`.
|
||||
fn statement_grid(stmt: &EvalStmt) -> Result<Arc<EvalGrid>> {
|
||||
///
|
||||
/// Visible to the planner because the binding a [`NativeTree`] is planned against has to name the
|
||||
/// same grid the tree itself was compiled on, and [`compile`]'s callers therefore read it here
|
||||
/// rather than deriving the grid a second time.
|
||||
pub(crate) fn statement_grid(stmt: &EvalStmt) -> Result<Arc<EvalGrid>> {
|
||||
let start_ms = signed_millis_since_epoch(stmt.start)?;
|
||||
let end_ms = signed_millis_since_epoch(stmt.end)?;
|
||||
// The step is narrowed before [`GridSpec`] sees it, so an interval to which no millisecond
|
||||
|
||||
@@ -16,6 +16,8 @@ mod at_modifier;
|
||||
mod function_plans;
|
||||
mod island;
|
||||
mod matching_filters;
|
||||
pub(crate) mod native;
|
||||
pub use native::NATIVE_PROMQL_EXTENSION_KEY;
|
||||
|
||||
mod set_operator;
|
||||
|
||||
@@ -292,6 +294,16 @@ pub struct PromPlanner {
|
||||
ctx: PromPlannerContext,
|
||||
/// Optional collector passed to native histogram UDFs.
|
||||
promql_annotations: Option<PromqlAnnotationCollector>,
|
||||
/// Whether a statement is first offered to the native fused-series compiler.
|
||||
///
|
||||
/// Read per query from the context extension
|
||||
/// [`NATIVE_PROMQL_EXTENSION_KEY`](super::planner::native::NATIVE_PROMQL_EXTENSION_KEY), and
|
||||
/// `false` when the query says nothing: with it off, planning takes the existing path unchanged
|
||||
/// and nothing below is reached.
|
||||
///
|
||||
/// The engine-wide `QueryOptions::experimental_native_promql` is applied while `QueryEngineState`
|
||||
/// is built; this planner reads the per-query request extension instead.
|
||||
native_fused: bool,
|
||||
}
|
||||
|
||||
type BinaryFieldPair<'a> = (&'a String, &'a String);
|
||||
@@ -302,22 +314,38 @@ impl PromPlanner {
|
||||
stmt: &EvalStmt,
|
||||
query_engine_state: &QueryEngineState,
|
||||
) -> Result<LogicalPlan> {
|
||||
Self::stmt_to_plan_with_annotations(table_provider, stmt, query_engine_state, None).await
|
||||
Self::stmt_to_plan_with_annotations(table_provider, stmt, query_engine_state, None, false)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Plans a PromQL statement and passes the optional collector to histogram UDFs.
|
||||
///
|
||||
/// `native_fused` enables the native fused-series route: the statement is offered to
|
||||
/// [`compile`](crate::promql::native_compile::compile) first, and whatever that compiler does
|
||||
/// not accept is planned by the existing path, unchanged. See
|
||||
/// [`Self::try_plan_native_fused`].
|
||||
pub async fn stmt_to_plan_with_annotations(
|
||||
table_provider: DfTableSourceProvider,
|
||||
stmt: &EvalStmt,
|
||||
query_engine_state: &QueryEngineState,
|
||||
promql_annotations: Option<PromqlAnnotationCollector>,
|
||||
native_fused: bool,
|
||||
) -> Result<LogicalPlan> {
|
||||
let mut planner = Self {
|
||||
table_provider,
|
||||
ctx: PromPlannerContext::from_eval_stmt(stmt)?,
|
||||
promql_annotations,
|
||||
native_fused,
|
||||
};
|
||||
|
||||
// The whole statement is offered to the native compiler first, because that compiler folds
|
||||
// a range function and the math functions over it into one tree: offered any later, the
|
||||
// expression has already been split into the nodes those functions would have shared. A
|
||||
// shape it does not accept is planned by the existing path exactly as if it were not here.
|
||||
if let Some(plan) = planner.try_plan_native_fused(&stmt.expr).await? {
|
||||
return planner.strip_tsid_column(plan);
|
||||
}
|
||||
|
||||
let plan = planner
|
||||
.prom_root_expr_to_plan(&stmt.expr, query_engine_state)
|
||||
.await?;
|
||||
@@ -2394,9 +2422,20 @@ impl PromPlanner {
|
||||
// let promql_parser::parser::ast::Extension { expr } = ext_expr;
|
||||
let expr = &ext_expr.expr;
|
||||
let children = expr.children();
|
||||
let plan = self
|
||||
.prom_root_expr_to_plan(&children[0], query_engine_state)
|
||||
.await?;
|
||||
let native_plan = match expr.name() {
|
||||
ANALYZE_NODE_NAME
|
||||
| ANALYZE_VERBOSE_NODE_NAME
|
||||
| EXPLAIN_NODE_NAME
|
||||
| EXPLAIN_VERBOSE_NODE_NAME => self.try_plan_native_fused(&children[0]).await?,
|
||||
_ => None,
|
||||
};
|
||||
let plan = match native_plan {
|
||||
Some(plan) => plan,
|
||||
None => {
|
||||
self.prom_root_expr_to_plan(&children[0], query_engine_state)
|
||||
.await?
|
||||
}
|
||||
};
|
||||
// Wrapper for the explanation/analyze of the existing plan
|
||||
// https://docs.rs/datafusion-expr/latest/datafusion_expr/logical_plan/builder/struct.LogicalPlanBuilder.html#method.explain
|
||||
// if `analyze` is true, runs the actual plan and produces
|
||||
|
||||
@@ -0,0 +1,528 @@
|
||||
// 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.
|
||||
|
||||
//! The native fused-series route of the PromQL planner.
|
||||
//!
|
||||
//! [`PromPlanner::try_plan_native_fused`] offers a small experimental subset to the native
|
||||
//! compiler: a vector selector (optionally parenthesized), or nested unary `rate`, `increase`,
|
||||
//! `sum_over_time`, and `abs` calls over a selector. It replaces that expression with one
|
||||
//! [`FusedSeries`] node over the selector's table scan. Other expression shapes keep the existing
|
||||
//! planner's route.
|
||||
//!
|
||||
//! What the node needs from here is three things the existing planner already resolves: the raw
|
||||
//! series schema the tree's input was compiled against (the selector's selected field columns), the
|
||||
//! statement's evaluation grid, and the logical tag names the output carries. The output schema is
|
||||
//! built to the shape the existing unary pipeline produces — the time index, then the series
|
||||
//! fields, then the tags — so a node planned here can be consumed by the same code that consumes a
|
||||
//! `SeriesNormalize` output.
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, UNIX_EPOCH};
|
||||
|
||||
use common_catalog::consts::{METRIC_ENGINE, MITO_ENGINE, MITO2_ENGINE};
|
||||
use common_query::prelude::OTLP_AGGREGATION_TEMPORALITY_LABEL;
|
||||
use datafusion::arrow::datatypes::{DataType as ArrowDataType, Field, TimeUnit as ArrowTimeUnit};
|
||||
use datafusion::common::{DFSchema, DFSchemaRef};
|
||||
use datafusion::error::DataFusionError;
|
||||
use datafusion::logical_expr::{Extension, LogicalPlan, TableScanBuilder};
|
||||
use datafusion_expr::Expr;
|
||||
use datatypes::series::SeriesOrdering;
|
||||
use datatypes::series::samples::{FieldSpec, SeriesSchema};
|
||||
use promql::extension_plan::{FusedSeries, FusedSeriesBinding};
|
||||
use promql_parser::parser::value::ValueType;
|
||||
use promql_parser::parser::{
|
||||
EvalStmt, Expr as PromExpr, MatrixSelector, ParenExpr, VectorSelector,
|
||||
};
|
||||
use session::context::QueryContextRef;
|
||||
use snafu::ResultExt;
|
||||
use store_api::metric_engine_consts::{DATA_SCHEMA_TABLE_ID_COLUMN_NAME, METRIC_ENGINE_NAME};
|
||||
|
||||
use crate::promql::error::{DataFusionPlanningSnafu, Result};
|
||||
use crate::promql::native_compile::{compile, statement_grid};
|
||||
use crate::promql::planner::PromPlanner;
|
||||
|
||||
/// The query-context key that turns the native fused-series route on for one query.
|
||||
///
|
||||
/// Named after the engine-wide option it mirrors ([`QueryOptions::experimental_native_promql`]),
|
||||
/// and read exactly the way the other per-query override in this crate is read
|
||||
/// ([`remote_dyn_filter_pushdown_enabled_from_extensions`]): a query-context extension, which a
|
||||
/// client can set per request over the HTTP hints header or as a gRPC context key. The engine-wide
|
||||
/// option is applied while building query-engine state; this planner reads the per-query request
|
||||
/// extension at the call site; see [`native_promql_requested`].
|
||||
///
|
||||
/// [`QueryOptions::experimental_native_promql`]: crate::options::QueryOptions::experimental_native_promql
|
||||
/// [`remote_dyn_filter_pushdown_enabled_from_extensions`]:
|
||||
/// crate::options::remote_dyn_filter_pushdown_enabled_from_extensions
|
||||
pub const NATIVE_PROMQL_EXTENSION_KEY: &str = "query.experimental_native_promql";
|
||||
|
||||
/// Whether `query_ctx` asks for the native fused-series route.
|
||||
///
|
||||
/// The route is off unless the query asks for it, so a query that says nothing keeps the plan the
|
||||
/// existing planner builds. The value is compared literally (`true` in any case) rather than through
|
||||
/// a shared boolean parser: an unrecognized value leaves the route off, which is the same plan as
|
||||
/// the flag never having been read, so a typo costs a slow query and never a wrong one.
|
||||
pub(crate) fn native_promql_requested(query_ctx: &QueryContextRef) -> bool {
|
||||
query_ctx
|
||||
.extension(NATIVE_PROMQL_EXTENSION_KEY)
|
||||
.is_some_and(|value| value.eq_ignore_ascii_case("true"))
|
||||
}
|
||||
|
||||
/// The ordering this route compiles its trees against.
|
||||
///
|
||||
/// Which ordering a raw-series scan declares is decided when the region scanner is prepared, during
|
||||
/// physical planning, and the logical planner cannot read it: a legacy scan declares `KeyThenTime`
|
||||
/// and a two-phase scan `Contiguous`. So the route names the weakest ordering both provide —
|
||||
/// `Contiguous`, every series contiguous and the series themselves unordered — rather than
|
||||
/// overclaiming the legacy scan's stronger `KeyThenTime`. A native evaluation only refuses
|
||||
/// `Interleaved`, so a tree compiled here runs on either scanner.
|
||||
const NATIVE_ORDERING: SeriesOrdering = SeriesOrdering::Contiguous;
|
||||
|
||||
impl PromPlanner {
|
||||
/// Plans `prom_expr` as one [`FusedSeries`] node, or returns `None` to keep the existing plan.
|
||||
///
|
||||
/// The route is offered the whole expression so its tree computes the requested expression
|
||||
/// over the planner's resolved selector context: table name, time index, tag columns, and
|
||||
/// field columns (after any `__field__` matcher). These are the raw series schema the input
|
||||
/// tree is compiled against.
|
||||
///
|
||||
/// `None` is not a failure: unsupported shapes and queries this route cannot bind faithfully
|
||||
/// continue through the existing planner. Admission is limited to selectors and unary calls
|
||||
/// named above, plus the scan, schema, matcher, and grid constraints checked below.
|
||||
pub(super) async fn try_plan_native_fused(
|
||||
&mut self,
|
||||
prom_expr: &PromExpr,
|
||||
) -> Result<Option<LogicalPlan>> {
|
||||
if !self.native_fused {
|
||||
return Ok(None);
|
||||
}
|
||||
// Only the statement root is offered: the tree is compiled against the statement's grid,
|
||||
// and a subquery's expression is evaluated on its own, smaller one. A range-vector root is
|
||||
// the instant query's raw sample result, which the existing path reports as a matrix rather
|
||||
// than as one point per grid position.
|
||||
if prom_expr.value_type() == ValueType::Matrix {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
// Admit only a direct selector or the selected unary calls applied to one selector.
|
||||
let Some(selector) = admitted_selector(prom_expr) else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
// These matchers are alternatives the one scan below the node cannot represent, so the
|
||||
// query keeps the plan the existing path builds for it.
|
||||
if !selector.matchers.or_matchers.is_empty() {
|
||||
return Ok(None);
|
||||
}
|
||||
// Resolve the selector the way the existing planner does, and take what is left as the
|
||||
// rows it selects. The scan below the node is the selector's own scan, unfiltered, so a
|
||||
// matcher that survives preprocessing selects fewer rows here than the scan reads.
|
||||
let matchers = self.preprocess_label_matchers(&selector.matchers, &selector.name)?;
|
||||
if !matchers.matchers.is_empty() {
|
||||
return Ok(None);
|
||||
}
|
||||
// An `__field__` matcher narrows the field columns the existing planner builds its scan
|
||||
// with; folding every field of the table instead would evaluate a different window.
|
||||
if self.ctx.field_column_matcher.is_some() {
|
||||
return Ok(None);
|
||||
}
|
||||
// A metric without a table reads as an empty metric, which the existing planner plans.
|
||||
if self.setup_context().await?.is_some() {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let Some(table_name) = self.ctx.table_name.clone() else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(time_index) = self.ctx.time_index_column.clone() else {
|
||||
return Ok(None);
|
||||
};
|
||||
let fields = self.ctx.field_columns.clone();
|
||||
let tags = self.ctx.tag_columns.clone();
|
||||
// A series with no field has nothing to fold, and the tree's input would be an empty
|
||||
// schema, which no window can read.
|
||||
if fields.is_empty() {
|
||||
return Ok(None);
|
||||
}
|
||||
// On a delta-temporality table a series' samples are increments, which the existing
|
||||
// planner sums over the window for `increase` and divides by the range for `rate`
|
||||
// ([`PromPlanner::select_delta_range_math`]), per series because the tag can differ between
|
||||
// them. A compiled tree computes both for a cumulative series, so any `rate`/`increase`
|
||||
// over such a table keeps the existing plan; a bare selector still folds.
|
||||
if tags
|
||||
.iter()
|
||||
.any(|tag| tag == OTLP_AGGREGATION_TEMPORALITY_LABEL)
|
||||
&& calls_rate_or_increase(prom_expr)
|
||||
{
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
// The selector's own table scan, built by the same method the existing path uses: the
|
||||
// physical planner turns it into the region scan the node reads, and building a scan here
|
||||
// instead would let the two paths disagree about the columns they read.
|
||||
let table_ref = self.table_ref()?;
|
||||
let mut input = self.create_table_scan_plan(table_ref.clone()).await?;
|
||||
let logical_scan = metric_scan_and_table_id(&input);
|
||||
let has_verified_metric_rewrite = logical_scan.is_some();
|
||||
let scan = match logical_scan {
|
||||
Some((mut scan, predicate)) => {
|
||||
// Preserve logical field order and the table restriction on the direct physical scan.
|
||||
let source_schema = scan.source.schema();
|
||||
let Some(mut ordered_projection) = fields
|
||||
.iter()
|
||||
.map(|field| source_schema.index_of(field).ok())
|
||||
.collect::<Option<Vec<_>>>()
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
let original_projection = scan.projection.as_ref();
|
||||
let selected_indices = original_projection
|
||||
.cloned()
|
||||
.unwrap_or_else(|| (0..source_schema.fields().len()).collect());
|
||||
for index in selected_indices {
|
||||
let is_field = fields
|
||||
.iter()
|
||||
.any(|field| field == source_schema.field(index).name());
|
||||
if !is_field && !ordered_projection.contains(&index) {
|
||||
ordered_projection.push(index);
|
||||
}
|
||||
}
|
||||
scan.filters.push(predicate);
|
||||
scan = TableScanBuilder::from(scan)
|
||||
.with_projection(Some(ordered_projection))
|
||||
.build()
|
||||
.context(DataFusionPlanningSnafu)?;
|
||||
input = LogicalPlan::TableScan(scan.clone());
|
||||
scan
|
||||
}
|
||||
None => match &input {
|
||||
LogicalPlan::TableScan(scan) => scan.clone(),
|
||||
_ => return Ok(None),
|
||||
},
|
||||
};
|
||||
let table = self.table_from_source(&scan.source)?;
|
||||
let table_meta = &table.table_info().meta;
|
||||
// An engine without a raw-series scanner has no frame to give this node, and by the time
|
||||
// the physical planner asks, the logical plan is already this node: the refusal there is a
|
||||
// planning error, not a fallback. So the engine is checked here, where the existing plan is
|
||||
// still the answer. A file-engine table is the concrete case: its scanner is a
|
||||
// `SinglePartitionScanner`, which leaves `raw_series_schema` at the default `None`.
|
||||
if !supports_native_raw_series(&table_meta.engine)
|
||||
|| (table_meta.engine == METRIC_ENGINE_NAME && !has_verified_metric_rewrite)
|
||||
{
|
||||
return Ok(None);
|
||||
}
|
||||
// The region's raw-series preflight refuses these, and the physical layer turns that refusal into an error, so decline this route here instead.
|
||||
let time_index_type = table_meta
|
||||
.schema
|
||||
.timestamp_column()
|
||||
.map(|column| &column.data_type);
|
||||
if time_index_type
|
||||
!= Some(&datatypes::prelude::ConcreteDataType::timestamp_millisecond_datatype())
|
||||
|| table_meta
|
||||
.options
|
||||
.extra_options
|
||||
.get(store_api::mito_engine_options::APPEND_MODE_KEY)
|
||||
.is_some_and(|value| value.eq_ignore_ascii_case("true"))
|
||||
{
|
||||
return Ok(None);
|
||||
}
|
||||
// FusedSeries evaluates each region independently, so partition keys must identify a series.
|
||||
if !table
|
||||
.table_info()
|
||||
.meta
|
||||
.partition_column_names()
|
||||
.all(|column| tags.iter().any(|tag| tag == column))
|
||||
{
|
||||
return Ok(None);
|
||||
}
|
||||
// The legacy selector converts an offset the same way; sharing that conversion keeps the
|
||||
// two scan paths from disagreeing about an extreme duration.
|
||||
let offset = Self::offset_millis(&selector.offset);
|
||||
let (saved_start, saved_end, saved_range) = (self.ctx.start, self.ctx.end, self.ctx.range);
|
||||
self.ctx.range = selector_range_millis(prom_expr);
|
||||
// Keep the fused scan's bounds identical to the legacy selector scan, including @ and offset.
|
||||
let time_filter = (|| {
|
||||
let anchored = self.at_modifier_offset(&selector.at, &selector.offset)?;
|
||||
if anchored.is_some() {
|
||||
self.ctx.end = saved_start;
|
||||
}
|
||||
self.build_time_index_filter(anchored.unwrap_or(offset), input.schema())
|
||||
})();
|
||||
self.ctx.start = saved_start;
|
||||
self.ctx.end = saved_end;
|
||||
self.ctx.range = saved_range;
|
||||
let time_filter = time_filter?;
|
||||
if let Some(filter) = time_filter {
|
||||
let LogicalPlan::TableScan(scan) = &mut input else {
|
||||
unreachable!();
|
||||
};
|
||||
scan.filters.push(filter);
|
||||
}
|
||||
// The region's columnar series-index samples, one of the readers this scan may be served
|
||||
// by, store `Float64` fields only, and the logical plan cannot tell which reader arrives.
|
||||
// So every field has to be `Float64`: an integer, a `Float32`, a decimal, a string or a
|
||||
// native histogram field keeps the existing plan, which the physical layer still serves.
|
||||
if !fields.iter().all(|field| {
|
||||
Self::field_column_type(input.schema(), field) == Some(&ArrowDataType::Float64)
|
||||
}) {
|
||||
return Ok(None);
|
||||
}
|
||||
let Some(stmt) = self.native_statement(prom_expr) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Ok(grid) = statement_grid(&stmt) else {
|
||||
// The statement's grid is not one a tree can be evaluated on, so there is no grid to
|
||||
// address the node's output by.
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let series_schema = Arc::new(SeriesSchema {
|
||||
fields: fields
|
||||
.iter()
|
||||
.map(|name| FieldSpec {
|
||||
name: Arc::from(name.as_str()),
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.into(),
|
||||
eval_grid: None,
|
||||
});
|
||||
let tree = match compile(&stmt, Arc::clone(&series_schema), NATIVE_ORDERING) {
|
||||
Ok(tree) => tree,
|
||||
// The compiler refusing a shape is this route declining it: the existing planner gets
|
||||
// the query exactly as it would have without this flag.
|
||||
Err(DataFusionError::NotImplemented(_)) => return Ok(None),
|
||||
Err(error) => return Err(error).context(DataFusionPlanningSnafu),
|
||||
};
|
||||
let binding = FusedSeriesBinding {
|
||||
series_schema,
|
||||
// The tree's own output is evaluated on the statement's grid, and the binding names
|
||||
// the same one, so a chunk's absolute positions are read on the grid it was computed
|
||||
// on.
|
||||
grid,
|
||||
metric_name: Some(Arc::from(table_name.as_str())),
|
||||
time_index: time_index.clone(),
|
||||
tags: tags.clone(),
|
||||
};
|
||||
let output_schema = fused_output_schema(&fields, &tags, &time_index)?;
|
||||
let node = FusedSeries::try_new(
|
||||
Arc::from(prom_expr.to_string()),
|
||||
tree,
|
||||
binding,
|
||||
input,
|
||||
output_schema,
|
||||
)
|
||||
.context(DataFusionPlanningSnafu)?;
|
||||
|
||||
Ok(Some(LogicalPlan::Extension(Extension {
|
||||
node: Arc::new(node),
|
||||
})))
|
||||
}
|
||||
|
||||
/// The statement being planned, rebuilt from the evaluation numbers the planner holds.
|
||||
///
|
||||
/// The compiler reads the statement's grid, lookback and `@` anchors rather than the planner's
|
||||
/// mutable state. The route is only offered at the statement root, where `ctx.start` and
|
||||
/// `ctx.end` are the statement's own bounds: they are rewritten while an anchored window or a
|
||||
/// subquery is folded, so a nested expression would need the bounds it was actually entered
|
||||
/// with, which is why no nested expression is offered here.
|
||||
///
|
||||
/// `None` means the bounds are not representable as a [`SystemTime`], which is not a grid a
|
||||
/// tree can be evaluated on either.
|
||||
fn native_statement(&self, prom_expr: &PromExpr) -> Option<EvalStmt> {
|
||||
let start = UNIX_EPOCH.checked_add(Duration::from_millis(self.ctx.stmt_start as u64))?;
|
||||
let end = UNIX_EPOCH.checked_add(Duration::from_millis(self.ctx.stmt_end as u64))?;
|
||||
Some(EvalStmt {
|
||||
expr: prom_expr.clone(),
|
||||
start,
|
||||
end,
|
||||
interval: Duration::from_millis(self.ctx.interval as u64),
|
||||
lookback_delta: Duration::from_millis(self.ctx.lookback_delta as u64),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// The columns a [`FusedSeries`] node produces.
|
||||
///
|
||||
/// The shape is the one the existing unary pipeline consumes — the time index, then the series
|
||||
/// fields, then the tags — and each column is declared the way the node's own conversion reads it
|
||||
/// back:
|
||||
/// the time column is a non-nullable millisecond timestamp with no timezone, because every
|
||||
/// evaluated chunk is addressed on the statement's grid and carries one; a field is a nullable
|
||||
/// `Float64`, because a grid point inside the statement's interval can have no sample; and a tag is
|
||||
/// a nullable `Utf8`, because a range function drops the metric name and a series can be missing a
|
||||
/// label entirely.
|
||||
fn fused_output_schema(
|
||||
fields: &[String],
|
||||
tags: &[String],
|
||||
time_index: &str,
|
||||
) -> Result<DFSchemaRef> {
|
||||
let mut columns = Vec::with_capacity(fields.len() + tags.len() + 1);
|
||||
columns.push((
|
||||
None,
|
||||
Arc::new(Field::new(
|
||||
time_index,
|
||||
ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None),
|
||||
false,
|
||||
)),
|
||||
));
|
||||
for name in fields {
|
||||
columns.push((
|
||||
None,
|
||||
Arc::new(Field::new(name, ArrowDataType::Float64, true)),
|
||||
));
|
||||
}
|
||||
for name in tags {
|
||||
columns.push((None, Arc::new(Field::new(name, ArrowDataType::Utf8, true))));
|
||||
}
|
||||
|
||||
Ok(Arc::new(
|
||||
DFSchema::new_with_metadata(columns, Default::default())
|
||||
.context(DataFusionPlanningSnafu)?,
|
||||
))
|
||||
}
|
||||
|
||||
/// The one selector `expr` folds, if it names exactly one.
|
||||
///
|
||||
/// The walk follows the grammar the native compiler accepts — a call, a range selector, a selector
|
||||
/// and parentheses — and stops at every other shape, because a shape the compiler cannot compile
|
||||
/// needs no selector resolved. Two distinct selectors are not one input, and a tree built over the
|
||||
/// wrong one of them would evaluate a window the query did not write.
|
||||
fn selector_range_millis(expr: &PromExpr) -> Option<i64> {
|
||||
match expr {
|
||||
PromExpr::MatrixSelector(matrix) => Some(matrix.range.as_millis() as i64),
|
||||
PromExpr::Paren(paren) => selector_range_millis(&paren.expr),
|
||||
PromExpr::Call(call) => call
|
||||
.args
|
||||
.args
|
||||
.iter()
|
||||
.find_map(|arg| selector_range_millis(arg)),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Whether `engine` can serve raw series frames to the fused-series node.
|
||||
///
|
||||
/// A whitelist avoids probing a region from logical planning; declining only costs the native route.
|
||||
fn supports_native_raw_series(engine: &str) -> bool {
|
||||
matches!(engine, MITO_ENGINE | MITO2_ENGINE | METRIC_ENGINE)
|
||||
}
|
||||
|
||||
/// Recognizes only the metric logical-table scan emitted by `create_table_scan_plan`.
|
||||
fn metric_scan_and_table_id(
|
||||
plan: &LogicalPlan,
|
||||
) -> Option<(datafusion::logical_expr::TableScan, Expr)> {
|
||||
let plan = match plan {
|
||||
LogicalPlan::Projection(projection) => projection.input.as_ref(),
|
||||
plan => plan,
|
||||
};
|
||||
let plan = match plan {
|
||||
LogicalPlan::SubqueryAlias(alias) => alias.input.as_ref(),
|
||||
plan => plan,
|
||||
};
|
||||
let LogicalPlan::Filter(filter) = plan else {
|
||||
return None;
|
||||
};
|
||||
let Expr::BinaryExpr(binary) = &filter.predicate else {
|
||||
return None;
|
||||
};
|
||||
if binary.op != datafusion::logical_expr::Operator::Eq
|
||||
|| !matches!(binary.left.as_ref(), Expr::Column(column) if column.name == DATA_SCHEMA_TABLE_ID_COLUMN_NAME)
|
||||
{
|
||||
return None;
|
||||
}
|
||||
if !matches!(
|
||||
binary.right.as_ref(),
|
||||
Expr::Literal(datafusion::scalar::ScalarValue::UInt32(Some(_)), _)
|
||||
) {
|
||||
return None;
|
||||
}
|
||||
let LogicalPlan::TableScan(scan) = filter.input.as_ref() else {
|
||||
return None;
|
||||
};
|
||||
Some((scan.clone(), filter.predicate.clone()))
|
||||
}
|
||||
|
||||
fn admitted_selector(expr: &PromExpr) -> Option<&VectorSelector> {
|
||||
match expr {
|
||||
PromExpr::VectorSelector(selector) => Some(selector),
|
||||
// Range selectors are admitted only as arguments within one of these supported calls.
|
||||
PromExpr::MatrixSelector(MatrixSelector { vs, .. }) => Some(vs),
|
||||
PromExpr::Paren(ParenExpr { expr }) => admitted_selector(expr),
|
||||
PromExpr::Call(call)
|
||||
if matches!(
|
||||
call.func.name,
|
||||
"rate" | "increase" | "sum_over_time" | "abs"
|
||||
) && call.args.args.len() == 1 =>
|
||||
{
|
||||
admitted_selector(&call.args.args[0])
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Whether `expr` calls `rate`/`increase`, the functions the existing planner rewrites per series
|
||||
/// for a delta-temporality table. Only nested calls and parentheses can hold one in an expression
|
||||
/// this route accepts; every other shape names no single selector or reads across series.
|
||||
fn calls_rate_or_increase(expr: &PromExpr) -> bool {
|
||||
match expr {
|
||||
PromExpr::Call(call) => {
|
||||
matches!(call.func.name, "rate" | "increase")
|
||||
|| call.args.args.iter().any(|arg| calls_rate_or_increase(arg))
|
||||
}
|
||||
PromExpr::Paren(ParenExpr { expr }) => calls_rate_or_increase(expr),
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use session::context::QueryContextBuilder;
|
||||
|
||||
use super::{NATIVE_PROMQL_EXTENSION_KEY, native_promql_requested};
|
||||
|
||||
#[test]
|
||||
fn native_promql_requested_is_false_when_extension_is_absent() {
|
||||
let query_ctx = Arc::new(QueryContextBuilder::default().build());
|
||||
|
||||
assert!(!native_promql_requested(&query_ctx));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn native_promql_requested_accepts_true_case_insensitively() {
|
||||
for value in ["true", "TRUE", "TrUe"] {
|
||||
let query_ctx = Arc::new(
|
||||
QueryContextBuilder::default()
|
||||
.set_extension(NATIVE_PROMQL_EXTENSION_KEY.to_string(), value.to_string())
|
||||
.build(),
|
||||
);
|
||||
|
||||
assert!(native_promql_requested(&query_ctx), "value: {value}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn native_promql_requested_rejects_false_and_unrecognized_values() {
|
||||
for value in ["false", "FALSE", "yes", "1", " true ", ""] {
|
||||
let query_ctx = Arc::new(
|
||||
QueryContextBuilder::default()
|
||||
.set_extension(NATIVE_PROMQL_EXTENSION_KEY.to_string(), value.to_string())
|
||||
.build(),
|
||||
);
|
||||
|
||||
assert!(!native_promql_requested(&query_ctx), "value: {value:?}");
|
||||
}
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -411,6 +411,7 @@ async fn delta_mixed_ranges_drop_and_float_ranges_sum() {
|
||||
.await,
|
||||
ctx: PromPlannerContext::default(),
|
||||
promql_annotations: Some(collector.clone()),
|
||||
native_fused: false,
|
||||
};
|
||||
let left_context = direct_or_context("lhs", &["job", marker], "v");
|
||||
let right_context = direct_or_context("rhs", &["job", marker], "v");
|
||||
@@ -712,6 +713,7 @@ async fn binary_joins_align_only_the_temporality_marker() {
|
||||
.await,
|
||||
ctx: PromPlannerContext::default(),
|
||||
promql_annotations: None,
|
||||
native_fused: false,
|
||||
};
|
||||
let joined = planner
|
||||
.join_on_non_field_columns(
|
||||
@@ -751,6 +753,7 @@ async fn binary_joins_align_only_the_temporality_marker() {
|
||||
.await,
|
||||
ctx: PromPlannerContext::default(),
|
||||
promql_annotations: None,
|
||||
native_fused: false,
|
||||
};
|
||||
let set = planner
|
||||
.set_op_on_non_field_columns(
|
||||
@@ -806,6 +809,7 @@ async fn binary_joins_align_only_the_temporality_marker() {
|
||||
.await,
|
||||
ctx: PromPlannerContext::default(),
|
||||
promql_annotations: None,
|
||||
native_fused: false,
|
||||
};
|
||||
let set = planner
|
||||
.set_op_on_non_field_columns(
|
||||
@@ -918,6 +922,7 @@ async fn binary_joins_align_dictionary_temporality_marker_with_tagless_vector()
|
||||
.await,
|
||||
ctx: PromPlannerContext::default(),
|
||||
promql_annotations: None,
|
||||
native_fused: false,
|
||||
};
|
||||
let scan = |name, table: Arc<MemTable>| {
|
||||
LogicalPlanBuilder::scan(name, provider_as_source(table), None)
|
||||
|
||||
@@ -189,7 +189,7 @@ impl QueryEngineFactory {
|
||||
plugins: Plugins,
|
||||
options: QueryOptions,
|
||||
) -> datafusion::error::Result<Self> {
|
||||
let state = Arc::new(QueryEngineState::try_new(
|
||||
Self::try_new_with_plugins_and_region_transport(
|
||||
catalog_manager,
|
||||
partition_rule_manager,
|
||||
region_query_handler,
|
||||
@@ -197,6 +197,40 @@ impl QueryEngineFactory {
|
||||
procedure_service_handler,
|
||||
flow_service_handler,
|
||||
with_dist_planner,
|
||||
false,
|
||||
plugins,
|
||||
options,
|
||||
)
|
||||
}
|
||||
|
||||
/// [`QueryEngineFactory::try_new_with_plugins`], additionally stating how the engine reaches
|
||||
/// the regions it queries.
|
||||
///
|
||||
/// `in_process_region_transport` reaches `DistPlannerAnalyzer`: when it is `true`, a region
|
||||
/// input may travel as a `LogicalPlan` instead of Substrait bytes, which only a region server in
|
||||
/// the same process can accept.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn try_new_with_plugins_and_region_transport(
|
||||
catalog_manager: CatalogManagerRef,
|
||||
partition_rule_manager: Option<PartitionRuleManagerRef>,
|
||||
region_query_handler: Option<RegionQueryHandlerRef>,
|
||||
table_mutation_handler: Option<TableMutationHandlerRef>,
|
||||
procedure_service_handler: Option<ProcedureServiceHandlerRef>,
|
||||
flow_service_handler: Option<FlowServiceHandlerRef>,
|
||||
with_dist_planner: bool,
|
||||
in_process_region_transport: bool,
|
||||
plugins: Plugins,
|
||||
options: QueryOptions,
|
||||
) -> datafusion::error::Result<Self> {
|
||||
let state = Arc::new(QueryEngineState::try_new_with_region_transport(
|
||||
catalog_manager,
|
||||
partition_rule_manager,
|
||||
region_query_handler,
|
||||
table_mutation_handler,
|
||||
procedure_service_handler,
|
||||
flow_service_handler,
|
||||
with_dist_planner,
|
||||
in_process_region_transport,
|
||||
plugins.clone(),
|
||||
options,
|
||||
)?);
|
||||
|
||||
@@ -99,6 +99,10 @@ pub struct QueryEngineState {
|
||||
table_functions: Arc<RwLock<HashMap<String, Arc<TableFunction>>>>,
|
||||
extension_rules: Vec<Arc<dyn ExtensionAnalyzerRule + Send + Sync>>,
|
||||
plugins: Plugins,
|
||||
experimental_native_promql: bool,
|
||||
/// Whether this engine's regions are reached in process, which is what lets a region take a
|
||||
/// `LogicalPlan` instead of Substrait bytes.
|
||||
in_process_region_transport: bool,
|
||||
}
|
||||
|
||||
impl fmt::Debug for QueryEngineState {
|
||||
@@ -147,6 +151,39 @@ impl QueryEngineState {
|
||||
with_dist_planner: bool,
|
||||
plugins: Plugins,
|
||||
options: QueryOptionsNew,
|
||||
) -> DfResult<Self> {
|
||||
Self::try_new_with_region_transport(
|
||||
catalog_list,
|
||||
partition_rule_manager,
|
||||
region_query_handler,
|
||||
table_mutation_handler,
|
||||
procedure_service_handler,
|
||||
flow_service_handler,
|
||||
with_dist_planner,
|
||||
false,
|
||||
plugins,
|
||||
options,
|
||||
)
|
||||
}
|
||||
|
||||
/// [`QueryEngineState::try_new`], additionally stating how the engine reaches its regions.
|
||||
///
|
||||
/// `in_process_region_transport` reaches [`DistPlannerAnalyzer`]: when it is `true` the
|
||||
/// distributed planner may hand a region a `LogicalPlan` instead of Substrait bytes, which only
|
||||
/// the in-process read path accepts. Every other caller builds an engine whose regions are
|
||||
/// reached over the wire, which is what [`QueryEngineState::try_new`] states.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) fn try_new_with_region_transport(
|
||||
catalog_list: CatalogManagerRef,
|
||||
partition_rule_manager: Option<PartitionRuleManagerRef>,
|
||||
region_query_handler: Option<RegionQueryHandlerRef>,
|
||||
table_mutation_handler: Option<TableMutationHandlerRef>,
|
||||
procedure_service_handler: Option<ProcedureServiceHandlerRef>,
|
||||
flow_service_handler: Option<FlowServiceHandlerRef>,
|
||||
with_dist_planner: bool,
|
||||
in_process_region_transport: bool,
|
||||
plugins: Plugins,
|
||||
options: QueryOptionsNew,
|
||||
) -> DfResult<Self> {
|
||||
let total_memory = get_total_memory_bytes().max(0) as u64;
|
||||
let memory_pool_size = options.memory_pool_size.resolve(total_memory) as usize;
|
||||
@@ -209,7 +246,9 @@ impl QueryEngineState {
|
||||
)),
|
||||
);
|
||||
if with_dist_planner {
|
||||
analyzer.rules.push(Arc::new(DistPlannerAnalyzer));
|
||||
analyzer.rules.push(Arc::new(DistPlannerAnalyzer::new(
|
||||
in_process_region_transport,
|
||||
)));
|
||||
analyzer.rules.push(Arc::new(JsonSchemaConcretizeRule));
|
||||
}
|
||||
analyzer.rules.push(Arc::new(FixStateUdafOrderingAnalyzer));
|
||||
@@ -297,6 +336,8 @@ impl QueryEngineState {
|
||||
table_functions: Arc::new(RwLock::new(HashMap::new())),
|
||||
extension_rules,
|
||||
plugins,
|
||||
experimental_native_promql: options.experimental_native_promql,
|
||||
in_process_region_transport,
|
||||
scalar_functions: Arc::new(RwLock::new(HashMap::new())),
|
||||
})
|
||||
}
|
||||
@@ -445,6 +486,15 @@ impl QueryEngineState {
|
||||
self.df_context.register_udwf(func);
|
||||
}
|
||||
|
||||
pub(crate) fn experimental_native_promql(&self) -> bool {
|
||||
self.experimental_native_promql
|
||||
}
|
||||
|
||||
/// Whether a region read is served in this process, so a region can take a `LogicalPlan`.
|
||||
pub(crate) fn in_process_region_transport(&self) -> bool {
|
||||
self.in_process_region_transport
|
||||
}
|
||||
|
||||
pub fn catalog_manager(&self) -> &CatalogManagerRef {
|
||||
&self.catalog_manager
|
||||
}
|
||||
|
||||
+139
-5
@@ -12,7 +12,7 @@
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
use std::collections::HashSet;
|
||||
use std::collections::{BinaryHeap, HashSet};
|
||||
use std::pin::Pin;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::task::{Context, Poll};
|
||||
@@ -109,6 +109,66 @@ impl std::fmt::Debug for RegionScanExec {
|
||||
}
|
||||
}
|
||||
|
||||
/// Distribute [`PartitionRange`]s to each partition.
|
||||
///
|
||||
/// Currently we assign ranges to partitions according to their rows so each partition
|
||||
/// has similar number of rows. This method always return `expected_partition_num` partitions.
|
||||
pub fn assign_partition_ranges(
|
||||
mut ranges: Vec<PartitionRange>,
|
||||
expected_partition_num: usize,
|
||||
) -> Vec<Vec<PartitionRange>> {
|
||||
if ranges.is_empty() {
|
||||
// Returns a single partition with no range.
|
||||
return vec![vec![]; expected_partition_num];
|
||||
}
|
||||
|
||||
if ranges.len() == 1 {
|
||||
let mut vec = vec![vec![]; expected_partition_num];
|
||||
vec[0] = ranges;
|
||||
return vec;
|
||||
}
|
||||
|
||||
// Sort ranges by number of rows in descending order.
|
||||
ranges.sort_by_key(|b| std::cmp::Reverse(b.num_rows));
|
||||
let mut partition_ranges = vec![vec![]; expected_partition_num];
|
||||
|
||||
#[derive(Eq, PartialEq)]
|
||||
struct HeapNode {
|
||||
num_rows: usize,
|
||||
partition_idx: usize,
|
||||
}
|
||||
|
||||
impl Ord for HeapNode {
|
||||
fn cmp(&self, other: &Self) -> std::cmp::Ordering {
|
||||
// Reverse for min-heap.
|
||||
self.num_rows.cmp(&other.num_rows).reverse()
|
||||
}
|
||||
}
|
||||
|
||||
impl PartialOrd for HeapNode {
|
||||
fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
|
||||
Some(self.cmp(other))
|
||||
}
|
||||
}
|
||||
|
||||
let mut part_heap =
|
||||
BinaryHeap::from_iter((0..expected_partition_num).map(|partition_idx| HeapNode {
|
||||
num_rows: 0,
|
||||
partition_idx,
|
||||
}));
|
||||
|
||||
// Assigns the range to the partition with the smallest number of rows.
|
||||
for range in ranges {
|
||||
let mut node = part_heap.pop().unwrap();
|
||||
let partition_idx = node.partition_idx;
|
||||
node.num_rows += range.num_rows;
|
||||
partition_ranges[partition_idx].push(range);
|
||||
part_heap.push(node);
|
||||
}
|
||||
|
||||
partition_ranges
|
||||
}
|
||||
|
||||
impl RegionScanExec {
|
||||
pub fn new(
|
||||
scanner: RegionScannerRef,
|
||||
@@ -434,6 +494,7 @@ impl RegionScanExec {
|
||||
|
||||
/// Resolves authorized logical names against this prepared scanner's physical metadata,
|
||||
/// then selects a native raw-series leaf and returns the binding it was validated against.
|
||||
/// Prepared scans may contain multiple partitions; the producer keeps each series in one partition.
|
||||
///
|
||||
/// Approval must be based on the returned [`RawSeriesRequest`], because that binding, not the
|
||||
/// requested names, describes the approved output. See [`Self::try_to_named_series_scan`]
|
||||
@@ -474,10 +535,6 @@ impl RegionScanExec {
|
||||
}
|
||||
let scanner = self.scanner.lock().unwrap();
|
||||
let metadata = scanner.metadata();
|
||||
if scanner.properties().num_partitions().max(1) != 1 {
|
||||
return None;
|
||||
}
|
||||
|
||||
let primary_key_ids = metadata
|
||||
.primary_key_columns()
|
||||
.map(|column| column.column_id)
|
||||
@@ -1302,6 +1359,83 @@ mod test {
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn assign_partition_ranges_balances_rows_deterministically() {
|
||||
let ranges: Vec<_> = [100, 200, 150, 250]
|
||||
.into_iter()
|
||||
.enumerate()
|
||||
.map(|(identifier, num_rows)| PartitionRange {
|
||||
start: common_time::Timestamp::new(
|
||||
identifier as i64,
|
||||
common_time::timestamp::TimeUnit::Second,
|
||||
),
|
||||
end: common_time::Timestamp::new(
|
||||
identifier as i64 + 1,
|
||||
common_time::timestamp::TimeUnit::Second,
|
||||
),
|
||||
num_rows,
|
||||
identifier: identifier + 1,
|
||||
})
|
||||
.collect();
|
||||
let assigned = assign_partition_ranges(ranges.clone(), 2);
|
||||
assert_eq!(
|
||||
assigned
|
||||
.iter()
|
||||
.map(|partition| {
|
||||
partition
|
||||
.iter()
|
||||
.map(|range| range.identifier)
|
||||
.collect::<Vec<_>>()
|
||||
})
|
||||
.collect::<Vec<_>>(),
|
||||
vec![vec![4, 1], vec![2, 3]]
|
||||
);
|
||||
let mut assigned_identifiers = assigned
|
||||
.iter()
|
||||
.flatten()
|
||||
.map(|range| range.identifier)
|
||||
.collect::<Vec<_>>();
|
||||
assigned_identifiers.sort_unstable();
|
||||
assert_eq!(assigned_identifiers, vec![1, 2, 3, 4]);
|
||||
|
||||
let assigned = assign_partition_ranges(ranges.clone(), 5);
|
||||
assert_eq!(
|
||||
assigned
|
||||
.iter()
|
||||
.map(|partition| {
|
||||
partition
|
||||
.iter()
|
||||
.map(|range| range.identifier)
|
||||
.collect::<Vec<_>>()
|
||||
})
|
||||
.collect::<Vec<_>>(),
|
||||
vec![vec![4], vec![1], vec![2], vec![], vec![3]]
|
||||
);
|
||||
let mut assigned_identifiers = assigned
|
||||
.iter()
|
||||
.flatten()
|
||||
.map(|range| range.identifier)
|
||||
.collect::<Vec<_>>();
|
||||
assigned_identifiers.sort_unstable();
|
||||
assert_eq!(assigned_identifiers, vec![1, 2, 3, 4]);
|
||||
|
||||
assert_eq!(assign_partition_ranges(vec![], 5), vec![vec![]; 5]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn named_series_binding_accepts_multiple_prepared_partitions() {
|
||||
let (_, plan) = dynamic_filter_fixture(5687);
|
||||
let prepared = plan
|
||||
.with_new_partitions(vec![vec![], vec![]], 2)
|
||||
.expect("prepare two scanner partitions");
|
||||
assert!(prepared.is_partition_set());
|
||||
assert!(
|
||||
prepared
|
||||
.resolve_named_series_binding("cpu", &[], &["a".to_string()])
|
||||
.is_some()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn named_series_scan_requires_prepared_single_partition() {
|
||||
let (_, plan) = dynamic_filter_fixture(5686);
|
||||
|
||||
@@ -91,6 +91,7 @@ pub struct GreptimeDbStandaloneBuilder {
|
||||
logical_batcher: Option<BatcherOptions>,
|
||||
table_batcher: BatcherOptions,
|
||||
mito_config: Option<mito2::config::MitoConfig>,
|
||||
query_options: query::options::QueryOptions,
|
||||
}
|
||||
|
||||
impl GreptimeDbStandaloneBuilder {
|
||||
@@ -112,9 +113,17 @@ impl GreptimeDbStandaloneBuilder {
|
||||
logical_batcher: None,
|
||||
table_batcher: BatcherOptions::default(),
|
||||
mito_config: None,
|
||||
query_options: query::options::QueryOptions::default(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Overrides query options for this standalone test instance.
|
||||
#[must_use]
|
||||
pub fn with_query_options(mut self, options: query::options::QueryOptions) -> Self {
|
||||
self.query_options = options;
|
||||
self
|
||||
}
|
||||
|
||||
/// Overrides the Mito configuration for this test instance.
|
||||
#[must_use]
|
||||
pub fn with_mito_config(mut self, config: mito2::config::MitoConfig) -> Self {
|
||||
@@ -316,6 +325,7 @@ impl GreptimeDbStandaloneBuilder {
|
||||
procedure_executor.clone(),
|
||||
Arc::new(ProcessManager::new(server_addr, None)),
|
||||
)
|
||||
.with_in_process_region_transport()
|
||||
.with_local_file_access(local_file_access)
|
||||
.with_plugin(plugins.clone())
|
||||
.try_build()
|
||||
@@ -389,6 +399,7 @@ impl GreptimeDbStandaloneBuilder {
|
||||
slow_query: self.slow_query_options.clone(),
|
||||
event_recorder: self.event_recorder_options.clone(),
|
||||
auto_create_table: self.auto_create_table,
|
||||
query: self.query_options.clone(),
|
||||
pending_rows_batcher: PendingRowsBatcherOptions {
|
||||
logical_table: self.logical_batcher.clone(),
|
||||
table: self.table_batcher.clone(),
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user