fix(mito2): share raw key buffers and validate legacy key order

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

src/mito2/src/read/series_scan.rs

@@ -22,27 +22,34 @@ use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::{Arc, Mutex};
 use std::time::{Duration, Instant};
 
+use api::v1::SemanticType;
 use async_stream::try_stream;
 use common_error::ext::BoxedError;
+use common_query::native_histogram::is_native_histogram_value_type;
 use common_recordbatch::util::ChainedRecordBatchStream;
 use common_recordbatch::{RecordBatchStreamWrapper, SendableRecordBatchStream};
 use common_telemetry::tracing::{self, Instrument};
-use common_telemetry::warn;
+use common_telemetry::{debug, warn};
 use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet;
 use datafusion::physical_plan::{DisplayAs, DisplayFormatType};
+use datafusion_common::tree_node::TreeNode;
+use datafusion_expr::{Expr, Volatility};
 use datatypes::arrow::array::BinaryArray;
 use datatypes::arrow::datatypes::DataType as ArrowDataType;
 use datatypes::arrow::record_batch::RecordBatch;
 use datatypes::schema::SchemaRef;
+use datatypes::series::samples::{FieldSpec, SeriesSchema};
 use datatypes::types::TimestampType;
 use futures::{StreamExt, TryStreamExt};
 use smallvec::SmallVec;
 use snafu::{OptionExt, ResultExt, ensure};
 use store_api::metadata::RegionMetadataRef;
+use store_api::metric_engine_consts::is_metric_engine_internal_column;
 use store_api::region_engine::{
     PartitionRange, PrepareRequest, QueryScanContext, RawLabelSource, RegionScanner,
     ScannerProperties, SendableRawSeriesStream,
 };
+use store_api::storage::consts::{ReservedColumnId, is_internal_column};
 use table::predicate::Predicate;
 use tokio::sync::Semaphore;
 use tokio::sync::mpsc::error::{SendTimeoutError, TrySendError};
@@ -578,12 +585,6 @@ impl RegionScanner for SeriesScan {
         partition: usize,
         label_source: &RawLabelSource,
     ) -> Result<Option<SendableRawSeriesStream>, BoxedError> {
-        use api::v1::SemanticType;
-        use common_query::native_histogram::is_native_histogram_value_type;
-        use datatypes::series::samples::{FieldSpec, SeriesSchema};
-        use store_api::metric_engine_consts::is_metric_engine_internal_column;
-        use store_api::storage::consts::{ReservedColumnId, is_internal_column};
-
         let input = &self.stream_ctx.input;
         let metadata = input.mapper.metadata();
         match label_source {
@@ -593,6 +594,13 @@ impl RegionScanner for SeriesScan {
                 ..
             } => {
                 if !Arc::ptr_eq(metadata, physical_metadata) {
+                    debug!(
+                        "Raw series scan preflight rejected metadata snapshot: region_id={}, physical_region_id={}, schema_version={}, physical_schema_version={}",
+                        metadata.region_id,
+                        physical_metadata.region_id,
+                        metadata.schema_version,
+                        physical_metadata.schema_version,
+                    );
                     return Err(BoxedError::new(
                         InvalidRecordBatchSnafu {
                             reason: "raw label source metadata does not match scanner snapshot"
@@ -615,7 +623,13 @@ impl RegionScanner for SeriesScan {
                             && !is_metric_engine_internal_column(tag.name.as_ref())
                             && !ReservedColumnId::is_reserved(tag.column_id)
                     });
-                    if !valid || !ids.insert(tag.column_id) || !names.insert(tag.name.as_ref()) {
+                    let unique_id = valid && ids.insert(tag.column_id);
+                    let unique_name = unique_id && names.insert(tag.name.as_ref());
+                    if !valid || !unique_id || !unique_name {
+                        debug!(
+                            "Raw series scan preflight rejected label binding: region_id={}, column_id={}, binding_valid={}, unique_id={}, unique_name={}",
+                            metadata.region_id, tag.column_id, valid, unique_id, unique_name,
+                        );
                         return Err(BoxedError::new(
                             InvalidRecordBatchSnafu {
                                 reason: "invalid raw label binding".to_string(),
@@ -625,7 +639,13 @@ impl RegionScanner for SeriesScan {
                     }
                 }
             }
-            RawLabelSource::Explicit => return Ok(None),
+            RawLabelSource::Explicit => {
+                debug!(
+                    "Raw series scan preflight rejected explicit labels: region_id={}",
+                    metadata.region_id,
+                );
+                return Ok(None);
+            }
         }
         let selected = input.mapper.output_schema();
         let mut fields = Vec::new();
@@ -644,6 +664,10 @@ impl RegionScanner for SeriesScan {
                 .enumerate()
                 .find(|(_, (id, _))| *id == column.column_id)
             else {
+                debug!(
+                    "Raw series scan preflight missing mapper field: region_id={}, column_id={}",
+                    metadata.region_id, column.column_id,
+                );
                 return Err(BoxedError::new(
                     InvalidRecordBatchSnafu {
                         reason: format!(
@@ -655,6 +679,10 @@ impl RegionScanner for SeriesScan {
                 ));
             };
             if !data_type.is_numeric() && !is_native_histogram_value_type(data_type) {
+                debug!(
+                    "Raw series scan preflight rejected field type: region_id={}, column_id={}",
+                    metadata.region_id, column.column_id,
+                );
                 return Ok(None);
             }
             fields.push((index, data_type.clone()));
@@ -668,6 +696,14 @@ impl RegionScanner for SeriesScan {
             || input.series_row_selector.is_some()
             || !time_is_millis
         {
+            debug!(
+                "Raw series scan preflight rejected scan options: region_id={}, selected_fields={}, append_mode={}, row_selector={}, time_is_millis={}",
+                metadata.region_id,
+                fields.len(),
+                input.append_mode,
+                input.series_row_selector.is_some(),
+                time_is_millis,
+            );
             return Ok(None);
         }
         let schema = Arc::new(SeriesSchema {
@@ -677,9 +713,6 @@ impl RegionScanner for SeriesScan {
         let typed_schema = input.mapper.input_arrow_schema(false);
         let mut predicates = Vec::new();
         if let Some(predicate) = input.predicate_group().predicate_without_region() {
-            use datafusion_common::tree_node::TreeNode;
-            use datafusion_expr::{Expr, Volatility};
-
             for expr in predicate.exprs() {
                 // Fresh execution properties cannot preserve Stable or query-dependent semantics.
                 if expr
@@ -690,16 +723,28 @@ impl RegionScanner for SeriesScan {
                     })
                     .unwrap_or(true)
                 {
+                    debug!(
+                        "Raw series scan preflight rejected predicate volatility or variable: region_id={}",
+                        metadata.region_id,
+                    );
                     return Ok(None);
                 }
                 // Evaluate residuals with original types before raw conversion to Float64.
                 let Ok(physical) = Predicate::to_physical_expr(expr, &typed_schema) else {
+                    debug!(
+                        "Raw series scan preflight failed predicate compilation: region_id={}",
+                        metadata.region_id,
+                    );
                     return Ok(None);
                 };
                 if !matches!(
                     physical.data_type(&typed_schema),
                     Ok(ArrowDataType::Boolean | ArrowDataType::Null)
                 ) {
+                    debug!(
+                        "Raw series scan preflight rejected predicate type: region_id={}",
+                        metadata.region_id,
+                    );
                     return Ok(None);
                 }
                 predicates.push(physical);

src/mito2/src/read/series_scan/raw.rs

@@ -28,11 +28,12 @@ use common_query::native_histogram::{
 use common_recordbatch::filter;
 use datafusion::physical_plan::PhysicalExpr;
 use datatypes::arrow::array::{
-    Array, BinaryArray, DictionaryArray, Float64Array, StructArray, TimestampMillisecondArray,
+    Array, BinaryArray, Float64Array, StructArray, TimestampMillisecondArray,
 };
 use datatypes::arrow::buffer::Buffer;
 use datatypes::arrow::compute::cast;
-use datatypes::arrow::datatypes::{DataType, UInt32Type};
+use datatypes::arrow::datatypes::ArrowNativeType;
+use datatypes::arrow::datatypes::DataType;
 use datatypes::arrow::record_batch::RecordBatch;
 use datatypes::prelude::ConcreteDataType;
 use datatypes::series::raw::{RawSeriesChunk, RawSeriesOrdering, RawTimeAxis};
@@ -50,6 +51,7 @@ use crate::read::scan_util::PartitionMetrics;
 use crate::read::series_scan::SeriesBatch;
 use crate::read::stream::{ScanBatch, ScanBatchStream};
 use crate::sst::parquet::flat_format::{primary_key_column_index, time_index_column_index};
+use crate::sst::parquet::format::PrimaryKeyArray;
 
 fn boxed(error: Error) -> BoxedError {
     BoxedError::new(error)
@@ -64,41 +66,14 @@ fn invalid(reason: impl Into<String>) -> BoxedError {
     )
 }
 
-fn read_key(batch: &RecordBatch, row: usize) -> Result<&[u8], BoxedError> {
-    let column = batch.column(primary_key_column_index(batch.num_columns()));
-    if let Some(binary) = column.as_any().downcast_ref::<BinaryArray>() {
-        return binary
-            .is_valid(row)
-            .then(|| binary.value(row))
-            .ok_or_else(|| invalid("null primary key"));
-    }
-    let dictionary = column
-        .as_any()
-        .downcast_ref::<DictionaryArray<UInt32Type>>()
-        .ok_or_else(|| invalid("invalid primary-key array"))?;
-    let values = dictionary
-        .values()
-        .as_any()
-        .downcast_ref::<BinaryArray>()
-        .ok_or_else(|| invalid("invalid primary-key dictionary"))?;
-    if !dictionary.keys().is_valid(row) {
-        return Err(invalid("null primary key"));
-    }
-    let index = dictionary.keys().value(row) as usize;
-    if index >= values.len() || !values.is_valid(index) {
-        return Err(invalid("invalid primary-key dictionary value"));
-    }
-    Ok(values.value(index))
-}
-
 fn chunk(
     batch: &RecordBatch,
     start: usize,
     end: usize,
+    key: Buffer,
     schema: &SeriesSchema,
     fields: &[(usize, ConcreteDataType)],
 ) -> Result<RawSeriesChunk, BoxedError> {
-    let key = read_key(batch, start)?;
     let timestamp = batch
         .column(time_index_column_index(batch.num_columns()))
         .as_any()
@@ -144,16 +119,9 @@ fn chunk(
     let fields = SeriesFields::try_from_vec(samples)
         .context(DataTypeMismatchSnafu)
         .map_err(BoxedError::new)?;
-    RawSeriesChunk::try_new(
-        schema,
-        Buffer::from(key.to_vec()),
-        None,
-        time,
-        fields,
-        false,
-    )
-    .context(DataTypeMismatchSnafu)
-    .map_err(BoxedError::new)
+    RawSeriesChunk::try_new(schema, key, None, time, fields, false)
+        .context(DataTypeMismatchSnafu)
+        .map_err(BoxedError::new)
 }
 
 struct RawStream {
@@ -194,9 +162,9 @@ pub(crate) fn make_stream(
         let mut input = inner;
         let mut pending: Option<RawSeriesChunk> = None;
         while let Some(item) = input.next().await {
-            let records = match item {
-                Ok(ScanBatch::Series(SeriesBatch::Flat(flat))) => flat.batches.into_vec(),
-                Ok(ScanBatch::RecordBatch(batch)) => vec![batch],
+            let (records, is_legacy) = match item {
+                Ok(ScanBatch::Series(SeriesBatch::Flat(flat))) => (flat.batches.into_vec(), true),
+                Ok(ScanBatch::RecordBatch(batch)) => (vec![batch], false),
                 Err(error) => { yield Err(BoxedError::new(error)); return; }
             };
             for record in records {
@@ -222,22 +190,59 @@ pub(crate) fn make_stream(
                 if record.num_rows() == 0 {
                     continue;
                 }
+                let pk_column = record.column(primary_key_column_index(record.num_columns()));
+                let primary_key = match pk_column.as_any().downcast_ref::<PrimaryKeyArray>() {
+                    Some(primary_key) => primary_key,
+                    None => { yield Err(invalid("invalid primary-key array")); return; }
+                };
+                let values = match primary_key.values().as_any().downcast_ref::<BinaryArray>() {
+                    Some(values) => values,
+                    None => { yield Err(invalid("invalid primary-key dictionary")); return; }
+                };
                 let mut start = 0;
                 while start < record.num_rows() {
-                    let key = match read_key(&record, start) { Ok(key) => key, Err(error) => { yield Err(error); return; } };
+                    let index = if primary_key.keys().is_valid(start) {
+                        primary_key.keys().value(start).as_usize()
+                    } else {
+                        yield Err(invalid("null primary key")); return;
+                    };
+                    if index >= values.len() || !values.is_valid(index) {
+                        yield Err(invalid("invalid primary-key dictionary value")); return;
+                    }
                     let mut end = start + 1;
                     while end < record.num_rows() {
-                        match read_key(&record, end) {
-                            Ok(next) if next == key => end += 1,
-                            Ok(_) => break,
-                            Err(error) => { yield Err(error); return; }
+                        let next_index = if primary_key.keys().is_valid(end) {
+                            primary_key.keys().value(end).as_usize()
+                        } else {
+                            yield Err(invalid("null primary key")); return;
+                        };
+                        if next_index >= values.len() || !values.is_valid(next_index) {
+                            yield Err(invalid("invalid primary-key dictionary value")); return;
+                        }
+                        if super::primary_key_at(primary_key, values, end)
+                            == super::primary_key_at(primary_key, values, start)
+                        {
+                            end += 1;
+                        } else {
+                            break;
                         }
                     }
+                    let offsets = values.value_offsets();
+                    let value_start = offsets[index].as_usize();
+                    let value_len = offsets[index + 1].as_usize() - value_start;
+                    let key = values.values().slice_with_length(value_start, value_len);
                     let begun = Instant::now();
-                    let current = match chunk(&record, start, end, &output_schema, &fields) {
+                    let current = match chunk(&record, start, end, key, &output_schema, &fields) {
                         Ok(chunk) => chunk, Err(error) => { yield Err(error); return; }
                     };
                     metrics.inc_convert_batch_cost(begun.elapsed());
+                    if let Some(old) = pending.as_ref()
+                        && is_legacy
+                        && old.read_key.as_slice() != current.read_key.as_slice()
+                        && current.read_key.as_slice() <= old.read_key.as_slice()
+                    {
+                        yield Err(invalid("legacy series primary keys are not strictly increasing")); return;
+                    }
                     if let Some(old) = pending.as_ref()
                         && old.read_key.as_slice() == current.read_key.as_slice()
                         && current.time.timestamp(0).zip(old.time.timestamp(old.time.len() - 1))

src/mito2/src/read/series_scan/raw/tests.rs

@@ -49,6 +49,8 @@ use crate::error::{Error, InvalidRecordBatchSnafu, Result};
 use crate::read::scan_util::PartitionMetrics;
 use crate::read::series_scan::{FlatSeriesBatch, SeriesBatch};
 use crate::read::stream::{ScanBatch, ScanBatchStream};
+use crate::sst::parquet::flat_format::wrap_pk_binary_to_dict;
+use crate::sst::parquet::format::PrimaryKeyArray;
 
 fn metrics() -> PartitionMetrics {
     PartitionMetrics::new(
@@ -71,7 +73,7 @@ fn schema() -> Arc<SeriesSchema> {
     })
 }
 
-fn batch(keys: &[&[u8]], times: &[i64], values: &[Option<i64>]) -> RecordBatch {
+fn binary_batch(keys: &[&[u8]], times: &[i64], values: &[Option<i64>]) -> RecordBatch {
     assert_eq!(keys.len(), times.len());
     assert_eq!(keys.len(), values.len());
     let columns: Vec<ArrayRef> = vec![
@@ -95,6 +97,33 @@ fn batch(keys: &[&[u8]], times: &[i64], values: &[Option<i64>]) -> RecordBatch {
     RecordBatch::try_new(schema, columns).unwrap()
 }
 
+fn dictionary_batch(keys: &[&[u8]], times: &[i64], values: &[Option<i64>]) -> RecordBatch {
+    let binary = binary_batch(keys, times, values);
+    let schema = Arc::new(Schema::new(
+        binary
+            .schema()
+            .fields()
+            .iter()
+            .enumerate()
+            .map(|(index, field)| {
+                if index == binary.num_columns() - 3 {
+                    Field::new(
+                        field.name(),
+                        DataType::Dictionary(
+                            Box::new(DataType::UInt32),
+                            Box::new(DataType::Binary),
+                        ),
+                        field.is_nullable(),
+                    )
+                } else {
+                    field.as_ref().clone()
+                }
+            })
+            .collect::<Vec<_>>(),
+    ));
+    wrap_pk_binary_to_dict(binary, &schema).unwrap()
+}
+
 fn histogram_input() -> (RecordBatch, Vec<NativeHistogram>) {
     let expected = vec![NativeHistogram {
         schema: 3,
@@ -145,7 +174,29 @@ fn histogram_input() -> (RecordBatch, Vec<NativeHistogram>) {
         Field::new("__op_type", DataType::UInt8, false),
     ];
     (
-        RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap(),
+        wrap_pk_binary_to_dict(
+            RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap(),
+            &Arc::new(Schema::new(vec![
+                Field::new(
+                    "histogram",
+                    native_histogram_value_type().as_arrow_type(),
+                    true,
+                ),
+                Field::new(
+                    "ts",
+                    DataType::Timestamp(TimeUnit::Millisecond, None),
+                    false,
+                ),
+                Field::new(
+                    "__primary_key",
+                    DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Binary)),
+                    false,
+                ),
+                Field::new("__sequence", DataType::UInt64, false),
+                Field::new("__op_type", DataType::UInt8, false),
+            ])),
+        )
+        .unwrap(),
         expected,
     )
 }
@@ -251,12 +302,12 @@ fn values(chunk: &RawSeriesChunk) -> Vec<Option<f64>> {
 #[tokio::test]
 async fn chunks_preserve_keys_values_times_and_final_flags_for_both_input_shapes() {
     let schema = schema();
-    let first = batch(
+    let first = dictionary_batch(
         &[b"prefix-full-key-a", b"prefix-full-key-a"],
         &[1, 2],
         &[Some(1), None],
     );
-    let second = batch(
+    let second = dictionary_batch(
         &[
             b"prefix-full-key-a",
             b"prefix-full-key-a",
@@ -328,7 +379,7 @@ async fn typed_predicate_filters_before_numeric_samples_are_cast() {
         physical_lit(ScalarValue::Int64(Some(threshold))),
     ));
     let mut output = make_stream(
-        input(vec![Ok(ScanBatch::RecordBatch(batch(
+        input(vec![Ok(ScanBatch::RecordBatch(dictionary_batch(
             &[b"a", b"a", b"a", b"b", b"c"],
             &[1, 2, 3, 1, 1],
             &[
@@ -363,7 +414,7 @@ async fn typed_predicate_filters_before_numeric_samples_are_cast() {
 async fn predicate_error_is_terminal() {
     let predicate: Arc<dyn PhysicalExpr> = Arc::new(Column::new("value", 0));
     let mut output = make_stream(
-        input(vec![Ok(ScanBatch::RecordBatch(batch(
+        input(vec![Ok(ScanBatch::RecordBatch(dictionary_batch(
             &[b"a"],
             &[1],
             &[Some(1)],
@@ -386,9 +437,21 @@ async fn upstream_error_is_terminal_after_yielded_prefix_without_flushing_pendin
     .build();
     let mut output = convert(
         input(vec![
-            Ok(ScanBatch::RecordBatch(batch(&[b"a"], &[1], &[Some(1)]))),
-            Ok(ScanBatch::RecordBatch(batch(&[b"a"], &[2], &[Some(2)]))),
-            Ok(ScanBatch::RecordBatch(batch(&[b"a"], &[3], &[Some(3)]))),
+            Ok(ScanBatch::RecordBatch(dictionary_batch(
+                &[b"a"],
+                &[1],
+                &[Some(1)],
+            ))),
+            Ok(ScanBatch::RecordBatch(dictionary_batch(
+                &[b"a"],
+                &[2],
+                &[Some(2)],
+            ))),
+            Ok(ScanBatch::RecordBatch(dictionary_batch(
+                &[b"a"],
+                &[3],
+                &[Some(3)],
+            ))),
             Err(upstream_error),
         ]),
         schema(),
@@ -418,9 +481,21 @@ async fn upstream_error_is_terminal_after_yielded_prefix_without_flushing_pendin
 async fn dropping_raw_stream_drops_unfinished_input() {
     let drops = Arc::new(AtomicUsize::new(0));
     let upstream = input(vec![
-        Ok(ScanBatch::RecordBatch(batch(&[b"a"], &[1], &[Some(1)]))),
-        Ok(ScanBatch::RecordBatch(batch(&[b"b"], &[1], &[Some(2)]))),
-        Ok(ScanBatch::RecordBatch(batch(&[b"c"], &[1], &[Some(3)]))),
+        Ok(ScanBatch::RecordBatch(dictionary_batch(
+            &[b"a"],
+            &[1],
+            &[Some(1)],
+        ))),
+        Ok(ScanBatch::RecordBatch(dictionary_batch(
+            &[b"b"],
+            &[1],
+            &[Some(2)],
+        ))),
+        Ok(ScanBatch::RecordBatch(dictionary_batch(
+            &[b"c"],
+            &[1],
+            &[Some(3)],
+        ))),
     ]);
     let spy = DropSpy {
         inner: upstream,
@@ -433,47 +508,107 @@ async fn dropping_raw_stream_drops_unfinished_input() {
 }
 
 #[tokio::test]
-async fn raw_stream_owns_schema_and_bound_source_metadata() {
-    let schema = schema();
-    let schema_weak = Arc::downgrade(&schema);
-    let metadata = metadata();
-    let metadata_weak = Arc::downgrade(&metadata);
-    let mut raw = convert_with_source(
-        input(vec![]),
-        schema.clone(),
-        primary_key_source(metadata.clone()),
-        ConcreteDataType::int64_datatype(),
+async fn legacy_series_requires_increasing_keys_but_record_batches_do_not() {
+    let mut batches = smallvec::SmallVec::new();
+    batches.push(dictionary_batch(&[b"b"], &[1], &[Some(1)]));
+    batches.push(dictionary_batch(&[b"a"], &[2], &[Some(2)]));
+    let mut output = convert(
+        input(vec![Ok(ScanBatch::Series(SeriesBatch::Flat(
+            FlatSeriesBatch { batches },
+        )))]),
+        schema(),
     );
-    drop(schema);
-    drop(metadata);
-    assert!(schema_weak.upgrade().is_some());
-    assert!(metadata_weak.upgrade().is_some());
-    assert_eq!(raw.schema().fields[0].name.as_ref(), "value");
-    match raw.label_source() {
-        RawLabelSource::PrimaryKey {
-            metric_name,
-            physical_metadata,
-            logical_tags,
-        } => {
-            assert_eq!(metric_name.as_ref(), "metric");
-            assert_eq!(physical_metadata.region_id, RegionId::new(1, 0));
-            assert_eq!(
-                physical_metadata.column_metadatas[0].column_schema.name,
-                "tag"
-            );
-            assert_eq!(logical_tags[0].column_id, 0);
-            assert_eq!(logical_tags[0].name.as_ref(), "tag");
-            assert!(Arc::ptr_eq(
-                physical_metadata,
-                &metadata_weak.upgrade().unwrap()
-            ));
-        }
-        RawLabelSource::Explicit => panic!("expected bound primary-key labels"),
-    }
-    assert!(raw.next().await.is_none());
-    drop(raw);
-    assert!(schema_weak.upgrade().is_none());
-    assert!(metadata_weak.upgrade().is_none());
+    assert!(output.next().await.unwrap().is_err());
+    assert!(output.next().await.is_none());
+
+    let mut batches = smallvec::SmallVec::new();
+    batches.push(dictionary_batch(&[b"a"], &[1], &[Some(1)]));
+    batches.push(dictionary_batch(&[b"b"], &[1], &[Some(2)]));
+    batches.push(dictionary_batch(&[b"a"], &[2], &[Some(3)]));
+    let mut output = convert(
+        input(vec![Ok(ScanBatch::Series(SeriesBatch::Flat(
+            FlatSeriesBatch { batches },
+        )))]),
+        schema(),
+    );
+    let prefix = output.next().await.unwrap().unwrap();
+    assert_eq!(prefix.read_key.as_slice(), b"a");
+    assert_eq!(values(&prefix), vec![Some(1.0)]);
+    assert_eq!(prefix.time.timestamp(0), Some(1));
+    assert!(prefix.is_last);
+    assert!(output.next().await.unwrap().is_err());
+    assert!(output.next().await.is_none());
+
+    let mut batches = smallvec::SmallVec::new();
+    batches.push(dictionary_batch(&[b"a"], &[1], &[Some(1)]));
+    batches.push(dictionary_batch(&[b"b"], &[1], &[Some(2)]));
+    let mut output = convert(
+        input(vec![Ok(ScanBatch::Series(SeriesBatch::Flat(
+            FlatSeriesBatch { batches },
+        )))]),
+        schema(),
+    );
+    let first = output.next().await.unwrap().unwrap();
+    assert_eq!(first.read_key.as_slice(), b"a");
+    assert_eq!(values(&first), vec![Some(1.0)]);
+    assert_eq!(first.time.timestamp(0), Some(1));
+    assert!(first.is_last);
+    let second = output.next().await.unwrap().unwrap();
+    assert_eq!(second.read_key.as_slice(), b"b");
+    assert_eq!(values(&second), vec![Some(2.0)]);
+    assert_eq!(second.time.timestamp(0), Some(1));
+    assert!(second.is_last);
+    assert!(output.next().await.is_none());
+
+    let mut output = convert(
+        input(vec![
+            Ok(ScanBatch::RecordBatch(dictionary_batch(
+                &[b"b"],
+                &[1],
+                &[Some(1)],
+            ))),
+            Ok(ScanBatch::RecordBatch(dictionary_batch(
+                &[b"a"],
+                &[1],
+                &[Some(2)],
+            ))),
+        ]),
+        schema(),
+    );
+    let first = output.next().await.unwrap().unwrap();
+    assert_eq!(first.read_key.as_slice(), b"b");
+    assert_eq!(values(&first), vec![Some(1.0)]);
+    assert_eq!(first.time.timestamp(0), Some(1));
+    assert!(first.is_last);
+    let second = output.next().await.unwrap().unwrap();
+    assert_eq!(second.read_key.as_slice(), b"a");
+    assert_eq!(values(&second), vec![Some(2.0)]);
+    assert_eq!(second.time.timestamp(0), Some(1));
+    assert!(second.is_last);
+    assert!(output.next().await.is_none());
+}
+
+#[tokio::test]
+async fn raw_key_shares_dictionary_values_and_survives_batch_drop() {
+    let record = dictionary_batch(
+        &[b"unused", b"prefix-key", b"other"],
+        &[0, 1, 1],
+        &[Some(0), Some(1), Some(2)],
+    )
+    .slice(1, 2);
+    let pk = record
+        .column(record.num_columns() - 3)
+        .as_any()
+        .downcast_ref::<PrimaryKeyArray>()
+        .unwrap();
+    let values = pk.values().as_any().downcast_ref::<BinaryArray>().unwrap();
+    let expected_ptr = values.value(1).as_ptr();
+    let mut output = convert(input(vec![Ok(ScanBatch::RecordBatch(record))]), schema());
+    let chunk = output.next().await.unwrap().unwrap();
+    assert_eq!(chunk.read_key.as_slice(), b"prefix-key");
+    assert_eq!(chunk.read_key.as_ptr(), expected_ptr);
+    drop(output);
+    assert_eq!(chunk.read_key.as_slice(), b"prefix-key");
 }
 
 #[tokio::test]
@@ -532,12 +667,16 @@ async fn duplicate_or_descending_cross_batch_time_is_terminal() {
     for next in [2, 1] {
         let mut output = convert(
             input(vec![
-                Ok(ScanBatch::RecordBatch(batch(
+                Ok(ScanBatch::RecordBatch(dictionary_batch(
                     &[b"k", b"k"],
                     &[1, 2],
                     &[Some(1), Some(2)],
                 ))),
-                Ok(ScanBatch::RecordBatch(batch(&[b"k"], &[next], &[Some(3)]))),
+                Ok(ScanBatch::RecordBatch(dictionary_batch(
+                    &[b"k"],
+                    &[next],
+                    &[Some(3)],
+                ))),
             ]),
             schema(),
         );