From 5b649daaeaf019616cb6e6b05798e73a54b73227 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Sat, 17 Oct 2026 11:30:55 +0800 Subject: [PATCH 1/4] feat(mito2): produce typed raw series from RegionScanner with truthful ordering Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- Cargo.lock | 1 + src/mito2/Cargo.toml | 1 + src/mito2/src/engine/scan_test.rs | 918 +++++++++++++++++++- src/mito2/src/read/series_scan.rs | 383 ++++---- src/mito2/src/read/series_scan/raw.rs | 13 +- src/mito2/src/read/series_scan/raw/tests.rs | 76 +- 6 files changed, 1211 insertions(+), 181 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index ea2a530d6b5..43d65a03bd8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -8742,6 +8742,7 @@ dependencies = [ "partition", "pin-project", "prometheus 0.14.0", + "promql", "prost 0.14.1", "puffin", "rand 0.9.4", diff --git a/src/mito2/Cargo.toml b/src/mito2/Cargo.toml index 02c134e964c..151694c40d1 100644 --- a/src/mito2/Cargo.toml +++ b/src/mito2/Cargo.toml @@ -101,6 +101,7 @@ dotenv.workspace = true log-store = { workspace = true, features = ["testing"] } mito-codec = { workspace = true, features = ["testing"] } object-store = { workspace = true, features = ["services-memory"] } +promql.workspace = true rskafka.workspace = true rstest.workspace = true rstest_reuse.workspace = true diff --git a/src/mito2/src/engine/scan_test.rs b/src/mito2/src/engine/scan_test.rs index a893a234adc..3ad13fac11e 100644 --- a/src/mito2/src/engine/scan_test.rs +++ b/src/mito2/src/engine/scan_test.rs @@ -4772,6 +4772,7 @@ async fn test_mito_handle_query_raw_series_candidate_projection_and_mvcc() { use datatypes::series::samples::SampleRef; use datatypes::value::Value as DatumValue; use mito_codec::row_converter::build_primary_key_codec; + use promql::source::materialize_raw_labels; use store_api::region_engine::{BoundTag, QueryScanContext, RawLabelSource}; use store_api::region_request::RegionDeleteRequest; use table::predicate::Predicate; @@ -4907,7 +4908,7 @@ async fn test_mito_handle_query_raw_series_candidate_projection_and_mvcc() { .scan_series_partition(&QueryScanContext::default(), &metrics, partition, &source) .unwrap() .expect("supported millisecond PerSeries scan"); - assert_eq!(SeriesOrdering::Contiguous, stream.ordering()); + assert_eq!(SeriesOrdering::KeyThenTime, stream.ordering()); assert!(stream.schema().eval_grid.is_none()); assert_eq!( vec!["field_1", "field_0"], @@ -4925,7 +4926,7 @@ async fn test_mito_handle_query_raw_series_candidate_projection_and_mvcc() { !closed_keys.contains(chunk.primary_key.as_slice()), "no chunk may follow a key's is_last chunk" ); - let times = chunk.time.as_buffer().as_ref(); + let times = chunk.time.as_buffer(); assert_eq!(2, chunk.fields.len()); assert_eq!(times.len(), chunk.fields.get(0).unwrap().len()); assert_eq!(times.len(), chunk.fields.get(1).unwrap().len()); @@ -4994,7 +4995,7 @@ async fn test_mito_handle_query_raw_series_candidate_projection_and_mvcc() { } // Compare keyed visible tuples through a separately opened legacy scanner. // A partition cannot be consumed through both paths. - let mut legacy_request = request; + let mut legacy_request = request.clone(); legacy_request.projection = Some(vec![0, 3, 2, 1, 4]); legacy_request.preserve_pk_dictionary_encoding = false; let legacy = engine @@ -5090,6 +5091,133 @@ async fn test_mito_handle_query_raw_series_candidate_projection_and_mvcc() { .collect::>(); raw_visible.sort_by(|a, b| a.0.cmp(&b.0).then(a.1.cmp(&b.1))); assert_eq!(raw_visible, visible); + + // Exercise the production table physical consumer over a fresh real Mito scanner. + use common_query::series_frame::unpack_raw_series; + use datafusion::execution::TaskContext; + use datafusion::physical_plan::ExecutionPlan; + use table::table::scan::RegionScanExec; + let native_scanner = engine + .handle_query(region_id, request.clone()) + .await + .unwrap(); + let native_metadata = native_scanner.metadata(); + let native_tags = native_metadata + .column_metadatas + .iter() + .filter(|c| c.semantic_type == SemanticType::Tag) + .map(|c| BoundTag { + column_id: c.column_id, + name: Arc::from(c.column_schema.name.as_str()), + }) + .collect::>(); + let native_source = RawLabelSource::PrimaryKey { + metric_name: Arc::from("fixture_metric"), + physical_metadata: native_metadata, + logical_tags: native_tags.into(), + }; + let native_plan = Arc::new(RegionScanExec::new(native_scanner, request.clone(), None).unwrap()); + let native_plan = native_plan + .try_to_series_scan(native_source.clone()) + .unwrap() + .expect("Mito raw-series capability"); + let native_schema = native_plan.series_schema().clone(); + let native_frames = native_plan + .execute(0, Arc::new(TaskContext::default())) + .unwrap(); + let unpacked = unpack_raw_series( + native_frames, + native_schema, + native_source, + SeriesOrdering::Contiguous, + ) + .unwrap(); + let native_schema = unpacked.schema().clone(); + let materialized = materialize_raw_labels(unpacked).unwrap(); + assert!(matches!( + materialized.label_source(), + RawLabelSource::Explicit + )); + let mut native = materialized.try_collect::>().await.unwrap(); + assert_eq!( + vec!["field_1", "field_0"], + native_schema + .fields + .iter() + .map(|field| field.name.as_ref()) + .collect::>() + ); + let mut native_visible = native + .drain(..) + .flat_map(|chunk| { + let metric = chunk.metric.as_ref().expect("materialized complete metric"); + assert_eq!(Some("fixture_metric"), metric.name()); + let labels = metric + .tags() + .iter() + .map(|label| (label.name.as_ref(), label.value.as_ref())) + .collect::>(); + assert_eq!(1, labels.len()); + let tag = labels["tag_0"].to_string(); + assert!(tag == "A" || tag == "B"); + let times = chunk.time.as_buffer(); + let key = chunk.primary_key.to_vec(); + let is_last = chunk.is_last; + let mut rows = Vec::with_capacity(times.len()); + for (i, time) in times.iter().enumerate() { + let SampleRef::Float(f1) = chunk.fields.get(0).unwrap().sample(i) else { + panic!("expected f1 float") + }; + let f0 = match chunk.fields.get(1).unwrap().sample(i) { + SampleRef::Float(value) => Some(value), + SampleRef::Missing => None, + _ => panic!("expected f0 float or null"), + }; + rows.push(( + key.clone(), + tag.clone(), + *time, + f1, + f0, + is_last && i + 1 == times.len(), + )); + } + rows + }) + .collect::>(); + let mut expected_native = raw + .iter() + .map(|(key, time, f1, f0, _)| (key.to_vec(), *time, *f1, *f0)) + .collect::>(); + type NativeRow = (Vec, String, i64, f64, Option, bool); + let sort_native = |rows: &mut Vec| { + rows.sort_by(|a, b| a.0.cmp(&b.0).then(a.2.cmp(&b.2))); + }; + sort_native(&mut native_visible); + expected_native.sort_by(|a, b| a.0.cmp(&b.0).then(a.1.cmp(&b.1))); + assert_eq!( + expected_native, + native_visible + .iter() + .map(|(key, _, time, f1, f0, _)| (key.clone(), *time, *f1, *f0)) + .collect::>() + ); + for key in &expected_keys { + assert_eq!( + 1, + native_visible + .iter() + .filter(|(row_key, _, _, _, _, is_last)| row_key == key && *is_last) + .count() + ); + } + let mut decoded_visible = native_visible + .iter() + .map(|(_, tag, time, f1, f0, _)| (tag.clone(), *time, *f1, *f0)) + .collect::>(); + decoded_visible.sort_by(|a, b| a.0.cmp(&b.0).then(a.1.cmp(&b.1))); + assert_eq!(visible, decoded_visible); + // Apply the typed predicate and final projection through the ordinary legacy filter path. assert_eq!(3, visible.len()); assert_eq!( @@ -5102,6 +5230,602 @@ async fn test_mito_handle_query_raw_series_candidate_projection_and_mvcc() { ); } +#[tokio::test] +async fn test_mito_regional_raw_merge_preserves_global_key_order_and_labels() { + use api::v1::{Row, Value}; + use common_query::series_frame::{pack_raw_series, unpack_raw_series}; + use datatypes::arrow::array::Array; + use datatypes::series::SeriesOrdering; + use datatypes::series::samples::{SampleRef, SeriesSchema}; + use datatypes::value::Value as DatumValue; + use mito_codec::row_converter::build_primary_key_codec; + use promql::source::{materialize_raw_labels, merge_ordered_raw_series}; + use store_api::region_engine::{BoundTag, QueryScanContext, RawLabelSource}; + use store_api::region_request::RegionDeleteRequest; + + let mut env = TestEnv::with_prefix("test_mito_regional_raw_merge").await; + let engine = env.create_engine(MitoConfig::default()).await; + let regions = [RegionId::new(1, 1), RegionId::new(1, 2)]; + let mut request = CreateRequestBuilder::new().field_num(3).build(); + request + .options + .insert("memtable.type".into(), "bulk".into()); + let schema = test_util::rows_schema(&request); + for region in regions { + engine + .handle_request(region, RegionRequest::Create(request.clone())) + .await + .unwrap(); + } + let make_row = |tag: &str, f0: Option, f1: f64, gate: f64, time: i64| Row { + values: vec![ + Some(ValueData::StringValue(tag.into())), + f0.map(ValueData::F64Value), + Some(ValueData::F64Value(f1)), + Some(ValueData::F64Value(gate)), + Some(ValueData::TimestampMillisecondValue(time)), + ] + .into_iter() + .map(|value_data| Value { value_data }) + .collect(), + }; + for (region, rows) in [ + ( + regions[0], + vec![ + make_row("A", Some(10.0), 1.0, 1.0, 10), + make_row("A", None, 3.0, 1.0, 30), + make_row("B", Some(-0.0), 5.0, 1.0, 10), + ], + ), + ( + regions[1], + vec![ + make_row("A", Some(20.0), 2.0, 1.0, 20), + make_row("A", Some(40.0), 4.0, 1.0, 40), + make_row("C", Some(60.0), 6.0, 1.0, 10), + ], + ), + ] { + test_util::put_rows( + &engine, + region, + Rows { + schema: schema.clone(), + rows, + }, + ) + .await; + } + test_util::flush_region(&engine, regions[0], None).await; + test_util::flush_region(&engine, regions[1], None).await; + test_util::put_rows( + &engine, + regions[0], + Rows { + schema: schema.clone(), + rows: vec![ + make_row("A", Some(11.0), 11.0, 1.0, 10), + make_row("A", Some(50.0), 5.0, 0.0, 50), + make_row("B", Some(60.0), 6.0, 1.0, 60), + ], + }, + ) + .await; + engine + .handle_request( + regions[0], + RegionRequest::Delete(RegionDeleteRequest { + rows: Rows { + schema: test_util::delete_rows_schema( + &CreateRequestBuilder::new().field_num(3).build(), + ), + rows: vec![row(vec![ + ValueData::StringValue("B".into()), + ValueData::TimestampMillisecondValue(60), + ])], + }, + hint: None, + partition_expr_version: None, + }), + ) + .await + .unwrap(); + let request = ScanRequest { + projection: Some(vec![2, 1, 4]), + filters: vec![col("field_2").gt(lit(0.0_f64))], + distribution: Some(TimeSeriesDistribution::PerSeries), + preserve_pk_dictionary_encoding: true, + ..Default::default() + }; + let mut streams = Vec::new(); + let mut metadata_for_codec = None; + let mut expected_schema: Option> = None; + let metrics = ExecutionPlanMetricsSet::default(); + for region in regions { + let selected = engine.scanner(region, request.clone()).await.unwrap(); + let mut scanner: Box = match selected { + Scanner::Series(series) => { + assert_eq!("legacy", series.mode()); + Box::new(series) + } + _ => panic!("PerSeries request must select a series scanner"), + }; + assert_eq!(SeriesOrdering::KeyThenTime, scanner.raw_series_ordering()); + let metadata = scanner.metadata(); + metadata_for_codec.get_or_insert_with(|| metadata.clone()); + let source = RawLabelSource::PrimaryKey { + metric_name: Arc::from("fixture_metric"), + physical_metadata: metadata.clone(), + logical_tags: metadata + .column_metadatas + .iter() + .filter(|column| { + column.semantic_type == SemanticType::Tag + && !store_api::storage::consts::is_internal_column( + &column.column_schema.name, + ) + && !store_api::metric_engine_consts::is_metric_engine_internal_column( + &column.column_schema.name, + ) + }) + .map(|column| BoundTag { + column_id: column.column_id, + name: Arc::from(column.column_schema.name.as_str()), + }) + .collect::>() + .into(), + }; + let ranges = scanner + .properties() + .partitions + .clone() + .into_iter() + .flatten() + .collect::>(); + assert!(!ranges.is_empty(), "region scan needs an input range"); + let partition_count = 2; + let mut partitions = vec![Vec::new(); partition_count]; + for (index, range) in ranges.into_iter().enumerate() { + partitions[index % partition_count].push(range); + } + scanner + .prepare(PrepareRequest { + ranges: Some(partitions), + target_partitions: Some(partition_count), + ..Default::default() + }) + .unwrap(); + let partition_count = scanner.properties().num_partitions(); + assert!( + partition_count > 1, + "prepare must use multiple supplied native range partitions" + ); + for partition in 0..partition_count { + let stream = scanner + .scan_series_partition(&QueryScanContext::default(), &metrics, partition, &source) + .unwrap() + .expect("native raw input"); + assert_eq!(scanner.raw_series_ordering(), stream.ordering()); + assert_eq!(SeriesOrdering::KeyThenTime, stream.ordering()); + assert_eq!( + vec!["field_1", "field_0"], + stream + .schema() + .fields + .iter() + .map(|field| field.name.as_ref()) + .collect::>() + ); + if let Some(schema) = &expected_schema { + assert_eq!( + schema + .fields + .iter() + .map(|field| field.name.as_ref()) + .collect::>(), + stream + .schema() + .fields + .iter() + .map(|field| field.name.as_ref()) + .collect::>() + ); + } else { + expected_schema = Some(Arc::clone(stream.schema())); + } + streams.push(stream); + } + // Each scanner owns the same schema and explicitly bound source. + } + let merged = merge_ordered_raw_series(streams).unwrap(); + let merged_schema = Arc::clone(merged.schema()); + let binding = merged.label_source().clone(); + let ordering = merged.ordering(); + assert_eq!(SeriesOrdering::KeyThenTime, ordering); + let packed = pack_raw_series(merged).unwrap(); + let unpacked = unpack_raw_series(packed, merged_schema, binding, ordering).unwrap(); + let mut materialized = materialize_raw_labels(unpacked).unwrap(); + let mut actual = Vec::new(); + while let Some(chunk) = materialized.try_next().await.unwrap() { + let metric = chunk.metric.as_ref().unwrap(); + assert_eq!(Some("fixture_metric"), metric.name()); + let tag = metric.get("tag_0").unwrap().to_string(); + for index in 0..chunk.time.len() { + let SampleRef::Float(f1) = chunk.fields.get(0).unwrap().sample(index) else { + panic!("field_1 float") + }; + let f0 = match chunk.fields.get(1).unwrap().sample(index) { + SampleRef::Float(value) => Some(value.to_bits()), + SampleRef::Missing => None, + _ => panic!("unexpected field_0 sample type"), + }; + actual.push(( + chunk.primary_key.to_vec(), + tag.clone(), + chunk.time.timestamp(index).unwrap(), + f1.to_bits(), + f0, + chunk.is_last && index + 1 == chunk.time.len(), + )); + } + } + assert!( + actual + .windows(2) + .all(|pair| pair[0].0 < pair[1].0 || (pair[0].0 == pair[1].0 && pair[0].2 < pair[1].2)), + "merged stream must be ordered by full read key then time" + ); + let metadata = metadata_for_codec.unwrap(); + let codec = build_primary_key_codec(&metadata); + let mut expected = Vec::new(); + for (tag, time, f1, f0) in [ + ("A", 10, 11.0_f64, Some(11.0_f64)), + ("A", 20, 2.0, Some(20.0)), + ("A", 30, 3.0, None), + ("A", 40, 4.0, Some(40.0)), + ("B", 10, 5.0, Some(-0.0_f64)), + ("C", 10, 6.0, Some(60.0)), + ] { + let mut key = Vec::new(); + codec + .encode_values( + &[(metadata.primary_key[0], DatumValue::String(tag.into()))], + &mut key, + ) + .unwrap(); + expected.push((key, tag, time, f1.to_bits(), f0.map(f64::to_bits))); + } + let duplicate_key = expected[0].0.clone(); + assert_eq!( + expected + .iter() + .map(|(key, _, _, _, _)| key.clone()) + .collect::>(), + actual.iter().map(|row| row.0.clone()).collect::>() + ); + let visible = actual + .iter() + .map(|(_, tag, time, f1, f0, _)| (tag.as_str(), *time, *f1, *f0)) + .collect::>(); + assert_eq!( + expected + .iter() + .map(|(_, tag, time, f1, f0)| (*tag, *time, *f1, *f0)) + .collect::>(), + visible + ); + let mut legacy_batches = Vec::new(); + for region in regions { + let mut legacy_request = request.clone(); + legacy_request.projection = Some(vec![0, 1, 2, 3, 4]); + legacy_request.filters.clear(); + legacy_request.preserve_pk_dictionary_encoding = false; + let legacy = engine.handle_query(region, legacy_request).await.unwrap(); + for partition in 0..legacy.properties().num_partitions() { + legacy_batches.extend( + legacy + .scan_partition(&QueryScanContext::default(), &metrics, partition) + .unwrap() + .try_collect::>() + .await + .unwrap(), + ); + } + } + let mut legacy_visible = Vec::new(); + for batch in &legacy_batches { + let batch = batch.df_record_batch(); + let tags = batch + .column_by_name("tag_0") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let f0 = batch + .column_by_name("field_0") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let f1 = batch + .column_by_name("field_1") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let times = batch + .column_by_name("ts") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let gate = batch + .column_by_name("field_2") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + for row in 0..batch.num_rows() { + if gate.value(row) > 0.0 { + legacy_visible.push(( + tags.value(row).to_string(), + times.value(row), + f1.value(row).to_bits(), + if f0.is_null(row) { + None + } else { + Some(f0.value(row).to_bits()) + }, + )); + } + } + } + legacy_visible.sort_by(|left, right| left.0.cmp(&right.0).then(left.1.cmp(&right.1))); + let mut raw_visible = actual + .iter() + .map(|(_, tag, time, f1, f0, _)| (tag.clone(), *time, *f1, *f0)) + .collect::>(); + raw_visible.sort_by(|left, right| left.0.cmp(&right.0).then(left.1.cmp(&right.1))); + assert_eq!( + legacy_visible, raw_visible, + "raw candidates after field_2 predicate must match legacy scan" + ); + test_util::put_rows( + &engine, + regions[1], + Rows { + schema: schema.clone(), + rows: vec![make_row("A", Some(99.0), 99.0, 1.0, 30)], + }, + ) + .await; + let mut duplicate_inputs = Vec::new(); + for region in regions { + let scanner = engine.handle_query(region, request.clone()).await.unwrap(); + let metadata = scanner.metadata(); + let source = RawLabelSource::PrimaryKey { + metric_name: Arc::from("fixture_metric"), + physical_metadata: metadata.clone(), + logical_tags: metadata + .column_metadatas + .iter() + .filter(|column| { + column.semantic_type == SemanticType::Tag + && !store_api::storage::consts::is_internal_column( + &column.column_schema.name, + ) + && !store_api::metric_engine_consts::is_metric_engine_internal_column( + &column.column_schema.name, + ) + }) + .map(|column| BoundTag { + column_id: column.column_id, + name: Arc::from(column.column_schema.name.as_str()), + }) + .collect::>() + .into(), + }; + let mut scanner = scanner; + let ranges = scanner + .properties() + .partitions + .clone() + .into_iter() + .flatten() + .collect::>(); + let partition_count = 2; + let mut partitions = vec![Vec::new(); partition_count]; + for (index, range) in ranges.into_iter().enumerate() { + partitions[index % partition_count].push(range); + } + scanner + .prepare(PrepareRequest { + ranges: Some(partitions), + target_partitions: Some(partition_count), + ..Default::default() + }) + .unwrap(); + for partition in 0..scanner.properties().num_partitions() { + duplicate_inputs.push( + scanner + .scan_series_partition( + &QueryScanContext::default(), + &metrics, + partition, + &source, + ) + .unwrap() + .unwrap(), + ); + } + } + let mut duplicate_merge = merge_ordered_raw_series(duplicate_inputs).unwrap(); + let mut affected_key_last = Vec::new(); + let duplicate_error = loop { + match duplicate_merge.try_next().await { + Ok(Some(chunk)) => { + if chunk.primary_key.as_slice() == duplicate_key.as_slice() { + affected_key_last.push(chunk.is_last); + } + } + Ok(None) => panic!("same full key/timestamp across regions must fail"), + Err(error) => break error.to_string(), + } + }; + assert!( + duplicate_error.contains("timestamp") || duplicate_error.contains("time"), + "unexpected duplicate timestamp error: {duplicate_error}" + ); + assert!(!affected_key_last.is_empty()); + assert!(affected_key_last.iter().all(|is_last| !is_last)); + assert!( + duplicate_merge.try_next().await.unwrap().is_none(), + "merge error must fuse stream" + ); + for tag in ["A", "B", "C"] { + let key_rows = actual.iter().filter(|row| row.1 == tag).collect::>(); + assert_eq!( + 1, + key_rows.iter().filter(|row| row.5).count(), + "global finality for {tag}" + ); + assert!(key_rows.last().unwrap().5); + assert!(key_rows[..key_rows.len() - 1].iter().all(|row| !row.5)); + } +} + +#[tokio::test] +async fn test_two_phase_native_raw_stream_is_rejected_by_ordered_merge() { + use datatypes::series::SeriesOrdering; + use store_api::region_engine::{BoundTag, QueryScanContext, RawLabelSource}; + use store_api::region_request::RegionPutRequest; + + let mut env = TestEnv::with_prefix("test_two_phase_native_raw_merge_rejection").await; + let engine = env + .create_engine(MitoConfig { + experimental_series_scan_v2: true, + ..Default::default() + }) + .await; + let region_id = RegionId::new(1, 1); + let metadata = Arc::new(sst_region_metadata_with_encoding( + PrimaryKeyEncoding::Sparse, + )); + let mut create = CreateRequestBuilder::new().build(); + create.column_metadatas = metadata.column_metadatas.clone(); + create.primary_key = metadata.primary_key.clone(); + create + .options + .insert(PRIMARY_KEY_ENCODING.to_string(), "sparse".to_string()); + create + .options + .insert("memtable.type".to_string(), "bulk".to_string()); + create + .options + .insert("sst_format".to_string(), "flat".to_string()); + let full_schema = test_util::rows_schema(&create); + let mut key_schema = full_schema[0].clone(); + key_schema.column_name = PRIMARY_KEY_COLUMN_NAME.to_string(); + key_schema.datatype = ColumnDataType::Binary.into(); + key_schema.semantic_type = SemanticType::Tag.into(); + let row_schema = vec![key_schema, full_schema[5].clone(), full_schema[4].clone()]; + engine + .handle_request(region_id, RegionRequest::Create(create)) + .await + .unwrap(); + let sparse_rows = Rows { + schema: row_schema, + rows: vec![ + row(vec![ + ValueData::BinaryValue(new_sparse_primary_key(&["a", "x"], &metadata, 10, 0)), + ValueData::TimestampMillisecondValue(10), + ValueData::U64Value(1), + ]), + row(vec![ + ValueData::BinaryValue(new_sparse_primary_key(&["b", "y"], &metadata, 20, 0)), + ValueData::TimestampMillisecondValue(20), + ValueData::U64Value(2), + ]), + ], + }; + engine + .handle_request( + region_id, + RegionRequest::Put(RegionPutRequest { + rows: sparse_rows, + hint: Some(WriteHint { + primary_key_encoding: api::v1::PrimaryKeyEncoding::Sparse.into(), + }), + skip_wal: false, + partition_expr_version: None, + }), + ) + .await + .unwrap(); + test_util::flush_region(&engine, region_id, None).await; + let request = ScanRequest { + distribution: Some(TimeSeriesDistribution::PerSeries), + preserve_pk_dictionary_encoding: true, + ..Default::default() + }; + let selected = engine.scanner(region_id, request.clone()).await.unwrap(); + let Scanner::Series(series) = selected else { + panic!("sparse PerSeries request must use SeriesScan"); + }; + assert_eq!("two_phase", series.mode()); + let mut scanner = engine.handle_query(region_id, request).await.unwrap(); + assert_eq!(SeriesOrdering::Contiguous, scanner.raw_series_ordering()); + let metadata = scanner.metadata(); + let source = RawLabelSource::PrimaryKey { + metric_name: Arc::from("fixture_metric"), + physical_metadata: metadata.clone(), + logical_tags: metadata + .column_metadatas + .iter() + .filter(|column| { + column.semantic_type == SemanticType::Tag + && !store_api::storage::consts::is_internal_column(&column.column_schema.name) + && !store_api::metric_engine_consts::is_metric_engine_internal_column( + &column.column_schema.name, + ) + }) + .map(|column| BoundTag { + column_id: column.column_id, + name: Arc::from(column.column_schema.name.as_str()), + }) + .collect::>() + .into(), + }; + let ranges = scanner.properties().partitions.clone(); + scanner + .prepare(PrepareRequest { + ranges: Some(ranges.clone()), + target_partitions: Some(ranges.len()), + ..Default::default() + }) + .unwrap(); + let metrics = ExecutionPlanMetricsSet::default(); + let mut streams = Vec::new(); + for partition in 0..ranges.len() { + let stream = scanner + .scan_series_partition(&QueryScanContext::default(), &metrics, partition, &source) + .unwrap() + .unwrap(); + assert_eq!(scanner.raw_series_ordering(), stream.ordering()); + assert_eq!(SeriesOrdering::Contiguous, stream.ordering()); + streams.push(stream); + } + let error = match promql::source::merge_ordered_raw_series(streams) { + Ok(_) => panic!("Contiguous two-phase input cannot satisfy ordered merge ordering"), + Err(error) => error, + }; + assert!( + error.to_string().contains("ordering") || error.to_string().contains("ordered"), + "expected ordering rejection, got: {error}" + ); +} + #[tokio::test] async fn test_mito_raw_series_int64_predicate_uses_original_precision() { use datatypes::series::samples::SampleRef; @@ -5199,7 +5923,7 @@ async fn test_mito_raw_series_int64_predicate_uses_original_precision() { ); let chunks = stream.try_collect::>().await.unwrap(); for chunk in chunks { - let axis = chunk.time.as_buffer().as_ref(); + let axis = chunk.time.as_buffer(); times.extend(axis.iter().copied()); for (index, &time) in axis.iter().enumerate() { output_values.push(( @@ -5214,6 +5938,112 @@ async fn test_mito_raw_series_int64_predicate_uses_original_precision() { } assert_eq!(vec![2], times); assert_eq!(vec![(2, 20.0)], output_values); + + // Validate the typed integer gate through RegionScanExec's actual physical consumer. + use common_query::series_frame::unpack_raw_series; + use datafusion::execution::TaskContext; + use datafusion::physical_plan::ExecutionPlan; + use datatypes::series::SeriesOrdering; + use table::table::scan::RegionScanExec; + let native_request = ScanRequest { + projection: Some(vec![2, 3]), + filters: vec![ + col("tag_0").eq(lit("series")), + col("field_0").gt(lit(boundary)), + ], + distribution: Some(TimeSeriesDistribution::PerSeries), + preserve_pk_dictionary_encoding: true, + ..Default::default() + }; + let native_scanner = engine + .handle_query(region_id, native_request.clone()) + .await + .unwrap(); + let native_metadata = native_scanner.metadata(); + let native_source = RawLabelSource::PrimaryKey { + metric_name: Arc::from("fixture_metric"), + physical_metadata: native_metadata.clone(), + logical_tags: Arc::from([BoundTag { + column_id: native_metadata.primary_key[0], + name: Arc::from("tag_0"), + }]), + }; + let native_plan = Arc::new(RegionScanExec::new(native_scanner, native_request, None).unwrap()) + .try_to_series_scan(native_source.clone()) + .unwrap() + .expect("typed int64 gate is supported"); + let native_schema = native_plan.series_schema().clone(); + let frames = native_plan + .execute(0, Arc::new(TaskContext::default())) + .unwrap(); + let chunks = unpack_raw_series( + frames, + native_schema, + native_source, + SeriesOrdering::Contiguous, + ) + .unwrap() + .try_collect::>() + .await + .unwrap(); + let native_rows = chunks + .iter() + .flat_map(|chunk| { + let times = chunk.time.as_buffer(); + (0..times.len()) + .map(|i| { + let SampleRef::Float(value) = chunk.fields.get(0).unwrap().sample(i) else { + panic!("expected float output field") + }; + (times[i], value) + }) + .collect::>() + }) + .collect::>(); + assert_eq!(vec![(2, 20.0)], native_rows); + + let empty_request = ScanRequest { + projection: Some(vec![2, 3]), + filters: vec![ + col("tag_0").eq(lit("series")), + col("field_0").gt(lit(boundary + 1)), + ], + distribution: Some(TimeSeriesDistribution::PerSeries), + preserve_pk_dictionary_encoding: true, + ..Default::default() + }; + let empty_scanner = engine + .handle_query(region_id, empty_request.clone()) + .await + .unwrap(); + let empty_metadata = empty_scanner.metadata(); + let empty_source = RawLabelSource::PrimaryKey { + metric_name: Arc::from("fixture_metric"), + physical_metadata: empty_metadata.clone(), + logical_tags: Arc::from([BoundTag { + column_id: empty_metadata.primary_key[0], + name: Arc::from("tag_0"), + }]), + }; + let empty_plan = Arc::new(RegionScanExec::new(empty_scanner, empty_request, None).unwrap()) + .try_to_series_scan(empty_source.clone()) + .unwrap() + .expect("statically empty scan still has native capability"); + let empty_schema = empty_plan.series_schema().clone(); + let empty_frames = empty_plan + .execute(0, Arc::new(TaskContext::default())) + .unwrap(); + let empty_chunks = unpack_raw_series( + empty_frames, + empty_schema, + empty_source, + SeriesOrdering::Contiguous, + ) + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert!(empty_chunks.is_empty()); } /// Invalid predicate bindings decline before consumption; explicit labels decline and legacy still reads. @@ -5304,6 +6134,11 @@ async fn test_mito_raw_series_invalid_binding_before_legacy_read() { name: Arc::from("tag_0"), }]), }; + assert_eq!( + supported, + scanner.raw_series_schema(&valid_source).unwrap().is_some(), + "schema preflight case: {name}" + ); let raw = scanner .scan_series_partition(&ctx, &metrics, 0, &valid_source) .unwrap(); @@ -5313,7 +6148,7 @@ async fn test_mito_raw_series_invalid_binding_before_legacy_read() { let tuples = chunks .iter() .map(|chunk| { - let times = chunk.time.as_buffer().as_ref(); + let times = chunk.time.as_buffer(); let SampleRef::Float(value) = chunk.fields.get(0).unwrap().sample(0) else { panic!("expected raw float sample"); }; @@ -5387,6 +6222,13 @@ async fn test_mito_raw_series_invalid_binding_before_legacy_read() { .unwrap() .is_none() ); + assert!( + scanner + .raw_series_schema(&RawLabelSource::Explicit) + .unwrap() + .is_none() + ); + assert!(scanner.raw_series_schema(&valid_shape_bad_tag).is_err()); assert!( scanner .scan_series_partition(&ctx, &metrics, 0, &valid_shape_bad_tag) @@ -5427,4 +6269,70 @@ async fn test_mito_raw_series_invalid_binding_before_legacy_read() { vec![("0".to_string(), 0.0, 0), ("1".to_string(), 1.0, 1000)], rows ); + + // Runtime pruning wrappers are deliberately ineligible for native planning, and + // wrappers installed after planning must be caught at execute time. + use datafusion::execution::TaskContext; + use datafusion::physical_plan::ExecutionPlan; + use table::table::scan::RegionScanExec; + let dynamic_request = ScanRequest { + distribution: Some(TimeSeriesDistribution::PerSeries), + preserve_pk_dictionary_encoding: true, + ..Default::default() + }; + let dynamic_source_for = |scanner: &dyn RegionScanner| { + let metadata = scanner.metadata(); + RawLabelSource::PrimaryKey { + metric_name: Arc::from("fixture_metric"), + physical_metadata: metadata.clone(), + logical_tags: Arc::from([BoundTag { + column_id: metadata.primary_key[0], + name: Arc::from("tag_0"), + }]), + } + }; + let make_dynamic = || { + Arc::new(DynamicFilterPhysicalExpr::new( + vec![Arc::new(Column::new("field_0", 1))], + physical_lit(true), + )) + }; + let before = engine + .handle_query(region_id, dynamic_request.clone()) + .await + .unwrap(); + let before_source = dynamic_source_for(before.as_ref()); + let before_scan = Arc::new(RegionScanExec::new(before, dynamic_request.clone(), None).unwrap()); + assert_eq!( + vec![true], + before_scan.add_dyn_filters_to_predicate(vec![make_dynamic()]) + ); + assert!( + before_scan + .clone() + .try_to_series_scan(before_source) + .unwrap() + .is_none() + ); + + let after = engine + .handle_query(region_id, dynamic_request.clone()) + .await + .unwrap(); + let after_source = dynamic_source_for(after.as_ref()); + let after_scan = Arc::new(RegionScanExec::new(after, dynamic_request, None).unwrap()); + let after_plan = after_scan + .clone() + .try_to_series_scan(after_source) + .unwrap() + .expect("native plan before dynamic wrapper installation"); + assert_eq!( + vec![true], + after_scan.add_dyn_filters_to_predicate(vec![make_dynamic()]) + ); + assert!( + after_plan + .execute(0, Arc::new(TaskContext::default())) + .is_err() + ); } diff --git a/src/mito2/src/read/series_scan.rs b/src/mito2/src/read/series_scan.rs index eb5c1048ab5..8b4ad2acf20 100644 --- a/src/mito2/src/read/series_scan.rs +++ b/src/mito2/src/read/series_scan.rs @@ -31,13 +31,15 @@ use common_recordbatch::{RecordBatchStreamWrapper, SendableRecordBatchStream}; use common_telemetry::tracing::{self, Instrument}; use common_telemetry::{debug, warn}; use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; -use datafusion::physical_plan::{DisplayAs, DisplayFormatType}; +use datafusion::physical_plan::{DisplayAs, DisplayFormatType, PhysicalExpr}; 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::prelude::ConcreteDataType; use datatypes::schema::SchemaRef; +use datatypes::series::SeriesOrdering; use datatypes::series::samples::{FieldSpec, SeriesSchema}; use datatypes::types::TimestampType; use futures::{StreamExt, TryStreamExt}; @@ -88,6 +90,12 @@ type CandidateReceiverList = Vec>; /// Candidate assignment senders for output partitions. type CandidateSenderList = Vec>>>; +struct RawScanPreparation { + schema: Arc, + fields: Vec<(usize, ConcreteDataType)>, + predicates: Vec>, +} + struct CandidateReceiver { receiver: mpsc::UnboundedReceiver>, active_receivers: Arc, @@ -194,6 +202,189 @@ impl SeriesScan { } } + fn prepare_raw_scan( + &self, + label_source: &RawLabelSource, + ) -> Result, BoxedError> { + let input = &self.stream_ctx.input; + let metadata = input.mapper.metadata(); + match label_source { + RawLabelSource::PrimaryKey { + physical_metadata, + logical_tags, + .. + } => { + 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" + .to_string(), + } + .build(), + )); + } + let mut ids = HashSet::new(); + let mut names = HashSet::new(); + for tag in logical_tags.iter() { + let column = metadata.column_by_id(tag.column_id).filter(|column| { + column.semantic_type == SemanticType::Tag + && metadata.primary_key.contains(&tag.column_id) + }); + let valid = column.is_some_and(|column| { + column.column_schema.name == tag.name.as_ref() + && tag.name.as_ref() != "__name__" + && !is_internal_column(tag.name.as_ref()) + && !is_metric_engine_internal_column(tag.name.as_ref()) + && !ReservedColumnId::is_reserved(tag.column_id) + }); + 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(), + } + .build(), + )); + } + } + } + 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(); + let mut specs = Vec::new(); + for output in selected.column_schemas() { + let Some(column) = metadata.column_metadatas.iter().find(|column| { + column.column_schema.name == output.name + && column.semantic_type == SemanticType::Field + }) else { + continue; + }; + let Some((index, (_, data_type))) = input + .mapper + .batch_schema() + .iter() + .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!( + "selected field {} is missing from mapper batch schema", + output.name + ), + } + .build(), + )); + }; + 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())); + specs.push(FieldSpec { + name: Arc::from(output.name.as_str()), + }); + } + let time_is_millis = matches!(metadata.time_index_type(), TimestampType::Millisecond(_)); + if fields.is_empty() + || input.append_mode + || 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 { + fields: specs.into(), + eval_grid: None, + }); + let typed_schema = input.mapper.input_arrow_schema(false); + let mut predicates = Vec::new(); + if let Some(predicate) = input.predicate_group().predicate_without_region() { + if !predicate.dyn_filters().is_empty() { + debug!( + "Raw series scan preflight rejected dynamic predicate filters: region_id={}", + metadata.region_id, + ); + return Ok(None); + } + for expr in predicate.exprs() { + // Fresh execution properties cannot preserve Stable or query-dependent semantics. + if expr + .exists(|node| { + Ok(matches!(node, Expr::ScalarFunction(function) + if function.func.signature().volatility != Volatility::Immutable) + || matches!(node, Expr::ScalarVariable(..))) + }) + .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); + } + } + Ok(Some(RawScanPreparation { + schema, + fields, + predicates, + })) + } + fn supports_two_phase(input: &ScanInput) -> bool { if !is_sparse_metric_metadata(input.region_metadata()) { return false; @@ -578,6 +769,22 @@ impl RegionScanner for SeriesScan { .map_err(BoxedError::new) } + fn raw_series_ordering(&self) -> SeriesOrdering { + match self.mode { + SeriesScanMode::Legacy => SeriesOrdering::KeyThenTime, + SeriesScanMode::TwoPhase => SeriesOrdering::Contiguous, + } + } + + fn raw_series_schema( + &self, + label_source: &RawLabelSource, + ) -> Result>, BoxedError> { + Ok(self + .prepare_raw_scan(label_source)? + .map(|preparation| preparation.schema)) + } + fn scan_series_partition( &self, ctx: &QueryScanContext, @@ -585,171 +792,9 @@ impl RegionScanner for SeriesScan { partition: usize, label_source: &RawLabelSource, ) -> Result, BoxedError> { - let input = &self.stream_ctx.input; - let metadata = input.mapper.metadata(); - match label_source { - RawLabelSource::PrimaryKey { - physical_metadata, - logical_tags, - .. - } => { - 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" - .to_string(), - } - .build(), - )); - } - let mut ids = HashSet::new(); - let mut names = HashSet::new(); - for tag in logical_tags.iter() { - let column = metadata.column_by_id(tag.column_id).filter(|column| { - column.semantic_type == SemanticType::Tag - && metadata.primary_key.contains(&tag.column_id) - }); - let valid = column.is_some_and(|column| { - column.column_schema.name == tag.name.as_ref() - && tag.name.as_ref() != "__name__" - && !is_internal_column(tag.name.as_ref()) - && !is_metric_engine_internal_column(tag.name.as_ref()) - && !ReservedColumnId::is_reserved(tag.column_id) - }); - 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(), - } - .build(), - )); - } - } - } - 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(); - let mut specs = Vec::new(); - for output in selected.column_schemas() { - let Some(column) = metadata.column_metadatas.iter().find(|column| { - column.column_schema.name == output.name - && column.semantic_type == SemanticType::Field - }) else { - continue; - }; - let Some((index, (_, data_type))) = input - .mapper - .batch_schema() - .iter() - .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!( - "selected field {} is missing from mapper batch schema", - output.name - ), - } - .build(), - )); - }; - 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())); - specs.push(FieldSpec { - name: Arc::from(output.name.as_str()), - }); - } - let time_is_millis = matches!(metadata.time_index_type(), TimestampType::Millisecond(_)); - if fields.is_empty() - || input.append_mode - || 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, - ); + let Some(preparation) = self.prepare_raw_scan(label_source)? else { return Ok(None); - } - let schema = Arc::new(SeriesSchema { - fields: specs.into(), - eval_grid: None, - }); - let typed_schema = input.mapper.input_arrow_schema(false); - let mut predicates = Vec::new(); - if let Some(predicate) = input.predicate_group().predicate_without_region() { - for expr in predicate.exprs() { - // Fresh execution properties cannot preserve Stable or query-dependent semantics. - if expr - .exists(|node| { - Ok(matches!(node, Expr::ScalarFunction(function) - if function.func.signature().volatility != Volatility::Immutable) - || matches!(node, Expr::ScalarVariable(..))) - }) - .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); - } - } + }; let metrics = new_partition_metrics( &self.stream_ctx, ctx.explain_verbose, @@ -760,13 +805,15 @@ impl RegionScanner for SeriesScan { let batches = self .scan_batch_in_partition(ctx, partition, metrics.clone(), metrics_set) .map_err(BoxedError::new)?; + let ordering = self.raw_series_ordering(); Ok(Some(raw::make_stream( batches, - schema, + preparation.schema, label_source.clone(), - fields, + preparation.fields, metrics, - predicates, + preparation.predicates, + ordering, ))) } diff --git a/src/mito2/src/read/series_scan/raw.rs b/src/mito2/src/read/series_scan/raw.rs index 6ed41bf8c14..4adac4263ae 100644 --- a/src/mito2/src/read/series_scan/raw.rs +++ b/src/mito2/src/read/series_scan/raw.rs @@ -128,6 +128,7 @@ fn chunk( struct RawStream { schema: Arc, label_source: RawLabelSource, + ordering: SeriesOrdering, inner: BoxStream<'static, Result>, } @@ -146,7 +147,7 @@ impl RawSeriesStream for RawStream { &self.label_source } fn ordering(&self) -> SeriesOrdering { - SeriesOrdering::Contiguous + self.ordering } } @@ -157,15 +158,16 @@ pub(crate) fn make_stream( fields: Vec<(usize, ConcreteDataType)>, metrics: PartitionMetrics, predicates: Vec>, + ordering: SeriesOrdering, ) -> SendableRawSeriesStream { let output_schema = schema.clone(); let stream = async_stream::stream! { let mut input = inner; let mut pending: Option = None; while let Some(item) = input.next().await { - let (records, is_legacy) = match item { - Ok(ScanBatch::Series(SeriesBatch::Flat(flat))) => (flat.batches.into_vec(), true), - Ok(ScanBatch::RecordBatch(batch)) => (vec![batch], false), + let records = match item { + Ok(ScanBatch::Series(SeriesBatch::Flat(flat))) => flat.batches.into_vec(), + Ok(ScanBatch::RecordBatch(batch)) => vec![batch], Err(error) => { yield Err(BoxedError::new(error)); return; } }; for record in records { @@ -238,7 +240,7 @@ pub(crate) fn make_stream( }; metrics.inc_convert_batch_cost(begun.elapsed()); if let Some(old) = pending.as_ref() - && is_legacy + && ordering == SeriesOrdering::KeyThenTime && old.primary_key.as_slice() != current.primary_key.as_slice() && current.primary_key.as_slice() <= old.primary_key.as_slice() { @@ -265,6 +267,7 @@ pub(crate) fn make_stream( Box::pin(RawStream { schema, label_source, + ordering, inner: Box::pin(stream), }) } diff --git a/src/mito2/src/read/series_scan/raw/tests.rs b/src/mito2/src/read/series_scan/raw/tests.rs index 22dd3eefa15..1fdd703b00d 100644 --- a/src/mito2/src/read/series_scan/raw/tests.rs +++ b/src/mito2/src/read/series_scan/raw/tests.rs @@ -37,6 +37,7 @@ use datatypes::arrow::record_batch::RecordBatch; use datatypes::data_type::DataType as _; use datatypes::prelude::ConcreteDataType; use datatypes::schema::ColumnSchema; +use datatypes::series::SeriesOrdering; use datatypes::series::raw::RawSeriesChunk; use datatypes::series::samples::{FieldSpec, SampleRef, SeriesSchema}; use futures::{Stream, StreamExt, stream}; @@ -239,6 +240,7 @@ fn convert_with_source( schema: Arc, source: RawLabelSource, field_type: ConcreteDataType, + ordering: SeriesOrdering, ) -> SendableRawSeriesStream { make_stream( inner, @@ -247,6 +249,7 @@ fn convert_with_source( vec![(0, field_type)], metrics(), vec![], + ordering, ) } @@ -267,6 +270,17 @@ fn convert(inner: ScanBatchStream, schema: Arc) -> SendableRawSeri schema, primary_key_source(metadata()), ConcreteDataType::int64_datatype(), + SeriesOrdering::Contiguous, + ) +} + +fn convert_ordered(inner: ScanBatchStream, schema: Arc) -> SendableRawSeriesStream { + convert_with_source( + inner, + schema, + primary_key_source(metadata()), + ConcreteDataType::int64_datatype(), + SeriesOrdering::KeyThenTime, ) } @@ -370,6 +384,59 @@ async fn chunks_preserve_keys_values_times_and_final_flags_for_both_input_shapes } } +#[tokio::test] +async fn raw_stream_reports_and_enforces_declared_ordering() { + let stream = convert(input(vec![]), schema()); + assert_eq!(stream.ordering(), SeriesOrdering::Contiguous); + drop(stream); + let stream = convert_ordered(input(vec![]), schema()); + assert_eq!(stream.ordering(), SeriesOrdering::KeyThenTime); + drop(stream); + + let mut output = convert_ordered( + input(vec![ + Ok(ScanBatch::RecordBatch(dictionary_batch( + &[b"z"], + &[1], + &[Some(1)], + ))), + Ok(ScanBatch::RecordBatch(dictionary_batch( + &[b"a"], + &[1], + &[Some(2)], + ))), + ]), + schema(), + ); + assert!(output.next().await.unwrap().is_err()); + assert!(output.next().await.is_none()); + + let mut output = convert( + input(vec![ + Ok(ScanBatch::RecordBatch(dictionary_batch( + &[b"z"], + &[1], + &[Some(1)], + ))), + Ok(ScanBatch::RecordBatch(dictionary_batch( + &[b"a"], + &[1], + &[Some(2)], + ))), + ]), + schema(), + ); + assert_eq!( + output.next().await.unwrap().unwrap().primary_key.as_slice(), + b"z" + ); + assert_eq!( + output.next().await.unwrap().unwrap().primary_key.as_slice(), + b"a" + ); + assert!(output.next().await.is_none()); +} + #[tokio::test] async fn typed_predicate_filters_before_numeric_samples_are_cast() { let threshold = 1_i64 << 53; @@ -395,6 +462,7 @@ async fn typed_predicate_filters_before_numeric_samples_are_cast() { vec![(0, ConcreteDataType::int64_datatype())], metrics(), vec![predicate], + SeriesOrdering::Contiguous, ); let chunk = output.next().await.unwrap().unwrap(); assert!(chunk.is_last); @@ -424,6 +492,7 @@ async fn predicate_error_is_terminal() { vec![(0, ConcreteDataType::int64_datatype())], metrics(), vec![predicate], + SeriesOrdering::Contiguous, ); assert!(output.next().await.unwrap().is_err()); assert!(output.next().await.is_none()); @@ -515,7 +584,7 @@ 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( + let mut output = convert_ordered( input(vec![Ok(ScanBatch::Series(SeriesBatch::Flat( FlatSeriesBatch { batches }, )))]), @@ -528,7 +597,7 @@ async fn legacy_series_requires_increasing_keys_but_record_batches_do_not() { 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( + let mut output = convert_ordered( input(vec![Ok(ScanBatch::Series(SeriesBatch::Flat( FlatSeriesBatch { batches }, )))]), @@ -545,7 +614,7 @@ async fn legacy_series_requires_increasing_keys_but_record_batches_do_not() { 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( + let mut output = convert_ordered( input(vec![Ok(ScanBatch::Series(SeriesBatch::Flat( FlatSeriesBatch { batches }, )))]), @@ -636,6 +705,7 @@ async fn native_histogram_adapter_preserves_complete_payload() { native_histogram_value_type().clone(), )), native_histogram_value_type().clone(), + SeriesOrdering::Contiguous, ); assert!(matches!( raw.label_source(), From 8e90db600d10ee226b42d66c64856766ddb259d2 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Tue, 20 Oct 2026 20:27:44 +0800 Subject: [PATCH 2/4] test(mito2): cover raw partition bounds and contiguity across assignments The raw-series path had two untested behaviours. An out-of-range partition must fail without consuming the scanner's partition state: the range check runs before the receiver is taken, so a later in-bounds scan on the same scanner still reads its complete data, and a repeat scan of that partition then reports the one-shot violation. A two-phase stream must stay contiguous in the sense the raw contract advertises, including across candidate assignments: chunks are checked per partition for a non-empty, disjoint key set, the whole stream is routed through the frame codec so its own sequence validator runs, and a series may appear in only one run so a non-adjacent recurrence would fail. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/mito2/src/engine/scan_test.rs | 501 ++++++++++++++++++++++++++++++ 1 file changed, 501 insertions(+) diff --git a/src/mito2/src/engine/scan_test.rs b/src/mito2/src/engine/scan_test.rs index 3ad13fac11e..6fe07d63ecd 100644 --- a/src/mito2/src/engine/scan_test.rs +++ b/src/mito2/src/engine/scan_test.rs @@ -5826,6 +5826,507 @@ async fn test_two_phase_native_raw_stream_is_rejected_by_ordered_merge() { ); } +#[tokio::test] +async fn test_two_phase_raw_partition_out_of_range_keeps_scanner_usable() { + use datatypes::series::samples::SampleRef; + use store_api::region_engine::QueryScanContext; + + let mut env = TestEnv::with_prefix("test_two_phase_raw_partition_out_of_range").await; + let engine = env + .create_engine(MitoConfig { + experimental_series_scan_v2: true, + ..Default::default() + }) + .await; + let region_id = RegionId::new(1, 1); + let (metadata, row_schema) = create_two_phase_sparse_region(&engine, region_id).await; + let expected = [ + (new_sparse_primary_key(&["a", "x"], &metadata, 10, 0), 10, 1), + (new_sparse_primary_key(&["b", "x"], &metadata, 10, 0), 20, 2), + ]; + put_sparse_rows( + &engine, + region_id, + &row_schema, + expected + .iter() + .map(|(key, time, value)| (key.clone(), *time, *value)) + .collect(), + ) + .await; + test_util::flush_region(&engine, region_id, None).await; + + let request = ScanRequest { + distribution: Some(TimeSeriesDistribution::PerSeries), + preserve_pk_dictionary_encoding: true, + ..Default::default() + }; + let Scanner::Series(selected) = engine.scanner(region_id, request.clone()).await.unwrap() + else { + panic!("sparse PerSeries request must use SeriesScan"); + }; + assert_eq!("two_phase", selected.mode()); + let mut scanner = engine.handle_query(region_id, request).await.unwrap(); + let source = primary_key_label_source(&scanner.metadata()); + let ranges = scanner.properties().partitions.clone(); + scanner + .prepare(PrepareRequest { + ranges: Some(ranges.clone()), + target_partitions: Some(ranges.len()), + ..Default::default() + }) + .unwrap(); + let partitions = scanner.properties().num_partitions(); + assert!(partitions > 0, "prepared scan needs an output partition"); + let metrics = ExecutionPlanMetricsSet::default(); + let ctx = QueryScanContext::default(); + + // `scan_batch_in_partition` validates the index before it takes that partition's receiver, so + // the first index past the last partition must fail without disturbing scanner state. + let error = match scanner.scan_series_partition(&ctx, &metrics, partitions, &source) { + Ok(_) => { + panic!("partition {partitions} must be rejected by a {partitions}-partitions scan") + } + Err(error) => error, + }; + let message = error.to_string(); + assert!( + message.contains("out of range") && message.contains(&partitions.to_string()), + "expected a partition range error, got: {message}" + ); + + // The in-range partition must still read the complete data set. + let stream = scanner + .scan_series_partition(&ctx, &metrics, 0, &source) + .unwrap() + .expect("the raw series scan stays available after a rejected partition"); + let mut actual = HashMap::new(); + for chunk in stream.try_collect::>().await.unwrap() { + let times = chunk.time.as_buffer(); + for (index, time) in times.iter().enumerate() { + let SampleRef::Float(value) = chunk.fields.get(0).unwrap().sample(index) else { + panic!("raw fields are converted to float samples"); + }; + actual.insert((chunk.primary_key.to_vec(), *time), value); + } + } + assert_eq!( + expected.len(), + actual.len(), + "the rejected call must not consume rows" + ); + for (key, time, value) in expected { + assert_eq!(Some(&(value as f64)), actual.get(&(key, time))); + } + + // A partition can only be scanned once, so this proves the rejected index never took the + // in-range partition's receiver: the error above is raised by the range check alone. + let repeated = match scanner.scan_series_partition(&ctx, &metrics, 0, &source) { + Ok(_) => panic!("partition 0 was already scanned"), + Err(error) => error.to_string(), + }; + assert!( + repeated.contains("scan multiple times"), + "expected a repeated-scan error, got: {repeated}" + ); +} + +#[tokio::test] +async fn test_two_phase_raw_stream_stays_contiguous_across_assignments() { + use common_query::series_frame::{pack_raw_series, unpack_raw_series}; + use datatypes::series::SeriesOrdering; + use datatypes::series::raw::RawSeriesChunk; + use datatypes::series::samples::{SampleRef, SeriesSchema}; + use store_api::region_engine::QueryScanContext; + + let mut env = TestEnv::with_prefix("test_two_phase_raw_contiguous_assignments").await; + let engine = env + .create_engine(MitoConfig { + experimental_series_scan_v2: true, + ..Default::default() + }) + .await; + let region_id = RegionId::new(1, 1); + let (metadata, row_schema) = create_two_phase_sparse_region(&engine, region_id).await; + + // A sparse primary key carries the TSID, and `SeriesRange::partition_for` splits the u64 TSID + // domain over the output partitions, so the fixture splits into two disjoint groups of series: + // tsids below 2^63 own partition 0 and tsids at or above 2^63 own partition 1. + let series = [ + ("a0", 0_u64), + ("a1", 1), + ("b0", 1_u64 << 63), + ("b1", (1_u64 << 63) + 1), + ]; + let keys = series + .iter() + .map(|(tag, tsid)| new_sparse_primary_key(&[tag, "x"], &metadata, 10, *tsid)) + .collect::>(); + + // Three write rounds put every series into two SST ranges plus one memtable range, and the + // large rounds make a series span more than one read batch, so one series is emitted as + // multiple chunks instead of exactly one. + const LARGE: i64 = 20_000; + const SMALL: i64 = 3; + let rounds: [&[(usize, i64, i64)]; 3] = [ + &[(0, 0, LARGE), (1, 0, SMALL), (2, 0, SMALL), (3, 0, LARGE)], + &[ + (0, LARGE, 5), + (1, SMALL, SMALL), + (2, SMALL, SMALL), + (3, LARGE, 5), + ], + &[ + (0, LARGE + 5, 2), + (1, 2 * SMALL, 2), + (2, 2 * SMALL, 2), + (3, LARGE + 5, 2), + ], + ]; + let mut expected_times = HashMap::, Vec>::new(); + for spec in rounds { + for &(index, first, count) in spec { + expected_times + .entry(keys[index].clone()) + .or_default() + .extend(first..first + count); + } + } + for (round, spec) in rounds.iter().enumerate() { + let rows = spec + .iter() + .flat_map(|&(index, first, count)| { + let key = keys[index].clone(); + (0..count).map(move |offset| { + let time = first + offset; + (key.clone(), time, (time + 1) as u64) + }) + }) + .collect::>(); + put_sparse_rows(&engine, region_id, &row_schema, rows).await; + if round + 1 < rounds.len() { + test_util::flush_region(&engine, region_id, None).await; + } + } + + let request = ScanRequest { + distribution: Some(TimeSeriesDistribution::PerSeries), + preserve_pk_dictionary_encoding: true, + ..Default::default() + }; + let mut scanner = engine.handle_query(region_id, request).await.unwrap(); + assert_eq!(SeriesOrdering::Contiguous, scanner.raw_series_ordering()); + let source = primary_key_label_source(&scanner.metadata()); + let ranges = scanner + .properties() + .partitions + .clone() + .into_iter() + .flatten() + .collect::>(); + assert!(!ranges.is_empty(), "region scan needs an input range"); + // Two output partitions make the candidate distributor create one reader assignment per + // partition, so the scan runs more than one assignment. + let partition_count = 2; + let mut partitions = vec![Vec::new(); partition_count]; + for (index, range) in ranges.into_iter().enumerate() { + partitions[index % partition_count].push(range); + } + scanner + .prepare(PrepareRequest { + ranges: Some(partitions), + target_partitions: Some(partition_count), + ..Default::default() + }) + .unwrap(); + assert_eq!(partition_count, scanner.properties().num_partitions()); + + let metrics = ExecutionPlanMetricsSet::default(); + let ctx = QueryScanContext::default(); + let mut per_partition = Vec::with_capacity(partition_count); + let mut schema: Option> = None; + for partition in 0..partition_count { + let stream = scanner + .scan_series_partition(&ctx, &metrics, partition, &source) + .unwrap() + .expect("two-phase raw scan is supported"); + assert_eq!(SeriesOrdering::Contiguous, stream.ordering()); + assert_eq!( + vec!["field_0"], + stream + .schema() + .fields + .iter() + .map(|field| field.name.as_ref()) + .collect::>() + ); + schema.get_or_insert_with(|| Arc::clone(stream.schema())); + let chunks = stream.try_collect::>().await.unwrap(); + assert!( + !chunks.is_empty(), + "partition {partition} must own a reader assignment" + ); + // Each partition owns exactly the series the TSID split assigned to it, so the two + // assignment streams are non-empty and disjoint. + let mut observed = chunks + .iter() + .map(|chunk| chunk.primary_key.to_vec()) + .collect::>(); + observed.sort(); + observed.dedup(); + let mut assigned = vec![keys[2 * partition].clone(), keys[2 * partition + 1].clone()]; + assigned.sort(); + assert_eq!( + assigned, observed, + "partition {partition} must own exactly its assigned series" + ); + per_partition.push(chunks); + } + + // The complete raw series stream of the multi-assignment scan: every assignment's chunks in + // partition order. + let complete = per_partition + .into_iter() + .flatten() + .collect::>(); + let schema = schema.unwrap(); + let complete_rows = complete.iter().map(|chunk| chunk.time.len()).sum::(); + assert_eq!( + expected_times.values().map(Vec::len).sum::(), + complete_rows, + "the complete stream must carry every written row" + ); + + // Reuse the frame adapter's cross-chunk validator (private to + // `common_query::series_frame`) by packing and unpacking the complete stream. It rejects a key + // change that was not preceded by that key's final chunk, a key that recurs right after its + // final chunk, timestamps that do not increase across chunks, and a stream whose last chunk is + // not final. + let packed = pack_raw_series(Box::pin(ChainedRawSeriesStream::new( + Arc::clone(&schema), + source.clone(), + SeriesOrdering::Contiguous, + complete.clone(), + ))) + .unwrap(); + let unpacked = unpack_raw_series( + packed, + Arc::clone(&schema), + source, + SeriesOrdering::Contiguous, + ) + .unwrap(); + let validated = unpacked.try_collect::>().await.unwrap(); + assert_eq!( + complete_rows, + validated + .iter() + .map(|chunk| chunk.time.len()) + .sum::(), + "packing and unpacking must preserve the complete stream" + ); + + // `SequenceValidator` checks adjacent chunks only, so also reject non-adjacent recurrence: a + // series interrupted by another series and resumed later would break `Contiguous`. + let mut runs = Vec::>::new(); + let mut chunks_per_series = HashMap::, usize>::new(); + for chunk in &complete { + let key = chunk.primary_key.to_vec(); + *chunks_per_series.entry(key.clone()).or_default() += 1; + if runs.last() != Some(&key) { + assert!( + !runs.contains(&key), + "a series' chunks must stay contiguous across assignments" + ); + runs.push(key); + } + } + assert_eq!(keys.len(), runs.len(), "every series appears exactly once"); + assert!( + chunks_per_series.values().copied().max().unwrap() > 1, + "the fixture must emit multiple chunks for one series, got {chunks_per_series:?}" + ); + assert_eq!(keys.len(), chunks_per_series.len()); + for (key, times) in &expected_times { + let observed = complete + .iter() + .filter(|chunk| chunk.primary_key.as_slice() == key.as_slice()) + .flat_map(|chunk| chunk.time.as_buffer().iter().copied()) + .collect::>(); + assert_eq!( + times, &observed, + "sample positions must increase per series" + ); + } + for (index, chunk) in complete.iter().enumerate() { + let next_is_same_series = complete + .get(index + 1) + .is_some_and(|next| next.primary_key.as_slice() == chunk.primary_key.as_slice()); + assert_eq!( + !next_is_same_series, chunk.is_last, + "is_last must mark the final chunk of a series" + ); + let times = chunk.time.as_buffer(); + for (sample, time) in times.iter().enumerate() { + let SampleRef::Float(value) = chunk.fields.get(0).unwrap().sample(sample) else { + panic!("raw fields are converted to float samples"); + }; + assert_eq!((*time + 1) as f64, value); + } + } +} + +/// Creates the sparse-encoded flat region shared by the two-phase raw scan tests. +/// +/// Returns the region metadata and the row schema accepted by the sparse write path. +async fn create_two_phase_sparse_region( + engine: &crate::engine::MitoEngine, + region_id: RegionId, +) -> ( + Arc, + Vec, +) { + let metadata = Arc::new(sst_region_metadata_with_encoding( + PrimaryKeyEncoding::Sparse, + )); + let mut create = CreateRequestBuilder::new().build(); + create.column_metadatas = metadata.column_metadatas.clone(); + create.primary_key = metadata.primary_key.clone(); + create + .options + .insert(PRIMARY_KEY_ENCODING.to_string(), "sparse".to_string()); + create + .options + .insert("memtable.type".to_string(), "bulk".to_string()); + create + .options + .insert("sst_format".to_string(), "flat".to_string()); + let full_schema = test_util::rows_schema(&create); + let mut key_schema = full_schema[0].clone(); + key_schema.column_name = PRIMARY_KEY_COLUMN_NAME.to_string(); + key_schema.datatype = ColumnDataType::Binary.into(); + key_schema.semantic_type = SemanticType::Tag.into(); + let row_schema = vec![key_schema, full_schema[5].clone(), full_schema[4].clone()]; + engine + .handle_request(region_id, RegionRequest::Create(create)) + .await + .unwrap(); + (metadata, row_schema) +} + +/// Writes `(primary key, timestamp, field_0)` rows through the sparse write path. +async fn put_sparse_rows( + engine: &crate::engine::MitoEngine, + region_id: RegionId, + schema: &[api::v1::ColumnSchema], + rows: Vec<(Vec, i64, u64)>, +) { + let rows = rows + .into_iter() + .map(|(key, time, value)| { + row(vec![ + ValueData::BinaryValue(key), + ValueData::TimestampMillisecondValue(time), + ValueData::U64Value(value), + ]) + }) + .collect(); + engine + .handle_request( + region_id, + RegionRequest::Put(RegionPutRequest { + rows: Rows { + schema: schema.to_vec(), + rows, + }, + hint: Some(WriteHint { + primary_key_encoding: api::v1::PrimaryKeyEncoding::Sparse.into(), + }), + skip_wal: false, + partition_expr_version: None, + }), + ) + .await + .unwrap(); +} + +/// Binds every non-internal tag column of `metadata` as the raw label source. +fn primary_key_label_source( + metadata: &store_api::metadata::RegionMetadataRef, +) -> store_api::region_engine::RawLabelSource { + use store_api::region_engine::{BoundTag, RawLabelSource}; + + RawLabelSource::PrimaryKey { + metric_name: Arc::from("fixture_metric"), + physical_metadata: Arc::clone(metadata), + logical_tags: metadata + .column_metadatas + .iter() + .filter(|column| { + column.semantic_type == SemanticType::Tag + && !store_api::storage::consts::is_internal_column(&column.column_schema.name) + && !store_api::metric_engine_consts::is_metric_engine_internal_column( + &column.column_schema.name, + ) + }) + .map(|column| BoundTag { + column_id: column.column_id, + name: Arc::from(column.column_schema.name.as_str()), + }) + .collect::>() + .into(), + } +} + +/// Concatenates the raw chunks of every reader assignment of one scan, in partition order. +struct ChainedRawSeriesStream { + schema: Arc, + label_source: store_api::region_engine::RawLabelSource, + ordering: datatypes::series::SeriesOrdering, + chunks: std::vec::IntoIter, +} + +impl ChainedRawSeriesStream { + fn new( + schema: Arc, + label_source: store_api::region_engine::RawLabelSource, + ordering: datatypes::series::SeriesOrdering, + chunks: Vec, + ) -> Self { + Self { + schema, + label_source, + ordering, + chunks: chunks.into_iter(), + } + } +} + +impl futures::Stream for ChainedRawSeriesStream { + type Item = Result; + + fn poll_next( + mut self: std::pin::Pin<&mut Self>, + _cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + std::task::Poll::Ready(self.chunks.next().map(Ok)) + } +} + +impl store_api::region_engine::RawSeriesStream for ChainedRawSeriesStream { + fn schema(&self) -> &Arc { + &self.schema + } + + fn label_source(&self) -> &store_api::region_engine::RawLabelSource { + &self.label_source + } + + fn ordering(&self) -> datatypes::series::SeriesOrdering { + self.ordering + } +} + #[tokio::test] async fn test_mito_raw_series_int64_predicate_uses_original_precision() { use datatypes::series::samples::SampleRef; From 38f728ee3fa97e4405a1a4466022f7e63fe04b43 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Tue, 20 Oct 2026 21:42:34 +0800 Subject: [PATCH 3/4] test(mito2): cover a second assignment reaching the same partition The cross-assignment test wrote four candidate series, gave each of two output partitions one assignment, and concatenated the two streams. It proved that two disjoint non-empty assignments stay contiguous once joined, but production only sends a second round to the *same* receiver when one collection round holds `CANDIDATE_SERIES_ASSIGNMENT_THRESHOLD` (1,000,000) candidates, so that boundary was never crossed. The new test lowers the threshold through a `#[cfg(test)]` override instead of writing a million series, then writes 501 series -- one candidate batch plus a trailing remainder -- into a single output partition and asserts on that partition's one stream: two or more assignment rounds were sent, every series appears exactly once in key order with no recurrence, timestamps increase within a key, `is_last` is correct at the boundary, and every value matches. The stream then goes through `pack_raw_series`/`unpack_raw_series` so `SequenceValidator` runs on it as well. Dropping a later round leaves a validator-clean stream, so the data assertions are what catch it. The layout assertions now check the composition rather than non-emptiness: the scan input reports two SST sources whose row counts are the two flushed write rounds and one memtable whose rows are the unflushed round. Compaction is disabled in both tests so a background merge cannot move that composition. The older test is renamed to say what it covers: two disjoint partitions, one assignment each, concatenated. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/mito2/src/engine/scan_test.rs | 369 +++++++++++++++++++++++++++++- src/mito2/src/read/series_scan.rs | 48 +++- 2 files changed, 408 insertions(+), 9 deletions(-) diff --git a/src/mito2/src/engine/scan_test.rs b/src/mito2/src/engine/scan_test.rs index 6fe07d63ecd..4a720ef73f2 100644 --- a/src/mito2/src/engine/scan_test.rs +++ b/src/mito2/src/engine/scan_test.rs @@ -5932,7 +5932,7 @@ async fn test_two_phase_raw_partition_out_of_range_keeps_scanner_usable() { } #[tokio::test] -async fn test_two_phase_raw_stream_stays_contiguous_across_assignments() { +async fn test_two_phase_raw_stream_stays_contiguous_across_disjoint_partition_assignments() { use common_query::series_frame::{pack_raw_series, unpack_raw_series}; use datatypes::series::SeriesOrdering; use datatypes::series::raw::RawSeriesChunk; @@ -5943,6 +5943,10 @@ async fn test_two_phase_raw_stream_stays_contiguous_across_assignments() { let engine = env .create_engine(MitoConfig { experimental_series_scan_v2: true, + // Automatic compaction could merge the two flushed SSTs before the scan pins its + // input and make the source-composition assertions below racy. + schedule_compaction_after_edit: false, + min_compaction_interval: Duration::from_secs(60 * 60), ..Default::default() }) .await; @@ -5963,9 +5967,11 @@ async fn test_two_phase_raw_stream_stays_contiguous_across_assignments() { .map(|(tag, tsid)| new_sparse_primary_key(&[tag, "x"], &metadata, 10, *tsid)) .collect::>(); - // Three write rounds put every series into two SST ranges plus one memtable range, and the - // large rounds make a series span more than one read batch, so one series is emitted as - // multiple chunks instead of exactly one. + // Three write rounds put every series into two SST sources (the first two rounds are flushed) + // plus one memtable source (the last round stays in memory), and the large rounds make a + // series span more than one read batch, so one series is emitted as multiple chunks instead + // of exactly one. The scan input assertions below check that composition instead of only + // checking that something was read. const LARGE: i64 = 20_000; const SMALL: i64 = 3; let rounds: [&[(usize, i64, i64)]; 3] = [ @@ -5992,6 +5998,7 @@ async fn test_two_phase_raw_stream_stays_contiguous_across_assignments() { .extend(first..first + count); } } + let mut round_rows = Vec::with_capacity(rounds.len()); for (round, spec) in rounds.iter().enumerate() { let rows = spec .iter() @@ -6003,6 +6010,7 @@ async fn test_two_phase_raw_stream_stays_contiguous_across_assignments() { }) }) .collect::>(); + round_rows.push(rows.len()); put_sparse_rows(&engine, region_id, &row_schema, rows).await; if round + 1 < rounds.len() { test_util::flush_region(&engine, region_id, None).await; @@ -6014,8 +6022,50 @@ async fn test_two_phase_raw_stream_stays_contiguous_across_assignments() { preserve_pk_dictionary_encoding: true, ..Default::default() }; - let mut scanner = engine.handle_query(region_id, request).await.unwrap(); + let selected = engine.scanner(region_id, request).await.unwrap(); + let Scanner::Series(mut scanner) = selected else { + panic!("sparse PerSeries request must use SeriesScan"); + }; + assert_eq!("two_phase", scanner.mode()); assert_eq!(SeriesOrdering::Contiguous, scanner.raw_series_ordering()); + // Composition of the pinned scan input: the two flushes pin two SST sources and the third + // write round stays in one memtable source. Each source reports the rows its write round + // wrote, so this asserts the layout instead of only asserting that something was read. + { + let input = scanner.input(); + let mut file_rows = input + .files + .iter() + .map(|file| file.num_rows()) + .collect::>(); + file_rows.sort_unstable(); + let mut expected_file_rows = round_rows[..2].to_vec(); + expected_file_rows.sort_unstable(); + assert_eq!( + expected_file_rows, file_rows, + "each flushed write round must be one SST source" + ); + assert_eq!(2, input.num_files()); + assert_eq!( + 1, + input.num_memtables(), + "the last round is one memtable range" + ); + assert_eq!( + round_rows[2], + input + .memtables + .iter() + .map(|memtable| memtable.stats().num_rows) + .sum::(), + "the unflushed round must be the memtable's contribution" + ); + assert_eq!( + round_rows.iter().sum::(), + input.total_rows(), + "the sources must add up to every written row" + ); + } let source = primary_key_label_source(&scanner.metadata()); let ranges = scanner .properties() @@ -6083,8 +6133,9 @@ async fn test_two_phase_raw_stream_stays_contiguous_across_assignments() { per_partition.push(chunks); } - // The complete raw series stream of the multi-assignment scan: every assignment's chunks in - // partition order. + // The complete raw series stream of the two-partition scan: each partition's one assignment, + // concatenated in partition order. `test_two_phase_raw_same_partition_multi_round_assignment_ + // stays_contiguous` covers two assignments handed to the same partition. let complete = per_partition .into_iter() .flatten() @@ -6136,7 +6187,7 @@ async fn test_two_phase_raw_stream_stays_contiguous_across_assignments() { if runs.last() != Some(&key) { assert!( !runs.contains(&key), - "a series' chunks must stay contiguous across assignments" + "a series' chunks must stay contiguous across the concatenated assignments" ); runs.push(key); } @@ -6176,6 +6227,308 @@ async fn test_two_phase_raw_stream_stays_contiguous_across_assignments() { } } +/// Lowers the candidate-series assignment threshold and restores it when dropped. +/// +/// The override is process-global, so it is restored from `Drop` to keep a failing assertion or a +/// panic from leaking a weakened threshold into later tests of the same binary. `cargo nextest` +/// (the documented runner for this crate) additionally runs every test in its own process. +struct AssignmentThresholdScope; + +impl AssignmentThresholdScope { + fn lower_to(threshold: usize) -> Self { + crate::read::series_scan::set_candidate_series_assignment_threshold(threshold); + Self + } +} + +impl Drop for AssignmentThresholdScope { + fn drop(&mut self) { + crate::read::series_scan::set_candidate_series_assignment_threshold(0); + } +} + +/// One partition receiving two candidate assignments must still emit one contiguous stream. +/// +/// Without the threshold override a second assignment reaches the same partition only when a +/// single collection round holds `CANDIDATE_SERIES_ASSIGNMENT_THRESHOLD` (1,000,000) candidate +/// series, so this test lowers the threshold instead of writing a million series. +#[tokio::test] +async fn test_two_phase_raw_same_partition_multi_round_assignment_stays_contiguous() { + use common_query::series_frame::{pack_raw_series, unpack_raw_series}; + use datatypes::series::SeriesOrdering; + use datatypes::series::samples::SampleRef; + use store_api::region_engine::QueryScanContext; + + // The candidate distributor starts a new assignment as soon as the collected series reach the + // threshold, and the candidate stream hands series over in globally sorted batches of at most + // `METRIC_SERIES_ID_BATCH_SIZE` (500) entries, yielding a trailing batch for the remainder. + // With the threshold lowered to 1 every one of those batches becomes one assignment round, so + // 501 series produce exactly two non-empty rounds: the first 500 sorted series, then the last + // one. Both rounds are addressed to the single output partition, i.e. to one receiver. + const CANDIDATE_BATCH: usize = 500; + const SERIES: usize = CANDIDATE_BATCH + 1; + let _scope = AssignmentThresholdScope::lower_to(1); + + let mut env = TestEnv::with_prefix("test_two_phase_raw_same_partition_multi_round").await; + let engine = env + .create_engine(MitoConfig { + experimental_series_scan_v2: true, + // The source-composition assertions below pin the flushed SST until the scan opens. + schedule_compaction_after_edit: false, + min_compaction_interval: Duration::from_secs(60 * 60), + ..Default::default() + }) + .await; + let region_id = RegionId::new(1, 1); + let (metadata, row_schema) = create_two_phase_sparse_region(&engine, region_id).await; + + let tags = (0..SERIES) + .map(|index| format!("s{index}")) + .collect::>(); + let keys = tags + .iter() + .enumerate() + .map(|(index, tag)| { + new_sparse_primary_key(&[tag.as_str(), "x"], &metadata, 10, index as u64) + }) + .collect::>(); + // The round split below relies on the candidate stream ordering series by primary key, which + // is the ascending TSID order for this fixture. + let mut sorted_keys = keys.clone(); + sorted_keys.sort(); + assert_eq!(keys, sorted_keys, "primary keys must sort in TSID order"); + + // One flushed SST source plus one memtable source. The wide series spans more than one read + // batch, so a series is emitted as several chunks and the round boundary sits next to chunks + // that must stay grouped. + const WIDE: i64 = 20_000; + const SMALL: i64 = 3; + const SECOND_ROUND_FIRST: i64 = 100_000; + let mut expected_times = HashMap::, Vec>::new(); + expected_times.insert( + keys[0].clone(), + (0..WIDE) + .chain(SECOND_ROUND_FIRST..SECOND_ROUND_FIRST + 3) + .collect(), + ); + for key in &keys[1..] { + expected_times.insert( + key.clone(), + (0..SMALL) + .chain(SECOND_ROUND_FIRST..SECOND_ROUND_FIRST + 3) + .collect(), + ); + } + + let first_round = std::iter::once((keys[0].clone(), 0_i64, WIDE)) + .chain(keys[1..].iter().map(|key| (key.clone(), 0_i64, SMALL))) + .flat_map(|(key, first, count)| { + (0..count).map(move |offset| { + let time = first + offset; + (key.clone(), time, (time + 1) as u64) + }) + }) + .collect::>(); + let second_round = keys + .iter() + .flat_map(|key| { + (0..SMALL).map(move |offset| { + let time = SECOND_ROUND_FIRST + offset; + (key.clone(), time, (time + 1) as u64) + }) + }) + .collect::>(); + let (sst_rows, memtable_rows) = (first_round.len(), second_round.len()); + put_sparse_rows(&engine, region_id, &row_schema, first_round).await; + test_util::flush_region(&engine, region_id, None).await; + put_sparse_rows(&engine, region_id, &row_schema, second_round).await; + + let request = ScanRequest { + distribution: Some(TimeSeriesDistribution::PerSeries), + preserve_pk_dictionary_encoding: true, + ..Default::default() + }; + let selected = engine.scanner(region_id, request).await.unwrap(); + let Scanner::Series(mut scanner) = selected else { + panic!("sparse PerSeries request must use SeriesScan"); + }; + assert_eq!("two_phase", scanner.mode()); + assert_eq!(SeriesOrdering::Contiguous, scanner.raw_series_ordering()); + // The fixture must really be one SST source plus one memtable source, or the two rounds below + // would be explained by something else. + { + let input = scanner.input(); + assert_eq!(1, input.num_files()); + assert_eq!( + vec![sst_rows], + input + .files + .iter() + .map(|file| file.num_rows()) + .collect::>() + ); + assert_eq!(1, input.num_memtables()); + assert_eq!( + memtable_rows, + input + .memtables + .iter() + .map(|memtable| memtable.stats().num_rows) + .sum::() + ); + } + + let source = primary_key_label_source(&scanner.metadata()); + let ranges = scanner + .properties() + .partitions + .clone() + .into_iter() + .flatten() + .collect::>(); + assert!(!ranges.is_empty(), "region scan needs an input range"); + // One output partition means every series the distributor assigns belongs to this receiver, + // so both rounds are delivered to partition 0. + scanner + .prepare(PrepareRequest { + ranges: Some(vec![ranges]), + target_partitions: Some(1), + ..Default::default() + }) + .unwrap(); + assert_eq!(1, scanner.properties().num_partitions()); + + let metrics = ExecutionPlanMetricsSet::default(); + let ctx = QueryScanContext::default(); + let stream = scanner + .scan_series_partition(&ctx, &metrics, 0, &source) + .unwrap() + .expect("two-phase raw scan is supported"); + assert_eq!(SeriesOrdering::Contiguous, stream.ordering()); + let schema = Arc::clone(stream.schema()); + let chunks = stream.try_collect::>().await.unwrap(); + // Two collection rounds really happened; the fixture never reaches the production threshold. + let rounds = crate::read::series_scan::candidate_assignment_rounds(); + assert!( + rounds >= 2, + "the partition must consume more than one assignment round, got {rounds}" + ); + assert!( + !chunks.is_empty(), + "partition 0 must own reader assignments" + ); + + // Partition 0's single stream carries every series exactly once, in ascending key order, so + // the stream crosses the round boundary: series 0..500 come from the first assignment and + // series 500 (the last of the globally sorted candidate stream) from the second one. + let mut runs = Vec::>::new(); + let mut chunks_per_series = HashMap::, usize>::new(); + for chunk in &chunks { + let key = chunk.primary_key.to_vec(); + *chunks_per_series.entry(key.clone()).or_default() += 1; + if runs.last() != Some(&key) { + assert!( + !runs.contains(&key), + "a series must not resume after another series across a round boundary" + ); + runs.push(key); + } + } + assert_eq!(keys, runs, "every series appears exactly once in key order"); + assert!( + chunks_per_series.values().copied().max().unwrap() > 1, + "the fixture must emit multiple chunks for one series, got {chunks_per_series:?}" + ); + // The round boundary: the first run belongs to the first assignment, the last run to the + // second assignment, and both are chunks of the same partition stream. + assert_eq!( + &keys[0], + runs.first().unwrap(), + "first round data is streamed" + ); + assert_eq!( + &keys[SERIES - 1], + runs.last().unwrap(), + "second round data is streamed" + ); + let last_round_chunks = chunks + .iter() + .filter(|chunk| chunk.primary_key.as_slice() == keys[SERIES - 1].as_slice()) + .count(); + assert!( + last_round_chunks > 0, + "the series of the second assignment must be streamed" + ); + + let stream_rows = chunks.iter().map(|chunk| chunk.time.len()).sum::(); + assert_eq!( + expected_times.values().map(Vec::len).sum::(), + stream_rows, + "the stream must carry every written row of both rounds" + ); + for (key, times) in &expected_times { + let observed = chunks + .iter() + .filter(|chunk| chunk.primary_key.as_slice() == key.as_slice()) + .flat_map(|chunk| chunk.time.as_buffer().iter().copied()) + .collect::>(); + assert_eq!(times, &observed, "timestamps must increase per series"); + assert!( + observed.windows(2).all(|window| window[0] < window[1]), + "timestamps must strictly increase within a key" + ); + } + for (index, chunk) in chunks.iter().enumerate() { + let next_is_same_series = chunks + .get(index + 1) + .is_some_and(|next| next.primary_key.as_slice() == chunk.primary_key.as_slice()); + assert_eq!( + !next_is_same_series, chunk.is_last, + "is_last must mark the final chunk of a series, including at the round boundary" + ); + let times = chunk.time.as_buffer(); + for (sample, time) in times.iter().enumerate() { + let SampleRef::Float(value) = chunk.fields.get(0).unwrap().sample(sample) else { + panic!("raw fields are converted to float samples"); + }; + assert_eq!((*time + 1) as f64, value); + } + } + + // Route the same partition stream through the frame codec so its `SequenceValidator` checks + // the chunks this reader produced, including the chunks the second assignment added. + let packed = pack_raw_series(Box::pin(ChainedRawSeriesStream::new( + Arc::clone(&schema), + source.clone(), + SeriesOrdering::Contiguous, + chunks.clone(), + ))) + .unwrap(); + let unpacked = unpack_raw_series( + packed, + Arc::clone(&schema), + source, + SeriesOrdering::Contiguous, + ) + .unwrap(); + let validated = unpacked.try_collect::>().await.unwrap(); + assert_eq!( + stream_rows, + validated + .iter() + .map(|chunk| chunk.time.len()) + .sum::(), + "packing and unpacking must preserve the multi-round stream" + ); + let mut schema_fields = schema + .fields + .iter() + .map(|field| field.name.as_ref()) + .collect::>(); + schema_fields.dedup(); + assert_eq!(vec!["field_0"], schema_fields); +} + /// Creates the sparse-encoded flat region shared by the two-phase raw scan tests. /// /// Returns the region metadata and the row schema accepted by the sparse write path. diff --git a/src/mito2/src/read/series_scan.rs b/src/mito2/src/read/series_scan.rs index 8b4ad2acf20..971ff93a742 100644 --- a/src/mito2/src/read/series_scan.rs +++ b/src/mito2/src/read/series_scan.rs @@ -81,6 +81,49 @@ const SEND_TIMEOUT: Duration = Duration::from_micros(100); /// Maximum number of candidate series retained before distributing assignments. const CANDIDATE_SERIES_ASSIGNMENT_THRESHOLD: usize = 1_000_000; +/// Test-only override of [`CANDIDATE_SERIES_ASSIGNMENT_THRESHOLD`], where `0` means "use the +/// production constant". +/// +/// The production build compiles this static, the setter below and the private accessor out, so +/// the shipped threshold is always the constant above. +#[cfg(test)] +static CANDIDATE_SERIES_ASSIGNMENT_THRESHOLD_OVERRIDE: AtomicUsize = AtomicUsize::new(0); + +/// Test-only number of candidate assignment rounds sent since the override was last set. +/// +/// A round is one `send_assignments` call, so a partition that received more than one assignment +/// sees `>= 2` rounds here. Increments are compiled out of production builds. +#[cfg(test)] +static CANDIDATE_ASSIGNMENT_ROUNDS: AtomicUsize = AtomicUsize::new(0); + +/// Returns the candidate-series assignment threshold, honouring the test-only override. +fn candidate_series_assignment_threshold() -> usize { + #[cfg(test)] + { + let threshold = CANDIDATE_SERIES_ASSIGNMENT_THRESHOLD_OVERRIDE.load(Ordering::Relaxed); + if threshold > 0 { + return threshold; + } + } + CANDIDATE_SERIES_ASSIGNMENT_THRESHOLD +} + +/// Test-only: overrides the candidate-series assignment threshold and resets the round counter. +/// +/// Pass `0` to restore the production threshold. Tests must reset the override before they +/// return; a scope guard in `engine::scan_test` does that even on panic. +#[cfg(test)] +pub(crate) fn set_candidate_series_assignment_threshold(threshold: usize) { + CANDIDATE_SERIES_ASSIGNMENT_THRESHOLD_OVERRIDE.store(threshold, Ordering::Relaxed); + CANDIDATE_ASSIGNMENT_ROUNDS.store(0, Ordering::Relaxed); +} + +/// Test-only: returns the number of assignment rounds sent since the override was last set. +#[cfg(test)] +pub(crate) fn candidate_assignment_rounds() -> usize { + CANDIDATE_ASSIGNMENT_ROUNDS.load(Ordering::Relaxed) +} + /// Legacy data receivers for output partitions. type LegacyReceiverList = Vec>>>; @@ -1143,7 +1186,7 @@ impl SeriesCandidateDistributor { return Ok(()); } collector.push(batch); - if collector.len() >= CANDIDATE_SERIES_ASSIGNMENT_THRESHOLD { + if collector.len() >= candidate_series_assignment_threshold() { chunked = true; self.send_assignments(collector.finish(false), &partition_pruner); if !self.should_fetch_candidates() { @@ -1167,6 +1210,9 @@ impl SeriesCandidateDistributor { assignments: Vec, partition_pruner: &Arc, ) { + // One round of candidate collection; a partition that receives two assignments sees two. + #[cfg(test)] + CANDIDATE_ASSIGNMENT_ROUNDS.fetch_add(1, Ordering::Relaxed); for (partition, assigned_series) in assignments.into_iter().enumerate() { if assigned_series.series().is_empty() { continue; From e1dba11eea857c2b1b9825753262fece81fffcfa Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Tue, 20 Oct 2026 22:46:23 +0800 Subject: [PATCH 4/4] test(mito2): check the ordering the producer declares, and drop the promql dev-dep MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The three end-to-end sites unpacked every frame as `Contiguous` while the producer declared `KeyThenTime`, and `SequenceValidator` only checks key order under the ordering it is given. A producer that declared key-then-time and emitted keys out of order passed all three. Each site now reads the ordering the frame declares, asserts it is `KeyThenTime` — so a future producer that quietly downgrades its declaration fails instead of skipping the check — and unpacks under it. Reversing two series' key order in the fixture now fails the test; passing `Contiguous` for the same frames still passes, which is what the old code did. The two-phase test asserted key order for an ordering that declares only `Contiguous`: a reader that stayed contiguous and stopped sorting keys is within contract and failed the test for it. It now asserts one run per key and that the two assignments hold the first and last key, which is what `Contiguous` plus the assignment composition actually promise. `promql` leaves mito2's dev-dependencies. The storage contract needs only `pack`/`unpack_raw_series` from `common-query`, which mito2 already depends on; the cross-region merge composition moves to promql, where it is a stream-level test rather than a two-region fixture: `merges_two_unpacked_frames_through_the_frame_codec`. The merge rejection that the two-phase test used to exercise is already covered by `rejects_unsupported_ordering_and_empty_input_without_polling`, which the remaining test now cites instead of duplicating. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- Cargo.lock | 1 - src/mito2/Cargo.toml | 1 - src/mito2/src/engine/scan_test.rs | 637 +++++---------------------- src/promql/src/source/merge/tests.rs | 62 +++ 4 files changed, 164 insertions(+), 537 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 43d65a03bd8..ea2a530d6b5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -8742,7 +8742,6 @@ dependencies = [ "partition", "pin-project", "prometheus 0.14.0", - "promql", "prost 0.14.1", "puffin", "rand 0.9.4", diff --git a/src/mito2/Cargo.toml b/src/mito2/Cargo.toml index 151694c40d1..02c134e964c 100644 --- a/src/mito2/Cargo.toml +++ b/src/mito2/Cargo.toml @@ -101,7 +101,6 @@ dotenv.workspace = true log-store = { workspace = true, features = ["testing"] } mito-codec = { workspace = true, features = ["testing"] } object-store = { workspace = true, features = ["services-memory"] } -promql.workspace = true rskafka.workspace = true rstest.workspace = true rstest_reuse.workspace = true diff --git a/src/mito2/src/engine/scan_test.rs b/src/mito2/src/engine/scan_test.rs index 4a720ef73f2..c0b155f7311 100644 --- a/src/mito2/src/engine/scan_test.rs +++ b/src/mito2/src/engine/scan_test.rs @@ -4772,7 +4772,6 @@ async fn test_mito_handle_query_raw_series_candidate_projection_and_mvcc() { use datatypes::series::samples::SampleRef; use datatypes::value::Value as DatumValue; use mito_codec::row_converter::build_primary_key_codec; - use promql::source::materialize_raw_labels; use store_api::region_engine::{BoundTag, QueryScanContext, RawLabelSource}; use store_api::region_request::RegionDeleteRequest; use table::predicate::Predicate; @@ -5093,7 +5092,7 @@ async fn test_mito_handle_query_raw_series_candidate_projection_and_mvcc() { assert_eq!(raw_visible, visible); // Exercise the production table physical consumer over a fresh real Mito scanner. - use common_query::series_frame::unpack_raw_series; + use common_query::series_frame::{raw_series_frame_ordering, unpack_raw_series}; use datafusion::execution::TaskContext; use datafusion::physical_plan::ExecutionPlan; use table::table::scan::RegionScanExec; @@ -5125,20 +5124,18 @@ async fn test_mito_handle_query_raw_series_candidate_projection_and_mvcc() { let native_frames = native_plan .execute(0, Arc::new(TaskContext::default())) .unwrap(); - let unpacked = unpack_raw_series( - native_frames, - native_schema, - native_source, - SeriesOrdering::Contiguous, - ) - .unwrap(); + // Unpack under the frame's own declaration. A legacy producer declares `KeyThenTime`, and + // only that declaration makes the decoder check that keys increase. + let frame_ordering = raw_series_frame_ordering(native_frames.schema().as_ref()).unwrap(); + assert_eq!( + SeriesOrdering::KeyThenTime, + frame_ordering, + "the legacy producer must declare key-then-time frames" + ); + let unpacked = + unpack_raw_series(native_frames, native_schema, native_source, frame_ordering).unwrap(); let native_schema = unpacked.schema().clone(); - let materialized = materialize_raw_labels(unpacked).unwrap(); - assert!(matches!( - materialized.label_source(), - RawLabelSource::Explicit - )); - let mut native = materialized.try_collect::>().await.unwrap(); + let mut native = unpacked.try_collect::>().await.unwrap(); assert_eq!( vec!["field_1", "field_0"], native_schema @@ -5150,18 +5147,17 @@ async fn test_mito_handle_query_raw_series_candidate_projection_and_mvcc() { let mut native_visible = native .drain(..) .flat_map(|chunk| { - let metric = chunk.metric.as_ref().expect("materialized complete metric"); - assert_eq!(Some("fixture_metric"), metric.name()); - let labels = metric - .tags() - .iter() - .map(|label| (label.name.as_ref(), label.value.as_ref())) - .collect::>(); - assert_eq!(1, labels.len()); - let tag = labels["tag_0"].to_string(); - assert!(tag == "A" || tag == "B"); - let times = chunk.time.as_buffer(); let key = chunk.primary_key.to_vec(); + // The frame binding is the raw primary key, so decode it back to the fixture tag to + // keep the raw rows comparable with the legacy path. + let tag = if key == expected_keys[0] { + "A" + } else if key == expected_keys[1] { + "B" + } else { + panic!("unexpected raw primary key {key:?}") + }; + let times = chunk.time.as_buffer(); let is_last = chunk.is_last; let mut rows = Vec::with_capacity(times.len()); for (i, time) in times.iter().enumerate() { @@ -5175,7 +5171,7 @@ async fn test_mito_handle_query_raw_series_candidate_projection_and_mvcc() { }; rows.push(( key.clone(), - tag.clone(), + tag.to_string(), *time, f1, f0, @@ -5230,478 +5226,22 @@ async fn test_mito_handle_query_raw_series_candidate_projection_and_mvcc() { ); } +/// A `TwoPhase` scan declares [`SeriesOrdering::Contiguous`] on the scanner, on every partition +/// stream, and therefore in every frame it packs; the frame adapter validates that declaration. +/// +/// The ordered merge that rejects a non-`KeyThenTime` input is covered by +/// `promql::source::merge::tests::rejects_unsupported_ordering_and_empty_input_without_polling`. #[tokio::test] -async fn test_mito_regional_raw_merge_preserves_global_key_order_and_labels() { - use api::v1::{Row, Value}; - use common_query::series_frame::{pack_raw_series, unpack_raw_series}; - use datatypes::arrow::array::Array; - use datatypes::series::SeriesOrdering; - use datatypes::series::samples::{SampleRef, SeriesSchema}; - use datatypes::value::Value as DatumValue; - use mito_codec::row_converter::build_primary_key_codec; - use promql::source::{materialize_raw_labels, merge_ordered_raw_series}; - use store_api::region_engine::{BoundTag, QueryScanContext, RawLabelSource}; - use store_api::region_request::RegionDeleteRequest; - - let mut env = TestEnv::with_prefix("test_mito_regional_raw_merge").await; - let engine = env.create_engine(MitoConfig::default()).await; - let regions = [RegionId::new(1, 1), RegionId::new(1, 2)]; - let mut request = CreateRequestBuilder::new().field_num(3).build(); - request - .options - .insert("memtable.type".into(), "bulk".into()); - let schema = test_util::rows_schema(&request); - for region in regions { - engine - .handle_request(region, RegionRequest::Create(request.clone())) - .await - .unwrap(); - } - let make_row = |tag: &str, f0: Option, f1: f64, gate: f64, time: i64| Row { - values: vec![ - Some(ValueData::StringValue(tag.into())), - f0.map(ValueData::F64Value), - Some(ValueData::F64Value(f1)), - Some(ValueData::F64Value(gate)), - Some(ValueData::TimestampMillisecondValue(time)), - ] - .into_iter() - .map(|value_data| Value { value_data }) - .collect(), +async fn test_two_phase_native_raw_stream_declares_contiguous_frames() { + use common_query::series_frame::{ + pack_raw_series, raw_series_frame_ordering, unpack_raw_series, }; - for (region, rows) in [ - ( - regions[0], - vec![ - make_row("A", Some(10.0), 1.0, 1.0, 10), - make_row("A", None, 3.0, 1.0, 30), - make_row("B", Some(-0.0), 5.0, 1.0, 10), - ], - ), - ( - regions[1], - vec![ - make_row("A", Some(20.0), 2.0, 1.0, 20), - make_row("A", Some(40.0), 4.0, 1.0, 40), - make_row("C", Some(60.0), 6.0, 1.0, 10), - ], - ), - ] { - test_util::put_rows( - &engine, - region, - Rows { - schema: schema.clone(), - rows, - }, - ) - .await; - } - test_util::flush_region(&engine, regions[0], None).await; - test_util::flush_region(&engine, regions[1], None).await; - test_util::put_rows( - &engine, - regions[0], - Rows { - schema: schema.clone(), - rows: vec![ - make_row("A", Some(11.0), 11.0, 1.0, 10), - make_row("A", Some(50.0), 5.0, 0.0, 50), - make_row("B", Some(60.0), 6.0, 1.0, 60), - ], - }, - ) - .await; - engine - .handle_request( - regions[0], - RegionRequest::Delete(RegionDeleteRequest { - rows: Rows { - schema: test_util::delete_rows_schema( - &CreateRequestBuilder::new().field_num(3).build(), - ), - rows: vec![row(vec![ - ValueData::StringValue("B".into()), - ValueData::TimestampMillisecondValue(60), - ])], - }, - hint: None, - partition_expr_version: None, - }), - ) - .await - .unwrap(); - let request = ScanRequest { - projection: Some(vec![2, 1, 4]), - filters: vec![col("field_2").gt(lit(0.0_f64))], - distribution: Some(TimeSeriesDistribution::PerSeries), - preserve_pk_dictionary_encoding: true, - ..Default::default() - }; - let mut streams = Vec::new(); - let mut metadata_for_codec = None; - let mut expected_schema: Option> = None; - let metrics = ExecutionPlanMetricsSet::default(); - for region in regions { - let selected = engine.scanner(region, request.clone()).await.unwrap(); - let mut scanner: Box = match selected { - Scanner::Series(series) => { - assert_eq!("legacy", series.mode()); - Box::new(series) - } - _ => panic!("PerSeries request must select a series scanner"), - }; - assert_eq!(SeriesOrdering::KeyThenTime, scanner.raw_series_ordering()); - let metadata = scanner.metadata(); - metadata_for_codec.get_or_insert_with(|| metadata.clone()); - let source = RawLabelSource::PrimaryKey { - metric_name: Arc::from("fixture_metric"), - physical_metadata: metadata.clone(), - logical_tags: metadata - .column_metadatas - .iter() - .filter(|column| { - column.semantic_type == SemanticType::Tag - && !store_api::storage::consts::is_internal_column( - &column.column_schema.name, - ) - && !store_api::metric_engine_consts::is_metric_engine_internal_column( - &column.column_schema.name, - ) - }) - .map(|column| BoundTag { - column_id: column.column_id, - name: Arc::from(column.column_schema.name.as_str()), - }) - .collect::>() - .into(), - }; - let ranges = scanner - .properties() - .partitions - .clone() - .into_iter() - .flatten() - .collect::>(); - assert!(!ranges.is_empty(), "region scan needs an input range"); - let partition_count = 2; - let mut partitions = vec![Vec::new(); partition_count]; - for (index, range) in ranges.into_iter().enumerate() { - partitions[index % partition_count].push(range); - } - scanner - .prepare(PrepareRequest { - ranges: Some(partitions), - target_partitions: Some(partition_count), - ..Default::default() - }) - .unwrap(); - let partition_count = scanner.properties().num_partitions(); - assert!( - partition_count > 1, - "prepare must use multiple supplied native range partitions" - ); - for partition in 0..partition_count { - let stream = scanner - .scan_series_partition(&QueryScanContext::default(), &metrics, partition, &source) - .unwrap() - .expect("native raw input"); - assert_eq!(scanner.raw_series_ordering(), stream.ordering()); - assert_eq!(SeriesOrdering::KeyThenTime, stream.ordering()); - assert_eq!( - vec!["field_1", "field_0"], - stream - .schema() - .fields - .iter() - .map(|field| field.name.as_ref()) - .collect::>() - ); - if let Some(schema) = &expected_schema { - assert_eq!( - schema - .fields - .iter() - .map(|field| field.name.as_ref()) - .collect::>(), - stream - .schema() - .fields - .iter() - .map(|field| field.name.as_ref()) - .collect::>() - ); - } else { - expected_schema = Some(Arc::clone(stream.schema())); - } - streams.push(stream); - } - // Each scanner owns the same schema and explicitly bound source. - } - let merged = merge_ordered_raw_series(streams).unwrap(); - let merged_schema = Arc::clone(merged.schema()); - let binding = merged.label_source().clone(); - let ordering = merged.ordering(); - assert_eq!(SeriesOrdering::KeyThenTime, ordering); - let packed = pack_raw_series(merged).unwrap(); - let unpacked = unpack_raw_series(packed, merged_schema, binding, ordering).unwrap(); - let mut materialized = materialize_raw_labels(unpacked).unwrap(); - let mut actual = Vec::new(); - while let Some(chunk) = materialized.try_next().await.unwrap() { - let metric = chunk.metric.as_ref().unwrap(); - assert_eq!(Some("fixture_metric"), metric.name()); - let tag = metric.get("tag_0").unwrap().to_string(); - for index in 0..chunk.time.len() { - let SampleRef::Float(f1) = chunk.fields.get(0).unwrap().sample(index) else { - panic!("field_1 float") - }; - let f0 = match chunk.fields.get(1).unwrap().sample(index) { - SampleRef::Float(value) => Some(value.to_bits()), - SampleRef::Missing => None, - _ => panic!("unexpected field_0 sample type"), - }; - actual.push(( - chunk.primary_key.to_vec(), - tag.clone(), - chunk.time.timestamp(index).unwrap(), - f1.to_bits(), - f0, - chunk.is_last && index + 1 == chunk.time.len(), - )); - } - } - assert!( - actual - .windows(2) - .all(|pair| pair[0].0 < pair[1].0 || (pair[0].0 == pair[1].0 && pair[0].2 < pair[1].2)), - "merged stream must be ordered by full read key then time" - ); - let metadata = metadata_for_codec.unwrap(); - let codec = build_primary_key_codec(&metadata); - let mut expected = Vec::new(); - for (tag, time, f1, f0) in [ - ("A", 10, 11.0_f64, Some(11.0_f64)), - ("A", 20, 2.0, Some(20.0)), - ("A", 30, 3.0, None), - ("A", 40, 4.0, Some(40.0)), - ("B", 10, 5.0, Some(-0.0_f64)), - ("C", 10, 6.0, Some(60.0)), - ] { - let mut key = Vec::new(); - codec - .encode_values( - &[(metadata.primary_key[0], DatumValue::String(tag.into()))], - &mut key, - ) - .unwrap(); - expected.push((key, tag, time, f1.to_bits(), f0.map(f64::to_bits))); - } - let duplicate_key = expected[0].0.clone(); - assert_eq!( - expected - .iter() - .map(|(key, _, _, _, _)| key.clone()) - .collect::>(), - actual.iter().map(|row| row.0.clone()).collect::>() - ); - let visible = actual - .iter() - .map(|(_, tag, time, f1, f0, _)| (tag.as_str(), *time, *f1, *f0)) - .collect::>(); - assert_eq!( - expected - .iter() - .map(|(_, tag, time, f1, f0)| (*tag, *time, *f1, *f0)) - .collect::>(), - visible - ); - let mut legacy_batches = Vec::new(); - for region in regions { - let mut legacy_request = request.clone(); - legacy_request.projection = Some(vec![0, 1, 2, 3, 4]); - legacy_request.filters.clear(); - legacy_request.preserve_pk_dictionary_encoding = false; - let legacy = engine.handle_query(region, legacy_request).await.unwrap(); - for partition in 0..legacy.properties().num_partitions() { - legacy_batches.extend( - legacy - .scan_partition(&QueryScanContext::default(), &metrics, partition) - .unwrap() - .try_collect::>() - .await - .unwrap(), - ); - } - } - let mut legacy_visible = Vec::new(); - for batch in &legacy_batches { - let batch = batch.df_record_batch(); - let tags = batch - .column_by_name("tag_0") - .unwrap() - .as_any() - .downcast_ref::() - .unwrap(); - let f0 = batch - .column_by_name("field_0") - .unwrap() - .as_any() - .downcast_ref::() - .unwrap(); - let f1 = batch - .column_by_name("field_1") - .unwrap() - .as_any() - .downcast_ref::() - .unwrap(); - let times = batch - .column_by_name("ts") - .unwrap() - .as_any() - .downcast_ref::() - .unwrap(); - let gate = batch - .column_by_name("field_2") - .unwrap() - .as_any() - .downcast_ref::() - .unwrap(); - for row in 0..batch.num_rows() { - if gate.value(row) > 0.0 { - legacy_visible.push(( - tags.value(row).to_string(), - times.value(row), - f1.value(row).to_bits(), - if f0.is_null(row) { - None - } else { - Some(f0.value(row).to_bits()) - }, - )); - } - } - } - legacy_visible.sort_by(|left, right| left.0.cmp(&right.0).then(left.1.cmp(&right.1))); - let mut raw_visible = actual - .iter() - .map(|(_, tag, time, f1, f0, _)| (tag.clone(), *time, *f1, *f0)) - .collect::>(); - raw_visible.sort_by(|left, right| left.0.cmp(&right.0).then(left.1.cmp(&right.1))); - assert_eq!( - legacy_visible, raw_visible, - "raw candidates after field_2 predicate must match legacy scan" - ); - test_util::put_rows( - &engine, - regions[1], - Rows { - schema: schema.clone(), - rows: vec![make_row("A", Some(99.0), 99.0, 1.0, 30)], - }, - ) - .await; - let mut duplicate_inputs = Vec::new(); - for region in regions { - let scanner = engine.handle_query(region, request.clone()).await.unwrap(); - let metadata = scanner.metadata(); - let source = RawLabelSource::PrimaryKey { - metric_name: Arc::from("fixture_metric"), - physical_metadata: metadata.clone(), - logical_tags: metadata - .column_metadatas - .iter() - .filter(|column| { - column.semantic_type == SemanticType::Tag - && !store_api::storage::consts::is_internal_column( - &column.column_schema.name, - ) - && !store_api::metric_engine_consts::is_metric_engine_internal_column( - &column.column_schema.name, - ) - }) - .map(|column| BoundTag { - column_id: column.column_id, - name: Arc::from(column.column_schema.name.as_str()), - }) - .collect::>() - .into(), - }; - let mut scanner = scanner; - let ranges = scanner - .properties() - .partitions - .clone() - .into_iter() - .flatten() - .collect::>(); - let partition_count = 2; - let mut partitions = vec![Vec::new(); partition_count]; - for (index, range) in ranges.into_iter().enumerate() { - partitions[index % partition_count].push(range); - } - scanner - .prepare(PrepareRequest { - ranges: Some(partitions), - target_partitions: Some(partition_count), - ..Default::default() - }) - .unwrap(); - for partition in 0..scanner.properties().num_partitions() { - duplicate_inputs.push( - scanner - .scan_series_partition( - &QueryScanContext::default(), - &metrics, - partition, - &source, - ) - .unwrap() - .unwrap(), - ); - } - } - let mut duplicate_merge = merge_ordered_raw_series(duplicate_inputs).unwrap(); - let mut affected_key_last = Vec::new(); - let duplicate_error = loop { - match duplicate_merge.try_next().await { - Ok(Some(chunk)) => { - if chunk.primary_key.as_slice() == duplicate_key.as_slice() { - affected_key_last.push(chunk.is_last); - } - } - Ok(None) => panic!("same full key/timestamp across regions must fail"), - Err(error) => break error.to_string(), - } - }; - assert!( - duplicate_error.contains("timestamp") || duplicate_error.contains("time"), - "unexpected duplicate timestamp error: {duplicate_error}" - ); - assert!(!affected_key_last.is_empty()); - assert!(affected_key_last.iter().all(|is_last| !is_last)); - assert!( - duplicate_merge.try_next().await.unwrap().is_none(), - "merge error must fuse stream" - ); - for tag in ["A", "B", "C"] { - let key_rows = actual.iter().filter(|row| row.1 == tag).collect::>(); - assert_eq!( - 1, - key_rows.iter().filter(|row| row.5).count(), - "global finality for {tag}" - ); - assert!(key_rows.last().unwrap().5); - assert!(key_rows[..key_rows.len() - 1].iter().all(|row| !row.5)); - } -} - -#[tokio::test] -async fn test_two_phase_native_raw_stream_is_rejected_by_ordered_merge() { use datatypes::series::SeriesOrdering; + use datatypes::series::samples::SeriesSchema; use store_api::region_engine::{BoundTag, QueryScanContext, RawLabelSource}; use store_api::region_request::RegionPutRequest; - let mut env = TestEnv::with_prefix("test_two_phase_native_raw_merge_rejection").await; + let mut env = TestEnv::with_prefix("test_two_phase_native_raw_contiguous_frames").await; let engine = env .create_engine(MitoConfig { experimental_series_scan_v2: true, @@ -5806,23 +5346,47 @@ async fn test_two_phase_native_raw_stream_is_rejected_by_ordered_merge() { }) .unwrap(); let metrics = ExecutionPlanMetricsSet::default(); - let mut streams = Vec::new(); + let mut schema: Option> = None; + let mut chunks = Vec::new(); for partition in 0..ranges.len() { let stream = scanner .scan_series_partition(&QueryScanContext::default(), &metrics, partition, &source) .unwrap() .unwrap(); + // The storage-side contract: every partition stream reports the scanner's declaration. assert_eq!(scanner.raw_series_ordering(), stream.ordering()); assert_eq!(SeriesOrdering::Contiguous, stream.ordering()); - streams.push(stream); + schema.get_or_insert_with(|| Arc::clone(stream.schema())); + chunks.extend(stream.try_collect::>().await.unwrap()); } - let error = match promql::source::merge_ordered_raw_series(streams) { - Ok(_) => panic!("Contiguous two-phase input cannot satisfy ordered merge ordering"), - Err(error) => error, - }; - assert!( - error.to_string().contains("ordering") || error.to_string().contains("ordered"), - "expected ordering rejection, got: {error}" + assert_eq!( + 2, + chunks.iter().map(|chunk| chunk.time.len()).sum::() + ); + + // The frames the adapter receives carry that declaration, and the adapter validates it + // through `pack`/`unpack_raw_series` without ever seeing a sorted key guarantee. + let schema = schema.expect("the scan must emit at least one partition stream"); + let packed = pack_raw_series(Box::pin(ChainedRawSeriesStream::new( + Arc::clone(&schema), + source.clone(), + SeriesOrdering::Contiguous, + chunks, + ))) + .unwrap(); + let ordering = raw_series_frame_ordering(packed.schema().as_ref()).unwrap(); + assert_eq!(SeriesOrdering::Contiguous, ordering); + let validated = unpack_raw_series(packed, schema, source, ordering) + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!( + 2, + validated + .iter() + .map(|chunk| chunk.time.len()) + .sum::() ); } @@ -6418,9 +5982,10 @@ async fn test_two_phase_raw_same_partition_multi_round_assignment_stays_contiguo "partition 0 must own reader assignments" ); - // Partition 0's single stream carries every series exactly once, in ascending key order, so - // the stream crosses the round boundary: series 0..500 come from the first assignment and - // series 500 (the last of the globally sorted candidate stream) from the second one. + // Partition 0's single stream carries every series in one contiguous run each. `TwoPhase` + // declares only `Contiguous`, so the runs are not contracted to be sorted; the assignment + // composition is: the first assignment holds the first 500 candidate series (which includes + // `keys[0]`) and the trailing assignment holds the last candidate series. let mut runs = Vec::>::new(); let mut chunks_per_series = HashMap::, usize>::new(); for chunk in &chunks { @@ -6434,22 +5999,21 @@ async fn test_two_phase_raw_same_partition_multi_round_assignment_stays_contiguo runs.push(key); } } - assert_eq!(keys, runs, "every series appears exactly once in key order"); + assert_eq!(keys.len(), runs.len(), "every series appears in one run"); assert!( chunks_per_series.values().copied().max().unwrap() > 1, "the fixture must emit multiple chunks for one series, got {chunks_per_series:?}" ); // The round boundary: the first run belongs to the first assignment, the last run to the // second assignment, and both are chunks of the same partition stream. - assert_eq!( - &keys[0], - runs.first().unwrap(), - "first round data is streamed" + let (first_assignment, second_assignment) = runs.split_at(SERIES - 1); + assert!( + first_assignment.contains(&keys[0]), + "the first assignment must stream the first key" ); - assert_eq!( - &keys[SERIES - 1], - runs.last().unwrap(), - "second round data is streamed" + assert!( + second_assignment.contains(&keys[SERIES - 1]), + "the second assignment must stream the last key" ); let last_round_chunks = chunks .iter() @@ -6794,7 +6358,7 @@ async fn test_mito_raw_series_int64_predicate_uses_original_precision() { assert_eq!(vec![(2, 20.0)], output_values); // Validate the typed integer gate through RegionScanExec's actual physical consumer. - use common_query::series_frame::unpack_raw_series; + use common_query::series_frame::{raw_series_frame_ordering, unpack_raw_series}; use datafusion::execution::TaskContext; use datafusion::physical_plan::ExecutionPlan; use datatypes::series::SeriesOrdering; @@ -6830,16 +6394,18 @@ async fn test_mito_raw_series_int64_predicate_uses_original_precision() { let frames = native_plan .execute(0, Arc::new(TaskContext::default())) .unwrap(); - let chunks = unpack_raw_series( - frames, - native_schema, - native_source, - SeriesOrdering::Contiguous, - ) - .unwrap() - .try_collect::>() - .await - .unwrap(); + // Unpack under the frame's own declaration so the decoder checks the legacy key order. + let ordering = raw_series_frame_ordering(frames.schema().as_ref()).unwrap(); + assert_eq!( + SeriesOrdering::KeyThenTime, + ordering, + "the legacy producer must declare key-then-time frames" + ); + let chunks = unpack_raw_series(frames, native_schema, native_source, ordering) + .unwrap() + .try_collect::>() + .await + .unwrap(); let native_rows = chunks .iter() .flat_map(|chunk| { @@ -6887,16 +6453,17 @@ async fn test_mito_raw_series_int64_predicate_uses_original_precision() { let empty_frames = empty_plan .execute(0, Arc::new(TaskContext::default())) .unwrap(); - let empty_chunks = unpack_raw_series( - empty_frames, - empty_schema, - empty_source, - SeriesOrdering::Contiguous, - ) - .unwrap() - .try_collect::>() - .await - .unwrap(); + let empty_ordering = raw_series_frame_ordering(empty_frames.schema().as_ref()).unwrap(); + assert_eq!( + SeriesOrdering::KeyThenTime, + empty_ordering, + "a statically empty legacy scan still declares key-then-time frames" + ); + let empty_chunks = unpack_raw_series(empty_frames, empty_schema, empty_source, empty_ordering) + .unwrap() + .try_collect::>() + .await + .unwrap(); assert!(empty_chunks.is_empty()); } diff --git a/src/promql/src/source/merge/tests.rs b/src/promql/src/source/merge/tests.rs index 4ad5d17cbb7..853da99e823 100644 --- a/src/promql/src/source/merge/tests.rs +++ b/src/promql/src/source/merge/tests.rs @@ -23,6 +23,7 @@ use std::task::{Context, Poll, Waker}; use api::v1::SemanticType; use common_error::ext::{BoxedError, ErrorExt, PlainError}; use common_error::status_code::StatusCode; +use common_query::series_frame::{pack_raw_series, raw_series_frame_ordering, unpack_raw_series}; use datatypes::arrow::array::Float64Array; use datatypes::arrow::buffer::{Buffer, ScalarBuffer}; use datatypes::data_type::ConcreteDataType; @@ -431,6 +432,67 @@ async fn merges_full_encoded_keys_and_interleaved_chunks_with_global_finality() assert!(series[2].3.last().unwrap()); } +/// Two inputs that each went through the frame codec merge like plain streams: unpacking keeps +/// the declared ordering, and the merge restores global key order and per-key finality across +/// the unpacked frames. +#[tokio::test] +async fn merges_two_unpacked_frames_through_the_frame_codec() { + let rack_a_sensor_1 = read_key("rack-a", "sensor-1"); + let rack_b_sensor_1 = read_key("rack-b", "sensor-1"); + let framed = |labels: RawLabelSource, chunks: Vec| { + let (stream, _, _) = stream( + labels, + SeriesOrdering::KeyThenTime, + chunks.into_iter().map(ready), + ); + let schema = Arc::clone(stream.schema()); + let binding = stream.label_source().clone(); + let packed = pack_raw_series(stream).unwrap(); + let ordering = raw_series_frame_ordering(packed.schema().as_ref()).unwrap(); + assert_eq!(SeriesOrdering::KeyThenTime, ordering); + unpack_raw_series(packed, schema, binding, ordering).unwrap() + }; + let left = framed( + label_source(1), + vec![ + chunk(rack_a_sensor_1.clone(), &[10, 30], &[1.0, 3.0], true), + chunk(rack_b_sensor_1.clone(), &[5], &[5.0], true), + ], + ); + let right = framed( + label_source(2), + vec![chunk(rack_a_sensor_1.clone(), &[20, 40], &[2.0, 4.0], true)], + ); + + let mut merged = merge_ordered_raw_series(vec![left, right]).unwrap(); + assert_eq!(SeriesOrdering::KeyThenTime, merged.ordering()); + let mut series = Vec::<(Buffer, Vec, Vec)>::new(); + while let Some(item) = merged.next().await { + let item = item.unwrap(); + if series + .last() + .is_none_or(|current| current.0 != item.primary_key) + { + series.push((item.primary_key.clone(), Vec::new(), Vec::new())); + } + let current = series.last_mut().unwrap(); + current.1.extend(item.time.as_buffer().to_vec()); + current.2.push(item.is_last); + } + assert_eq!(2, series.len()); + assert_eq!(rack_a_sensor_1, series[0].0); + assert_eq!(vec![10, 20, 30, 40], series[0].1); + assert!(*series[0].2.last().unwrap()); + assert!( + series[0].2[..series[0].2.len() - 1] + .iter() + .all(|last| !last) + ); + assert_eq!(rack_b_sensor_1, series[1].0); + assert_eq!(vec![5], series[1].1); + assert!(*series[1].2.last().unwrap()); +} + #[tokio::test] async fn rejects_equal_timestamp_from_separate_heads() { let key = read_key("rack-a", "sensor-1");