diff --git a/src/mito2/src/read/series_scan.rs b/src/mito2/src/read/series_scan.rs index 85d95cf172b..eb5c1048ab5 100644 --- a/src/mito2/src/read/series_scan.rs +++ b/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, 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); diff --git a/src/mito2/src/read/series_scan/raw.rs b/src/mito2/src/read/series_scan/raw.rs index ffc28564303..2792eb24e81 100644 --- a/src/mito2/src/read/series_scan/raw.rs +++ b/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) -> 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::() { - return binary - .is_valid(row) - .then(|| binary.value(row)) - .ok_or_else(|| invalid("null primary key")); - } - let dictionary = column - .as_any() - .downcast_ref::>() - .ok_or_else(|| invalid("invalid primary-key array"))?; - let values = dictionary - .values() - .as_any() - .downcast_ref::() - .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 { - 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 = 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::() { + Some(primary_key) => primary_key, + None => { yield Err(invalid("invalid primary-key array")); return; } + }; + let values = match primary_key.values().as_any().downcast_ref::() { + 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)) diff --git a/src/mito2/src/read/series_scan/raw/tests.rs b/src/mito2/src/read/series_scan/raw/tests.rs index e70bc43f405..5f581fd1acb 100644 --- a/src/mito2/src/read/series_scan/raw/tests.rs +++ b/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 { }) } -fn batch(keys: &[&[u8]], times: &[i64], values: &[Option]) -> RecordBatch { +fn binary_batch(keys: &[&[u8]], times: &[i64], values: &[Option]) -> RecordBatch { assert_eq!(keys.len(), times.len()); assert_eq!(keys.len(), values.len()); let columns: Vec = vec![ @@ -95,6 +97,33 @@ fn batch(keys: &[&[u8]], times: &[i64], values: &[Option]) -> RecordBatch { RecordBatch::try_new(schema, columns).unwrap() } +fn dictionary_batch(keys: &[&[u8]], times: &[i64], values: &[Option]) -> 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::>(), + )); + wrap_pk_binary_to_dict(binary, &schema).unwrap() +} + fn histogram_input() -> (RecordBatch, Vec) { let expected = vec![NativeHistogram { schema: 3, @@ -145,7 +174,29 @@ fn histogram_input() -> (RecordBatch, Vec) { 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> { #[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 = 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::() + .unwrap(); + let values = pk.values().as_any().downcast_ref::().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(), );