diff --git a/src/promql/src/extension_plan/instant_manipulate.rs b/src/promql/src/extension_plan/instant_manipulate.rs index 0210b7f6558..85b08f5ca41 100644 --- a/src/promql/src/extension_plan/instant_manipulate.rs +++ b/src/promql/src/extension_plan/instant_manipulate.rs @@ -12,6 +12,9 @@ // See the License for the specific language governing permissions and // limitations under the License. +mod native; +mod native_exec; + use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; diff --git a/src/promql/src/extension_plan/instant_manipulate/native.rs b/src/promql/src/extension_plan/instant_manipulate/native.rs new file mode 100644 index 00000000000..f7df2a5feeb --- /dev/null +++ b/src/promql/src/extension_plan/instant_manipulate/native.rs @@ -0,0 +1,302 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; + +use datafusion::arrow::array::{ArrayRef, Float64Array, StringArray, TimestampMillisecondArray}; +use datafusion::arrow::datatypes::{DataType, SchemaRef, TimeUnit}; +use datafusion::arrow::record_batch::RecordBatch; +use datafusion::error::{DataFusionError, Result as DataFusionResult}; +use datatypes::series::MetricName; +use datatypes::series::raw::RawTimestamps; +use datatypes::series::samples::{SampleRef, Samples, SeriesFields, SeriesSchema}; + +#[cfg(test)] +#[path = "native_tests.rs"] +mod tests; + +use crate::extension_plan::instant_manipulate::InstantManipulate; +use crate::instant_eval::{MAX_INSTANT_MANIPULATE_OUTPUT_POINTS, select_instants}; + +#[derive(Clone)] +enum Column { + Time, + Field(usize), + Tag(String), +} + +/// Preflighted mapping from a complete typed series to the original query schema. +pub(crate) struct NativeInstantMapping { + schema: SchemaRef, + series_schema: Arc, + columns: Vec, + fields: Vec, + start: i64, + end: i64, + step: i64, + lookback: i64, +} + +impl NativeInstantMapping { + /// Returns `None` when the node/schema is outside the bounded native subset. + pub(crate) fn try_new( + node: &InstantManipulate, + schema: SchemaRef, + series_schema: Arc, + logical_tags: &[String], + ) -> DataFusionResult> { + if node.offset != 0 || node.unfix.is_some() || node.start >= node.end { + return Ok(None); + } + if node.interval <= 0 { + return Err(DataFusionError::Execution( + "InstantManipulate interval must be positive".into(), + )); + } + if node.lookback_delta < 0 { + return Err(DataFusionError::Execution( + "InstantManipulate lookback must be non-negative".into(), + )); + } + if series_schema.eval_grid.is_some() { + return Ok(None); + } + + let mut names = std::collections::HashSet::new(); + for field in series_schema.fields.iter() { + if !names.insert(field.name.as_ref()) { + return Ok(None); + } + } + if logical_tags.iter().any(|name| !names.insert(name.as_str())) { + return Ok(None); + } + + let stale_field_names = node.staleness_field_columns().collect::>(); + if stale_field_names.iter().any(|name| { + series_schema + .fields + .iter() + .all(|field| field.name.as_ref() != *name) + }) { + return Ok(None); + } + + let mut columns = Vec::with_capacity(schema.fields().len()); + let mut seen_fields = vec![false; series_schema.fields.len()]; + let mut seen_tags = vec![false; logical_tags.len()]; + let mut seen_time = false; + for field in schema.fields() { + let name = field.name(); + let column = if name == node.time_index_column.as_str() { + if seen_time + || field.data_type() != &DataType::Timestamp(TimeUnit::Millisecond, None) + { + return Ok(None); + } + seen_time = true; + Column::Time + } else if let Some(index) = series_schema + .fields + .iter() + .position(|v| v.name.as_ref() == name) + { + if seen_fields[index] || field.data_type() != &DataType::Float64 { + return Ok(None); + } + seen_fields[index] = true; + Column::Field(index) + } else if let Some(index) = logical_tags.iter().position(|tag| tag == name) { + if seen_tags[index] || !is_string_type(field.data_type()) { + return Ok(None); + } + seen_tags[index] = true; + Column::Tag(name.clone()) + } else { + return Ok(None); + }; + columns.push(column); + } + if !seen_time || seen_fields.iter().any(|seen| !seen) || seen_tags.iter().any(|seen| !seen) + { + return Ok(None); + } + + let fields = stale_field_names + .iter() + .filter_map(|name| { + series_schema + .fields + .iter() + .position(|field| field.name.as_ref() == *name) + }) + .collect(); + Ok(Some(Self { + schema, + series_schema, + columns, + fields, + start: node.start, + end: node.end, + step: node.interval, + lookback: node.lookback_delta, + })) + } + + /// The original output schema retained by this mapping. + pub(crate) fn schema(&self) -> SchemaRef { + self.schema.clone() + } + + /// Estimates the clipped evaluation grid for complete source timestamp bounds. + pub(crate) fn estimated_points(&self, first: i64, last: i64) -> DataFusionResult { + let first = first as i128; + let last = last as i128; + let useful_last = if self.lookback == 0 { + last + } else { + last + self.lookback as i128 - 1 + }; + let query_start = self.start as i128; + let query_end = self.end as i128; + let max_start = first.max(query_start); + let min_end = useful_last.min(query_end); + let (aligned_start, aligned_end) = if max_start > min_end { + (1_i128, 0_i128) + } else { + ( + query_start + (max_start - query_start) / self.step as i128 * self.step as i128, + query_end - (query_end - min_end) / self.step as i128 * self.step as i128, + ) + }; + let points = if aligned_end >= aligned_start { + (aligned_end - aligned_start) / self.step as i128 + 1 + } else { + 0 + }; + if points > MAX_INSTANT_MANIPULATE_OUTPUT_POINTS as i128 { + return Err(DataFusionError::Execution(format!( + "InstantManipulate output points exceed limit: {points} > {MAX_INSTANT_MANIPULATE_OUTPUT_POINTS}" + ))); + } + usize::try_from(points).map_err(|_| { + DataFusionError::ResourcesExhausted("native instant point count overflow".into()) + }) + } + + /// Selects instants against typed samples, then maps only selected values to Arrow. + pub(crate) fn evaluate( + &self, + metric: &MetricName, + time: &RawTimestamps, + fields: &SeriesFields, + ) -> DataFusionResult { + if fields.len() != self.series_schema.fields.len() { + return Err(DataFusionError::Execution( + "native instant field count mismatch".into(), + )); + } + let checked_fields = self + .fields + .iter() + .map(|i| fields.as_slice()[*i].clone()) + .collect::>(); + let checked_fields = if checked_fields.is_empty() { + vec![Samples::from_float(Float64Array::new_null(time.len()))] + } else { + checked_fields + }; + let checked = SeriesFields::try_from_vec(checked_fields) + .map_err(|e| DataFusionError::Execution(e.to_string()))?; + let (indices, timestamps) = select_instants( + time, + &checked, + self.start..=self.end, + self.step, + self.lookback, + 0, + )?; + if indices.is_empty() { + return Ok(RecordBatch::new_empty(self.schema.clone())); + } + + let arrays: DataFusionResult> = self + .columns + .iter() + .enumerate() + .map(|(output_index, column)| match column { + Column::Time => { + Ok(Arc::new(TimestampMillisecondArray::from(timestamps.clone())) as ArrayRef) + } + Column::Field(index) => { + let source = fields.get(*index).ok_or_else(|| { + DataFusionError::Execution("native instant field missing".into()) + })?; + let mut samples = source.iter(); + let mut cursor = 0usize; + let mut selected_index = None; + let mut selected_value = None; + let mut values = Vec::with_capacity(indices.len()); + for index in &indices { + let index = *index as usize; + if selected_index != Some(index) { + while cursor <= index { + selected_value = Some( + match samples.next().ok_or_else(|| { + DataFusionError::Execution( + "native instant sample iterator ended before selected index" + .into(), + ) + })? { + SampleRef::Float(value) => Some(value), + SampleRef::Missing => None, + SampleRef::Histogram(_) => { + return Err(DataFusionError::Execution( + "native instant expected Float64 sample".into(), + )); + } + }, + ); + selected_index = Some(cursor); + cursor += 1; + } + } + values.push(selected_value.ok_or_else(|| { + DataFusionError::Execution( + "native instant selected sample was not initialized".into(), + ) + })?); + } + Ok(Arc::new(Float64Array::from(values)) as ArrayRef) + } + Column::Tag(name) => { + let values = std::iter::repeat_with(|| metric.get(name)) + .take(indices.len()) + .collect::(); + let target = self.schema.field(output_index).data_type(); + let cast = datafusion::arrow::compute::cast(&values, target)?; + Ok(cast) + } + }) + .collect(); + Ok(RecordBatch::try_new(self.schema.clone(), arrays?)?) + } +} + +fn is_string_type(data_type: &DataType) -> bool { + matches!( + data_type, + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View + ) || matches!(data_type, DataType::Dictionary(key, value) if key.as_ref() == &DataType::UInt32 && value.as_ref() == &DataType::Utf8) +} diff --git a/src/promql/src/extension_plan/instant_manipulate/native_exec.rs b/src/promql/src/extension_plan/instant_manipulate/native_exec.rs new file mode 100644 index 00000000000..383d341b259 --- /dev/null +++ b/src/promql/src/extension_plan/instant_manipulate/native_exec.rs @@ -0,0 +1,501 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use common_query::series_frame::{raw_series_frame_ordering, unpack_raw_series}; +use datafusion::arrow::array::Array; +use datafusion::arrow::datatypes::SchemaRef; +use datafusion::arrow::record_batch::RecordBatch; +use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::error::{DataFusionError, Result as DfResult}; +use datafusion::execution::context::TaskContext; +use datafusion::execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use datafusion::physical_expr::EquivalenceProperties; +use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, PhysicalExpr, PlanProperties, RecordBatchStream, + SendableRecordBatchStream, +}; +use datatypes::series::SeriesOrdering; +use datatypes::series::raw::{RawSeriesChunk, RawTimestamps}; +use datatypes::series::samples::{DenseSamplesRef, Samples, SeriesFields, SeriesSchema}; +use futures::{Stream, ready}; +use store_api::region_engine::{RawLabelSource, SendableRawSeriesStream}; + +use crate::extension_plan::instant_manipulate::InstantManipulate; +use crate::extension_plan::instant_manipulate::native::NativeInstantMapping; +use crate::source::materialize_raw_labels; + +#[cfg(test)] +#[path = "native_exec_tests.rs"] +mod tests; + +impl InstantManipulate { + /// Returns whether the bounded native mapping accepts this output schema, series schema, and + /// complete logical tag set. + /// + /// This is the local half of [`Self::try_to_native_execution_plan`]: the same decision that + /// method makes before it inspects its input, so a caller can decline a plan without building + /// one. `Err` is the same error that method propagates for an invalid interval or lookback. + pub fn is_native_mapping_supported( + &self, + output_schema: SchemaRef, + series_schema: Arc, + logical_tags: &[String], + ) -> DfResult { + Ok( + NativeInstantMapping::try_new(self, output_schema, series_schema, logical_tags)? + .is_some(), + ) + } + + /// Builds the bounded native execution path for a prepared raw-series frame. + pub fn try_to_native_execution_plan( + &self, + input: Arc, + output_schema: SchemaRef, + label_source: RawLabelSource, + series_schema: Arc, + logical_tags: &[String], + ) -> DfResult>> { + let Some(mapping) = NativeInstantMapping::try_new( + self, + output_schema, + series_schema.clone(), + logical_tags, + )? + else { + return Ok(None); + }; + if input.properties().partitioning.partition_count() != 1 { + return Ok(None); + } + let declared_schema = input.schema(); + let ordering = raw_series_frame_ordering(declared_schema.as_ref())?; + // Validate the declared raw-frame schema without executing the child. + unpack_raw_series( + Box::pin(RecordBatchStreamAdapter::new( + declared_schema.clone(), + futures::stream::empty(), + )), + series_schema.clone(), + label_source.clone(), + ordering, + )?; + Ok(Some(Arc::new(NativeInstantManipulateExec::new( + input, + Arc::new(mapping), + declared_schema, + series_schema, + label_source, + ordering, + )))) + } +} + +struct NativeInstantManipulateExec { + input: Arc, + mapping: Arc, + raw_schema: SchemaRef, + series_schema: Arc, + label_source: RawLabelSource, + ordering: SeriesOrdering, + properties: Arc, +} +impl NativeInstantManipulateExec { + fn new( + input: Arc, + mapping: Arc, + raw_schema: SchemaRef, + series_schema: Arc, + label_source: RawLabelSource, + ordering: SeriesOrdering, + ) -> Self { + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(mapping.schema()), + datafusion::physical_plan::Partitioning::UnknownPartitioning(1), + EmissionType::Incremental, + Boundedness::Bounded, + )); + Self { + input, + mapping, + raw_schema, + series_schema, + label_source, + ordering, + properties, + } + } +} +impl std::fmt::Debug for NativeInstantManipulateExec { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("NativeInstantManipulateExec").finish() + } +} +impl DisplayAs for NativeInstantManipulateExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!(f, "NativeInstantManipulateExec") + } +} +impl ExecutionPlan for NativeInstantManipulateExec { + fn name(&self) -> &str { + "NativeInstantManipulateExec" + } + fn schema(&self) -> SchemaRef { + self.mapping.schema() + } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> DfResult, + ) -> DfResult { + Ok(TreeNodeRecursion::Continue) + } + fn properties(&self) -> &Arc { + &self.properties + } + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> DfResult> { + if children.len() != 1 + || children[0].schema() != self.raw_schema + || children[0].properties().partitioning.partition_count() != 1 + { + return Err(DataFusionError::Plan( + "NativeInstantManipulateExec requires one unchanged one-partition raw child".into(), + )); + } + Ok(Arc::new(Self::new( + children[0].clone(), + self.mapping.clone(), + self.raw_schema.clone(), + self.series_schema.clone(), + self.label_source.clone(), + self.ordering, + ))) + } + fn execute( + &self, + partition: usize, + context: Arc, + ) -> DfResult { + if partition != 0 { + return Err(DataFusionError::Execution( + "native instant partition must be zero".into(), + )); + } + let frame = self.input.execute(partition, context.clone())?; + let source = unpack_raw_series( + frame, + self.series_schema.clone(), + self.label_source.clone(), + self.ordering, + )?; + let source = materialize_raw_labels(source) + .map_err(|e| DataFusionError::Execution(e.to_string()))?; + let reservation = + MemoryConsumer::new("NativeInstantStream").register(&context.runtime_env().memory_pool); + Ok(Box::pin(NativeInstantStream { + source: Some(source), + mapping: self.mapping.clone(), + schema: self.mapping.schema(), + series_schema: self.series_schema.clone(), + pending: Vec::new(), + reservation, + })) + } +} + +struct NativeInstantStream { + source: Option, + mapping: Arc, + schema: SchemaRef, + series_schema: Arc, + pending: Vec, + reservation: MemoryReservation, +} +impl NativeInstantStream { + fn fail(&mut self, error: DataFusionError) -> Poll>> { + self.source = None; + self.pending = Vec::new(); + self.reservation.free(); + Poll::Ready(Some(Err(error))) + } + fn finish_series(&mut self) -> DfResult { + let chunks = std::mem::take(&mut self.pending); + let rows = chunks + .iter() + .try_fold(0usize, |n, c| n.checked_add(c.time.len())) + .ok_or_else(|| { + DataFusionError::ResourcesExhausted("native instant row count overflow".into()) + })?; + let first = chunks + .first() + .and_then(|c| c.time.timestamp(0)) + .ok_or_else(|| { + DataFusionError::Execution("native instant series has no timestamps".into()) + })?; + let last = chunks + .last() + .and_then(|c| c.time.timestamp(c.time.len().saturating_sub(1))) + .ok_or_else(|| { + DataFusionError::Execution("native instant series has no timestamps".into()) + })?; + let points = self.mapping.estimated_points(first, last)?; + let label_bytes = chunks + .first() + .and_then(|c| c.metric.as_ref()) + .map(|m| { + m.name().map_or(0, str::len) + + m.tags() + .iter() + .map(|l| l.name.len() + l.value.len() + 16) + .sum::() + }) + .unwrap_or(0); + // Reserve timestamp concatenation by row count, but field descriptors by their + // actual chunk/segment counts: a dense chunk is still one field descriptor. + let field_count = self.series_schema.fields.len(); + let chunk_count = chunks.len(); + let segment_count = chunks + .iter() + .flat_map(|chunk| chunk.fields.as_slice()) + .try_fold(0usize, |count, field| { + count.checked_add(field.segments().len()) + }) + .ok_or_else(|| { + DataFusionError::ResourcesExhausted("native instant segment count overflow".into()) + })?; + let chunk_field_descriptors = chunk_count + .checked_mul(field_count) + .and_then(|n| n.checked_mul(std::mem::size_of::())) + .ok_or_else(|| { + DataFusionError::ResourcesExhausted( + "native instant descriptor estimate overflow".into(), + ) + })?; + // Flattening and converting to Arc can transiently retain the target descriptors and + // two source/flattened vectors; three segment-counts conservatively cover those owners. + let flattened_descriptors = segment_count + .checked_mul(3) + .and_then(|n| n.checked_mul(std::mem::size_of::())) + .ok_or_else(|| { + DataFusionError::ResourcesExhausted( + "native instant segment estimate overflow".into(), + ) + })?; + let scratch = rows + .checked_mul(16) + .and_then(|n| n.checked_add(chunk_field_descriptors)) + .and_then(|n| n.checked_add(flattened_descriptors)) + .and_then(|n| { + n.checked_add(field_count.checked_mul(std::mem::size_of::>())?) + }) + .and_then(|n| { + points + .checked_mul( + self.schema + .fields() + .len() + .checked_mul(24)? + .checked_add(label_bytes)? + .checked_add(48)?, + )? + .checked_add(n) + }) + .ok_or_else(|| { + DataFusionError::ResourcesExhausted( + "native instant scratch estimate overflow".into(), + ) + })?; + self.reservation + .try_grow(scratch) + .map_err(|e| DataFusionError::ResourcesExhausted(e.to_string()))?; + let metric = chunks + .first() + .and_then(|c| c.metric.clone()) + .ok_or_else(|| { + DataFusionError::Execution("native instant series has no metric labels".into()) + })?; + let mut times = Vec::new(); + times + .try_reserve_exact(rows) + .map_err(|e| DataFusionError::ResourcesExhausted(e.to_string()))?; + let mut field_chunks = (0..field_count).map(|_| Vec::new()).collect::>(); + for field in &mut field_chunks { + field + .try_reserve_exact(chunk_count) + .map_err(|e| DataFusionError::ResourcesExhausted(e.to_string()))?; + } + for chunk in chunks { + for i in 0..chunk.time.len() { + times.push(chunk.time.timestamp(i).ok_or_else(|| { + DataFusionError::Execution("native instant timestamp missing".into()) + })?); + } + for (target, field) in field_chunks.iter_mut().zip(chunk.fields.as_slice()) { + target.push(field.clone()); + } + } + let time = RawTimestamps::try_new(times.into()) + .map_err(|e| DataFusionError::Execution(e.to_string()))?; + let fields = field_chunks + .into_iter() + .map(Samples::try_from_chunks) + .collect::, _>>() + .map_err(|e| DataFusionError::Execution(e.to_string()))?; + let fields = SeriesFields::try_from_vec(fields) + .map_err(|e| DataFusionError::Execution(e.to_string()))?; + self.mapping.evaluate(&metric, &time, &fields) + } +} +fn retained_chunk_bytes(chunk: &RawSeriesChunk) -> Option { + let timestamp_bytes = chunk.time.as_buffer().inner().capacity(); + let labels = chunk + .metric + .as_ref() + .map(|metric| { + metric + .tags() + .iter() + .try_fold(metric.name().map_or(0, str::len), |total, label| { + total + .checked_add(label.name.len())? + .checked_add(label.value.len())? + .checked_add(std::mem::size_of::())? + .checked_add(2 * std::mem::size_of::()) + }) + }) + .unwrap_or(Some(0))?; + let field_descriptors = std::mem::size_of::() + .checked_add(std::mem::size_of::())? + .checked_add( + chunk + .fields + .as_slice() + .len() + .checked_mul(std::mem::size_of::())?, + )?; + let (segment_descriptors, sample_buffers) = chunk + .fields + .as_slice() + .iter() + .flat_map(|field| field.segments()) + .try_fold((0usize, 0usize), |(descriptors, buffers), segment| { + let bytes = match segment.dense()? { + DenseSamplesRef::Float(array) => array.get_buffer_memory_size(), + _ => return None, + }; + Some(( + descriptors.checked_add(std::mem::size_of::())?, + buffers.checked_add(bytes)?, + )) + })?; + timestamp_bytes + .checked_add(chunk.primary_key.capacity())? + .checked_add(labels)? + .checked_add(field_descriptors)? + .checked_add(segment_descriptors)? + .checked_add(sample_buffers) +} + +impl Stream for NativeInstantStream { + type Item = DfResult; + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + for _ in 0..64 { + let Some(source) = self.source.as_mut() else { + self.pending = Vec::new(); + self.reservation.free(); + return Poll::Ready(None); + }; + match ready!(source.as_mut().poll_next(cx)) { + Some(Err(error)) => { + return self.fail(DataFusionError::External(Box::new(error))); + } + None if !self.pending.is_empty() => { + return self.fail(DataFusionError::Execution( + "raw series ended without final chunk".into(), + )); + } + None => { + self.source = None; + self.pending = Vec::new(); + self.reservation.free(); + return Poll::Ready(None); + } + Some(Ok(chunk)) => { + let Some(bytes) = retained_chunk_bytes(&chunk) else { + return self.fail(DataFusionError::ResourcesExhausted( + "native instant retained-byte estimate overflow".into(), + )); + }; + if self.pending.len() == self.pending.capacity() { + let old_capacity = self.pending.capacity(); + let new_capacity = old_capacity.saturating_mul(2).max(4); + let slot_growth = + new_capacity.checked_sub(old_capacity).and_then(|slots| { + slots.checked_mul(std::mem::size_of::()) + }); + let Some(slot_growth) = slot_growth else { + return self.fail(DataFusionError::ResourcesExhausted( + "native instant pending capacity overflow".into(), + )); + }; + let Some(growth) = bytes.checked_add(slot_growth) else { + return self.fail(DataFusionError::ResourcesExhausted( + "native instant chunk estimate overflow".into(), + )); + }; + if let Err(e) = self.reservation.try_grow(growth) { + return self.fail(DataFusionError::ResourcesExhausted(e.to_string())); + } + let additional = new_capacity - self.pending.capacity(); + if let Err(e) = self.pending.try_reserve_exact(additional) { + return self.fail(DataFusionError::ResourcesExhausted(e.to_string())); + } + } else if let Err(e) = self.reservation.try_grow(bytes) { + return self.fail(DataFusionError::ResourcesExhausted(e.to_string())); + } + let last = chunk.is_last; + self.pending.push(chunk); + if last { + match self.finish_series() { + Ok(batch) => { + self.pending = Vec::new(); + self.reservation.free(); + return Poll::Ready(Some(Ok(batch))); + } + Err(e) => return self.fail(e), + } + } + } + } + } + cx.waker().wake_by_ref(); + Poll::Pending + } +} +impl RecordBatchStream for NativeInstantStream { + fn schema(&self) -> SchemaRef { + self.schema.clone() + } +} diff --git a/src/promql/src/extension_plan/instant_manipulate/native_exec_tests.rs b/src/promql/src/extension_plan/instant_manipulate/native_exec_tests.rs new file mode 100644 index 00000000000..4250fc44ddb --- /dev/null +++ b/src/promql/src/extension_plan/instant_manipulate/native_exec_tests.rs @@ -0,0 +1,427 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::VecDeque; +use std::pin::Pin; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::task::{Context, Poll}; + +use common_error::ext::{BoxedError, ErrorExt, PlainError}; +use common_error::status_code::StatusCode; +use datafusion::arrow::array::{Array, Float64Array, TimestampMillisecondArray}; +use datafusion::arrow::datatypes::{DataType, Field, Schema, TimeUnit}; +use datafusion::common::DFSchema; +use datafusion::error::DataFusionError; +use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryConsumer, MemoryPool}; +use datafusion::logical_expr::{EmptyRelation, LogicalPlan}; +use datatypes::arrow::buffer::ScalarBuffer; +use datatypes::series::raw::{RawSeriesChunk, RawTimestamps}; +use datatypes::series::samples::{FieldSpec, Samples, SeriesFields, SeriesSchema}; +use datatypes::series::{MetricName, SeriesOrdering}; +use futures::Stream; +use store_api::region_engine::{RawLabelSource, RawSeriesStream, SendableRawSeriesStream}; + +use super::*; + +fn raw_schema() -> Arc { + Arc::new(SeriesSchema { + fields: Arc::from([FieldSpec { + name: Arc::from("value"), + }]), + eval_grid: None, + }) +} + +fn output_schema() -> Arc { + Arc::new(Schema::new(vec![ + Field::new("value", DataType::Float64, true), + Field::new( + "time", + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + ])) +} + +fn mapping_for( + start: i64, + end: i64, + step: i64, + lookback: i64, +) -> (Arc, Arc) { + let logical = LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: Arc::new(DFSchema::empty()), + }); + let node = InstantManipulate::new( + start, + end, + lookback, + step, + 0, + "time".into(), + vec![], + Some("value".into()), + logical, + ); + let series = raw_schema(); + let mapping = NativeInstantMapping::try_new(&node, output_schema(), series.clone(), &[]) + .unwrap() + .unwrap(); + (Arc::new(mapping), series) +} + +fn chunk(time: &[i64], values: &[f64], last: bool) -> RawSeriesChunk { + RawSeriesChunk::try_new( + &raw_schema(), + datatypes::arrow::buffer::Buffer::from(vec![1, 2, 3]), + Some(MetricName::try_new(None, vec![]).unwrap()), + RawTimestamps::try_new(ScalarBuffer::from(time.to_vec())).unwrap(), + SeriesFields::one(Samples::from_float(Float64Array::from(values.to_vec()))), + last, + ) + .unwrap() +} + +struct RawMock { + schema: Arc, + labels: RawLabelSource, + items: VecDeque>>>, + polls: Arc, + drops: Arc, +} +impl Stream for RawMock { + type Item = std::result::Result; + fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + self.polls.fetch_add(1, Ordering::SeqCst); + self.items.pop_front().unwrap_or(Poll::Ready(None)) + } +} +impl Drop for RawMock { + fn drop(&mut self) { + self.drops.fetch_add(1, Ordering::SeqCst); + } +} +impl RawSeriesStream for RawMock { + fn schema(&self) -> &Arc { + &self.schema + } + fn label_source(&self) -> &RawLabelSource { + &self.labels + } + fn ordering(&self) -> SeriesOrdering { + SeriesOrdering::Contiguous + } +} + +fn raw( + items: impl IntoIterator>>>, +) -> (SendableRawSeriesStream, Arc, Arc) { + let polls = Arc::new(AtomicUsize::new(0)); + let drops = Arc::new(AtomicUsize::new(0)); + ( + Box::pin(RawMock { + schema: raw_schema(), + labels: RawLabelSource::Explicit, + items: items.into_iter().collect(), + polls: polls.clone(), + drops: drops.clone(), + }), + polls, + drops, + ) +} + +fn stream(source: SendableRawSeriesStream, pool: &Arc) -> NativeInstantStream { + stream_for(source, pool, 10, 30, 10, 10) +} + +fn stream_for( + source: SendableRawSeriesStream, + pool: &Arc, + start: i64, + end: i64, + step: i64, + lookback: i64, +) -> NativeInstantStream { + let (mapping, schema) = mapping_for(start, end, step, lookback); + NativeInstantStream { + source: Some(source), + mapping: mapping.clone(), + schema: mapping.schema(), + series_schema: schema, + pending: Vec::new(), + reservation: MemoryConsumer::new("NativeInstantStreamTest").register(pool), + } +} + +fn poll(stream: &mut NativeInstantStream) -> Poll>> { + let waker = futures::task::noop_waker(); + let mut cx = Context::from_waker(&waker); + Pin::new(stream).poll_next(&mut cx) +} + +#[test] +fn emits_complete_series_before_polling_later_pending_input() { + let (input, polls, _) = raw([ + Poll::Ready(Some(Ok(chunk(&[10], &[1.0], false)))), + Poll::Ready(Some(Ok(chunk(&[20], &[2.0], true)))), + Poll::Pending, + ]); + let pool = Arc::new(GreedyMemoryPool::new(1_000_000)) as Arc; + let mut stream = stream(input, &pool); + let Poll::Ready(Some(Ok(batch))) = poll(&mut stream) else { + panic!("completed series should emit") + }; + assert_eq!(polls.load(Ordering::SeqCst), 2); + assert_eq!( + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values(), + &[1.0, 2.0] + ); + assert_eq!( + batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .values(), + &[10, 20] + ); + assert!(matches!(poll(&mut stream), Poll::Pending)); + assert_eq!(polls.load(Ordering::SeqCst), 3); +} + +#[test] +fn source_error_after_prefix_is_returned_once_and_clears_reservation_and_input() { + let failure = BoxedError::new(PlainError::new( + "mock source error".into(), + StatusCode::Internal, + )); + let (input, polls, drops) = raw([ + Poll::Ready(Some(Ok(chunk(&[10], &[1.0], true)))), + Poll::Ready(Some(Err(failure))), + ]); + let pool = Arc::new(GreedyMemoryPool::new(1_000_000)) as Arc; + let mut stream = stream(input, &pool); + assert!(matches!(poll(&mut stream), Poll::Ready(Some(Ok(_))))); + let Poll::Ready(Some(Err(error))) = poll(&mut stream) else { + panic!("source error must be returned once") + }; + assert_eq!( + common_query::error::datafusion_status_code::(&error, None), + StatusCode::Internal + ); + let DataFusionError::External(source) = error else { + panic!("source error must stay a typed DataFusionError::External payload") + }; + let source = source + .downcast_ref::() + .expect("source error must still be the original BoxedError"); + assert_eq!(source.status_code(), StatusCode::Internal); + assert_eq!(source.to_string(), "mock source error"); + assert!(matches!(poll(&mut stream), Poll::Ready(None))); + assert!(matches!(poll(&mut stream), Poll::Ready(None))); + assert_eq!(polls.load(Ordering::SeqCst), 2); + assert_eq!(drops.load(Ordering::SeqCst), 1); + assert_eq!(pool.reserved(), 0); +} + +#[test] +fn unfinished_series_eof_is_error_once_and_drops_partial_chunks() { + let (input, _, drops) = raw([Poll::Ready(Some(Ok(chunk(&[10], &[1.0], false))))]); + let pool = Arc::new(GreedyMemoryPool::new(1_000_000)) as Arc; + let mut stream = stream(input, &pool); + assert!(matches!(poll(&mut stream), Poll::Ready(Some(Err(_))))); + assert_eq!(stream.pending.len(), 0); + assert_eq!(drops.load(Ordering::SeqCst), 1); + assert_eq!(pool.reserved(), 0); + assert!(matches!(poll(&mut stream), Poll::Ready(None))); +} + +#[test] +fn dropping_pending_stream_releases_source_chunks_and_memory() { + let array = Arc::new(Float64Array::from(vec![1.0])); + let samples = Samples::from_float(array.as_ref().clone()); + drop(array); + let pending_chunk = RawSeriesChunk::try_new( + &raw_schema(), + datatypes::arrow::buffer::Buffer::from(vec![1]), + Some(MetricName::try_new(None, vec![]).unwrap()), + RawTimestamps::try_new(ScalarBuffer::from(vec![10])).unwrap(), + SeriesFields::one(samples), + false, + ) + .unwrap(); + let (input, _, drops) = raw([Poll::Ready(Some(Ok(pending_chunk))), Poll::Pending]); + let pool = Arc::new(GreedyMemoryPool::new(1_000_000)) as Arc; + let mut stream = stream(input, &pool); + assert!(matches!(poll(&mut stream), Poll::Pending)); + assert!(pool.reserved() > 0); + drop(stream); + assert_eq!(drops.load(Ordering::SeqCst), 1); + assert_eq!(pool.reserved(), 0); +} + +#[test] +fn reservation_failure_is_terminal_and_releases_input_and_partial_series() { + let (input, _, drops) = raw([Poll::Ready(Some(Ok(chunk(&[10], &[1.0], false))))]); + let pool = Arc::new(GreedyMemoryPool::new(1)) as Arc; + let mut stream = stream(input, &pool); + assert!(matches!( + poll(&mut stream), + Poll::Ready(Some(Err(DataFusionError::ResourcesExhausted(_)))) + )); + assert_eq!(drops.load(Ordering::SeqCst), 1); + assert_eq!(stream.pending.len(), 0); + assert_eq!(stream.pending.capacity(), 0); + assert_eq!(pool.reserved(), 0); + assert!(matches!(poll(&mut stream), Poll::Ready(None))); +} + +#[test] +fn source_error_drops_unfinished_series_without_emitting_it() { + let failure = BoxedError::new(PlainError::new( + "mock source error".into(), + StatusCode::Internal, + )); + let (input, polls, drops) = raw([ + Poll::Ready(Some(Ok(chunk(&[10], &[1.0], false)))), + Poll::Ready(Some(Err(failure))), + ]); + let pool = Arc::new(GreedyMemoryPool::new(1_000_000)) as Arc; + let mut stream = stream(input, &pool); + assert!(matches!(poll(&mut stream), Poll::Ready(Some(Err(_))))); + assert!(stream.pending.is_empty()); + assert_eq!(stream.pending.capacity(), 0); + assert!(matches!(poll(&mut stream), Poll::Ready(None))); + assert_eq!(polls.load(Ordering::SeqCst), 2); + assert_eq!(drops.load(Ordering::SeqCst), 1); + assert_eq!(pool.reserved(), 0); +} + +#[test] +fn descriptor_heavy_series_fails_reservation_before_accumulating_unbounded_chunks() { + let items = (0..512) + .map(|i| Poll::Ready(Some(Ok(chunk(&[i], &[i as f64], false))))) + .chain([Poll::Pending]); + let (input, polls, drops) = raw(items); + let pool = Arc::new(GreedyMemoryPool::new(32 * 1024)) as Arc; + let mut stream = stream(input, &pool); + assert!(matches!( + poll(&mut stream), + Poll::Ready(Some(Err(DataFusionError::ResourcesExhausted(_)))) + )); + assert!(stream.pending.is_empty()); + assert_eq!(stream.pending.capacity(), 0); + assert!(matches!(poll(&mut stream), Poll::Ready(None))); + assert!(polls.load(Ordering::SeqCst) < 512); + assert_eq!(drops.load(Ordering::SeqCst), 1); + assert_eq!(pool.reserved(), 0); +} + +#[test] +fn dense_large_chunk_uses_segment_count_not_row_count_for_descriptors() { + let count = 100_000usize; + let timestamps = (0..count).map(|i| i as i64).collect::>(); + let values = Float64Array::from(vec![7.0; count]); + let dense_chunk = RawSeriesChunk::try_new( + &raw_schema(), + datatypes::arrow::buffer::Buffer::from(vec![1]), + Some(MetricName::try_new(None, vec![]).unwrap()), + RawTimestamps::try_new(ScalarBuffer::from(timestamps)).unwrap(), + SeriesFields::one(Samples::from_float(values)), + true, + ) + .unwrap(); + let (input, _, _) = raw([Poll::Ready(Some(Ok(dense_chunk)))]); + // Input timestamp/field buffers and concatenation scratch fit; row-scaled Samples + // descriptors would exceed this pool despite there being only one dense segment. + let pool = Arc::new(GreedyMemoryPool::new(4 * 1024 * 1024)) as Arc; + let mut stream = stream_for(input, &pool, 10, 11, 1, 0); + let Poll::Ready(Some(Ok(batch))) = poll(&mut stream) else { + panic!("dense chunk should fit its bounded pool") + }; + assert_eq!(batch.num_rows(), 2); + assert_eq!( + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values(), + &[7.0, 7.0] + ); + assert_eq!( + batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .values(), + &[10, 11] + ); + assert_eq!(pool.reserved(), 0); +} + +#[test] +fn output_grid_scratch_is_reserved_before_large_evaluation() { + let (input, polls, drops) = raw([Poll::Ready(Some(Ok(chunk(&[0], &[7.0], true))))]); + let pool = Arc::new(GreedyMemoryPool::new(128 * 1024)) as Arc; + // One source point with a long lookback produces a large grid despite the tiny source. + let mut stream = stream_for(input, &pool, 0, 900_000, 1, 900_001); + assert!(matches!( + poll(&mut stream), + Poll::Ready(Some(Err(DataFusionError::ResourcesExhausted(_)))) + )); + assert_eq!(polls.load(Ordering::SeqCst), 1); + assert_eq!(drops.load(Ordering::SeqCst), 1); + assert!(stream.pending.is_empty()); + assert_eq!(pool.reserved(), 0); + assert!(matches!(poll(&mut stream), Poll::Ready(None))); +} + +#[test] +fn larger_pool_completes_one_point_many_evaluation_grid() { + let (input, _, _) = raw([Poll::Ready(Some(Ok(chunk(&[0], &[7.0], true))))]); + let pool = Arc::new(GreedyMemoryPool::new(32 * 1024 * 1024)) as Arc; + let mut stream = stream_for(input, &pool, 0, 4, 1, 5); + let Poll::Ready(Some(Ok(batch))) = poll(&mut stream) else { + panic!("sufficient reservation should evaluate") + }; + assert_eq!( + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values(), + &[7.0; 5] + ); + assert_eq!( + batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .values(), + &[0, 1, 2, 3, 4] + ); + assert_eq!(pool.reserved(), 0); +} diff --git a/src/promql/src/extension_plan/instant_manipulate/native_tests.rs b/src/promql/src/extension_plan/instant_manipulate/native_tests.rs new file mode 100644 index 00000000000..2193d13e695 --- /dev/null +++ b/src/promql/src/extension_plan/instant_manipulate/native_tests.rs @@ -0,0 +1,597 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::HashMap; +use std::sync::Arc; + +use common_query::prometheus::PROMETHEUS_STALE_NAN_BITS; +use datafusion::arrow::array::{Array, Float64Array, StringArray, TimestampMillisecondArray}; +use datafusion::arrow::datatypes::{DataType, Field, Schema, TimeUnit}; +use datafusion::common::DFSchema; +use datafusion::logical_expr::{EmptyRelation, LogicalPlan}; +use datatypes::arrow::buffer::ScalarBuffer; +use datatypes::series::raw::RawTimestamps; +use datatypes::series::samples::{FieldSpec, Samples, SeriesFields, SeriesSchema}; +use datatypes::series::{Label, MetricName}; + +use super::*; + +fn node() -> InstantManipulate { + let input = LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: Arc::new(DFSchema::empty()), + }); + InstantManipulate::new( + 10, + 20, + 10, + 10, + 0, + "time".into(), + vec![], + Some("value".into()), + input, + ) +} + +fn schema(tag_type: DataType) -> SchemaRef { + let mut metadata = HashMap::new(); + metadata.insert("schema".to_string(), "kept".to_string()); + Arc::new(Schema::new_with_metadata( + vec![ + Field::new("value", DataType::Float64, true) + .with_metadata([("field-key".to_string(), "field-value".to_string())].into()), + Field::new( + "time", + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new("zone", tag_type, true), + ], + metadata, + )) +} + +fn series_schema() -> Arc { + Arc::new(SeriesSchema { + fields: Arc::from([FieldSpec { + name: Arc::from("value"), + }]), + eval_grid: None, + }) +} + +fn metric(zone: Option<&str>) -> MetricName { + let labels = zone + .map(|value| { + vec![Label { + name: Arc::from("zone"), + value: Arc::from(value), + }] + }) + .unwrap_or_default(); + MetricName::try_new(None, labels).unwrap() +} + +/// The ordered, deliberately non-lexical field names of the multifloat regression. +const MULTIFLOAT_FIELDS: [&str; 3] = ["value", "a_second", "m_third"]; + +/// The strict millisecond timestamps of every multifloat case. +const MULTIFLOAT_TIMESTAMPS: [i64; 5] = [1000, 2000, 3000, 4000, 5000]; + +/// The kept source rows of each configured staleness field: the stale row of the configured +/// field is suppressed, so the first field drops row 3 (4000) and the second field drops row 2 +/// (3000) without ever retrying an older sample. +const MULTIFLOAT_KEPT: [[usize; 4]; 2] = [[0, 1, 2, 4], [0, 1, 3, 4]]; + +/// An ordinary, non-stale NaN payload that must survive the mapping as exact bits. +const ORDINARY_NAN_BITS: u64 = 0x7ff8_0000_0000_0001; + +/// The output schema of the regression: the declared field order, then the time index and the +/// tag, with metadata that every case must retain. +/// +/// Every declared field is nullable because every fixture lane keeps a NULL row in each of them: +/// the first row of the second and third fields and the second row of the first field. +fn multifloat_schema(field_names: &[&str], tag_type: DataType) -> SchemaRef { + let mut metadata = HashMap::new(); + metadata.insert("schema".to_string(), "kept".to_string()); + let mut fields = field_names + .iter() + .enumerate() + .map(|(index, name)| { + Field::new(*name, DataType::Float64, true) + .with_metadata([("field-key".to_string(), format!("field-value-{index}"))].into()) + }) + .collect::>(); + fields.push(Field::new( + "time", + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + )); + fields.push(Field::new("zone", tag_type, true)); + Arc::new(Schema::new_with_metadata(fields, metadata)) +} + +/// The series schema of the same declared field order. +fn multifloat_series_schema(field_names: &[&str]) -> Arc { + Arc::new(SeriesSchema { + fields: field_names + .iter() + .map(|name| FieldSpec { + name: Arc::from(*name), + }) + .collect::>() + .into(), + eval_grid: None, + }) +} + +/// The five source rows of each multifloat field. +/// +/// Row `i` is the strict millisecond timestamp `MULTIFLOAT_TIMESTAMPS[i]`. The first field is +/// stale at 4000 and carries an ordinary NaN payload at 5000; the unrelated second field is +/// stale at 3000; the fields carry NULLs and a signed zero as well. +fn multifloat_source(field_names: &[&str]) -> Vec { + let stale = f64::from_bits(PROMETHEUS_STALE_NAN_BITS); + let nan = f64::from_bits(ORDINARY_NAN_BITS); + let rows: [&[Option]; 3] = [ + &[Some(1.0), None, Some(-0.0), Some(stale), Some(nan)], + &[None, Some(-0.0), Some(stale), Some(2.0), Some(4.0)], + &[None, Some(3.0), Some(nan), Some(3.0), Some(5.0)], + ]; + field_names + .iter() + .enumerate() + .map(|(index, _)| Float64Array::from(rows[index].to_vec())) + .collect() +} + +/// The chunked samples of the multifloat input: the first field arrives in three chunks, two of +/// which are slices with non-zero offsets, and every other field arrives as one dense array. +fn multifloat_fields(source: &[Float64Array]) -> SeriesFields { + let fields = source + .iter() + .enumerate() + .map(|(index, values)| { + if index > 0 { + return Samples::from_float(values.clone()); + } + Samples::try_from_chunks(vec![ + Samples::from_float(values.slice(0, 1)), + Samples::from_float(values.slice(1, 3)), + Samples::from_float(values.slice(4, 1)), + ]) + .unwrap() + }) + .collect::>(); + SeriesFields::try_from_vec(fields).unwrap() +} + +/// The bounded instant declaration of the regression: one output point per input millisecond +/// with one lookback step, whose staleness field is exactly one declared field. +fn multifloat_node(field_column: &str) -> InstantManipulate { + let mut query = node(); + query.start = MULTIFLOAT_TIMESTAMPS[0]; + query.end = MULTIFLOAT_TIMESTAMPS[4]; + query.lookback_delta = 1000; + query.interval = 1000; + query.field_column = Some(field_column.to_string()); + query +} + +#[test] +fn evaluates_sliced_chunked_samples_and_preserves_original_schema() { + let output_schema = schema(DataType::Dictionary( + Box::new(DataType::UInt32), + Box::new(DataType::Utf8), + )); + let mapping = NativeInstantMapping::try_new( + &node(), + output_schema.clone(), + series_schema(), + &["zone".into()], + ) + .unwrap() + .unwrap(); + let a = Float64Array::from(vec![Some(99.0), Some(1.0), None]); + let b = Float64Array::from(vec![Some(2.0), Some(3.0)]); + let samples = Samples::try_from_chunks(vec![ + Samples::from_float( + a.slice(1, 2) + .as_any() + .downcast_ref::() + .unwrap() + .clone(), + ), + Samples::from_float(b), + ]) + .unwrap(); + let fields = SeriesFields::one(samples); + let time = RawTimestamps::try_new(ScalarBuffer::from(vec![0, 10, 20, 30])).unwrap(); + let batch = mapping.evaluate(&metric(Some("")), &time, &fields).unwrap(); + assert_eq!(batch.num_rows(), 2); + assert_eq!(batch.schema(), output_schema); + let output_values = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert!(output_values.is_null(0)); + assert_eq!(output_values.value(1), 2.0); + assert_eq!( + batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .values(), + &[10, 20] + ); + let tag_values = datafusion::arrow::compute::cast(batch.column(2), &DataType::Utf8).unwrap(); + let tag_values = tag_values.as_any().downcast_ref::().unwrap(); + assert_eq!(tag_values.value(0), ""); + assert_eq!(tag_values.value(1), ""); +} + +#[test] +fn selected_repeated_indices_across_sliced_chunks_preserve_missing_and_nan() { + let mut query = node(); + query.end = 30; + query.lookback_delta = 20; + let mapping = NativeInstantMapping::try_new( + &query, + schema(DataType::Utf8), + series_schema(), + &["zone".into()], + ) + .unwrap() + .unwrap(); + let first = Float64Array::from(vec![Some(99.0), None]).slice(1, 1); + let second = Float64Array::from(vec![Some(f64::NAN), Some(88.0)]).slice(0, 1); + let samples = Samples::try_from_chunks(vec![ + Samples::from_float( + first + .as_any() + .downcast_ref::() + .unwrap() + .clone(), + ), + Samples::from_float( + second + .as_any() + .downcast_ref::() + .unwrap() + .clone(), + ), + ]) + .unwrap(); + let fields = SeriesFields::one(samples); + let time = RawTimestamps::try_new(ScalarBuffer::from(vec![10, 30])).unwrap(); + + let batch = mapping.evaluate(&metric(None), &time, &fields).unwrap(); + let values = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(batch.num_rows(), 3); + assert!(values.is_null(0)); + assert!(values.is_null(1)); + assert!(!values.is_null(2)); + assert!(values.value(2).is_nan()); +} + +#[test] +fn maps_all_supported_string_types_and_distinguishes_missing_from_empty() { + let tag_types = [ + DataType::Utf8, + DataType::LargeUtf8, + DataType::Utf8View, + DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8)), + ]; + let fields = SeriesFields::one(Samples::from_float(Float64Array::from(vec![ + Some(1.0), + Some(2.0), + ]))); + let time = RawTimestamps::try_new(ScalarBuffer::from(vec![10, 20])).unwrap(); + + for tag_type in tag_types { + let output_schema = schema(tag_type); + let mapping = NativeInstantMapping::try_new( + &node(), + output_schema.clone(), + series_schema(), + &["zone".into()], + ) + .unwrap() + .unwrap(); + + for (label, expected) in [(Some(""), Some("")), (None, None)] { + let batch = mapping.evaluate(&metric(label), &time, &fields).unwrap(); + assert_eq!(batch.schema(), output_schema); + let values = + datafusion::arrow::compute::cast(batch.column(2), &DataType::Utf8).unwrap(); + let values = values.as_any().downcast_ref::().unwrap(); + match expected { + Some(value) => { + assert_eq!(values.value(0), value); + assert_eq!(values.value(1), value); + assert_eq!(values.null_count(), 0); + } + None => assert_eq!(values.null_count(), 2), + } + } + } +} + +#[test] +fn selected_field_staleness_and_no_field_semantics_match_existing_selector_contract() { + let stale = f64::from_bits(PROMETHEUS_STALE_NAN_BITS); + let times = RawTimestamps::try_new(ScalarBuffer::from(vec![10, 20])).unwrap(); + let values = SeriesFields::one(Samples::from_float(Float64Array::from(vec![ + Some(1.0), + Some(stale), + ]))); + let mapping = NativeInstantMapping::try_new( + &node(), + schema(DataType::Utf8), + series_schema(), + &["zone".into()], + ) + .unwrap() + .unwrap(); + let batch = mapping.evaluate(&metric(None), ×, &values).unwrap(); + assert_eq!(batch.num_rows(), 1); + assert_eq!( + batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + 10 + ); + + let mut no_staleness_node = node(); + no_staleness_node.field_column = None; + let no_staleness_mapping = NativeInstantMapping::try_new( + &no_staleness_node, + schema(DataType::Utf8), + series_schema(), + &["zone".into()], + ) + .unwrap() + .unwrap(); + let batch = no_staleness_mapping + .evaluate(&metric(None), ×, &values) + .unwrap(); + let output = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(output.len(), 2); + assert_eq!(output.value(1).to_bits(), PROMETHEUS_STALE_NAN_BITS); +} + +#[test] +fn rejects_unmapped_missing_duplicate_and_non_millisecond_schema_columns() { + let valid = schema(DataType::Utf8); + let mut cases = Vec::new(); + + cases.push(Arc::new(Schema::new(vec![ + valid.field(0).clone(), + valid.field(1).clone(), + valid.field(2).clone(), + Field::new("helper", DataType::UInt64, true), + ]))); + cases.push(Arc::new(Schema::new(vec![ + valid.field(0).clone(), + valid.field(1).clone(), + valid.field(2).clone(), + Field::new("__tsid", DataType::UInt64, true), + ]))); + cases.push(Arc::new(Schema::new(vec![ + valid.field(0).clone(), + valid.field(1).clone(), + ]))); + cases.push(Arc::new(Schema::new(vec![ + valid.field(1).clone(), + valid.field(2).clone(), + ]))); + cases.push(Arc::new(Schema::new(vec![ + valid.field(0).clone(), + valid.field(1).clone(), + valid.field(2).clone(), + valid.field(0).clone(), + ]))); + cases.push(Arc::new(Schema::new(vec![ + valid.field(0).clone(), + Field::new( + "time", + DataType::Timestamp(TimeUnit::Microsecond, None), + false, + ), + valid.field(2).clone(), + ]))); + cases.push(Arc::new(Schema::new(vec![ + Field::new("value", DataType::Int64, true), + valid.field(1).clone(), + valid.field(2).clone(), + ]))); + + for output_schema in cases { + assert!( + NativeInstantMapping::try_new( + &node(), + output_schema, + series_schema(), + &["zone".into()], + ) + .unwrap() + .is_none() + ); + } +} + +#[test] +fn absent_tag_is_null_and_nonzero_offset_is_rejected() { + let output_schema = schema(DataType::Utf8); + let mapping = + NativeInstantMapping::try_new(&node(), output_schema, series_schema(), &["zone".into()]) + .unwrap() + .unwrap(); + let fields = SeriesFields::one(Samples::from_float(Float64Array::from(vec![ + Some(1.0), + Some(2.0), + ]))); + let time = RawTimestamps::try_new(ScalarBuffer::from(vec![10, 20])).unwrap(); + let batch = mapping.evaluate(&metric(None), &time, &fields).unwrap(); + assert_eq!(batch.column(2).null_count(), 2); + + let mut unsupported = node(); + unsupported.offset = 1; + assert!( + NativeInstantMapping::try_new( + &unsupported, + schema(DataType::Utf8), + series_schema(), + &["zone".into()], + ) + .unwrap() + .is_none() + ); +} + +#[test] +fn multifloat_native_mapping_keeps_ordered_fields_and_checks_only_the_configured_field() { + let time = RawTimestamps::try_new(ScalarBuffer::from(MULTIFLOAT_TIMESTAMPS.to_vec())).unwrap(); + + for field_names in [&MULTIFLOAT_FIELDS[..2], &MULTIFLOAT_FIELDS[..3]] { + for (configured, kept) in MULTIFLOAT_KEPT.iter().enumerate() { + let output_schema = multifloat_schema(field_names, DataType::Utf8); + let source = multifloat_source(field_names); + let fields = multifloat_fields(&source); + let mapping = NativeInstantMapping::try_new( + &multifloat_node(field_names[configured]), + output_schema.clone(), + multifloat_series_schema(field_names), + &["zone".into()], + ) + .unwrap() + .expect("the complete ordered Float64 fields are mapped instead of declined"); + let batch = mapping + .evaluate(&metric(Some("multifloat")), &time, &fields) + .unwrap(); + + // The complete output schema, including its metadata and nullability, is retained. + assert_eq!(batch.schema(), output_schema); + assert_eq!(batch.num_rows(), kept.len()); + let kept_timestamps = batch + .column(field_names.len()) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec(); + assert_eq!( + kept_timestamps, + kept.iter() + .map(|row| MULTIFLOAT_TIMESTAMPS[*row]) + .collect::>() + ); + let suppressed = if configured == 0 { 3 } else { 2 }; + let unrelated = if configured == 0 { 2 } else { 3 }; + assert!( + !kept_timestamps.contains(&MULTIFLOAT_TIMESTAMPS[suppressed]), + "the stale row of the configured field is suppressed instead of retried" + ); + assert!( + kept_timestamps.contains(&MULTIFLOAT_TIMESTAMPS[unrelated]), + "the stale marker of another field does not suppress its row" + ); + + // Every output field keeps the nullness and the exact bits of its own source row. + for (index, values) in source.iter().enumerate() { + let output = batch + .column(index) + .as_any() + .downcast_ref::() + .unwrap(); + for (row, source_row) in kept.iter().enumerate() { + if values.is_null(*source_row) { + assert!(output.is_null(row)); + } else { + assert!(!output.is_null(row)); + assert_eq!( + output.value(row).to_bits(), + values.value(*source_row).to_bits(), + "field {index} row {row} must keep its exact source bits" + ); + } + } + } + + // The unchecked stale marker travels through as exact bits, so the row of another + // field is kept with the marker intact instead of being compared or re-encoded. + let (unchecked_field, unchecked_row) = if configured == 0 { (1, 2) } else { (0, 3) }; + let unchecked_output_row = kept + .iter() + .position(|row| *row == unchecked_row) + .expect("the unchecked stale row is kept"); + let marker = batch + .column(unchecked_field) + .as_any() + .downcast_ref::() + .unwrap() + .value(unchecked_output_row); + assert_eq!(marker.to_bits(), PROMETHEUS_STALE_NAN_BITS); + + // The ordinary NaN payload and the signed zero of the first two fields survive as + // exact bits as well. + let nan_row = kept + .iter() + .position(|row| *row == 4) + .expect("the ordinary NaN row is kept"); + let nan = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(nan_row); + assert_eq!(nan.to_bits(), ORDINARY_NAN_BITS); + let signed_zero_row = kept + .iter() + .position(|row| *row == 1) + .expect("the signed zero row is kept"); + let signed_zero = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .value(signed_zero_row); + assert_eq!(signed_zero.to_bits(), (-0.0_f64).to_bits()); + + // The series key of every kept row is the requested label value. + let tags = datafusion::arrow::compute::cast( + batch.column(field_names.len() + 1), + &DataType::Utf8, + ) + .unwrap(); + let tags = tags.as_any().downcast_ref::().unwrap(); + assert_eq!(tags.null_count(), 0); + assert!(tags.iter().all(|value| value == Some("multifloat"))); + } + } +}