mirror of
https://github.com/GreptimeTeam/greptimedb.git
synced 2026-11-14 06:46:56 +00:00
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>
This commit is contained in:
discord9
@@ -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::<Vec<_>>().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::<Vec<_>>();
|
||||
|
||||
// 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<u8>, Vec<i64>>::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::<Vec<_>>();
|
||||
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::<Vec<_>>();
|
||||
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<Arc<SeriesSchema>> = 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::<Vec<_>>()
|
||||
);
|
||||
schema.get_or_insert_with(|| Arc::clone(stream.schema()));
|
||||
let chunks = stream.try_collect::<Vec<_>>().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::<Vec<_>>();
|
||||
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::<Vec<RawSeriesChunk>>();
|
||||
let schema = schema.unwrap();
|
||||
let complete_rows = complete.iter().map(|chunk| chunk.time.len()).sum::<usize>();
|
||||
assert_eq!(
|
||||
expected_times.values().map(Vec::len).sum::<usize>(),
|
||||
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::<Vec<_>>().await.unwrap();
|
||||
assert_eq!(
|
||||
complete_rows,
|
||||
validated
|
||||
.iter()
|
||||
.map(|chunk| chunk.time.len())
|
||||
.sum::<usize>(),
|
||||
"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::<Vec<u8>>::new();
|
||||
let mut chunks_per_series = HashMap::<Vec<u8>, 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::<Vec<_>>();
|
||||
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<store_api::metadata::RegionMetadata>,
|
||||
Vec<api::v1::ColumnSchema>,
|
||||
) {
|
||||
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<u8>, 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::<Vec<_>>()
|
||||
.into(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Concatenates the raw chunks of every reader assignment of one scan, in partition order.
|
||||
struct ChainedRawSeriesStream {
|
||||
schema: Arc<datatypes::series::samples::SeriesSchema>,
|
||||
label_source: store_api::region_engine::RawLabelSource,
|
||||
ordering: datatypes::series::SeriesOrdering,
|
||||
chunks: std::vec::IntoIter<datatypes::series::raw::RawSeriesChunk>,
|
||||
}
|
||||
|
||||
impl ChainedRawSeriesStream {
|
||||
fn new(
|
||||
schema: Arc<datatypes::series::samples::SeriesSchema>,
|
||||
label_source: store_api::region_engine::RawLabelSource,
|
||||
ordering: datatypes::series::SeriesOrdering,
|
||||
chunks: Vec<datatypes::series::raw::RawSeriesChunk>,
|
||||
) -> Self {
|
||||
Self {
|
||||
schema,
|
||||
label_source,
|
||||
ordering,
|
||||
chunks: chunks.into_iter(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl futures::Stream for ChainedRawSeriesStream {
|
||||
type Item = Result<datatypes::series::raw::RawSeriesChunk, common_error::ext::BoxedError>;
|
||||
|
||||
fn poll_next(
|
||||
mut self: std::pin::Pin<&mut Self>,
|
||||
_cx: &mut std::task::Context<'_>,
|
||||
) -> std::task::Poll<Option<Self::Item>> {
|
||||
std::task::Poll::Ready(self.chunks.next().map(Ok))
|
||||
}
|
||||
}
|
||||
|
||||
impl store_api::region_engine::RawSeriesStream for ChainedRawSeriesStream {
|
||||
fn schema(&self) -> &Arc<datatypes::series::samples::SeriesSchema> {
|
||||
&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;
|
||||
|
||||
Reference in New Issue
Block a user