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;