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