feat(common-query): declare raw series frame ordering in schema metadata

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

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<ArrowSchemaRef> {
+    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<SeriesOrdering> {
+    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<Option<datafusion_common::Result<DfRecordBatch>>> {
+        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)))),
                     };

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<SeriesSchema> {
     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<String, String>,
+        ordering: Option<String>,
+    }
+
+    let cases: Vec<Case> = 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::<Vec<_>>();
+            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());
+            }
+        }
+    }
+}

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.

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
+  }
+]