diff --git a/config/config.md b/config/config.md
index 930608aa4f5..b5aa8488929 100644
--- a/config/config.md
+++ b/config/config.md
@@ -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.
The maximum number of procedures that can be running at the same time.
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.
Default: false |
+| `query.experimental_native_promql` | Bool | `false` | Enable experimental native PromQL execution for compatible queries.
Requires the per-request query.experimental_native_promql=true extension.
Default: false |
| `query.parallelism` | Integer | `0` | Parallelism of the query engine.
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).
Supports absolute size (e.g., "2GB", "4GB") or percentage of system memory (e.g., "20%").
Setting it to 0 disables the limit. Defaults to 0 (unlimited); this example sets it to `50%`.
When this limit is reached, queries will fail with ResourceExhausted error.
NOTE: This does NOT limit memory used by table scans. |
| `query.experimental_memory_pool_policy` | String | `greedy` | Experimental memory pool allocation policy:
- "greedy" (default): first-come-first-served allocation; preserves current behavior.
- "fair": divides available memory among spillable operators and may spill earlier.
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.
Default: false |
+| `query.experimental_native_promql` | Bool | `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 |
| `query.parallelism` | Integer | `0` | Parallelism of the query engine.
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.
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).
Supports absolute size (e.g., "4GB", "8GB") or percentage of system memory (e.g., "30%").
Setting it to 0 disables the limit. Defaults to 0 (unlimited); this example sets it to `50%`.
When this limit is reached, queries will fail with ResourceExhausted error.
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.
**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.
- `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).
- `fail`: the read fails, so the region does not open.
**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.
Default: false |
| `query.parallelism` | Integer | `0` | Parallelism of the query engine.
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).
Supports absolute size (e.g., "2GB", "4GB") or percentage of system memory (e.g., "20%").
Setting it to 0 disables the limit. Defaults to 0 (unlimited); this example sets it to `50%`.
When this limit is reached, queries will fail with ResourceExhausted error.
NOTE: This does NOT limit memory used by table scans. |
| `query.experimental_memory_pool_policy` | String | `greedy` | Experimental memory pool allocation policy:
- "greedy" (default): first-come-first-served allocation; preserves current behavior.
- "fair": divides available memory among spillable operators and may spill earlier.
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.
Default: false |
| `query.parallelism` | Integer | `1` | Parallelism of the query engine for query sent by flownode.
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).
Supports absolute size (e.g., "1GB", "2GB") or percentage of system memory (e.g., "20%").
Setting it to 0 disables the limit. Defaults to 0 (unlimited); this example sets it to `50%`.
When this limit is reached, queries will fail with ResourceExhausted error.
NOTE: This does NOT limit memory used by table scans. |
| `query.experimental_memory_pool_policy` | String | `greedy` | Experimental memory pool allocation policy:
- "greedy" (default): first-come-first-served allocation; preserves current behavior.
- "fair": divides available memory among spillable operators and may spill earlier.
Only effective when `memory_pool_size` is bounded (>0). |
diff --git a/config/datanode.example.toml b/config/datanode.example.toml
index 2af83991638..0a01fa31973 100644
--- a/config/datanode.example.toml
+++ b/config/datanode.example.toml
@@ -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
diff --git a/config/flownode.example.toml b/config/flownode.example.toml
index 4827e576258..fcd5282cb7c 100644
--- a/config/flownode.example.toml
+++ b/config/flownode.example.toml
@@ -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
diff --git a/config/frontend.example.toml b/config/frontend.example.toml
index 8f0122a04ff..d798f67e56e 100644
--- a/config/frontend.example.toml
+++ b/config/frontend.example.toml
@@ -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.
diff --git a/config/standalone.example.toml b/config/standalone.example.toml
index 4346b4daeb8..ea39a2e5a6b 100644
--- a/config/standalone.example.toml
+++ b/config/standalone.example.toml
@@ -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.
diff --git a/src/common/query/src/request.rs b/src/common/query/src/request.rs
index 77971440447..86c6758aa25 100644
--- a/src/common/query/src/request.rs
+++ b/src/common/query/src/request.rs
@@ -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;
diff --git a/src/common/query/src/request/raw_series_request.rs b/src/common/query/src/request/raw_series_request.rs
deleted file mode 100644
index 0dd02b5290d..00000000000
--- a/src/common/query/src/request/raw_series_request.rs
+++ /dev/null
@@ -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,
- /// Physical fields in requested output order.
- pub field_ids: Vec,
-}
-
-impl RawSeriesRequest {
- /// Binds expected output fields and labels to the actual opened scanner metadata.
- pub fn bind(
- &self,
- actual: RegionMetadataRef,
- ) -> DataFusionResult<(Arc, 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::>>()?;
- 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::>();
- 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::>(),
- ["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());
- }
-}
diff --git a/src/common/query/src/series_frame.rs b/src/common/query/src/series_frame.rs
index 04096c2b903..94022e4eb3a 100644
--- a/src/common/query/src/series_frame.rs
+++ b/src/common/query/src/series_frame.rs
@@ -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;
diff --git a/src/common/query/src/series_frame/record_batches.rs b/src/common/query/src/series_frame/record_batches.rs
deleted file mode 100644
index fa16f16deb3..00000000000
--- a/src/common/query/src/series_frame/record_batches.rs
+++ /dev/null
@@ -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>,
- /// Complete logical tag mapping, as `(tag_name, source_column_index)` pairs.
- pub tags: Vec<(Arc, 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 {
- 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, 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::>()
- .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,
- arrow_schema: SchemaRef,
- mapping: InputMapping,
- schema: Arc,
- converter: RowConverter,
- batch: Option<(DfRecordBatch, Rows)>,
- row: usize,
- pending: Option,
- label_source: RawLabelSource,
-}
-impl RawSeriesStream for InputStream {
- fn schema(&self) -> &Arc {
- &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 {
- 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,
- ) -> DfResult {
- let len = end - start;
- let time = batch
- .column(self.mapping.time_index)
- .as_any()
- .downcast_ref::()
- .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::()
- .ok_or_else(|| invalid("invalid Float64 array"))?;
- fields.push(Samples::from_float(values.slice(start, len)));
- } else {
- let hist = array
- .as_any()
- .downcast_ref::()
- .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;
- fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> {
- 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);
- }
- }
-}
diff --git a/src/common/query/src/series_frame/record_batches/tests.rs b/src/common/query/src/series_frame/record_batches/tests.rs
deleted file mode 100644
index 7d4948e7d18..00000000000
--- a/src/common/query/src/series_frame/record_batches/tests.rs
+++ /dev/null
@@ -1,1263 +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, PlainError};
-use common_error::status_code::StatusCode;
-use common_recordbatch::DfSendableRecordBatchStream;
-use datafusion::datasource::memory::MemorySourceConfig;
-use datafusion::physical_expr::expressions::{BinaryExpr, Column, Literal};
-use datafusion::physical_expr::{LexOrdering, PhysicalExpr, PhysicalSortExpr};
-use datafusion::physical_plan::ExecutionPlan;
-use datafusion::physical_plan::RecordBatchStream;
-use datafusion::physical_plan::filter::FilterExec;
-use datafusion::physical_plan::sorts::sort::SortExec;
-use datafusion::prelude::SessionContext;
-use datafusion_common::ScalarValue;
-use datafusion_expr::Operator;
-use datatypes::arrow::array::{
- Array, Float64Array, Int64Array, StringArray, StringDictionaryBuilder, StructArray,
- TimestampMillisecondArray,
-};
-use datatypes::arrow::compute::cast;
-use datatypes::arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit};
-use datatypes::arrow::record_batch::RecordBatch;
-use datatypes::series::SeriesOrdering;
-use datatypes::series::raw::RawSeriesChunk;
-use datatypes::series::samples::{DenseSamplesRef, FieldSpec, SampleRef};
-use futures_util::{Stream, StreamExt};
-
-use super::{InputMapping, record_batches_to_raw_series};
-
-struct CountWake(Arc);
-
-impl std::task::Wake for CountWake {
- fn wake(self: Arc) {
- self.0.fetch_add(1, Ordering::SeqCst);
- }
-
- fn wake_by_ref(self: &Arc) {
- self.0.fetch_add(1, Ordering::SeqCst);
- }
-}
-
-enum InputItem {
- Batch(datafusion_common::Result),
- Pending,
-}
-
-struct BatchStream {
- schema: SchemaRef,
- batches: VecDeque,
- polls: Arc,
- drops: Arc,
-}
-
-impl Drop for BatchStream {
- fn drop(&mut self) {
- self.drops.fetch_add(1, Ordering::SeqCst);
- }
-}
-
-impl Stream for BatchStream {
- type Item = datafusion_common::Result;
-
- fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> {
- self.polls.fetch_add(1, Ordering::SeqCst);
- match self.batches.pop_front() {
- Some(InputItem::Batch(item)) => Poll::Ready(Some(item)),
- Some(InputItem::Pending) => {
- cx.waker().wake_by_ref();
- Poll::Pending
- }
- None => Poll::Ready(None),
- }
- }
-}
-
-impl RecordBatchStream for BatchStream {
- fn schema(&self) -> SchemaRef {
- self.schema.clone()
- }
-}
-
-fn schema() -> SchemaRef {
- Arc::new(Schema::new(vec![
- Field::new(
- "time",
- DataType::Timestamp(TimeUnit::Millisecond, None),
- false,
- ),
- Field::new("host", DataType::Utf8, true),
- Field::new("value", DataType::Float64, true),
- Field::new("predicate", DataType::Int64, false),
- ]))
-}
-
-fn batch(
- times: &[i64],
- hosts: &[Option<&str>],
- values: &[Option],
- predicates: &[i64],
-) -> RecordBatch {
- RecordBatch::try_new(
- schema(),
- vec![
- Arc::new(TimestampMillisecondArray::from(times.to_vec())),
- Arc::new(StringArray::from(hosts.to_vec())),
- Arc::new(Float64Array::from(values.to_vec())),
- Arc::new(Int64Array::from(predicates.to_vec())),
- ],
- )
- .unwrap()
-}
-
-fn mapping() -> InputMapping {
- InputMapping {
- time_index: 0,
- metric_name: Some(Arc::from("cpu")),
- tags: vec![(Arc::from("host"), 1)],
- fields: vec![(
- FieldSpec {
- name: Arc::from("value"),
- },
- 2,
- )],
- }
-}
-
-fn source(
- batches: Vec,
-) -> (
- DfSendableRecordBatchStream,
- Arc,
- Arc,
-) {
- let polls = Arc::new(AtomicUsize::new(0));
- let drops = Arc::new(AtomicUsize::new(0));
- let input = Box::pin(BatchStream {
- schema: schema(),
- batches: batches
- .into_iter()
- .map(|batch| InputItem::Batch(Ok(batch)))
- .collect(),
- polls: polls.clone(),
- drops: drops.clone(),
- });
- (input, polls, drops)
-}
-
-async fn collect(
- mut stream: store_api::region_engine::SendableRawSeriesStream,
-) -> Vec {
- let mut chunks = Vec::new();
- while let Some(item) = stream.next().await {
- chunks.push(item.unwrap());
- }
- chunks
-}
-
-fn key(chunk: &RawSeriesChunk) -> &[u8] {
- chunk.primary_key.as_slice()
-}
-
-#[tokio::test]
-async fn maps_identity_and_payload_columns_across_batches() {
- let (input, _, _) = source(vec![
- batch(
- &[1, 2],
- &[Some("a"), Some("a")],
- &[Some(10.0), None],
- &[9, 9],
- ),
- RecordBatch::new_empty(schema()),
- batch(
- &[3, 1],
- &[Some("a"), Some("b")],
- &[None, Some(11.0)],
- &[9, 9],
- ),
- ]);
- let output = record_batches_to_raw_series(input, mapping()).unwrap();
- assert!(matches!(
- output.label_source(),
- store_api::region_engine::RawLabelSource::Explicit
- ));
- assert_eq!(output.ordering(), SeriesOrdering::Contiguous);
- assert_eq!(
- output
- .schema()
- .fields
- .iter()
- .map(|field| field.name.as_ref())
- .collect::>(),
- ["value"]
- );
- let chunks = collect(output).await;
- assert_eq!(chunks.len(), 3);
- assert_eq!(chunks[0].time.as_buffer().as_ref(), &[1, 2]);
- assert_eq!(chunks[1].time.timestamp(0), Some(3));
- assert_eq!(chunks[2].time.timestamp(0), Some(1));
- assert_eq!(key(&chunks[0]), key(&chunks[1]));
- assert_ne!(key(&chunks[1]), key(&chunks[2]));
- assert_eq!(
- chunks.iter().map(|chunk| chunk.is_last).collect::>(),
- [false, true, true]
- );
- for chunk in &chunks {
- assert_eq!(chunk.metric.as_ref().unwrap().name(), Some("cpu"));
- assert_eq!(chunk.fields.as_slice().len(), 1);
- }
- let DenseSamplesRef::Float(a) = chunks[0].fields.as_slice()[0].segments()[0]
- .dense()
- .unwrap()
- else {
- panic!("float field expected")
- };
- assert_eq!(a.value(0), 10.0);
- assert!(a.is_null(1));
- let DenseSamplesRef::Float(a) = chunks[1].fields.as_slice()[0].segments()[0]
- .dense()
- .unwrap()
- else {
- panic!("float field expected")
- };
- assert!(a.is_null(0));
- let DenseSamplesRef::Float(a) = chunks[2].fields.as_slice()[0].segments()[0]
- .dense()
- .unwrap()
- else {
- panic!("float field expected")
- };
- assert_eq!(a.value(0), 11.0);
-}
-
-#[tokio::test]
-async fn identity_uses_complete_labels_and_null_is_not_empty() {
- let schema = Arc::new(Schema::new(vec![
- Field::new(
- "time",
- DataType::Timestamp(TimeUnit::Millisecond, None),
- false,
- ),
- Field::new("left", DataType::Utf8, true),
- Field::new("right", DataType::Utf8, true),
- Field::new("value", DataType::Float64, false),
- Field::new("unmapped", DataType::Int64, false),
- ]));
- let make = |hosts: Vec>, rights: Vec>| {
- let len = hosts.len();
- RecordBatch::try_new(
- schema.clone(),
- vec![
- Arc::new(TimestampMillisecondArray::from(vec![1; len])),
- Arc::new(StringArray::from(hosts)),
- Arc::new(StringArray::from(rights)),
- Arc::new(Float64Array::from(vec![Some(1.0); len])),
- Arc::new(Int64Array::from(vec![i64::MAX; len])),
- ],
- )
- .unwrap()
- };
- let mapping = InputMapping {
- time_index: 0,
- metric_name: None,
- tags: vec![(Arc::from("right"), 2), (Arc::from("left"), 1)],
- fields: vec![(
- FieldSpec {
- name: Arc::from("v"),
- },
- 3,
- )],
- };
- let (input, _, _) = source_custom(
- schema.clone(),
- vec![make(
- vec![None, Some(""), Some("ab")],
- vec![Some("x"), None, Some("c")],
- )],
- );
- let chunks = collect(record_batches_to_raw_series(input, mapping).unwrap()).await;
- assert_eq!(chunks.len(), 3);
- assert!(chunks.iter().all(|chunk| chunk.is_last));
- assert!(
- chunks
- .iter()
- .all(|chunk| chunk.metric.as_ref().unwrap().name().is_none())
- );
- let labels = chunks
- .iter()
- .map(|chunk| {
- chunk
- .metric
- .as_ref()
- .unwrap()
- .tags()
- .iter()
- .map(|l| (l.name.as_ref(), l.value.as_ref()))
- .collect::>()
- })
- .collect::>();
- assert_eq!(labels[0], vec![("right", "x")]);
- assert_eq!(labels[1], vec![("left", "")]);
- assert_eq!(labels[2], vec![("left", "ab"), ("right", "c")]);
- assert_ne!(key(&chunks[0]), key(&chunks[1]));
- assert_ne!(key(&chunks[1]), key(&chunks[2]));
-}
-
-fn scripted_source(
- schema: SchemaRef,
- batches: Vec,
-) -> (
- DfSendableRecordBatchStream,
- Arc,
- Arc,
-) {
- let polls = Arc::new(AtomicUsize::new(0));
- let drops = Arc::new(AtomicUsize::new(0));
- (
- Box::pin(BatchStream {
- schema,
- batches: batches.into(),
- polls: polls.clone(),
- drops: drops.clone(),
- }),
- polls,
- drops,
- )
-}
-
-fn source_custom(
- schema: SchemaRef,
- batches: Vec,
-) -> (
- DfSendableRecordBatchStream,
- Arc,
- Arc,
-) {
- let polls = Arc::new(AtomicUsize::new(0));
- let drops = Arc::new(AtomicUsize::new(0));
- (
- Box::pin(BatchStream {
- schema,
- batches: batches
- .into_iter()
- .map(|batch| InputItem::Batch(Ok(batch)))
- .collect(),
- polls: polls.clone(),
- drops: drops.clone(),
- }),
- polls,
- drops,
- )
-}
-
-#[tokio::test]
-async fn preserves_sliced_timestamp_and_float_buffers_and_exact_values() {
- let schema = schema();
- let times_owner = TimestampMillisecondArray::from(vec![99, 4, 5, 6, 100]);
- let floats_owner = Float64Array::from(vec![
- Some(88.0),
- Some(-0.0),
- None,
- Some(f64::from_bits(0x7ff8_0000_0000_0042)),
- Some(77.0),
- ]);
- let times = times_owner.slice(1, 3);
- let floats = floats_owner.slice(1, 3);
- let input_time_ptr = times.values().as_ptr();
- let input_float_ptr = floats.values().as_ptr();
- let batch = RecordBatch::try_new(
- schema.clone(),
- vec![
- Arc::new(times),
- Arc::new(StringArray::from(vec![Some("h"); 3])),
- Arc::new(floats),
- Arc::new(Int64Array::from(vec![1, 2, 3])),
- ],
- )
- .unwrap();
- let (input, _, _) = source_custom(schema, vec![batch]);
- let chunks = collect(record_batches_to_raw_series(input, mapping()).unwrap()).await;
- assert_eq!(chunks.len(), 1);
- let time = chunks[0].time.as_buffer();
- assert_eq!(time.as_ptr(), input_time_ptr);
- let DenseSamplesRef::Float(values) = chunks[0].fields.as_slice()[0].segments()[0]
- .dense()
- .unwrap()
- else {
- panic!("float field expected")
- };
- assert_eq!(values.values().as_ptr(), input_float_ptr);
- assert_eq!(time.as_ref(), &[4, 5, 6]);
- assert_eq!(values.value(0).to_bits(), (-0.0_f64).to_bits());
- assert!(values.is_null(1));
- assert_eq!(values.value(2).to_bits(), 0x7ff8_0000_0000_0042);
- drop(times_owner);
- drop(floats_owner);
- assert_eq!(time.as_ref(), &[4, 5, 6]);
- assert_eq!(values.value(2).to_bits(), 0x7ff8_0000_0000_0042);
-}
-
-#[tokio::test]
-async fn native_histogram_payload_and_null_positions_are_preserved() {
- use crate::native_histogram::{
- NativeHistogram, Span, build_histogram_array, native_histogram_arrow_type,
- };
-
- let histogram = NativeHistogram {
- schema: -53,
- zero_threshold: 0.25,
- sum: -17.5,
- reset_hint: crate::native_histogram::CounterResetHint::Unrecognized(37),
- start_timestamp: Some(-123),
- custom_values: vec![1.5, 4.5],
- positive_spans: vec![Span {
- offset: 1,
- length: 2,
- }],
- negative_spans: vec![Span {
- offset: -2,
- length: 1,
- }],
- count: 17.25,
- zero_count: 2.5,
- positive_buckets: vec![3.25, 7.5],
- negative_buckets: vec![0.75],
- };
- let values = build_histogram_array(&[Some(histogram.clone()), None]);
- let schema = Arc::new(Schema::new(vec![
- Field::new(
- "time",
- DataType::Timestamp(TimeUnit::Millisecond, None),
- false,
- ),
- Field::new("host", DataType::Utf8, false),
- Field::new("hist", native_histogram_arrow_type(), true),
- ]));
- let batch = RecordBatch::try_new(
- schema.clone(),
- vec![
- Arc::new(TimestampMillisecondArray::from(vec![1, 2])),
- Arc::new(StringArray::from(vec!["h", "h"])),
- values,
- ],
- )
- .unwrap();
- let mapping = InputMapping {
- time_index: 0,
- metric_name: Some(Arc::from("cpu")),
- tags: vec![(Arc::from("host"), 1)],
- fields: vec![(
- FieldSpec {
- name: Arc::from("hist"),
- },
- 2,
- )],
- };
- let (input, _, _) = source_custom(schema, vec![batch]);
- let chunks = collect(record_batches_to_raw_series(input, mapping).unwrap()).await;
- assert_eq!(chunks.len(), 1);
- assert_eq!(chunks[0].time.as_buffer().as_ref(), &[1, 2]);
- let samples = &chunks[0].fields.as_slice()[0].segments()[0];
- let SampleRef::Histogram(view) = samples.sample(0) else {
- panic!("histogram expected")
- };
- assert_eq!(view.schema, histogram.schema);
- assert_eq!(view.zero_threshold, histogram.zero_threshold);
- assert_eq!(view.sum.to_bits(), histogram.sum.to_bits());
- assert_eq!(view.count, histogram.count);
- assert_eq!(view.zero_count, histogram.zero_count);
- assert_eq!(view.reset_hint, histogram.reset_hint);
- assert_eq!(view.start_timestamp, histogram.start_timestamp);
- assert_eq!(view.custom_values, histogram.custom_values);
- assert_eq!(view.positive.span_offsets, [1]);
- assert_eq!(view.positive.span_lengths, [2]);
- assert_eq!(view.positive.buckets, [3.25, 7.5]);
- assert_eq!(view.negative.span_offsets, [-2]);
- assert_eq!(view.negative.span_lengths, [1]);
- assert_eq!(view.negative.buckets, [0.75]);
- assert!(matches!(samples.sample(1), SampleRef::Missing));
-}
-
-#[tokio::test]
-async fn malformed_native_histogram_is_terminal_after_completed_prefix() {
- use crate::native_histogram::{build_histogram_array, native_histogram_arrow_type};
- let valid = build_histogram_array(&[None]);
- let valid = valid
- .as_any()
- .downcast_ref::()
- .unwrap()
- .clone();
- let mut columns = valid.columns().to_vec();
- let bad_values = Arc::new(Float64Array::new_null(1)) as Arc;
- columns[1] = bad_values;
- let malformed = StructArray::try_new(valid.fields().clone(), columns, None).unwrap();
- let hist_type = native_histogram_arrow_type();
- let schema = Arc::new(Schema::new(vec![
- Field::new(
- "time",
- DataType::Timestamp(TimeUnit::Millisecond, None),
- false,
- ),
- Field::new("host", DataType::Utf8, false),
- Field::new("hist", hist_type, true),
- ]));
- let good = build_histogram_array(&[None]);
- let good_batch = RecordBatch::try_new(
- schema.clone(),
- vec![
- Arc::new(TimestampMillisecondArray::from(vec![10])),
- Arc::new(StringArray::from(vec!["a"])),
- good,
- ],
- )
- .unwrap();
- let bad_batch = RecordBatch::try_new(
- schema.clone(),
- vec![
- Arc::new(TimestampMillisecondArray::from(vec![1])),
- Arc::new(StringArray::from(vec!["b"])),
- Arc::new(malformed),
- ],
- )
- .unwrap();
- let mapping = InputMapping {
- time_index: 0,
- metric_name: None,
- tags: vec![(Arc::from("host"), 1)],
- fields: vec![(
- FieldSpec {
- name: Arc::from("hist"),
- },
- 2,
- )],
- };
- let (input, _, _) = source_custom(schema, vec![good_batch, bad_batch]);
- let mut stream = record_batches_to_raw_series(input, mapping).unwrap();
- let prefix = stream.next().await.unwrap().unwrap();
- assert_eq!(prefix.metric.as_ref().unwrap().get("host"), Some("a"));
- assert_eq!(prefix.time.timestamp(0), Some(10));
- assert!(prefix.is_last);
- assert!(stream.next().await.unwrap().is_err());
- assert!(stream.next().await.is_none());
-}
-
-#[tokio::test]
-async fn string_view_and_dictionary_tags_use_logical_values_for_identity() {
- let encodings = [
- DataType::Utf8,
- DataType::LargeUtf8,
- DataType::Utf8View,
- DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
- ];
- let mut output_chunks = Vec::new();
- for tag_type in encodings {
- let schema = Arc::new(Schema::new(vec![
- Field::new(
- "time",
- DataType::Timestamp(TimeUnit::Millisecond, None),
- false,
- ),
- Field::new("tag", tag_type.clone(), false),
- Field::new("value", DataType::Float64, false),
- ]));
- let tag_values: Arc =
- Arc::new(StringArray::from(vec![Some("same"), Some("same")]));
- let tag = if tag_type == DataType::Utf8 {
- tag_values
- } else {
- cast(tag_values.as_ref(), &tag_type).unwrap()
- };
- let batch = RecordBatch::try_new(
- schema.clone(),
- vec![
- Arc::new(TimestampMillisecondArray::from(vec![1, 2])),
- tag,
- Arc::new(Float64Array::from(vec![1.0, 2.0])),
- ],
- )
- .unwrap();
- let mapping = InputMapping {
- time_index: 0,
- metric_name: None,
- tags: vec![(Arc::from("tag"), 1)],
- fields: vec![(
- FieldSpec {
- name: Arc::from("value"),
- },
- 2,
- )],
- };
- let (input, _, _) = source_custom(schema.clone(), vec![batch]);
- let chunks = collect(record_batches_to_raw_series(input, mapping).unwrap()).await;
- assert_eq!(chunks.len(), 1);
- assert!(chunks[0].is_last);
- assert_eq!(chunks[0].metric.as_ref().unwrap().get("tag"), Some("same"));
- output_chunks.push(chunks.into_iter().next().unwrap());
- }
- for chunk in &output_chunks {
- assert_eq!(chunk.metric.as_ref().unwrap().get("tag"), Some("same"));
- assert_eq!(chunk.time.timestamp(0), Some(1));
- }
-}
-
-#[tokio::test]
-async fn dictionary_index_changes_preserve_same_identity_across_batches() {
- use datatypes::arrow::datatypes::Int8Type;
-
- let tag_type = DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8));
- let schema = Arc::new(Schema::new(vec![
- Field::new(
- "time",
- DataType::Timestamp(TimeUnit::Millisecond, None),
- false,
- ),
- Field::new("tag", tag_type, false),
- Field::new("value", DataType::Float64, false),
- ]));
- let mut first_dict = StringDictionaryBuilder::::new();
- first_dict.append("other").unwrap();
- first_dict.append("same").unwrap();
- let first_tag = Arc::new(first_dict.finish().slice(1, 1)) as Arc;
- let mut second_dict = StringDictionaryBuilder::::new();
- second_dict.append("same").unwrap();
- second_dict.append("other").unwrap();
- let second_tag = Arc::new(second_dict.finish().slice(0, 1)) as Arc;
- let batch_with = |time, tag| {
- RecordBatch::try_new(
- schema.clone(),
- vec![
- Arc::new(TimestampMillisecondArray::from(vec![time])),
- tag,
- Arc::new(Float64Array::from(vec![time as f64])),
- ],
- )
- .unwrap()
- };
- let mapping = InputMapping {
- time_index: 0,
- metric_name: None,
- tags: vec![(Arc::from("tag"), 1)],
- fields: vec![(
- FieldSpec {
- name: Arc::from("value"),
- },
- 2,
- )],
- };
- let (input, _, _) = source_custom(
- schema.clone(),
- vec![batch_with(1, first_tag), batch_with(2, second_tag)],
- );
- let chunks = collect(record_batches_to_raw_series(input, mapping).unwrap()).await;
- assert_eq!(chunks.len(), 2);
- assert_eq!(key(&chunks[0]), key(&chunks[1]));
- assert_eq!(chunks[0].time.timestamp(0), Some(1));
- assert_eq!(chunks[1].time.timestamp(0), Some(2));
- assert_eq!(chunks[0].metric.as_ref().unwrap().get("tag"), Some("same"));
- assert_eq!(chunks[1].metric.as_ref().unwrap().get("tag"), Some("same"));
- assert_eq!((chunks[0].is_last, chunks[1].is_last), (false, true));
-}
-
-#[tokio::test]
-async fn constructor_rejects_bad_mapping_without_polling_and_drops_source() {
- let alternate = Arc::new(Schema::new(vec![
- Field::new(
- "time",
- DataType::Timestamp(TimeUnit::Millisecond, None),
- false,
- ),
- Field::new("host", DataType::Utf8, true),
- Field::new("other_tag", DataType::Utf8, true),
- Field::new("value", DataType::Float64, true),
- Field::new("other_value", DataType::Float64, true),
- Field::new("predicate", DataType::Int64, false),
- ]));
- let alternate_batch = RecordBatch::new_empty(alternate.clone());
- let duplicate_tag = InputMapping {
- time_index: 0,
- metric_name: None,
- tags: vec![(Arc::from("dup"), 1), (Arc::from("dup"), 2)],
- fields: vec![(
- FieldSpec {
- name: Arc::from("value"),
- },
- 3,
- )],
- };
- let duplicate_fields = InputMapping {
- time_index: 0,
- metric_name: None,
- tags: vec![],
- fields: vec![
- (
- FieldSpec {
- name: Arc::from("same"),
- },
- 3,
- ),
- (
- FieldSpec {
- name: Arc::from("same"),
- },
- 4,
- ),
- ],
- };
- for invalid in [duplicate_tag, duplicate_fields] {
- let (input, polls, drops) = source_custom(alternate.clone(), vec![alternate_batch.clone()]);
- assert!(record_batches_to_raw_series(input, invalid).is_err());
- assert_eq!(polls.load(Ordering::SeqCst), 0);
- assert_eq!(drops.load(Ordering::SeqCst), 1);
- }
- let (wrong_time, time_mapping) = {
- let schema = schema();
- let mut fields = schema
- .fields()
- .iter()
- .map(|field| field.as_ref().clone())
- .collect::>();
- fields[0] = Field::new("time", DataType::Int64, false);
- let schema = Arc::new(Schema::new(fields));
- let mapping = InputMapping {
- time_index: 0,
- ..mapping()
- };
- (schema, mapping)
- };
- let (wrong_tag, tag_mapping) = {
- let schema = schema();
- let mut fields = schema
- .fields()
- .iter()
- .map(|field| field.as_ref().clone())
- .collect::>();
- fields[1] = Field::new("host", DataType::Boolean, true);
- let schema = Arc::new(Schema::new(fields));
- let mapping = mapping();
- (schema, mapping)
- };
- let (wrong_field, field_mapping) = {
- let schema = schema();
- let mut fields = schema
- .fields()
- .iter()
- .map(|field| field.as_ref().clone())
- .collect::>();
- fields[2] = Field::new("value", DataType::Int64, true);
- let schema = Arc::new(Schema::new(fields));
- let mapping = mapping();
- (schema, mapping)
- };
- for (schema, invalid) in [
- (wrong_time, time_mapping),
- (wrong_tag, tag_mapping),
- (wrong_field, field_mapping),
- ] {
- let (input, polls, drops) =
- source_custom(schema.clone(), vec![RecordBatch::new_empty(schema)]);
- assert!(record_batches_to_raw_series(input, invalid).is_err());
- assert_eq!(polls.load(Ordering::SeqCst), 0);
- assert_eq!(drops.load(Ordering::SeqCst), 1);
- }
- let invalids = [
- InputMapping {
- time_index: 99,
- ..mapping()
- },
- InputMapping {
- tags: vec![(Arc::from("host"), 99)],
- ..mapping()
- },
- InputMapping {
- fields: vec![(
- FieldSpec {
- name: Arc::from("value"),
- },
- 99,
- )],
- ..mapping()
- },
- InputMapping {
- tags: vec![(Arc::from("__name__"), 1)],
- ..mapping()
- },
- InputMapping {
- fields: vec![],
- ..mapping()
- },
- ];
- for invalid in invalids {
- let (input, polls, drops) = source(vec![]);
- assert!(record_batches_to_raw_series(input, invalid).is_err());
- assert_eq!(polls.load(Ordering::SeqCst), 0);
- assert_eq!(drops.load(Ordering::SeqCst), 1);
- }
-}
-
-#[tokio::test]
-async fn pending_upstream_does_not_flush_buffered_series() {
- let first = batch(&[1], &[Some("a")], &[Some(10.0)], &[1]);
- let second = batch(&[2], &[Some("a")], &[Some(20.0)], &[2]);
- let (input, polls, drops) = scripted_source(
- schema(),
- vec![
- InputItem::Batch(Ok(first)),
- InputItem::Pending,
- InputItem::Batch(Ok(second)),
- ],
- );
- let mut stream = record_batches_to_raw_series(input, mapping()).unwrap();
- let wakes = Arc::new(AtomicUsize::new(0));
- let waker = std::task::Waker::from(Arc::new(CountWake(wakes.clone())));
- let mut cx = Context::from_waker(&waker);
- assert!(matches!(
- Pin::new(&mut stream).poll_next(&mut cx),
- Poll::Pending
- ));
- assert_eq!(wakes.load(Ordering::SeqCst), 1);
- assert_eq!(polls.load(Ordering::SeqCst), 2);
- assert_eq!(drops.load(Ordering::SeqCst), 0);
- let Poll::Ready(Some(Ok(first))) = Pin::new(&mut stream).poll_next(&mut cx) else {
- panic!("expected completed prefix")
- };
- assert!(!first.is_last);
- assert_eq!(first.time.timestamp(0), Some(1));
- let Poll::Ready(Some(Ok(last))) = Pin::new(&mut stream).poll_next(&mut cx) else {
- panic!("expected final chunk")
- };
- assert!(last.is_last);
- assert_eq!(last.time.timestamp(0), Some(2));
- assert!(matches!(
- Pin::new(&mut stream).poll_next(&mut cx),
- Poll::Ready(None)
- ));
- assert_eq!(drops.load(Ordering::SeqCst), 1);
-}
-
-#[tokio::test]
-async fn upstream_error_drops_unfinished_buffer_without_emitting_it() {
- let first = batch(&[1], &[Some("a")], &[Some(10.0)], &[1]);
- let error = datafusion_common::DataFusionError::External(Box::new(BoxedError::new(
- PlainError::new("upstream failure".into(), StatusCode::RequestOutdated),
- )));
- let (input, polls, drops) = scripted_source(
- schema(),
- vec![
- InputItem::Batch(Ok(first)),
- InputItem::Pending,
- InputItem::Batch(Err(error)),
- ],
- );
- let mut stream = record_batches_to_raw_series(input, mapping()).unwrap();
- let wakes = Arc::new(AtomicUsize::new(0));
- let waker = std::task::Waker::from(Arc::new(CountWake(wakes.clone())));
- let mut cx = Context::from_waker(&waker);
- assert!(matches!(
- Pin::new(&mut stream).poll_next(&mut cx),
- Poll::Pending
- ));
- assert_eq!(wakes.load(Ordering::SeqCst), 1);
- let Poll::Ready(Some(Err(error))) = Pin::new(&mut stream).poll_next(&mut cx) else {
- panic!("expected upstream failure")
- };
- assert_eq!(
- common_error::ext::ErrorExt::status_code(&error),
- StatusCode::RequestOutdated
- );
- assert_eq!(drops.load(Ordering::SeqCst), 1);
- assert!(matches!(
- Pin::new(&mut stream).poll_next(&mut cx),
- Poll::Ready(None)
- ));
- assert_eq!(polls.load(Ordering::SeqCst), 3);
-}
-
-#[tokio::test]
-async fn completed_key_is_emitted_before_next_key_validation_fails() {
- let (input, _, drops) = source(vec![
- batch(&[10], &[Some("a")], &[Some(10.0)], &[1]),
- batch(
- &[3, 2],
- &[Some("b"), Some("b")],
- &[Some(30.0), Some(20.0)],
- &[2, 2],
- ),
- ]);
- let mut stream = record_batches_to_raw_series(input, mapping()).unwrap();
- let prefix = stream.next().await.unwrap().unwrap();
- assert_eq!(prefix.metric.as_ref().unwrap().get("host"), Some("a"));
- assert_eq!(prefix.time.timestamp(0), Some(10));
- assert!(prefix.is_last);
- let DenseSamplesRef::Float(values) = prefix.fields.as_slice()[0].segments()[0].dense().unwrap()
- else {
- panic!("float expected")
- };
- assert_eq!(values.value(0), 10.0);
- assert!(stream.next().await.unwrap().is_err());
- assert!(stream.next().await.is_none());
- assert_eq!(drops.load(Ordering::SeqCst), 1);
-}
-
-#[tokio::test]
-async fn eof_releases_source_batch_and_emitted_chunks_own_sliced_buffers() {
- let times = TimestampMillisecondArray::from(vec![1, 2]);
- let floats = Float64Array::from(vec![Some(10.0), Some(20.0)]);
- let time_buffer = times.values().inner().clone();
- let float_buffer = floats.values().inner().clone();
- let input_batch = RecordBatch::try_new(
- schema(),
- vec![
- Arc::new(times),
- Arc::new(StringArray::from(vec![Some("a"), Some("b")])),
- Arc::new(floats),
- Arc::new(Int64Array::from(vec![1, 2])),
- ],
- )
- .unwrap();
- let (input, _, drops) = source(vec![input_batch]);
- let mut stream = record_batches_to_raw_series(input, mapping()).unwrap();
- let first = stream.next().await.unwrap().unwrap();
- assert!(first.is_last);
- assert_eq!(drops.load(Ordering::SeqCst), 0);
- assert!(time_buffer.strong_count() > 1);
- assert!(float_buffer.strong_count() > 1);
- drop(first);
- assert!(stream.next().await.unwrap().unwrap().is_last);
- assert_eq!(drops.load(Ordering::SeqCst), 1);
- assert_eq!(time_buffer.strong_count(), 1);
- assert_eq!(float_buffer.strong_count(), 1);
- assert!(stream.next().await.is_none());
-}
-
-#[tokio::test]
-async fn canceling_pending_stream_releases_buffered_batch_and_owner() {
- let time_owner = TimestampMillisecondArray::from(vec![1]);
- let float_owner = Float64Array::from(vec![Some(10.0)]);
- let time_buffer = time_owner.values().inner().clone();
- let float_buffer = float_owner.values().inner().clone();
- let input_batch = RecordBatch::try_new(
- schema(),
- vec![
- Arc::new(time_owner),
- Arc::new(StringArray::from(vec![Some("a")])),
- Arc::new(float_owner),
- Arc::new(Int64Array::from(vec![1])),
- ],
- )
- .unwrap();
- let (input, _, drops) = scripted_source(
- schema(),
- vec![InputItem::Batch(Ok(input_batch)), InputItem::Pending],
- );
- let mut stream = record_batches_to_raw_series(input, mapping()).unwrap();
- let wakes = Arc::new(AtomicUsize::new(0));
- let waker = std::task::Waker::from(Arc::new(CountWake(wakes.clone())));
- let mut cx = Context::from_waker(&waker);
- assert!(matches!(
- Pin::new(&mut stream).poll_next(&mut cx),
- Poll::Pending
- ));
- assert_eq!(wakes.load(Ordering::SeqCst), 1);
- assert!(time_buffer.strong_count() > 1);
- assert!(float_buffer.strong_count() > 1);
- drop(stream);
- assert_eq!(drops.load(Ordering::SeqCst), 1);
- assert_eq!(time_buffer.strong_count(), 1);
- assert_eq!(float_buffer.strong_count(), 1);
-}
-
-#[tokio::test]
-async fn runtime_errors_terminate_once_and_keep_no_unfinished_identity() {
- let (input, polls, drops) = source(vec![
- batch(
- &[1, 1],
- &[Some("a"), Some("b")],
- &[Some(1.0), Some(2.0)],
- &[1, 2],
- ),
- batch(&[1], &[Some("b")], &[Some(3.0)], &[3]),
- batch(&[3], &[Some("c")], &[Some(4.0)], &[4]),
- ]);
- let mut stream = record_batches_to_raw_series(input, mapping()).unwrap();
- assert!(stream.next().await.unwrap().is_ok());
- assert!(stream.next().await.unwrap().is_err());
- assert!(stream.next().await.is_none());
- assert_eq!(drops.load(Ordering::SeqCst), 1);
- assert!(polls.load(Ordering::SeqCst) <= 3);
-}
-
-#[tokio::test]
-async fn empty_batch_with_changed_schema_is_terminal_before_skip() {
- let declared = schema();
- let changed = Arc::new(Schema::new(vec![Field::new(
- "different",
- DataType::Int64,
- false,
- )]));
- let (input, polls, drops) = scripted_source(
- declared,
- vec![InputItem::Batch(Ok(RecordBatch::new_empty(changed)))],
- );
- let mut stream = record_batches_to_raw_series(input, mapping()).unwrap();
- assert!(matches!(stream.next().await, Some(Err(_))));
- assert!(stream.next().await.is_none());
- assert_eq!(polls.load(Ordering::SeqCst), 1);
- assert_eq!(drops.load(Ordering::SeqCst), 1);
-}
-
-#[tokio::test]
-async fn null_timestamp_is_terminal() {
- let nullable = Arc::new(Schema::new(vec![
- Field::new(
- "time",
- DataType::Timestamp(TimeUnit::Millisecond, None),
- true,
- ),
- Field::new("host", DataType::Utf8, true),
- Field::new("value", DataType::Float64, true),
- Field::new("predicate", DataType::Int64, false),
- ]));
- let input_batch = RecordBatch::try_new(
- nullable.clone(),
- vec![
- Arc::new(TimestampMillisecondArray::from(vec![None])),
- Arc::new(StringArray::from(vec![Some("h")])),
- Arc::new(Float64Array::from(vec![Some(1.0)])),
- Arc::new(Int64Array::from(vec![1])),
- ],
- )
- .unwrap();
- let (input, _, _) = source_custom(nullable, vec![input_batch]);
- let mut stream = record_batches_to_raw_series(input, mapping()).unwrap();
- assert!(stream.next().await.unwrap().is_err());
- assert!(stream.next().await.is_none());
-}
-
-#[tokio::test]
-async fn rejects_descending_time_inside_source_batch() {
- let (input, _, _) = source(vec![batch(
- &[2, 1],
- &[Some("h"), Some("h")],
- &[Some(1.0), Some(2.0)],
- &[1, 1],
- )]);
- let mut output = record_batches_to_raw_series(input, mapping()).unwrap();
- assert!(output.next().await.unwrap().is_err());
- assert!(output.next().await.is_none());
-}
-
-#[tokio::test]
-async fn exact_int64_filter_sort_adapter_frame_pipeline() {
- let schema = Arc::new(Schema::new(vec![
- Field::new(
- "time",
- DataType::Timestamp(TimeUnit::Millisecond, None),
- false,
- ),
- Field::new("host", DataType::Utf8, false),
- Field::new("value", DataType::Float64, false),
- Field::new("gate", DataType::Int64, false),
- ]));
- let threshold = 9_007_199_254_740_993_i64;
- let batch = RecordBatch::try_new(
- schema.clone(),
- vec![
- Arc::new(TimestampMillisecondArray::from(vec![4, 2, 1, 3])),
- Arc::new(StringArray::from(vec!["h", "h", "h", "h"])),
- Arc::new(Float64Array::from(vec![40.0, 20.0, 10.0, 30.0])),
- Arc::new(Int64Array::from(vec![
- threshold - 1,
- threshold,
- threshold,
- threshold + 1,
- ])),
- ],
- )
- .unwrap();
- let source = MemorySourceConfig::try_new_exec(&[vec![batch]], schema.clone(), None).unwrap();
- let gate: Arc = Arc::new(Column::new("gate", 3));
- let predicate = Arc::new(BinaryExpr::new(
- gate,
- Operator::Eq,
- Arc::new(Literal::new(ScalarValue::Int64(Some(threshold)))),
- ));
- let filter = Arc::new(FilterExec::try_new(predicate, source).unwrap());
- let time: Arc = Arc::new(Column::new("time", 0));
- let host: Arc = Arc::new(Column::new("host", 1));
- let ordering = LexOrdering::new(vec![
- PhysicalSortExpr::new_default(host),
- PhysicalSortExpr::new_default(time),
- ])
- .unwrap();
- let sort = SortExec::new(ordering, filter);
- let ctx = SessionContext::new();
- let input = sort.execute(0, ctx.task_ctx()).unwrap();
- let mapping = InputMapping {
- time_index: 0,
- metric_name: Some(Arc::from("cpu")),
- tags: vec![(Arc::from("host"), 1)],
- fields: vec![(
- FieldSpec {
- name: Arc::from("value"),
- },
- 2,
- )],
- };
- let raw = record_batches_to_raw_series(input, mapping).unwrap();
- let frames = super::super::pack_raw_series(raw).unwrap();
- let expected = Arc::new(datatypes::series::samples::SeriesSchema {
- fields: vec![FieldSpec {
- name: Arc::from("value"),
- }]
- .into(),
- eval_grid: None,
- });
- let mut decoded = super::super::unpack_raw_series(frames, Some(expected)).unwrap();
- let chunk = decoded.next().await.unwrap().unwrap();
- assert!(chunk.is_last);
- assert_eq!(chunk.time.as_buffer().as_ref(), &[1, 2]);
- assert_eq!(chunk.metric.as_ref().unwrap().name(), Some("cpu"));
- assert_eq!(chunk.metric.as_ref().unwrap().get("host"), Some("h"));
- let DenseSamplesRef::Float(values) = chunk.fields.as_slice()[0].segments()[0].dense().unwrap()
- else {
- panic!("float expected")
- };
- assert_eq!(values.value(0), 10.0);
- assert_eq!(values.value(1), 20.0);
- assert!(decoded.next().await.is_none());
-}
-
-#[tokio::test]
-async fn empty_input_finishes_with_explicit_contiguous_metadata() {
- let (input, polls, drops) = source(vec![RecordBatch::new_empty(schema())]);
- let mut stream = record_batches_to_raw_series(input, mapping()).unwrap();
- assert!(matches!(
- stream.label_source(),
- store_api::region_engine::RawLabelSource::Explicit
- ));
- assert_eq!(stream.ordering(), SeriesOrdering::Contiguous);
- assert!(stream.next().await.is_none());
- assert_eq!(polls.load(Ordering::SeqCst), 2);
- assert_eq!(drops.load(Ordering::SeqCst), 1);
-}
-
-#[test]
-fn empty_batches_are_bounded_per_poll_and_eventually_progress() {
- let mut batches = vec![RecordBatch::new_empty(schema()); 65];
- batches.push(batch(&[7], &[Some("h")], &[Some(7.0)], &[7]));
- let (input, polls, _) = source(batches);
- let mut stream = record_batches_to_raw_series(input, mapping()).unwrap();
- let wakes = Arc::new(AtomicUsize::new(0));
- let waker = std::task::Waker::from(Arc::new(CountWake(wakes.clone())));
- let mut cx = Context::from_waker(&waker);
- assert!(matches!(
- Pin::new(&mut stream).poll_next(&mut cx),
- Poll::Pending
- ));
- assert_eq!(polls.load(Ordering::SeqCst), 32);
- assert_eq!(wakes.load(Ordering::SeqCst), 1);
- assert!(matches!(
- Pin::new(&mut stream).poll_next(&mut cx),
- Poll::Pending
- ));
- assert_eq!(polls.load(Ordering::SeqCst), 64);
- assert_eq!(wakes.load(Ordering::SeqCst), 2);
- let first = Pin::new(&mut stream).poll_next(&mut cx);
- let Poll::Ready(Some(Ok(chunk))) = first else {
- panic!("expected row chunk")
- };
- assert_eq!(chunk.time.timestamp(0), Some(7));
- assert!(chunk.is_last);
- assert!(matches!(
- Pin::new(&mut stream).poll_next(&mut cx),
- Poll::Ready(None)
- ));
- assert_eq!(polls.load(Ordering::SeqCst), 67);
-}
-
-#[tokio::test]
-async fn output_schema_and_field_order_follow_logical_mapping() {
- let schema = Arc::new(Schema::new(vec![
- Field::new(
- "time",
- DataType::Timestamp(TimeUnit::Millisecond, None),
- false,
- ),
- Field::new("host", DataType::Utf8, true),
- Field::new("first", DataType::Float64, true),
- Field::new("second", DataType::Float64, true),
- Field::new("predicate", DataType::Int64, false),
- ]));
- let batch = RecordBatch::try_new(
- schema.clone(),
- vec![
- Arc::new(TimestampMillisecondArray::from(vec![1])),
- Arc::new(StringArray::from(vec![Some("h")])),
- Arc::new(Float64Array::from(vec![Some(11.0)])),
- Arc::new(Float64Array::from(vec![Some(22.0)])),
- Arc::new(Int64Array::from(vec![7])),
- ],
- )
- .unwrap();
- let input_mapping = InputMapping {
- time_index: 0,
- metric_name: Some(Arc::from("cpu")),
- tags: vec![(Arc::from("host"), 1)],
- fields: vec![
- (
- FieldSpec {
- name: Arc::from("logical_second"),
- },
- 3,
- ),
- (
- FieldSpec {
- name: Arc::from("logical_first"),
- },
- 2,
- ),
- ],
- };
- let (input, _, _) = source_custom(schema, vec![batch]);
- let mut stream = record_batches_to_raw_series(input, input_mapping).unwrap();
- assert_eq!(
- stream
- .schema()
- .fields
- .iter()
- .map(|field| field.name.as_ref())
- .collect::>(),
- ["logical_second", "logical_first"]
- );
- let chunk = stream.next().await.unwrap().unwrap();
- let DenseSamplesRef::Float(second) = chunk.fields.as_slice()[0].segments()[0].dense().unwrap()
- else {
- panic!("float expected")
- };
- let DenseSamplesRef::Float(first) = chunk.fields.as_slice()[1].segments()[0].dense().unwrap()
- else {
- panic!("float expected")
- };
- assert_eq!(second.value(0), 22.0);
- assert_eq!(first.value(0), 11.0);
- assert_eq!(chunk.metric.as_ref().unwrap().name(), Some("cpu"));
- assert!(chunk.is_last);
-}
diff --git a/src/common/query/src/series_frame/tests/evaluated.rs b/src/common/query/src/series_frame/tests/evaluated.rs
index dba22e8d702..b936247be64 100644
--- a/src/common/query/src/series_frame/tests/evaluated.rs
+++ b/src/common/query/src/series_frame/tests/evaluated.rs
@@ -388,7 +388,16 @@ async fn evaluated_pack_drains_large_sparse_chunk_before_polling_again() {
let expected: Vec> = (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 = (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::(),
- 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 = rows.iter().flat_map(|(row, _)| positions(row)).collect();
assert_eq!(
diff --git a/src/common/recordbatch/src/adapter.rs b/src/common/recordbatch/src/adapter.rs
index 19cb403307f..12a617d4858 100644
--- a/src/common/recordbatch/src/adapter.rs
+++ b/src/common/recordbatch/src/adapter.rs
@@ -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,
diff --git a/src/common/recordbatch/src/error.rs b/src/common/recordbatch/src/error.rs
index bddb8ae8032..df706f5d9b8 100644
--- a/src/common/recordbatch/src/error.rs
+++ b/src/common/recordbatch/src/error.rs
@@ -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"));
+ }
}
}
}
diff --git a/src/datatypes/src/series/frame.rs b/src/datatypes/src/series/frame.rs
index 4271c6c6f02..3bdcf581c65 100644
--- a/src/datatypes/src/series/frame.rs
+++ b/src/datatypes/src/series/frame.rs
@@ -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 {
+ 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> {
+ fn next_row(&mut self, available_bytes: Option) -> Result> {
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> {
+ /// 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) -> Result> {
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 = 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::>>()?;
+ }
+ *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 {
+ 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 {
+ 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::()
+ .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 {
+ samples.iter().try_fold(0usize, |bytes, sample| {
+ add_checked(bytes, sample_charge(sample)?)
+ })
+}
+
+fn sample_charge(sample: &Samples) -> Result {
+ 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 {
+ 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 {
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 {
fn histogram_charge(values: &HistogramSamples) -> Result {
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::()
- .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.
diff --git a/src/datatypes/src/series/frame/tests.rs b/src/datatypes/src/series/frame/tests.rs
index aa748a528f4..bf6e5e35d4c 100644
--- a/src/datatypes/src/series/frame/tests.rs
+++ b/src/datatypes/src/series/frame/tests.rs
@@ -249,6 +249,28 @@ fn assert_float_bits(actual: &Samples, expected: &[Option]) {
}
}
+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, 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![(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::>(),
+ );
+ let table = histogram_batch();
+ let indices = UInt32Array::from(
+ (0..9)
+ .map(|index| (index != 3 && index != 6).then_some(0))
+ .collect::>(),
+ );
+ let mixed_indices = UInt32Array::from(
+ (0..9)
+ .map(|index| (index == 3 || index == 6).then_some(0))
+ .collect::>(),
+ );
+ 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::(),
+ expected_rows
+ );
+ let decoded = batches
+ .iter()
+ .flat_map(|batch| {
+ SeriesFrameDecoder::try_new(&schema, FrameLabelMode::Explicit, None)
+ .unwrap()
+ .decode_evaluated(batch)
+ .unwrap()
+ })
+ .collect::>();
+ assert_eq!(
+ decoded
+ .iter()
+ .map(|(series, _)| series.time.len())
+ .sum::(),
+ 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::>(),
+ [(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::>();
+ 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::>();
+ 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::>();
+ let histograms = (0..128)
+ .map(|index| (!missing(index)).then_some(0))
+ .collect::>();
+ let mixed_floats = (0..128)
+ .map(|index| (!missing(index) && index % 2 == 0).then(|| floats[index].unwrap()))
+ .collect::>();
+ let mixed_histograms = (0..128)
+ .map(|index| (!missing(index) && index % 2 == 1).then_some(0))
+ .collect::>();
+ let prefix_float = |values: Vec>| {
+ Float64Array::from(
+ std::iter::once(Some(999.0))
+ .chain(values)
+ .collect::>(),
+ )
+ .slice(1, 128)
+ };
+ let prefix_indices = |values: Vec>| {
+ UInt32Array::from(std::iter::once(Some(999)).chain(values).collect::>())
+ .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::>(),
+ spans
+ .iter()
+ .map(|span| (span.start(), span.len()))
+ .collect::>()
+ );
+ 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![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![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::>();
+ assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 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![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::>();
+ assert_eq!(
+ decoded
+ .iter()
+ .flat_map(|(series, _)| series.fields.get(0).unwrap().iter().map(
+ |sample| match sample {
+ SampleRef::Float(value) => value,
+ _ => panic!(),
+ }
+ ))
+ .collect::>(),
+ 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![(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;
diff --git a/src/mito-codec/src/labels.rs b/src/mito-codec/src/labels.rs
index 70df9c447c3..8a6de5b8d56 100644
--- a/src/mito-codec/src/labels.rs
+++ b/src/mito-codec/src/labels.rs
@@ -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]>),
-}
-
-/// 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