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

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

Browse Source 
Signed-off-by: Dennis Zhuang <killme2008@gmail.com>
This commit is contained in:
Dennis Zhuang
2026-10-27 18:20:21 +08:00
parent ba8156be47 b29a471575
commit 02eb3aa60c
54 changed files with 2269 additions and 9689 deletions
Download Patch File Download Diff File Expand all files Collapse all files
config/config.md
+2 -4
View File
@@ -140,7 +140,7 @@
| `procedure.retry_delay` | String | `500ms` | Initial retry delay of procedures, increases exponentially |
| `procedure.max_running_procedures` | Integer | `128` | Max running procedures.<br/>The maximum number of procedures that can be running at the same time.<br/>If the number of running procedures exceeds this limit, the procedure will be rejected. |
| `query` | -- | -- | The query engine options. |
| `query.experimental_native_promql` | Bool | `false` | Whether to enable experimental native PromQL execution.<br/>Default: false |
| `query.experimental_native_promql` | Bool | `false` | Enable experimental native PromQL execution for compatible queries.<br/>Requires the per-request query.experimental_native_promql=true extension.<br/>Default: false |
| `query.parallelism` | Integer | `0` | Parallelism of the query engine.<br/>Default to 0, which means the number of CPU cores. |
| `query.memory_pool_size` | String | 0 | Memory pool size for query execution operators (aggregation, sorting, join).<br/>Supports absolute size (e.g., "2GB", "4GB") or percentage of system memory (e.g., "20%").<br/>Setting it to 0 disables the limit. Defaults to 0 (unlimited); this example sets it to `50%`.<br/>When this limit is reached, queries will fail with ResourceExhausted error.<br/>NOTE: This does NOT limit memory used by table scans. |
| `query.experimental_memory_pool_policy` | String | `greedy` | Experimental memory pool allocation policy:<br/>- "greedy" (default): first-come-first-served allocation; preserves current behavior.<br/>- "fair": divides available memory among spillable operators and may spill earlier.<br/>Only effective when `memory_pool_size` is bounded (>0). |
@@ -385,7 +385,7 @@
| `meta_client.metadata_cache_ttl` | String | `10m` | TTL of the metadata cache. |
| `meta_client.metadata_cache_tti` | String | `5m` | -- |
| `query` | -- | -- | The query engine options. |
| `query.experimental_native_promql` | Bool | `false` | Whether to enable experimental native PromQL execution.<br/>Default: false |
| `query.experimental_native_promql` | Bool | `false` | Enable experimental native PromQL execution for compatible standalone queries.<br/>Requires in-process region transport and the per-request<br/>query.experimental_native_promql=true extension. Default: false |
| `query.parallelism` | Integer | `0` | Parallelism of the query engine.<br/>Default to 0, which means the number of CPU cores. |
| `query.allow_query_fallback` | Bool | `false` | Whether to allow query fallback when push down optimize fails.<br/>Default to false, meaning when push down optimize failed, return error msg |
| `query.memory_pool_size` | String | 0 | Memory pool size for query execution operators (aggregation, sorting, join).<br/>Supports absolute size (e.g., "4GB", "8GB") or percentage of system memory (e.g., "30%").<br/>Setting it to 0 disables the limit. Defaults to 0 (unlimited); this example sets it to `50%`.<br/>When this limit is reached, queries will fail with ResourceExhausted error.<br/>NOTE: This does NOT limit memory used by table scans (only applies to datanodes). |
@@ -616,7 +616,6 @@
| `wal.max_unpersisted_age` | String | `8s` | The age of the oldest unpersisted entry at which new appends stall until an upload completes, in `enqueued` mode.<br/>**It's only used when the provider is `experimental_object_store` and `ack_mode` is `enqueued`**. |
| `wal.on_corrupted_segment` | String | `skip` | What a read does with a segment that still does not decode after a second fetch, because its checksum does not match or its content disagrees with its footer entry.<br/>- `skip`: the segment is skipped and recorded as a WAL hole of its region, a metric is incremented and a warning is logged; the other regions of the object are unaffected (the default).<br/>- `fail`: the read fails, so the region does not open.<br/>**It's only used when the provider is `experimental_object_store`**. |
| `query` | -- | -- | The query engine options. |
| `query.experimental_native_promql` | Bool | `false` | Whether to enable experimental native PromQL execution.<br/>Default: false |
| `query.parallelism` | Integer | `0` | Parallelism of the query engine.<br/>Default to 0, which means the number of CPU cores. |
| `query.memory_pool_size` | String | 0 | Memory pool size for query execution operators (aggregation, sorting, join).<br/>Supports absolute size (e.g., "2GB", "4GB") or percentage of system memory (e.g., "20%").<br/>Setting it to 0 disables the limit. Defaults to 0 (unlimited); this example sets it to `50%`.<br/>When this limit is reached, queries will fail with ResourceExhausted error.<br/>NOTE: This does NOT limit memory used by table scans. |
| `query.experimental_memory_pool_policy` | String | `greedy` | Experimental memory pool allocation policy:<br/>- "greedy" (default): first-come-first-served allocation; preserves current behavior.<br/>- "fair": divides available memory among spillable operators and may spill earlier.<br/>Only effective when `memory_pool_size` is bounded (>0). |
@@ -804,7 +803,6 @@
| `tracing` | -- | -- | The tracing options. Only effect when compiled with `tokio-console` feature. |
| `tracing.tokio_console_addr` | String | Unset | The tokio console address. |
| `query` | -- | -- | -- |
| `query.experimental_native_promql` | Bool | `false` | Whether to enable experimental native PromQL execution.<br/>Default: false |
| `query.parallelism` | Integer | `1` | Parallelism of the query engine for query sent by flownode.<br/>Default to 1, so it won't use too much cpu or memory |
| `query.memory_pool_size` | String | 0 | Memory pool size for query execution operators (aggregation, sorting, join).<br/>Supports absolute size (e.g., "1GB", "2GB") or percentage of system memory (e.g., "20%").<br/>Setting it to 0 disables the limit. Defaults to 0 (unlimited); this example sets it to `50%`.<br/>When this limit is reached, queries will fail with ResourceExhausted error.<br/>NOTE: This does NOT limit memory used by table scans. |
| `query.experimental_memory_pool_policy` | String | `greedy` | Experimental memory pool allocation policy:<br/>- "greedy" (default): first-come-first-served allocation; preserves current behavior.<br/>- "fair": divides available memory among spillable operators and may spill earlier.<br/>Only effective when `memory_pool_size` is bounded (>0). |
config/datanode.example.toml
-3
View File
@@ -327,9 +327,6 @@ overwrite_entry_start_id = false
## The query engine options.
[query]
## Whether to enable experimental native PromQL execution.
## Default: false
#+ experimental_native_promql = false
## Parallelism of the query engine.
## Default to 0, which means the number of CPU cores.
parallelism = 0
config/flownode.example.toml
-3
View File
@@ -152,9 +152,6 @@ default_ratio = 1.0
[query]
## Whether to enable experimental native PromQL execution.
## Default: false
#+ experimental_native_promql = false
## Parallelism of the query engine for query sent by flownode.
## Default to 1, so it won't use too much cpu or memory
parallelism = 1
config/frontend.example.toml
+3 -2
View File
@@ -360,8 +360,9 @@ metadata_cache_tti = "5m"
## The query engine options.
[query]
## Whether to enable experimental native PromQL execution.
## Default: false
## Enable experimental native PromQL execution for compatible standalone queries.
## Requires in-process region transport and the per-request
## query.experimental_native_promql=true extension. Default: false
#+ experimental_native_promql = false
## Parallelism of the query engine.
## Default to 0, which means the number of CPU cores.
config/standalone.example.toml
+2 -1
View File
@@ -538,7 +538,8 @@ max_running_procedures = 128
## The query engine options.
[query]
## Whether to enable experimental native PromQL execution.
## Enable experimental native PromQL execution for compatible queries.
## Requires the per-request query.experimental_native_promql=true extension.
## Default: false
#+ experimental_native_promql = false
## Parallelism of the query engine.
src/common/query/src/request.rs
-3
View File
@@ -14,9 +14,6 @@
mod base64_serde;
mod initial_remote_dyn_filter_reg;
mod raw_series_request;
pub use self::raw_series_request::{RawSeriesRequest, RawSeriesTag};
use std::cell::Cell;
use std::sync::Arc;
src/common/query/src/request/raw_series_request.rs
-249
View File
@@ -1,249 +0,0 @@
// Copyright 2023 Greptime Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use std::collections::HashSet;
use std::sync::Arc;
use api::v1::SemanticType;
use datafusion_common::{DataFusionError, Result as DataFusionResult};
use datatypes::series::samples::{FieldSpec, SeriesSchema};
use serde::{Deserialize, Serialize};
use store_api::metadata::{RegionMetadata, RegionMetadataRef};
use store_api::region_engine::{BoundTag, RawLabelSource};
use store_api::storage::ColumnId;
/// A logical label mapped to its physical region column.
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct RawSeriesTag {
/// Physical region column identifier.
pub column_id: ColumnId,
/// Logical tag name.
pub name: String,
}
/// Explicit versioned binding for a single-region raw-series request.
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct RawSeriesRequest {
/// Control format version (currently 1).
pub version: u32,
/// Expected physical metadata snapshot.
pub physical_metadata: RegionMetadata,
/// Logical metric name.
pub metric_name: String,
/// Logical tags in output order.
pub logical_tags: Vec<RawSeriesTag>,
/// Physical fields in requested output order.
pub field_ids: Vec<ColumnId>,
}
impl RawSeriesRequest {
/// Binds expected output fields and labels to the actual opened scanner metadata.
pub fn bind(
&self,
actual: RegionMetadataRef,
) -> DataFusionResult<(Arc<SeriesSchema>, RawLabelSource)> {
self.validate()?;
if !same_encoding_identity(&self.physical_metadata, &actual) {
return Err(DataFusionError::Plan("Raw-series metadata mismatch".into()));
}
let fields = self
.field_ids
.iter()
.map(|id| {
actual
.column_by_id(*id)
.filter(|column| column.semantic_type == SemanticType::Field)
.map(|column| FieldSpec {
name: Arc::from(column.column_schema.name.as_str()),
})
.ok_or_else(|| DataFusionError::Plan("Invalid raw-series fields".into()))
})
.collect::<DataFusionResult<Vec<_>>>()?;
let schema = Arc::new(SeriesSchema {
fields: fields.into(),
eval_grid: None,
});
let tags = self
.logical_tags
.iter()
.map(|tag| BoundTag {
column_id: tag.column_id,
name: Arc::from(tag.name.as_str()),
})
.collect::<Vec<_>>();
let source = RawLabelSource::PrimaryKey {
metric_name: Arc::from(self.metric_name.as_str()),
physical_metadata: actual,
logical_tags: tags.into(),
};
Ok((schema, source))
}
fn validate(&self) -> DataFusionResult<()> {
if self.version != 1 || self.field_ids.is_empty() {
return Err(DataFusionError::Plan("Invalid raw-series request".into()));
}
let mut ids = HashSet::new();
for id in &self.field_ids {
let Some(column) = self.physical_metadata.column_by_id(*id) else {
return Err(DataFusionError::Plan("Invalid raw-series fields".into()));
};
if column.semantic_type != SemanticType::Field || !ids.insert(*id) {
return Err(DataFusionError::Plan("Invalid raw-series fields".into()));
}
}
Ok(())
}
}
fn same_encoding_identity(left: &RegionMetadata, right: &RegionMetadata) -> bool {
left.region_id == right.region_id
&& left.schema_version == right.schema_version
&& left.column_metadatas == right.column_metadatas
&& left.primary_key == right.primary_key
&& left.primary_key_encoding == right.primary_key_encoding
&& left.time_index_column().column_id == right.time_index_column().column_id
&& left.time_index_type() == right.time_index_type()
}
#[cfg(test)]
mod tests {
use datatypes::data_type::ConcreteDataType;
use datatypes::schema::ColumnSchema;
use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder};
use store_api::storage::RegionId;
use super::*;
/// Builds the physical metadata of a one-tag, two-field region.
fn fixture_metadata() -> RegionMetadata {
let mut builder = RegionMetadataBuilder::new(RegionId::new(1, 1));
builder
.push_column_metadata(ColumnMetadata {
column_schema: ColumnSchema::new(
"host",
ConcreteDataType::string_datatype(),
false,
),
semantic_type: SemanticType::Tag,
column_id: 1,
})
.push_column_metadata(ColumnMetadata {
column_schema: ColumnSchema::new(
"value",
ConcreteDataType::int32_datatype(),
false,
),
semantic_type: SemanticType::Field,
column_id: 2,
})
.push_column_metadata(ColumnMetadata {
column_schema: ColumnSchema::new("aux", ConcreteDataType::int64_datatype(), false),
semantic_type: SemanticType::Field,
column_id: 3,
})
.push_column_metadata(ColumnMetadata {
column_schema: ColumnSchema::new(
"ts",
ConcreteDataType::timestamp_millisecond_datatype(),
false,
),
semantic_type: SemanticType::Timestamp,
column_id: 4,
})
.primary_key(vec![1]);
builder.build().unwrap()
}
/// Builds a request asking for fields in descending column order.
fn fixture_request() -> RawSeriesRequest {
RawSeriesRequest {
version: 1,
physical_metadata: fixture_metadata(),
metric_name: "cpu".to_string(),
logical_tags: vec![RawSeriesTag {
column_id: 1,
name: "host".to_string(),
}],
field_ids: vec![3, 2],
}
}
#[test]
fn rejects_unsupported_version() {
let mut request = fixture_request();
request.version = 2;
assert!(request.validate().is_err());
}
#[test]
fn rejects_empty_duplicate_and_non_field_ids() {
let mut request = fixture_request();
request.field_ids.clear();
assert!(request.validate().is_err());
let mut request = fixture_request();
request.field_ids.push(3);
assert!(request.validate().is_err());
let mut request = fixture_request();
request.field_ids = vec![1];
assert!(request.validate().is_err());
}
#[test]
fn binds_schema_order_and_actual_metadata_arc() {
let request = fixture_request();
let mut reordered = request.clone();
reordered.field_ids.reverse();
let actual = Arc::new(reordered.physical_metadata.clone());
let (schema, source) = reordered.bind(actual.clone()).unwrap();
assert_eq!(
schema
.fields
.iter()
.map(|field| field.name.as_ref())
.collect::<Vec<_>>(),
["value", "aux"]
);
let RawLabelSource::PrimaryKey {
physical_metadata, ..
} = source
else {
panic!("expected key labels")
};
assert!(Arc::ptr_eq(&actual, &physical_metadata));
let mut mismatch = request.clone();
mismatch.physical_metadata.schema_version += 1;
assert!(mismatch.bind(actual.clone()).is_err());
let mut mismatch = request.clone();
mismatch.physical_metadata.region_id = RegionId::new(1, 2);
assert!(mismatch.bind(actual.clone()).is_err());
let mut mismatch = request.clone();
mismatch.physical_metadata.primary_key.push(2);
assert!(mismatch.bind(actual.clone()).is_err());
let mut mismatch = request.clone();
mismatch.physical_metadata.column_metadatas.swap(0, 1);
assert!(mismatch.bind(actual.clone()).is_err());
let mut mismatch = request.clone();
mismatch.physical_metadata.column_metadatas[1].column_schema =
ColumnSchema::new("value", ConcreteDataType::int64_datatype(), false);
assert!(mismatch.bind(actual.clone()).is_err());
let mut partition_only = request;
partition_only
.physical_metadata
.set_partition_expr(Some("ignored".into()));
assert!(partition_only.bind(actual).is_ok());
}
}
src/common/query/src/series_frame.rs
-4
View File
@@ -57,10 +57,6 @@ use store_api::region_engine::{
NO_BOUND_LABEL_SOURCE, RawLabelSource, RawSeriesStream, SendableRawSeriesStream,
};
mod record_batches;
pub use record_batches::{InputMapping, record_batches_to_raw_series};
#[cfg(test)]
mod tests;
src/common/query/src/series_frame/record_batches.rs
-393
View File
@@ -1,393 +0,0 @@
// Copyright 2023 Greptime Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#[cfg(test)]
mod tests;
use std::collections::HashSet;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};
use common_error::ext::BoxedError;
use common_recordbatch::{DfRecordBatch, DfSendableRecordBatchStream};
use datafusion::arrow::array::{
Array, ArrayRef, Float64Array, StringArray, StructArray, TimestampMillisecondArray,
};
use datafusion::arrow::buffer::Buffer;
use datafusion::arrow::datatypes::{DataType, SchemaRef, TimeUnit};
use datafusion::arrow::row::{RowConverter, Rows, SortField};
use datafusion_common::{DataFusionError, Result as DfResult};
use datatypes::arrow_array::string_array_value_at_index;
use datatypes::series::raw::{RawSeriesChunk, RawTimestamps};
use datatypes::series::samples::{FieldSpec, Samples, SeriesFields, SeriesSchema};
use datatypes::series::{Label, SeriesLabels, SeriesOrdering};
use futures_util::Stream;
use snafu::Location;
use store_api::region_engine::{RawLabelSource, RawSeriesStream, SendableRawSeriesStream};
use crate::error::Error;
use crate::native_histogram::{histogram_samples_from_arrow, native_histogram_arrow_type};
use crate::series_frame::{codec_error, invalid};
/// Explicit mapping from selected record-batch columns to raw-series identity and payload.
///
/// `time_index` selects a timezone-naive millisecond timestamp column. `metric_name` is an
/// optional constant logical metric name, not a source column. `tags` binds all logical tag
/// names to their selected source columns; supported tag columns are Utf8, LargeUtf8, Utf8View,
/// or integer-key string dictionaries. `fields` binds selected source columns in output field
/// order; supported field columns are Float64 and the native-histogram struct type.
#[derive(Clone, Debug)]
pub struct InputMapping {
/// Source column containing timezone-naive millisecond timestamps.
pub time_index: usize,
/// Optional constant metric name; `None` omits the metric name.
pub metric_name: Option<Arc<str>>,
/// Complete logical tag mapping, as `(tag_name, source_column_index)` pairs.
pub tags: Vec<(Arc<str>, usize)>,
/// Selected payload columns in output order, paired with their logical field specs.
pub fields: Vec<(FieldSpec, usize)>,
}
/// Converts ordinary record batches to contiguous raw-series chunks.
///
/// The already-executed child must keep each complete identity contiguous and timestamps
/// ascending. The adapter checks adjacent runs only; it cannot prove global absence of A/B/A.
/// Encoded read keys are opaque identities comparable only within this converter stream; they
/// are not canonical cross-region or wire encodings. Float fields share their source buffers;
/// native-histogram struct fields use the existing histogram bridge.
pub fn record_batches_to_raw_series(
source: DfSendableRecordBatchStream,
mapping: InputMapping,
) -> DfResult<SendableRawSeriesStream> {
let arrow_schema = source.schema();
let (mapping, schema, converter) = validate_mapping(&arrow_schema, mapping)?;
Ok(Box::pin(InputStream {
source: Some(source),
arrow_schema,
mapping,
schema,
converter,
batch: None,
row: 0,
pending: None,
label_source: RawLabelSource::Explicit,
}))
}
fn validate_mapping(
schema: &SchemaRef,
mut mapping: InputMapping,
) -> DfResult<(InputMapping, Arc<SeriesSchema>, RowConverter)> {
let mut used = HashSet::new();
let mut use_index = |index: usize| -> DfResult<()> {
if index >= schema.fields().len() {
return Err(invalid("mapped column index is out of bounds"));
}
if !used.insert(index) {
return Err(invalid("source column is mapped to multiple roles"));
}
Ok(())
};
use_index(mapping.time_index)?;
if schema.field(mapping.time_index).data_type()
!= &DataType::Timestamp(TimeUnit::Millisecond, None)
{
return Err(invalid("time column must be Timestamp(Millisecond, None)"));
}
for (_, index) in &mapping.tags {
use_index(*index)?;
let ty = schema.field(*index).data_type();
if !string_type(ty) && !integer_key_type(ty) {
return Err(invalid(
"tag column must be string or integer-key string dictionary",
));
}
if matches!(ty, DataType::Dictionary(_, value) if !string_type(value)) {
return Err(invalid("tag dictionary values must be strings"));
}
}
if mapping.fields.is_empty() {
return Err(invalid("at least one field is required"));
}
let mut field_names = HashSet::new();
for (spec, index) in &mapping.fields {
use_index(*index)?;
if !field_names.insert(spec.name.clone()) {
return Err(invalid("duplicate logical field name"));
}
let ty = schema.field(*index).data_type();
if ty != &DataType::Float64 && ty != &native_histogram_arrow_type() {
return Err(invalid("field must be Float64 or native histogram"));
}
}
// SeriesLabels validates logical label names and duplicate tags through its canonical path.
SeriesLabels::try_new(
mapping.metric_name.clone(),
mapping
.tags
.iter()
.map(|(name, _)| Label {
name: name.clone(),
value: Arc::from(""),
})
.collect(),
)
.map_err(codec_error)?;
mapping.tags.sort_by(|a, b| a.0.cmp(&b.0));
let raw_schema = Arc::new(SeriesSchema {
fields: mapping
.fields
.iter()
.map(|(spec, _)| spec.clone())
.collect::<Vec<_>>()
.into(),
eval_grid: None,
});
let mut sorts = Vec::with_capacity(1 + mapping.tags.len() * 2);
sorts.push(SortField::new(DataType::Utf8));
for (_, index) in &mapping.tags {
sorts.push(SortField::new(DataType::Utf8));
sorts.push(SortField::new(schema.field(*index).data_type().clone()));
}
let converter = RowConverter::new(sorts)?;
Ok((mapping, raw_schema, converter))
}
fn string_type(ty: &DataType) -> bool {
matches!(
ty,
DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View
)
}
fn integer_key_type(ty: &DataType) -> bool {
matches!(ty, DataType::Dictionary(key, value) if matches!(key.as_ref(), DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 | DataType::UInt8 | DataType::UInt16 | DataType::UInt32 | DataType::UInt64) && string_type(value))
}
fn stream_error(error: DataFusionError) -> BoxedError {
BoxedError::new(Error::GeneralDataFusion {
error,
location: Location::default(),
})
}
struct InputStream {
source: Option<DfSendableRecordBatchStream>,
arrow_schema: SchemaRef,
mapping: InputMapping,
schema: Arc<SeriesSchema>,
converter: RowConverter,
batch: Option<(DfRecordBatch, Rows)>,
row: usize,
pending: Option<RawSeriesChunk>,
label_source: RawLabelSource,
}
impl RawSeriesStream for InputStream {
fn schema(&self) -> &Arc<SeriesSchema> {
&self.schema
}
fn label_source(&self) -> &RawLabelSource {
&self.label_source
}
fn ordering(&self) -> SeriesOrdering {
SeriesOrdering::Contiguous
}
}
impl InputStream {
fn finish(&mut self) {
self.source = None;
self.batch = None;
self.row = 0;
self.pending = None;
}
fn batch_key_columns(&self, batch: &DfRecordBatch) -> Vec<ArrayRef> {
let n = batch.num_rows();
let metric: ArrayRef = match &self.mapping.metric_name {
Some(name) => Arc::new(StringArray::from_iter_values(std::iter::repeat_n(
name.as_ref(),
n,
))),
None => Arc::new(StringArray::new_null(n)),
};
let mut cols = vec![metric];
for (tag, index) in &self.mapping.tags {
cols.push(Arc::new(StringArray::from_iter_values(
std::iter::repeat_n(tag.as_ref(), n),
)));
cols.push(batch.column(*index).clone());
}
cols
}
fn read_chunk(
&self,
batch: &DfRecordBatch,
start: usize,
end: usize,
key: Vec<u8>,
) -> DfResult<RawSeriesChunk> {
let len = end - start;
let time = batch
.column(self.mapping.time_index)
.as_any()
.downcast_ref::<TimestampMillisecondArray>()
.ok_or_else(|| invalid("invalid timestamp array"))?;
if time.slice(start, len).null_count() != 0 {
return Err(invalid("time values must not be null"));
}
let axis = RawTimestamps::try_new(time.values().slice(start, len)).map_err(codec_error)?;
let mut fields = Vec::with_capacity(self.mapping.fields.len());
for (_, index) in &self.mapping.fields {
let array = batch.column(*index);
if array.data_type() == &DataType::Float64 {
let values = array
.as_any()
.downcast_ref::<Float64Array>()
.ok_or_else(|| invalid("invalid Float64 array"))?;
fields.push(Samples::from_float(values.slice(start, len)));
} else {
let hist = array
.as_any()
.downcast_ref::<StructArray>()
.ok_or_else(|| invalid("invalid histogram array"))?;
fields.push(histogram_samples_from_arrow(&hist.slice(start, len))?);
}
}
let fields = SeriesFields::try_from_vec(fields).map_err(codec_error)?;
let metric_name = self.mapping.metric_name.clone();
let mut labels = Vec::new();
for (name, index) in &self.mapping.tags {
if let Some(value) = string_array_value_at_index(batch.column(*index), start) {
labels.push(Label {
name: name.clone(),
value: Arc::from(value),
});
}
}
let metric = SeriesLabels::try_new(metric_name, labels).map_err(codec_error)?;
RawSeriesChunk::try_new(
&self.schema,
Buffer::from(key),
Some(metric),
axis,
fields,
true,
)
.map_err(codec_error)
}
}
impl Stream for InputStream {
type Item = Result<RawSeriesChunk, BoxedError>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
const MAX_EMPTY_BATCHES: usize = 32;
let this = &mut *self;
let mut empties = 0;
loop {
if this.batch.is_none() {
match this.source.as_mut().map(|s| s.as_mut().poll_next(cx)) {
None | Some(Poll::Ready(None)) => {
if let Some(mut final_chunk) = this.pending.take() {
final_chunk.is_last = true;
this.finish();
return Poll::Ready(Some(Ok(final_chunk)));
}
this.finish();
return Poll::Ready(None);
}
Some(Poll::Pending) => return Poll::Pending,
Some(Poll::Ready(Some(Err(error)))) => {
this.finish();
return Poll::Ready(Some(Err(stream_error(error))));
}
Some(Poll::Ready(Some(Ok(batch)))) => {
if batch.schema() != this.arrow_schema {
this.finish();
return Poll::Ready(Some(Err(stream_error(invalid(
"record batch schema mismatch",
)))));
}
if batch.num_rows() == 0 {
empties += 1;
if empties >= MAX_EMPTY_BATCHES {
cx.waker().wake_by_ref();
return Poll::Pending;
}
continue;
}
let columns = this.batch_key_columns(&batch);
let rows = match this.converter.convert_columns(&columns) {
Ok(rows) => rows,
Err(error) => {
this.finish();
return Poll::Ready(Some(Err(stream_error(error.into()))));
}
};
this.batch = Some((batch, rows));
this.row = 0;
}
}
}
let Some((batch, rows)) = this.batch.as_ref() else {
continue;
};
let start = this.row;
let key = rows.row(start).as_ref().to_vec();
if this
.pending
.as_ref()
.is_some_and(|previous| previous.primary_key.as_slice() != key.as_slice())
&& let Some(mut complete) = this.pending.take()
{
complete.is_last = true;
return Poll::Ready(Some(Ok(complete)));
}
let mut end = start + 1;
while end < batch.num_rows() && rows.row(end).as_ref() == key {
end += 1;
}
let chunk = match this.read_chunk(batch, start, end, key) {
Ok(chunk) => chunk,
Err(error) => {
this.finish();
return Poll::Ready(Some(Err(stream_error(error))));
}
};
this.row = end;
if end == batch.num_rows() {
this.batch = None;
}
if let Some(mut previous) = this.pending.take() {
let Some(last) = previous.time.timestamp(previous.time.len() - 1) else {
this.finish();
return Poll::Ready(Some(Err(stream_error(invalid(
"empty pending raw chunk",
)))));
};
let Some(first) = chunk.time.timestamp(0) else {
this.finish();
return Poll::Ready(Some(Err(stream_error(invalid("empty raw chunk")))));
};
if first <= last {
this.finish();
return Poll::Ready(Some(Err(stream_error(invalid(
"timestamps are not increasing across adjacent batches",
)))));
}
previous.is_last = false;
this.pending = Some(chunk);
return Poll::Ready(Some(Ok(previous)));
}
this.pending = Some(chunk);
}
}
}
src/common/query/src/series_frame/record_batches/tests.rs
-1263
View File
File diff suppressed because it is too large Load Diff
src/common/query/src/series_frame/tests/evaluated.rs
+29 -16
View File
@@ -388,7 +388,16 @@ async fn evaluated_pack_drains_large_sparse_chunk_before_polling_again() {
let expected: Vec<Option<f64>> = (0..positions.len())
.map(|index| Some(index as f64))
.collect();
let chunk = series_on(&schema, &spans, expected.clone(), labels("large"));
let mut chunk = series_on(&schema, &spans, expected.clone(), labels("large"));
chunk.fields = SeriesFields::one(
Samples::try_from_chunks(
expected
.iter()
.map(|value| Samples::from_float(Float64Array::from(vec![*value])))
.collect(),
)
.unwrap(),
);
let polls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let dropped = Arc::new(std::sync::atomic::AtomicBool::new(false));
let upstream = observed_source(vec![Ok((chunk, true))], polls.clone(), dropped.clone());
@@ -430,7 +439,15 @@ async fn evaluated_pack_drop_releases_unfinished_large_chunk() {
let schema = large_evaluated_schema();
let positions: Vec<usize> = (0..1600).step_by(2).collect();
let spans: Vec<(usize, usize)> = positions.iter().map(|position| (*position, 1)).collect();
let chunk = series_on(&schema, &spans, vec![Some(1.0); spans.len()], metric);
let mut chunk = series_on(&schema, &spans, vec![Some(1.0); spans.len()], metric);
chunk.fields = SeriesFields::one(
Samples::try_from_chunks(
(0..positions.len())
.map(|_| Samples::from_float(Float64Array::from(vec![Some(1.0)])))
.collect(),
)
.unwrap(),
);
drop(identity);
let polls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let dropped = Arc::new(std::sync::atomic::AtomicBool::new(false));
@@ -475,8 +492,7 @@ async fn evaluated_pack_error_discards_buffer_and_releases_source() {
assert!(stream.next().await.is_none());
}
/// Rows split across batches keep their order and finality: every position appears once and in
/// order, and only the last row of the source's final chunk closes the chain.
/// Dense source chunks coalesce to one frame row each while preserving order and finality across batches.
#[tokio::test]
async fn multi_chunk_multi_batch_stream_preserves_order_and_finality() {
let metric = labels("d");
@@ -493,13 +509,11 @@ async fn multi_chunk_multi_batch_stream_preserves_order_and_finality() {
),
])
.await;
assert!(
batches.len() > 1,
"the default row bound splits 310 rows into several batches"
);
assert_eq!(batches.len(), 1, "both dense chunks fit in one batch");
assert_eq!(
batches.iter().map(RecordBatch::num_rows).sum::<usize>(),
310
2,
"one frame row per dense chunk"
);
let declared = batches[0].schema();
@@ -1045,10 +1059,10 @@ async fn multi_field_partial_nulls_are_not_nan() {
)],
)
.await;
// One row per carried run, so the hole between the runs is not what the frame trades away.
// Sparse runs share one frame row; the grid hole is still represented by separate spans.
assert_eq!(batches.len(), 1);
assert_eq!(batches[0].num_rows(), 2);
// Both fields are columns of every row, and both carry float values of their own.
assert_eq!(batches[0].num_rows(), 1);
// Both fields are columns of the row, and both carry float values of their own.
for row in 0..batches[0].num_rows() {
for field in [0, 1] {
let payload = frame_field_payload(&batches[0], field);
@@ -1061,13 +1075,12 @@ async fn multi_field_partial_nulls_are_not_nan() {
let declared = batches[0].schema();
let rows = drain(unpack_as(batches, declared, schema)).await;
assert_eq!(rows.len(), 2, "one decoded chunk per frame row");
assert_eq!(rows.len(), 1, "one decoded chunk per frame row");
assert!(
rows.iter().all(|(row, _)| row.metric == metric),
"both rows name the same series"
"the row names the same series"
);
assert!(!rows[0].1, "only the last row closes the chain");
assert!(rows[1].1);
assert!(rows[0].1, "the sole coalesced row closes the chain");
// Concatenated in stream order, exactly as the producer's single chunk was laid out.
let decoded_positions: Vec<usize> = rows.iter().flat_map(|(row, _)| positions(row)).collect();
assert_eq!(
src/common/recordbatch/src/adapter.rs
+12 -1
View File
@@ -1425,7 +1425,12 @@ mod test {
let waker = futures::task::noop_waker_ref();
let mut cx = Context::from_waker(waker);
for status in [StatusCode::Internal, StatusCode::RequestOutdated] {
for status in [
StatusCode::InvalidArguments,
StatusCode::Internal,
StatusCode::Unknown,
StatusCode::RequestOutdated,
] {
// Unique per case, so a match proves the *original* error survived.
let marker = format!("df-record-batch-stream-adapter-typed-error-marker-{status}");
@@ -1491,6 +1496,12 @@ mod test {
panic!("expected a poll stream error, got {roundtrip_error:?}");
};
assert_eq!(roundtrip_error.status_code(), status);
let expected_output = match status {
StatusCode::Internal => "Internal error: 1003".to_string(),
StatusCode::Unknown => "Internal error: 1000".to_string(),
_ => marker.clone(),
};
assert_eq!(roundtrip_error.output_msg(), expected_output);
assert_typed_external(
df_error,
status,
src/common/recordbatch/src/error.rs
+14
View File
@@ -314,6 +314,11 @@ mod tests {
location: Location::default(),
};
assert_eq!(error.status_code(), status);
let expected_output = match status {
StatusCode::Unknown => "Internal error: 1000",
_ => "neutral error",
};
assert!(error.output_msg().contains(expected_output));
}
}
@@ -361,6 +366,15 @@ mod tests {
location: Location::default(),
};
assert_eq!(error.status_code(), StatusCode::EngineExecuteQuery);
if matches!(
&error,
Error::PollStream {
error: DataFusionError::Internal(_),
..
}
) {
assert!(error.output_msg().contains("neutral internal error"));
}
}
}
}
src/datatypes/src/series/frame.rs
+202 -110
View File
@@ -810,6 +810,15 @@ impl SeriesFrameBuilder {
self.rows.len() >= self.max_rows || self.bytes >= self.max_bytes
}
/// Bytes available to an evaluated row that is being assembled before cursor advancement.
fn available_bytes(&self) -> Option<usize> {
if self.rows.len() >= self.max_rows || self.bytes >= self.max_bytes {
None
} else {
Some(self.max_bytes - self.bytes)
}
}
/// Whether `row` may join the buffered rows without crossing either bound.
///
/// An empty builder accepts any individually measured row; a group total that would overflow
@@ -1001,7 +1010,7 @@ impl SeriesFrameCursor {
loop {
let row = match self.pending.take() {
Some(row) => row,
None => match self.next_row()? {
None => match self.next_row(builder.available_bytes())? {
Some(row) => {
self.pending = Some(row);
continue;
@@ -1027,10 +1036,10 @@ impl SeriesFrameCursor {
}
/// Builds and measures the next row, or `None` when the chunk is exhausted.
fn next_row(&mut self) -> Result<Option<FrameRow>> {
fn next_row(&mut self, available_bytes: Option<usize>) -> Result<Option<FrameRow>> {
match &mut self.walk {
ChunkWalk::Raw(walk) => walk.next_row(),
ChunkWalk::Evaluated(walk) => walk.next_row(),
ChunkWalk::Evaluated(walk) => walk.next_row(available_bytes),
}
}
}
@@ -1092,11 +1101,10 @@ impl RawWalk {
impl EvaluatedWalk {
/// Builds the next row, or `None` when the chunk is exhausted.
///
/// A row covers one `(run of carried positions, sample segment)` intersection, which is why its
/// positions are always one span: the holes between runs are not rows, so nothing is ever filled
/// in for a grid position the chunk does not hold. The `is_last` a producer declared is carried
/// by the chunk's final row only, so a consumer that sees it knows the chain ended exactly there.
fn next_row(&mut self) -> Result<Option<FrameRow>> {
/// A row may cover several run/segment intersections, but never crosses a dense segment
/// boundary. Grid holes are represented by separate spans while the sample slice stays
/// contiguous. Cursor state is committed only after the complete row has been measured.
fn next_row(&mut self, available_bytes: Option<usize>) -> Result<Option<FrameRow>> {
let Self {
series,
is_last,
@@ -1108,7 +1116,6 @@ impl EvaluatedWalk {
closing,
} = self;
if *closing {
*closing = false;
// A chunk that carries no position still owes its consumer one row: the row that says the
// chain ends here. It holds no grid position and no sample, and each field is kept as the
// producer built it - an empty float stays float, an empty histogram stays histogram - so a
@@ -1127,53 +1134,101 @@ impl EvaluatedWalk {
bytes: 0,
};
row.bytes = measure_row(&row)?;
*closing = false;
return Ok(Some(row));
}
let Some(current) = runs.get(*run) else {
let Some(_) = runs.get(*run) else {
return Ok(None);
};
let remaining = current.len - *offset;
let mut take = remaining;
for (index, field) in series.fields.iter().enumerate() {
let part = segment(field, segments[index])?;
take = take.min(part.len().saturating_sub(offsets[index]));
}
if take == 0 {
return invalid_series("empty sample segment");
}
let row_offset = current.row + *offset;
let mut samples = Vec::with_capacity(series.fields.len());
for (index, field) in series.fields.iter().enumerate() {
let part = segment(field, segments[index])?;
samples.push(part.slice(offsets[index], take)?);
}
let mut row = FrameRow {
key: None,
metric: Some(series.metric.clone()),
is_last: *is_last && *run + 1 == runs.len() && take == remaining,
time: RowTime::Spans(vec![GridSpan::new(current.position + *offset, take)]),
samples,
bytes: 0,
};
row.bytes = measure_row(&row)?;
*offset += take;
if *offset == current.len {
*run += 1;
*offset = 0;
}
for (index, field) in series.fields.iter().enumerate() {
let length = segment(field, segments[index])?.len();
offsets[index] += take;
if offsets[index] == length {
segments[index] += 1;
offsets[index] = 0;
let mut next_run = *run;
let mut next_offset = *offset;
let mut next_segments = segments.clone();
let mut next_offsets = offsets.clone();
let mut total = 0usize;
let mut row: Option<FrameRow> = None;
while let Some(current) = runs.get(next_run) {
let remaining = current.len - next_offset;
let mut take = remaining;
let mut segment_remaining_lengths = Vec::with_capacity(series.fields.len());
for (index, field) in series.fields.iter().enumerate() {
let part = segment(field, next_segments[index])?;
let segment_remaining = part.len().saturating_sub(next_offsets[index]);
segment_remaining_lengths.push(segment_remaining);
take = take.min(segment_remaining);
}
let reaches_segment_boundary = segment_remaining_lengths.contains(&take);
if take == 0 {
return invalid_series("empty sample segment");
}
let proposed_total = add_checked(total, take)?;
let sample_growth = sample_range_growth(&series.fields, segments, total, take)?;
if let Some(row) = &mut row {
let RowTime::Spans(row_spans) = &mut row.time else {
unreachable!("evaluated row has spans")
};
let proposed_span_count = row_spans
.len()
.checked_add(1)
.ok_or_else(|| err("raw series frame offset overflow"))?;
checked_offset(proposed_span_count)?;
checked_offset(proposed_total)?;
let proposed_bytes = add_checked(
add_checked(row.bytes, add_checked(SPAN_CHARGE, 2)?)?,
sample_growth,
)?;
if available_bytes.is_some_and(|available| proposed_bytes > available) {
break;
}
row_spans.push(GridSpan::new(current.position + next_offset, take));
row.is_last = *is_last && next_run + 1 == runs.len() && take == remaining;
row.bytes = proposed_bytes;
} else {
let mut candidate = FrameRow {
key: None,
metric: Some(series.metric.clone()),
is_last: *is_last && next_run + 1 == runs.len() && take == remaining,
time: RowTime::Spans(vec![GridSpan::new(current.position + next_offset, take)]),
samples: Vec::new(),
bytes: 0,
};
candidate.bytes = add_checked(measure_row(&candidate)?, sample_growth)?;
row = Some(candidate);
}
total = proposed_total;
next_offset += take;
if next_offset == current.len {
next_run += 1;
next_offset = 0;
}
for (index, field) in series.fields.iter().enumerate() {
let length = segment(field, next_segments[index])?.len();
next_offsets[index] += take;
if next_offsets[index] == length {
next_segments[index] += 1;
next_offsets[index] = 0;
}
}
// An individually oversized atomic intersection must be emitted alone.
if available_bytes.is_none()
|| available_bytes.is_some_and(|available| row.as_ref().unwrap().bytes > available)
|| reaches_segment_boundary
{
break;
}
}
debug_assert!(
row_offset + take <= series.time.len(),
"an evaluated row exceeds the axis it walks"
);
Ok(Some(row))
if let Some(row) = &mut row {
row.samples = series
.fields
.iter()
.enumerate()
.map(|(index, field)| segment(field, segments[index])?.slice(offsets[index], total))
.collect::<Result<Vec<_>>>()?;
}
*run = next_run;
*offset = next_offset;
segments.clone_from(&next_segments);
offsets.clone_from(&next_offsets);
Ok(row)
}
}
@@ -1260,17 +1315,101 @@ impl RowTime {
}
}
/// Per-row encoded-size measure, computed once per row when the cursor produces it.
///
/// The measure dominates every cumulative offset domain the encoder writes: binary/UTF-8 child
/// bytes, list element counts, bitmap bytes, and the visible extent of each histogram table list.
/// A group total of `min(max_bytes, i32::MAX)` therefore guarantees that all Arrow offsets and
/// lengths stay representable, without tracking a separate offset counter. This is a bound used
/// for batching work, not the exact encoded size.
/// Measures the added dense positions without slicing Arrow arrays or rescanning histogram tables.
fn sample_range_growth(
fields: &SeriesFields,
segments: &[usize],
previous: usize,
added: usize,
) -> Result<usize> {
let end = previous
.checked_add(added)
.ok_or_else(|| err("raw series frame offset overflow"))?;
checked_offset(end)?;
let bitmap_growth = bitmap_bytes(end).saturating_sub(bitmap_bytes(previous));
fields
.iter()
.enumerate()
.try_fold(0usize, |bytes, (index, field)| {
let part = segment(field, segments[index])?;
let (value_width, bitmap_count, table_charge) = match part
.dense()
.ok_or_else(|| err("sample segment is not dense"))?
{
DenseSamplesRef::Float(_) => (8, 1, None),
DenseSamplesRef::Histogram(values) => (4, 1, Some(values)),
DenseSamplesRef::Mixed(values) => (12, 2, Some(values.histograms())),
};
let mut charge = mul_checked(added, value_width)?;
charge = add_checked(charge, mul_checked(bitmap_growth, bitmap_count)?)?;
if previous == 0 {
charge = add_checked(charge, FIELD_CHARGE)?;
if let Some(values) = table_charge {
charge = add_checked(charge, histogram_table_charge(values)?)?;
}
}
add_checked(bytes, charge)
})
}
fn histogram_table_charge(values: &HistogramSamples) -> Result<usize> {
let table = values.batch().as_arrow();
checked_offset(table.len())?;
let mut bytes = add_checked(
HISTOGRAM_TABLE_CHARGE,
mul_checked(table.len(), HISTOGRAM_ROW_CHARGE)?,
)?;
for (position, width) in HISTOGRAM_LIST_WIDTHS.iter().enumerate() {
let list = table
.column(7 + position)
.as_any()
.downcast_ref::<ListArray>()
.ok_or_else(|| err("histogram table schema mismatch"))?;
// Charge the visible offset range, not the backing child array: the table may be sliced.
let offsets = list.value_offsets();
let first = offsets
.first()
.copied()
.ok_or_else(|| err("histogram table schema mismatch"))?;
let last = offsets
.get(list.len())
.copied()
.ok_or_else(|| err("histogram table schema mismatch"))?;
let span = last
.checked_sub(first)
.and_then(|span| usize::try_from(span).ok())
.ok_or_else(|| err("histogram table offsets are invalid"))?;
bytes = add_checked(bytes, mul_checked(span, *width)?)?;
}
Ok(bytes)
}
fn sample_charges(samples: &[Samples]) -> Result<usize> {
samples.iter().try_fold(0usize, |bytes, sample| {
add_checked(bytes, sample_charge(sample)?)
})
}
fn sample_charge(sample: &Samples) -> Result<usize> {
let mut bytes = FIELD_CHARGE;
match sample
.dense()
.ok_or_else(|| err("sample segment is not dense"))?
{
DenseSamplesRef::Float(values) => add_checked(bytes, float_charge(values.len())?),
DenseSamplesRef::Histogram(values) => add_checked(bytes, histogram_charge(values)?),
DenseSamplesRef::Mixed(values) => {
bytes = add_checked(bytes, float_charge(values.floats().len())?)?;
add_checked(bytes, histogram_charge(values.histograms())?)
}
}
}
fn measure_row(row: &FrameRow) -> Result<usize> {
let time_len = row.time.len()?;
let key_len = row.key.as_ref().map(Buffer::len).unwrap_or(0);
checked_offset(key_len)?;
checked_offset(row.time.len()?)?;
checked_offset(time_len)?;
let mut bytes = add_checked(ROW_CHARGE, key_len)?;
bytes = match &row.time {
RowTime::Timestamps(timestamps) => add_checked(bytes, mul_checked(timestamps.len(), 8)?)?,
@@ -1280,24 +1419,7 @@ fn measure_row(row: &FrameRow) -> Result<usize> {
mul_checked(spans.len(), 2)?,
)?,
};
for samples in &row.samples {
bytes = add_checked(bytes, FIELD_CHARGE)?;
match samples
.dense()
.ok_or_else(|| err("sample segment is not dense"))?
{
DenseSamplesRef::Float(values) => {
bytes = add_checked(bytes, float_charge(values.len())?)?
}
DenseSamplesRef::Histogram(values) => {
bytes = add_checked(bytes, histogram_charge(values)?)?
}
DenseSamplesRef::Mixed(values) => {
bytes = add_checked(bytes, float_charge(values.floats().len())?)?;
bytes = add_checked(bytes, histogram_charge(values.histograms())?)?;
}
}
}
bytes = add_checked(bytes, sample_charges(&row.samples)?)?;
if let Some(metric) = &row.metric {
if let Some(name) = metric.name() {
checked_offset(name.len())?;
@@ -1329,40 +1451,10 @@ fn float_charge(len: usize) -> Result<usize> {
fn histogram_charge(values: &HistogramSamples) -> Result<usize> {
let indices = values.indices();
checked_offset(indices.len())?;
let mut bytes = add_checked(mul_checked(indices.len(), 4)?, bitmap_bytes(indices.len()))?;
let table = values.batch().as_arrow();
checked_offset(table.len())?;
bytes = add_checked(
bytes,
add_checked(
HISTOGRAM_TABLE_CHARGE,
mul_checked(table.len(), HISTOGRAM_ROW_CHARGE)?,
)?,
)?;
for (position, width) in HISTOGRAM_LIST_WIDTHS.iter().enumerate() {
let list = table
.column(7 + position)
.as_any()
.downcast_ref::<ListArray>()
.ok_or_else(|| err("histogram table schema mismatch"))?;
// `value_offsets` is already a view of the visible slice; the retained child array may be
// wider than the offsets span, so charge the visible extent rather than its length.
let offsets = list.value_offsets();
let first = offsets
.first()
.copied()
.ok_or_else(|| err("histogram table schema mismatch"))?;
let last = offsets
.get(list.len())
.copied()
.ok_or_else(|| err("histogram table schema mismatch"))?;
let span = last
.checked_sub(first)
.and_then(|span| usize::try_from(span).ok())
.ok_or_else(|| err("histogram table offsets are invalid"))?;
bytes = add_checked(bytes, mul_checked(span, *width)?)?;
}
Ok(bytes)
add_checked(
add_checked(mul_checked(indices.len(), 4)?, bitmap_bytes(indices.len()))?,
histogram_table_charge(values)?,
)
}
/// Encodes buffered rows as one batch with direct N-row Arrow array construction.
src/datatypes/src/series/frame/tests.rs
+515 -1
View File
@@ -249,6 +249,28 @@ fn assert_float_bits(actual: &Samples, expected: &[Option<f64>]) {
}
}
fn assert_evaluated_sample_equal(actual: SampleRef<'_>, expected: SampleRef<'_>, index: usize) {
match (actual, expected) {
(SampleRef::Missing, SampleRef::Missing) => {}
(SampleRef::Float(actual), SampleRef::Float(expected)) => {
assert_eq!(actual.to_bits(), expected.to_bits(), "position {index}");
}
(SampleRef::Histogram(actual), SampleRef::Histogram(expected)) => {
assert_eq!(
actual.sum.to_bits(),
expected.sum.to_bits(),
"position {index}"
);
assert_eq!(
actual.count.to_bits(),
expected.count.to_bits(),
"position {index}"
);
}
_ => panic!("sample kind changed at position {index}"),
}
}
/// Packs one owned chunk through the public cursor/builder pair, mirroring the explicit
/// single-chunk wrapper: every batch the walk produces, then the final flush.
fn pack_chunk_with_budget(
@@ -1261,6 +1283,498 @@ fn histogram_payload_and_present_all_null_representation_survive() {
);
}
fn evaluated_series(schema: &SeriesSchema, spans: Vec<GridSpan>, fields: SeriesFields) -> Series {
Series::try_new(
schema,
metric(Some("cpu"), &[("host", "a")]),
TimeAxis::from_grid_spans(schema.eval_grid.clone().unwrap(), spans).unwrap(),
fields,
)
.unwrap()
}
#[test]
fn evaluated_sparse_runs_share_row_preserve_gaps_and_finality() {
let base = test_schema(&["value"]);
let schema = SeriesSchema {
fields: base.fields,
eval_grid: Some(Arc::new(
crate::series::time::EvalGrid::try_new(
crate::series::time::GridSpec::try_new(0, 10, 10).unwrap(),
)
.unwrap(),
)),
};
let values = Float64Array::from(vec![Some(1.0), None, Some(3.0), Some(4.0)]);
let series = evaluated_series(
&schema,
vec![GridSpan::new(0, 2), GridSpan::new(5, 2)],
SeriesFields::one(Samples::from_float(values.clone())),
);
let mut cursor = SeriesFrameCursor::try_new_evaluated(&schema, series, true).unwrap();
let mut builder =
SeriesFrameBuilder::try_new(&schema, FrameLabelMode::Explicit, None, 8, 4096).unwrap();
assert!(cursor.append_to(&mut builder).unwrap().is_none());
let batch = builder.flush().unwrap().unwrap();
assert_eq!(batch.num_rows(), 1);
assert_eq!(
batch.column(1).len(),
1,
"labels are encoded once for the row"
);
let decoded = SeriesFrameDecoder::try_new(&schema, FrameLabelMode::Explicit, None)
.unwrap()
.decode_evaluated(&batch)
.unwrap();
assert_eq!(decoded.len(), 1);
let (decoded, last) = &decoded[0];
assert!(*last);
assert_eq!(
decoded
.time
.runs()
.map(|run| (run.position, run.len))
.collect::<Vec<_>>(),
vec![(0, 2), (5, 2)]
);
let values = decoded.fields.get(0).unwrap().dense().unwrap();
let DenseSamplesRef::Float(values) = values else {
panic!("expected float values")
};
assert_eq!(
values,
&Float64Array::from(vec![Some(1.0), None, Some(3.0), Some(4.0)])
);
assert_eq!(decoded.time.row_of(3), None);
}
#[test]
fn evaluated_incremental_measure_matches_full_measure_at_exact_byte_budget() {
let base = test_schema(&["value"]);
let schema = SeriesSchema {
fields: base.fields,
eval_grid: Some(Arc::new(
crate::series::time::EvalGrid::try_new(
crate::series::time::GridSpec::try_new(0, 1, 16).unwrap(),
)
.unwrap(),
)),
};
let metric = metric(Some("cpu"), &[("host", "a")]);
let floats = Float64Array::from(
(0..9)
.map(|index| (index != 3 && index != 6).then_some(index as f64))
.collect::<Vec<_>>(),
);
let table = histogram_batch();
let indices = UInt32Array::from(
(0..9)
.map(|index| (index != 3 && index != 6).then_some(0))
.collect::<Vec<_>>(),
);
let mixed_indices = UInt32Array::from(
(0..9)
.map(|index| (index == 3 || index == 6).then_some(0))
.collect::<Vec<_>>(),
);
let cases = vec![
Samples::from_float(floats.clone()),
Samples::try_from_histogram(indices.clone(), table.clone()).unwrap(),
Samples::try_from_mixed(
floats,
crate::series::samples::HistogramSamples::try_new(mixed_indices, table).unwrap(),
)
.unwrap(),
];
for samples in cases {
let whole = FrameRow {
key: None,
metric: Some(metric.clone()),
is_last: true,
time: RowTime::Spans(vec![GridSpan::new(0, 8), GridSpan::new(10, 1)]),
samples: vec![samples.clone()],
bytes: 0,
};
let exact_budget = measure_row(&whole).unwrap();
for (budget, expected_rows) in [(exact_budget, 1), (exact_budget - 1, 2)] {
let series = evaluated_series(
&schema,
vec![GridSpan::new(0, 8), GridSpan::new(10, 1)],
SeriesFields::one(samples.clone()),
);
let mut cursor = SeriesFrameCursor::try_new_evaluated(&schema, series, true).unwrap();
let mut builder =
SeriesFrameBuilder::try_new(&schema, FrameLabelMode::Explicit, None, 8, budget)
.unwrap();
let mut batches = Vec::new();
while let Some(batch) = cursor.append_to(&mut builder).unwrap() {
batches.push(batch);
}
if let Some(batch) = builder.flush().unwrap() {
batches.push(batch);
}
assert_eq!(
batches.iter().map(RecordBatch::num_rows).sum::<usize>(),
expected_rows
);
let decoded = batches
.iter()
.flat_map(|batch| {
SeriesFrameDecoder::try_new(&schema, FrameLabelMode::Explicit, None)
.unwrap()
.decode_evaluated(batch)
.unwrap()
})
.collect::<Vec<_>>();
assert_eq!(
decoded
.iter()
.map(|(series, _)| series.time.len())
.sum::<usize>(),
9
);
assert_eq!(decoded.iter().filter(|(_, is_last)| *is_last).count(), 1);
assert!(decoded.last().unwrap().1);
assert_eq!(
decoded
.iter()
.flat_map(|(series, _)| series.time.runs())
.map(|run| (run.position, run.len))
.collect::<Vec<_>>(),
[(0, 8), (10, 1)]
);
let decoded_samples = decoded
.iter()
.flat_map(|(series, _)| {
(0..series.time.len())
.map(move |index| series.fields.get(0).unwrap().sample(index))
})
.collect::<Vec<_>>();
assert_eq!(decoded_samples.len(), 9);
for (index, actual) in decoded_samples.into_iter().enumerate() {
assert_evaluated_sample_equal(actual, samples.sample(index), index);
}
}
}
}
#[test]
fn evaluated_many_sparse_runs_form_one_row_without_losing_positions() {
let base = test_schema(&["value"]);
let spans = (0..128)
.map(|run| GridSpan::new(run * 2, 1))
.collect::<Vec<_>>();
let schema = SeriesSchema {
fields: base.fields,
eval_grid: Some(Arc::new(
crate::series::time::EvalGrid::try_new(
crate::series::time::GridSpec::try_new(0, 1, 256).unwrap(),
)
.unwrap(),
)),
};
let missing = |index| matches!(index, 7 | 8 | 63 | 64 | 127);
let floats = (0..128)
.map(|index| match index {
index if missing(index) => None,
0 => Some(f64::from_bits(0x7ff8_0000_dead_beef)),
1 => Some(-0.0),
_ => Some(index as f64),
})
.collect::<Vec<_>>();
let histograms = (0..128)
.map(|index| (!missing(index)).then_some(0))
.collect::<Vec<_>>();
let mixed_floats = (0..128)
.map(|index| (!missing(index) && index % 2 == 0).then(|| floats[index].unwrap()))
.collect::<Vec<_>>();
let mixed_histograms = (0..128)
.map(|index| (!missing(index) && index % 2 == 1).then_some(0))
.collect::<Vec<_>>();
let prefix_float = |values: Vec<Option<f64>>| {
Float64Array::from(
std::iter::once(Some(999.0))
.chain(values)
.collect::<Vec<_>>(),
)
.slice(1, 128)
};
let prefix_indices = |values: Vec<Option<u32>>| {
UInt32Array::from(std::iter::once(Some(999)).chain(values).collect::<Vec<_>>())
.slice(1, 128)
};
let table = histogram_batch();
let cases = vec![
Samples::from_float(prefix_float(floats)),
Samples::try_from_histogram(prefix_indices(histograms), table.clone()).unwrap(),
Samples::try_from_mixed(
prefix_float(mixed_floats),
crate::series::samples::HistogramSamples::try_new(
prefix_indices(mixed_histograms),
table,
)
.unwrap(),
)
.unwrap(),
];
for samples in cases {
let series = evaluated_series(&schema, spans.clone(), SeriesFields::one(samples.clone()));
let mut cursor = SeriesFrameCursor::try_new_evaluated(&schema, series, true).unwrap();
let mut builder =
SeriesFrameBuilder::try_new(&schema, FrameLabelMode::Explicit, None, 8, 1 << 20)
.unwrap();
while cursor.append_to(&mut builder).unwrap().is_some() {}
let batch = builder.flush().unwrap().unwrap();
assert_eq!(batch.num_rows(), 1);
let decoded = SeriesFrameDecoder::try_new(&schema, FrameLabelMode::Explicit, None)
.unwrap()
.decode_evaluated(&batch)
.unwrap();
assert_eq!(decoded.len(), 1);
assert!(decoded[0].1);
assert_eq!(
decoded[0]
.0
.time
.runs()
.map(|run| (run.position, run.len))
.collect::<Vec<_>>(),
spans
.iter()
.map(|span| (span.start(), span.len()))
.collect::<Vec<_>>()
);
let decoded_samples = decoded[0].0.fields.get(0).unwrap();
for index in 0..128 {
assert_evaluated_sample_equal(
decoded_samples.sample(index),
samples.sample(index),
index,
);
}
}
}
#[test]
fn evaluated_rows_stop_at_earliest_field_segment_and_builder_bounds() {
let base = test_schema(&["a", "b"]);
let schema = SeriesSchema {
fields: base.fields,
eval_grid: Some(Arc::new(
crate::series::time::EvalGrid::try_new(
crate::series::time::GridSpec::try_new(0, 1, 12).unwrap(),
)
.unwrap(),
)),
};
let first = Samples::try_from_chunks(vec![
Samples::from_float(Float64Array::from(vec![Some(1.0), Some(2.0)])),
Samples::from_float(Float64Array::from(vec![Some(3.0), Some(4.0)])),
])
.unwrap();
let second = Samples::from_float(Float64Array::from(vec![Some(10.0); 4]));
let series = evaluated_series(
&schema,
vec![
GridSpan::new(0, 1),
GridSpan::new(3, 1),
GridSpan::new(6, 1),
GridSpan::new(9, 1),
],
SeriesFields::try_from_vec(vec![first, second]).unwrap(),
);
let mut cursor = SeriesFrameCursor::try_new_evaluated(&schema, series, true).unwrap();
let mut builder =
SeriesFrameBuilder::try_new(&schema, FrameLabelMode::Explicit, None, 1, 4096).unwrap();
let mut rows = Vec::new();
while let Some(batch) = cursor.append_to(&mut builder).unwrap() {
rows.extend(
SeriesFrameDecoder::try_new(&schema, FrameLabelMode::Explicit, None)
.unwrap()
.decode_evaluated(&batch)
.unwrap(),
);
}
if let Some(batch) = builder.flush().unwrap() {
rows.extend(
SeriesFrameDecoder::try_new(&schema, FrameLabelMode::Explicit, None)
.unwrap()
.decode_evaluated(&batch)
.unwrap(),
);
}
assert_eq!(
rows.iter().map(|(s, _)| s.time.len()).collect::<Vec<_>>(),
vec![2, 2]
);
assert_eq!(
rows.iter()
.flat_map(|(s, _)| s.fields.get(0).unwrap().iter().map(|v| match v {
SampleRef::Float(v) => v,
_ => panic!(),
}))
.collect::<Vec<_>>(),
vec![1.0, 2.0, 3.0, 4.0]
);
assert!(rows.last().unwrap().1);
}
#[test]
fn evaluated_available_bytes_preserve_pending_intersection_and_oversized_is_atomic() {
let base = test_schema(&["value"]);
let schema = SeriesSchema {
fields: base.fields,
eval_grid: Some(Arc::new(
crate::series::time::EvalGrid::try_new(
crate::series::time::GridSpec::try_new(0, 1, 8).unwrap(),
)
.unwrap(),
)),
};
let series = evaluated_series(
&schema,
vec![
GridSpan::new(0, 1),
GridSpan::new(2, 1),
GridSpan::new(4, 1),
GridSpan::new(6, 1),
],
SeriesFields::one(Samples::from_float(Float64Array::from(vec![
Some(1.0),
Some(2.0),
Some(3.0),
Some(4.0),
]))),
);
let mut cursor = SeriesFrameCursor::try_new_evaluated(&schema, series, true).unwrap();
let mut builder =
SeriesFrameBuilder::try_new(&schema, FrameLabelMode::Explicit, None, 8, 190).unwrap();
let batches = {
let mut out = Vec::new();
while let Some(batch) = cursor.append_to(&mut builder).unwrap() {
out.push(batch);
}
if let Some(batch) = builder.flush().unwrap() {
out.push(batch);
}
out
};
let rows = batches
.iter()
.flat_map(|batch| {
SeriesFrameDecoder::try_new(&schema, FrameLabelMode::Explicit, None)
.unwrap()
.decode_evaluated(batch)
.unwrap()
})
.collect::<Vec<_>>();
assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 4);
assert_eq!(
rows.iter()
.flat_map(|(s, _)| s.fields.get(0).unwrap().iter().map(|v| match v {
SampleRef::Float(v) => v,
_ => panic!(),
}))
.collect::<Vec<_>>(),
vec![1.0, 2.0, 3.0, 4.0]
);
assert!(rows.last().unwrap().1);
let oversized = evaluated_series(
&schema,
vec![GridSpan::new(0, 1), GridSpan::new(2, 1)],
SeriesFields::one(Samples::from_float(Float64Array::from(vec![
Some(1.0),
Some(2.0),
]))),
);
let mut cursor = SeriesFrameCursor::try_new_evaluated(&schema, oversized, true).unwrap();
let mut tiny =
SeriesFrameBuilder::try_new(&schema, FrameLabelMode::Explicit, None, 8, 1).unwrap();
let first = cursor.append_to(&mut tiny).unwrap().unwrap();
assert_eq!(
first.num_rows(),
1,
"oversized atomic intersection cannot absorb another run"
);
let second = cursor.append_to(&mut tiny).unwrap().unwrap();
assert_eq!(second.num_rows(), 1);
assert!(cursor.append_to(&mut tiny).unwrap().is_none());
let decoded = [first, second]
.iter()
.flat_map(|batch| {
SeriesFrameDecoder::try_new(&schema, FrameLabelMode::Explicit, None)
.unwrap()
.decode_evaluated(batch)
.unwrap()
})
.collect::<Vec<_>>();
assert_eq!(
decoded
.iter()
.flat_map(|(series, _)| series.fields.get(0).unwrap().iter().map(
|sample| match sample {
SampleRef::Float(value) => value,
_ => panic!(),
}
))
.collect::<Vec<_>>(),
vec![1.0, 2.0]
);
assert!(decoded.last().unwrap().1);
}
#[test]
fn evaluated_histogram_sparse_row_retains_full_table_and_indices() {
let base = test_schema(&["hist"]);
let schema = SeriesSchema {
fields: base.fields,
eval_grid: Some(Arc::new(
crate::series::time::EvalGrid::try_new(
crate::series::time::GridSpec::try_new(0, 1, 8).unwrap(),
)
.unwrap(),
)),
};
let table = histogram_batch();
let indices = UInt32Array::from(vec![Some(0), None, Some(1)]);
let samples = Samples::try_from_histogram(indices.clone(), table.clone()).unwrap();
let series = evaluated_series(
&schema,
vec![GridSpan::new(0, 1), GridSpan::new(3, 2)],
SeriesFields::one(samples),
);
let mut cursor = SeriesFrameCursor::try_new_evaluated(&schema, series, true).unwrap();
let mut builder =
SeriesFrameBuilder::try_new(&schema, FrameLabelMode::Explicit, None, 8, 4096).unwrap();
while cursor.append_to(&mut builder).unwrap().is_some() {}
let batch = builder.flush().unwrap().unwrap();
assert_eq!(batch.num_rows(), 1);
let decoded = SeriesFrameDecoder::try_new(&schema, FrameLabelMode::Explicit, None)
.unwrap()
.decode_evaluated(&batch)
.unwrap();
let DenseSamplesRef::Histogram(hist) = decoded[0].0.fields.get(0).unwrap().dense().unwrap()
else {
panic!("expected histogram")
};
assert_eq!(
hist.batch().len(),
table.len(),
"unused histogram table entries remain present"
);
assert_eq!(hist.indices(), &indices);
assert_eq!(
decoded[0]
.0
.time
.runs()
.map(|run| (run.position, run.len))
.collect::<Vec<_>>(),
vec![(0, 1), (3, 2)]
);
}
#[test]
fn direct_decoder_validates_ordering_for_raw_and_evaluated_batches() {
let raw_schema = test_schema(&["value"]);
@@ -3318,7 +3832,7 @@ fn byte_budget_flushes_before_exceeding_and_retries_the_pending_row() {
for chunk in &chunks {
let mut cursor =
SeriesFrameCursor::try_new_raw(&schema, FrameLabelMode::Key, chunk.clone()).unwrap();
measured.push(cursor.next_row().unwrap().unwrap().bytes);
measured.push(cursor.next_row(None).unwrap().unwrap().bytes);
}
let budget = measured[0] + measured[1] - 1;
src/mito-codec/src/labels.rs
+35 -135
View File
@@ -15,16 +15,14 @@
use std::collections::HashSet;
use std::sync::Arc;
use api::v1::SemanticType;
use common_telemetry::debug;
use datatypes::data_type::ConcreteDataType;
use datatypes::series::frame::{LabelDecodeEncoding, LabelDecodeSpec};
use datatypes::series::{Label, SeriesLabels};
use datatypes::value::Value;
use snafu::{OptionExt, ResultExt};
use snafu::ResultExt;
use store_api::codec::PrimaryKeyEncoding;
use store_api::metric_engine_consts::is_metric_engine_internal_column;
use store_api::region_engine::RawLabelSource;
use store_api::storage::{
ColumnId,
consts::{ReservedColumnId, is_internal_column},
@@ -32,31 +30,12 @@ use store_api::storage::{
use crate::error::{
InvalidDensePrimaryKeySnafu, InvalidLabelDecodeDescriptionSnafu, InvalidSparsePrimaryKeySnafu,
Result, UnsupportedOperationSnafu,
Result,
};
use crate::row_converter::SortField;
use crate::row_converter::dense::DensePrimaryKeyCodec;
use crate::row_converter::sparse::{SparseOffsetsCache, SparsePrimaryKeyView};
/// Selects labels materialized from a primary key.
///
/// `All` produces a complete metric name and label set. `Only` produces a projected
/// label vector, including `__name__` only when explicitly requested. Unexplained sparse
/// column IDs are errors with `All`; `Only` skips them after checking their framing.
#[derive(Clone, Debug, Default)]
pub enum LabelProjection {
#[default]
All,
Only(Arc<[Arc<str>]>),
}
/// Decoded labels: either a complete metric name or a possibly partial label vector.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum DecodedLabels {
Complete(SeriesLabels),
Projected(Vec<Label>),
}
enum Layout {
Sparse(SparseOffsetsCache),
Dense {
@@ -65,6 +44,13 @@ enum Layout {
},
}
fn invalid_binding_error(encoding: PrimaryKeyEncoding, reason: String) -> crate::error::Error {
match encoding {
PrimaryKeyEncoding::Dense => InvalidDensePrimaryKeySnafu { reason }.build(),
PrimaryKeyEncoding::Sparse => InvalidSparsePrimaryKeySnafu { reason }.build(),
}
}
struct BoundLabel {
id: ColumnId,
name: Arc<str>,
@@ -81,8 +67,8 @@ struct KeyColumn {
/// Minimal primary-key layout labels can be decoded from.
///
/// It holds only what decoding reads: the metric name, the primary-key encoding, and the ordered
/// primary-key columns with their physical types. A scan-bound [`RawLabelSource::PrimaryKey`] and a
/// frame schema's [`LabelDecodeSpec`] both reduce to this view, so both bind the same decoder.
/// primary-key columns with their physical types. A frame schema's [`LabelDecodeSpec`] reduces
/// to this view, so decoding binds directly to the described layout.
/// Every key column is a `Tag` column, which is checked when a decoder is bound, so the layout
/// carries no per-column role.
struct KeyLayout {
@@ -106,69 +92,13 @@ impl KeyLayout {
/// Decodes labels on demand using the minimal primary-key layout bound to a raw-label source.
pub struct PrimaryKeyLabelDecoder {
key: KeyLayout,
complete: bool,
include_name: bool,
labels: Vec<BoundLabel>,
value_scratch: Vec<u8>,
layout: Layout,
}
impl PrimaryKeyLabelDecoder {
/// Binds a decoder to a primary-key source snapshot and projection.
///
/// Explicit sources already contain complete labels and are rejected. Dense decoding
/// requires the full key layout described by the bound snapshot; compatibility with
/// older key layouts must be normalized upstream.
///
/// Every primary-key column must be a `Tag` column; [`RegionMetadata::validate`] enforces it, so
/// this only rejects a snapshot built without validation.
pub fn new(source: &RawLabelSource, projection: LabelProjection) -> Result<Self> {
let RawLabelSource::PrimaryKey {
metric_name,
physical_metadata,
logical_tags,
} = source
else {
return UnsupportedOperationSnafu {
err_msg: "explicit raw labels are already complete",
}
.fail();
};
let metadata = physical_metadata.as_ref();
if let Some(column) = metadata
.primary_key_columns()
.find(|column| column.semantic_type != SemanticType::Tag)
{
return InvalidSparsePrimaryKeySnafu {
reason: format!(
"primary-key column {} is not a tag column",
column.column_schema.name
),
}
.fail();
}
let key = KeyLayout {
metric_name: metric_name.clone(),
encoding: metadata.primary_key_encoding,
columns: metadata
.primary_key_columns()
.map(|column| KeyColumn {
id: column.column_id,
name: column.column_schema.name.clone(),
data_type: column.column_schema.data_type.clone(),
})
.collect(),
};
Self::bind(
key,
logical_tags
.iter()
.map(|tag| (tag.column_id, tag.name.as_ref())),
projection,
)
}
/// Binds a decoder to a frame schema's label decode description and projection.
/// Binds a decoder to a frame schema's label decode description.
///
/// This is the decoding path for a consumer that received only the raw frame schema: the
/// description carries the metric name, the primary-key encoding, the ordered key columns with
@@ -180,7 +110,7 @@ impl PrimaryKeyLabelDecoder {
/// is an encoding position. The logical tags are still checked here, one by one, exactly as
/// for a typed binding, which is the only place the decoder inspects them (a listed column no
/// tag selects is not examined).
pub fn from_decode_spec(spec: &LabelDecodeSpec, projection: LabelProjection) -> Result<Self> {
pub fn from_decode_spec(spec: &LabelDecodeSpec) -> Result<Self> {
spec.validate()
.context(InvalidLabelDecodeDescriptionSnafu)?;
let key = KeyLayout {
@@ -204,12 +134,11 @@ impl PrimaryKeyLabelDecoder {
spec.logical_tags
.iter()
.map(|tag| (tag.column_id, tag.name.as_str())),
projection,
)
}
/// Binds the ordered logical tags of `key` to one projection and one decode layout.
fn bind<'a, I>(key: KeyLayout, tags: I, projection: LabelProjection) -> Result<Self>
/// Binds the ordered logical tags of `key` to one decode layout.
fn bind<'a, I>(key: KeyLayout, tags: I) -> Result<Self>
where
I: IntoIterator<Item = (ColumnId, &'a str)>,
{
@@ -229,16 +158,17 @@ impl PrimaryKeyLabelDecoder {
.column(column_id)
.is_some_and(|column| column.name == name)
{
return InvalidSparsePrimaryKeySnafu {
reason: format!("invalid logical tag binding {name}"),
}
.fail();
return Err(invalid_binding_error(
key.encoding,
format!("invalid logical tag binding {name}"),
));
}
let position = key
.position(column_id)
.context(InvalidSparsePrimaryKeySnafu {
reason: format!("logical tag {name} is not a primary-key column"),
})?;
let position = key.position(column_id).ok_or_else(|| {
invalid_binding_error(
key.encoding,
format!("logical tag {name} is not a primary-key column"),
)
})?;
labels.push(BoundLabel {
id: column_id,
name: Arc::from(name),
@@ -246,17 +176,6 @@ impl PrimaryKeyLabelDecoder {
});
}
labels.sort_unstable_by(|a, b| a.name.cmp(&b.name));
let complete = matches!(&projection, LabelProjection::All);
let include_name = match &projection {
LabelProjection::All => false,
LabelProjection::Only(requested) => {
requested.iter().any(|name| name.as_ref() == "__name__")
}
};
if let LabelProjection::Only(requested) = &projection {
let requested: HashSet<_> = requested.iter().map(AsRef::as_ref).collect();
labels.retain(|label| requested.contains(label.name.as_ref()));
}
let layout = match key.encoding {
PrimaryKeyEncoding::Sparse => Layout::Sparse(SparseOffsetsCache::new()),
PrimaryKeyEncoding::Dense => Layout::Dense {
@@ -273,18 +192,14 @@ impl PrimaryKeyLabelDecoder {
PrimaryKeyEncoding::Sparse => "sparse",
PrimaryKeyEncoding::Dense => "dense",
};
let projection_mode = if complete { "all" } else { "only" };
debug!(
encoding,
projection_mode,
logical_tag_count,
selected_tag_count = labels.len(),
"Constructed primary key label decoder"
);
Ok(Self {
key,
complete,
include_name,
labels,
value_scratch: Vec::new(),
layout,
@@ -293,16 +208,14 @@ impl PrimaryKeyLabelDecoder {
/// Decodes labels without modifying or consuming `read_key`.
///
/// `All` returns a complete [`SeriesLabels`]; `Only` returns a partial vector. Dense
/// keys are fully checked against the exact snapshot layout even when decoding only
/// selected fields. Sparse keys receive a full structural audit before values are read.
pub fn decode(&mut self, read_key: &[u8]) -> Result<DecodedLabels> {
let mut values = Vec::with_capacity(self.labels.len() + usize::from(self.include_name));
/// Returns complete [`SeriesLabels`]. Dense keys are fully checked against the exact
/// described layout. Sparse keys receive a full structural audit before values are read.
pub fn decode(&mut self, read_key: &[u8]) -> Result<SeriesLabels> {
let mut values = Vec::with_capacity(self.labels.len());
let Self {
layout,
key,
labels,
complete,
value_scratch,
..
} = self;
@@ -310,7 +223,7 @@ impl PrimaryKeyLabelDecoder {
Layout::Sparse(cache) => {
let mut view = SparsePrimaryKeyView::new(read_key, cache)?;
view.visit_label_ids(|id| {
if *complete && key.column(id).is_none() {
if key.column(id).is_none() {
return InvalidSparsePrimaryKeySnafu {
reason: format!("unrecognized sparse tag id {id}"),
}
@@ -351,25 +264,12 @@ impl PrimaryKeyLabelDecoder {
}
}
}
if *complete {
Ok(DecodedLabels::Complete(
SeriesLabels::try_new(Some(key.metric_name.clone()), values).map_err(|error| {
InvalidSparsePrimaryKeySnafu {
reason: error.to_string(),
}
.build()
})?,
))
} else {
if self.include_name {
values.push(Label {
name: Arc::from("__name__"),
value: self.key.metric_name.clone(),
});
values.sort_unstable_by(|a, b| a.name.cmp(&b.name));
SeriesLabels::try_new(Some(key.metric_name.clone()), values).map_err(|error| {
InvalidSparsePrimaryKeySnafu {
reason: error.to_string(),
}
Ok(DecodedLabels::Projected(values))
}
.build()
})
}
}
src/mito-codec/src/labels/tests.rs
+73 -570
View File
@@ -15,14 +15,13 @@
use std::sync::Arc;
use api::v1::SemanticType;
use common_time::timestamp::TimeUnit;
use common_time::{Date, Timestamp};
use datatypes::data_type::ConcreteDataType;
use datatypes::prelude::Value;
use datatypes::schema::ColumnSchema;
use store_api::codec::PrimaryKeyEncoding;
use store_api::metadata::{ColumnMetadata, RegionMetadata, RegionMetadataBuilder};
use store_api::region_engine::BoundTag;
use store_api::region_engine::RawLabelSource;
use store_api::storage::RegionId;
use crate::row_converter::{
@@ -122,591 +121,95 @@ fn values() -> Vec<(u32, Value)> {
}
#[test]
fn dense_and_sparse_decode_complete_or_projected_labels() {
for encoding in [PrimaryKeyEncoding::Dense, PrimaryKeyEncoding::Sparse] {
let source = source(
string_metadata(encoding),
"requests_total",
&[(8, "zeta"), (41, "alpha")],
);
let key = encoded(encoding, &values());
let mut all = PrimaryKeyLabelDecoder::new(&source, LabelProjection::All).unwrap();
let complete = all.decode(&key).unwrap();
let DecodedLabels::Complete(metric) = complete else {
panic!("All must return a complete metric identity")
};
assert_eq!(metric.name(), Some("requests_total"));
assert_eq!(metric.get("alpha"), Some("first"));
assert_eq!(metric.get("zeta"), Some("last"));
assert_eq!(metric.get("private"), None);
let mut projected = PrimaryKeyLabelDecoder::new(
&source,
LabelProjection::Only(Arc::from([
Arc::<str>::from("zeta"),
Arc::from("alpha"),
Arc::from("__name__"),
Arc::from("missing"),
Arc::from("alpha"),
])),
)
.unwrap();
assert_eq!(
projected.decode(&key).unwrap(),
DecodedLabels::Projected(vec![
Label {
name: Arc::from("__name__"),
value: Arc::from("requests_total")
},
Label {
name: Arc::from("alpha"),
value: Arc::from("first")
},
Label {
name: Arc::from("zeta"),
value: Arc::from("last")
},
])
);
assert_eq!(
projected.decode(&key).unwrap(),
DecodedLabels::Projected(vec![
Label {
name: Arc::from("__name__"),
value: Arc::from("requests_total")
},
Label {
name: Arc::from("alpha"),
value: Arc::from("first")
},
Label {
name: Arc::from("zeta"),
value: Arc::from("last")
},
])
);
let mut without_name = PrimaryKeyLabelDecoder::new(
&source,
LabelProjection::Only(Arc::from([Arc::<str>::from("alpha")])),
)
.unwrap();
assert_eq!(
without_name.decode(&key).unwrap(),
DecodedLabels::Projected(vec![Label {
name: Arc::from("alpha"),
value: Arc::from("first")
}])
);
let mut empty_projection =
PrimaryKeyLabelDecoder::new(&source, LabelProjection::Only(Arc::from([]))).unwrap();
assert_eq!(
empty_projection.decode(&key).unwrap(),
DecodedLabels::Projected(vec![])
);
}
}
#[test]
fn projection_snapshots_its_layout_and_distinguishes_missing_and_empty() {
let original = metadata(
&[
(21, "old", ConcreteDataType::string_datatype()),
(4, "added_later", ConcreteDataType::string_datatype()),
],
PrimaryKeyEncoding::Sparse,
);
let old_source = source(
original.clone(),
"old_metric",
&[(21, "old"), (4, "added_later")],
);
let weak_metadata = Arc::downgrade(&original);
let replacement = metadata(
&[(21, "renamed", ConcreteDataType::string_datatype())],
PrimaryKeyEncoding::Sparse,
);
let new_source = source(replacement, "new_metric", &[(21, "renamed")]);
assert!(PrimaryKeyLabelDecoder::new(&new_source, LabelProjection::All).is_ok());
let codec = SparsePrimaryKeyCodec::with_fields(vec![(
21,
SortField::new(ConcreteDataType::string_datatype()),
)]);
let mut key = Vec::new();
codec.encode_internal(42, 9, &mut key).unwrap();
codec
.encode_values(&[(21, Value::String("".into()))], &mut key)
.unwrap();
let mut decoder = PrimaryKeyLabelDecoder::new(
&old_source,
LabelProjection::Only(Arc::from([
Arc::<str>::from("old"),
Arc::from("added_later"),
])),
)
.unwrap();
let mut all = PrimaryKeyLabelDecoder::new(&old_source, LabelProjection::All).unwrap();
drop(old_source);
drop(original);
assert!(
weak_metadata.upgrade().is_none(),
"the decoder keeps its own key layout (ids, positions, types), not the region metadata"
);
assert_eq!(
decoder.decode(&key).unwrap(),
DecodedLabels::Projected(vec![Label {
name: Arc::from("old"),
value: Arc::from("")
},])
);
let mut absent = Vec::new();
codec.encode_internal(42, 9, &mut absent).unwrap();
assert_eq!(
decoder.decode(&absent).unwrap(),
DecodedLabels::Projected(vec![])
);
let DecodedLabels::Complete(absent_tags) = all.decode(&absent).unwrap() else {
panic!("All must return complete identity")
};
assert_eq!(absent_tags.get("old"), None);
let DecodedLabels::Complete(empty_tag) = all.decode(&key).unwrap() else {
panic!("All must return complete identity")
};
assert_eq!(empty_tag.get("old"), Some(""));
let mut complete = PrimaryKeyLabelDecoder::new(
&source(
metadata(
&[(21, "old", ConcreteDataType::string_datatype())],
PrimaryKeyEncoding::Sparse,
),
"",
&[(21, "old")],
),
LabelProjection::All,
)
.unwrap();
let DecodedLabels::Complete(empty_name) = complete.decode(&absent).unwrap() else {
panic!("All must return complete identity")
};
assert_eq!(empty_name.name(), Some(""));
let named_source = source(
metadata(
&[(21, "old", ConcreteDataType::string_datatype())],
PrimaryKeyEncoding::Sparse,
),
"old_metric",
&[(21, "old")],
);
let mut complete = PrimaryKeyLabelDecoder::new(&named_source, LabelProjection::All).unwrap();
let DecodedLabels::Complete(named) = complete.decode(&absent).unwrap() else {
panic!("All must return complete identity")
};
assert_eq!(named.name(), Some("old_metric"));
// A null value is represented by its sparse column ID and null marker.
let mut null = Vec::new();
codec.encode_internal(42, 9, &mut null).unwrap();
null.extend_from_slice(&21_u32.to_be_bytes());
null.push(0);
assert_eq!(
decoder.decode(&null).unwrap(),
DecodedLabels::Projected(vec![])
);
drop(decoder);
drop(all);
assert!(weak_metadata.upgrade().is_none());
}
#[test]
fn projection_does_not_hide_bad_framing_or_materialize_unselected_values() {
let meta = metadata(
&[
(11, "wanted", ConcreteDataType::string_datatype()),
(35, "unwanted", ConcreteDataType::string_datatype()),
],
PrimaryKeyEncoding::Sparse,
);
let bound = source(meta, "m", &[(11, "wanted"), (35, "unwanted")]);
let codec = SparsePrimaryKeyCodec::with_fields(vec![
(11, SortField::new(ConcreteDataType::string_datatype())),
(35, SortField::new(ConcreteDataType::string_datatype())),
]);
let mut key = Vec::new();
codec.encode_internal(42, 9, &mut key).unwrap();
codec
.encode_raw_tag_value(
[(11, b"ok".as_slice()), (35, &[0xff])].into_iter(),
&mut key,
)
.unwrap();
let mut only = PrimaryKeyLabelDecoder::new(
&bound,
LabelProjection::Only(Arc::from([Arc::<str>::from("wanted")])),
)
.unwrap();
assert_eq!(
only.decode(&key).unwrap(),
DecodedLabels::Projected(vec![Label {
name: Arc::from("wanted"),
value: Arc::from("ok")
},])
);
let mut all = PrimaryKeyLabelDecoder::new(&bound, LabelProjection::All).unwrap();
assert!(all.decode(&key).is_err());
let mut malformed = key.clone();
malformed.pop();
assert!(only.decode(&malformed).is_err());
assert!(all.decode(&malformed).is_err());
let mut empty_projection =
PrimaryKeyLabelDecoder::new(&bound, LabelProjection::Only(Arc::from([]))).unwrap();
assert!(empty_projection.decode(&malformed).is_err());
}
#[test]
fn dense_projection_still_rejects_short_and_extra_key_bytes() {
let meta = metadata(
&[(17, "label", ConcreteDataType::string_datatype())],
PrimaryKeyEncoding::Dense,
);
let bound = source(meta, "m", &[(17, "label")]);
let key = encoded(
PrimaryKeyEncoding::Dense,
&[(17, Value::String("ok".into()))],
);
let mut decoder =
PrimaryKeyLabelDecoder::new(&bound, LabelProjection::Only(Arc::from([]))).unwrap();
assert!(decoder.decode(&key[..key.len() - 1]).is_err());
let mut extra = key;
extra.push(0);
assert!(decoder.decode(&extra).is_err());
}
#[test]
fn dense_unselected_invalid_utf8_is_not_materialized() {
let meta = metadata(
&[
(17, "wanted", ConcreteDataType::string_datatype()),
(18, "unwanted", ConcreteDataType::string_datatype()),
],
PrimaryKeyEncoding::Dense,
);
let bound = source(meta, "m", &[(17, "wanted"), (18, "unwanted")]);
let first_value = encoded(
PrimaryKeyEncoding::Dense,
&[(17, Value::String("ok".into()))],
);
let mut key = encoded(
PrimaryKeyEncoding::Dense,
&[
(17, Value::String("ok".into())),
(18, Value::String("x".into())),
],
);
key[first_value.len() + 2] = 0xff;
let mut only = PrimaryKeyLabelDecoder::new(
&bound,
LabelProjection::Only(Arc::from([Arc::<str>::from("wanted")])),
)
.unwrap();
assert_eq!(
only.decode(&key).unwrap(),
DecodedLabels::Projected(vec![Label {
name: Arc::from("wanted"),
value: Arc::from("ok")
}])
);
let mut selected_invalid = PrimaryKeyLabelDecoder::new(
&bound,
LabelProjection::Only(Arc::from([Arc::<str>::from("unwanted")])),
)
.unwrap();
assert!(selected_invalid.decode(&key).is_err());
let mut all = PrimaryKeyLabelDecoder::new(&bound, LabelProjection::All).unwrap();
assert!(all.decode(&key).is_err());
}
#[test]
fn unknown_sparse_ids_are_skipped_only_for_projection_and_full_keys_stay_distinct() {
let meta = metadata(
&[(11, "wanted", ConcreteDataType::string_datatype())],
PrimaryKeyEncoding::Sparse,
);
let bound = source(meta, "m", &[(11, "wanted")]);
let codec = SparsePrimaryKeyCodec::with_fields(vec![(
11,
SortField::new(ConcreteDataType::string_datatype()),
)]);
let mut first = Vec::new();
codec.encode_internal(42, 9, &mut first).unwrap();
codec
.encode_raw_tag_value(
[
(10, b"before".as_slice()),
(11, b"one".as_slice()),
(12, b"between".as_slice()),
(90, b"after".as_slice()),
]
.into_iter(),
&mut first,
)
.unwrap();
let original_first = first.clone();
let mut second = Vec::new();
codec.encode_internal(42, 9, &mut second).unwrap();
codec
.encode_raw_tag_value(
[
(10, b"before".as_slice()),
(11, b"two".as_slice()),
(12, b"between".as_slice()),
(90, b"after".as_slice()),
]
.into_iter(),
&mut second,
)
.unwrap();
let mut only = PrimaryKeyLabelDecoder::new(
&bound,
LabelProjection::Only(Arc::from([Arc::<str>::from("wanted")])),
)
.unwrap();
assert_eq!(
only.decode(&first).unwrap(),
DecodedLabels::Projected(vec![Label {
name: Arc::from("wanted"),
value: Arc::from("one")
},])
);
assert_eq!(
only.decode(&second).unwrap(),
DecodedLabels::Projected(vec![Label {
name: Arc::from("wanted"),
value: Arc::from("two")
},])
);
assert_ne!(first, second);
assert_eq!(first, original_first);
let original_second = second.clone();
let mut all = PrimaryKeyLabelDecoder::new(&bound, LabelProjection::All).unwrap();
assert!(all.decode(&first).is_err());
assert_eq!(first, original_first);
assert_eq!(second, original_second);
}
#[test]
fn unknown_sparse_ids_before_between_and_after_known_tags_are_projected() {
let meta = metadata(
&[
(11, "wanted", ConcreteDataType::string_datatype()),
(13, "other", ConcreteDataType::string_datatype()),
],
PrimaryKeyEncoding::Sparse,
);
let bound = source(meta, "m", &[(11, "wanted"), (13, "other")]);
let codec = SparsePrimaryKeyCodec::schemaless();
let cases: [&[(u32, &[u8])]; 3] = [
&[(10, b"before"), (11, b"one"), (12, b"two")],
&[(11, b"one"), (12, b"between"), (13, b"two")],
&[(11, b"one"), (13, b"two"), (90, b"after")],
];
for tags in cases {
let mut key = Vec::new();
codec.encode_internal(42, 9, &mut key).unwrap();
codec
.encode_raw_tag_value(tags.iter().copied(), &mut key)
.unwrap();
let mut only = PrimaryKeyLabelDecoder::new(
&bound,
LabelProjection::Only(Arc::from([Arc::<str>::from("wanted")])),
)
.unwrap();
assert_eq!(
only.decode(&key).unwrap(),
DecodedLabels::Projected(vec![Label {
name: Arc::from("wanted"),
value: Arc::from("one")
}])
);
let mut all = PrimaryKeyLabelDecoder::new(&bound, LabelProjection::All).unwrap();
assert!(all.decode(&key).is_err());
}
}
#[test]
fn constructor_rejects_invalid_bindings_but_accepts_valid_control() {
let meta = metadata(
&[
(11, "wanted", ConcreteDataType::string_datatype()),
(35, "other", ConcreteDataType::string_datatype()),
],
PrimaryKeyEncoding::Sparse,
);
for bindings in [
vec![(11, "wrong")],
vec![(11, "same"), (35, "same")],
vec![(11, "wanted"), (11, "wanted")],
vec![(0, "ts")],
vec![(90, "not_a_tag")],
vec![(99, "field")],
vec![(11, "__name__")],
vec![(u32::MAX, "internal")],
] {
let invalid = source(meta.clone(), "m", &bindings);
assert!(PrimaryKeyLabelDecoder::new(&invalid, LabelProjection::All).is_err());
}
let valid = source(meta, "m", &[(11, "wanted"), (35, "other")]);
assert!(PrimaryKeyLabelDecoder::new(&valid, LabelProjection::All).is_ok());
assert!(PrimaryKeyLabelDecoder::new(&RawLabelSource::Explicit, LabelProjection::All).is_err());
}
#[test]
fn dense_typed_values_keep_exact_integer_string_representation() {
let cases = [
(
"count",
ConcreteDataType::int64_datatype(),
Value::Int64(9_007_199_254_740_993),
),
(
"binary",
ConcreteDataType::binary_datatype(),
Value::Binary(vec![0xab, 0xcd].into()),
),
(
"date",
ConcreteDataType::date_datatype(),
Value::Date(Date::new(0)),
),
(
"timestamp",
ConcreteDataType::timestamp_millisecond_datatype(),
Value::Timestamp(Timestamp::new(1234, TimeUnit::Millisecond)),
),
];
for (name, data_type, value) in cases {
let expected = value.to_string();
let meta = metadata(&[(17, name, data_type)], PrimaryKeyEncoding::Dense);
let bound = source(meta, "typed", &[(17, name)]);
let key = encoded(PrimaryKeyEncoding::Dense, &[(17, value)]);
let mut decoder = PrimaryKeyLabelDecoder::new(&bound, LabelProjection::All).unwrap();
let DecodedLabels::Complete(metric) = decoder.decode(&key).unwrap() else {
panic!("All must return a complete metric identity")
};
assert_eq!(metric.get(name), Some(expected.as_str()));
}
assert_eq!(
Value::Int64(9_007_199_254_740_993).to_string(),
"9007199254740993"
);
}
/// One field is enough: the frame schema only carries ordered payload fields and metadata.
fn frame_series_schema() -> Arc<datatypes::series::samples::SeriesSchema> {
Arc::new(datatypes::series::samples::SeriesSchema {
fields: vec![datatypes::series::samples::FieldSpec {
name: Arc::from("value"),
}]
.into(),
eval_grid: None,
})
}
/// A key-encoded frame's schema describes its own labels: recovering the description from the
/// schema and decoding with it matches decoding with the original bound source, for both primary
/// key encodings and both projections.
#[test]
fn frame_label_decode_description_decodes_like_the_bound_source() {
fn dense_and_sparse_decode_complete_labels() {
for encoding in [PrimaryKeyEncoding::Dense, PrimaryKeyEncoding::Sparse] {
let bound = source(
string_metadata(encoding),
"requests_total",
&[(8, "zeta"), (41, "alpha")],
);
let description = common_query::series_frame::label_decode_spec(&bound)
.expect("a primary-key binding describes its key layout");
let frame_schema = datatypes::series::frame::series_frame_schema(
&frame_series_schema(),
datatypes::series::frame::FrameLabelMode::Key,
Some(&description),
)
.unwrap();
let recovered = datatypes::series::frame::series_frame_label_decode(&frame_schema)
.unwrap()
.expect("the frame declares its decode description");
assert_eq!(recovered, description);
let description = common_query::series_frame::label_decode_spec(&bound).unwrap();
let key = encoded(encoding, &values());
for projection in [
LabelProjection::All,
LabelProjection::Only(Arc::from([Arc::<str>::from("alpha")])),
LabelProjection::Only(Arc::from([])),
] {
let mut direct = PrimaryKeyLabelDecoder::new(&bound, projection.clone()).unwrap();
let mut from_frame =
PrimaryKeyLabelDecoder::from_decode_spec(&recovered, projection).unwrap();
assert_eq!(
from_frame.decode(&key).unwrap(),
direct.decode(&key).unwrap()
);
}
let mut decoder = PrimaryKeyLabelDecoder::from_decode_spec(&description).unwrap();
let metric = decoder.decode(&key).unwrap();
assert_eq!(metric.name(), Some("requests_total"));
assert_eq!(metric.get("alpha"), Some("first"));
assert_eq!(metric.get("zeta"), Some("last"));
assert_eq!(metric.get("private"), None);
}
}
/// A description this revision cannot honour is rejected when a decoder is bound, not decoded
/// best-effort.
#[test]
fn frame_label_decode_description_rejects_unsupported_version() {
fn dense_malformed_key_reports_dense_error() {
let bound = source(
string_metadata(PrimaryKeyEncoding::Dense),
"requests_total",
&[(8, "zeta")],
metadata(
&[(17, "label", ConcreteDataType::string_datatype())],
PrimaryKeyEncoding::Dense,
),
"m",
&[(17, "label")],
);
let description = common_query::series_frame::label_decode_spec(&bound).unwrap();
let mut future = description.clone();
future.version += 1;
assert!(PrimaryKeyLabelDecoder::from_decode_spec(&future, LabelProjection::All).is_err());
// The description must list the column a logical tag reads.
let mut unbound = description;
unbound.primary_key.retain(|column| column.column_id != 8);
assert!(PrimaryKeyLabelDecoder::from_decode_spec(&unbound, LabelProjection::All).is_err());
let mut decoder = PrimaryKeyLabelDecoder::from_decode_spec(&description).unwrap();
let key = encoded(
PrimaryKeyEncoding::Dense,
&[(17, Value::String("ok".into()))],
);
let mut extra = key.clone();
extra.push(0);
let error = decoder.decode(&extra).unwrap_err();
assert!(matches!(
error,
crate::error::Error::InvalidDensePrimaryKey { .. }
));
let error = decoder.decode(&key[..key.len() - 1]).unwrap_err();
assert!(matches!(
error,
crate::error::Error::DeserializeField { .. }
));
}
/// A primary-key column that is not a `Tag` column is rejected when a decoder is bound, because the
/// layout records no per-column role and decodes every listed key column as a tag.
///
/// `RegionMetadata::validate` rejects such metadata, so this snapshot is built without validation to
/// exercise the check the decoder performs itself.
#[test]
fn binding_rejects_a_primary_key_column_that_is_not_a_tag() {
let mut builder = RegionMetadataBuilder::new(RegionId::new(7, 1));
builder.push_column_metadata(ColumnMetadata {
column_schema: ColumnSchema::new(
"ts",
ConcreteDataType::timestamp_millisecond_datatype(),
false,
fn dense_invalid_logical_binding_reports_dense_error() {
let bound = source(
metadata(
&[(17, "label", ConcreteDataType::string_datatype())],
PrimaryKeyEncoding::Dense,
),
semantic_type: SemanticType::Timestamp,
column_id: 0,
});
builder.push_column_metadata(ColumnMetadata {
column_schema: ColumnSchema::new("not_a_tag", ConcreteDataType::string_datatype(), true),
semantic_type: SemanticType::Field,
column_id: 8,
});
builder.primary_key(vec![8]);
builder.primary_key_encoding(PrimaryKeyEncoding::Sparse);
let invalid = Arc::new(builder.build_without_validation().unwrap());
let bound = source(invalid, "m", &[(8, "not_a_tag")]);
let error = PrimaryKeyLabelDecoder::new(&bound, LabelProjection::All)
"m",
&[(17, "label")],
);
let mut description = common_query::series_frame::label_decode_spec(&bound).unwrap();
description.logical_tags[0].name = "wrong_name".to_string();
let error = PrimaryKeyLabelDecoder::from_decode_spec(&description)
.err()
.expect("a non-tag primary-key column must be rejected")
.to_string();
assert!(error.contains("not a tag column"), "{error}");
.unwrap();
assert!(matches!(
error,
crate::error::Error::InvalidDensePrimaryKey { .. }
));
}
#[test]
fn sparse_malformed_key_reports_sparse_error() {
let bound = source(
metadata(
&[(17, "label", ConcreteDataType::string_datatype())],
PrimaryKeyEncoding::Sparse,
),
"m",
&[(17, "label")],
);
let description = common_query::series_frame::label_decode_spec(&bound).unwrap();
let mut decoder = PrimaryKeyLabelDecoder::from_decode_spec(&description).unwrap();
let mut malformed = encoded(
PrimaryKeyEncoding::Sparse,
&[(17, Value::String("ok".into()))],
);
malformed.pop();
let error = decoder.decode(&malformed).unwrap_err();
assert!(matches!(
error,
crate::error::Error::InvalidSparsePrimaryKey { .. }
));
}
src/mito2/src/engine/scan_test.rs
+18 -15
View File
@@ -5229,8 +5229,8 @@ async fn test_mito_handle_query_raw_series_candidate_projection_and_mvcc() {
/// A `TwoPhase` scan declares [`SeriesOrdering::Contiguous`] on the scanner, on every partition
/// stream, and therefore in every frame it packs; the frame adapter validates that declaration.
///
/// The ordered merge that rejects a non-`KeyThenTime` input is covered by
/// `promql::source::merge::tests::rejects_unsupported_ordering_and_empty_input_without_polling`.
/// An ordering the frame cannot carry is rejected by the frame constructor before any poll; that
/// refusal is covered by `common_query::series_frame::tests::interleaved_ordering_is_rejected_before_any_poll`.
#[tokio::test]
async fn test_two_phase_native_raw_stream_declares_contiguous_frames() {
use common_query::series_frame::{pack_raw_series, series_frame_ordering, unpack_raw_series};
@@ -5313,7 +5313,10 @@ async fn test_two_phase_native_raw_stream_declares_contiguous_frames() {
};
assert_eq!("two_phase", series.mode());
let mut scanner = engine.handle_query(region_id, request).await.unwrap();
assert_eq!(SeriesOrdering::Contiguous, scanner.raw_series_ordering());
assert_eq!(
SeriesOrdering::Contiguous,
scanner.raw_series_ordering().unwrap()
);
let metadata = scanner.metadata();
let source = RawLabelSource::PrimaryKey {
metric_name: Arc::from("fixture_metric"),
@@ -5352,7 +5355,7 @@ async fn test_two_phase_native_raw_stream_declares_contiguous_frames() {
.unwrap()
.unwrap();
// The storage-side contract: every partition stream reports the scanner's declaration.
assert_eq!(scanner.raw_series_ordering(), stream.ordering());
assert_eq!(scanner.raw_series_ordering().unwrap(), stream.ordering());
assert_eq!(SeriesOrdering::Contiguous, stream.ordering());
schema.get_or_insert_with(|| Arc::clone(stream.schema()));
chunks.extend(stream.try_collect::<Vec<_>>().await.unwrap());
@@ -5589,7 +5592,10 @@ async fn test_two_phase_raw_stream_stays_contiguous_across_disjoint_partition_as
panic!("sparse PerSeries request must use SeriesScan");
};
assert_eq!("two_phase", scanner.mode());
assert_eq!(SeriesOrdering::Contiguous, scanner.raw_series_ordering());
assert_eq!(
SeriesOrdering::Contiguous,
scanner.raw_series_ordering().unwrap()
);
// Composition of the pinned scan input: the two flushes pin two SST sources and the third
// write round stays in one memtable source. Each source reports the rows its write round
// wrote, so this asserts the layout instead of only asserting that something was read.
@@ -5909,7 +5915,10 @@ async fn test_two_phase_raw_same_partition_multi_round_assignment_stays_contiguo
panic!("sparse PerSeries request must use SeriesScan");
};
assert_eq!("two_phase", scanner.mode());
assert_eq!(SeriesOrdering::Contiguous, scanner.raw_series_ordering());
assert_eq!(
SeriesOrdering::Contiguous,
scanner.raw_series_ordering().unwrap()
);
// The fixture must really be one SST source plus one memtable source, or the two rounds below
// would be explained by something else.
{
@@ -6244,10 +6253,9 @@ fn decode_declared_labels(
description: &datatypes::series::frame::LabelDecodeSpec,
chunks: &[datatypes::series::raw::RawSeriesChunk],
) -> Vec<datatypes::series::SeriesLabels> {
use mito_codec::labels::{DecodedLabels, LabelProjection, PrimaryKeyLabelDecoder};
use mito_codec::labels::PrimaryKeyLabelDecoder;
let mut decoder =
PrimaryKeyLabelDecoder::from_decode_spec(description, LabelProjection::All).unwrap();
let mut decoder = PrimaryKeyLabelDecoder::from_decode_spec(description).unwrap();
chunks
.iter()
.map(|chunk| {
@@ -6255,12 +6263,7 @@ fn decode_declared_labels(
chunk.metric.is_none(),
"a key-encoded chunk leaves its labels undecoded"
);
match decoder.decode(chunk.primary_key.as_slice()).unwrap() {
DecodedLabels::Complete(metric) => metric,
DecodedLabels::Projected(_) => {
panic!("decoding every label must yield complete labels")
}
}
decoder.decode(chunk.primary_key.as_slice()).unwrap()
})
.collect()
}
src/mito2/src/read/scan_util.rs
-10
View File
@@ -202,8 +202,6 @@ pub(crate) struct ScanMetricsSet {
/// Elapsed time before the first poll operation.
first_poll: Duration,
/// Number of send timeout in SeriesScan.
num_series_send_timeout: usize,
/// Number of send full in SeriesScan.
num_series_send_full: usize,
/// Number of rows the series distributor scanned.
@@ -325,7 +323,6 @@ impl fmt::Debug for ScanMetricsSet {
num_sst_batches,
num_sst_rows,
first_poll,
num_series_send_timeout,
num_series_send_full,
num_distributor_rows,
num_distributor_batches,
@@ -459,9 +456,6 @@ impl fmt::Debug for ScanMetricsSet {
}
// Write non-zero distributor metrics
if *num_series_send_timeout > 0 {
write!(f, ", \"num_series_send_timeout\":{num_series_send_timeout}")?;
}
if *num_series_send_full > 0 {
write!(f, ", \"num_series_send_full\":{num_series_send_full}")?;
}
@@ -775,7 +769,6 @@ impl ScanMetricsSet {
/// Sets distributor metrics.
fn set_distributor_metrics(&mut self, distributor_metrics: &SeriesDistributorMetrics) {
let SeriesDistributorMetrics {
num_series_send_timeout,
num_series_send_full,
num_rows,
num_batches,
@@ -784,7 +777,6 @@ impl ScanMetricsSet {
divider_cost,
} = distributor_metrics;
self.num_series_send_timeout += *num_series_send_timeout;
self.num_series_send_full += *num_series_send_full;
self.num_distributor_rows += *num_rows;
self.num_distributor_batches += *num_batches;
@@ -1144,8 +1136,6 @@ impl fmt::Debug for PartitionMetrics {
/// Metrics for the series distributor.
#[derive(Default)]
pub(crate) struct SeriesDistributorMetrics {
/// Number of send timeout in SeriesScan.
pub(crate) num_series_send_timeout: usize,
/// Number of send full in SeriesScan.
pub(crate) num_series_send_full: usize,
/// Number of rows the series distributor scanned.
src/mito2/src/read/series_scan.rs
+60 -50
View File
@@ -21,7 +21,7 @@ use std::collections::HashSet;
use std::fmt;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use std::time::Instant;
use api::v1::SemanticType;
use async_stream::try_stream;
@@ -55,7 +55,7 @@ use store_api::region_engine::{
use store_api::storage::consts::{ReservedColumnId, is_internal_column};
use table::predicate::Predicate;
use tokio::sync::Semaphore;
use tokio::sync::mpsc::error::{SendTimeoutError, TrySendError};
use tokio::sync::mpsc::error::TrySendError;
use tokio::sync::mpsc::{self, Receiver, Sender};
use crate::error::{
@@ -76,9 +76,6 @@ use crate::read::stream::{ConvertBatchStream, ScanBatch, ScanBatchStream};
use crate::sst::parquet::flat_format::primary_key_column_index;
use crate::sst::parquet::format::PrimaryKeyArray;
/// Timeout to send a batch to a sender.
const SEND_TIMEOUT: Duration = Duration::from_micros(100);
/// Maximum number of candidate series retained before distributing assignments.
const CANDIDATE_SERIES_ASSIGNMENT_THRESHOLD: usize = 1_000_000;
@@ -871,13 +868,10 @@ impl RegionScanner for SeriesScan {
.map_err(BoxedError::new)
}
fn raw_series_ordering(&self) -> SeriesOrdering {
// 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),
}
fn raw_series_ordering(&self) -> Result<SeriesOrdering, BoxedError> {
self.uses_prepared_columnar()
.map(|prepared_columnar| self.raw_ordering(prepared_columnar))
.map_err(BoxedError::new)
}
fn raw_series_schema(
@@ -1268,7 +1262,6 @@ impl SeriesDistributor {
}
metrics.scan_cost += fetch_start.elapsed();
metrics.num_series_send_timeout = self.senders.num_timeout;
metrics.num_series_send_full = self.senders.num_full;
part_metrics.set_distributor_metrics(&metrics);
@@ -1423,8 +1416,6 @@ struct SenderList {
num_nones: usize,
/// Index of the current partition to send.
sender_idx: usize,
/// Number of timeout.
num_timeout: usize,
/// Number of full senders.
num_full: usize,
}
@@ -1436,7 +1427,6 @@ impl SenderList {
senders,
num_nones,
sender_idx: 0,
num_timeout: 0,
num_full: 0,
}
}
@@ -1472,46 +1462,35 @@ impl SenderList {
}
/// Finds a partition and sends the batch to the partition.
async fn send_batch(&mut self, mut batch: SeriesBatch) -> Result<()> {
// Sends the batch without blocking first.
match self.try_send_batch(batch)? {
Some(b) => {
// Unable to send batch to partition.
batch = b;
}
None => {
return Ok(());
}
}
async fn send_batch(&mut self, batch: SeriesBatch) -> Result<()> {
let Some(batch) = self.try_send_batch(batch)? else {
return Ok(());
};
loop {
ensure!(self.num_nones < self.senders.len(), InvalidSenderSnafu);
let sender_idx = self.fetch_add_sender_idx();
let Some(sender) = &self.senders[sender_idx] else {
continue;
// Waits on every partition at once and sends to whichever frees a slot first. A
// consumer may leave partitions unpolled for a long time, e.g. one that drains them one
// after another, and waiting on such a partition would stall the scan.
let closed = {
let reservations = self.senders.iter().enumerate().filter_map(|(idx, sender)| {
sender
.as_ref()
.map(|sender| Box::pin(async move { (idx, sender.reserve().await) }))
});
let ((idx, reserved), _, _) = futures::future::select_all(reservations).await;
match reserved {
Ok(permit) => {
permit.send(Ok(batch));
return Ok(());
}
Err(_) => idx,
}
};
// Adds a timeout to avoid blocking indefinitely and sending
// the batch in a round-robin fashion when some partitions
// don't poll their inputs. This may happen if we have a
// node like sort merging. But it is rare when we are using SeriesScan.
match sender.send_timeout(Ok(batch), SEND_TIMEOUT).await {
Ok(()) => break,
Err(SendTimeoutError::Timeout(res)) => {
self.num_timeout += 1;
// Safety: we send Ok.
batch = res.unwrap();
}
Err(SendTimeoutError::Closed(res)) => {
self.senders[sender_idx] = None;
self.num_nones += 1;
// Safety: we send Ok.
batch = res.unwrap();
}
}
self.senders[closed] = None;
self.num_nones += 1;
}
Ok(())
}
async fn send_error(&self, error: Error) {
@@ -1721,6 +1700,37 @@ mod tests {
assert_eq!(0, active_receivers.load(Ordering::Relaxed));
}
#[tokio::test]
async fn distributor_does_not_stall_on_partitions_read_one_after_another() {
const NUM_BATCHES: usize = 2000;
let (mut senders, receivers) = new_legacy_channel_list(16);
let producer = tokio::spawn(async move {
for _ in 0..NUM_BATCHES {
senders
.send_batch(SeriesBatch::Flat(FlatSeriesBatch::default()))
.await
.unwrap();
}
});
// Reads partition 0 to its end before touching the next one, so all other partitions stay
// full while the producer still has batches to send.
let consume = async {
let mut received = 0;
for mut receiver in receivers.into_iter().flatten() {
while let Some(batch) = receiver.recv().await {
batch.unwrap();
received += 1;
}
}
received
};
let received = tokio::time::timeout(std::time::Duration::from_secs(1), consume)
.await
.expect("the distributor stalled on partitions the consumer was not reading");
producer.await.unwrap();
assert_eq!(NUM_BATCHES, received);
}
fn new_test_record_batch(
primary_keys: &[&[u8]],
timestamps: &[i64],
src/mito2/src/read/series_scan/prepared/tests.rs
+28 -3
View File
@@ -479,7 +479,10 @@ async fn native_scan_falls_back_to_raw_sources_until_columnar_coverage(#[case] v
.await
.unwrap();
let binding = labels(raw_scanner.metadata());
assert_eq!(fallback_ordering, raw_scanner.raw_series_ordering());
assert_eq!(
fallback_ordering,
raw_scanner.raw_series_ordering().unwrap()
);
assert!(raw_scanner.raw_series_schema(&binding).unwrap().is_some());
assert_eq!(
vec![(1000, 1.0), (2000, 1.0)],
@@ -522,7 +525,7 @@ async fn native_scan_falls_back_to_raw_sources_until_columnar_coverage(#[case] v
let binding = labels(raw_scanner.metadata());
assert_eq!(
SeriesOrdering::KeyThenTime,
raw_scanner.raw_series_ordering()
raw_scanner.raw_series_ordering().unwrap()
);
let source = raw_scanner
.prepare_series(&ctx, &metrics, &binding)
@@ -564,7 +567,10 @@ async fn native_scan_falls_back_to_raw_sources_until_columnar_coverage(#[case] v
.await
.unwrap();
let binding = labels(raw_scanner.metadata());
assert_eq!(fallback_ordering, raw_scanner.raw_series_ordering());
assert_eq!(
fallback_ordering,
raw_scanner.raw_series_ordering().unwrap()
);
assert!(raw_scanner.raw_series_schema(&binding).unwrap().is_some());
assert!(
raw_scanner
@@ -1218,6 +1224,25 @@ async fn preflight_rejects_unsupported_assignments_without_artifact_io() {
.to_string()
.contains("invalid prepared scan range")
);
// The public ordering API must propagate malformed scanner assignments, rather than
// silently describing them as an ordinary scan.
let mut scan = SeriesScan::new(make_context().input, false);
scan.prepare(PrepareRequest::default()).unwrap();
let malformed = PartitionRange {
identifier: ranges[0].identifier + ranges.len(),
..ranges[0]
};
// Inject after valid preparation: prepare() also checks scan limits by indexing ranges,
// whereas this test targets the ordering API's prepared-input validation.
scan.properties.partitions = vec![vec![malformed]];
assert!(
scan.raw_series_ordering()
.unwrap_err()
.to_string()
.contains("invalid prepared scan range")
);
let scan = SeriesScan::new(context.input, false);
let binding = labels(scan.metadata());
// The raw path stays available for the ordinary scan and preparation is skipped.
src/promql/AGENTS.md
+59 -88
View File
@@ -3,111 +3,82 @@
Navigation map for `src/promql`. Repo-wide invariants:
[`.agents/architecture-invariants.md`](../../.agents/architecture-invariants.md).
`promql` owns PromQL evaluation: the series/evaluated data model, range and
instant evaluation, function kernels, vector matching, and the PromQL extension
operators. Statement-to-plan translation lives in `src/query/src/promql/`.
`promql` owns PromQL evaluation: series/evaluated data, range and instant
functions, vector matching and extension operators. Statement planning lives in
`src/query/src/promql/`.
## Module map
| Area | Path | Entry point |
| --- | --- | --- |
| Data model | `src/promql/src/grid_series.rs`, `src/promql/src/evaluated.rs` | Grid series, evaluated chunks |
| Data model | `src/promql/src/grid_series.rs`, `src/promql/src/evaluated.rs` | Grid series and evaluated chunks |
| Evaluation | `src/promql/src/range_eval.rs`, `src/promql/src/instant_eval.rs` | Range and instant evaluation |
| 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 and execution |
| Extension operators | `src/promql/src/extension_plan.rs`, `src/promql/src/extension_plan/` | Logical and physical extension nodes |
| Native model | `src/promql/src/native/` | Bound tree, execution, metrics and display |
| Native compiler | `src/query/src/promql/native_compile.rs` | PromQL AST to native tree |
## Evaluation paths
## Native execution boundary
- **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.
Native PromQL is an experimental, standalone evaluation route alongside the
existing DataFusion PromQL path. It operates on one storage region, which may
supply multiple scan partitions. It can evaluate each independent input series
in those partitions, but does not merge one logical series split across regions
or partitions; cross-region `SeriesMerge`/`Exchange` and global final-label
semantics are not implemented. Do not describe partition-local work as a global
series merge.
Vector operations use the existing vector-matching semantics; do not describe
them as pairing only corresponding child series. A single-input operation can
independently map N series to N outputs. Aggregation placement must respect its
stage: `Complete` is valid only when the execution sees all contributions to
each group. Partitioned aggregation requires `Partial` state per input partition
followed by `Final` aggregation of that state. Do not claim that a complete
aggregate on each scan partition is globally complete.
The native `timestamp` operation evaluates on its input grid, and literal
scalar constants (including foldable constant expressions) bind on the
statement grid. This does not imply support for time-varying `time()` or
`scalar(v)`. Native `count` produces Float64 values, unlike the legacy PromQL
path's Int64 count. Native SQL/TQL output is not necessarily schema-identical
to legacy output.
The native tree and its compiler describe the implemented model, not the whole
PromQL surface. Keep comments aligned with executable nodes and kernels rather
than listing an exhaustive function allowlist. In range functions,
`double_exponential_smoothing`, `mad_over_time`, and `ts_of_*_over_time` are
experimental in Prometheus; `first_over_time` is not.
## Change coupling
- 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` and
`RawSeriesSource` are the exception: they are not serialized to a datanode. A
`FusedSeries` is placed in the region only over direct scans of one table, with
a bound metric name and a tree that reads one series at a time. A tree over two
tables stays on the frontend over one merge scan per table, each carrying a
`RawSeriesSource` that answers explicit-label raw frames; the frame schema is
compared exactly at the merge scan, so its declaration and the region's output
change together. 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`.
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 native nodes must
not be treated as serializable or distributed operators. New logical extension
nodes also need a physical planner reachable from
`DfQueryPlanner::create_physical_plan`, assembled in
`src/query/src/query_engine/state.rs`.
A `FusedSeries` over two tables stays on the frontend over one merge scan per
table, each carrying a `RawSeriesSource` that answers explicit-label raw frames.
The frame schema is compared exactly at the merge scan, so its declaration and
the region's output change together.
## Gotchas
- **One operation, two implementations; build on the live one.** `InstantManipulate`
routes to `InstantManipulateExec`. The native alternative
`NativeInstantManipulateExec` is crate-private and
`InstantManipulate::try_to_native_execution_plan` has no caller on any branch —
its doc is written for a caller that does not exist.
- **A native window requires non-interleaved input.** Only `Contiguous` and
`KeyThenTime` are accepted; `Interleaved` is refused when compiled.
- **A fused tree runs ordinary raw children and a literal Scalar root.** Recursive native
execution keeps `source_index` routing for one or more raw children; only a validated
zero-source Scalar root is materialized on its bound grid without a scanner. It must not
displace the current multi-input vector path or the single-input partial aggregate path.
- **A series closes exactly once**, on the chunk that says so — per input series,
not per label set. A node that rewrites labels can hand two input series over
under one label set, and each closes on its own chunk, so a label set may be
closed more than once; a merge above them must not read the second close as a
defect. That closing chunk is the only way a consumer distinguishes "nothing more
here" from "no more input": dropping it leaves a series open, and a child that
merges regions must recompute finality globally rather than pass a region-local
flag through.
- **Raw-frame labels come from the binding, not the schema.** A raw series frame
carries an opaque encoded primary key while its schema metadata describes only
the format: `greptime.series.layout`, `.content`, `.labels` and `.time`, plus
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.
- **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
is not part of that check.
A series closes exactly once per input series, on the chunk that says so—not
once per final label set. A label-rewriting operation may cause distinct input
series to share final labels, and each still closes on its own chunk. The
closing chunk is the only way a consumer distinguishes “nothing more here”
from “no more input”; dropping it leaves a series open. Consumers must
preserve this per-series finality, including sparse trailing empty chunks. A
merge above them must not mistake a second close for a defect. A child that
merges regions must recompute finality globally rather than pass through a
region-local flag.
## Testing
## Validation
```bash
cargo nextest run -p promql
src/promql/src/extension_plan/aggregate_finalize.rs
+29 -1
View File
@@ -926,7 +926,10 @@ async fn merge_into_batches(
reservation.try_resize(merger.state_bytes())?;
}
// Only a successful end of the input answers: a failed read dropped the state above instead.
let series = merger.finish_budgeted(|total| reservation.try_resize(total))?;
let mut series = merger.finish_budgeted(|total| reservation.try_resize(total))?;
// Partition arrival order is not a statement-level output order. Match Prometheus label
// comparison before presenting the finalized groups so equivalent executions answer alike.
series.sort_by(compare_series_labels);
// The merge is consumed by its answer, which replaces what it kept. The whole answer is charged at
// once and held until the loop ends: the vector and every series still waiting in it stay live
// while the loop folds them one by one, so releasing them one at a time would leave the rest
@@ -976,6 +979,31 @@ async fn merge_into_batches(
Ok(batches)
}
/// Compares the complete Prometheus label sets in label-name order, with no temporary allocation.
fn compare_series_labels(a: &Series, b: &Series) -> Ordering {
fn labels(series: &Series) -> impl Iterator<Item = (&str, &str)> {
let tags = series.metric.tags();
let name_index = tags.partition_point(|tag| tag.name.as_ref() < "__name__");
tags[..name_index]
.iter()
.map(|tag| (tag.name.as_ref(), tag.value.as_ref()))
.chain(
series
.metric
.name()
.into_iter()
.map(|name| ("__name__", name)),
)
.chain(
tags[name_index..]
.iter()
.map(|tag| (tag.name.as_ref(), tag.value.as_ref())),
)
}
labels(a).cmp(labels(b))
}
/// Hands over every chunk the suffix has settled, as rows.
///
/// Each chunk is admitted before it is converted, by a coarse estimate of the storage its rows need -
src/promql/src/extension_plan/aggregate_finalize/exec_tests.rs
+81
View File
@@ -485,6 +485,87 @@ async fn the_state_of_two_producers_merges_into_one_series() {
);
}
/// The metric name participates in the same sorted label comparison as ordinary tags.
#[test]
fn finalized_label_comparison_includes_metric_name() {
let schema = cut_schema();
let make = |name: Option<&str>| {
Series::try_new(
&schema,
SeriesLabels::try_new(
name.map(Arc::from),
vec![Label {
name: K0.into(),
value: LABEL.into(),
}],
)
.unwrap(),
TimeAxis::Grid(GridRange::try_new(grid(), POINT, 1).unwrap()),
SeriesFields::try_from_vec(vec![Samples::from_float(Float64Array::from(vec![1.0]))])
.unwrap(),
)
.unwrap()
};
let absent = make(None);
let a_name = make(Some("a"));
let z_name = make(Some("z"));
assert_eq!(
compare_series_labels(&absent, &a_name),
std::cmp::Ordering::Greater
);
assert_eq!(
compare_series_labels(&a_name, &z_name),
std::cmp::Ordering::Less
);
}
/// Final groups are presented by label order, not by the order partitions return their states.
#[tokio::test]
async fn finalized_groups_are_ordered_independently_of_partition_arrival() {
let kind = AggregateKind::Avg;
let schema = cut_schema();
let make_state = |value: &str, sample: f64| {
let metric = SeriesLabels::try_new(
Some(METRIC.into()),
vec![Label {
name: K0.into(),
value: value.into(),
}],
)
.unwrap();
let series = Series::try_new(
&schema,
metric,
TimeAxis::Grid(GridRange::try_new(grid(), POINT, 1).unwrap()),
SeriesFields::try_from_vec(vec![Samples::from_float(Float64Array::from(vec![sample]))])
.unwrap(),
)
.unwrap();
state_of_series(&schema, cut_op(kind), &[K0], vec![series])
};
let a = make_state("a", 2.0);
let z = make_state("z", 8.0);
let run = |states| async move {
let plan = finalize(avg_tree(kind), output_schema(VALUE, &[K0]))
.to_execution_plan(state_input(states))
.unwrap();
assert_eq!(plan.properties().emission_type, EmissionType::Final);
collect(plan, SessionContext::default().task_ctx())
.await
.unwrap()
};
let forward = rows(&run(vec![a.clone(), z.clone()]).await);
let reverse = rows(&run(vec![z, a]).await);
assert_eq!(forward, reverse);
assert_eq!(
forward
.iter()
.map(|row| row.2.as_deref())
.collect::<Vec<_>>(),
[Some("a"), Some("z")]
);
}
/// Every fixed-width accumulator merges the producers' state, whichever fields it keeps.
#[tokio::test]
async fn every_fixed_width_kind_merges_across_producers() {
src/promql/src/extension_plan/fused_series.rs
+8 -3
View File
@@ -1044,9 +1044,14 @@ impl FusedSeriesExec {
} else {
partition_count
}),
inputs.first().map_or(EmissionType::Final, |(input, _, _)| {
input.properties().emission_type
}),
if partial_fold.is_some() {
// Partial aggregate state is only emitted after the scan reaches EOF.
EmissionType::Final
} else {
inputs.first().map_or(EmissionType::Final, |(input, _, _)| {
input.properties().emission_type
})
},
inputs
.first()
.map_or(Boundedness::Bounded, |(input, _, _)| {
src/promql/src/extension_plan/fused_series/tests.rs
+5
View File
@@ -3630,6 +3630,11 @@ async fn a_partial_fold_exports_state_from_every_partition() {
)
.expect("a rate tree with a partial count is executable");
assert_eq!(
datafusion::physical_plan::execution_plan::EmissionType::Final,
exec.properties().emission_type,
"the partial state is emitted only after the input reaches EOF"
);
assert_eq!(
2,
exec.properties().output_partitioning().partition_count(),
src/promql/src/extension_plan/instant_manipulate.rs
+131 -20
View File
@@ -12,9 +12,6 @@
// See the License for the specific language governing permissions and
// limitations under the License.
mod native;
mod native_exec;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};
@@ -22,7 +19,7 @@ use std::task::{Context, Poll};
use common_query::prelude::{greptime_native_histogram, greptime_value};
use datafusion::arrow::array::{
Array, ArrayRef, Float64Array, LargeListArray, LargeStringArray, ListArray, StringArray,
StructArray, TimestampMillisecondArray, UInt64Array,
StringViewArray, StructArray, TimestampMillisecondArray, UInt64Array,
};
use datafusion::arrow::datatypes::{DataType, SchemaRef};
use datafusion::arrow::record_batch::RecordBatch;
@@ -80,14 +77,12 @@ fn mixed_sample_fields(field: Option<&str>) -> [Option<&str>; 2] {
/// `start` and `end` with step `interval`. Find in the `lookback` range if data
/// is missing at the given timestamp.
///
/// The memory pool bounds one instant query, not a point count: the execution reserves the
/// workspace its selection allocates before allocating it, and a pool that cannot meet the
/// reservation refuses the query with `ResourcesExhausted` (see
/// [`InstantManipulateStream::manipulate_reserved`]). That bound exists only where the pool does:
/// a resolved `query.memory_pool_size` of zero leaves DataFusion's unbounded pool attached, whose
/// growth always succeeds, so an operator that leaves the option unset has no bound on this path
/// at all. What one query can cost still grows with the number of series, which no reservation of
/// a single series accounts for.
/// The memory pool admits this operator's per-series selection workspace: the execution reserves
/// the estimate before allocating that workspace and a pool that cannot meet it refuses the
/// operation with `ResourcesExhausted` (see [`InstantManipulateStream::manipulate_reserved`]).
/// The reservation is released as soon as the output batch is built; it does not bound the child's
/// input, the returned output batch, or total memory retained by an instant query. Its effect also
/// depends on the attached pool: an unbounded pool always accepts the reservation.
#[derive(Debug, PartialEq, Eq, Hash)]
pub struct InstantManipulate {
start: Millisecond,
@@ -906,9 +901,9 @@ impl InstantManipulateStream {
/// What one selected series' workspace holds, in bytes, before it is allocated.
///
/// The figure follows the reasoning of `InstantRuntime::gather_bytes` in `native::execute` and
/// of `NativeInstantStream::finish_series` in `native_exec`: a selection decides one point of
/// its grid per sample position it reads, so what it allocates grows with the points of the
/// The figure follows the reasoning of `InstantRuntime::gather_bytes` in `native::execute`: a
/// selection decides one point of its grid per sample position it reads, so what it allocates
/// grows with the points of the
/// clipped grid and not with the samples of the series. Each term below is one allocation, and
/// the ones this path shares with those are named as they are there.
///
@@ -952,7 +947,20 @@ impl InstantManipulateStream {
continue;
}
bytes = bytes.saturating_add(match row_bytes(column.as_ref()) {
Some(width) => points.saturating_mul(width.saturating_add(1)),
Some(width) => {
let output_offsets = match column.data_type() {
DataType::Utf8 => points
.saturating_add(1)
.saturating_mul(std::mem::size_of::<i32>()),
DataType::LargeUtf8 => points
.saturating_add(1)
.saturating_mul(std::mem::size_of::<i64>()),
_ => 0,
};
points
.saturating_mul(width.saturating_add(1))
.saturating_add(output_offsets)
}
None => column.get_buffer_memory_size(),
});
}
@@ -1053,9 +1061,12 @@ impl InstantManipulateStream {
/// carry.
///
/// `take` allocates one row of the column per point, so one point costs the widest row the column
/// holds: the width of its value for a fixed-width column, the longest value for a byte column, and
/// the whole row for a struct such as a native histogram -- its scalar columns, the offsets of its
/// list columns and the values those lists hold for the widest row of the column.
/// holds: the width of its value for a fixed-width column; value bytes for Utf8 or LargeUtf8 (the
/// output offsets buffer is accounted once in `workspace_bytes`), a view descriptor plus value
/// bytes for Utf8View (a conservative allowance because short views may point into shared buffers),
/// dictionary key bytes for string dictionaries (whose values remain shared),
/// and the whole row for a struct such as a native histogram -- its scalar columns, the offsets of
/// its list columns and their widest values.
fn row_bytes(column: &dyn Array) -> Option<usize> {
if let Some(width) = column.data_type().primitive_width() {
return Some(width);
@@ -1070,6 +1081,28 @@ fn row_bytes(column: &dyn Array) -> Option<usize> {
if let Some(values) = column.as_any().downcast_ref::<LargeStringArray>() {
return Some(values.offsets().lengths().max().unwrap_or(0));
}
if let Some(values) = column.as_any().downcast_ref::<StringViewArray>() {
return Some(
std::mem::size_of::<u128>().saturating_add(
values
.iter()
.map(|value| value.map_or(0, str::len))
.max()
.unwrap_or(0),
),
);
}
if let DataType::Dictionary(key_type, value_type) = column.data_type() {
// `take` copies dictionary keys per output row and shares the dictionary values.
// Instant tags supported by the producer use string dictionary values.
if matches!(
value_type.as_ref(),
DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View
) {
return key_type.primitive_width();
}
return None;
}
let values = column.as_any().downcast_ref::<StructArray>()?;
Some(
values
@@ -1135,7 +1168,8 @@ mod test {
use common_query::native_histogram::build_histogram_array;
use common_query::prometheus::PROMETHEUS_STALE_NAN_BITS;
use datafusion::arrow::array::{
Float64Array, TimestampMicrosecondArray, TimestampNanosecondArray, TimestampSecondArray,
Float64Array, LargeStringArray, StringArray, StringViewArray, TimestampMicrosecondArray,
TimestampNanosecondArray, TimestampSecondArray,
};
use datafusion::arrow::buffer::NullBuffer;
use datafusion::arrow::datatypes::{DataType, Field, Schema, TimeUnit};
@@ -1951,6 +1985,83 @@ mod test {
);
}
#[test]
fn row_bytes_accounts_for_string_representation_storage() {
use datafusion::arrow::array::{DictionaryArray, Int8Array};
use datafusion::arrow::datatypes::Int8Type;
let utf8 = StringArray::from(vec!["a", "long"]);
// `row_bytes` is the per-row value width; workspace offsets are charged separately.
assert_eq!(row_bytes(&utf8), Some(4));
let large_utf8 = LargeStringArray::from(vec!["a", "long"]);
assert_eq!(row_bytes(&large_utf8), Some(4));
let view = StringViewArray::from(vec!["a", "long"]);
assert_eq!(row_bytes(&view), Some(16 + 4));
let dictionary = DictionaryArray::<Int8Type>::try_new(
Int8Array::from(vec![Some(0), Some(1)]),
Arc::new(StringArray::from(vec!["a", "long"])),
)
.unwrap();
assert_eq!(row_bytes(&dictionary), Some(std::mem::size_of::<i8>()));
}
#[test]
fn workspace_bytes_charges_utf8_output_offsets_once() {
let schema = Arc::new(Schema::new(vec![
Field::new(
TIME_INDEX_COLUMN,
DataType::Timestamp(TimeUnit::Millisecond, None),
false,
),
Field::new("utf8", DataType::Utf8, false),
]));
let batch = RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(TimestampMillisecondArray::from(vec![0, 1])),
Arc::new(StringArray::from(vec!["payload", "payload"])),
],
)
.unwrap();
let stream = InstantManipulateStream {
offset: 0,
start: 0,
end: 1,
lookback_delta: 1,
interval: 1,
time_index: 0,
time_unit: TimeUnit::Millisecond,
field_indices: [None, None],
tsid_index: None,
reuse_tsid_column: false,
schema,
input: Box::pin(
datafusion::physical_plan::memory::MemoryStream::try_new(
vec![],
batch.schema(),
None,
)
.unwrap(),
),
metric: BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0),
num_series: Count::new(),
reservation: unbounded(),
};
let points = stream.clipped_points(&batch);
assert_eq!(points, 2);
let estimate = stream.workspace_bytes(&batch);
let output = stream.manipulate(batch.clone()).unwrap();
let output_buffers = output
.columns()
.iter()
.map(|column| column.get_buffer_memory_size())
.sum::<usize>();
assert!(estimate >= output_buffers);
}
/// The figure `manipulate_reserved` grows is the clipped one: 262 bytes for a series that ends
/// at its first sample, 413 337 bytes for one that spans the same seven-day grid of 10 081
/// points at a 60 s step. Per point it is the four `u64` of the two index vectors and the
src/promql/src/extension_plan/instant_manipulate/native.rs
-301
View File
@@ -1,301 +0,0 @@
// Copyright 2023 Greptime Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use std::sync::Arc;
use datafusion::arrow::array::{ArrayRef, Float64Array, StringArray, TimestampMillisecondArray};
use datafusion::arrow::datatypes::{DataType, SchemaRef, TimeUnit};
use datafusion::arrow::record_batch::RecordBatch;
use datafusion::error::{DataFusionError, Result as DataFusionResult};
use datatypes::series::SeriesLabels;
use datatypes::series::raw::RawTimestamps;
use datatypes::series::samples::{SampleRef, Samples, SeriesFields, SeriesSchema};
#[cfg(test)]
#[path = "native_tests.rs"]
mod tests;
use crate::extension_plan::instant_manipulate::InstantManipulate;
use crate::instant_eval::select_instants;
#[derive(Clone)]
enum Column {
Time,
Field(usize),
Tag(String),
}
/// Preflighted mapping from a complete typed series to the original query schema.
pub(crate) struct NativeInstantMapping {
schema: SchemaRef,
series_schema: Arc<SeriesSchema>,
columns: Vec<Column>,
fields: Vec<usize>,
start: i64,
end: i64,
step: i64,
lookback: i64,
}
impl NativeInstantMapping {
/// Returns `None` when the node/schema is outside the bounded native subset.
pub(crate) fn try_new(
node: &InstantManipulate,
schema: SchemaRef,
series_schema: Arc<SeriesSchema>,
logical_tags: &[String],
) -> DataFusionResult<Option<Self>> {
if node.offset != 0 || node.unfix.is_some() || node.start >= node.end {
return Ok(None);
}
if node.interval <= 0 {
return Err(DataFusionError::Execution(
"InstantManipulate interval must be positive".into(),
));
}
if node.lookback_delta < 0 {
return Err(DataFusionError::Execution(
"InstantManipulate lookback must be non-negative".into(),
));
}
if series_schema.eval_grid.is_some() {
return Ok(None);
}
let mut names = std::collections::HashSet::new();
for field in series_schema.fields.iter() {
if !names.insert(field.name.as_ref()) {
return Ok(None);
}
}
if logical_tags.iter().any(|name| !names.insert(name.as_str())) {
return Ok(None);
}
let stale_field_names = node.staleness_field_columns().collect::<Vec<_>>();
if stale_field_names.iter().any(|name| {
series_schema
.fields
.iter()
.all(|field| field.name.as_ref() != *name)
}) {
return Ok(None);
}
let mut columns = Vec::with_capacity(schema.fields().len());
let mut seen_fields = vec![false; series_schema.fields.len()];
let mut seen_tags = vec![false; logical_tags.len()];
let mut seen_time = false;
for field in schema.fields() {
let name = field.name();
let column = if name == node.time_index_column.as_str() {
if seen_time
|| field.data_type() != &DataType::Timestamp(TimeUnit::Millisecond, None)
{
return Ok(None);
}
seen_time = true;
Column::Time
} else if let Some(index) = series_schema
.fields
.iter()
.position(|v| v.name.as_ref() == name)
{
if seen_fields[index] || field.data_type() != &DataType::Float64 {
return Ok(None);
}
seen_fields[index] = true;
Column::Field(index)
} else if let Some(index) = logical_tags.iter().position(|tag| tag == name) {
if seen_tags[index] || !is_string_type(field.data_type()) {
return Ok(None);
}
seen_tags[index] = true;
Column::Tag(name.clone())
} else {
return Ok(None);
};
columns.push(column);
}
if !seen_time || seen_fields.iter().any(|seen| !seen) || seen_tags.iter().any(|seen| !seen)
{
return Ok(None);
}
let fields = stale_field_names
.iter()
.filter_map(|name| {
series_schema
.fields
.iter()
.position(|field| field.name.as_ref() == *name)
})
.collect();
Ok(Some(Self {
schema,
series_schema,
columns,
fields,
start: node.start,
end: node.end,
step: node.interval,
lookback: node.lookback_delta,
}))
}
/// The original output schema retained by this mapping.
pub(crate) fn schema(&self) -> SchemaRef {
self.schema.clone()
}
/// Estimates the clipped evaluation grid for complete source timestamp bounds.
///
/// This count is what the caller reserves the workspace of a selection for before it selects,
/// so nothing but the reservation bounds the grid: a grid too long to index is refused as
/// exhausted rather than by a point limit of its own.
pub(crate) fn estimated_points(&self, first: i64, last: i64) -> DataFusionResult<usize> {
let first = first as i128;
let last = last as i128;
let useful_last = if self.lookback == 0 {
last
} else {
last + self.lookback as i128 - 1
};
let query_start = self.start as i128;
let query_end = self.end as i128;
let max_start = first.max(query_start);
let min_end = useful_last.min(query_end);
let (aligned_start, aligned_end) = if max_start > min_end {
(1_i128, 0_i128)
} else {
(
query_start + (max_start - query_start) / self.step as i128 * self.step as i128,
query_end - (query_end - min_end) / self.step as i128 * self.step as i128,
)
};
let points = if aligned_end >= aligned_start {
(aligned_end - aligned_start) / self.step as i128 + 1
} else {
0
};
usize::try_from(points).map_err(|_| {
DataFusionError::ResourcesExhausted("native instant point count overflow".into())
})
}
/// Selects instants against typed samples, then maps only selected values to Arrow.
pub(crate) fn evaluate(
&self,
metric: &SeriesLabels,
time: &RawTimestamps,
fields: &SeriesFields,
) -> DataFusionResult<RecordBatch> {
if fields.len() != self.series_schema.fields.len() {
return Err(DataFusionError::Execution(
"native instant field count mismatch".into(),
));
}
let checked_fields = self
.fields
.iter()
.map(|i| fields.as_slice()[*i].clone())
.collect::<Vec<_>>();
let checked_fields = if checked_fields.is_empty() {
vec![Samples::from_float(Float64Array::new_null(time.len()))]
} else {
checked_fields
};
let checked = SeriesFields::try_from_vec(checked_fields)
.map_err(|e| DataFusionError::Execution(e.to_string()))?;
let (indices, timestamps) = select_instants(
time,
&checked,
self.start..=self.end,
self.step,
self.lookback,
0,
)?;
if indices.is_empty() {
return Ok(RecordBatch::new_empty(self.schema.clone()));
}
let arrays: DataFusionResult<Vec<ArrayRef>> = self
.columns
.iter()
.enumerate()
.map(|(output_index, column)| match column {
Column::Time => {
Ok(Arc::new(TimestampMillisecondArray::from(timestamps.clone())) as ArrayRef)
}
Column::Field(index) => {
let source = fields.get(*index).ok_or_else(|| {
DataFusionError::Execution("native instant field missing".into())
})?;
let mut samples = source.iter();
let mut cursor = 0usize;
let mut selected_index = None;
let mut selected_value = None;
let mut values = Vec::with_capacity(indices.len());
for index in &indices {
let index = *index as usize;
if selected_index != Some(index) {
while cursor <= index {
selected_value = Some(
match samples.next().ok_or_else(|| {
DataFusionError::Execution(
"native instant sample iterator ended before selected index"
.into(),
)
})? {
SampleRef::Float(value) => Some(value),
SampleRef::Missing => None,
SampleRef::Histogram(_) => {
return Err(DataFusionError::Execution(
"native instant expected Float64 sample".into(),
));
}
},
);
selected_index = Some(cursor);
cursor += 1;
}
}
values.push(selected_value.ok_or_else(|| {
DataFusionError::Execution(
"native instant selected sample was not initialized".into(),
)
})?);
}
Ok(Arc::new(Float64Array::from(values)) as ArrayRef)
}
Column::Tag(name) => {
let values = std::iter::repeat_with(|| metric.get(name))
.take(indices.len())
.collect::<StringArray>();
let target = self.schema.field(output_index).data_type();
let cast = datafusion::arrow::compute::cast(&values, target)?;
Ok(cast)
}
})
.collect();
Ok(RecordBatch::try_new(self.schema.clone(), arrays?)?)
}
}
fn is_string_type(data_type: &DataType) -> bool {
matches!(
data_type,
DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View
) || matches!(data_type, DataType::Dictionary(key, value) if key.as_ref() == &DataType::UInt32 && value.as_ref() == &DataType::Utf8)
}
src/promql/src/extension_plan/instant_manipulate/native_exec.rs
-500
View File
@@ -1,500 +0,0 @@
// Copyright 2023 Greptime Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use std::pin::Pin;
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;
use datafusion::arrow::datatypes::SchemaRef;
use datafusion::arrow::record_batch::RecordBatch;
use datafusion::common::tree_node::TreeNodeRecursion;
use datafusion::error::{DataFusionError, Result as DfResult};
use datafusion::execution::context::TaskContext;
use datafusion::execution::memory_pool::{MemoryConsumer, MemoryReservation};
use datafusion::physical_expr::EquivalenceProperties;
use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
use datafusion::physical_plan::{
DisplayAs, DisplayFormatType, ExecutionPlan, PhysicalExpr, PlanProperties, RecordBatchStream,
SendableRecordBatchStream,
};
use datatypes::series::raw::{RawSeriesChunk, RawTimestamps};
use datatypes::series::samples::{DenseSamplesRef, Samples, SeriesFields, SeriesSchema};
use futures::{Stream, ready};
use store_api::region_engine::{RawLabelSource, SendableRawSeriesStream};
use crate::extension_plan::instant_manipulate::InstantManipulate;
use crate::extension_plan::instant_manipulate::native::NativeInstantMapping;
use crate::source::materialize_raw_labels;
#[cfg(test)]
#[path = "native_exec_tests.rs"]
mod tests;
impl InstantManipulate {
/// Returns whether the bounded native mapping accepts this output schema, series schema, and
/// complete logical tag set.
///
/// This is the local half of [`Self::try_to_native_execution_plan`]: the same decision that
/// method makes before it inspects its input, so a caller can decline a plan without building
/// one. `Err` is the same error that method propagates for an invalid interval or lookback.
pub fn is_native_mapping_supported(
&self,
output_schema: SchemaRef,
series_schema: Arc<SeriesSchema>,
logical_tags: &[String],
) -> DfResult<bool> {
Ok(
NativeInstantMapping::try_new(self, output_schema, series_schema, logical_tags)?
.is_some(),
)
}
/// Builds the bounded native execution path for a prepared raw-series frame.
pub fn try_to_native_execution_plan(
&self,
input: Arc<dyn ExecutionPlan>,
output_schema: SchemaRef,
label_source: RawLabelSource,
series_schema: Arc<SeriesSchema>,
logical_tags: &[String],
) -> DfResult<Option<Arc<dyn ExecutionPlan>>> {
let Some(mapping) = NativeInstantMapping::try_new(
self,
output_schema,
series_schema.clone(),
logical_tags,
)?
else {
return Ok(None);
};
if input.properties().partitioning.partition_count() != 1 {
return Ok(None);
}
let declared_schema = input.schema();
// Validate the declared raw-frame schema without executing the child. The frame states its
// own label layout, so the plan's binding is checked against that declaration rather than
// supplied to the decoder.
ensure_frame_matches_label_source(declared_schema.as_ref(), &label_source)?;
unpack_raw_series(
Box::pin(RecordBatchStreamAdapter::new(
declared_schema.clone(),
futures::stream::empty(),
)),
Some(series_schema.clone()),
)?;
Ok(Some(Arc::new(NativeInstantManipulateExec::new(
input,
Arc::new(mapping),
declared_schema,
series_schema,
label_source,
))))
}
}
struct NativeInstantManipulateExec {
input: Arc<dyn ExecutionPlan>,
mapping: Arc<NativeInstantMapping>,
raw_schema: SchemaRef,
series_schema: Arc<SeriesSchema>,
label_source: RawLabelSource,
properties: Arc<PlanProperties>,
}
impl NativeInstantManipulateExec {
fn new(
input: Arc<dyn ExecutionPlan>,
mapping: Arc<NativeInstantMapping>,
raw_schema: SchemaRef,
series_schema: Arc<SeriesSchema>,
label_source: RawLabelSource,
) -> Self {
let properties = Arc::new(PlanProperties::new(
EquivalenceProperties::new(mapping.schema()),
datafusion::physical_plan::Partitioning::UnknownPartitioning(1),
EmissionType::Incremental,
Boundedness::Bounded,
));
Self {
input,
mapping,
raw_schema,
series_schema,
label_source,
properties,
}
}
}
impl std::fmt::Debug for NativeInstantManipulateExec {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("NativeInstantManipulateExec").finish()
}
}
impl DisplayAs for NativeInstantManipulateExec {
fn fmt_as(&self, _t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result {
write!(f, "NativeInstantManipulateExec")
}
}
impl ExecutionPlan for NativeInstantManipulateExec {
fn name(&self) -> &str {
"NativeInstantManipulateExec"
}
fn schema(&self) -> SchemaRef {
self.mapping.schema()
}
fn apply_expressions(
&self,
_f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> DfResult<TreeNodeRecursion>,
) -> DfResult<TreeNodeRecursion> {
Ok(TreeNodeRecursion::Continue)
}
fn properties(&self) -> &Arc<PlanProperties> {
&self.properties
}
fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
vec![&self.input]
}
fn with_new_children(
self: Arc<Self>,
children: Vec<Arc<dyn ExecutionPlan>>,
) -> DfResult<Arc<dyn ExecutionPlan>> {
if children.len() != 1
|| children[0].schema() != self.raw_schema
|| children[0].properties().partitioning.partition_count() != 1
{
return Err(DataFusionError::Plan(
"NativeInstantManipulateExec requires one unchanged one-partition raw child".into(),
));
}
Ok(Arc::new(Self::new(
children[0].clone(),
self.mapping.clone(),
self.raw_schema.clone(),
self.series_schema.clone(),
self.label_source.clone(),
)))
}
fn execute(
&self,
partition: usize,
context: Arc<TaskContext>,
) -> DfResult<SendableRecordBatchStream> {
if partition != 0 {
return Err(DataFusionError::Execution(
"native instant partition must be zero".into(),
));
}
let frame = self.input.execute(partition, context.clone())?;
// The child's frames state their own label layout; the plan's binding is an expectation
// checked against that declaration, not a decode input.
ensure_frame_matches_label_source(frame.schema().as_ref(), &self.label_source)?;
let source = unpack_raw_series(frame, Some(self.series_schema.clone()))?;
let source = materialize_raw_labels(source)
.map_err(|e| DataFusionError::Execution(e.to_string()))?;
let reservation =
MemoryConsumer::new("NativeInstantStream").register(&context.runtime_env().memory_pool);
Ok(Box::pin(NativeInstantStream {
source: Some(source),
mapping: self.mapping.clone(),
schema: self.mapping.schema(),
series_schema: self.series_schema.clone(),
pending: Vec::new(),
reservation,
}))
}
}
struct NativeInstantStream {
source: Option<SendableRawSeriesStream>,
mapping: Arc<NativeInstantMapping>,
schema: SchemaRef,
series_schema: Arc<SeriesSchema>,
pending: Vec<RawSeriesChunk>,
reservation: MemoryReservation,
}
impl NativeInstantStream {
fn fail(&mut self, error: DataFusionError) -> Poll<Option<DfResult<RecordBatch>>> {
self.source = None;
self.pending = Vec::new();
self.reservation.free();
Poll::Ready(Some(Err(error)))
}
fn finish_series(&mut self) -> DfResult<RecordBatch> {
let chunks = std::mem::take(&mut self.pending);
let rows = chunks
.iter()
.try_fold(0usize, |n, c| n.checked_add(c.time.len()))
.ok_or_else(|| {
DataFusionError::ResourcesExhausted("native instant row count overflow".into())
})?;
let first = chunks
.first()
.and_then(|c| c.time.timestamp(0))
.ok_or_else(|| {
DataFusionError::Execution("native instant series has no timestamps".into())
})?;
let last = chunks
.last()
.and_then(|c| c.time.timestamp(c.time.len().saturating_sub(1)))
.ok_or_else(|| {
DataFusionError::Execution("native instant series has no timestamps".into())
})?;
let points = self.mapping.estimated_points(first, last)?;
let label_bytes = chunks
.first()
.and_then(|c| c.metric.as_ref())
.map(|m| {
m.name().map_or(0, str::len)
+ m.tags()
.iter()
.map(|l| l.name.len() + l.value.len() + 16)
.sum::<usize>()
})
.unwrap_or(0);
// Reserve timestamp concatenation by row count, but field descriptors by their
// actual chunk/segment counts: a dense chunk is still one field descriptor.
let field_count = self.series_schema.fields.len();
let chunk_count = chunks.len();
let segment_count = chunks
.iter()
.flat_map(|chunk| chunk.fields.as_slice())
.try_fold(0usize, |count, field| {
count.checked_add(field.segments().len())
})
.ok_or_else(|| {
DataFusionError::ResourcesExhausted("native instant segment count overflow".into())
})?;
let chunk_field_descriptors = chunk_count
.checked_mul(field_count)
.and_then(|n| n.checked_mul(std::mem::size_of::<Samples>()))
.ok_or_else(|| {
DataFusionError::ResourcesExhausted(
"native instant descriptor estimate overflow".into(),
)
})?;
// Flattening and converting to Arc can transiently retain the target descriptors and
// two source/flattened vectors; three segment-counts conservatively cover those owners.
let flattened_descriptors = segment_count
.checked_mul(3)
.and_then(|n| n.checked_mul(std::mem::size_of::<Samples>()))
.ok_or_else(|| {
DataFusionError::ResourcesExhausted(
"native instant segment estimate overflow".into(),
)
})?;
let scratch = rows
.checked_mul(16)
.and_then(|n| n.checked_add(chunk_field_descriptors))
.and_then(|n| n.checked_add(flattened_descriptors))
.and_then(|n| {
n.checked_add(field_count.checked_mul(std::mem::size_of::<Vec<Samples>>())?)
})
.and_then(|n| {
// Per point of the clipped grid, the most `select_instants` and the mapping of its
// selection can allocate: 24 bytes per output column, the labels of the row, and
// 48 for the two vectors the selection returns and the time column it is mapped
// into -- a `u64` source row index, an `i64` evaluation timestamp and the `i64`
// of the output column, at 8 bytes each. The point count is what stands in for
// the point limit `estimated_points` used to carry, so the selection arrays are
// reserved for here rather than refused there.
points
.checked_mul(
self.schema
.fields()
.len()
.checked_mul(24)?
.checked_add(label_bytes)?
.checked_add(48)?,
)?
.checked_add(n)
})
.ok_or_else(|| {
DataFusionError::ResourcesExhausted(
"native instant scratch estimate overflow".into(),
)
})?;
self.reservation
.try_grow(scratch)
.map_err(|e| DataFusionError::ResourcesExhausted(e.to_string()))?;
let metric = chunks
.first()
.and_then(|c| c.metric.clone())
.ok_or_else(|| {
DataFusionError::Execution("native instant series has no metric labels".into())
})?;
let mut times = Vec::new();
times
.try_reserve_exact(rows)
.map_err(|e| DataFusionError::ResourcesExhausted(e.to_string()))?;
let mut field_chunks = (0..field_count).map(|_| Vec::new()).collect::<Vec<_>>();
for field in &mut field_chunks {
field
.try_reserve_exact(chunk_count)
.map_err(|e| DataFusionError::ResourcesExhausted(e.to_string()))?;
}
for chunk in chunks {
for i in 0..chunk.time.len() {
times.push(chunk.time.timestamp(i).ok_or_else(|| {
DataFusionError::Execution("native instant timestamp missing".into())
})?);
}
for (target, field) in field_chunks.iter_mut().zip(chunk.fields.as_slice()) {
target.push(field.clone());
}
}
let time = RawTimestamps::try_new(times.into())
.map_err(|e| DataFusionError::Execution(e.to_string()))?;
let fields = field_chunks
.into_iter()
.map(Samples::try_from_chunks)
.collect::<std::result::Result<Vec<_>, _>>()
.map_err(|e| DataFusionError::Execution(e.to_string()))?;
let fields = SeriesFields::try_from_vec(fields)
.map_err(|e| DataFusionError::Execution(e.to_string()))?;
self.mapping.evaluate(&metric, &time, &fields)
}
}
fn retained_chunk_bytes(chunk: &RawSeriesChunk) -> Option<usize> {
let timestamp_bytes = chunk.time.as_buffer().inner().capacity();
let labels = chunk
.metric
.as_ref()
.map(|metric| {
metric
.tags()
.iter()
.try_fold(metric.name().map_or(0, str::len), |total, label| {
total
.checked_add(label.name.len())?
.checked_add(label.value.len())?
.checked_add(std::mem::size_of::<datatypes::series::Label>())?
.checked_add(2 * std::mem::size_of::<usize>())
})
})
.unwrap_or(Some(0))?;
let field_descriptors = std::mem::size_of::<RawSeriesChunk>()
.checked_add(std::mem::size_of::<SeriesFields>())?
.checked_add(
chunk
.fields
.as_slice()
.len()
.checked_mul(std::mem::size_of::<Samples>())?,
)?;
let (segment_descriptors, sample_buffers) = chunk
.fields
.as_slice()
.iter()
.flat_map(|field| field.segments())
.try_fold((0usize, 0usize), |(descriptors, buffers), segment| {
let bytes = match segment.dense()? {
DenseSamplesRef::Float(array) => array.get_buffer_memory_size(),
_ => return None,
};
Some((
descriptors.checked_add(std::mem::size_of::<Samples>())?,
buffers.checked_add(bytes)?,
))
})?;
timestamp_bytes
.checked_add(chunk.primary_key.capacity())?
.checked_add(labels)?
.checked_add(field_descriptors)?
.checked_add(segment_descriptors)?
.checked_add(sample_buffers)
}
impl Stream for NativeInstantStream {
type Item = DfResult<RecordBatch>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
for _ in 0..64 {
let Some(source) = self.source.as_mut() else {
self.pending = Vec::new();
self.reservation.free();
return Poll::Ready(None);
};
match ready!(source.as_mut().poll_next(cx)) {
Some(Err(error)) => {
return self.fail(DataFusionError::External(Box::new(error)));
}
None if !self.pending.is_empty() => {
return self.fail(DataFusionError::Execution(
"raw series ended without final chunk".into(),
));
}
None => {
self.source = None;
self.pending = Vec::new();
self.reservation.free();
return Poll::Ready(None);
}
Some(Ok(chunk)) => {
let Some(bytes) = retained_chunk_bytes(&chunk) else {
return self.fail(DataFusionError::ResourcesExhausted(
"native instant retained-byte estimate overflow".into(),
));
};
if self.pending.len() == self.pending.capacity() {
let old_capacity = self.pending.capacity();
let new_capacity = old_capacity.saturating_mul(2).max(4);
let slot_growth =
new_capacity.checked_sub(old_capacity).and_then(|slots| {
slots.checked_mul(std::mem::size_of::<RawSeriesChunk>())
});
let Some(slot_growth) = slot_growth else {
return self.fail(DataFusionError::ResourcesExhausted(
"native instant pending capacity overflow".into(),
));
};
let Some(growth) = bytes.checked_add(slot_growth) else {
return self.fail(DataFusionError::ResourcesExhausted(
"native instant chunk estimate overflow".into(),
));
};
if let Err(e) = self.reservation.try_grow(growth) {
return self.fail(DataFusionError::ResourcesExhausted(e.to_string()));
}
let additional = new_capacity - self.pending.capacity();
if let Err(e) = self.pending.try_reserve_exact(additional) {
return self.fail(DataFusionError::ResourcesExhausted(e.to_string()));
}
} else if let Err(e) = self.reservation.try_grow(bytes) {
return self.fail(DataFusionError::ResourcesExhausted(e.to_string()));
}
let last = chunk.is_last;
self.pending.push(chunk);
if last {
match self.finish_series() {
Ok(batch) => {
self.pending = Vec::new();
self.reservation.free();
return Poll::Ready(Some(Ok(batch)));
}
Err(e) => return self.fail(e),
}
}
}
}
}
cx.waker().wake_by_ref();
Poll::Pending
}
}
impl RecordBatchStream for NativeInstantStream {
fn schema(&self) -> SchemaRef {
self.schema.clone()
}
}
src/promql/src/extension_plan/instant_manipulate/native_exec_tests.rs
-427
View File
@@ -1,427 +0,0 @@
// Copyright 2023 Greptime Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use std::collections::VecDeque;
use std::pin::Pin;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::task::{Context, Poll};
use common_error::ext::{BoxedError, ErrorExt, PlainError};
use common_error::status_code::StatusCode;
use datafusion::arrow::array::{Array, Float64Array, TimestampMillisecondArray};
use datafusion::arrow::datatypes::{DataType, Field, Schema, TimeUnit};
use datafusion::common::DFSchema;
use datafusion::error::DataFusionError;
use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryConsumer, MemoryPool};
use datafusion::logical_expr::{EmptyRelation, LogicalPlan};
use datatypes::arrow::buffer::ScalarBuffer;
use datatypes::series::raw::{RawSeriesChunk, RawTimestamps};
use datatypes::series::samples::{FieldSpec, Samples, SeriesFields, SeriesSchema};
use datatypes::series::{SeriesLabels, SeriesOrdering};
use futures::Stream;
use store_api::region_engine::{RawLabelSource, RawSeriesStream, SendableRawSeriesStream};
use super::*;
fn raw_schema() -> Arc<SeriesSchema> {
Arc::new(SeriesSchema {
fields: Arc::from([FieldSpec {
name: Arc::from("value"),
}]),
eval_grid: None,
})
}
fn output_schema() -> Arc<Schema> {
Arc::new(Schema::new(vec![
Field::new("value", DataType::Float64, true),
Field::new(
"time",
DataType::Timestamp(TimeUnit::Millisecond, None),
false,
),
]))
}
fn mapping_for(
start: i64,
end: i64,
step: i64,
lookback: i64,
) -> (Arc<NativeInstantMapping>, Arc<SeriesSchema>) {
let logical = LogicalPlan::EmptyRelation(EmptyRelation {
produce_one_row: false,
schema: Arc::new(DFSchema::empty()),
});
let node = InstantManipulate::new(
start,
end,
lookback,
step,
0,
"time".into(),
vec![],
Some("value".into()),
logical,
);
let series = raw_schema();
let mapping = NativeInstantMapping::try_new(&node, output_schema(), series.clone(), &[])
.unwrap()
.unwrap();
(Arc::new(mapping), series)
}
fn chunk(time: &[i64], values: &[f64], last: bool) -> RawSeriesChunk {
RawSeriesChunk::try_new(
&raw_schema(),
datatypes::arrow::buffer::Buffer::from(vec![1, 2, 3]),
Some(SeriesLabels::try_new(None, vec![]).unwrap()),
RawTimestamps::try_new(ScalarBuffer::from(time.to_vec())).unwrap(),
SeriesFields::one(Samples::from_float(Float64Array::from(values.to_vec()))),
last,
)
.unwrap()
}
struct RawMock {
schema: Arc<SeriesSchema>,
labels: RawLabelSource,
items: VecDeque<Poll<Option<std::result::Result<RawSeriesChunk, BoxedError>>>>,
polls: Arc<AtomicUsize>,
drops: Arc<AtomicUsize>,
}
impl Stream for RawMock {
type Item = std::result::Result<RawSeriesChunk, BoxedError>;
fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
self.polls.fetch_add(1, Ordering::SeqCst);
self.items.pop_front().unwrap_or(Poll::Ready(None))
}
}
impl Drop for RawMock {
fn drop(&mut self) {
self.drops.fetch_add(1, Ordering::SeqCst);
}
}
impl RawSeriesStream for RawMock {
fn schema(&self) -> &Arc<SeriesSchema> {
&self.schema
}
fn label_source(&self) -> &RawLabelSource {
&self.labels
}
fn ordering(&self) -> SeriesOrdering {
SeriesOrdering::Contiguous
}
}
fn raw(
items: impl IntoIterator<Item = Poll<Option<std::result::Result<RawSeriesChunk, BoxedError>>>>,
) -> (SendableRawSeriesStream, Arc<AtomicUsize>, Arc<AtomicUsize>) {
let polls = Arc::new(AtomicUsize::new(0));
let drops = Arc::new(AtomicUsize::new(0));
(
Box::pin(RawMock {
schema: raw_schema(),
labels: RawLabelSource::Explicit,
items: items.into_iter().collect(),
polls: polls.clone(),
drops: drops.clone(),
}),
polls,
drops,
)
}
fn stream(source: SendableRawSeriesStream, pool: &Arc<dyn MemoryPool>) -> NativeInstantStream {
stream_for(source, pool, 10, 30, 10, 10)
}
fn stream_for(
source: SendableRawSeriesStream,
pool: &Arc<dyn MemoryPool>,
start: i64,
end: i64,
step: i64,
lookback: i64,
) -> NativeInstantStream {
let (mapping, schema) = mapping_for(start, end, step, lookback);
NativeInstantStream {
source: Some(source),
mapping: mapping.clone(),
schema: mapping.schema(),
series_schema: schema,
pending: Vec::new(),
reservation: MemoryConsumer::new("NativeInstantStreamTest").register(pool),
}
}
fn poll(stream: &mut NativeInstantStream) -> Poll<Option<DfResult<RecordBatch>>> {
let waker = futures::task::noop_waker();
let mut cx = Context::from_waker(&waker);
Pin::new(stream).poll_next(&mut cx)
}
#[test]
fn emits_complete_series_before_polling_later_pending_input() {
let (input, polls, _) = raw([
Poll::Ready(Some(Ok(chunk(&[10], &[1.0], false)))),
Poll::Ready(Some(Ok(chunk(&[20], &[2.0], true)))),
Poll::Pending,
]);
let pool = Arc::new(GreedyMemoryPool::new(1_000_000)) as Arc<dyn MemoryPool>;
let mut stream = stream(input, &pool);
let Poll::Ready(Some(Ok(batch))) = poll(&mut stream) else {
panic!("completed series should emit")
};
assert_eq!(polls.load(Ordering::SeqCst), 2);
assert_eq!(
batch
.column(0)
.as_any()
.downcast_ref::<Float64Array>()
.unwrap()
.values(),
&[1.0, 2.0]
);
assert_eq!(
batch
.column(1)
.as_any()
.downcast_ref::<TimestampMillisecondArray>()
.unwrap()
.values(),
&[10, 20]
);
assert!(matches!(poll(&mut stream), Poll::Pending));
assert_eq!(polls.load(Ordering::SeqCst), 3);
}
#[test]
fn source_error_after_prefix_is_returned_once_and_clears_reservation_and_input() {
let failure = BoxedError::new(PlainError::new(
"mock source error".into(),
StatusCode::Internal,
));
let (input, polls, drops) = raw([
Poll::Ready(Some(Ok(chunk(&[10], &[1.0], true)))),
Poll::Ready(Some(Err(failure))),
]);
let pool = Arc::new(GreedyMemoryPool::new(1_000_000)) as Arc<dyn MemoryPool>;
let mut stream = stream(input, &pool);
assert!(matches!(poll(&mut stream), Poll::Ready(Some(Ok(_)))));
let Poll::Ready(Some(Err(error))) = poll(&mut stream) else {
panic!("source error must be returned once")
};
assert_eq!(
common_query::error::datafusion_status_code::<common_query::error::Error>(&error, None),
StatusCode::Internal
);
let DataFusionError::External(source) = error else {
panic!("source error must stay a typed DataFusionError::External payload")
};
let source = source
.downcast_ref::<BoxedError>()
.expect("source error must still be the original BoxedError");
assert_eq!(source.status_code(), StatusCode::Internal);
assert_eq!(source.to_string(), "mock source error");
assert!(matches!(poll(&mut stream), Poll::Ready(None)));
assert!(matches!(poll(&mut stream), Poll::Ready(None)));
assert_eq!(polls.load(Ordering::SeqCst), 2);
assert_eq!(drops.load(Ordering::SeqCst), 1);
assert_eq!(pool.reserved(), 0);
}
#[test]
fn unfinished_series_eof_is_error_once_and_drops_partial_chunks() {
let (input, _, drops) = raw([Poll::Ready(Some(Ok(chunk(&[10], &[1.0], false))))]);
let pool = Arc::new(GreedyMemoryPool::new(1_000_000)) as Arc<dyn MemoryPool>;
let mut stream = stream(input, &pool);
assert!(matches!(poll(&mut stream), Poll::Ready(Some(Err(_)))));
assert_eq!(stream.pending.len(), 0);
assert_eq!(drops.load(Ordering::SeqCst), 1);
assert_eq!(pool.reserved(), 0);
assert!(matches!(poll(&mut stream), Poll::Ready(None)));
}
#[test]
fn dropping_pending_stream_releases_source_chunks_and_memory() {
let array = Arc::new(Float64Array::from(vec![1.0]));
let samples = Samples::from_float(array.as_ref().clone());
drop(array);
let pending_chunk = RawSeriesChunk::try_new(
&raw_schema(),
datatypes::arrow::buffer::Buffer::from(vec![1]),
Some(SeriesLabels::try_new(None, vec![]).unwrap()),
RawTimestamps::try_new(ScalarBuffer::from(vec![10])).unwrap(),
SeriesFields::one(samples),
false,
)
.unwrap();
let (input, _, drops) = raw([Poll::Ready(Some(Ok(pending_chunk))), Poll::Pending]);
let pool = Arc::new(GreedyMemoryPool::new(1_000_000)) as Arc<dyn MemoryPool>;
let mut stream = stream(input, &pool);
assert!(matches!(poll(&mut stream), Poll::Pending));
assert!(pool.reserved() > 0);
drop(stream);
assert_eq!(drops.load(Ordering::SeqCst), 1);
assert_eq!(pool.reserved(), 0);
}
#[test]
fn reservation_failure_is_terminal_and_releases_input_and_partial_series() {
let (input, _, drops) = raw([Poll::Ready(Some(Ok(chunk(&[10], &[1.0], false))))]);
let pool = Arc::new(GreedyMemoryPool::new(1)) as Arc<dyn MemoryPool>;
let mut stream = stream(input, &pool);
assert!(matches!(
poll(&mut stream),
Poll::Ready(Some(Err(DataFusionError::ResourcesExhausted(_))))
));
assert_eq!(drops.load(Ordering::SeqCst), 1);
assert_eq!(stream.pending.len(), 0);
assert_eq!(stream.pending.capacity(), 0);
assert_eq!(pool.reserved(), 0);
assert!(matches!(poll(&mut stream), Poll::Ready(None)));
}
#[test]
fn source_error_drops_unfinished_series_without_emitting_it() {
let failure = BoxedError::new(PlainError::new(
"mock source error".into(),
StatusCode::Internal,
));
let (input, polls, drops) = raw([
Poll::Ready(Some(Ok(chunk(&[10], &[1.0], false)))),
Poll::Ready(Some(Err(failure))),
]);
let pool = Arc::new(GreedyMemoryPool::new(1_000_000)) as Arc<dyn MemoryPool>;
let mut stream = stream(input, &pool);
assert!(matches!(poll(&mut stream), Poll::Ready(Some(Err(_)))));
assert!(stream.pending.is_empty());
assert_eq!(stream.pending.capacity(), 0);
assert!(matches!(poll(&mut stream), Poll::Ready(None)));
assert_eq!(polls.load(Ordering::SeqCst), 2);
assert_eq!(drops.load(Ordering::SeqCst), 1);
assert_eq!(pool.reserved(), 0);
}
#[test]
fn descriptor_heavy_series_fails_reservation_before_accumulating_unbounded_chunks() {
let items = (0..512)
.map(|i| Poll::Ready(Some(Ok(chunk(&[i], &[i as f64], false)))))
.chain([Poll::Pending]);
let (input, polls, drops) = raw(items);
let pool = Arc::new(GreedyMemoryPool::new(32 * 1024)) as Arc<dyn MemoryPool>;
let mut stream = stream(input, &pool);
assert!(matches!(
poll(&mut stream),
Poll::Ready(Some(Err(DataFusionError::ResourcesExhausted(_))))
));
assert!(stream.pending.is_empty());
assert_eq!(stream.pending.capacity(), 0);
assert!(matches!(poll(&mut stream), Poll::Ready(None)));
assert!(polls.load(Ordering::SeqCst) < 512);
assert_eq!(drops.load(Ordering::SeqCst), 1);
assert_eq!(pool.reserved(), 0);
}
#[test]
fn dense_large_chunk_uses_segment_count_not_row_count_for_descriptors() {
let count = 100_000usize;
let timestamps = (0..count).map(|i| i as i64).collect::<Vec<_>>();
let values = Float64Array::from(vec![7.0; count]);
let dense_chunk = RawSeriesChunk::try_new(
&raw_schema(),
datatypes::arrow::buffer::Buffer::from(vec![1]),
Some(SeriesLabels::try_new(None, vec![]).unwrap()),
RawTimestamps::try_new(ScalarBuffer::from(timestamps)).unwrap(),
SeriesFields::one(Samples::from_float(values)),
true,
)
.unwrap();
let (input, _, _) = raw([Poll::Ready(Some(Ok(dense_chunk)))]);
// Input timestamp/field buffers and concatenation scratch fit; row-scaled Samples
// descriptors would exceed this pool despite there being only one dense segment.
let pool = Arc::new(GreedyMemoryPool::new(4 * 1024 * 1024)) as Arc<dyn MemoryPool>;
let mut stream = stream_for(input, &pool, 10, 11, 1, 0);
let Poll::Ready(Some(Ok(batch))) = poll(&mut stream) else {
panic!("dense chunk should fit its bounded pool")
};
assert_eq!(batch.num_rows(), 2);
assert_eq!(
batch
.column(0)
.as_any()
.downcast_ref::<Float64Array>()
.unwrap()
.values(),
&[7.0, 7.0]
);
assert_eq!(
batch
.column(1)
.as_any()
.downcast_ref::<TimestampMillisecondArray>()
.unwrap()
.values(),
&[10, 11]
);
assert_eq!(pool.reserved(), 0);
}
#[test]
fn output_grid_scratch_is_reserved_before_large_evaluation() {
let (input, polls, drops) = raw([Poll::Ready(Some(Ok(chunk(&[0], &[7.0], true))))]);
let pool = Arc::new(GreedyMemoryPool::new(128 * 1024)) as Arc<dyn MemoryPool>;
// One source point with a long lookback produces a large grid despite the tiny source.
let mut stream = stream_for(input, &pool, 0, 900_000, 1, 900_001);
assert!(matches!(
poll(&mut stream),
Poll::Ready(Some(Err(DataFusionError::ResourcesExhausted(_))))
));
assert_eq!(polls.load(Ordering::SeqCst), 1);
assert_eq!(drops.load(Ordering::SeqCst), 1);
assert!(stream.pending.is_empty());
assert_eq!(pool.reserved(), 0);
assert!(matches!(poll(&mut stream), Poll::Ready(None)));
}
#[test]
fn larger_pool_completes_one_point_many_evaluation_grid() {
let (input, _, _) = raw([Poll::Ready(Some(Ok(chunk(&[0], &[7.0], true))))]);
let pool = Arc::new(GreedyMemoryPool::new(32 * 1024 * 1024)) as Arc<dyn MemoryPool>;
let mut stream = stream_for(input, &pool, 0, 4, 1, 5);
let Poll::Ready(Some(Ok(batch))) = poll(&mut stream) else {
panic!("sufficient reservation should evaluate")
};
assert_eq!(
batch
.column(0)
.as_any()
.downcast_ref::<Float64Array>()
.unwrap()
.values(),
&[7.0; 5]
);
assert_eq!(
batch
.column(1)
.as_any()
.downcast_ref::<TimestampMillisecondArray>()
.unwrap()
.values(),
&[0, 1, 2, 3, 4]
);
assert_eq!(pool.reserved(), 0);
}
src/promql/src/extension_plan/instant_manipulate/native_tests.rs
-597
View File
@@ -1,597 +0,0 @@
// Copyright 2023 Greptime Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use std::collections::HashMap;
use std::sync::Arc;
use common_query::prometheus::PROMETHEUS_STALE_NAN_BITS;
use datafusion::arrow::array::{Array, Float64Array, StringArray, TimestampMillisecondArray};
use datafusion::arrow::datatypes::{DataType, Field, Schema, TimeUnit};
use datafusion::common::DFSchema;
use datafusion::logical_expr::{EmptyRelation, LogicalPlan};
use datatypes::arrow::buffer::ScalarBuffer;
use datatypes::series::raw::RawTimestamps;
use datatypes::series::samples::{FieldSpec, Samples, SeriesFields, SeriesSchema};
use datatypes::series::{Label, SeriesLabels};
use super::*;
fn node() -> InstantManipulate {
let input = LogicalPlan::EmptyRelation(EmptyRelation {
produce_one_row: false,
schema: Arc::new(DFSchema::empty()),
});
InstantManipulate::new(
10,
20,
10,
10,
0,
"time".into(),
vec![],
Some("value".into()),
input,
)
}
fn schema(tag_type: DataType) -> SchemaRef {
let mut metadata = HashMap::new();
metadata.insert("schema".to_string(), "kept".to_string());
Arc::new(Schema::new_with_metadata(
vec![
Field::new("value", DataType::Float64, true)
.with_metadata([("field-key".to_string(), "field-value".to_string())].into()),
Field::new(
"time",
DataType::Timestamp(TimeUnit::Millisecond, None),
false,
),
Field::new("zone", tag_type, true),
],
metadata,
))
}
fn series_schema() -> Arc<SeriesSchema> {
Arc::new(SeriesSchema {
fields: Arc::from([FieldSpec {
name: Arc::from("value"),
}]),
eval_grid: None,
})
}
fn metric(zone: Option<&str>) -> SeriesLabels {
let labels = zone
.map(|value| {
vec![Label {
name: Arc::from("zone"),
value: Arc::from(value),
}]
})
.unwrap_or_default();
SeriesLabels::try_new(None, labels).unwrap()
}
/// The ordered, deliberately non-lexical field names of the multifloat regression.
const MULTIFLOAT_FIELDS: [&str; 3] = ["value", "a_second", "m_third"];
/// The strict millisecond timestamps of every multifloat case.
const MULTIFLOAT_TIMESTAMPS: [i64; 5] = [1000, 2000, 3000, 4000, 5000];
/// The kept source rows of each configured staleness field: the stale row of the configured
/// field is suppressed, so the first field drops row 3 (4000) and the second field drops row 2
/// (3000) without ever retrying an older sample.
const MULTIFLOAT_KEPT: [[usize; 4]; 2] = [[0, 1, 2, 4], [0, 1, 3, 4]];
/// An ordinary, non-stale NaN payload that must survive the mapping as exact bits.
const ORDINARY_NAN_BITS: u64 = 0x7ff8_0000_0000_0001;
/// The output schema of the regression: the declared field order, then the time index and the
/// tag, with metadata that every case must retain.
///
/// Every declared field is nullable because every fixture lane keeps a NULL row in each of them:
/// the first row of the second and third fields and the second row of the first field.
fn multifloat_schema(field_names: &[&str], tag_type: DataType) -> SchemaRef {
let mut metadata = HashMap::new();
metadata.insert("schema".to_string(), "kept".to_string());
let mut fields = field_names
.iter()
.enumerate()
.map(|(index, name)| {
Field::new(*name, DataType::Float64, true)
.with_metadata([("field-key".to_string(), format!("field-value-{index}"))].into())
})
.collect::<Vec<_>>();
fields.push(Field::new(
"time",
DataType::Timestamp(TimeUnit::Millisecond, None),
false,
));
fields.push(Field::new("zone", tag_type, true));
Arc::new(Schema::new_with_metadata(fields, metadata))
}
/// The series schema of the same declared field order.
fn multifloat_series_schema(field_names: &[&str]) -> Arc<SeriesSchema> {
Arc::new(SeriesSchema {
fields: field_names
.iter()
.map(|name| FieldSpec {
name: Arc::from(*name),
})
.collect::<Vec<_>>()
.into(),
eval_grid: None,
})
}
/// The five source rows of each multifloat field.
///
/// Row `i` is the strict millisecond timestamp `MULTIFLOAT_TIMESTAMPS[i]`. The first field is
/// stale at 4000 and carries an ordinary NaN payload at 5000; the unrelated second field is
/// stale at 3000; the fields carry NULLs and a signed zero as well.
fn multifloat_source(field_names: &[&str]) -> Vec<Float64Array> {
let stale = f64::from_bits(PROMETHEUS_STALE_NAN_BITS);
let nan = f64::from_bits(ORDINARY_NAN_BITS);
let rows: [&[Option<f64>]; 3] = [
&[Some(1.0), None, Some(-0.0), Some(stale), Some(nan)],
&[None, Some(-0.0), Some(stale), Some(2.0), Some(4.0)],
&[None, Some(3.0), Some(nan), Some(3.0), Some(5.0)],
];
field_names
.iter()
.enumerate()
.map(|(index, _)| Float64Array::from(rows[index].to_vec()))
.collect()
}
/// The chunked samples of the multifloat input: the first field arrives in three chunks, two of
/// which are slices with non-zero offsets, and every other field arrives as one dense array.
fn multifloat_fields(source: &[Float64Array]) -> SeriesFields {
let fields = source
.iter()
.enumerate()
.map(|(index, values)| {
if index > 0 {
return Samples::from_float(values.clone());
}
Samples::try_from_chunks(vec![
Samples::from_float(values.slice(0, 1)),
Samples::from_float(values.slice(1, 3)),
Samples::from_float(values.slice(4, 1)),
])
.unwrap()
})
.collect::<Vec<_>>();
SeriesFields::try_from_vec(fields).unwrap()
}
/// The bounded instant declaration of the regression: one output point per input millisecond
/// with one lookback step, whose staleness field is exactly one declared field.
fn multifloat_node(field_column: &str) -> InstantManipulate {
let mut query = node();
query.start = MULTIFLOAT_TIMESTAMPS[0];
query.end = MULTIFLOAT_TIMESTAMPS[4];
query.lookback_delta = 1000;
query.interval = 1000;
query.field_column = Some(field_column.to_string());
query
}
#[test]
fn evaluates_sliced_chunked_samples_and_preserves_original_schema() {
let output_schema = schema(DataType::Dictionary(
Box::new(DataType::UInt32),
Box::new(DataType::Utf8),
));
let mapping = NativeInstantMapping::try_new(
&node(),
output_schema.clone(),
series_schema(),
&["zone".into()],
)
.unwrap()
.unwrap();
let a = Float64Array::from(vec![Some(99.0), Some(1.0), None]);
let b = Float64Array::from(vec![Some(2.0), Some(3.0)]);
let samples = Samples::try_from_chunks(vec![
Samples::from_float(
a.slice(1, 2)
.as_any()
.downcast_ref::<Float64Array>()
.unwrap()
.clone(),
),
Samples::from_float(b),
])
.unwrap();
let fields = SeriesFields::one(samples);
let time = RawTimestamps::try_new(ScalarBuffer::from(vec![0, 10, 20, 30])).unwrap();
let batch = mapping.evaluate(&metric(Some("")), &time, &fields).unwrap();
assert_eq!(batch.num_rows(), 2);
assert_eq!(batch.schema(), output_schema);
let output_values = batch
.column(0)
.as_any()
.downcast_ref::<Float64Array>()
.unwrap();
assert!(output_values.is_null(0));
assert_eq!(output_values.value(1), 2.0);
assert_eq!(
batch
.column(1)
.as_any()
.downcast_ref::<TimestampMillisecondArray>()
.unwrap()
.values(),
&[10, 20]
);
let tag_values = datafusion::arrow::compute::cast(batch.column(2), &DataType::Utf8).unwrap();
let tag_values = tag_values.as_any().downcast_ref::<StringArray>().unwrap();
assert_eq!(tag_values.value(0), "");
assert_eq!(tag_values.value(1), "");
}
#[test]
fn selected_repeated_indices_across_sliced_chunks_preserve_missing_and_nan() {
let mut query = node();
query.end = 30;
query.lookback_delta = 20;
let mapping = NativeInstantMapping::try_new(
&query,
schema(DataType::Utf8),
series_schema(),
&["zone".into()],
)
.unwrap()
.unwrap();
let first = Float64Array::from(vec![Some(99.0), None]).slice(1, 1);
let second = Float64Array::from(vec![Some(f64::NAN), Some(88.0)]).slice(0, 1);
let samples = Samples::try_from_chunks(vec![
Samples::from_float(
first
.as_any()
.downcast_ref::<Float64Array>()
.unwrap()
.clone(),
),
Samples::from_float(
second
.as_any()
.downcast_ref::<Float64Array>()
.unwrap()
.clone(),
),
])
.unwrap();
let fields = SeriesFields::one(samples);
let time = RawTimestamps::try_new(ScalarBuffer::from(vec![10, 30])).unwrap();
let batch = mapping.evaluate(&metric(None), &time, &fields).unwrap();
let values = batch
.column(0)
.as_any()
.downcast_ref::<Float64Array>()
.unwrap();
assert_eq!(batch.num_rows(), 3);
assert!(values.is_null(0));
assert!(values.is_null(1));
assert!(!values.is_null(2));
assert!(values.value(2).is_nan());
}
#[test]
fn maps_all_supported_string_types_and_distinguishes_missing_from_empty() {
let tag_types = [
DataType::Utf8,
DataType::LargeUtf8,
DataType::Utf8View,
DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8)),
];
let fields = SeriesFields::one(Samples::from_float(Float64Array::from(vec![
Some(1.0),
Some(2.0),
])));
let time = RawTimestamps::try_new(ScalarBuffer::from(vec![10, 20])).unwrap();
for tag_type in tag_types {
let output_schema = schema(tag_type);
let mapping = NativeInstantMapping::try_new(
&node(),
output_schema.clone(),
series_schema(),
&["zone".into()],
)
.unwrap()
.unwrap();
for (label, expected) in [(Some(""), Some("")), (None, None)] {
let batch = mapping.evaluate(&metric(label), &time, &fields).unwrap();
assert_eq!(batch.schema(), output_schema);
let values =
datafusion::arrow::compute::cast(batch.column(2), &DataType::Utf8).unwrap();
let values = values.as_any().downcast_ref::<StringArray>().unwrap();
match expected {
Some(value) => {
assert_eq!(values.value(0), value);
assert_eq!(values.value(1), value);
assert_eq!(values.null_count(), 0);
}
None => assert_eq!(values.null_count(), 2),
}
}
}
}
#[test]
fn selected_field_staleness_and_no_field_semantics_match_existing_selector_contract() {
let stale = f64::from_bits(PROMETHEUS_STALE_NAN_BITS);
let times = RawTimestamps::try_new(ScalarBuffer::from(vec![10, 20])).unwrap();
let values = SeriesFields::one(Samples::from_float(Float64Array::from(vec![
Some(1.0),
Some(stale),
])));
let mapping = NativeInstantMapping::try_new(
&node(),
schema(DataType::Utf8),
series_schema(),
&["zone".into()],
)
.unwrap()
.unwrap();
let batch = mapping.evaluate(&metric(None), &times, &values).unwrap();
assert_eq!(batch.num_rows(), 1);
assert_eq!(
batch
.column(1)
.as_any()
.downcast_ref::<TimestampMillisecondArray>()
.unwrap()
.value(0),
10
);
let mut no_staleness_node = node();
no_staleness_node.field_column = None;
let no_staleness_mapping = NativeInstantMapping::try_new(
&no_staleness_node,
schema(DataType::Utf8),
series_schema(),
&["zone".into()],
)
.unwrap()
.unwrap();
let batch = no_staleness_mapping
.evaluate(&metric(None), &times, &values)
.unwrap();
let output = batch
.column(0)
.as_any()
.downcast_ref::<Float64Array>()
.unwrap();
assert_eq!(output.len(), 2);
assert_eq!(output.value(1).to_bits(), PROMETHEUS_STALE_NAN_BITS);
}
#[test]
fn rejects_unmapped_missing_duplicate_and_non_millisecond_schema_columns() {
let valid = schema(DataType::Utf8);
let mut cases = Vec::new();
cases.push(Arc::new(Schema::new(vec![
valid.field(0).clone(),
valid.field(1).clone(),
valid.field(2).clone(),
Field::new("helper", DataType::UInt64, true),
])));
cases.push(Arc::new(Schema::new(vec![
valid.field(0).clone(),
valid.field(1).clone(),
valid.field(2).clone(),
Field::new("__tsid", DataType::UInt64, true),
])));
cases.push(Arc::new(Schema::new(vec![
valid.field(0).clone(),
valid.field(1).clone(),
])));
cases.push(Arc::new(Schema::new(vec![
valid.field(1).clone(),
valid.field(2).clone(),
])));
cases.push(Arc::new(Schema::new(vec![
valid.field(0).clone(),
valid.field(1).clone(),
valid.field(2).clone(),
valid.field(0).clone(),
])));
cases.push(Arc::new(Schema::new(vec![
valid.field(0).clone(),
Field::new(
"time",
DataType::Timestamp(TimeUnit::Microsecond, None),
false,
),
valid.field(2).clone(),
])));
cases.push(Arc::new(Schema::new(vec![
Field::new("value", DataType::Int64, true),
valid.field(1).clone(),
valid.field(2).clone(),
])));
for output_schema in cases {
assert!(
NativeInstantMapping::try_new(
&node(),
output_schema,
series_schema(),
&["zone".into()],
)
.unwrap()
.is_none()
);
}
}
#[test]
fn absent_tag_is_null_and_nonzero_offset_is_rejected() {
let output_schema = schema(DataType::Utf8);
let mapping =
NativeInstantMapping::try_new(&node(), output_schema, series_schema(), &["zone".into()])
.unwrap()
.unwrap();
let fields = SeriesFields::one(Samples::from_float(Float64Array::from(vec![
Some(1.0),
Some(2.0),
])));
let time = RawTimestamps::try_new(ScalarBuffer::from(vec![10, 20])).unwrap();
let batch = mapping.evaluate(&metric(None), &time, &fields).unwrap();
assert_eq!(batch.column(2).null_count(), 2);
let mut unsupported = node();
unsupported.offset = 1;
assert!(
NativeInstantMapping::try_new(
&unsupported,
schema(DataType::Utf8),
series_schema(),
&["zone".into()],
)
.unwrap()
.is_none()
);
}
#[test]
fn multifloat_native_mapping_keeps_ordered_fields_and_checks_only_the_configured_field() {
let time = RawTimestamps::try_new(ScalarBuffer::from(MULTIFLOAT_TIMESTAMPS.to_vec())).unwrap();
for field_names in [&MULTIFLOAT_FIELDS[..2], &MULTIFLOAT_FIELDS[..3]] {
for (configured, kept) in MULTIFLOAT_KEPT.iter().enumerate() {
let output_schema = multifloat_schema(field_names, DataType::Utf8);
let source = multifloat_source(field_names);
let fields = multifloat_fields(&source);
let mapping = NativeInstantMapping::try_new(
&multifloat_node(field_names[configured]),
output_schema.clone(),
multifloat_series_schema(field_names),
&["zone".into()],
)
.unwrap()
.expect("the complete ordered Float64 fields are mapped instead of declined");
let batch = mapping
.evaluate(&metric(Some("multifloat")), &time, &fields)
.unwrap();
// The complete output schema, including its metadata and nullability, is retained.
assert_eq!(batch.schema(), output_schema);
assert_eq!(batch.num_rows(), kept.len());
let kept_timestamps = batch
.column(field_names.len())
.as_any()
.downcast_ref::<TimestampMillisecondArray>()
.unwrap()
.values()
.to_vec();
assert_eq!(
kept_timestamps,
kept.iter()
.map(|row| MULTIFLOAT_TIMESTAMPS[*row])
.collect::<Vec<_>>()
);
let suppressed = if configured == 0 { 3 } else { 2 };
let unrelated = if configured == 0 { 2 } else { 3 };
assert!(
!kept_timestamps.contains(&MULTIFLOAT_TIMESTAMPS[suppressed]),
"the stale row of the configured field is suppressed instead of retried"
);
assert!(
kept_timestamps.contains(&MULTIFLOAT_TIMESTAMPS[unrelated]),
"the stale marker of another field does not suppress its row"
);
// Every output field keeps the nullness and the exact bits of its own source row.
for (index, values) in source.iter().enumerate() {
let output = batch
.column(index)
.as_any()
.downcast_ref::<Float64Array>()
.unwrap();
for (row, source_row) in kept.iter().enumerate() {
if values.is_null(*source_row) {
assert!(output.is_null(row));
} else {
assert!(!output.is_null(row));
assert_eq!(
output.value(row).to_bits(),
values.value(*source_row).to_bits(),
"field {index} row {row} must keep its exact source bits"
);
}
}
}
// The unchecked stale marker travels through as exact bits, so the row of another
// field is kept with the marker intact instead of being compared or re-encoded.
let (unchecked_field, unchecked_row) = if configured == 0 { (1, 2) } else { (0, 3) };
let unchecked_output_row = kept
.iter()
.position(|row| *row == unchecked_row)
.expect("the unchecked stale row is kept");
let marker = batch
.column(unchecked_field)
.as_any()
.downcast_ref::<Float64Array>()
.unwrap()
.value(unchecked_output_row);
assert_eq!(marker.to_bits(), PROMETHEUS_STALE_NAN_BITS);
// The ordinary NaN payload and the signed zero of the first two fields survive as
// exact bits as well.
let nan_row = kept
.iter()
.position(|row| *row == 4)
.expect("the ordinary NaN row is kept");
let nan = batch
.column(0)
.as_any()
.downcast_ref::<Float64Array>()
.unwrap()
.value(nan_row);
assert_eq!(nan.to_bits(), ORDINARY_NAN_BITS);
let signed_zero_row = kept
.iter()
.position(|row| *row == 1)
.expect("the signed zero row is kept");
let signed_zero = batch
.column(1)
.as_any()
.downcast_ref::<Float64Array>()
.unwrap()
.value(signed_zero_row);
assert_eq!(signed_zero.to_bits(), (-0.0_f64).to_bits());
// The series key of every kept row is the requested label value.
let tags = datafusion::arrow::compute::cast(
batch.column(field_names.len() + 1),
&DataType::Utf8,
)
.unwrap();
let tags = tags.as_any().downcast_ref::<StringArray>().unwrap();
assert_eq!(tags.null_count(), 0);
assert!(tags.iter().all(|value| value == Some("multifloat")));
}
}
}
src/promql/src/instant_eval.rs
+3 -4
View File
@@ -180,10 +180,9 @@ pub fn select_instants(
///
/// `timestamp(x @ t)` selects at `t` for every point, so its selection repeats one row. This
/// kernel reports the sample the caller selected and does not decide the selector's time
/// modifiers: the native Instant caller passes its offset into the selection, so
/// `timestamp(x @ t offset d)` reports the sample at `t - d`, matching GreptimeDB's existing
/// planner rather than Prometheus' special case for #8433, which ignores the offset here. `info()`
/// applies it too.
/// modifiers: the native Instant caller selects the sample under the intersection of the ordinary
/// anchored selector window and the timestamp function's unshifted window. `info()` retains ordinary
/// selector offset behavior.
pub fn timestamp_of_selection(
schema: &Arc<SeriesSchema>,
metric: &SeriesLabels,
src/promql/src/instant_eval/tests.rs
+1
View File
@@ -535,6 +535,7 @@ fn anchored_instant_uses_lookback_offset_and_staleness_rules() {
select_instant_at(&raw_axis, &fields, 499, 0, 0).unwrap(),
None
);
assert!(select_instant_at(&raw_axis, &fields, 500, -1, 0).is_err());
let stale = Samples::from_float(Float64Array::from(vec![
Some(7.0),
Some(f64::from_bits(
src/promql/src/native/aggregate.rs
+3 -5
View File
@@ -160,11 +160,9 @@ impl AggregateRuntime {
// chunk of another field count or another grid would be folded into the wrong places.
validate_float_series(&self.schema, &chunk.series)
.map_err(|error| folded_chunk(&self.describe(), error))?;
let aggregate = self.describe();
let accumulator = self
.accumulator
.as_mut()
.ok_or_else(|| no_fold(&aggregate))?;
let Some(accumulator) = self.accumulator.as_mut() else {
return Err(no_fold(&self.describe()));
};
// Keep the chunk covered while the fold decides whether it needs a new group or block.
accumulator.update_budgeted(&chunk.series, |fold_bytes| {
buffer.resize(pending.saturating_add(fold_bytes))
src/promql/src/native/execute.rs
+203 -60
View File
@@ -434,11 +434,13 @@ pub(super) struct NodeBuffer {
impl NodeBuffer {
/// Registers a reservation for the node called `name` with `pool`.
///
/// The name is what the operator is — `native.input`, `native.window` — rather than where it
/// sits in the tree: a consumer is named once, when the node owning it is built, so a tree
/// executed twice registers the same consumers again instead of a second set of them.
fn new(name: &'static str, metrics: Arc<NodeMetrics>, pool: &Arc<dyn MemoryPool>) -> Self {
let consumer = MemoryConsumer::new(name);
/// The name identifies both the operator and its tree location, such as
/// `native.window at root.left`, so a failure names the node that hit it rather than only its
/// kind. It labels the reservation, it does not key it: the pool accounts a consumer by its own
/// id, so two nodes sharing a name stay separate reservations. This name is what the error
/// context and the pool's top-consumer report print.
fn new(name: impl Into<String>, metrics: Arc<NodeMetrics>, pool: &Arc<dyn MemoryPool>) -> Self {
let consumer = MemoryConsumer::new(name.into());
Self {
metrics,
reservation: consumer.register(pool),
@@ -965,12 +967,36 @@ impl InstantRuntime {
// The grid's own points are what the selection decides, so the range and the step it is
// asked for are the grid's: `start..=start + (len - 1) * step`, step by step.
if let Some(at_ms) = self.selection.at_ms {
let timestamp_window = if self.selection.timestamp && self.selection.lookback_ms >= 0 {
let anchor = at_ms as i128;
let shifted = anchor - self.selection.offset_ms as i128;
let lower = anchor.max(shifted) - self.selection.lookback_ms as i128;
let upper = anchor.min(shifted);
if lower > upper || (lower == upper && self.selection.lookback_ms != 0) {
None
} else {
let effective_lookback = i64::try_from(upper - lower).ok();
effective_lookback.and_then(|lookback| {
i64::try_from(upper).ok().map(|upper| (upper, lookback))
})
}
} else {
Some((at_ms, self.selection.lookback_ms))
};
let Some((selection_at, selection_lookback)) = timestamp_window else {
let metric = drop_metric_name(&metric)?;
return no_row_chunk(&self.schema, metric);
};
let Some(index) = select_instant_at(
&time,
&fields,
at_ms,
self.selection.lookback_ms,
self.selection.offset_ms,
selection_at,
selection_lookback,
if self.selection.timestamp {
0
} else {
self.selection.offset_ms
},
)?
else {
let metric = if self.selection.timestamp {
@@ -2206,7 +2232,15 @@ pub struct NativeExecutionState {
}
impl NodeState {
fn new(name: &'static str, metrics: Arc<NodeMetrics>, pool: &Arc<dyn MemoryPool>) -> Self {
fn context(&self, error: DataFusionError) -> DataFusionError {
DataFusionError::Context(
self.buffer.reservation.consumer().name().to_string(),
Box::new(error),
)
}
fn new(name: impl Into<String>, metrics: Arc<NodeMetrics>, pool: &Arc<dyn MemoryPool>) -> Self {
let name = name.into();
Self {
queued: VecDeque::new(),
queued_tables: 0,
@@ -2221,6 +2255,7 @@ impl RuntimeNode {
node: &NativeNode,
pool: &Arc<dyn MemoryPool>,
inputs: &mut Vec<InputSlot>,
location: String,
) -> Result<(Self, bool)> {
let metrics = Arc::new(NodeMetrics::default());
let (op, shared) = match &node.expr {
@@ -2239,13 +2274,21 @@ impl RuntimeNode {
source_index: bound.source_index,
schema,
open_source_series: false,
state: NodeState::new("native.input", Arc::clone(&metrics), pool),
state: NodeState::new(
format!(
"native.input at {location} (source_index={})",
bound.source_index
),
Arc::clone(&metrics),
pool,
),
});
(RuntimeOp::Input(id), false)
}
NativeExpr::Scalar(_) => (RuntimeOp::Scalar, false),
NativeExpr::Instant { input, op } => {
let (child, shared) = Self::build(input, pool, inputs)?;
let (child, shared) =
Self::build(input, pool, inputs, format!("{location}.input"))?;
let input_schema = input.output.samples().cloned().ok_or_else(|| {
DataFusionError::Execution(
"instant selection consumes non-sample input".to_string(),
@@ -2256,8 +2299,11 @@ impl RuntimeNode {
input: Box::new(child),
instant: Some(InstantRuntime::try_new(&input_schema, op)?),
state: {
let mut state =
NodeState::new("native.instant", Arc::clone(&metrics), pool);
let mut state = NodeState::new(
format!("native.instant at {location}"),
Arc::clone(&metrics),
pool,
);
if op.timestamp {
state.emitted = Some(Emitted::default());
}
@@ -2268,11 +2314,15 @@ impl RuntimeNode {
)
}
NativeExpr::Merge { input } => {
let (child, _) = Self::build(input, pool, inputs)?;
let (child, _) = Self::build(input, pool, inputs, format!("{location}.input"))?;
let schema = node.output.samples().cloned().ok_or_else(|| {
DataFusionError::Execution("merge does not produce samples".to_string())
})?;
let state = NodeState::new("native.merge", Arc::clone(&metrics), pool);
let state = NodeState::new(
format!("native.merge at {location}"),
Arc::clone(&metrics),
pool,
);
(
RuntimeOp::Merge {
input: Box::new(child),
@@ -2283,13 +2333,18 @@ impl RuntimeNode {
)
}
NativeExpr::Window { input, op } => {
let (child, child_shared) = Self::build(input, pool, inputs)?;
let (child, child_shared) =
Self::build(input, pool, inputs, format!("{location}.input"))?;
let input_schema = input.output.samples().cloned().ok_or_else(|| {
DataFusionError::Execution("window consumes non-sample input".to_string())
})?;
let merge = child_shared.then(|| MergeInput::new(Arc::clone(&input_schema)));
let shared = op.function.drops_metric_name() || !child_shared;
let mut state = NodeState::new("native.window", Arc::clone(&metrics), pool);
let mut state = NodeState::new(
format!("native.window at {location}"),
Arc::clone(&metrics),
pool,
);
state.emitted = Some(Emitted::default());
(
RuntimeOp::Window {
@@ -2307,11 +2362,15 @@ impl RuntimeNode {
)
}
NativeExpr::Map { input, op } => {
let (child, _) = Self::build(input, pool, inputs)?;
let (child, _) = Self::build(input, pool, inputs, format!("{location}.input"))?;
let schema = node.output.samples().cloned().ok_or_else(|| {
DataFusionError::Execution("map does not produce samples".to_string())
})?;
let mut state = NodeState::new("native.map", Arc::clone(&metrics), pool);
let mut state = NodeState::new(
format!("native.map at {location}"),
Arc::clone(&metrics),
pool,
);
state.emitted = Some(Emitted::default());
(
RuntimeOp::Map {
@@ -2326,11 +2385,15 @@ impl RuntimeNode {
)
}
NativeExpr::Timestamp { input } => {
let (child, _) = Self::build(input, pool, inputs)?;
let (child, _) = Self::build(input, pool, inputs, format!("{location}.input"))?;
let schema = node.output.samples().cloned().ok_or_else(|| {
DataFusionError::Execution("timestamp does not produce samples".to_string())
})?;
let mut state = NodeState::new("native.timestamp", Arc::clone(&metrics), pool);
let mut state = NodeState::new(
format!("native.timestamp at {location}"),
Arc::clone(&metrics),
pool,
);
state.emitted = Some(Emitted::default());
(
RuntimeOp::Timestamp {
@@ -2362,15 +2425,21 @@ impl RuntimeNode {
}
_ => (None, Some(broadcast_runtime(left, right, matching)?)),
};
let (left_node, left_shared) = Self::build(left, pool, inputs)?;
let (right_node, right_shared) = Self::build(right, pool, inputs)?;
let (left_node, left_shared) =
Self::build(left, pool, inputs, format!("{location}.left"))?;
let (right_node, right_shared) =
Self::build(right, pool, inputs, format!("{location}.right"))?;
match (binary, broadcast) {
(Some(binary), None) => (
RuntimeOp::Binary {
left: Box::new(left_node),
right: Box::new(right_node),
binary: Some(binary),
state: NodeState::new("native.binary", Arc::clone(&metrics), pool),
state: NodeState::new(
format!("native.binary at {location}"),
Arc::clone(&metrics),
pool,
),
},
false,
),
@@ -2381,7 +2450,11 @@ impl RuntimeNode {
} else {
right_shared
};
let mut state = NodeState::new("native.binary", Arc::clone(&metrics), pool);
let mut state = NodeState::new(
format!("native.binary at {location}"),
Arc::clone(&metrics),
pool,
);
if broadcast.drops_labels() {
state.emitted = Some(Emitted::default());
}
@@ -2414,15 +2487,20 @@ impl RuntimeNode {
PairKernel::Set(op.clone()),
[Arc::clone(left_schema), Arc::clone(right_schema)],
);
let (left_node, _) = Self::build(left, pool, inputs)?;
let (right_node, _) = Self::build(right, pool, inputs)?;
let (left_node, _) = Self::build(left, pool, inputs, format!("{location}.left"))?;
let (right_node, _) =
Self::build(right, pool, inputs, format!("{location}.right"))?;
// The kernel files one series per label set, so its output shares none.
(
RuntimeOp::Binary {
left: Box::new(left_node),
right: Box::new(right_node),
binary: Some(set),
state: NodeState::new("native.set", Arc::clone(&metrics), pool),
state: NodeState::new(
format!("native.set at {location}"),
Arc::clone(&metrics),
pool,
),
},
false,
)
@@ -2436,13 +2514,17 @@ impl RuntimeNode {
.ok_or_else(|| {
DataFusionError::Execution("aggregate does not produce samples".to_string())
})?;
let mut state = NodeState::new("native.aggregate", Arc::clone(&metrics), pool);
let mut state = NodeState::new(
format!("native.aggregate at {location}"),
Arc::clone(&metrics),
pool,
);
let aggregate = Some(AggregateRuntime::new(
op.clone(),
schema,
&mut state.buffer,
)?);
let (child, _) = Self::build(input, pool, inputs)?;
let (child, _) = Self::build(input, pool, inputs, format!("{location}.input"))?;
(
RuntimeOp::Aggregate {
input: Box::new(child),
@@ -2471,7 +2553,8 @@ impl RuntimeNode {
));
}
let mut vector = VectorRuntime::new(op.clone(), schema);
let (child, child_shared) = Self::build(input, pool, inputs)?;
let (child, child_shared) =
Self::build(input, pool, inputs, format!("{location}.input"))?;
let limitk = matches!(op, BoundVectorOp::LimitK { .. });
if limitk && child_shared {
vector.merge_input = true;
@@ -2487,7 +2570,11 @@ impl RuntimeNode {
RuntimeOp::Vector {
input: Box::new(child),
vector: Some(vector),
state: NodeState::new("native.vector", Arc::clone(&metrics), pool),
state: NodeState::new(
format!("native.vector at {location}"),
Arc::clone(&metrics),
pool,
),
},
child_shared && !whole,
)
@@ -2592,7 +2679,10 @@ impl RuntimeNode {
let Some(chunk) = input.pull(inputs, inputs_finished)? else {
return Ok(None);
};
if let Some(output) = state.instant(selection, chunk)? {
if let Some(output) = state
.instant(selection, chunk)
.map_err(|error| state.context(error))?
{
return Ok(Some(NativeChunk::Evaluated(output)));
}
}
@@ -2629,7 +2719,7 @@ impl RuntimeNode {
// A chunk this window computed while flushing its collected input leaves first: it
// was produced before the window was drained, and the pull that hands it up is the
// one that empties the queue.
if let Some(chunk) = state.take_output()? {
if let Some(chunk) = state.take_output().map_err(|error| state.context(error))? {
return Ok(Some(chunk));
}
let Some(evaluator) = window.as_mut() else {
@@ -2641,7 +2731,10 @@ impl RuntimeNode {
let Some(chunk) = input.pull(inputs, inputs_finished)? else {
return Ok(None);
};
if let Some(output) = state.window(evaluator, merge, chunk)? {
if let Some(output) = state
.window(evaluator, merge, chunk)
.map_err(|error| state.context(error))?
{
return Ok(Some(NativeChunk::Evaluated(output)));
}
}
@@ -2661,7 +2754,10 @@ impl RuntimeNode {
"the map node consumes an evaluated chunk, not raw samples".to_string(),
));
};
if let Some(output) = state.map(operation, chunk)? {
if let Some(output) = state
.map(operation, chunk)
.map_err(|error| state.context(error))?
{
return Ok(Some(NativeChunk::Evaluated(output)));
}
}
@@ -2683,7 +2779,10 @@ impl RuntimeNode {
"timestamp consumes evaluated samples, not raw samples".into(),
));
};
if let Some(output) = state.timestamp(schema, chunk)? {
if let Some(output) = state
.timestamp(schema, chunk)
.map_err(|error| state.context(error))?
{
return Ok(Some(NativeChunk::Evaluated(output)));
}
}
@@ -2701,10 +2800,14 @@ impl RuntimeNode {
// binary node only retains the chunks, because a pair is not decided until both
// sides are whole, and after it ends this is the drain that empties the two.
while let Some(chunk) = left.pull(inputs, inputs_finished)? {
state.retain_binary(operation, 0, chunk)?;
state
.retain_binary(operation, 0, chunk)
.map_err(|error| state.context(error))?;
}
while let Some(chunk) = right.pull(inputs, inputs_finished)? {
state.retain_binary(operation, 1, chunk)?;
state
.retain_binary(operation, 1, chunk)
.map_err(|error| state.context(error))?;
}
// Nothing has settled for a point that has an operand still arriving, and an
// operand below a collecting window has not arrived until that window is flushed:
@@ -2712,7 +2815,9 @@ impl RuntimeNode {
if !inputs_finished || !left.settled(inputs) || !right.settled(inputs) {
return Ok(None);
}
state.compute_binary(operation)?;
state
.compute_binary(operation)
.map_err(|error| state.context(error))?;
// A result holds one value per field, and a field may hold none where another
// does, so its values are counted rather than read as its positions.
state.pop_queued(evaluated_values)
@@ -2740,7 +2845,10 @@ impl RuntimeNode {
let Some(chunk) = vector.pull(inputs, inputs_finished)? else {
return Ok(None);
};
if let Some(output) = state.broadcast(operation, chunk)? {
if let Some(output) = state
.broadcast(operation, chunk)
.map_err(|error| state.context(error))?
{
return Ok(Some(NativeChunk::Evaluated(output)));
}
}
@@ -2758,12 +2866,16 @@ impl RuntimeNode {
// bytes. Nothing is answered before the input ends: until then a group's value at a
// point is still missing the members that have not arrived.
while let Some(chunk) = input.pull(inputs, inputs_finished)? {
state.retain_aggregate(fold, chunk)?;
state
.retain_aggregate(fold, chunk)
.map_err(|error| state.context(error))?;
}
if !inputs_finished || !input.settled(inputs) {
return Ok(None);
}
state.compute_aggregate(fold)?;
state
.compute_aggregate(fold)
.map_err(|error| state.context(error))?;
// A group's series holds one value per field and point it decided, so its values
// are counted rather than read as its positions, as a binary result's are.
state.pop_queued(evaluated_values)
@@ -2779,12 +2891,16 @@ impl RuntimeNode {
// Which samples are kept at a point is not decided until the whole input has
// arrived, so the input is read to the end before anything is handed over.
while let Some(chunk) = input.pull(inputs, inputs_finished)? {
state.retain_vector(operation, chunk)?;
state
.retain_vector(operation, chunk)
.map_err(|error| state.context(error))?;
}
if !inputs_finished || !input.settled(inputs) {
return Ok(None);
}
state.compute_vector(operation)?;
state
.compute_vector(operation)
.map_err(|error| state.context(error))?;
state.pop_queued(evaluated_values)
}
}
@@ -2927,7 +3043,9 @@ impl RuntimeNode {
let Some(evaluator) = window.as_mut() else {
return Err(no_window());
};
state.flush_merge(evaluator, merge)?;
state
.flush_merge(evaluator, merge)
.map_err(|error| state.context(error))?;
flushed = true;
}
Ok(flushed)
@@ -2972,10 +3090,10 @@ impl RuntimeNode {
left,
right,
binary,
..
state,
} => {
if let Some(operation) = binary {
operation.closed()?;
operation.closed().map_err(|error| state.context(error))?;
}
left.unclosed()?;
right.unclosed()
@@ -2984,9 +3102,16 @@ impl RuntimeNode {
left.unclosed()?;
right.unclosed()
}
RuntimeOp::Vector { input, vector, .. } => {
RuntimeOp::Vector {
input,
vector,
state,
} => {
if let Some(operation) = vector {
operation.input.ensure_closed(Reader::Vector)?;
operation
.input
.ensure_closed(Reader::Vector)
.map_err(|error| state.context(error))?;
}
input.unclosed()
}
@@ -3015,9 +3140,17 @@ impl RuntimeNode {
let second = right.finish_windows();
first.and(second)
}
RuntimeOp::Window { input, window, .. } => {
RuntimeOp::Window {
input,
window,
state,
..
} => {
let child = input.finish_windows();
let own = window.take().map_or(Ok(()), RangeEvaluator::finish);
let own = window
.take()
.map_or(Ok(()), RangeEvaluator::finish)
.map_err(|error| state.context(error));
child.and(own)
}
}
@@ -3124,13 +3257,17 @@ impl RuntimeNode {
impl NativeExecutionState {
pub fn try_new(tree: &NativeTree, pool: Arc<dyn MemoryPool>) -> Result<Self> {
let mut inputs = Vec::new();
let (root, _) = RuntimeNode::build(tree.root(), &pool, &mut inputs)?;
let (root, _) = RuntimeNode::build(tree.root(), &pool, &mut inputs, "root".to_string())?;
let scalar_root = if inputs.is_empty() {
match (&tree.root().expr, &tree.root().output) {
(NativeExpr::Scalar(BoundScalar { value }), NativeType::Scalar(grid)) => Some((
Arc::clone(grid),
*value,
NodeState::new("native.scalar", Arc::clone(&root.metrics), &pool),
NodeState::new(
"native.scalar at root".to_string(),
Arc::clone(&root.metrics),
&pool,
),
)),
_ => {
return Err(DataFusionError::Execution(
@@ -3224,8 +3361,11 @@ impl NativeExecutionState {
.ok_or_else(|| {
DataFusionError::Execution("scalar allocation size overflow".into())
})?;
state.buffer.grow(bytes)?;
let chunk = scalar_chunk(grid, value)?;
state
.buffer
.grow(bytes)
.map_err(|error| state.context(error))?;
let chunk = scalar_chunk(grid, value).map_err(|error| state.context(error))?;
state.buffer.metrics.record_completed_series(1);
state
.queued
@@ -3243,7 +3383,8 @@ impl NativeExecutionState {
/// terms.
fn unclosed(&self) -> Result<()> {
if let Some(index) = self.inputs.iter().position(|slot| slot.open_source_series) {
return Err(truncated_input(InputId::new(index)));
let slot = &self.inputs[index];
return Err(slot.state.context(truncated_input(InputId::new(index))));
}
self.root.unclosed()
}
@@ -4094,7 +4235,7 @@ impl InputSlot {
let bound = self.schema.eval_grid.is_some();
let evaluated = matches!(chunk, NativeChunk::Evaluated(_));
if evaluated != bound {
return Err(DataFusionError::Execution(format!(
return Err(self.state.context(DataFusionError::Execution(format!(
"the native source is bound to {}, but {} chunk was queued at it",
if bound {
"an evaluated series"
@@ -4102,9 +4243,11 @@ impl InputSlot {
"raw samples"
},
if evaluated { "an evaluated" } else { "a raw" },
)));
))));
}
self.state.enqueue(chunk)
self.state
.enqueue(chunk)
.map_err(|error| self.state.context(error))
}
/// Hands over the next queued source chunk, if there is one.
src/promql/src/native/execute/tests.rs
+139 -32
View File
@@ -32,6 +32,14 @@ use datatypes::series::samples::{
use datatypes::series::time::{EvalGrid, GridRange, GridSpan, GridSpec, TimeAxis};
use datatypes::series::{Label, SeriesLabels, SeriesOrdering};
fn is_resources_exhausted(error: &DataFusionError) -> bool {
match error {
DataFusionError::ResourcesExhausted(_) => true,
DataFusionError::Context(_, inner) => is_resources_exhausted(inner),
_ => false,
}
}
use super::{
BinaryRuntime, Emitted, InputId, InputSlot, NativeChunk, NativeExecutionState, NativeExecutor,
NodeState, PairKernel, RuntimeNode, RuntimeOp, absorb_rebuilds, evaluated_bytes, grid_end_ms,
@@ -265,6 +273,35 @@ fn instant(
}
}
#[test]
fn timestamp_instant_rejects_negative_lookback() {
let output_grid = grid(500, 1_000, 1);
let mut root = instant(input(), -1, 0, Arc::clone(&output_grid));
let NativeExpr::Instant { op, .. } = &mut root.expr else {
unreachable!();
};
op.timestamp = true;
op.at_ms = Some(500);
let tree = NativeTree::try_new(root).unwrap();
let pool = pool();
let mut execution = executor(&tree, Arc::clone(&pool));
execution
.push_raw(
InputId::new(0),
raw_chunk(&raw_schema(), &[500], &[1.0], true),
)
.unwrap();
let error = execution
.finish()
.expect_err("negative lookback was accepted");
assert!(
error.to_string().contains("lookback must be non-negative"),
"unexpected error: {error}"
);
assert_eq!(pool.reserved(), 0);
}
fn merge(inner: NativeNode) -> NativeNode {
let output = inner.output.clone();
NativeNode {
@@ -521,7 +558,11 @@ fn source_runtime(pool: &Arc<dyn MemoryPool>) -> InputSlot {
source_index: 0,
schema: raw_schema(),
open_source_series: false,
state: NodeState::new("native.input", Arc::new(NodeMetrics::default()), pool),
state: NodeState::new(
"native.input at root (source_index=0)",
Arc::new(NodeMetrics::default()),
pool,
),
}
}
@@ -535,7 +576,11 @@ fn binary_runtime(
schema: &Arc<SeriesSchema>,
) -> (NodeState, BinaryRuntime) {
(
NodeState::new("native.binary", Arc::new(NodeMetrics::default()), pool),
NodeState::new(
"native.binary at root",
Arc::new(NodeMetrics::default()),
pool,
),
BinaryRuntime::new(
PairKernel::Binary(Arc::new(binary_matching(BinaryOp::Add))),
[Arc::clone(schema), Arc::clone(schema)],
@@ -1139,10 +1184,11 @@ fn timestamp_projection_workspace_failure_releases_runtime() {
evaluated_chunk(&evaluated_schema(grid(0, 1, 10_000)), 0, vec![1.0], true),
)
.unwrap();
assert!(matches!(
execution.next_chunk(),
Err(DataFusionError::ResourcesExhausted(_))
));
assert!(
execution
.next_chunk()
.is_err_and(|error| is_resources_exhausted(&error))
);
assert_eq!(pool.reserved(), 0);
assert_eq!(execution.state().reserved_bytes(), 0);
}
@@ -1725,6 +1771,10 @@ fn an_error_releases_the_reservation() {
.unwrap();
let error = execution.next_chunk().unwrap_err();
assert!(error.to_string().contains("not increasing"), "{error}");
assert!(
error.to_string().contains("native.window at root"),
"{error}"
);
// Everything went back to the pool while the execution is alive: waiting for the drop would
// keep the query's bytes in use for as long as the caller holds the executor.
assert_eq!(pool.reserved(), 0);
@@ -1815,13 +1865,60 @@ fn a_fanout_enqueue_failure_releases_every_recipient() {
let pool = bounded_pool(first_charge);
let mut execution = executor(&tree, Arc::clone(&pool));
let ids: Vec<_> = execution.state().input_ids(0).collect();
let names = &execution.state().inputs;
assert_ne!(
names[0].state.buffer.reservation.consumer().name(),
names[1].state.buffer.reservation.consumer().name()
);
assert!(
names[0]
.state
.buffer
.reservation
.consumer()
.name()
.contains("root.left.input")
);
assert!(
names[1]
.state
.buffer
.reservation
.consumer()
.name()
.contains("root.right.input")
);
assert!(
names[0]
.state
.buffer
.reservation
.consumer()
.name()
.contains("source_index=0")
);
assert!(
names[1]
.state
.buffer
.reservation
.consumer()
.name()
.contains("source_index=0")
);
let first = ids[0];
let second = ids[1];
let error = execution
.push_raw(first, chunk.clone())
.and_then(|()| execution.push_raw(second, chunk))
.expect_err("second fanout recipient exceeded the shared pool");
assert!(matches!(error, DataFusionError::ResourcesExhausted(_)));
match &error {
DataFusionError::Context(context, inner) => {
assert_eq!(context, "native.input at root.right.input (source_index=0)");
assert!(is_resources_exhausted(inner), "{inner}");
}
other => panic!("expected input context around resource exhaustion, got {other}"),
}
assert_eq!(execution.state().reserved_bytes(), 0);
assert_eq!(pool.reserved(), 0);
assert!(
@@ -1870,10 +1967,7 @@ fn a_pool_that_cannot_admit_a_chunk_releases_what_it_held() {
raw_chunk(&schema, &timestamps, &values, true),
)
.expect_err("a chunk larger than the pool was admitted");
assert!(
matches!(error, DataFusionError::ResourcesExhausted(_)),
"{error}"
);
assert!(is_resources_exhausted(&error), "{error}");
// The chunk never entered the queue, and everything the execution held went back to the pool
// while the executor is still alive.
@@ -2090,6 +2184,12 @@ fn an_input_that_ends_inside_a_collected_series_is_reported() {
.contains("ended before the last chunk of a series"),
"the truncation was reported as {error}"
);
assert!(
error
.to_string()
.contains("native.input at root.input.input (source_index=0)"),
"{error}"
);
// Nothing was merged and nothing was answered, and everything went back while the executor is
// alive.
assert_eq!(execution.state().reserved_bytes(), 0);
@@ -2685,10 +2785,7 @@ fn a_pool_that_cannot_admit_a_collected_fragment_releases_what_it_held() {
let error = execution
.next_chunk()
.expect_err("the chain was admitted by a pool that cannot hold what it makes of the chunk");
assert!(
matches!(error, DataFusionError::ResourcesExhausted(_)),
"{error}"
);
assert!(is_resources_exhausted(&error), "{error}");
assert_eq!(execution.state().reserved_bytes(), 0);
assert_eq!(pool.reserved(), 0);
assert!(
@@ -2762,10 +2859,7 @@ fn a_pool_that_cannot_admit_the_merge_workspace_releases_what_it_held() {
let error = execution
.finish()
.expect_err("the merge ran without its workspace admitted");
assert!(
matches!(error, DataFusionError::ResourcesExhausted(_)),
"{error}"
);
assert!(is_resources_exhausted(&error), "{error}");
// Everything the chain held went back while the executor is alive: the fragments, the
// evaluator and both queues.
assert_eq!(execution.state().reserved_bytes(), 0);
@@ -5166,7 +5260,7 @@ fn an_instant_selection_reserves_the_workspace_it_gathers() {
let error = execution.next_chunk().unwrap_err();
assert!(
matches!(error, DataFusionError::ResourcesExhausted(_)),
is_resources_exhausted(&error),
"the gather's workspace was taken without asking the pool: {error}"
);
assert!(
@@ -5274,7 +5368,7 @@ fn a_selection_reserves_the_points_it_decides() {
.unwrap();
let error = execution.next_chunk().unwrap_err();
assert!(
matches!(error, DataFusionError::ResourcesExhausted(_)),
is_resources_exhausted(&error),
"a series covering the grid was gathered without asking the pool: {error}"
);
assert_eq!(pool.reserved(), 0);
@@ -5306,7 +5400,7 @@ fn an_oversized_selection_is_refused_by_the_pool() {
let error = execution.next_chunk().unwrap_err();
assert!(
matches!(error, DataFusionError::ResourcesExhausted(_)),
is_resources_exhausted(&error),
"the gather's workspace was taken without asking the pool: {error}"
);
assert!(
@@ -6078,7 +6172,7 @@ fn an_oversized_binary_workspace_is_refused_by_the_pool() {
.unwrap();
let error = execution.finish().unwrap_err();
assert!(
matches!(error, DataFusionError::ResourcesExhausted(_)),
is_resources_exhausted(&error),
"a point limit answered in the pool's place: {error}"
);
assert_eq!(pool.reserved(), 0);
@@ -6105,7 +6199,7 @@ fn an_oversized_binary_operand_is_refused_by_the_pool() {
.push_evaluated(right, evaluated_chunk(&schema, 0, values, true))
.unwrap_err();
assert!(
matches!(error, DataFusionError::ResourcesExhausted(_)),
is_resources_exhausted(&error),
"an operand limit answered in the pool's place: {error}"
);
assert_eq!(pool.reserved(), 0);
@@ -6389,7 +6483,7 @@ fn a_pool_that_holds_one_histogram_table_refuses_a_second() {
.retain_binary(&mut refused_binary, 0, NativeChunk::Evaluated(second_own))
.unwrap_err();
assert!(
matches!(error, DataFusionError::ResourcesExhausted(_)),
is_resources_exhausted(&error),
"a second chunk holding a table of its own was admitted beside the first: {error}"
);
assert_eq!(
@@ -6676,7 +6770,7 @@ fn an_assembly_the_pool_cannot_hold_is_refused_before_it_allocates() {
assert_eq!(state.buffer.reservation.size(), retained);
let error = state.compute_binary(&mut binary).unwrap_err();
assert!(
matches!(error, DataFusionError::ResourcesExhausted(_)),
is_resources_exhausted(&error),
"the assembly was not reserved before it allocated: {error}"
);
// The assembly never ran, so its charge is gone: what the node holds is what it retained, and
@@ -6745,7 +6839,7 @@ fn an_assembly_the_executors_pool_cannot_hold_stops_it() {
let error = execution.finish().unwrap_err();
assert!(
matches!(error, DataFusionError::ResourcesExhausted(_)),
is_resources_exhausted(&error),
"the assembly was not reserved before it allocated: {error}"
);
assert_eq!(
@@ -7813,7 +7907,7 @@ fn an_oversized_set_workspace_is_refused_by_the_pool() {
}
let error = execution.finish().unwrap_err();
assert!(
matches!(error, DataFusionError::ResourcesExhausted(_)),
matches!(error.find_root(), DataFusionError::ResourcesExhausted(_)),
"the set workspace was not refused by the pool: {error}"
);
assert_eq!(pool.reserved(), 0);
@@ -7839,6 +7933,13 @@ fn queued_set_results_hold_their_labels_until_handed_on() {
.unwrap();
let start = |pool: &Arc<dyn MemoryPool>| {
let mut execution = executor(&tree, Arc::clone(pool));
match &execution.state().root.op {
RuntimeOp::Binary { state, .. } => assert_eq!(
state.buffer.reservation.consumer().name(),
"native.set at root"
),
_ => panic!("expected set binary runtime"),
}
let (left, right) = (
execution.state().input_ids(0).next().unwrap(),
execution.state().input_ids(1).next().unwrap(),
@@ -9281,10 +9382,16 @@ fn an_aggregate_refuses_a_chunk_of_another_layout_and_returns_its_reservation()
let error = execution
.next_chunk()
.expect_err("a chunk of another field count was folded");
assert!(
matches!(error, DataFusionError::Execution(_)),
"the refusal is not an execution error: {error}"
);
match &error {
DataFusionError::Context(context, inner) => {
assert_eq!(context, "native.aggregate at root");
assert!(
matches!(inner.as_ref(), DataFusionError::Execution(_)),
"{inner}"
);
}
other => panic!("the refusal is not aggregate-contextualized: {other}"),
}
let message = error.to_string();
assert!(
message.contains("the `Sum` aggregate"),
src/promql/src/range_eval/function.rs
+7 -2
View File
@@ -46,17 +46,22 @@ pub enum RangeFunction {
Deriv,
/// `predict_linear` with a constant horizon in seconds.
PredictLinear(f64),
/// `double_exponential_smoothing` (`holt_winters`) with constant factors in `(0, 1)`.
/// Experimental in Prometheus: `double_exponential_smoothing` (`holt_winters`) with constant
/// factors in `(0, 1)`.
DoubleExponentialSmoothing {
smoothing: f64,
trend: f64,
},
// Prometheus marks the following as experimental.
FirstOverTime,
/// Experimental in Prometheus.
MadOverTime,
/// Experimental in Prometheus.
TsOfFirstOverTime,
/// Experimental in Prometheus.
TsOfLastOverTime,
/// Experimental in Prometheus.
TsOfMinOverTime,
/// Experimental in Prometheus.
TsOfMaxOverTime,
}
src/promql/src/source.rs
+8 -23
View File
@@ -12,7 +12,6 @@
// See the License for the specific language governing permissions and
// limitations under the License.
mod merge;
#[cfg(test)]
mod tests;
@@ -27,11 +26,9 @@ use datatypes::series::frame::LabelDecodeSpec;
use datatypes::series::raw::RawSeriesChunk;
use datatypes::series::samples::SeriesSchema;
use futures::Stream;
use mito_codec::labels::{DecodedLabels, LabelProjection, PrimaryKeyLabelDecoder};
use mito_codec::labels::PrimaryKeyLabelDecoder;
use store_api::region_engine::{RawLabelSource, RawSeriesStream, SendableRawSeriesStream};
pub use merge::merge_ordered_raw_series;
/// Materializes complete metric labels from a raw-series frame stream.
///
/// The labels are decoded from the stream's own label decode description, so a consumer that holds
@@ -64,10 +61,9 @@ pub fn materialize_raw_labels(
));
}
let decoder = match &decode {
Some(description) => Some(
PrimaryKeyLabelDecoder::from_decode_spec(description, LabelProjection::All)
.map_err(BoxedError::new)?,
),
Some(description) => {
Some(PrimaryKeyLabelDecoder::from_decode_spec(description).map_err(BoxedError::new)?)
}
None => None,
};
@@ -143,22 +139,11 @@ impl RawSeriesStream for MaterializedRawSeriesStream {
impl MaterializedRawSeriesStream {
fn materialize(&mut self, chunk: RawSeriesChunk) -> Result<RawSeriesChunk, BoxedError> {
let metric = match (&mut self.decoder, chunk.metric) {
(Some(decoder), None) => {
match decoder
(Some(decoder), None) => Some(
decoder
.decode(&chunk.primary_key)
.map_err(BoxedError::new)?
{
DecodedLabels::Complete(metric) => Some(metric),
DecodedLabels::Projected(_) => {
return Err(BoxedError::new(
InvalidSeriesSnafu {
reason: "primary-key decoder returned projected labels",
}
.build(),
));
}
}
}
.map_err(BoxedError::new)?,
),
(None, Some(metric)) => Some(metric),
(Some(_), Some(_)) => {
return Err(BoxedError::new(
src/promql/src/source/merge.rs
-459
View File
@@ -1,459 +0,0 @@
// Copyright 2023 Greptime Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#[cfg(test)]
mod tests;
use std::cmp::Ordering;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};
use common_error::ext::BoxedError;
use datatypes::error::InvalidSeriesSnafu;
use datatypes::series::SeriesOrdering;
use datatypes::series::frame::LabelDecodeSpec;
use datatypes::series::raw::RawSeriesChunk;
use datatypes::series::samples::{SeriesFields, SeriesSchema};
use futures::Stream;
use mito_codec::labels::{LabelProjection, PrimaryKeyLabelDecoder};
use store_api::region_engine::{
NO_BOUND_LABEL_SOURCE, RawLabelSource, RawSeriesStream, SendableRawSeriesStream,
};
/// Merge compatible, explicitly supplied raw streams by complete read key and time.
///
/// Callers provide authorized bindings and own snapshot/storage policy. This does not select an
/// MVCC view, resolve overlapping physical copies, or roll back emitted prefixes after errors.
/// `is_last` is global only over these inputs; there is no fallback, replay, or input sorting.
///
/// The merge contract is the label layout, not a binding: every input must declare the same one —
/// from its frame if it arrived as a frame, otherwise from its own binding — and the merged stream
/// advertises that one declaration. A consumer therefore decodes the merged stream from the frame;
/// the identity it reports is none, see [`RawSeriesStream::label_source`].
pub fn merge_ordered_raw_series(
inputs: Vec<SendableRawSeriesStream>,
) -> Result<SendableRawSeriesStream, BoxedError> {
let mut inputs = inputs.into_iter();
let first = inputs
.next()
.ok_or_else(|| invalid("raw merge requires at least one input"))?;
let schema = Arc::clone(first.schema());
let decode = merged_decode_spec(first.as_ref().get_ref())?;
let mut streams = vec![Input::new(first)];
for stream in inputs {
streams.push(Input::new(stream));
}
for input in &streams {
let stream = input
.stream
.as_ref()
.ok_or_else(|| invalid("raw merge input disappeared"))?;
validate_input(stream.as_ref().get_ref(), &schema)?;
validate_decode(stream.as_ref().get_ref(), &decode)?;
}
Ok(Box::pin(OrderedRawMerge {
streams,
schema,
decode,
candidate: None,
}))
}
fn invalid(reason: &'static str) -> BoxedError {
BoxedError::new(InvalidSeriesSnafu { reason }.build())
}
/// Returns the label layout a stream declares, with no binding required to state it.
///
/// Every key-encoded stream declares its own description — a producer derives it from the binding it
/// read under and states it before the first poll, and a stream reassembled from a frame carries the
/// description its producer wrote — so the answer is exactly what decoding that stream's chunks
/// needs. A merge can therefore compare inputs that were framed and inputs that were not in the same
/// terms. `None` means the stream carries complete labels and declares no key layout at all.
fn declared_layout(stream: &dyn RawSeriesStream) -> Option<LabelDecodeSpec> {
stream.label_decode_spec().cloned()
}
/// Returns the one label decode description the merged stream advertises for all inputs.
///
/// The first input's declaration is forwarded unchanged, so a merge advertises exactly what its
/// inputs advertise, whether they are frame-decoded or straight from a producer. The declaration is
/// validated as a decode layout before it is advertised, so a downstream consumer holding only the
/// merged frame schema can decode it.
fn merged_decode_spec(stream: &dyn RawSeriesStream) -> Result<LabelDecodeSpec, BoxedError> {
let description = declared_layout(stream).ok_or_else(|| {
invalid("raw merge requires inputs that describe a primary-key label layout")
})?;
PrimaryKeyLabelDecoder::from_decode_spec(&description, LabelProjection::All)
.map_err(BoxedError::new)?;
Ok(description)
}
/// Requires every input to declare the advertised label layout.
///
/// The merged stream advertises one description, so an input declaring another would be decoded as
/// if it had the advertised encoding. Every input is compared through [`declared_layout`], which
/// reads the input's own declaration, so no input needs a binding to be present.
fn validate_decode(
stream: &dyn RawSeriesStream,
decode: &LabelDecodeSpec,
) -> Result<(), BoxedError> {
if declared_layout(stream).as_ref() != Some(decode) {
return Err(invalid(
"raw merge inputs describe different label decode layouts",
));
}
Ok(())
}
fn validate_input(stream: &dyn RawSeriesStream, schema: &SeriesSchema) -> Result<(), BoxedError> {
let other = stream.schema();
if stream.ordering() != SeriesOrdering::KeyThenTime
|| other.eval_grid.is_some()
|| other.fields.is_empty()
|| other.fields.len() != schema.fields.len()
|| other
.fields
.iter()
.zip(schema.fields.iter())
.any(|(a, b)| a.name != b.name)
{
return Err(invalid(
"raw merge requires equal raw schemas and ReadKeyThenTime ordering",
));
}
Ok(())
}
struct Input {
stream: Option<SendableRawSeriesStream>,
head: Option<RawSeriesChunk>,
previous_key: Option<Vec<u8>>,
previous_time: Option<i64>,
ended_series: bool,
}
impl Input {
fn new(stream: SendableRawSeriesStream) -> Self {
Self {
stream: Some(stream),
head: None,
previous_key: None,
previous_time: None,
ended_series: false,
}
}
fn validate(
&mut self,
chunk: &RawSeriesChunk,
schema: &SeriesSchema,
) -> Result<(), BoxedError> {
if chunk.metric.is_some() {
return Err(invalid("raw merge requires undecoded primary-key chunks"));
}
RawSeriesChunk::try_new(
schema,
chunk.primary_key.clone(),
None,
chunk.time.clone(),
chunk.fields.clone(),
chunk.is_last,
)
.map_err(BoxedError::new)?;
let times = chunk.time.as_buffer();
if let Some(previous_key) = &self.previous_key {
match previous_key.as_slice().cmp(chunk.primary_key.as_slice()) {
Ordering::Greater => {
return Err(invalid("raw merge input read keys are not ordered"));
}
Ordering::Equal
if self.ended_series
|| self.previous_time.is_some_and(|time| times[0] <= time) =>
{
return Err(invalid(
"raw merge input has repeated, non-increasing timestamps or invalid is_last",
));
}
Ordering::Less if !self.ended_series => {
return Err(invalid("raw merge input changed read keys before is_last"));
}
_ => {}
}
}
self.previous_key = Some(chunk.primary_key.as_slice().to_vec());
self.previous_time = times.last().copied();
self.ended_series = chunk.is_last;
Ok(())
}
/// Returns whether this input can no longer contribute chunks for `key`.
///
/// A held head beyond `key` proves completion. Without a head, only an ended input or `key`'s
/// own final chunk proves completion: an earlier completed key never rules out unknown later
/// or intervening keys.
fn finished_at(&self, key: &[u8]) -> bool {
match self.head.as_ref() {
Some(head) => head.primary_key.as_slice() > key,
None => {
self.stream.is_none()
|| self.previous_key.as_deref().is_some_and(|previous| {
previous > key || (previous == key && self.ended_series)
})
}
}
}
}
struct OrderedRawMerge {
streams: Vec<Input>,
schema: Arc<SeriesSchema>,
/// The description the merged stream advertises for all inputs; see [`merged_decode_spec`].
decode: LabelDecodeSpec,
candidate: Option<RawSeriesChunk>,
}
impl OrderedRawMerge {
fn fail(&mut self, error: BoxedError) -> Poll<Option<Result<RawSeriesChunk, BoxedError>>> {
self.streams.clear();
self.candidate = None;
Poll::Ready(Some(Err(error)))
}
/// Polls every input that has no head once, returning upstream errors before progress.
fn poll_inputs(&mut self, cx: &mut Context<'_>) -> Result<(), BoxedError> {
for input in &mut self.streams {
if input.stream.is_none() || input.head.is_some() {
continue;
}
match input
.stream
.as_mut()
.map(|stream| stream.as_mut().poll_next(cx))
{
Some(Poll::Pending) => {}
Some(Poll::Ready(Some(Err(error)))) => return Err(error),
Some(Poll::Ready(Some(Ok(chunk)))) => {
input.validate(&chunk, &self.schema)?;
input.head = Some(chunk);
}
Some(Poll::Ready(None)) => {
if input.previous_key.is_some() && !input.ended_series {
return Err(invalid("raw merge input ended before its final chunk"));
}
input.stream = None;
input.previous_key = None;
input.previous_time = None;
}
None => {}
}
}
Ok(())
}
/// Returns the index of the smallest available head, by read key then first timestamp.
fn smallest_head_index(&self) -> Option<usize> {
self.streams
.iter()
.enumerate()
.filter_map(|(index, input)| input.head.as_ref().map(|head| (index, head)))
.min_by(|(_, a), (_, b)| {
a.primary_key
.as_slice()
.cmp(b.primary_key.as_slice())
.then_with(|| a.time.timestamp(0).cmp(&b.time.timestamp(0)))
})
.map(|(index, _)| index)
}
/// Returns whether every input accounts for `key`: it holds a head, or completion is proven.
///
/// A headless input that cannot prove completion may still add chunks to `key`, so the merge
/// must wait instead of deciding `key`.
fn all_inputs_ready_at(&self, key: &[u8]) -> bool {
self.streams
.iter()
.all(|input| input.head.is_some() || input.finished_at(key))
}
fn take_next(&mut self) -> Result<Option<RawSeriesChunk>, BoxedError> {
let Some(index) = self.smallest_head_index() else {
return Ok(None);
};
let chunk = self.streams[index]
.head
.take()
.ok_or_else(|| invalid("raw merge lost selected head"))?;
let key = chunk.primary_key.as_slice();
let start = chunk
.time
.timestamp(0)
.ok_or_else(|| invalid("raw merge chunk has no timestamp"))?;
let mut end = chunk.time.len();
for input in &self.streams {
if let Some(head) = &input.head
&& head.primary_key.as_slice() == key
{
let next = head
.time
.timestamp(0)
.ok_or_else(|| invalid("raw merge chunk has no timestamp"))?;
if next == start {
return Err(invalid(
"raw merge inputs contain duplicate timestamps for a complete read key",
));
}
if next > start {
end = end.min(
chunk
.time
.as_buffer()
.partition_point(|timestamp| *timestamp < next),
);
}
}
}
if end == chunk.time.len() {
return Ok(Some(RawSeriesChunk {
primary_key: chunk.primary_key,
metric: None,
time: chunk.time,
fields: chunk.fields,
is_last: false,
}));
}
let fields = SeriesFields::try_from_vec(
chunk
.fields
.iter()
.map(|samples| samples.slice(0, end).map_err(BoxedError::new))
.collect::<Result<Vec<_>, _>>()?,
)
.map_err(BoxedError::new)?;
let time = chunk.time.slice(0, end).map_err(BoxedError::new)?;
if end < chunk.time.len() {
let remainder = SeriesFields::try_from_vec(
chunk
.fields
.iter()
.map(|samples| {
samples
.slice(end, chunk.time.len() - end)
.map_err(BoxedError::new)
})
.collect::<Result<Vec<_>, _>>()?,
)
.map_err(BoxedError::new)?;
self.streams[index].head = Some(RawSeriesChunk {
primary_key: chunk.primary_key.clone(),
metric: None,
time: chunk
.time
.slice(end, chunk.time.len() - end)
.map_err(BoxedError::new)?,
fields: remainder,
is_last: chunk.is_last,
});
}
Ok(Some(RawSeriesChunk {
primary_key: chunk.primary_key,
metric: None,
time,
fields,
is_last: false,
}))
}
}
impl Stream for OrderedRawMerge {
type Item = Result<RawSeriesChunk, BoxedError>;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let this = self.get_mut();
if this.streams.is_empty() {
return Poll::Ready(None);
}
if this.candidate.is_none() {
if let Err(error) = this.poll_inputs(cx) {
return this.fail(error);
}
let Some(head) = this
.smallest_head_index()
.and_then(|index| this.streams[index].head.as_ref())
else {
// Only a fully drained input set is EOF: a live input may still reveal future keys
// or a late error, so its state stays owned instead of being cleared.
if this.streams.iter().all(|input| input.stream.is_none()) {
this.streams.clear();
return Poll::Ready(None);
}
return Poll::Pending;
};
// A headless input that cannot prove completion at the smallest available key may
// still add chunks to it, so the merge waits before consuming that head.
if !this.all_inputs_ready_at(head.primary_key.as_slice()) {
return Poll::Pending;
}
match this.take_next() {
Ok(candidate) => this.candidate = candidate,
Err(error) => return this.fail(error),
}
}
if let Err(error) = this.poll_inputs(cx) {
return this.fail(error);
}
let Some(candidate) = this.candidate.take() else {
this.streams.clear();
return Poll::Ready(None);
};
if !this.all_inputs_ready_at(candidate.primary_key.as_slice()) {
// The candidate's key may still be extended or interleaved by a headless input: an
// already-decided prefix is held back instead of being emitted as if it were final.
this.candidate = Some(candidate);
return Poll::Pending;
}
let is_last = this
.streams
.iter()
.all(|input| input.finished_at(candidate.primary_key.as_slice()));
Poll::Ready(Some(Ok(RawSeriesChunk {
is_last,
..candidate
})))
}
}
impl RawSeriesStream for OrderedRawMerge {
fn schema(&self) -> &Arc<SeriesSchema> {
&self.schema
}
/// Returns [`NO_BOUND_LABEL_SOURCE`] unconditionally.
///
/// A merge is not a region scan, so it has no scan-bound identity to report, and it does not
/// invent one from its inputs. Its decode contract is [`Self::label_decode_spec`]; identity is
/// not something a merged stream can report.
fn label_source(&self) -> &RawLabelSource {
&NO_BOUND_LABEL_SOURCE
}
/// Returns the description this merge advertises for every input, so a consumer of the merged
/// stream decodes labels from the frame instead of from the binding.
fn label_decode_spec(&self) -> Option<&LabelDecodeSpec> {
Some(&self.decode)
}
fn ordering(&self) -> SeriesOrdering {
SeriesOrdering::KeyThenTime
}
}
src/promql/src/source/merge/tests.rs
-2020
View File
File diff suppressed because it is too large Load Diff
src/promql/src/source/merge/tests/payload.rs
-654
View File
@@ -1,654 +0,0 @@
// Copyright 2023 Greptime Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use std::collections::VecDeque;
use std::pin::Pin;
use std::sync::{Arc, Weak};
use std::task::{Context, Poll};
use common_error::ext::BoxedError;
use common_query::series_frame::{
ensure_frame_matches_label_source, label_decode_spec, pack_raw_series, series_frame_ordering,
unpack_raw_series,
};
use datatypes::arrow::array::{Array, Float64Array, ListArray, UInt32Array};
use datatypes::arrow::buffer::{Buffer, ScalarBuffer};
use datatypes::series::SeriesOrdering;
use datatypes::series::frame::LabelDecodeSpec;
use datatypes::series::histogram::{
BucketSideView, CounterResetHint, HistogramBatch, HistogramBatchBuilder, HistogramView,
};
use datatypes::series::raw::{RawSeriesChunk, RawTimestamps};
use datatypes::series::samples::{
DenseSamplesRef, FieldSpec, HistogramSamples, SampleRef, Samples, SeriesFields, SeriesSchema,
};
use futures::{Stream, StreamExt, TryStreamExt};
use store_api::region_engine::{RawLabelSource, RawSeriesStream, SendableRawSeriesStream};
use super::{label_source, read_key};
use crate::source::{materialize_raw_labels, merge_ordered_raw_series};
fn payload_schema() -> Arc<SeriesSchema> {
Arc::new(SeriesSchema {
fields: vec!["temperature", "latency", "observations"]
.into_iter()
.map(|name| FieldSpec {
name: Arc::from(name),
})
.collect::<Vec<_>>()
.into(),
eval_grid: None,
})
}
fn histograms(sums: &[f64]) -> Arc<HistogramBatch> {
let mut builder = HistogramBatchBuilder::new();
for &sum in sums {
builder
.push(HistogramView {
schema: -19,
zero_threshold: 0.125,
sum,
count: 23.5,
zero_count: 2.25,
reset_hint: CounterResetHint::Unrecognized(37),
start_timestamp: Some(-456),
custom_values: &[1.5, 4.5],
positive: BucketSideView {
span_offsets: &[1],
span_lengths: &[2],
buckets: &[3.25, 7.5],
},
negative: BucketSideView {
span_offsets: &[-2],
span_lengths: &[1],
buckets: &[0.75],
},
})
.unwrap();
}
builder.finish().unwrap()
}
fn payload_chunk(
key: Buffer,
times: &[i64],
floats: &[Option<f64>],
latency: &[Option<f64>],
mixed_sums: &[Option<f64>],
histogram_sums: &[Option<f64>],
last: bool,
) -> RawSeriesChunk {
assert_eq!(times.len(), floats.len());
assert_eq!(times.len(), latency.len());
assert_eq!(times.len(), mixed_sums.len());
assert_eq!(times.len(), histogram_sums.len());
let mixed_floats = Float64Array::from(latency.to_vec());
let mut mixed_index = 0;
let mixed_indices = UInt32Array::from(
mixed_sums
.iter()
.map(|sum| {
sum.map(|_| {
let index = mixed_index;
mixed_index += 1;
index
})
})
.collect::<Vec<_>>(),
);
let mixed = Samples::try_from_mixed(
mixed_floats,
HistogramSamples::try_new(
mixed_indices,
histograms(&mixed_sums.iter().flatten().copied().collect::<Vec<_>>()),
)
.unwrap(),
)
.unwrap();
let mixed = if times.len() > 1 {
Samples::try_from_chunks(vec![
mixed.slice(0, 1).unwrap(),
mixed.slice(1, times.len() - 1).unwrap(),
])
.unwrap()
} else {
mixed
};
let mut histogram_index = 0;
let histogram_indices = UInt32Array::from(
histogram_sums
.iter()
.map(|sum| {
sum.map(|_| {
let index = histogram_index;
histogram_index += 1;
index
})
})
.collect::<Vec<_>>(),
);
let fields = SeriesFields::try_from_vec(vec![
Samples::from_float(Float64Array::from(floats.to_vec())),
mixed,
{
let table = histograms(&histogram_sums.iter().flatten().copied().collect::<Vec<_>>());
let samples = Samples::try_from_histogram(histogram_indices, table).unwrap();
if times.len() > 1 {
Samples::try_from_chunks(vec![
samples.slice(0, 1).unwrap(),
samples.slice(1, times.len() - 1).unwrap(),
])
.unwrap()
} else {
samples
}
},
])
.unwrap();
RawSeriesChunk::try_new(
&payload_schema(),
key,
None,
RawTimestamps::try_new(ScalarBuffer::from(times.to_vec())).unwrap(),
fields,
last,
)
.unwrap()
}
struct PayloadStream {
schema: Arc<SeriesSchema>,
labels: RawLabelSource,
/// The description this fixture declares, derived once from `labels`, as a producer does.
decode: Option<LabelDecodeSpec>,
items: VecDeque<Result<RawSeriesChunk, BoxedError>>,
}
impl Stream for PayloadStream {
type Item = Result<RawSeriesChunk, BoxedError>;
fn poll_next(mut self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Option<Self::Item>> {
Poll::Ready(self.items.pop_front())
}
}
impl RawSeriesStream for PayloadStream {
fn schema(&self) -> &Arc<SeriesSchema> {
&self.schema
}
fn label_source(&self) -> &RawLabelSource {
&self.labels
}
fn label_decode_spec(&self) -> Option<&LabelDecodeSpec> {
self.decode.as_ref()
}
fn ordering(&self) -> SeriesOrdering {
SeriesOrdering::KeyThenTime
}
}
fn payload_stream(labels: RawLabelSource, chunks: Vec<RawSeriesChunk>) -> SendableRawSeriesStream {
Box::pin(PayloadStream {
schema: payload_schema(),
decode: label_decode_spec(&labels),
labels,
items: chunks.into_iter().map(Ok).collect(),
})
}
fn assert_histogram(value: HistogramView<'_>, sum: f64) {
assert_eq!(value.schema, -19);
assert_eq!(value.zero_threshold, 0.125);
assert_eq!(value.sum.to_bits(), sum.to_bits());
assert_eq!(value.count, 23.5);
assert_eq!(value.zero_count, 2.25);
assert_eq!(value.reset_hint, CounterResetHint::Unrecognized(37));
assert_eq!(value.start_timestamp, Some(-456));
assert_eq!(value.custom_values, &[1.5, 4.5]);
assert_eq!(value.positive.span_offsets, &[1]);
assert_eq!(value.positive.span_lengths, &[2]);
assert_eq!(value.positive.buckets, &[3.25, 7.5]);
assert_eq!(value.negative.span_offsets, &[-2]);
assert_eq!(value.negative.span_lengths, &[1]);
assert_eq!(value.negative.buckets, &[0.75]);
}
#[tokio::test]
async fn merges_typed_multifield_payload_through_frame_and_label_roundtrip() {
let key = read_key("rack-payload", "sensor-7");
let a = payload_stream(
label_source(1),
vec![payload_chunk(
key.clone(),
&[10, 30],
&[Some(f64::from_bits(0x7ff8_0000_0000_0042)), Some(-0.0)],
&[Some(1.25), None],
&[None, Some(11.0)],
&[Some(101.0), None],
true,
)],
);
let b = payload_stream(
label_source(2),
vec![payload_chunk(
key.clone(),
&[20, 40],
&[None, Some(f64::from_bits(0x7ff0_0000_0000_0002))],
&[None, None],
&[None, Some(21.0)],
&[Some(202.0), Some(303.0)],
true,
)],
);
let merged = merge_ordered_raw_series(vec![a, b]).unwrap();
let frame_schema = merged.schema().clone();
let frame_ordering = merged.ordering();
assert!(
matches!(merged.label_source(), RawLabelSource::Explicit),
"a merge is not a region scan, so it reports no scan-bound identity"
);
let framed = pack_raw_series(merged).unwrap();
// The frame states the layout the decoder uses: it matches the description an input's binding
// declares, since the merge only accepts inputs that declare one layout.
ensure_frame_matches_label_source(framed.schema().as_ref(), &label_source(1)).unwrap();
assert_eq!(
series_frame_ordering(framed.schema().as_ref()).unwrap(),
frame_ordering,
"the packed frame declares the merge's ordering"
);
let raw = unpack_raw_series(framed, Some(frame_schema)).unwrap();
assert_eq!(
raw.ordering(),
frame_ordering,
"unpacking derives the ordering from the frame"
);
let rows = materialize_raw_labels(raw)
.unwrap()
.try_collect::<Vec<_>>()
.await
.unwrap();
assert_eq!(
rows.iter()
.map(|r| r.time.timestamp(0).unwrap())
.collect::<Vec<_>>(),
[10, 20, 30, 40]
);
assert!(rows.iter().all(|r| r.primary_key == key));
for row in &rows {
let metric = row.metric.as_ref().unwrap();
assert_eq!(metric.name(), Some("room_temperature"));
assert_eq!(metric.get("rack"), Some("rack-payload"));
assert_eq!(metric.get("sensor"), Some("sensor-7"));
}
let floats: Vec<_> = rows
.iter()
.map(|r| match r.fields.get(0).unwrap().sample(0) {
SampleRef::Float(v) => Some(v.to_bits()),
SampleRef::Missing => None,
_ => panic!("unexpected temperature sample variant"),
})
.collect();
assert_eq!(
floats,
[
Some(0x7ff8_0000_0000_0042),
None,
Some((-0.0_f64).to_bits()),
Some(0x7ff0_0000_0000_0002)
]
);
let mixed: Vec<_> = rows
.iter()
.map(|r| match r.fields.get(1).unwrap().sample(0) {
SampleRef::Float(v) => Some(("float", v.to_bits())),
SampleRef::Missing => None,
SampleRef::Histogram(h) => {
assert_histogram(
h,
if r.time.timestamp(0) == Some(30) {
11.0
} else {
21.0
},
);
Some(("hist", 0))
}
})
.collect();
assert_eq!(
mixed,
[
Some(("float", 1.25_f64.to_bits())),
None,
Some(("hist", 0)),
Some(("hist", 0))
]
);
let histograms: Vec<_> = rows
.iter()
.map(|r| match r.fields.get(2).unwrap().sample(0) {
SampleRef::Histogram(h) => {
let expected = match r.time.timestamp(0).unwrap() {
10 => 101.0,
20 => 202.0,
40 => 303.0,
_ => panic!("unexpected histogram position"),
};
assert_histogram(h, expected);
Some(h.sum.to_bits())
}
SampleRef::Missing => None,
_ => panic!("unexpected observation sample variant"),
})
.collect();
assert_eq!(
histograms,
[
Some(101.0_f64.to_bits()),
Some(202.0_f64.to_bits()),
None,
Some(303.0_f64.to_bits())
]
);
assert!(!rows[0].is_last && !rows[1].is_last && !rows[2].is_last && rows[3].is_last);
}
#[tokio::test]
async fn merge_retains_sliced_float_time_and_histogram_owners_until_output_drop() {
let key = read_key("rack-slice", "sensor-9");
let times_owner = ScalarBuffer::from(vec![1, 11, 22, 99]);
let times = times_owner.slice(1, 2);
let floats_owner = Float64Array::from(vec![
Some(90.0),
Some(f64::from_bits(0x7ff8_0000_0000_0042)),
Some(-0.0),
Some(80.0),
]);
let floats = floats_owner.slice(1, 2);
let hist_owner = histograms(&[99.0, 12.5, 88.0]);
let hist_weak = Arc::downgrade(&hist_owner);
let sliced_hist =
Arc::new(HistogramBatch::try_from_arrow(hist_owner.as_arrow().slice(1, 1)).unwrap());
let hist_indices_owner = UInt32Array::from(vec![Some(9), Some(0), Some(0), Some(9)]);
let hist_indices = hist_indices_owner.slice(1, 2);
let mixed_floats_owner = Float64Array::from(vec![Some(90.0), Some(7.0), None, Some(80.0)]);
let mixed_floats = mixed_floats_owner.slice(1, 2);
let mixed_indices_owner = UInt32Array::from(vec![Some(9), None, Some(0), Some(9)]);
let mixed_indices = mixed_indices_owner.slice(1, 2);
let fields = SeriesFields::try_from_vec(vec![
Samples::from_float(floats.clone()),
Samples::try_from_mixed(
mixed_floats.clone(),
HistogramSamples::try_new(mixed_indices.clone(), sliced_hist.clone()).unwrap(),
)
.unwrap(),
Samples::try_from_histogram(hist_indices.clone(), sliced_hist.clone()).unwrap(),
])
.unwrap();
let chunk = RawSeriesChunk::try_new(
&payload_schema(),
key.clone(),
None,
RawTimestamps::try_new(times.clone()).unwrap(),
fields,
true,
)
.unwrap();
let time_ptr = times.as_ptr();
let key_ptr = key.as_ptr();
let float_ptr = floats.values().as_ptr();
let indices_ptr = hist_indices.values().as_ptr();
let mixed_float_ptr = mixed_floats.values().as_ptr();
let mixed_indices_ptr = mixed_indices.values().as_ptr();
let scalar_ptr = sliced_hist
.as_arrow()
.column(2)
.as_any()
.downcast_ref::<Float64Array>()
.unwrap()
.values()
.as_ptr();
let custom_ptr = sliced_hist
.as_arrow()
.column(7)
.as_any()
.downcast_ref::<ListArray>()
.unwrap()
.values()
.as_any()
.downcast_ref::<Float64Array>()
.unwrap()
.values()
.as_ptr();
let bucket_ptr = sliced_hist
.as_arrow()
.column(10)
.as_any()
.downcast_ref::<ListArray>()
.unwrap()
.values()
.as_any()
.downcast_ref::<Float64Array>()
.unwrap()
.values()
.as_ptr();
let batch_weak: Weak<HistogramBatch> = Arc::downgrade(&sliced_hist);
let interleaved = payload_chunk(
key.clone(),
&[15],
&[Some(6.0)],
&[None],
&[None],
&[Some(7.5)],
true,
);
let mut merged = merge_ordered_raw_series(vec![
payload_stream(label_source(1), vec![chunk]),
payload_stream(label_source(2), vec![interleaved]),
])
.unwrap();
assert_eq!(
merged
.schema()
.fields
.iter()
.map(|field| field.name.as_ref())
.collect::<Vec<_>>(),
["temperature", "latency", "observations"]
);
assert!(
matches!(merged.label_source(), RawLabelSource::Explicit),
"a merge is not a region scan, so it reports no scan-bound identity"
);
assert_eq!(merged.ordering(), SeriesOrdering::KeyThenTime);
let prefix = merged.next().await.unwrap().unwrap();
assert_eq!(prefix.time.as_buffer().as_ref(), &[11]);
assert_eq!(prefix.time.as_buffer().as_ptr(), time_ptr);
assert_eq!(prefix.primary_key.as_ptr(), key_ptr);
let interleaved = merged.next().await.unwrap().unwrap();
assert_eq!(interleaved.time.as_buffer().as_ref(), &[15]);
let remainder = merged.next().await.unwrap().unwrap();
assert_eq!(remainder.primary_key.as_ptr(), key_ptr);
assert_eq!(remainder.time.as_buffer().as_ref(), &[22]);
assert_eq!(
remainder.time.as_buffer().as_ptr(),
time_ptr.wrapping_add(1)
);
assert!(
matches!(prefix.fields.get(0).unwrap().dense().unwrap(), DenseSamplesRef::Float(a) if a.values().as_ptr() == float_ptr)
);
assert!(
matches!(remainder.fields.get(0).unwrap().dense().unwrap(), DenseSamplesRef::Float(a) if a.values().as_ptr() == float_ptr.wrapping_add(1))
);
let DenseSamplesRef::Mixed(prefix_mixed) = prefix.fields.get(1).unwrap().dense().unwrap()
else {
panic!("expected mixed payload")
};
assert_eq!(prefix_mixed.floats().values().as_ptr(), mixed_float_ptr);
assert_eq!(
prefix_mixed.histograms().indices().values().as_ptr(),
mixed_indices_ptr
);
assert!(
matches!(prefix.fields.get(1).unwrap().sample(0), SampleRef::Float(v) if v.to_bits() == 7.0_f64.to_bits())
);
let DenseSamplesRef::Mixed(remainder_mixed) = remainder.fields.get(1).unwrap().dense().unwrap()
else {
panic!("expected mixed payload")
};
assert_eq!(
remainder_mixed.floats().values().as_ptr(),
mixed_float_ptr.wrapping_add(1)
);
assert_eq!(
remainder_mixed.histograms().indices().values().as_ptr(),
mixed_indices_ptr.wrapping_add(1)
);
assert!(
matches!(remainder.fields.get(1).unwrap().sample(0), SampleRef::Histogram(h) if h.sum.to_bits() == 12.5_f64.to_bits())
);
let DenseSamplesRef::Histogram(prefix_hist) = prefix.fields.get(2).unwrap().dense().unwrap()
else {
panic!("expected histogram payload")
};
let DenseSamplesRef::Histogram(remainder_hist) =
remainder.fields.get(2).unwrap().dense().unwrap()
else {
panic!("expected histogram payload")
};
assert_eq!(prefix_hist.indices().values().as_ptr(), indices_ptr);
assert_eq!(prefix_hist.indices().values(), &[0]);
assert_histogram(prefix_hist.view_at(0).unwrap(), 12.5);
assert!(Arc::ptr_eq(prefix_hist.batch(), remainder_hist.batch()));
let output_arrow = prefix_hist.batch().as_arrow();
assert_eq!(
output_arrow
.column(2)
.as_any()
.downcast_ref::<Float64Array>()
.unwrap()
.values()
.as_ptr(),
scalar_ptr
);
assert_eq!(
output_arrow
.column(7)
.as_any()
.downcast_ref::<ListArray>()
.unwrap()
.values()
.as_any()
.downcast_ref::<Float64Array>()
.unwrap()
.values()
.as_ptr(),
custom_ptr
);
assert_eq!(
output_arrow
.column(10)
.as_any()
.downcast_ref::<ListArray>()
.unwrap()
.values()
.as_any()
.downcast_ref::<Float64Array>()
.unwrap()
.values()
.as_ptr(),
bucket_ptr
);
let negative_bucket_ptr = sliced_hist
.as_arrow()
.column(13)
.as_any()
.downcast_ref::<ListArray>()
.unwrap()
.values()
.as_any()
.downcast_ref::<Float64Array>()
.unwrap()
.values()
.as_ptr();
assert_eq!(
output_arrow
.column(13)
.as_any()
.downcast_ref::<ListArray>()
.unwrap()
.values()
.as_any()
.downcast_ref::<Float64Array>()
.unwrap()
.values()
.as_ptr(),
negative_bucket_ptr
);
assert_eq!(
remainder_hist.indices().values().as_ptr(),
indices_ptr.wrapping_add(1)
);
assert_eq!(remainder_hist.indices().values(), &[0]);
assert_histogram(remainder_hist.view_at(0).unwrap(), 12.5);
drop(merged);
drop(times_owner);
drop(times);
drop(floats_owner);
drop(floats);
drop(mixed_floats_owner);
drop(mixed_floats);
drop(mixed_indices_owner);
drop(mixed_indices);
drop(hist_owner);
drop(sliced_hist);
drop(hist_indices_owner);
drop(hist_indices);
assert!(hist_weak.upgrade().is_none());
assert!(batch_weak.upgrade().is_some());
assert_eq!(prefix.time.as_buffer().as_ref(), &[11]);
assert_eq!(remainder.time.as_buffer().as_ref(), &[22]);
let SampleRef::Histogram(prefix_histogram) = prefix.fields.get(2).unwrap().sample(0) else {
panic!("expected histogram");
};
assert_histogram(prefix_histogram, 12.5);
let SampleRef::Histogram(remainder_histogram) = remainder.fields.get(2).unwrap().sample(0)
else {
panic!("expected histogram");
};
assert_histogram(remainder_histogram, 12.5);
let SampleRef::Histogram(mixed_histogram) = remainder.fields.get(1).unwrap().sample(0) else {
panic!("expected histogram");
};
assert_histogram(mixed_histogram, 12.5);
assert!(
matches!(prefix.fields.get(0).unwrap().sample(0), SampleRef::Float(v) if v.to_bits() == 0x7ff8_0000_0000_0042)
);
assert!(
matches!(remainder.fields.get(0).unwrap().sample(0), SampleRef::Float(v) if v.to_bits() == (-0.0_f64).to_bits())
);
drop(prefix);
assert!(batch_weak.upgrade().is_some());
let SampleRef::Histogram(remainder_histogram) = remainder.fields.get(2).unwrap().sample(0)
else {
panic!("expected histogram");
};
assert_histogram(remainder_histogram, 12.5);
assert!(
matches!(remainder.fields.get(1).unwrap().sample(0), SampleRef::Histogram(h) if h.sum.to_bits() == 12.5_f64.to_bits() && h.custom_values == [1.5, 4.5] && h.positive.span_offsets == [1] && h.positive.span_lengths == [2] && h.positive.buckets == [3.25, 7.5] && h.negative.span_offsets == [-2] && h.negative.span_lengths == [1] && h.negative.buckets == [0.75])
);
drop(remainder);
drop(interleaved);
assert!(batch_weak.upgrade().is_none());
}
src/query/src/datafusion.rs
+133 -13
View File
@@ -30,7 +30,9 @@ use common_error::ext::BoxedError;
use common_function::function::FunctionContext;
use common_function::function_factory::ScalarFunctionFactory;
use common_query::{Output, OutputData, OutputMeta};
use common_recordbatch::adapter::{RecordBatchStreamAdapter, RegionQueryStatCounters};
use common_recordbatch::adapter::{
RecordBatchStreamAdapter, RegionQueryStatCounters, SERIES_SCAN_EXEC_NAME,
};
use common_recordbatch::{EmptyRecordBatchStream, SendableRecordBatchStream};
use common_telemetry::tracing;
use datafusion::catalog::TableFunction;
@@ -50,7 +52,7 @@ use snafu::{OptionExt, ResultExt, ensure};
use sqlparser::ast::AnalyzeFormat;
use table::TableRef;
use table::requests::{DeleteRequest, InsertRequest};
use table::table::scan::{REGION_SCAN_EXEC_NAME, RegionScanExec};
use table::table::scan::{REGION_SCAN_EXEC_NAME, RegionScanExec, SeriesScanExec};
use tracing::Span;
use crate::analyze::DistAnalyzeExec;
@@ -87,10 +89,16 @@ fn query_load_region_id(plan: &Arc<dyn ExecutionPlan>) -> Option<u64> {
let mut stack = vec![plan.clone()];
while let Some(plan) = stack.pop() {
if plan.name() == REGION_SCAN_EXEC_NAME
&& let Some(scan) = plan.downcast_ref::<RegionScanExec>()
&& let Some(scan_region_id) = scan.query_load_region_id()
{
let scan_region_id = if plan.name() == REGION_SCAN_EXEC_NAME {
plan.downcast_ref::<RegionScanExec>()
.and_then(RegionScanExec::query_load_region_id)
} else if plan.name() == SERIES_SCAN_EXEC_NAME {
plan.downcast_ref::<SeriesScanExec>()
.and_then(SeriesScanExec::query_load_region_id)
} else {
None
};
if let Some(scan_region_id) = scan_region_id {
match region_id {
Some(region_id) if region_id != scan_region_id => return None,
Some(_) => {}
@@ -108,7 +116,7 @@ fn query_load_region_id(plan: &Arc<dyn ExecutionPlan>) -> Option<u64> {
// Unlike the Prometheus read-load reporting in `MergeScanExec`, the heartbeat
// counters must be updated before metrics leave the datanode process. The
// `RecordBatchStreamAdapter` that resolves `RecordBatchMetrics` does not know
// the owning `MitoRegion`, so we extract the counters from `RegionScanExec` and
// the owning `MitoRegion`, so we extract the counters from the scan execs and
// pass them to the adapter. If a plan contains scans from different regions,
// return `None` to avoid charging the whole plan metrics to one region.
fn query_stat_counters(plan: &Arc<dyn ExecutionPlan>) -> Option<RegionQueryStatCounters> {
@@ -116,10 +124,16 @@ fn query_stat_counters(plan: &Arc<dyn ExecutionPlan>) -> Option<RegionQueryStatC
let mut stack = vec![plan.clone()];
while let Some(plan) = stack.pop() {
if plan.name() == REGION_SCAN_EXEC_NAME
&& let Some(scan) = plan.downcast_ref::<RegionScanExec>()
&& let Some(scan_counters) = scan.query_stat_counters()
{
let scan_counters = if plan.name() == REGION_SCAN_EXEC_NAME {
plan.downcast_ref::<RegionScanExec>()
.and_then(RegionScanExec::query_stat_counters)
} else if plan.name() == SERIES_SCAN_EXEC_NAME {
plan.downcast_ref::<SeriesScanExec>()
.and_then(SeriesScanExec::query_stat_counters)
} else {
None
};
if let Some(scan_counters) = scan_counters {
match &counters {
Some(counters)
if !Arc::ptr_eq(&counters.query_cpu_time, &scan_counters.query_cpu_time)
@@ -758,7 +772,7 @@ mod tests {
};
use datafusion::physical_plan::display::{DisplayAs, DisplayFormatType};
use datafusion::physical_plan::expressions::PhysicalSortExpr;
use datafusion::physical_plan::joins::{HashJoinExec, JoinOn, PartitionMode};
use datafusion::physical_plan::joins::{CrossJoinExec, HashJoinExec, JoinOn, PartitionMode};
use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet;
use datafusion::physical_plan::{ExecutionPlan, PhysicalExpr};
use datafusion::prelude::{col, lit};
@@ -770,7 +784,8 @@ mod tests {
use session::context::{QueryContext, QueryContextBuilder};
use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder, RegionMetadataRef};
use store_api::region_engine::{
PartitionRange, PrepareRequest, QueryScanContext, RegionScanner, ScannerProperties,
PartitionRange, PrepareRequest, QueryScanContext, RawLabelSource, RegionScanner,
ScannerProperties,
};
use store_api::storage::{RegionId, ScanRequest};
use table::metadata::{TableInfoBuilder, TableMetaBuilder};
@@ -844,6 +859,17 @@ mod tests {
true
}
fn raw_series_schema(
&self,
_label_source: &RawLabelSource,
) -> std::result::Result<Option<Arc<datatypes::series::samples::SeriesSchema>>, BoxedError>
{
Ok(Some(Arc::new(datatypes::series::samples::SeriesSchema {
fields: Arc::from([datatypes::series::samples::FieldSpec { name: "ts".into() }]),
eval_grid: None,
})))
}
fn add_dyn_filter_to_predicate(
&mut self,
filter_exprs: Vec<Arc<dyn PhysicalExpr>>,
@@ -885,6 +911,13 @@ mod tests {
query_load_region_id: Option<RegionId>,
query_stat_counters: Option<RegionQueryStatCounters>,
) -> Arc<dyn ExecutionPlan> {
build_region_scan_exec(query_load_region_id, query_stat_counters)
}
fn build_region_scan_exec(
query_load_region_id: Option<RegionId>,
query_stat_counters: Option<RegionQueryStatCounters>,
) -> Arc<RegionScanExec> {
let schema = Arc::new(datatypes::schema::Schema::new(vec![ColumnSchema::new(
"ts",
ConcreteDataType::timestamp_millisecond_datatype(),
@@ -921,6 +954,23 @@ mod tests {
Arc::new(RegionScanExec::new(Box::new(scanner), ScanRequest::default(), None).unwrap())
}
fn build_series_scan(
query_load_region_id: Option<RegionId>,
query_stat_counters: Option<RegionQueryStatCounters>,
) -> Arc<dyn ExecutionPlan> {
build_region_scan_exec(query_load_region_id, query_stat_counters)
.try_to_series_scan(RawLabelSource::Explicit)
.unwrap()
.expect("RecordingScanner supports a native series scan")
}
fn build_mixed_scan_plan(
left: Arc<dyn ExecutionPlan>,
right: Arc<dyn ExecutionPlan>,
) -> Arc<dyn ExecutionPlan> {
Arc::new(CrossJoinExec::new(left, right))
}
fn query_stat_counters_for_test() -> RegionQueryStatCounters {
RegionQueryStatCounters {
query_cpu_time: Arc::new(AtomicU64::new(0)),
@@ -928,6 +978,33 @@ mod tests {
}
}
#[test]
fn query_load_region_id_reads_native_series_scan() {
let region_id = RegionId::new(1024, 42);
let scan = build_series_scan(Some(region_id), None);
assert_eq!(super::query_load_region_id(&scan), Some(region_id.as_u64()));
}
#[test]
fn query_load_region_id_accepts_mixed_scans_for_same_region() {
let region_id = RegionId::new(1024, 42);
let left = build_query_load_region_scan(Some(region_id));
let right = build_series_scan(Some(region_id), None);
let plan = build_mixed_scan_plan(left, right);
assert_eq!(super::query_load_region_id(&plan), Some(region_id.as_u64()));
}
#[test]
fn query_load_region_id_rejects_mixed_region_scans() {
let left = build_query_load_region_scan(Some(RegionId::new(1024, 42)));
let right = build_series_scan(Some(RegionId::new(1024, 43)), None);
let plan = build_mixed_scan_plan(left, right);
assert!(super::query_load_region_id(&plan).is_none());
}
#[test]
fn query_load_region_id_ignores_scans_without_region_id() {
let query_load_region_id = RegionId::new(1024, 42);
@@ -958,6 +1035,49 @@ mod tests {
);
}
#[test]
fn query_stat_counters_reads_native_series_scan() {
let counters = query_stat_counters_for_test();
let scan = build_series_scan(None, Some(counters.clone()));
let actual = super::query_stat_counters(&scan).unwrap();
assert!(Arc::ptr_eq(
&actual.query_cpu_time,
&counters.query_cpu_time
));
assert!(Arc::ptr_eq(
&actual.query_scanned_bytes,
&counters.query_scanned_bytes
));
}
#[test]
fn query_stat_counters_returns_shared_counter_for_mixed_scan_plan() {
let counters = query_stat_counters_for_test();
let left = build_query_stat_counter_region_scan(counters.clone());
let right = build_series_scan(None, Some(counters.clone()));
let plan = build_mixed_scan_plan(left, right);
let actual = super::query_stat_counters(&plan).unwrap();
assert!(Arc::ptr_eq(
&actual.query_cpu_time,
&counters.query_cpu_time
));
assert!(Arc::ptr_eq(
&actual.query_scanned_bytes,
&counters.query_scanned_bytes
));
}
#[test]
fn query_stat_counters_rejects_mixed_counter_plan_with_series_scan() {
let left = build_query_stat_counter_region_scan(query_stat_counters_for_test());
let right = build_series_scan(None, Some(query_stat_counters_for_test()));
let plan = build_mixed_scan_plan(left, right);
assert!(super::query_stat_counters(&plan).is_none());
}
#[test]
fn query_stat_counters_returns_shared_counter_for_multi_scan_plan() {
let counters = query_stat_counters_for_test();
src/query/src/dist_plan/merge_scan.rs
+2 -289
View File
@@ -69,7 +69,6 @@ use crate::dist_plan::dyn_filter_bridge::{
query_context_with_refreshed_initial_dyn_filter_regs,
register_dyn_filter_subscribers_for_region, register_remote_dyn_filters,
};
use crate::dist_plan::planner::CaptureNativeRawRoute;
use crate::dist_plan::{
FilterId, RemoteDynFilterProducerId, RemoteDynFilterRegistryLease, Subscriber,
};
@@ -495,32 +494,6 @@ pub struct MergeScanExec {
captured_remote_dyn_filters: Arc<Mutex<Vec<CapturedDynFilter>>>,
target_partition: usize,
enable_per_region_metrics: bool,
/// Optional route provenance captured by the distributed planner from the
/// same physical route that selected `regions`. Ordinary execution never
/// reads it, so `None` keeps every existing plan unchanged.
native_raw_route: Option<CaptureNativeRawRoute>,
}
/// Borrowed handoff view over the ordinary plan state a native raw consumer
/// needs to read the captured route.
///
/// It borrows the existing [`MergeScanExec`] fields instead of duplicating them;
/// the output schema is available from the public `ExecutionPlan::schema()` of
/// the plan this source came from.
// TODO(narthex): remove this allowance once the remote native PromQL planner
// (#9670) consumes the captured route; this PR only captures it.
#[allow(dead_code)]
pub(crate) struct NativeRawSource<'a> {
/// Frozen route provenance in ordinary region order.
pub(crate) route: &'a CaptureNativeRawRoute,
/// The exact ordinary input logical plan sent to every region.
pub(crate) plan: &'a LogicalPlan,
/// The table the ordinary plan scans.
pub(crate) table: &'a TableName,
/// The handler that performs the region reads.
pub(crate) region_query_handler: &'a RegionQueryHandlerRef,
/// The query context of the ordinary execution.
pub(crate) query_ctx: &'a QueryContextRef,
}
impl std::fmt::Debug for MergeScanExec {
@@ -604,43 +577,6 @@ impl MergeScanExec {
captured_remote_dyn_filters: Arc::default(),
target_partition,
enable_per_region_metrics,
native_raw_route: None,
})
}
/// Attaches the native raw route captured from the same physical route that
/// selected this plan's regions.
///
/// The artifact is provenance for a native raw consumer; ordinary execution
/// never reads it.
pub(crate) fn with_native_raw_route(mut self, native_raw_route: CaptureNativeRawRoute) -> Self {
self.native_raw_route = Some(native_raw_route);
self
}
/// Returns the frozen region targets captured for this plan together with
/// the ordinary execution state a native raw consumer needs to read them.
///
/// Returns `None` unless the planner proved the route and ordinary region
/// dispatch stays unchanged: a scan that already captured remote dynamic
/// filters, or a flow sink scan, keeps its ordinary path, so a consumer must
/// fall back as well. An assigned producer id alone is not active intent,
/// because the analyzer assigns one to every merge scan.
// TODO(narthex): remove this allowance once the remote native PromQL planner
// (#9670) consumes the captured route; this PR only captures it.
#[allow(dead_code)]
pub(crate) fn native_raw_source(&self) -> Option<NativeRawSource<'_>> {
if !self.captured_remote_dyn_filters.lock().unwrap().is_empty() || self.is_flow_sink_scan()
{
return None;
}
let route = self.native_raw_route.as_ref()?;
Some(NativeRawSource {
route,
plan: &self.plan,
table: &self.table,
region_query_handler: &self.region_query_handler,
query_ctx: &self.query_ctx,
})
}
@@ -1431,7 +1367,7 @@ mod tests {
use datatypes::vectors::{Int64Vector, StringVector, TimestampMillisecondVector};
use futures_util::{Stream, TryStreamExt};
use session::ReadPreference;
use session::context::{QueryContext, QueryContextBuilder};
use session::context::QueryContext;
use session::query_id::QueryId;
use snafu::IntoError;
use table::table::scan::REGION_SCAN_EXEC_NAME;
@@ -1441,7 +1377,7 @@ mod tests {
use super::*;
use crate::dist_plan::DynFilterRegistryManager;
use crate::options::{FLOW_SINK_TABLE_ID, QueryOptions};
use crate::options::QueryOptions;
use crate::query_engine::{QueryEngineContext, QueryEngineState};
use crate::region_query::RegionQueryHandler;
@@ -1610,217 +1546,6 @@ mod tests {
);
}
#[test]
fn merge_scan_ordinary_construction_has_no_native_raw_route() {
// The artifact is provenance only: an ordinary plan carries none and
// executes exactly as before.
let exec = merge_scan_exec_with_sorted_input(1, 1);
assert!(exec.native_raw_source().is_none());
}
#[test]
fn merge_scan_keeps_captured_native_raw_route() {
let region_id = RegionId::new(1024, 1);
let exec =
merge_scan_exec_with_sorted_input(1, 1).with_native_raw_route(CaptureNativeRawRoute {
targets: vec![(region_id, test_target(1))],
});
let source = exec.native_raw_source().unwrap();
assert_eq!(vec![(region_id, test_target(1))], source.route.targets);
// The bundle borrows the ordinary plan state, not a copy of it.
assert_eq!(1, source.plan.schema().fields().len());
assert_eq!("table", source.table.table_name);
assert_eq!(1, exec.partition_count());
}
#[test]
fn merge_scan_declines_native_raw_route_for_captured_dyn_filters() {
// An assigned remote dynamic-filter producer id is identity, not intent:
// the analyzer assigns one to every merge scan, so a plan that captured no
// filters keeps the handoff view.
let region_id = RegionId::new(1024, 1);
let exec = remote_dyn_filter_test_exec(
Arc::new(TestRegionQueryHandler::default()),
QueryContext::arc(),
)
.with_native_raw_route(CaptureNativeRawRoute {
targets: vec![(region_id, test_target(1))],
});
assert!(exec.captured_remote_dyn_filters().is_empty());
assert!(exec.native_raw_source().is_some());
// A filter this producer accepted is really captured, and only then does the
// plan keep its ordinary dispatch.
install_remote_dyn_filter(&exec);
assert_eq!(1, exec.captured_remote_dyn_filters().len());
assert!(exec.native_raw_source().is_none());
// A rewrite of the same scan keeps the frozen route and the captured filters,
// so the rewritten plan still declines the handoff view.
let plan: Arc<dyn ExecutionPlan> = Arc::new(exec);
let replaced = plan
.replace_children(
vec![],
datafusion::physical_plan::ReplaceChildrenOptions::new(
datafusion::physical_plan::ChildrenPropertiesMode::Recompute,
),
)
.unwrap();
let exec = replaced.downcast_ref::<MergeScanExec>().unwrap();
assert_eq!(
vec![(region_id, test_target(1))],
exec.native_raw_route.as_ref().unwrap().targets
);
assert_eq!(1, exec.captured_remote_dyn_filters().len());
assert!(exec.native_raw_source().is_none());
}
#[tokio::test]
async fn merge_scan_execution_ignores_native_raw_route() {
// Attaching the artifact must not change what an ordinary execution
// reads: every region is still dispatched through the handler's own
// target selection, with the ordinary read preference, and the same
// batches come back in the same order.
let regions = vec![
RegionId::new(1024, 2),
RegionId::new(1024, 1),
RegionId::new(1024, 3),
];
let batch_for = |region_number: u32| {
let base = i64::from(region_number) * 10;
record_batch(
int64_schema(&["a", "b"]),
vec![
Arc::new(Int64Vector::from_slice([base + 1])) as _,
Arc::new(Int64Vector::from_slice([base + 2])) as _,
],
)
};
let responses = || -> Vec<(RegionId, RecordBatch)> {
regions
.iter()
.map(|region_id| (*region_id, batch_for(region_id.region_number())))
.collect()
};
// Region striping keeps the selection order, so this is the whole output.
let expected = [(21, 22), (11, 12), (31, 32)];
let baseline_handler = Arc::new(TestRegionQueryHandler::new(responses()));
let baseline = collect_merge_scan(merge_scan_exec_with_handler(
regions.clone(),
expected_int64_schema(),
baseline_handler.clone(),
1,
))
.await
.unwrap();
assert_eq!(expected.len(), baseline.len());
for (batch, values) in baseline.iter().zip(expected) {
assert_int64_batch(batch, values);
}
// Deliberately different peers, and in a deliberately different order:
// an execution that read the artifact would dispatch to these.
let handler = Arc::new(TestRegionQueryHandler::new(responses()));
let with_route = collect_merge_scan(
merge_scan_exec_with_handler(
regions.clone(),
expected_int64_schema(),
handler.clone(),
1,
)
.with_native_raw_route(CaptureNativeRawRoute {
targets: vec![
(regions[2], test_target(7)),
(regions[0], test_target(8)),
(regions[1], test_target(9)),
],
}),
)
.await
.unwrap();
assert_eq!(baseline, with_route);
assert_eq!(expected.len(), with_route.len());
for (batch, values) in with_route.iter().zip(expected) {
assert_int64_batch(batch, values);
}
for exec_handler in [&baseline_handler, &handler] {
assert_eq!(regions.len(), exec_handler.select_target_calls());
assert_eq!(
vec![ReadPreference::Leader; regions.len()],
exec_handler.read_preferences()
);
// The ordinary plan dispatches to the target its handler selected,
// never to a peer from the artifact.
assert_eq!(vec![1; regions.len()], exec_handler.do_get_targets());
}
}
#[test]
fn merge_scan_with_new_children_preserves_native_raw_route() {
let region_id = RegionId::new(1024, 1);
let plan: Arc<dyn ExecutionPlan> = Arc::new(
merge_scan_exec_with_sorted_input(1, 1).with_native_raw_route(CaptureNativeRawRoute {
targets: vec![(region_id, test_target(3))],
}),
);
let replaced = plan
.replace_children(
vec![],
datafusion::physical_plan::ReplaceChildrenOptions::new(
datafusion::physical_plan::ChildrenPropertiesMode::Recompute,
),
)
.unwrap();
let exec = replaced.downcast_ref::<MergeScanExec>().unwrap();
assert_eq!(
vec![(region_id, test_target(3))],
exec.native_raw_source().unwrap().route.targets
);
}
#[test]
fn merge_scan_declines_native_raw_route_for_flow_sink_scan() {
// A flow sink read keeps its ordinary dispatch, so the handoff view is
// not offered even though the artifact exists for these regions.
let region_id = RegionId::new(1024, 1);
let query_ctx = Arc::new(
QueryContextBuilder::default()
.set_extension(
FLOW_SINK_TABLE_ID.to_string(),
region_id.table_id().to_string(),
)
.build(),
);
let plan = LogicalPlanBuilder::empty(true).build().unwrap();
let schema = plan.schema().as_arrow().clone();
let exec = MergeScanExec::new(
&SessionStateBuilder::new().build(),
TableName::new("catalog", "schema", "table"),
vec![region_id],
plan,
&schema,
Arc::new(TestRegionQueryHandler::default()),
query_ctx,
1,
None,
false,
)
.unwrap()
.with_native_raw_route(CaptureNativeRawRoute {
targets: vec![(region_id, test_target(1))],
});
assert!(exec.is_flow_sink_scan());
assert!(exec.native_raw_source().is_none());
}
#[test]
fn merge_scan_advertises_ordering_when_each_partition_reads_at_most_one_region() {
let exec = merge_scan_exec_with_sorted_input(3, 3);
@@ -2398,18 +2123,6 @@ mod tests {
..Default::default()
}
}
fn select_target_calls(&self) -> usize {
self.select_target_calls.load(Ordering::SeqCst)
}
fn read_preferences(&self) -> Vec<ReadPreference> {
self.read_preferences.lock().unwrap().clone()
}
fn do_get_targets(&self) -> Vec<u64> {
self.do_get_targets.lock().unwrap().clone()
}
}
struct TestRecordBatchStream {
src/query/src/dist_plan/planner.rs
+18 -1027
View File
File diff suppressed because it is too large Load Diff
src/query/src/options.rs
+4 -2
View File
@@ -110,8 +110,10 @@ pub struct QueryOptions {
/// - `fair`: Fair divides memory available after unspillable reservations
/// evenly among spillable reservations and may trigger earlier spills.
pub experimental_memory_pool_policy: QueryMemoryPoolPolicy,
/// Enables the experimental native PromQL execution path for compatible
/// instant queries over shared Metric raw series. Default: `false`.
/// Enables the experimental native PromQL execution path for compatible queries. This
/// path currently applies only to standalone queries using in-process region transport,
/// and also requires the per-request `query.experimental_native_promql=true` extension.
/// Default: `false`.
pub experimental_native_promql: bool,
}
src/query/src/promql/native_compile.rs
+5 -7
View File
@@ -15,8 +15,8 @@
//! Compiles a PromQL AST into a [`NativeTree`].
//!
//! This is the front half of the native evaluation model: it reads the parsed statement and writes
//! down what it means as bound nodes. It runs nothing, and nothing calls it yet, so the subset it
//! accepts is one recursive rule:
//! down what it means as bound nodes. The native planner uses this compiler to build a tree for
//! standalone execution. Its accepted expressions are described by the recursive rule below:
//!
//! ```text
//! <expression> one of, recursively:
@@ -718,11 +718,9 @@ fn binary_op(op: &TokenType) -> Option<BinaryOp> {
/// Compiles an aggregation into an aggregate node over the expression it groups.
///
/// Every aggregator that folds a group's values into one value per grid point binds, whether or not
/// its kernel exists yet: `sum`, `avg`, `count`, `min`, `max`, `group`, `stddev`, `stdvar` and
/// `quantile`. Binding a shape the compiler can describe and letting the runtime refuse a kind it
/// cannot execute keeps this the one place that says what an aggregate means, so a kernel landing
/// later changes no binding.
/// The aggregators that fold a group's values into one value per grid point bind to an aggregate
/// node: `sum`, `avg`, `count`, `min`, `max`, `group`, `stddev`, `stdvar` and `quantile`. The
/// runtime validates whether the resulting tree can execute.
///
/// The rest of the parser's aggregates are a different shape and bind as a vector operation
/// ([`vector_aggregate`]) rather than as a kind that would compute something else: `topk`,
src/query/src/promql/native_compile/tests.rs
+39
View File
@@ -4645,6 +4645,45 @@ fn timestamp_multi_field_input_is_not_supported_natively() {
.unwrap_or_else(|error| panic!("`{query}` over one field was refused: {error}"));
}
}
#[test]
fn timestamp_anchored_selector_retains_offset_and_ordinary_selection_keeps_it() {
for (query, expected_offset, expected_anchor, expected_bounds) in [
(
"timestamp(x @ 50 offset 60s)",
60_000,
Some(50_000),
(-15_000, -10_000),
),
("timestamp(x offset 2s)", 2_000, None, (-7_000, 98_000)),
(
"x @ 50 offset 60s",
60_000,
Some(50_000),
(-15_000, -10_000),
),
] {
let compiled = compile_selectors(
&eval_stmt(query),
&|_| Ok(raw_schema()),
SeriesOrdering::KeyThenTime,
)
.unwrap_or_else(|error| panic!("{query}: {error}"));
let requirement = &compiled.inputs[0].requirement;
assert_eq!(requirement.offset_ms, expected_offset, "{query}");
assert_eq!(requirement.at_ms, expected_anchor, "{query}");
let (start, end, _) = requirement.raw_scan_bounds().unwrap();
assert_eq!((start, end), expected_bounds, "{query}");
if query.starts_with("timestamp(") && expected_anchor.is_some() {
let NativeExpr::Instant { op, .. } = &compiled.tree.root().expr else {
panic!("{query} did not compile as instant timestamp");
};
assert_eq!(op.offset_ms, expected_offset, "{query}");
assert_eq!(op.at_ms, expected_anchor, "{query}");
}
}
}
#[test]
fn selector_requirements_capture_bounds_and_source_occurrence() {
let compiled = compile_selectors(
src/query/src/promql/planner/native.rs
+36 -4
View File
@@ -969,11 +969,13 @@ fn collect_selectors<'a>(expr: &'a PromExpr, out: &mut Vec<&'a VectorSelector>)
match expr {
PromExpr::VectorSelector(selector) => out.push(selector),
PromExpr::MatrixSelector(MatrixSelector { vs, .. }) => out.push(vs),
PromExpr::Paren(ParenExpr { expr }) => return collect_selectors(expr, out),
PromExpr::Paren(ParenExpr { expr }) => {
return collect_selectors(expr, out);
}
PromExpr::Call(call)
if matches!(
call.func.name,
"rate" | "increase" | "sum_over_time" | "abs"
"rate" | "increase" | "sum_over_time" | "abs" | "timestamp"
) && call.args.args.len() == 1 =>
{
return collect_selectors(&call.args.args[0], out);
@@ -983,7 +985,9 @@ fn collect_selectors<'a>(expr: &'a PromExpr, out: &mut Vec<&'a VectorSelector>)
}
// Literals are permitted only as compiler-recognized broadcasts in a bounded binary tree.
PromExpr::NumberLiteral(_) => return true,
PromExpr::Unary(unary) => return collect_selectors(&unary.expr, out),
PromExpr::Unary(unary) => {
return collect_selectors(&unary.expr, out);
}
PromExpr::Call(call) => {
if matches!(call.func.name, "label_replace" | "label_join") {
return false;
@@ -994,7 +998,9 @@ fn collect_selectors<'a>(expr: &'a PromExpr, out: &mut Vec<&'a VectorSelector>)
}
}
}
PromExpr::Aggregate(aggregate) => return collect_selectors(&aggregate.expr, out),
PromExpr::Aggregate(aggregate) => {
return collect_selectors(&aggregate.expr, out);
}
PromExpr::Subquery(_) => return false,
_ => {}
}
@@ -1133,6 +1139,32 @@ mod tests {
.collect()
}
#[test]
fn selector_walk_keeps_plain_selector_identity() {
for (query, expected) in [
("timestamp((x @ 2 offset -1s))", vec!["x"]),
("abs(timestamp(x @ 2 offset -1s))", vec!["x"]),
("timestamp(x @ start() offset 1s)", vec!["x"]),
("timestamp(x @ end() offset -1s)", vec!["x"]),
(
"timestamp(x @ 2 offset 1s) + y @ 2 offset 1s",
vec!["x", "y"],
),
] {
let expr = promql_parser::parser::parse(query).unwrap();
let mut selectors = Vec::new();
assert!(collect_selectors(&expr, &mut selectors));
assert_eq!(
selectors
.iter()
.map(|s| s.name.as_deref().unwrap())
.collect::<Vec<_>>(),
expected,
"{query}"
);
}
}
#[test]
fn collect_selectors_descends_into_aggregate_inputs() {
for query in [
src/query/src/query_engine/state.rs
-3
View File
@@ -313,7 +313,6 @@ impl QueryEngineState {
partition_rule_manager,
region_query_handler.clone(),
options.enable_per_region_metrics,
options.experimental_native_promql,
)))
.with_optimizer_rules(optimizer.rules)
.with_physical_optimizer_rules(physical_optimizer.rules)
@@ -626,7 +625,6 @@ impl DfQueryPlanner {
partition_rule_manager: Option<PartitionRuleManagerRef>,
region_query_handler: Option<RegionQueryHandlerRef>,
enable_per_region_metrics: bool,
capture_native_raw_route: bool,
) -> Self {
let mut planners: Vec<Arc<dyn ExtensionPlanner + Send + Sync>> = vec![
Arc::new(PromExtensionPlanner),
@@ -641,7 +639,6 @@ impl DfQueryPlanner {
partition_rule_manager,
region_query_handler,
enable_per_region_metrics,
capture_native_raw_route,
)));
planners.push(Arc::new(MergeSortExtensionPlanner {}));
}
src/query/tests/native_instant_at.rs
+129
View File
@@ -179,6 +179,33 @@ fn native_instant_at_timestamp_values_and_labels() {
run("timestamp(abs(a @ 1))", 0, 4000, 3000, &values, &times, 1)[0].0,
[0.0, 1.0, 2.0, 3.0, 4.0]
);
// The raw source axis is deliberately unpruned: timestamp selection must enforce both
// windows itself. At @11s offset -2s with 5s lookback, 10.5s is eligible but 11.5s is not.
assert_eq!(
run(
"timestamp(a @ 11 offset -2s)",
11_000,
11_000,
5_000,
&[Some(10_500.0), Some(11_500.0)],
&[10_500, 11_500],
1,
)[0]
.0,
[10.5]
);
assert!(
run(
"timestamp(a @ 50 offset 60s)",
50_000,
50_000,
30_000,
&[Some(10_000.0)],
&[10_000],
1,
)
.is_empty()
);
assert_eq!(
run("a @ -1.5", 2000, 6000, 1000, &values, &times, 1)[0].0,
[7.0; 5]
@@ -204,6 +231,108 @@ fn native_instant_at_timestamp_values_and_labels() {
);
}
#[test]
fn native_anchored_timestamp_selection_boundaries() {
let stale = f64::from_bits(common_query::prometheus::PROMETHEUS_STALE_NAN_BITS);
let all_times = [-1500, 500, 1500, 3500];
let all_values = [Some(7.0), Some(1.0), Some(2.0), Some(3.0)];
type Case<'a> = (&'a str, u64, &'a [i64], &'a [Option<f64>], &'a [f64]);
let cases: &[Case<'_>] = &[
(
"timestamp(a @ 2 offset 1s)",
5_000,
&all_times,
&all_values,
&[0.5],
),
(
"timestamp(a @ 2 offset -1s)",
5_000,
&all_times,
&all_values,
&[1.5],
),
(
"timestamp(a @ 2 offset 2s)",
750,
&all_times,
&all_values,
&[],
),
("timestamp(a @ 0.5)", 0, &all_times, &all_values, &[0.5]),
("timestamp(a @ 1)", 0, &all_times, &all_values, &[]),
(
"timestamp(a @ 2 offset 1s)",
1_000,
&all_times,
&all_values,
&[],
),
// The lower endpoint is exclusive for either sign of offset.
(
"timestamp(a @ 2 offset 1s)",
1_500,
&[500],
&[Some(1.0)],
&[],
),
(
"timestamp(a @ 2 offset -1s)",
1_500,
&[1500],
&[Some(2.0)],
&[],
),
// The upper endpoint is inclusive.
(
"timestamp(a @ 2 offset 1s)",
5_000,
&[1000],
&[Some(4.0)],
&[1.0],
),
(
"timestamp(a @ 2 offset -1s)",
5_000,
&[2000, 3000],
&[Some(2.0), Some(3.0)],
&[2.0],
),
// A stale marker at the eligible upper endpoint suppresses the preceding value.
(
"timestamp(a @ 2 offset 1s)",
5_000,
&[500, 1000],
&[Some(1.0), Some(stale)],
&[],
),
// A stale marker above the intersection's upper endpoint is ignored.
(
"timestamp(a @ 2 offset 1s)",
5_000,
&[1000, 1500],
&[Some(4.0), Some(stale)],
&[1.0],
),
// A stale marker inside the eligible window prevents fallback to an older finite sample.
(
"timestamp(a @ 2 offset -1s)",
5_000,
&[500, 1500],
&[Some(1.0), Some(stale)],
&[],
),
];
for (query, lookback, times, values, expected) in cases {
let actual = run(query, 2_000, 2_000, *lookback, values, times, 1)
.into_iter()
.flat_map(|series| series.0)
.collect::<Vec<_>>();
assert_eq!(actual, *expected, "{query} with lookback {lookback}");
}
}
/// Two source series that differ only by a metadata label `timestamp()` drops must not be emitted
/// as two rows of one label set: the second is refused at the position it collides on, which is the
/// same rule a map follows when it drops the metric name.
src/store-api/src/region_engine.rs
+6 -3
View File
@@ -655,8 +655,8 @@ pub trait RegionScanner: Debug + DisplayAs + Send {
///
/// Legacy raw sources have no ordering declaration, so `Contiguous` is the conservative
/// default; `Interleaved` is not representable on the raw frame wire.
fn raw_series_ordering(&self) -> SeriesOrdering {
SeriesOrdering::Contiguous
fn raw_series_ordering(&self) -> Result<SeriesOrdering, BoxedError> {
Ok(SeriesOrdering::Contiguous)
}
/// Scans a partition as raw series when this source supports that path.
@@ -1018,7 +1018,10 @@ mod tests {
.unwrap()
.is_none()
);
assert_eq!(scanner.raw_series_ordering(), SeriesOrdering::Contiguous);
assert_eq!(
scanner.raw_series_ordering().unwrap(),
SeriesOrdering::Contiguous
);
// SinglePartitionScanner uses the default raw-scan implementation, which panics.
let raw_scan = catch_unwind(AssertUnwindSafe(|| {
src/table/src/table/scan.rs
+96 -143
View File
@@ -20,7 +20,6 @@ use std::time::Instant;
use api::v1::SemanticType;
use common_error::ext::BoxedError;
use common_query::request::{RawSeriesRequest, RawSeriesTag};
use common_query::series_frame::{pack_raw_series, raw_series_frame_schema_with_ordering};
use common_recordbatch::adapter::{
RecordBatchStreamAdapter, RegionQueryStatCounters, SERIES_SCAN_EXEC_NAME,
@@ -57,11 +56,11 @@ use datafusion_physical_expr::{
use datatypes::arrow::datatypes::SchemaRef as ArrowSchemaRef;
use datatypes::compute::SortOptions;
use datatypes::series::SeriesOrdering;
use datatypes::series::samples::SeriesSchema;
use datatypes::series::samples::{FieldSpec, SeriesSchema};
use futures::{Stream, StreamExt};
use store_api::metric_engine_consts::DATA_SCHEMA_TSID_COLUMN_NAME;
use store_api::region_engine::{
PartitionRange, PrepareRequest, QueryScanContext, RawLabelSource, RegionScannerRef,
BoundTag, PartitionRange, PrepareRequest, QueryScanContext, RawLabelSource, RegionScannerRef,
};
use store_api::storage::{ScanRequest, TimeSeriesDistribution};
@@ -476,153 +475,117 @@ impl RegionScanExec {
}
/// Resolves authorized logical names against this prepared scanner's physical metadata,
/// then selects a native raw-series leaf through the existing binding path.
/// then selects a native raw-series leaf.
///
/// The caller is responsible for authorizing the logical table and supplying its complete
/// logical tag names and requested field names. This method does not discover logical schema
/// or authorize names; it resolves them only against this scanner's metadata snapshot.
/// logical tag names and requested field names. This method resolves them only against this
/// scanner's metadata snapshot.
pub fn try_to_named_series_scan(
self: Arc<Self>,
metric_name: &str,
logical_tags: &[String],
field_names: &[String],
) -> DfResult<Option<Arc<SeriesScanExec>>> {
Ok(self
.try_to_named_series_binding(metric_name, logical_tags, field_names)?
.map(|(_, native)| native))
}
/// 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`]
/// for the authorization the caller owes this method.
pub fn try_to_named_series_binding(
self: Arc<Self>,
metric_name: &str,
logical_tags: &[String],
field_names: &[String],
) -> DfResult<Option<(RawSeriesRequest, Arc<SeriesScanExec>)>> {
let Some(binding) =
self.resolve_named_series_binding(metric_name, logical_tags, field_names)
else {
return Ok(None);
};
let Some(native) = self.try_to_bound_series_scan(&binding)? else {
return Ok(None);
};
Ok(Some((binding, native)))
}
/// Builds the raw-series binding of authorized logical names from this scanner's metadata.
///
/// The binding copies this scanner's metadata snapshot, and every requested name must resolve
/// to a distinct logical column of that snapshot. `None` means at least one name does not.
fn resolve_named_series_binding(
&self,
metric_name: &str,
logical_tags: &[String],
field_names: &[String],
) -> Option<RawSeriesRequest> {
if !self.is_partition_set || field_names.is_empty() {
return None;
return Ok(None);
}
if !self.pushed_dyn_filters.lock().unwrap().is_empty() {
return None;
return Ok(None);
}
let scanner = self.scanner.lock().unwrap();
let metadata = scanner.metadata();
let primary_key_ids = metadata
.primary_key_columns()
.map(|column| column.column_id)
.collect::<HashSet<_>>();
let mut names = HashSet::new();
let mut column_ids = HashSet::new();
let mut tags = Vec::with_capacity(logical_tags.len());
for name in logical_tags {
if !names.insert(name.as_str())
|| name == "__name__"
|| store_api::metric_engine_consts::is_metric_engine_internal_column(name)
|| store_api::storage::consts::is_internal_column(name)
{
return None;
}
let column = metadata.column_by_name(name)?;
if column.semantic_type != SemanticType::Tag
|| store_api::storage::consts::ReservedColumnId::is_reserved(column.column_id)
|| !primary_key_ids.contains(&column.column_id)
|| !column_ids.insert(column.column_id)
{
return None;
}
tags.push(RawSeriesTag {
column_id: column.column_id,
name: name.clone(),
});
}
let mut fields = Vec::with_capacity(field_names.len());
for name in field_names {
if !names.insert(name.as_str()) {
return None;
}
let column = metadata.column_by_name(name)?;
if column.semantic_type != SemanticType::Field
|| store_api::storage::consts::ReservedColumnId::is_reserved(column.column_id)
|| !column_ids.insert(column.column_id)
{
return None;
}
fields.push(column.column_id);
}
Some(RawSeriesRequest {
version: 1,
physical_metadata: metadata.as_ref().clone(),
metric_name: metric_name.to_string(),
logical_tags: tags,
field_ids: fields,
})
}
/// Selects a native raw-series leaf only when the prepared scanner matches the binding.
pub fn try_to_bound_series_scan(
self: Arc<Self>,
binding: &RawSeriesRequest,
) -> DfResult<Option<Arc<SeriesScanExec>>> {
let actual = {
let (metadata, expected_field_ids, fields, tags) = {
let scanner = self.scanner.lock().unwrap();
let metadata = scanner.metadata();
let output_ids = self
.arrow_schema
.fields()
.iter()
.filter_map(|field| {
metadata
.column_by_name(field.name())
.filter(|column| column.semantic_type == SemanticType::Field)
.map(|column| column.column_id)
})
.collect::<Vec<_>>();
(metadata, output_ids)
let primary_key_ids = metadata
.primary_key_columns()
.map(|column| column.column_id)
.collect::<HashSet<_>>();
let mut names = HashSet::new();
let mut column_ids = HashSet::new();
let mut tags = Vec::with_capacity(logical_tags.len());
for name in logical_tags {
if !names.insert(name.as_str())
|| name == "__name__"
|| store_api::metric_engine_consts::is_metric_engine_internal_column(name)
|| store_api::storage::consts::is_internal_column(name)
{
return Ok(None);
}
let Some(column) = metadata.column_by_name(name) else {
return Ok(None);
};
if column.semantic_type != SemanticType::Tag
|| store_api::storage::consts::ReservedColumnId::is_reserved(column.column_id)
|| !primary_key_ids.contains(&column.column_id)
|| !column_ids.insert(column.column_id)
{
return Ok(None);
}
tags.push(BoundTag {
column_id: column.column_id,
name: Arc::from(name.as_str()),
});
}
let mut fields = Vec::with_capacity(field_names.len());
let mut expected_field_ids = Vec::with_capacity(field_names.len());
for name in field_names {
if !names.insert(name.as_str())
|| name == "__name__"
|| store_api::metric_engine_consts::is_metric_engine_internal_column(name)
|| store_api::storage::consts::is_internal_column(name)
{
return Ok(None);
}
let Some(column) = metadata.column_by_name(name) else {
return Ok(None);
};
if column.semantic_type != SemanticType::Field
|| store_api::storage::consts::ReservedColumnId::is_reserved(column.column_id)
|| !column_ids.insert(column.column_id)
{
return Ok(None);
}
expected_field_ids.push(column.column_id);
fields.push(FieldSpec {
name: Arc::from(name.as_str()),
});
}
(metadata, expected_field_ids, fields, tags)
};
if actual.1 != binding.field_ids {
// The requested field sequence must exactly match the scanner's projected field IDs.
let projected = self
.arrow_schema
.fields()
.iter()
.filter_map(|field| {
metadata
.column_by_name(field.name())
.filter(|column| column.semantic_type == SemanticType::Field)
.map(|column| column.column_id)
})
.collect::<Vec<_>>();
if expected_field_ids != projected {
return Ok(None);
}
let (expected_schema, source) = binding.bind(actual.0)?;
let Some(native) = self.clone().try_to_series_scan(source)? else {
let source = RawLabelSource::PrimaryKey {
metric_name: Arc::from(metric_name),
physical_metadata: metadata,
logical_tags: tags.into(),
};
let schema = Arc::new(SeriesSchema {
fields: fields.into(),
eval_grid: None,
});
let Some(native) = self.try_to_series_scan(source)? else {
return Ok(None);
};
if native.series_schema.eval_grid.is_some()
|| native.series_schema.fields.len() != expected_schema.fields.len()
|| native.series_schema.fields.len() != schema.fields.len()
|| native
.series_schema
.fields
.iter()
.zip(expected_schema.fields.iter())
.zip(schema.fields.iter())
.any(|(actual, expected)| actual.name != expected.name)
{
return Ok(None);
@@ -648,7 +611,9 @@ impl RegionScanExec {
else {
return Ok(None);
};
let ordering = scanner.raw_series_ordering();
let ordering = scanner
.raw_series_ordering()
.map_err(|e| DataFusionError::External(Box::new(e)))?;
let arrow_schema = raw_series_frame_schema_with_ordering(&schema, &label_source, ordering)?;
let partition_count = scanner.properties().num_partitions().max(1);
let properties = Arc::new(PlanProperties::new(
@@ -829,7 +794,9 @@ impl ExecutionPlan for SeriesScanExec {
.ok_or_else(|| {
DataFusionError::Execution("raw series capability changed after planning".into())
})?;
let current_ordering = scanner.raw_series_ordering();
let current_ordering = scanner
.raw_series_ordering()
.map_err(|e| DataFusionError::External(Box::new(e)))?;
let current_frame =
raw_series_frame_schema_with_ordering(&current, &self.label_source, current_ordering)?;
if current_frame != self.arrow_schema
@@ -1423,21 +1390,7 @@ mod test {
}
#[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() {
fn named_series_scan_requires_prepared_scanner_and_valid_names() {
let (_, plan) = dynamic_filter_fixture(5686);
let plan = Arc::new(plan);
assert!(
src/table/src/table/scan/series_tests.rs
+6 -165
View File
@@ -19,7 +19,6 @@ use std::task::{Context, Poll};
use api::v1::SemanticType;
use common_error::ext::BoxedError;
use common_query::request::RawSeriesRequest;
use common_query::series_frame::series_frame_ordering;
use common_recordbatch::{
QueryMemoryTracker, RecordBatch, RecordBatches, SendableRecordBatchStream,
@@ -170,12 +169,13 @@ impl RegionScanner for TestScanner {
fn metadata(&self) -> store_api::metadata::RegionMetadataRef {
self.metadata.clone()
}
fn raw_series_ordering(&self) -> SeriesOrdering {
self.state
fn raw_series_ordering(&self) -> Result<SeriesOrdering, BoxedError> {
Ok(self
.state
.lock()
.unwrap()
.planned_ordering
.unwrap_or(SeriesOrdering::KeyThenTime)
.unwrap_or(SeriesOrdering::KeyThenTime))
}
fn prepare(&mut self, request: PrepareRequest) -> Result<(), BoxedError> {
self.properties.prepare(request);
@@ -504,9 +504,11 @@ fn named_series_scan_declines_invalid_names_before_raw_preflight() {
(vec!["aux"], vec!["value"]),
(vec!["ts"], vec!["value"]),
(vec!["host", "host"], vec!["value"]),
(vec!["__name__"], vec!["value"]),
(vec!["__sequence"], vec!["value"]),
(vec!["__table_id"], vec!["value"]),
(vec!["host"], vec!["missing"]),
(vec![], vec!["__name__"]),
(vec!["host"], vec!["ts"]),
(vec!["host"], vec!["value", "value"]),
] {
@@ -518,12 +520,6 @@ fn named_series_scan_declines_invalid_names_before_raw_preflight() {
.unwrap()
.is_none()
);
assert!(
scan.clone()
.try_to_named_series_binding("cpu", &tags, &fields)
.unwrap()
.is_none()
);
}
let state = state.lock().unwrap();
assert_eq!(state.preflights, 0);
@@ -531,109 +527,6 @@ fn named_series_scan_declines_invalid_names_before_raw_preflight() {
assert_eq!(state.legacy_calls, 0);
}
#[test]
fn named_series_binding_reports_the_binding_of_the_actual_metadata() {
let (scan, state, metadata) = named_fixture();
state.lock().unwrap().planned_ordering = Some(SeriesOrdering::Contiguous);
let scan = Arc::new(scan.with_new_partitions(vec![vec![]], 1).unwrap());
let (binding, native) = scan
.try_to_named_series_binding(
"cpu",
&["host".to_string()],
&["aux".to_string(), "value".to_string()],
)
.unwrap()
.expect("complete logical names must resolve against the actual metadata");
assert_eq!(binding.version, 1);
assert_eq!(binding.metric_name, "cpu");
assert_eq!(binding.physical_metadata.region_id, metadata.region_id);
assert_eq!(
binding.physical_metadata.schema_version,
metadata.schema_version
);
assert_eq!(
binding
.physical_metadata
.column_metadatas
.iter()
.map(|column| column.column_id)
.collect::<Vec<_>>(),
metadata
.column_metadatas
.iter()
.map(|column| column.column_id)
.collect::<Vec<_>>()
);
assert_eq!(binding.logical_tags.len(), 1);
assert_eq!(binding.logical_tags[0].name, "host");
assert_eq!(binding.logical_tags[0].column_id, 7);
assert_eq!(binding.field_ids, [11, 12]);
assert_eq!(native.raw_series_ordering(), SeriesOrdering::Contiguous);
assert_eq!(
native
.series_schema
.fields
.iter()
.map(|field| field.name.as_ref())
.collect::<Vec<_>>(),
["aux", "value"]
);
let RawLabelSource::PrimaryKey {
physical_metadata, ..
} = native.label_source()
else {
panic!("expected primary-key label source");
};
assert!(Arc::ptr_eq(&metadata, physical_metadata));
let state = state.lock().unwrap();
assert_eq!(state.preflights, 1);
assert_eq!(state.raw_calls, 0);
assert_eq!(state.legacy_calls, 0);
}
#[test]
fn named_series_binding_never_reorders_or_omits_requested_fields() {
let (scan, state, _) = named_fixture();
let scan = Arc::new(scan.with_new_partitions(vec![vec![]], 1).unwrap());
let (binding, native) = scan
.clone()
.try_to_named_series_binding(
"cpu",
&["host".to_string()],
&["aux".to_string(), "value".to_string()],
)
.unwrap()
.expect("the requested fields are the projected fields");
assert_eq!(binding.field_ids, [11, 12]);
assert_eq!(
native
.series_schema
.fields
.iter()
.map(|field| field.name.as_ref())
.collect::<Vec<_>>(),
["aux", "value"]
);
for fields in [
vec!["value", "aux"],
vec!["aux"],
vec!["value"],
vec!["aux", "value", "value"],
] {
let fields = fields.into_iter().map(str::to_string).collect::<Vec<_>>();
assert!(
scan.clone()
.try_to_named_series_binding("cpu", &["host".to_string()], &fields)
.unwrap()
.is_none()
);
}
let state = state.lock().unwrap();
assert_eq!(state.preflights, 1);
assert_eq!(state.raw_calls, 0);
assert_eq!(state.legacy_calls, 0);
}
#[tokio::test]
async fn named_series_scan_raw_capability_declines_and_legacy_scan_is_unchanged() {
let (scan, state, _) = named_fixture();
@@ -1450,58 +1343,6 @@ fn native_series_scan_forwards_query_load_region_id() {
);
}
#[test]
fn bound_series_scan_uses_open_scanner_metadata_without_execution() {
let schema = Arc::new(SeriesSchema {
fields: vec![FieldSpec {
name: Arc::from("value"),
}]
.into(),
eval_grid: None,
});
let (scan, state) = fixture(Some(schema), vec![]);
let actual = scan.scanner.lock().unwrap().metadata();
let request = RawSeriesRequest {
version: 1,
physical_metadata: (*actual).clone(),
metric_name: "cpu".to_string(),
logical_tags: vec![],
field_ids: vec![1],
};
let native = scan
.clone()
.try_to_bound_series_scan(&request)
.unwrap()
.unwrap();
assert_eq!(native.series_schema.fields[0].name.as_ref(), "value");
let RawLabelSource::PrimaryKey {
physical_metadata, ..
} = native.label_source()
else {
panic!("expected key labels")
};
assert!(Arc::ptr_eq(&actual, physical_metadata));
assert_eq!(state.lock().unwrap().preflights, 1);
let mut wrong_projection = request.clone();
wrong_projection.field_ids = vec![2];
assert!(
scan.clone()
.try_to_bound_series_scan(&wrong_projection)
.unwrap()
.is_none()
);
let mut wrong_metadata = request.clone();
wrong_metadata.physical_metadata.schema_version += 1;
assert!(
scan.clone()
.try_to_bound_series_scan(&wrong_metadata)
.is_err()
);
assert_eq!(state.lock().unwrap().preflights, 1);
assert_eq!(state.lock().unwrap().raw_calls, 0);
assert_eq!(state.lock().unwrap().legacy_calls, 0);
}
#[tokio::test]
async fn raw_series_admission_captures_scanner_ordering_before_execution() {
for ordering in [SeriesOrdering::KeyThenTime, SeriesOrdering::Contiguous] {
tests-integration/src/tests/promql_test.rs
+125 -1
View File
@@ -2153,6 +2153,65 @@ async fn native_promql_executes_fused_node_and_matches_legacy() {
.await;
assert_eq!(native_rows, expected);
assert_eq!(native_rows, legacy_rows);
// Exercise the split aggregate's fixed-width min, group and stdvar states through the same
// native/legacy comparison harness. Sorting is part of the statement output contract.
for aggregate in [
"min by (host)(native_e2e)",
"group by (host)(native_e2e)",
"stdvar by (host)(native_e2e)",
] {
let query = PromQuery {
query: aggregate.to_string(),
start: "0".to_string(),
end: "0".to_string(),
step: "1s".to_string(),
..PromQuery::default()
};
let native_rows = native_result_rows(
instance
.do_promql_query(&query, native_query_context(true, 4))
.await
.remove(0)
.unwrap(),
)
.await;
let legacy_rows = native_result_rows(
instance
.do_promql_query(&query, native_query_context(false, 4))
.await
.remove(0)
.unwrap(),
)
.await;
assert_eq!(native_rows, legacy_rows, "{aggregate}");
let expected_value = if aggregate.starts_with("min") {
None
} else if aggregate.starts_with("group") {
Some(1.0)
} else {
Some(0.0)
};
assert_eq!(native_rows.len(), 16, "{aggregate}");
for (index, (time, host, value)) in native_rows.iter().enumerate() {
assert_eq!(*time, 0, "{aggregate}");
assert_eq!(
host,
&char::from(b'a' + index as u8).to_string(),
"{aggregate}"
);
if let Some(expected) = expected_value {
assert_eq!(*value, expected, "{aggregate}, host={host}");
}
}
if aggregate.starts_with("min") {
assert_eq!(
native_rows.iter().map(|row| row.2).collect::<Vec<_>>(),
(1..=16).map(|value| value as f64 * 2.0).collect::<Vec<_>>(),
"{aggregate}"
);
}
}
}
/// Native instant selectors preserve anchored evaluation and sample timestamps in standalone.
@@ -2166,7 +2225,7 @@ async fn native_promql_instant_at_uses_sample_time_and_anchor() {
let instance = standalone.fe_instance().clone();
execute_all(
&instance,
"CREATE TABLE native_at (host STRING, ts TIMESTAMP TIME INDEX, val DOUBLE, PRIMARY KEY(host)); INSERT INTO native_at (host, ts, val) VALUES ('a', -1500, 1), ('a', 500, 2), ('a', 1500, 3), ('a', 3500, 4), ('b', -1500, 9), ('missing', 500, 8), ('missing', 1500, NULL), ('stale', 500, 7)",
"CREATE TABLE native_at (host STRING, ts TIMESTAMP TIME INDEX, val DOUBLE, PRIMARY KEY(host)); INSERT INTO native_at (host, ts, val) VALUES ('a', -1500, 1), ('a', 500, 2), ('a', 1500, 3), ('a', 3500, 4), ('neg', 10500, 10), ('neg', 11500, 11), ('b', -1500, 9), ('missing', 500, 8), ('missing', 1500, NULL), ('stale', 500, 7)",
QueryContext::arc(),
)
.await;
@@ -2378,6 +2437,63 @@ async fn native_promql_instant_at_uses_sample_time_and_anchor() {
"a",
vec![1.5],
),
(
"timestamp(native_at{host=\"neg\"} @ 11 offset -2s)",
11_000,
11_000,
5_000,
"neg",
vec![10.5],
),
// The anchored and offset-adjusted selector windows are disjoint here.
(
"timestamp(native_at{host=\"a\"} @ 2 offset 2s)",
2_000,
2_000,
750,
"a",
vec![],
),
(
"timestamp((native_at{host=\"a\"} @ 2 offset -1s))",
2_000,
2_000,
5_000,
"a",
vec![1.5],
),
(
"abs(timestamp(native_at{host=\"a\"} @ 2 offset -1s))",
2_000,
2_000,
5_000,
"a",
vec![1.5],
),
(
"timestamp(native_at{host=\"a\"} @ end() offset -1s)",
0,
2_000,
5_000,
"a",
vec![1.5; 3],
),
(
"timestamp(native_at{host=\"a\"} @ 2 offset 1s) + native_at{host=\"a\"} @ 2 offset 1s",
2_000,
2_000,
5_000,
"a",
vec![2.5],
),
(
"timestamp(native_at{host=\"a\"} offset 1s)",
2_000,
2_000,
5_000,
"a",
vec![0.5],
),
];
for (expr, start_ms, end_ms, lookback_ms, host, values) in cases {
@@ -2475,6 +2591,14 @@ async fn native_promql_instant_at_uses_sample_time_and_anchor() {
assert_eq!(timestamps.value(0), 2_000);
assert_eq!(hosts.value(0), "missing");
}
} else if expr == "timestamp(native_at{host=\"neg\"} @ 11 offset -2s)" {
// Legacy uses the shifted instant for its lookback window and selects 11.5s; the
// anchored native selection intersects that window with the unshifted anchor.
assert_eq!(
native_result_rows(legacy).await,
expected_rows(11_000, "neg", &[11.5]),
"legacy known mismatch: {expr}"
);
} else {
assert_eq!(
native_rows,