From 535923629b8418348cee15a3e4bed9fc6a8ef265 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Sat, 17 Oct 2026 13:23:54 +0800 Subject: [PATCH] feat(promql): merge ordered raw series inputs Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/promql/src/source.rs | 3 + src/promql/src/source/merge.rs | 450 +++++ src/promql/src/source/merge/tests.rs | 1822 ++++++++++++++++++ src/promql/src/source/merge/tests/payload.rs | 627 ++++++ 4 files changed, 2902 insertions(+) create mode 100644 src/promql/src/source/merge.rs create mode 100644 src/promql/src/source/merge/tests.rs create mode 100644 src/promql/src/source/merge/tests/payload.rs diff --git a/src/promql/src/source.rs b/src/promql/src/source.rs index 6feefa79e86..acd06605f19 100644 --- a/src/promql/src/source.rs +++ b/src/promql/src/source.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +mod merge; #[cfg(test)] mod tests; @@ -28,6 +29,8 @@ use futures::Stream; use mito_codec::labels::{DecodedLabels, LabelProjection, PrimaryKeyLabelDecoder}; use store_api::region_engine::{RawLabelSource, RawSeriesStream, SendableRawSeriesStream}; +pub use merge::merge_ordered_raw_series; + /// Materializes complete metric labels from a scan-bound raw-series source. /// /// The caller supplies an already-bound source. Adjacent-chunk pack/unpack validation remains diff --git a/src/promql/src/source/merge.rs b/src/promql/src/source/merge.rs new file mode 100644 index 00000000000..a734d63c51b --- /dev/null +++ b/src/promql/src/source/merge.rs @@ -0,0 +1,450 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#[cfg(test)] +mod tests; + +use std::cmp::Ordering; +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use common_error::ext::BoxedError; +use datatypes::error::InvalidSeriesSnafu; +use datatypes::series::SeriesOrdering; +use datatypes::series::raw::RawSeriesChunk; +use datatypes::series::samples::{SeriesFields, SeriesSchema}; +use futures::Stream; +use mito_codec::labels::{LabelProjection, PrimaryKeyLabelDecoder}; +use store_api::region_engine::{RawLabelSource, RawSeriesStream, SendableRawSeriesStream}; + +/// Merge compatible, explicitly supplied raw streams by complete read key and time. +/// +/// Callers provide authorized bindings and own snapshot/storage policy. This does not select an +/// MVCC view, resolve overlapping physical copies, or roll back emitted prefixes after errors. +/// `is_last` is global only over these inputs; there is no fallback, replay, or input sorting. +pub fn merge_ordered_raw_series( + inputs: Vec, +) -> Result { + let mut inputs = inputs.into_iter(); + let first = inputs + .next() + .ok_or_else(|| invalid("raw merge requires at least one input"))?; + let schema = Arc::clone(first.schema()); + let binding = first.label_source().clone(); + let mut streams = vec![Input::new(first)]; + for stream in inputs { + streams.push(Input::new(stream)); + } + for input in &streams { + let stream = input + .stream + .as_ref() + .ok_or_else(|| invalid("raw merge input disappeared"))?; + validate_binding(stream.label_source())?; + validate_input(stream.as_ref().get_ref(), &schema, &binding)?; + } + Ok(Box::pin(OrderedRawMerge { + streams, + schema, + binding, + candidate: None, + })) +} + +fn invalid(reason: &'static str) -> BoxedError { + BoxedError::new(InvalidSeriesSnafu { reason }.build()) +} + +fn validate_binding(binding: &RawLabelSource) -> Result<(), BoxedError> { + PrimaryKeyLabelDecoder::new(binding, LabelProjection::All) + .map(|_| ()) + .map_err(BoxedError::new) +} + +fn validate_input( + stream: &dyn RawSeriesStream, + schema: &SeriesSchema, + binding: &RawLabelSource, +) -> Result<(), BoxedError> { + let other = stream.schema(); + if stream.ordering() != SeriesOrdering::KeyThenTime + || other.eval_grid.is_some() + || other.fields.is_empty() + || other.fields.len() != schema.fields.len() + || other + .fields + .iter() + .zip(schema.fields.iter()) + .any(|(a, b)| a.name != b.name) + || !bindings_compatible(stream.label_source(), binding) + { + return Err(invalid( + "raw merge requires equal raw schemas, primary-key bindings, and ReadKeyThenTime ordering", + )); + } + Ok(()) +} + +fn bindings_compatible(left: &RawLabelSource, right: &RawLabelSource) -> bool { + let ( + RawLabelSource::PrimaryKey { + metric_name: lm, + physical_metadata: lmeta, + logical_tags: lt, + }, + RawLabelSource::PrimaryKey { + metric_name: rm, + physical_metadata: rmeta, + logical_tags: rt, + }, + ) = (left, right) + else { + return false; + }; + if lm != rm + || lt.len() != rt.len() + || lt.iter().any(|tag| { + !rt.iter() + .any(|other| tag.column_id == other.column_id && tag.name == other.name) + }) + { + return false; + } + let lm = lmeta.as_ref(); + let rm = rmeta.as_ref(); + lm.primary_key == rm.primary_key + && lm.primary_key_encoding == rm.primary_key_encoding + && lm + .primary_key + .iter() + .all(|id| match (lm.column_by_id(*id), rm.column_by_id(*id)) { + (Some(a), Some(b)) => { + a.column_id == b.column_id + && a.column_schema.name == b.column_schema.name + && a.column_schema.data_type == b.column_schema.data_type + && a.semantic_type == b.semantic_type + } + _ => false, + }) +} + +struct Input { + stream: Option, + head: Option, + previous_key: Option>, + previous_time: Option, + ended_series: bool, +} + +impl Input { + fn new(stream: SendableRawSeriesStream) -> Self { + Self { + stream: Some(stream), + head: None, + previous_key: None, + previous_time: None, + ended_series: false, + } + } + + fn validate( + &mut self, + chunk: &RawSeriesChunk, + schema: &SeriesSchema, + ) -> Result<(), BoxedError> { + if chunk.metric.is_some() { + return Err(invalid("raw merge requires undecoded primary-key chunks")); + } + RawSeriesChunk::try_new( + schema, + chunk.primary_key.clone(), + None, + chunk.time.clone(), + chunk.fields.clone(), + chunk.is_last, + ) + .map_err(BoxedError::new)?; + let times = chunk.time.as_buffer(); + if let Some(previous_key) = &self.previous_key { + match previous_key.as_slice().cmp(chunk.primary_key.as_slice()) { + Ordering::Greater => { + return Err(invalid("raw merge input read keys are not ordered")); + } + Ordering::Equal + if self.ended_series + || self.previous_time.is_some_and(|time| times[0] <= time) => + { + return Err(invalid( + "raw merge input has repeated, non-increasing timestamps or invalid is_last", + )); + } + Ordering::Less if !self.ended_series => { + return Err(invalid("raw merge input changed read keys before is_last")); + } + _ => {} + } + } + self.previous_key = Some(chunk.primary_key.as_slice().to_vec()); + self.previous_time = times.last().copied(); + self.ended_series = chunk.is_last; + Ok(()) + } + + /// Returns whether this input can no longer contribute chunks for `key`. + /// + /// A held head beyond `key` proves completion. Without a head, only an ended input or `key`'s + /// own final chunk proves completion: an earlier completed key never rules out unknown later + /// or intervening keys. + fn finished_at(&self, key: &[u8]) -> bool { + match self.head.as_ref() { + Some(head) => head.primary_key.as_slice() > key, + None => { + self.stream.is_none() + || self.previous_key.as_deref().is_some_and(|previous| { + previous > key || (previous == key && self.ended_series) + }) + } + } + } +} + +struct OrderedRawMerge { + streams: Vec, + schema: Arc, + binding: RawLabelSource, + candidate: Option, +} + +impl OrderedRawMerge { + fn fail(&mut self, error: BoxedError) -> Poll>> { + self.streams.clear(); + self.candidate = None; + Poll::Ready(Some(Err(error))) + } + + /// Polls every input that has no head once, returning upstream errors before progress. + fn poll_inputs(&mut self, cx: &mut Context<'_>) -> Result<(), BoxedError> { + for input in &mut self.streams { + if input.stream.is_none() || input.head.is_some() { + continue; + } + match input + .stream + .as_mut() + .map(|stream| stream.as_mut().poll_next(cx)) + { + Some(Poll::Pending) => {} + Some(Poll::Ready(Some(Err(error)))) => return Err(error), + Some(Poll::Ready(Some(Ok(chunk)))) => { + input.validate(&chunk, &self.schema)?; + input.head = Some(chunk); + } + Some(Poll::Ready(None)) => { + if input.previous_key.is_some() && !input.ended_series { + return Err(invalid("raw merge input ended before its final chunk")); + } + input.stream = None; + input.previous_key = None; + input.previous_time = None; + } + None => {} + } + } + Ok(()) + } + + /// Returns the index of the smallest available head, by read key then first timestamp. + fn smallest_head_index(&self) -> Option { + self.streams + .iter() + .enumerate() + .filter_map(|(index, input)| input.head.as_ref().map(|head| (index, head))) + .min_by(|(_, a), (_, b)| { + a.primary_key + .as_slice() + .cmp(b.primary_key.as_slice()) + .then_with(|| a.time.timestamp(0).cmp(&b.time.timestamp(0))) + }) + .map(|(index, _)| index) + } + + /// Returns whether every input accounts for `key`: it holds a head, or completion is proven. + /// + /// A headless input that cannot prove completion may still add chunks to `key`, so the merge + /// must wait instead of deciding `key`. + fn all_inputs_ready_at(&self, key: &[u8]) -> bool { + self.streams + .iter() + .all(|input| input.head.is_some() || input.finished_at(key)) + } + + fn take_next(&mut self) -> Result, BoxedError> { + let Some(index) = self.smallest_head_index() else { + return Ok(None); + }; + let chunk = self.streams[index] + .head + .take() + .ok_or_else(|| invalid("raw merge lost selected head"))?; + let key = chunk.primary_key.as_slice(); + let start = chunk + .time + .timestamp(0) + .ok_or_else(|| invalid("raw merge chunk has no timestamp"))?; + let mut end = chunk.time.len(); + for input in &self.streams { + if let Some(head) = &input.head + && head.primary_key.as_slice() == key + { + let next = head + .time + .timestamp(0) + .ok_or_else(|| invalid("raw merge chunk has no timestamp"))?; + if next == start { + return Err(invalid( + "raw merge inputs contain duplicate timestamps for a complete read key", + )); + } + if next > start { + end = end.min( + chunk + .time + .as_buffer() + .partition_point(|timestamp| *timestamp < next), + ); + } + } + } + if end == chunk.time.len() { + return Ok(Some(RawSeriesChunk { + primary_key: chunk.primary_key, + metric: None, + time: chunk.time, + fields: chunk.fields, + is_last: false, + })); + } + let fields = SeriesFields::try_from_vec( + chunk + .fields + .iter() + .map(|samples| samples.slice(0, end).map_err(BoxedError::new)) + .collect::, _>>()?, + ) + .map_err(BoxedError::new)?; + let time = chunk.time.slice(0, end).map_err(BoxedError::new)?; + if end < chunk.time.len() { + let remainder = SeriesFields::try_from_vec( + chunk + .fields + .iter() + .map(|samples| { + samples + .slice(end, chunk.time.len() - end) + .map_err(BoxedError::new) + }) + .collect::, _>>()?, + ) + .map_err(BoxedError::new)?; + self.streams[index].head = Some(RawSeriesChunk { + primary_key: chunk.primary_key.clone(), + metric: None, + time: chunk + .time + .slice(end, chunk.time.len() - end) + .map_err(BoxedError::new)?, + fields: remainder, + is_last: chunk.is_last, + }); + } + Ok(Some(RawSeriesChunk { + primary_key: chunk.primary_key, + metric: None, + time, + fields, + is_last: false, + })) + } +} + +impl Stream for OrderedRawMerge { + type Item = Result; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + if this.streams.is_empty() { + return Poll::Ready(None); + } + if this.candidate.is_none() { + if let Err(error) = this.poll_inputs(cx) { + return this.fail(error); + } + let Some(head) = this + .smallest_head_index() + .and_then(|index| this.streams[index].head.as_ref()) + else { + // Only a fully drained input set is EOF: a live input may still reveal future keys + // or a late error, so its state stays owned instead of being cleared. + if this.streams.iter().all(|input| input.stream.is_none()) { + this.streams.clear(); + return Poll::Ready(None); + } + return Poll::Pending; + }; + // A headless input that cannot prove completion at the smallest available key may + // still add chunks to it, so the merge waits before consuming that head. + if !this.all_inputs_ready_at(head.primary_key.as_slice()) { + return Poll::Pending; + } + match this.take_next() { + Ok(candidate) => this.candidate = candidate, + Err(error) => return this.fail(error), + } + } + if let Err(error) = this.poll_inputs(cx) { + return this.fail(error); + } + let Some(candidate) = this.candidate.take() else { + this.streams.clear(); + return Poll::Ready(None); + }; + if !this.all_inputs_ready_at(candidate.primary_key.as_slice()) { + // The candidate's key may still be extended or interleaved by a headless input: an + // already-decided prefix is held back instead of being emitted as if it were final. + this.candidate = Some(candidate); + return Poll::Pending; + } + let is_last = this + .streams + .iter() + .all(|input| input.finished_at(candidate.primary_key.as_slice())); + Poll::Ready(Some(Ok(RawSeriesChunk { + is_last, + ..candidate + }))) + } +} + +impl RawSeriesStream for OrderedRawMerge { + fn schema(&self) -> &Arc { + &self.schema + } + fn label_source(&self) -> &RawLabelSource { + &self.binding + } + fn ordering(&self) -> SeriesOrdering { + SeriesOrdering::KeyThenTime + } +} diff --git a/src/promql/src/source/merge/tests.rs b/src/promql/src/source/merge/tests.rs new file mode 100644 index 00000000000..4ad5d17cbb7 --- /dev/null +++ b/src/promql/src/source/merge/tests.rs @@ -0,0 +1,1822 @@ +// 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. + +mod payload; + +use std::collections::VecDeque; +use std::pin::Pin; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; +use std::task::{Context, Poll, Waker}; + +use api::v1::SemanticType; +use common_error::ext::{BoxedError, ErrorExt, PlainError}; +use common_error::status_code::StatusCode; +use datatypes::arrow::array::Float64Array; +use datatypes::arrow::buffer::{Buffer, ScalarBuffer}; +use datatypes::data_type::ConcreteDataType; +use datatypes::prelude::Value; +use datatypes::schema::ColumnSchema; +use datatypes::series::raw::{RawSeriesChunk, RawTimestamps}; +use datatypes::series::samples::{FieldSpec, SampleRef, Samples, SeriesFields, SeriesSchema}; +use datatypes::series::time::{EvalGrid, GridSpec}; +use datatypes::series::{MetricName, SeriesOrdering}; +use futures::{Stream, StreamExt, TryStreamExt}; +use mito_codec::row_converter::{DensePrimaryKeyCodec, PrimaryKeyCodec, SortField}; +use store_api::codec::PrimaryKeyEncoding; +use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder}; +use store_api::region_engine::{ + BoundTag, RawLabelSource, RawSeriesStream, SendableRawSeriesStream, +}; +use store_api::storage::RegionId; + +use super::{Input, merge_ordered_raw_series}; + +fn schema() -> Arc { + Arc::new(SeriesSchema { + fields: vec![FieldSpec { + name: Arc::from("temperature"), + }] + .into(), + eval_grid: None, + }) +} + +fn label_source(region: u32) -> RawLabelSource { + label_source_with(region, PrimaryKeyEncoding::Dense, "room_temperature") +} + +fn label_source_with(region: u32, encoding: PrimaryKeyEncoding, metric: &str) -> RawLabelSource { + label_source_with_tags(region, encoding, metric, vec![(8, "rack"), (9, "sensor")]) +} + +fn label_source_with_tags( + region: u32, + encoding: PrimaryKeyEncoding, + metric: &str, + tags: Vec<(u32, &str)>, +) -> RawLabelSource { + label_source_with_schema_version(region, encoding, metric, tags, 0) +} + +fn label_source_with_schema_version( + region: u32, + encoding: PrimaryKeyEncoding, + metric: &str, + tags: Vec<(u32, &str)>, + schema_version: u64, +) -> RawLabelSource { + let mut builder = RegionMetadataBuilder::new(RegionId::new(region, 1)); + for _ in 0..schema_version { + builder.bump_version(); + } + let columns = [ + ( + 0, + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + SemanticType::Timestamp, + ), + ( + 20, + "temperature", + ConcreteDataType::float64_datatype(), + SemanticType::Field, + ), + ]; + for (column_id, name, ty, semantic_type) in + columns.into_iter().chain(tags.iter().map(|(id, name)| { + ( + *id, + *name, + ConcreteDataType::string_datatype(), + SemanticType::Tag, + ) + })) + { + builder.push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new(name, ty, semantic_type != SemanticType::Timestamp), + semantic_type, + column_id, + }); + } + builder.primary_key(tags.iter().map(|(id, _)| *id).collect()); + builder.primary_key_encoding(encoding); + RawLabelSource::PrimaryKey { + metric_name: Arc::from(metric), + physical_metadata: Arc::new(builder.build().unwrap()), + logical_tags: tags + .iter() + .map(|(column_id, name)| BoundTag { + column_id: *column_id, + name: Arc::from(*name), + }) + .collect::>() + .into(), + } +} + +// Keys use the real DensePrimaryKeyCodec and the same primary-key field layout as metadata. +fn read_key(rack: &str, sensor: &str) -> Buffer { + let codec = DensePrimaryKeyCodec::with_fields(vec![ + (8, SortField::new(ConcreteDataType::string_datatype())), + (9, SortField::new(ConcreteDataType::string_datatype())), + ]); + let mut bytes = Vec::new(); + codec + .encode_values( + &[ + (8, Value::String(rack.into())), + (9, Value::String(sensor.into())), + ], + &mut bytes, + ) + .unwrap(); + Buffer::from(bytes) +} + +fn chunk(key: Buffer, times: &[i64], values: &[f64], is_last: bool) -> RawSeriesChunk { + RawSeriesChunk::try_new( + &schema(), + key, + None, + RawTimestamps::try_new(ScalarBuffer::from(times.to_vec())).unwrap(), + SeriesFields::one(Samples::from_float(Float64Array::from(values.to_vec()))), + is_last, + ) + .unwrap() +} + +struct TestStream { + schema: Arc, + labels: RawLabelSource, + ordering: SeriesOrdering, + items: VecDeque>>>, + // Delivered once the gate opens, so a blocked body can resume with later chunks. + deferred: VecDeque>>>, + polls: Arc, + drops: Arc, + gate: Option>>>, + release: Option>, +} + +impl Stream for TestStream { + type Item = Result; + + fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + self.polls.fetch_add(1, Ordering::SeqCst); + if let Some(gate) = &self.gate + && self.items.is_empty() + && self.release.as_ref().unwrap().load(Ordering::SeqCst) == 0 + { + *gate.lock().unwrap() = Some(_cx.waker().clone()); + return Poll::Pending; + } + self.items + .pop_front() + .or_else(|| self.deferred.pop_front()) + .unwrap_or(Poll::Ready(None)) + } +} + +impl Drop for TestStream { + fn drop(&mut self) { + self.drops.fetch_add(1, Ordering::SeqCst); + } +} + +impl RawSeriesStream for TestStream { + fn schema(&self) -> &Arc { + &self.schema + } + fn label_source(&self) -> &RawLabelSource { + &self.labels + } + fn ordering(&self) -> SeriesOrdering { + self.ordering + } +} + +fn stream( + labels: RawLabelSource, + ordering: SeriesOrdering, + items: impl IntoIterator>>>, +) -> (SendableRawSeriesStream, Arc, Arc) { + stream_with(schema(), labels, ordering, items) +} + +fn stream_with( + schema: Arc, + labels: RawLabelSource, + ordering: SeriesOrdering, + items: impl IntoIterator>>>, +) -> (SendableRawSeriesStream, Arc, Arc) { + let polls = Arc::new(AtomicUsize::new(0)); + let drops = Arc::new(AtomicUsize::new(0)); + ( + Box::pin(TestStream { + schema, + labels, + ordering, + items: items.into_iter().collect(), + deferred: VecDeque::new(), + polls: polls.clone(), + drops: drops.clone(), + gate: None, + release: None, + }), + polls, + drops, + ) +} + +fn gated_stream( + labels: RawLabelSource, + first: RawSeriesChunk, + gate: Arc>>, + release: Arc, +) -> (SendableRawSeriesStream, Arc, Arc) { + let polls = Arc::new(AtomicUsize::new(0)); + let drops = Arc::new(AtomicUsize::new(0)); + ( + Box::pin(TestStream { + schema: schema(), + labels, + ordering: SeriesOrdering::KeyThenTime, + items: [ready(first)].into(), + deferred: VecDeque::new(), + polls: polls.clone(), + drops: drops.clone(), + gate: Some(gate), + release: Some(release), + }), + polls, + drops, + ) +} + +// Yields `first` immediately, then blocks while `release` is zero, then yields `resumed` and EOF. +fn gated_resuming_stream( + labels: RawLabelSource, + first: RawSeriesChunk, + resumed: impl IntoIterator>>>, + gate: Arc>>, + release: Arc, +) -> (SendableRawSeriesStream, Arc, Arc) { + let polls = Arc::new(AtomicUsize::new(0)); + let drops = Arc::new(AtomicUsize::new(0)); + ( + Box::pin(TestStream { + schema: schema(), + labels, + ordering: SeriesOrdering::KeyThenTime, + items: [ready(first)].into(), + deferred: resumed.into_iter().collect(), + polls: polls.clone(), + drops: drops.clone(), + gate: Some(gate), + release: Some(release), + }), + polls, + drops, + ) +} + +struct KeyOwner { + bytes: Vec, + drops: Arc, +} + +impl Drop for KeyOwner { + fn drop(&mut self) { + self.drops.fetch_add(1, Ordering::SeqCst); + } +} + +fn owned_key(bytes: Vec, drops: Arc) -> (Buffer, std::sync::Weak) { + let mut owner = KeyOwner { bytes, drops }; + let pointer = std::ptr::NonNull::new(owner.bytes.as_mut_ptr()).unwrap(); + let len = owner.bytes.len(); + let owner = Arc::new(owner); + let weak = Arc::downgrade(&owner); + // SAFETY: Buffer keeps the Arc owner alive for the entire allocation lifetime. + let buffer = unsafe { Buffer::from_custom_allocation(pointer, len, owner) }; + (buffer, weak) +} + +fn gated_pending_stream( + labels: RawLabelSource, + gate: Arc>>, + release: Arc, +) -> (SendableRawSeriesStream, Arc, Arc) { + let polls = Arc::new(AtomicUsize::new(0)); + let drops = Arc::new(AtomicUsize::new(0)); + ( + Box::pin(TestStream { + schema: schema(), + labels, + ordering: SeriesOrdering::KeyThenTime, + items: VecDeque::new(), + deferred: VecDeque::new(), + polls: polls.clone(), + drops: drops.clone(), + gate: Some(gate), + release: Some(release), + }), + polls, + drops, + ) +} + +fn ready(chunk: RawSeriesChunk) -> Poll>> { + Poll::Ready(Some(Ok(chunk))) +} + +fn float_sample(chunk: &RawSeriesChunk, index: usize) -> f64 { + match chunk.fields.get(0).unwrap().sample(index) { + SampleRef::Float(value) => value, + _ => panic!("expected float sample"), + } +} + +#[tokio::test] +async fn merges_full_encoded_keys_and_interleaved_chunks_with_global_finality() { + let rack_a_sensor_1 = read_key("rack-a", "sensor-1"); + let rack_a_sensor_2 = read_key("rack-a", "sensor-2"); + let rack_b_sensor_1 = read_key("rack-b", "sensor-1"); + let (left, _, _) = stream( + label_source(1), + SeriesOrdering::KeyThenTime, + [ + ready(chunk(rack_a_sensor_1.clone(), &[10, 30], &[1.0, 3.0], true)), + ready(chunk(rack_b_sensor_1.clone(), &[5, 25], &[5.0, 25.0], true)), + ], + ); + let (right, _, _) = stream( + label_source(2), + SeriesOrdering::KeyThenTime, + [ + ready(chunk(rack_a_sensor_1.clone(), &[20, 40], &[2.0, 4.0], true)), + ready(chunk(rack_a_sensor_2.clone(), &[7, 17], &[7.0, 17.0], true)), + ], + ); + + let mut merged = merge_ordered_raw_series(vec![left, right]).unwrap(); + assert_eq!(merged.ordering(), SeriesOrdering::KeyThenTime); + assert!(merged.schema().eval_grid.is_none()); + assert!(matches!( + merged.label_source(), + RawLabelSource::PrimaryKey { .. } + )); + let mut observed = Vec::new(); + while let Some(item) = merged.next().await { + let item = item.unwrap(); + let values = (0..item.fields.get(0).unwrap().len()) + .map(|index| match item.fields.get(0).unwrap().sample(index) { + SampleRef::Float(value) => value.to_bits(), + _ => panic!("expected float sample"), + }) + .collect::>(); + observed.push(( + item.primary_key, + item.time.as_buffer().to_vec(), + values, + item.is_last, + )); + } + let mut series = Vec::<(Buffer, Vec, Vec, Vec)>::new(); + for (key, times, values, is_last) in observed { + if series.last().is_none_or(|current| current.0 != key) { + series.push((key.clone(), Vec::new(), Vec::new(), Vec::new())); + } + let current = series.last_mut().unwrap(); + current.1.extend(times); + current.2.extend(values); + current.3.push(is_last); + } + assert_eq!(series.len(), 3); + assert_eq!(series[0].0, rack_a_sensor_1); + assert_eq!(series[0].1, [10, 20, 30, 40]); + assert_eq!( + series[0].2, + [ + 1.0_f64.to_bits(), + 2.0_f64.to_bits(), + 3.0_f64.to_bits(), + 4.0_f64.to_bits() + ] + ); + assert!(series[0].3.last().unwrap()); + assert!( + series[0].3[..series[0].3.len() - 1] + .iter() + .all(|last| !last) + ); + assert_eq!(series[1].0, rack_a_sensor_2); + assert_eq!(series[1].1, [7, 17]); + assert!(series[1].3.last().unwrap()); + assert_eq!(series[2].0, rack_b_sensor_1); + assert_eq!(series[2].1, [5, 25]); + assert!(series[2].3.last().unwrap()); +} + +#[tokio::test] +async fn rejects_equal_timestamp_from_separate_heads() { + let key = read_key("rack-a", "sensor-1"); + let (left, _, _) = stream( + label_source(1), + SeriesOrdering::KeyThenTime, + [ready(chunk(key.clone(), &[10], &[1.0], true))], + ); + let (right, _, _) = stream( + label_source(2), + SeriesOrdering::KeyThenTime, + [ready(chunk(key, &[10], &[1.0], true))], + ); + let mut merged = merge_ordered_raw_series(vec![left, right]).unwrap(); + let error = merged.next().await.unwrap().unwrap_err().to_string(); + assert!( + error.contains("timestamp") || error.contains("time"), + "unexpected error: {error}" + ); + assert!(merged.next().await.is_none()); +} + +#[tokio::test] +async fn rejects_unsupported_ordering_and_empty_input_without_polling() { + for ordering in [SeriesOrdering::Contiguous, SeriesOrdering::Interleaved] { + let (input, polls, drops) = stream(label_source(1), ordering, []); + let error = match merge_ordered_raw_series(vec![input]) { + Ok(_) => panic!("{ordering:?} ordering must be rejected"), + Err(error) => error.to_string(), + }; + assert!( + error.contains("ordering") || error.contains("ordered"), + "{ordering:?}: {error}" + ); + assert_eq!(polls.load(Ordering::SeqCst), 0, "{ordering:?}"); + assert_eq!(drops.load(Ordering::SeqCst), 1, "{ordering:?}"); + } + + let error = match merge_ordered_raw_series(vec![]) { + Ok(_) => panic!("empty inputs must be rejected"), + Err(error) => error.to_string(), + }; + assert!( + error.contains("empty") || error.contains("input"), + "unexpected error: {error}" + ); +} + +#[tokio::test] +async fn preflight_rejects_binding_and_schema_mismatches_without_polling() { + let (explicit, explicit_polls, explicit_drops) = + stream(RawLabelSource::Explicit, SeriesOrdering::KeyThenTime, []); + let error = match merge_ordered_raw_series(vec![explicit]) { + Ok(_) => panic!("Explicit label binding must be rejected"), + Err(error) => error.to_string(), + }; + assert!( + error.to_lowercase().contains("primary") || error.to_lowercase().contains("label"), + "{error}" + ); + assert_eq!(explicit_polls.load(Ordering::SeqCst), 0); + assert_eq!(explicit_drops.load(Ordering::SeqCst), 1); + + let (single, single_polls, single_drops) = + stream(label_source(1), SeriesOrdering::KeyThenTime, []); + let _merged = merge_ordered_raw_series(vec![single]).unwrap(); + assert_eq!( + single_polls.load(Ordering::SeqCst), + 0, + "single input construction must not poll eagerly" + ); + drop(_merged); + assert_eq!(single_drops.load(Ordering::SeqCst), 1); + + let cases = [ + ( + "different field", + vec![FieldSpec { + name: Arc::from("humidity"), + }], + None, + ), + ("zero fields", vec![], None), + ( + "evaluation grid", + vec![FieldSpec { + name: Arc::from("temperature"), + }], + Some(Arc::new( + EvalGrid::try_new(GridSpec::try_new(0, 1, 2).unwrap()).unwrap(), + )), + ), + ]; + let (invalid_mode, _, invalid_mode_drops) = + stream(label_source(1), SeriesOrdering::Contiguous, []); + let ordering_error = match merge_ordered_raw_series(vec![invalid_mode]) { + Ok(_) => panic!("Contiguous ordering must be rejected"), + Err(error) => error.to_string(), + }; + assert!( + ordering_error.to_lowercase().contains("ordering"), + "{ordering_error}" + ); + assert_eq!(invalid_mode_drops.load(Ordering::SeqCst), 1); + + for (description, fields, eval_grid) in cases { + let (first, first_polls, first_drops) = + stream(label_source(1), SeriesOrdering::KeyThenTime, []); + let incompatible = Arc::new(SeriesSchema { + fields: fields.into(), + eval_grid, + }); + let (second, second_polls, second_drops) = stream_with( + incompatible, + label_source(2), + SeriesOrdering::KeyThenTime, + [], + ); + let error = match merge_ordered_raw_series(vec![first, second]) { + Ok(_) => panic!("{description} must be rejected"), + Err(error) => error.to_string(), + }; + assert!( + error.to_lowercase().contains("schema") || error.to_lowercase().contains("field"), + "{description}: {error}" + ); + assert_eq!(first_polls.load(Ordering::SeqCst), 0); + assert_eq!(second_polls.load(Ordering::SeqCst), 0); + assert_eq!(first_drops.load(Ordering::SeqCst), 1); + assert_eq!(second_drops.load(Ordering::SeqCst), 1); + } +} + +#[tokio::test] +async fn later_error_drops_held_candidate_and_heads_before_returning_error() { + let key_a = read_key("rack-a", "sensor-1"); + let key_b = read_key("rack-b", "sensor-1"); + let key_a_drops = Arc::new(AtomicUsize::new(0)); + let key_b_drops = Arc::new(AtomicUsize::new(0)); + let (owned_a, weak_a) = owned_key(key_a.as_slice().to_vec(), key_a_drops.clone()); + let (owned_b, weak_b) = owned_key(key_b.as_slice().to_vec(), key_b_drops.clone()); + let (first, _, first_drops) = stream( + label_source(1), + SeriesOrdering::KeyThenTime, + [ + ready(chunk(owned_a, &[10], &[1.0], false)), + Poll::Pending, + Poll::Ready(Some(Err(BoxedError::new(PlainError::new( + "upstream boom".to_string(), + StatusCode::Internal, + ))))), + ], + ); + let (second, _, second_drops) = stream( + label_source(2), + SeriesOrdering::KeyThenTime, + [ready(chunk(owned_b, &[20], &[2.0], true))], + ); + let mut merged = merge_ordered_raw_series(vec![first, second]).unwrap(); + let waker = futures::task::noop_waker(); + let mut cx = Context::from_waker(&waker); + assert!(matches!( + Pin::new(&mut merged).poll_next(&mut cx), + Poll::Pending + )); + assert!(weak_a.upgrade().is_some() && weak_b.upgrade().is_some()); + let error = match Pin::new(&mut merged).poll_next(&mut cx) { + Poll::Ready(Some(Err(error))) => error.to_string(), + _ => panic!("upstream failure should be returned"), + }; + assert!(error.contains("upstream boom"), "{error}"); + assert!(weak_a.upgrade().is_none() && weak_b.upgrade().is_none()); + assert_eq!(key_a_drops.load(Ordering::SeqCst), 1); + assert_eq!(key_b_drops.load(Ordering::SeqCst), 1); + assert_eq!(first_drops.load(Ordering::SeqCst), 1); + assert_eq!(second_drops.load(Ordering::SeqCst), 1); + assert!(merged.next().await.is_none()); +} + +// A chunk marked `is_last` proves its input cannot extend that read key, so a body that never +// becomes ready again must postpone neither the already-decided prefix nor the late error waiting +// behind another input's greater head. Global finality is still recomputed over all inputs. +#[test] +fn completed_key_pending_body_does_not_starve_ready_prefix_or_late_error() { + for failing_input_first in [false, true] { + let key = read_key("rack-a", "sensor-1"); + let marker = "late alpha body behind the completed key never unblocked"; + let key_drops = Arc::new(AtomicUsize::new(0)); + let (owned_lower_key, weak_lower_key) = + owned_key(key.as_slice().to_vec(), key_drops.clone()); + let gate = Arc::new(Mutex::new(None::)); + let release = Arc::new(AtomicUsize::new(0)); + // The lower chunk completes the read key; its body stays unreleased, so every later poll of + // that input is a genuine, permanent `Pending`. + let (completed, _, completed_drops) = gated_stream( + label_source(1), + chunk(owned_lower_key, &[10], &[1.0], true), + gate, + release, + ); + let (failing, _, failing_drops) = stream( + label_source(2), + SeriesOrdering::KeyThenTime, + [ + ready(chunk(key.clone(), &[30], &[3.0], true)), + Poll::Ready(Some(Err(BoxedError::new(PlainError::new( + marker.to_string(), + StatusCode::Internal, + ))))), + ], + ); + let inputs = if failing_input_first { + vec![failing, completed] + } else { + vec![completed, failing] + }; + let mut merged = merge_ordered_raw_series(inputs).unwrap(); + let waker = futures::task::noop_waker(); + let mut cx = Context::from_waker(&waker); + + // The decided lower chunk must not wait for the pending body; the greater same-key head + // still holds the key open, so the prefix is not final. + let prefix = match Pin::new(&mut merged).poll_next(&mut cx) { + Poll::Ready(Some(Ok(prefix))) => prefix, + other => panic!( + "completed lower chunk must be ready before the pending body settles (failing \ + input first: {failing_input_first}): {other:?}" + ), + }; + assert_eq!(prefix.primary_key, key); + assert_eq!(prefix.time.as_buffer().as_ref(), &[10]); + assert!( + !prefix.is_last, + "the greater same-key head keeps the prefix non-final" + ); + match prefix.fields.get(0).unwrap().sample(0) { + SampleRef::Float(value) => assert_eq!(value, 1.0), + _ => panic!("expected float sample"), + } + assert_eq!( + key_drops.load(Ordering::SeqCst), + 0, + "the emitted prefix retains its key allocation" + ); + drop(prefix); + assert!( + weak_lower_key.upgrade().is_none(), + "dropping the prefix releases its key allocation" + ); + assert_eq!(key_drops.load(Ordering::SeqCst), 1); + + // The late error waits behind the held greater-time head, and that head's own input is + // polled again as soon as the head is consumed, so the error must surface on this very + // demand: neither the pending body nor the already-decided successor chunk may delay it. + let error = match Pin::new(&mut merged).poll_next(&mut cx) { + Poll::Ready(Some(Err(error))) => error, + other => panic!( + "the late error must be returned before the greater-time successor is emitted \ + (failing input first: {failing_input_first}): {other:?}" + ), + }; + assert_eq!(error.status_code(), StatusCode::Internal); + assert!(error.to_string().contains(marker), "{error}"); + assert_eq!( + completed_drops.load(Ordering::SeqCst), + 1, + "both inputs drop before the error is observed" + ); + assert_eq!(failing_drops.load(Ordering::SeqCst), 1); + assert!(weak_lower_key.upgrade().is_none()); + assert_eq!(key_drops.load(Ordering::SeqCst), 1); + assert!(matches!( + Pin::new(&mut merged).poll_next(&mut cx), + Poll::Ready(None) + )); + } +} + +// A `is_last` chunk completes its read key even while the input stays live, so the final chunk is +// emitted without waiting for the body; the live input must not be dropped or treated as EOF. +#[test] +fn live_input_completed_at_its_last_chunk_emits_final_chunk_without_dropping_input() { + let key = read_key("rack-a", "sensor-1"); + let gate = Arc::new(Mutex::new(None::)); + let release = Arc::new(AtomicUsize::new(0)); + let (input, _, drops) = gated_stream( + label_source(1), + chunk(key.clone(), &[10], &[1.0], true), + gate, + release.clone(), + ); + let mut merged = merge_ordered_raw_series(vec![input]).unwrap(); + let waker = futures::task::noop_waker(); + let mut cx = Context::from_waker(&waker); + + let output = match Pin::new(&mut merged).poll_next(&mut cx) { + Poll::Ready(Some(Ok(output))) => output, + other => panic!("a completed series must not wait for its live body: {other:?}"), + }; + assert_eq!(output.primary_key, key); + assert_eq!(output.time.as_buffer().as_ref(), &[10]); + assert_eq!(float_sample(&output, 0), 1.0); + assert!( + output.is_last, + "the last chunk of a completed series is final" + ); + drop(output); + + // The input is live, so nothing may be dropped or reported as EOF yet. + assert!(matches!( + Pin::new(&mut merged).poll_next(&mut cx), + Poll::Pending + )); + assert_eq!( + drops.load(Ordering::SeqCst), + 0, + "a live input stays owned after its final chunk" + ); + release.store(1, Ordering::SeqCst); + assert!(matches!( + Pin::new(&mut merged).poll_next(&mut cx), + Poll::Ready(None) + )); + assert_eq!(drops.load(Ordering::SeqCst), 1); +} + +// A live input that completed one read key says nothing about other keys, so a greater head waits +// until that input resolves, and is still emitted afterwards. +#[test] +fn live_completed_input_withholds_other_key_head_until_it_resolves() { + let completed_key = read_key("rack-a", "sensor-1"); + let later_key = read_key("rack-b", "sensor-1"); + let gate = Arc::new(Mutex::new(None::)); + let release = Arc::new(AtomicUsize::new(0)); + let (completed, _, completed_drops) = gated_stream( + label_source(1), + chunk(completed_key.clone(), &[10], &[1.0], true), + gate, + release.clone(), + ); + let (later, _, later_drops) = stream( + label_source(2), + SeriesOrdering::KeyThenTime, + [ready(chunk(later_key.clone(), &[50], &[5.0], true))], + ); + let mut merged = merge_ordered_raw_series(vec![completed, later]).unwrap(); + let waker = futures::task::noop_waker(); + let mut cx = Context::from_waker(&waker); + + let first = match Pin::new(&mut merged).poll_next(&mut cx) { + Poll::Ready(Some(Ok(first))) => first, + other => panic!("the completed key must be emitted despite the live body: {other:?}"), + }; + assert_eq!(first.primary_key, completed_key); + assert_eq!(float_sample(&first, 0), 1.0); + assert!( + first.is_last, + "the greater head proves the completed key is globally final" + ); + drop(first); + + // The live input may still reveal keys before `rack-b`, so that head is not provable yet. + assert!(matches!( + Pin::new(&mut merged).poll_next(&mut cx), + Poll::Pending + )); + assert_eq!(completed_drops.load(Ordering::SeqCst), 0); + release.store(1, Ordering::SeqCst); + let second = match Pin::new(&mut merged).poll_next(&mut cx) { + Poll::Ready(Some(Ok(second))) => second, + other => panic!("the held head must be emitted once the live input resolves: {other:?}"), + }; + assert_eq!(second.primary_key, later_key); + assert_eq!(second.time.as_buffer().as_ref(), &[50]); + assert_eq!(float_sample(&second, 0), 5.0); + assert!(second.is_last); + drop(second); + assert!(matches!( + Pin::new(&mut merged).poll_next(&mut cx), + Poll::Ready(None) + )); + assert_eq!(completed_drops.load(Ordering::SeqCst), 1); + assert_eq!(later_drops.load(Ordering::SeqCst), 1); +} + +// An input that has not finished a read key may still add chunks to it, so even its already-decided +// slice waits; its own final chunk then unsticks the exact ordered sequence. +#[test] +fn unfinished_same_key_pending_input_blocks_its_slice_until_its_final_chunk() { + let key = read_key("rack-a", "sensor-1"); + let gate = Arc::new(Mutex::new(None::)); + let release = Arc::new(AtomicUsize::new(0)); + let (unfinished, _, unfinished_drops) = gated_resuming_stream( + label_source(1), + chunk(key.clone(), &[10], &[1.0], false), + [ready(chunk(key.clone(), &[20], &[2.0], true))], + gate, + release.clone(), + ); + let (later, _, later_drops) = stream( + label_source(2), + SeriesOrdering::KeyThenTime, + [ready(chunk(key.clone(), &[30], &[3.0], true))], + ); + let mut merged = merge_ordered_raw_series(vec![unfinished, later]).unwrap(); + let waker = futures::task::noop_waker(); + let mut cx = Context::from_waker(&waker); + + assert!( + matches!(Pin::new(&mut merged).poll_next(&mut cx), Poll::Pending), + "an unfinished same-key input holds its whole key back" + ); + assert_eq!(unfinished_drops.load(Ordering::SeqCst), 0); + release.store(1, Ordering::SeqCst); + + let mut observed = Vec::new(); + loop { + match Pin::new(&mut merged).poll_next(&mut cx) { + Poll::Ready(Some(Ok(next))) => observed.push(( + next.primary_key.clone(), + next.time.as_buffer().to_vec(), + float_sample(&next, 0).to_bits(), + next.is_last, + )), + Poll::Ready(None) => break, + other => panic!("expected the ordered slices or EOF, got {other:?}"), + } + } + assert_eq!( + observed, + vec![ + (key.clone(), vec![10], 1.0_f64.to_bits(), false), + (key.clone(), vec![20], 2.0_f64.to_bits(), false), + (key, vec![30], 3.0_f64.to_bits(), true), + ] + ); + assert_eq!(unfinished_drops.load(Ordering::SeqCst), 1); + assert_eq!(later_drops.load(Ordering::SeqCst), 1); +} + +// After completing a series, a still-live input keeps being polled: once its blocked body resumes, +// its later keys are emitted in order with complete values and global finality, and no input is +// dropped or treated as EOF while the merge is merely blocked. +#[test] +fn completed_input_resumes_with_later_keys_after_being_blocked() { + let completed_key = read_key("rack-a", "sensor-1"); + let later_key = read_key("rack-c", "sensor-1"); + let gate_a = Arc::new(Mutex::new(None::)); + let gate_b = Arc::new(Mutex::new(None::)); + let release_a = Arc::new(AtomicUsize::new(0)); + let release_b = Arc::new(AtomicUsize::new(0)); + let (resuming, _, resuming_drops) = gated_resuming_stream( + label_source(1), + chunk(completed_key.clone(), &[10], &[1.0], true), + [ready(chunk(later_key.clone(), &[40], &[4.0], true))], + gate_a, + release_a.clone(), + ); + let (completed, _, completed_drops) = gated_stream( + label_source(2), + chunk(completed_key.clone(), &[30], &[3.0], true), + gate_b, + release_b.clone(), + ); + let mut merged = merge_ordered_raw_series(vec![resuming, completed]).unwrap(); + let waker = futures::task::noop_waker(); + let mut cx = Context::from_waker(&waker); + + let first = match Pin::new(&mut merged).poll_next(&mut cx) { + Poll::Ready(Some(Ok(first))) => first, + other => panic!("the completed prefix must be emitted despite live bodies: {other:?}"), + }; + assert_eq!(first.primary_key, completed_key); + assert_eq!(first.time.as_buffer().as_ref(), &[10]); + assert_eq!(float_sample(&first, 0), 1.0); + assert!(!first.is_last); + drop(first); + + let second = match Pin::new(&mut merged).poll_next(&mut cx) { + Poll::Ready(Some(Ok(second))) => second, + other => panic!("the remaining completed chunk must be emitted before waiting: {other:?}"), + }; + assert_eq!(second.primary_key, completed_key); + assert_eq!(second.time.as_buffer().as_ref(), &[30]); + assert_eq!(float_sample(&second, 0), 3.0); + assert!(second.is_last); + drop(second); + + // Both bodies are blocked: the merge waits and keeps both inputs (and their data) owned. + assert!(matches!( + Pin::new(&mut merged).poll_next(&mut cx), + Poll::Pending + )); + assert_eq!(resuming_drops.load(Ordering::SeqCst), 0); + assert_eq!(completed_drops.load(Ordering::SeqCst), 0); + release_a.store(1, Ordering::SeqCst); + release_b.store(1, Ordering::SeqCst); + let third = match Pin::new(&mut merged).poll_next(&mut cx) { + Poll::Ready(Some(Ok(third))) => third, + other => panic!("the resumed key must be emitted after the bodies unblock: {other:?}"), + }; + assert_eq!(third.primary_key, later_key); + assert_eq!(third.time.as_buffer().as_ref(), &[40]); + assert_eq!(float_sample(&third, 0), 4.0); + assert!(third.is_last); + drop(third); + assert!(matches!( + Pin::new(&mut merged).poll_next(&mut cx), + Poll::Ready(None) + )); + assert_eq!(resuming_drops.load(Ordering::SeqCst), 1); + assert_eq!(completed_drops.load(Ordering::SeqCst), 1); +} + +// A chunk split against a later same-key head leaves a remainder that is locally final, but its key +// stays open until that remainder itself is emitted. +#[test] +fn split_remainder_head_keeps_its_key_open_despite_local_final_chunk() { + let key = read_key("rack-a", "sensor-1"); + let (left, _, left_drops) = stream( + label_source(1), + SeriesOrdering::KeyThenTime, + [ready(chunk(key.clone(), &[10, 40], &[1.0, 4.0], true))], + ); + let (right, _, right_drops) = stream( + label_source(2), + SeriesOrdering::KeyThenTime, + [ready(chunk(key.clone(), &[30], &[3.0], true))], + ); + let mut merged = merge_ordered_raw_series(vec![left, right]).unwrap(); + let waker = futures::task::noop_waker(); + let mut cx = Context::from_waker(&waker); + + let first = match Pin::new(&mut merged).poll_next(&mut cx) { + Poll::Ready(Some(Ok(first))) => first, + other => panic!("expected the split prefix, got {other:?}"), + }; + assert_eq!(first.time.as_buffer().as_ref(), &[10]); + assert_eq!(float_sample(&first, 0), 1.0); + assert!( + !first.is_last, + "the locally final remainder head still holds the key open" + ); + drop(first); + + let mut observed = Vec::new(); + loop { + match Pin::new(&mut merged).poll_next(&mut cx) { + Poll::Ready(Some(Ok(next))) => observed.push(( + next.time.as_buffer().to_vec(), + float_sample(&next, 0).to_bits(), + next.is_last, + )), + Poll::Ready(None) => break, + other => panic!("expected the remaining slices or EOF, got {other:?}"), + } + } + assert_eq!( + observed, + vec![ + (vec![30], 3.0_f64.to_bits(), false), + (vec![40], 4.0_f64.to_bits(), true), + ] + ); + assert_eq!(left_drops.load(Ordering::SeqCst), 1); + assert_eq!(right_drops.load(Ordering::SeqCst), 1); +} + +#[tokio::test] +async fn rejects_malformed_public_chunk_and_unfinished_series_at_eof() { + let malformed = RawSeriesChunk { + primary_key: read_key("rack-a", "sensor-1"), + metric: None, + time: RawTimestamps::try_new(ScalarBuffer::from(vec![10, 20])).unwrap(), + fields: SeriesFields::one(Samples::from_float(Float64Array::from(vec![1.0]))), + is_last: true, + }; + let (input, _, drops) = stream( + label_source(1), + SeriesOrdering::KeyThenTime, + [ready(malformed)], + ); + let mut merged = merge_ordered_raw_series(vec![input]).unwrap(); + let error = merged.next().await.unwrap().unwrap_err().to_string(); + assert!( + error.to_lowercase().contains("field") || error.to_lowercase().contains("length"), + "{error}" + ); + assert_eq!(drops.load(Ordering::SeqCst), 1); + assert!(merged.next().await.is_none()); + + for (malformed, reason) in [ + ( + RawSeriesChunk { + primary_key: read_key("rack-a", "sensor-2"), + metric: None, + time: RawTimestamps::try_new(ScalarBuffer::from(Vec::::new())).unwrap(), + fields: SeriesFields::one(Samples::from_float(Float64Array::from( + Vec::::new(), + ))), + is_last: true, + }, + "raw series chunks must not be empty", + ), + ( + RawSeriesChunk { + primary_key: read_key("rack-a", "sensor-3"), + metric: None, + time: RawTimestamps::try_new(ScalarBuffer::from(vec![30])).unwrap(), + fields: SeriesFields::try_from_vec(vec![ + Samples::from_float(Float64Array::from(vec![3.0])), + Samples::from_float(Float64Array::from(vec![4.0])), + ]) + .unwrap(), + is_last: true, + }, + "field count", + ), + ] { + let (input, _, drops) = stream( + label_source(1), + SeriesOrdering::KeyThenTime, + [ready(malformed)], + ); + let mut merged = merge_ordered_raw_series(vec![input]).unwrap(); + let error = merged.next().await.unwrap().unwrap_err().to_string(); + assert!( + error.to_lowercase().contains(reason), + "expected {reason:?}: {error}" + ); + assert_eq!(drops.load(Ordering::SeqCst), 1); + assert!(merged.next().await.is_none()); + } + + let unfinished = RawSeriesChunk { + primary_key: read_key("rack-a", "sensor-1"), + metric: None, + time: RawTimestamps::try_new(ScalarBuffer::from(vec![30])).unwrap(), + fields: SeriesFields::one(Samples::from_float(Float64Array::from(vec![3.0]))), + is_last: false, + }; + let (input, _, drops) = stream( + label_source(1), + SeriesOrdering::KeyThenTime, + [ready(unfinished), Poll::Ready(None)], + ); + let mut merged = merge_ordered_raw_series(vec![input]).unwrap(); + let error = merged.next().await.unwrap().unwrap_err().to_string(); + assert!( + error.to_lowercase().contains("last") || error.to_lowercase().contains("final"), + "{error}" + ); + assert_eq!(drops.load(Ordering::SeqCst), 1); + assert!(merged.next().await.is_none()); +} + +#[tokio::test] +async fn equal_timestamp_at_successor_boundary_is_an_error() { + let key = read_key("rack-a", "sensor-1"); + let (left, _, drops_a) = stream( + label_source(1), + SeriesOrdering::KeyThenTime, + [ready(chunk(key.clone(), &[10, 30], &[1.0, 3.0], true))], + ); + let (right, _, drops_b) = stream( + label_source(2), + SeriesOrdering::KeyThenTime, + [ready(chunk(key, &[20, 30], &[2.0, 30.0], true))], + ); + let mut merged = merge_ordered_raw_series(vec![left, right]).unwrap(); + let first = merged.next().await.unwrap().unwrap(); + let second = merged.next().await.unwrap().unwrap(); + assert_eq!(first.time.as_buffer().as_ref(), &[10]); + assert_eq!(second.time.as_buffer().as_ref(), &[20]); + assert!(!first.is_last && !second.is_last); + let error = merged.next().await.unwrap().unwrap_err().to_string(); + assert!( + error.to_lowercase().contains("duplicate") || error.to_lowercase().contains("timestamp"), + "{error}" + ); + assert_eq!(drops_a.load(Ordering::SeqCst), 1); + assert_eq!(drops_b.load(Ordering::SeqCst), 1); + assert!(merged.next().await.is_none()); +} + +#[tokio::test] +async fn greater_key_head_allows_finished_series_before_future_eof() { + let key_a = read_key("rack-a", "sensor-1"); + let key_b = read_key("rack-b", "sensor-1"); + let gate = Arc::new(Mutex::new(None::)); + let release = Arc::new(AtomicUsize::new(0)); + let (right, right_polls, right_drops) = gated_stream( + label_source(2), + chunk(key_b, &[5], &[5.0], true), + gate.clone(), + release.clone(), + ); + let (left, _, left_drops) = stream( + label_source(1), + SeriesOrdering::KeyThenTime, + [ready(chunk(key_a.clone(), &[10], &[1.0], true))], + ); + let mut merged = merge_ordered_raw_series(vec![left, right]).unwrap(); + let waker = futures::task::noop_waker(); + let mut cx = Context::from_waker(&waker); + let item = match Pin::new(&mut merged).poll_next(&mut cx) { + Poll::Ready(Some(Ok(item))) => item, + _ => panic!("expected finished smaller key"), + }; + assert_eq!(item.primary_key, key_a); + assert!(item.is_last); + assert!( + gate.lock().unwrap().is_none(), + "known greater head makes an EOF poll unnecessary" + ); + assert_eq!( + right_polls.load(Ordering::SeqCst), + 1, + "known greater head must not be polled again" + ); + drop(item); + drop(merged); + assert_eq!(left_drops.load(Ordering::SeqCst), 1); + assert_eq!(right_drops.load(Ordering::SeqCst), 1); +} + +#[tokio::test] +async fn preflight_rejects_incompatible_primary_key_bindings_without_polling() { + let full_baseline = label_source(1); + let physical_comparison_baseline = { + let RawLabelSource::PrimaryKey { + metric_name, + physical_metadata, + .. + } = label_source(1) + else { + unreachable!() + }; + RawLabelSource::PrimaryKey { + metric_name, + physical_metadata, + logical_tags: Arc::from([BoundTag { + column_id: 9, + name: Arc::from("sensor"), + }]), + } + }; + let cases = [ + ( + "metric name", + label_source_with(2, PrimaryKeyEncoding::Dense, "other_metric"), + "binding", + ), + ( + "encoding", + label_source_with(2, PrimaryKeyEncoding::Sparse, "room_temperature"), + "binding", + ), + ( + "primary key column ids", + { + let RawLabelSource::PrimaryKey { + metric_name, + physical_metadata, + logical_tags, + } = label_source(2) + else { + unreachable!() + }; + let mut changed = (*physical_metadata).clone(); + changed.primary_key = vec![9, 8]; + RawLabelSource::PrimaryKey { + metric_name, + physical_metadata: Arc::new(changed), + logical_tags, + } + }, + "binding", + ), + ( + "primary key physical name mismatch", + { + let RawLabelSource::PrimaryKey { + metric_name, + physical_metadata, + .. + } = label_source(2) + else { + unreachable!() + }; + let mut changed = (*physical_metadata).clone(); + changed + .column_metadatas + .iter_mut() + .find(|column| column.column_id == 8) + .unwrap() + .column_schema + .name = "shelf".to_string(); + RawLabelSource::PrimaryKey { + metric_name, + physical_metadata: Arc::new(changed), + logical_tags: Arc::from([BoundTag { + column_id: 9, + name: Arc::from("sensor"), + }]), + } + }, + "merge compatibility", + ), + ( + "primary key semantic role", + { + let RawLabelSource::PrimaryKey { + metric_name, + physical_metadata, + .. + } = label_source(2) + else { + unreachable!() + }; + let mut changed = (*physical_metadata).clone(); + changed + .column_metadatas + .iter_mut() + .find(|column| column.column_id == 8) + .unwrap() + .semantic_type = SemanticType::Field; + RawLabelSource::PrimaryKey { + metric_name, + physical_metadata: Arc::new(changed), + logical_tags: Arc::from([BoundTag { + column_id: 9, + name: Arc::from("sensor"), + }]), + } + }, + "merge compatibility", + ), + ( + "primary key physical type", + { + let RawLabelSource::PrimaryKey { + metric_name, + physical_metadata, + logical_tags, + } = label_source(2) + else { + unreachable!() + }; + let mut changed = (*physical_metadata).clone(); + changed + .column_metadatas + .iter_mut() + .find(|column| column.column_id == 8) + .unwrap() + .column_schema + .data_type = ConcreteDataType::int64_datatype(); + RawLabelSource::PrimaryKey { + metric_name, + physical_metadata: Arc::new(changed), + logical_tags, + } + }, + "binding", + ), + ( + "logical tag mapping", + { + let RawLabelSource::PrimaryKey { + metric_name, + physical_metadata, + .. + } = label_source(2) + else { + unreachable!() + }; + RawLabelSource::PrimaryKey { + metric_name, + physical_metadata, + logical_tags: Arc::from([BoundTag { + column_id: 8, + name: Arc::from("rack"), + }]), + } + }, + "binding", + ), + ]; + for (description, incompatible, reason) in cases { + let baseline = if reason == "merge compatibility" { + &physical_comparison_baseline + } else { + &full_baseline + }; + let (first, polls_a, drops_a) = + stream((*baseline).clone(), SeriesOrdering::KeyThenTime, []); + let (second, polls_b, drops_b) = stream(incompatible, SeriesOrdering::KeyThenTime, []); + let error = match merge_ordered_raw_series(vec![first, second]) { + Ok(_) => panic!("{description} mismatch must be rejected"), + Err(error) => error.to_string(), + }; + if reason == "merge compatibility" { + assert_eq!( + error, + "Invalid series: raw merge requires equal raw schemas, primary-key bindings, and ReadKeyThenTime ordering", + "{description}" + ); + } else { + assert!( + error.to_lowercase().contains(reason), + "{description}: {error}" + ); + } + assert_eq!(polls_a.load(Ordering::SeqCst), 0); + assert_eq!(polls_b.load(Ordering::SeqCst), 0); + assert_eq!(drops_a.load(Ordering::SeqCst), 1); + assert_eq!(drops_b.load(Ordering::SeqCst), 1); + } +} + +#[tokio::test] +async fn metadata_snapshots_with_same_layout_are_accepted() { + let base = label_source(1); + let RawLabelSource::PrimaryKey { + metric_name, + physical_metadata, + logical_tags, + } = base + else { + unreachable!() + }; + let mut builder = RegionMetadataBuilder::from_existing((*physical_metadata).clone()); + builder.bump_version(); + builder.push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "unrelated_note", + ConcreteDataType::string_datatype(), + true, + ), + semantic_type: SemanticType::Field, + column_id: 99, + }); + let reordered_tags = Arc::from([ + BoundTag { + column_id: 9, + name: Arc::from("sensor"), + }, + BoundTag { + column_id: 8, + name: Arc::from("rack"), + }, + ]); + let mut second_metadata = builder.build().unwrap(); + second_metadata.region_id = RegionId::new(2, 1); + let second_binding = RawLabelSource::PrimaryKey { + metric_name, + physical_metadata: Arc::new(second_metadata), + logical_tags: reordered_tags, + }; + assert_ne!( + physical_metadata.region_id, + match &second_binding { + RawLabelSource::PrimaryKey { + physical_metadata, .. + } => physical_metadata.region_id, + _ => unreachable!(), + } + ); + let first_binding = RawLabelSource::PrimaryKey { + metric_name: Arc::from("room_temperature"), + physical_metadata, + logical_tags, + }; + let (left, _, _) = stream( + first_binding, + SeriesOrdering::KeyThenTime, + [ready(chunk( + read_key("rack-a", "sensor-1"), + &[1], + &[1.0], + true, + ))], + ); + let (right, _, _) = stream( + second_binding, + SeriesOrdering::KeyThenTime, + [ready(chunk( + read_key("rack-b", "sensor-1"), + &[2], + &[2.0], + true, + ))], + ); + let mut merged = merge_ordered_raw_series(vec![left, right]).unwrap(); + let first = merged.next().await.unwrap().unwrap(); + assert_eq!(first.time.as_buffer().as_ref(), &[1]); + drop(first); + assert_eq!( + merged + .next() + .await + .unwrap() + .unwrap() + .time + .as_buffer() + .as_ref(), + &[2] + ); + assert!(merged.next().await.is_none()); +} + +#[tokio::test] +async fn pending_inputs_are_all_polled_and_wakers_registered_without_spinning() { + use futures::task::{ArcWake, waker_ref}; + + #[derive(Default)] + struct WakeCounter(AtomicUsize); + impl ArcWake for WakeCounter { + fn wake_by_ref(arc_self: &Arc) { + arc_self.0.fetch_add(1, Ordering::SeqCst); + } + } + + let gate_a = Arc::new(Mutex::new(None::)); + let gate_b = Arc::new(Mutex::new(None::)); + let release_a = Arc::new(AtomicUsize::new(0)); + let release_b = Arc::new(AtomicUsize::new(0)); + let (left, left_polls, left_drops) = + gated_pending_stream(label_source(1), gate_a.clone(), release_a.clone()); + let (right, right_polls, right_drops) = + gated_pending_stream(label_source(2), gate_b.clone(), release_b.clone()); + let mut merged = merge_ordered_raw_series(vec![left, right]).unwrap(); + let counter = Arc::new(WakeCounter::default()); + let waker = waker_ref(&counter); + let mut cx = Context::from_waker(&waker); + assert!(matches!( + Pin::new(&mut merged).poll_next(&mut cx), + Poll::Pending + )); + assert!(gate_a.lock().unwrap().is_some()); + assert!(gate_b.lock().unwrap().is_some()); + let before = ( + left_polls.load(Ordering::SeqCst), + right_polls.load(Ordering::SeqCst), + ); + assert_eq!( + before, + (1, 1), + "poll each missing head once in this outer poll" + ); + assert_eq!( + counter.0.load(Ordering::SeqCst), + 0, + "Pending inputs must not self-wake" + ); + assert!(matches!( + Pin::new(&mut merged).poll_next(&mut cx), + Poll::Pending + )); + assert_eq!( + ( + left_polls.load(Ordering::SeqCst), + right_polls.load(Ordering::SeqCst) + ), + (2, 2), + "each new outer poll checks each pending input once" + ); + release_a.store(1, Ordering::SeqCst); + release_b.store(1, Ordering::SeqCst); + gate_a.lock().unwrap().take().unwrap().wake(); + gate_b.lock().unwrap().take().unwrap().wake(); + assert_eq!(counter.0.load(Ordering::SeqCst), 2); + assert!(matches!( + Pin::new(&mut merged).poll_next(&mut cx), + Poll::Ready(None) + )); + assert_eq!(left_drops.load(Ordering::SeqCst), 1); + assert_eq!(right_drops.load(Ordering::SeqCst), 1); + drop(merged); +} + +#[tokio::test] +async fn cancel_and_eof_drop_owned_inputs_and_candidate_key_buffers() { + let key = read_key("rack-a", "sensor-1"); + let bytes = key.as_slice().to_vec(); + let key_drops = Arc::new(AtomicUsize::new(0)); + let (owned, weak_owner) = owned_key(bytes, key_drops.clone()); + let (input, _, drops) = stream( + label_source(1), + SeriesOrdering::KeyThenTime, + [ready(chunk(owned, &[1], &[1.0], true))], + ); + let mut merged = merge_ordered_raw_series(vec![input]).unwrap(); + let waker = futures::task::noop_waker(); + let mut cx = Context::from_waker(&waker); + let output = match Pin::new(&mut merged).poll_next(&mut cx) { + Poll::Ready(Some(Ok(output))) => output, + _ => panic!("expected emitted output"), + }; + assert_eq!( + key_drops.load(Ordering::SeqCst), + 0, + "output retains its actual key allocation" + ); + drop(merged); + assert_eq!(drops.load(Ordering::SeqCst), 1); + assert_eq!( + key_drops.load(Ordering::SeqCst), + 0, + "output remains the owner after upstream drop" + ); + assert!(weak_owner.upgrade().is_some()); + drop(output); + assert!( + weak_owner.upgrade().is_none(), + "dropping output releases the allocation" + ); + assert_eq!(key_drops.load(Ordering::SeqCst), 1); + + let gate = Arc::new(Mutex::new(None::)); + let release = Arc::new(AtomicUsize::new(0)); + let pending_drops = Arc::new(AtomicUsize::new(0)); + let (pending_key, pending_weak) = owned_key( + read_key("rack-z", "sensor-9").as_slice().to_vec(), + pending_drops.clone(), + ); + let (input, _, drops) = gated_stream( + label_source(2), + chunk(pending_key, &[1], &[1.0], false), + gate.clone(), + release, + ); + let mut merged = merge_ordered_raw_series(vec![input]).unwrap(); + assert!(matches!( + Pin::new(&mut merged).poll_next(&mut cx), + Poll::Pending + )); + assert!(gate.lock().unwrap().is_some()); + assert!( + pending_weak.upgrade().is_some(), + "pending merge retains its candidate payload" + ); + drop(merged); + assert_eq!(drops.load(Ordering::SeqCst), 1); + assert!( + pending_weak.upgrade().is_none(), + "cancellation releases candidate key owner" + ); + assert_eq!(pending_drops.load(Ordering::SeqCst), 1); +} + +#[test] +fn previous_key_validation_retains_only_key_bytes_not_input_allocation() { + let parent = (0..32).map(|n| (n % 251) as u8).collect::>(); + let owner_drops = Arc::new(AtomicUsize::new(0)); + let (parent_key, weak_owner) = owned_key(parent, owner_drops.clone()); + let offset = 7; + let key_len = read_key("rack-a", "sensor-1").len(); + let owner_key = parent_key.slice_with_length(offset, key_len); + drop(parent_key); + let input_chunk = chunk(owner_key, &[10], &[1.0], true); + let (upstream, _, _) = stream(label_source(1), SeriesOrdering::KeyThenTime, []); + let mut input = Input::new(upstream); + input.validate(&input_chunk, &schema()).unwrap(); + drop(input_chunk); + assert!( + weak_owner.upgrade().is_none(), + "previous_key should copy key bytes, not retain its allocation owner" + ); + assert_eq!(owner_drops.load(Ordering::SeqCst), 1); +} + +#[tokio::test] +async fn ordered_input_contract_rejects_bad_key_time_and_finality_transitions() { + let a = read_key("rack-a", "sensor-1"); + let b = read_key("rack-b", "sensor-1"); + // `RawTimestamps` rejects descending axes at construction; a merge only observes + // chunks that already satisfy the strictly increasing time contract. + assert!(RawTimestamps::try_new(ScalarBuffer::from(vec![20, 10])).is_err()); + let cases = [ + ( + "series recurs after another key", + vec![ + chunk(a.clone(), &[1], &[1.0], true), + chunk(b.clone(), &[2], &[2.0], true), + chunk(a.clone(), &[3], &[3.0], true), + ], + "key", + ), + ( + "unfinished series at key change", + vec![ + chunk(a.clone(), &[1], &[1.0], false), + chunk(b.clone(), &[2], &[2.0], true), + ], + "last", + ), + ( + "series recurs after final chunk", + vec![ + chunk(a.clone(), &[1], &[1.0], true), + chunk(a.clone(), &[2], &[2.0], true), + ], + "last", + ), + ( + "non-increasing timestamps across chunks", + vec![ + chunk(a.clone(), &[10], &[1.0], false), + chunk(a.clone(), &[10], &[2.0], true), + ], + "time", + ), + ( + "explicit metric in primary-key mode", + vec![RawSeriesChunk { + primary_key: a.clone(), + metric: Some(MetricName::try_new(None, vec![]).unwrap()), + time: RawTimestamps::try_new(ScalarBuffer::from(vec![1])).unwrap(), + fields: SeriesFields::one(Samples::from_float(Float64Array::from(vec![1.0]))), + is_last: true, + }], + "primary-key", + ), + ]; + for (description, chunks, reason) in cases { + let (input, _, drops) = stream( + label_source(1), + SeriesOrdering::KeyThenTime, + chunks.into_iter().map(ready), + ); + let mut merged = merge_ordered_raw_series(vec![input]).unwrap(); + let error = loop { + match merged.next().await { + Some(Ok(_)) => {} + Some(Err(error)) => break error.to_string(), + None => panic!("{description} should fail"), + } + }; + assert!( + error.to_lowercase().contains(reason), + "{description}: {error}" + ); + assert_eq!(drops.load(Ordering::SeqCst), 1, "{description}"); + assert!(merged.next().await.is_none(), "{description} must be fused"); + } +} + +#[tokio::test] +async fn sparse_complete_keys_keep_tag_and_internal_identity_distinct() { + use mito_codec::row_converter::SparsePrimaryKeyCodec; + + let sparse_source = label_source_with_tags( + 1, + PrimaryKeyEncoding::Sparse, + "room_temperature", + vec![(42, "rack"), (43, "sensor")], + ); + let codec = SparsePrimaryKeyCodec::with_fields(vec![ + (42, SortField::new(ConcreteDataType::string_datatype())), + (43, SortField::new(ConcreteDataType::string_datatype())), + ]); + let sparse_key = |table_id, tsid, rack: &str, sensor: &str| { + let mut bytes = Vec::new(); + codec.encode_internal(table_id, tsid, &mut bytes).unwrap(); + codec + .encode_values( + &[ + (42, Value::String(rack.into())), + (43, Value::String(sensor.into())), + ], + &mut bytes, + ) + .unwrap(); + Buffer::from(bytes) + }; + let key_one = sparse_key(42, 9, "rack-a", "sensor-1"); + let key_two = sparse_key(42, 9, "rack-a", "sensor-2"); + let key_three = sparse_key(43, 9, "rack-a", "sensor-1"); + let mut key_records = [ + (key_one.clone(), 42, 9, "sensor-1"), + (key_two.clone(), 42, 9, "sensor-2"), + (key_three.clone(), 43, 9, "sensor-1"), + ]; + key_records.sort_by(|a, b| a.0.as_slice().cmp(b.0.as_slice())); + let left_chunks = [key_records[0].0.clone(), key_records[2].0.clone()] + .into_iter() + .enumerate() + .map(|(index, key)| { + let (time, value) = if index == 0 { (1, 10.0) } else { (3, 30.0) }; + ready(chunk(key, &[time], &[value], true)) + }); + let right_chunks = [key_records[1].0.clone()] + .into_iter() + .map(|key| ready(chunk(key, &[2], &[20.0], true))); + let (left, _, _) = stream( + sparse_source.clone(), + SeriesOrdering::KeyThenTime, + left_chunks, + ); + let (right, _, _) = stream(sparse_source, SeriesOrdering::KeyThenTime, right_chunks); + let merged = merge_ordered_raw_series(vec![left, right]).unwrap(); + let rows = super::super::materialize_raw_labels(merged) + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!(rows.len(), 3); + let observed = rows + .iter() + .map(|chunk| chunk.primary_key.clone()) + .collect::>(); + assert_eq!( + observed, + key_records + .iter() + .map(|record| record.0.clone()) + .collect::>() + ); + assert!(rows.iter().all(|chunk| chunk.is_last)); + let expected_times_and_values = [ + (1, 10.0_f64.to_bits()), + (2, 20.0_f64.to_bits()), + (3, 30.0_f64.to_bits()), + ]; + for (row, (expected_time, expected_value)) in rows.iter().zip(expected_times_and_values) { + assert_eq!(row.time.timestamp(0), Some(expected_time)); + assert!( + matches!(row.fields.get(0).unwrap().sample(0), SampleRef::Float(value) if value.to_bits() == expected_value) + ); + } + let metrics = rows + .iter() + .map(|row| row.metric.as_ref().unwrap()) + .collect::>(); + assert_eq!(metrics[0].get("rack"), Some("rack-a")); + assert_eq!(metrics[1].get("rack"), Some("rack-a")); + assert_eq!(metrics[2].get("rack"), Some("rack-a")); + assert_eq!(metrics[0].tags().len(), 2); + assert_eq!(metrics[0].get("sensor"), Some(key_records[0].3)); + assert_eq!(metrics[1].get("sensor"), Some(key_records[1].3)); + assert_eq!(metrics[2].get("sensor"), Some(key_records[2].3)); + assert_eq!(metrics[0].name(), Some("room_temperature")); + assert_eq!( + metrics[0], metrics[2], + "different complete internal keys with same labels remain separate" + ); + assert_ne!(rows[0].primary_key, rows[2].primary_key); +} + +#[tokio::test] +async fn primary_key_codec_identity_survives_merge_then_label_materialization() { + let key = read_key("rack-z", "sensor-4"); + let (input, _, _) = stream( + label_source(1), + SeriesOrdering::KeyThenTime, + [ready(chunk( + key.clone(), + &[1, 2], + &[18.5, f64::from_bits(0x7ff8_0000_0000_0042)], + true, + ))], + ); + let merged = merge_ordered_raw_series(vec![input]).unwrap(); + let mut materialized = super::super::materialize_raw_labels(merged).unwrap(); + let output = materialized.next().await.unwrap().unwrap(); + assert_eq!(output.primary_key, key); + assert_eq!( + output.metric.as_ref().unwrap().name(), + Some("room_temperature") + ); + assert_eq!(output.metric.as_ref().unwrap().get("rack"), Some("rack-z")); + assert_eq!( + output.metric.as_ref().unwrap().get("sensor"), + Some("sensor-4") + ); + assert_eq!(output.time.as_buffer().as_ref(), &[1, 2]); + match output.fields.get(0).unwrap().sample(0) { + SampleRef::Float(value) => assert_eq!(value, 18.5), + _ => panic!("expected float sample"), + } + match output.fields.get(0).unwrap().sample(1) { + SampleRef::Float(value) => { + assert_eq!(value.to_bits(), 0x7ff8_0000_0000_0042) + } + _ => panic!("expected NaN sample"), + } + assert!(output.is_last); +} diff --git a/src/promql/src/source/merge/tests/payload.rs b/src/promql/src/source/merge/tests/payload.rs new file mode 100644 index 00000000000..2aff48a3c14 --- /dev/null +++ b/src/promql/src/source/merge/tests/payload.rs @@ -0,0 +1,627 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::VecDeque; +use std::pin::Pin; +use std::sync::{Arc, Weak}; +use std::task::{Context, Poll}; + +use common_error::ext::BoxedError; +use common_query::series_frame::{pack_raw_series, unpack_raw_series}; +use datatypes::arrow::array::{Array, Float64Array, ListArray, UInt32Array}; +use datatypes::arrow::buffer::{Buffer, ScalarBuffer}; +use datatypes::series::SeriesOrdering; +use datatypes::series::histogram::{ + BucketSideView, CounterResetHint, HistogramBatch, HistogramBatchBuilder, HistogramView, +}; +use datatypes::series::raw::{RawSeriesChunk, RawTimestamps}; +use datatypes::series::samples::{ + DenseSamplesRef, FieldSpec, HistogramSamples, SampleRef, Samples, SeriesFields, SeriesSchema, +}; +use futures::{Stream, StreamExt, TryStreamExt}; +use store_api::region_engine::{RawLabelSource, RawSeriesStream, SendableRawSeriesStream}; + +use super::{label_source, read_key}; +use crate::source::{materialize_raw_labels, merge_ordered_raw_series}; + +fn payload_schema() -> Arc { + Arc::new(SeriesSchema { + fields: vec!["temperature", "latency", "observations"] + .into_iter() + .map(|name| FieldSpec { + name: Arc::from(name), + }) + .collect::>() + .into(), + eval_grid: None, + }) +} + +fn histograms(sums: &[f64]) -> Arc { + let mut builder = HistogramBatchBuilder::new(); + for &sum in sums { + builder + .push(HistogramView { + schema: -19, + zero_threshold: 0.125, + sum, + count: 23.5, + zero_count: 2.25, + reset_hint: CounterResetHint::Unrecognized(37), + start_timestamp: Some(-456), + custom_values: &[1.5, 4.5], + positive: BucketSideView { + span_offsets: &[1], + span_lengths: &[2], + buckets: &[3.25, 7.5], + }, + negative: BucketSideView { + span_offsets: &[-2], + span_lengths: &[1], + buckets: &[0.75], + }, + }) + .unwrap(); + } + builder.finish().unwrap() +} + +fn payload_chunk( + key: Buffer, + times: &[i64], + floats: &[Option], + latency: &[Option], + mixed_sums: &[Option], + histogram_sums: &[Option], + last: bool, +) -> RawSeriesChunk { + assert_eq!(times.len(), floats.len()); + assert_eq!(times.len(), latency.len()); + assert_eq!(times.len(), mixed_sums.len()); + assert_eq!(times.len(), histogram_sums.len()); + let mixed_floats = Float64Array::from(latency.to_vec()); + let mut mixed_index = 0; + let mixed_indices = UInt32Array::from( + mixed_sums + .iter() + .map(|sum| { + sum.map(|_| { + let index = mixed_index; + mixed_index += 1; + index + }) + }) + .collect::>(), + ); + let mixed = Samples::try_from_mixed( + mixed_floats, + HistogramSamples::try_new( + mixed_indices, + histograms(&mixed_sums.iter().flatten().copied().collect::>()), + ) + .unwrap(), + ) + .unwrap(); + let mixed = if times.len() > 1 { + Samples::try_from_chunks(vec![ + mixed.slice(0, 1).unwrap(), + mixed.slice(1, times.len() - 1).unwrap(), + ]) + .unwrap() + } else { + mixed + }; + let mut histogram_index = 0; + let histogram_indices = UInt32Array::from( + histogram_sums + .iter() + .map(|sum| { + sum.map(|_| { + let index = histogram_index; + histogram_index += 1; + index + }) + }) + .collect::>(), + ); + let fields = SeriesFields::try_from_vec(vec![ + Samples::from_float(Float64Array::from(floats.to_vec())), + mixed, + { + let table = histograms(&histogram_sums.iter().flatten().copied().collect::>()); + let samples = Samples::try_from_histogram(histogram_indices, table).unwrap(); + if times.len() > 1 { + Samples::try_from_chunks(vec![ + samples.slice(0, 1).unwrap(), + samples.slice(1, times.len() - 1).unwrap(), + ]) + .unwrap() + } else { + samples + } + }, + ]) + .unwrap(); + RawSeriesChunk::try_new( + &payload_schema(), + key, + None, + RawTimestamps::try_new(ScalarBuffer::from(times.to_vec())).unwrap(), + fields, + last, + ) + .unwrap() +} + +struct PayloadStream { + schema: Arc, + labels: RawLabelSource, + items: VecDeque>, +} + +impl Stream for PayloadStream { + type Item = Result; + fn poll_next(mut self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { + Poll::Ready(self.items.pop_front()) + } +} + +impl RawSeriesStream for PayloadStream { + fn schema(&self) -> &Arc { + &self.schema + } + fn label_source(&self) -> &RawLabelSource { + &self.labels + } + fn ordering(&self) -> SeriesOrdering { + SeriesOrdering::KeyThenTime + } +} + +fn payload_stream(labels: RawLabelSource, chunks: Vec) -> SendableRawSeriesStream { + Box::pin(PayloadStream { + schema: payload_schema(), + labels, + items: chunks.into_iter().map(Ok).collect(), + }) +} + +fn assert_histogram(value: HistogramView<'_>, sum: f64) { + assert_eq!(value.schema, -19); + assert_eq!(value.zero_threshold, 0.125); + assert_eq!(value.sum.to_bits(), sum.to_bits()); + assert_eq!(value.count, 23.5); + assert_eq!(value.zero_count, 2.25); + assert_eq!(value.reset_hint, CounterResetHint::Unrecognized(37)); + assert_eq!(value.start_timestamp, Some(-456)); + assert_eq!(value.custom_values, &[1.5, 4.5]); + assert_eq!(value.positive.span_offsets, &[1]); + assert_eq!(value.positive.span_lengths, &[2]); + assert_eq!(value.positive.buckets, &[3.25, 7.5]); + assert_eq!(value.negative.span_offsets, &[-2]); + assert_eq!(value.negative.span_lengths, &[1]); + assert_eq!(value.negative.buckets, &[0.75]); +} + +#[tokio::test] +async fn merges_typed_multifield_payload_through_frame_and_label_roundtrip() { + let key = read_key("rack-payload", "sensor-7"); + let a = payload_stream( + label_source(1), + vec![payload_chunk( + key.clone(), + &[10, 30], + &[Some(f64::from_bits(0x7ff8_0000_0000_0042)), Some(-0.0)], + &[Some(1.25), None], + &[None, Some(11.0)], + &[Some(101.0), None], + true, + )], + ); + let b = payload_stream( + label_source(2), + vec![payload_chunk( + key.clone(), + &[20, 40], + &[None, Some(f64::from_bits(0x7ff0_0000_0000_0002))], + &[None, None], + &[None, Some(21.0)], + &[Some(202.0), Some(303.0)], + true, + )], + ); + let merged = merge_ordered_raw_series(vec![a, b]).unwrap(); + let frame_schema = merged.schema().clone(); + let frame_binding = merged.label_source().clone(); + let frame_ordering = merged.ordering(); + let framed = pack_raw_series(merged).unwrap(); + let raw = unpack_raw_series(framed, frame_schema, frame_binding, frame_ordering).unwrap(); + let rows = materialize_raw_labels(raw) + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!( + rows.iter() + .map(|r| r.time.timestamp(0).unwrap()) + .collect::>(), + [10, 20, 30, 40] + ); + assert!(rows.iter().all(|r| r.primary_key == key)); + for row in &rows { + let metric = row.metric.as_ref().unwrap(); + assert_eq!(metric.name(), Some("room_temperature")); + assert_eq!(metric.get("rack"), Some("rack-payload")); + assert_eq!(metric.get("sensor"), Some("sensor-7")); + } + let floats: Vec<_> = rows + .iter() + .map(|r| match r.fields.get(0).unwrap().sample(0) { + SampleRef::Float(v) => Some(v.to_bits()), + SampleRef::Missing => None, + _ => panic!("unexpected temperature sample variant"), + }) + .collect(); + assert_eq!( + floats, + [ + Some(0x7ff8_0000_0000_0042), + None, + Some((-0.0_f64).to_bits()), + Some(0x7ff0_0000_0000_0002) + ] + ); + let mixed: Vec<_> = rows + .iter() + .map(|r| match r.fields.get(1).unwrap().sample(0) { + SampleRef::Float(v) => Some(("float", v.to_bits())), + SampleRef::Missing => None, + SampleRef::Histogram(h) => { + assert_histogram( + h, + if r.time.timestamp(0) == Some(30) { + 11.0 + } else { + 21.0 + }, + ); + Some(("hist", 0)) + } + }) + .collect(); + assert_eq!( + mixed, + [ + Some(("float", 1.25_f64.to_bits())), + None, + Some(("hist", 0)), + Some(("hist", 0)) + ] + ); + let histograms: Vec<_> = rows + .iter() + .map(|r| match r.fields.get(2).unwrap().sample(0) { + SampleRef::Histogram(h) => { + let expected = match r.time.timestamp(0).unwrap() { + 10 => 101.0, + 20 => 202.0, + 40 => 303.0, + _ => panic!("unexpected histogram position"), + }; + assert_histogram(h, expected); + Some(h.sum.to_bits()) + } + SampleRef::Missing => None, + _ => panic!("unexpected observation sample variant"), + }) + .collect(); + assert_eq!( + histograms, + [ + Some(101.0_f64.to_bits()), + Some(202.0_f64.to_bits()), + None, + Some(303.0_f64.to_bits()) + ] + ); + assert!(!rows[0].is_last && !rows[1].is_last && !rows[2].is_last && rows[3].is_last); +} + +#[tokio::test] +async fn merge_retains_sliced_float_time_and_histogram_owners_until_output_drop() { + let key = read_key("rack-slice", "sensor-9"); + let times_owner = ScalarBuffer::from(vec![1, 11, 22, 99]); + let times = times_owner.slice(1, 2); + let floats_owner = Float64Array::from(vec![ + Some(90.0), + Some(f64::from_bits(0x7ff8_0000_0000_0042)), + Some(-0.0), + Some(80.0), + ]); + let floats = floats_owner.slice(1, 2); + let hist_owner = histograms(&[99.0, 12.5, 88.0]); + let hist_weak = Arc::downgrade(&hist_owner); + let sliced_hist = + Arc::new(HistogramBatch::try_from_arrow(hist_owner.as_arrow().slice(1, 1)).unwrap()); + let hist_indices_owner = UInt32Array::from(vec![Some(9), Some(0), Some(0), Some(9)]); + let hist_indices = hist_indices_owner.slice(1, 2); + let mixed_floats_owner = Float64Array::from(vec![Some(90.0), Some(7.0), None, Some(80.0)]); + let mixed_floats = mixed_floats_owner.slice(1, 2); + let mixed_indices_owner = UInt32Array::from(vec![Some(9), None, Some(0), Some(9)]); + let mixed_indices = mixed_indices_owner.slice(1, 2); + let fields = SeriesFields::try_from_vec(vec![ + Samples::from_float(floats.clone()), + Samples::try_from_mixed( + mixed_floats.clone(), + HistogramSamples::try_new(mixed_indices.clone(), sliced_hist.clone()).unwrap(), + ) + .unwrap(), + Samples::try_from_histogram(hist_indices.clone(), sliced_hist.clone()).unwrap(), + ]) + .unwrap(); + let chunk = RawSeriesChunk::try_new( + &payload_schema(), + key.clone(), + None, + RawTimestamps::try_new(times.clone()).unwrap(), + fields, + true, + ) + .unwrap(); + let time_ptr = times.as_ptr(); + let key_ptr = key.as_ptr(); + let float_ptr = floats.values().as_ptr(); + let indices_ptr = hist_indices.values().as_ptr(); + let mixed_float_ptr = mixed_floats.values().as_ptr(); + let mixed_indices_ptr = mixed_indices.values().as_ptr(); + let scalar_ptr = sliced_hist + .as_arrow() + .column(2) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .as_ptr(); + let custom_ptr = sliced_hist + .as_arrow() + .column(7) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .as_ptr(); + let bucket_ptr = sliced_hist + .as_arrow() + .column(10) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .as_ptr(); + let batch_weak: Weak = Arc::downgrade(&sliced_hist); + let interleaved = payload_chunk( + key.clone(), + &[15], + &[Some(6.0)], + &[None], + &[None], + &[Some(7.5)], + true, + ); + let mut merged = merge_ordered_raw_series(vec![ + payload_stream(label_source(1), vec![chunk]), + payload_stream(label_source(2), vec![interleaved]), + ]) + .unwrap(); + assert_eq!( + merged + .schema() + .fields + .iter() + .map(|field| field.name.as_ref()) + .collect::>(), + ["temperature", "latency", "observations"] + ); + assert!( + matches!(merged.label_source(), RawLabelSource::PrimaryKey { metric_name, .. } if metric_name.as_ref() == "room_temperature") + ); + assert_eq!(merged.ordering(), SeriesOrdering::KeyThenTime); + let prefix = merged.next().await.unwrap().unwrap(); + assert_eq!(prefix.time.as_buffer().as_ref(), &[11]); + assert_eq!(prefix.time.as_buffer().as_ptr(), time_ptr); + assert_eq!(prefix.primary_key.as_ptr(), key_ptr); + let interleaved = merged.next().await.unwrap().unwrap(); + assert_eq!(interleaved.time.as_buffer().as_ref(), &[15]); + let remainder = merged.next().await.unwrap().unwrap(); + assert_eq!(remainder.primary_key.as_ptr(), key_ptr); + assert_eq!(remainder.time.as_buffer().as_ref(), &[22]); + assert_eq!( + remainder.time.as_buffer().as_ptr(), + time_ptr.wrapping_add(1) + ); + assert!( + matches!(prefix.fields.get(0).unwrap().dense().unwrap(), DenseSamplesRef::Float(a) if a.values().as_ptr() == float_ptr) + ); + assert!( + matches!(remainder.fields.get(0).unwrap().dense().unwrap(), DenseSamplesRef::Float(a) if a.values().as_ptr() == float_ptr.wrapping_add(1)) + ); + let DenseSamplesRef::Mixed(prefix_mixed) = prefix.fields.get(1).unwrap().dense().unwrap() + else { + panic!("expected mixed payload") + }; + assert_eq!(prefix_mixed.floats().values().as_ptr(), mixed_float_ptr); + assert_eq!( + prefix_mixed.histograms().indices().values().as_ptr(), + mixed_indices_ptr + ); + assert!( + matches!(prefix.fields.get(1).unwrap().sample(0), SampleRef::Float(v) if v.to_bits() == 7.0_f64.to_bits()) + ); + let DenseSamplesRef::Mixed(remainder_mixed) = remainder.fields.get(1).unwrap().dense().unwrap() + else { + panic!("expected mixed payload") + }; + assert_eq!( + remainder_mixed.floats().values().as_ptr(), + mixed_float_ptr.wrapping_add(1) + ); + assert_eq!( + remainder_mixed.histograms().indices().values().as_ptr(), + mixed_indices_ptr.wrapping_add(1) + ); + assert!( + matches!(remainder.fields.get(1).unwrap().sample(0), SampleRef::Histogram(h) if h.sum.to_bits() == 12.5_f64.to_bits()) + ); + let DenseSamplesRef::Histogram(prefix_hist) = prefix.fields.get(2).unwrap().dense().unwrap() + else { + panic!("expected histogram payload") + }; + let DenseSamplesRef::Histogram(remainder_hist) = + remainder.fields.get(2).unwrap().dense().unwrap() + else { + panic!("expected histogram payload") + }; + assert_eq!(prefix_hist.indices().values().as_ptr(), indices_ptr); + assert_eq!(prefix_hist.indices().values(), &[0]); + assert_histogram(prefix_hist.view_at(0).unwrap(), 12.5); + assert!(Arc::ptr_eq(prefix_hist.batch(), remainder_hist.batch())); + let output_arrow = prefix_hist.batch().as_arrow(); + assert_eq!( + output_arrow + .column(2) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .as_ptr(), + scalar_ptr + ); + assert_eq!( + output_arrow + .column(7) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .as_ptr(), + custom_ptr + ); + assert_eq!( + output_arrow + .column(10) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .as_ptr(), + bucket_ptr + ); + let negative_bucket_ptr = sliced_hist + .as_arrow() + .column(13) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .as_ptr(); + assert_eq!( + output_arrow + .column(13) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .as_ptr(), + negative_bucket_ptr + ); + assert_eq!( + remainder_hist.indices().values().as_ptr(), + indices_ptr.wrapping_add(1) + ); + assert_eq!(remainder_hist.indices().values(), &[0]); + assert_histogram(remainder_hist.view_at(0).unwrap(), 12.5); + drop(merged); + drop(times_owner); + drop(times); + drop(floats_owner); + drop(floats); + drop(mixed_floats_owner); + drop(mixed_floats); + drop(mixed_indices_owner); + drop(mixed_indices); + drop(hist_owner); + drop(sliced_hist); + drop(hist_indices_owner); + drop(hist_indices); + assert!(hist_weak.upgrade().is_none()); + assert!(batch_weak.upgrade().is_some()); + assert_eq!(prefix.time.as_buffer().as_ref(), &[11]); + assert_eq!(remainder.time.as_buffer().as_ref(), &[22]); + let SampleRef::Histogram(prefix_histogram) = prefix.fields.get(2).unwrap().sample(0) else { + panic!("expected histogram"); + }; + assert_histogram(prefix_histogram, 12.5); + let SampleRef::Histogram(remainder_histogram) = remainder.fields.get(2).unwrap().sample(0) + else { + panic!("expected histogram"); + }; + assert_histogram(remainder_histogram, 12.5); + let SampleRef::Histogram(mixed_histogram) = remainder.fields.get(1).unwrap().sample(0) else { + panic!("expected histogram"); + }; + assert_histogram(mixed_histogram, 12.5); + assert!( + matches!(prefix.fields.get(0).unwrap().sample(0), SampleRef::Float(v) if v.to_bits() == 0x7ff8_0000_0000_0042) + ); + assert!( + matches!(remainder.fields.get(0).unwrap().sample(0), SampleRef::Float(v) if v.to_bits() == (-0.0_f64).to_bits()) + ); + drop(prefix); + assert!(batch_weak.upgrade().is_some()); + let SampleRef::Histogram(remainder_histogram) = remainder.fields.get(2).unwrap().sample(0) + else { + panic!("expected histogram"); + }; + assert_histogram(remainder_histogram, 12.5); + assert!( + matches!(remainder.fields.get(1).unwrap().sample(0), SampleRef::Histogram(h) if h.sum.to_bits() == 12.5_f64.to_bits() && h.custom_values == [1.5, 4.5] && h.positive.span_offsets == [1] && h.positive.span_lengths == [2] && h.positive.buckets == [3.25, 7.5] && h.negative.span_offsets == [-2] && h.negative.span_lengths == [1] && h.negative.buckets == [0.75]) + ); + drop(remainder); + drop(interleaved); + assert!(batch_weak.upgrade().is_none()); +}