From 1dbc2c15ce3c718359cf4cd9958875cc686da483 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Tue, 27 Oct 2026 03:42:27 +0000 Subject: [PATCH 1/2] fix(mito2): stop the series distributor from stalling on unread partitions (#9991) When every partition channel is full, the distributor walked the senders round-robin with a 100us send timeout each. Tokio timers fire on a 1ms tick, so a consumer that leaves some partitions unpolled costs about one millisecond per unpolled partition for every series batch. SeriesScan::build_stream and scan_all_partitions drain partitions one after another, and SortPreservingMergeExec above a multi-partition SeriesScan only pulls the partition holding the current minimum key. With 16 partitions and one batch per series after a merge, a 40-series scan read partition by partition took ~480ms (measured on the native PromQL two-input route). Wait on all senders at once and send to the first one that frees a slot. The send-timeout metric goes away with the timeout. Signed-off-by: Dennis Zhuang (cherry picked from commit e6b60644b3ce10b96d8be3d8eb11c43c8d18d0ed) Signed-off-by: Dennis Zhuang --- src/mito2/src/read/scan_util.rs | 10 ---- src/mito2/src/read/series_scan.rs | 99 +++++++++++++++++-------------- 2 files changed, 56 insertions(+), 53 deletions(-) diff --git a/src/mito2/src/read/scan_util.rs b/src/mito2/src/read/scan_util.rs index 7224143f2c4..a932942e2d1 100644 --- a/src/mito2/src/read/scan_util.rs +++ b/src/mito2/src/read/scan_util.rs @@ -202,8 +202,6 @@ pub(crate) struct ScanMetricsSet { /// Elapsed time before the first poll operation. first_poll: Duration, - /// Number of send timeout in SeriesScan. - num_series_send_timeout: usize, /// Number of send full in SeriesScan. num_series_send_full: usize, /// Number of rows the series distributor scanned. @@ -325,7 +323,6 @@ impl fmt::Debug for ScanMetricsSet { num_sst_batches, num_sst_rows, first_poll, - num_series_send_timeout, num_series_send_full, num_distributor_rows, num_distributor_batches, @@ -459,9 +456,6 @@ impl fmt::Debug for ScanMetricsSet { } // Write non-zero distributor metrics - if *num_series_send_timeout > 0 { - write!(f, ", \"num_series_send_timeout\":{num_series_send_timeout}")?; - } if *num_series_send_full > 0 { write!(f, ", \"num_series_send_full\":{num_series_send_full}")?; } @@ -775,7 +769,6 @@ impl ScanMetricsSet { /// Sets distributor metrics. fn set_distributor_metrics(&mut self, distributor_metrics: &SeriesDistributorMetrics) { let SeriesDistributorMetrics { - num_series_send_timeout, num_series_send_full, num_rows, num_batches, @@ -784,7 +777,6 @@ impl ScanMetricsSet { divider_cost, } = distributor_metrics; - self.num_series_send_timeout += *num_series_send_timeout; self.num_series_send_full += *num_series_send_full; self.num_distributor_rows += *num_rows; self.num_distributor_batches += *num_batches; @@ -1144,8 +1136,6 @@ impl fmt::Debug for PartitionMetrics { /// Metrics for the series distributor. #[derive(Default)] pub(crate) struct SeriesDistributorMetrics { - /// Number of send timeout in SeriesScan. - pub(crate) num_series_send_timeout: usize, /// Number of send full in SeriesScan. pub(crate) num_series_send_full: usize, /// Number of rows the series distributor scanned. diff --git a/src/mito2/src/read/series_scan.rs b/src/mito2/src/read/series_scan.rs index 8fb06782240..b8c7533efa7 100644 --- a/src/mito2/src/read/series_scan.rs +++ b/src/mito2/src/read/series_scan.rs @@ -21,7 +21,7 @@ use std::collections::HashSet; use std::fmt; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; -use std::time::{Duration, Instant}; +use std::time::Instant; use api::v1::SemanticType; use async_stream::try_stream; @@ -55,7 +55,7 @@ use store_api::region_engine::{ use store_api::storage::consts::{ReservedColumnId, is_internal_column}; use table::predicate::Predicate; use tokio::sync::Semaphore; -use tokio::sync::mpsc::error::{SendTimeoutError, TrySendError}; +use tokio::sync::mpsc::error::TrySendError; use tokio::sync::mpsc::{self, Receiver, Sender}; use crate::error::{ @@ -76,9 +76,6 @@ use crate::read::stream::{ConvertBatchStream, ScanBatch, ScanBatchStream}; use crate::sst::parquet::flat_format::primary_key_column_index; use crate::sst::parquet::format::PrimaryKeyArray; -/// Timeout to send a batch to a sender. -const SEND_TIMEOUT: Duration = Duration::from_micros(100); - /// Maximum number of candidate series retained before distributing assignments. const CANDIDATE_SERIES_ASSIGNMENT_THRESHOLD: usize = 1_000_000; @@ -1268,7 +1265,6 @@ impl SeriesDistributor { } metrics.scan_cost += fetch_start.elapsed(); - metrics.num_series_send_timeout = self.senders.num_timeout; metrics.num_series_send_full = self.senders.num_full; part_metrics.set_distributor_metrics(&metrics); @@ -1423,8 +1419,6 @@ struct SenderList { num_nones: usize, /// Index of the current partition to send. sender_idx: usize, - /// Number of timeout. - num_timeout: usize, /// Number of full senders. num_full: usize, } @@ -1436,7 +1430,6 @@ impl SenderList { senders, num_nones, sender_idx: 0, - num_timeout: 0, num_full: 0, } } @@ -1472,46 +1465,35 @@ impl SenderList { } /// Finds a partition and sends the batch to the partition. - async fn send_batch(&mut self, mut batch: SeriesBatch) -> Result<()> { - // Sends the batch without blocking first. - match self.try_send_batch(batch)? { - Some(b) => { - // Unable to send batch to partition. - batch = b; - } - None => { - return Ok(()); - } - } + async fn send_batch(&mut self, batch: SeriesBatch) -> Result<()> { + let Some(batch) = self.try_send_batch(batch)? else { + return Ok(()); + }; loop { ensure!(self.num_nones < self.senders.len(), InvalidSenderSnafu); - let sender_idx = self.fetch_add_sender_idx(); - let Some(sender) = &self.senders[sender_idx] else { - continue; + // Waits on every partition at once and sends to whichever frees a slot first. A + // consumer may leave partitions unpolled for a long time, e.g. one that drains them one + // after another, and waiting on such a partition would stall the scan. + let closed = { + let reservations = self.senders.iter().enumerate().filter_map(|(idx, sender)| { + sender + .as_ref() + .map(|sender| Box::pin(async move { (idx, sender.reserve().await) })) + }); + let ((idx, reserved), _, _) = futures::future::select_all(reservations).await; + match reserved { + Ok(permit) => { + permit.send(Ok(batch)); + return Ok(()); + } + Err(_) => idx, + } }; - // Adds a timeout to avoid blocking indefinitely and sending - // the batch in a round-robin fashion when some partitions - // don't poll their inputs. This may happen if we have a - // node like sort merging. But it is rare when we are using SeriesScan. - match sender.send_timeout(Ok(batch), SEND_TIMEOUT).await { - Ok(()) => break, - Err(SendTimeoutError::Timeout(res)) => { - self.num_timeout += 1; - // Safety: we send Ok. - batch = res.unwrap(); - } - Err(SendTimeoutError::Closed(res)) => { - self.senders[sender_idx] = None; - self.num_nones += 1; - // Safety: we send Ok. - batch = res.unwrap(); - } - } + self.senders[closed] = None; + self.num_nones += 1; } - - Ok(()) } async fn send_error(&self, error: Error) { @@ -1721,6 +1703,37 @@ mod tests { assert_eq!(0, active_receivers.load(Ordering::Relaxed)); } + #[tokio::test] + async fn distributor_does_not_stall_on_partitions_read_one_after_another() { + const NUM_BATCHES: usize = 2000; + let (mut senders, receivers) = new_legacy_channel_list(16); + let producer = tokio::spawn(async move { + for _ in 0..NUM_BATCHES { + senders + .send_batch(SeriesBatch::Flat(FlatSeriesBatch::default())) + .await + .unwrap(); + } + }); + // Reads partition 0 to its end before touching the next one, so all other partitions stay + // full while the producer still has batches to send. + let consume = async { + let mut received = 0; + for mut receiver in receivers.into_iter().flatten() { + while let Some(batch) = receiver.recv().await { + batch.unwrap(); + received += 1; + } + } + received + }; + let received = tokio::time::timeout(std::time::Duration::from_secs(1), consume) + .await + .expect("the distributor stalled on partitions the consumer was not reading"); + producer.await.unwrap(); + assert_eq!(NUM_BATCHES, received); + } + fn new_test_record_batch( primary_keys: &[&[u8]], timestamps: &[i64], From b29a47157550e1c216c911705744d5f545c1b77d Mon Sep 17 00:00:00 2001 From: discord9 Date: Tue, 27 Oct 2026 17:41:38 +0800 Subject: [PATCH 2/2] fix(promql): repair native execution contracts and remove obsolete paths (#9980) * fix(promql): restore native scan statistics and remove unused raw merger Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(promql): repair native execution contracts and remove obsolete paths Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * refactor(promql): leave fused frame output repairs to separate PR Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * docs: restore native contracts and regenerate config reference Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix: slice nullable sparse frame fields after span selection Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(promql): intersect anchored timestamp selection windows Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(query): name anchored timestamp boundary case type Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(promql): align partitioned pair checks with hybrid operands Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> Signed-off-by: discord9 --- config/config.md | 6 +- config/datanode.example.toml | 3 - config/flownode.example.toml | 3 - config/frontend.example.toml | 5 +- config/standalone.example.toml | 3 +- src/common/query/src/request.rs | 3 - .../query/src/request/raw_series_request.rs | 249 -- src/common/query/src/series_frame.rs | 4 - .../query/src/series_frame/record_batches.rs | 393 ---- .../src/series_frame/record_batches/tests.rs | 1263 ----------- .../query/src/series_frame/tests/evaluated.rs | 45 +- src/common/recordbatch/src/adapter.rs | 13 +- src/common/recordbatch/src/error.rs | 14 + src/datatypes/src/series/frame.rs | 312 ++- src/datatypes/src/series/frame/tests.rs | 516 ++++- src/mito-codec/src/labels.rs | 170 +- src/mito-codec/src/labels/tests.rs | 643 +----- src/mito2/src/engine/scan_test.rs | 33 +- src/mito2/src/read/series_scan.rs | 11 +- .../src/read/series_scan/prepared/tests.rs | 31 +- src/promql/AGENTS.md | 138 +- .../src/extension_plan/aggregate_finalize.rs | 30 +- .../aggregate_finalize/exec_tests.rs | 81 + src/promql/src/extension_plan/fused_series.rs | 11 +- .../src/extension_plan/fused_series/tests.rs | 5 + .../src/extension_plan/instant_manipulate.rs | 151 +- .../instant_manipulate/native.rs | 301 --- .../instant_manipulate/native_exec.rs | 500 ---- .../instant_manipulate/native_exec_tests.rs | 427 ---- .../instant_manipulate/native_tests.rs | 597 ----- src/promql/src/instant_eval.rs | 7 +- src/promql/src/instant_eval/tests.rs | 1 + src/promql/src/native/aggregate.rs | 8 +- src/promql/src/native/execute.rs | 263 ++- src/promql/src/native/execute/tests.rs | 171 +- src/promql/src/range_eval/function.rs | 9 +- src/promql/src/source.rs | 31 +- src/promql/src/source/merge.rs | 459 ---- src/promql/src/source/merge/tests.rs | 2020 ----------------- src/promql/src/source/merge/tests/payload.rs | 654 ------ src/query/src/datafusion.rs | 146 +- src/query/src/dist_plan/merge_scan.rs | 291 +-- src/query/src/dist_plan/planner.rs | 1045 +-------- src/query/src/options.rs | 6 +- src/query/src/promql/native_compile.rs | 12 +- src/query/src/promql/native_compile/tests.rs | 39 + src/query/src/promql/planner/native.rs | 40 +- src/query/src/query_engine/state.rs | 3 - src/query/tests/native_instant_at.rs | 129 ++ src/store-api/src/region_engine.rs | 9 +- src/table/src/table/scan.rs | 239 +- src/table/src/table/scan/series_tests.rs | 171 +- tests-integration/src/tests/promql_test.rs | 126 +- 53 files changed, 2208 insertions(+), 9632 deletions(-) delete mode 100644 src/common/query/src/request/raw_series_request.rs delete mode 100644 src/common/query/src/series_frame/record_batches.rs delete mode 100644 src/common/query/src/series_frame/record_batches/tests.rs delete mode 100644 src/promql/src/extension_plan/instant_manipulate/native.rs delete mode 100644 src/promql/src/extension_plan/instant_manipulate/native_exec.rs delete mode 100644 src/promql/src/extension_plan/instant_manipulate/native_exec_tests.rs delete mode 100644 src/promql/src/extension_plan/instant_manipulate/native_tests.rs delete mode 100644 src/promql/src/source/merge.rs delete mode 100644 src/promql/src/source/merge/tests.rs delete mode 100644 src/promql/src/source/merge/tests/payload.rs 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