diff --git a/src/promql/AGENTS.md b/src/promql/AGENTS.md index 6af954ee08e..437c58b1315 100644 --- a/src/promql/AGENTS.md +++ b/src/promql/AGENTS.md @@ -23,10 +23,11 @@ functions, vector matching and extension operators. Statement planning lives in Native PromQL is an experimental, standalone evaluation route alongside the existing DataFusion PromQL path. It operates on one storage region, which may -supply multiple scan partitions. It can evaluate each independent input series -in those partitions, but does not merge one logical series split across regions -or partitions; cross-region `SeriesMerge`/`Exchange` and global final-label -semantics are not implemented. Do not describe partition-local work as a global +supply multiple scan partitions. Where a boundary must answer one logical +series once, `SeriesRepartitionExec` merges that identity's fragments across +scan partitions and recomputes a single close for it. Cross-region +`SeriesMerge`/`Exchange` are still not implemented, so a logical series split +across regions is not merged. Do not describe partition-local work as a global series merge. Vector operations use the existing vector-matching semantics; do not describe diff --git a/src/promql/src/extension_plan.rs b/src/promql/src/extension_plan.rs index ba4cf65e0d5..c8c281a22ff 100644 --- a/src/promql/src/extension_plan.rs +++ b/src/promql/src/extension_plan.rs @@ -57,7 +57,7 @@ pub use range_manipulate::{RangeManipulate, RangeManipulateExec, RangeManipulate pub use raw_series_source::{RawSeriesSource, RawSeriesSourceExec}; pub use scalar_calculate::ScalarCalculate; pub use series_divide::{SeriesDivide, SeriesDivideExec, SeriesDivideStream}; -pub use series_repartition::SeriesRepartitionExec; +pub use series_repartition::{SeriesRepartitionExec, SeriesRouting}; pub use union_distinct_on::{UnionDistinctOn, UnionDistinctOnExec, UnionDistinctOnStream}; pub type Millisecond = ::Native; diff --git a/src/promql/src/extension_plan/series_repartition.rs b/src/promql/src/extension_plan/series_repartition.rs index af53ef796d9..fc3c89d68cb 100644 --- a/src/promql/src/extension_plan/series_repartition.rs +++ b/src/promql/src/extension_plan/series_repartition.rs @@ -12,12 +12,12 @@ // See the License for the specific language governing permissions and // limitations under the License. -//! Routes evaluated series chains by canonical label set across DataFusion input partitions. +//! Routes evaluated and raw series across DataFusion input partitions. //! -//! This node routes chains by canonical label-set hash, then merges exact metric-name/tag -//! identities at each owner after their close/EOF barriers. This is owner-local grouping, not a claim -//! of global canonical completeness. It operates on DataFusion partitions only; arbitrary `MatchKey` -//! and `GroupKey` routing is unsupported. +//! Canonical routing hashes complete labels; matching routing hashes the evaluator's own PromQL +//! matching key. Routing groups are not merge identities: evaluated series merge only by exact name +//! and tags. Raw ownership is the complete primary key plus exact label agreement; each input's +//! complete chain or EOF establishes its contribution to a raw identity. //! //! Chains are transported intact to owners with their input-partition ID, followed by one charged //! normal-EOF message per live output from each completed input. Owners group exact identities and @@ -121,7 +121,8 @@ use std::task::{Context, Poll}; use common_error::ext::BoxedError; use common_query::series_frame::{ - SendableEvaluatedSeriesStream, pack_evaluated_series, unpack_evaluated_series, + SendableEvaluatedSeriesStream, pack_evaluated_series, pack_raw_series, unpack_evaluated_series, + unpack_raw_series, }; use datafusion::arrow::datatypes::SchemaRef; use datafusion::arrow::record_batch::RecordBatch; @@ -135,18 +136,64 @@ use datafusion::physical_plan::{ DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, InputDistributionRequirements, Partitioning, PlanProperties, RecordBatchStream, SendableRecordBatchStream, }; -use datatypes::series::SeriesLabels; -use datatypes::series::frame::{FrameContent, SeriesFrameDeclaration}; +use datatypes::series::frame::{ + FrameContent, FrameLabelMode, LabelDecodeSpec, SeriesFrameDeclaration, +}; +use datatypes::series::raw::RawSeriesChunk; use datatypes::series::samples::{Series, SeriesSchema}; +use datatypes::series::{SeriesLabels, SeriesOrdering}; use futures::{Stream, StreamExt}; +use store_api::region_engine::{NO_BOUND_LABEL_SOURCE, RawSeriesStream, SendableRawSeriesStream}; use tokio::sync::mpsc; use twox_hash::XxHash64; use crate::grid_series::{label_set, merge_same_labelset}; -use crate::native::execute::merge_workspace_bytes; -use crate::native::execute::operand_chunk_bytes; +use crate::native::execute::{merge_workspace_bytes, operand_chunk_bytes}; +use crate::source::merge_raw_identity_charging; use datatypes::series::samples::DenseSamplesRef; +/// Selects the key used to route a series to an output partition. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum SeriesRouting { + /// Hash the canonical label set. + CanonicalLabels, + /// Hash the exact PromQL matching key, so both sides of a vector operation meet. + Matching(crate::vector_match::MatchingLabels), + /// Route raw frames by their complete stored primary key. + RawPrimaryKey, +} + +struct EmptyRawStream { + schema: Arc, + ordering: SeriesOrdering, +} + +impl Stream for EmptyRawStream { + type Item = Result; + + fn poll_next(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(None) + } +} + +impl RawSeriesStream for EmptyRawStream { + fn schema(&self) -> &Arc { + &self.schema + } + + fn label_source(&self) -> &store_api::region_engine::RawLabelSource { + &NO_BOUND_LABEL_SOURCE + } + + fn label_decode_spec(&self) -> Option<&LabelDecodeSpec> { + None + } + + fn ordering(&self) -> SeriesOrdering { + self.ordering + } +} + /// The seed behind every routing decision. /// /// Fixed rather than drawn per execution, so which output a label set belongs on is a function of @@ -165,6 +212,9 @@ pub struct SeriesRepartitionExec { series_schema: Arc, /// The one hash configuration of this node's execution. routing: Routing, + routing_policy: SeriesRouting, + content: FrameContent, + ordering: SeriesOrdering, /// Started by the first `execute` call, since the output partitions share one read of the child. shared: Mutex>>, properties: Arc, @@ -184,6 +234,25 @@ impl SeriesRepartitionExec { /// and a frame whose declaration does not survive the round trip would be published under a /// schema its batches do not match. pub fn try_new(input: Arc, outputs: usize) -> DataFusionResult { + let declaration = + SeriesFrameDeclaration::parse(input.schema().as_ref()).map_err(|error| { + DataFusionError::Plan(format!( + "SeriesRepartition input is not a series frame: {error}" + )) + })?; + let policy = match declaration.content() { + FrameContent::Evaluated => SeriesRouting::CanonicalLabels, + FrameContent::Raw => SeriesRouting::RawPrimaryKey, + }; + Self::try_new_with_routing(input, outputs, policy) + } + + /// Creates the node with an explicit content-compatible routing policy. + pub fn try_new_with_routing( + input: Arc, + outputs: usize, + routing_policy: SeriesRouting, + ) -> DataFusionResult { if outputs == 0 { return Err(DataFusionError::Plan( "SeriesRepartition requires at least one output partition".to_string(), @@ -200,17 +269,45 @@ impl SeriesRepartitionExec { "SeriesRepartition input is not a series frame: {error}" )) })?; - if declaration.content() != FrameContent::Evaluated { + let content = declaration.content(); + match (content, &routing_policy) { + (FrameContent::Evaluated, SeriesRouting::RawPrimaryKey) => { + return Err(DataFusionError::Plan( + "SeriesRepartition evaluated frames cannot use RawPrimaryKey routing".into(), + )); + } + (FrameContent::Raw, SeriesRouting::CanonicalLabels) => { + return Err(DataFusionError::Plan( + "SeriesRepartition raw frames cannot use CanonicalLabels routing".into(), + )); + } + _ => {} + } + if content == FrameContent::Raw && declaration.labels() != FrameLabelMode::Explicit { return Err(DataFusionError::Plan( - "SeriesRepartition routes evaluated series fragments, not raw samples".to_string(), + "SeriesRepartition raw frames require explicit labels".into(), )); } let series_schema = Arc::new(declaration.series().clone()); - let declared = pack_evaluated_series( - Box::pin(futures::stream::empty()), - Arc::clone(&series_schema), - )? - .schema(); + let ordering = if content == FrameContent::Raw { + common_query::series_frame::series_frame_ordering(input.schema().as_ref())? + } else { + SeriesOrdering::Contiguous + }; + let declared = match content { + FrameContent::Evaluated => pack_evaluated_series( + Box::pin(futures::stream::empty()), + Arc::clone(&series_schema), + )? + .schema(), + FrameContent::Raw => { + let source = EmptyRawStream { + schema: Arc::clone(&series_schema), + ordering, + }; + pack_raw_series(Box::pin(source))?.schema() + } + }; if declared != input.schema() { return Err(DataFusionError::Plan( "SeriesRepartition input frame carries a schema this node cannot re-declare" @@ -231,14 +328,22 @@ impl SeriesRepartitionExec { Ok(Self { input, series_schema, - routing: Routing::new(outputs), + routing: Routing::with_policy(outputs, routing_policy.clone()), + routing_policy, + content, + ordering, shared: Mutex::new(None), properties, }) } + /// The routing policy used by this node. + pub fn routing_policy(&self) -> &SeriesRouting { + &self.routing_policy + } + /// How many output partitions this node has. - fn outputs(&self) -> usize { + pub fn outputs(&self) -> usize { self.properties.partitioning.partition_count() } } @@ -313,9 +418,10 @@ impl ExecutionPlan for SeriesRepartitionExec { } // Rebuilt rather than patched, so a child whose schema changed is refused here instead of // routed under a frame declaration the new child does not make. - Ok(Arc::new(Self::try_new( + Ok(Arc::new(Self::try_new_with_routing( children.swap_remove(0), self.outputs(), + self.routing_policy.clone(), )?)) } @@ -340,7 +446,7 @@ impl ExecutionPlan for SeriesRepartitionExec { // emitted twice or read by a scan that was never meant to run twice. *state = Some(SharedExecution::start( &self.input, - self.routing, + self.routing.clone(), Arc::clone(&self.series_schema), &context, )?); @@ -359,9 +465,25 @@ impl ExecutionPlan for SeriesRepartitionExec { Arc::clone(&self.series_schema), &context.runtime_env().memory_pool, self.input.properties().partitioning.partition_count(), + self.ordering, ); let handoff = Arc::clone(&inner.handoff); - let output = pack_evaluated_series(Box::pin(inner), Arc::clone(&self.series_schema))?; + let output: SendableRecordBatchStream = match self.content { + FrameContent::Evaluated => pack_evaluated_series( + Box::pin(EvaluatedOutputStream { + inner: Some(Box::pin(inner)), + }), + Arc::clone(&self.series_schema), + )?, + FrameContent::Raw => { + let raw = RawOutputStream { + inner: Some(Box::pin(inner)), + schema: Arc::clone(&self.series_schema), + ordering: self.ordering, + }; + pack_raw_series(Box::pin(raw))? + } + }; let schema = output.schema(); let output = Box::pin(PackedHandoffStream { inner: Some(output), @@ -433,6 +555,11 @@ impl SharedExecution { series_schema: Arc, context: &Arc, ) -> DataFusionResult> { + // What the child carries is what it declares, not what the policy suggests: a raw frame + // can be routed by a match key too, so the policy no longer names the content. + let content = SeriesFrameDeclaration::parse(input.schema().as_ref()) + .map_err(|error| DataFusionError::Plan(error.to_string()))? + .content(); let partitions = input.properties().partitioning.partition_count(); let mut streams = Vec::with_capacity(partitions); for child_partition in 0..partitions { @@ -440,10 +567,15 @@ impl SharedExecution { // The expected schema is a check and never an input: the decoder reads the frame's own // declaration, and a partition declaring another frame is refused here rather than // routed under labels it does not carry. - streams.push(unpack_evaluated_series( - frames, - Some(Arc::clone(&series_schema)), - )?); + streams.push(match content { + FrameContent::Evaluated => RoutedInput::Evaluated(unpack_evaluated_series( + frames, + Some(Arc::clone(&series_schema)), + )?), + FrameContent::Raw => { + RoutedInput::Raw(unpack_raw_series(frames, Some(Arc::clone(&series_schema)))?) + } + }); } let mut senders = Vec::with_capacity(routing.outputs); let mut receivers = Vec::with_capacity(routing.outputs); @@ -487,7 +619,7 @@ impl SharedExecution { let task = common_runtime::spawn_query(route_partition( stream, partition, - routing, + routing.clone(), senders.clone(), // What this partition reads ahead is charged to its own consumer, so the read-ahead // of one partition is attributable to the task that did it. @@ -839,7 +971,7 @@ impl Inputs { /// queued chains remain charged to the pool until consumed or dropped. Each task records its outcome /// before returning and dropping its senders, so output streams observe the final read outcome. async fn route_partition( - mut stream: SendableEvaluatedSeriesStream, + mut stream: RoutedInput, partition: usize, routing: Routing, senders: Vec>, @@ -847,7 +979,15 @@ async fn route_partition( inputs: Arc, mut report: ReportIfUnreported, ) { - let stop = read_chains(&mut stream, partition, &routing, &senders, &read_ahead).await; + // The variant is the content: decoding the input's declared frame is what produced it. + let stop = match &mut stream { + RoutedInput::Evaluated(stream) => { + read_evaluated_chains(stream, partition, &routing, &senders, &read_ahead).await + } + RoutedInput::Raw(stream) => { + read_raw_chains(stream, partition, &routing, &senders, &read_ahead).await + } + }; if matches!(&stop, Stop::Ended) { for sender in &senders { if sender.is_closed() { @@ -963,7 +1103,7 @@ pub(crate) fn histogram_table_slice_bytes( /// to the pool. Output queues are unbounded, so routing does not wait for output polling. Each queued /// chain retains its pool charge until the owner consumes it or the queue is dropped. This function /// otherwise waits only on its own child. -async fn read_chains( +async fn read_evaluated_chains( stream: &mut SendableEvaluatedSeriesStream, partition: usize, routing: &Routing, @@ -973,7 +1113,7 @@ async fn read_chains( loop { // One whole chain, assembled here: the chunk that closes the series ends it, and so does a // stream that ends between two chains. - let mut chain: Chain = Vec::new(); + let mut chain: Vec<(Series, bool, MemoryReservation)> = Vec::new(); let mut chain_output: Option = None; let mut discarding = false; let mut ended = false; @@ -1000,7 +1140,7 @@ async fn read_chains( let output = match chain_output { Some(output) => output, None => { - let output = routing.output(&series.metric); + let output = routing.output_labels(&series.metric); chain_output = Some(output); output } @@ -1095,7 +1235,7 @@ async fn read_chains( senders[output] .send(Queued::Chain { input: partition, - chain, + chain: RoutedChain::Evaluated(chain), _metadata: metadata, }) .is_err() @@ -1106,6 +1246,112 @@ async fn read_chains( } } +async fn read_raw_chains( + stream: &mut SendableRawSeriesStream, + partition: usize, + routing: &Routing, + senders: &[mpsc::UnboundedSender], + read_ahead: &MemoryReservation, +) -> Stop { + loop { + let mut chain: Vec<(RawSeriesChunk, MemoryReservation)> = Vec::new(); + let mut destination: Option = None; + let mut discarding = false; + loop { + tokio::task::coop::consume_budget().await; + let next = match (destination, discarding) { + (Some(output), false) if !chain.is_empty() => tokio::select! { + biased; + _ = senders[output].closed() => { chain.clear(); discarding = true; continue; } + item = stream.next() => item, + }, + _ => stream.next().await, + }; + match next { + Some(Ok(chunk)) => { + let output = match destination { + Some(output) => output, + None => { + let output = match &routing.policy { + SeriesRouting::Matching(_) => { + let Some(labels) = chunk.metric.as_ref() else { + return Stop::Failed(Arc::new(DataFusionError::Execution( + "a matching-routed raw chain carries no labels".into(), + ))); + }; + routing.output_labels(labels) + } + _ => routing.output_key(chunk.primary_key.as_slice()), + }; + *destination.insert(output) + } + }; + let final_chunk = chunk.is_last; + if senders[output].is_closed() { + chain.clear(); + discarding = true; + if final_chunk { + break; + } + continue; + } + if discarding { + if final_chunk { + break; + } + continue; + } + let charge = read_ahead.new_empty(); + let bytes = match raw_chunk_slice_bytes(&chunk) { + Ok(bytes) => bytes, + Err(error) => return Stop::Failed(Arc::new(error)), + }; + if let Err(error) = charge.try_grow(bytes) { + return Stop::Failed(Arc::new(error)); + } + chain.push((chunk, charge)); + if final_chunk { + break; + } + } + Some(Err(error)) => { + return Stop::Failed(Arc::new(DataFusionError::External(Box::new(error)))); + } + None => { + return if destination.is_none() { + Stop::Ended + } else { + Stop::Failed(Arc::new(DataFusionError::Execution( + "raw series ended without is_last".into(), + ))) + }; + } + } + } + let Some(output) = destination else { + return Stop::Failed(Arc::new(DataFusionError::Internal( + "raw chain has no destination".into(), + ))); + }; + if !discarding { + let metadata = read_ahead.new_empty(); + if let Err(error) = metadata.try_grow(std::mem::size_of::()) { + return Stop::Failed(Arc::new(error)); + } + if senders[output] + .send(Queued::Chain { + input: partition, + chain: RoutedChain::Raw(chain), + _metadata: metadata, + }) + .is_err() + { + continue; + } + } + } +} + fn exact_identity(series: &Series) -> ExactIdentity { ( series.metric.name().map(Arc::::from), @@ -1225,13 +1471,21 @@ enum Ending { /// /// The transport unit is a chain so queued source contributions remain intact. Owners collect all /// chains before locally merging exact identities and producing their final output. -type Chain = Vec<(Series, bool, MemoryReservation)>; +enum RoutedInput { + Evaluated(SendableEvaluatedSeriesStream), + Raw(SendableRawSeriesStream), +} + +enum RoutedChain { + Evaluated(Vec<(Series, bool, MemoryReservation)>), + Raw(Vec<(RawSeriesChunk, MemoryReservation)>), +} /// What one output queue carries: a whole chain and its memory charge. enum Queued { Chain { input: usize, - chain: Chain, + chain: RoutedChain, _metadata: MemoryReservation, }, Eof { @@ -1242,14 +1496,199 @@ enum Queued { type ExactIdentity = (Option>, Vec<(Arc, Arc)>); +/// Turns an owner's evaluated publications into the evaluated stream the packer consumes. +struct EvaluatedOutputStream { + inner: Option>>, +} + +impl Stream for EvaluatedOutputStream { + type Item = Result<(Series, bool), BoxedError>; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + let Some(inner) = this.inner.as_mut() else { + return Poll::Ready(None); + }; + match inner.as_mut().poll_next(cx) { + Poll::Pending => Poll::Pending, + Poll::Ready(None) => { + this.inner = None; + Poll::Ready(None) + } + Poll::Ready(Some(Err(error))) => { + this.inner = None; + Poll::Ready(Some(Err(error))) + } + // A merged identity is final: the owner publishes it once, already closed. + Poll::Ready(Some(Ok(ReadyOutput::Evaluated(series)))) => { + Poll::Ready(Some(Ok((series, true)))) + } + Poll::Ready(Some(Ok(ReadyOutput::Raw(_)))) => { + this.inner = None; + Poll::Ready(Some(Err(boxed(DataFusionError::Internal( + "an evaluated repartition output carried a raw series".into(), + ))))) + } + } + } +} + +/// Turns an owner's raw publications into the raw frame stream the packer consumes. +struct RawOutputStream { + inner: Option>>, + schema: Arc, + ordering: SeriesOrdering, +} + +impl Stream for RawOutputStream { + type Item = Result; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + let Some(inner) = this.inner.as_mut() else { + return Poll::Ready(None); + }; + match inner.as_mut().poll_next(cx) { + Poll::Pending => Poll::Pending, + Poll::Ready(None) => { + this.inner = None; + Poll::Ready(None) + } + Poll::Ready(Some(Err(error))) => { + this.inner = None; + Poll::Ready(Some(Err(error))) + } + Poll::Ready(Some(Ok(ReadyOutput::Raw(chunk)))) => Poll::Ready(Some(Ok(chunk))), + Poll::Ready(Some(Ok(ReadyOutput::Evaluated(_)))) => { + this.inner = None; + Poll::Ready(Some(Err(boxed(DataFusionError::Internal( + "a raw repartition output carried an evaluated series".into(), + ))))) + } + } + } +} + +impl RawSeriesStream for RawOutputStream { + fn schema(&self) -> &Arc { + &self.schema + } + + fn label_source(&self) -> &store_api::region_engine::RawLabelSource { + &store_api::region_engine::NO_BOUND_LABEL_SOURCE + } + + fn label_decode_spec(&self) -> Option<&LabelDecodeSpec> { + None + } + + fn ordering(&self) -> SeriesOrdering { + self.ordering + } +} + +/// What one owner groups by: an evaluated exact identity, or a complete raw primary key. +/// +/// The two are separate variants rather than one tuple, so a raw key can never be confused with an +/// evaluated identity that happens to look similar. +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)] +enum OwnerKey { + Evaluated(ExactIdentity), + Raw(Vec), +} + +/// The raw identity of one chunk: its complete stored primary key and nothing else. +/// +/// Labels are deliberately not part of the key. Two pieces of one physical key must always meet in +/// the same owner so a label disagreement can be rejected there, instead of being split across +/// owners that would each look consistent on its own. +fn raw_chunk_slice_bytes(chunk: &RawSeriesChunk) -> DataFusionResult { + use datafusion::arrow::array::Array; + + let metric_bytes = chunk.metric.as_ref().map_or(0, |metric| { + metric + .name() + .map_or(0, |name| name.len()) + .saturating_add(metric.tags().iter().fold(0usize, |bytes, tag| { + bytes + .saturating_add(tag.name.len()) + .saturating_add(tag.value.len()) + })) + }); + // A histogram field is an index array naming rows of a shared table; both are held per sample. + let mut fields_bytes = 0usize; + for samples in chunk.fields.iter() { + let field_bytes = match samples.dense() { + Some(DenseSamplesRef::Float(values)) => values.to_data().get_slice_memory_size()?, + Some(DenseSamplesRef::Histogram(histograms)) => histograms + .indices() + .to_data() + .get_slice_memory_size()? + .saturating_add(histogram_table_slice_bytes(histograms.batch())?), + Some(DenseSamplesRef::Mixed(mixed)) => mixed + .floats() + .to_data() + .get_slice_memory_size()? + .saturating_add( + mixed + .histograms() + .indices() + .to_data() + .get_slice_memory_size()?, + ) + .saturating_add(histogram_table_slice_bytes(mixed.histograms().batch())?), + None => 0, + }; + fields_bytes = fields_bytes.saturating_add(field_bytes); + } + Ok(std::mem::size_of::() + .saturating_add(chunk.primary_key.len()) + .saturating_add(metric_bytes) + .saturating_add( + chunk + .time + .as_buffer() + .len() + .saturating_mul(std::mem::size_of::()), + ) + .saturating_add(fields_bytes)) +} + +fn raw_identity(chunk: &RawSeriesChunk) -> OwnerKey { + OwnerKey::Raw(chunk.primary_key.as_slice().to_vec()) +} + struct PendingGroup { fragments: Vec, + raw_chunks: Vec, charges: Vec, received: Vec, metadata: MemoryReservation, } -fn group_ready(group: &PendingGroup, eof: &[bool]) -> bool { +/// Whether this owner may publish `group` now. +/// +/// A key-then-time raw frame promises its keys arrive in order, and this owner publishes in key +/// order only when it publishes them together: an identity that is complete while a smaller one is +/// still waiting would go out first and break that promise. Watermarks let it publish once every +/// input has advanced past the identity, without waiting for every input to end. +fn group_ready( + group: &PendingGroup, + eof: &[bool], + last_key_from: &[Option>], + ordering: SeriesOrdering, + key: &OwnerKey, +) -> bool { + if ordering == SeriesOrdering::KeyThenTime { + let OwnerKey::Raw(key) = key else { + return false; + }; + return group.received.iter().zip(eof).zip(last_key_from).all( + |((received, eof), last_key)| { + *received || *eof || last_key.as_ref().is_some_and(|last| last > key) + }, + ); + } group .received .iter() @@ -1257,8 +1696,13 @@ fn group_ready(group: &PendingGroup, eof: &[bool]) -> bool { .all(|(received, eof)| *received || *eof) } +enum ReadyOutput { + Evaluated(Series), + Raw(RawSeriesChunk), +} + struct ReadySeries { - series: Series, + series: ReadyOutput, charges: ReadyCharges, } @@ -1274,8 +1718,11 @@ struct RepartitionStream { lease: OutputLease, schema: Arc, pool: Arc, - groups: BTreeMap, + /// What the frame this stream emits declares about the order of its keys. + ordering: SeriesOrdering, + groups: BTreeMap, eof: Vec, + last_key_from: Vec>>, eof_metadata: Option, inputs: usize, output: VecDeque, @@ -1288,7 +1735,7 @@ struct RepartitionStream { /// every identity it releases. rescan: bool, /// The last identity a pass examined, exclusive. Retains at most one identity descriptor. - resume_after: Option, + resume_after: Option, reported: bool, } @@ -1346,14 +1793,17 @@ impl RepartitionStream { schema: Arc, pool: &Arc, inputs: usize, + ordering: SeriesOrdering, ) -> Self { Self { receiver, lease, schema, + ordering, pool: Arc::clone(pool), groups: BTreeMap::new(), eof: Vec::new(), + last_key_from: vec![None; inputs], eof_metadata: None, inputs, output: VecDeque::new(), @@ -1364,10 +1814,13 @@ impl RepartitionStream { } } - fn finish_group(&mut self, key: &ExactIdentity) -> DataFusionResult<()> { + fn finish_group(&mut self, key: &OwnerKey) -> DataFusionResult<()> { let Some(group) = self.groups.remove(key) else { return Ok(()); }; + if matches!(key, OwnerKey::Raw(_)) { + return self.finish_raw_group(group); + } let workspace = if group.fragments.len() > 1 { let rows = group.fragments.iter().fold(0usize, |rows, series| { rows.saturating_add(series.time.len()) @@ -1395,7 +1848,7 @@ impl RepartitionStream { )); }; self.output.push_back(ReadySeries { - series, + series: ReadyOutput::Evaluated(series), charges: ReadyCharges { _charges: group.charges, _metadata: group.metadata, @@ -1406,8 +1859,78 @@ impl RepartitionStream { } } +impl RepartitionStream { + /// Publishes one complete raw identity. + /// + /// Raw pieces are only published once every input reached its normal EOF, because a piece's + /// `is_last` closes that producer's stream, not the series. The workspace this reconstruction + /// needs is charged before it runs and rides with the output until the packer hands the batch + /// off, like any other owner output. + fn finish_raw_group(&mut self, group: PendingGroup) -> DataFusionResult<()> { + let labels = group + .raw_chunks + .first() + .and_then(|chunk| chunk.metric.clone()) + .ok_or_else(|| { + DataFusionError::Execution("raw identity has no explicit labels".to_string()) + })?; + if group + .raw_chunks + .iter() + .any(|chunk| chunk.metric.as_ref() != Some(&labels)) + { + return Err(DataFusionError::Execution( + "raw identity pieces disagree on their explicit labels".to_string(), + )); + } + if group.raw_chunks.len() == 1 { + let mut chunk = group.raw_chunks.into_iter().next().expect("one raw chunk"); + chunk.is_last = true; + self.output.push_back(ReadySeries { + series: ReadyOutput::Raw(chunk), + charges: ReadyCharges { + _charges: group.charges, + _metadata: group.metadata, + _workspace: None, + }, + }); + return Ok(()); + } + // Reconstruction's working set depends on how the pieces interleave, which is only known + // while it runs: charge each allocation as it is made rather than reserving a worst case + // that a well-ordered identity never approaches. + let workspace = MemoryConsumer::new("SeriesRepartitionRawOwner").register(&self.pool); + let mut reserve = |bytes: usize| workspace.try_grow(bytes); + let mut shrink = |bytes: usize| { + let result = workspace.try_shrink(bytes); + debug_assert!( + result.is_ok(), + "raw merge workspace shrink failed: {result:?}" + ); + // The merged data is already correct; tolerate accounting inconsistency rather than + // turning a valid query result into an error. + let _ = result; + }; + let chunk = merge_raw_identity_charging( + &self.schema, + &group.raw_chunks, + &mut reserve, + &mut shrink, + )?; + self.output.push_back(ReadySeries { + series: ReadyOutput::Raw(chunk), + charges: ReadyCharges { + _charges: group.charges, + _metadata: group.metadata, + _workspace: Some(workspace), + }, + }); + Ok(()) + } +} + impl Stream for RepartitionStream { - type Item = Result<(Series, bool), BoxedError>; + type Item = Result; fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let this = self.get_mut(); @@ -1456,19 +1979,34 @@ impl Stream for RepartitionStream { return Poll::Ready(Some(Err(boxed(DataFusionError::from(&error))))); } if this.rescan { - let next = match &this.resume_after { - Some(after) => this - .groups - .range(( - std::ops::Bound::Excluded(after.clone()), - std::ops::Bound::Unbounded, - )) - .find_map(|(key, group)| { - group_ready(group, &this.eof).then(|| key.clone()) + let next = if this.ordering == SeriesOrdering::KeyThenTime { + this.groups.first_key_value().and_then(|(key, group)| { + group_ready(group, &this.eof, &this.last_key_from, this.ordering, key) + .then(|| key.clone()) + }) + } else { + match &this.resume_after { + Some(after) => this + .groups + .range(( + std::ops::Bound::Excluded(after.clone()), + std::ops::Bound::Unbounded, + )) + .find_map(|(key, group)| { + group_ready( + group, + &this.eof, + &this.last_key_from, + this.ordering, + key, + ) + .then(|| key.clone()) + }), + None => this.groups.iter().find_map(|(key, group)| { + group_ready(group, &this.eof, &this.last_key_from, this.ordering, key) + .then(|| key.clone()) }), - None => this.groups.iter().find_map(|(key, group)| { - group_ready(group, &this.eof).then(|| key.clone()) - }), + } }; match next { Some(key) => { @@ -1493,7 +2031,7 @@ impl Stream for RepartitionStream { .lock() .unwrap_or_else(|e| e.into_inner()) .current = Some(ready.charges); - return Poll::Ready(Some(Ok((ready.series, true)))); + return Poll::Ready(Some(Ok(ready.series))); } if this.reported { return Poll::Ready(None); @@ -1509,36 +2047,67 @@ impl Stream for RepartitionStream { _metadata, })) => { drained += 1; - let Some(first) = chain.first() else { - continue; + let (key, pieces) = match chain { + RoutedChain::Evaluated(chain) => { + let Some(first) = chain.first() else { + continue; + }; + let key = exact_identity(&first.0); + if chain + .iter() + .any(|(series, _, _)| exact_identity(series) != key) + { + this.lease.fail(Arc::new(DataFusionError::Execution( + "one series chain contains multiple exact identities".into(), + ))); + continue; + } + (OwnerKey::Evaluated(key), RoutedChain::Evaluated(chain)) + } + RoutedChain::Raw(chain) => { + let Some(first) = chain.first() else { + continue; + }; + let key = raw_identity(&first.0); + if chain.iter().any(|(chunk, _)| raw_identity(chunk) != key) { + this.lease.fail(Arc::new(DataFusionError::Execution( + "one raw chain contains multiple primary keys".into(), + ))); + continue; + } + if let OwnerKey::Raw(raw_key) = &key + && this.ordering == SeriesOrdering::KeyThenTime + { + let last = &mut this.last_key_from[input]; + if last.as_ref().is_none_or(|last| raw_key > last) { + *last = Some(raw_key.clone()); + } + } + (key, RoutedChain::Raw(chain)) + } }; - let key = exact_identity(&first.0); - if chain - .iter() - .any(|(series, _, _)| exact_identity(series) != key) - { - this.lease.fail(Arc::new(DataFusionError::Execution( - "one series chain contains multiple exact identities".into(), - ))); - continue; - } if !this.groups.contains_key(&key) { let metadata = MemoryConsumer::new("SeriesRepartitionIdentity").register(&this.pool); + let key_bytes = + match &key { + OwnerKey::Evaluated((name, tags)) => { + std::mem::size_of::() + .saturating_add(name.as_ref().map_or(0, |name| name.len())) + .saturating_add(tags.capacity().saturating_mul( + std::mem::size_of::<(Arc, Arc)>(), + )) + .saturating_add( + tags.iter() + .map(|(n, v)| n.len().saturating_add(v.len())) + .sum::(), + ) + } + OwnerKey::Raw(key) => key.capacity(), + }; let bytes = std::mem::size_of::() - .saturating_add(std::mem::size_of::()) - .saturating_add(key.0.as_ref().map_or(0, |name| name.len())) - .saturating_add(key.1.capacity().saturating_mul(std::mem::size_of::<( - Arc, - Arc, - )>( - ))) - .saturating_add( - key.1 - .iter() - .map(|(n, v)| n.len().saturating_add(v.len())) - .sum::(), - ) + .saturating_add(std::mem::size_of::()) + .saturating_add(key_bytes) .saturating_add(this.eof.len().saturating_mul(2)); if let Err(error) = metadata.try_grow(bytes) { this.lease.fail(Arc::new(error)); @@ -1548,6 +2117,7 @@ impl Stream for RepartitionStream { key.clone(), PendingGroup { fragments: Vec::new(), + raw_chunks: Vec::new(), charges: Vec::new(), received: vec![false; this.eof.len()], metadata, @@ -1567,12 +2137,30 @@ impl Stream for RepartitionStream { continue; } group.received[input] = true; - for (series, _, charge) in chain { - group.fragments.push(series); - group.charges.push(charge); + match pieces { + RoutedChain::Evaluated(chain) => { + for (series, _, charge) in chain { + group.fragments.push(series); + group.charges.push(charge); + } + } + RoutedChain::Raw(chain) => { + for (chunk, charge) in chain { + group.raw_chunks.push(chunk); + group.charges.push(charge); + } + } } - if group_ready(group, &this.eof) - && let Err(error) = this.finish_group(&key) + if this.ordering == SeriesOrdering::KeyThenTime { + this.rescan = true; + this.resume_after = None; + } else if group_ready( + group, + &this.eof, + &this.last_key_from, + this.ordering, + &key, + ) && let Err(error) = this.finish_group(&key) { this.lease.fail(Arc::new(error)); } @@ -1612,7 +2200,7 @@ impl Stream for RepartitionStream { .lock() .unwrap_or_else(|e| e.into_inner()) .current = Some(ready.charges); - return Poll::Ready(Some(Ok((ready.series, true)))); + return Poll::Ready(Some(Ok(ready.series))); } return Poll::Ready(None); } @@ -1631,33 +2219,57 @@ fn boxed(error: DataFusionError) -> BoxedError { /// /// The seed is what an execution builds once; a [`XxHash64`] accumulates, so a fresh hasher starts /// from this configuration for each chunk, and that is the only thing each chunk builds. -#[derive(Clone, Copy, Debug)] +#[derive(Clone, Debug)] struct Routing { - /// The number of output partitions a hash is reduced modulo. outputs: usize, - /// The seed every chunk of this execution is hashed with. seed: u64, + policy: SeriesRouting, } impl Routing { - /// The routing of an execution with `outputs` output partitions. + /// The canonical routing of an execution with `outputs` output partitions. + #[cfg(test)] fn new(outputs: usize) -> Self { + Self::with_policy(outputs, SeriesRouting::CanonicalLabels) + } + + /// The routing of an execution with `outputs` output partitions and an explicit policy. + fn with_policy(outputs: usize, policy: SeriesRouting) -> Self { Self { outputs, seed: ROUTING_SEED, + policy, } } - /// The output partition `metric`'s canonical label set belongs on. - /// - /// The whole label set is hashed, since that is what a merge groups by: pieces of two series - /// whose label sets differ have to be able to land apart, or the split this node exists to make - /// would not happen. [`label_set`] is what canonicalizes the labels - dropping empty values and - /// keeping the metric name - so two pieces that merge hash the same. + /// The output a label-carrying chunk belongs on. + #[cfg(test)] fn output(&self, metric: &SeriesLabels) -> usize { + self.output_labels(metric) + } + + fn output_labels(&self, metric: &SeriesLabels) -> usize { let mut hasher = XxHash64::with_seed(self.seed); - label_set(metric).hash(&mut hasher); - // `outputs` is at least one, which [`SeriesRepartitionExec::try_new`] is what guarantees. + match &self.policy { + SeriesRouting::CanonicalLabels => label_set(metric).hash(&mut hasher), + SeriesRouting::Matching(policy) => { + crate::vector_match::group_key(metric, policy).hash(&mut hasher) + } + // A raw frame routed by its stored key never reaches here: its destination is the + // stored key itself, not a label hash. Evaluated frames never use this policy. + SeriesRouting::RawPrimaryKey => { + debug_assert!( + false, + "RawPrimaryKey routing hashes stored keys, not labels" + ); + } + } + (hasher.finish() % self.outputs as u64) as usize + } + + fn output_key(&self, key: &[u8]) -> usize { + let mut hasher = XxHash64::with_seed(self.seed); + key.hash(&mut hasher); (hasher.finish() % self.outputs as u64) as usize } } diff --git a/src/promql/src/extension_plan/series_repartition/tests.rs b/src/promql/src/extension_plan/series_repartition/tests.rs index 6dcbff6ef61..46aa8a511cd 100644 --- a/src/promql/src/extension_plan/series_repartition/tests.rs +++ b/src/promql/src/extension_plan/series_repartition/tests.rs @@ -38,7 +38,8 @@ use std::sync::{Arc, Mutex, Weak}; use std::task::{Poll, Waker}; use std::time::Duration; -use datafusion::arrow::array::{Array, Float64Array, ListArray, StructArray}; +use datafusion::arrow::array::{Array, Float64Array, ListArray, StructArray, UInt32Array}; +use datafusion::arrow::buffer::{Buffer, ScalarBuffer}; use datafusion::common::config::ConfigOptions; use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryPool, UnboundedMemoryPool}; use datafusion::execution::runtime_env::RuntimeEnvBuilder; @@ -54,6 +55,7 @@ use datatypes::series::Label; use datatypes::series::histogram::{ BucketSideView, CounterResetHint, HistogramBatchBuilder, HistogramView, }; +use datatypes::series::raw::RawTimestamps; use datatypes::series::samples::{FieldSpec, Samples, SeriesFields}; use datatypes::series::time::{EvalGrid, GridSpan, GridSpec, TimeAxis}; use futures::TryStreamExt; @@ -519,6 +521,236 @@ fn node(child: Arc, outputs: usize) -> Arc Arc { + Arc::new(SeriesSchema { + fields: Arc::from([FieldSpec { + name: "value".into(), + }]), + eval_grid: None, + }) +} + +/// One chunk of primary key `key`, carrying `metric`'s labels, over `times`/`values`. +fn raw_chunk( + key: &str, + metric: Option, + times: &[i64], + values: Vec>, + last: bool, +) -> RawSeriesChunk { + let schema = raw_schema(); + let fields = vec![Samples::from_float(Float64Array::from(values))]; + RawSeriesChunk::try_new( + &schema, + Buffer::from_slice_ref(key.as_bytes()), + metric, + RawTimestamps::try_new(ScalarBuffer::from(times.to_vec())).expect("ascending timestamps"), + SeriesFields::try_from_vec(fields).expect("one sample per timestamp per field"), + last, + ) + .expect("a legal raw chunk") +} + +/// `{__name__="m", host=""}`. +fn raw_metric(host: &str) -> SeriesLabels { + SeriesLabels::try_new( + Some(Arc::from("m")), + vec![Label { + name: Arc::from("host"), + value: Arc::from(host), + }], + ) + .expect("legal labels") +} + +/// One raw partition: its chunks, and an optional gate that keeps the stream open afterwards. +#[derive(Clone)] +struct RawPartition { + chunks: Vec, + gate: Option>, + gate_chunks: bool, +} + +struct RawScript { + chunks: std::vec::IntoIter, + schema: Arc, + binding: store_api::region_engine::RawLabelSource, + ordering: SeriesOrdering, + gate: Option>, + gate_chunks: bool, +} + +impl Stream for RawScript { + type Item = Result; + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + if self.gate_chunks + && let Some(gate) = &self.gate + && !gate.is_open() + { + gate.wait(cx); + return Poll::Pending; + } + if let Some(chunk) = self.chunks.next() { + return Poll::Ready(Some(Ok(chunk))); + } + if let Some(gate) = &self.gate + && !gate.is_open() + { + gate.wait(cx); + return Poll::Pending; + } + Poll::Ready(None) + } +} + +impl RawSeriesStream for RawScript { + fn schema(&self) -> &Arc { + &self.schema + } + fn label_source(&self) -> &store_api::region_engine::RawLabelSource { + &self.binding + } + fn ordering(&self) -> SeriesOrdering { + self.ordering + } +} + +/// A child whose partitions hand over raw chunks, one partition per list. +struct RawScripted { + partitions: Vec, + series_schema: Arc, + ordering: SeriesOrdering, + properties: Arc, +} + +impl RawScripted { + fn plan(partitions: Vec) -> Arc { + Self::plan_with_ordering(partitions, SeriesOrdering::Contiguous) + } + + /// The child over `partitions` declaring `ordering` for its keys. + fn plan_with_ordering( + partitions: Vec, + ordering: SeriesOrdering, + ) -> Arc { + let series_schema = raw_schema(); + let frame = common_query::series_frame::raw_series_frame_schema_with_ordering( + &series_schema, + &store_api::region_engine::RawLabelSource::Explicit, + ordering, + ) + .expect("a raw schema declares an explicit-label frame"); + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(frame), + Partitioning::UnknownPartitioning(partitions.len()), + EmissionType::Incremental, + Boundedness::Bounded, + )); + Arc::new(Self { + partitions, + series_schema, + ordering, + properties, + }) + } +} + +impl fmt::Debug for RawScripted { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + f.debug_struct("RawScripted") + .field("partitions", &self.partitions.len()) + .finish() + } +} + +impl DisplayAs for RawScripted { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> fmt::Result { + write!(f, "RawScripted: partitions={}", self.partitions.len()) + } +} + +impl ExecutionPlan for RawScripted { + fn name(&self) -> &str { + "RawScripted" + } + fn properties(&self) -> &Arc { + &self.properties + } + fn children(&self) -> Vec<&Arc> { + Vec::new() + } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> DataFusionResult, + ) -> DataFusionResult { + Ok(TreeNodeRecursion::Continue) + } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> DataFusionResult> { + assert!(children.is_empty(), "a raw scripted child has no children"); + Ok(self) + } + fn execute( + &self, + partition: usize, + _context: Arc, + ) -> DataFusionResult { + let spec = self.partitions.get(partition).cloned().ok_or_else(|| { + DataFusionError::Execution(format!("raw child has no partition {partition}")) + })?; + pack_raw_series(Box::pin(RawScript { + chunks: spec.chunks.into_iter(), + schema: Arc::clone(&self.series_schema), + binding: store_api::region_engine::RawLabelSource::Explicit, + // The packed batches have to declare what this child's own plan declares: the + // sequence validator and the receiver both read the ordering off the frame, so a + // stream that reported a different one would be validated under the wrong contract. + ordering: self.ordering, + gate: spec.gate, + gate_chunks: spec.gate_chunks, + })) + } +} + +/// The raw owner over `child` routing by an explicit policy. +fn raw_node_with_routing( + child: Arc, + outputs: usize, + routing: SeriesRouting, +) -> DataFusionResult { + SeriesRepartitionExec::try_new_with_routing(child, outputs, routing) +} + +/// The raw owner over `child` grouping by complete stored primary key. +fn raw_node(child: Arc, outputs: usize) -> Arc { + Arc::new( + SeriesRepartitionExec::try_new_with_routing(child, outputs, SeriesRouting::RawPrimaryKey) + .expect("the raw owner is built"), + ) +} + +/// Every chunk of `plan`'s one output, decoded, plus the error it ended with if any. +async fn drain_raw( + plan: &Arc, + context: &Arc, +) -> (Vec, Option) { + let frames = plan + .execute(0, Arc::clone(context)) + .expect("partition 0 exists"); + let mut decoded = unpack_raw_series(frames, Some(raw_schema())).expect("the frame decodes"); + let mut chunks = Vec::new(); + while let Some(next) = decoded.next().await { + match next { + Ok(chunk) => chunks.push(chunk), + Err(error) => return (chunks, Some(error)), + } + } + (chunks, None) +} + /// Every output partition of `plan`, decoded, and the error it ended with if it ended with one. /// /// All partitions are read at once, as in normal execution. Draining them concurrently makes this @@ -545,6 +777,14 @@ async fn drain( } /// The output partition of `plan` at `partition`, as the frame stream a consumer reads. +/// The evaluated view of one owner publication, for tests that drive the internal stream. +fn evaluated(output: ReadyOutput) -> (Series, bool) { + match output { + ReadyOutput::Evaluated(series) => (series, true), + ReadyOutput::Raw(_) => panic!("an evaluated output carried a raw chunk"), + } +} + fn output( plan: &Arc, partition: usize, @@ -1957,6 +2197,7 @@ async fn eof_readies_two_groups_but_poll_merges_one_at_a_time() { Arc::clone(&schema()), &(Arc::clone(&pool) as Arc), 1, + SeriesOrdering::Contiguous, ); wait_until("both identities and EOF are queued", || { producer_stopped(&shared) @@ -1967,8 +2208,9 @@ async fn eof_readies_two_groups_but_poll_merges_one_at_a_time() { .await .expect("first ready identity") .expect("first identity succeeds"); - assert_eq!(first_metric, first.0.metric); - assert!(first.1); + let (first_series, first_closed) = evaluated(first); + assert_eq!(first_metric, first_series.metric); + assert!(first_closed); assert!( stream.output.is_empty(), "no eagerly merged second result is buffered" @@ -1978,11 +2220,12 @@ async fn eof_readies_two_groups_but_poll_merges_one_at_a_time() { .await .expect("second ready identity") .expect("second identity succeeds"); - assert_eq!(second_metric, second.0.metric); - assert!(second.1); + let (second_series, second_closed) = evaluated(second); + assert_eq!(second_metric, second_series.metric); + assert!(second_closed); assert!(stream.groups.is_empty()); - drop(first); - drop(second); + drop(first_series); + drop(second_series); assert!( futures::future::poll_fn(|cx| Pin::new(&mut stream).poll_next(cx)) .await @@ -3877,7 +4120,7 @@ async fn a_task_attached_after_the_failure_is_aborted_and_not_counted() { ); let task = tokio::spawn(route_partition( - Box::pin(futures::stream::pending()), + RoutedInput::Evaluated(Box::pin(futures::stream::pending())), 1, Routing::new(1), Vec::new(), @@ -3959,12 +4202,12 @@ async fn a_recorded_failure_reaches_every_output_once() { .flat_map(|(zero, one)| [zero, one].map(|metric| Step::chunk(chunk(metric, 0, 1, true)))) .collect::>(); steps.push(Step::ErrorSaying(MESSAGE)); + // The producer is gated so it cannot reach its error before this test owns both queues. A + // failure drains every queue nobody has claimed yet, so a producer that ran first would release + // the very charges this test reads, and the retention assertion below could never hold. + let switch = Arc::new(Switch::default()); let plan = node( - Fragile::plan( - vec![Behaviour::OneAtATime(steps)], - Arc::new(Switch::default()), - ) - .plan, + Fragile::plan(vec![Behaviour::Gated(steps)], Arc::clone(&switch)).plan, 2, ); let pool = Arc::new(GreedyMemoryPool::new(usize::MAX)); @@ -3976,27 +4219,35 @@ async fn a_recorded_failure_reaches_every_output_once() { let mut streams = Vec::new(); for partition in 0..2 { let (receiver, lease) = SharedExecution::claim(&shared, partition).expect("output claim"); - streams.push(Box::pin(RepartitionStream::new( - receiver, - lease, - Arc::clone(&schema()), - &(Arc::clone(&pool) as Arc), - 2, - )) as SendableEvaluatedSeriesStream); + streams.push(Box::pin(EvaluatedOutputStream { + inner: Some(Box::pin(RepartitionStream::new( + receiver, + lease, + Arc::clone(&schema()), + &(Arc::clone(&pool) as Arc), + 2, + SeriesOrdering::Contiguous, + ))), + }) as SendableEvaluatedSeriesStream); } let execution = shared; + // Both queues are claimed, so the gates below are the only thing between the queued prefix and + // the failure: what is retained now is retained until this test drains it. + assert!( + !switch.is_open(), + "the producer is held before its error, so nothing drained the queues yet" + ); + wait_until("queued output chains retain memory before drain", || { + pool.reserved() > 0 + }) + .await; + switch.open(); wait_until( "the read records its failure while no output is drained", || matches!(execution.ending(), Ending::Failed(_)), ) .await; - // The producer records the failure before its queued chains are all visible, so wait for the - // retention this assertion is about instead of reading the pool once. - wait_until("queued output chains retain memory before drain", || { - pool.reserved() > 0 - }) - .await; let outputs = tokio::time::timeout( Duration::from_secs(10), futures::future::join_all(streams.iter_mut().map(until_end)), @@ -4303,13 +4554,16 @@ async fn dropping_unpolled_route_future_ends_outputs_and_cancels_peer() { let (receiver, lease) = SharedExecution::claim(&shared, partition).expect("output claim"); let frames = pack_evaluated_series( - Box::pin(RepartitionStream::new( - receiver, - lease, - Arc::clone(&schema()), - &(Arc::clone(&pool) as Arc), - 2, - )), + Box::pin(EvaluatedOutputStream { + inner: Some(Box::pin(RepartitionStream::new( + receiver, + lease, + Arc::clone(&schema()), + &(Arc::clone(&pool) as Arc), + 2, + SeriesOrdering::Contiguous, + ))), + }), Arc::clone(&series_schema), ) .expect("output frame builds"); @@ -4326,7 +4580,7 @@ async fn dropping_unpolled_route_future_ends_outputs_and_cancels_peer() { Poll::Pending }); let peer = route_partition( - Box::pin(peer_stream), + RoutedInput::Evaluated(Box::pin(peer_stream)), 1, Routing::new(2), senders.clone(), @@ -4351,7 +4605,7 @@ async fn dropping_unpolled_route_future_ends_outputs_and_cancels_peer() { reported: false, }; let unpolled = route_partition( - Box::pin(futures::stream::pending()), + RoutedInput::Evaluated(Box::pin(futures::stream::pending())), 0, Routing::new(2), senders.clone(), @@ -4410,13 +4664,16 @@ async fn unpolled_route_drop_preserves_an_already_elected_failure() { for partition in 0..2 { let (receiver, lease) = SharedExecution::claim(&shared, partition).expect("output claim"); let frames = pack_evaluated_series( - Box::pin(RepartitionStream::new( - receiver, - lease, - Arc::clone(&schema()), - &(Arc::clone(&pool) as Arc), - 2, - )), + Box::pin(EvaluatedOutputStream { + inner: Some(Box::pin(RepartitionStream::new( + receiver, + lease, + Arc::clone(&schema()), + &(Arc::clone(&pool) as Arc), + 2, + SeriesOrdering::Contiguous, + ))), + }), Arc::clone(&schema()), ) .expect("output frame builds"); @@ -4424,7 +4681,7 @@ async fn unpolled_route_drop_preserves_an_already_elected_failure() { } inputs.failed(1, &Arc::new(DataFusionError::Execution(FIRST.to_string()))); let unpolled = route_partition( - Box::pin(futures::stream::pending()), + RoutedInput::Evaluated(Box::pin(futures::stream::pending())), 0, Routing::new(2), senders.clone(), @@ -4491,7 +4748,7 @@ async fn ready_chain_yields_after_a_charged_prefix_and_drop_releases_it() { let (sender, receiver) = mpsc::unbounded_channel(); let routing = Routing::new(1); let senders = [sender]; - let mut read = Box::pin(read_chains( + let mut read = Box::pin(read_evaluated_chains( &mut source, 0, &routing, @@ -4567,15 +4824,18 @@ async fn sibling_failure_aborts_budget_yielded_route_and_releases_prefix() { ending: Arc::clone(&ending), }); let (receiver, lease) = SharedExecution::claim(&shared, 0).expect("claim output"); - let mut output: SendableEvaluatedSeriesStream = Box::pin(RepartitionStream::new( - receiver, - lease, - Arc::clone(&schema()), - &(Arc::clone(&pool) as Arc), - 1, - )); + let mut output: SendableEvaluatedSeriesStream = Box::pin(EvaluatedOutputStream { + inner: Some(Box::pin(RepartitionStream::new( + receiver, + lease, + Arc::clone(&schema()), + &(Arc::clone(&pool) as Arc), + 1, + SeriesOrdering::Contiguous, + ))), + }); let future = route_partition( - source, + RoutedInput::Evaluated(source), 0, Routing::new(1), vec![sender], @@ -4893,7 +5153,8 @@ async fn discarded_open_chain_eof_is_still_an_error() { Box::pin(futures::stream::iter(vec![Ok((series, false))])); let pool: Arc = Arc::new(GreedyMemoryPool::new(1 << 20)); let reservation = MemoryConsumer::new("discarded-open-eof-test").register(&pool); - let result = read_chains(&mut input, 0, &Routing::new(1), &[sender], &reservation).await; + let result = + read_evaluated_chains(&mut input, 0, &Routing::new(1), &[sender], &reservation).await; assert!( matches!(result, Stop::Failed(_)), "open discarded chain EOF is not clean" @@ -5162,6 +5423,7 @@ async fn internal_stream_releases_each_chunk_charge_only_when_polled() { Arc::clone(&schema()), &(Arc::clone(&pool) as Arc), 1, + SeriesOrdering::Contiguous, ); wait_until("internal producer finishes", || producer_stopped(&shared)).await; // Queue, EOF, and identity-group metadata are charged alongside the retained chunks; the @@ -5181,10 +5443,11 @@ async fn internal_stream_releases_each_chunk_charge_only_when_polled() { .await .unwrap() .unwrap(); - assert!(merged.1, "the logical owner closes once after global EOF"); + let (merged_series, closed) = evaluated(merged); + assert!(closed, "the logical owner closes once after global EOF"); assert_eq!( 4, - merged.0.time.len(), + merged_series.time.len(), "both slices retain their two positions in the merged logical owner" ); assert!( @@ -5197,7 +5460,7 @@ async fn internal_stream_releases_each_chunk_charge_only_when_polled() { .is_none(), "the logical owner has one stable end" ); - drop(merged); + drop(merged_series); drop(stream); assert_eq!( 0, @@ -5457,6 +5720,7 @@ async fn histogram_chain_slices_fit_a_budget_below_the_full_table_charge() { Arc::clone(&schema()), &(Arc::clone(&pool) as Arc), 1, + SeriesOrdering::Contiguous, ); wait_until("histogram producer queues the whole chain", || { producer_stopped(&shared) @@ -5483,12 +5747,9 @@ async fn histogram_chain_slices_fit_a_budget_below_the_full_table_charge() { chunks.len(), "one owner series spans all physical slices" ); - assert!(chunks[0].1, "the complete owner closes once"); - assert_eq!( - 256, - chunks[0].0.time.len(), - "all histogram rows survive merge" - ); + let (series, closed) = evaluated(chunks.remove(0)); + assert!(closed, "the complete owner closes once"); + assert_eq!(256, series.time.len(), "all histogram rows survive merge"); assert!(matches!(shared.ending(), Ending::Finished)); assert!( pool.reserved() >= base_bytes + independent_slice_lower_bound, @@ -5763,3 +6024,751 @@ async fn a_chain_that_outgrows_the_pool_ends_every_output_and_releases_it() { }) .await; } + +#[tokio::test] +async fn dropping_a_raw_open_chain_cancels_and_releases_its_reservation() { + let switch = Arc::new(Switch::default()); + let child = RawScripted::plan(vec![RawPartition { + chunks: vec![ + raw_chunk("k", Some(raw_metric("a")), &[1], vec![Some(1.0)], false), + raw_chunk("k", Some(raw_metric("a")), &[2], vec![Some(2.0)], false), + ], + gate: Some(Arc::clone(&switch)), + // The switch stays closed while this partial chunk is delivered and charged; the producer + // then parks at EOF until the destination is dropped. + gate_chunks: false, + }]); + let plan = raw_node(child, 1); + let pool = Arc::new(GreedyMemoryPool::new(8 << 20)); + let context = task_context(&pool); + let frames = plan.execute(0, Arc::clone(&context)).unwrap(); + let stream = unpack_raw_series(frames, Some(raw_schema())).unwrap(); + wait_until("raw open chain is reserved", || pool.reserved() > 0).await; + drop(stream); + tokio::time::timeout(Duration::from_secs(10), async { + while !producer_stopped(&execution(&plan)) { + tokio::task::yield_now().await; + } + }) + .await + .expect("dropping the raw stream cancels the open-chain read"); + assert_eq!(0, pool.reserved()); +} + +#[tokio::test] +async fn dropping_a_raw_destination_releases_its_partial_chain() { + let switch = Arc::new(Switch::default()); + let child = RawScripted::plan(vec![RawPartition { + chunks: vec![raw_chunk( + "k", + Some(raw_metric("a")), + &[1], + vec![Some(1.0)], + false, + )], + gate: Some(Arc::clone(&switch)), + // The switch stays closed while this partial chunk is delivered and charged; the producer + // then parks at EOF until the destination is dropped. + gate_chunks: false, + }]); + let plan = raw_node(child, 1); + let pool = Arc::new(GreedyMemoryPool::new(8 << 20)); + let context = task_context(&pool); + let frames = plan.execute(0, Arc::clone(&context)).unwrap(); + let stream = unpack_raw_series(frames, Some(raw_schema())).unwrap(); + wait_until("raw partial chain is reserved", || pool.reserved() > 0).await; + drop(stream); + tokio::time::timeout(Duration::from_secs(10), async { + while !producer_stopped(&execution(&plan)) { + tokio::task::yield_now().await; + } + }) + .await + .expect("dropping the destination stops its raw producer"); + assert_eq!(0, pool.reserved()); +} + +#[tokio::test] +async fn raw_chain_pool_refusal_releases_partial_reservation() { + let first = raw_chunk("k", Some(raw_metric("a")), &[1], vec![Some(1.0)], false); + let second = raw_chunk("k", Some(raw_metric("a")), &[2], vec![Some(2.0)], false); + // Measure what the first chunk reserves in this exact path, then leave no room for a second. + let probe_pool = Arc::new(GreedyMemoryPool::new(8 << 20)); + let probe_switch = Arc::new(Switch::default()); + let probe_plan = raw_node( + RawScripted::plan(vec![RawPartition { + chunks: vec![first.clone()], + gate: Some(Arc::clone(&probe_switch)), + gate_chunks: false, + }]), + 1, + ); + let probe_context = task_context(&probe_pool); + let probe_frames = probe_plan.execute(0, Arc::clone(&probe_context)).unwrap(); + let probe_stream = unpack_raw_series(probe_frames, Some(raw_schema())).unwrap(); + wait_until("probe raw chunk is reserved", || probe_pool.reserved() > 0).await; + let budget = probe_pool.reserved(); + drop(probe_stream); + probe_switch.open(); + let pool = Arc::new(GreedyMemoryPool::new(budget)); + let child = RawScripted::plan(vec![RawPartition { + chunks: vec![first, second], + gate: None, + gate_chunks: false, + }]); + let plan = raw_node(child, 1); + let context = task_context(&pool); + let (_, error) = tokio::time::timeout(Duration::from_secs(10), drain_raw(&plan, &context)) + .await + .expect("raw pool refusal terminates the execution"); + assert!( + error.is_some(), + "second chunk must be refused after first reservation" + ); + assert_eq!(0, pool.reserved()); +} + +#[tokio::test] +async fn raw_matching_routing_groups_by_match_key() { + use crate::vector_match::MatchingLabels; + + let with_zone = |zone: &str| { + SeriesLabels::try_new( + Some(Arc::from("m")), + vec![ + Label { + name: Arc::from("host"), + value: Arc::from("h"), + }, + Label { + name: Arc::from("zone"), + value: Arc::from(zone), + }, + ], + ) + .expect("legal labels") + }; + let child = RawScripted::plan(vec![ + RawPartition { + chunks: vec![raw_chunk( + "ka", + Some(with_zone("a")), + &[1], + vec![Some(1.0)], + true, + )], + gate: None, + gate_chunks: false, + }, + RawPartition { + chunks: vec![raw_chunk( + "kb", + Some(with_zone("b")), + &[1], + vec![Some(2.0)], + true, + )], + gate: None, + gate_chunks: false, + }, + ]); + let plan = Arc::new( + raw_node_with_routing( + child, + 3, + SeriesRouting::Matching(MatchingLabels::On(vec![Arc::from("host")])), + ) + .expect("raw matching routing is accepted"), + ); + let context = task_context(&Arc::new(GreedyMemoryPool::new(8 << 20))); + let matching = Routing::with_policy( + 3, + SeriesRouting::Matching(MatchingLabels::On(vec![Arc::from("host")])), + ); + let primary = Routing::with_policy(3, SeriesRouting::RawPrimaryKey); + let expected_owner = matching.output_labels(&with_zone("a")); + // Both match keys have host=h, while the primary keys ka/kb route separately. Comparing + // to the positive match-key owner therefore fails if this accidentally routes by primary key. + assert_eq!(expected_owner, matching.output_labels(&with_zone("b"))); + assert_ne!(expected_owner, primary.output_key(b"ka")); + assert_ne!(expected_owner, primary.output_key(b"kb")); + let mut owner = None; + let mut identities = std::collections::HashSet::new(); + for output in 0..plan.outputs() { + let frames = plan + .execute(output, Arc::clone(&context)) + .expect("output exists"); + let mut stream = unpack_raw_series(frames, Some(raw_schema())).expect("frame decodes"); + while let Some(chunk) = stream.next().await { + let chunk = chunk.expect("matching-routed raw stream succeeds"); + identities.insert(chunk.primary_key.as_slice().to_vec()); + if !chunk.is_last { + continue; + } + if [b"ka".to_vec(), b"kb".to_vec()].contains(&chunk.primary_key.as_slice().to_vec()) { + assert_eq!( + output, expected_owner, + "identity must use its match-key owner" + ); + assert!( + owner + .replace(output) + .is_none_or(|previous| previous == output), + "both match-key identities share an output partition" + ); + } + } + } + assert_eq!( + identities, + [b"ka".to_vec(), b"kb".to_vec()].into_iter().collect() + ); + assert!( + owner.is_some(), + "both identities were published by their common owner" + ); +} + +#[test] +fn raw_canonical_labels_routing_is_refused() { + let child = RawScripted::plan(vec![RawPartition { + chunks: vec![], + gate: None, + gate_chunks: false, + }]); + let error = raw_node_with_routing(child, 1, SeriesRouting::CanonicalLabels) + .expect_err("raw canonical labels routing must remain unsupported"); + assert!(error.to_string().contains("cannot use CanonicalLabels")); +} + +#[tokio::test] +async fn raw_identity_publishes_once_every_input_contributed() { + let switch = Arc::new(Switch::default()); + let child = RawScripted::plan(vec![ + RawPartition { + chunks: vec![raw_chunk( + "k", + Some(raw_metric("a")), + &[1, 3], + vec![Some(1.0), Some(3.0)], + true, + )], + gate: None, + gate_chunks: false, + }, + RawPartition { + chunks: vec![raw_chunk( + "k", + Some(raw_metric("a")), + &[2, 4], + vec![Some(2.0), Some(4.0)], + true, + )], + gate: Some(Arc::clone(&switch)), + gate_chunks: false, + }, + ]); + let plan = raw_node(child, 1); + let context = task_context(&Arc::new(GreedyMemoryPool::new(8 << 20))); + let frames = plan.execute(0, context).expect("partition 0 exists"); + let mut stream = unpack_raw_series(frames, Some(raw_schema())).expect("the frame decodes"); + let mut chunks = Vec::new(); + while let Some(item) = tokio::time::timeout(Duration::from_secs(2), stream.next()) + .await + .expect("published while EOF gate remains closed") + { + chunks.push(item.expect("raw merge succeeds")); + if chunks.last().is_some_and(|chunk| chunk.is_last) { + break; + } + } + let timestamps: Vec = chunks + .iter() + .flat_map(|chunk| chunk.time.as_buffer().iter().copied()) + .collect(); + assert_eq!(timestamps, [1, 2, 3, 4]); + assert!(chunks.last().is_some_and(|chunk| chunk.is_last)); + switch.open(); +} + +#[tokio::test] +async fn raw_identity_waits_for_an_input_that_has_not_contributed() { + let switch = Arc::new(Switch::default()); + let child = RawScripted::plan(vec![ + RawPartition { + chunks: vec![raw_chunk( + "k", + Some(raw_metric("a")), + &[1, 3], + vec![Some(1.0), Some(3.0)], + true, + )], + gate: None, + gate_chunks: false, + }, + RawPartition { + chunks: vec![raw_chunk( + "k", + Some(raw_metric("a")), + &[2, 4], + vec![Some(2.0), Some(4.0)], + true, + )], + gate: Some(Arc::clone(&switch)), + gate_chunks: true, + }, + ]); + let plan = raw_node(child, 1); + let context = task_context(&Arc::new(GreedyMemoryPool::new(8 << 20))); + let frames = plan.execute(0, context).expect("partition 0 exists"); + let mut stream = unpack_raw_series(frames, Some(raw_schema())).expect("the frame decodes"); + assert!( + tokio::time::timeout(Duration::from_millis(200), stream.next()) + .await + .is_err() + ); + switch.open(); + let chunks: Vec<_> = stream.try_collect().await.expect("raw merge succeeds"); + let timestamps: Vec = chunks + .iter() + .flat_map(|chunk| chunk.time.as_buffer().iter().copied()) + .collect(); + assert_eq!(timestamps, [1, 2, 3, 4]); + assert!(chunks.last().is_some_and(|chunk| chunk.is_last)); +} + +/// Chunks decoded from one frame batch share its Arrow buffers, so charging each chunk the whole +/// allocation would charge the batch's timestamps once per chunk. +/// +/// A second input that stays open and contributes nothing holds every group unpublished, which is +/// what keeps the charges of the first input's chunks alive to be measured. +#[tokio::test] +async fn raw_chunks_of_one_batch_are_charged_their_own_slice() { + const SERIES: usize = 64; + const ROWS: usize = 64; + let switch = Arc::new(Switch::default()); + let child = RawScripted::plan(vec![ + RawPartition { + chunks: (0..SERIES) + .map(|series| { + let times: Vec = (0..ROWS as i64) + .map(|row| series as i64 * 1000 + row) + .collect(); + let values: Vec> = (0..ROWS).map(|row| Some(row as f64)).collect(); + raw_chunk( + &format!("k{series}"), + Some(raw_metric("a")), + ×, + values, + true, + ) + }) + .collect(), + gate: None, + gate_chunks: false, + }, + RawPartition { + chunks: Vec::new(), + gate: Some(Arc::clone(&switch)), + gate_chunks: false, + }, + ]); + let plan = raw_node(child, 1); + let pool = Arc::new(GreedyMemoryPool::new(64 << 20)); + let context = task_context(&pool); + let mut frames = plan + .execute(0, Arc::clone(&context)) + .expect("partition 0 exists"); + // Everything the first input has is delivered; the second input has contributed nothing, so no + // group is publishable and every chunk charge of the first input is still held. + // + // Wait for the first charge before sampling: under load the input task may not have run yet, + // and a fixed number of idle ticks would then measure nothing. + wait_until("the first chunk is charged", || pool.reserved() > 0).await; + // Keep the largest charge seen while the rest of the partition is routed, so a charge that + // grows per chunk is measured rather than missed. + let mut peak = 0usize; + let mut idle = 0usize; + while idle < 100 { + let state = futures::future::poll_fn(|cx| Poll::Ready(frames.as_mut().poll_next(cx))).await; + peak = peak.max(pool.reserved()); + match state { + Poll::Ready(Some(Ok(_))) => idle = 0, + Poll::Ready(Some(Err(error))) => panic!("the raw read failed: {error}"), + Poll::Ready(None) => panic!("the read ended before the gated input was opened"), + Poll::Pending => idle += 1, + } + tokio::time::sleep(Duration::from_millis(1)).await; + } + switch.open(); + drop(frames); + // The timestamps of every chunk in the batch come from one buffer of at most SERIES * ROWS + // values. Charging that buffer per chunk would retain about SERIES times its size. + let batch_timestamps = SERIES * ROWS * std::mem::size_of::(); + assert!( + peak < SERIES * batch_timestamps / 2, + "peak {peak} must stay well below the {SERIES}x over-charge of {batch_timestamps} bytes" + ); +} + +/// A histogram field is an index array naming rows of a shared table. A query over few distinct +/// histograms holds mostly indices, and those are charged alongside the table rather than instead +/// of it. +#[tokio::test] +async fn raw_histogram_index_arrays_are_charged() { + const ROWS: usize = 4096; + let side = BucketSideView { + span_offsets: &[], + span_lengths: &[], + buckets: &[], + }; + let mut builder = HistogramBatchBuilder::new(); + builder + .push(HistogramView { + schema: 1, + zero_threshold: 0.5, + sum: 1.0, + count: 1.0, + zero_count: 0.0, + reset_hint: CounterResetHint::Unknown, + start_timestamp: None, + custom_values: &[], + positive: side, + negative: side, + }) + .expect("one histogram row"); + // Every position names the same row, so the table stays small while the index array grows with + // the number of samples. + let samples = Samples::try_from_histogram( + UInt32Array::from(vec![Some(0u32); ROWS]), + builder.finish().expect("the table is built"), + ) + .expect("one row is referenced by every position"); + let chunk = RawSeriesChunk::try_new( + &raw_schema(), + Buffer::from_slice_ref(b"k"), + Some(raw_metric("a")), + RawTimestamps::try_new(ScalarBuffer::from((0..ROWS as i64).collect::>())) + .expect("ascending timestamps"), + SeriesFields::try_from_vec(vec![samples]).expect("one sample per timestamp"), + true, + ) + .expect("one series' samples are a legal raw chunk"); + let switch = Arc::new(Switch::default()); + let child = RawScripted::plan(vec![ + RawPartition { + chunks: vec![chunk], + gate: None, + gate_chunks: false, + }, + // An input that has contributed nothing keeps the group unpublished, which is what holds + // the first input's charges alive to be measured. + RawPartition { + chunks: Vec::new(), + gate: Some(Arc::clone(&switch)), + gate_chunks: false, + }, + ]); + let plan = raw_node(child, 1); + let pool = Arc::new(GreedyMemoryPool::new(64 << 20)); + let context = task_context(&pool); + let frames = plan + .execute(0, Arc::clone(&context)) + .expect("partition 0 exists"); + // The chunk is charged as the input task routes it, which happens without this stream being + // polled; wait for the charge instead of sampling it for a fixed number of ticks. + wait_until("the routed chunk is charged", || pool.reserved() > 0).await; + let peak = pool.reserved(); + switch.open(); + drop(frames); + let timestamps = ROWS * std::mem::size_of::(); + let indices = ROWS * std::mem::size_of::(); + assert!( + peak >= timestamps + indices, + "peak {peak} must cover the {timestamps} timestamp bytes and the {indices} index bytes" + ); +} + +/// A key-then-time frame promises its keys arrive in order, so an identity whose inputs have all +/// contributed is published without waiting for any of them to end. +#[tokio::test] +async fn an_ordered_frame_publishes_a_key_its_inputs_have_passed() { + let switch = Arc::new(Switch::default()); + let child = RawScripted::plan_with_ordering( + vec![ + RawPartition { + chunks: vec![ + raw_chunk("a", Some(raw_metric("a")), &[1], vec![Some(1.0)], true), + raw_chunk("b", Some(raw_metric("b")), &[1], vec![Some(2.0)], true), + ], + gate: None, + gate_chunks: false, + }, + // Contributes `a` and then stays open: `a` is complete, and nothing has been seen past + // `b`, so `b` waits. + RawPartition { + chunks: vec![raw_chunk( + "a", + Some(raw_metric("a")), + &[3], + vec![Some(3.0)], + true, + )], + gate: Some(Arc::clone(&switch)), + gate_chunks: false, + }, + ], + SeriesOrdering::KeyThenTime, + ); + let plan = raw_node(child, 1); + let pool = Arc::new(GreedyMemoryPool::new(8 << 20)); + let context = task_context(&pool); + let frames = plan + .execute(0, Arc::clone(&context)) + .expect("partition exists"); + let mut stream = unpack_raw_series(frames, Some(raw_schema())).expect("frame decodes"); + // A frame may split one identity's rows across chunks, so read until `a` closes. + let mut rows = std::collections::BTreeMap::, Vec>::new(); + let mut closed = std::collections::HashSet::new(); + while !closed.contains(b"a".as_slice()) { + let chunk = tokio::time::timeout(Duration::from_secs(2), stream.next()) + .await + .expect("`a` publishes without EOF") + .expect("`a` arrives") + .expect("valid frame"); + rows.entry(chunk.primary_key.to_vec()) + .or_default() + .extend(chunk.time.as_buffer().iter().copied()); + if chunk.is_last { + closed.insert(chunk.primary_key.to_vec()); + } + } + assert_eq!( + vec![1, 3], + rows[b"a".as_slice()], + "both inputs' samples are merged and `a` closed without any input ending" + ); + assert!( + !rows.contains_key(b"b".as_slice()), + "`b` waits: no input has been seen past it and none has ended" + ); + switch.open(); + while stream.next().await.is_some() {} +} + +/// A watermark lets `a` publish after another input has passed it, without that input ending. +#[tokio::test] +async fn an_ordered_watermark_lets_an_identity_publish_without_input_eof() { + let switch = Arc::new(Switch::default()); + let child = RawScripted::plan_with_ordering( + vec![ + RawPartition { + chunks: vec![raw_chunk( + "a", + Some(raw_metric("a")), + &[1], + vec![Some(1.0)], + true, + )], + gate: None, + gate_chunks: false, + }, + // Delivers `b`, which proves this input passed `a`, then stays open so neither input + // ends before `a` is published. + RawPartition { + chunks: vec![raw_chunk( + "b", + Some(raw_metric("b")), + &[2], + vec![Some(2.0)], + true, + )], + gate: Some(Arc::clone(&switch)), + gate_chunks: false, + }, + ], + SeriesOrdering::KeyThenTime, + ); + let plan = raw_node(child, 1); + let pool = Arc::new(GreedyMemoryPool::new(8 << 20)); + let context = task_context(&pool); + let frames = plan + .execute(0, Arc::clone(&context)) + .expect("partition exists"); + let mut stream = unpack_raw_series(frames, Some(raw_schema())).expect("frame decodes"); + let first = tokio::time::timeout(Duration::from_secs(2), stream.next()) + .await + .expect("`a` publishes although no input ended") + .expect("`a` arrives") + .expect("valid frame"); + assert_eq!(b"a", first.primary_key.as_slice()); + switch.open(); + while stream.next().await.is_some() {} +} + +/// Without a watermark from the silent input, `a` is not safe to publish until that input ends. +#[tokio::test] +async fn an_ordered_identity_waits_for_an_input_that_has_not_passed_the_key() { + let switch = Arc::new(Switch::default()); + let child = RawScripted::plan_with_ordering( + vec![ + RawPartition { + chunks: vec![raw_chunk( + "a", + Some(raw_metric("a")), + &[1], + vec![Some(1.0)], + true, + )], + gate: None, + gate_chunks: false, + }, + RawPartition { + chunks: Vec::new(), + gate: Some(Arc::clone(&switch)), + gate_chunks: true, + }, + ], + SeriesOrdering::KeyThenTime, + ); + let plan = raw_node(child, 1); + let context = task_context(&Arc::new(GreedyMemoryPool::new(8 << 20))); + let frames = plan + .execute(0, Arc::clone(&context)) + .expect("partition exists"); + let mut stream = unpack_raw_series(frames, Some(raw_schema())).expect("frame decodes"); + // The bounded wait itself is the assertion: a publication here would mean the silent input + // was incorrectly treated as having passed `a` (or ended). + assert!( + tokio::time::timeout(Duration::from_millis(100), stream.next()) + .await + .is_err(), + "`a` stays held while another input is silent and open" + ); + switch.open(); + let chunk = tokio::time::timeout(Duration::from_secs(2), stream.next()) + .await + .expect("the opened input lets the output proceed") + .expect("`a` publishes") + .expect("valid frame"); + assert_eq!(b"a", chunk.primary_key.as_slice()); + assert!(chunk.is_last); +} + +#[tokio::test] +async fn raw_identity_groups_by_primary_key_not_labels() { + let child = RawScripted::plan(vec![ + RawPartition { + chunks: vec![raw_chunk( + "k1", + Some(raw_metric("a")), + &[1], + vec![Some(1.0)], + true, + )], + gate: None, + gate_chunks: false, + }, + RawPartition { + chunks: vec![raw_chunk( + "k2", + Some(raw_metric("a")), + &[2], + vec![Some(2.0)], + true, + )], + gate: None, + gate_chunks: false, + }, + ]); + let plan = raw_node(child, 1); + let pool = Arc::new(GreedyMemoryPool::new(8 << 20)); + let context = task_context(&pool); + let (chunks, error) = drain_raw(&plan, &context).await; + assert!( + error.is_none(), + "different keys with identical labels merge independently" + ); + let keys: std::collections::HashSet> = chunks + .iter() + .map(|chunk| chunk.primary_key.as_slice().to_vec()) + .collect(); + assert_eq!(chunks.len(), 2, "one chunk is emitted per primary key"); + assert_eq!(keys, [b"k1".to_vec(), b"k2".to_vec()].into_iter().collect()); +} + +#[tokio::test] +async fn raw_pieces_of_one_key_must_agree_on_labels() { + let child = RawScripted::plan(vec![ + RawPartition { + chunks: vec![raw_chunk( + "k", + Some(raw_metric("a")), + &[1], + vec![Some(1.0)], + true, + )], + gate: None, + gate_chunks: false, + }, + RawPartition { + chunks: vec![raw_chunk( + "k", + Some(raw_metric("b")), + &[2], + vec![Some(2.0)], + true, + )], + gate: None, + gate_chunks: false, + }, + ]); + let plan = raw_node(child, 1); + let pool = Arc::new(GreedyMemoryPool::new(8 << 20)); + let context = task_context(&pool); + let (_, error) = drain_raw(&plan, &context).await; + let error = error.expect("label disagreement is surfaced"); + assert!( + error + .to_string() + .contains("disagree on their explicit labels") + ); +} + +#[tokio::test] +async fn a_raw_merge_conflict_fails_the_execution() { + let child = RawScripted::plan(vec![ + RawPartition { + chunks: vec![raw_chunk( + "k", + Some(raw_metric("a")), + &[1], + vec![Some(1.0)], + true, + )], + gate: None, + gate_chunks: false, + }, + RawPartition { + chunks: vec![raw_chunk( + "k", + Some(raw_metric("a")), + &[1], + vec![Some(2.0)], + true, + )], + gate: None, + gate_chunks: false, + }, + ]); + let plan = raw_node(child, 1); + let pool = Arc::new(GreedyMemoryPool::new(8 << 20)); + let context = task_context(&pool); + let (_, error) = drain_raw(&plan, &context).await; + let error = error.expect("conflicting duplicate samples fail the execution"); + assert!( + error + .to_string() + .contains("duplicate raw timestamp has conflicting samples") + ); +} diff --git a/src/promql/src/source.rs b/src/promql/src/source.rs index 7d882c7251d..8ad9c086964 100644 --- a/src/promql/src/source.rs +++ b/src/promql/src/source.rs @@ -12,9 +12,12 @@ // See the License for the specific language governing permissions and // limitations under the License. +mod merge; #[cfg(test)] mod tests; +pub(crate) use merge::merge_raw_identity_charging; + use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; diff --git a/src/promql/src/source/merge.rs b/src/promql/src/source/merge.rs new file mode 100644 index 00000000000..cb0fd0d587a --- /dev/null +++ b/src/promql/src/source/merge.rs @@ -0,0 +1,252 @@ +// 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 datafusion::arrow::buffer::ScalarBuffer; +use datafusion_common::DataFusionError; +use datatypes::error::InvalidSeriesSnafu; +use datatypes::series::histogram::HistogramView; +use datatypes::series::raw::{RawSeriesChunk, RawTimestamps}; +use datatypes::series::samples::{SampleRef, Samples, SeriesFields, SeriesSchema}; + +/// Rebuilds one raw series from every piece of one complete physical identity, reporting each +/// allocation before it happens. +/// +/// Reconstruction holds three intermediate vectors whose sizes depend on how the pieces interleave, +/// so what it needs cannot be known up front. The caller charges `reserve` as the work grows +/// instead of reserving a worst case it may never reach. +pub(crate) fn merge_raw_identity_charging( + schema: &SeriesSchema, + chunks: &[RawSeriesChunk], + reserve: &mut dyn FnMut(usize) -> std::result::Result<(), DataFusionError>, + shrink: &mut dyn FnMut(usize), +) -> std::result::Result { + let invalid = |reason: &str| { + DataFusionError::External(Box::new( + InvalidSeriesSnafu { + reason: reason.to_string(), + } + .build(), + )) + }; + let Some(first) = chunks.first() else { + return Err(invalid("cannot merge an empty raw identity")); + }; + if schema.eval_grid.is_some() { + return Err(invalid("raw series schema cannot have an evaluation grid")); + } + let key = first.primary_key.clone(); + let labels = first + .metric + .as_ref() + .ok_or_else(|| invalid("raw identity chunk has no explicit labels"))? + .clone(); + for chunk in chunks { + if chunk.primary_key != key { + return Err(invalid("raw identity chunks have different primary keys")); + } + if chunk.metric.as_ref() != Some(&labels) { + return Err(invalid( + "raw identity chunks have different or missing labels", + )); + } + if chunk.fields.len() != schema.fields.len() + || chunk + .fields + .iter() + .any(|field| field.len() != chunk.time.len()) + { + return Err(invalid("raw identity chunk shape does not match schema")); + } + if chunk + .time + .as_buffer() + .windows(2) + .any(|pair| pair[0] >= pair[1]) + { + return Err(invalid( + "raw identity chunk timestamps are not strictly increasing", + )); + } + } + + let mut row_count = 0usize; + for chunk in chunks.iter() { + row_count = row_count.saturating_add(chunk.time.len()); + } + let chunks_are_ordered_and_disjoint = chunks.windows(2).all(|pair| { + pair[0] + .time + .as_buffer() + .last() + .zip(pair[1].time.as_buffer().first()) + .is_some_and(|(last, first)| last < first) + }); + let mut transient_workspace_bytes = 0usize; + let mut output_fields = Vec::with_capacity(schema.fields.len()); + let times = if chunks_are_ordered_and_disjoint { + reserve(row_count.saturating_mul(std::mem::size_of::()))?; + let mut times = Vec::with_capacity(row_count); + for chunk in chunks { + times.extend_from_slice(chunk.time.as_buffer()); + } + for field_index in 0..schema.fields.len() { + let mut parts = Vec::with_capacity(chunks.len()); + for chunk in chunks { + let source = chunk + .fields + .get(field_index) + .ok_or_else(|| invalid("field index out of bounds"))?; + reserve(std::mem::size_of::())?; + parts.push( + source + .slice(0, chunk.time.len()) + .map_err(|error| DataFusionError::External(Box::new(error)))?, + ); + } + output_fields.push( + Samples::try_from_chunks(parts) + .map_err(|error| DataFusionError::External(Box::new(error)))?, + ); + } + times + } else { + transient_workspace_bytes = row_count + .saturating_mul(std::mem::size_of::<(i64, usize, usize)>()) + .saturating_mul(2); + reserve(row_count.saturating_mul(std::mem::size_of::<(i64, usize, usize)>()))?; + let mut rows = Vec::with_capacity(row_count); + for (chunk_index, chunk) in chunks.iter().enumerate() { + for row_index in 0..chunk.time.len() { + rows.push((chunk.time.as_buffer()[row_index], chunk_index, row_index)); + } + } + rows.sort_unstable(); + + reserve( + rows.len() + .saturating_mul(std::mem::size_of::<(i64, usize, usize)>()), + )?; + let mut selected = Vec::with_capacity(rows.len()); + let mut index = 0; + while index < rows.len() { + let mut end = index + 1; + while end < rows.len() && rows[end].0 == rows[index].0 { + end += 1; + } + for duplicate in &rows[index + 1..end] { + for field_index in 0..schema.fields.len() { + if !samples_equal( + chunks[rows[index].1] + .fields + .get(field_index) + .ok_or_else(|| invalid("field index out of bounds"))? + .sample(rows[index].2), + chunks[duplicate.1] + .fields + .get(field_index) + .ok_or_else(|| invalid("field index out of bounds"))? + .sample(duplicate.2), + ) { + return Err(invalid("duplicate raw timestamp has conflicting samples")); + } + } + } + selected.push(rows[index]); + index = end; + } + + reserve(selected.len().saturating_mul(std::mem::size_of::()))?; + let mut times = Vec::with_capacity(selected.len()); + for row in &selected { + times.push(row.0); + } + for field_index in 0..schema.fields.len() { + let mut parts = Vec::new(); + let mut start = 0; + while start < selected.len() { + let (chunk_index, row_index) = (selected[start].1, selected[start].2); + let mut end = start + 1; + while end < selected.len() + && selected[end].1 == chunk_index + && selected[end].2 == selected[end - 1].2 + 1 + { + end += 1; + } + let source = chunks[chunk_index] + .fields + .get(field_index) + .ok_or_else(|| invalid("field index out of bounds"))?; + reserve(std::mem::size_of::())?; + parts.push( + source + .slice(row_index, end - start) + .map_err(|error| DataFusionError::External(Box::new(error)))?, + ); + start = end; + } + output_fields.push( + Samples::try_from_chunks(parts) + .map_err(|error| DataFusionError::External(Box::new(error)))?, + ); + } + times + }; + let result = RawSeriesChunk::try_new( + schema, + key, + Some(labels), + RawTimestamps::try_new(ScalarBuffer::from(times)) + .map_err(|error| DataFusionError::External(Box::new(error)))?, + SeriesFields::try_from_vec(output_fields) + .map_err(|error| DataFusionError::External(Box::new(error)))?, + true, + ) + .map_err(|error| DataFusionError::External(Box::new(error)))?; + // Only the general merge path allocates transient row-index workspaces. + shrink(transient_workspace_bytes); + Ok(result) +} + +fn samples_equal(left: SampleRef<'_>, right: SampleRef<'_>) -> bool { + match (left, right) { + (SampleRef::Missing, SampleRef::Missing) => true, + (SampleRef::Float(a), SampleRef::Float(b)) => a.to_bits() == b.to_bits(), + (SampleRef::Histogram(a), SampleRef::Histogram(b)) => histograms_equal(a, b), + _ => false, + } +} + +fn histograms_equal(a: HistogramView<'_>, b: HistogramView<'_>) -> bool { + a.schema == b.schema + && a.zero_threshold.to_bits() == b.zero_threshold.to_bits() + && a.sum.to_bits() == b.sum.to_bits() + && a.count.to_bits() == b.count.to_bits() + && a.zero_count.to_bits() == b.zero_count.to_bits() + && a.reset_hint == b.reset_hint + && a.start_timestamp == b.start_timestamp + && float_slices_equal(a.custom_values, b.custom_values) + && a.positive.span_offsets == b.positive.span_offsets + && a.positive.span_lengths == b.positive.span_lengths + && float_slices_equal(a.positive.buckets, b.positive.buckets) + && a.negative.span_offsets == b.negative.span_offsets + && a.negative.span_lengths == b.negative.span_lengths + && float_slices_equal(a.negative.buckets, b.negative.buckets) +} + +fn float_slices_equal(a: &[f64], b: &[f64]) -> bool { + a.len() == b.len() && a.iter().zip(b).all(|(a, b)| a.to_bits() == b.to_bits()) +} diff --git a/src/promql/src/source/merge/tests.rs b/src/promql/src/source/merge/tests.rs new file mode 100644 index 00000000000..1bbc927cbf0 --- /dev/null +++ b/src/promql/src/source/merge/tests.rs @@ -0,0 +1,702 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; + +use datafusion::arrow::array::{Float64Array, UInt32Array}; +use datafusion::arrow::buffer::Buffer; +use datatypes::series::histogram::{BucketSideView, CounterResetHint, HistogramBatchBuilder}; +use datatypes::series::raw::{RawSeriesChunk, RawTimestamps}; +use datatypes::series::samples::{Samples, SeriesFields, SeriesSchema}; +use datatypes::series::{Label, SeriesLabels}; + +use super::merge_raw_identity_charging; + +/// The same merge, with accounting the tests do not measure. +fn merge_raw_identity( + schema: &SeriesSchema, + chunks: &[RawSeriesChunk], +) -> datafusion::common::Result { + merge_raw_identity_charging(schema, chunks, &mut |_| Ok(()), &mut |_| {}) +} + +fn schema(count: usize) -> SeriesSchema { + SeriesSchema { + fields: (0..count) + .map(|i| datatypes::series::samples::FieldSpec { + name: Arc::from(format!("f{i}")), + }) + .collect::>() + .into(), + eval_grid: None, + } +} + +fn histogram(sum: f64, threshold: f64) -> Samples { + let mut builder = HistogramBatchBuilder::new(); + let side = BucketSideView { + span_offsets: &[], + span_lengths: &[], + buckets: &[], + }; + let index = builder + .push(datatypes::series::histogram::HistogramView { + schema: 1, + zero_threshold: threshold, + sum, + count: 1.0, + zero_count: 0.0, + reset_hint: CounterResetHint::Unknown, + start_timestamp: None, + custom_values: &[], + positive: side, + negative: side, + }) + .unwrap(); + Samples::try_from_histogram( + UInt32Array::from(vec![Some(index)]), + builder.finish().unwrap(), + ) + .unwrap() +} + +fn labels(value: &str) -> SeriesLabels { + SeriesLabels::try_new( + Some(Arc::from("metric")), + vec![Label { + name: Arc::from("tag"), + value: Arc::from(value), + }], + ) + .unwrap() +} + +fn chunk( + key: &[u8], + label: &str, + times: Vec, + fields: Vec>>, +) -> RawSeriesChunk { + RawSeriesChunk::try_new( + &schema(fields.len()), + Buffer::from(key.to_vec()), + Some(labels(label)), + RawTimestamps::try_new(times.into()).unwrap(), + SeriesFields::try_from_vec( + fields + .into_iter() + .map(|values| Samples::from_float(Float64Array::from(values))) + .collect(), + ) + .unwrap(), + false, + ) + .unwrap() +} + +#[test] +fn interleaved_pieces_merge_in_time_order() { + let schema = schema(2); + let result = merge_raw_identity( + &schema, + &[ + chunk( + b"k", + "v", + vec![1, 3], + vec![vec![Some(10.), Some(30.)], vec![Some(1.), Some(3.)]], + ), + chunk( + b"k", + "v", + vec![2, 4], + vec![vec![Some(20.), Some(40.)], vec![Some(2.), Some(4.)]], + ), + ], + ) + .unwrap(); + assert_eq!(result.time.as_buffer().as_ref(), &[1, 2, 3, 4]); + assert_eq!( + result + .fields + .get(0) + .unwrap() + .iter() + .map(|s| match s { + datatypes::series::samples::SampleRef::Float(v) => v, + _ => 0., + }) + .collect::>(), + vec![10., 20., 30., 40.] + ); + assert!(result.is_last); +} + +#[test] +fn ordered_disjoint_pieces_concatenate_without_row_workspace() { + let schema = schema(1); + let chunks = [ + chunk(b"k", "v", vec![1, 2], vec![vec![Some(10.), Some(20.)]]), + chunk(b"k", "v", vec![3, 4], vec![vec![Some(30.), Some(40.)]]), + ]; + let result = merge_raw_identity(&schema, &chunks).unwrap(); + assert_eq!(result.time.as_buffer().as_ref(), &[1, 2, 3, 4]); + assert_eq!(result.time.len(), 4); + assert_eq!( + result + .fields + .get(0) + .unwrap() + .iter() + .map(|sample| match sample { + datatypes::series::samples::SampleRef::Float(value) => value, + _ => 0., + }) + .collect::>(), + vec![10., 20., 30., 40.] + ); + + let mut reserved = 0usize; + let mut shrunk = 0usize; + merge_raw_identity_charging( + &schema, + &chunks, + &mut |bytes| { + reserved += bytes; + Ok(()) + }, + &mut |bytes| shrunk += bytes, + ) + .unwrap(); + assert_eq!( + reserved, + 4 * std::mem::size_of::() + 2 * std::mem::size_of::() + ); + assert_eq!(shrunk, 0); + assert!(shrunk <= reserved); +} + +#[test] +fn touching_ordered_chunks_use_general_duplicate_path() { + let schema = schema(1); + let chunks = [ + chunk(b"k", "v", vec![1, 2], vec![vec![Some(10.), Some(20.)]]), + chunk(b"k", "v", vec![2, 3], vec![vec![Some(20.), Some(30.)]]), + ]; + let mut reserved = 0usize; + let mut shrunk = 0usize; + let result = merge_raw_identity_charging( + &schema, + &chunks, + &mut |bytes| { + reserved += bytes; + Ok(()) + }, + &mut |bytes| shrunk += bytes, + ) + .unwrap(); + assert_eq!(result.time.as_buffer().as_ref(), &[1, 2, 3]); + assert_eq!(shrunk, 2 * 4 * std::mem::size_of::<(i64, usize, usize)>()); + assert!(shrunk <= reserved); +} + +#[test] +fn equal_duplicates_collapse_but_bit_differences_conflict() { + assert!( + merge_raw_identity( + &schema(1), + &[ + chunk(b"k", "v", vec![1], vec![vec![Some(1.)]]), + chunk(b"k", "v", vec![1], vec![vec![Some(1.)]]) + ] + ) + .is_ok() + ); + for (left, right) in [ + (-0.0, 0.0), + ( + f64::from_bits(0x7ff8_0000_0000_0001), + f64::from_bits(0x7ff8_0000_0000_0002), + ), + ] { + assert!( + merge_raw_identity( + &schema(1), + &[ + chunk(b"k", "v", vec![1], vec![vec![Some(left)]]), + chunk(b"k", "v", vec![1], vec![vec![Some(right)]]) + ] + ) + .is_err() + ); + } + assert!( + merge_raw_identity( + &schema(1), + &[ + chunk(b"k", "v", vec![1], vec![vec![None]]), + chunk(b"k", "v", vec![1], vec![vec![Some(1.)]]) + ] + ) + .is_err() + ); +} + +#[test] +fn rejects_different_keys_labels_and_internal_duplicates() { + assert!( + merge_raw_identity( + &schema(1), + &[ + chunk(b"a", "v", vec![1], vec![vec![Some(1.)]]), + chunk(b"b", "v", vec![2], vec![vec![Some(2.)]]) + ] + ) + .is_err() + ); + assert!( + merge_raw_identity( + &schema(1), + &[ + chunk(b"a", "v", vec![1], vec![vec![Some(1.)]]), + chunk(b"a", "other", vec![2], vec![vec![Some(2.)]]) + ] + ) + .is_err() + ); + assert!(RawTimestamps::try_new(vec![1, 1].into()).is_err()); + let mut no_labels = chunk(b"a", "v", vec![1], vec![vec![Some(1.)]]); + no_labels.metric = None; + assert!(merge_raw_identity(&schema(1), &[no_labels]).is_err()); +} + +#[test] +fn histogram_duplicates_compare_every_float_property() { + let schema = schema(1); + let make = |sample| RawSeriesChunk { + primary_key: Buffer::from(b"k".to_vec()), + metric: Some(labels("v")), + time: RawTimestamps::try_new(vec![1].into()).unwrap(), + fields: SeriesFields::one(sample), + is_last: false, + }; + assert!( + merge_raw_identity( + &schema, + &[make(histogram(2.0, 0.5)), make(histogram(2.0, 0.5))] + ) + .is_ok() + ); + assert!( + merge_raw_identity( + &schema, + &[make(histogram(2.0, 0.5)), make(histogram(3.0, 0.5))] + ) + .is_err() + ); + assert!( + merge_raw_identity( + &schema, + &[make(histogram(2.0, 0.5)), make(histogram(2.0, -0.0))] + ) + .is_err() + ); +} + +/// Builds a histogram with configurable comparison properties for duplicate-conflict tests. +fn histogram_with(view: datatypes::series::histogram::HistogramView<'_>) -> Samples { + let mut builder = HistogramBatchBuilder::new(); + let index = builder.push(view).unwrap(); + Samples::try_from_histogram( + UInt32Array::from(vec![Some(index)]), + builder.finish().unwrap(), + ) + .unwrap() +} + +/// Two histograms are duplicates only when every property matches. +/// +/// Bucket values are independently free within a valid layout and are compared directly. The +/// span/cursor layout cannot be varied independently of bucket count: each layout must have +/// nonnegative lengths summing to the bucket count and a valid cursor. The cases below vary each +/// offset and length while preserving valid layouts and equal bucket counts. +#[test] +fn histogram_duplicates_compare_every_property() { + let schema = schema(1); + let make = |sample| RawSeriesChunk { + primary_key: Buffer::from(b"k".to_vec()), + metric: Some(labels("v")), + time: RawTimestamps::try_new(vec![1].into()).unwrap(), + fields: SeriesFields::one(sample), + is_last: false, + }; + let positive = BucketSideView { + span_offsets: &[0, 0], + span_lengths: &[1, 1], + buckets: &[1.0, 2.0], + }; + let negative = BucketSideView { + span_offsets: &[0, 0], + span_lengths: &[1, 1], + buckets: &[3.0, 4.0], + }; + let build = |schema, + count, + zero_count, + sum, + zero_threshold, + positive, + negative, + reset_hint, + start_timestamp, + custom_values| { + histogram_with(datatypes::series::histogram::HistogramView { + schema, + count, + zero_count, + sum, + zero_threshold, + positive, + negative, + reset_hint, + start_timestamp, + custom_values, + }) + }; + let baseline = || { + build( + 1, + 1.0, + 0.0, + 2.0, + 0.5, + positive, + negative, + CounterResetHint::Unknown, + None, + &[], + ) + }; + assert!( + merge_raw_identity(&schema, &[make(baseline()), make(baseline())]).is_ok(), + "identical histograms are duplicates" + ); + + let positive_buckets_changed = BucketSideView { + buckets: &[1.0, 2.5], + ..positive + }; + let negative_buckets_changed = BucketSideView { + buckets: &[3.0, 4.5], + ..negative + }; + // Same span count and lengths as the baseline, different cursors: only the offsets differ. + let positive_offsets_changed = BucketSideView { + span_offsets: &[1, 0], + ..positive + }; + let negative_offsets_changed = BucketSideView { + span_offsets: &[1, 0], + ..negative + }; + // Same cursors and same total bucket count (2), different split into spans. A zero-length + // span is a valid layout, so only the lengths differ here. + let positive_lengths_changed = BucketSideView { + span_lengths: &[2, 0], + ..positive + }; + let negative_lengths_changed = BucketSideView { + span_lengths: &[2, 0], + ..negative + }; + let cases: Vec<(&str, Samples)> = vec![ + ( + "schema (1 vs 2)", + build( + 2, + 1.0, + 0.0, + 2.0, + 0.5, + positive, + negative, + CounterResetHint::Unknown, + None, + &[], + ), + ), + ( + "count (1.0 vs 2.0)", + build( + 1, + 2.0, + 0.0, + 2.0, + 0.5, + positive, + negative, + CounterResetHint::Unknown, + None, + &[], + ), + ), + ( + "zero_count (0.0 vs 1.0)", + build( + 1, + 1.0, + 1.0, + 2.0, + 0.5, + positive, + negative, + CounterResetHint::Unknown, + None, + &[], + ), + ), + ( + "sum (2.0 vs 3.0)", + build( + 1, + 1.0, + 0.0, + 3.0, + 0.5, + positive, + negative, + CounterResetHint::Unknown, + None, + &[], + ), + ), + ( + "zero_threshold (0.5 vs -0.0)", + build( + 1, + 1.0, + 0.0, + 2.0, + -0.0, + positive, + negative, + CounterResetHint::Unknown, + None, + &[], + ), + ), + ( + "positive bucket value (2.0 vs 2.5)", + build( + 1, + 1.0, + 0.0, + 2.0, + 0.5, + positive_buckets_changed, + negative, + CounterResetHint::Unknown, + None, + &[], + ), + ), + ( + "negative bucket value (4.0 vs 4.5)", + build( + 1, + 1.0, + 0.0, + 2.0, + 0.5, + positive, + negative_buckets_changed, + CounterResetHint::Unknown, + None, + &[], + ), + ), + ( + "positive span offset ([0, 0] vs [1, 0])", + build( + 1, + 1.0, + 0.0, + 2.0, + 0.5, + positive_offsets_changed, + negative, + CounterResetHint::Unknown, + None, + &[], + ), + ), + ( + "negative span offset ([0, 0] vs [1, 0])", + build( + 1, + 1.0, + 0.0, + 2.0, + 0.5, + positive, + negative_offsets_changed, + CounterResetHint::Unknown, + None, + &[], + ), + ), + ( + "positive span lengths ([1, 1] vs [2, 0])", + build( + 1, + 1.0, + 0.0, + 2.0, + 0.5, + positive_lengths_changed, + negative, + CounterResetHint::Unknown, + None, + &[], + ), + ), + ( + "negative span lengths ([1, 1] vs [2, 0])", + build( + 1, + 1.0, + 0.0, + 2.0, + 0.5, + positive, + negative_lengths_changed, + CounterResetHint::Unknown, + None, + &[], + ), + ), + ( + "reset hint (Unknown vs CounterReset)", + build( + 1, + 1.0, + 0.0, + 2.0, + 0.5, + positive, + negative, + CounterResetHint::CounterReset, + None, + &[], + ), + ), + ( + "start timestamp (None vs Some(7))", + build( + 1, + 1.0, + 0.0, + 2.0, + 0.5, + positive, + negative, + CounterResetHint::Unknown, + Some(7), + &[], + ), + ), + ( + "custom values ([] vs [1.5])", + build( + 1, + 1.0, + 0.0, + 2.0, + 0.5, + positive, + negative, + CounterResetHint::Unknown, + None, + &[1.5], + ), + ), + ]; + for (what, changed) in cases { + assert!( + merge_raw_identity(&schema, &[make(baseline()), make(changed)]).is_err(), + "two histograms differing only in {what} are a conflict, not a duplicate" + ); + } +} + +#[test] +fn rejects_empty_and_shape_mismatch() { + assert!(merge_raw_identity(&schema(1), &[]).is_err()); + let valid = chunk(b"k", "v", vec![1], vec![vec![Some(1.)]]); + assert!(merge_raw_identity(&schema(2), &[valid]).is_err()); + let malformed = RawSeriesChunk { + primary_key: Buffer::from(b"k".to_vec()), + metric: Some(labels("v")), + time: RawTimestamps::try_new(vec![1].into()).unwrap(), + fields: SeriesFields::one(Samples::from_float(Float64Array::from(vec![ + Some(1.), + Some(2.), + ]))), + is_last: false, + }; + assert!(merge_raw_identity(&schema(1), &[malformed]).is_err()); +} + +#[test] +fn reservation_error_is_preserved() { + let error = merge_raw_identity_charging( + &schema(1), + &[chunk(b"k", "v", vec![1], vec![vec![Some(1.)]])], + &mut |_| { + Err(datafusion::common::DataFusionError::Execution( + "reservation refused".into(), + )) + }, + &mut |_| {}, + ) + .unwrap_err(); + assert!(matches!( + error, + datafusion::common::DataFusionError::Execution(_) + )); +} + +#[test] +fn merge_releases_transient_row_workspace() { + let mut reserved = 0usize; + let mut shrunk = 0usize; + let result = merge_raw_identity_charging( + &schema(1), + &[ + chunk(b"k", "v", vec![1, 3], vec![vec![Some(1.), Some(3.)]]), + chunk(b"k", "v", vec![1, 2], vec![vec![Some(1.), Some(2.)]]), + ], + &mut |bytes| { + reserved = reserved.saturating_add(bytes); + Ok(()) + }, + &mut |bytes| shrunk = shrunk.saturating_add(bytes), + ) + .unwrap(); + assert_eq!(result.time.len(), 3); + assert_eq!(shrunk, 2 * 4 * std::mem::size_of::<(i64, usize, usize)>()); + let retained = reserved - shrunk; + assert!(retained < reserved); + assert!(retained >= result.time.len() * std::mem::size_of::()); +} diff --git a/src/promql/src/vector_match.rs b/src/promql/src/vector_match.rs index f68cc50f454..2196893e3a9 100644 --- a/src/promql/src/vector_match.rs +++ b/src/promql/src/vector_match.rs @@ -1528,9 +1528,9 @@ fn same_grid(lhs: &SeriesSchema, rhs: &SeriesSchema) -> Result> { /// Label pairs identifying a match group or a result; empty values are absent, as in /// Prometheus. -type GroupKey = Vec<(Arc, Arc)>; +pub(crate) type GroupKey = Vec<(Arc, Arc)>; -fn group_key(metric: &SeriesLabels, labels: &MatchingLabels) -> GroupKey { +pub(crate) fn group_key(metric: &SeriesLabels, labels: &MatchingLabels) -> GroupKey { let mut key: GroupKey = match labels { MatchingLabels::On(names) => names .iter() diff --git a/tests-integration/tests/native_prefix.rs b/tests-integration/tests/native_prefix.rs index 9b7685e793c..4cd3b9246db 100644 --- a/tests-integration/tests/native_prefix.rs +++ b/tests-integration/tests/native_prefix.rs @@ -55,13 +55,15 @@ fn evaluated_frame_schema(tree: &NativeTree) -> datafusion::common::DFSchemaRef ) } -// Reproduction from the review: 16 hosts × 20k samples, `rate[10s]`, and 16 DN -// partitions produces four batches per chain and fails while unpacking because -// `CoalescePartitionsExec` interleaves their batches, so a chain is cut and the frames no -// longer decode. This test asserts that failure, which is why it runs rather than being -// ignored: it keeps the boundary honest until the frontend region merge becomes -// chain-atomic, and it is the test to change when that fix lands. The parallelism-one -// path is exercised by `native_prefix_executes_over_real_mito_region`. +// Sends `rate[10s]` on a one-point evaluation grid for 16 hosts over 16 region-scan +// partitions to the datanode's remote-read entry point, then checks that the evaluated frames +// decode and each host reports the rate its samples produce. The former decode failure no longer +// reproduces because #10023 makes region package completion close its pending final raw chunk, +// so a chain's value and close are no longer latched apart in separate batches by the +// batch-interleaving `CoalescePartitionsExec`. That coalescing is still not chain-atomic: a chain +// spanning batches can still be cut. This one-point grid gives each host one evaluated value plus +// a possible empty closing chunk, so it does not exercise a chain spanning many batches. The +// parallelism-one path is exercised by `native_prefix_executes_over_real_mito_region`. #[tokio::test(flavor = "multi_thread")] async fn native_prefix_multi_partition_series_chains() { let instance = GreptimeDbStandaloneBuilder::new("native_prefix_partitioned_chains") @@ -205,20 +207,48 @@ async fn native_prefix_multi_partition_series_chains() { } let schema = batches.first().expect("evaluated frames").schema.clone(); let batches = RecordBatches::try_new(schema, batches).unwrap(); - let error = common_query::series_frame::unpack_evaluated_series( + // Package completion closes each region's pending final raw chunk, so the evaluated frames + // decode; ordinary partition coalescing still does not keep chains atomic. Check that every + // host's chain closes exactly once and carries the rate its own samples produce. + let decoded = common_query::series_frame::unpack_evaluated_series( Box::pin(DfRecordBatchStreamAdapter::new(batches.as_stream())), None, ) .unwrap() .try_collect::>() .await - .unwrap_err(); - assert!( - error - .to_string() - .contains("evaluated series identity changed before its chain closed"), - "{error}" - ); + .expect("a chain-atomic merge decodes every chain"); + + let mut closed = std::collections::BTreeMap::)>::new(); + for (series, is_last) in decoded { + let host = series + .metric + .get("host") + .unwrap_or_else(|| panic!("a decoded series keeps its host label: {:?}", series.metric)) + .to_string(); + let entry = closed.entry(host).or_default(); + if is_last { + entry.0 += 1; + } + for field in series.fields.iter() { + entry + .1 + .extend(field.iter().filter_map(|sample| match sample { + datatypes::series::samples::SampleRef::Float(value) => Some(value), + datatypes::series::samples::SampleRef::Histogram(_) + | datatypes::series::samples::SampleRef::Missing => None, + })); + } + } + assert_eq!(16, closed.len(), "one chain per host: {closed:?}"); + for (host, (closes, values)) in &closed { + assert_eq!(1, *closes, "`{host}` closes its chain exactly once"); + assert_eq!( + vec![1.0], + *values, + "`{host}` reports the rate of its own samples, which a cut chain would lose" + ); + } } #[tokio::test(flavor = "multi_thread")]