From 69c73740ecc898c30238a49c878476a1691130ea Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Thu, 15 Oct 2026 15:35:59 +0800 Subject: [PATCH 01/11] feat(datatypes): add raw series frame codec Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/datatypes/src/series.rs | 1 + src/datatypes/src/series/frame.rs | 552 ++++++ src/datatypes/src/series/frame/tests.rs | 2213 +++++++++++++++++++++++ src/datatypes/src/series/samples.rs | 40 + 4 files changed, 2806 insertions(+) create mode 100644 src/datatypes/src/series/frame.rs create mode 100644 src/datatypes/src/series/frame/tests.rs diff --git a/src/datatypes/src/series.rs b/src/datatypes/src/series.rs index ad3bf52aa52..f00051ce46f 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..2fd242c9b76 --- /dev/null +++ b/src/datatypes/src/series/frame.rs @@ -0,0 +1,552 @@ +// 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 are ordinal (`f0`, `f1`, ...); +//! callers bind those positions to their logical names, which this frame cannot authenticate. + +#[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::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, RawTimeAxis}; +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 using ordinal domain fields. +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() + .enumerate() + .map(|(i, _)| Field::new(format!("f{i}"), 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))) +} + +/// Packs one raw chunk into one-row batches at the union of field segment boundaries. +pub fn pack_raw_series_chunk( + schema: &SeriesSchema, + mode: RawLabelMode, + chunk: &RawSeriesChunk, +) -> Result> { + let checked = RawSeriesChunk::try_new( + schema, + chunk.read_key.clone(), + chunk.metric.clone(), + chunk.time.clone(), + chunk.fields.clone(), + chunk.is_last, + )?; + match (mode, &checked.metric) { + (RawLabelMode::Key, None) | (RawLabelMode::Explicit, Some(_)) => {} + _ => return invalid_series("raw chunk label mode does not match frame mode"), + } + let frame_schema = raw_series_frame_schema(schema, mode)?; + let RawTimeAxis::Explicit(times) = &checked.time; + let all: Vec<&[Samples]> = checked.fields.iter().map(Samples::segments).collect(); + let mut segment = vec![0usize; all.len()]; + let mut offset = vec![0usize; all.len()]; + let mut pos = 0usize; + let mut batches = Vec::new(); + while pos < times.len() { + let take = all + .iter() + .enumerate() + .map(|(i, parts)| parts[segment[i]].len() - offset[i]) + .min() + .unwrap_or(times.len()) + .min(times.len() - pos); + if take == 0 { + return invalid_series("empty sample segment"); + } + let pieces = all + .iter() + .enumerate() + .map(|(i, parts)| parts[segment[i]].slice(offset[i], take)) + .collect::>>()?; + batches.push(encode_row( + &frame_schema, + mode, + &checked.read_key, + checked.metric.as_ref(), + checked.is_last && pos + take == times.len(), + times.slice(pos, take), + &pieces, + )?); + pos += take; + for (i, parts) in all.iter().enumerate() { + offset[i] += take; + if offset[i] == parts[segment[i]].len() { + segment[i] += 1; + offset[i] = 0; + } + } + } + Ok(batches) +} + +/// 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> { + let expected = raw_series_frame_schema(schema, mode)?; + 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)?); + let hs = HistogramSamples::try_new(indices, batch)?; + if has_float { + Samples::try_from_mixed(floats, hs)? + } else { + Samples::try_from_histogram(hs.indices().clone(), hs.batch().clone())? + } + }; + decoded.push(samples); + } + let fields = SeriesFields::try_from_vec(decoded)?; + let key = binary_slice(keys, row)?; + let raw_time = RawTimeAxis::try_new(time)?; + out.push(RawSeriesChunk::try_new( + schema, + key, + metric, + raw_time, + fields, + lasts.value(row), + )?); + } + Ok(out) +} + +fn encode_row( + schema: &SchemaRef, + mode: RawLabelMode, + key: &Buffer, + metric: Option<&MetricName>, + last: bool, + times: ScalarBuffer, + samples: &[Samples], +) -> Result { + let mut arrays: Vec = vec![ + Arc::new(binary_one(key.as_slice())?), + Arc::new(BooleanArray::from(vec![last])), + ]; + let ts = wrap_list( + Arc::new(Int64Array::new(times, None)), + DataType::Int64, + false, + )?; + let DataType::Struct(time_fields) = schema.field(2).data_type() else { + return invalid_series("time schema mismatch"); + }; + arrays.push(Arc::new(struct_array( + time_fields.clone(), + vec![Arc::new(ts)], + )?)); + let DataType::Struct(field_schema) = schema.field(3).data_type() else { + return invalid_series("fields schema mismatch"); + }; + let DataType::Struct(first_payload) = field_schema[0].data_type() else { + return invalid_series("payload schema mismatch"); + }; + let DataType::Struct(histogram_fields) = first_payload[1].data_type() else { + return invalid_series("histogram schema mismatch"); + }; + let mut field_arrays = Vec::new(); + for sample in samples { + let DataType::Struct(dtype) = field_schema[field_arrays.len()].data_type() else { + return invalid_series("payload schema mismatch"); + }; + let (floats, hist) = encode_payload(sample, histogram_fields)?; + field_arrays.push(Arc::new(struct_array(dtype.clone(), vec![floats, hist])?) as ArrayRef); + } + arrays.push(Arc::new(struct_array(field_schema.clone(), field_arrays)?)); + if mode == RawLabelMode::Explicit { + let metric = metric.ok_or_else(|| err("explicit labels missing"))?; + checked_offset(metric.name().map_or(0, str::len))?; + arrays.push(Arc::new(StringArray::from(vec![metric.name()]))); + checked_offset(metric.tags().len())?; + checked_offset( + metric + .tags() + .iter() + .try_fold(0usize, |n, l| n.checked_add(l.name.len())) + .ok_or_else(|| err("label name length overflow"))?, + )?; + checked_offset( + metric + .tags() + .iter() + .try_fold(0usize, |n, l| n.checked_add(l.value.len())) + .ok_or_else(|| err("label value length overflow"))?, + )?; + let names: StringArray = metric + .tags() + .iter() + .map(|l| Some(l.name.as_ref())) + .collect(); + let values: StringArray = metric + .tags() + .iter() + .map(|l| Some(l.value.as_ref())) + .collect(); + let DataType::List(label_list) = schema.field(5).data_type() else { + return invalid_series("labels schema mismatch"); + }; + let DataType::Struct(label_fields) = label_list.data_type() else { + return invalid_series("label struct schema mismatch"); + }; + let labels = struct_array( + label_fields.clone(), + vec![Arc::new(names), Arc::new(values)], + )?; + arrays.push(Arc::new(wrap_list( + Arc::new(labels), + DataType::Struct(label_fields.clone()), + false, + )?)); + } + arrow(RecordBatch::try_new_with_options( + schema.clone(), + arrays, + &RecordBatchOptions::new().with_row_count(Some(1)), + )) +} + +fn encode_payload(samples: &Samples, histogram_fields: &Fields) -> Result<(ArrayRef, ArrayRef)> { + let (floats, histograms) = match samples + .dense() + .ok_or_else(|| err("sample segment is not dense"))? + { + DenseSamplesRef::Float(values) => (Some(values), None), + DenseSamplesRef::Histogram(values) => (None, Some(values)), + DenseSamplesRef::Mixed(values) => (Some(values.floats()), Some(values.histograms())), + }; + let float_array = match floats { + Some(values) => Arc::new(wrap_list( + Arc::new(values.clone()), + DataType::Float64, + true, + )?) as ArrayRef, + None => Arc::new(ListArray::new_null( + Arc::new(Field::new("item", DataType::Float64, true)), + 1, + )) as ArrayRef, + }; + let hist_array = match histograms { + Some(values) => { + let DataType::List(indices_field) = histogram_fields[0].data_type() else { + return invalid_series("histogram index schema mismatch"); + }; + let DataType::List(table_field) = histogram_fields[1].data_type() else { + return invalid_series("histogram table schema mismatch"); + }; + let indices = wrap_list( + Arc::new(values.indices().clone()), + indices_field.data_type().clone(), + true, + )?; + let table = wrap_list( + Arc::new(values.batch().as_arrow().clone()), + table_field.data_type().clone(), + false, + )?; + Arc::new(struct_array( + histogram_fields.clone(), + vec![Arc::new(indices), Arc::new(table)], + )?) as ArrayRef + } + None => null_histogram(histogram_fields)?, + }; + Ok((float_array, hist_array)) +} + +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 binary_one(value: &[u8]) -> Result { + checked_offset(value.len())?; + Ok(BinaryArray::from(vec![Some(value)])) +} +fn checked_offset(len: usize) -> Result { + i32::try_from(len).map_err(|_| err("Arrow offset limit exceeded")) +} +fn offsets(len: usize) -> Result> { + let len = checked_offset(len)?; + Ok(OffsetBuffer::new(ScalarBuffer::from(vec![0, len]))) +} +fn wrap_list(values: ArrayRef, item: DataType, nullable: bool) -> Result { + let len = values.len(); + arrow(ListArray::try_new( + Arc::new(Field::new("item", item, nullable)), + offsets(len)?, + values, + None, + )) +} +fn null_histogram(fields: &Fields) -> Result { + let DataType::List(indices_field) = fields[0].data_type() else { + return invalid_series("histogram index schema mismatch"); + }; + let DataType::List(table_field) = fields[1].data_type() else { + return invalid_series("histogram table schema mismatch"); + }; + let DataType::Struct(table_fields) = table_field.data_type() else { + return invalid_series("histogram table schema mismatch"); + }; + let indices = wrap_list( + Arc::new(UInt32Array::new_null(0)), + indices_field.data_type().clone(), + true, + )?; + let columns = table_fields + .iter() + .map(|f| arrow::array::new_null_array(f.data_type(), 0)) + .collect(); + let table = arrow(StructArray::try_new(table_fields.clone(), columns, None))?; + let table = wrap_list(Arc::new(table), table_field.data_type().clone(), false)?; + Ok(Arc::new(arrow(StructArray::try_new( + fields.clone(), + vec![Arc::new(indices), Arc::new(table)], + Some(NullBuffer::new_null(1)), + ))?)) +} +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..beb5f9ae0ec --- /dev/null +++ b/src/datatypes/src/series/frame/tests.rs @@ -0,0 +1,2213 @@ +// 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, Float64Array, Int64Array, ListArray, StringArray, StructArray, UInt32Array, +}; +use arrow::buffer::{Buffer, OffsetBuffer, ScalarBuffer}; +use arrow::compute::concat_batches; +use arrow::datatypes::{DataType, Field, Schema}; + +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, + RawTimeAxis::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}"), + } + } +} + +#[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("f0", payload, false)].into()), + false, + ), + ], + expected_key_metadata, + ); + assert_eq!(key_schema.as_ref(), &expected_key); + assert_eq!(key_schema.metadata().len(), 4); + assert_eq!( + key_schema.metadata().get("greptime.series.layout").unwrap(), + "v1" + ); + assert_eq!( + key_schema + .metadata() + .get("greptime.series.content") + .unwrap(), + "raw" + ); + assert_eq!( + key_schema.metadata().get("greptime.series.labels").unwrap(), + "key" + ); + assert_eq!( + key_schema.metadata().get("greptime.series.time").unwrap(), + "timestamps" + ); + 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); + + assert_eq!( + expected_key + .metadata() + .get("greptime.series.labels") + .unwrap(), + "key" + ); + assert_eq!( + expected_key.metadata().get("greptime.series.time").unwrap(), + "timestamps" + ); + + 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].read_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].read_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, + RawTimeAxis::try_new(vec![10_i64, 20, 30, 40].into()).unwrap(), + SeriesFields::try_from_vec(vec![first, second]).unwrap(), + true, + ) + .unwrap(); + let frames = pack_raw_series_chunk(&schema, RawLabelMode::Key, &chunk).unwrap(); + assert_eq!(frames.len(), 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(); + for (frame_index, frame) in frames.iter().enumerate() { + let decoded = unpack_raw_series_frame(&schema, RawLabelMode::Key, frame).unwrap(); + assert_eq!(decoded.len(), 1); + assert_eq!(decoded[0].is_last, frame_index == frames.len() - 1); + observed_markers.push(decoded[0].is_last); + assert_eq!(decoded[0].read_key.as_slice(), &[0, 0x91, 0xfe]); + observed_keys.push(decoded[0].read_key.as_slice().to_vec()); + for position in 0..decoded[0].time.len() { + observed_times.push(decoded[0].time.timestamp(position).unwrap()); + match decoded[0].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 decoded[0].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, + RawTimeAxis::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(); + 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, batch) in aligned_frames.iter().enumerate() { + let decoded = unpack_raw_series_frame(&aligned_schema, RawLabelMode::Key, batch).unwrap(); + assert_eq!(decoded[0].is_last, i == 2); + interval_lengths.push(decoded[0].time.len()); + let start = aligned_times.len(); + let DenseSamplesRef::Float(decoded_a) = decoded[0].fields.get(0).unwrap().dense().unwrap() + else { + panic!("expected float field"); + }; + let DenseSamplesRef::Float(decoded_b) = decoded[0].fields.get(1).unwrap().dense().unwrap() + else { + panic!("expected float field"); + }; + assert_eq!( + decoded_a.values().as_ptr(), + a_owner.values().as_ptr().wrapping_add(start) + ); + assert_eq!( + decoded_b.values().as_ptr(), + b_owner.values().as_ptr().wrapping_add(start) + ); + for (decoded_values, owner) in [(&decoded_a, &a_owner), (&decoded_b, &b_owner)] { + match (decoded_values.nulls(), owner.nulls()) { + (Some(decoded_nulls), Some(owner_nulls)) => { + assert_eq!( + decoded_nulls.buffer().as_ptr(), + owner_nulls.buffer().as_ptr() + ); + assert_eq!(decoded_nulls.offset(), owner_nulls.offset() + start); + } + (None, None) => {} + _ => panic!("fragment changed validity representation"), + } + } + for position in 0..decoded[0].time.len() { + aligned_times.push(decoded[0].time.timestamp(position).unwrap()); + aligned_a.push(match decoded[0].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 decoded[0].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); + 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()) + ] + ); + + 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, + RawTimeAxis::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].read_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.read_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, + RawTimeAxis::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, + RawTimeAxis::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.batch().len(), 2); + assert_eq!(histograms.indices().len(), 2); + assert!(histograms.indices().is_valid(0)); + assert!(histograms.indices().is_valid(1)); + assert_eq!(histograms.indices().value(0), 0); + assert_eq!(histograms.indices().value(1), 0); + assert_histogram(histograms.batch().view(0).unwrap(), -9.5); + assert_unreferenced_histogram(histograms.batch().view(1).unwrap()); + + 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.batch().len(), 2); + assert!((0..2).all(|i| !all_null_histograms.indices().is_valid(i))); + assert_histogram(all_null_histograms.batch().view(0).unwrap(), -9.5); + assert_unreferenced_histogram(all_null_histograms.batch().view(1).unwrap()); + 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, + RawTimeAxis::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().batch().len(), 2); + assert!( + (0..2).all(|i| !mixed.floats().is_valid(i) && !mixed.histograms().indices().is_valid(i)) + ); + assert_histogram(mixed.histograms().batch().view(0).unwrap(), -9.5); + assert_unreferenced_histogram(mixed.histograms().batch().view(1).unwrap()); +} + +#[test] +fn rejects_grid_and_invalid_schema_even_for_empty_batch_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()); + } + + 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 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()); + + 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()); + + 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, + RawTimeAxis::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, + RawTimeAxis::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 crate::series::raw::RawTimeAxis::Explicit(decoded_time) = &decoded[0].time; + assert_eq!(decoded_time.as_ptr(), time_values.values().as_ptr()); + let crate::series::raw::RawTimeAxis::Explicit(source_time) = &chunk.time; + 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].read_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, + RawTimeAxis::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 crate::series::raw::RawTimeAxis::Explicit(source_times) = &chunk.time; + + let frames = pack_raw_series_chunk(&schema, RawLabelMode::Key, &chunk).unwrap(); + 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 crate::series::raw::RawTimeAxis::Explicit(decoded_time_values) = &decoded[0].time; + 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(_)) + )); +} diff --git a/src/datatypes/src/series/samples.rs b/src/datatypes/src/series/samples.rs index a19329c62e0..5bc5b5e3edc 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 { From 3717cc2e2a14f2e53ea286e4c1d6ac015f57ccd4 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Thu, 15 Oct 2026 16:31:56 +0800 Subject: [PATCH 02/11] fix(datatypes): retain histogram table views across frame densification Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/datatypes/src/series/frame/tests.rs | 203 ++++++++------ src/datatypes/src/series/samples.rs | 15 +- src/datatypes/src/series/samples/tests.rs | 306 +++++++++++++++++++++- 3 files changed, 439 insertions(+), 85 deletions(-) diff --git a/src/datatypes/src/series/frame/tests.rs b/src/datatypes/src/series/frame/tests.rs index beb5f9ae0ec..6247c4924ee 100644 --- a/src/datatypes/src/series/frame/tests.rs +++ b/src/datatypes/src/series/frame/tests.rs @@ -247,6 +247,98 @@ fn assert_float_bits(actual: &Samples, expected: &[Option]) { } } +#[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, + RawTimeAxis::try_new((0_i64..64).collect::>().into()).unwrap(), + SeriesFields::one(samples), + true, + ) + .unwrap(); + + let frames = pack_raw_series_chunk(&schema, RawLabelMode::Key, &chunk).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() + ); + + 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 schema_and_key_explicit_round_trips_preserve_identity_and_bits() { let schema = test_schema(&["value"]); @@ -315,26 +407,6 @@ fn schema_and_key_explicit_round_trips_preserve_identity_and_bits() { expected_key_metadata, ); assert_eq!(key_schema.as_ref(), &expected_key); - assert_eq!(key_schema.metadata().len(), 4); - assert_eq!( - key_schema.metadata().get("greptime.series.layout").unwrap(), - "v1" - ); - assert_eq!( - key_schema - .metadata() - .get("greptime.series.content") - .unwrap(), - "raw" - ); - assert_eq!( - key_schema.metadata().get("greptime.series.labels").unwrap(), - "key" - ); - assert_eq!( - key_schema.metadata().get("greptime.series.time").unwrap(), - "timestamps" - ); let mut explicit_fields = expected_key.fields().iter().cloned().collect::>(); explicit_fields.push(Field::new("name", DataType::Utf8, true).into()); explicit_fields.push( @@ -360,18 +432,6 @@ fn schema_and_key_explicit_round_trips_preserve_identity_and_bits() { let expected_explicit = Schema::new_with_metadata(explicit_fields, explicit_metadata); assert_eq!(explicit_schema.as_ref(), &expected_explicit); - assert_eq!( - expected_key - .metadata() - .get("greptime.series.labels") - .unwrap(), - "key" - ); - assert_eq!( - expected_key.metadata().get("greptime.series.time").unwrap(), - "timestamps" - ); - 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)]; @@ -628,26 +688,6 @@ fn segmented_mixed_fields_emit_every_aligned_position_and_marker() { ); drop(aligned_chunk); drop(aligned_frames); - 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()) - ] - ); let mut distinct_rows = frames.clone(); let alternate_schema = test_schema(&["first", "second"]); @@ -817,24 +857,19 @@ fn histogram_payload_and_present_all_null_representation_survive() { else { panic!("histogram representation was normalized"); }; - assert_eq!(histograms.batch().len(), 2); assert_eq!(histograms.indices().len(), 2); assert!(histograms.indices().is_valid(0)); assert!(histograms.indices().is_valid(1)); - assert_eq!(histograms.indices().value(0), 0); - assert_eq!(histograms.indices().value(1), 0); - assert_histogram(histograms.batch().view(0).unwrap(), -9.5); - assert_unreferenced_histogram(histograms.batch().view(1).unwrap()); + 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.batch().len(), 2); + assert_eq!(all_null_histograms.indices().len(), 2); assert!((0..2).all(|i| !all_null_histograms.indices().is_valid(i))); - assert_histogram(all_null_histograms.batch().view(0).unwrap(), -9.5); - assert_unreferenced_histogram(all_null_histograms.batch().view(1).unwrap()); let DenseSamplesRef::Float(all_null_float) = decoded[0].fields.get(2).unwrap().dense().unwrap() else { panic!("all-null float representation was normalized"); @@ -877,16 +912,14 @@ fn histogram_payload_and_present_all_null_representation_survive() { panic!("all-null mixed representation was normalized"); }; assert_eq!(mixed.floats().len(), 2); - assert_eq!(mixed.histograms().batch().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)) ); - assert_histogram(mixed.histograms().batch().view(0).unwrap(), -9.5); - assert_unreferenced_histogram(mixed.histograms().batch().view(1).unwrap()); } #[test] -fn rejects_grid_and_invalid_schema_even_for_empty_batch_and_accepts_valid_empty() { +fn rejects_invalid_schema_and_metadata_and_accepts_valid_empty() { let schema = test_schema(&["value"]); let grid_schema = SeriesSchema { fields: schema.fields.clone(), @@ -1000,7 +1033,15 @@ fn rejects_grid_and_invalid_schema_even_for_empty_batch_and_accepts_valid_empty( .unwrap(); assert!(unpack_raw_series_frame(&schema, RawLabelMode::Key, &malformed).is_err()); } +} +#[test] +fn rejects_invalid_timestamps_and_empty_raw_payloads() { + 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()); @@ -1107,6 +1148,23 @@ fn rejects_grid_and_invalid_schema_even_for_empty_batch_and_accepts_valid_empty( // 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, @@ -1213,21 +1271,10 @@ fn rejects_grid_and_invalid_schema_even_for_empty_batch_and_accepts_valid_empty( 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()); +} - 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_malformed_histogram_and_payload_encodings() { let hist_schema = test_schema(&["hist"]); let valid_hist = pack_raw_series_chunk( &hist_schema, diff --git a/src/datatypes/src/series/samples.rs b/src/datatypes/src/series/samples.rs index eb02ccfd4db..6b2056df691 100644 --- a/src/datatypes/src/series/samples.rs +++ b/src/datatypes/src/series/samples.rs @@ -355,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); @@ -390,9 +389,15 @@ 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 shared_batch = histogram_batches.next().filter(|first| { + histogram_batches.all(|batch| { + Arc::ptr_eq(first, batch) + || 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 3f672ed418a..b623937945d 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]); From 696259a1ce3050b9656a427d3e10ff11ce99637a Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Thu, 15 Oct 2026 17:16:16 +0800 Subject: [PATCH 03/11] feat(query): add incremental raw series frame streams Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/common/query/Cargo.toml | 2 +- src/common/query/src/lib.rs | 1 + src/common/query/src/series_frame.rs | 363 +++++ src/common/query/src/series_frame/tests.rs | 1186 +++++++++++++++++ .../query/src/series_frame/tests/lifecycle.rs | 322 +++++ 5 files changed, 1873 insertions(+), 1 deletion(-) create mode 100644 src/common/query/src/series_frame.rs create mode 100644 src/common/query/src/series_frame/tests.rs create mode 100644 src/common/query/src/series_frame/tests/lifecycle.rs 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..47c81cc49fc --- /dev/null +++ b/src/common/query/src/series_frame.rs @@ -0,0 +1,363 @@ +// 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 codec layout and ordinal payload fields; 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. + +use std::collections::HashSet; +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use common_error::ext::BoxedError; +use common_recordbatch::{DfRecordBatch, DfSendableRecordBatchStream}; +use datafusion_common::DataFusionError; +use datatypes::error::InvalidSeriesSnafu; +use datatypes::series::MetricName; +use datatypes::series::frame::{ + RawLabelMode, pack_raw_series_chunk, raw_series_frame_schema, unpack_raw_series_frame, +}; +use datatypes::series::raw::{RawSeriesChunk, RawSeriesOrdering}; +use datatypes::series::samples::SeriesSchema; +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. +pub fn pack_raw_series( + source: SendableRawSeriesStream, +) -> datafusion_common::Result { + let schema = source.schema().clone(); + let label_source = source.label_source().clone(); + let ordering = source.ordering(); + let mode = label_mode(&label_source); + let frame_schema = raw_series_frame_schema(&schema, mode).map_err(codec_error)?; + Ok(Box::pin(PackStream { + source: Some(source), + schema, + frame_schema, + mode, + validator: SequenceValidator::new(ordering), + pending: Vec::::new().into_iter(), + })) +} + +/// Unpacks DataFusion batches using the caller's explicit scan binding and 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: RawSeriesOrdering, +) -> datafusion_common::Result { + let mode = label_mode(&label_source); + let frame_schema = source.schema(); + // Validate even empty streams' declared schema and layout before returning the stream. + let empty = DfRecordBatch::new_empty(frame_schema.clone()); + unpack_raw_series_frame(&expected, mode, &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(), + })) +} + +fn label_mode(source: &RawLabelSource) -> RawLabelMode { + match source { + RawLabelSource::PrimaryKey { .. } => RawLabelMode::Key, + RawLabelSource::Explicit => RawLabelMode::Explicit, + } +} + +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. Contiguous ordering retains complete owned keys for recurrence +/// detection, so memory is O(total distinct complete-key bytes + distinct keys), plus current-key +/// and metric state, until stream termination; there is no eviction or cap. This owns key bytes, +/// never Arrow buffers that might pin larger payloads. +struct SequenceValidator { + ordering: RawSeriesOrdering, + previous_key: Option>, + previous_timestamp: Option, + previous_metric: Option, + previous_is_last: bool, + closed: HashSet>, +} + +impl SequenceValidator { + fn new(ordering: RawSeriesOrdering) -> Self { + Self { + ordering, + previous_key: None, + previous_timestamp: None, + previous_metric: None, + previous_is_last: false, + closed: HashSet::new(), + } + } + + 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.read_key.as_slice(); + let timestamp = chunk + .time + .timestamp(chunk.time.len() - 1) + .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 == RawSeriesOrdering::ReadKeyThenTime + && previous.as_slice() > key + { + return Err(invalid("raw series keys are not ordered")); + } + if self.ordering == RawSeriesOrdering::Contiguous && self.closed.contains(key) { + return Err(invalid("raw series key recurred after its contiguous run")); + } + } + } + } + let owned_key = key.to_vec(); + if chunk.is_last && self.ordering == RawSeriesOrdering::Contiguous { + self.closed.insert(owned_key.clone()); + } + self.previous_key = Some(owned_key); + 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, + frame_schema: datatypes::arrow::datatypes::SchemaRef, + mode: RawLabelMode, + validator: SequenceValidator, + pending: std::vec::IntoIter, +} + +impl PackStream { + fn finish(&mut self) { + self.source = None; + self.pending = Vec::new().into_iter(); + self.validator = SequenceValidator::new(self.validator.ordering); + } +} + +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 { + if let Some(batch) = this.pending.next() { + return Poll::Ready(Some(Ok(batch))); + } + 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(Err)); + } + Poll::Ready(Some(Err(error))) => { + this.finish(); + return Poll::Ready(Some(Err(upstream_error(error)))); + } + Poll::Ready(Some(Ok(chunk))) => { + let result = (|| { + let batches = pack_raw_series_chunk(&this.schema, this.mode, &chunk) + .map_err(codec_error)?; + this.validator.validate(&chunk)?; + Ok(batches) + })(); + match result { + Ok(batches) => this.pending = batches.into_iter(), + Err(error) => { + this.finish(); + return Poll::Ready(Some(Err(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: RawSeriesOrdering, + validator: SequenceValidator, + pending: std::vec::IntoIter, +} + +impl RawSeriesStream for UnpackStream { + fn schema(&self) -> &Arc { + &self.schema + } + + fn label_source(&self) -> &RawLabelSource { + &self.label_source + } + + fn ordering(&self) -> RawSeriesOrdering { + 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> { + let this = &mut *self; + 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")); + } + unpack_raw_series_frame( + &this.schema, + label_mode(&this.label_source), + &batch, + ) + .map_err(codec_error) + .and_then(|chunks| { + for chunk in &chunks { + this.validator.validate(chunk)?; + } + Ok(chunks) + }) + })(); + match result { + Ok(chunks) => 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..a2c8a8a5ffd --- /dev/null +++ b/src/common/query/src/series_frame/tests.rs @@ -0,0 +1,1186 @@ +// 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::pin::Pin; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::task::{Context, Poll}; + +use common_error::ext::{BoxedError, PlainError}; +use common_error::status_code::StatusCode; +use common_recordbatch::DfSendableRecordBatchStream; +use datafusion::physical_plan::RecordBatchStream; +use datatypes::arrow::array::{Array, Float64Array}; +use datatypes::arrow::buffer::{Buffer, ScalarBuffer}; +use datatypes::arrow::compute::concat_batches; +use datatypes::arrow::datatypes::SchemaRef; +use datatypes::arrow::record_batch::RecordBatch; +use datatypes::series::Label; +use datatypes::series::MetricName; +use datatypes::series::raw::{RawSeriesChunk, RawSeriesOrdering, RawTimeAxis}; +use datatypes::series::samples::{FieldSpec, 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, + RawTimeAxis::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: RawSeriesOrdering, + polls: Arc, + drops: Arc, + pending: bool, + emitted: bool, +} +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> { + self.polls.fetch_add(1, Ordering::SeqCst); + if self.pending && self.emitted { + self.pending = false; + cx.waker().wake_by_ref(); + return Poll::Pending; + } + let item = self.items.pop_front(); + if item.is_some() { + self.emitted = true; + } + Poll::Ready(item) + } +} +impl RawSeriesStream for RawMock { + fn schema(&self) -> &Arc { + &self.schema + } + fn label_source(&self) -> &RawLabelSource { + &self.binding + } + fn ordering(&self) -> RawSeriesOrdering { + self.ordering + } +} +struct DfMock { + schema: SchemaRef, + items: VecDeque>, + polls: Arc, + drops: Arc, + pending: bool, + emitted: bool, +} +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> { + self.polls.fetch_add(1, Ordering::SeqCst); + if self.pending && self.emitted { + self.pending = false; + cx.waker().wake_by_ref(); + return Poll::Pending; + } + let item = self.items.pop_front(); + if item.is_some() { + self.emitted = true; + } + Poll::Ready(item) + } +} +impl RecordBatchStream for DfMock { + fn schema(&self) -> SchemaRef { + self.schema.clone() + } +} +fn raw( + items: Vec, + ordering: RawSeriesOrdering, +) -> (SendableRawSeriesStream, Arc, Arc) { + raw_custom(items, ordering, RawLabelSource::Explicit, false) +} +fn raw_custom( + items: Vec, + ordering: RawSeriesOrdering, + binding: RawLabelSource, + pending: bool, +) -> (SendableRawSeriesStream, Arc, Arc) { + let polls = Arc::new(AtomicUsize::new(0)); + let drops = Arc::new(AtomicUsize::new(0)); + ( + Box::pin(RawMock { + schema: schema(), + items: items.into_iter().map(Ok).collect(), + binding, + ordering, + polls: polls.clone(), + drops: drops.clone(), + pending, + emitted: false, + }), + polls, + drops, + ) +} +fn df_source( + batches: Vec, + declared: SchemaRef, + pending: bool, +) -> ( + DfSendableRecordBatchStream, + Arc, + Arc, +) { + let polls = Arc::new(AtomicUsize::new(0)); + let drops = Arc::new(AtomicUsize::new(0)); + ( + Box::pin(DfMock { + schema: declared, + items: batches.into_iter().map(Ok).collect(), + polls: polls.clone(), + drops: drops.clone(), + pending, + emitted: false, + }), + 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 +} +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.read_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, RawSeriesOrdering::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!( + polls.load(Ordering::SeqCst), + 1, + "first frame must precede drain" + ); + let mut batches = vec![first]; + while let Some(b) = packed.next().await { + batches.push(b.unwrap()); + } + 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: false, + emitted: false, + }); + let mut output = unpack_raw_series( + source, + schema(), + RawLabelSource::Explicit, + RawSeriesOrdering::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 read_key_then_time_valid_roundtrip() { + let (input, _, _) = raw( + vec![chunk(b"a", &[1], true), chunk(b"b", &[2], true)], + RawSeriesOrdering::ReadKeyThenTime, + ); + let batches = pack_all(input).await; + 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: false, + emitted: false, + }); + let mut out = unpack_raw_series( + source, + schema(), + RawLabelSource::Explicit, + RawSeriesOrdering::ReadKeyThenTime, + ) + .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: RawSeriesOrdering, + expected_prefix_count: usize, + 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: RawSeriesOrdering::Contiguous, + expected_prefix_count: 1, + 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: RawSeriesOrdering::Contiguous, + expected_prefix_count: 1, + prefix: vec![(b"a", 1, false, None)], + }, + Case { + name: "nonlocal_recurrence", + items: vec![ + chunk(b"a", &[1], true), + chunk(b"b", &[2], true), + chunk(b"a", &[3], true), + ], + order: RawSeriesOrdering::Contiguous, + expected_prefix_count: 2, + 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: RawSeriesOrdering::ReadKeyThenTime, + expected_prefix_count: 1, + 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: RawSeriesOrdering::Contiguous, + expected_prefix_count: 1, + 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: RawSeriesOrdering::Contiguous, + expected_prefix_count: 1, + 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: RawSeriesOrdering::Contiguous, + expected_prefix_count: 1, + prefix: vec![(b"a", 1, false, Some("x"))], + }, + ]; + for case in cases { + let _case_name = case.name; + assert_eq!(case.prefix.len(), case.expected_prefix_count); + let (input, _, _) = raw(case.items.clone(), case.order); + let mut s = pack_raw_series(input).unwrap(); + for (key, time, last, host) in &case.prefix { + let batch = s.next().await.unwrap().unwrap(); + let chunks = datatypes::series::frame::unpack_raw_series_frame( + &schema(), + datatypes::series::frame::RawLabelMode::Explicit, + &batch, + ) + .unwrap(); + assert_eq!(chunks.len(), 1); + assert_chunk(&chunks[0], key, &[*time], &[Some(*time as f64)], *last); + assert_metric(&chunks[0], *host); + } + 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()); + + // 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, false); + 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()); + } + let (input, _, _) = raw( + vec![chunk(b"unfinished", &[7], false)], + RawSeriesOrdering::Contiguous, + ); + let mut s = pack_raw_series(input).unwrap(); + assert!(s.next().await.unwrap().is_ok()); + assert!(s.next().await.unwrap().is_err()); + assert!(s.next().await.is_none()); + 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, false); + let mut decoded = unpack_raw_series( + source, + schema(), + RawLabelSource::Explicit, + RawSeriesOrdering::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()); +} + +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, + RawTimeAxis::try_new(ScalarBuffer::from(vec![1])).unwrap(), + SeriesFields::one(Samples::from_float(Float64Array::from(vec![Some(3.0)]))), + true, + ) + .unwrap(), + ], + RawSeriesOrdering::Contiguous, + binding.clone(), + false, + ); + 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: false, + emitted: false, + }); + let mut out = + unpack_raw_series(source, schema(), binding, RawSeriesOrdering::Contiguous).unwrap(); + let decoded = out.next().await.unwrap().unwrap(); + assert_eq!(decoded.read_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, + RawTimeAxis::try_new(ScalarBuffer::from(vec![1])).unwrap(), + SeriesFields::one(Samples::from_float(Float64Array::from(vec![Some(1.0)]))), + true, + ) + .unwrap(), + ], + RawSeriesOrdering::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)], + RawSeriesOrdering::Contiguous, + primary_binding(metadata), + false, + ); + 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() { + let (input, _, _) = raw( + 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), + ], + RawSeriesOrdering::Contiguous, + ); + let mut s = pack_raw_series(input).unwrap(); + assert!(s.next().await.unwrap().is_ok()); + assert!(s.next().await.unwrap().is_err()); + 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: RawSeriesOrdering::Contiguous, + polls, + drops: Arc::new(AtomicUsize::new(0)), + pending: false, + emitted: false, + }); + 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![], RawSeriesOrdering::Contiguous); + let mut packed = pack_raw_series(input).unwrap(); + let frame_schema = packed.schema(); + assert!(packed.next().await.is_none()); + let (source, polls, _) = df_source(vec![], frame_schema.clone(), false); + let mut unpacked = unpack_raw_series( + source, + schema(), + RawLabelSource::Explicit, + RawSeriesOrdering::Contiguous, + ) + .unwrap(); + assert_eq!(polls.load(Ordering::SeqCst), 0); + assert_eq!(unpacked.schema().fields.len(), 1); + 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: RawSeriesOrdering::Contiguous, + polls, + drops: drops.clone(), + pending: false, + emitted: false, + }); + 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!(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: false, + emitted: false, + }); + let mut unpacked = unpack_raw_series( + source, + schema(), + RawLabelSource::Explicit, + RawSeriesOrdering::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_custom( + vec![chunk(b"a", &[1], true), chunk(b"b", &[2], true)], + RawSeriesOrdering::Contiguous, + RawLabelSource::Explicit, + true, + ); + let mut packed = pack_raw_series(input).unwrap(); + assert!(packed.next().await.unwrap().is_ok()); + assert_eq!( + polls.load(Ordering::SeqCst), + 1, + "first item is emitted immediately" + ); + assert!( + packed.next().await.unwrap().is_ok(), + "Pending is resumed before the next item" + ); + assert_eq!(polls.load(Ordering::SeqCst), 3); + 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, false); + assert!( + unpack_raw_series( + source, + schema(), + RawLabelSource::Explicit, + RawSeriesOrdering::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: RawSeriesOrdering::Contiguous, + polls: polls.clone(), + drops: Arc::new(AtomicUsize::new(0)), + pending: false, + emitted: false, + }); + 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)], + RawSeriesOrdering::Contiguous, + ); + batches.extend(pack_all(input).await); + } + let multi = concat_batches(&frame_schema, &batches).unwrap(); + let (source, _, _) = df_source(vec![multi], frame_schema, false); + let mut unpacked = unpack_raw_series( + source, + schema(), + RawLabelSource::Explicit, + RawSeriesOrdering::Contiguous, + ) + .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, false); + let mut unpacked = unpack_raw_series( + source, + schema(), + RawLabelSource::Explicit, + RawSeriesOrdering::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(), false); + let mut stable = unpack_raw_series( + source, + schema(), + RawLabelSource::Explicit, + RawSeriesOrdering::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(), false); + let mut unpacked = unpack_raw_series( + source, + schema(), + RawLabelSource::Explicit, + RawSeriesOrdering::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()); + } +} + +#[tokio::test] +async fn unpack_drains_many_empty_batches_before_the_next_row() { + let (input, _, _) = raw( + vec![chunk(b"after-empty", &[9], true)], + RawSeriesOrdering::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, false); + let mut unpacked = unpack_raw_series( + source, + schema(), + RawLabelSource::Explicit, + RawSeriesOrdering::Contiguous, + ) + .unwrap(); + assert_chunk( + &unpacked.next().await.unwrap().unwrap(), + b"after-empty", + &[9], + &[Some(9.0)], + true, + ); + assert_eq!(polls.load(Ordering::SeqCst), 10_001); + assert!(unpacked.next().await.is_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)], + RawSeriesOrdering::ReadKeyThenTime, + ); + 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, false); + let mut unpacked = unpack_raw_series( + source, + schema(), + RawLabelSource::Explicit, + RawSeriesOrdering::ReadKeyThenTime, + ) + .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, false); + assert!( + unpack_raw_series( + source, + Arc::new(expected), + RawLabelSource::Explicit, + RawSeriesOrdering::Contiguous, + ) + .is_err() + ); + assert_eq!(polls.load(Ordering::SeqCst), 0); + assert_eq!(drops.load(Ordering::SeqCst), 1); + } +} 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..1c487fc73f2 --- /dev/null +++ b/src/common/query/src/series_frame/tests/lifecycle.rs @@ -0,0 +1,322 @@ +// 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::histogram::{BucketSideView, CounterResetHint, HistogramBatchBuilder}; +use datatypes::series::raw::{RawSeriesChunk, RawSeriesOrdering, RawTimeAxis}; +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) -> RawSeriesOrdering { + RawSeriesOrdering::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 = RawTimeAxis::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 = match &chunk.time { + RawTimeAxis::Explicit(values) => values.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, + RawSeriesOrdering::Contiguous, + ) + .unwrap(); + let mut output = Vec::new(); + while let Some(item) = decoded.next().await { + output.push(item.unwrap()); + } + assert_eq!( + output.len(), + 2, + "the float field's segment boundary yields two frames" + ); + + // Validate buffer sharing while the source owners are still alive. + let DenseSamplesRef::Float(decoded_float) = output[0].fields.get(0).unwrap().segments()[0] + .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()) + ); + 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()) + ); + let RawTimeAxis::Explicit(first_times) = &output[0].time; + assert_eq!(first_times.as_ptr(), input_times_ptr); + + 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.read_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.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]); +} From b774c320752eb16832550e6c6c71ee6b764edf6f Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Thu, 15 Oct 2026 19:18:07 +0800 Subject: [PATCH 04/11] fix(query): bound raw frame validation history and empty polling Keep only adjacent sequence state instead of retaining every completed key. The producer remains responsible for contiguous key runs; ordered streams still reject key regressions. Yield after 32 validated empty batches per unpack poll without losing pending rows or terminal errors. Preserve whole-batch atomicity and document the producer contract. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/common/query/src/series_frame.rs | 39 +-- src/common/query/src/series_frame/tests.rs | 263 ++++++++++++++++++++- src/datatypes/src/series/raw.rs | 5 +- src/store-api/src/region_engine.rs | 6 +- 4 files changed, 281 insertions(+), 32 deletions(-) diff --git a/src/common/query/src/series_frame.rs b/src/common/query/src/series_frame.rs index 47c81cc49fc..c9e02646540 100644 --- a/src/common/query/src/series_frame.rs +++ b/src/common/query/src/series_frame.rs @@ -20,7 +20,6 @@ //! candidates only: it is not a residual filter, authorization check, remote transport, or //! compaction boundary. -use std::collections::HashSet; use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; @@ -42,6 +41,9 @@ use store_api::region_engine::{RawLabelSource, RawSeriesStream, SendableRawSerie mod tests; /// Packs a storage raw-series stream into DataFusion record batches incrementally. +/// +/// Producers must keep each read key's 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 { @@ -62,6 +64,9 @@ pub fn pack_raw_series( /// Unpacks DataFusion batches using the caller's explicit scan binding and ordering. /// +/// Producers must keep each read key's 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, @@ -109,17 +114,15 @@ fn invalid(reason: impl Into) -> DataFusionError { ))) } -/// Cross-chunk validation state. Contiguous ordering retains complete owned keys for recurrence -/// detection, so memory is O(total distinct complete-key bytes + distinct keys), plus current-key -/// and metric state, until stream termination; there is no eviction or cap. This owns key bytes, -/// never Arrow buffers that might pin larger payloads. +/// Cross-chunk validation state. The producer must keep each key contiguous; this validator +/// checks adjacent chunks only. It retains only the previous complete owned key (not an Arrow +/// buffer), timestamp, metric, and finality. struct SequenceValidator { ordering: RawSeriesOrdering, previous_key: Option>, previous_timestamp: Option, previous_metric: Option, previous_is_last: bool, - closed: HashSet>, } impl SequenceValidator { @@ -130,7 +133,6 @@ impl SequenceValidator { previous_timestamp: None, previous_metric: None, previous_is_last: false, - closed: HashSet::new(), } } @@ -175,17 +177,10 @@ impl SequenceValidator { { return Err(invalid("raw series keys are not ordered")); } - if self.ordering == RawSeriesOrdering::Contiguous && self.closed.contains(key) { - return Err(invalid("raw series key recurred after its contiguous run")); - } } } } - let owned_key = key.to_vec(); - if chunk.is_last && self.ordering == RawSeriesOrdering::Contiguous { - self.closed.insert(owned_key.clone()); - } - self.previous_key = Some(owned_key); + 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; @@ -296,7 +291,9 @@ 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))); @@ -345,7 +342,17 @@ impl Stream for UnpackStream { }) })(); match result { - Ok(chunks) => this.pending = chunks.into_iter(), + 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( diff --git a/src/common/query/src/series_frame/tests.rs b/src/common/query/src/series_frame/tests.rs index a2c8a8a5ffd..10664e85d3c 100644 --- a/src/common/query/src/series_frame/tests.rs +++ b/src/common/query/src/series_frame/tests.rs @@ -18,7 +18,7 @@ use std::collections::VecDeque; use std::pin::Pin; use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; -use std::task::{Context, Poll}; +use std::task::{Context, Poll, Wake, Waker}; use common_error::ext::{BoxedError, PlainError}; use common_error::status_code::StatusCode; @@ -227,6 +227,41 @@ async fn pack_all(input: SendableRawSeriesStream) -> Vec { } 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, RawSeriesOrdering::Contiguous); + let batches = pack_all(source).await; + let frame_schema = batches[0].schema(); + let (source, _, _) = df_source(batches, frame_schema, false); + let mut decoded = unpack_raw_series( + source, + schema(), + RawLabelSource::Explicit, + RawSeriesOrdering::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 @@ -405,13 +440,13 @@ async fn invalid_sequence_reports_valid_prefix_then_one_terminal_error() { prefix: vec![(b"a", 1, false, None)], }, Case { - name: "nonlocal_recurrence", + name: "key_decrease_after_intervening_keys", items: vec![ chunk(b"a", &[1], true), chunk(b"b", &[2], true), chunk(b"a", &[3], true), ], - order: RawSeriesOrdering::Contiguous, + order: RawSeriesOrdering::ReadKeyThenTime, expected_prefix_count: 2, prefix: vec![(b"a", 1, true, None), (b"b", 2, true, None)], }, @@ -896,7 +931,7 @@ async fn unpack_batch_sequence_error_is_atomic_for_all_rows() { ] { let (input, _, _) = raw( vec![chunk(key, &[time], true)], - RawSeriesOrdering::Contiguous, + RawSeriesOrdering::ReadKeyThenTime, ); batches.extend(pack_all(input).await); } @@ -906,7 +941,7 @@ async fn unpack_batch_sequence_error_is_atomic_for_all_rows() { source, schema(), RawLabelSource::Explicit, - RawSeriesOrdering::Contiguous, + RawSeriesOrdering::ReadKeyThenTime, ) .unwrap(); assert!(unpacked.next().await.unwrap().is_err()); @@ -1077,8 +1112,28 @@ async fn frame_application_metadata_is_stable_and_schema_change_rejected() { } } +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)) +} + #[tokio::test] -async fn unpack_drains_many_empty_batches_before_the_next_row() { +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)], RawSeriesOrdering::Contiguous, @@ -1096,15 +1151,197 @@ async fn unpack_drains_many_empty_batches_before_the_next_row() { RawSeriesOrdering::Contiguous, ) .unwrap(); - assert_chunk( - &unpacked.next().await.unwrap().unwrap(), - b"after-empty", - &[9], - &[Some(9.0)], + 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, false); + let mut unpacked = unpack_raw_series( + source, + schema(), + RawLabelSource::Explicit, + RawSeriesOrdering::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_source( + vec![RecordBatch::new_empty(frame_schema.clone())], + frame_schema, true, ); - assert_eq!(polls.load(Ordering::SeqCst), 10_001); - assert!(unpacked.next().await.is_none()); + let mut unpacked = unpack_raw_series( + source, + schema(), + RawLabelSource::Explicit, + RawSeriesOrdering::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, false); + let mut unpacked = unpack_raw_series( + source, + schema(), + RawLabelSource::Explicit, + RawSeriesOrdering::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: false, + emitted: false, + }); + let mut unpacked = unpack_raw_series( + source, + schema(), + RawLabelSource::Explicit, + RawSeriesOrdering::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] diff --git a/src/datatypes/src/series/raw.rs b/src/datatypes/src/series/raw.rs index 1ff8cc873e0..56baeca1d30 100644 --- a/src/datatypes/src/series/raw.rs +++ b/src/datatypes/src/series/raw.rs @@ -73,7 +73,10 @@ impl RawTimeAxis { } } -/// Ordering guarantee for chunks in one input stream. +/// Producer ordering guarantee for chunks in one input stream. Producers MUST emit all chunks +/// for each read key contiguously; adapters validate adjacent chunk conditions but do not certify +/// nonadjacent recurrence for [`Self::Contiguous`]. An `A/B/A` sequence violates the producer +/// contract even when a consumer cannot detect that recurrence. #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum RawSeriesOrdering { /// Each read key's chunks are contiguous; distinct read keys need not be ordered. diff --git a/src/store-api/src/region_engine.rs b/src/store-api/src/region_engine.rs index 82183f77ec3..a30f0cbd701 100644 --- a/src/store-api/src/region_engine.rs +++ b/src/store-api/src/region_engine.rs @@ -498,8 +498,10 @@ pub enum RawLabelSource { /// A stream of raw series chunks with stable schema, label binding, and ordering. /// -/// Chunks for a read key are contiguous, with strictly increasing timestamps within and across -/// chunks. `is_last` marks the last chunk of that key in this stream, not across streams. +/// Producers MUST emit chunks for each read key contiguously, with strictly increasing timestamps +/// within and across chunks. Adapters validate adjacent chunk conditions but do not certify +/// nonadjacent recurrence for [`RawSeriesOrdering::Contiguous`]; an A/B/A sequence violates the +/// producer contract. `is_last` marks the last chunk of that key in this stream, not across streams. pub trait RawSeriesStream: Stream> + Send { /// Returns the schema shared by all chunks; raw schemas have no evaluation grid. fn schema(&self) -> &Arc; From d4a3f67f65fedc9b4c8bf0733059c8ec3fa6d779 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Thu, 15 Oct 2026 20:37:16 +0800 Subject: [PATCH 05/11] fix(datatypes): preserve series frame field names Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/datatypes/src/series/frame.rs | 8 ++-- src/datatypes/src/series/frame/tests.rs | 54 ++++++++++++++++++++++++- 2 files changed, 56 insertions(+), 6 deletions(-) diff --git a/src/datatypes/src/series/frame.rs b/src/datatypes/src/series/frame.rs index 2fd242c9b76..d9f9c1ee347 100644 --- a/src/datatypes/src/series/frame.rs +++ b/src/datatypes/src/series/frame.rs @@ -12,8 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -//! Finite Arrow batches for raw series chunks. Field names are ordinal (`f0`, `f1`, ...); -//! callers bind those positions to their logical names, which this frame cannot authenticate. +//! Finite Arrow batches for raw series chunks. Field names preserve the logical series schema in order. #[cfg(test)] mod tests; @@ -50,7 +49,7 @@ pub enum RawLabelMode { Explicit, } -/// Builds the typed schema for a raw frame using ordinal domain fields. +/// 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"); @@ -88,8 +87,7 @@ pub fn raw_series_frame_schema(schema: &SeriesSchema, mode: RawLabelMode) -> Res let fields: Fields = schema .fields .iter() - .enumerate() - .map(|(i, _)| Field::new(format!("f{i}"), payload.clone(), false)) + .map(|field| Field::new(field.name.as_ref(), payload.clone(), false)) .collect(); let mut columns = vec![ Field::new("read_key", DataType::Binary, false), diff --git a/src/datatypes/src/series/frame/tests.rs b/src/datatypes/src/series/frame/tests.rs index 6247c4924ee..7edc38168b5 100644 --- a/src/datatypes/src/series/frame/tests.rs +++ b/src/datatypes/src/series/frame/tests.rs @@ -339,6 +339,58 @@ fn sliced_histogram_frames_densify_without_duplicating_shared_table() { ); } +#[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, + RawTimeAxis::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].read_key.as_slice(), &[0xaa]); + assert!(decoded[0].is_last); + let RawTimeAxis::Explicit(times) = &decoded[0].time; + 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"]); @@ -400,7 +452,7 @@ fn schema_and_key_explicit_round_trips_preserve_identity_and_bits() { ), Field::new( "fields", - DataType::Struct(vec![Field::new("f0", payload, false)].into()), + DataType::Struct(vec![Field::new("value", payload, false)].into()), false, ), ], From 1284c043e5a07b8d42536ee1d89396c552e34716 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Fri, 16 Oct 2026 15:00:04 +0800 Subject: [PATCH 06/11] feat(storage): separate prepared series selection from raw readers Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/datatypes/src/series/raw.rs | 101 ++--- src/datatypes/src/series/raw/tests.rs | 51 +-- src/store-api/src/region_engine.rs | 75 +++- .../region_engine/prepared_series_tests.rs | 405 ++++++++++++++++++ 4 files changed, 546 insertions(+), 86 deletions(-) create mode 100644 src/store-api/src/region_engine/prepared_series_tests.rs diff --git a/src/datatypes/src/series/raw.rs b/src/datatypes/src/series/raw.rs index 6c5b04ceac7..9c18ae2941a 100644 --- a/src/datatypes/src/series/raw.rs +++ b/src/datatypes/src/series/raw.rs @@ -22,59 +22,51 @@ use crate::error::{InvalidSeriesSnafu, Result}; use crate::series::samples::{SeriesFields, SeriesSchema}; use crate::series::{MetricName, invalid_series}; -/// Explicit source timestamps in milliseconds for a raw series chunk. -#[non_exhaustive] +/// Strictly increasing source timestamps in milliseconds for a raw series chunk. +/// +/// Its private buffer ensures timestamps are validated once at construction. Direct +/// construction is intentionally unavailable: +/// +/// ```compile_fail +/// use arrow::buffer::ScalarBuffer; +/// use datatypes::series::raw::RawTimestamps; +/// +/// let timestamps = RawTimestamps { +/// values: ScalarBuffer::from(vec![1_i64, 2]), +/// }; +/// ``` #[derive(Clone, Debug)] -pub enum RawTimeAxis { - /// Source timestamps in strictly increasing order. - Explicit(ScalarBuffer), +pub struct RawTimestamps { + values: ScalarBuffer, } -impl RawTimeAxis { - /// Validates and wraps strictly increasing timestamps; an empty axis is valid on its own. - pub fn try_new(timestamps: ScalarBuffer) -> Result { - let axis = Self::Explicit(timestamps); - axis.try_len()?; - Ok(axis) +impl RawTimestamps { + /// Validates and wraps strictly increasing timestamps; an empty wrapper is valid. + pub fn try_new(values: ScalarBuffer) -> Result { + if values.windows(2).any(|pair| pair[0] >= pair[1]) { + return invalid_series("raw timestamps must be strictly increasing"); + } + Ok(Self { values }) } - /// Returns the explicit timestamp buffer, if this axis uses explicit timestamps. - /// - /// Future axis variants may not have an explicit buffer. - pub fn as_explicit(&self) -> Option<&ScalarBuffer> { - match self { - Self::Explicit(values) => Some(values), - } - } - - fn try_len(&self) -> Result { - match self { - Self::Explicit(values) => { - if values.windows(2).any(|pair| pair[0] >= pair[1]) { - return invalid_series("raw timestamps must be strictly increasing"); - } - Ok(values.len()) - } - } + /// Returns the underlying timestamp buffer. + pub fn as_buffer(&self) -> &ScalarBuffer { + &self.values } /// Returns the number of timestamps. pub fn len(&self) -> usize { - match self { - Self::Explicit(values) => values.len(), - } + self.values.len() } - /// Returns whether the axis has no timestamps. + /// Returns whether there are no timestamps. pub fn is_empty(&self) -> bool { - self.len() == 0 + self.values.is_empty() } /// Returns the timestamp at `index`, or `None` when out of bounds. pub fn timestamp(&self, index: usize) -> Option { - match self { - Self::Explicit(values) => values.get(index).copied(), - } + self.values.get(index).copied() } /// Returns a zero-copy slice, reporting overflowing or out-of-range bounds as an error. @@ -82,23 +74,23 @@ impl RawTimeAxis { let end = offset .checked_add(len) .with_context(|| InvalidSeriesSnafu { - reason: "raw time axis slice bounds overflow".to_string(), + reason: "raw timestamps slice bounds overflow".to_string(), })?; if end > self.len() { - return invalid_series("raw time axis slice is out of bounds"); - } - match self { - Self::Explicit(values) => Ok(Self::Explicit(values.slice(offset, len))), + return invalid_series("raw timestamps slice is out of bounds"); } + Ok(Self { + values: self.values.slice(offset, len), + }) } } /// Ordering guarantee for chunks in one input stream. #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum RawSeriesOrdering { - /// Each read key's chunks are contiguous; distinct read keys need not be ordered. + /// Each complete primary key's chunks are contiguous; distinct keys need not be ordered. Contiguous, - /// Chunks are ordered by read key and then by timestamp. + /// Chunks are ordered by complete primary key and then by timestamp. ReadKeyThenTime, } @@ -109,27 +101,29 @@ pub enum RawSeriesOrdering { #[derive(Clone, Debug)] pub struct RawSeriesChunk { /// Complete encoded identity of the source series, not a TSID or hash alone. - pub read_key: Buffer, - /// `None` leaves labels unmaterialized in the complete read key; `Some` contains - /// complete explicit labels. + pub primary_key: Buffer, + /// `None` denotes an encoded key under bound-label mode; `Some` contains complete + /// explicit labels and retains their identity. pub metric: Option, /// Source timestamps aligned with every field. - pub time: RawTimeAxis, + pub time: RawTimestamps, /// Field values aligned with the source timestamps. pub fields: SeriesFields, - /// Whether this is the final chunk for this input stream's series, not all regions. + /// Whether this is the final chunk for this series in the input stream exposing it. + /// Merge code must recompute finalization globally across all relevant inputs before + /// downstream finalization. pub is_last: bool, } impl RawSeriesChunk { - /// Creates a nonempty chunk after checking its explicit time axis and schema alignment. + /// Creates a nonempty chunk after checking its time axis and schema alignment. /// /// Validation is local to this chunk and does not verify ordering between chunks. pub fn try_new( schema: &SeriesSchema, - read_key: Buffer, + primary_key: Buffer, metric: Option, - time: RawTimeAxis, + time: RawTimestamps, fields: SeriesFields, is_last: bool, ) -> Result { @@ -145,9 +139,8 @@ impl RawSeriesChunk { if time.is_empty() { return invalid_series("raw series chunks must not be empty"); } - time.try_len()?; Ok(Self { - read_key, + primary_key, metric, time, fields, diff --git a/src/datatypes/src/series/raw/tests.rs b/src/datatypes/src/series/raw/tests.rs index ee982148d73..ea2d4f80193 100644 --- a/src/datatypes/src/series/raw/tests.rs +++ b/src/datatypes/src/series/raw/tests.rs @@ -46,8 +46,8 @@ fn metric() -> MetricName { .unwrap() } -fn times(values: Vec) -> RawTimeAxis { - RawTimeAxis::try_new(values.into()).unwrap() +fn times(values: Vec) -> RawTimestamps { + RawTimestamps::try_new(values.into()).unwrap() } fn float_fields(first: Vec>, second: Option>>) -> SeriesFields { @@ -64,8 +64,13 @@ fn constructors_reject_invalid_chunk_boundaries_and_unordered_timestamps() { let key = Buffer::from(vec![1_u8]); let fields = float_fields(vec![Some(1.0), Some(2.0)], None); - assert!(RawTimeAxis::try_new(vec![1_i64, 1].into()).is_err()); - assert!(RawTimeAxis::try_new(vec![2_i64, 1].into()).is_err()); + assert!(RawTimestamps::try_new(vec![1_i64, 1].into()).is_err()); + assert!(RawTimestamps::try_new(vec![2_i64, 1].into()).is_err()); + assert!( + RawTimestamps::try_new(ScalarBuffer::from(Vec::::new())) + .unwrap() + .is_empty() + ); assert!( RawSeriesChunk::try_new( @@ -101,22 +106,6 @@ fn constructors_reject_invalid_chunk_boundaries_and_unordered_timestamps() { .is_err() ); - // Public enum construction bypasses RawTimeAxis::try_new, so chunk construction - // must independently reject duplicate and descending timestamps. - for unordered in [vec![1_i64, 1], vec![2_i64, 1]] { - assert!( - RawSeriesChunk::try_new( - &schema, - key.clone(), - None, - RawTimeAxis::Explicit(unordered.into()), - fields.clone(), - true, - ) - .is_err() - ); - } - let grid = Arc::new(EvalGrid::try_new(GridSpec::try_new(10, 5, 2).unwrap()).unwrap()); let grid_schema = SeriesSchema { fields: schema.fields.clone(), @@ -131,9 +120,9 @@ fn constructors_reject_invalid_chunk_boundaries_and_unordered_timestamps() { #[test] fn sliced_axes_and_chunks_keep_shared_buffers_and_values_alive() { let time_owner: ScalarBuffer = vec![0, 10, 20, 30, 40, 80].into(); - let time_source = RawTimeAxis::try_new(time_owner.clone()).unwrap(); + let time_source = RawTimestamps::try_new(time_owner.clone()).unwrap(); let time = time_source.slice(1, 4).unwrap(); - let time_buffer = time.as_explicit().unwrap(); + let time_buffer = time.as_buffer(); assert_eq!(time_buffer.as_ptr(), time_owner.as_ptr().wrapping_add(1)); let nan_payload = f64::from_bits(0x7ff8_0000_0000_0042); @@ -160,12 +149,12 @@ fn sliced_axes_and_chunks_keep_shared_buffers_and_values_alive() { ]) .unwrap(); let key_source = Buffer::from(vec![0_u8, 0xA1, 0xB2, 0xC3, 0xD4, 0xFF]); - let read_key = key_source.slice_with_length(1, 4); - assert_eq!(read_key.as_ptr(), key_source.as_ptr().wrapping_add(1)); + let primary_key = key_source.slice_with_length(1, 4); + assert_eq!(primary_key.as_ptr(), key_source.as_ptr().wrapping_add(1)); let schema = schema(&["first", "second"]); let chunk = - RawSeriesChunk::try_new(&schema, read_key, Some(metric()), time, fields, true).unwrap(); + RawSeriesChunk::try_new(&schema, primary_key, Some(metric()), time, fields, true).unwrap(); let Some(crate::series::samples::DenseSamplesRef::Float(first)) = chunk.fields.get(0).unwrap().dense() else { @@ -181,7 +170,7 @@ fn sliced_axes_and_chunks_keep_shared_buffers_and_values_alive() { drop(time_owner); drop(float_owner); drop(second_owner); - assert_eq!(chunk.read_key.as_slice(), &[0xA1, 0xB2, 0xC3, 0xD4]); + assert_eq!(chunk.primary_key.as_slice(), &[0xA1, 0xB2, 0xC3, 0xD4]); assert_eq!( (0..4) .map(|i| chunk.time.timestamp(i).unwrap()) @@ -224,9 +213,9 @@ fn sliced_axes_and_chunks_keep_shared_buffers_and_values_alive() { } #[test] -fn raw_time_axis_checks_bounds_and_retains_exact_end_slices() { +fn raw_timestamps_check_bounds_and_retain_exact_end_slices() { let source: ScalarBuffer = vec![4, 8, 12].into(); - let axis = RawTimeAxis::try_new(source.clone()).unwrap(); + let axis = RawTimestamps::try_new(source.clone()).unwrap(); assert!(axis.slice(usize::MAX, 1).is_err()); assert!(axis.slice(2, usize::MAX).is_err()); assert!(axis.slice(3, 1).is_err()); @@ -272,11 +261,11 @@ fn raw_chunks_preserve_native_unmaterialized_and_explicit_metrics() { false, ) .unwrap(); - assert_eq!(native.read_key.as_slice(), &[0, 0x12, 0x34, 0x56]); + assert_eq!(native.primary_key.as_slice(), &[0, 0x12, 0x34, 0x56]); assert!(native.metric.is_none()); - assert!(tagless.read_key.is_empty()); + assert!(tagless.primary_key.is_empty()); assert!(tagless.metric.is_none()); - assert_eq!(explicit.read_key.as_slice(), &[0, 0xAB, 0xCD, 0xEF]); + assert_eq!(explicit.primary_key.as_slice(), &[0, 0xAB, 0xCD, 0xEF]); assert_eq!(explicit.metric.as_ref().unwrap().name(), Some("cpu")); assert_eq!( explicit.metric.as_ref().unwrap().get("host"), diff --git a/src/store-api/src/region_engine.rs b/src/store-api/src/region_engine.rs index 48017c94f0f..d9c4fa477f6 100644 --- a/src/store-api/src/region_engine.rs +++ b/src/store-api/src/region_engine.rs @@ -499,7 +499,11 @@ pub enum RawLabelSource { /// A stream of raw series chunks with stable schema, label binding, and ordering. /// /// Chunks for a read key are contiguous, with strictly increasing timestamps within and across -/// chunks. `is_last` marks the last chunk of that key in this stream, not across streams. +/// chunks. `is_last` marks the last chunk of that key in this stream, not across streams; merge +/// code must recompute finalization globally across all relevant inputs before downstream +/// per-input-series finalization (not aggregate or vector-match state). `RawLabelSource` keeps +/// the stream's label interpretation self-contained through transport rather than relying on a +/// local source handle. pub trait RawSeriesStream: Stream> + Send { /// Returns the schema shared by all chunks; raw schemas have no evaluation grid. fn schema(&self) -> &Arc; @@ -514,6 +518,48 @@ pub trait RawSeriesStream: Stream> + S /// A pinned raw-series stream that can be sent between threads. pub type SendableRawSeriesStream = Pin>; +/// Stable identity metadata for one ordinal in a prepared series source. +#[derive(Clone, Debug)] +pub struct SeriesInfo { + /// Complete encoded primary key, not a TSID or hash alone. + pub primary_key: datatypes::arrow::buffer::Buffer, + /// Complete explicit metric and labels, or `None` for primary-key label binding. + pub metric: Option, +} + +/// An immutable, finite selected set of series prepared for independent range reads. +/// +/// The count may include series with no samples after filtering. Ordinals are stable local +/// indexes, not protocol identifiers; each has exactly one complete identity returned by +/// [`PreparedSeriesSource::series_info`], and that identity matches chunks for that ordinal. +/// Every reader exposes the same fixed schema, label binding, and ordering before polling, +/// including for an empty range, and describes labels independently of this source's lifetime. `PrimaryKey` output has `metric: None`; `Explicit` output has a +/// complete `Some(metric)` matching each chunk. [`RawSeriesOrdering::Contiguous`] and +/// [`RawSeriesOrdering::ReadKeyThenTime`] describe stream order, not ordinal enumeration order. +/// +/// Preparation, identity lookup, and reader creation do not load sample payloads or perform +/// hidden blocking I/O. The selected set is bound to the partition's immutable MVCC view, +/// predicates, authorization, projection, labels, and existing scan resources. Readers own +/// resources needed to continue after this source or scanner is dropped. Dropping one reader +/// does not cancel sibling readers; normal ownership cleanup applies without detached work. +/// +/// A range must satisfy `start <= end <= num_series()`. Invalid ranges and invalid ordinals +/// return errors; an empty in-bounds range returns an ordinary empty stream. Repeated and +/// overlapping reads are independent and do not consume shared iteration state. +pub trait PreparedSeriesSource: Send + Sync { + /// Returns the exact, stable number of selected local series ordinals. + fn num_series(&self) -> usize; + + /// Returns identity metadata for an ordinal without polling sample data. + fn series_info(&self, ordinal: usize) -> Result; + + /// Creates an independent reader for the half-open ordinal range. + fn read_series( + &self, + range: std::ops::Range, + ) -> Result; +} + /// A scanner that provides a way to scan the region concurrently. /// /// The scanner splits the region into partitions so that each partition can be scanned concurrently. @@ -566,6 +612,22 @@ pub trait RegionScanner: Debug + DisplayAs + Send { todo!() } + /// Prepares a finite, stable set of series for independent ordinal-range reads. + /// + /// `None` means unsupported before consumption, so callers may use the legacy record-batch + /// scan. `Some`, including a source with zero series, selects this path; errors propagate + /// without replay. Choose exactly one entrypoint per partition per execution: this prepared + /// source, [`Self::scan_series_partition`], or [`Self::scan_partition`]. In particular, do + /// not call the legacy raw-scan default merely because this returns `None`. + fn prepare_series_partition( + &self, + _ctx: &QueryScanContext, + _metrics_set: &ExecutionPlanMetricsSet, + _partition: usize, + ) -> Result>, BoxedError> { + Ok(None) + } + /// Check if there is any predicate exclude region partition exprs that may be executed in this scanner. fn has_predicate_without_region(&self) -> bool; @@ -804,6 +866,10 @@ impl RegionStatistic { } } +#[cfg(test)] +#[path = "region_engine/prepared_series_tests.rs"] +mod prepared_series_tests; + #[cfg(test)] mod tests { use common_recordbatch::{RecordBatch, RecordBatches}; @@ -859,6 +925,13 @@ mod tests { let ctx = QueryScanContext::default(); let metrics = ExecutionPlanMetricsSet::new(); + assert!( + scanner + .prepare_series_partition(&ctx, &metrics, 0) + .unwrap() + .is_none() + ); + // SinglePartitionScanner uses the default raw-scan implementation, which panics. let raw_scan = catch_unwind(AssertUnwindSafe(|| { scanner.scan_series_partition(&ctx, &metrics, 0) diff --git a/src/store-api/src/region_engine/prepared_series_tests.rs b/src/store-api/src/region_engine/prepared_series_tests.rs new file mode 100644 index 00000000000..05882ddcc70 --- /dev/null +++ b/src/store-api/src/region_engine/prepared_series_tests.rs @@ -0,0 +1,405 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::VecDeque; +use std::pin::Pin; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; +use std::task::{Context, Poll}; + +use api::v1::SemanticType; +use common_error::ext::BoxedError; +use datatypes::arrow::array::Float64Array; +use datatypes::arrow::buffer::Buffer; +use datatypes::error::InvalidSeriesSnafu; +use datatypes::prelude::ConcreteDataType; +use datatypes::schema::ColumnSchema; +use datatypes::series::raw::{RawSeriesChunk, RawSeriesOrdering, RawTimestamps}; +use datatypes::series::samples::{FieldSpec, Samples, SeriesFields, SeriesSchema}; +use datatypes::series::{Label, MetricName}; +use futures::{Stream, StreamExt}; + +use crate::region_engine::{ + PreparedSeriesSource, RawLabelSource, RawSeriesStream, SendableRawSeriesStream, SeriesInfo, +}; + +#[derive(Clone)] +struct TestSeries { + info: SeriesInfo, + chunks: Vec, +} + +struct MemorySource { + series: Arc<[TestSeries]>, + schema: Arc, + label_source: RawLabelSource, + polls: Arc, + drops: Arc, +} + +impl PreparedSeriesSource for MemorySource { + fn num_series(&self) -> usize { + self.series.len() + } + + fn series_info(&self, ordinal: usize) -> Result { + self.series + .get(ordinal) + .map(|series| series.info.clone()) + .ok_or_else(|| { + BoxedError::new( + InvalidSeriesSnafu { + reason: "invalid series ordinal".to_string(), + } + .build(), + ) + }) + } + + fn read_series( + &self, + range: std::ops::Range, + ) -> Result { + if range.start > range.end || range.end > self.series.len() { + return Err(BoxedError::new( + InvalidSeriesSnafu { + reason: "invalid series range".to_string(), + } + .build(), + )); + } + let chunks = self.series[range] + .iter() + .flat_map(|series| series.chunks.iter().cloned()) + .collect(); + Ok(Box::pin(MemoryReader { + chunks, + schema: self.schema.clone(), + label_source: self.label_source.clone(), + polls: self.polls.clone(), + drops: self.drops.clone(), + })) + } +} + +struct MemoryReader { + chunks: VecDeque, + schema: Arc, + label_source: RawLabelSource, + polls: Arc, + drops: Arc, +} + +impl Stream for MemoryReader { + type Item = Result; + + fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + let chunk = self.chunks.pop_front(); + if chunk.is_some() { + self.polls.fetch_add(1, Ordering::SeqCst); + } + Poll::Ready(chunk.map(Ok)) + } +} + +impl Drop for MemoryReader { + fn drop(&mut self) { + self.drops.fetch_add(1, Ordering::SeqCst); + } +} + +impl RawSeriesStream for MemoryReader { + fn schema(&self) -> &Arc { + &self.schema + } + + fn label_source(&self) -> &RawLabelSource { + &self.label_source + } + + fn ordering(&self) -> RawSeriesOrdering { + RawSeriesOrdering::Contiguous + } +} + +fn metric(host: &str) -> MetricName { + MetricName::try_new( + Some(Arc::from("cpu")), + vec![Label { + name: Arc::from("host"), + value: Arc::from(host), + }], + ) + .unwrap() +} + +fn schema() -> Arc { + Arc::new(SeriesSchema { + fields: vec![FieldSpec { + name: Arc::from("usage"), + }] + .into(), + eval_grid: None, + }) +} + +fn chunk( + key: &[u8], + metric: Option, + timestamp: i64, + value: f64, + last: bool, +) -> RawSeriesChunk { + RawSeriesChunk::try_new( + &schema(), + Buffer::from(key), + metric, + RawTimestamps::try_new(vec![timestamp].into()).unwrap(), + SeriesFields::one(Samples::from_float(Float64Array::from(vec![value]))), + last, + ) + .unwrap() +} + +fn primary_key_metadata() -> crate::metadata::RegionMetadata { + let mut builder = + crate::metadata::RegionMetadataBuilder::new(crate::storage::RegionId::new(1, 1)); + builder.push_column_metadata(crate::metadata::ColumnMetadata { + column_schema: ColumnSchema::new("host", ConcreteDataType::string_datatype(), false), + semantic_type: SemanticType::Tag, + column_id: 1, + }); + builder.push_column_metadata(crate::metadata::ColumnMetadata { + column_schema: ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + semantic_type: SemanticType::Timestamp, + column_id: 2, + }); + builder.primary_key(vec![1]); + builder.build().unwrap() +} + +fn empty_source() -> Arc { + Arc::new(MemorySource { + series: Arc::from([]), + schema: schema(), + label_source: RawLabelSource::Explicit, + polls: Arc::new(AtomicUsize::new(0)), + drops: Arc::new(AtomicUsize::new(0)), + }) +} + +fn source(explicit: bool) -> Arc { + source_from_keys( + explicit, + vec![ + Buffer::from(vec![0x10, 0x01]), + Buffer::from(vec![0x10, 0x02]), + ] + .into(), + ) +} + +fn source_from_keys(explicit: bool, keys: Arc<[Buffer]>) -> Arc { + assert_eq!(keys.len(), 2); + let first_metric = explicit.then(|| metric("a")); + let second_metric = explicit.then(|| metric("b")); + let series = vec![ + TestSeries { + info: SeriesInfo { + primary_key: keys[0].clone(), + metric: first_metric.clone(), + }, + chunks: vec![chunk(keys[0].as_slice(), first_metric, 10, 1.0, true)], + }, + TestSeries { + info: SeriesInfo { + primary_key: keys[1].clone(), + metric: second_metric.clone(), + }, + chunks: vec![ + chunk(keys[1].as_slice(), second_metric.clone(), 10, 2.0, false), + chunk(keys[1].as_slice(), second_metric, 20, 3.0, true), + ], + }, + ]; + Arc::new(MemorySource { + series: series.into(), + schema: schema(), + label_source: if explicit { + RawLabelSource::Explicit + } else { + RawLabelSource::PrimaryKey { + metric_name: Arc::from("cpu"), + physical_metadata: Arc::new(primary_key_metadata()), + logical_tags: Arc::from([]), + } + }, + polls: Arc::new(AtomicUsize::new(0)), + drops: Arc::new(AtomicUsize::new(0)), + }) +} + +#[tokio::test] +async fn prepared_identity_matches_chunks_in_both_label_modes_including_empty_reader() { + for explicit in [false, true] { + let source = source(explicit); + assert_eq!(source.num_series(), 2); + let first = source.series_info(0).unwrap(); + let second = source.series_info(1).unwrap(); + assert_eq!(first.primary_key, Buffer::from(vec![0x10, 0x01])); + assert_eq!(second.primary_key, Buffer::from(vec![0x10, 0x02])); + assert_eq!( + first.metric.as_ref().and_then(MetricName::name), + if explicit { Some("cpu") } else { None } + ); + assert_eq!( + first.metric.as_ref().and_then(|m| m.get("host")), + explicit.then_some("a") + ); + assert_eq!( + second.metric.as_ref().and_then(|m| m.get("host")), + explicit.then_some("b") + ); + + let empty = source.read_series(1..1).unwrap(); + assert_eq!(empty.schema().fields.len(), 1); + assert_eq!(empty.ordering(), RawSeriesOrdering::Contiguous); + assert_eq!( + matches!(empty.label_source(), RawLabelSource::Explicit), + explicit + ); + assert_eq!(source.polls.load(Ordering::SeqCst), 0); + assert!( + futures::TryStreamExt::try_collect::>(empty) + .await + .unwrap() + .is_empty() + ); + assert_eq!(source.polls.load(Ordering::SeqCst), 0); + let end = source.read_series(2..2).unwrap(); + assert_eq!(end.schema().fields.len(), 1); + assert_eq!( + matches!(end.label_source(), RawLabelSource::Explicit), + explicit + ); + + let chunks = + futures::TryStreamExt::try_collect::>(source.read_series(0..2).unwrap()) + .await + .unwrap(); + assert_eq!(chunks.len(), 3); + assert_eq!(source.polls.load(Ordering::SeqCst), 3); + assert_eq!(chunks[1].time.timestamp(0), Some(10)); + assert_eq!(chunks[2].time.timestamp(0), Some(20)); + assert!( + matches!(chunks[1].fields.get(0).unwrap().sample(0), datatypes::series::samples::SampleRef::Float(value) if value == 2.0) + ); + assert!( + matches!(chunks[2].fields.get(0).unwrap().sample(0), datatypes::series::samples::SampleRef::Float(value) if value == 3.0) + ); + for chunk in &chunks { + let ordinal = if chunk.primary_key == first.primary_key { + 0 + } else { + 1 + }; + let info = source.series_info(ordinal).unwrap(); + assert_eq!(chunk.primary_key, info.primary_key); + assert_eq!(chunk.metric, info.metric); + } + assert_eq!(chunks[0].metric.is_some(), explicit); + assert!(!chunks[1].is_last); + assert!(chunks[2].is_last); + } +} + +#[tokio::test] +async fn empty_prepared_source_supports_its_empty_range() { + let source = empty_source(); + assert_eq!(source.num_series(), 0); + let stream = source.read_series(0..0).unwrap(); + assert_eq!(stream.schema().fields.len(), 1); + assert!( + futures::TryStreamExt::try_collect::>(stream) + .await + .unwrap() + .is_empty() + ); + assert_eq!(source.polls.load(Ordering::SeqCst), 0); +} + +#[test] +fn invalid_ordinals_and_ranges_fail_without_polling_samples() { + let source = source(false); + assert!(source.series_info(2).is_err()); + assert!(source.read_series(2..3).is_err()); + assert!(source.read_series(2..1).is_err()); + assert_eq!(source.polls.load(Ordering::SeqCst), 0); +} + +#[tokio::test] +async fn overlapping_readers_are_independent_across_source_and_reader_drops() { + let source = source(false); + let drops = source.drops.clone(); + let mut left = source.read_series(0..2).unwrap(); + let mut right = source.read_series(1..2).unwrap(); + drop(source); + assert!(left.next().await.unwrap().unwrap().is_last); + assert_eq!( + right.next().await.unwrap().unwrap().primary_key, + Buffer::from(vec![0x10, 0x02]) + ); + drop(left); + assert_eq!( + futures::TryStreamExt::try_collect::>(right) + .await + .unwrap() + .len(), + 1 + ); + assert_eq!(drops.load(Ordering::SeqCst), 2); +} + +#[tokio::test] +async fn prepared_selection_is_independent_of_later_mutable_selection_changes() { + let selected = Arc::new(Mutex::new(vec![ + Buffer::from(vec![1_u8]), + Buffer::from(vec![2_u8]), + ])); + let prepared = source_from_keys(false, selected.lock().unwrap().clone().into()); + selected.lock().unwrap().push(Buffer::from(vec![3_u8])); + + assert_eq!(prepared.num_series(), 2); + assert_eq!( + prepared.series_info(0).unwrap().primary_key, + Buffer::from(vec![1_u8]) + ); + assert_eq!( + prepared.series_info(1).unwrap().primary_key, + Buffer::from(vec![2_u8]) + ); + let chunks = futures::TryStreamExt::try_collect::>(prepared.read_series(0..2).unwrap()) + .await + .unwrap(); + assert_eq!(chunks.len(), 3); + assert!( + chunks + .iter() + .all(|chunk| chunk.primary_key != Buffer::from(vec![3_u8])) + ); +} From 1fc0c1c7d665f673285fb6bab8f4302b73cff496 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Fri, 16 Oct 2026 15:00:04 +0800 Subject: [PATCH 07/11] feat(storage): separate prepared series selection from raw readers Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/datatypes/src/series/raw.rs | 101 ++--- src/datatypes/src/series/raw/tests.rs | 51 +-- src/store-api/src/region_engine.rs | 79 +++- .../region_engine/prepared_series_tests.rs | 405 ++++++++++++++++++ 4 files changed, 547 insertions(+), 89 deletions(-) create mode 100644 src/store-api/src/region_engine/prepared_series_tests.rs diff --git a/src/datatypes/src/series/raw.rs b/src/datatypes/src/series/raw.rs index 2bad48e7b3b..130f59565a9 100644 --- a/src/datatypes/src/series/raw.rs +++ b/src/datatypes/src/series/raw.rs @@ -22,59 +22,51 @@ use crate::error::{InvalidSeriesSnafu, Result}; use crate::series::samples::{SeriesFields, SeriesSchema}; use crate::series::{MetricName, invalid_series}; -/// Explicit source timestamps in milliseconds for a raw series chunk. -#[non_exhaustive] +/// Strictly increasing source timestamps in milliseconds for a raw series chunk. +/// +/// Its private buffer ensures timestamps are validated once at construction. Direct +/// construction is intentionally unavailable: +/// +/// ```compile_fail +/// use arrow::buffer::ScalarBuffer; +/// use datatypes::series::raw::RawTimestamps; +/// +/// let timestamps = RawTimestamps { +/// values: ScalarBuffer::from(vec![1_i64, 2]), +/// }; +/// ``` #[derive(Clone, Debug)] -pub enum RawTimeAxis { - /// Source timestamps in strictly increasing order. - Explicit(ScalarBuffer), +pub struct RawTimestamps { + values: ScalarBuffer, } -impl RawTimeAxis { - /// Validates and wraps strictly increasing timestamps; an empty axis is valid on its own. - pub fn try_new(timestamps: ScalarBuffer) -> Result { - let axis = Self::Explicit(timestamps); - axis.try_len()?; - Ok(axis) +impl RawTimestamps { + /// Validates and wraps strictly increasing timestamps; an empty wrapper is valid. + pub fn try_new(values: ScalarBuffer) -> Result { + if values.windows(2).any(|pair| pair[0] >= pair[1]) { + return invalid_series("raw timestamps must be strictly increasing"); + } + Ok(Self { values }) } - /// Returns the explicit timestamp buffer, if this axis uses explicit timestamps. - /// - /// Future axis variants may not have an explicit buffer. - pub fn as_explicit(&self) -> Option<&ScalarBuffer> { - match self { - Self::Explicit(values) => Some(values), - } - } - - fn try_len(&self) -> Result { - match self { - Self::Explicit(values) => { - if values.windows(2).any(|pair| pair[0] >= pair[1]) { - return invalid_series("raw timestamps must be strictly increasing"); - } - Ok(values.len()) - } - } + /// Returns the underlying timestamp buffer. + pub fn as_buffer(&self) -> &ScalarBuffer { + &self.values } /// Returns the number of timestamps. pub fn len(&self) -> usize { - match self { - Self::Explicit(values) => values.len(), - } + self.values.len() } - /// Returns whether the axis has no timestamps. + /// Returns whether there are no timestamps. pub fn is_empty(&self) -> bool { - self.len() == 0 + self.values.is_empty() } /// Returns the timestamp at `index`, or `None` when out of bounds. pub fn timestamp(&self, index: usize) -> Option { - match self { - Self::Explicit(values) => values.get(index).copied(), - } + self.values.get(index).copied() } /// Returns a zero-copy slice, reporting overflowing or out-of-range bounds as an error. @@ -82,14 +74,14 @@ impl RawTimeAxis { let end = offset .checked_add(len) .with_context(|| InvalidSeriesSnafu { - reason: "raw time axis slice bounds overflow".to_string(), + reason: "raw timestamps slice bounds overflow".to_string(), })?; if end > self.len() { - return invalid_series("raw time axis slice is out of bounds"); - } - match self { - Self::Explicit(values) => Ok(Self::Explicit(values.slice(offset, len))), + return invalid_series("raw timestamps slice is out of bounds"); } + Ok(Self { + values: self.values.slice(offset, len), + }) } } @@ -99,9 +91,9 @@ impl RawTimeAxis { /// contract even when a consumer cannot detect that recurrence. #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum RawSeriesOrdering { - /// Each read key's chunks are contiguous; distinct read keys need not be ordered. + /// Each complete primary key's chunks are contiguous; distinct keys need not be ordered. Contiguous, - /// Chunks are ordered by read key and then by timestamp. + /// Chunks are ordered by complete primary key and then by timestamp. ReadKeyThenTime, } @@ -112,27 +104,29 @@ pub enum RawSeriesOrdering { #[derive(Clone, Debug)] pub struct RawSeriesChunk { /// Complete encoded identity of the source series, not a TSID or hash alone. - pub read_key: Buffer, - /// `None` leaves labels unmaterialized in the complete read key; `Some` contains - /// complete explicit labels. + pub primary_key: Buffer, + /// `None` denotes an encoded key under bound-label mode; `Some` contains complete + /// explicit labels and retains their identity. pub metric: Option, /// Source timestamps aligned with every field. - pub time: RawTimeAxis, + pub time: RawTimestamps, /// Field values aligned with the source timestamps. pub fields: SeriesFields, - /// Whether this is the final chunk for this input stream's series, not all regions. + /// Whether this is the final chunk for this series in the input stream exposing it. + /// Merge code must recompute finalization globally across all relevant inputs before + /// downstream finalization. pub is_last: bool, } impl RawSeriesChunk { - /// Creates a nonempty chunk after checking its explicit time axis and schema alignment. + /// Creates a nonempty chunk after checking its time axis and schema alignment. /// /// Validation is local to this chunk and does not verify ordering between chunks. pub fn try_new( schema: &SeriesSchema, - read_key: Buffer, + primary_key: Buffer, metric: Option, - time: RawTimeAxis, + time: RawTimestamps, fields: SeriesFields, is_last: bool, ) -> Result { @@ -148,9 +142,8 @@ impl RawSeriesChunk { if time.is_empty() { return invalid_series("raw series chunks must not be empty"); } - time.try_len()?; Ok(Self { - read_key, + primary_key, metric, time, fields, diff --git a/src/datatypes/src/series/raw/tests.rs b/src/datatypes/src/series/raw/tests.rs index ee982148d73..ea2d4f80193 100644 --- a/src/datatypes/src/series/raw/tests.rs +++ b/src/datatypes/src/series/raw/tests.rs @@ -46,8 +46,8 @@ fn metric() -> MetricName { .unwrap() } -fn times(values: Vec) -> RawTimeAxis { - RawTimeAxis::try_new(values.into()).unwrap() +fn times(values: Vec) -> RawTimestamps { + RawTimestamps::try_new(values.into()).unwrap() } fn float_fields(first: Vec>, second: Option>>) -> SeriesFields { @@ -64,8 +64,13 @@ fn constructors_reject_invalid_chunk_boundaries_and_unordered_timestamps() { let key = Buffer::from(vec![1_u8]); let fields = float_fields(vec![Some(1.0), Some(2.0)], None); - assert!(RawTimeAxis::try_new(vec![1_i64, 1].into()).is_err()); - assert!(RawTimeAxis::try_new(vec![2_i64, 1].into()).is_err()); + assert!(RawTimestamps::try_new(vec![1_i64, 1].into()).is_err()); + assert!(RawTimestamps::try_new(vec![2_i64, 1].into()).is_err()); + assert!( + RawTimestamps::try_new(ScalarBuffer::from(Vec::::new())) + .unwrap() + .is_empty() + ); assert!( RawSeriesChunk::try_new( @@ -101,22 +106,6 @@ fn constructors_reject_invalid_chunk_boundaries_and_unordered_timestamps() { .is_err() ); - // Public enum construction bypasses RawTimeAxis::try_new, so chunk construction - // must independently reject duplicate and descending timestamps. - for unordered in [vec![1_i64, 1], vec![2_i64, 1]] { - assert!( - RawSeriesChunk::try_new( - &schema, - key.clone(), - None, - RawTimeAxis::Explicit(unordered.into()), - fields.clone(), - true, - ) - .is_err() - ); - } - let grid = Arc::new(EvalGrid::try_new(GridSpec::try_new(10, 5, 2).unwrap()).unwrap()); let grid_schema = SeriesSchema { fields: schema.fields.clone(), @@ -131,9 +120,9 @@ fn constructors_reject_invalid_chunk_boundaries_and_unordered_timestamps() { #[test] fn sliced_axes_and_chunks_keep_shared_buffers_and_values_alive() { let time_owner: ScalarBuffer = vec![0, 10, 20, 30, 40, 80].into(); - let time_source = RawTimeAxis::try_new(time_owner.clone()).unwrap(); + let time_source = RawTimestamps::try_new(time_owner.clone()).unwrap(); let time = time_source.slice(1, 4).unwrap(); - let time_buffer = time.as_explicit().unwrap(); + let time_buffer = time.as_buffer(); assert_eq!(time_buffer.as_ptr(), time_owner.as_ptr().wrapping_add(1)); let nan_payload = f64::from_bits(0x7ff8_0000_0000_0042); @@ -160,12 +149,12 @@ fn sliced_axes_and_chunks_keep_shared_buffers_and_values_alive() { ]) .unwrap(); let key_source = Buffer::from(vec![0_u8, 0xA1, 0xB2, 0xC3, 0xD4, 0xFF]); - let read_key = key_source.slice_with_length(1, 4); - assert_eq!(read_key.as_ptr(), key_source.as_ptr().wrapping_add(1)); + let primary_key = key_source.slice_with_length(1, 4); + assert_eq!(primary_key.as_ptr(), key_source.as_ptr().wrapping_add(1)); let schema = schema(&["first", "second"]); let chunk = - RawSeriesChunk::try_new(&schema, read_key, Some(metric()), time, fields, true).unwrap(); + RawSeriesChunk::try_new(&schema, primary_key, Some(metric()), time, fields, true).unwrap(); let Some(crate::series::samples::DenseSamplesRef::Float(first)) = chunk.fields.get(0).unwrap().dense() else { @@ -181,7 +170,7 @@ fn sliced_axes_and_chunks_keep_shared_buffers_and_values_alive() { drop(time_owner); drop(float_owner); drop(second_owner); - assert_eq!(chunk.read_key.as_slice(), &[0xA1, 0xB2, 0xC3, 0xD4]); + assert_eq!(chunk.primary_key.as_slice(), &[0xA1, 0xB2, 0xC3, 0xD4]); assert_eq!( (0..4) .map(|i| chunk.time.timestamp(i).unwrap()) @@ -224,9 +213,9 @@ fn sliced_axes_and_chunks_keep_shared_buffers_and_values_alive() { } #[test] -fn raw_time_axis_checks_bounds_and_retains_exact_end_slices() { +fn raw_timestamps_check_bounds_and_retain_exact_end_slices() { let source: ScalarBuffer = vec![4, 8, 12].into(); - let axis = RawTimeAxis::try_new(source.clone()).unwrap(); + let axis = RawTimestamps::try_new(source.clone()).unwrap(); assert!(axis.slice(usize::MAX, 1).is_err()); assert!(axis.slice(2, usize::MAX).is_err()); assert!(axis.slice(3, 1).is_err()); @@ -272,11 +261,11 @@ fn raw_chunks_preserve_native_unmaterialized_and_explicit_metrics() { false, ) .unwrap(); - assert_eq!(native.read_key.as_slice(), &[0, 0x12, 0x34, 0x56]); + assert_eq!(native.primary_key.as_slice(), &[0, 0x12, 0x34, 0x56]); assert!(native.metric.is_none()); - assert!(tagless.read_key.is_empty()); + assert!(tagless.primary_key.is_empty()); assert!(tagless.metric.is_none()); - assert_eq!(explicit.read_key.as_slice(), &[0, 0xAB, 0xCD, 0xEF]); + assert_eq!(explicit.primary_key.as_slice(), &[0, 0xAB, 0xCD, 0xEF]); assert_eq!(explicit.metric.as_ref().unwrap().name(), Some("cpu")); assert_eq!( explicit.metric.as_ref().unwrap().get("host"), diff --git a/src/store-api/src/region_engine.rs b/src/store-api/src/region_engine.rs index 315133cdcf8..d9c4fa477f6 100644 --- a/src/store-api/src/region_engine.rs +++ b/src/store-api/src/region_engine.rs @@ -498,10 +498,12 @@ pub enum RawLabelSource { /// A stream of raw series chunks with stable schema, label binding, and ordering. /// -/// Producers MUST emit chunks for each read key contiguously, with strictly increasing timestamps -/// within and across chunks. Adapters validate adjacent chunk conditions but do not certify -/// nonadjacent recurrence for [`RawSeriesOrdering::Contiguous`]; an A/B/A sequence violates the -/// producer contract. `is_last` marks the last chunk of that key in this stream, not across streams. +/// Chunks for a read key are contiguous, with strictly increasing timestamps within and across +/// chunks. `is_last` marks the last chunk of that key in this stream, not across streams; merge +/// code must recompute finalization globally across all relevant inputs before downstream +/// per-input-series finalization (not aggregate or vector-match state). `RawLabelSource` keeps +/// the stream's label interpretation self-contained through transport rather than relying on a +/// local source handle. pub trait RawSeriesStream: Stream> + Send { /// Returns the schema shared by all chunks; raw schemas have no evaluation grid. fn schema(&self) -> &Arc; @@ -516,6 +518,48 @@ pub trait RawSeriesStream: Stream> + S /// A pinned raw-series stream that can be sent between threads. pub type SendableRawSeriesStream = Pin>; +/// Stable identity metadata for one ordinal in a prepared series source. +#[derive(Clone, Debug)] +pub struct SeriesInfo { + /// Complete encoded primary key, not a TSID or hash alone. + pub primary_key: datatypes::arrow::buffer::Buffer, + /// Complete explicit metric and labels, or `None` for primary-key label binding. + pub metric: Option, +} + +/// An immutable, finite selected set of series prepared for independent range reads. +/// +/// The count may include series with no samples after filtering. Ordinals are stable local +/// indexes, not protocol identifiers; each has exactly one complete identity returned by +/// [`PreparedSeriesSource::series_info`], and that identity matches chunks for that ordinal. +/// Every reader exposes the same fixed schema, label binding, and ordering before polling, +/// including for an empty range, and describes labels independently of this source's lifetime. `PrimaryKey` output has `metric: None`; `Explicit` output has a +/// complete `Some(metric)` matching each chunk. [`RawSeriesOrdering::Contiguous`] and +/// [`RawSeriesOrdering::ReadKeyThenTime`] describe stream order, not ordinal enumeration order. +/// +/// Preparation, identity lookup, and reader creation do not load sample payloads or perform +/// hidden blocking I/O. The selected set is bound to the partition's immutable MVCC view, +/// predicates, authorization, projection, labels, and existing scan resources. Readers own +/// resources needed to continue after this source or scanner is dropped. Dropping one reader +/// does not cancel sibling readers; normal ownership cleanup applies without detached work. +/// +/// A range must satisfy `start <= end <= num_series()`. Invalid ranges and invalid ordinals +/// return errors; an empty in-bounds range returns an ordinary empty stream. Repeated and +/// overlapping reads are independent and do not consume shared iteration state. +pub trait PreparedSeriesSource: Send + Sync { + /// Returns the exact, stable number of selected local series ordinals. + fn num_series(&self) -> usize; + + /// Returns identity metadata for an ordinal without polling sample data. + fn series_info(&self, ordinal: usize) -> Result; + + /// Creates an independent reader for the half-open ordinal range. + fn read_series( + &self, + range: std::ops::Range, + ) -> Result; +} + /// A scanner that provides a way to scan the region concurrently. /// /// The scanner splits the region into partitions so that each partition can be scanned concurrently. @@ -568,6 +612,22 @@ pub trait RegionScanner: Debug + DisplayAs + Send { todo!() } + /// Prepares a finite, stable set of series for independent ordinal-range reads. + /// + /// `None` means unsupported before consumption, so callers may use the legacy record-batch + /// scan. `Some`, including a source with zero series, selects this path; errors propagate + /// without replay. Choose exactly one entrypoint per partition per execution: this prepared + /// source, [`Self::scan_series_partition`], or [`Self::scan_partition`]. In particular, do + /// not call the legacy raw-scan default merely because this returns `None`. + fn prepare_series_partition( + &self, + _ctx: &QueryScanContext, + _metrics_set: &ExecutionPlanMetricsSet, + _partition: usize, + ) -> Result>, BoxedError> { + Ok(None) + } + /// Check if there is any predicate exclude region partition exprs that may be executed in this scanner. fn has_predicate_without_region(&self) -> bool; @@ -806,6 +866,10 @@ impl RegionStatistic { } } +#[cfg(test)] +#[path = "region_engine/prepared_series_tests.rs"] +mod prepared_series_tests; + #[cfg(test)] mod tests { use common_recordbatch::{RecordBatch, RecordBatches}; @@ -861,6 +925,13 @@ mod tests { let ctx = QueryScanContext::default(); let metrics = ExecutionPlanMetricsSet::new(); + assert!( + scanner + .prepare_series_partition(&ctx, &metrics, 0) + .unwrap() + .is_none() + ); + // SinglePartitionScanner uses the default raw-scan implementation, which panics. let raw_scan = catch_unwind(AssertUnwindSafe(|| { scanner.scan_series_partition(&ctx, &metrics, 0) diff --git a/src/store-api/src/region_engine/prepared_series_tests.rs b/src/store-api/src/region_engine/prepared_series_tests.rs new file mode 100644 index 00000000000..05882ddcc70 --- /dev/null +++ b/src/store-api/src/region_engine/prepared_series_tests.rs @@ -0,0 +1,405 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::VecDeque; +use std::pin::Pin; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; +use std::task::{Context, Poll}; + +use api::v1::SemanticType; +use common_error::ext::BoxedError; +use datatypes::arrow::array::Float64Array; +use datatypes::arrow::buffer::Buffer; +use datatypes::error::InvalidSeriesSnafu; +use datatypes::prelude::ConcreteDataType; +use datatypes::schema::ColumnSchema; +use datatypes::series::raw::{RawSeriesChunk, RawSeriesOrdering, RawTimestamps}; +use datatypes::series::samples::{FieldSpec, Samples, SeriesFields, SeriesSchema}; +use datatypes::series::{Label, MetricName}; +use futures::{Stream, StreamExt}; + +use crate::region_engine::{ + PreparedSeriesSource, RawLabelSource, RawSeriesStream, SendableRawSeriesStream, SeriesInfo, +}; + +#[derive(Clone)] +struct TestSeries { + info: SeriesInfo, + chunks: Vec, +} + +struct MemorySource { + series: Arc<[TestSeries]>, + schema: Arc, + label_source: RawLabelSource, + polls: Arc, + drops: Arc, +} + +impl PreparedSeriesSource for MemorySource { + fn num_series(&self) -> usize { + self.series.len() + } + + fn series_info(&self, ordinal: usize) -> Result { + self.series + .get(ordinal) + .map(|series| series.info.clone()) + .ok_or_else(|| { + BoxedError::new( + InvalidSeriesSnafu { + reason: "invalid series ordinal".to_string(), + } + .build(), + ) + }) + } + + fn read_series( + &self, + range: std::ops::Range, + ) -> Result { + if range.start > range.end || range.end > self.series.len() { + return Err(BoxedError::new( + InvalidSeriesSnafu { + reason: "invalid series range".to_string(), + } + .build(), + )); + } + let chunks = self.series[range] + .iter() + .flat_map(|series| series.chunks.iter().cloned()) + .collect(); + Ok(Box::pin(MemoryReader { + chunks, + schema: self.schema.clone(), + label_source: self.label_source.clone(), + polls: self.polls.clone(), + drops: self.drops.clone(), + })) + } +} + +struct MemoryReader { + chunks: VecDeque, + schema: Arc, + label_source: RawLabelSource, + polls: Arc, + drops: Arc, +} + +impl Stream for MemoryReader { + type Item = Result; + + fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + let chunk = self.chunks.pop_front(); + if chunk.is_some() { + self.polls.fetch_add(1, Ordering::SeqCst); + } + Poll::Ready(chunk.map(Ok)) + } +} + +impl Drop for MemoryReader { + fn drop(&mut self) { + self.drops.fetch_add(1, Ordering::SeqCst); + } +} + +impl RawSeriesStream for MemoryReader { + fn schema(&self) -> &Arc { + &self.schema + } + + fn label_source(&self) -> &RawLabelSource { + &self.label_source + } + + fn ordering(&self) -> RawSeriesOrdering { + RawSeriesOrdering::Contiguous + } +} + +fn metric(host: &str) -> MetricName { + MetricName::try_new( + Some(Arc::from("cpu")), + vec![Label { + name: Arc::from("host"), + value: Arc::from(host), + }], + ) + .unwrap() +} + +fn schema() -> Arc { + Arc::new(SeriesSchema { + fields: vec![FieldSpec { + name: Arc::from("usage"), + }] + .into(), + eval_grid: None, + }) +} + +fn chunk( + key: &[u8], + metric: Option, + timestamp: i64, + value: f64, + last: bool, +) -> RawSeriesChunk { + RawSeriesChunk::try_new( + &schema(), + Buffer::from(key), + metric, + RawTimestamps::try_new(vec![timestamp].into()).unwrap(), + SeriesFields::one(Samples::from_float(Float64Array::from(vec![value]))), + last, + ) + .unwrap() +} + +fn primary_key_metadata() -> crate::metadata::RegionMetadata { + let mut builder = + crate::metadata::RegionMetadataBuilder::new(crate::storage::RegionId::new(1, 1)); + builder.push_column_metadata(crate::metadata::ColumnMetadata { + column_schema: ColumnSchema::new("host", ConcreteDataType::string_datatype(), false), + semantic_type: SemanticType::Tag, + column_id: 1, + }); + builder.push_column_metadata(crate::metadata::ColumnMetadata { + column_schema: ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + semantic_type: SemanticType::Timestamp, + column_id: 2, + }); + builder.primary_key(vec![1]); + builder.build().unwrap() +} + +fn empty_source() -> Arc { + Arc::new(MemorySource { + series: Arc::from([]), + schema: schema(), + label_source: RawLabelSource::Explicit, + polls: Arc::new(AtomicUsize::new(0)), + drops: Arc::new(AtomicUsize::new(0)), + }) +} + +fn source(explicit: bool) -> Arc { + source_from_keys( + explicit, + vec![ + Buffer::from(vec![0x10, 0x01]), + Buffer::from(vec![0x10, 0x02]), + ] + .into(), + ) +} + +fn source_from_keys(explicit: bool, keys: Arc<[Buffer]>) -> Arc { + assert_eq!(keys.len(), 2); + let first_metric = explicit.then(|| metric("a")); + let second_metric = explicit.then(|| metric("b")); + let series = vec![ + TestSeries { + info: SeriesInfo { + primary_key: keys[0].clone(), + metric: first_metric.clone(), + }, + chunks: vec![chunk(keys[0].as_slice(), first_metric, 10, 1.0, true)], + }, + TestSeries { + info: SeriesInfo { + primary_key: keys[1].clone(), + metric: second_metric.clone(), + }, + chunks: vec![ + chunk(keys[1].as_slice(), second_metric.clone(), 10, 2.0, false), + chunk(keys[1].as_slice(), second_metric, 20, 3.0, true), + ], + }, + ]; + Arc::new(MemorySource { + series: series.into(), + schema: schema(), + label_source: if explicit { + RawLabelSource::Explicit + } else { + RawLabelSource::PrimaryKey { + metric_name: Arc::from("cpu"), + physical_metadata: Arc::new(primary_key_metadata()), + logical_tags: Arc::from([]), + } + }, + polls: Arc::new(AtomicUsize::new(0)), + drops: Arc::new(AtomicUsize::new(0)), + }) +} + +#[tokio::test] +async fn prepared_identity_matches_chunks_in_both_label_modes_including_empty_reader() { + for explicit in [false, true] { + let source = source(explicit); + assert_eq!(source.num_series(), 2); + let first = source.series_info(0).unwrap(); + let second = source.series_info(1).unwrap(); + assert_eq!(first.primary_key, Buffer::from(vec![0x10, 0x01])); + assert_eq!(second.primary_key, Buffer::from(vec![0x10, 0x02])); + assert_eq!( + first.metric.as_ref().and_then(MetricName::name), + if explicit { Some("cpu") } else { None } + ); + assert_eq!( + first.metric.as_ref().and_then(|m| m.get("host")), + explicit.then_some("a") + ); + assert_eq!( + second.metric.as_ref().and_then(|m| m.get("host")), + explicit.then_some("b") + ); + + let empty = source.read_series(1..1).unwrap(); + assert_eq!(empty.schema().fields.len(), 1); + assert_eq!(empty.ordering(), RawSeriesOrdering::Contiguous); + assert_eq!( + matches!(empty.label_source(), RawLabelSource::Explicit), + explicit + ); + assert_eq!(source.polls.load(Ordering::SeqCst), 0); + assert!( + futures::TryStreamExt::try_collect::>(empty) + .await + .unwrap() + .is_empty() + ); + assert_eq!(source.polls.load(Ordering::SeqCst), 0); + let end = source.read_series(2..2).unwrap(); + assert_eq!(end.schema().fields.len(), 1); + assert_eq!( + matches!(end.label_source(), RawLabelSource::Explicit), + explicit + ); + + let chunks = + futures::TryStreamExt::try_collect::>(source.read_series(0..2).unwrap()) + .await + .unwrap(); + assert_eq!(chunks.len(), 3); + assert_eq!(source.polls.load(Ordering::SeqCst), 3); + assert_eq!(chunks[1].time.timestamp(0), Some(10)); + assert_eq!(chunks[2].time.timestamp(0), Some(20)); + assert!( + matches!(chunks[1].fields.get(0).unwrap().sample(0), datatypes::series::samples::SampleRef::Float(value) if value == 2.0) + ); + assert!( + matches!(chunks[2].fields.get(0).unwrap().sample(0), datatypes::series::samples::SampleRef::Float(value) if value == 3.0) + ); + for chunk in &chunks { + let ordinal = if chunk.primary_key == first.primary_key { + 0 + } else { + 1 + }; + let info = source.series_info(ordinal).unwrap(); + assert_eq!(chunk.primary_key, info.primary_key); + assert_eq!(chunk.metric, info.metric); + } + assert_eq!(chunks[0].metric.is_some(), explicit); + assert!(!chunks[1].is_last); + assert!(chunks[2].is_last); + } +} + +#[tokio::test] +async fn empty_prepared_source_supports_its_empty_range() { + let source = empty_source(); + assert_eq!(source.num_series(), 0); + let stream = source.read_series(0..0).unwrap(); + assert_eq!(stream.schema().fields.len(), 1); + assert!( + futures::TryStreamExt::try_collect::>(stream) + .await + .unwrap() + .is_empty() + ); + assert_eq!(source.polls.load(Ordering::SeqCst), 0); +} + +#[test] +fn invalid_ordinals_and_ranges_fail_without_polling_samples() { + let source = source(false); + assert!(source.series_info(2).is_err()); + assert!(source.read_series(2..3).is_err()); + assert!(source.read_series(2..1).is_err()); + assert_eq!(source.polls.load(Ordering::SeqCst), 0); +} + +#[tokio::test] +async fn overlapping_readers_are_independent_across_source_and_reader_drops() { + let source = source(false); + let drops = source.drops.clone(); + let mut left = source.read_series(0..2).unwrap(); + let mut right = source.read_series(1..2).unwrap(); + drop(source); + assert!(left.next().await.unwrap().unwrap().is_last); + assert_eq!( + right.next().await.unwrap().unwrap().primary_key, + Buffer::from(vec![0x10, 0x02]) + ); + drop(left); + assert_eq!( + futures::TryStreamExt::try_collect::>(right) + .await + .unwrap() + .len(), + 1 + ); + assert_eq!(drops.load(Ordering::SeqCst), 2); +} + +#[tokio::test] +async fn prepared_selection_is_independent_of_later_mutable_selection_changes() { + let selected = Arc::new(Mutex::new(vec![ + Buffer::from(vec![1_u8]), + Buffer::from(vec![2_u8]), + ])); + let prepared = source_from_keys(false, selected.lock().unwrap().clone().into()); + selected.lock().unwrap().push(Buffer::from(vec![3_u8])); + + assert_eq!(prepared.num_series(), 2); + assert_eq!( + prepared.series_info(0).unwrap().primary_key, + Buffer::from(vec![1_u8]) + ); + assert_eq!( + prepared.series_info(1).unwrap().primary_key, + Buffer::from(vec![2_u8]) + ); + let chunks = futures::TryStreamExt::try_collect::>(prepared.read_series(0..2).unwrap()) + .await + .unwrap(); + assert_eq!(chunks.len(), 3); + assert!( + chunks + .iter() + .all(|chunk| chunk.primary_key != Buffer::from(vec![3_u8])) + ); +} From d0180318237eae243a35beabd549f59c548256a2 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Fri, 16 Oct 2026 15:45:42 +0800 Subject: [PATCH 08/11] feat(datatypes): share series ordering and expose source label binding Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/datatypes/src/series.rs | 21 ++++ src/datatypes/src/series/raw.rs | 15 +-- src/store-api/src/region_engine.rs | 40 +++++--- .../region_engine/prepared_series_tests.rs | 97 +++++++++++++++---- 4 files changed, 129 insertions(+), 44 deletions(-) diff --git a/src/datatypes/src/series.rs b/src/datatypes/src/series.rs index f00051ce46f..0de711e2da3 100644 --- a/src/datatypes/src/series.rs +++ b/src/datatypes/src/series.rs @@ -25,6 +25,27 @@ use std::sync::Arc; use crate::error::{InvalidSeriesSnafu, Result}; +/// Ordering guarantee within one stream and each of its output partitions. +/// +/// Every series has strictly increasing sample positions, including across chunks; positions +/// are raw timestamps or, for evaluated series, current-stage evaluation times or grid positions. +/// They need not all be explicit times on the grid. No mode permits +/// duplicate positions. `Interleaved` permits A1, B1, A2; only the latter two modes keep each +/// series contiguous. `KeyThenTime` also orders by this stream's series identity, then time: +/// raw streams use the complete bound-compatible identity (not just a TSID or hash), while +/// evaluated streams use current name/label identity, not storage primary keys or object address/hash. +/// This ordering +/// is not a query-global ordering, partition-ownership rule, or finality guarantee. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum SeriesOrdering { + /// Series chunks may interleave; positions within each series remain strictly increasing. + Interleaved, + /// Each series is contiguous; distinct series need not be ordered. + Contiguous, + /// Series are ordered by stream identity and then by strictly increasing position. + KeyThenTime, +} + /// A label with shared UTF-8 name and value strings. #[derive(Clone, Debug, PartialEq, Eq)] pub struct Label { diff --git a/src/datatypes/src/series/raw.rs b/src/datatypes/src/series/raw.rs index 130f59565a9..da4ba487204 100644 --- a/src/datatypes/src/series/raw.rs +++ b/src/datatypes/src/series/raw.rs @@ -85,20 +85,11 @@ impl RawTimestamps { } } -/// Producer ordering guarantee for chunks in one input stream. Producers MUST emit all chunks -/// for each read key contiguously; adapters validate adjacent chunk conditions but do not certify -/// nonadjacent recurrence for [`Self::Contiguous`]. An `A/B/A` sequence violates the producer -/// contract even when a consumer cannot detect that recurrence. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum RawSeriesOrdering { - /// Each complete primary key's chunks are contiguous; distinct keys need not be ordered. - Contiguous, - /// Chunks are ordered by complete primary key and then by timestamp. - ReadKeyThenTime, -} - /// One raw source batch, with fields aligned to its timestamps. /// +/// `metric: None` means labels are not yet decoded because identity is bound by the source; it +/// does not mean the metric name is absent. An explicit `Some(metric)` is always complete. +/// /// Public fields allow direct access and construction, which bypasses the checks in /// [`RawSeriesChunk::try_new`]; callers crossing a trust boundary must validate input. #[derive(Clone, Debug)] diff --git a/src/store-api/src/region_engine.rs b/src/store-api/src/region_engine.rs index d9c4fa477f6..59a959c09f5 100644 --- a/src/store-api/src/region_engine.rs +++ b/src/store-api/src/region_engine.rs @@ -30,7 +30,8 @@ use common_time::Timestamp; use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; use datafusion_physical_plan::{DisplayAs, DisplayFormatType, PhysicalExpr}; use datatypes::schema::SchemaRef; -use datatypes::series::raw::{RawSeriesChunk, RawSeriesOrdering}; +use datatypes::series::SeriesOrdering; +use datatypes::series::raw::RawSeriesChunk; use datatypes::series::samples::SeriesSchema; use futures::Stream; use futures::future::join_all; @@ -492,27 +493,33 @@ pub enum RawLabelSource { /// Logical tag names mapped to physical column identifiers. logical_tags: Arc<[BoundTag]>, }, - /// The stream supplies a complete metric name and label set in each chunk. + /// The stream supplies a complete metric name and label set in each chunk, not a partial + /// label projection; `Some(metric)` is a complete name/label set (the name itself may be + /// absent). Explicit, } /// A stream of raw series chunks with stable schema, label binding, and ordering. /// -/// Chunks for a read key are contiguous, with strictly increasing timestamps within and across -/// chunks. `is_last` marks the last chunk of that key in this stream, not across streams; merge -/// code must recompute finalization globally across all relevant inputs before downstream -/// per-input-series finalization (not aggregate or vector-match state). `RawLabelSource` keeps -/// the stream's label interpretation self-contained through transport rather than relying on a -/// local source handle. +/// Every series has strictly increasing timestamps within and across chunks; timestamps are +/// unique per series. `Interleaved` permits a series to resume after another series; `Contiguous` +/// and `KeyThenTime` guarantee each series' chunks are contiguous. `KeyThenTime` uses the complete +/// stream identity compatible with its label binding, not just a TSID or hash. These are +/// per-stream guarantees, not query-global order, partition ownership, or finality. `is_last` +/// marks the last chunk of that series in this stream only; merging requires recomputing +/// per-input-series finalization across all related inputs, not aggregate-group finality. +/// `RawLabelSource` keeps label interpretation self-contained through transport. Stream finality +/// applies to all related input series, not aggregate groups or vector-match state. pub trait RawSeriesStream: Stream> + Send { /// Returns the schema shared by all chunks; raw schemas have no evaluation grid. fn schema(&self) -> &Arc; /// Returns scan-bound label metadata. `PrimaryKey` chunks leave `metric` as `None`, meaning /// labels are undecoded, not that the metric name is missing; `Explicit` chunks supply - /// `Some` with a complete metric name and labels. + /// `Some` with the complete name/label set, not a partial projection. The metric name itself + /// may be absent. fn label_source(&self) -> &RawLabelSource; /// Returns this stream's ordering guarantee. - fn ordering(&self) -> RawSeriesOrdering; + fn ordering(&self) -> SeriesOrdering; } /// A pinned raw-series stream that can be sent between threads. @@ -523,7 +530,8 @@ pub type SendableRawSeriesStream = Pin>; pub struct SeriesInfo { /// Complete encoded primary key, not a TSID or hash alone. pub primary_key: datatypes::arrow::buffer::Buffer, - /// Complete explicit metric and labels, or `None` for primary-key label binding. + /// Complete explicit metric name and labels, not a partial label projection; or `None` for + /// primary-key label binding (undecoded, not nameless). The metric name itself may be absent. pub metric: Option, } @@ -533,9 +541,9 @@ pub struct SeriesInfo { /// indexes, not protocol identifiers; each has exactly one complete identity returned by /// [`PreparedSeriesSource::series_info`], and that identity matches chunks for that ordinal. /// Every reader exposes the same fixed schema, label binding, and ordering before polling, -/// including for an empty range, and describes labels independently of this source's lifetime. `PrimaryKey` output has `metric: None`; `Explicit` output has a -/// complete `Some(metric)` matching each chunk. [`RawSeriesOrdering::Contiguous`] and -/// [`RawSeriesOrdering::ReadKeyThenTime`] describe stream order, not ordinal enumeration order. +/// including for an empty range, and describes labels independently of this source's lifetime. +/// `PrimaryKey` output has `metric: None` (undecoded); `Explicit` output has a complete +/// `Some(metric)` matching each chunk. Ordering describes the stream, not ordinal enumeration. /// /// Preparation, identity lookup, and reader creation do not load sample payloads or perform /// hidden blocking I/O. The selected set is bound to the partition's immutable MVCC view, @@ -550,6 +558,10 @@ pub trait PreparedSeriesSource: Send + Sync { /// Returns the exact, stable number of selected local series ordinals. fn num_series(&self) -> usize; + /// Returns label binding shared by series metadata and every reader, without opening or + /// polling a reader; the binding is self-contained for remote consumers. + fn label_source(&self) -> &RawLabelSource; + /// Returns identity metadata for an ordinal without polling sample data. fn series_info(&self, ordinal: usize) -> Result; diff --git a/src/store-api/src/region_engine/prepared_series_tests.rs b/src/store-api/src/region_engine/prepared_series_tests.rs index 05882ddcc70..a489598b567 100644 --- a/src/store-api/src/region_engine/prepared_series_tests.rs +++ b/src/store-api/src/region_engine/prepared_series_tests.rs @@ -25,7 +25,8 @@ use datatypes::arrow::buffer::Buffer; use datatypes::error::InvalidSeriesSnafu; use datatypes::prelude::ConcreteDataType; use datatypes::schema::ColumnSchema; -use datatypes::series::raw::{RawSeriesChunk, RawSeriesOrdering, RawTimestamps}; +use datatypes::series::SeriesOrdering; +use datatypes::series::raw::{RawSeriesChunk, RawTimestamps}; use datatypes::series::samples::{FieldSpec, Samples, SeriesFields, SeriesSchema}; use datatypes::series::{Label, MetricName}; use futures::{Stream, StreamExt}; @@ -53,6 +54,10 @@ impl PreparedSeriesSource for MemorySource { self.series.len() } + fn label_source(&self) -> &RawLabelSource { + &self.label_source + } + fn series_info(&self, ordinal: usize) -> Result { self.series .get(ordinal) @@ -128,8 +133,8 @@ impl RawSeriesStream for MemoryReader { &self.label_source } - fn ordering(&self) -> RawSeriesOrdering { - RawSeriesOrdering::Contiguous + fn ordering(&self) -> SeriesOrdering { + SeriesOrdering::Contiguous } } @@ -144,6 +149,33 @@ fn metric(host: &str) -> MetricName { .unwrap() } +fn assert_same_label_source(expected: &RawLabelSource, actual: &RawLabelSource) { + match (expected, actual) { + (RawLabelSource::Explicit, RawLabelSource::Explicit) => {} + ( + RawLabelSource::PrimaryKey { + metric_name: expected_name, + physical_metadata: expected_metadata, + logical_tags: expected_tags, + }, + RawLabelSource::PrimaryKey { + metric_name: actual_name, + physical_metadata: actual_metadata, + logical_tags: actual_tags, + }, + ) => { + assert_eq!(expected_name, actual_name); + assert!(Arc::ptr_eq(expected_metadata, actual_metadata)); + assert_eq!(expected_tags.len(), actual_tags.len()); + for (expected, actual) in expected_tags.iter().zip(actual_tags.iter()) { + assert_eq!(expected.name, actual.name); + assert_eq!(expected.column_id, actual.column_id); + } + } + _ => panic!("label source differs"), + } +} + fn schema() -> Arc { Arc::new(SeriesSchema { fields: vec![FieldSpec { @@ -246,7 +278,11 @@ fn source_from_keys(explicit: bool, keys: Arc<[Buffer]>) -> Arc { RawLabelSource::PrimaryKey { metric_name: Arc::from("cpu"), physical_metadata: Arc::new(primary_key_metadata()), - logical_tags: Arc::from([]), + logical_tags: vec![crate::region_engine::BoundTag { + column_id: 1, + name: Arc::from("host"), + }] + .into(), } }, polls: Arc::new(AtomicUsize::new(0)), @@ -255,10 +291,36 @@ fn source_from_keys(explicit: bool, keys: Arc<[Buffer]>) -> Arc { } #[tokio::test] -async fn prepared_identity_matches_chunks_in_both_label_modes_including_empty_reader() { +async fn prepared_source_label_binding_is_available_before_readers_and_matches_chunks() { for explicit in [false, true] { let source = source(explicit); assert_eq!(source.num_series(), 2); + match (explicit, source.label_source()) { + (true, RawLabelSource::Explicit) => {} + ( + false, + RawLabelSource::PrimaryKey { + metric_name, + physical_metadata, + logical_tags, + }, + ) => { + assert_eq!(metric_name.as_ref(), "cpu"); + assert_eq!( + physical_metadata + .column_by_id(1) + .unwrap() + .column_schema + .name, + "host" + ); + assert_eq!(logical_tags.len(), 1); + assert_eq!(logical_tags[0].column_id, 1); + assert_eq!(logical_tags[0].name.as_ref(), "host"); + } + _ => panic!("unexpected source label binding"), + } + assert_eq!(source.polls.load(Ordering::SeqCst), 0); let first = source.series_info(0).unwrap(); let second = source.series_info(1).unwrap(); assert_eq!(first.primary_key, Buffer::from(vec![0x10, 0x01])); @@ -278,11 +340,8 @@ async fn prepared_identity_matches_chunks_in_both_label_modes_including_empty_re let empty = source.read_series(1..1).unwrap(); assert_eq!(empty.schema().fields.len(), 1); - assert_eq!(empty.ordering(), RawSeriesOrdering::Contiguous); - assert_eq!( - matches!(empty.label_source(), RawLabelSource::Explicit), - explicit - ); + assert_eq!(empty.ordering(), SeriesOrdering::Contiguous); + assert_same_label_source(source.label_source(), empty.label_source()); assert_eq!(source.polls.load(Ordering::SeqCst), 0); assert!( futures::TryStreamExt::try_collect::>(empty) @@ -293,15 +352,16 @@ async fn prepared_identity_matches_chunks_in_both_label_modes_including_empty_re assert_eq!(source.polls.load(Ordering::SeqCst), 0); let end = source.read_series(2..2).unwrap(); assert_eq!(end.schema().fields.len(), 1); - assert_eq!( - matches!(end.label_source(), RawLabelSource::Explicit), - explicit - ); + assert_eq!(end.ordering(), SeriesOrdering::Contiguous); + assert_same_label_source(source.label_source(), end.label_source()); + assert_eq!(source.polls.load(Ordering::SeqCst), 0); - let chunks = - futures::TryStreamExt::try_collect::>(source.read_series(0..2).unwrap()) - .await - .unwrap(); + let reader = source.read_series(0..2).unwrap(); + assert_same_label_source(source.label_source(), reader.label_source()); + assert_eq!(source.polls.load(Ordering::SeqCst), 0); + let chunks = futures::TryStreamExt::try_collect::>(reader) + .await + .unwrap(); assert_eq!(chunks.len(), 3); assert_eq!(source.polls.load(Ordering::SeqCst), 3); assert_eq!(chunks[1].time.timestamp(0), Some(10)); @@ -331,6 +391,7 @@ async fn prepared_identity_matches_chunks_in_both_label_modes_including_empty_re #[tokio::test] async fn empty_prepared_source_supports_its_empty_range() { let source = empty_source(); + assert!(matches!(source.label_source(), RawLabelSource::Explicit)); assert_eq!(source.num_series(), 0); let stream = source.read_series(0..0).unwrap(); assert_eq!(stream.schema().fields.len(), 1); From 3406a5334574a518924c67c319e2d92b6cac943b Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Fri, 16 Oct 2026 16:32:50 +0800 Subject: [PATCH 09/11] perf(datatypes): encode bounded multirow series frames Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/datatypes/src/series/frame.rs | 964 ++++++++++++++++++------ src/datatypes/src/series/frame/tests.rs | 882 ++++++++++++++++++++-- src/datatypes/src/series/samples.rs | 12 +- 3 files changed, 1544 insertions(+), 314 deletions(-) diff --git a/src/datatypes/src/series/frame.rs b/src/datatypes/src/series/frame.rs index d9f9c1ee347..27902ecf805 100644 --- a/src/datatypes/src/series/frame.rs +++ b/src/datatypes/src/series/frame.rs @@ -13,6 +13,12 @@ // 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; @@ -24,12 +30,13 @@ use arrow::array::{ 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, RawTimeAxis}; +use crate::series::raw::{RawSeriesChunk, RawTimestamps}; use crate::series::samples::{ DenseSamplesRef, HistogramSamples, Samples, SeriesFields, SeriesSchema, }; @@ -130,66 +137,714 @@ pub fn raw_series_frame_schema(schema: &SeriesSchema, mode: RawLabelMode) -> Res Ok(Arc::new(Schema::new_with_metadata(columns, metadata))) } -/// Packs one raw chunk into one-row batches at the union of field segment boundaries. +/// 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() { + if 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 checked = RawSeriesChunk::try_new( + let mut builder = RawSeriesFrameBuilder::try_new( schema, - chunk.read_key.clone(), - chunk.metric.clone(), - chunk.time.clone(), - chunk.fields.clone(), - chunk.is_last, + mode, + RawSeriesFrameBuilder::DEFAULT_MAX_ROWS, + RawSeriesFrameBuilder::DEFAULT_MAX_BYTES, )?; - match (mode, &checked.metric) { - (RawLabelMode::Key, None) | (RawLabelMode::Explicit, Some(_)) => {} - _ => return invalid_series("raw chunk label mode does not match frame mode"), - } - let frame_schema = raw_series_frame_schema(schema, mode)?; - let RawTimeAxis::Explicit(times) = &checked.time; - let all: Vec<&[Samples]> = checked.fields.iter().map(Samples::segments).collect(); - let mut segment = vec![0usize; all.len()]; - let mut offset = vec![0usize; all.len()]; - let mut pos = 0usize; + let mut cursor = RawSeriesChunkCursor::try_new(schema, mode, chunk.clone())?; let mut batches = Vec::new(); - while pos < times.len() { - let take = all - .iter() - .enumerate() - .map(|(i, parts)| parts[segment[i]].len() - offset[i]) - .min() - .unwrap_or(times.len()) - .min(times.len() - pos); - if take == 0 { - return invalid_series("empty sample segment"); - } - let pieces = all - .iter() - .enumerate() - .map(|(i, parts)| parts[segment[i]].slice(offset[i], take)) - .collect::>>()?; - batches.push(encode_row( - &frame_schema, - mode, - &checked.read_key, - checked.metric.as_ref(), - checked.is_last && pos + take == times.len(), - times.slice(pos, take), - &pieces, - )?); - pos += take; - for (i, parts) in all.iter().enumerate() { - offset[i] += take; - if offset[i] == parts[segment[i]].len() { - segment[i] += 1; - offset[i] = 0; + 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())?)?; } } } - Ok(batches) + 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. @@ -198,7 +853,15 @@ pub fn unpack_raw_series_frame( mode: RawLabelMode, batch: &RecordBatch, ) -> Result> { - let expected = raw_series_frame_schema(schema, mode)?; + 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) { @@ -307,18 +970,17 @@ pub fn unpack_raw_series_frame( } let table = down::(&tables.value(row), "histogram table")?.clone(); let batch = Arc::new(HistogramBatch::try_from_arrow(table)?); - let hs = HistogramSamples::try_new(indices, batch)?; if has_float { - Samples::try_from_mixed(floats, hs)? + Samples::try_from_mixed(floats, HistogramSamples::try_new(indices, batch)?)? } else { - Samples::try_from_histogram(hs.indices().clone(), hs.batch().clone())? + 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 = RawTimeAxis::try_new(time)?; + let raw_time = RawTimestamps::try_new(time)?; out.push(RawSeriesChunk::try_new( schema, key, @@ -331,149 +993,6 @@ pub fn unpack_raw_series_frame( Ok(out) } -fn encode_row( - schema: &SchemaRef, - mode: RawLabelMode, - key: &Buffer, - metric: Option<&MetricName>, - last: bool, - times: ScalarBuffer, - samples: &[Samples], -) -> Result { - let mut arrays: Vec = vec![ - Arc::new(binary_one(key.as_slice())?), - Arc::new(BooleanArray::from(vec![last])), - ]; - let ts = wrap_list( - Arc::new(Int64Array::new(times, None)), - DataType::Int64, - false, - )?; - let DataType::Struct(time_fields) = schema.field(2).data_type() else { - return invalid_series("time schema mismatch"); - }; - arrays.push(Arc::new(struct_array( - time_fields.clone(), - vec![Arc::new(ts)], - )?)); - let DataType::Struct(field_schema) = schema.field(3).data_type() else { - return invalid_series("fields schema mismatch"); - }; - let DataType::Struct(first_payload) = field_schema[0].data_type() else { - return invalid_series("payload schema mismatch"); - }; - let DataType::Struct(histogram_fields) = first_payload[1].data_type() else { - return invalid_series("histogram schema mismatch"); - }; - let mut field_arrays = Vec::new(); - for sample in samples { - let DataType::Struct(dtype) = field_schema[field_arrays.len()].data_type() else { - return invalid_series("payload schema mismatch"); - }; - let (floats, hist) = encode_payload(sample, histogram_fields)?; - field_arrays.push(Arc::new(struct_array(dtype.clone(), vec![floats, hist])?) as ArrayRef); - } - arrays.push(Arc::new(struct_array(field_schema.clone(), field_arrays)?)); - if mode == RawLabelMode::Explicit { - let metric = metric.ok_or_else(|| err("explicit labels missing"))?; - checked_offset(metric.name().map_or(0, str::len))?; - arrays.push(Arc::new(StringArray::from(vec![metric.name()]))); - checked_offset(metric.tags().len())?; - checked_offset( - metric - .tags() - .iter() - .try_fold(0usize, |n, l| n.checked_add(l.name.len())) - .ok_or_else(|| err("label name length overflow"))?, - )?; - checked_offset( - metric - .tags() - .iter() - .try_fold(0usize, |n, l| n.checked_add(l.value.len())) - .ok_or_else(|| err("label value length overflow"))?, - )?; - let names: StringArray = metric - .tags() - .iter() - .map(|l| Some(l.name.as_ref())) - .collect(); - let values: StringArray = metric - .tags() - .iter() - .map(|l| Some(l.value.as_ref())) - .collect(); - let DataType::List(label_list) = schema.field(5).data_type() else { - return invalid_series("labels schema mismatch"); - }; - let DataType::Struct(label_fields) = label_list.data_type() else { - return invalid_series("label struct schema mismatch"); - }; - let labels = struct_array( - label_fields.clone(), - vec![Arc::new(names), Arc::new(values)], - )?; - arrays.push(Arc::new(wrap_list( - Arc::new(labels), - DataType::Struct(label_fields.clone()), - false, - )?)); - } - arrow(RecordBatch::try_new_with_options( - schema.clone(), - arrays, - &RecordBatchOptions::new().with_row_count(Some(1)), - )) -} - -fn encode_payload(samples: &Samples, histogram_fields: &Fields) -> Result<(ArrayRef, ArrayRef)> { - let (floats, histograms) = match samples - .dense() - .ok_or_else(|| err("sample segment is not dense"))? - { - DenseSamplesRef::Float(values) => (Some(values), None), - DenseSamplesRef::Histogram(values) => (None, Some(values)), - DenseSamplesRef::Mixed(values) => (Some(values.floats()), Some(values.histograms())), - }; - let float_array = match floats { - Some(values) => Arc::new(wrap_list( - Arc::new(values.clone()), - DataType::Float64, - true, - )?) as ArrayRef, - None => Arc::new(ListArray::new_null( - Arc::new(Field::new("item", DataType::Float64, true)), - 1, - )) as ArrayRef, - }; - let hist_array = match histograms { - Some(values) => { - let DataType::List(indices_field) = histogram_fields[0].data_type() else { - return invalid_series("histogram index schema mismatch"); - }; - let DataType::List(table_field) = histogram_fields[1].data_type() else { - return invalid_series("histogram table schema mismatch"); - }; - let indices = wrap_list( - Arc::new(values.indices().clone()), - indices_field.data_type().clone(), - true, - )?; - let table = wrap_list( - Arc::new(values.batch().as_arrow().clone()), - table_field.data_type().clone(), - false, - )?; - Arc::new(struct_array( - histogram_fields.clone(), - vec![Arc::new(indices), Arc::new(table)], - )?) as ArrayRef - } - None => null_histogram(histogram_fields)?, - }; - Ok((float_array, hist_array)) -} - 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"))?; @@ -483,68 +1002,29 @@ fn binary_slice(array: &BinaryArray, row: usize) -> Result { .ok_or_else(|| err("invalid binary offsets"))?; Ok(array.values().slice_with_length(start, len)) } -fn binary_one(value: &[u8]) -> Result { - checked_offset(value.len())?; - Ok(BinaryArray::from(vec![Some(value)])) -} + fn checked_offset(len: usize) -> Result { i32::try_from(len).map_err(|_| err("Arrow offset limit exceeded")) } -fn offsets(len: usize) -> Result> { - let len = checked_offset(len)?; - Ok(OffsetBuffer::new(ScalarBuffer::from(vec![0, len]))) -} -fn wrap_list(values: ArrayRef, item: DataType, nullable: bool) -> Result { - let len = values.len(); - arrow(ListArray::try_new( - Arc::new(Field::new("item", item, nullable)), - offsets(len)?, - values, - None, - )) -} -fn null_histogram(fields: &Fields) -> Result { - let DataType::List(indices_field) = fields[0].data_type() else { - return invalid_series("histogram index schema mismatch"); - }; - let DataType::List(table_field) = fields[1].data_type() else { - return invalid_series("histogram table schema mismatch"); - }; - let DataType::Struct(table_fields) = table_field.data_type() else { - return invalid_series("histogram table schema mismatch"); - }; - let indices = wrap_list( - Arc::new(UInt32Array::new_null(0)), - indices_field.data_type().clone(), - true, - )?; - let columns = table_fields - .iter() - .map(|f| arrow::array::new_null_array(f.data_type(), 0)) - .collect(); - let table = arrow(StructArray::try_new(table_fields.clone(), columns, None))?; - let table = wrap_list(Arc::new(table), table_field.data_type().clone(), false)?; - Ok(Arc::new(arrow(StructArray::try_new( - fields.clone(), - vec![Arc::new(indices), Arc::new(table)], - Some(NullBuffer::new_null(1)), - ))?)) -} + 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 index 7edc38168b5..976a433ee27 100644 --- a/src/datatypes/src/series/frame/tests.rs +++ b/src/datatypes/src/series/frame/tests.rs @@ -15,11 +15,13 @@ use std::sync::Arc; use arrow::array::{ - Array, Float64Array, Int64Array, ListArray, StringArray, StructArray, UInt32Array, + 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}; @@ -67,7 +69,7 @@ fn float_chunk( schema, Buffer::from(key.to_vec()), metric, - RawTimeAxis::try_new(times.into()).unwrap(), + RawTimestamps::try_new(times.into()).unwrap(), SeriesFields::one(Samples::from_float(Float64Array::from(values))), is_last, ) @@ -247,6 +249,107 @@ fn assert_float_bits(actual: &Samples, expected: &[Option]) { } } +/// 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"]); @@ -263,13 +366,22 @@ fn sliced_histogram_frames_densify_without_duplicating_shared_table() { &schema, Buffer::from(vec![0x42_u8]), None, - RawTimeAxis::try_new((0_i64..64).collect::>().into()).unwrap(), + RawTimestamps::try_new((0_i64..64).collect::>().into()).unwrap(), SeriesFields::one(samples), true, ) .unwrap(); - let frames = pack_raw_series_chunk(&schema, RawLabelMode::Key, &chunk).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 { @@ -305,6 +417,24 @@ fn sliced_histogram_frames_densify_without_duplicating_shared_table() { 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() @@ -354,7 +484,7 @@ fn raw_frame_field_names_bind_payloads_to_their_schema() { &schema, Buffer::from(vec![0xaa_u8]), metric, - RawTimeAxis::try_new(vec![100, 200].into()).unwrap(), + RawTimestamps::try_new(vec![100, 200].into()).unwrap(), fields.clone(), true, ) @@ -375,9 +505,9 @@ fn raw_frame_field_names_bind_payloads_to_their_schema() { ); let decoded = unpack_raw_series_frame(&schema, mode, batch).unwrap(); - assert_eq!(decoded[0].read_key.as_slice(), &[0xaa]); + assert_eq!(decoded[0].primary_key.as_slice(), &[0xaa]); assert!(decoded[0].is_last); - let RawTimeAxis::Explicit(times) = &decoded[0].time; + 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)]); @@ -499,7 +629,7 @@ fn schema_and_key_explicit_round_trips_preserve_identity_and_bits() { 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].read_key.as_slice(), &[0, 0xff, 0, 0x80]); + 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!( @@ -520,7 +650,7 @@ fn schema_and_key_explicit_round_trips_preserve_identity_and_bits() { ); 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].read_key.as_slice(), &[0x80, 0, 0xff]); + 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); @@ -577,28 +707,31 @@ fn segmented_mixed_fields_emit_every_aligned_position_and_marker() { &schema, Buffer::from(vec![0_u8, 0x91, 0xfe]), None, - RawTimeAxis::try_new(vec![10_i64, 20, 30, 40].into()).unwrap(), + 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(), 4); + 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(); - for (frame_index, frame) in frames.iter().enumerate() { - let decoded = unpack_raw_series_frame(&schema, RawLabelMode::Key, frame).unwrap(); - assert_eq!(decoded.len(), 1); - assert_eq!(decoded[0].is_last, frame_index == frames.len() - 1); - observed_markers.push(decoded[0].is_last); - assert_eq!(decoded[0].read_key.as_slice(), &[0, 0x91, 0xfe]); - observed_keys.push(decoded[0].read_key.as_slice().to_vec()); - for position in 0..decoded[0].time.len() { - observed_times.push(decoded[0].time.timestamp(position).unwrap()); - match decoded[0].fields.get(0).unwrap().sample(position) { + 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) => { @@ -606,7 +739,7 @@ fn segmented_mixed_fields_emit_every_aligned_position_and_marker() { observed_first.push(("h", histogram.sum.to_bits())); } } - match decoded[0].fields.get(1).unwrap().sample(position) { + 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) => { @@ -657,59 +790,59 @@ fn segmented_mixed_fields_emit_every_aligned_position_and_marker() { &aligned_schema, Buffer::from(vec![0xa1_u8]), None, - RawTimeAxis::try_new(vec![101_i64, 102, 103, 104, 105].into()).unwrap(), + 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, batch) in aligned_frames.iter().enumerate() { - let decoded = unpack_raw_series_frame(&aligned_schema, RawLabelMode::Key, batch).unwrap(); - assert_eq!(decoded[0].is_last, i == 2); - interval_lengths.push(decoded[0].time.len()); + 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) = decoded[0].fields.get(0).unwrap().dense().unwrap() - else { + let DenseSamplesRef::Float(decoded_a) = row.fields.get(0).unwrap().dense().unwrap() else { panic!("expected float field"); }; - let DenseSamplesRef::Float(decoded_b) = decoded[0].fields.get(1).unwrap().dense().unwrap() - else { + let DenseSamplesRef::Float(decoded_b) = row.fields.get(1).unwrap().dense().unwrap() else { panic!("expected float field"); }; - assert_eq!( - decoded_a.values().as_ptr(), - a_owner.values().as_ptr().wrapping_add(start) - ); - assert_eq!( - decoded_b.values().as_ptr(), - b_owner.values().as_ptr().wrapping_add(start) - ); + // 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)] { - match (decoded_values.nulls(), owner.nulls()) { - (Some(decoded_nulls), Some(owner_nulls)) => { - assert_eq!( - decoded_nulls.buffer().as_ptr(), - owner_nulls.buffer().as_ptr() - ); - assert_eq!(decoded_nulls.offset(), owner_nulls.offset() + start); + 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"), } - (None, None) => {} - _ => panic!("fragment changed validity representation"), } } - for position in 0..decoded[0].time.len() { - aligned_times.push(decoded[0].time.timestamp(position).unwrap()); - aligned_a.push(match decoded[0].fields.get(0).unwrap().sample(position) { + 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 decoded[0].fields.get(1).unwrap().sample(position) { + 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"), @@ -757,7 +890,7 @@ fn segmented_mixed_fields_emit_every_aligned_position_and_marker() { &alternate_schema, Buffer::from(vec![0xde_u8, 0xad]), None, - RawTimeAxis::try_new(vec![100_i64].into()).unwrap(), + RawTimestamps::try_new(vec![100_i64].into()).unwrap(), SeriesFields::try_from_vec(vec![first_distinct, second_distinct]).unwrap(), false, ) @@ -781,7 +914,7 @@ fn segmented_mixed_fields_emit_every_aligned_position_and_marker() { .collect::>(), [10, 20, 30, 40, 100] ); - assert_eq!(decoded[4].read_key.as_slice(), &[0xde, 0xad]); + 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()) ); @@ -807,7 +940,7 @@ fn segmented_mixed_fields_emit_every_aligned_position_and_marker() { assert_eq!( sliced_rows .iter() - .map(|row| row.read_key.as_slice()) + .map(|row| row.primary_key.as_slice()) .collect::>(), [&[0xee][..], &[0xef][..], &[0xed][..]] ); @@ -837,7 +970,7 @@ fn segmented_mixed_fields_emit_every_aligned_position_and_marker() { &schema, Buffer::from(vec![0_u8, 0x91, 0xfe]), None, - RawTimeAxis::try_new(vec![50_i64, 60, 70, 80].into()).unwrap(), + RawTimestamps::try_new(vec![50_i64, 60, 70, 80].into()).unwrap(), chunk.fields.clone(), false, ) @@ -898,7 +1031,7 @@ fn histogram_payload_and_present_all_null_representation_survive() { &schema, Buffer::from(vec![0x12_u8, 0x00]), None, - RawTimeAxis::try_new(vec![1_i64, 2].into()).unwrap(), + RawTimestamps::try_new(vec![1_i64, 2].into()).unwrap(), fields, false, ) @@ -950,7 +1083,7 @@ fn histogram_payload_and_present_all_null_representation_survive() { &mixed_schema, Buffer::from(vec![0x42_u8]), None, - RawTimeAxis::try_new(vec![5_i64, 6].into()).unwrap(), + RawTimestamps::try_new(vec![5_i64, 6].into()).unwrap(), SeriesFields::one(mixed), false, ) @@ -1089,6 +1222,9 @@ fn rejects_invalid_schema_and_metadata_and_accepts_valid_empty() { #[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) @@ -1335,7 +1471,7 @@ fn rejects_malformed_histogram_and_payload_encodings() { &hist_schema, Buffer::from(vec![1_u8]), None, - RawTimeAxis::try_new(vec![1_i64].into()).unwrap(), + RawTimestamps::try_new(vec![1_i64].into()).unwrap(), SeriesFields::one( Samples::try_from_histogram(UInt32Array::from(vec![Some(0)]), histogram_batch()) .unwrap(), @@ -1705,7 +1841,7 @@ fn decoded_key_and_payload_slices_keep_buffers_alive() { &schema, Buffer::from(vec![0xaa_u8, 0x10, 0x20, 0xbb]), None, - RawTimeAxis::try_new(vec![0_i64, 11, 22, 33].into()) + RawTimestamps::try_new(vec![0_i64, 11, 22, 33].into()) .unwrap() .slice(1, 2) .unwrap(), @@ -1776,9 +1912,9 @@ fn decoded_key_and_payload_slices_keep_buffers_alive() { .as_any() .downcast_ref::() .unwrap(); - let crate::series::raw::RawTimeAxis::Explicit(decoded_time) = &decoded[0].time; + let decoded_time = decoded[0].time.as_buffer(); assert_eq!(decoded_time.as_ptr(), time_values.values().as_ptr()); - let crate::series::raw::RawTimeAxis::Explicit(source_time) = &chunk.time; + let source_time = chunk.time.as_buffer(); assert_eq!(decoded_time.as_ptr(), source_time.as_ptr()); drop(owner); drop(chunk); @@ -1790,7 +1926,7 @@ fn decoded_key_and_payload_slices_keep_buffers_alive() { decoded[0].fields.get(0).unwrap().sample(1), SampleRef::Missing )); - assert_eq!(decoded[0].read_key.as_slice(), &[0xaa, 0x10, 0x20, 0xbb]); + assert_eq!(decoded[0].primary_key.as_slice(), &[0xaa, 0x10, 0x20, 0xbb]); } #[test] @@ -1855,7 +1991,7 @@ fn sliced_canonical_mixed_histogram_round_trip_keeps_buffers_alive() { &schema, Buffer::from(vec![0x42_u8]), None, - RawTimeAxis::try_new(vec![0_i64, 11, 22, 33].into()) + RawTimestamps::try_new(vec![0_i64, 11, 22, 33].into()) .unwrap() .slice(1, 2) .unwrap(), @@ -1948,9 +2084,19 @@ fn sliced_canonical_mixed_histogram_round_trip_keeps_buffers_alive() { .values() .as_ptr() ); - let crate::series::raw::RawTimeAxis::Explicit(source_times) = &chunk.time; + let source_times = chunk.time.as_buffer(); - let frames = pack_raw_series_chunk(&schema, RawLabelMode::Key, &chunk).unwrap(); + // 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) @@ -2064,7 +2210,7 @@ fn sliced_canonical_mixed_histogram_round_trip_keeps_buffers_alive() { else { panic!("histogram representation was normalized"); }; - let crate::series::raw::RawTimeAxis::Explicit(decoded_time_values) = &decoded[0].time; + 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 [ @@ -2310,3 +2456,605 @@ fn public_segments_are_flat_ordered_and_retain_empty_dense_kinds() { 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 12e3e17eaac..42b1d4bba0b 100644 --- a/src/datatypes/src/series/samples.rs +++ b/src/datatypes/src/series/samples.rs @@ -389,13 +389,15 @@ impl Samples { SampleStorage::Mixed(mixed) => Some(&mixed.histograms.batch), SampleStorage::Float(_) | SampleStorage::Chunks(_, _) => None, }); + let mut first_data = None; let shared_batch = histogram_batches.next().filter(|first| { histogram_batches.all(|batch| { - Arc::ptr_eq(first, batch) - || first - .as_arrow() - .to_data() - .ptr_eq(&batch.as_arrow().to_data()) + 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| { From eddd595b82888452b3c066678fabc6c1200c74f7 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Fri, 16 Oct 2026 16:38:09 +0800 Subject: [PATCH 10/11] perf(query): batch raw series across stream chunks Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/common/query/src/series_frame.rs | 191 ++- src/common/query/src/series_frame/tests.rs | 1287 ++++++++++++++--- .../query/src/series_frame/tests/lifecycle.rs | 251 +++- 3 files changed, 1485 insertions(+), 244 deletions(-) diff --git a/src/common/query/src/series_frame.rs b/src/common/query/src/series_frame.rs index 081a8be9b28..e2ebc8f03eb 100644 --- a/src/common/query/src/series_frame.rs +++ b/src/common/query/src/series_frame.rs @@ -18,7 +18,9 @@ //! 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. +//! 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; @@ -28,12 +30,12 @@ use common_error::ext::BoxedError; use common_recordbatch::{DfRecordBatch, DfSendableRecordBatchStream}; use datafusion_common::DataFusionError; use datatypes::error::InvalidSeriesSnafu; -use datatypes::series::MetricName; use datatypes::series::frame::{ - RawLabelMode, pack_raw_series_chunk, raw_series_frame_schema, unpack_raw_series_frame, + RawLabelMode, RawSeriesChunkCursor, RawSeriesFrameBuilder, RawSeriesFrameDecoder, }; -use datatypes::series::raw::{RawSeriesChunk, RawSeriesOrdering}; +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}; @@ -42,43 +44,63 @@ mod tests; /// Packs a storage raw-series stream into DataFusion record batches incrementally. /// -/// Producers must keep each read key's chunks contiguous. Validation checks adjacent chunk +/// 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 ordering = source.ordering(); let mode = label_mode(&label_source); - let frame_schema = raw_series_frame_schema(&schema, mode).map_err(codec_error)?; + 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, - frame_schema, mode, + frame_schema, + ordering, validator: SequenceValidator::new(ordering), - pending: Vec::::new().into_iter(), + builder: Some(builder), + cursor: None, })) } /// Unpacks DataFusion batches using the caller's explicit scan binding and ordering. /// -/// Producers must keep each read key's chunks contiguous. Validation checks adjacent chunk -/// conditions but does not certify nonadjacent recurrence for `Contiguous` 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: RawSeriesOrdering, + 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()); - unpack_raw_series_frame(&expected, mode, &empty).map_err(codec_error)?; + decoder.decode(&empty).map_err(codec_error)?; Ok(Box::pin(UnpackStream { source: Some(source), schema: expected, @@ -87,6 +109,7 @@ pub fn unpack_raw_series( ordering, validator: SequenceValidator::new(ordering), pending: Vec::new().into_iter(), + decoder, })) } @@ -97,6 +120,16 @@ fn label_mode(source: &RawLabelSource) -> RawLabelMode { } } +/// 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))) } @@ -114,11 +147,11 @@ fn invalid(reason: impl Into) -> DataFusionError { ))) } -/// Cross-chunk validation state. The producer must keep each key contiguous; this validator -/// checks adjacent chunks only. It retains only the previous complete owned key (not an Arrow -/// buffer), timestamp, metric, and finality. +/// 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: RawSeriesOrdering, + ordering: SeriesOrdering, previous_key: Option>, previous_timestamp: Option, previous_metric: Option, @@ -126,7 +159,7 @@ struct SequenceValidator { } impl SequenceValidator { - fn new(ordering: RawSeriesOrdering) -> Self { + fn new(ordering: SeriesOrdering) -> Self { Self { ordering, previous_key: None, @@ -144,10 +177,12 @@ impl SequenceValidator { } fn validate(&mut self, chunk: &RawSeriesChunk) -> datafusion_common::Result<()> { - let key = chunk.read_key.as_slice(); + let key = chunk.primary_key.as_slice(); let timestamp = chunk .time - .timestamp(chunk.time.len() - 1) + .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) { @@ -172,9 +207,7 @@ impl SequenceValidator { if !self.previous_is_last { return Err(invalid("raw series key changed before final chunk")); } - if self.ordering == RawSeriesOrdering::ReadKeyThenTime - && previous.as_slice() > key - { + if self.ordering == SeriesOrdering::KeyThenTime && previous.as_slice() > key { return Err(invalid("raw series keys are not ordered")); } } @@ -191,17 +224,46 @@ impl SequenceValidator { struct PackStream { source: Option, schema: Arc, - frame_schema: datatypes::arrow::datatypes::SchemaRef, mode: RawLabelMode, + frame_schema: datatypes::arrow::datatypes::SchemaRef, + ordering: SeriesOrdering, validator: SequenceValidator, - pending: std::vec::IntoIter, + 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.pending = Vec::new().into_iter(); - self.validator = SequenceValidator::new(self.validator.ordering); + 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), + } } } @@ -211,36 +273,57 @@ impl Stream for PackStream { fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let this = &mut *self; loop { - if let Some(batch) = this.pending.next() { - return Poll::Ready(Some(Ok(batch))); + 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 => return Poll::Pending, + 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) => { - let result = this.validator.finish(); + if let Err(error) = this.validator.finish() { + this.finish(); + return Poll::Ready(Some(Err(error))); + } + let flushed = this.flush_buffered(); this.finish(); - return Poll::Ready(result.err().map(Err)); + 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))) => { - let result = (|| { - let batches = pack_raw_series_chunk(&this.schema, this.mode, &chunk) - .map_err(codec_error)?; - this.validator.validate(&chunk)?; - Ok(batches) - })(); - match result { - Ok(batches) => this.pending = batches.into_iter(), + 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(error))); + return Poll::Ready(Some(Err(codec_error(error)))); } } } @@ -260,9 +343,10 @@ struct UnpackStream { schema: Arc, frame_schema: datatypes::arrow::datatypes::SchemaRef, label_source: RawLabelSource, - ordering: RawSeriesOrdering, + ordering: SeriesOrdering, validator: SequenceValidator, pending: std::vec::IntoIter, + decoder: RawSeriesFrameDecoder, } impl RawSeriesStream for UnpackStream { @@ -274,7 +358,7 @@ impl RawSeriesStream for UnpackStream { &self.label_source } - fn ordering(&self) -> RawSeriesOrdering { + fn ordering(&self) -> SeriesOrdering { self.ordering } } @@ -328,18 +412,15 @@ impl Stream for UnpackStream { if batch.schema() != this.frame_schema { return Err(invalid("raw frame schema mismatch")); } - unpack_raw_series_frame( - &this.schema, - label_mode(&this.label_source), - &batch, - ) - .map_err(codec_error) - .and_then(|chunks| { - for chunk in &chunks { - this.validator.validate(chunk)?; - } - Ok(chunks) - }) + 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) => { diff --git a/src/common/query/src/series_frame/tests.rs b/src/common/query/src/series_frame/tests.rs index c282a546e7c..85d840ab4f2 100644 --- a/src/common/query/src/series_frame/tests.rs +++ b/src/common/query/src/series_frame/tests.rs @@ -25,7 +25,7 @@ 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}; +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}; @@ -33,8 +33,12 @@ 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::raw::{RawSeriesChunk, RawSeriesOrdering, RawTimeAxis}; -use datatypes::series::samples::{DenseSamplesRef, FieldSpec, Samples, SeriesFields, SeriesSchema}; +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, @@ -75,7 +79,7 @@ fn chunk_with( &schema(), Buffer::from(key.to_vec()), metric, - RawTimeAxis::try_new(ScalarBuffer::from(times.to_vec())).unwrap(), + RawTimestamps::try_new(ScalarBuffer::from(times.to_vec())).unwrap(), SeriesFields::one(Samples::from_float(Float64Array::from(values))), last, ) @@ -96,11 +100,16 @@ struct RawMock { schema: Arc, items: VecDeque>, binding: RawLabelSource, - ordering: RawSeriesOrdering, + ordering: SeriesOrdering, polls: Arc, drops: Arc, - pending: bool, - emitted: bool, + /// 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) { @@ -110,15 +119,19 @@ impl Drop for RawMock { 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 && self.emitted { - self.pending = false; + 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 = true; + self.emitted += 1; } Poll::Ready(item) } @@ -130,7 +143,7 @@ impl RawSeriesStream for RawMock { fn label_source(&self) -> &RawLabelSource { &self.binding } - fn ordering(&self) -> RawSeriesOrdering { + fn ordering(&self) -> SeriesOrdering { self.ordering } } @@ -139,8 +152,13 @@ struct DfMock { items: VecDeque>, polls: Arc, drops: Arc, - pending: bool, - emitted: bool, + /// 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) { @@ -150,15 +168,19 @@ impl Drop for DfMock { 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 && self.emitted { - self.pending = false; + 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 = true; + self.emitted += 1; } Poll::Ready(item) } @@ -170,28 +192,85 @@ impl RecordBatchStream for DfMock { } fn raw( items: Vec, - ordering: RawSeriesOrdering, + ordering: SeriesOrdering, ) -> (SendableRawSeriesStream, Arc, Arc) { - raw_custom(items, ordering, RawLabelSource::Explicit, false) + 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: RawSeriesOrdering, + ordering: SeriesOrdering, binding: RawLabelSource, - pending: bool, +) -> (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: items.into_iter().map(Ok).collect(), + 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, - emitted: false, + pending_after, + pending_reported: false, + panic_on_poll: false, + emitted: 0, }), polls, drops, @@ -200,7 +279,30 @@ fn raw_custom( fn df_source( batches: Vec, declared: SchemaRef, - pending: bool, +) -> ( + 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, @@ -211,11 +313,39 @@ fn df_source( ( Box::pin(DfMock { schema: declared, - items: batches.into_iter().map(Ok).collect(), + items: VecDeque::new(), polls: polls.clone(), drops: drops.clone(), - pending, - emitted: false, + 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, @@ -245,15 +375,17 @@ async fn contiguous_adapter_does_not_certify_nonadjacent_recurrence() { chunk_with(key, &[*time], vec![Some(*value)], Some(metric()), true) }) .collect(); - let (source, _, _) = raw(input, RawSeriesOrdering::Contiguous); + 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, false); + let (source, _, _) = df_source(batches, frame_schema); let mut decoded = unpack_raw_series( source, schema(), RawLabelSource::Explicit, - RawSeriesOrdering::Contiguous, + SeriesOrdering::Contiguous, ) .unwrap(); @@ -284,7 +416,7 @@ fn assert_chunk( values: &[Option], last: bool, ) { - assert_eq!(actual.read_key.as_slice(), key); + assert_eq!(actual.primary_key.as_slice(), key); assert_eq!( (0..actual.time.len()) .map(|i| actual.time.timestamp(i).unwrap()) @@ -324,33 +456,38 @@ async fn roundtrip_explicit_unordered_contiguous_preserves_complete_chunk_conten ), chunk_with(b"complete-a", &[4], vec![Some(9.25)], Some(metric()), true), ]; - let (input, polls, _) = raw(chunks, RawSeriesOrdering::Contiguous); + 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!( - polls.load(Ordering::SeqCst), - 1, - "first frame must precede drain" + first.num_rows(), + 3, + "three dense chunks coalesce into three frame rows, containing four positions" ); - let mut batches = vec![first]; - while let Some(b) = packed.next().await { - batches.push(b.unwrap()); - } + 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: false, - emitted: false, + pending_after: None, + pending_reported: false, + panic_on_poll: false, + emitted: 0, }); let mut output = unpack_raw_series( source, schema(), RawLabelSource::Explicit, - RawSeriesOrdering::Contiguous, + SeriesOrdering::Contiguous, ) .unwrap(); assert_chunk( @@ -378,26 +515,30 @@ async fn roundtrip_explicit_unordered_contiguous_preserves_complete_chunk_conten } #[tokio::test] -async fn read_key_then_time_valid_roundtrip() { +async fn key_then_time_valid_roundtrip() { let (input, _, _) = raw( vec![chunk(b"a", &[1], true), chunk(b"b", &[2], true)], - RawSeriesOrdering::ReadKeyThenTime, + 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: false, - emitted: false, + pending_after: None, + pending_reported: false, + panic_on_poll: false, + emitted: 0, }); let mut out = unpack_raw_series( source, schema(), RawLabelSource::Explicit, - RawSeriesOrdering::ReadKeyThenTime, + SeriesOrdering::KeyThenTime, ) .unwrap(); assert_chunk( @@ -422,23 +563,20 @@ async fn invalid_sequence_reports_valid_prefix_then_one_terminal_error() { struct Case { name: &'static str, items: Vec, - order: RawSeriesOrdering, - expected_prefix_count: usize, + 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: RawSeriesOrdering::Contiguous, - expected_prefix_count: 1, + 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: RawSeriesOrdering::Contiguous, - expected_prefix_count: 1, + order: SeriesOrdering::Contiguous, prefix: vec![(b"a", 1, false, None)], }, Case { @@ -448,29 +586,25 @@ async fn invalid_sequence_reports_valid_prefix_then_one_terminal_error() { chunk(b"b", &[2], true), chunk(b"a", &[3], true), ], - order: RawSeriesOrdering::ReadKeyThenTime, - expected_prefix_count: 2, + 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: RawSeriesOrdering::ReadKeyThenTime, - expected_prefix_count: 1, + 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: RawSeriesOrdering::Contiguous, - expected_prefix_count: 1, + 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: RawSeriesOrdering::Contiguous, - expected_prefix_count: 1, + order: SeriesOrdering::Contiguous, prefix: vec![(b"a", 3, false, None)], }, Case { @@ -485,34 +619,39 @@ async fn invalid_sequence_reports_valid_prefix_then_one_terminal_error() { ), chunk_with(b"a", &[2], vec![Some(2.0)], Some(labeled_metric("y")), true), ], - order: RawSeriesOrdering::Contiguous, - expected_prefix_count: 1, + order: SeriesOrdering::Contiguous, prefix: vec![(b"a", 1, false, Some("x"))], }, ]; for case in cases { - let _case_name = case.name; - assert_eq!(case.prefix.len(), case.expected_prefix_count); - let (input, _, _) = raw(case.items.clone(), case.order); + // 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(); - for (key, time, last, host) in &case.prefix { + 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(), 1); - assert_chunk(&chunks[0], key, &[*time], &[Some(*time as f64)], *last); - assert_metric(&chunks[0], *host); + 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 + StatusCode::InvalidArguments, + "{}", + case.name ); - assert!(s.next().await.is_none()); + 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. @@ -528,7 +667,7 @@ async fn invalid_sequence_reports_valid_prefix_then_one_terminal_error() { ); } let frame_schema = frames[0].schema(); - let (source, _, _) = df_source(frames, frame_schema, false); + 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 { @@ -543,14 +682,8 @@ async fn invalid_sequence_reports_valid_prefix_then_one_terminal_error() { ); assert!(decoded.next().await.is_none()); } - let (input, _, _) = raw( - vec![chunk(b"unfinished", &[7], false)], - RawSeriesOrdering::Contiguous, - ); - let mut s = pack_raw_series(input).unwrap(); - assert!(s.next().await.unwrap().is_ok()); - assert!(s.next().await.unwrap().is_err()); - assert!(s.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, @@ -558,12 +691,12 @@ async fn invalid_sequence_reports_valid_prefix_then_one_terminal_error() { ) .unwrap(); let frame_schema = frame[0].schema(); - let (source, _, _) = df_source(frame, frame_schema, false); + let (source, _, _) = df_source(frame, frame_schema); let mut decoded = unpack_raw_series( source, schema(), RawLabelSource::Explicit, - RawSeriesOrdering::Contiguous, + SeriesOrdering::Contiguous, ) .unwrap(); let first = decoded.next().await.unwrap().unwrap(); @@ -575,6 +708,36 @@ async fn invalid_sequence_reports_valid_prefix_then_one_terminal_error() { 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"), @@ -621,15 +784,14 @@ async fn primary_key_binding_is_retained_exactly() { &schema(), Buffer::from(b"encoded-complete-key".to_vec()), None, - RawTimeAxis::try_new(ScalarBuffer::from(vec![1])).unwrap(), + RawTimestamps::try_new(ScalarBuffer::from(vec![1])).unwrap(), SeriesFields::one(Samples::from_float(Float64Array::from(vec![Some(3.0)]))), true, ) .unwrap(), ], - RawSeriesOrdering::Contiguous, + SeriesOrdering::Contiguous, binding.clone(), - false, ); let mut packed = pack_raw_series(input).unwrap(); let batch = packed.next().await.unwrap().unwrap(); @@ -638,13 +800,14 @@ async fn primary_key_binding_is_retained_exactly() { items: VecDeque::from([Ok(batch)]), polls: Arc::new(AtomicUsize::new(0)), drops: Arc::new(AtomicUsize::new(0)), - pending: false, - emitted: false, + pending_after: None, + pending_reported: false, + panic_on_poll: false, + emitted: 0, }); - let mut out = - unpack_raw_series(source, schema(), binding, RawSeriesOrdering::Contiguous).unwrap(); + let mut out = unpack_raw_series(source, schema(), binding, SeriesOrdering::Contiguous).unwrap(); let decoded = out.next().await.unwrap().unwrap(); - assert_eq!(decoded.read_key.as_slice(), b"encoded-complete-key"); + assert_eq!(decoded.primary_key.as_slice(), b"encoded-complete-key"); assert!(decoded.metric.is_none()); let retained = out.label_source(); match retained { @@ -671,13 +834,13 @@ async fn explicit_mode_rejects_key_payload_and_primary_mode_rejects_explicit_met &schema(), Buffer::from(b"a".to_vec()), None, - RawTimeAxis::try_new(ScalarBuffer::from(vec![1])).unwrap(), + RawTimestamps::try_new(ScalarBuffer::from(vec![1])).unwrap(), SeriesFields::one(Samples::from_float(Float64Array::from(vec![Some(1.0)]))), true, ) .unwrap(), ], - RawSeriesOrdering::Contiguous, + SeriesOrdering::Contiguous, ); let mut s = pack_raw_series(input).unwrap(); assert!(s.next().await.unwrap().is_err()); @@ -714,9 +877,8 @@ async fn primary_key_mode_rejects_explicit_metric_payload() { let metadata = Arc::new(builder.build().unwrap()); let (input, _, _) = raw_custom( vec![chunk(b"a", &[1], true)], - RawSeriesOrdering::Contiguous, + SeriesOrdering::Contiguous, primary_binding(metadata), - false, ); let mut stream = pack_raw_series(input).unwrap(); assert_eq!( @@ -731,7 +893,9 @@ async fn primary_key_mode_rejects_explicit_metric_payload() { #[tokio::test] async fn explicit_metric_must_remain_stable_within_key() { - let (input, _, _) = raw( + // 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", @@ -742,11 +906,30 @@ async fn explicit_metric_must_remain_stable_within_key() { ), chunk_with(b"a", &[2], vec![Some(2.0)], Some(labeled_metric("y")), true), ], - RawSeriesOrdering::Contiguous, + SeriesOrdering::Contiguous, + 1, ); let mut s = pack_raw_series(input).unwrap(); - assert!(s.next().await.unwrap().is_ok()); - assert!(s.next().await.unwrap().is_err()); + 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()); } @@ -762,11 +945,13 @@ async fn locally_invalid_public_raw_chunk_is_rejected() { schema: schema(), items: VecDeque::from([Ok(bad)]), binding: RawLabelSource::Explicit, - ordering: RawSeriesOrdering::Contiguous, + ordering: SeriesOrdering::Contiguous, polls, drops: Arc::new(AtomicUsize::new(0)), - pending: false, - emitted: false, + 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(); @@ -779,20 +964,28 @@ async fn locally_invalid_public_raw_chunk_is_rejected() { #[tokio::test] async fn empty_streams_keep_self_describing_schema() { - let (input, _, _) = raw(vec![], RawSeriesOrdering::Contiguous); + 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()); - let (source, polls, _) = df_source(vec![], frame_schema.clone(), false); + // 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, - RawSeriesOrdering::Contiguous, + 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()); } @@ -806,11 +999,13 @@ async fn upstream_error_status_is_preserved_once_and_releases_source() { schema: schema(), items: VecDeque::from([Err(e)]), binding: RawLabelSource::Explicit, - ordering: RawSeriesOrdering::Contiguous, - polls, + ordering: SeriesOrdering::Contiguous, + polls: polls.clone(), drops: drops.clone(), - pending: false, - emitted: false, + 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(); @@ -819,6 +1014,7 @@ async fn upstream_error_status_is_preserved_once_and_releases_source() { 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(), @@ -835,14 +1031,16 @@ async fn upstream_error_status_is_preserved_once_and_releases_source() { )))]), polls: polls.clone(), drops: drops.clone(), - pending: false, - emitted: false, + pending_after: None, + pending_reported: false, + panic_on_poll: false, + emitted: 0, }); let mut unpacked = unpack_raw_series( source, schema(), RawLabelSource::Explicit, - RawSeriesOrdering::Contiguous, + SeriesOrdering::Contiguous, ) .unwrap(); let error = unpacked.next().await.unwrap().unwrap_err(); @@ -855,24 +1053,40 @@ async fn upstream_error_status_is_preserved_once_and_releases_source() { #[tokio::test] async fn dropping_adapter_releases_upstream_and_pending_item_is_resumed() { - let (input, polls, drops) = raw_custom( + let (input, polls, drops) = raw_pending( vec![chunk(b"a", &[1], true), chunk(b"b", &[2], true)], - RawSeriesOrdering::Contiguous, - RawLabelSource::Explicit, - true, + SeriesOrdering::Contiguous, + 1, ); let mut packed = pack_raw_series(input).unwrap(); - assert!(packed.next().await.unwrap().is_ok()); + let first = packed.next().await.unwrap().unwrap(); + assert_eq!(first.num_rows(), 1); assert_eq!( polls.load(Ordering::SeqCst), - 1, - "first item is emitted immediately" + 2, + "the buffered row is flushed at Pending before the next item" ); - assert!( - packed.next().await.unwrap().is_ok(), - "Pending is resumed 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" ); - assert_eq!(polls.load(Ordering::SeqCst), 3); + 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); } @@ -880,13 +1094,13 @@ async fn dropping_adapter_releases_upstream_and_pending_item_is_resumed() { #[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, false); + let (source, polls, drops) = df_source(vec![], wrong); assert!( unpack_raw_series( source, schema(), RawLabelSource::Explicit, - RawSeriesOrdering::Contiguous + SeriesOrdering::Contiguous ) .is_err() ); @@ -908,11 +1122,13 @@ async fn invalid_grid_schema_rejected_for_empty_raw_stream() { schema: Arc::new(invalid), items: VecDeque::new(), binding: RawLabelSource::Explicit, - ordering: RawSeriesOrdering::Contiguous, + ordering: SeriesOrdering::Contiguous, polls: polls.clone(), drops: Arc::new(AtomicUsize::new(0)), - pending: false, - emitted: false, + 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); @@ -931,19 +1147,16 @@ async fn unpack_batch_sequence_error_is_atomic_for_all_rows() { (b"b".as_slice(), 2), (b"a".as_slice(), 3), ] { - let (input, _, _) = raw( - vec![chunk(key, &[time], true)], - RawSeriesOrdering::ReadKeyThenTime, - ); + 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, false); + let (source, _, _) = df_source(vec![multi], frame_schema); let mut unpacked = unpack_raw_series( source, schema(), RawLabelSource::Explicit, - RawSeriesOrdering::ReadKeyThenTime, + SeriesOrdering::KeyThenTime, ) .unwrap(); assert!(unpacked.next().await.unwrap().is_err()); @@ -1001,12 +1214,12 @@ async fn unpack_rejects_late_invalid_frame_row_before_prefix() { ); 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, false); + let (source, _, _) = df_source(vec![multi], frame_schema); let mut unpacked = unpack_raw_series( source, schema(), RawLabelSource::Explicit, - RawSeriesOrdering::Contiguous, + SeriesOrdering::Contiguous, ) .unwrap(); let error = unpacked.next().await.unwrap().unwrap_err(); @@ -1049,12 +1262,12 @@ async fn frame_application_metadata_is_stable_and_schema_change_rejected() { .remove(0) .with_schema(declared.clone()) .unwrap(); - let (source, _, _) = df_source(vec![first.clone(), stable_second], declared.clone(), false); + let (source, _, _) = df_source(vec![first.clone(), stable_second], declared.clone()); let mut stable = unpack_raw_series( source, schema(), RawLabelSource::Explicit, - RawSeriesOrdering::Contiguous, + SeriesOrdering::Contiguous, ) .unwrap(); assert_chunk( @@ -1090,12 +1303,12 @@ async fn frame_application_metadata_is_stable_and_schema_change_rejected() { .remove(0) .with_schema(changed_schema) .unwrap(); - let (source, _, _) = df_source(vec![first.clone(), bad_batch], declared.clone(), false); + let (source, _, _) = df_source(vec![first.clone(), bad_batch], declared.clone()); let mut unpacked = unpack_raw_series( source, schema(), RawLabelSource::Explicit, - RawSeriesOrdering::Contiguous, + SeriesOrdering::Contiguous, ) .unwrap(); assert_chunk( @@ -1133,24 +1346,31 @@ fn poll_unpacked( 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)], - RawSeriesOrdering::Contiguous, + 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, false); + let (source, polls, _) = df_source(batches, frame_schema); let mut unpacked = unpack_raw_series( source, schema(), RawLabelSource::Explicit, - RawSeriesOrdering::Contiguous, + SeriesOrdering::Contiguous, ) .unwrap(); let wake = Arc::new(CountingWake(AtomicUsize::new(0))); @@ -1205,12 +1425,12 @@ fn unpack_yields_at_empty_batch_budget_before_observing_eof() { ) .unwrap(); let batches = vec![RecordBatch::new_empty(frame_schema.clone()); BUDGET]; - let (source, polls, _) = df_source(batches, frame_schema, false); + let (source, polls, _) = df_source(batches, frame_schema); let mut unpacked = unpack_raw_series( source, schema(), RawLabelSource::Explicit, - RawSeriesOrdering::Contiguous, + SeriesOrdering::Contiguous, ) .unwrap(); let wake = Arc::new(CountingWake(AtomicUsize::new(0))); @@ -1236,16 +1456,16 @@ fn unpack_propagates_upstream_pending_without_extra_wake() { datatypes::series::frame::RawLabelMode::Explicit, ) .unwrap(); - let (source, polls, _) = df_source( + let (source, polls, _) = df_pending( vec![RecordBatch::new_empty(frame_schema.clone())], frame_schema, - true, + 1, ); let mut unpacked = unpack_raw_series( source, schema(), RawLabelSource::Explicit, - RawSeriesOrdering::Contiguous, + SeriesOrdering::Contiguous, ) .unwrap(); let wake = Arc::new(CountingWake(AtomicUsize::new(0))); @@ -1271,12 +1491,12 @@ fn unpack_reports_bad_batch_at_budget_boundary_before_yielding() { batches.push(RecordBatch::new_empty( datatypes::arrow::datatypes::Schema::empty().into(), )); - let (source, polls, _) = df_source(batches, frame_schema, false); + let (source, polls, _) = df_source(batches, frame_schema); let mut unpacked = unpack_raw_series( source, schema(), RawLabelSource::Explicit, - RawSeriesOrdering::Contiguous, + SeriesOrdering::Contiguous, ) .unwrap(); let wake = Arc::new(CountingWake(AtomicUsize::new(0))); @@ -1315,14 +1535,16 @@ fn unpack_surfaces_upstream_error_after_budget_yield_once() { items: items.into(), polls: polls.clone(), drops: Arc::new(AtomicUsize::new(0)), - pending: false, - emitted: false, + pending_after: None, + pending_reported: false, + panic_on_poll: false, + emitted: 0, }); let mut unpacked = unpack_raw_series( source, schema(), RawLabelSource::Explicit, - RawSeriesOrdering::Contiguous, + SeriesOrdering::Contiguous, ) .unwrap(); let wake = Arc::new(CountingWake(AtomicUsize::new(0))); @@ -1350,17 +1572,17 @@ fn unpack_surfaces_upstream_error_after_budget_yield_once() { async fn unpack_drains_decoded_rows_before_polling_another_batch() { let (input, _, _) = raw( vec![chunk(b"a", &[1], true), chunk(b"b", &[2], true)], - RawSeriesOrdering::ReadKeyThenTime, + 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, false); + let (source, polls, _) = df_source(vec![multi], frame_schema); let mut unpacked = unpack_raw_series( source, schema(), RawLabelSource::Explicit, - RawSeriesOrdering::ReadKeyThenTime, + SeriesOrdering::KeyThenTime, ) .unwrap(); assert_chunk( @@ -1409,13 +1631,13 @@ async fn unpack_preflights_empty_and_grid_schemas_without_polling() { datatypes::series::frame::RawLabelMode::Explicit, ) .unwrap(); - let (source, polls, drops) = df_source(vec![], frame_schema, false); + let (source, polls, drops) = df_source(vec![], frame_schema); assert!( unpack_raw_series( source, Arc::new(expected), RawLabelSource::Explicit, - RawSeriesOrdering::Contiguous, + SeriesOrdering::Contiguous, ) .is_err() ); @@ -1441,7 +1663,7 @@ fn named_fields_chunk(schema: &SeriesSchema) -> RawSeriesChunk { schema, Buffer::from(b"series-key".to_vec()), Some(metric()), - RawTimeAxis::try_new(ScalarBuffer::from(vec![7_i64, 8])).unwrap(), + 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)])), @@ -1460,11 +1682,13 @@ async fn named_frame_fields_roundtrip_and_reject_wrong_bindings_before_polling() schema: expected.clone(), items: VecDeque::from([Ok(input_chunk)]), binding: RawLabelSource::Explicit, - ordering: RawSeriesOrdering::ReadKeyThenTime, + ordering: SeriesOrdering::KeyThenTime, polls: Arc::new(AtomicUsize::new(0)), drops: Arc::new(AtomicUsize::new(0)), - pending: false, - emitted: false, + pending_after: None, + pending_reported: false, + panic_on_poll: false, + emitted: 0, }); let batches = pack_all(input).await; let mut ipc = Vec::new(); @@ -1490,12 +1714,12 @@ async fn named_frame_fields_roundtrip_and_reject_wrong_bindings_before_polling() ["cpu_usage", "cpu_idle"] ); - let (source, _, _) = df_source(ipc_batches.clone(), frame_schema.clone(), false); + let (source, _, _) = df_source(ipc_batches.clone(), frame_schema.clone()); let mut decoded = unpack_raw_series( source, expected.clone(), RawLabelSource::Explicit, - RawSeriesOrdering::ReadKeyThenTime, + SeriesOrdering::KeyThenTime, ) .unwrap(); let actual = decoded.next().await.unwrap().unwrap(); @@ -1524,13 +1748,13 @@ async fn named_frame_fields_roundtrip_and_reject_wrong_bindings_before_polling() ] { 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(), false); + let (source, polls, drops) = df_source(batches, frame_schema.clone()); assert!( unpack_raw_series( source, wrong_schema, RawLabelSource::Explicit, - RawSeriesOrdering::ReadKeyThenTime, + SeriesOrdering::KeyThenTime, ) .is_err() ); @@ -1538,3 +1762,734 @@ async fn named_frame_fields_roundtrip_and_reject_wrong_bindings_before_polling() 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 index 0e467730403..b8558c89b03 100644 --- a/src/common/query/src/series_frame/tests/lifecycle.rs +++ b/src/common/query/src/series_frame/tests/lifecycle.rs @@ -19,8 +19,9 @@ 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, RawSeriesOrdering, RawTimeAxis}; +use datatypes::series::raw::{RawSeriesChunk, RawTimestamps}; use datatypes::series::samples::{ DenseSamplesRef, FieldSpec, HistogramSamples, SampleRef, Samples, SeriesFields, SeriesSchema, }; @@ -53,8 +54,8 @@ impl RawSeriesStream for RawFixture { &self.binding } - fn ordering(&self) -> RawSeriesOrdering { - RawSeriesOrdering::Contiguous + fn ordering(&self) -> SeriesOrdering { + SeriesOrdering::Contiguous } } @@ -128,7 +129,7 @@ async fn streaming_round_trip_preserves_sliced_multifield_payload_and_owners() { eval_grid: None, }); let time_owner = ScalarBuffer::from(vec![0, 11, 22, 33, 44]); - let times = RawTimeAxis::try_new(time_owner.slice(1, 3)).unwrap(); + let times = RawTimestamps::try_new(time_owner.slice(1, 3)).unwrap(); let float_owner = Float64Array::from(vec![ Some(90.0), None, @@ -180,7 +181,7 @@ async fn streaming_round_trip_preserves_sliced_multifield_payload_and_owners() { true, ) .unwrap(); - let input_times_ptr = chunk.time.as_explicit().unwrap().as_ptr(); + 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(); @@ -196,7 +197,7 @@ async fn streaming_round_trip_preserves_sliced_multifield_payload_and_owners() { frames, schema, RawLabelSource::Explicit, - RawSeriesOrdering::Contiguous, + SeriesOrdering::Contiguous, ) .unwrap(); let mut output = Vec::new(); @@ -206,55 +207,86 @@ async fn streaming_round_trip_preserves_sliced_multifield_payload_and_owners() { assert_eq!( output.len(), 2, - "the float field's segment boundary yields two frames" + "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" ); - // Validate buffer sharing while the source owners are still alive. - let DenseSamplesRef::Float(decoded_float) = output[0].fields.get(0).unwrap().segments()[0] - .dense() - .unwrap() + // 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_eq!(decoded_float.values().as_ptr(), input_float_ptr); + 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_eq!( + assert_ne!( decoded_hist.indices().values().as_ptr(), - input_hist_indices_ptr + input_hist_indices_ptr, + "multi-row concatenation copies histogram indices" ); assert!( - decoded_hist + !decoded_hist .batch() .as_arrow() .to_data() - .ptr_eq(&sliced_table.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_eq!( + assert_ne!( decoded_mixed.floats().values().as_ptr(), input_mixed_float_ptr ); - assert_eq!( + assert_ne!( decoded_mixed.histograms().indices().values().as_ptr(), input_mixed_indices_ptr ); assert!( - decoded_mixed + !decoded_mixed .histograms() .batch() .as_arrow() .to_data() - .ptr_eq(&sliced_table.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_explicit().unwrap(); - assert_eq!(first_times.as_ptr(), input_times_ptr); + 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); @@ -277,9 +309,10 @@ async fn streaming_round_trip_preserves_sliced_multifield_payload_and_owners() { let mut histogram_seen = Vec::new(); let mut mixed_seen = Vec::new(); for chunk in &output { - assert_eq!(chunk.read_key.as_slice(), [0xaa, 0x10, 0xbb]); + 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(); @@ -318,3 +351,175 @@ async fn streaming_round_trip_preserves_sliced_multifield_payload_and_owners() { 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"), + } +} From ba88f817cdcffeae6594943c55aaced2ffa769e5 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Fri, 16 Oct 2026 17:44:58 +0800 Subject: [PATCH 11/11] fix(datatypes): keep frame flush guard clippy clean Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/datatypes/src/series/frame.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/datatypes/src/series/frame.rs b/src/datatypes/src/series/frame.rs index 27902ecf805..da6ff38d1ae 100644 --- a/src/datatypes/src/series/frame.rs +++ b/src/datatypes/src/series/frame.rs @@ -303,10 +303,10 @@ impl RawSeriesChunkCursor { return Ok(Some(batch)); } builder.push(row)?; - if builder.is_full() { - if let Some(batch) = builder.flush()? { - return Ok(Some(batch)); - } + if builder.is_full() + && let Some(batch) = builder.flush()? + { + return Ok(Some(batch)); } } }