rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-11-14 06:46:56 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(mito2): produce typed raw series from RegionScanner with truthful ordering

Browse Source 
Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>
This commit is contained in:
discord9
2026-10-20 20:27:55 +08:00
parent 3ad5f96328
commit 5b649daaea
6 changed files with 1211 additions and 181 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Cargo.lock
Generated
+1
View File
@@ -8742,6 +8742,7 @@ dependencies = [
"partition",
"pin-project",
"prometheus 0.14.0",
"promql",
"prost 0.14.1",
"puffin",
"rand 0.9.4",
src/mito2/Cargo.toml
+1
View File
@@ -101,6 +101,7 @@ dotenv.workspace = true
log-store = { workspace = true, features = ["testing"] }
mito-codec = { workspace = true, features = ["testing"] }
object-store = { workspace = true, features = ["services-memory"] }
promql.workspace = true
rskafka.workspace = true
rstest.workspace = true
rstest_reuse.workspace = true
src/mito2/src/engine/scan_test.rs
+913 -5
View File
@@ -4772,6 +4772,7 @@ async fn test_mito_handle_query_raw_series_candidate_projection_and_mvcc() {
use datatypes::series::samples::SampleRef;
use datatypes::value::Value as DatumValue;
use mito_codec::row_converter::build_primary_key_codec;
use promql::source::materialize_raw_labels;
use store_api::region_engine::{BoundTag, QueryScanContext, RawLabelSource};
use store_api::region_request::RegionDeleteRequest;
use table::predicate::Predicate;
@@ -4907,7 +4908,7 @@ async fn test_mito_handle_query_raw_series_candidate_projection_and_mvcc() {
.scan_series_partition(&QueryScanContext::default(), &metrics, partition, &source)
.unwrap()
.expect("supported millisecond PerSeries scan");
assert_eq!(SeriesOrdering::Contiguous, stream.ordering());
assert_eq!(SeriesOrdering::KeyThenTime, stream.ordering());
assert!(stream.schema().eval_grid.is_none());
assert_eq!(
vec!["field_1", "field_0"],
@@ -4925,7 +4926,7 @@ async fn test_mito_handle_query_raw_series_candidate_projection_and_mvcc() {
!closed_keys.contains(chunk.primary_key.as_slice()),
"no chunk may follow a key's is_last chunk"
);
let times = chunk.time.as_buffer().as_ref();
let times = chunk.time.as_buffer();
assert_eq!(2, chunk.fields.len());
assert_eq!(times.len(), chunk.fields.get(0).unwrap().len());
assert_eq!(times.len(), chunk.fields.get(1).unwrap().len());
@@ -4994,7 +4995,7 @@ async fn test_mito_handle_query_raw_series_candidate_projection_and_mvcc() {
}
// Compare keyed visible tuples through a separately opened legacy scanner.
// A partition cannot be consumed through both paths.
let mut legacy_request = request;
let mut legacy_request = request.clone();
legacy_request.projection = Some(vec![0, 3, 2, 1, 4]);
legacy_request.preserve_pk_dictionary_encoding = false;
let legacy = engine
@@ -5090,6 +5091,133 @@ async fn test_mito_handle_query_raw_series_candidate_projection_and_mvcc() {
.collect::<Vec<_>>();
raw_visible.sort_by(|a, b| a.0.cmp(&b.0).then(a.1.cmp(&b.1)));
assert_eq!(raw_visible, visible);
// Exercise the production table physical consumer over a fresh real Mito scanner.
use common_query::series_frame::unpack_raw_series;
use datafusion::execution::TaskContext;
use datafusion::physical_plan::ExecutionPlan;
use table::table::scan::RegionScanExec;
let native_scanner = engine
.handle_query(region_id, request.clone())
.await
.unwrap();
let native_metadata = native_scanner.metadata();
let native_tags = native_metadata
.column_metadatas
.iter()
.filter(|c| c.semantic_type == SemanticType::Tag)
.map(|c| BoundTag {
column_id: c.column_id,
name: Arc::from(c.column_schema.name.as_str()),
})
.collect::<Vec<_>>();
let native_source = RawLabelSource::PrimaryKey {
metric_name: Arc::from("fixture_metric"),
physical_metadata: native_metadata,
logical_tags: native_tags.into(),
};
let native_plan = Arc::new(RegionScanExec::new(native_scanner, request.clone(), None).unwrap());
let native_plan = native_plan
.try_to_series_scan(native_source.clone())
.unwrap()
.expect("Mito raw-series capability");
let native_schema = native_plan.series_schema().clone();
let native_frames = native_plan
.execute(0, Arc::new(TaskContext::default()))
.unwrap();
let unpacked = unpack_raw_series(
native_frames,
native_schema,
native_source,
SeriesOrdering::Contiguous,
)
.unwrap();
let native_schema = unpacked.schema().clone();
let materialized = materialize_raw_labels(unpacked).unwrap();
assert!(matches!(
materialized.label_source(),
RawLabelSource::Explicit
));
let mut native = materialized.try_collect::<Vec<_>>().await.unwrap();
assert_eq!(
vec!["field_1", "field_0"],
native_schema
.fields
.iter()
.map(|field| field.name.as_ref())
.collect::<Vec<_>>()
);
let mut native_visible = native
.drain(..)
.flat_map(|chunk| {
let metric = chunk.metric.as_ref().expect("materialized complete metric");
assert_eq!(Some("fixture_metric"), metric.name());
let labels = metric
.tags()
.iter()
.map(|label| (label.name.as_ref(), label.value.as_ref()))
.collect::<HashMap<_, _>>();
assert_eq!(1, labels.len());
let tag = labels["tag_0"].to_string();
assert!(tag == "A" || tag == "B");
let times = chunk.time.as_buffer();
let key = chunk.primary_key.to_vec();
let is_last = chunk.is_last;
let mut rows = Vec::with_capacity(times.len());
for (i, time) in times.iter().enumerate() {
let SampleRef::Float(f1) = chunk.fields.get(0).unwrap().sample(i) else {
panic!("expected f1 float")
};
let f0 = match chunk.fields.get(1).unwrap().sample(i) {
SampleRef::Float(value) => Some(value),
SampleRef::Missing => None,
_ => panic!("expected f0 float or null"),
};
rows.push((
key.clone(),
tag.clone(),
*time,
f1,
f0,
is_last && i + 1 == times.len(),
));
}
rows
})
.collect::<Vec<_>>();
let mut expected_native = raw
.iter()
.map(|(key, time, f1, f0, _)| (key.to_vec(), *time, *f1, *f0))
.collect::<Vec<_>>();
type NativeRow = (Vec<u8>, String, i64, f64, Option<f64>, bool);
let sort_native = |rows: &mut Vec<NativeRow>| {
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::<Vec<_>>()
);
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::<Vec<_>>();
decoded_visible.sort_by(|a, b| a.0.cmp(&b.0).then(a.1.cmp(&b.1)));
assert_eq!(visible, decoded_visible);
// Apply the typed predicate and final projection through the ordinary legacy filter path.
assert_eq!(3, visible.len());
assert_eq!(
@@ -5102,6 +5230,602 @@ async fn test_mito_handle_query_raw_series_candidate_projection_and_mvcc() {
);
}
#[tokio::test]
async fn test_mito_regional_raw_merge_preserves_global_key_order_and_labels() {
use api::v1::{Row, Value};
use common_query::series_frame::{pack_raw_series, unpack_raw_series};
use datatypes::arrow::array::Array;
use datatypes::series::SeriesOrdering;
use datatypes::series::samples::{SampleRef, SeriesSchema};
use datatypes::value::Value as DatumValue;
use mito_codec::row_converter::build_primary_key_codec;
use promql::source::{materialize_raw_labels, merge_ordered_raw_series};
use store_api::region_engine::{BoundTag, QueryScanContext, RawLabelSource};
use store_api::region_request::RegionDeleteRequest;
let mut env = TestEnv::with_prefix("test_mito_regional_raw_merge").await;
let engine = env.create_engine(MitoConfig::default()).await;
let regions = [RegionId::new(1, 1), RegionId::new(1, 2)];
let mut request = CreateRequestBuilder::new().field_num(3).build();
request
.options
.insert("memtable.type".into(), "bulk".into());
let schema = test_util::rows_schema(&request);
for region in regions {
engine
.handle_request(region, RegionRequest::Create(request.clone()))
.await
.unwrap();
}
let make_row = |tag: &str, f0: Option<f64>, f1: f64, gate: f64, time: i64| Row {
values: vec![
Some(ValueData::StringValue(tag.into())),
f0.map(ValueData::F64Value),
Some(ValueData::F64Value(f1)),
Some(ValueData::F64Value(gate)),
Some(ValueData::TimestampMillisecondValue(time)),
]
.into_iter()
.map(|value_data| Value { value_data })
.collect(),
};
for (region, rows) in [
(
regions[0],
vec![
make_row("A", Some(10.0), 1.0, 1.0, 10),
make_row("A", None, 3.0, 1.0, 30),
make_row("B", Some(-0.0), 5.0, 1.0, 10),
],
),
(
regions[1],
vec![
make_row("A", Some(20.0), 2.0, 1.0, 20),
make_row("A", Some(40.0), 4.0, 1.0, 40),
make_row("C", Some(60.0), 6.0, 1.0, 10),
],
),
] {
test_util::put_rows(
&engine,
region,
Rows {
schema: schema.clone(),
rows,
},
)
.await;
}
test_util::flush_region(&engine, regions[0], None).await;
test_util::flush_region(&engine, regions[1], None).await;
test_util::put_rows(
&engine,
regions[0],
Rows {
schema: schema.clone(),
rows: vec![
make_row("A", Some(11.0), 11.0, 1.0, 10),
make_row("A", Some(50.0), 5.0, 0.0, 50),
make_row("B", Some(60.0), 6.0, 1.0, 60),
],
},
)
.await;
engine
.handle_request(
regions[0],
RegionRequest::Delete(RegionDeleteRequest {
rows: Rows {
schema: test_util::delete_rows_schema(
&CreateRequestBuilder::new().field_num(3).build(),
),
rows: vec![row(vec![
ValueData::StringValue("B".into()),
ValueData::TimestampMillisecondValue(60),
])],
},
hint: None,
partition_expr_version: None,
}),
)
.await
.unwrap();
let request = ScanRequest {
projection: Some(vec![2, 1, 4]),
filters: vec![col("field_2").gt(lit(0.0_f64))],
distribution: Some(TimeSeriesDistribution::PerSeries),
preserve_pk_dictionary_encoding: true,
..Default::default()
};
let mut streams = Vec::new();
let mut metadata_for_codec = None;
let mut expected_schema: Option<Arc<SeriesSchema>> = None;
let metrics = ExecutionPlanMetricsSet::default();
for region in regions {
let selected = engine.scanner(region, request.clone()).await.unwrap();
let mut scanner: Box<dyn RegionScanner> = match selected {
Scanner::Series(series) => {
assert_eq!("legacy", series.mode());
Box::new(series)
}
_ => panic!("PerSeries request must select a series scanner"),
};
assert_eq!(SeriesOrdering::KeyThenTime, scanner.raw_series_ordering());
let metadata = scanner.metadata();
metadata_for_codec.get_or_insert_with(|| metadata.clone());
let source = RawLabelSource::PrimaryKey {
metric_name: Arc::from("fixture_metric"),
physical_metadata: metadata.clone(),
logical_tags: metadata
.column_metadatas
.iter()
.filter(|column| {
column.semantic_type == SemanticType::Tag
&& !store_api::storage::consts::is_internal_column(
&column.column_schema.name,
)
&& !store_api::metric_engine_consts::is_metric_engine_internal_column(
&column.column_schema.name,
)
})
.map(|column| BoundTag {
column_id: column.column_id,
name: Arc::from(column.column_schema.name.as_str()),
})
.collect::<Vec<_>>()
.into(),
};
let ranges = scanner
.properties()
.partitions
.clone()
.into_iter()
.flatten()
.collect::<Vec<_>>();
assert!(!ranges.is_empty(), "region scan needs an input range");
let partition_count = 2;
let mut partitions = vec![Vec::new(); partition_count];
for (index, range) in ranges.into_iter().enumerate() {
partitions[index % partition_count].push(range);
}
scanner
.prepare(PrepareRequest {
ranges: Some(partitions),
target_partitions: Some(partition_count),
..Default::default()
})
.unwrap();
let partition_count = scanner.properties().num_partitions();
assert!(
partition_count > 1,
"prepare must use multiple supplied native range partitions"
);
for partition in 0..partition_count {
let stream = scanner
.scan_series_partition(&QueryScanContext::default(), &metrics, partition, &source)
.unwrap()
.expect("native raw input");
assert_eq!(scanner.raw_series_ordering(), stream.ordering());
assert_eq!(SeriesOrdering::KeyThenTime, stream.ordering());
assert_eq!(
vec!["field_1", "field_0"],
stream
.schema()
.fields
.iter()
.map(|field| field.name.as_ref())
.collect::<Vec<_>>()
);
if let Some(schema) = &expected_schema {
assert_eq!(
schema
.fields
.iter()
.map(|field| field.name.as_ref())
.collect::<Vec<_>>(),
stream
.schema()
.fields
.iter()
.map(|field| field.name.as_ref())
.collect::<Vec<_>>()
);
} else {
expected_schema = Some(Arc::clone(stream.schema()));
}
streams.push(stream);
}
// Each scanner owns the same schema and explicitly bound source.
}
let merged = merge_ordered_raw_series(streams).unwrap();
let merged_schema = Arc::clone(merged.schema());
let binding = merged.label_source().clone();
let ordering = merged.ordering();
assert_eq!(SeriesOrdering::KeyThenTime, ordering);
let packed = pack_raw_series(merged).unwrap();
let unpacked = unpack_raw_series(packed, merged_schema, binding, ordering).unwrap();
let mut materialized = materialize_raw_labels(unpacked).unwrap();
let mut actual = Vec::new();
while let Some(chunk) = materialized.try_next().await.unwrap() {
let metric = chunk.metric.as_ref().unwrap();
assert_eq!(Some("fixture_metric"), metric.name());
let tag = metric.get("tag_0").unwrap().to_string();
for index in 0..chunk.time.len() {
let SampleRef::Float(f1) = chunk.fields.get(0).unwrap().sample(index) else {
panic!("field_1 float")
};
let f0 = match chunk.fields.get(1).unwrap().sample(index) {
SampleRef::Float(value) => Some(value.to_bits()),
SampleRef::Missing => None,
_ => panic!("unexpected field_0 sample type"),
};
actual.push((
chunk.primary_key.to_vec(),
tag.clone(),
chunk.time.timestamp(index).unwrap(),
f1.to_bits(),
f0,
chunk.is_last && index + 1 == chunk.time.len(),
));
}
}
assert!(
actual
.windows(2)
.all(|pair| pair[0].0 < pair[1].0 || (pair[0].0 == pair[1].0 && pair[0].2 < pair[1].2)),
"merged stream must be ordered by full read key then time"
);
let metadata = metadata_for_codec.unwrap();
let codec = build_primary_key_codec(&metadata);
let mut expected = Vec::new();
for (tag, time, f1, f0) in [
("A", 10, 11.0_f64, Some(11.0_f64)),
("A", 20, 2.0, Some(20.0)),
("A", 30, 3.0, None),
("A", 40, 4.0, Some(40.0)),
("B", 10, 5.0, Some(-0.0_f64)),
("C", 10, 6.0, Some(60.0)),
] {
let mut key = Vec::new();
codec
.encode_values(
&[(metadata.primary_key[0], DatumValue::String(tag.into()))],
&mut key,
)
.unwrap();
expected.push((key, tag, time, f1.to_bits(), f0.map(f64::to_bits)));
}
let duplicate_key = expected[0].0.clone();
assert_eq!(
expected
.iter()
.map(|(key, _, _, _, _)| key.clone())
.collect::<Vec<_>>(),
actual.iter().map(|row| row.0.clone()).collect::<Vec<_>>()
);
let visible = actual
.iter()
.map(|(_, tag, time, f1, f0, _)| (tag.as_str(), *time, *f1, *f0))
.collect::<Vec<_>>();
assert_eq!(
expected
.iter()
.map(|(_, tag, time, f1, f0)| (*tag, *time, *f1, *f0))
.collect::<Vec<_>>(),
visible
);
let mut legacy_batches = Vec::new();
for region in regions {
let mut legacy_request = request.clone();
legacy_request.projection = Some(vec![0, 1, 2, 3, 4]);
legacy_request.filters.clear();
legacy_request.preserve_pk_dictionary_encoding = false;
let legacy = engine.handle_query(region, legacy_request).await.unwrap();
for partition in 0..legacy.properties().num_partitions() {
legacy_batches.extend(
legacy
.scan_partition(&QueryScanContext::default(), &metrics, partition)
.unwrap()
.try_collect::<Vec<_>>()
.await
.unwrap(),
);
}
}
let mut legacy_visible = Vec::new();
for batch in &legacy_batches {
let batch = batch.df_record_batch();
let tags = batch
.column_by_name("tag_0")
.unwrap()
.as_any()
.downcast_ref::<StringArray>()
.unwrap();
let f0 = batch
.column_by_name("field_0")
.unwrap()
.as_any()
.downcast_ref::<Float64Array>()
.unwrap();
let f1 = batch
.column_by_name("field_1")
.unwrap()
.as_any()
.downcast_ref::<Float64Array>()
.unwrap();
let times = batch
.column_by_name("ts")
.unwrap()
.as_any()
.downcast_ref::<TimestampMillisecondArray>()
.unwrap();
let gate = batch
.column_by_name("field_2")
.unwrap()
.as_any()
.downcast_ref::<Float64Array>()
.unwrap();
for row in 0..batch.num_rows() {
if gate.value(row) > 0.0 {
legacy_visible.push((
tags.value(row).to_string(),
times.value(row),
f1.value(row).to_bits(),
if f0.is_null(row) {
None
} else {
Some(f0.value(row).to_bits())
},
));
}
}
}
legacy_visible.sort_by(|left, right| left.0.cmp(&right.0).then(left.1.cmp(&right.1)));
let mut raw_visible = actual
.iter()
.map(|(_, tag, time, f1, f0, _)| (tag.clone(), *time, *f1, *f0))
.collect::<Vec<_>>();
raw_visible.sort_by(|left, right| left.0.cmp(&right.0).then(left.1.cmp(&right.1)));
assert_eq!(
legacy_visible, raw_visible,
"raw candidates after field_2 predicate must match legacy scan"
);
test_util::put_rows(
&engine,
regions[1],
Rows {
schema: schema.clone(),
rows: vec![make_row("A", Some(99.0), 99.0, 1.0, 30)],
},
)
.await;
let mut duplicate_inputs = Vec::new();
for region in regions {
let scanner = engine.handle_query(region, request.clone()).await.unwrap();
let metadata = scanner.metadata();
let source = RawLabelSource::PrimaryKey {
metric_name: Arc::from("fixture_metric"),
physical_metadata: metadata.clone(),
logical_tags: metadata
.column_metadatas
.iter()
.filter(|column| {
column.semantic_type == SemanticType::Tag
&& !store_api::storage::consts::is_internal_column(
&column.column_schema.name,
)
&& !store_api::metric_engine_consts::is_metric_engine_internal_column(
&column.column_schema.name,
)
})
.map(|column| BoundTag {
column_id: column.column_id,
name: Arc::from(column.column_schema.name.as_str()),
})
.collect::<Vec<_>>()
.into(),
};
let mut scanner = scanner;
let ranges = scanner
.properties()
.partitions
.clone()
.into_iter()
.flatten()
.collect::<Vec<_>>();
let partition_count = 2;
let mut partitions = vec![Vec::new(); partition_count];
for (index, range) in ranges.into_iter().enumerate() {
partitions[index % partition_count].push(range);
}
scanner
.prepare(PrepareRequest {
ranges: Some(partitions),
target_partitions: Some(partition_count),
..Default::default()
})
.unwrap();
for partition in 0..scanner.properties().num_partitions() {
duplicate_inputs.push(
scanner
.scan_series_partition(
&QueryScanContext::default(),
&metrics,
partition,
&source,
)
.unwrap()
.unwrap(),
);
}
}
let mut duplicate_merge = merge_ordered_raw_series(duplicate_inputs).unwrap();
let mut affected_key_last = Vec::new();
let duplicate_error = loop {
match duplicate_merge.try_next().await {
Ok(Some(chunk)) => {
if chunk.primary_key.as_slice() == duplicate_key.as_slice() {
affected_key_last.push(chunk.is_last);
}
}
Ok(None) => panic!("same full key/timestamp across regions must fail"),
Err(error) => break error.to_string(),
}
};
assert!(
duplicate_error.contains("timestamp") || duplicate_error.contains("time"),
"unexpected duplicate timestamp error: {duplicate_error}"
);
assert!(!affected_key_last.is_empty());
assert!(affected_key_last.iter().all(|is_last| !is_last));
assert!(
duplicate_merge.try_next().await.unwrap().is_none(),
"merge error must fuse stream"
);
for tag in ["A", "B", "C"] {
let key_rows = actual.iter().filter(|row| row.1 == tag).collect::<Vec<_>>();
assert_eq!(
1,
key_rows.iter().filter(|row| row.5).count(),
"global finality for {tag}"
);
assert!(key_rows.last().unwrap().5);
assert!(key_rows[..key_rows.len() - 1].iter().all(|row| !row.5));
}
}
#[tokio::test]
async fn test_two_phase_native_raw_stream_is_rejected_by_ordered_merge() {
use datatypes::series::SeriesOrdering;
use store_api::region_engine::{BoundTag, QueryScanContext, RawLabelSource};
use store_api::region_request::RegionPutRequest;
let mut env = TestEnv::with_prefix("test_two_phase_native_raw_merge_rejection").await;
let engine = env
.create_engine(MitoConfig {
experimental_series_scan_v2: true,
..Default::default()
})
.await;
let region_id = RegionId::new(1, 1);
let metadata = Arc::new(sst_region_metadata_with_encoding(
PrimaryKeyEncoding::Sparse,
));
let mut create = CreateRequestBuilder::new().build();
create.column_metadatas = metadata.column_metadatas.clone();
create.primary_key = metadata.primary_key.clone();
create
.options
.insert(PRIMARY_KEY_ENCODING.to_string(), "sparse".to_string());
create
.options
.insert("memtable.type".to_string(), "bulk".to_string());
create
.options
.insert("sst_format".to_string(), "flat".to_string());
let full_schema = test_util::rows_schema(&create);
let mut key_schema = full_schema[0].clone();
key_schema.column_name = PRIMARY_KEY_COLUMN_NAME.to_string();
key_schema.datatype = ColumnDataType::Binary.into();
key_schema.semantic_type = SemanticType::Tag.into();
let row_schema = vec![key_schema, full_schema[5].clone(), full_schema[4].clone()];
engine
.handle_request(region_id, RegionRequest::Create(create))
.await
.unwrap();
let sparse_rows = Rows {
schema: row_schema,
rows: vec![
row(vec![
ValueData::BinaryValue(new_sparse_primary_key(&["a", "x"], &metadata, 10, 0)),
ValueData::TimestampMillisecondValue(10),
ValueData::U64Value(1),
]),
row(vec![
ValueData::BinaryValue(new_sparse_primary_key(&["b", "y"], &metadata, 20, 0)),
ValueData::TimestampMillisecondValue(20),
ValueData::U64Value(2),
]),
],
};
engine
.handle_request(
region_id,
RegionRequest::Put(RegionPutRequest {
rows: sparse_rows,
hint: Some(WriteHint {
primary_key_encoding: api::v1::PrimaryKeyEncoding::Sparse.into(),
}),
skip_wal: false,
partition_expr_version: None,
}),
)
.await
.unwrap();
test_util::flush_region(&engine, region_id, None).await;
let request = ScanRequest {
distribution: Some(TimeSeriesDistribution::PerSeries),
preserve_pk_dictionary_encoding: true,
..Default::default()
};
let selected = engine.scanner(region_id, request.clone()).await.unwrap();
let Scanner::Series(series) = selected else {
panic!("sparse PerSeries request must use SeriesScan");
};
assert_eq!("two_phase", series.mode());
let mut scanner = engine.handle_query(region_id, request).await.unwrap();
assert_eq!(SeriesOrdering::Contiguous, scanner.raw_series_ordering());
let metadata = scanner.metadata();
let source = RawLabelSource::PrimaryKey {
metric_name: Arc::from("fixture_metric"),
physical_metadata: metadata.clone(),
logical_tags: metadata
.column_metadatas
.iter()
.filter(|column| {
column.semantic_type == SemanticType::Tag
&& !store_api::storage::consts::is_internal_column(&column.column_schema.name)
&& !store_api::metric_engine_consts::is_metric_engine_internal_column(
&column.column_schema.name,
)
})
.map(|column| BoundTag {
column_id: column.column_id,
name: Arc::from(column.column_schema.name.as_str()),
})
.collect::<Vec<_>>()
.into(),
};
let ranges = scanner.properties().partitions.clone();
scanner
.prepare(PrepareRequest {
ranges: Some(ranges.clone()),
target_partitions: Some(ranges.len()),
..Default::default()
})
.unwrap();
let metrics = ExecutionPlanMetricsSet::default();
let mut streams = Vec::new();
for partition in 0..ranges.len() {
let stream = scanner
.scan_series_partition(&QueryScanContext::default(), &metrics, partition, &source)
.unwrap()
.unwrap();
assert_eq!(scanner.raw_series_ordering(), stream.ordering());
assert_eq!(SeriesOrdering::Contiguous, stream.ordering());
streams.push(stream);
}
let error = match promql::source::merge_ordered_raw_series(streams) {
Ok(_) => panic!("Contiguous two-phase input cannot satisfy ordered merge ordering"),
Err(error) => error,
};
assert!(
error.to_string().contains("ordering") || error.to_string().contains("ordered"),
"expected ordering rejection, got: {error}"
);
}
#[tokio::test]
async fn test_mito_raw_series_int64_predicate_uses_original_precision() {
use datatypes::series::samples::SampleRef;
@@ -5199,7 +5923,7 @@ async fn test_mito_raw_series_int64_predicate_uses_original_precision() {
);
let chunks = stream.try_collect::<Vec<_>>().await.unwrap();
for chunk in chunks {
let axis = chunk.time.as_buffer().as_ref();
let axis = chunk.time.as_buffer();
times.extend(axis.iter().copied());
for (index, &time) in axis.iter().enumerate() {
output_values.push((
@@ -5214,6 +5938,112 @@ async fn test_mito_raw_series_int64_predicate_uses_original_precision() {
}
assert_eq!(vec![2], times);
assert_eq!(vec![(2, 20.0)], output_values);
// Validate the typed integer gate through RegionScanExec's actual physical consumer.
use common_query::series_frame::unpack_raw_series;
use datafusion::execution::TaskContext;
use datafusion::physical_plan::ExecutionPlan;
use datatypes::series::SeriesOrdering;
use table::table::scan::RegionScanExec;
let native_request = ScanRequest {
projection: Some(vec![2, 3]),
filters: vec![
col("tag_0").eq(lit("series")),
col("field_0").gt(lit(boundary)),
],
distribution: Some(TimeSeriesDistribution::PerSeries),
preserve_pk_dictionary_encoding: true,
..Default::default()
};
let native_scanner = engine
.handle_query(region_id, native_request.clone())
.await
.unwrap();
let native_metadata = native_scanner.metadata();
let native_source = RawLabelSource::PrimaryKey {
metric_name: Arc::from("fixture_metric"),
physical_metadata: native_metadata.clone(),
logical_tags: Arc::from([BoundTag {
column_id: native_metadata.primary_key[0],
name: Arc::from("tag_0"),
}]),
};
let native_plan = Arc::new(RegionScanExec::new(native_scanner, native_request, None).unwrap())
.try_to_series_scan(native_source.clone())
.unwrap()
.expect("typed int64 gate is supported");
let native_schema = native_plan.series_schema().clone();
let frames = native_plan
.execute(0, Arc::new(TaskContext::default()))
.unwrap();
let chunks = unpack_raw_series(
frames,
native_schema,
native_source,
SeriesOrdering::Contiguous,
)
.unwrap()
.try_collect::<Vec<_>>()
.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::<Vec<_>>()
})
.collect::<Vec<_>>();
assert_eq!(vec![(2, 20.0)], native_rows);
let empty_request = ScanRequest {
projection: Some(vec![2, 3]),
filters: vec![
col("tag_0").eq(lit("series")),
col("field_0").gt(lit(boundary + 1)),
],
distribution: Some(TimeSeriesDistribution::PerSeries),
preserve_pk_dictionary_encoding: true,
..Default::default()
};
let empty_scanner = engine
.handle_query(region_id, empty_request.clone())
.await
.unwrap();
let empty_metadata = empty_scanner.metadata();
let empty_source = RawLabelSource::PrimaryKey {
metric_name: Arc::from("fixture_metric"),
physical_metadata: empty_metadata.clone(),
logical_tags: Arc::from([BoundTag {
column_id: empty_metadata.primary_key[0],
name: Arc::from("tag_0"),
}]),
};
let empty_plan = Arc::new(RegionScanExec::new(empty_scanner, empty_request, None).unwrap())
.try_to_series_scan(empty_source.clone())
.unwrap()
.expect("statically empty scan still has native capability");
let empty_schema = empty_plan.series_schema().clone();
let empty_frames = empty_plan
.execute(0, Arc::new(TaskContext::default()))
.unwrap();
let empty_chunks = unpack_raw_series(
empty_frames,
empty_schema,
empty_source,
SeriesOrdering::Contiguous,
)
.unwrap()
.try_collect::<Vec<_>>()
.await
.unwrap();
assert!(empty_chunks.is_empty());
}
/// Invalid predicate bindings decline before consumption; explicit labels decline and legacy still reads.
@@ -5304,6 +6134,11 @@ async fn test_mito_raw_series_invalid_binding_before_legacy_read() {
name: Arc::from("tag_0"),
}]),
};
assert_eq!(
supported,
scanner.raw_series_schema(&valid_source).unwrap().is_some(),
"schema preflight case: {name}"
);
let raw = scanner
.scan_series_partition(&ctx, &metrics, 0, &valid_source)
.unwrap();
@@ -5313,7 +6148,7 @@ async fn test_mito_raw_series_invalid_binding_before_legacy_read() {
let tuples = chunks
.iter()
.map(|chunk| {
let times = chunk.time.as_buffer().as_ref();
let times = chunk.time.as_buffer();
let SampleRef::Float(value) = chunk.fields.get(0).unwrap().sample(0) else {
panic!("expected raw float sample");
};
@@ -5387,6 +6222,13 @@ async fn test_mito_raw_series_invalid_binding_before_legacy_read() {
.unwrap()
.is_none()
);
assert!(
scanner
.raw_series_schema(&RawLabelSource::Explicit)
.unwrap()
.is_none()
);
assert!(scanner.raw_series_schema(&valid_shape_bad_tag).is_err());
assert!(
scanner
.scan_series_partition(&ctx, &metrics, 0, &valid_shape_bad_tag)
@@ -5427,4 +6269,70 @@ async fn test_mito_raw_series_invalid_binding_before_legacy_read() {
vec![("0".to_string(), 0.0, 0), ("1".to_string(), 1.0, 1000)],
rows
);
// Runtime pruning wrappers are deliberately ineligible for native planning, and
// wrappers installed after planning must be caught at execute time.
use datafusion::execution::TaskContext;
use datafusion::physical_plan::ExecutionPlan;
use table::table::scan::RegionScanExec;
let dynamic_request = ScanRequest {
distribution: Some(TimeSeriesDistribution::PerSeries),
preserve_pk_dictionary_encoding: true,
..Default::default()
};
let dynamic_source_for = |scanner: &dyn RegionScanner| {
let metadata = scanner.metadata();
RawLabelSource::PrimaryKey {
metric_name: Arc::from("fixture_metric"),
physical_metadata: metadata.clone(),
logical_tags: Arc::from([BoundTag {
column_id: metadata.primary_key[0],
name: Arc::from("tag_0"),
}]),
}
};
let make_dynamic = || {
Arc::new(DynamicFilterPhysicalExpr::new(
vec![Arc::new(Column::new("field_0", 1))],
physical_lit(true),
))
};
let before = engine
.handle_query(region_id, dynamic_request.clone())
.await
.unwrap();
let before_source = dynamic_source_for(before.as_ref());
let before_scan = Arc::new(RegionScanExec::new(before, dynamic_request.clone(), None).unwrap());
assert_eq!(
vec![true],
before_scan.add_dyn_filters_to_predicate(vec![make_dynamic()])
);
assert!(
before_scan
.clone()
.try_to_series_scan(before_source)
.unwrap()
.is_none()
);
let after = engine
.handle_query(region_id, dynamic_request.clone())
.await
.unwrap();
let after_source = dynamic_source_for(after.as_ref());
let after_scan = Arc::new(RegionScanExec::new(after, dynamic_request, None).unwrap());
let after_plan = after_scan
.clone()
.try_to_series_scan(after_source)
.unwrap()
.expect("native plan before dynamic wrapper installation");
assert_eq!(
vec![true],
after_scan.add_dyn_filters_to_predicate(vec![make_dynamic()])
);
assert!(
after_plan
.execute(0, Arc::new(TaskContext::default()))
.is_err()
);
}
src/mito2/src/read/series_scan.rs
+215 -168
View File
@@ -31,13 +31,15 @@ use common_recordbatch::{RecordBatchStreamWrapper, SendableRecordBatchStream};
use common_telemetry::tracing::{self, Instrument};
use common_telemetry::{debug, warn};
use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet;
use datafusion::physical_plan::{DisplayAs, DisplayFormatType};
use datafusion::physical_plan::{DisplayAs, DisplayFormatType, PhysicalExpr};
use datafusion_common::tree_node::TreeNode;
use datafusion_expr::{Expr, Volatility};
use datatypes::arrow::array::BinaryArray;
use datatypes::arrow::datatypes::DataType as ArrowDataType;
use datatypes::arrow::record_batch::RecordBatch;
use datatypes::prelude::ConcreteDataType;
use datatypes::schema::SchemaRef;
use datatypes::series::SeriesOrdering;
use datatypes::series::samples::{FieldSpec, SeriesSchema};
use datatypes::types::TimestampType;
use futures::{StreamExt, TryStreamExt};
@@ -88,6 +90,12 @@ type CandidateReceiverList = Vec<Option<CandidateReceiver>>;
/// Candidate assignment senders for output partitions.
type CandidateSenderList = Vec<Option<mpsc::UnboundedSender<Result<SeriesReaderInput>>>>;
struct RawScanPreparation {
schema: Arc<SeriesSchema>,
fields: Vec<(usize, ConcreteDataType)>,
predicates: Vec<Arc<dyn PhysicalExpr>>,
}
struct CandidateReceiver {
receiver: mpsc::UnboundedReceiver<Result<SeriesReaderInput>>,
active_receivers: Arc<AtomicUsize>,
@@ -194,6 +202,189 @@ impl SeriesScan {
}
}
fn prepare_raw_scan(
&self,
label_source: &RawLabelSource,
) -> Result<Option<RawScanPreparation>, BoxedError> {
let input = &self.stream_ctx.input;
let metadata = input.mapper.metadata();
match label_source {
RawLabelSource::PrimaryKey {
physical_metadata,
logical_tags,
..
} => {
if !Arc::ptr_eq(metadata, physical_metadata) {
debug!(
"Raw series scan preflight rejected metadata snapshot: region_id={}, physical_region_id={}, schema_version={}, physical_schema_version={}",
metadata.region_id,
physical_metadata.region_id,
metadata.schema_version,
physical_metadata.schema_version,
);
return Err(BoxedError::new(
InvalidRecordBatchSnafu {
reason: "raw label source metadata does not match scanner snapshot"
.to_string(),
}
.build(),
));
}
let mut ids = HashSet::new();
let mut names = HashSet::new();
for tag in logical_tags.iter() {
let column = metadata.column_by_id(tag.column_id).filter(|column| {
column.semantic_type == SemanticType::Tag
&& metadata.primary_key.contains(&tag.column_id)
});
let valid = column.is_some_and(|column| {
column.column_schema.name == tag.name.as_ref()
&& tag.name.as_ref() != "__name__"
&& !is_internal_column(tag.name.as_ref())
&& !is_metric_engine_internal_column(tag.name.as_ref())
&& !ReservedColumnId::is_reserved(tag.column_id)
});
let unique_id = valid && ids.insert(tag.column_id);
let unique_name = unique_id && names.insert(tag.name.as_ref());
if !valid || !unique_id || !unique_name {
debug!(
"Raw series scan preflight rejected label binding: region_id={}, column_id={}, binding_valid={}, unique_id={}, unique_name={}",
metadata.region_id, tag.column_id, valid, unique_id, unique_name,
);
return Err(BoxedError::new(
InvalidRecordBatchSnafu {
reason: "invalid raw label binding".to_string(),
}
.build(),
));
}
}
}
RawLabelSource::Explicit => {
debug!(
"Raw series scan preflight rejected explicit labels: region_id={}",
metadata.region_id,
);
return Ok(None);
}
}
let selected = input.mapper.output_schema();
let mut fields = Vec::new();
let mut specs = Vec::new();
for output in selected.column_schemas() {
let Some(column) = metadata.column_metadatas.iter().find(|column| {
column.column_schema.name == output.name
&& column.semantic_type == SemanticType::Field
}) else {
continue;
};
let Some((index, (_, data_type))) = input
.mapper
.batch_schema()
.iter()
.enumerate()
.find(|(_, (id, _))| *id == column.column_id)
else {
debug!(
"Raw series scan preflight missing mapper field: region_id={}, column_id={}",
metadata.region_id, column.column_id,
);
return Err(BoxedError::new(
InvalidRecordBatchSnafu {
reason: format!(
"selected field {} is missing from mapper batch schema",
output.name
),
}
.build(),
));
};
if !data_type.is_numeric() && !is_native_histogram_value_type(data_type) {
debug!(
"Raw series scan preflight rejected field type: region_id={}, column_id={}",
metadata.region_id, column.column_id,
);
return Ok(None);
}
fields.push((index, data_type.clone()));
specs.push(FieldSpec {
name: Arc::from(output.name.as_str()),
});
}
let time_is_millis = matches!(metadata.time_index_type(), TimestampType::Millisecond(_));
if fields.is_empty()
|| input.append_mode
|| input.series_row_selector.is_some()
|| !time_is_millis
{
debug!(
"Raw series scan preflight rejected scan options: region_id={}, selected_fields={}, append_mode={}, row_selector={}, time_is_millis={}",
metadata.region_id,
fields.len(),
input.append_mode,
input.series_row_selector.is_some(),
time_is_millis,
);
return Ok(None);
}
let schema = Arc::new(SeriesSchema {
fields: specs.into(),
eval_grid: None,
});
let typed_schema = input.mapper.input_arrow_schema(false);
let mut predicates = Vec::new();
if let Some(predicate) = input.predicate_group().predicate_without_region() {
if !predicate.dyn_filters().is_empty() {
debug!(
"Raw series scan preflight rejected dynamic predicate filters: region_id={}",
metadata.region_id,
);
return Ok(None);
}
for expr in predicate.exprs() {
// Fresh execution properties cannot preserve Stable or query-dependent semantics.
if expr
.exists(|node| {
Ok(matches!(node, Expr::ScalarFunction(function)
if function.func.signature().volatility != Volatility::Immutable)
|| matches!(node, Expr::ScalarVariable(..)))
})
.unwrap_or(true)
{
debug!(
"Raw series scan preflight rejected predicate volatility or variable: region_id={}",
metadata.region_id,
);
return Ok(None);
}
// Evaluate residuals with original types before raw conversion to Float64.
let Ok(physical) = Predicate::to_physical_expr(expr, &typed_schema) else {
debug!(
"Raw series scan preflight failed predicate compilation: region_id={}",
metadata.region_id,
);
return Ok(None);
};
if !matches!(
physical.data_type(&typed_schema),
Ok(ArrowDataType::Boolean | ArrowDataType::Null)
) {
debug!(
"Raw series scan preflight rejected predicate type: region_id={}",
metadata.region_id,
);
return Ok(None);
}
predicates.push(physical);
}
}
Ok(Some(RawScanPreparation {
schema,
fields,
predicates,
}))
}
fn supports_two_phase(input: &ScanInput) -> bool {
if !is_sparse_metric_metadata(input.region_metadata()) {
return false;
@@ -578,6 +769,22 @@ impl RegionScanner for SeriesScan {
.map_err(BoxedError::new)
}
fn raw_series_ordering(&self) -> SeriesOrdering {
match self.mode {
SeriesScanMode::Legacy => SeriesOrdering::KeyThenTime,
SeriesScanMode::TwoPhase => SeriesOrdering::Contiguous,
}
}
fn raw_series_schema(
&self,
label_source: &RawLabelSource,
) -> Result<Option<Arc<SeriesSchema>>, BoxedError> {
Ok(self
.prepare_raw_scan(label_source)?
.map(|preparation| preparation.schema))
}
fn scan_series_partition(
&self,
ctx: &QueryScanContext,
@@ -585,171 +792,9 @@ impl RegionScanner for SeriesScan {
partition: usize,
label_source: &RawLabelSource,
) -> Result<Option<SendableRawSeriesStream>, BoxedError> {
let input = &self.stream_ctx.input;
let metadata = input.mapper.metadata();
match label_source {
RawLabelSource::PrimaryKey {
physical_metadata,
logical_tags,
..
} => {
if !Arc::ptr_eq(metadata, physical_metadata) {
debug!(
"Raw series scan preflight rejected metadata snapshot: region_id={}, physical_region_id={}, schema_version={}, physical_schema_version={}",
metadata.region_id,
physical_metadata.region_id,
metadata.schema_version,
physical_metadata.schema_version,
);
return Err(BoxedError::new(
InvalidRecordBatchSnafu {
reason: "raw label source metadata does not match scanner snapshot"
.to_string(),
}
.build(),
));
}
let mut ids = HashSet::new();
let mut names = HashSet::new();
for tag in logical_tags.iter() {
let column = metadata.column_by_id(tag.column_id).filter(|column| {
column.semantic_type == SemanticType::Tag
&& metadata.primary_key.contains(&tag.column_id)
});
let valid = column.is_some_and(|column| {
column.column_schema.name == tag.name.as_ref()
&& tag.name.as_ref() != "__name__"
&& !is_internal_column(tag.name.as_ref())
&& !is_metric_engine_internal_column(tag.name.as_ref())
&& !ReservedColumnId::is_reserved(tag.column_id)
});
let unique_id = valid && ids.insert(tag.column_id);
let unique_name = unique_id && names.insert(tag.name.as_ref());
if !valid || !unique_id || !unique_name {
debug!(
"Raw series scan preflight rejected label binding: region_id={}, column_id={}, binding_valid={}, unique_id={}, unique_name={}",
metadata.region_id, tag.column_id, valid, unique_id, unique_name,
);
return Err(BoxedError::new(
InvalidRecordBatchSnafu {
reason: "invalid raw label binding".to_string(),
}
.build(),
));
}
}
}
RawLabelSource::Explicit => {
debug!(
"Raw series scan preflight rejected explicit labels: region_id={}",
metadata.region_id,
);
return Ok(None);
}
}
let selected = input.mapper.output_schema();
let mut fields = Vec::new();
let mut specs = Vec::new();
for output in selected.column_schemas() {
let Some(column) = metadata.column_metadatas.iter().find(|column| {
column.column_schema.name == output.name
&& column.semantic_type == SemanticType::Field
}) else {
continue;
};
let Some((index, (_, data_type))) = input
.mapper
.batch_schema()
.iter()
.enumerate()
.find(|(_, (id, _))| *id == column.column_id)
else {
debug!(
"Raw series scan preflight missing mapper field: region_id={}, column_id={}",
metadata.region_id, column.column_id,
);
return Err(BoxedError::new(
InvalidRecordBatchSnafu {
reason: format!(
"selected field {} is missing from mapper batch schema",
output.name
),
}
.build(),
));
};
if !data_type.is_numeric() && !is_native_histogram_value_type(data_type) {
debug!(
"Raw series scan preflight rejected field type: region_id={}, column_id={}",
metadata.region_id, column.column_id,
);
return Ok(None);
}
fields.push((index, data_type.clone()));
specs.push(FieldSpec {
name: Arc::from(output.name.as_str()),
});
}
let time_is_millis = matches!(metadata.time_index_type(), TimestampType::Millisecond(_));
if fields.is_empty()
|| input.append_mode
|| input.series_row_selector.is_some()
|| !time_is_millis
{
debug!(
"Raw series scan preflight rejected scan options: region_id={}, selected_fields={}, append_mode={}, row_selector={}, time_is_millis={}",
metadata.region_id,
fields.len(),
input.append_mode,
input.series_row_selector.is_some(),
time_is_millis,
);
let Some(preparation) = self.prepare_raw_scan(label_source)? else {
return Ok(None);
}
let schema = Arc::new(SeriesSchema {
fields: specs.into(),
eval_grid: None,
});
let typed_schema = input.mapper.input_arrow_schema(false);
let mut predicates = Vec::new();
if let Some(predicate) = input.predicate_group().predicate_without_region() {
for expr in predicate.exprs() {
// Fresh execution properties cannot preserve Stable or query-dependent semantics.
if expr
.exists(|node| {
Ok(matches!(node, Expr::ScalarFunction(function)
if function.func.signature().volatility != Volatility::Immutable)
|| matches!(node, Expr::ScalarVariable(..)))
})
.unwrap_or(true)
{
debug!(
"Raw series scan preflight rejected predicate volatility or variable: region_id={}",
metadata.region_id,
);
return Ok(None);
}
// Evaluate residuals with original types before raw conversion to Float64.
let Ok(physical) = Predicate::to_physical_expr(expr, &typed_schema) else {
debug!(
"Raw series scan preflight failed predicate compilation: region_id={}",
metadata.region_id,
);
return Ok(None);
};
if !matches!(
physical.data_type(&typed_schema),
Ok(ArrowDataType::Boolean | ArrowDataType::Null)
) {
debug!(
"Raw series scan preflight rejected predicate type: region_id={}",
metadata.region_id,
);
return Ok(None);
}
predicates.push(physical);
}
}
};
let metrics = new_partition_metrics(
&self.stream_ctx,
ctx.explain_verbose,
@@ -760,13 +805,15 @@ impl RegionScanner for SeriesScan {
let batches = self
.scan_batch_in_partition(ctx, partition, metrics.clone(), metrics_set)
.map_err(BoxedError::new)?;
let ordering = self.raw_series_ordering();
Ok(Some(raw::make_stream(
batches,
schema,
preparation.schema,
label_source.clone(),
fields,
preparation.fields,
metrics,
predicates,
preparation.predicates,
ordering,
)))
}
src/mito2/src/read/series_scan/raw.rs
+8 -5
View File
@@ -128,6 +128,7 @@ fn chunk(
struct RawStream {
schema: Arc<SeriesSchema>,
label_source: RawLabelSource,
ordering: SeriesOrdering,
inner: BoxStream<'static, Result<RawSeriesChunk, BoxedError>>,
}
@@ -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<Arc<dyn PhysicalExpr>>,
ordering: SeriesOrdering,
) -> SendableRawSeriesStream {
let output_schema = schema.clone();
let stream = async_stream::stream! {
let mut input = inner;
let mut pending: Option<RawSeriesChunk> = 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),
})
}
src/mito2/src/read/series_scan/raw/tests.rs
+73 -3
View File
@@ -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<SeriesSchema>,
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<SeriesSchema>) -> SendableRawSeri
schema,
primary_key_source(metadata()),
ConcreteDataType::int64_datatype(),
SeriesOrdering::Contiguous,
)
}
fn convert_ordered(inner: ScanBatchStream, schema: Arc<SeriesSchema>) -> 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(),