diff --git a/src/common/query/src/series_frame.rs b/src/common/query/src/series_frame.rs index e2ebc8f03eb..847e18499b5 100644 --- a/src/common/query/src/series_frame.rs +++ b/src/common/query/src/series_frame.rs @@ -29,9 +29,11 @@ use std::task::{Context, Poll}; use common_error::ext::BoxedError; use common_recordbatch::{DfRecordBatch, DfSendableRecordBatchStream}; use datafusion_common::DataFusionError; +use datatypes::arrow::datatypes::{Schema as ArrowSchema, SchemaRef as ArrowSchemaRef}; use datatypes::error::InvalidSeriesSnafu; use datatypes::series::frame::{ RawLabelMode, RawSeriesChunkCursor, RawSeriesFrameBuilder, RawSeriesFrameDecoder, + raw_series_frame_schema, }; use datatypes::series::raw::RawSeriesChunk; use datatypes::series::samples::SeriesSchema; @@ -42,6 +44,48 @@ use store_api::region_engine::{RawLabelSource, RawSeriesStream, SendableRawSerie #[cfg(test)] mod tests; +const ORDERING_KEY: &str = "greptime.series.ordering"; + +/// Builds the canonical raw frame schema with the producer's guaranteed ordering. +/// +/// Ordering records a producer guarantee; it is not inferred from frame contents. `Interleaved` +/// has no v1 wire representation and is rejected before the schema is assembled. +pub fn raw_series_frame_schema_with_ordering( + schema: &SeriesSchema, + mode: RawLabelMode, + ordering: SeriesOrdering, +) -> datafusion_common::Result { + let ordering = match ordering { + SeriesOrdering::Contiguous => "contiguous", + SeriesOrdering::KeyThenTime => "read_key_then_time", + SeriesOrdering::Interleaved => { + return Err(invalid( + "interleaved raw series ordering has no wire representation", + )); + } + }; + let schema = raw_series_frame_schema(schema, mode).map_err(codec_error)?; + let mut metadata = schema.metadata().clone(); + metadata.insert(ORDERING_KEY.to_string(), ordering.to_string()); + Ok(Arc::new(ArrowSchema::new_with_metadata( + schema.fields().clone(), + metadata, + ))) +} + +/// Reads the declared producer ordering, conservatively treating legacy frames as contiguous. +/// +/// Only the two orderings representable on the wire can be returned. +pub fn raw_series_frame_ordering( + schema: &ArrowSchema, +) -> datafusion_common::Result { + match schema.metadata().get(ORDERING_KEY).map(String::as_str) { + None | Some("contiguous") => Ok(SeriesOrdering::Contiguous), + Some("read_key_then_time") => Ok(SeriesOrdering::KeyThenTime), + Some(_) => Err(invalid("unknown raw series ordering")), + } +} + /// Packs a storage raw-series stream into DataFusion record batches incrementally. /// /// Only [`SeriesOrdering::Contiguous`] and [`SeriesOrdering::KeyThenTime`] streams are @@ -59,6 +103,9 @@ pub fn pack_raw_series( let schema = source.schema().clone(); let label_source = source.label_source().clone(); let mode = label_mode(&label_source); + // The declared frame schema carries the producer ordering; the shared builder's schema + // intentionally omits it, so every ready batch is declared before it is yielded. + let frame_schema = raw_series_frame_schema_with_ordering(&schema, mode, ordering)?; let builder = RawSeriesFrameBuilder::try_new( &schema, mode, @@ -66,7 +113,6 @@ pub fn pack_raw_series( RawSeriesFrameBuilder::DEFAULT_MAX_BYTES, ) .map_err(codec_error)?; - let frame_schema = builder.schema().clone(); Ok(Box::pin(PackStream { source: Some(source), schema, @@ -265,6 +311,22 @@ impl PackStream { None => Ok(None), } } + + /// Publishes one ready builder batch under the declared frame schema, which is the shared + /// builder's schema plus the producer's ordering metadata. A batch whose fields cannot be + /// declared releases the stream instead of yielding rows with an undeclared layout. + fn declare_batch( + &mut self, + batch: DfRecordBatch, + ) -> Poll>> { + match batch.with_schema(self.frame_schema.clone()) { + Ok(batch) => Poll::Ready(Some(Ok(batch))), + Err(error) => { + self.finish(); + Poll::Ready(Some(Err(DataFusionError::from(error)))) + } + } + } } impl Stream for PackStream { @@ -274,7 +336,7 @@ impl Stream for PackStream { let this = &mut *self; loop { match this.advance_cursor() { - Ok(Some(batch)) => return Poll::Ready(Some(Ok(batch))), + Ok(Some(batch)) => return this.declare_batch(batch), Ok(None) => {} Err(error) => { this.finish(); @@ -289,7 +351,7 @@ impl Stream for PackStream { Poll::Pending => { let flushed = this.flush_buffered(); return match flushed { - Ok(Some(batch)) => Poll::Ready(Some(Ok(batch))), + Ok(Some(batch)) => this.declare_batch(batch), Ok(None) => Poll::Pending, Err(error) => { this.finish(); @@ -305,7 +367,7 @@ impl Stream for PackStream { let flushed = this.flush_buffered(); this.finish(); return match flushed { - Ok(Some(batch)) => Poll::Ready(Some(Ok(batch))), + Ok(Some(batch)) => this.declare_batch(batch), Ok(None) => Poll::Ready(None), Err(error) => Poll::Ready(Some(Err(codec_error(error)))), }; diff --git a/src/common/query/src/series_frame/tests.rs b/src/common/query/src/series_frame/tests.rs index 85d840ab4f2..feabbe9d019 100644 --- a/src/common/query/src/series_frame/tests.rs +++ b/src/common/query/src/series_frame/tests.rs @@ -14,7 +14,7 @@ mod lifecycle; -use std::collections::VecDeque; +use std::collections::{HashMap, VecDeque}; use std::io::Cursor; use std::pin::Pin; use std::sync::Arc; @@ -34,6 +34,7 @@ use datatypes::arrow::record_batch::RecordBatch; use datatypes::series::Label; use datatypes::series::MetricName; use datatypes::series::SeriesOrdering; +use datatypes::series::frame::{RawLabelMode, raw_series_frame_schema}; use datatypes::series::histogram::{BucketSideView, CounterResetHint, HistogramBatchBuilder}; use datatypes::series::raw::{RawSeriesChunk, RawTimestamps}; use datatypes::series::samples::{ @@ -44,7 +45,7 @@ use store_api::region_engine::{ BoundTag, RawLabelSource, RawSeriesStream, SendableRawSeriesStream, }; -use super::{pack_raw_series, unpack_raw_series}; +use super::{pack_raw_series, raw_series_frame_ordering, unpack_raw_series}; fn schema() -> Arc { Arc::new(SeriesSchema { @@ -2383,9 +2384,20 @@ async fn public_chunk_with_empty_times_is_rejected_without_panic() { } /// `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. +/// an empty unpack stream; the wire-schema helper rejects it before assembling a schema, since no +/// v1 metadata value represents it. No adjacent-only fallback is installed. #[tokio::test] async fn interleaved_ordering_is_rejected_before_any_poll() { + assert!( + super::raw_series_frame_schema_with_ordering( + &schema(), + datatypes::series::frame::RawLabelMode::Explicit, + SeriesOrdering::Interleaved, + ) + .is_err(), + "interleaved cannot be declared on the wire schema" + ); + let (source, polls, drops) = raw_never_polled(SeriesOrdering::Interleaved); let error = match pack_raw_series(source) { Ok(_) => panic!("interleaved pack stream must be rejected"), @@ -2493,3 +2505,126 @@ async fn unpack_rejects_unrelated_metadata_drift_and_empty_schema_changes() { assert!(unpacked.next().await.is_none(), "{label}"); } } + +#[test] +fn raw_series_ordering_metadata_fixture() { + #[derive(serde::Deserialize)] + struct Case { + name: String, + metadata: HashMap, + ordering: Option, + } + + let cases: Vec = serde_json::from_str(include_str!( + "../../../../../tests/compatibility/wire/raw_series_ordering/v1_metadata.json" + )) + .unwrap(); + for case in cases { + let literal_ordering = match case.ordering.as_deref() { + Some("contiguous") => Some(SeriesOrdering::Contiguous), + Some("read_key_then_time") => Some(SeriesOrdering::KeyThenTime), + Some(value) => panic!("{} has invalid fixture ordering {value}", case.name), + None => None, + }; + let frame_schema = raw_series_frame_schema(&schema(), RawLabelMode::Explicit).unwrap(); + let frame_schema = frame_schema.as_ref().clone().with_metadata(case.metadata); + match literal_ordering { + Some(expected) => assert_eq!( + raw_series_frame_ordering(&frame_schema) + .unwrap_or_else(|error| panic!("{}: {error}", case.name)), + expected, + "{}", + case.name + ), + None => assert!( + raw_series_frame_ordering(&frame_schema).is_err(), + "{}", + case.name + ), + } + } +} + +#[tokio::test] +async fn pack_declares_ordering_on_every_batch_and_ipc_roundtrips_it() { + for ordering in [SeriesOrdering::Contiguous, SeriesOrdering::KeyThenTime] { + for chunks in [ + vec![ + chunk_with(b"a", &[1], vec![Some(1.0)], Some(metric()), false), + chunk_with(b"a", &[2], vec![None], Some(metric()), true), + chunk_with(b"b", &[3], vec![Some(3.0)], Some(metric()), true), + ], + vec![], + ] { + let is_empty = chunks.is_empty(); + // The bounded multi-row builder coalesces these three rows into one batch; the + // replaced single-row packing emitted one batch per chunk. + let expected_batch_count = if is_empty { 0 } else { 1 }; + let canonical_schema = + raw_series_frame_schema(&schema(), RawLabelMode::Explicit).unwrap(); + let mut expected_metadata = canonical_schema.metadata().clone(); + expected_metadata.insert( + "greptime.series.ordering".to_string(), + match ordering { + SeriesOrdering::Contiguous => "contiguous", + SeriesOrdering::KeyThenTime => "read_key_then_time", + // The v1 wire carries exactly these two spellings, so the loop above + // admits only the two representable modes; `Interleaved` is rejected + // before any poll (see `interleaved_ordering_is_rejected_before_any_poll`). + SeriesOrdering::Interleaved => { + unreachable!("interleaved raw series ordering has no wire spelling") + } + } + .to_string(), + ); + let (input, _, _) = raw(chunks, ordering); + let mut packed = pack_raw_series(input).unwrap(); + let declared = packed.schema(); + assert_eq!(declared.fields(), canonical_schema.fields()); + assert_eq!(declared.metadata(), &expected_metadata); + assert_eq!(raw_series_frame_ordering(&declared).unwrap(), ordering); + let mut batches = Vec::new(); + while let Some(batch) = packed.next().await { + let batch = batch.unwrap(); + assert_eq!(batch.schema(), declared); + batches.push(batch); + } + + let mut ipc = Vec::new(); + { + let mut writer = StreamWriter::try_new(&mut ipc, &declared).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(), declared); + assert_eq!( + raw_series_frame_ordering(&reader.schema()).unwrap(), + ordering + ); + let ipc_batches = reader.map(Result::unwrap).collect::>(); + assert_eq!(ipc_batches.len(), batches.len()); + assert_eq!(ipc_batches.len(), expected_batch_count); + assert!(ipc_batches.iter().all(|batch| batch.schema() == declared)); + let (source, _, _) = df_source(ipc_batches, declared); + let mut decoded = + unpack_raw_series(source, schema(), RawLabelSource::Explicit, ordering).unwrap(); + if is_empty { + assert!(decoded.next().await.is_none()); + } else { + for (key, time, value, last) in [ + (b"a".as_slice(), 1, Some(1.0), false), + (b"a".as_slice(), 2, None, true), + (b"b".as_slice(), 3, Some(3.0), true), + ] { + let actual = decoded.next().await.unwrap().unwrap(); + assert_chunk(&actual, key, &[time], &[value], last); + assert_metric(&actual, None); + } + assert!(decoded.next().await.is_none()); + } + } + } +} diff --git a/tests/compatibility/wire/raw_series_ordering/README.md b/tests/compatibility/wire/raw_series_ordering/README.md new file mode 100644 index 00000000000..630dc25e2a2 --- /dev/null +++ b/tests/compatibility/wire/raw_series_ordering/README.md @@ -0,0 +1,19 @@ +# Raw-series ordering schema metadata compatibility + +`v1_metadata.json` is a hand-authored, frozen set of Arrow schema metadata maps +for the v1 raw-series frame. It records the established layout/content/labels/ +time metadata and the new optional ordering entry. The Rust unit test parses +these literal maps through `raw_series_frame_ordering`; the fixture is not +produced by the schema builder. + +Run the focused check: + +```shell +cargo nextest run -p common-query -E 'test(raw_series_ordering_metadata_fixture)' +``` + +The fixture proves that absent ordering metadata defaults to `Contiguous`, the +two known values decode, and an unknown value is rejected. It does not prove +interoperability with an old binary, nor does it test an end-to-end client or +service. Separate pack/IPC tests check the declared ordering schema and batch +schema agreement for both orderings, including empty streams. diff --git a/tests/compatibility/wire/raw_series_ordering/v1_metadata.json b/tests/compatibility/wire/raw_series_ordering/v1_metadata.json new file mode 100644 index 00000000000..c5b3c0ea69b --- /dev/null +++ b/tests/compatibility/wire/raw_series_ordering/v1_metadata.json @@ -0,0 +1,45 @@ +[ + { + "name": "missing_defaults_to_contiguous", + "metadata": { + "greptime.series.layout": "v1", + "greptime.series.content": "raw", + "greptime.series.labels": "explicit", + "greptime.series.time": "timestamps" + }, + "ordering": "contiguous" + }, + { + "name": "known_contiguous", + "metadata": { + "greptime.series.layout": "v1", + "greptime.series.content": "raw", + "greptime.series.labels": "explicit", + "greptime.series.time": "timestamps", + "greptime.series.ordering": "contiguous" + }, + "ordering": "contiguous" + }, + { + "name": "known_read_key_then_time", + "metadata": { + "greptime.series.layout": "v1", + "greptime.series.content": "raw", + "greptime.series.labels": "explicit", + "greptime.series.time": "timestamps", + "greptime.series.ordering": "read_key_then_time" + }, + "ordering": "read_key_then_time" + }, + { + "name": "unknown_ordering_rejected", + "metadata": { + "greptime.series.layout": "v1", + "greptime.series.content": "raw", + "greptime.series.labels": "explicit", + "greptime.series.time": "timestamps", + "greptime.series.ordering": "future_value" + }, + "ordering": null + } +]