diff --git a/src/common/query/Cargo.toml b/src/common/query/Cargo.toml index 4acdfafcb6c..75db66bb12b 100644 --- a/src/common/query/Cargo.toml +++ b/src/common/query/Cargo.toml @@ -25,6 +25,7 @@ datafusion-common.workspace = true datafusion-expr.workspace = true datafusion-proto.workspace = true datatypes.workspace = true +futures-util.workspace = true once_cell.workspace = true prost.workspace = true serde.workspace = true @@ -36,7 +37,6 @@ store-api.workspace = true [dev-dependencies] common-base.workspace = true -futures-util.workspace = true serde_json.workspace = true tokio.workspace = true diff --git a/src/common/query/src/lib.rs b/src/common/query/src/lib.rs index 8151dfc5731..ee5cf92767b 100644 --- a/src/common/query/src/lib.rs +++ b/src/common/query/src/lib.rs @@ -20,6 +20,7 @@ pub mod prelude; pub mod prometheus; pub mod promql_annotations; pub mod request; +pub mod series_frame; pub mod stream; #[cfg(any(test, feature = "testing"))] pub mod test_util; diff --git a/src/common/query/src/series_frame.rs b/src/common/query/src/series_frame.rs new file mode 100644 index 00000000000..e2ebc8f03eb --- /dev/null +++ b/src/common/query/src/series_frame.rs @@ -0,0 +1,451 @@ +// 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. + +//! Streaming boundary between finite raw-series chunks and DataFusion record batches. +//! +//! Schema fields identify the ordered logical payload fields and codec layout; additional +//! application metadata is permitted when stable across the stream. The caller-provided +//! [`RawLabelSource`] is the label binding, not authenticated by the frame schema or field +//! provenance. This adapter preserves scan candidates only; it is not a residual filter, +//! authorization check, remote transport, or compaction boundary. Only +//! [`SeriesOrdering::Contiguous`] and [`SeriesOrdering::KeyThenTime`] are accepted; +//! `Interleaved` is rejected before any poll. + +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use common_error::ext::BoxedError; +use common_recordbatch::{DfRecordBatch, DfSendableRecordBatchStream}; +use datafusion_common::DataFusionError; +use datatypes::error::InvalidSeriesSnafu; +use datatypes::series::frame::{ + RawLabelMode, RawSeriesChunkCursor, RawSeriesFrameBuilder, RawSeriesFrameDecoder, +}; +use datatypes::series::raw::RawSeriesChunk; +use datatypes::series::samples::SeriesSchema; +use datatypes::series::{MetricName, SeriesOrdering}; +use futures_util::Stream; +use store_api::region_engine::{RawLabelSource, RawSeriesStream, SendableRawSeriesStream}; + +#[cfg(test)] +mod tests; + +/// Packs a storage raw-series stream into DataFusion record batches incrementally. +/// +/// Only [`SeriesOrdering::Contiguous`] and [`SeriesOrdering::KeyThenTime`] streams are +/// accepted; an `Interleaved` declaration is rejected before the first poll, including for an +/// empty stream. Rows accumulate across chunks into batches bounded by the builder's row and +/// estimated payload limits; one legal oversized row may be emitted alone, and these limits are a +/// work bound, not a resident-memory cap. A ready batch is yielded without polling the source +/// again. Producers must keep each series' chunks contiguous. Validation checks adjacent chunk +/// conditions but does not certify nonadjacent recurrence for `Contiguous` ordering. +pub fn pack_raw_series( + source: SendableRawSeriesStream, +) -> datafusion_common::Result { + let ordering = source.ordering(); + ensure_supported_ordering(ordering)?; + let schema = source.schema().clone(); + let label_source = source.label_source().clone(); + let mode = label_mode(&label_source); + let builder = RawSeriesFrameBuilder::try_new( + &schema, + mode, + RawSeriesFrameBuilder::DEFAULT_MAX_ROWS, + RawSeriesFrameBuilder::DEFAULT_MAX_BYTES, + ) + .map_err(codec_error)?; + let frame_schema = builder.schema().clone(); + Ok(Box::pin(PackStream { + source: Some(source), + schema, + mode, + frame_schema, + ordering, + validator: SequenceValidator::new(ordering), + builder: Some(builder), + cursor: None, + })) +} + +/// Unpacks DataFusion batches using the caller's explicit scan binding and ordering. +/// +/// Only [`SeriesOrdering::Contiguous`] and [`SeriesOrdering::KeyThenTime`] streams are +/// accepted; an `Interleaved` declaration is rejected before the stream is returned, including +/// for an empty stream. Producers must keep each series' chunks contiguous. Validation checks +/// adjacent chunk conditions but does not certify nonadjacent recurrence for `Contiguous` +/// ordering. +/// +/// The supplied binding is retained unchanged and is not validated as field provenance. +pub fn unpack_raw_series( + source: DfSendableRecordBatchStream, + expected: Arc, + label_source: RawLabelSource, + ordering: SeriesOrdering, +) -> datafusion_common::Result { + ensure_supported_ordering(ordering)?; + let mode = label_mode(&label_source); + let frame_schema = source.schema(); + let decoder = RawSeriesFrameDecoder::try_new(&expected, mode).map_err(codec_error)?; + // Validate even empty streams' declared schema and layout before returning the stream. + let empty = DfRecordBatch::new_empty(frame_schema.clone()); + decoder.decode(&empty).map_err(codec_error)?; + Ok(Box::pin(UnpackStream { + source: Some(source), + schema: expected, + frame_schema, + label_source, + ordering, + validator: SequenceValidator::new(ordering), + pending: Vec::new().into_iter(), + decoder, + })) +} + +fn label_mode(source: &RawLabelSource) -> RawLabelMode { + match source { + RawLabelSource::PrimaryKey { .. } => RawLabelMode::Key, + RawLabelSource::Explicit => RawLabelMode::Explicit, + } +} + +/// Rejects orderings this adapter cannot validate; a bad declaration is an error, not a fallback. +fn ensure_supported_ordering(ordering: SeriesOrdering) -> datafusion_common::Result<()> { + match ordering { + SeriesOrdering::Contiguous | SeriesOrdering::KeyThenTime => Ok(()), + SeriesOrdering::Interleaved => { + Err(invalid("interleaved raw series ordering is not supported")) + } + } +} + +fn codec_error(error: datatypes::error::Error) -> DataFusionError { + DataFusionError::External(Box::new(BoxedError::new(error))) +} + +fn upstream_error(error: BoxedError) -> DataFusionError { + DataFusionError::External(Box::new(error)) +} + +fn invalid(reason: impl Into) -> DataFusionError { + DataFusionError::External(Box::new(BoxedError::new( + InvalidSeriesSnafu { + reason: reason.into(), + } + .build(), + ))) +} + +/// Cross-chunk validation state. The producer must keep each series' chunks contiguous; this +/// validator checks adjacent chunks only and never sees `Interleaved` streams. It retains only +/// the previous complete owned key (not an Arrow buffer), timestamp, metric, and finality. +struct SequenceValidator { + ordering: SeriesOrdering, + previous_key: Option>, + previous_timestamp: Option, + previous_metric: Option, + previous_is_last: bool, +} + +impl SequenceValidator { + fn new(ordering: SeriesOrdering) -> Self { + Self { + ordering, + previous_key: None, + previous_timestamp: None, + previous_metric: None, + previous_is_last: false, + } + } + + fn finish(&self) -> datafusion_common::Result<()> { + if self.previous_key.is_some() && !self.previous_is_last { + return Err(invalid("raw series stream ended before key was final")); + } + Ok(()) + } + + fn validate(&mut self, chunk: &RawSeriesChunk) -> datafusion_common::Result<()> { + let key = chunk.primary_key.as_slice(); + let timestamp = chunk + .time + .len() + .checked_sub(1) + .and_then(|last| chunk.time.timestamp(last)) + .ok_or_else(|| invalid("empty raw series chunk"))?; + if let Some(previous) = &self.previous_key { + match previous.as_slice().cmp(key) { + std::cmp::Ordering::Equal => { + if self.previous_is_last { + return Err(invalid("raw series key recurred after final chunk")); + } + if self.previous_metric != chunk.metric { + return Err(invalid("raw series metric changed within key")); + } + let first = chunk + .time + .timestamp(0) + .ok_or_else(|| invalid("empty raw series chunk"))?; + if self.previous_timestamp.is_some_and(|last| first <= last) { + return Err(invalid( + "raw series timestamps are not increasing across chunks", + )); + } + } + _ => { + if !self.previous_is_last { + return Err(invalid("raw series key changed before final chunk")); + } + if self.ordering == SeriesOrdering::KeyThenTime && previous.as_slice() > key { + return Err(invalid("raw series keys are not ordered")); + } + } + } + } + self.previous_key = Some(key.to_vec()); + self.previous_timestamp = Some(timestamp); + self.previous_metric = chunk.metric.clone(); + self.previous_is_last = chunk.is_last; + Ok(()) + } +} + +struct PackStream { + source: Option, + schema: Arc, + mode: RawLabelMode, + frame_schema: datatypes::arrow::datatypes::SchemaRef, + ordering: SeriesOrdering, + validator: SequenceValidator, + builder: Option, + cursor: Option, +} + +impl PackStream { + /// Releases the source and drops every buffered owner, leaving the stream terminal; batches + /// already yielded are unaffected. + fn finish(&mut self) { + self.source = None; + self.cursor = None; + self.builder = None; + self.validator = SequenceValidator::new(self.ordering); + } + + /// Advances the installed cursor until it yields one ready batch or is exhausted, dropping an + /// exhausted cursor. `None` means no batch is ready, so the source may be polled again. + fn advance_cursor(&mut self) -> datatypes::error::Result> { + let ready = match (self.cursor.as_mut(), self.builder.as_mut()) { + (Some(cursor), Some(builder)) => cursor.append_to(builder)?, + _ => return Ok(None), + }; + match ready { + Some(batch) => Ok(Some(batch)), + None => { + self.cursor = None; + Ok(None) + } + } + } + + /// Flushes the rows buffered so far, if any, so a pending source can still deliver data. + fn flush_buffered(&mut self) -> datatypes::error::Result> { + match self.builder.as_mut() { + Some(builder) => builder.flush(), + None => Ok(None), + } + } +} + +impl Stream for PackStream { + type Item = datafusion_common::Result; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = &mut *self; + loop { + match this.advance_cursor() { + Ok(Some(batch)) => return Poll::Ready(Some(Ok(batch))), + Ok(None) => {} + Err(error) => { + this.finish(); + return Poll::Ready(Some(Err(codec_error(error)))); + } + } + let polled = match this.source.as_mut() { + Some(source) => source.as_mut().poll_next(cx), + None => return Poll::Ready(None), + }; + match polled { + Poll::Pending => { + let flushed = this.flush_buffered(); + return match flushed { + Ok(Some(batch)) => Poll::Ready(Some(Ok(batch))), + Ok(None) => Poll::Pending, + Err(error) => { + this.finish(); + Poll::Ready(Some(Err(codec_error(error)))) + } + }; + } + Poll::Ready(None) => { + if let Err(error) = this.validator.finish() { + this.finish(); + return Poll::Ready(Some(Err(error))); + } + let flushed = this.flush_buffered(); + this.finish(); + return match flushed { + Ok(Some(batch)) => Poll::Ready(Some(Ok(batch))), + Ok(None) => Poll::Ready(None), + Err(error) => Poll::Ready(Some(Err(codec_error(error)))), + }; + } + Poll::Ready(Some(Err(error))) => { + this.finish(); + return Poll::Ready(Some(Err(upstream_error(error)))); + } + Poll::Ready(Some(Ok(chunk))) => { + if let Err(error) = this.validator.validate(&chunk) { + this.finish(); + return Poll::Ready(Some(Err(error))); + } + match RawSeriesChunkCursor::try_new(&this.schema, this.mode, chunk) { + Ok(cursor) => this.cursor = Some(cursor), + Err(error) => { + this.finish(); + return Poll::Ready(Some(Err(codec_error(error)))); + } + } + } + } + } + } +} + +impl datafusion::physical_plan::RecordBatchStream for PackStream { + fn schema(&self) -> datatypes::arrow::datatypes::SchemaRef { + self.frame_schema.clone() + } +} + +struct UnpackStream { + source: Option, + schema: Arc, + frame_schema: datatypes::arrow::datatypes::SchemaRef, + label_source: RawLabelSource, + ordering: SeriesOrdering, + validator: SequenceValidator, + pending: std::vec::IntoIter, + decoder: RawSeriesFrameDecoder, +} + +impl RawSeriesStream for UnpackStream { + fn schema(&self) -> &Arc { + &self.schema + } + + fn label_source(&self) -> &RawLabelSource { + &self.label_source + } + + fn ordering(&self) -> SeriesOrdering { + self.ordering + } +} + +impl UnpackStream { + fn finish(&mut self) { + self.source = None; + self.pending = Vec::new().into_iter(); + self.validator = SequenceValidator::new(self.ordering); + } +} + +impl Stream for UnpackStream { + type Item = Result; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + const MAX_EMPTY_BATCHES_PER_POLL: usize = 32; + let this = &mut *self; + let mut empty_batches = 0; + loop { + if let Some(chunk) = this.pending.next() { + return Poll::Ready(Some(Ok(chunk))); + } + let polled = match this.source.as_mut() { + Some(source) => source.as_mut().poll_next(cx), + None => return Poll::Ready(None), + }; + match polled { + Poll::Pending => return Poll::Pending, + Poll::Ready(None) => { + let result = this.validator.finish(); + this.finish(); + return Poll::Ready(result.err().map(|error| { + Err(BoxedError::new(crate::error::Error::GeneralDataFusion { + error, + location: snafu::Location::default(), + })) + })); + } + Poll::Ready(Some(Err(error))) => { + this.finish(); + return Poll::Ready(Some(Err(BoxedError::new( + crate::error::Error::GeneralDataFusion { + error, + location: snafu::Location::default(), + }, + )))); + } + Poll::Ready(Some(Ok(batch))) => { + let result = (|| { + if batch.schema() != this.frame_schema { + return Err(invalid("raw frame schema mismatch")); + } + this.decoder + .decode(&batch) + .map_err(codec_error) + .and_then(|chunks| { + for chunk in &chunks { + this.validator.validate(chunk)?; + } + Ok(chunks) + }) + })(); + match result { + Ok(chunks) => { + if chunks.is_empty() { + empty_batches += 1; + if empty_batches == MAX_EMPTY_BATCHES_PER_POLL { + cx.waker().wake_by_ref(); + return Poll::Pending; + } + } else { + this.pending = chunks.into_iter(); + } + } + Err(error) => { + this.finish(); + return Poll::Ready(Some(Err(BoxedError::new( + crate::error::Error::GeneralDataFusion { + error, + location: snafu::Location::default(), + }, + )))); + } + } + } + } + } + } +} diff --git a/src/common/query/src/series_frame/tests.rs b/src/common/query/src/series_frame/tests.rs new file mode 100644 index 00000000000..85d840ab4f2 --- /dev/null +++ b/src/common/query/src/series_frame/tests.rs @@ -0,0 +1,2495 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +mod lifecycle; + +use std::collections::VecDeque; +use std::io::Cursor; +use std::pin::Pin; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::task::{Context, Poll, Wake, Waker}; + +use common_error::ext::{BoxedError, PlainError}; +use common_error::status_code::StatusCode; +use common_recordbatch::DfSendableRecordBatchStream; +use datafusion::physical_plan::RecordBatchStream; +use datatypes::arrow::array::{Array, Float64Array, UInt32Array}; +use datatypes::arrow::buffer::{Buffer, ScalarBuffer}; +use datatypes::arrow::compute::concat_batches; +use datatypes::arrow::datatypes::{DataType, SchemaRef}; +use datatypes::arrow::ipc::{reader::StreamReader, writer::StreamWriter}; +use datatypes::arrow::record_batch::RecordBatch; +use datatypes::series::Label; +use datatypes::series::MetricName; +use datatypes::series::SeriesOrdering; +use datatypes::series::histogram::{BucketSideView, CounterResetHint, HistogramBatchBuilder}; +use datatypes::series::raw::{RawSeriesChunk, RawTimestamps}; +use datatypes::series::samples::{ + DenseSamplesRef, FieldSpec, SampleRef, Samples, SeriesFields, SeriesSchema, +}; +use futures_util::{Stream, StreamExt}; +use store_api::region_engine::{ + BoundTag, RawLabelSource, RawSeriesStream, SendableRawSeriesStream, +}; + +use super::{pack_raw_series, unpack_raw_series}; + +fn schema() -> Arc { + Arc::new(SeriesSchema { + fields: vec![FieldSpec { + name: Arc::from("value"), + }] + .into(), + eval_grid: None, + }) +} +fn metric() -> MetricName { + MetricName::try_new(Some(Arc::from("cpu")), vec![]).unwrap() +} +fn labeled_metric(host: &str) -> MetricName { + MetricName::try_new( + Some(Arc::from("cpu")), + vec![Label { + name: Arc::from("host"), + value: Arc::from(host), + }], + ) + .unwrap() +} +fn chunk_with( + key: &[u8], + times: &[i64], + values: Vec>, + metric: Option, + last: bool, +) -> RawSeriesChunk { + RawSeriesChunk::try_new( + &schema(), + Buffer::from(key.to_vec()), + metric, + RawTimestamps::try_new(ScalarBuffer::from(times.to_vec())).unwrap(), + SeriesFields::one(Samples::from_float(Float64Array::from(values))), + last, + ) + .unwrap() +} +fn chunk(key: &[u8], times: &[i64], last: bool) -> RawSeriesChunk { + chunk_with( + key, + times, + times.iter().map(|t| Some(*t as f64)).collect(), + Some(metric()), + last, + ) +} + +#[derive(Debug)] +struct RawMock { + schema: Arc, + items: VecDeque>, + binding: RawLabelSource, + ordering: SeriesOrdering, + polls: Arc, + drops: Arc, + /// After this many rows were returned, the next poll reports `Pending` once and wakes the + /// task. `None` keeps the source always ready. + pending_after: Option, + pending_reported: bool, + /// Fails instead of polling, proving a constructor rejected the source before any poll. + panic_on_poll: bool, + emitted: usize, +} +impl Drop for RawMock { + fn drop(&mut self) { + self.drops.fetch_add(1, Ordering::SeqCst); + } +} +impl Stream for RawMock { + type Item = Result; + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + assert!( + !self.panic_on_poll, + "a source rejected before polling must never be polled" + ); + self.polls.fetch_add(1, Ordering::SeqCst); + if !self.pending_reported && self.pending_after == Some(self.emitted) { + self.pending_reported = true; + cx.waker().wake_by_ref(); + return Poll::Pending; + } + let item = self.items.pop_front(); + if item.is_some() { + self.emitted += 1; + } + Poll::Ready(item) + } +} +impl RawSeriesStream for RawMock { + fn schema(&self) -> &Arc { + &self.schema + } + fn label_source(&self) -> &RawLabelSource { + &self.binding + } + fn ordering(&self) -> SeriesOrdering { + self.ordering + } +} +struct DfMock { + schema: SchemaRef, + items: VecDeque>, + polls: Arc, + drops: Arc, + /// After this many batches were returned, the next poll reports `Pending` once and wakes + /// the task. `None` keeps the source always ready. + pending_after: Option, + pending_reported: bool, + /// Fails instead of polling, proving a constructor rejected the source before any poll. + panic_on_poll: bool, + emitted: usize, +} +impl Drop for DfMock { + fn drop(&mut self) { + self.drops.fetch_add(1, Ordering::SeqCst); + } +} +impl Stream for DfMock { + type Item = datafusion_common::Result; + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + assert!( + !self.panic_on_poll, + "a source rejected before polling must never be polled" + ); + self.polls.fetch_add(1, Ordering::SeqCst); + if !self.pending_reported && self.pending_after == Some(self.emitted) { + self.pending_reported = true; + cx.waker().wake_by_ref(); + return Poll::Pending; + } + let item = self.items.pop_front(); + if item.is_some() { + self.emitted += 1; + } + Poll::Ready(item) + } +} +impl RecordBatchStream for DfMock { + fn schema(&self) -> SchemaRef { + self.schema.clone() + } +} +fn raw( + items: Vec, + ordering: SeriesOrdering, +) -> (SendableRawSeriesStream, Arc, Arc) { + raw_items( + items.into_iter().map(Ok).collect(), + schema(), + RawLabelSource::Explicit, + ordering, + None, + ) +} +fn raw_pending( + items: Vec, + ordering: SeriesOrdering, + pending_after: usize, +) -> (SendableRawSeriesStream, Arc, Arc) { + raw_items( + items.into_iter().map(Ok).collect(), + schema(), + RawLabelSource::Explicit, + ordering, + Some(pending_after), + ) +} +fn raw_custom( + items: Vec, + ordering: SeriesOrdering, + binding: RawLabelSource, +) -> (SendableRawSeriesStream, Arc, Arc) { + raw_items( + items.into_iter().map(Ok).collect(), + schema(), + binding, + ordering, + None, + ) +} +fn raw_never_polled( + ordering: SeriesOrdering, +) -> (SendableRawSeriesStream, Arc, Arc) { + let polls = Arc::new(AtomicUsize::new(0)); + let drops = Arc::new(AtomicUsize::new(0)); + ( + Box::pin(RawMock { + schema: schema(), + items: VecDeque::new(), + binding: RawLabelSource::Explicit, + ordering, + polls: polls.clone(), + drops: drops.clone(), + pending_after: None, + pending_reported: false, + panic_on_poll: true, + emitted: 0, + }), + polls, + drops, + ) +} +fn raw_items( + items: Vec>, + schema: Arc, + binding: RawLabelSource, + ordering: SeriesOrdering, + pending_after: Option, +) -> (SendableRawSeriesStream, Arc, Arc) { + let polls = Arc::new(AtomicUsize::new(0)); + let drops = Arc::new(AtomicUsize::new(0)); + ( + Box::pin(RawMock { + schema, + items: items.into(), + binding, + ordering, + polls: polls.clone(), + drops: drops.clone(), + pending_after, + pending_reported: false, + panic_on_poll: false, + emitted: 0, + }), + polls, + drops, + ) +} +fn df_source( + batches: Vec, + declared: SchemaRef, +) -> ( + DfSendableRecordBatchStream, + Arc, + Arc, +) { + df_items(batches.into_iter().map(Ok).collect(), declared, None) +} +fn df_pending( + batches: Vec, + declared: SchemaRef, + pending_after: usize, +) -> ( + DfSendableRecordBatchStream, + Arc, + Arc, +) { + df_items( + batches.into_iter().map(Ok).collect(), + declared, + Some(pending_after), + ) +} +fn df_never_polled( + declared: SchemaRef, +) -> ( + DfSendableRecordBatchStream, + Arc, + Arc, +) { + let polls = Arc::new(AtomicUsize::new(0)); + let drops = Arc::new(AtomicUsize::new(0)); + ( + Box::pin(DfMock { + schema: declared, + items: VecDeque::new(), + polls: polls.clone(), + drops: drops.clone(), + pending_after: None, + pending_reported: false, + panic_on_poll: true, + emitted: 0, + }), + polls, + drops, + ) +} +fn df_items( + items: Vec>, + declared: SchemaRef, + pending_after: Option, +) -> ( + DfSendableRecordBatchStream, + Arc, + Arc, +) { + let polls = Arc::new(AtomicUsize::new(0)); + let drops = Arc::new(AtomicUsize::new(0)); + ( + Box::pin(DfMock { + schema: declared, + items: items.into(), + polls: polls.clone(), + drops: drops.clone(), + pending_after, + pending_reported: false, + panic_on_poll: false, + emitted: 0, + }), + polls, + drops, + ) +} +async fn pack_all(input: SendableRawSeriesStream) -> Vec { + let mut stream = pack_raw_series(input).unwrap(); + let mut out = vec![]; + while let Some(v) = stream.next().await { + out.push(v.unwrap()); + } + out +} + +#[tokio::test] +async fn contiguous_adapter_does_not_certify_nonadjacent_recurrence() { + // A/B/A violates the producer's contiguous-key contract; adapters validate adjacent chunks + // only and preserve the supplied sequence rather than sorting, merging, or dropping rows. + let expected = [ + (b"a".as_slice(), 1, 11.0), + (b"b".as_slice(), 2, 22.0), + (b"a".as_slice(), 3, 33.0), + ]; + let input = expected + .iter() + .map(|(key, time, value)| { + chunk_with(key, &[*time], vec![Some(*value)], Some(metric()), true) + }) + .collect(); + let (source, _, _) = raw(input, SeriesOrdering::Contiguous); + let batches = pack_all(source).await; + assert_eq!(batches.len(), 1, "ready chunks coalesce into one batch"); + assert_eq!(batches[0].num_rows(), 3); + let frame_schema = batches[0].schema(); + let (source, _, _) = df_source(batches, frame_schema); + let mut decoded = unpack_raw_series( + source, + schema(), + RawLabelSource::Explicit, + SeriesOrdering::Contiguous, + ) + .unwrap(); + + for (key, time, value) in expected { + let actual = decoded.next().await.unwrap().unwrap(); + assert_chunk(&actual, key, &[time], &[Some(value)], true); + assert_metric(&actual, None); + } + assert!(decoded.next().await.is_none()); +} +fn assert_metric(chunk: &RawSeriesChunk, host: Option<&str>) { + let actual = chunk + .metric + .as_ref() + .expect("explicit chunks retain metric"); + assert_eq!(actual.name(), Some("cpu")); + assert_eq!(actual.tags().len(), usize::from(host.is_some())); + if let Some(host) = host { + assert_eq!(actual.tags()[0].name.as_ref(), "host"); + assert_eq!(actual.tags()[0].value.as_ref(), host); + } +} + +fn assert_chunk( + actual: &RawSeriesChunk, + key: &[u8], + times: &[i64], + values: &[Option], + last: bool, +) { + assert_eq!(actual.primary_key.as_slice(), key); + assert_eq!( + (0..actual.time.len()) + .map(|i| actual.time.timestamp(i).unwrap()) + .collect::>(), + times + ); + assert_eq!(actual.is_last, last); + let arr = match actual.fields.as_slice()[0].dense().unwrap() { + datatypes::series::samples::DenseSamplesRef::Float(a) => a, + _ => panic!("expected float field"), + }; + assert_eq!(arr.len(), values.len()); + for (i, value) in values.iter().enumerate() { + assert_eq!(arr.is_valid(i), value.is_some()); + if let Some(value) = value { + assert_eq!(arr.value(i).to_bits(), value.to_bits()); + } + } +} + +#[tokio::test] +async fn roundtrip_explicit_unordered_contiguous_preserves_complete_chunk_contents() { + let chunks = vec![ + chunk_with( + b"complete-b", + &[1, 2], + vec![Some(-0.0), None], + Some(metric()), + false, + ), + chunk_with( + b"complete-b", + &[3], + vec![Some(f64::from_bits(0x7ff8_0000_0000_0042))], + Some(metric()), + true, + ), + chunk_with(b"complete-a", &[4], vec![Some(9.25)], Some(metric()), true), + ]; + let (input, polls, _) = raw(chunks, SeriesOrdering::Contiguous); + let mut packed = pack_raw_series(input).unwrap(); + assert_eq!(polls.load(Ordering::SeqCst), 0); + let frame_schema = packed.schema(); + let first = packed.next().await.unwrap().unwrap(); + assert_eq!( + first.num_rows(), + 3, + "three dense chunks coalesce into three frame rows, containing four positions" + ); + assert_eq!( + polls.load(Ordering::SeqCst), + 4, + "the complete batch is flushed once EOF is observed" + ); + assert!(packed.next().await.is_none()); + let batches = vec![first]; + let source: DfSendableRecordBatchStream = Box::pin(DfMock { + schema: frame_schema, + items: batches.into_iter().map(Ok).collect(), + polls: Arc::new(AtomicUsize::new(0)), + drops: Arc::new(AtomicUsize::new(0)), + pending_after: None, + pending_reported: false, + panic_on_poll: false, + emitted: 0, + }); + let mut output = unpack_raw_series( + source, + schema(), + RawLabelSource::Explicit, + SeriesOrdering::Contiguous, + ) + .unwrap(); + assert_chunk( + &output.next().await.unwrap().unwrap(), + b"complete-b", + &[1, 2], + &[Some(-0.0), None], + false, + ); + assert_chunk( + &output.next().await.unwrap().unwrap(), + b"complete-b", + &[3], + &[Some(f64::from_bits(0x7ff8_0000_0000_0042))], + true, + ); + assert_chunk( + &output.next().await.unwrap().unwrap(), + b"complete-a", + &[4], + &[Some(9.25)], + true, + ); + assert!(output.next().await.is_none()); +} + +#[tokio::test] +async fn key_then_time_valid_roundtrip() { + let (input, _, _) = raw( + vec![chunk(b"a", &[1], true), chunk(b"b", &[2], true)], + SeriesOrdering::KeyThenTime, + ); + let batches = pack_all(input).await; + assert_eq!(batches.len(), 1, "cross-chunk rows share one bounded batch"); + assert_eq!(batches[0].num_rows(), 2); + let frame_schema = batches[0].schema(); + let source: DfSendableRecordBatchStream = Box::pin(DfMock { + schema: frame_schema, + items: batches.into_iter().map(Ok).collect(), + polls: Arc::new(AtomicUsize::new(0)), + drops: Arc::new(AtomicUsize::new(0)), + pending_after: None, + pending_reported: false, + panic_on_poll: false, + emitted: 0, + }); + let mut out = unpack_raw_series( + source, + schema(), + RawLabelSource::Explicit, + SeriesOrdering::KeyThenTime, + ) + .unwrap(); + assert_chunk( + &out.next().await.unwrap().unwrap(), + b"a", + &[1], + &[Some(1.0)], + true, + ); + assert_chunk( + &out.next().await.unwrap().unwrap(), + b"b", + &[2], + &[Some(2.0)], + true, + ); + assert!(out.next().await.is_none()); +} + +#[tokio::test] +async fn invalid_sequence_reports_valid_prefix_then_one_terminal_error() { + struct Case { + name: &'static str, + items: Vec, + order: SeriesOrdering, + prefix: Vec<(&'static [u8], i64, bool, Option<&'static str>)>, + } + let cases = vec![ + Case { + name: "final_then_same_key", + items: vec![chunk(b"a", &[1], true), chunk(b"a", &[2], true)], + order: SeriesOrdering::Contiguous, + prefix: vec![(b"a", 1, true, None)], + }, + Case { + name: "unfinished_then_new_key", + items: vec![chunk(b"a", &[1], false), chunk(b"b", &[2], true)], + order: SeriesOrdering::Contiguous, + prefix: vec![(b"a", 1, false, None)], + }, + Case { + name: "key_decrease_after_intervening_keys", + items: vec![ + chunk(b"a", &[1], true), + chunk(b"b", &[2], true), + chunk(b"a", &[3], true), + ], + order: SeriesOrdering::KeyThenTime, + prefix: vec![(b"a", 1, true, None), (b"b", 2, true, None)], + }, + Case { + name: "sorted_key_decrease", + items: vec![chunk(b"b", &[1], true), chunk(b"a", &[2], true)], + order: SeriesOrdering::KeyThenTime, + prefix: vec![(b"b", 1, true, None)], + }, + Case { + name: "timestamp_equal_across_chunks", + items: vec![chunk(b"a", &[2], false), chunk(b"a", &[2], true)], + order: SeriesOrdering::Contiguous, + prefix: vec![(b"a", 2, false, None)], + }, + Case { + name: "timestamp_decreases_across_chunks", + items: vec![chunk(b"a", &[3], false), chunk(b"a", &[2], true)], + order: SeriesOrdering::Contiguous, + prefix: vec![(b"a", 3, false, None)], + }, + Case { + name: "metric_changes_within_key", + items: vec![ + chunk_with( + b"a", + &[1], + vec![Some(1.0)], + Some(labeled_metric("x")), + false, + ), + chunk_with(b"a", &[2], vec![Some(2.0)], Some(labeled_metric("y")), true), + ], + order: SeriesOrdering::Contiguous, + prefix: vec![(b"a", 1, false, Some("x"))], + }, + ]; + for case in cases { + // The valid prefix is flushed by an upstream `Pending`, so the test observes both the + // retained rows and the single terminal error honestly under the batching builder. + let (input, _, _) = raw_pending(case.items.clone(), case.order, case.prefix.len()); + let mut s = pack_raw_series(input).unwrap(); + if !case.prefix.is_empty() { + let batch = s.next().await.unwrap().unwrap(); + assert_eq!(batch.num_rows(), case.prefix.len(), "{}", case.name); + let chunks = datatypes::series::frame::unpack_raw_series_frame( + &schema(), + datatypes::series::frame::RawLabelMode::Explicit, + &batch, + ) + .unwrap(); + assert_eq!(chunks.len(), case.prefix.len(), "{}", case.name); + for (chunk, prefix) in chunks.iter().zip(&case.prefix) { + let (key, time, last, host) = *prefix; + assert_chunk(chunk, key, &[time], &[Some(time as f64)], last); + assert_metric(chunk, host); + } + } + let err = s.next().await.unwrap().unwrap_err(); + assert_eq!( + crate::error::datafusion_status_code::(&err, None), + StatusCode::InvalidArguments, + "{}", + case.name + ); + assert!(s.next().await.is_none(), "{}", case.name); + + // Encode chunks independently, bypassing stream validation, to exercise unpack's + // own cross-frame state machine against precisely the same invalid sequence. + let mut frames = Vec::new(); + for item in &case.items { + frames.extend( + datatypes::series::frame::pack_raw_series_chunk( + &schema(), + datatypes::series::frame::RawLabelMode::Explicit, + item, + ) + .unwrap(), + ); + } + let frame_schema = frames[0].schema(); + let (source, _, _) = df_source(frames, frame_schema); + let mut decoded = + unpack_raw_series(source, schema(), RawLabelSource::Explicit, case.order).unwrap(); + for (key, time, last, host) in &case.prefix { + let chunk = decoded.next().await.unwrap().unwrap(); + assert_chunk(&chunk, key, &[*time], &[Some(*time as f64)], *last); + assert_metric(&chunk, *host); + } + let error = decoded.next().await.unwrap().unwrap_err(); + assert_eq!( + common_error::ext::ErrorExt::status_code(&error), + StatusCode::InvalidArguments + ); + assert!(decoded.next().await.is_none()); + } + // The unpack side of a non-final stream: the row arrives with `is_last` false and the + // missing final key is reported once, at EOF. + let frame = datatypes::series::frame::pack_raw_series_chunk( + &schema(), + datatypes::series::frame::RawLabelMode::Explicit, + &chunk(b"unfinished", &[7], false), + ) + .unwrap(); + let frame_schema = frame[0].schema(); + let (source, _, _) = df_source(frame, frame_schema); + let mut decoded = unpack_raw_series( + source, + schema(), + RawLabelSource::Explicit, + SeriesOrdering::Contiguous, + ) + .unwrap(); + let first = decoded.next().await.unwrap().unwrap(); + assert_chunk(&first, b"unfinished", &[7], &[Some(7.0)], false); + assert_eq!( + common_error::ext::ErrorExt::status_code(&decoded.next().await.unwrap().unwrap_err()), + StatusCode::InvalidArguments + ); + assert!(decoded.next().await.is_none()); +} + +/// EOF before a non-final key is an error: buffered rows are not reported as a partial success. +#[tokio::test] +async fn pack_rejects_nonfinal_eof_without_buffered_partial_success() { + let (input, polls, _) = raw( + vec![chunk_with( + b"unfinished", + &[1, 2], + vec![Some(1.0), Some(2.0)], + Some(metric()), + false, + )], + SeriesOrdering::Contiguous, + ); + let mut packed = pack_raw_series(input).unwrap(); + let error = packed.next().await.unwrap().unwrap_err(); + assert_eq!( + crate::error::datafusion_status_code::(&error, None), + StatusCode::InvalidArguments + ); + assert!( + packed.next().await.is_none(), + "the buffered rows are discarded, never emitted as a successful batch" + ); + assert_eq!( + polls.load(Ordering::SeqCst), + 2, + "the row was buffered, then EOF rejected the missing final key" + ); +} + +fn primary_binding(metadata: store_api::metadata::RegionMetadataRef) -> RawLabelSource { + RawLabelSource::PrimaryKey { + metric_name: Arc::from("cpu"), + physical_metadata: metadata, + logical_tags: Arc::from([BoundTag { + column_id: 0, + name: Arc::from("host"), + }]), + } +} + +#[tokio::test] +async fn primary_key_binding_is_retained_exactly() { + use api::v1::SemanticType; + use datatypes::schema::ColumnSchema; + use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder}; + use store_api::storage::RegionId; + let mut builder = RegionMetadataBuilder::new(RegionId::new(1, 1)); + builder + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "host", + datatypes::data_type::ConcreteDataType::string_datatype(), + false, + ), + semantic_type: SemanticType::Tag, + column_id: 0, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "ts", + datatypes::data_type::ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + semantic_type: SemanticType::Timestamp, + column_id: 1, + }) + .primary_key(vec![0]); + let metadata = Arc::new(builder.build().unwrap()); + let binding = primary_binding(metadata.clone()); + let (input, _, _) = raw_custom( + vec![ + RawSeriesChunk::try_new( + &schema(), + Buffer::from(b"encoded-complete-key".to_vec()), + None, + RawTimestamps::try_new(ScalarBuffer::from(vec![1])).unwrap(), + SeriesFields::one(Samples::from_float(Float64Array::from(vec![Some(3.0)]))), + true, + ) + .unwrap(), + ], + SeriesOrdering::Contiguous, + binding.clone(), + ); + let mut packed = pack_raw_series(input).unwrap(); + let batch = packed.next().await.unwrap().unwrap(); + let source: DfSendableRecordBatchStream = Box::pin(DfMock { + schema: batch.schema(), + items: VecDeque::from([Ok(batch)]), + polls: Arc::new(AtomicUsize::new(0)), + drops: Arc::new(AtomicUsize::new(0)), + pending_after: None, + pending_reported: false, + panic_on_poll: false, + emitted: 0, + }); + let mut out = unpack_raw_series(source, schema(), binding, SeriesOrdering::Contiguous).unwrap(); + let decoded = out.next().await.unwrap().unwrap(); + assert_eq!(decoded.primary_key.as_slice(), b"encoded-complete-key"); + assert!(decoded.metric.is_none()); + let retained = out.label_source(); + match retained { + RawLabelSource::PrimaryKey { + metric_name, + physical_metadata, + logical_tags, + } => { + assert!(Arc::ptr_eq(physical_metadata, &metadata)); + assert_eq!(metric_name.as_ref(), "cpu"); + assert_eq!(logical_tags.as_ref().len(), 1); + assert_eq!(logical_tags[0].column_id, 0); + assert_eq!(logical_tags[0].name.as_ref(), "host"); + } + _ => panic!("binding changed"), + } +} + +#[tokio::test] +async fn explicit_mode_rejects_key_payload_and_primary_mode_rejects_explicit_metric() { + let (input, _, _) = raw( + vec![ + RawSeriesChunk::try_new( + &schema(), + Buffer::from(b"a".to_vec()), + None, + RawTimestamps::try_new(ScalarBuffer::from(vec![1])).unwrap(), + SeriesFields::one(Samples::from_float(Float64Array::from(vec![Some(1.0)]))), + true, + ) + .unwrap(), + ], + SeriesOrdering::Contiguous, + ); + let mut s = pack_raw_series(input).unwrap(); + assert!(s.next().await.unwrap().is_err()); + assert!(s.next().await.is_none()); +} + +#[tokio::test] +async fn primary_key_mode_rejects_explicit_metric_payload() { + let mut builder = + store_api::metadata::RegionMetadataBuilder::new(store_api::storage::RegionId::new(1, 2)); + use api::v1::SemanticType; + use datatypes::schema::ColumnSchema; + use store_api::metadata::ColumnMetadata; + builder + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "host", + datatypes::data_type::ConcreteDataType::string_datatype(), + false, + ), + semantic_type: SemanticType::Tag, + column_id: 0, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "ts", + datatypes::data_type::ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + semantic_type: SemanticType::Timestamp, + column_id: 1, + }) + .primary_key(vec![0]); + let metadata = Arc::new(builder.build().unwrap()); + let (input, _, _) = raw_custom( + vec![chunk(b"a", &[1], true)], + SeriesOrdering::Contiguous, + primary_binding(metadata), + ); + let mut stream = pack_raw_series(input).unwrap(); + assert_eq!( + crate::error::datafusion_status_code::( + &stream.next().await.unwrap().unwrap_err(), + None, + ), + StatusCode::InvalidArguments + ); + assert!(stream.next().await.is_none()); +} + +#[tokio::test] +async fn explicit_metric_must_remain_stable_within_key() { + // The first row is flushed by `Pending` so the retained frame and the later violation are + // both observed under the batching builder. + let (input, polls, _) = raw_pending( + vec![ + chunk_with( + b"a", + &[1], + vec![Some(1.0)], + Some(labeled_metric("x")), + false, + ), + chunk_with(b"a", &[2], vec![Some(2.0)], Some(labeled_metric("y")), true), + ], + SeriesOrdering::Contiguous, + 1, + ); + let mut s = pack_raw_series(input).unwrap(); + let batch = s.next().await.unwrap().unwrap(); + assert_eq!(batch.num_rows(), 1); + assert_eq!( + polls.load(Ordering::SeqCst), + 2, + "the stable row is flushed at Pending before the second chunk is read" + ); + let chunks = datatypes::series::frame::unpack_raw_series_frame( + &schema(), + datatypes::series::frame::RawLabelMode::Explicit, + &batch, + ) + .unwrap(); + assert_chunk(&chunks[0], b"a", &[1], &[Some(1.0)], false); + assert_metric(&chunks[0], Some("x")); + let err = s.next().await.unwrap().unwrap_err(); + assert_eq!( + crate::error::datafusion_status_code::(&err, None), + StatusCode::InvalidArguments + ); + assert!(s.next().await.is_none()); +} + +#[tokio::test] +async fn locally_invalid_public_raw_chunk_is_rejected() { + let mut bad = chunk(b"a", &[1], true); + bad.fields = SeriesFields::one(Samples::from_float(Float64Array::from(vec![ + Some(1.0), + Some(2.0), + ]))); + let polls = Arc::new(AtomicUsize::new(0)); + let source = Box::pin(RawMock { + schema: schema(), + items: VecDeque::from([Ok(bad)]), + binding: RawLabelSource::Explicit, + ordering: SeriesOrdering::Contiguous, + polls, + drops: Arc::new(AtomicUsize::new(0)), + pending_after: None, + pending_reported: false, + panic_on_poll: false, + emitted: 0, + }); + let mut s = pack_raw_series(source).unwrap(); + let err = s.next().await.unwrap().unwrap_err(); + assert_eq!( + crate::error::datafusion_status_code::(&err, None), + StatusCode::InvalidArguments + ); + assert!(s.next().await.is_none()); +} + +#[tokio::test] +async fn empty_streams_keep_self_describing_schema() { + let (input, _, _) = raw(vec![], SeriesOrdering::Contiguous); + let mut packed = pack_raw_series(input).unwrap(); + let frame_schema = packed.schema(); + assert_eq!( + frame_schema.field(0).name(), + "read_key", + "the v1 wire column name is unchanged" + ); + assert!(packed.next().await.is_none()); + // A declared-but-empty stream is still validated against the expected schema before any + // poll, and then ends without producing rows. + let (source, polls, _) = df_source(vec![], frame_schema.clone()); + let mut unpacked = unpack_raw_series( + source, + schema(), + RawLabelSource::Explicit, + SeriesOrdering::Contiguous, + ) + .unwrap(); + assert_eq!(polls.load(Ordering::SeqCst), 0); + assert_eq!(unpacked.schema().fields.len(), 1); + assert_eq!(unpacked.schema().fields[0].name.as_ref(), "value"); + assert!(unpacked.next().await.is_none()); +} + +#[tokio::test] +async fn upstream_error_status_is_preserved_once_and_releases_source() { + for status in [StatusCode::RequestOutdated, StatusCode::Unsupported] { + let polls = Arc::new(AtomicUsize::new(0)); + let drops = Arc::new(AtomicUsize::new(0)); + let e = BoxedError::new(PlainError::new("upstream".into(), status)); + let source = Box::pin(RawMock { + schema: schema(), + items: VecDeque::from([Err(e)]), + binding: RawLabelSource::Explicit, + ordering: SeriesOrdering::Contiguous, + polls: polls.clone(), + drops: drops.clone(), + pending_after: None, + pending_reported: false, + panic_on_poll: false, + emitted: 0, + }); + let mut packed = pack_raw_series(source).unwrap(); + let error = packed.next().await.unwrap().unwrap_err(); + assert_eq!( + crate::error::datafusion_status_code::(&error, None), + status + ); + assert!(packed.next().await.is_none()); + assert_eq!(polls.load(Ordering::SeqCst), 1); + assert_eq!(drops.load(Ordering::SeqCst), 1); + let batch_schema = datatypes::series::frame::raw_series_frame_schema( + &schema(), + datatypes::series::frame::RawLabelMode::Explicit, + ) + .unwrap(); + let plain = PlainError::new("upstream df".into(), status); + let polls = Arc::new(AtomicUsize::new(0)); + let drops = Arc::new(AtomicUsize::new(0)); + let source: DfSendableRecordBatchStream = Box::pin(DfMock { + schema: batch_schema, + items: VecDeque::from([Err(datafusion_common::DataFusionError::External(Box::new( + BoxedError::new(plain), + )))]), + polls: polls.clone(), + drops: drops.clone(), + pending_after: None, + pending_reported: false, + panic_on_poll: false, + emitted: 0, + }); + let mut unpacked = unpack_raw_series( + source, + schema(), + RawLabelSource::Explicit, + SeriesOrdering::Contiguous, + ) + .unwrap(); + let error = unpacked.next().await.unwrap().unwrap_err(); + assert_eq!(common_error::ext::ErrorExt::status_code(&error), status); + assert!(unpacked.next().await.is_none()); + assert_eq!(polls.load(Ordering::SeqCst), 1); + assert_eq!(drops.load(Ordering::SeqCst), 1); + } +} + +#[tokio::test] +async fn dropping_adapter_releases_upstream_and_pending_item_is_resumed() { + let (input, polls, drops) = raw_pending( + vec![chunk(b"a", &[1], true), chunk(b"b", &[2], true)], + SeriesOrdering::Contiguous, + 1, + ); + let mut packed = pack_raw_series(input).unwrap(); + let first = packed.next().await.unwrap().unwrap(); + assert_eq!(first.num_rows(), 1); + assert_eq!( + polls.load(Ordering::SeqCst), + 2, + "the buffered row is flushed at Pending before the next item" + ); + let chunks = datatypes::series::frame::unpack_raw_series_frame( + &schema(), + datatypes::series::frame::RawLabelMode::Explicit, + &first, + ) + .unwrap(); + assert_chunk(&chunks[0], b"a", &[1], &[Some(1.0)], true); + let second = packed.next().await.unwrap().unwrap(); + assert_eq!(second.num_rows(), 1); + assert_eq!( + polls.load(Ordering::SeqCst), + 4, + "Pending is resumed and the remaining row is flushed at EOF" + ); + let chunks = datatypes::series::frame::unpack_raw_series_frame( + &schema(), + datatypes::series::frame::RawLabelMode::Explicit, + &second, + ) + .unwrap(); + assert_chunk(&chunks[0], b"b", &[2], &[Some(2.0)], true); + drop(packed); + assert_eq!(drops.load(Ordering::SeqCst), 1); +} + +#[tokio::test] +async fn unpack_rejects_mismatched_batch_schema_before_polling() { + let wrong = datatypes::arrow::datatypes::Schema::empty().into(); + let (source, polls, drops) = df_source(vec![], wrong); + assert!( + unpack_raw_series( + source, + schema(), + RawLabelSource::Explicit, + SeriesOrdering::Contiguous + ) + .is_err() + ); + assert_eq!(polls.load(Ordering::SeqCst), 0); + assert_eq!(drops.load(Ordering::SeqCst), 1); +} + +#[tokio::test] +async fn invalid_grid_schema_rejected_for_empty_raw_stream() { + let mut invalid = (*schema()).clone(); + invalid.eval_grid = Some(Arc::new( + datatypes::series::time::EvalGrid::try_new( + datatypes::series::time::GridSpec::try_new(0, 10, 1).unwrap(), + ) + .unwrap(), + )); + let polls = Arc::new(AtomicUsize::new(0)); + let source = Box::pin(RawMock { + schema: Arc::new(invalid), + items: VecDeque::new(), + binding: RawLabelSource::Explicit, + ordering: SeriesOrdering::Contiguous, + polls: polls.clone(), + drops: Arc::new(AtomicUsize::new(0)), + pending_after: None, + pending_reported: false, + panic_on_poll: false, + emitted: 0, + }); + assert!(pack_raw_series(source).is_err()); + assert_eq!(polls.load(Ordering::SeqCst), 0); +} + +#[tokio::test] +async fn unpack_batch_sequence_error_is_atomic_for_all_rows() { + let mut batches = Vec::new(); + let frame_schema = datatypes::series::frame::raw_series_frame_schema( + &schema(), + datatypes::series::frame::RawLabelMode::Explicit, + ) + .unwrap(); + for (key, time) in [ + (b"a".as_slice(), 1), + (b"b".as_slice(), 2), + (b"a".as_slice(), 3), + ] { + let (input, _, _) = raw(vec![chunk(key, &[time], true)], SeriesOrdering::KeyThenTime); + batches.extend(pack_all(input).await); + } + let multi = concat_batches(&frame_schema, &batches).unwrap(); + let (source, _, _) = df_source(vec![multi], frame_schema); + let mut unpacked = unpack_raw_series( + source, + schema(), + RawLabelSource::Explicit, + SeriesOrdering::KeyThenTime, + ) + .unwrap(); + assert!(unpacked.next().await.unwrap().is_err()); + assert!(unpacked.next().await.is_none()); +} + +#[tokio::test] +async fn unpack_rejects_late_invalid_frame_row_before_prefix() { + let frame_schema = datatypes::series::frame::raw_series_frame_schema( + &schema(), + datatypes::series::frame::RawLabelMode::Explicit, + ) + .unwrap(); + let first = datatypes::series::frame::pack_raw_series_chunk( + &schema(), + datatypes::series::frame::RawLabelMode::Explicit, + &chunk(b"a", &[1], true), + ) + .unwrap() + .remove(0); + let mut second = datatypes::series::frame::pack_raw_series_chunk( + &schema(), + datatypes::series::frame::RawLabelMode::Explicit, + &chunk(b"b", &[2], true), + ) + .unwrap() + .remove(0); + let mut columns = second.columns().to_vec(); + let time_struct = columns[2] + .as_any() + .downcast_ref::() + .unwrap(); + let lists = time_struct + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let datatypes::arrow::datatypes::DataType::List(list_field) = lists.data_type() else { + panic!("expected list") + }; + let invalid_times = datatypes::arrow::array::Int64Array::from(vec![2, 1]); + let invalid_list = datatypes::arrow::array::ListArray::new( + list_field.clone(), + datatypes::arrow::buffer::OffsetBuffer::new(ScalarBuffer::from(vec![0, 2])), + Arc::new(invalid_times), + None, + ); + columns[2] = Arc::new( + datatypes::arrow::array::StructArray::try_new( + time_struct.fields().clone(), + vec![Arc::new(invalid_list)], + None, + ) + .unwrap(), + ); + second = RecordBatch::try_new(frame_schema.clone(), columns).unwrap(); + let multi = concat_batches(&frame_schema, &[first, second]).unwrap(); + let (source, _, _) = df_source(vec![multi], frame_schema); + let mut unpacked = unpack_raw_series( + source, + schema(), + RawLabelSource::Explicit, + SeriesOrdering::Contiguous, + ) + .unwrap(); + let error = unpacked.next().await.unwrap().unwrap_err(); + assert_eq!( + common_error::ext::ErrorExt::status_code(&error), + StatusCode::InvalidArguments + ); + assert!(unpacked.next().await.is_none()); +} + +#[tokio::test] +async fn frame_application_metadata_is_stable_and_schema_change_rejected() { + let base_schema = datatypes::series::frame::raw_series_frame_schema( + &schema(), + datatypes::series::frame::RawLabelMode::Explicit, + ) + .unwrap(); + let mut application_metadata = base_schema.metadata().clone(); + application_metadata.insert("app".into(), "v1".into()); + let declared = Arc::new( + (*base_schema) + .clone() + .with_metadata(application_metadata.clone()), + ); + let first = datatypes::series::frame::pack_raw_series_chunk( + &schema(), + datatypes::series::frame::RawLabelMode::Explicit, + &chunk(b"a", &[1], true), + ) + .unwrap() + .remove(0) + .with_schema(declared.clone()) + .unwrap(); + let stable_second = datatypes::series::frame::pack_raw_series_chunk( + &schema(), + datatypes::series::frame::RawLabelMode::Explicit, + &chunk(b"b", &[2], true), + ) + .unwrap() + .remove(0) + .with_schema(declared.clone()) + .unwrap(); + let (source, _, _) = df_source(vec![first.clone(), stable_second], declared.clone()); + let mut stable = unpack_raw_series( + source, + schema(), + RawLabelSource::Explicit, + SeriesOrdering::Contiguous, + ) + .unwrap(); + assert_chunk( + &stable.next().await.unwrap().unwrap(), + b"a", + &[1], + &[Some(1.0)], + true, + ); + assert_chunk( + &stable.next().await.unwrap().unwrap(), + b"b", + &[2], + &[Some(2.0)], + true, + ); + assert!(stable.next().await.is_none()); + + for changed_value in [Some("v2"), None] { + let mut changed = application_metadata.clone(); + if let Some(value) = changed_value { + changed.insert("app".into(), value.into()); + } else { + changed.remove("app"); + } + let changed_schema = Arc::new((*base_schema).clone().with_metadata(changed)); + let bad_batch = datatypes::series::frame::pack_raw_series_chunk( + &schema(), + datatypes::series::frame::RawLabelMode::Explicit, + &chunk(b"b", &[2], true), + ) + .unwrap() + .remove(0) + .with_schema(changed_schema) + .unwrap(); + let (source, _, _) = df_source(vec![first.clone(), bad_batch], declared.clone()); + let mut unpacked = unpack_raw_series( + source, + schema(), + RawLabelSource::Explicit, + SeriesOrdering::Contiguous, + ) + .unwrap(); + assert_chunk( + &unpacked.next().await.unwrap().unwrap(), + b"a", + &[1], + &[Some(1.0)], + true, + ); + let error = unpacked.next().await.unwrap().unwrap_err(); + assert_eq!( + common_error::ext::ErrorExt::status_code(&error), + StatusCode::InvalidArguments + ); + assert!(unpacked.next().await.is_none()); + } +} + +struct CountingWake(AtomicUsize); + +impl Wake for CountingWake { + fn wake(self: Arc) { + self.0.fetch_add(1, Ordering::SeqCst); + } + + fn wake_by_ref(self: &Arc) { + self.0.fetch_add(1, Ordering::SeqCst); + } +} + +fn poll_unpacked( + stream: &mut SendableRawSeriesStream, + waker: &Waker, +) -> Poll>> { + stream.as_mut().poll_next(&mut Context::from_waker(waker)) +} + +fn poll_packed( + stream: &mut DfSendableRecordBatchStream, + waker: &Waker, +) -> Poll>> { + stream.as_mut().poll_next(&mut Context::from_waker(waker)) +} + +#[tokio::test] +async fn unpack_bounds_empty_batches_per_poll_and_eventually_yields_row() { + const BUDGET: usize = 32; + let (input, _, _) = raw( + vec![chunk(b"after-empty", &[9], true)], + SeriesOrdering::Contiguous, + ); + let mut packed = pack_raw_series(input).unwrap(); + let row = packed.next().await.unwrap().unwrap(); + let frame_schema = row.schema(); + let mut batches = vec![RecordBatch::new_empty(frame_schema.clone()); 10_000]; + batches.push(row); + let (source, polls, _) = df_source(batches, frame_schema); + let mut unpacked = unpack_raw_series( + source, + schema(), + RawLabelSource::Explicit, + SeriesOrdering::Contiguous, + ) + .unwrap(); + let wake = Arc::new(CountingWake(AtomicUsize::new(0))); + let waker = Waker::from(wake.clone()); + + let before = polls.load(Ordering::SeqCst); + assert!(matches!( + poll_unpacked(&mut unpacked, &waker), + Poll::Pending + )); + assert_eq!(polls.load(Ordering::SeqCst) - before, BUDGET); + assert_eq!(wake.0.load(Ordering::SeqCst), 1); + + let mut yielded = false; + let mut expected_wakes = 1; + for _ in 0..=10_000 / BUDGET { + let before = polls.load(Ordering::SeqCst); + let wakes_before = wake.0.load(Ordering::SeqCst); + match poll_unpacked(&mut unpacked, &waker) { + Poll::Pending => { + assert_eq!(polls.load(Ordering::SeqCst) - before, BUDGET); + expected_wakes += 1; + assert_eq!(wake.0.load(Ordering::SeqCst) - wakes_before, 1); + assert_eq!(wake.0.load(Ordering::SeqCst), expected_wakes); + } + Poll::Ready(Some(Ok(chunk))) => { + assert!(polls.load(Ordering::SeqCst) - before <= BUDGET + 1); + assert_eq!(polls.load(Ordering::SeqCst), 10_001); + assert_eq!(wake.0.load(Ordering::SeqCst), wakes_before); + assert_eq!(wake.0.load(Ordering::SeqCst), 10_000 / BUDGET); + assert_chunk(&chunk, b"after-empty", &[9], &[Some(9.0)], true); + yielded = true; + break; + } + other => panic!("unexpected poll result: {other:?}"), + } + } + assert!(yielded, "bounded polling should eventually yield the row"); + assert!(matches!( + poll_unpacked(&mut unpacked, &waker), + Poll::Ready(None) + )); + assert_eq!(wake.0.load(Ordering::SeqCst), 10_000 / BUDGET); +} + +#[test] +fn unpack_yields_at_empty_batch_budget_before_observing_eof() { + const BUDGET: usize = 32; + let frame_schema = datatypes::series::frame::raw_series_frame_schema( + &schema(), + datatypes::series::frame::RawLabelMode::Explicit, + ) + .unwrap(); + let batches = vec![RecordBatch::new_empty(frame_schema.clone()); BUDGET]; + let (source, polls, _) = df_source(batches, frame_schema); + let mut unpacked = unpack_raw_series( + source, + schema(), + RawLabelSource::Explicit, + SeriesOrdering::Contiguous, + ) + .unwrap(); + let wake = Arc::new(CountingWake(AtomicUsize::new(0))); + let waker = Waker::from(wake.clone()); + + assert!(matches!( + poll_unpacked(&mut unpacked, &waker), + Poll::Pending + )); + assert_eq!(polls.load(Ordering::SeqCst), BUDGET); + assert_eq!(wake.0.load(Ordering::SeqCst), 1); + assert!(matches!( + poll_unpacked(&mut unpacked, &waker), + Poll::Ready(None) + )); + assert_eq!(wake.0.load(Ordering::SeqCst), 1); +} + +#[test] +fn unpack_propagates_upstream_pending_without_extra_wake() { + let frame_schema = datatypes::series::frame::raw_series_frame_schema( + &schema(), + datatypes::series::frame::RawLabelMode::Explicit, + ) + .unwrap(); + let (source, polls, _) = df_pending( + vec![RecordBatch::new_empty(frame_schema.clone())], + frame_schema, + 1, + ); + let mut unpacked = unpack_raw_series( + source, + schema(), + RawLabelSource::Explicit, + SeriesOrdering::Contiguous, + ) + .unwrap(); + let wake = Arc::new(CountingWake(AtomicUsize::new(0))); + let waker = Waker::from(wake.clone()); + + assert!(matches!( + poll_unpacked(&mut unpacked, &waker), + Poll::Pending + )); + assert_eq!(polls.load(Ordering::SeqCst), 2); + assert_eq!(wake.0.load(Ordering::SeqCst), 1, "only upstream wakes"); +} + +#[test] +fn unpack_reports_bad_batch_at_budget_boundary_before_yielding() { + const BUDGET: usize = 32; + let frame_schema = datatypes::series::frame::raw_series_frame_schema( + &schema(), + datatypes::series::frame::RawLabelMode::Explicit, + ) + .unwrap(); + let mut batches = vec![RecordBatch::new_empty(frame_schema.clone()); BUDGET - 1]; + batches.push(RecordBatch::new_empty( + datatypes::arrow::datatypes::Schema::empty().into(), + )); + let (source, polls, _) = df_source(batches, frame_schema); + let mut unpacked = unpack_raw_series( + source, + schema(), + RawLabelSource::Explicit, + SeriesOrdering::Contiguous, + ) + .unwrap(); + let wake = Arc::new(CountingWake(AtomicUsize::new(0))); + let waker = Waker::from(wake.clone()); + + assert!(matches!( + poll_unpacked(&mut unpacked, &waker), + Poll::Ready(Some(Err(_))) + )); + assert_eq!(polls.load(Ordering::SeqCst), BUDGET); + assert_eq!(wake.0.load(Ordering::SeqCst), 0); + assert!(matches!( + poll_unpacked(&mut unpacked, &waker), + Poll::Ready(None) + )); +} + +#[test] +fn unpack_surfaces_upstream_error_after_budget_yield_once() { + const BUDGET: usize = 32; + let frame_schema = datatypes::series::frame::raw_series_frame_schema( + &schema(), + datatypes::series::frame::RawLabelMode::Explicit, + ) + .unwrap(); + let error = datafusion_common::DataFusionError::External(Box::new(BoxedError::new( + PlainError::new("after empty batches".into(), StatusCode::Unsupported), + ))); + let mut items = (0..BUDGET) + .map(|_| Ok(RecordBatch::new_empty(frame_schema.clone()))) + .collect::>(); + items.push(Err(error)); + let polls = Arc::new(AtomicUsize::new(0)); + let source = Box::pin(DfMock { + schema: frame_schema, + items: items.into(), + polls: polls.clone(), + drops: Arc::new(AtomicUsize::new(0)), + pending_after: None, + pending_reported: false, + panic_on_poll: false, + emitted: 0, + }); + let mut unpacked = unpack_raw_series( + source, + schema(), + RawLabelSource::Explicit, + SeriesOrdering::Contiguous, + ) + .unwrap(); + let wake = Arc::new(CountingWake(AtomicUsize::new(0))); + let waker = Waker::from(wake.clone()); + + assert!(matches!( + poll_unpacked(&mut unpacked, &waker), + Poll::Pending + )); + assert_eq!(polls.load(Ordering::SeqCst), BUDGET); + let Poll::Ready(Some(Err(error))) = poll_unpacked(&mut unpacked, &waker) else { + panic!("upstream error should be returned after budget yield"); + }; + assert_eq!( + common_error::ext::ErrorExt::status_code(&error), + StatusCode::Unsupported + ); + assert!(matches!( + poll_unpacked(&mut unpacked, &waker), + Poll::Ready(None) + )); +} + +#[tokio::test] +async fn unpack_drains_decoded_rows_before_polling_another_batch() { + let (input, _, _) = raw( + vec![chunk(b"a", &[1], true), chunk(b"b", &[2], true)], + SeriesOrdering::KeyThenTime, + ); + let batches = pack_all(input).await; + let frame_schema = batches[0].schema(); + let multi = concat_batches(&frame_schema, &batches).unwrap(); + let (source, polls, _) = df_source(vec![multi], frame_schema); + let mut unpacked = unpack_raw_series( + source, + schema(), + RawLabelSource::Explicit, + SeriesOrdering::KeyThenTime, + ) + .unwrap(); + assert_chunk( + &unpacked.next().await.unwrap().unwrap(), + b"a", + &[1], + &[Some(1.0)], + true, + ); + assert_eq!(polls.load(Ordering::SeqCst), 1); + assert_chunk( + &unpacked.next().await.unwrap().unwrap(), + b"b", + &[2], + &[Some(2.0)], + true, + ); + assert_eq!( + polls.load(Ordering::SeqCst), + 1, + "decoded pending row precedes upstream polling" + ); + assert!(unpacked.next().await.is_none()); +} + +#[tokio::test] +async fn unpack_preflights_empty_and_grid_schemas_without_polling() { + let invalid_schemas = [ + SeriesSchema { + fields: Arc::from([]), + eval_grid: None, + }, + SeriesSchema { + fields: schema().fields.clone(), + eval_grid: Some(Arc::new( + datatypes::series::time::EvalGrid::try_new( + datatypes::series::time::GridSpec::try_new(0, 10, 1).unwrap(), + ) + .unwrap(), + )), + }, + ]; + for expected in invalid_schemas { + let frame_schema = datatypes::series::frame::raw_series_frame_schema( + &schema(), + datatypes::series::frame::RawLabelMode::Explicit, + ) + .unwrap(); + let (source, polls, drops) = df_source(vec![], frame_schema); + assert!( + unpack_raw_series( + source, + Arc::new(expected), + RawLabelSource::Explicit, + SeriesOrdering::Contiguous, + ) + .is_err() + ); + assert_eq!(polls.load(Ordering::SeqCst), 0); + assert_eq!(drops.load(Ordering::SeqCst), 1); + } +} + +fn named_fields_schema(names: [&str; 2]) -> Arc { + Arc::new(SeriesSchema { + fields: names + .into_iter() + .map(|name| FieldSpec { + name: Arc::from(name), + }) + .collect(), + eval_grid: None, + }) +} + +fn named_fields_chunk(schema: &SeriesSchema) -> RawSeriesChunk { + RawSeriesChunk::try_new( + schema, + Buffer::from(b"series-key".to_vec()), + Some(metric()), + RawTimestamps::try_new(ScalarBuffer::from(vec![7_i64, 8])).unwrap(), + SeriesFields::try_from_vec(vec![ + Samples::from_float(Float64Array::from(vec![Some(1.0), Some(2.0)])), + Samples::from_float(Float64Array::from(vec![Some(10.0), Some(20.0)])), + ]) + .unwrap(), + true, + ) + .unwrap() +} + +#[tokio::test] +async fn named_frame_fields_roundtrip_and_reject_wrong_bindings_before_polling() { + let expected = named_fields_schema(["cpu_usage", "cpu_idle"]); + let input_chunk = named_fields_chunk(&expected); + let input: SendableRawSeriesStream = Box::pin(RawMock { + schema: expected.clone(), + items: VecDeque::from([Ok(input_chunk)]), + binding: RawLabelSource::Explicit, + ordering: SeriesOrdering::KeyThenTime, + polls: Arc::new(AtomicUsize::new(0)), + drops: Arc::new(AtomicUsize::new(0)), + pending_after: None, + pending_reported: false, + panic_on_poll: false, + emitted: 0, + }); + let batches = pack_all(input).await; + let mut ipc = Vec::new(); + { + let mut writer = StreamWriter::try_new(&mut ipc, &batches[0].schema()).unwrap(); + for batch in &batches { + writer.write(batch).unwrap(); + } + writer.finish().unwrap(); + } + let reader = StreamReader::try_new(Cursor::new(ipc), None).unwrap(); + let frame_schema = reader.schema(); + let ipc_batches = reader.map(Result::unwrap).collect::>(); + let payload_fields = frame_schema.field_with_name("fields").unwrap().data_type(); + let DataType::Struct(payload_fields) = payload_fields else { + panic!("expected struct payload field"); + }; + assert_eq!( + payload_fields + .iter() + .map(|field| field.name()) + .collect::>(), + ["cpu_usage", "cpu_idle"] + ); + + let (source, _, _) = df_source(ipc_batches.clone(), frame_schema.clone()); + let mut decoded = unpack_raw_series( + source, + expected.clone(), + RawLabelSource::Explicit, + SeriesOrdering::KeyThenTime, + ) + .unwrap(); + let actual = decoded.next().await.unwrap().unwrap(); + assert_chunk( + &actual, + b"series-key", + &[7, 8], + &[Some(1.0), Some(2.0)], + true, + ); + assert_eq!(actual.fields.len(), 2); + let second = match actual.fields.as_slice()[1].dense().unwrap() { + DenseSamplesRef::Float(values) => values, + _ => panic!("expected float field"), + }; + assert_eq!(second.null_count(), 0); + assert_eq!(second.values(), &[10.0, 20.0]); + assert_metric(&actual, None); + assert!(decoded.next().await.is_none()); + + for (wrong_names, empty) in [ + (["cpu_idle", "cpu_usage"], false), + (["cpu_usage", "renamed_idle"], false), + (["cpu_idle", "cpu_usage"], true), + (["cpu_usage", "renamed_idle"], true), + ] { + let wrong_schema = named_fields_schema(wrong_names); + let batches = if empty { vec![] } else { ipc_batches.clone() }; + let (source, polls, drops) = df_source(batches, frame_schema.clone()); + assert!( + unpack_raw_series( + source, + wrong_schema, + RawLabelSource::Explicit, + SeriesOrdering::KeyThenTime, + ) + .is_err() + ); + assert_eq!(polls.load(Ordering::SeqCst), 0); + assert_eq!(drops.load(Ordering::SeqCst), 1); + } +} + +fn histogram_table() -> Arc { + let mut builder = HistogramBatchBuilder::new(); + for (sum, start_timestamp) in [(10.5, Some(-123)), (-9.5, Some(456)), (7.25, None)] { + builder + .push(datatypes::series::histogram::HistogramView { + schema: -7, + zero_threshold: 0.125, + sum, + count: 17.25, + zero_count: 2.5, + reset_hint: CounterResetHint::Unrecognized(37), + start_timestamp, + custom_values: &[1.5, 4.5], + positive: BucketSideView { + span_offsets: &[1], + span_lengths: &[2], + buckets: &[3.25, 7.5], + }, + negative: BucketSideView { + span_offsets: &[-2], + span_lengths: &[1], + buckets: &[0.75], + }, + }) + .unwrap(); + } + builder.finish().unwrap() +} + +/// Ready rows are bounded to the default row limit and emitted as soon as that limit is reached, +/// without reading ahead to EOF; the remainder is flushed when EOF is observed. +#[tokio::test] +async fn pack_bounds_ready_rows_at_default_row_limit_before_eof() { + const TOTAL: usize = 300; + const ROW_LIMIT: usize = 256; + let items = (0..TOTAL) + .map(|i| { + RawSeriesChunk::try_new( + &schema(), + Buffer::from(format!("s{i:03}").into_bytes()), + Some(metric()), + RawTimestamps::try_new(ScalarBuffer::from(vec![i as i64])).unwrap(), + SeriesFields::one(Samples::from_float(Float64Array::from(vec![Some( + i as f64 + 0.5, + )]))), + true, + ) + .unwrap() + }) + .collect::>(); + let (input, polls, drops) = raw(items, SeriesOrdering::KeyThenTime); + let mut packed = pack_raw_series(input).unwrap(); + + let first = packed.next().await.unwrap().unwrap(); + assert_eq!(first.num_rows(), ROW_LIMIT); + assert_eq!( + polls.load(Ordering::SeqCst), + ROW_LIMIT, + "the row bound emits a batch without observing EOF or polling a further chunk" + ); + assert_eq!(drops.load(Ordering::SeqCst), 0, "the source is still live"); + + let second = packed.next().await.unwrap().unwrap(); + assert_eq!(second.num_rows(), TOTAL - ROW_LIMIT); + assert_eq!(polls.load(Ordering::SeqCst), TOTAL + 1); + assert_eq!(drops.load(Ordering::SeqCst), 1, "EOF releases the source"); + assert!(packed.next().await.is_none()); + assert_eq!( + polls.load(Ordering::SeqCst), + TOTAL + 1, + "the terminal poll does not revisit the released source" + ); + + // Compare every series' full payload, not only the batch shape. + let mut seen = 0usize; + for batch in [&first, &second] { + let decoded = datatypes::series::frame::unpack_raw_series_frame( + &schema(), + datatypes::series::frame::RawLabelMode::Explicit, + batch, + ) + .unwrap(); + for chunk in decoded { + let key = format!("s{seen:03}"); + assert_chunk( + &chunk, + key.as_bytes(), + &[seen as i64], + &[Some(seen as f64 + 0.5)], + true, + ); + assert_metric(&chunk, None); + seen += 1; + } + } + assert_eq!(seen, TOTAL, "every ready series survives exactly once"); +} + +/// One chunk of `rows` single-sample float segments (so every row is a union-of-segment row), +/// plus a clone of every per-row value array, used to observe that a retained batch stays +/// readable once these source-side owners are dropped. +fn per_row_segment_chunk(rows: usize, key: &[u8]) -> (RawSeriesChunk, Vec) { + let owners = (0..rows) + .map(|i| Float64Array::from(vec![Some(i as f64)])) + .collect::>(); + let samples = owners + .clone() + .into_iter() + .map(Samples::from_float) + .collect::>(); + let chunk = RawSeriesChunk::try_new( + &schema(), + Buffer::from(key.to_vec()), + Some(metric()), + RawTimestamps::try_new(ScalarBuffer::from((0..rows as i64).collect::>())).unwrap(), + SeriesFields::one(Samples::try_from_chunks(samples).unwrap()), + true, + ) + .unwrap(); + (chunk, owners) +} + +/// One chunk with more rows than the limit yields exactly one bounded batch per poll: the first +/// batch needs no second source poll, and the remainder needs only EOF. The stream never builds a +/// chunk-wide list of batches. +#[tokio::test] +async fn pack_splits_one_chunk_into_bounded_batches_without_repolling_the_source() { + const ROWS: usize = 300; + let (chunk, _owners) = per_row_segment_chunk(ROWS, b"single"); + let (input, polls, drops) = raw(vec![chunk], SeriesOrdering::Contiguous); + let mut packed = pack_raw_series(input).unwrap(); + + let first = packed.next().await.unwrap().unwrap(); + assert_eq!(first.num_rows(), 256); + assert_eq!( + polls.load(Ordering::SeqCst), + 1, + "a single source poll fills the bounded first batch" + ); + let second = packed.next().await.unwrap().unwrap(); + assert_eq!(second.num_rows(), ROWS - 256); + assert_eq!( + polls.load(Ordering::SeqCst), + 2, + "the remainder needs only the EOF poll, not another chunk" + ); + assert_eq!(drops.load(Ordering::SeqCst), 1); + assert!(packed.next().await.is_none()); + + let mut seen = 0usize; + for batch in [&first, &second] { + let decoded = datatypes::series::frame::unpack_raw_series_frame( + &schema(), + datatypes::series::frame::RawLabelMode::Explicit, + batch, + ) + .unwrap(); + for chunk in decoded { + assert_chunk( + &chunk, + b"single", + &[seen as i64], + &[Some(seen as f64)], + seen + 1 == ROWS, + ); + seen += 1; + } + } + assert_eq!(seen, ROWS); +} + +/// Dropping the pack stream while its chunk cursor is still live: the upstream source and the +/// remaining rows go away without another poll, and the last emitted batch stays self-contained, +/// readable after every source-side owner is dropped. +#[tokio::test] +async fn dropping_pack_mid_chunk_keeps_last_batch_and_releases_source() { + const ROWS: usize = 300; + let (chunk, owners) = per_row_segment_chunk(ROWS, b"single"); + let (input, polls, drops) = raw(vec![chunk], SeriesOrdering::Contiguous); + let mut packed = pack_raw_series(input).unwrap(); + + let batch = packed.next().await.unwrap().unwrap(); + assert_eq!(batch.num_rows(), 256); + assert_eq!(polls.load(Ordering::SeqCst), 1); + assert_eq!( + drops.load(Ordering::SeqCst), + 0, + "the live cursor keeps its chunk and the source alive" + ); + + drop(packed); + assert_eq!( + drops.load(Ordering::SeqCst), + 1, + "the drop releases upstream" + ); + assert_eq!( + polls.load(Ordering::SeqCst), + 1, + "the drop neither polls upstream nor drains the rest of the chunk" + ); + + // The retained batch owns its payload: it decodes fully after the source-side owners go. + drop(owners); + let decoded = datatypes::series::frame::unpack_raw_series_frame( + &schema(), + datatypes::series::frame::RawLabelMode::Explicit, + &batch, + ) + .unwrap(); + assert_eq!(decoded.len(), 256); + for (index, chunk) in decoded.iter().enumerate() { + // `is_last` stays false: the final row of the chunk is not in this batch. + assert_chunk( + chunk, + b"single", + &[index as i64], + &[Some(index as f64)], + false, + ); + } + assert_eq!(polls.load(Ordering::SeqCst), 1); +} + +/// Two ready chunks of one key, with more rows buffered than `Pending` reports: the partial +/// batch is flushed complete at that `Pending` without EOF, keeping `is_last` false while the +/// series is not final in the source, and the remainder follows at EOF. +#[tokio::test] +async fn pack_flushes_ready_rows_at_pending_then_remainder_at_eof() { + let (input, polls, drops) = raw_pending( + vec![ + chunk_with(b"k", &[1], vec![Some(1.0)], Some(metric()), false), + chunk_with(b"k", &[2], vec![Some(2.0)], Some(metric()), false), + chunk_with(b"k", &[3], vec![Some(3.0)], Some(metric()), true), + ], + SeriesOrdering::Contiguous, + 2, + ); + let mut packed = pack_raw_series(input).unwrap(); + let wake = Arc::new(CountingWake(AtomicUsize::new(0))); + let waker = Waker::from(wake.clone()); + + let first = match poll_packed(&mut packed, &waker) { + Poll::Ready(Some(Ok(batch))) => batch, + other => panic!("expected a flush at Pending, got {other:?}"), + }; + assert_eq!(first.num_rows(), 2, "both ready rows are complete results"); + assert_eq!( + polls.load(Ordering::SeqCst), + 3, + "two ready rows and one Pending poll, without reaching EOF" + ); + assert_eq!(wake.0.load(Ordering::SeqCst), 1, "upstream wakes the task"); + let decoded = datatypes::series::frame::unpack_raw_series_frame( + &schema(), + datatypes::series::frame::RawLabelMode::Explicit, + &first, + ) + .unwrap(); + assert_eq!(decoded.len(), 2); + assert_chunk(&decoded[0], b"k", &[1], &[Some(1.0)], false); + assert_chunk(&decoded[1], b"k", &[2], &[Some(2.0)], false); + + let second = match poll_packed(&mut packed, &waker) { + Poll::Ready(Some(Ok(batch))) => batch, + other => panic!("expected the remainder at EOF, got {other:?}"), + }; + assert_eq!(second.num_rows(), 1); + assert_eq!(polls.load(Ordering::SeqCst), 5); + assert_eq!(drops.load(Ordering::SeqCst), 1); + let decoded = datatypes::series::frame::unpack_raw_series_frame( + &schema(), + datatypes::series::frame::RawLabelMode::Explicit, + &second, + ) + .unwrap(); + assert_chunk(&decoded[0], b"k", &[3], &[Some(3.0)], true); + assert!(matches!( + poll_packed(&mut packed, &waker), + Poll::Ready(None) + )); +} + +/// The byte bound is a work estimate, not a memory cap: a legal row larger than the estimate is +/// emitted whole as its own batch and never split or dropped. +#[tokio::test] +async fn pack_emits_oversized_legal_rows_whole() { + const SAMPLES: usize = 200_000; + let big = |key: &[u8], offset: i64| { + RawSeriesChunk::try_new( + &schema(), + Buffer::from(key.to_vec()), + Some(metric()), + RawTimestamps::try_new(ScalarBuffer::from((0..SAMPLES as i64).collect::>())) + .unwrap(), + SeriesFields::one(Samples::from_float(Float64Array::from( + (0..SAMPLES) + .map(|i| Some(offset as f64 + i as f64)) + .collect::>(), + ))), + true, + ) + .unwrap() + }; + let (input, polls, _) = raw( + vec![big(b"large-a", 0), big(b"large-b", SAMPLES as i64)], + SeriesOrdering::Contiguous, + ); + let mut packed = pack_raw_series(input).unwrap(); + + let first = packed.next().await.unwrap().unwrap(); + assert_eq!(first.num_rows(), 1, "an oversized row is not split"); + assert!( + polls.load(Ordering::SeqCst) <= 2, + "the oversized row is emitted before EOF, at most after one look-ahead poll" + ); + let second = packed.next().await.unwrap().unwrap(); + assert_eq!(second.num_rows(), 1); + assert!(packed.next().await.is_none()); + + for (batch, key, offset) in [ + (&first, b"large-a".as_slice(), 0_i64), + (&second, b"large-b", 200_000), + ] { + let chunks = datatypes::series::frame::unpack_raw_series_frame( + &schema(), + datatypes::series::frame::RawLabelMode::Explicit, + batch, + ) + .unwrap(); + assert_eq!(chunks.len(), 1); + let chunk = &chunks[0]; + assert_eq!(chunk.primary_key.as_slice(), key); + assert!(chunk.is_last); + assert_eq!(chunk.time.len(), SAMPLES); + assert_eq!(chunk.time.timestamp(0), Some(0)); + assert_eq!(chunk.time.timestamp(SAMPLES - 1), Some(SAMPLES as i64 - 1)); + let samples = chunk.fields.get(0).unwrap(); + assert_eq!(samples.len(), SAMPLES); + for index in [0, SAMPLES / 2, SAMPLES - 1] { + match samples.sample(index) { + SampleRef::Float(value) => { + assert_eq!(value.to_bits(), (offset as f64 + index as f64).to_bits()); + } + SampleRef::Missing | SampleRef::Histogram(_) => panic!("expected float samples"), + } + } + } +} + +/// Two chunks with different per-field segment boundaries coalesce into one frame, survive a real +/// Arrow IPC round trip, and decode back to every value, sample bit, row-local histogram index, +/// and finality flag. +#[tokio::test] +async fn multirow_ipc_roundtrip_preserves_cross_chunk_variable_segments() { + let schema = named_fields_schema(["value", "latency"]); + let table = histogram_table(); + let first = RawSeriesChunk::try_new( + &schema, + Buffer::from(b"series-key".to_vec()), + Some(metric()), + RawTimestamps::try_new(ScalarBuffer::from(vec![10_i64, 20, 30])).unwrap(), + SeriesFields::try_from_vec(vec![ + Samples::try_from_chunks(vec![ + Samples::from_float(Float64Array::from(vec![Some(1.0), None])), + Samples::from_float(Float64Array::from(vec![Some(f64::from_bits( + 0x7ff8_0000_0000_0042, + ))])), + ]) + .unwrap(), + Samples::try_from_histogram( + UInt32Array::from(vec![Some(1), None, Some(0)]), + table.clone(), + ) + .unwrap(), + ]) + .unwrap(), + false, + ) + .unwrap(); + let second = RawSeriesChunk::try_new( + &schema, + Buffer::from(b"series-key".to_vec()), + Some(metric()), + RawTimestamps::try_new(ScalarBuffer::from(vec![40_i64, 50])).unwrap(), + SeriesFields::try_from_vec(vec![ + Samples::try_from_chunks(vec![ + Samples::from_float(Float64Array::from(vec![Some(-0.0)])), + Samples::from_float(Float64Array::from(vec![Some(f64::from_bits( + crate::prometheus::PROMETHEUS_STALE_NAN_BITS, + ))])), + ]) + .unwrap(), + Samples::try_from_histogram(UInt32Array::from(vec![Some(0), Some(2)]), table).unwrap(), + ]) + .unwrap(), + true, + ) + .unwrap(); + + let (input, polls, _) = raw_items( + vec![Ok(first), Ok(second)], + schema.clone(), + RawLabelSource::Explicit, + SeriesOrdering::KeyThenTime, + None, + ); + let batches = pack_all(input).await; + assert_eq!(polls.load(Ordering::SeqCst), 3); + assert_eq!(batches.len(), 1, "cross-chunk rows share one bounded batch"); + assert_eq!(batches[0].num_rows(), 4); + + let mut ipc = Vec::new(); + { + let mut writer = StreamWriter::try_new(&mut ipc, &batches[0].schema()).unwrap(); + for batch in &batches { + writer.write(batch).unwrap(); + } + writer.finish().unwrap(); + } + let reader = StreamReader::try_new(Cursor::new(ipc), None).unwrap(); + assert_eq!( + reader.schema(), + batches[0].schema(), + "IPC preserves the v1 layout, metadata, and field order" + ); + let frame_schema = reader.schema(); + let ipc_batches = reader.map(Result::unwrap).collect::>(); + let (source, _, _) = df_source(ipc_batches, frame_schema); + let mut decoded = unpack_raw_series( + source, + schema.clone(), + RawLabelSource::Explicit, + SeriesOrdering::KeyThenTime, + ) + .unwrap(); + + struct Expected { + time: i64, + float_bits: Option, + hist_index: Option, + } + // Table rows: 0 => sum 10.5, 1 => sum -9.5, 2 => sum 7.25. Row-local indices stay + // unchanged, so equal indices in different rows keep the same sum. + let nan_bits = f64::from_bits(0x7ff8_0000_0000_0042).to_bits(); + let rows: [(Vec, bool); 4] = [ + ( + vec![ + Expected { + time: 10, + float_bits: Some(1.0_f64.to_bits()), + hist_index: Some(1), + }, + Expected { + time: 20, + float_bits: None, + hist_index: None, + }, + ], + false, + ), + ( + vec![Expected { + time: 30, + float_bits: Some(nan_bits), + hist_index: Some(0), + }], + false, + ), + ( + vec![Expected { + time: 40, + float_bits: Some((-0.0_f64).to_bits()), + hist_index: Some(0), + }], + false, + ), + ( + vec![Expected { + time: 50, + float_bits: Some(crate::prometheus::PROMETHEUS_STALE_NAN_BITS), + hist_index: Some(2), + }], + true, + ), + ]; + let sums = [10.5_f64, -9.5, 7.25]; + for (expected, last) in rows { + let chunk = decoded.next().await.unwrap().unwrap(); + assert_eq!(chunk.primary_key.as_slice(), b"series-key"); + assert_eq!(chunk.is_last, last); + assert_metric(&chunk, None); + assert_eq!(chunk.time.len(), expected.len()); + let DenseSamplesRef::Histogram(hist) = chunk.fields.get(1).unwrap().dense().unwrap() else { + panic!("expected a dense histogram field") + }; + for (index, expected) in expected.iter().enumerate() { + assert_eq!(chunk.time.timestamp(index), Some(expected.time)); + match chunk.fields.get(0).unwrap().sample(index) { + SampleRef::Float(value) => { + assert_eq!(Some(value.to_bits()), expected.float_bits); + } + SampleRef::Missing => assert_eq!(expected.float_bits, None), + SampleRef::Histogram(_) => panic!("expected a float field"), + } + match expected.hist_index { + Some(index_value) => { + assert!(hist.indices().is_valid(index)); + assert_eq!(hist.indices().value(index), index_value); + let view = hist.view_at(index).expect("histogram sample present"); + assert_eq!(view.sum.to_bits(), sums[index_value as usize].to_bits()); + assert_eq!(view.schema, -7); + assert_eq!(view.reset_hint, CounterResetHint::Unrecognized(37)); + assert_eq!(view.custom_values, [1.5, 4.5]); + } + None => { + assert!(hist.indices().is_null(index), "null index survives IPC"); + assert!(hist.view_at(index).is_none()); + } + } + } + } + assert!(decoded.next().await.is_none()); +} + +/// A `Pending` flush after ready rows and a later upstream error: the buffered row is never +/// reported, the original error surfaces immediately and once, and the source is released. +#[tokio::test] +async fn pack_discards_buffered_rows_and_reports_upstream_error_once() { + let status = StatusCode::Unsupported; + let (input, polls, drops) = raw_items( + vec![ + Ok(chunk(b"a", &[1], true)), + Ok(chunk(b"b", &[2], true)), + Err(BoxedError::new(PlainError::new( + "mid-stream".into(), + status, + ))), + ], + schema(), + RawLabelSource::Explicit, + SeriesOrdering::Contiguous, + Some(1), + ); + let mut packed = pack_raw_series(input).unwrap(); + let wake = Arc::new(CountingWake(AtomicUsize::new(0))); + let waker = Waker::from(wake.clone()); + + let first = match poll_packed(&mut packed, &waker) { + Poll::Ready(Some(Ok(batch))) => batch, + other => panic!("expected the buffered row at Pending, got {other:?}"), + }; + assert_eq!(first.num_rows(), 1); + assert_eq!(polls.load(Ordering::SeqCst), 2); + + let error = match poll_packed(&mut packed, &waker) { + Poll::Ready(Some(Err(error))) => error, + other => panic!("expected the upstream error, got {other:?}"), + }; + assert_eq!( + crate::error::datafusion_status_code::(&error, None), + status + ); + assert_eq!( + polls.load(Ordering::SeqCst), + 4, + "the second row is buffered, then the error is polled" + ); + assert_eq!( + drops.load(Ordering::SeqCst), + 1, + "the error releases the source" + ); + assert!( + matches!(poll_packed(&mut packed, &waker), Poll::Ready(None)), + "no buffered row may follow the error" + ); + assert_eq!( + polls.load(Ordering::SeqCst), + 4, + "the error is reported exactly once" + ); +} + +/// A malformed public chunk that bypasses `RawSeriesChunk::try_new` reports the existing +/// invalid-arguments error instead of panicking, and still releases the source. +#[tokio::test] +async fn public_chunk_with_empty_times_is_rejected_without_panic() { + let empty_times = || RawTimestamps::try_new(ScalarBuffer::from(Vec::::new())).unwrap(); + let mut aligned = chunk(b"a", &[1], true); + aligned.time = empty_times(); + aligned.fields = SeriesFields::one(Samples::from_float(Float64Array::new_null(0))); + let mut mismatched = chunk(b"a", &[1], true); + mismatched.time = empty_times(); + + for (label, bad) in [("empty", aligned), ("mismatched", mismatched)] { + let (input, polls, drops) = raw_items( + vec![Ok(bad)], + schema(), + RawLabelSource::Explicit, + SeriesOrdering::Contiguous, + None, + ); + let mut packed = pack_raw_series(input).unwrap(); + let error = match packed.next().await { + Some(Err(error)) => error, + other => panic!("{label}: expected a rejection, got {other:?}"), + }; + assert_eq!( + crate::error::datafusion_status_code::(&error, None), + StatusCode::InvalidArguments, + "{label}" + ); + assert!(packed.next().await.is_none(), "{label}"); + assert_eq!(polls.load(Ordering::SeqCst), 1, "{label}"); + assert_eq!(drops.load(Ordering::SeqCst), 1, "{label}"); + } +} + +/// `Interleaved` is rejected by the constructor before any source poll, for a pack source and for +/// an empty unpack stream; no adjacent-only fallback is installed. +#[tokio::test] +async fn interleaved_ordering_is_rejected_before_any_poll() { + let (source, polls, drops) = raw_never_polled(SeriesOrdering::Interleaved); + let error = match pack_raw_series(source) { + Ok(_) => panic!("interleaved pack stream must be rejected"), + Err(error) => error, + }; + assert_eq!( + crate::error::datafusion_status_code::(&error, None), + StatusCode::InvalidArguments + ); + assert_eq!( + polls.load(Ordering::SeqCst), + 0, + "a rejected source is never polled" + ); + assert_eq!(drops.load(Ordering::SeqCst), 1); + + let frame_schema = datatypes::series::frame::raw_series_frame_schema( + &schema(), + datatypes::series::frame::RawLabelMode::Explicit, + ) + .unwrap(); + let (source, polls, drops) = df_never_polled(frame_schema); + assert!( + unpack_raw_series( + source, + schema(), + RawLabelSource::Explicit, + SeriesOrdering::Interleaved, + ) + .is_err() + ); + assert_eq!(polls.load(Ordering::SeqCst), 0); + assert_eq!(drops.load(Ordering::SeqCst), 1); +} + +/// The unpack stream rejects a batch whose schema differs from the stream's declared schema, +/// including drift in metadata the codec itself does not interpret and an empty batch that +/// carries the drift; the reserved layout metadata is rejected as well. +#[tokio::test] +async fn unpack_rejects_unrelated_metadata_drift_and_empty_schema_changes() { + let base = datatypes::series::frame::raw_series_frame_schema( + &schema(), + datatypes::series::frame::RawLabelMode::Explicit, + ) + .unwrap(); + let first = datatypes::series::frame::pack_raw_series_chunk( + &schema(), + datatypes::series::frame::RawLabelMode::Explicit, + &chunk(b"a", &[1], true), + ) + .unwrap() + .remove(0); + + let mut unrelated = base.metadata().clone(); + unrelated.insert("app.unrelated".into(), "x".into()); + let mut reserved = base.metadata().clone(); + reserved.insert("greptime.series.layout".into(), "v2".into()); + let mut renamed = base.metadata().clone(); + renamed.insert("app.renamed".into(), "y".into()); + + let cases: [(&str, Arc, bool); 3] = [ + ( + "unrelated metadata added", + Arc::new((*base).clone().with_metadata(unrelated)), + false, + ), + ( + "reserved layout metadata changed", + Arc::new((*base).clone().with_metadata(reserved)), + false, + ), + ( + "unrelated metadata on an empty batch", + Arc::new((*base).clone().with_metadata(renamed)), + true, + ), + ]; + for (label, changed, empty) in cases { + let bad = if empty { + RecordBatch::new_empty(changed) + } else { + RecordBatch::try_new(changed, first.columns().to_vec()).unwrap() + }; + let (source, _, _) = df_source(vec![first.clone(), bad], base.clone()); + let mut unpacked = unpack_raw_series( + source, + schema(), + RawLabelSource::Explicit, + SeriesOrdering::Contiguous, + ) + .unwrap(); + assert_chunk( + &unpacked.next().await.unwrap().unwrap(), + b"a", + &[1], + &[Some(1.0)], + true, + ); + let error = unpacked.next().await.unwrap().unwrap_err(); + assert_eq!( + common_error::ext::ErrorExt::status_code(&error), + StatusCode::InvalidArguments, + "{label}" + ); + assert!(unpacked.next().await.is_none(), "{label}"); + } +} diff --git a/src/common/query/src/series_frame/tests/lifecycle.rs b/src/common/query/src/series_frame/tests/lifecycle.rs new file mode 100644 index 00000000000..b8558c89b03 --- /dev/null +++ b/src/common/query/src/series_frame/tests/lifecycle.rs @@ -0,0 +1,525 @@ +// 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_error::ext::BoxedError; +use datatypes::arrow::array::{Array, Float64Array, UInt32Array}; +use datatypes::arrow::buffer::{Buffer, ScalarBuffer}; +use datatypes::series::SeriesOrdering; +use datatypes::series::histogram::{BucketSideView, CounterResetHint, HistogramBatchBuilder}; +use datatypes::series::raw::{RawSeriesChunk, RawTimestamps}; +use datatypes::series::samples::{ + DenseSamplesRef, FieldSpec, HistogramSamples, SampleRef, Samples, SeriesFields, SeriesSchema, +}; +use datatypes::series::{Label, MetricName}; +use futures_util::{Stream, StreamExt}; +use store_api::region_engine::{RawLabelSource, RawSeriesStream, SendableRawSeriesStream}; + +use super::super::{pack_raw_series, unpack_raw_series}; + +struct RawFixture { + schema: Arc, + binding: RawLabelSource, + chunk: Option, +} + +impl Stream for RawFixture { + type Item = Result; + + fn poll_next(mut self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { + Poll::Ready(self.chunk.take().map(Ok)) + } +} + +impl RawSeriesStream for RawFixture { + fn schema(&self) -> &Arc { + &self.schema + } + + fn label_source(&self) -> &RawLabelSource { + &self.binding + } + + fn ordering(&self) -> SeriesOrdering { + SeriesOrdering::Contiguous + } +} + +fn histogram_batch() -> Arc { + let custom = [1.5, 4.5]; + let positive_offsets = [1]; + let positive_lengths = [2]; + let positive_buckets = [3.25, 7.5]; + let negative_offsets = [-2]; + let negative_lengths = [1]; + let negative_buckets = [0.75]; + let mut builder = HistogramBatchBuilder::new(); + for (sum, start_timestamp) in [(10.5, Some(-123)), (-9.5, Some(456)), (7.25, None)] { + builder + .push(datatypes::series::histogram::HistogramView { + schema: -7, + zero_threshold: 0.125, + sum, + count: 17.25, + zero_count: 2.5, + reset_hint: CounterResetHint::Unrecognized(37), + start_timestamp, + custom_values: &custom, + positive: BucketSideView { + span_offsets: &positive_offsets, + span_lengths: &positive_lengths, + buckets: &positive_buckets, + }, + negative: BucketSideView { + span_offsets: &negative_offsets, + span_lengths: &negative_lengths, + buckets: &negative_buckets, + }, + }) + .unwrap(); + } + builder.finish().unwrap() +} + +fn assert_histogram( + view: datatypes::series::histogram::HistogramView<'_>, + sum: f64, + start: Option, +) { + assert_eq!(view.schema, -7); + assert_eq!(view.zero_threshold, 0.125); + assert_eq!(view.sum.to_bits(), sum.to_bits()); + assert_eq!(view.count, 17.25); + assert_eq!(view.zero_count, 2.5); + assert_eq!(view.reset_hint, CounterResetHint::Unrecognized(37)); + assert_eq!(view.start_timestamp, start); + assert_eq!(view.custom_values, [1.5, 4.5]); + assert_eq!(view.positive.span_offsets, [1]); + assert_eq!(view.positive.span_lengths, [2]); + assert_eq!(view.positive.buckets, [3.25, 7.5]); + assert_eq!(view.negative.span_offsets, [-2]); + assert_eq!(view.negative.span_lengths, [1]); + assert_eq!(view.negative.buckets, [0.75]); +} + +#[tokio::test] +async fn streaming_round_trip_preserves_sliced_multifield_payload_and_owners() { + let schema = Arc::new(SeriesSchema { + fields: vec!["float", "histogram", "mixed"] + .into_iter() + .map(|name| FieldSpec { + name: Arc::from(name), + }) + .collect::>() + .into(), + eval_grid: None, + }); + let time_owner = ScalarBuffer::from(vec![0, 11, 22, 33, 44]); + let times = RawTimestamps::try_new(time_owner.slice(1, 3)).unwrap(); + let float_owner = Float64Array::from(vec![ + Some(90.0), + None, + Some(f64::from_bits(0x7ff8_0000_0000_0042)), + Some(-0.0), + Some(80.0), + ]); + let float_a = float_owner.slice(1, 2); + let float_b = float_owner.slice(3, 1); + let floats = Samples::try_from_chunks(vec![ + Samples::from_float(float_a.clone()), + Samples::from_float(float_b.clone()), + ]) + .unwrap(); + + let table_owner = histogram_batch(); + let table_weak = Arc::downgrade(&table_owner); + let source_table = table_owner.as_arrow().clone(); + let sliced_table = Arc::new( + datatypes::series::histogram::HistogramBatch::try_from_arrow(source_table.slice(1, 2)) + .unwrap(), + ); + let sliced_weak = Arc::downgrade(&sliced_table); + let index_owner = UInt32Array::from(vec![Some(2), Some(0), Some(1), Some(0), Some(1)]); + let indices = index_owner.slice(1, 3); + let hist = Samples::try_from_histogram(indices.clone(), sliced_table.clone()).unwrap(); + let mixed_float_owner = Float64Array::from(vec![Some(1.25), None, None]); + let mixed_index_owner = UInt32Array::from(vec![None, Some(0), None]); + let mixed = Samples::try_from_mixed( + mixed_float_owner.clone(), + HistogramSamples::try_new(mixed_index_owner.clone(), sliced_table.clone()).unwrap(), + ) + .unwrap(); + let chunk = RawSeriesChunk::try_new( + &schema, + Buffer::from(vec![0xaa_u8, 0x10, 0xbb]), + Some( + MetricName::try_new( + Some(Arc::from("cpu")), + vec![Label { + name: Arc::from("host"), + value: Arc::from("a"), + }], + ) + .unwrap(), + ), + times, + SeriesFields::try_from_vec(vec![floats, hist, mixed]).unwrap(), + true, + ) + .unwrap(); + let input_times_ptr = chunk.time.as_buffer().as_ptr(); + let input_float_ptr = float_a.values().as_ptr(); + let input_hist_indices_ptr = indices.values().as_ptr(); + let input_mixed_float_ptr = mixed_float_owner.values().as_ptr(); + let input_mixed_indices_ptr = mixed_index_owner.values().as_ptr(); + + let input: SendableRawSeriesStream = Box::pin(RawFixture { + schema: schema.clone(), + binding: RawLabelSource::Explicit, + chunk: Some(chunk), + }); + let frames = pack_raw_series(input).unwrap(); + let mut decoded = unpack_raw_series( + frames, + schema, + RawLabelSource::Explicit, + SeriesOrdering::Contiguous, + ) + .unwrap(); + let mut output = Vec::new(); + while let Some(item) = decoded.next().await { + output.push(item.unwrap()); + } + assert_eq!( + output.len(), + 2, + "one bounded batch still decodes one chunk per union-of-segment row" + ); + assert_eq!(output[0].time.len(), 2); + assert!( + !output[0].is_last, + "a chunk that stops at the field segment boundary is not final" + ); + assert_eq!(output[1].time.len(), 1); + assert!(output[1].is_last); + assert_eq!( + output[0].fields.get(0).unwrap().segments().len(), + 1, + "the decoded chunk is dense again after the frame round trip" + ); + + // The batch now holds both rows, so its per-row Arrow constituents are concatenated: multi-row + // values are copied rather than aliasing the source slices. Multi-row copying is expected and + // does not change contents, offsets, or owner cleanup; only genuinely single-row frames keep + // the buffers by pointer (see the single-row zero-copy test below). + let DenseSamplesRef::Float(decoded_float) = output[0].fields.get(0).unwrap().dense().unwrap() + else { + panic!("expected float") + }; + assert_ne!( + decoded_float.values().as_ptr(), + input_float_ptr, + "multi-row concatenation copies float values" + ); + assert_eq!(decoded_float.len(), 2); + assert!(!decoded_float.is_valid(0)); + assert_eq!( + decoded_float.value(1).to_bits(), + f64::from_bits(0x7ff8_0000_0000_0042).to_bits() + ); + let DenseSamplesRef::Histogram(decoded_hist) = + output[0].fields.get(1).unwrap().dense().unwrap() + else { + panic!("expected histogram") + }; + assert_ne!( + decoded_hist.indices().values().as_ptr(), + input_hist_indices_ptr, + "multi-row concatenation copies histogram indices" + ); + assert!( + !decoded_hist + .batch() + .as_arrow() + .to_data() + .ptr_eq(&sliced_table.as_arrow().to_data()), + "multi-row concatenation copies the histogram table" + ); + let DenseSamplesRef::Mixed(decoded_mixed) = output[0].fields.get(2).unwrap().dense().unwrap() + else { + panic!("expected mixed") + }; + assert_ne!( + decoded_mixed.floats().values().as_ptr(), + input_mixed_float_ptr + ); + assert_ne!( + decoded_mixed.histograms().indices().values().as_ptr(), + input_mixed_indices_ptr + ); + assert!( + !decoded_mixed + .histograms() + .batch() + .as_arrow() + .to_data() + .ptr_eq(&sliced_table.as_arrow().to_data()), + "multi-row concatenation copies the mixed histogram table" + ); + let first_times = output[0].time.as_buffer(); + assert_ne!( + first_times.as_ptr(), + input_times_ptr, + "multi-row concatenation copies timestamps" + ); + assert_eq!(&first_times[..], &[11_i64, 22]); + + drop(decoded); + drop(time_owner); + drop(float_owner); + drop(float_a); + drop(float_b); + drop(index_owner); + drop(indices); + drop(mixed_float_owner); + drop(mixed_index_owner); + drop(source_table); + drop(table_owner); + drop(sliced_table); + assert!(table_weak.upgrade().is_none()); + assert!(sliced_weak.upgrade().is_none()); + + let expected_histograms = [(-9.5, Some(456)), (7.25, None), (-9.5, Some(456))]; + let mut times_seen = Vec::new(); + let mut float_seen = Vec::new(); + let mut histogram_seen = Vec::new(); + let mut mixed_seen = Vec::new(); + for chunk in &output { + assert_eq!(chunk.primary_key.as_slice(), [0xaa, 0x10, 0xbb]); + assert_eq!(chunk.metric.as_ref().unwrap().name(), Some("cpu")); + assert_eq!(chunk.metric.as_ref().unwrap().get("host"), Some("a")); + // Only the row that ends the chunk in the source stream is final. + assert_eq!(chunk.is_last, times_seen.len() + chunk.time.len() == 3); + for i in 0..chunk.time.len() { + let position = times_seen.len(); + times_seen.push(chunk.time.timestamp(i).unwrap()); + match chunk.fields.get(0).unwrap().sample(i) { + SampleRef::Float(value) => float_seen.push(Some(value.to_bits())), + SampleRef::Missing => float_seen.push(None), + SampleRef::Histogram(_) => panic!("expected float field"), + } + let SampleRef::Histogram(value) = chunk.fields.get(1).unwrap().sample(i) else { + panic!("expected histogram") + }; + let (sum, start) = expected_histograms[position]; + assert_histogram(value, sum, start); + histogram_seen.push(value.sum); + match (position, chunk.fields.get(2).unwrap().sample(i)) { + (0, SampleRef::Float(value)) => mixed_seen.push(Some(value.to_bits())), + (1, SampleRef::Histogram(value)) => { + assert_histogram(value, -9.5, Some(456)); + mixed_seen.push(None); + } + (2, SampleRef::Missing) => mixed_seen.push(None), + _ => panic!("mixed payload variant did not match expected position {position}"), + } + } + } + assert_eq!(times_seen, [11, 22, 33]); + assert_eq!( + float_seen, + [ + None, + Some(f64::from_bits(0x7ff8_0000_0000_0042).to_bits()), + Some((-0.0_f64).to_bits()) + ] + ); + assert_eq!(histogram_seen, [-9.5, 7.25, -9.5]); + assert_eq!(mixed_seen, [Some(1.25_f64.to_bits()), None, None]); +} + +/// One row of one sample, with every array a nonzero-offset slice of a larger owner: the decoded +/// chunk keeps the source buffers by pointer, and the values stay readable and bit-exact after +/// every source owner is dropped. +#[tokio::test] +async fn single_row_sliced_payload_keeps_zero_copy_buffers_and_releases_owners() { + let schema = Arc::new(SeriesSchema { + fields: vec!["float", "histogram", "mixed"] + .into_iter() + .map(|name| FieldSpec { + name: Arc::from(name), + }) + .collect::>() + .into(), + eval_grid: None, + }); + let metric = MetricName::try_new( + Some(Arc::from("cpu")), + vec![Label { + name: Arc::from("host"), + value: Arc::from("a"), + }], + ) + .unwrap(); + + let time_owner = ScalarBuffer::from(vec![0_i64, 11, 22]); + let times = RawTimestamps::try_new(time_owner.slice(1, 1)).unwrap(); + let float_owner = Float64Array::from(vec![ + Some(90.0), + Some(f64::from_bits(0x7ff8_0000_0000_0042)), + ]); + let float = float_owner.slice(1, 1); + let table_owner = histogram_batch(); + let table_weak = Arc::downgrade(&table_owner); + let source_table = table_owner.as_arrow().clone(); + let sliced_table = Arc::new( + datatypes::series::histogram::HistogramBatch::try_from_arrow(source_table.slice(1, 1)) + .unwrap(), + ); + let sliced_weak = Arc::downgrade(&sliced_table); + let index_owner = UInt32Array::from(vec![Some(3), Some(0)]); + let indices = index_owner.slice(1, 1); + let hist = Samples::try_from_histogram(indices.clone(), sliced_table.clone()).unwrap(); + let mixed_float_owner = Float64Array::from(vec![Some(0.0), Some(-0.0)]); + let mixed_floats = mixed_float_owner.slice(1, 1); + let mixed_index_owner = UInt32Array::from(vec![Some(0), None]); + let mixed_indices = mixed_index_owner.slice(1, 1); + let mixed = Samples::try_from_mixed( + mixed_floats.clone(), + HistogramSamples::try_new(mixed_indices.clone(), sliced_table.clone()).unwrap(), + ) + .unwrap(); + let chunk = RawSeriesChunk::try_new( + &schema, + Buffer::from(vec![0xaa_u8, 0x10, 0xbb]), + Some(metric), + times, + SeriesFields::try_from_vec(vec![Samples::from_float(float.clone()), hist, mixed]).unwrap(), + true, + ) + .unwrap(); + + let input_times_ptr = chunk.time.as_buffer().as_ptr(); + let input_float_ptr = float.values().as_ptr(); + let input_hist_indices_ptr = indices.values().as_ptr(); + let input_mixed_float_ptr = mixed_floats.values().as_ptr(); + let input_mixed_indices_ptr = mixed_indices.values().as_ptr(); + + let input: SendableRawSeriesStream = Box::pin(RawFixture { + schema: schema.clone(), + binding: RawLabelSource::Explicit, + chunk: Some(chunk), + }); + let frames = pack_raw_series(input).unwrap(); + let mut decoded = unpack_raw_series( + frames, + schema, + RawLabelSource::Explicit, + SeriesOrdering::Contiguous, + ) + .unwrap(); + let mut output = Vec::new(); + while let Some(item) = decoded.next().await { + output.push(item.unwrap()); + } + assert_eq!(output.len(), 1, "one sample packs as one row"); + assert_eq!(output[0].time.len(), 1); + assert!(output[0].is_last); + + let DenseSamplesRef::Float(decoded_float) = output[0].fields.get(0).unwrap().dense().unwrap() + else { + panic!("expected float") + }; + assert_eq!(decoded_float.values().as_ptr(), input_float_ptr); + let DenseSamplesRef::Histogram(decoded_hist) = + output[0].fields.get(1).unwrap().dense().unwrap() + else { + panic!("expected histogram") + }; + assert_eq!( + decoded_hist.indices().values().as_ptr(), + input_hist_indices_ptr + ); + assert!( + decoded_hist + .batch() + .as_arrow() + .to_data() + .ptr_eq(&sliced_table.as_arrow().to_data()), + "a single row keeps the sliced histogram table by pointer" + ); + let DenseSamplesRef::Mixed(decoded_mixed) = output[0].fields.get(2).unwrap().dense().unwrap() + else { + panic!("expected mixed") + }; + assert_eq!( + decoded_mixed.floats().values().as_ptr(), + input_mixed_float_ptr + ); + assert_eq!( + decoded_mixed.histograms().indices().values().as_ptr(), + input_mixed_indices_ptr + ); + assert!( + decoded_mixed + .histograms() + .batch() + .as_arrow() + .to_data() + .ptr_eq(&sliced_table.as_arrow().to_data()) + ); + assert_eq!(output[0].time.as_buffer().as_ptr(), input_times_ptr); + + // The decoded chunk owns the slices, so the source owners can go while the values remain + // valid and bit-exact. + drop(decoded); + drop(time_owner); + drop(float_owner); + drop(float); + drop(index_owner); + drop(indices); + drop(mixed_float_owner); + drop(mixed_floats); + drop(mixed_index_owner); + drop(mixed_indices); + drop(source_table); + drop(table_owner); + drop(sliced_table); + assert!(table_weak.upgrade().is_none()); + assert!(sliced_weak.upgrade().is_none()); + + let chunk = &output[0]; + assert_eq!(chunk.primary_key.as_slice(), [0xaa, 0x10, 0xbb]); + assert_eq!(chunk.metric.as_ref().unwrap().name(), Some("cpu")); + assert_eq!(chunk.metric.as_ref().unwrap().get("host"), Some("a")); + assert_eq!(chunk.time.timestamp(0), Some(11)); + match chunk.fields.get(0).unwrap().sample(0) { + SampleRef::Float(value) => assert_eq!( + value.to_bits(), + f64::from_bits(0x7ff8_0000_0000_0042).to_bits() + ), + SampleRef::Missing | SampleRef::Histogram(_) => panic!("expected float field"), + } + let SampleRef::Histogram(value) = chunk.fields.get(1).unwrap().sample(0) else { + panic!("expected histogram") + }; + assert_histogram(value, -9.5, Some(456)); + match chunk.fields.get(2).unwrap().sample(0) { + SampleRef::Float(value) => assert_eq!(value.to_bits(), (-0.0_f64).to_bits()), + SampleRef::Missing | SampleRef::Histogram(_) => panic!("expected mixed float"), + } +} diff --git a/src/datatypes/src/series.rs b/src/datatypes/src/series.rs index 860358eff20..0de711e2da3 100644 --- a/src/datatypes/src/series.rs +++ b/src/datatypes/src/series.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +pub mod frame; pub mod histogram; pub mod raw; pub mod samples; diff --git a/src/datatypes/src/series/frame.rs b/src/datatypes/src/series/frame.rs new file mode 100644 index 00000000000..da6ff38d1ae --- /dev/null +++ b/src/datatypes/src/series/frame.rs @@ -0,0 +1,1030 @@ +// 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. + +//! Finite Arrow batches for raw series chunks. Field names preserve the logical series schema in order. +//! +//! Packing accumulates bounded multi-row batches: [`RawSeriesChunkCursor`] walks one validated +//! chunk and [`RawSeriesFrameBuilder`] emits rows directly as Arrow arrays once a row or estimated +//! byte bound is reached. Decoding reuses one prepared schema per stream through +//! [`RawSeriesFrameDecoder`]. Column layout, metadata, and wire strings stay on the frozen `v1` +//! contract. + +#[cfg(test)] +mod tests; + +use std::sync::Arc; + +use arrow::array::{ + Array, ArrayRef, BinaryArray, BooleanArray, Float64Array, Int64Array, ListArray, StringArray, + StructArray, UInt32Array, +}; +use arrow::buffer::{Buffer, NullBuffer, OffsetBuffer, ScalarBuffer}; +use arrow::compute::concat; +use arrow::datatypes::{DataType, Field, Fields, Schema, SchemaRef}; +use arrow::record_batch::{RecordBatch, RecordBatchOptions}; + +use crate::error::{InvalidSeriesSnafu, Result}; +use crate::series::histogram::{HistogramBatch, fields as histogram_fields}; +use crate::series::raw::{RawSeriesChunk, RawTimestamps}; +use crate::series::samples::{ + DenseSamplesRef, HistogramSamples, Samples, SeriesFields, SeriesSchema, +}; +use crate::series::{Label, MetricName, invalid_series}; + +const LAYOUT: &str = "greptime.series.layout"; +const CONTENT: &str = "greptime.series.content"; +const LABELS: &str = "greptime.series.labels"; +const TIME: &str = "greptime.series.time"; + +/// Whether a raw frame carries only the complete key or explicit metric labels as well. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum RawLabelMode { + /// Labels remain encoded solely in `read_key`. + Key, + /// Complete metric name and ordinary labels are present in the frame. + Explicit, +} + +/// Builds the typed schema for a raw frame, preserving logical field names and order. +pub fn raw_series_frame_schema(schema: &SeriesSchema, mode: RawLabelMode) -> Result { + if schema.eval_grid.is_some() { + return invalid_series("raw series schema cannot have an evaluation grid"); + } + if schema.fields.is_empty() { + return invalid_series("raw series schema must have fields"); + } + let list = |ty, nullable| DataType::List(Arc::new(Field::new("item", ty, nullable))); + let time = DataType::Struct( + vec![Field::new( + "timestamps", + list(DataType::Int64, false), + false, + )] + .into(), + ); + let hist = DataType::Struct( + vec![ + Field::new("indices", list(DataType::UInt32, true), false), + Field::new( + "table", + list(DataType::Struct(histogram_fields()), false), + false, + ), + ] + .into(), + ); + let payload = DataType::Struct( + vec![ + Field::new("float_values", list(DataType::Float64, true), true), + Field::new("histogram", hist, true), + ] + .into(), + ); + let fields: Fields = schema + .fields + .iter() + .map(|field| Field::new(field.name.as_ref(), payload.clone(), false)) + .collect(); + let mut columns = vec![ + Field::new("read_key", DataType::Binary, false), + Field::new("is_last", DataType::Boolean, false), + Field::new("time", time, false), + Field::new("fields", DataType::Struct(fields), false), + ]; + if mode == RawLabelMode::Explicit { + columns.push(Field::new("name", DataType::Utf8, true)); + columns.push(Field::new( + "labels", + list( + DataType::Struct( + vec![ + Field::new("name", DataType::Utf8, false), + Field::new("value", DataType::Utf8, false), + ] + .into(), + ), + false, + ), + false, + )); + } + let metadata = [ + (LAYOUT.to_string(), "v1".to_string()), + (CONTENT.to_string(), "raw".to_string()), + ( + LABELS.to_string(), + match mode { + RawLabelMode::Key => "key", + RawLabelMode::Explicit => "explicit", + } + .to_string(), + ), + (TIME.to_string(), "timestamps".to_string()), + ] + .into_iter() + .collect(); + Ok(Arc::new(Schema::new_with_metadata(columns, metadata))) +} + +/// Bounded multi-row builder for finite raw frame batches. +/// +/// Rows accumulate until the row limit or the estimated byte limit is reached; the next +/// [`RawSeriesChunkCursor::append_to`] call then emits them as one batch. `max_bytes` is stored as +/// the effective group cap `min(requested, i32::MAX)`, and each row carries a measure that +/// dominates every cumulative Arrow offset domain, so any admitted group can be encoded. The +/// measure is an internal work estimate, not an exact IPC size and not a resident-memory cap: one +/// individually representable row that exceeds the requested cap is still accepted and emitted +/// alone, because that slice cannot be split into smaller legal rows. +pub struct RawSeriesFrameBuilder { + schema: SeriesSchema, + frame_schema: SchemaRef, + mode: RawLabelMode, + max_rows: usize, + max_bytes: usize, + rows: Vec, + bytes: usize, +} + +impl RawSeriesFrameBuilder { + /// Default maximum rows in one packed batch when no explicit bound is needed. + pub const DEFAULT_MAX_ROWS: usize = 256; + /// Default maximum estimated encoded bytes in one packed batch when no explicit bound is + /// needed. + pub const DEFAULT_MAX_BYTES: usize = 1024 * 1024; + + /// Creates an empty builder for the given schema, label mode, and nonzero bounds. + /// + /// The frame schema is built once here and reused for every emitted batch. + pub fn try_new( + schema: &SeriesSchema, + mode: RawLabelMode, + max_rows: usize, + max_bytes: usize, + ) -> Result { + if max_rows == 0 || max_bytes == 0 { + return invalid_series("raw series frame limits must be nonzero"); + } + Ok(Self { + schema: schema.clone(), + frame_schema: raw_series_frame_schema(schema, mode)?, + mode, + max_rows, + max_bytes: max_bytes.min(i32::MAX as usize), + rows: Vec::new(), + bytes: 0, + }) + } + + /// The frame schema of every batch this builder emits. + pub fn schema(&self) -> &SchemaRef { + &self.frame_schema + } + + /// Emits the accumulated rows as one batch, or `None` when nothing is buffered. + pub fn flush(&mut self) -> Result> { + if self.rows.is_empty() { + return Ok(None); + } + let batch = encode_rows(&self.frame_schema, self.mode, &self.rows)?; + self.rows.clear(); + self.bytes = 0; + Ok(Some(batch)) + } + + /// Whether the buffered rows reached either bound. + fn is_full(&self) -> bool { + self.rows.len() >= self.max_rows || self.bytes >= self.max_bytes + } + + /// Whether `row` may join the buffered rows without crossing either bound. + /// + /// An empty builder accepts any individually measured row; a group total that would overflow + /// or exceed the cap flushes the buffered rows before this row is consumed, never dropping or + /// splitting the row. + fn accepts(&self, row: &RawFrameRow) -> bool { + self.rows.is_empty() + || (self.bytes < self.max_bytes + && self.rows.len() < self.max_rows + && self + .bytes + .checked_add(row.bytes) + .is_some_and(|total| total <= self.max_bytes)) + } + + /// Buffers one individually measured row. + fn push(&mut self, row: RawFrameRow) -> Result<()> { + self.bytes = add_checked(self.bytes, row.bytes)?; + self.rows.push(row); + Ok(()) + } +} + +/// One validated chunk walked segment by segment into a bounded builder. +pub struct RawSeriesChunkCursor { + schema: SeriesSchema, + mode: RawLabelMode, + chunk: RawSeriesChunk, + segments: Vec, + offsets: Vec, + pos: usize, + pending: Option, +} + +impl RawSeriesChunkCursor { + /// Validates one public chunk's shape and label mode once, taking ownership of it. + pub fn try_new( + schema: &SeriesSchema, + mode: RawLabelMode, + chunk: RawSeriesChunk, + ) -> Result { + let RawSeriesChunk { + primary_key, + metric, + time, + fields, + is_last, + } = chunk; + let chunk = RawSeriesChunk::try_new(schema, primary_key, metric, time, fields, is_last)?; + match (mode, &chunk.metric) { + (RawLabelMode::Key, None) | (RawLabelMode::Explicit, Some(_)) => {} + _ => return invalid_series("raw chunk label mode does not match frame mode"), + } + let segments = vec![0usize; chunk.fields.len()]; + let offsets = vec![0usize; chunk.fields.len()]; + Ok(Self { + schema: schema.clone(), + mode, + chunk, + segments, + offsets, + pos: 0, + pending: None, + }) + } + + /// Appends rows to the builder until one batch is ready or this chunk is exhausted. + /// + /// `None` means the cursor is exhausted; calling again returns `None` without touching the + /// builder. The single pending row is measured once by [`Self::next_row`] and retried after the + /// buffered rows are emitted. + pub fn append_to( + &mut self, + builder: &mut RawSeriesFrameBuilder, + ) -> Result> { + self.compatible(builder)?; + loop { + let row = match self.pending.take() { + Some(row) => row, + None => match self.next_row()? { + Some(row) => { + self.pending = Some(row); + continue; + } + None => return Ok(None), + }, + }; + if !builder.accepts(&row) { + // Reinstate the row before flushing so a flush error cannot consume it. + self.pending = Some(row); + let batch = builder + .flush()? + .ok_or_else(|| err("raw frame builder rejected a row without buffered rows"))?; + return Ok(Some(batch)); + } + builder.push(row)?; + if builder.is_full() + && let Some(batch) = builder.flush()? + { + return Ok(Some(batch)); + } + } + } + + /// The cursor and builder must describe the same fields and label mode. + fn compatible(&self, builder: &RawSeriesFrameBuilder) -> Result<()> { + let same_fields = self.schema.fields.len() == builder.schema.fields.len() + && self + .schema + .fields + .iter() + .zip(builder.schema.fields.iter()) + .all(|(left, right)| left.name == right.name); + if self.mode != builder.mode + || self.schema.eval_grid.is_some() != builder.schema.eval_grid.is_some() + || !same_fields + { + return invalid_series("raw frame cursor and builder schema mismatch"); + } + Ok(()) + } + + /// Builds and measures the next union-of-segment row, or `None` when the chunk is exhausted. + /// + /// The chunk fields are walked three times so no per-row temporary vectors are allocated; the + /// offsets advance only after the row measure succeeded. + fn next_row(&mut self) -> Result> { + if self.pos >= self.chunk.time.len() { + return Ok(None); + } + let time_axis = self.chunk.time.as_buffer(); + let mut take = time_axis.len() - self.pos; + for (index, field) in self.chunk.fields.iter().enumerate() { + let part = segment(field, self.segments[index])?; + take = take.min(part.len().saturating_sub(self.offsets[index])); + } + if take == 0 { + return invalid_series("empty sample segment"); + } + let mut samples = Vec::with_capacity(self.chunk.fields.len()); + for (index, field) in self.chunk.fields.iter().enumerate() { + let part = segment(field, self.segments[index])?; + samples.push(part.slice(self.offsets[index], take)?); + } + let mut row = RawFrameRow { + primary_key: self.chunk.primary_key.clone(), + metric: self.chunk.metric.clone(), + is_last: self.chunk.is_last && self.pos + take == self.chunk.time.len(), + times: time_axis.slice(self.pos, take), + samples, + bytes: 0, + }; + row.bytes = measure_row(&row)?; + self.pos += take; + for (index, field) in self.chunk.fields.iter().enumerate() { + let length = segment(field, self.segments[index])?.len(); + self.offsets[index] += take; + if self.offsets[index] == length { + self.segments[index] += 1; + self.offsets[index] = 0; + } + } + Ok(Some(row)) + } +} + +/// The next dense segment of one field, or an error when the walk left the segment list. +fn segment(field: &Samples, index: usize) -> Result<&Samples> { + field + .segments() + .get(index) + .ok_or_else(|| err("empty sample segment")) +} + +/// Packs one raw chunk, retaining the `Vec` compatibility signature. +/// +/// The chunk packs into bounded multi-row batches; callers that need streaming behavior should +/// drive [`RawSeriesChunkCursor`] and [`RawSeriesFrameBuilder`] directly instead of collecting. +pub fn pack_raw_series_chunk( + schema: &SeriesSchema, + mode: RawLabelMode, + chunk: &RawSeriesChunk, +) -> Result> { + let mut builder = RawSeriesFrameBuilder::try_new( + schema, + mode, + RawSeriesFrameBuilder::DEFAULT_MAX_ROWS, + RawSeriesFrameBuilder::DEFAULT_MAX_BYTES, + )?; + let mut cursor = RawSeriesChunkCursor::try_new(schema, mode, chunk.clone())?; + let mut batches = Vec::new(); + while let Some(batch) = cursor.append_to(&mut builder)? { + batches.push(batch); + } + if let Some(batch) = builder.flush()? { + batches.push(batch); + } + Ok(batches) +} + +/// One pending output row: a complete key with one dense sample slice per field and its checked +/// encoded-size measure. +struct RawFrameRow { + primary_key: Buffer, + metric: Option, + is_last: bool, + times: ScalarBuffer, + samples: Vec, + bytes: usize, +} + +/// Per-row encoded-size measure, computed once per row when the cursor produces it. +/// +/// The measure dominates every cumulative offset domain the encoder writes: binary/UTF-8 child +/// bytes, list element counts, bitmap bytes, and the visible extent of each histogram table list. +/// A group total of `min(max_bytes, i32::MAX)` therefore guarantees that all Arrow offsets and +/// lengths stay representable, without tracking a separate offset counter. This is a bound used +/// for batching work, not the exact encoded size. +fn measure_row(row: &RawFrameRow) -> Result { + checked_offset(row.primary_key.len())?; + checked_offset(row.times.len())?; + let mut bytes = add_checked(ROW_CHARGE, row.primary_key.len())?; + bytes = add_checked(bytes, mul_checked(row.times.len(), 8)?)?; + for samples in &row.samples { + bytes = add_checked(bytes, FIELD_CHARGE)?; + match samples + .dense() + .ok_or_else(|| err("sample segment is not dense"))? + { + DenseSamplesRef::Float(values) => { + bytes = add_checked(bytes, float_charge(values.len())?)? + } + DenseSamplesRef::Histogram(values) => { + bytes = add_checked(bytes, histogram_charge(values)?)? + } + DenseSamplesRef::Mixed(values) => { + bytes = add_checked(bytes, float_charge(values.floats().len())?)?; + bytes = add_checked(bytes, histogram_charge(values.histograms())?)?; + } + } + } + if let Some(metric) = &row.metric { + if let Some(name) = metric.name() { + checked_offset(name.len())?; + bytes = add_checked(bytes, name.len())?; + } + checked_offset(metric.tags().len())?; + bytes = add_checked(bytes, mul_checked(metric.tags().len(), 16)?)?; + let mut names = 0usize; + let mut values = 0usize; + for label in metric.tags() { + names = add_checked(names, label.name.len())?; + values = add_checked(values, label.value.len())?; + } + checked_offset(names)?; + checked_offset(values)?; + bytes = add_checked(bytes, names)?; + bytes = add_checked(bytes, values)?; + } + Ok(bytes) +} + +/// Float values plus their validity bitmap. +fn float_charge(len: usize) -> Result { + checked_offset(len)?; + add_checked(mul_checked(len, 8)?, bitmap_bytes(len)) +} + +/// Reference indices plus their bitmap and the full referenced histogram table. +fn histogram_charge(values: &HistogramSamples) -> Result { + let indices = values.indices(); + checked_offset(indices.len())?; + let mut bytes = add_checked(mul_checked(indices.len(), 4)?, bitmap_bytes(indices.len()))?; + let table = values.batch().as_arrow(); + checked_offset(table.len())?; + bytes = add_checked( + bytes, + add_checked( + HISTOGRAM_TABLE_CHARGE, + mul_checked(table.len(), HISTOGRAM_ROW_CHARGE)?, + )?, + )?; + for (position, width) in HISTOGRAM_LIST_WIDTHS.iter().enumerate() { + let list = table + .column(7 + position) + .as_any() + .downcast_ref::() + .ok_or_else(|| err("histogram table schema mismatch"))?; + // `value_offsets` is already a view of the visible slice; the retained child array may be + // wider than the offsets span, so charge the visible extent rather than its length. + let offsets = list.value_offsets(); + let first = offsets + .first() + .copied() + .ok_or_else(|| err("histogram table schema mismatch"))?; + let last = offsets + .get(list.len()) + .copied() + .ok_or_else(|| err("histogram table schema mismatch"))?; + let span = last + .checked_sub(first) + .and_then(|span| usize::try_from(span).ok()) + .ok_or_else(|| err("histogram table offsets are invalid"))?; + bytes = add_checked(bytes, mul_checked(span, *width)?)?; + } + Ok(bytes) +} + +/// Encodes buffered rows as one batch with direct N-row Arrow array construction. +/// +/// Row-local histogram indices stay exactly as encoded: physical tables are concatenated under +/// cumulative row list offsets, so no index is remapped. Each histogram-bearing row carries its +/// complete table, including unreferenced entries; many segments referencing one large table can +/// therefore repeat that table in the encoded batch. A single constituent child array is reused +/// by clone (sharing its buffers); multiple constituents are concatenated, which copies. +fn encode_rows( + schema: &SchemaRef, + mode: RawLabelMode, + rows: &[RawFrameRow], +) -> Result { + let count = rows.len(); + let DataType::Struct(time_fields) = schema.field(2).data_type() else { + return invalid_series("time schema mismatch"); + }; + let DataType::List(time_item) = time_fields[0].data_type() else { + return invalid_series("time schema mismatch"); + }; + let DataType::Struct(field_types) = schema.field(3).data_type() else { + return invalid_series("fields schema mismatch"); + }; + let mut time_lengths = Vec::with_capacity(count); + let mut time_parts: Vec = Vec::with_capacity(count); + for row in rows { + time_lengths.push(row.times.len()); + time_parts.push(Arc::new(Int64Array::new(row.times.clone(), None))); + } + let time_child = materialize_parts(&time_parts, || { + Ok(Arc::new(Int64Array::new_null(0)) as ArrayRef) + })?; + let times = list_array(&time_lengths, time_child, time_item.as_ref().clone(), None)?; + let time = struct_array(time_fields.clone(), vec![Arc::new(times)])?; + + let mut empty_table: Option = None; + let mut field_arrays = Vec::with_capacity(field_types.len()); + for (index, field) in field_types.iter().enumerate() { + let DataType::Struct(payload) = field.data_type() else { + return invalid_series("payload schema mismatch"); + }; + let DataType::List(float_item) = payload[0].data_type() else { + return invalid_series("float values schema mismatch"); + }; + let DataType::Struct(histogram_fields) = payload[1].data_type() else { + return invalid_series("histogram schema mismatch"); + }; + let DataType::List(index_item) = histogram_fields[0].data_type() else { + return invalid_series("histogram index schema mismatch"); + }; + let DataType::List(table_item) = histogram_fields[1].data_type() else { + return invalid_series("histogram table schema mismatch"); + }; + let DataType::Struct(table_fields) = table_item.data_type() else { + return invalid_series("histogram table schema mismatch"); + }; + let mut float_lengths = Vec::with_capacity(count); + let mut float_parts: Vec = Vec::new(); + let mut float_valid = Vec::with_capacity(count); + let mut index_lengths = Vec::with_capacity(count); + let mut index_parts: Vec = Vec::new(); + let mut table_lengths = Vec::with_capacity(count); + let mut table_parts: Vec = Vec::new(); + let mut histogram_valid = Vec::with_capacity(count); + for row in rows { + let dense = row + .samples + .get(index) + .ok_or_else(|| err("raw frame row field count mismatch"))? + .dense() + .ok_or_else(|| err("sample segment is not dense"))?; + match dense { + DenseSamplesRef::Float(values) => { + float_lengths.push(values.len()); + float_parts.push(Arc::new(values.clone())); + float_valid.push(true); + index_lengths.push(0); + table_lengths.push(0); + histogram_valid.push(false); + } + DenseSamplesRef::Histogram(values) => { + float_lengths.push(0); + float_valid.push(false); + index_lengths.push(values.len()); + index_parts.push(Arc::new(values.indices().clone())); + table_lengths.push(values.batch().len()); + table_parts.push(Arc::new(values.batch().as_arrow().clone())); + histogram_valid.push(true); + } + DenseSamplesRef::Mixed(values) => { + float_lengths.push(values.floats().len()); + float_parts.push(Arc::new(values.floats().clone())); + float_valid.push(true); + index_lengths.push(values.histograms().len()); + index_parts.push(Arc::new(values.histograms().indices().clone())); + table_lengths.push(values.histograms().batch().len()); + table_parts.push(Arc::new(values.histograms().batch().as_arrow().clone())); + histogram_valid.push(true); + } + } + } + let float_child = materialize_parts(&float_parts, || { + Ok(Arc::new(Float64Array::new_null(0)) as ArrayRef) + })?; + let floats = list_array( + &float_lengths, + float_child, + float_item.as_ref().clone(), + validity(&float_valid), + )?; + let index_child = materialize_parts(&index_parts, || { + Ok(Arc::new(UInt32Array::new_null(0)) as ArrayRef) + })?; + let indices = list_array( + &index_lengths, + index_child, + index_item.as_ref().clone(), + None, + )?; + // A batch without histogram tables shares one empty schema-typed table, built once here + // instead of per float-only row. + let table_child = if table_parts.is_empty() { + match &empty_table { + Some(table) => table.clone(), + None => { + let table = empty_histogram_table(table_fields)?; + empty_table = Some(table.clone()); + table + } + } + } else { + materialize_parts(&table_parts, || Err(err("missing histogram table values")))? + }; + let tables = list_array( + &table_lengths, + table_child, + table_item.as_ref().clone(), + None, + )?; + let histogram = arrow(StructArray::try_new( + histogram_fields.clone(), + vec![Arc::new(indices), Arc::new(tables)], + validity(&histogram_valid), + ))?; + field_arrays.push(Arc::new(struct_array( + payload.clone(), + vec![Arc::new(floats), Arc::new(histogram)], + )?) as ArrayRef); + } + let fields = struct_array(field_types.clone(), field_arrays)?; + + let mut columns: Vec = vec![ + Arc::new(binary_array(rows)?), + Arc::new(BooleanArray::from( + rows.iter().map(|row| row.is_last).collect::>(), + )), + Arc::new(time), + Arc::new(fields), + ]; + if mode == RawLabelMode::Explicit { + columns.push(Arc::new( + rows.iter() + .map(|row| row.metric.as_ref().and_then(MetricName::name)) + .collect::(), + )); + let DataType::List(label_item) = schema.field(5).data_type() else { + return invalid_series("labels schema mismatch"); + }; + let DataType::Struct(label_fields) = label_item.data_type() else { + return invalid_series("labels schema mismatch"); + }; + let mut label_lengths = Vec::with_capacity(count); + let mut tag_names: Vec<&str> = Vec::new(); + let mut tag_values: Vec<&str> = Vec::new(); + for row in rows { + let metric = row + .metric + .as_ref() + .ok_or_else(|| err("explicit labels missing"))?; + label_lengths.push(metric.tags().len()); + for label in metric.tags() { + tag_names.push(label.name.as_ref()); + tag_values.push(label.value.as_ref()); + } + } + let labels = struct_array( + label_fields.clone(), + vec![ + Arc::new(StringArray::from_iter_values(tag_names)), + Arc::new(StringArray::from_iter_values(tag_values)), + ], + )?; + columns.push(Arc::new(list_array( + &label_lengths, + Arc::new(labels), + label_item.as_ref().clone(), + None, + )?)); + } + + arrow(RecordBatch::try_new_with_options( + schema.clone(), + columns, + &RecordBatchOptions::new().with_row_count(Some(count)), + )) +} + +/// Row-level binary keys; keys are copied exactly like the previous single-row encoder. +fn binary_array(rows: &[RawFrameRow]) -> Result { + let mut offsets = Vec::with_capacity(rows.len() + 1); + let mut total = 0usize; + offsets.push(0); + for row in rows { + total = add_checked(total, row.primary_key.len())?; + offsets.push(checked_offset(total)?); + } + let mut values = Vec::with_capacity(total); + for row in rows { + values.extend_from_slice(row.primary_key.as_slice()); + } + arrow(BinaryArray::try_new( + OffsetBuffer::new(ScalarBuffer::from(offsets)), + Buffer::from_vec(values), + None, + )) +} + +/// Per-row charge for key/time/nullable overhead shared by every row. +const ROW_CHARGE: usize = 64; +/// Per-field charge for the payload struct, nested list and validity overhead. +const FIELD_CHARGE: usize = 64; +/// Conservative per-table-row charge: the fixed fourteen-column histogram schema has seven +/// primitive columns plus seven list columns whose item widths are [`HISTOGRAM_LIST_WIDTHS`]. +const HISTOGRAM_ROW_CHARGE: usize = 128; +/// Fixed per-histogram charge for the seven list columns' offset and validity overhead. +const HISTOGRAM_TABLE_CHARGE: usize = 28; +/// Item widths of histogram table list columns `7..=13`, in that order. +const HISTOGRAM_LIST_WIDTHS: [usize; 7] = [8, 4, 4, 8, 4, 4, 8]; + +fn list_array( + lengths: &[usize], + values: ArrayRef, + item: Field, + nulls: Option, +) -> Result { + arrow(ListArray::try_new( + Arc::new(item), + offset_buffer(lengths)?, + values, + nulls, + )) +} + +fn offset_buffer(lengths: &[usize]) -> Result> { + let mut offsets = Vec::with_capacity(lengths.len() + 1); + let mut total = 0usize; + offsets.push(0); + for length in lengths { + total = add_checked(total, *length)?; + offsets.push(checked_offset(total)?); + } + Ok(OffsetBuffer::new(ScalarBuffer::from(offsets))) +} + +/// Materializes one list child: a single constituent is cloned to keep its buffers shared. +fn materialize_parts( + parts: &[ArrayRef], + empty: impl FnOnce() -> Result, +) -> Result { + match parts { + [] => empty(), + [only] => Ok(only.clone()), + _ => arrow(concat( + &parts.iter().map(|array| &**array).collect::>(), + )), + } +} + +fn empty_histogram_table(fields: &Fields) -> Result { + let columns = fields + .iter() + .map(|field| arrow::array::new_null_array(field.data_type(), 0)) + .collect(); + arrow(StructArray::try_new(fields.clone(), columns, None)) + .map(|table| Arc::new(table) as ArrayRef) +} + +fn validity(valid: &[bool]) -> Option { + if valid.iter().all(|valid| *valid) { + None + } else { + Some(NullBuffer::from(valid.to_vec())) + } +} + +fn bitmap_bytes(len: usize) -> usize { + len.div_ceil(8) +} + +fn add_checked(total: usize, amount: usize) -> Result { + total + .checked_add(amount) + .ok_or_else(|| err("raw series frame offset overflow")) +} + +fn mul_checked(len: usize, width: usize) -> Result { + len.checked_mul(width) + .ok_or_else(|| err("raw series frame offset overflow")) +} + +/// Reuses one prepared frame schema while decoding all batches of a stream. +pub struct RawSeriesFrameDecoder { + schema: SeriesSchema, + mode: RawLabelMode, + expected: SchemaRef, +} + +impl RawSeriesFrameDecoder { + /// Builds the expected frame schema once from the series schema and label mode. + pub fn try_new(schema: &SeriesSchema, mode: RawLabelMode) -> Result { + Ok(Self { + schema: schema.clone(), + mode, + expected: raw_series_frame_schema(schema, mode)?, + }) + } + + /// Decodes all rows of one batch, checking its reserved metadata, declared fields, and rows. + pub fn decode(&self, batch: &RecordBatch) -> Result> { + decode_raw_series_frame(&self.schema, self.mode, &self.expected, batch) + } +} + +/// Decodes all rows of a raw frame; any malformed row makes the entire call fail. +pub fn unpack_raw_series_frame( + schema: &SeriesSchema, + mode: RawLabelMode, + batch: &RecordBatch, +) -> Result> { + RawSeriesFrameDecoder::try_new(schema, mode)?.decode(batch) +} + +fn decode_raw_series_frame( + schema: &SeriesSchema, + mode: RawLabelMode, + expected: &SchemaRef, + batch: &RecordBatch, +) -> Result> { + let actual = batch.schema(); + for key in [LAYOUT, CONTENT, LABELS, TIME] { + if actual.metadata().get(key) != expected.metadata().get(key) { + return invalid_series("raw frame metadata mismatch"); + } + } + if actual + .metadata() + .keys() + .any(|key| key == "greptime.series.grid" || key.starts_with("greptime.series.grid.")) + || actual.fields() != expected.fields() + { + return invalid_series("raw frame schema mismatch"); + } + let cols = batch.columns(); + let keys = down::(&cols[0], "read_key")?; + let lasts = down::(&cols[1], "is_last")?; + let time_struct = down::(&cols[2], "time")?; + let times = down::(time_struct.column(0), "timestamps")?; + let fs = down::(&cols[3], "fields")?; + let (names, labels) = if mode == RawLabelMode::Explicit { + ( + Some(down::(&cols[4], "name")?), + Some(down::(&cols[5], "labels")?), + ) + } else { + (None, None) + }; + let mut out = Vec::with_capacity(batch.num_rows()); + for row in 0..batch.num_rows() { + if keys.is_null(row) + || lasts.is_null(row) + || time_struct.is_null(row) + || times.is_null(row) + || fs.is_null(row) + { + return invalid_series("null required raw frame parent"); + } + let timestamp_values = times.value(row); + let tv = down::(×tamp_values, "timestamp values")?; + if tv.null_count() != 0 { + return invalid_series("null timestamp"); + } + let time = tv.values().clone(); + let metric = match (names, labels) { + (None, None) => None, + (Some(ns), Some(ls)) => { + if ls.is_null(row) { + return invalid_series("null labels list"); + } + let label_values = ls.value(row); + let label_struct = down::(&label_values, "labels")?; + if label_struct.null_count() != 0 { + return invalid_series("null label struct element"); + } + let ln = down::(label_struct.column(0), "label names")?; + let lv = down::(label_struct.column(1), "label values")?; + if ln.null_count() != 0 || lv.null_count() != 0 { + return invalid_series("null label element"); + } + let tags = (0..ln.len()) + .map(|i| Label { + name: Arc::from(ln.value(i)), + value: Arc::from(lv.value(i)), + }) + .collect(); + Some(MetricName::try_new( + (!ns.is_null(row)).then(|| Arc::from(ns.value(row))), + tags, + )?) + } + _ => return invalid_series("incomplete explicit label columns"), + }; + let mut decoded = Vec::with_capacity(fs.num_columns()); + for col in fs.columns() { + let payload = down::(col, "field payload")?; + if payload.is_null(row) { + return invalid_series("null field payload"); + } + let floats_list = down::(payload.column(0), "float values")?; + let hist = down::(payload.column(1), "histogram payload")?; + let has_float = !floats_list.is_null(row); + let has_hist = !hist.is_null(row); + if !has_float && !has_hist { + return invalid_series("field has neither float nor histogram payload"); + } + let floats = if has_float { + down::(&floats_list.value(row), "float values")?.clone() + } else { + Float64Array::new_null(0) + }; + if has_float && floats.len() != time.len() { + return invalid_series("float payload length mismatch"); + } + let samples = if !has_hist { + Samples::from_float(floats) + } else { + let inds = down::(hist.column(0), "histogram indices")?; + let tables = down::(hist.column(1), "histogram table")?; + if inds.is_null(row) || tables.is_null(row) { + return invalid_series("null histogram indices or table"); + } + let indices = down::(&inds.value(row), "histogram indices")?.clone(); + if indices.len() != time.len() { + return invalid_series("histogram index length mismatch"); + } + let table = down::(&tables.value(row), "histogram table")?.clone(); + let batch = Arc::new(HistogramBatch::try_from_arrow(table)?); + if has_float { + Samples::try_from_mixed(floats, HistogramSamples::try_new(indices, batch)?)? + } else { + Samples::try_from_histogram(indices, batch)? + } + }; + decoded.push(samples); + } + let fields = SeriesFields::try_from_vec(decoded)?; + let key = binary_slice(keys, row)?; + let raw_time = RawTimestamps::try_new(time)?; + out.push(RawSeriesChunk::try_new( + schema, + key, + metric, + raw_time, + fields, + lasts.value(row), + )?); + } + Ok(out) +} + +fn binary_slice(array: &BinaryArray, row: usize) -> Result { + let offsets = array.value_offsets(); + let start = usize::try_from(offsets[row]).map_err(|_| err("negative binary offset"))?; + let end = usize::try_from(offsets[row + 1]).map_err(|_| err("negative binary offset"))?; + let len = end + .checked_sub(start) + .ok_or_else(|| err("invalid binary offsets"))?; + Ok(array.values().slice_with_length(start, len)) +} + +fn checked_offset(len: usize) -> Result { + i32::try_from(len).map_err(|_| err("Arrow offset limit exceeded")) +} + +fn struct_array(fields: Fields, arrays: Vec) -> Result { + arrow(StructArray::try_new(fields, arrays, None)) +} + +fn down<'a, T: Array + 'static>(array: &'a ArrayRef, label: &str) -> Result<&'a T> { + array + .as_any() + .downcast_ref() + .ok_or_else(|| err(format!("{label} type mismatch"))) +} + +fn err(reason: impl Into) -> crate::error::Error { + InvalidSeriesSnafu { + reason: reason.into(), + } + .build() +} + +fn arrow(result: std::result::Result) -> Result { + result.map_err(|e| err(e.to_string())) +} diff --git a/src/datatypes/src/series/frame/tests.rs b/src/datatypes/src/series/frame/tests.rs new file mode 100644 index 00000000000..976a433ee27 --- /dev/null +++ b/src/datatypes/src/series/frame/tests.rs @@ -0,0 +1,3060 @@ +// 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 arrow::array::{ + Array, BinaryArray, BooleanArray, Float64Array, Int64Array, ListArray, StringArray, + StructArray, UInt32Array, +}; +use arrow::buffer::{Buffer, OffsetBuffer, ScalarBuffer}; +use arrow::compute::concat_batches; +use arrow::datatypes::{DataType, Field, Schema}; +use arrow::ipc::{reader::StreamReader, writer::StreamWriter}; + +use super::*; +use crate::series::histogram::{BucketSideView, CounterResetHint, HistogramBatchBuilder}; +use crate::series::samples::{ + DenseSamplesRef, FieldSpec, SampleRef, Samples, SeriesFields, SeriesSchema, +}; +use crate::series::{Label, MetricName}; + +fn test_schema(names: &[&str]) -> SeriesSchema { + SeriesSchema { + fields: names + .iter() + .map(|name| FieldSpec { + name: Arc::from(*name), + }) + .collect::>() + .into(), + eval_grid: None, + } +} + +fn metric(name: Option<&str>, labels: &[(&str, &str)]) -> MetricName { + MetricName::try_new( + name.map(Arc::from), + labels + .iter() + .map(|(name, value)| Label { + name: Arc::from(*name), + value: Arc::from(*value), + }) + .collect(), + ) + .unwrap() +} + +fn float_chunk( + schema: &SeriesSchema, + key: &[u8], + metric: Option, + times: Vec, + values: Vec>, + is_last: bool, +) -> RawSeriesChunk { + RawSeriesChunk::try_new( + schema, + Buffer::from(key.to_vec()), + metric, + RawTimestamps::try_new(times.into()).unwrap(), + SeriesFields::one(Samples::from_float(Float64Array::from(values))), + is_last, + ) + .unwrap() +} + +fn histogram_batch() -> Arc { + histogram_batch_with_sum(-9.5) +} + +fn histogram_batch_with_sum(sum: f64) -> Arc { + let custom = [1.5, 4.5]; + let positive_offsets = [1]; + let positive_lengths = [2]; + let positive_buckets = [3.25, 7.5]; + let negative_offsets = [-2]; + let negative_lengths = [1]; + let negative_buckets = [0.75]; + let view = |sum, start_timestamp| crate::series::histogram::HistogramView { + schema: -7, + zero_threshold: 0.125, + sum, + count: 17.25, + zero_count: 2.5, + reset_hint: CounterResetHint::Unrecognized(37), + start_timestamp, + custom_values: &custom, + positive: BucketSideView { + span_offsets: &positive_offsets, + span_lengths: &positive_lengths, + buckets: &positive_buckets, + }, + negative: BucketSideView { + span_offsets: &negative_offsets, + span_lengths: &negative_lengths, + buckets: &negative_buckets, + }, + }; + let mut builder = HistogramBatchBuilder::new(); + builder.push(view(sum, Some(-123))).unwrap(); + builder.push(view(sum + 100.0, None)).unwrap(); + builder.finish().unwrap() +} + +fn hist_schema_frame_field(schema: &SeriesSchema) -> DataType { + raw_series_frame_schema(schema, RawLabelMode::Key) + .unwrap() + .field(3) + .data_type() + .clone() +} + +fn with_timestamp_values( + batch: &arrow::record_batch::RecordBatch, + values: Vec, +) -> arrow::record_batch::RecordBatch { + let time = batch + .column(2) + .as_any() + .downcast_ref::() + .unwrap(); + let timestamp_item = match batch.schema().field(2).data_type() { + DataType::Struct(fields) => match fields[0].data_type() { + DataType::List(item) => item.clone(), + _ => panic!("timestamp field is not a List"), + }, + _ => panic!("time field is not a Struct"), + }; + let list = ListArray::try_new( + timestamp_item, + OffsetBuffer::new(ScalarBuffer::from(vec![0_i32, values.len() as i32])), + Arc::new(Int64Array::from(values)), + None, + ) + .unwrap(); + let time = StructArray::try_new(time.fields().clone(), vec![Arc::new(list)], None).unwrap(); + let mut columns = batch.columns().to_vec(); + columns[2] = Arc::new(time); + arrow::record_batch::RecordBatch::try_new(batch.schema(), columns).unwrap() +} + +fn frame_field_payload(batch: &arrow::record_batch::RecordBatch, field: usize) -> &StructArray { + let fields = batch + .column(3) + .as_any() + .downcast_ref::() + .unwrap(); + fields + .column(field) + .as_any() + .downcast_ref::() + .unwrap() +} + +fn frame_histogram_table(batch: &arrow::record_batch::RecordBatch, field: usize) -> StructArray { + let histogram = frame_field_payload(batch, field) + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + let tables = histogram + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + tables + .value(0) + .as_any() + .downcast_ref::() + .unwrap() + .clone() +} + +fn assert_histogram(histogram: crate::series::histogram::HistogramView<'_>, sum: f64) { + assert_histogram_with_start(histogram, sum, Some(-123)); +} + +fn assert_histogram_with_start( + histogram: crate::series::histogram::HistogramView<'_>, + sum: f64, + start_timestamp: Option, +) { + assert_eq!(histogram.schema, -7); + assert_eq!(histogram.zero_threshold.to_bits(), 0.125_f64.to_bits()); + assert_eq!(histogram.sum.to_bits(), sum.to_bits()); + assert_eq!(histogram.count.to_bits(), 17.25_f64.to_bits()); + assert_eq!(histogram.zero_count.to_bits(), 2.5_f64.to_bits()); + assert_eq!(histogram.reset_hint, CounterResetHint::Unrecognized(37)); + assert_eq!(histogram.start_timestamp, start_timestamp); + assert_eq!( + histogram + .custom_values + .iter() + .map(|v| v.to_bits()) + .collect::>(), + [1.5_f64.to_bits(), 4.5_f64.to_bits()] + ); + assert_eq!(histogram.positive.span_offsets, [1]); + assert_eq!(histogram.positive.span_lengths, [2]); + assert_eq!( + histogram + .positive + .buckets + .iter() + .map(|v| v.to_bits()) + .collect::>(), + [3.25_f64.to_bits(), 7.5_f64.to_bits()] + ); + assert_eq!(histogram.negative.span_offsets, [-2]); + assert_eq!(histogram.negative.span_lengths, [1]); + assert_eq!( + histogram + .negative + .buckets + .iter() + .map(|v| v.to_bits()) + .collect::>(), + [0.75_f64.to_bits()] + ); +} + +fn assert_unreferenced_histogram(histogram: crate::series::histogram::HistogramView<'_>) { + assert_eq!(histogram.schema, -7); + assert_histogram_with_start(histogram, 90.5, None); +} + +fn assert_float_bits(actual: &Samples, expected: &[Option]) { + assert_eq!(actual.len(), expected.len()); + for (index, expected) in expected.iter().enumerate() { + match (actual.sample(index), expected) { + (SampleRef::Missing, None) => {} + (SampleRef::Float(actual), Some(expected)) => { + assert_eq!(actual.to_bits(), expected.to_bits(), "position {index}"); + } + _ => panic!("unexpected sample at position {index}"), + } + } +} + +/// Packs one owned chunk through the public cursor/builder pair, mirroring the explicit +/// single-chunk wrapper: every batch the walk produces, then the final flush. +fn pack_chunk_with_budget( + schema: &SeriesSchema, + mode: RawLabelMode, + chunk: &RawSeriesChunk, + max_rows: usize, + max_bytes: usize, +) -> crate::error::Result> { + let mut builder = RawSeriesFrameBuilder::try_new(schema, mode, max_rows, max_bytes)?; + let mut cursor = RawSeriesChunkCursor::try_new(schema, mode, chunk.clone())?; + let mut batches = Vec::new(); + while let Some(batch) = cursor.append_to(&mut builder)? { + batches.push(batch); + } + if let Some(batch) = builder.flush()? { + batches.push(batch); + } + Ok(batches) +} + +/// Advances one owned chunk through an existing builder. A partially filled buffer stays buffered +/// instead of becoming a batch of its own, which is what makes cross-chunk grouping observable. +fn append_chunk( + builder: &mut RawSeriesFrameBuilder, + schema: &SeriesSchema, + mode: RawLabelMode, + chunk: &RawSeriesChunk, +) -> crate::error::Result> { + let mut cursor = RawSeriesChunkCursor::try_new(schema, mode, chunk.clone())?; + let mut batches = Vec::new(); + while let Some(batch) = cursor.append_to(builder)? { + batches.push(batch); + } + Ok(batches) +} + +/// Decodes every batch with one reused decoder, so repeated calls exercise the cached schema. +fn decode_batches( + schema: &SeriesSchema, + mode: RawLabelMode, + batches: &[RecordBatch], +) -> Vec { + let decoder = RawSeriesFrameDecoder::try_new(schema, mode).unwrap(); + batches + .iter() + .flat_map(|batch| decoder.decode(batch).unwrap()) + .collect() +} + +/// Row-local histogram index and table lists of one field payload of a frame batch. +fn frame_histogram_lists(batch: &RecordBatch, field: usize) -> (ListArray, ListArray) { + let histogram = frame_field_payload(batch, field) + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + let indices = histogram + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .clone(); + let tables = histogram + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .clone(); + (indices, tables) +} + +/// Every row is expected to resolve index 0 inside its own table only, with the whole table +/// retained and the other index entry null. +fn assert_row_local_histogram_rows(rows: &[RawSeriesChunk]) { + assert_eq!(rows.len(), 2); + assert_eq!(rows[0].time.as_buffer().as_ref(), &[1_i64, 2]); + assert_eq!(rows[1].time.as_buffer().as_ref(), &[3_i64, 4]); + + let DenseSamplesRef::Histogram(first) = rows[0].fields.get(0).unwrap().dense().unwrap() else { + panic!("expected histogram samples"); + }; + assert_eq!(first.batch().len(), 2); + assert_eq!(first.indices().len(), 2); + assert_eq!(first.indices().value(0), 0); + assert!(!first.indices().is_valid(1)); + assert_histogram(first.view_at(0).unwrap(), -9.5); + assert_unreferenced_histogram(first.batch().view(1).unwrap()); + + let DenseSamplesRef::Histogram(second) = rows[1].fields.get(0).unwrap().dense().unwrap() else { + panic!("expected histogram samples"); + }; + assert_eq!(second.batch().len(), 2); + assert_eq!(second.indices().len(), 2); + assert!(!second.indices().is_valid(0)); + assert_eq!(second.indices().value(1), 0); + // A global index remap would resolve index 0 into the first row's table instead. + assert_histogram(second.view_at(1).unwrap(), 42.25); + assert_histogram_with_start(second.batch().view(1).unwrap(), 142.25, None); +} + +#[test] +fn sliced_histogram_frames_densify_without_duplicating_shared_table() { + let schema = test_schema(&["histogram"]); + let source_table = histogram_batch(); + let source_samples = + Samples::try_from_histogram(UInt32Array::from(vec![Some(0); 64]), source_table).unwrap(); + let samples = Samples::try_from_chunks( + (0..64) + .map(|position| source_samples.slice(position, 1).unwrap()) + .collect(), + ) + .unwrap(); + let chunk = RawSeriesChunk::try_new( + &schema, + Buffer::from(vec![0x42_u8]), + None, + RawTimestamps::try_new((0_i64..64).collect::>().into()).unwrap(), + SeriesFields::one(samples), + true, + ) + .unwrap(); + + // One row per batch keeps the original segment-boundary granularity, so the shared-table fast + // paths below stay observable; the default budget coalesces the same rows later on. + let frames = pack_chunk_with_budget( + &schema, + RawLabelMode::Key, + &chunk, + 1, + RawSeriesFrameBuilder::DEFAULT_MAX_BYTES, + ) + .unwrap(); + assert_eq!(frames.len(), 64); + let mut decoded_chunks = Vec::with_capacity(frames.len()); + for frame in &frames { + let decoded = unpack_raw_series_frame(&schema, RawLabelMode::Key, frame).unwrap(); + assert_eq!(decoded.len(), 1); + assert_eq!(decoded[0].time.len(), 1); + decoded_chunks.push(decoded.into_iter().next().unwrap()); + } + + let first_table = match decoded_chunks[0].fields.get(0).unwrap().dense().unwrap() { + DenseSamplesRef::Histogram(histograms) => histograms.batch().clone(), + _ => panic!("expected histogram samples"), + }; + let second_table = match decoded_chunks[1].fields.get(0).unwrap().dense().unwrap() { + DenseSamplesRef::Histogram(histograms) => histograms.batch(), + _ => panic!("expected histogram samples"), + }; + assert!(!Arc::ptr_eq(&first_table, second_table)); + let first_values = first_table + .as_arrow() + .column(2) + .as_any() + .downcast_ref::() + .unwrap(); + let second_values = second_table + .as_arrow() + .column(2) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!( + first_values.values().as_ptr(), + second_values.values().as_ptr() + ); + + // The default budget groups the same 64 rows into one multirow batch; every row still decodes + // to its own slice of the shared table with bit-exact values. + let coalesced = pack_raw_series_chunk(&schema, RawLabelMode::Key, &chunk).unwrap(); + assert_eq!(coalesced.len(), 1); + assert_eq!(coalesced[0].num_rows(), 64); + let coalesced_rows = decode_batches(&schema, RawLabelMode::Key, &coalesced); + assert_eq!(coalesced_rows.len(), 64); + for (position, row) in coalesced_rows.iter().enumerate() { + assert_eq!(row.primary_key.as_slice(), &[0x42]); + assert_eq!(row.time.timestamp(0), Some(position as i64)); + let DenseSamplesRef::Histogram(histograms) = row.fields.get(0).unwrap().dense().unwrap() + else { + panic!("expected histogram samples"); + }; + assert_eq!(histograms.batch().len(), first_table.len()); + assert_histogram(histograms.view_at(0).unwrap(), -9.5); + } + + let samples = Samples::try_from_chunks( + decoded_chunks + .iter() + .map(|decoded| decoded.fields.get(0).unwrap().clone()) + .collect(), + ) + .unwrap(); + drop(chunk); + drop(frames); + drop(decoded_chunks); + drop(source_samples); + + let dense = samples.into_dense().unwrap(); + let DenseSamplesRef::Histogram(histograms) = dense.dense().unwrap() else { + panic!("expected dense histogram samples"); + }; + assert_eq!(histograms.len(), 64); + assert_eq!(histograms.batch().len(), first_table.len()); + for position in 0..64 { + assert_histogram(histograms.view_at(position).unwrap(), -9.5); + } + let dense_values = histograms + .batch() + .as_arrow() + .column(2) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!( + dense_values.values().as_ptr(), + first_values.values().as_ptr() + ); +} + +#[test] +fn raw_frame_field_names_bind_payloads_to_their_schema() { + let schema = test_schema(&["cpu_usage", "cpu_idle"]); + let fields = SeriesFields::try_from_vec(vec![ + Samples::from_float(Float64Array::from(vec![Some(1.0), Some(2.0)])), + Samples::from_float(Float64Array::from(vec![Some(10.0), Some(20.0)])), + ]) + .unwrap(); + + for mode in [RawLabelMode::Key, RawLabelMode::Explicit] { + let metric = (mode == RawLabelMode::Explicit).then(|| metric(Some("host"), &[])); + let chunk = RawSeriesChunk::try_new( + &schema, + Buffer::from(vec![0xaa_u8]), + metric, + RawTimestamps::try_new(vec![100, 200].into()).unwrap(), + fields.clone(), + true, + ) + .unwrap(); + let batch = &pack_raw_series_chunk(&schema, mode, &chunk).unwrap()[0]; + let payloads = batch + .column(3) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!( + payloads + .fields() + .iter() + .map(|field| field.name().as_str()) + .collect::>(), + ["cpu_usage", "cpu_idle"] + ); + + let decoded = unpack_raw_series_frame(&schema, mode, batch).unwrap(); + assert_eq!(decoded[0].primary_key.as_slice(), &[0xaa]); + assert!(decoded[0].is_last); + let times = decoded[0].time.as_buffer(); + assert_eq!(times.as_ref(), &[100, 200]); + assert_float_bits(decoded[0].fields.get(0).unwrap(), &[Some(1.0), Some(2.0)]); + assert_float_bits(decoded[0].fields.get(1).unwrap(), &[Some(10.0), Some(20.0)]); + + assert!( + unpack_raw_series_frame(&test_schema(&["cpu_idle", "cpu_usage"]), mode, batch).is_err() + ); + assert!( + unpack_raw_series_frame(&test_schema(&["cpu_usage", "cpu_busy"]), mode, batch).is_err() + ); + } +} + +#[test] +fn schema_and_key_explicit_round_trips_preserve_identity_and_bits() { + let schema = test_schema(&["value"]); + let key_schema = raw_series_frame_schema(&schema, RawLabelMode::Key).unwrap(); + let explicit_schema = raw_series_frame_schema(&schema, RawLabelMode::Explicit).unwrap(); + let payload = DataType::Struct( + vec![ + Field::new( + "float_values", + DataType::List(Arc::new(Field::new("item", DataType::Float64, true))), + true, + ), + Field::new( + "histogram", + DataType::Struct( + vec![ + Field::new( + "indices", + DataType::List(Arc::new(Field::new("item", DataType::UInt32, true))), + false, + ), + Field::new( + "table", + DataType::List(Arc::new(Field::new( + "item", + DataType::Struct(crate::series::histogram::fields()), + false, + ))), + false, + ), + ] + .into(), + ), + true, + ), + ] + .into(), + ); + let mut expected_key_metadata = std::collections::HashMap::new(); + expected_key_metadata.insert("greptime.series.layout".to_string(), "v1".to_string()); + expected_key_metadata.insert("greptime.series.content".to_string(), "raw".to_string()); + expected_key_metadata.insert("greptime.series.labels".to_string(), "key".to_string()); + expected_key_metadata.insert("greptime.series.time".to_string(), "timestamps".to_string()); + let expected_key = Schema::new_with_metadata( + vec![ + Field::new("read_key", DataType::Binary, false), + Field::new("is_last", DataType::Boolean, false), + Field::new( + "time", + DataType::Struct( + vec![Field::new( + "timestamps", + DataType::List(Arc::new(Field::new("item", DataType::Int64, false))), + false, + )] + .into(), + ), + false, + ), + Field::new( + "fields", + DataType::Struct(vec![Field::new("value", payload, false)].into()), + false, + ), + ], + expected_key_metadata, + ); + assert_eq!(key_schema.as_ref(), &expected_key); + let mut explicit_fields = expected_key.fields().iter().cloned().collect::>(); + explicit_fields.push(Field::new("name", DataType::Utf8, true).into()); + explicit_fields.push( + Field::new( + "labels", + DataType::List(Arc::new(Field::new( + "item", + DataType::Struct( + vec![ + Field::new("name", DataType::Utf8, false), + Field::new("value", DataType::Utf8, false), + ] + .into(), + ), + false, + ))), + false, + ) + .into(), + ); + let mut explicit_metadata = expected_key.metadata().clone(); + explicit_metadata.insert("greptime.series.labels".to_string(), "explicit".to_string()); + let expected_explicit = Schema::new_with_metadata(explicit_fields, explicit_metadata); + assert_eq!(explicit_schema.as_ref(), &expected_explicit); + + let nan = f64::from_bits(0x7ff8_0000_0000_0042); + let stale_nan = f64::from_bits(0x7ff0_0000_0000_0002); + let values = [Some(-0.0), None, Some(nan), Some(stale_nan)]; + let native = float_chunk( + &schema, + &[0, 0xff, 0, 0x80], + None, + vec![10, 20, 30, 40], + values.to_vec(), + true, + ); + let packed = pack_raw_series_chunk(&schema, RawLabelMode::Key, &native).unwrap(); + assert_eq!(packed.len(), 1); + let decoded = unpack_raw_series_frame(&schema, RawLabelMode::Key, &packed[0]).unwrap(); + assert_eq!(decoded.len(), 1); + assert_eq!(decoded[0].primary_key.as_slice(), &[0, 0xff, 0, 0x80]); + assert!(decoded[0].metric.is_none()); + assert!(decoded[0].is_last); + assert_eq!( + (0..4) + .map(|i| decoded[0].time.timestamp(i).unwrap()) + .collect::>(), + [10, 20, 30, 40] + ); + assert_float_bits(decoded[0].fields.get(0).unwrap(), &values); + + let explicit = float_chunk( + &schema, + &[0x80, 0, 0xff], + Some(metric(None, &[("rack", ""), ("zone", "west")])), + vec![7], + vec![Some(1.25)], + false, + ); + let packed = pack_raw_series_chunk(&schema, RawLabelMode::Explicit, &explicit).unwrap(); + let decoded = unpack_raw_series_frame(&schema, RawLabelMode::Explicit, &packed[0]).unwrap(); + assert_eq!(decoded[0].primary_key.as_slice(), &[0x80, 0, 0xff]); + assert!(!decoded[0].is_last); + let labels = decoded[0].metric.as_ref().unwrap(); + assert_eq!(labels.name(), None); + assert_eq!(labels.get("rack"), Some("")); + assert_eq!(labels.get("zone"), Some("west")); + assert!(pack_raw_series_chunk(&schema, RawLabelMode::Key, &explicit).is_err()); + assert!(pack_raw_series_chunk(&schema, RawLabelMode::Explicit, &native).is_err()); + assert!(unpack_raw_series_frame(&schema, RawLabelMode::Key, &packed[0]).is_err()); + + let key_packed = pack_raw_series_chunk(&schema, RawLabelMode::Key, &native).unwrap(); + assert!(unpack_raw_series_frame(&schema, RawLabelMode::Explicit, &key_packed[0]).is_err()); + for name in [None, Some(""), Some("cpu")] { + let named = float_chunk( + &schema, + &[0xff, 0, 0xfe], + Some(metric(name, &[("empty", "")])), + vec![9], + vec![Some(2.0)], + false, + ); + let encoded = pack_raw_series_chunk(&schema, RawLabelMode::Explicit, &named).unwrap(); + let restored = + unpack_raw_series_frame(&schema, RawLabelMode::Explicit, &encoded[0]).unwrap(); + assert_eq!(restored[0].metric.as_ref().unwrap().name(), name); + assert_eq!(restored[0].metric.as_ref().unwrap().get("empty"), Some("")); + } +} + +#[test] +fn segmented_mixed_fields_emit_every_aligned_position_and_marker() { + let schema = test_schema(&["first", "second"]); + let first_table = histogram_batch(); + let second_table = histogram_batch_with_sum(42.25); + let first = Samples::try_from_chunks(vec![ + Samples::from_float(Float64Array::from(vec![Some(1.0), None])), + Samples::try_from_mixed( + Float64Array::from(vec![Some(3.5), None]), + crate::series::samples::HistogramSamples::try_new( + UInt32Array::from(vec![None, Some(0)]), + first_table.clone(), + ) + .unwrap(), + ) + .unwrap(), + ]) + .unwrap(); + let second = Samples::try_from_chunks(vec![ + Samples::from_float(Float64Array::from(vec![Some(5.0)])), + Samples::try_from_histogram(UInt32Array::from(vec![Some(0), None]), second_table).unwrap(), + Samples::from_float(Float64Array::from(vec![Some(8.0)])), + ]) + .unwrap(); + let chunk = RawSeriesChunk::try_new( + &schema, + Buffer::from(vec![0_u8, 0x91, 0xfe]), + None, + RawTimestamps::try_new(vec![10_i64, 20, 30, 40].into()).unwrap(), + SeriesFields::try_from_vec(vec![first, second]).unwrap(), + true, + ) + .unwrap(); + // Variable field boundaries now group every aligned interval into one multirow batch instead + // of one batch per boundary, so decode the actual rows that batch carries. + let frames = pack_raw_series_chunk(&schema, RawLabelMode::Key, &chunk).unwrap(); + assert_eq!(frames.len(), 1); + assert_eq!(frames[0].num_rows(), 4); + let mut observed_times = Vec::new(); + let mut observed_first = Vec::new(); + let mut observed_second = Vec::new(); + let mut observed_keys = Vec::new(); + let mut observed_markers = Vec::new(); + let rows = unpack_raw_series_frame(&schema, RawLabelMode::Key, &frames[0]).unwrap(); + assert_eq!(rows.len(), 4); + for (row_index, row) in rows.iter().enumerate() { + assert_eq!(row.is_last, row_index == rows.len() - 1); + observed_markers.push(row.is_last); + assert_eq!(row.primary_key.as_slice(), &[0, 0x91, 0xfe]); + observed_keys.push(row.primary_key.as_slice().to_vec()); + for position in 0..row.time.len() { + observed_times.push(row.time.timestamp(position).unwrap()); + match row.fields.get(0).unwrap().sample(position) { + SampleRef::Missing => observed_first.push(("m", 0.0_f64.to_bits())), + SampleRef::Float(value) => observed_first.push(("f", value.to_bits())), + SampleRef::Histogram(histogram) => { + assert_histogram(histogram, -9.5); + observed_first.push(("h", histogram.sum.to_bits())); + } + } + match row.fields.get(1).unwrap().sample(position) { + SampleRef::Missing => observed_second.push(("m", 0.0_f64.to_bits())), + SampleRef::Float(value) => observed_second.push(("f", value.to_bits())), + SampleRef::Histogram(histogram) => { + assert_histogram(histogram, 42.25); + observed_second.push(("h", histogram.sum.to_bits())); + } + } + } + } + assert_eq!(observed_times, [10, 20, 30, 40]); + assert_eq!(observed_keys, [[0, 0x91, 0xfe]; 4]); + assert_eq!(observed_markers, [false, false, false, true]); + assert_eq!( + observed_first, + [ + ("f", 1.0_f64.to_bits()), + ("m", 0.0_f64.to_bits()), + ("f", 3.5_f64.to_bits()), + ("h", (-9.5_f64).to_bits()) + ] + ); + assert_eq!( + observed_second, + [ + ("f", 5.0_f64.to_bits()), + ("h", 42.25_f64.to_bits()), + ("m", 0.0_f64.to_bits()), + ("f", 8.0_f64.to_bits()) + ] + ); + + // Different chunk boundaries (A: 2+3, B: 1+4) yield aligned + // emitted intervals of 1, 1, and 3 positions. + let aligned_schema = test_schema(&["a", "b"]); + let a_owner = Float64Array::from(vec![Some(1.0), Some(2.0), Some(3.0), None, Some(5.0)]); + let b_owner = Float64Array::from(vec![Some(10.0), None, Some(30.0), Some(40.0), Some(50.0)]); + let a = Samples::try_from_chunks(vec![ + Samples::from_float(a_owner.slice(0, 2)), + Samples::from_float(a_owner.slice(2, 3)), + ]) + .unwrap(); + let b = Samples::try_from_chunks(vec![ + Samples::from_float(b_owner.slice(0, 1)), + Samples::from_float(b_owner.slice(1, 4)), + ]) + .unwrap(); + let aligned_chunk = RawSeriesChunk::try_new( + &aligned_schema, + Buffer::from(vec![0xa1_u8]), + None, + RawTimestamps::try_new(vec![101_i64, 102, 103, 104, 105].into()).unwrap(), + SeriesFields::try_from_vec(vec![a, b]).unwrap(), + true, + ) + .unwrap(); + let aligned_frames = + pack_raw_series_chunk(&aligned_schema, RawLabelMode::Key, &aligned_chunk).unwrap(); + assert_eq!(aligned_frames.len(), 1); + assert_eq!(aligned_frames[0].num_rows(), 3); + let aligned_rows = + unpack_raw_series_frame(&aligned_schema, RawLabelMode::Key, &aligned_frames[0]).unwrap(); + assert_eq!(aligned_rows.len(), 3); + let mut interval_lengths = Vec::new(); + let mut aligned_times = Vec::new(); + let mut aligned_a = Vec::new(); + let mut aligned_b = Vec::new(); + for (i, row) in aligned_rows.iter().enumerate() { + assert_eq!(row.is_last, i == aligned_rows.len() - 1); + interval_lengths.push(row.time.len()); + let start = aligned_times.len(); + let DenseSamplesRef::Float(decoded_a) = row.fields.get(0).unwrap().dense().unwrap() else { + panic!("expected float field"); + }; + let DenseSamplesRef::Float(decoded_b) = row.fields.get(1).unwrap().dense().unwrap() else { + panic!("expected float field"); + }; + // Grouping several constituents into one row may copy their payloads, so compare values + // and validity bit-exactly instead of requiring the source buffers by pointer. + for (decoded_values, owner) in [(&decoded_a, &a_owner), (&decoded_b, &b_owner)] { + assert_eq!(decoded_values.len(), row.time.len()); + for position in 0..decoded_values.len() { + match ( + decoded_values.is_valid(position), + owner.is_valid(start + position), + ) { + (true, true) => assert_eq!( + decoded_values.value(position).to_bits(), + owner.value(start + position).to_bits(), + "grouping changed a value" + ), + (false, false) => {} + _ => panic!("grouping changed validity representation"), + } + } + } + for position in 0..row.time.len() { + aligned_times.push(row.time.timestamp(position).unwrap()); + aligned_a.push(match row.fields.get(0).unwrap().sample(position) { + SampleRef::Missing => None, + SampleRef::Float(value) => Some(value.to_bits()), + SampleRef::Histogram(_) => panic!("expected float field"), + }); + aligned_b.push(match row.fields.get(1).unwrap().sample(position) { + SampleRef::Missing => None, + SampleRef::Float(value) => Some(value.to_bits()), + SampleRef::Histogram(_) => panic!("expected float field"), + }); + } + } + assert_eq!(interval_lengths, [1, 1, 3]); + assert_eq!(aligned_times, [101, 102, 103, 104, 105]); + assert_eq!( + aligned_a, + [ + Some(1.0_f64.to_bits()), + Some(2.0_f64.to_bits()), + Some(3.0_f64.to_bits()), + None, + Some(5.0_f64.to_bits()) + ] + ); + assert_eq!( + aligned_b, + [ + Some(10.0_f64.to_bits()), + None, + Some(30.0_f64.to_bits()), + Some(40.0_f64.to_bits()), + Some(50.0_f64.to_bits()) + ] + ); + drop(aligned_chunk); + drop(aligned_frames); + + let mut distinct_rows = frames.clone(); + let alternate_schema = test_schema(&["first", "second"]); + let first_distinct = Samples::try_from_histogram( + UInt32Array::from(vec![Some(0)]), + histogram_batch_with_sum(123.0), + ) + .unwrap(); + let second_distinct = Samples::try_from_histogram( + UInt32Array::from(vec![Some(0)]), + histogram_batch_with_sum(456.0), + ) + .unwrap(); + let alternate = RawSeriesChunk::try_new( + &alternate_schema, + Buffer::from(vec![0xde_u8, 0xad]), + None, + RawTimestamps::try_new(vec![100_i64].into()).unwrap(), + SeriesFields::try_from_vec(vec![first_distinct, second_distinct]).unwrap(), + false, + ) + .unwrap(); + distinct_rows.push( + pack_raw_series_chunk(&alternate_schema, RawLabelMode::Key, &alternate) + .unwrap() + .remove(0), + ); + let multirow = concat_batches(&frames[0].schema(), &distinct_rows).unwrap(); + let decoded = unpack_raw_series_frame(&schema, RawLabelMode::Key, &multirow).unwrap(); + assert_eq!(decoded.len(), 5); + assert_eq!( + decoded.iter().map(|row| row.is_last).collect::>(), + [false, false, false, true, false] + ); + assert_eq!( + decoded + .iter() + .map(|row| row.time.timestamp(0).unwrap()) + .collect::>(), + [10, 20, 30, 40, 100] + ); + assert_eq!(decoded[4].primary_key.as_slice(), &[0xde, 0xad]); + assert!( + matches!(decoded[4].fields.get(0).unwrap().sample(0), SampleRef::Histogram(h) if h.sum.to_bits() == 123.0_f64.to_bits()) + ); + assert!( + matches!(decoded[4].fields.get(1).unwrap().sample(0), SampleRef::Histogram(h) if h.sum.to_bits() == 456.0_f64.to_bits()) + ); + let sliced = multirow.slice(1, 3); + let mut altered = sliced.columns().to_vec(); + altered[0] = Arc::new(arrow::array::BinaryArray::from(vec![ + Some(&[0xee][..]), + Some(&[0xef][..]), + Some(&[0xed][..]), + ])); + let sliced = arrow::record_batch::RecordBatch::try_new(sliced.schema(), altered).unwrap(); + let sliced_rows = unpack_raw_series_frame(&schema, RawLabelMode::Key, &sliced).unwrap(); + assert_eq!( + sliced_rows + .iter() + .map(|row| row.time.timestamp(0).unwrap()) + .collect::>(), + [20, 30, 40] + ); + assert_eq!( + sliced_rows + .iter() + .map(|row| row.primary_key.as_slice()) + .collect::>(), + [&[0xee][..], &[0xef][..], &[0xed][..]] + ); + assert!(matches!( + sliced_rows[0].fields.get(0).unwrap().sample(0), + SampleRef::Missing + )); + assert!(matches!( + sliced_rows[0].fields.get(1).unwrap().sample(0), + SampleRef::Histogram(h) if h.sum.to_bits() == 42.25_f64.to_bits() + )); + assert!( + matches!(sliced_rows[1].fields.get(0).unwrap().sample(0), SampleRef::Float(v) if v.to_bits() == 3.5_f64.to_bits()) + ); + assert!(matches!( + sliced_rows[1].fields.get(1).unwrap().sample(0), + SampleRef::Missing + )); + assert!( + matches!(sliced_rows[2].fields.get(0).unwrap().sample(0), SampleRef::Histogram(h) if h.sum.to_bits() == (-9.5_f64).to_bits()) + ); + assert!( + matches!(sliced_rows[2].fields.get(1).unwrap().sample(0), SampleRef::Float(v) if v.to_bits() == 8.0_f64.to_bits()) + ); + + let not_last = RawSeriesChunk::try_new( + &schema, + Buffer::from(vec![0_u8, 0x91, 0xfe]), + None, + RawTimestamps::try_new(vec![50_i64, 60, 70, 80].into()).unwrap(), + chunk.fields.clone(), + false, + ) + .unwrap(); + let not_last_frames = pack_raw_series_chunk(&schema, RawLabelMode::Key, ¬_last).unwrap(); + let markers = not_last_frames + .iter() + .flat_map(|frame| unpack_raw_series_frame(&schema, RawLabelMode::Key, frame).unwrap()) + .map(|row| row.is_last) + .collect::>(); + assert_eq!(markers, [false; 4]); +} + +#[test] +fn histogram_payload_and_present_all_null_representation_survive() { + let schema = test_schema(&["histogram", "missing_histogram", "missing_float"]); + let mut builder = HistogramBatchBuilder::new(); + let custom = [1.5, 4.5]; + let po = [1]; + let pl = [2]; + let pb = [3.25, 7.5]; + let no = [-2]; + let nl = [1]; + let nb = [0.75]; + for (sum, start_timestamp) in [(-9.5, Some(-123)), (90.5, None)] { + builder + .push(crate::series::histogram::HistogramView { + schema: -7, + zero_threshold: 0.125, + sum, + count: 17.25, + zero_count: 2.5, + reset_hint: CounterResetHint::Unrecognized(37), + start_timestamp, + custom_values: &custom, + positive: BucketSideView { + span_offsets: &po, + span_lengths: &pl, + buckets: &pb, + }, + negative: BucketSideView { + span_offsets: &no, + span_lengths: &nl, + buckets: &nb, + }, + }) + .unwrap(); + } + let table = builder.finish().unwrap(); + let fields = SeriesFields::try_from_vec(vec![ + Samples::try_from_histogram(UInt32Array::from(vec![Some(0), Some(0)]), table.clone()) + .unwrap(), + Samples::try_from_histogram(UInt32Array::from(vec![None, None]), table.clone()).unwrap(), + Samples::from_float(Float64Array::new_null(2)), + ]) + .unwrap(); + let chunk = RawSeriesChunk::try_new( + &schema, + Buffer::from(vec![0x12_u8, 0x00]), + None, + RawTimestamps::try_new(vec![1_i64, 2].into()).unwrap(), + fields, + false, + ) + .unwrap(); + let frame = pack_raw_series_chunk(&schema, RawLabelMode::Key, &chunk).unwrap(); + let decoded = unpack_raw_series_frame(&schema, RawLabelMode::Key, &frame[0]).unwrap(); + let DenseSamplesRef::Histogram(histograms) = decoded[0].fields.get(0).unwrap().dense().unwrap() + else { + panic!("histogram representation was normalized"); + }; + assert_eq!(histograms.indices().len(), 2); + assert!(histograms.indices().is_valid(0)); + assert!(histograms.indices().is_valid(1)); + assert_histogram(histograms.view_at(0).unwrap(), -9.5); + assert_histogram(histograms.view_at(1).unwrap(), -9.5); + + let DenseSamplesRef::Histogram(all_null_histograms) = + decoded[0].fields.get(1).unwrap().dense().unwrap() + else { + panic!("all-null histogram representation was normalized"); + }; + assert_eq!(all_null_histograms.indices().len(), 2); + assert!((0..2).all(|i| !all_null_histograms.indices().is_valid(i))); + let DenseSamplesRef::Float(all_null_float) = decoded[0].fields.get(2).unwrap().dense().unwrap() + else { + panic!("all-null float representation was normalized"); + }; + assert_eq!(all_null_float.len(), 2); + assert!((0..2).all(|i| all_null_float.is_null(i))); + + let histogram = match decoded[0].fields.get(0).unwrap().sample(0) { + SampleRef::Histogram(histogram) => histogram, + _ => panic!("histogram was lost"), + }; + assert_histogram(histogram, -9.5); + + let mixed_schema = test_schema(&["mixed"]); + let table = histogram_batch(); + let mixed = Samples::try_from_mixed( + Float64Array::new_null(2), + crate::series::samples::HistogramSamples::try_new( + UInt32Array::from(vec![None, None]), + table.clone(), + ) + .unwrap(), + ) + .unwrap(); + let mixed_chunk = RawSeriesChunk::try_new( + &mixed_schema, + Buffer::from(vec![0x42_u8]), + None, + RawTimestamps::try_new(vec![5_i64, 6].into()).unwrap(), + SeriesFields::one(mixed), + false, + ) + .unwrap(); + let mixed_frame = + pack_raw_series_chunk(&mixed_schema, RawLabelMode::Key, &mixed_chunk).unwrap(); + let mixed_decoded = + unpack_raw_series_frame(&mixed_schema, RawLabelMode::Key, &mixed_frame[0]).unwrap(); + let DenseSamplesRef::Mixed(mixed) = mixed_decoded[0].fields.get(0).unwrap().dense().unwrap() + else { + panic!("all-null mixed representation was normalized"); + }; + assert_eq!(mixed.floats().len(), 2); + assert_eq!(mixed.histograms().indices().len(), 2); + assert!( + (0..2).all(|i| !mixed.floats().is_valid(i) && !mixed.histograms().indices().is_valid(i)) + ); +} + +#[test] +fn rejects_invalid_schema_and_metadata_and_accepts_valid_empty() { + let schema = test_schema(&["value"]); + let grid_schema = SeriesSchema { + fields: schema.fields.clone(), + eval_grid: Some(Arc::new( + crate::series::time::EvalGrid::try_new( + crate::series::time::GridSpec::try_new(1, 2, 3).unwrap(), + ) + .unwrap(), + )), + }; + assert!(raw_series_frame_schema(&grid_schema, RawLabelMode::Key).is_err()); + let empty_schema = raw_series_frame_schema(&schema, RawLabelMode::Key).unwrap(); + let empty = arrow::record_batch::RecordBatch::new_empty(empty_schema.clone()); + assert!( + unpack_raw_series_frame(&schema, RawLabelMode::Key, &empty) + .unwrap() + .is_empty() + ); + let wrong = + arrow::record_batch::RecordBatch::new_empty(Arc::new(Schema::new(vec![Field::new( + "read_key", + DataType::Utf8, + false, + )]))); + assert!(unpack_raw_series_frame(&schema, RawLabelMode::Key, &wrong).is_err()); + let no_fields = SeriesSchema { + fields: Arc::from([]), + eval_grid: None, + }; + assert!(raw_series_frame_schema(&no_fields, RawLabelMode::Key).is_err()); + + let mut metadata = empty_schema.metadata().clone(); + metadata.insert( + "greptime.series.content".to_string(), + "evaluated".to_string(), + ); + let contaminated = Arc::new(Schema::new_with_metadata( + empty_schema.fields().clone(), + metadata, + )); + let empty_bad = arrow::record_batch::RecordBatch::new_empty(contaminated); + assert!(unpack_raw_series_frame(&schema, RawLabelMode::Key, &empty_bad).is_err()); + + let chunk = float_chunk(&schema, &[1], None, vec![1], vec![Some(2.0)], false); + assert!(pack_raw_series_chunk(&grid_schema, RawLabelMode::Key, &chunk).is_err()); + + let valid = pack_raw_series_chunk(&schema, RawLabelMode::Key, &chunk) + .unwrap() + .remove(0); + let mut extra_metadata = valid.schema().metadata().clone(); + extra_metadata.insert("application.note".to_string(), "ignored".to_string()); + let unrelated = arrow::record_batch::RecordBatch::try_new( + Arc::new(Schema::new_with_metadata( + valid.schema().fields().clone(), + extra_metadata, + )), + valid.columns().to_vec(), + ) + .unwrap(); + assert_eq!( + unpack_raw_series_frame(&schema, RawLabelMode::Key, &unrelated) + .unwrap() + .len(), + 1 + ); + + let mut wrong_metadata = valid.schema().metadata().clone(); + wrong_metadata.remove("greptime.series.time"); + let wrong = arrow::record_batch::RecordBatch::try_new( + Arc::new(Schema::new_with_metadata( + valid.schema().fields().clone(), + wrong_metadata, + )), + valid.columns().to_vec(), + ) + .unwrap(); + assert!(unpack_raw_series_frame(&schema, RawLabelMode::Key, &wrong).is_err()); + + let mut grid_metadata = valid.schema().metadata().clone(); + grid_metadata.insert("greptime.series.grid.start".to_string(), "0".to_string()); + let contaminated = arrow::record_batch::RecordBatch::try_new( + Arc::new(Schema::new_with_metadata( + valid.schema().fields().clone(), + grid_metadata, + )), + valid.columns().to_vec(), + ) + .unwrap(); + assert!(unpack_raw_series_frame(&schema, RawLabelMode::Key, &contaminated).is_err()); + + for metadata_case in ["wrong_layout", "wrong_labels", "wrong_time"] { + let mut metadata = valid.schema().metadata().clone(); + match metadata_case { + "wrong_layout" => { + metadata.insert("greptime.series.layout".to_string(), "v2".to_string()); + } + "wrong_labels" => { + metadata.insert("greptime.series.labels".to_string(), "explicit".to_string()); + } + _ => { + metadata.insert("greptime.series.time".to_string(), "grid".to_string()); + } + } + let malformed = arrow::record_batch::RecordBatch::try_new( + Arc::new(Schema::new_with_metadata( + valid.schema().fields().clone(), + metadata, + )), + valid.columns().to_vec(), + ) + .unwrap(); + assert!(unpack_raw_series_frame(&schema, RawLabelMode::Key, &malformed).is_err()); + } +} + +#[test] +fn rejects_invalid_timestamps_and_empty_raw_payloads() { + assert!(RawTimestamps::try_new(vec![1_i64, 1].into()).is_err()); + assert!(RawTimestamps::try_new(vec![2_i64, 1].into()).is_err()); + + let schema = test_schema(&["value"]); + let chunk = float_chunk(&schema, &[1], None, vec![1], vec![Some(2.0)], false); + let valid = pack_raw_series_chunk(&schema, RawLabelMode::Key, &chunk) + .unwrap() + .remove(0); + for bad_timestamps in [vec![], vec![1, 1], vec![2, 1]] { + let malformed = with_timestamp_values(&valid, bad_timestamps); + assert!(unpack_raw_series_frame(&schema, RawLabelMode::Key, &malformed).is_err()); + } + let mut empty_columns = valid.columns().to_vec(); + let empty_time = with_timestamp_values(&valid, vec![]); + empty_columns[2] = empty_time.column(2).clone(); + let fields = valid.columns()[3] + .as_any() + .downcast_ref::() + .unwrap(); + let payload = fields + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let empty_float_list = ListArray::try_new( + match payload.fields()[0].data_type() { + DataType::List(item) => item.clone(), + _ => unreachable!(), + }, + OffsetBuffer::new(ScalarBuffer::from(vec![0_i32, 0])), + Arc::new(Float64Array::new_null(0)), + None, + ) + .unwrap(); + let empty_payload = StructArray::try_new( + payload.fields().clone(), + vec![Arc::new(empty_float_list), payload.column(1).clone()], + None, + ) + .unwrap(); + let empty_fields = + StructArray::try_new(fields.fields().clone(), vec![Arc::new(empty_payload)], None).unwrap(); + empty_columns[3] = Arc::new(empty_fields); + let empty_time_and_payload = + arrow::record_batch::RecordBatch::try_new(valid.schema(), empty_columns).unwrap(); + assert!(unpack_raw_series_frame(&schema, RawLabelMode::Key, &empty_time_and_payload).is_err()); + + let two_samples = float_chunk( + &schema, + &[1], + None, + vec![1, 2], + vec![Some(2.0), Some(3.0)], + false, + ); + let valid_two_samples = pack_raw_series_chunk(&schema, RawLabelMode::Key, &two_samples) + .unwrap() + .remove(0); + for bad_timestamps in [vec![1, 1], vec![2, 1]] { + let malformed = with_timestamp_values(&valid_two_samples, bad_timestamps); + assert!(unpack_raw_series_frame(&schema, RawLabelMode::Key, &malformed).is_err()); + } + + let mut empty_payload_columns = valid.columns().to_vec(); + let empty_time = with_timestamp_values(&valid, vec![]); + empty_payload_columns[2] = empty_time.column(2).clone(); + let float_fields = valid.columns()[3] + .as_any() + .downcast_ref::() + .unwrap(); + let float_payload = float_fields + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let empty_float_list = ListArray::try_new( + match float_payload.fields()[0].data_type() { + DataType::List(item) => item.clone(), + _ => unreachable!(), + }, + OffsetBuffer::new(ScalarBuffer::from(vec![0_i32, 0])), + Arc::new(Float64Array::from(Vec::>::new())), + None, + ) + .unwrap(); + let empty_payload = StructArray::try_new( + float_payload.fields().clone(), + vec![Arc::new(empty_float_list), float_payload.column(1).clone()], + None, + ) + .unwrap(); + let empty_fields = StructArray::try_new( + float_fields.fields().clone(), + vec![Arc::new(empty_payload)], + None, + ) + .unwrap(); + empty_payload_columns[3] = Arc::new(empty_fields); + let empty_payload = + arrow::record_batch::RecordBatch::try_new(valid.schema(), empty_payload_columns).unwrap(); + assert!(unpack_raw_series_frame(&schema, RawLabelMode::Key, &empty_payload).is_err()); + + let mut columns = valid.columns().to_vec(); + columns[3] = Arc::new(StructArray::new_null( + match valid.schema().field(3).data_type() { + DataType::Struct(fields) => fields.clone(), + _ => unreachable!(), + }, + 1, + )); + // Arrow rejects nulls in this canonical non-nullable top-level field before + // the decoder can inspect it; keep this a constructor-invariant assertion. + assert!(arrow::record_batch::RecordBatch::try_new(valid.schema(), columns).is_err()); + + let mut bad = valid.columns().to_vec(); + bad[2] = Arc::new( + with_timestamp_values(&valid, vec![1, 1]) + .column(2) + .as_any() + .downcast_ref::() + .unwrap() + .clone(), + ); + let invalid_later_row = arrow::record_batch::RecordBatch::try_new(valid.schema(), bad).unwrap(); + let later_invalid = + concat_batches(&valid.schema(), &[valid.clone(), invalid_later_row]).unwrap(); + assert!(unpack_raw_series_frame(&schema, RawLabelMode::Key, &later_invalid).is_err()); +} + +#[test] +fn rejects_invalid_explicit_labels() { + let explicit_schema = test_schema(&["value"]); + let explicit_chunk = float_chunk( + &explicit_schema, + &[1], + Some(metric(None, &[])), + vec![1], + vec![Some(1.0)], + false, + ); + let explicit = pack_raw_series_chunk(&explicit_schema, RawLabelMode::Explicit, &explicit_chunk) + .unwrap() + .remove(0); + let name = explicit + .column(4) + .as_any() + .downcast_ref::() + .unwrap() + .clone(); + let null_label_struct = StructArray::new_null( + match explicit.schema().field(5).data_type() { + DataType::List(item) => match item.data_type() { + DataType::Struct(fields) => fields.clone(), + _ => unreachable!(), + }, + _ => unreachable!(), + }, + 1, + ); + // Arrow rejects null list items because the canonical label item is non-nullable; + // this malformed value cannot reach the decoder through a safe constructor. + assert!( + ListArray::try_new( + match explicit.schema().field(5).data_type() { + DataType::List(item) => item.clone(), + _ => unreachable!(), + }, + OffsetBuffer::new(ScalarBuffer::from(vec![0_i32, 1])), + Arc::new(null_label_struct), + None, + ) + .is_err() + ); + let mut columns = explicit.columns().to_vec(); + columns[4] = Arc::new(name); + + let mut duplicate_labels = explicit.columns().to_vec(); + let duplicate_items = StructArray::try_new( + match explicit.schema().field(5).data_type() { + DataType::List(item) => match item.data_type() { + DataType::Struct(fields) => fields.clone(), + _ => unreachable!(), + }, + _ => unreachable!(), + }, + vec![ + Arc::new(StringArray::from(vec!["x", "x"])), + Arc::new(StringArray::from(vec!["1", "2"])), + ], + None, + ) + .unwrap(); + let duplicate_list = ListArray::try_new( + match explicit.schema().field(5).data_type() { + DataType::List(item) => item.clone(), + _ => unreachable!(), + }, + OffsetBuffer::new(ScalarBuffer::from(vec![0_i32, 2])), + Arc::new(duplicate_items), + None, + ) + .unwrap(); + duplicate_labels[5] = Arc::new(duplicate_list); + let duplicate = + arrow::record_batch::RecordBatch::try_new(explicit.schema(), duplicate_labels).unwrap(); + assert!(unpack_raw_series_frame(&explicit_schema, RawLabelMode::Explicit, &duplicate).is_err()); + + let mut reserved_labels = explicit.columns().to_vec(); + let reserved_items = StructArray::try_new( + match explicit.schema().field(5).data_type() { + DataType::List(item) => match item.data_type() { + DataType::Struct(fields) => fields.clone(), + _ => unreachable!(), + }, + _ => unreachable!(), + }, + vec![ + Arc::new(StringArray::from(vec!["__name__"])), + Arc::new(StringArray::from(vec!["x"])), + ], + None, + ) + .unwrap(); + let reserved_list = ListArray::try_new( + match explicit.schema().field(5).data_type() { + DataType::List(item) => item.clone(), + _ => unreachable!(), + }, + OffsetBuffer::new(ScalarBuffer::from(vec![0_i32, 1])), + Arc::new(reserved_items), + None, + ) + .unwrap(); + reserved_labels[5] = Arc::new(reserved_list); + let reserved = + arrow::record_batch::RecordBatch::try_new(explicit.schema(), reserved_labels).unwrap(); + assert!(unpack_raw_series_frame(&explicit_schema, RawLabelMode::Explicit, &reserved).is_err()); +} + +#[test] +fn rejects_malformed_histogram_and_payload_encodings() { + let hist_schema = test_schema(&["hist"]); + let valid_hist = pack_raw_series_chunk( + &hist_schema, + RawLabelMode::Key, + &RawSeriesChunk::try_new( + &hist_schema, + Buffer::from(vec![1_u8]), + None, + RawTimestamps::try_new(vec![1_i64].into()).unwrap(), + SeriesFields::one( + Samples::try_from_histogram(UInt32Array::from(vec![Some(0)]), histogram_batch()) + .unwrap(), + ), + false, + ) + .unwrap(), + ) + .unwrap() + .remove(0); + let columns = valid_hist.columns().to_vec(); + let fields = columns[3].as_any().downcast_ref::().unwrap(); + let payload = fields + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let null_hist = payload + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .clone(); + let hidden_indices = ListArray::try_new( + match hist_schema_frame_field(&hist_schema) { + DataType::Struct(f) => match f[0].data_type() { + DataType::Struct(p) => match p[1].data_type() { + DataType::Struct(h) => match h[0].data_type() { + DataType::List(item) => item.clone(), + _ => unreachable!(), + }, + _ => unreachable!(), + }, + _ => unreachable!(), + }, + _ => unreachable!(), + }, + OffsetBuffer::new(ScalarBuffer::from(vec![0_i32, 2])), + Arc::new(UInt32Array::from(vec![Some(99), None])), + None, + ) + .unwrap(); + let null_hist = StructArray::try_new( + null_hist.fields().clone(), + vec![Arc::new(hidden_indices), null_hist.column(1).clone()], + Some(arrow::buffer::NullBuffer::new_null(1)), + ) + .unwrap(); + let visible_float = ListArray::try_new( + match hist_schema_frame_field(&hist_schema) { + DataType::Struct(f) => match f[0].data_type() { + DataType::Struct(p) => match p[0].data_type() { + DataType::List(item) => item.clone(), + _ => unreachable!(), + }, + _ => unreachable!(), + }, + _ => unreachable!(), + }, + OffsetBuffer::new(ScalarBuffer::from(vec![0_i32, 1])), + Arc::new(Float64Array::from(vec![Some(1.25)])), + None, + ) + .unwrap(); + let payload = StructArray::try_new( + payload.fields().clone(), + vec![Arc::new(visible_float), Arc::new(null_hist.clone())], + None, + ) + .unwrap(); + let fields = StructArray::try_new( + fields.fields().clone(), + vec![Arc::new(payload.clone())], + None, + ) + .unwrap(); + let mut absent_hist_columns = valid_hist.columns().to_vec(); + absent_hist_columns[3] = Arc::new(fields.clone()); + let absent_hist = + arrow::record_batch::RecordBatch::try_new(valid_hist.schema(), absent_hist_columns) + .unwrap(); + let decoded_absent_hist = + unpack_raw_series_frame(&hist_schema, RawLabelMode::Key, &absent_hist).unwrap(); + assert!( + matches!(decoded_absent_hist[0].fields.get(0).unwrap().sample(0), SampleRef::Float(v) if v == 1.25) + ); + + let mut absent_both_columns = valid_hist.columns().to_vec(); + let absent_both_fields = absent_both_columns[3] + .as_any() + .downcast_ref::() + .unwrap(); + let absent_both_payload = absent_both_fields + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let absent_both_float = ListArray::try_new( + match absent_both_payload.fields()[0].data_type() { + DataType::List(item) => item.clone(), + _ => unreachable!(), + }, + OffsetBuffer::new(ScalarBuffer::from(vec![0_i32, 1])), + Arc::new(Float64Array::from(vec![Some(1.25)])), + Some(arrow::buffer::NullBuffer::new_null(1)), + ) + .unwrap(); + let absent_both_hist = absent_both_payload + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + let absent_both_hist = StructArray::try_new( + absent_both_hist.fields().clone(), + absent_both_hist.columns().to_vec(), + Some(arrow::buffer::NullBuffer::new_null(1)), + ) + .unwrap(); + let absent_both_payload = StructArray::try_new( + absent_both_payload.fields().clone(), + vec![Arc::new(absent_both_float), Arc::new(absent_both_hist)], + None, + ) + .unwrap(); + let absent_both_fields = StructArray::try_new( + absent_both_fields.fields().clone(), + vec![Arc::new(absent_both_payload)], + None, + ) + .unwrap(); + absent_both_columns[3] = Arc::new(absent_both_fields); + let absent_both = + arrow::record_batch::RecordBatch::try_new(valid_hist.schema(), absent_both_columns) + .unwrap(); + assert!(unpack_raw_series_frame(&hist_schema, RawLabelMode::Key, &absent_both).is_err()); + + let DataType::List(index_item) = null_hist.fields()[0].data_type() else { + panic!("expected histogram indices list"); + }; + let bad_indices = ListArray::try_new( + index_item.clone(), + OffsetBuffer::new(ScalarBuffer::from(vec![0_i32, 1])), + Arc::new(UInt32Array::from(vec![Some(99)])), + None, + ) + .unwrap(); + let bad_hist = StructArray::try_new( + null_hist.fields().clone(), + vec![Arc::new(bad_indices), null_hist.column(1).clone()], + None, + ) + .unwrap(); + let bad_payload = StructArray::try_new( + payload.fields().clone(), + vec![payload.column(0).clone(), Arc::new(bad_hist)], + None, + ) + .unwrap(); + let bad_fields = + StructArray::try_new(fields.fields().clone(), vec![Arc::new(bad_payload)], None).unwrap(); + let mut bad_parent_columns = valid_hist.columns().to_vec(); + bad_parent_columns[3] = Arc::new(bad_fields); + let present_bad_parent = + arrow::record_batch::RecordBatch::try_new(valid_hist.schema(), bad_parent_columns).unwrap(); + assert!(unpack_raw_series_frame(&hist_schema, RawLabelMode::Key, &present_bad_parent).is_err()); + + let valid_float = ListArray::try_new( + match payload.fields()[0].data_type() { + DataType::List(item) => item.clone(), + _ => unreachable!(), + }, + OffsetBuffer::new(ScalarBuffer::from(vec![0_i32, 1])), + Arc::new(Float64Array::from(vec![Some(1.25)])), + None, + ) + .unwrap(); + let overlap_indices = ListArray::try_new( + index_item.clone(), + OffsetBuffer::new(ScalarBuffer::from(vec![0_i32, 1])), + Arc::new(UInt32Array::from(vec![Some(0)])), + None, + ) + .unwrap(); + let overlap_hist = StructArray::try_new( + null_hist.fields().clone(), + vec![Arc::new(overlap_indices), null_hist.column(1).clone()], + None, + ) + .unwrap(); + let overlap_payload = StructArray::try_new( + payload.fields().clone(), + vec![Arc::new(valid_float), Arc::new(overlap_hist)], + None, + ) + .unwrap(); + let overlap_fields = StructArray::try_new( + fields.fields().clone(), + vec![Arc::new(overlap_payload)], + None, + ) + .unwrap(); + let mut overlap_columns = valid_hist.columns().to_vec(); + overlap_columns[3] = Arc::new(overlap_fields); + let overlap = + arrow::record_batch::RecordBatch::try_new(valid_hist.schema(), overlap_columns).unwrap(); + assert!(unpack_raw_series_frame(&hist_schema, RawLabelMode::Key, &overlap).is_err()); + + let short_indices = ListArray::try_new( + index_item.clone(), + OffsetBuffer::new(ScalarBuffer::from(vec![0_i32, 0])), + Arc::new(UInt32Array::from(Vec::>::new())), + None, + ) + .unwrap(); + let short_hist = StructArray::try_new( + null_hist.fields().clone(), + vec![Arc::new(short_indices), null_hist.column(1).clone()], + None, + ) + .unwrap(); + let short_payload = StructArray::try_new( + payload.fields().clone(), + vec![payload.column(0).clone(), Arc::new(short_hist)], + None, + ) + .unwrap(); + let short_fields = + StructArray::try_new(fields.fields().clone(), vec![Arc::new(short_payload)], None).unwrap(); + let mut short_columns = valid_hist.columns().to_vec(); + short_columns[3] = Arc::new(short_fields); + let short_indices = + arrow::record_batch::RecordBatch::try_new(valid_hist.schema(), short_columns).unwrap(); + assert!(unpack_raw_series_frame(&hist_schema, RawLabelMode::Key, &short_indices).is_err()); + + let malformed_histogram = || { + let mut columns = valid_hist.columns().to_vec(); + let fields = columns[3].as_any().downcast_ref::().unwrap(); + let payload = fields + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let hist = payload + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + let table_list = hist.column(1).as_any().downcast_ref::().unwrap(); + let table_value = table_list.value(0); + let table = table_value.as_any().downcast_ref::().unwrap(); + let lengths = ListArray::try_new( + match table.fields()[9].data_type() { + DataType::List(item) => item.clone(), + _ => unreachable!(), + }, + OffsetBuffer::new(ScalarBuffer::from(vec![0_i32, 1, 1])), + Arc::new(arrow::array::Int32Array::from(vec![1])), + None, + ) + .unwrap(); + let buckets = ListArray::try_new( + match table.fields()[10].data_type() { + DataType::List(item) => item.clone(), + _ => unreachable!(), + }, + OffsetBuffer::new(ScalarBuffer::from(vec![0_i32, 2, 5])), + Arc::new(Float64Array::from(vec![3.25, 7.5, 3.0, 7.0, 8.0])), + None, + ) + .unwrap(); + let mut table_columns = table.columns().to_vec(); + table_columns[9] = Arc::new(lengths); + table_columns[10] = Arc::new(buckets); + let malformed_table = + StructArray::try_new(table.fields().clone(), table_columns, None).unwrap(); + let malformed_tables = ListArray::try_new( + match hist.fields()[1].data_type() { + DataType::List(item) => item.clone(), + _ => unreachable!(), + }, + OffsetBuffer::new(ScalarBuffer::from(vec![0_i32, 1])), + Arc::new(malformed_table), + None, + ) + .unwrap(); + let malformed_hist = StructArray::try_new( + hist.fields().clone(), + vec![hist.column(0).clone(), Arc::new(malformed_tables)], + None, + ) + .unwrap(); + let malformed_payload = StructArray::try_new( + payload.fields().clone(), + vec![payload.column(0).clone(), Arc::new(malformed_hist)], + None, + ) + .unwrap(); + let malformed_fields = StructArray::try_new( + fields.fields().clone(), + vec![Arc::new(malformed_payload)], + None, + ) + .unwrap(); + columns[3] = Arc::new(malformed_fields); + arrow::record_batch::RecordBatch::try_new(valid_hist.schema(), columns).unwrap() + }; + assert!( + unpack_raw_series_frame(&hist_schema, RawLabelMode::Key, &malformed_histogram()).is_err() + ); + + let mut hidden_float = valid_hist.columns().to_vec(); + let fields = hidden_float[3] + .as_any() + .downcast_ref::() + .unwrap(); + let payload = fields + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let float_field = match hist_schema_frame_field(&hist_schema) { + DataType::Struct(fields) => match fields[0].data_type() { + DataType::Struct(payload_fields) => payload_fields[0].clone(), + _ => unreachable!(), + }, + _ => unreachable!(), + }; + let floats = ListArray::try_new( + match float_field.data_type() { + DataType::List(field) => field.clone(), + _ => unreachable!(), + }, + OffsetBuffer::new(ScalarBuffer::from(vec![0_i32, 1])), + Arc::new(Float64Array::from(vec![Some(999.0)])), + Some(arrow::buffer::NullBuffer::new_null(1)), + ) + .unwrap(); + let payload = StructArray::try_new( + payload.fields().clone(), + vec![Arc::new(floats), payload.column(1).clone()], + None, + ) + .unwrap(); + let fields = + StructArray::try_new(fields.fields().clone(), vec![Arc::new(payload)], None).unwrap(); + hidden_float[3] = Arc::new(fields); + let absent_float = + arrow::record_batch::RecordBatch::try_new(valid_hist.schema(), hidden_float).unwrap(); + let decoded_absent_float = + unpack_raw_series_frame(&hist_schema, RawLabelMode::Key, &absent_float).unwrap(); + assert_histogram( + match decoded_absent_float[0].fields.get(0).unwrap().sample(0) { + SampleRef::Histogram(histogram) => histogram, + _ => panic!("valid histogram was masked"), + }, + -9.5, + ); +} + +#[test] +fn decoded_key_and_payload_slices_keep_buffers_alive() { + let schema = test_schema(&["value"]); + let (owner, chunk, frames, decoded) = { + let owner = Float64Array::from(vec![Some(99.0), Some(2.5), None, Some(8.0)]); + let values = Samples::from_float(owner.slice(1, 2)); + let chunk = RawSeriesChunk::try_new( + &schema, + Buffer::from(vec![0xaa_u8, 0x10, 0x20, 0xbb]), + None, + RawTimestamps::try_new(vec![0_i64, 11, 22, 33].into()) + .unwrap() + .slice(1, 2) + .unwrap(), + SeriesFields::one(values), + true, + ) + .unwrap(); + let frames = pack_raw_series_chunk(&schema, RawLabelMode::Key, &chunk).unwrap(); + let decoded = unpack_raw_series_frame(&schema, RawLabelMode::Key, &frames[0]).unwrap(); + (owner, chunk, frames, decoded) + }; + let batch = frames[0].clone(); + let fields = batch + .column(3) + .as_any() + .downcast_ref::() + .unwrap(); + let field_payload = fields + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let floats = field_payload + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .as_any() + .downcast_ref::() + .unwrap(); + let decoded_floats = decoded[0].fields.get(0).unwrap(); + let DenseSamplesRef::Float(source_floats) = chunk.fields.get(0).unwrap().dense().unwrap() + else { + panic!() + }; + let DenseSamplesRef::Float(decoded_floats) = decoded_floats.dense().unwrap() else { + panic!() + }; + assert_eq!(decoded_floats.values().as_ptr(), floats.values().as_ptr()); + assert_eq!( + decoded_floats.values().as_ptr(), + owner.values().as_ptr().wrapping_add(1) + ); + assert_eq!( + decoded_floats.values().as_ptr(), + source_floats.values().as_ptr() + ); + let decoded_nulls = decoded_floats.nulls().unwrap(); + let source_nulls = source_floats.nulls().unwrap(); + assert_eq!( + decoded_nulls.buffer().as_ptr(), + source_nulls.buffer().as_ptr() + ); + assert_eq!(decoded_nulls.offset(), source_nulls.offset()); + let time_struct = batch + .column(2) + .as_any() + .downcast_ref::() + .unwrap(); + let timestamps = time_struct + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let time_values = timestamps + .values() + .as_any() + .downcast_ref::() + .unwrap(); + let decoded_time = decoded[0].time.as_buffer(); + assert_eq!(decoded_time.as_ptr(), time_values.values().as_ptr()); + let source_time = chunk.time.as_buffer(); + assert_eq!(decoded_time.as_ptr(), source_time.as_ptr()); + drop(owner); + drop(chunk); + drop(frames); + drop(batch); + assert_eq!(&decoded_time[..], &[11, 22]); + assert!(matches!(decoded[0].fields.get(0).unwrap().sample(0), SampleRef::Float(v) if v == 2.5)); + assert!(matches!( + decoded[0].fields.get(0).unwrap().sample(1), + SampleRef::Missing + )); + assert_eq!(decoded[0].primary_key.as_slice(), &[0xaa, 0x10, 0x20, 0xbb]); +} + +#[test] +fn sliced_canonical_mixed_histogram_round_trip_keeps_buffers_alive() { + let schema = test_schema(&["mixed", "histogram"]); + let mut builder = HistogramBatchBuilder::new(); + let custom = [1.5, 4.5]; + let po = [1]; + let pl = [2]; + let pb = [3.25, 7.5]; + let no = [-2]; + let nl = [1]; + let nb = [0.75]; + for (sum, start_timestamp) in [(99.5, Some(-123)), (-9.5, Some(-123)), (90.5, None)] { + builder + .push(crate::series::histogram::HistogramView { + schema: -7, + zero_threshold: 0.125, + sum, + count: 17.25, + zero_count: 2.5, + reset_hint: CounterResetHint::Unrecognized(37), + start_timestamp, + custom_values: &custom, + positive: BucketSideView { + span_offsets: &po, + span_lengths: &pl, + buckets: &pb, + }, + negative: BucketSideView { + span_offsets: &no, + span_lengths: &nl, + buckets: &nb, + }, + }) + .unwrap(); + } + let original_table = builder.finish().unwrap(); + let source_table = original_table.as_arrow(); + let sliced_table = Arc::new( + crate::series::histogram::HistogramBatch::try_from_arrow(source_table.slice(1, 2)).unwrap(), + ); + let indices_owner = UInt32Array::from(vec![Some(1), Some(0), None, Some(1)]); + let indices = indices_owner.slice(1, 2); + let histogram_indices_owner = UInt32Array::from(vec![Some(0), Some(1), Some(1), None]); + let histogram_indices = histogram_indices_owner.slice(1, 2); + let float_owner = Float64Array::from(vec![ + Some(99.0), + None, + Some(f64::from_bits(0x400c_0000_0000_0001)), + Some(77.0), + ]); + let mixed = Samples::try_from_mixed( + float_owner.slice(1, 2), + crate::series::samples::HistogramSamples::try_new(indices.clone(), sliced_table.clone()) + .unwrap(), + ) + .unwrap(); + let histogram = + Samples::try_from_histogram(histogram_indices.clone(), sliced_table.clone()).unwrap(); + let chunk = RawSeriesChunk::try_new( + &schema, + Buffer::from(vec![0x42_u8]), + None, + RawTimestamps::try_new(vec![0_i64, 11, 22, 33].into()) + .unwrap() + .slice(1, 2) + .unwrap(), + SeriesFields::try_from_vec(vec![mixed, histogram]).unwrap(), + true, + ) + .unwrap(); + + let source_indices = match chunk.fields.get(0).unwrap().dense().unwrap() { + DenseSamplesRef::Mixed(mixed) => mixed.histograms().indices(), + _ => panic!("expected mixed source"), + }; + assert_eq!(source_indices.offset(), 0); + assert_eq!(source_indices.value(0), 0); + assert!(!source_indices.is_valid(1)); + assert_eq!( + source_indices.values().as_ptr(), + indices_owner.values().as_ptr().wrapping_add(1) + ); + assert_eq!( + source_indices.nulls().unwrap().buffer().as_ptr(), + indices_owner.nulls().unwrap().buffer().as_ptr() + ); + assert_eq!(source_indices.nulls().unwrap().offset(), 1); + assert_eq!(histogram_indices.value(0), 1); + assert_eq!(histogram_indices.value(1), 1); + assert_eq!( + histogram_indices.values().as_ptr(), + histogram_indices_owner.values().as_ptr().wrapping_add(1) + ); + assert_eq!( + histogram_indices.nulls().unwrap().buffer().as_ptr(), + histogram_indices_owner.nulls().unwrap().buffer().as_ptr() + ); + assert_eq!(histogram_indices.nulls().unwrap().offset(), 1); + assert_eq!(sliced_table.len(), 2); + let source_sum = source_table + .column(2) + .as_any() + .downcast_ref::() + .unwrap(); + let sliced_sum = sliced_table + .as_arrow() + .column(2) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!( + sliced_sum.values().as_ptr(), + source_sum.values().as_ptr().wrapping_add(1) + ); + let source_positive = source_table + .column(10) + .as_any() + .downcast_ref::() + .unwrap(); + let sliced_positive = sliced_table + .as_arrow() + .column(10) + .as_any() + .downcast_ref::() + .unwrap(); + let sliced_table_offset = sliced_table.as_arrow().offset(); + assert_eq!(sliced_positive.offset(), source_positive.offset()); + assert_eq!( + sliced_positive.value_offsets().as_ptr(), + source_positive.value_offsets().as_ptr().wrapping_add(1) + ); + assert_eq!( + sliced_positive.value_offsets()[sliced_table_offset], + source_positive.value_offsets()[1] + ); + assert_eq!( + sliced_positive.value_offsets()[sliced_table_offset + 1], + source_positive.value_offsets()[2] + ); + assert_eq!( + sliced_positive + .values() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .as_ptr(), + source_positive + .values() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .as_ptr() + ); + let source_times = chunk.time.as_buffer(); + + // One row per batch keeps this fixture's single constituents and their nonzero offsets + // observable: the default budget could coalesce several rows and copy their payloads. + let frames = pack_chunk_with_budget( + &schema, + RawLabelMode::Key, + &chunk, + 1, + RawSeriesFrameBuilder::DEFAULT_MAX_BYTES, + ) + .unwrap(); + assert_eq!(frames.len(), 1); + let packed = &frames[0]; + let packed_times = packed + .column(2) + .as_any() + .downcast_ref::() + .unwrap() + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(packed_times.values().as_ptr(), source_times.as_ptr()); + assert_eq!(packed_times.values(), source_times); + let packed_mixed_indices = frame_field_payload(packed, 0) + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .as_any() + .downcast_ref::() + .unwrap(); + let packed_hist_indices = frame_field_payload(packed, 1) + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .as_any() + .downcast_ref::() + .unwrap(); + for (packed_indices, source_indices) in [ + (packed_mixed_indices, source_indices), + (packed_hist_indices, &histogram_indices), + ] { + assert_eq!( + packed_indices.values().as_ptr(), + source_indices.values().as_ptr() + ); + assert_eq!( + packed_indices.nulls().unwrap().buffer().as_ptr(), + source_indices.nulls().unwrap().buffer().as_ptr() + ); + assert_eq!( + packed_indices.nulls().unwrap().offset(), + source_indices.nulls().unwrap().offset() + ); + } + let packed_mixed_table = frame_histogram_table(packed, 0); + let packed_hist_table = frame_histogram_table(packed, 1); + for packed_table in [&packed_mixed_table, &packed_hist_table] { + let packed_sum = packed_table + .column(2) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(packed_sum.values().as_ptr(), sliced_sum.values().as_ptr()); + let packed_positive = packed_table + .column(10) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(packed_positive.offset(), sliced_positive.offset()); + assert_eq!( + packed_positive.value_offsets().as_ptr(), + sliced_positive.value_offsets().as_ptr() + ); + assert_eq!( + packed_positive.value_offsets()[sliced_table_offset], + sliced_positive.value_offsets()[sliced_table_offset] + ); + assert_eq!( + packed_positive.value_offsets()[sliced_table_offset + 1], + sliced_positive.value_offsets()[sliced_table_offset + 1] + ); + assert_eq!( + packed_positive + .values() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .as_ptr(), + sliced_positive + .values() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .as_ptr() + ); + } + let decoded = unpack_raw_series_frame(&schema, RawLabelMode::Key, packed).unwrap(); + let DenseSamplesRef::Mixed(decoded_mixed) = decoded[0].fields.get(0).unwrap().dense().unwrap() + else { + panic!("mixed representation was normalized"); + }; + let DenseSamplesRef::Histogram(decoded_histogram) = + decoded[0].fields.get(1).unwrap().dense().unwrap() + else { + panic!("histogram representation was normalized"); + }; + let decoded_time_values = decoded[0].time.as_buffer(); + assert_eq!(decoded_time_values.as_ptr(), packed_times.values().as_ptr()); + assert_eq!(decoded_time_values, source_times); + for (decoded_indices, source_indices) in [ + (decoded_mixed.histograms().indices(), source_indices), + (decoded_histogram.indices(), &histogram_indices), + ] { + assert_eq!( + decoded_indices.values().as_ptr(), + source_indices.values().as_ptr() + ); + assert_eq!( + decoded_indices.nulls().unwrap().buffer().as_ptr(), + source_indices.nulls().unwrap().buffer().as_ptr() + ); + assert_eq!( + decoded_indices.nulls().unwrap().offset(), + source_indices.nulls().unwrap().offset() + ); + } + assert_eq!( + decoded_mixed.histograms().indices().values().as_ptr(), + indices_owner.values().as_ptr().wrapping_add(1) + ); + assert_eq!( + decoded_histogram.indices().values().as_ptr(), + histogram_indices_owner.values().as_ptr().wrapping_add(1) + ); + assert_eq!(decoded_mixed.histograms().indices().value(0), 0); + assert!(!decoded_mixed.histograms().indices().is_valid(1)); + assert_eq!( + decoded_mixed.floats().values().as_ptr(), + float_owner.values().as_ptr().wrapping_add(1) + ); + assert_eq!( + decoded_mixed.floats().nulls().unwrap().buffer().as_ptr(), + float_owner.nulls().unwrap().buffer().as_ptr() + ); + assert_eq!(decoded_mixed.floats().nulls().unwrap().offset(), 1); + assert_eq!(decoded_histogram.indices().value(0), 1); + assert_eq!(decoded_histogram.indices().value(1), 1); + assert_eq!(decoded_histogram.batch().len(), 2); + let decoded_table = decoded_histogram.batch().as_arrow(); + let decoded_sum = decoded_table + .column(2) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(decoded_sum.values().as_ptr(), sliced_sum.values().as_ptr()); + let decoded_mixed_table = decoded_mixed.histograms().batch().as_arrow(); + for decoded_table_view in [decoded_mixed_table, decoded_table] { + let primitive = decoded_table_view + .column(2) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!( + primitive.values().as_ptr(), + source_sum.values().as_ptr().wrapping_add(1) + ); + let positive = decoded_table_view + .column(10) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(positive.offset(), source_positive.offset()); + assert_eq!( + positive.value_offsets()[sliced_table_offset], + source_positive.value_offsets()[1] + ); + assert_eq!( + positive.value_offsets().as_ptr(), + sliced_positive.value_offsets().as_ptr() + ); + assert_eq!( + positive.value_offsets()[sliced_table_offset + 1], + source_positive.value_offsets()[2] + ); + assert_eq!( + positive + .values() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .as_ptr(), + source_positive + .values() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .as_ptr() + ); + } + let decoded_positive_buckets = decoded_table + .column(10) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .as_any() + .downcast_ref::() + .unwrap(); + let source_positive_buckets = source_table + .column(10) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!( + decoded_positive_buckets.values().as_ptr(), + source_positive_buckets.values().as_ptr() + ); + + drop(packed_mixed_table); + drop(packed_hist_table); + drop(original_table); + drop(indices); + drop(histogram_indices); + drop(indices_owner); + drop(histogram_indices_owner); + drop(float_owner); + drop(sliced_table); + drop(chunk); + drop(frames); + + assert_eq!(decoded_table.len(), 2); + assert_histogram(decoded_mixed.histograms().batch().view(0).unwrap(), -9.5); + assert_unreferenced_histogram(decoded_mixed.histograms().batch().view(1).unwrap()); + assert_histogram(decoded_histogram.batch().view(0).unwrap(), -9.5); + assert_unreferenced_histogram(decoded_histogram.batch().view(1).unwrap()); + assert!(matches!( + decoded[0].fields.get(0).unwrap().sample(0), + SampleRef::Histogram(_) + )); + assert!( + matches!(decoded[0].fields.get(0).unwrap().sample(1), SampleRef::Float(v) if v.to_bits() == 0x400c_0000_0000_0001) + ); +} + +#[test] +fn malformed_pack_input_is_revalidated_before_encoding() { + let schema = test_schema(&["value"]); + let chunk = float_chunk( + &schema, + &[1], + None, + vec![1, 2], + vec![Some(3.0), Some(4.0)], + true, + ); + let invalid = RawSeriesChunk { + fields: SeriesFields::one(Samples::from_float(Float64Array::from(vec![Some(3.0)]))), + ..chunk + }; + assert!(pack_raw_series_chunk(&schema, RawLabelMode::Key, &invalid).is_err()); +} + +#[test] +fn public_segments_are_flat_ordered_and_retain_empty_dense_kinds() { + let histogram = + Samples::try_from_histogram(UInt32Array::from(vec![Some(0)]), histogram_batch()).unwrap(); + let mixed = Samples::try_from_mixed( + Float64Array::from(vec![Some(2.5)]), + crate::series::samples::HistogramSamples::try_new( + UInt32Array::from(vec![None]), + histogram_batch(), + ) + .unwrap(), + ) + .unwrap(); + let empty_float = Samples::from_float(Float64Array::new_null(0)); + let empty_histogram = Samples::try_from_histogram( + UInt32Array::from(Vec::>::new()), + histogram_batch(), + ) + .unwrap(); + let empty_mixed = Samples::try_from_mixed( + Float64Array::new_null(0), + crate::series::samples::HistogramSamples::try_new( + UInt32Array::from(Vec::>::new()), + histogram_batch(), + ) + .unwrap(), + ) + .unwrap(); + + for (samples, expected_len) in [(&empty_float, 0), (&empty_histogram, 0), (&empty_mixed, 0)] { + assert_eq!(samples.segments().len(), 1); + assert_eq!(samples.segments()[0].len(), expected_len); + assert!(samples.segments()[0].dense().is_some()); + } + + let first_segment_owner = Float64Array::from(vec![1.0, 2.0]); + let nested = Samples::try_from_chunks(vec![ + Samples::try_from_chunks(vec![ + empty_float.clone(), + Samples::from_float(first_segment_owner.clone()), + ]) + .unwrap(), + empty_histogram.clone(), + Samples::try_from_chunks(vec![histogram, empty_mixed.clone(), mixed]).unwrap(), + ]) + .unwrap(); + assert_eq!(nested.len(), 4); + assert_eq!(nested.segments().len(), 3); + let DenseSamplesRef::Float(first_segment) = nested.segments()[0].dense().unwrap() else { + panic!("expected float segment"); + }; + assert_eq!( + first_segment.values().as_ptr(), + first_segment_owner.values().as_ptr() + ); + let mut start = 0; + let mut starts = Vec::new(); + for segment in nested.segments() { + starts.push(start); + assert!(segment.dense().is_some()); + start += segment.len(); + } + assert_eq!(starts, [0, 2, 3]); + assert_eq!(start, nested.len()); + assert!(nested.dense().is_none()); + + let discarded_empty = Samples::try_from_chunks(vec![ + empty_float.clone(), + Samples::from_float(Float64Array::from(vec![1.0])), + empty_histogram.clone(), + Samples::from_float(Float64Array::from(vec![2.0])), + ]) + .unwrap(); + assert_eq!(discarded_empty.segments().len(), 2); + assert_eq!(discarded_empty.segments()[0].len(), 1); + assert_eq!(discarded_empty.segments()[1].len(), 1); + let all_empty = Samples::try_from_chunks(vec![empty_histogram, empty_mixed]).unwrap(); + assert_eq!(all_empty.segments().len(), 1); + assert_eq!(all_empty.segments()[0].len(), 0); + assert!(matches!( + all_empty.segments()[0].dense(), + Some(DenseSamplesRef::Float(_)) + )); +} + +#[test] +fn cross_chunk_rows_coalesce_into_one_multirow_batch() { + let schema = test_schema(&["value"]); + let first = float_chunk( + &schema, + &[0xaa], + None, + vec![10, 20], + vec![Some(1.0), None], + false, + ); + let second = float_chunk( + &schema, + &[0xbb], + None, + vec![30, 40, 50], + vec![Some(3.0), Some(4.0), None], + true, + ); + assert_eq!(RawSeriesFrameBuilder::DEFAULT_MAX_ROWS, 256); + assert_eq!(RawSeriesFrameBuilder::DEFAULT_MAX_BYTES, 1 << 20); + let expected = raw_series_frame_schema(&schema, RawLabelMode::Key).unwrap(); + + let mut builder = RawSeriesFrameBuilder::try_new( + &schema, + RawLabelMode::Key, + RawSeriesFrameBuilder::DEFAULT_MAX_ROWS, + RawSeriesFrameBuilder::DEFAULT_MAX_BYTES, + ) + .unwrap(); + assert_eq!(builder.schema(), &expected); + + let mut batches = append_chunk(&mut builder, &schema, RawLabelMode::Key, &first).unwrap(); + batches.extend(append_chunk(&mut builder, &schema, RawLabelMode::Key, &second).unwrap()); + if let Some(batch) = builder.flush().unwrap() { + batches.push(batch); + } + // Rows of both chunks coalesce into one batch: a partially filled buffer must survive the + // cursor boundary instead of becoming a batch of its own. + assert_eq!(batches.len(), 1); + let batch = &batches[0]; + assert_eq!(batch.num_rows(), 2); + assert_eq!(batch.schema(), expected); + assert_eq!( + batch + .schema() + .fields() + .iter() + .map(|field| field.name().as_str()) + .collect::>(), + ["read_key", "is_last", "time", "fields"] + ); + for (key, value) in [ + ("greptime.series.layout", "v1"), + ("greptime.series.content", "raw"), + ("greptime.series.labels", "key"), + ("greptime.series.time", "timestamps"), + ] { + assert_eq!( + batch.schema().metadata().get(key).map(String::as_str), + Some(value) + ); + } + + let keys = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(keys.value(0), &[0xaa]); + assert_eq!(keys.value(1), &[0xbb]); + let lasts = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!( + (0..2).map(|row| lasts.value(row)).collect::>(), + [false, true] + ); + let times = batch + .column(2) + .as_any() + .downcast_ref::() + .unwrap() + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(times.value_offsets(), [0, 2, 5]); + assert_eq!( + times + .values() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .as_ref(), + &[10_i64, 20, 30, 40, 50] + ); + let floats = frame_field_payload(batch, 0) + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(floats.value_offsets(), [0, 2, 5]); + let float_values = floats + .values() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!( + (0..float_values.len()) + .map(|position| float_values.is_valid(position)) + .collect::>(), + [true, false, true, true, false] + ); + assert_eq!(float_values.value(0).to_bits(), 1.0_f64.to_bits()); + assert_eq!(float_values.value(2).to_bits(), 3.0_f64.to_bits()); + assert_eq!(float_values.value(3).to_bits(), 4.0_f64.to_bits()); + + let decoded = decode_batches(&schema, RawLabelMode::Key, &batches); + assert_eq!(decoded.len(), 2); + assert_eq!(decoded[0].primary_key.as_slice(), &[0xaa]); + assert!(!decoded[0].is_last); + assert_eq!( + (0..decoded[0].time.len()) + .map(|position| decoded[0].time.timestamp(position).unwrap()) + .collect::>(), + [10, 20] + ); + assert_float_bits(decoded[0].fields.get(0).unwrap(), &[Some(1.0), None]); + assert_eq!(decoded[1].primary_key.as_slice(), &[0xbb]); + assert!(decoded[1].is_last); + assert_eq!( + (0..decoded[1].time.len()) + .map(|position| decoded[1].time.timestamp(position).unwrap()) + .collect::>(), + [30, 40, 50] + ); + assert_float_bits( + decoded[1].fields.get(0).unwrap(), + &[Some(3.0), Some(4.0), None], + ); + assert_eq!(decoded.iter().map(|row| row.time.len()).sum::(), 5); +} + +#[test] +fn raw_frame_builder_emits_at_the_row_budget_without_another_cursor() { + let schema = test_schema(&["value"]); + let chunks = [ + float_chunk(&schema, &[1], None, vec![1], vec![Some(1.0)], false), + float_chunk(&schema, &[2], None, vec![2], vec![Some(2.0)], false), + float_chunk(&schema, &[3], None, vec![3], vec![Some(3.0)], true), + ]; + + for max_rows in [2usize, 3] { + let mut builder = RawSeriesFrameBuilder::try_new( + &schema, + RawLabelMode::Key, + max_rows, + RawSeriesFrameBuilder::DEFAULT_MAX_BYTES, + ) + .unwrap(); + let mut emitted_at = None; + let mut batches = Vec::new(); + for (index, chunk) in chunks.iter().enumerate() { + let produced = append_chunk(&mut builder, &schema, RawLabelMode::Key, chunk).unwrap(); + assert!(produced.len() <= 1, "one cursor emits at most one batch"); + if let Some(batch) = produced.first() { + assert!( + emitted_at.replace(index).is_none(), + "the row bound is reached once, without another cursor" + ); + batches.push(batch.clone()); + } + } + if let Some(batch) = builder.flush().unwrap() { + batches.push(batch); + } + // Three one-row chunks emit at the max_rows-th row: a budget of two is reached during the + // second chunk and a budget of three during the third, so no further cursor is needed. + assert_eq!(emitted_at, Some(max_rows - 1)); + assert_eq!(batches[0].num_rows(), max_rows); + assert_eq!( + batches.iter().map(|batch| batch.num_rows()).sum::(), + 3 + ); + assert!(batches.iter().all(|batch| batch.num_rows() <= max_rows)); + let keys = batches[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!( + (0..max_rows) + .map(|row| keys.value(row).to_vec()) + .collect::>(), + (1..=max_rows) + .map(|key| vec![key as u8]) + .collect::>() + ); + + let decoded = decode_batches(&schema, RawLabelMode::Key, &batches); + assert_eq!(decoded.len(), 3); + assert_eq!( + decoded + .iter() + .map(|row| row.time.timestamp(0).unwrap()) + .collect::>(), + [1, 2, 3] + ); + assert_eq!( + decoded.iter().map(|row| row.is_last).collect::>(), + [false, false, true] + ); + } +} + +#[test] +fn raw_frame_builder_rejects_zero_budgets_and_bounds_offsets() { + let schema = test_schema(&["value"]); + // Zero in either budget is not a legal configuration; the literals take the parameter's type. + assert!( + RawSeriesFrameBuilder::try_new( + &schema, + RawLabelMode::Key, + 0, + RawSeriesFrameBuilder::DEFAULT_MAX_BYTES, + ) + .is_err() + ); + assert!( + RawSeriesFrameBuilder::try_new( + &schema, + RawLabelMode::Key, + RawSeriesFrameBuilder::DEFAULT_MAX_ROWS, + 0, + ) + .is_err() + ); + assert!(RawSeriesFrameBuilder::try_new(&schema, RawLabelMode::Key, 0, 0).is_err()); + assert!( + RawSeriesFrameBuilder::try_new( + &schema, + RawLabelMode::Key, + 1, + RawSeriesFrameBuilder::DEFAULT_MAX_BYTES, + ) + .is_ok() + ); + + // Cumulative Arrow list offsets are bounded through the shared helper, so the limit is proven + // without allocating a two-gigabyte fixture. + assert_eq!(checked_offset(i32::MAX as usize).unwrap(), i32::MAX); + assert!(checked_offset(i32::MAX as usize + 1).is_err()); + assert!(checked_offset(usize::MAX).is_err()); +} + +#[test] +fn byte_budget_emits_legal_oversized_rows_without_dropping_them() { + let schema = test_schema(&["value"]); + let owner = Float64Array::from(vec![Some(1.0), Some(2.0), Some(3.0)]); + let samples = Samples::try_from_chunks( + (0..3) + .map(|position| Samples::from_float(owner.slice(position, 1))) + .collect(), + ) + .unwrap(); + let chunk = RawSeriesChunk::try_new( + &schema, + Buffer::from(vec![0x9a_u8]), + None, + RawTimestamps::try_new(vec![7_i64, 8, 9].into()).unwrap(), + SeriesFields::one(samples), + true, + ) + .unwrap(); + + // A one-byte budget cannot hold any row, yet every legal row is still emitted as its own + // batch instead of being dropped or merged. + let batches = pack_chunk_with_budget( + &schema, + RawLabelMode::Key, + &chunk, + RawSeriesFrameBuilder::DEFAULT_MAX_ROWS, + 1, + ) + .unwrap(); + assert_eq!(batches.len(), 3); + assert!( + batches + .iter() + .all(|batch| batch.num_rows() == 1 && batch.schema() == batches[0].schema()), + "an oversized row is emitted alone, never dropped" + ); + + let decoded = decode_batches(&schema, RawLabelMode::Key, &batches); + assert_eq!(decoded.len(), 3); + for (position, row) in decoded.iter().enumerate() { + assert_eq!(row.primary_key.as_slice(), &[0x9a]); + assert_eq!(row.time.timestamp(0), Some(7 + position as i64)); + assert_eq!(row.is_last, position == 2); + assert_float_bits(row.fields.get(0).unwrap(), &[Some(position as f64 + 1.0)]); + } +} + +#[test] +fn byte_budget_flushes_before_exceeding_and_retries_the_pending_row() { + let schema = test_schema(&["value"]); + let chunks = [ + float_chunk(&schema, &[1], None, vec![10], vec![Some(1.0)], false), + float_chunk(&schema, &[2], None, vec![20], vec![Some(2.0)], true), + ]; + // The bound comes from the real rows instead of a fabricated counter: one byte less than + // both rows together fits the first row alone, so the second one must not join it. + let mut measured = Vec::with_capacity(chunks.len()); + for chunk in &chunks { + let mut cursor = + RawSeriesChunkCursor::try_new(&schema, RawLabelMode::Key, chunk.clone()).unwrap(); + measured.push(cursor.next_row().unwrap().unwrap().bytes); + } + let budget = measured[0] + measured[1] - 1; + + let mut builder = RawSeriesFrameBuilder::try_new( + &schema, + RawLabelMode::Key, + RawSeriesFrameBuilder::DEFAULT_MAX_ROWS, + budget, + ) + .unwrap(); + let mut first = + RawSeriesChunkCursor::try_new(&schema, RawLabelMode::Key, chunks[0].clone()).unwrap(); + assert!(first.append_to(&mut builder).unwrap().is_none()); + + let mut second = + RawSeriesChunkCursor::try_new(&schema, RawLabelMode::Key, chunks[1].clone()).unwrap(); + let emitted = second.append_to(&mut builder).unwrap().unwrap(); + assert_eq!(emitted.num_rows(), 1, "the first row is flushed alone"); + // The second row stays pending across the flush and is retried, never dropped. + assert!(second.append_to(&mut builder).unwrap().is_none()); + let remainder = builder.flush().unwrap().unwrap(); + assert_eq!(remainder.num_rows(), 1); + + let decoded = decode_batches(&schema, RawLabelMode::Key, &[emitted, remainder]); + assert_eq!(decoded.len(), 2); + assert_eq!(decoded[0].primary_key.as_slice(), &[1]); + assert_eq!(decoded[1].primary_key.as_slice(), &[2]); + assert_eq!(decoded[0].time.timestamp(0), Some(10)); + assert_eq!(decoded[1].time.timestamp(0), Some(20)); + assert_eq!( + decoded.iter().map(|row| row.is_last).collect::>(), + [false, true] + ); + assert_float_bits(decoded[0].fields.get(0).unwrap(), &[Some(1.0)]); + assert_float_bits(decoded[1].fields.get(0).unwrap(), &[Some(2.0)]); + + // A requested byte bound is clamped to the representable cumulative offset range, and that + // range is rejected before any child array is allocated. + let clamped = RawSeriesFrameBuilder::try_new( + &schema, + RawLabelMode::Key, + RawSeriesFrameBuilder::DEFAULT_MAX_ROWS, + usize::MAX, + ) + .unwrap(); + assert_eq!(clamped.max_bytes, i32::MAX as usize); + assert!(offset_buffer(&[i32::MAX as usize, 1]).is_err()); +} + +#[test] +fn multirow_histogram_rows_keep_row_local_indices_and_tables() { + let schema = test_schema(&["histogram"]); + let first_table = histogram_batch(); + let second_table = histogram_batch_with_sum(42.25); + let first_row = + Samples::try_from_histogram(UInt32Array::from(vec![Some(0), None]), first_table).unwrap(); + let second_row = + Samples::try_from_histogram(UInt32Array::from(vec![None, Some(0)]), second_table).unwrap(); + let fields = SeriesFields::one(Samples::try_from_chunks(vec![first_row, second_row]).unwrap()); + let chunk = RawSeriesChunk::try_new( + &schema, + Buffer::from(vec![0x5a_u8]), + None, + RawTimestamps::try_new(vec![1_i64, 2, 3, 4].into()).unwrap(), + fields, + true, + ) + .unwrap(); + + // The public wrapper groups two independent rows into one multirow batch. + let packed = pack_raw_series_chunk(&schema, RawLabelMode::Key, &chunk).unwrap(); + assert_eq!(packed.len(), 1); + assert_eq!(packed[0].num_rows(), 2); + + let (indices, tables) = frame_histogram_lists(&packed[0], 0); + let encoded_first = indices.value(0); + let encoded_first = encoded_first + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(encoded_first.value(0), 0); + assert!(!encoded_first.is_valid(1)); + let encoded_second = indices.value(1); + let encoded_second = encoded_second + .as_any() + .downcast_ref::() + .unwrap(); + assert!(!encoded_second.is_valid(0)); + assert_eq!( + encoded_second.value(1), + 0, + "row-local indices must not be remapped" + ); + // The second row references its own table through a nonzero slice of the concatenated child. + assert_eq!(tables.value_offsets(), [0, 2, 4]); + let concatenated = tables + .values() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(concatenated.len(), 4); + + let decoded = decode_batches(&schema, RawLabelMode::Key, std::slice::from_ref(&packed[0])); + assert_row_local_histogram_rows(&decoded); + + // A real Arrow IPC round trip keeps the schema, the whole tables, the null index entries and + // the bit-exact payloads. + let mut ipc = Vec::new(); + { + let mut writer = StreamWriter::try_new(&mut ipc, &packed[0].schema()).unwrap(); + writer.write(&packed[0]).unwrap(); + writer.finish().unwrap(); + } + let reader = StreamReader::try_new(std::io::Cursor::new(ipc), None).unwrap(); + let restored = reader.map(|batch| batch.unwrap()).collect::>(); + assert_eq!(restored.len(), 1); + assert_eq!(restored[0].schema(), packed[0].schema()); + assert_eq!(restored[0].num_rows(), 2); + let (restored_indices, restored_tables) = frame_histogram_lists(&restored[0], 0); + assert_eq!(restored_tables.value_offsets(), tables.value_offsets()); + let restored_second = restored_indices.value(1); + assert_eq!( + restored_second + .as_any() + .downcast_ref::() + .unwrap() + .value(1), + 0 + ); + let restored_rows = decode_batches(&schema, RawLabelMode::Key, &restored); + assert_row_local_histogram_rows(&restored_rows); +} + +#[test] +fn decoder_reuses_cached_schema_and_rejects_schema_evolution() { + let schema = test_schema(&["value"]); + let expected = raw_series_frame_schema(&schema, RawLabelMode::Key).unwrap(); + let first = pack_raw_series_chunk( + &schema, + RawLabelMode::Key, + &float_chunk( + &schema, + &[1], + None, + vec![1, 2], + vec![Some(1.0), Some(2.0)], + false, + ), + ) + .unwrap() + .remove(0); + let second = pack_raw_series_chunk( + &schema, + RawLabelMode::Key, + &float_chunk(&schema, &[2], None, vec![3], vec![Some(3.0)], true), + ) + .unwrap() + .remove(0); + + let key_decoder = RawSeriesFrameDecoder::try_new(&schema, RawLabelMode::Key).unwrap(); + let decoded = key_decoder.decode(&first).unwrap(); + assert_eq!(decoded.len(), 1); + assert_eq!(decoded[0].primary_key.as_slice(), &[1]); + assert_eq!( + (0..decoded[0].time.len()) + .map(|position| decoded[0].time.timestamp(position).unwrap()) + .collect::>(), + [1, 2] + ); + // The same decoder accepts a later valid batch before the malformed ones below. + let decoded = key_decoder.decode(&second).unwrap(); + assert_eq!(decoded.len(), 1); + assert_eq!(decoded[0].primary_key.as_slice(), &[2]); + assert!(decoded[0].is_last); + + // Unrelated metadata is not part of the raw frame contract. + let mut unrelated_metadata = second.schema().metadata().clone(); + unrelated_metadata.insert("application.note".to_string(), "ignored".to_string()); + let unrelated = arrow::record_batch::RecordBatch::try_new( + Arc::new(Schema::new_with_metadata( + expected.fields().clone(), + unrelated_metadata, + )), + second.columns().to_vec(), + ) + .unwrap(); + assert_eq!(key_decoder.decode(&unrelated).unwrap().len(), 1); + + // Reserved metadata and evolved fields are rejected, including by a zero-row batch. + for (key, value) in [ + ("greptime.series.layout", "v2"), + ("greptime.series.content", "evaluated"), + ("greptime.series.labels", "explicit"), + ("greptime.series.time", "grid"), + ] { + let mut metadata = second.schema().metadata().clone(); + metadata.insert(key.to_string(), value.to_string()); + let contaminated = arrow::record_batch::RecordBatch::try_new( + Arc::new(Schema::new_with_metadata( + expected.fields().clone(), + metadata, + )), + second.columns().to_vec(), + ) + .unwrap(); + assert!( + key_decoder.decode(&contaminated).is_err(), + "{key} must be rejected" + ); + } + let evolved_schema = test_schema(&["other"]); + let evolved = pack_raw_series_chunk( + &evolved_schema, + RawLabelMode::Key, + &float_chunk(&evolved_schema, &[3], None, vec![4], vec![Some(4.0)], false), + ) + .unwrap() + .remove(0); + assert!(key_decoder.decode(&evolved).is_err()); + let empty_evolved = arrow::record_batch::RecordBatch::new_empty( + raw_series_frame_schema(&evolved_schema, RawLabelMode::Key).unwrap(), + ); + assert!(key_decoder.decode(&empty_evolved).is_err()); + // Rejections leave the cached schema usable, and a valid zero-row batch stays valid. + assert_eq!(key_decoder.decode(&first).unwrap().len(), 1); + assert!( + key_decoder + .decode(&arrow::record_batch::RecordBatch::new_empty( + expected.clone() + )) + .unwrap() + .is_empty() + ); + + // Explicit mode reuses one decoder for a present but nameless metric and a named one. + let explicit_decoder = RawSeriesFrameDecoder::try_new(&schema, RawLabelMode::Explicit).unwrap(); + let nameless = pack_raw_series_chunk( + &schema, + RawLabelMode::Explicit, + &float_chunk( + &schema, + &[7], + Some(metric(None, &[("zone", "west")])), + vec![5], + vec![Some(5.0)], + true, + ), + ) + .unwrap() + .remove(0); + let named = pack_raw_series_chunk( + &schema, + RawLabelMode::Explicit, + &float_chunk( + &schema, + &[8], + Some(metric(Some("cpu"), &[])), + vec![6], + vec![Some(6.0)], + false, + ), + ) + .unwrap() + .remove(0); + let decoded = explicit_decoder.decode(&nameless).unwrap(); + let labels = decoded[0].metric.as_ref().unwrap(); + assert_eq!(labels.name(), None); + assert_eq!(labels.get("zone"), Some("west")); + assert_eq!( + explicit_decoder.decode(&named).unwrap()[0] + .metric + .as_ref() + .unwrap() + .name(), + Some("cpu") + ); + // Label modes are not interchangeable. + assert!(explicit_decoder.decode(&first).is_err()); + assert!(key_decoder.decode(&nameless).is_err()); +} diff --git a/src/datatypes/src/series/samples.rs b/src/datatypes/src/series/samples.rs index dead85c6d68..42b1d4bba0b 100644 --- a/src/datatypes/src/series/samples.rs +++ b/src/datatypes/src/series/samples.rs @@ -199,6 +199,46 @@ impl Samples { storage: SampleStorage::Chunks(flat.into(), total), }) } + /// Returns the dense segments in logical order without allocating, copying payloads, or + /// densifying the samples. + /// + /// Chunked samples expose their retained segment boundaries; different fields may therefore + /// have different boundaries. The logical start of each segment is the sum of the lengths of + /// all preceding segments, and the total length is [`Self::len`]. For non-chunked samples, + /// including empty samples, this returns a singleton slice containing `self`. Chunk + /// construction flattens nested chunks and drops empty segments; an all-empty input becomes + /// one empty float segment. Each segment is dense: [`Self::dense`] returns `Some` with its + /// typed Arrow view. Clone the samples cheaply if they need to outlive this borrow. + /// + /// ``` + /// use arrow::array::Float64Array; + /// use datatypes::series::samples::{DenseSamplesRef, Samples}; + /// + /// let samples = Samples::try_from_chunks(vec![ + /// Samples::from_float(Float64Array::from(vec![1.0, 2.0])), + /// Samples::from_float(Float64Array::from(vec![3.0])), + /// ]) + /// .unwrap(); + /// + /// let mut logical_start = 0; + /// let mut starts = Vec::new(); + /// for segment in samples.segments() { + /// starts.push(logical_start); + /// match segment.dense().unwrap() { + /// DenseSamplesRef::Float(values) => assert_eq!(values.len(), segment.len()), + /// _ => unreachable!("these segments contain floats"), + /// } + /// logical_start += segment.len(); + /// } + /// assert_eq!(starts, vec![0, 2]); + /// assert_eq!(logical_start, samples.len()); + /// ``` + pub fn segments(&self) -> &[Samples] { + match &self.storage { + SampleStorage::Chunks(chunks, _) => chunks, + _ => std::slice::from_ref(self), + } + } /// Number of sample positions. pub fn len(&self) -> usize { match &self.storage { @@ -315,8 +355,7 @@ impl Samples { } } /// Materialize chunks into one dense array. When all histogram-bearing chunks reference the - /// same table, reuse it; otherwise concatenate tables per chunk and remap indices, including - /// unused rows. + /// same Arrow table view, reuse it; otherwise concatenate tables per chunk and remap indices. pub fn into_dense(self) -> Result { let SampleStorage::Chunks(chunks, total_len) = self.storage else { return Ok(self); @@ -350,9 +389,17 @@ impl Samples { SampleStorage::Mixed(mixed) => Some(&mixed.histograms.batch), SampleStorage::Float(_) | SampleStorage::Chunks(_, _) => None, }); - let shared_batch = histogram_batches - .next() - .filter(|first| histogram_batches.all(|batch| Arc::ptr_eq(first, batch))); + let mut first_data = None; + let shared_batch = histogram_batches.next().filter(|first| { + histogram_batches.all(|batch| { + if Arc::ptr_eq(first, batch) { + return true; + } + first_data + .get_or_insert_with(|| first.as_arrow().to_data()) + .ptr_eq(&batch.as_arrow().to_data()) + }) + }); let needs_float = chunks.iter().any(|chunk| { matches!( chunk.storage, diff --git a/src/datatypes/src/series/samples/tests.rs b/src/datatypes/src/series/samples/tests.rs index 7cfb56bcd03..62539f885fa 100644 --- a/src/datatypes/src/series/samples/tests.rs +++ b/src/datatypes/src/series/samples/tests.rs @@ -15,8 +15,8 @@ use super::*; use std::sync::atomic::{AtomicUsize, Ordering}; -use arrow::array::{Float64Array, UInt32Array}; -use arrow::buffer::{Buffer, ScalarBuffer}; +use arrow::array::{ArrayRef, Float64Array, Int64Array, StructArray, UInt32Array}; +use arrow::buffer::{BooleanBuffer, Buffer, NullBuffer, ScalarBuffer}; use crate::series::histogram::{BucketSideView, CounterResetHint, HistogramBatchBuilder}; @@ -235,6 +235,308 @@ fn chunk_iterator_is_sequential_exact_and_handles_many_singletons() { assert!(iter.next().is_none()); } +#[test] +fn densification_reuses_distinct_wrappers_of_the_same_arrow_table_view() { + let base_row = HistogramView { + schema: 2, + zero_threshold: 0.25, + sum: 7.0, + count: 8.0, + zero_count: 1.0, + reset_hint: CounterResetHint::Gauge, + start_timestamp: None, + custom_values: &[], + positive: BucketSideView { + span_offsets: &[0], + span_lengths: &[1], + buckets: &[9.0], + }, + negative: BucketSideView { + span_offsets: &[], + span_lengths: &[], + buckets: &[], + }, + }; + let rows = [ + HistogramView { + schema: 1, + zero_threshold: 0.0, + sum: 1.0, + count: 2.0, + zero_count: 0.0, + reset_hint: CounterResetHint::Unknown, + start_timestamp: Some(11), + custom_values: &[1.0], + positive: BucketSideView { + span_offsets: &[0], + span_lengths: &[1], + buckets: &[3.0], + }, + negative: BucketSideView { + span_offsets: &[], + span_lengths: &[], + buckets: &[], + }, + }, + HistogramView { + sum: -0.0, + count: f64::NAN, + start_timestamp: Some(42), + custom_values: &[f64::from_bits(0x7ff8_0000_0000_0042)], + ..base_row + }, + ]; + let backing = histogram_batch(&[rows[0], rows[1], rows[1]]); + let original = + Arc::new(HistogramBatch::try_from_arrow(backing.as_arrow().slice(1, 1)).unwrap()); + let wrapped = Arc::new(HistogramBatch::try_from_arrow(original.as_arrow().clone()).unwrap()); + assert!(!Arc::ptr_eq(&original, &wrapped)); + assert!( + original + .as_arrow() + .to_data() + .ptr_eq(&wrapped.as_arrow().to_data()) + ); + + let mixed = Samples::try_from_mixed( + Float64Array::from(vec![ + Some(-0.0), + Some(f64::from_bits(0x7ff8_0000_0000_0042)), + None, + ]), + HistogramSamples::try_new(UInt32Array::from(vec![None, None, Some(0)]), wrapped).unwrap(), + ) + .unwrap(); + let dense = Samples::try_from_chunks(vec![ + Samples::try_from_histogram(UInt32Array::from(vec![Some(0), None]), original.clone()) + .unwrap(), + mixed, + ]) + .unwrap() + .into_dense() + .unwrap(); + let DenseSamplesRef::Mixed(values) = dense.dense().unwrap() else { + panic!() + }; + assert_eq!(values.histograms.batch().len(), 1); + assert!(Arc::ptr_eq(values.histograms.batch(), &original)); + assert_eq!(values.histograms.indices().value(0), 0); + assert!(values.histograms.indices().is_null(1)); + assert!(values.histograms.indices().is_null(2)); + assert!(values.histograms.indices().is_null(3)); + assert_eq!(values.histograms.indices().value(4), 0); + let expected = [ + SampleRef::Histogram(original.view(0).unwrap()), + SampleRef::Missing, + SampleRef::Float(-0.0), + SampleRef::Float(f64::from_bits(0x7ff8_0000_0000_0042)), + SampleRef::Histogram(original.view(0).unwrap()), + ]; + for (index, expected) in expected.into_iter().enumerate() { + assert!( + same_sample(dense.sample(index), expected), + "position {index}" + ); + } +} + +#[test] +fn densification_does_not_merge_equal_content_distinct_buffers() { + let first = batch(2.0, 3.0, &[4.0]); + let equal = batch(2.0, 3.0, &[4.0]); + assert!( + !first + .as_arrow() + .to_data() + .ptr_eq(&equal.as_arrow().to_data()) + ); + let dense = Samples::try_from_chunks(vec![ + Samples::try_from_histogram(UInt32Array::from(vec![Some(0)]), first.clone()).unwrap(), + Samples::try_from_histogram(UInt32Array::from(vec![Some(0)]), equal.clone()).unwrap(), + ]) + .unwrap() + .into_dense() + .unwrap(); + let DenseSamplesRef::Histogram(values) = dense.dense().unwrap() else { + panic!() + }; + assert_eq!(values.batch().len(), 2); + assert_eq!(values.indices().value(0), 0); + assert_eq!(values.indices().value(1), 1); + assert!(same_sample( + dense.sample(0), + SampleRef::Histogram(first.view(0).unwrap()) + )); + assert!(same_sample( + dense.sample(1), + SampleRef::Histogram(equal.view(0).unwrap()) + )); +} + +#[test] +fn densification_does_not_merge_distinct_table_slices() { + let base_row = HistogramView { + schema: 1, + zero_threshold: 0.0, + sum: 2.0, + count: 3.0, + zero_count: 0.0, + reset_hint: CounterResetHint::Unknown, + start_timestamp: None, + custom_values: &[], + positive: BucketSideView { + span_offsets: &[0], + span_lengths: &[1], + buckets: &[4.0], + }, + negative: BucketSideView { + span_offsets: &[], + span_lengths: &[], + buckets: &[], + }, + }; + let rows = [ + base_row, + HistogramView { + start_timestamp: Some(42), + sum: 12.0, + ..base_row + }, + HistogramView { + start_timestamp: Some(99), + sum: 19.0, + ..base_row + }, + ]; + let backing = histogram_batch(&rows); + let first = Arc::new(HistogramBatch::try_from_arrow(backing.as_arrow().slice(0, 1)).unwrap()); + let equal_len_other_offset = + Arc::new(HistogramBatch::try_from_arrow(backing.as_arrow().slice(1, 1)).unwrap()); + let longer_other_offset = + Arc::new(HistogramBatch::try_from_arrow(backing.as_arrow().slice(1, 2)).unwrap()); + + for (second, expected_table_len, second_expected) in [ + ( + equal_len_other_offset, + 2, + SampleRef::Histogram(backing.view(1).unwrap()), + ), + ( + longer_other_offset, + 3, + SampleRef::Histogram(backing.view(1).unwrap()), + ), + ] { + let dense = Samples::try_from_chunks(vec![ + Samples::try_from_histogram(UInt32Array::from(vec![Some(0)]), first.clone()).unwrap(), + Samples::try_from_histogram(UInt32Array::from(vec![Some(0)]), second).unwrap(), + ]) + .unwrap() + .into_dense() + .unwrap(); + let DenseSamplesRef::Histogram(values) = dense.dense().unwrap() else { + panic!() + }; + assert_eq!(values.batch().len(), expected_table_len); + assert_eq!(values.indices().value(0), 0); + assert_eq!(values.indices().value(1), 1); + assert!(same_sample( + dense.sample(0), + SampleRef::Histogram(first.view(0).unwrap()) + )); + assert!(same_sample(dense.sample(1), second_expected)); + } +} + +#[test] +fn densification_does_not_merge_shared_value_buffer_with_different_null_bitmap_offsets() { + let source = batch(2.0, 3.0, &[4.0]); + let start_values = ScalarBuffer::from(vec![42_i64]); + let bitmap = Buffer::from(vec![0b10_u8]); + let missing_starts = Int64Array::new( + start_values.clone(), + Some(NullBuffer::new(BooleanBuffer::new(bitmap.clone(), 0, 1))), + ); + let present_starts = Int64Array::new( + start_values.clone(), + Some(NullBuffer::new(BooleanBuffer::new(bitmap.clone(), 1, 1))), + ); + assert_eq!( + missing_starts.values().inner().as_ptr(), + present_starts.values().inner().as_ptr() + ); + let missing_nulls = missing_starts.nulls().unwrap(); + let present_nulls = present_starts.nulls().unwrap(); + assert_eq!( + missing_nulls.buffer().as_ptr(), + present_nulls.buffer().as_ptr() + ); + assert_ne!(missing_nulls.offset(), present_nulls.offset()); + assert!(missing_starts.is_null(0)); + assert_eq!(present_starts.value(0), 42); + + let with_start = |starts: Int64Array| { + let mut columns = source.as_arrow().columns().to_vec(); + columns[6] = Arc::new(starts) as ArrayRef; + Arc::new( + HistogramBatch::try_from_arrow( + StructArray::try_new(source.as_arrow().fields().clone(), columns, None).unwrap(), + ) + .unwrap(), + ) + }; + let missing = with_start(missing_starts); + let present = with_start(present_starts); + let missing_starts = missing + .as_arrow() + .column(6) + .as_any() + .downcast_ref::() + .unwrap(); + let present_starts = present + .as_arrow() + .column(6) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!( + missing_starts.values().inner().as_ptr(), + present_starts.values().inner().as_ptr() + ); + assert_eq!( + missing_starts.nulls().unwrap().buffer().as_ptr(), + present_starts.nulls().unwrap().buffer().as_ptr() + ); + assert_ne!( + missing_starts.nulls().unwrap().offset(), + present_starts.nulls().unwrap().offset() + ); + assert!(missing_starts.is_null(0)); + assert_eq!(present_starts.value(0), 42); + assert!( + !missing + .as_arrow() + .to_data() + .ptr_eq(&present.as_arrow().to_data()) + ); + let dense = Samples::try_from_chunks(vec![ + Samples::try_from_histogram(UInt32Array::from(vec![Some(0)]), missing.clone()).unwrap(), + Samples::try_from_histogram(UInt32Array::from(vec![Some(0)]), present.clone()).unwrap(), + ]) + .unwrap() + .into_dense() + .unwrap(); + assert_eq!(dense.len(), 2); + assert!(same_sample( + dense.sample(0), + SampleRef::Histogram(missing.view(0).unwrap()) + )); + assert!(same_sample( + dense.sample(1), + SampleRef::Histogram(present.view(0).unwrap()) + )); +} + #[test] fn densifying_slices_keeps_shared_histogram_table() { let table = batch(2.0, 3.0, &[4.0, 5.0]);