diff --git a/src/mito2/src/engine/scan_test.rs b/src/mito2/src/engine/scan_test.rs index a893a234adc..c0b155f7311 100644 --- a/src/mito2/src/engine/scan_test.rs +++ b/src/mito2/src/engine/scan_test.rs @@ -4907,7 +4907,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 +4925,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 +4994,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 +5090,130 @@ 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::{raw_series_frame_ordering, 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(); + // 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 mut native = unpacked.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 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() { + 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.to_string(), + *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 +5226,1024 @@ 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_two_phase_native_raw_stream_declares_contiguous_frames() { + use common_query::series_frame::{ + pack_raw_series, raw_series_frame_ordering, unpack_raw_series, + }; + 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_contiguous_frames").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 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()); + schema.get_or_insert_with(|| Arc::clone(stream.schema())); + chunks.extend(stream.try_collect::>().await.unwrap()); + } + 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::() + ); +} + +#[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_disjoint_partition_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, + // 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; + 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 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] = [ + &[(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); + } + } + let mut round_rows = Vec::with_capacity(rounds.len()); + 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::>(); + 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; + } + } + + 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()); + // 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() + .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 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() + .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 the concatenated 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); + } + } +} + +/// 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 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 { + 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.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. + 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!( + second_assignment.contains(&keys[SERIES - 1]), + "the second assignment must stream the last key" + ); + 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. +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; @@ -5199,7 +6341,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 +6356,115 @@ 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::{raw_series_frame_ordering, 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(); + // 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| { + 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_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()); } /// Invalid predicate bindings decline before consumption; explicit labels decline and legacy still reads. @@ -5304,6 +6555,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 +6569,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 +6643,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 +6690,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..971ff93a742 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}; @@ -79,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>>>; @@ -88,6 +133,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 +245,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 +812,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 +835,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 +848,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, ))) } @@ -1096,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() { @@ -1120,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; 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(), 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");