1use std::collections::{HashMap, HashSet};
16use std::fmt;
17use std::path::PathBuf;
18use std::sync::Arc;
19use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
20
21use clap::Parser;
22use colored::Colorize;
23use common_base::Plugins;
24use common_error::ext::{BoxedError, PlainError};
25use common_error::status_code::StatusCode;
26use common_meta::cache::{new_schema_cache, new_table_schema_cache};
27use common_meta::key::SchemaMetadataManager;
28use common_meta::kv_backend::memory::MemoryKvBackend;
29use common_wal::config::DatanodeWalConfig;
30use common_wal::config::kafka::DatanodeKafkaConfig;
31use common_wal::config::raft_engine::RaftEngineConfig;
32use datafusion::execution::SessionStateBuilder;
33use datafusion::logical_expr::{BinaryExpr, Expr as DfExpr, ExprSchemable, Operator};
34use datafusion_common::tree_node::{Transformed, TreeNodeRewriter};
35use datafusion_common::{DFSchemaRef, ScalarValue, ToDFSchema};
36use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet;
37use datafusion_physical_plan::{DisplayAs, DisplayFormatType};
38use datatypes::arrow::compute;
39use futures::StreamExt;
40use futures::stream::FuturesUnordered;
41use log_store::kafka::log_store::KafkaLogStore;
42use log_store::noop::log_store::NoopLogStore;
43use log_store::raft_engine::log_store::RaftEngineLogStore;
44use mito2::config::MitoConfig;
45use mito2::engine::MitoEngine;
46use mito2::sst::file_ref::FileReferenceManager;
47use moka::future::CacheBuilder;
48use object_store::manager::ObjectStoreManager;
49use object_store::util::normalize_dir;
50use query::optimizer::parallelize_scan::ParallelizeScan;
51use serde::{Deserialize, Serialize};
52use snafu::{OptionExt, ResultExt};
53use sqlparser::ast::ExprWithAlias as SqlExprWithAlias;
54use sqlparser::dialect::GenericDialect;
55use sqlparser::parser::Parser as SqlParser;
56use store_api::metadata::RegionMetadata;
57use store_api::path_utils::WAL_DIR;
58use store_api::region_engine::{PrepareRequest, QueryScanContext, RegionEngine};
59use store_api::region_request::{RegionOpenRequest, RegionRequest};
60use store_api::storage::{RegionId, ScanRequest, TimeSeriesDistribution, TimeSeriesRowSelector};
61use tokio::fs;
62
63use crate::datanode::tool_util::{
64 build_object_store, format_bytes, parse_config, parse_path_type, parse_region_id,
65};
66use crate::error;
67
68struct VerboseScannerDisplay<'a, T: ?Sized>(&'a T);
70
71impl<T: DisplayAs + ?Sized> fmt::Display for VerboseScannerDisplay<'_, T> {
72 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
73 self.0.fmt_as(DisplayFormatType::Verbose, f)
74 }
75}
76
77struct PartitionScanStats {
78 partition: usize,
79 rows: u64,
80 batches: u64,
81 array_mem_size: u64,
82 estimated_size: u64,
83 first_batch_elapsed: Option<Duration>,
84 elapsed: Duration,
85}
86
87#[derive(Debug, Parser)]
89pub struct ScanbenchCommand {
90 #[clap(long, value_name = "FILE")]
92 config: PathBuf,
93
94 #[clap(long)]
96 region_id: String,
97
98 #[clap(long)]
100 table_dir: String,
101
102 #[clap(long, default_value = "seq")]
104 scanner: String,
105
106 #[clap(long, value_name = "FILE", conflicts_with = "scan_configs")]
108 scan_config: Option<PathBuf>,
109
110 #[clap(long, value_name = "FILE", conflicts_with = "scan_config")]
112 scan_configs: Option<PathBuf>,
113
114 #[clap(long, default_value = "1")]
116 parallelism: usize,
117
118 #[clap(long, default_value = "1")]
120 iterations: usize,
121
122 #[clap(long, default_value = "bare")]
124 path_type: String,
125
126 #[clap(short, long, default_value_t = false)]
128 verbose: bool,
129
130 #[clap(long, value_name = "FILE")]
132 pprof_file: Option<PathBuf>,
133
134 #[clap(long, value_name = "FILE")]
136 result_file: Option<PathBuf>,
137
138 #[clap(long, default_value_t = false)]
140 enable_wal: bool,
141
142 #[clap(long, default_value_t = false)]
144 pprof_after_warmup: bool,
145}
146
147#[derive(Debug, Clone, Deserialize, Default)]
149struct ScanConfig {
150 name: Option<String>,
151 projection: Option<Vec<usize>>,
152 projection_names: Option<Vec<String>>,
153 filters: Option<Vec<String>>,
154 series_row_selector: Option<String>,
155}
156
157struct ScanConfigSet {
158 configs: Vec<ScanConfig>,
159 is_suite: bool,
160}
161
162impl ScanConfigSet {
163 fn run_count(&self, iterations: usize) -> usize {
164 if self.is_suite {
165 self.configs.len()
166 } else {
167 iterations
168 }
169 }
170
171 fn query_index(&self, iteration: usize) -> usize {
172 if self.is_suite { iteration } else { 0 }
173 }
174}
175
176struct ResolvedScanConfig {
177 name: String,
178 projection: Option<Vec<usize>>,
179 filters: Vec<DfExpr>,
180 series_row_selector: Option<TimeSeriesRowSelector>,
181}
182
183struct QueryRunSummary {
184 name: String,
185 runs: u64,
186 total_rows: u64,
187 total_elapsed: Duration,
188}
189
190impl QueryRunSummary {
191 fn new(name: String) -> Self {
192 Self {
193 name,
194 runs: 0,
195 total_rows: 0,
196 total_elapsed: Duration::ZERO,
197 }
198 }
199
200 fn record(&mut self, rows: u64, elapsed: Duration) {
201 self.runs += 1;
202 self.total_rows += rows;
203 self.total_elapsed += elapsed;
204 }
205
206 fn mean_rows(&self) -> u64 {
207 self.total_rows.checked_div(self.runs).unwrap_or_default()
208 }
209
210 fn mean_elapsed(&self) -> Duration {
211 self.total_elapsed
212 .checked_div(self.runs as u32)
213 .unwrap_or_default()
214 }
215}
216
217const RESULT_FORMAT_VERSION: u32 = 1;
218
219#[derive(Debug, Serialize)]
220struct ScanbenchResult {
221 format_version: u32,
222 started_at_unix_ms: u64,
223 benchmark: BenchmarkMetadata,
224 runs: Vec<ScanRunResult>,
225 summary: BenchmarkResultSummary,
226}
227
228#[derive(Debug, Serialize)]
229struct BenchmarkMetadata {
230 scanner: String,
231 region_id: String,
232 region_id_u64: u64,
233 table_dir: String,
234 path_type: String,
235 parallelism: usize,
236 enable_wal: bool,
237 config_mode: String,
238 run_count: usize,
239}
240
241#[derive(Debug, Serialize)]
242struct NormalizedScanConfig {
243 name: String,
244 projection: Option<Vec<usize>>,
245 filters: Vec<String>,
246 series_row_selector: Option<String>,
247}
248
249#[derive(Debug, Serialize)]
250struct ScanRunResult {
251 iteration: usize,
252 query_index: usize,
253 name: String,
254 config: NormalizedScanConfig,
255 rows: u64,
256 batches: u64,
257 setup_elapsed_ns: u64,
258 scan_elapsed_ns: u64,
259 elapsed_ns: u64,
260 array_mem_size_bytes: u64,
261 estimated_size_bytes: u64,
262 partitions: Vec<PartitionResult>,
263 scanner_explain: String,
264}
265
266#[derive(Debug, Serialize)]
267struct PartitionResult {
268 partition: usize,
269 rows: u64,
270 batches: u64,
271 array_mem_size_bytes: u64,
272 estimated_size_bytes: u64,
273 first_batch_elapsed_ns: Option<u64>,
274 elapsed_ns: u64,
275}
276
277#[derive(Debug, Serialize)]
278struct BenchmarkResultSummary {
279 runs: u64,
280 total_rows: u64,
281 total_elapsed_ns: u64,
282 mean_rows: u64,
283 mean_elapsed_ns: u64,
284 queries: Vec<QueryResultSummary>,
285}
286
287#[derive(Debug, Serialize)]
288struct QueryResultSummary {
289 query_index: usize,
290 name: String,
291 runs: u64,
292 total_rows: u64,
293 total_elapsed_ns: u64,
294 mean_rows: u64,
295 mean_elapsed_ns: u64,
296}
297
298fn duration_ns(duration: Duration) -> u64 {
299 duration.as_nanos().try_into().unwrap_or(u64::MAX)
300}
301
302fn started_at_unix_ms() -> u64 {
303 SystemTime::now()
304 .duration_since(UNIX_EPOCH)
305 .unwrap_or_default()
306 .as_millis()
307 .try_into()
308 .unwrap_or(u64::MAX)
309}
310
311fn result_summary(
312 total_rows: u64,
313 total_elapsed: Duration,
314 summaries: &[QueryRunSummary],
315) -> BenchmarkResultSummary {
316 let runs = summaries.iter().map(|summary| summary.runs).sum::<u64>();
317 BenchmarkResultSummary {
318 runs,
319 total_rows,
320 total_elapsed_ns: duration_ns(total_elapsed),
321 mean_rows: total_rows.checked_div(runs).unwrap_or_default(),
322 mean_elapsed_ns: duration_ns(total_elapsed.checked_div(runs as u32).unwrap_or_default()),
323 queries: summaries
324 .iter()
325 .enumerate()
326 .map(|(index, summary)| QueryResultSummary {
327 query_index: index + 1,
328 name: summary.name.clone(),
329 runs: summary.runs,
330 total_rows: summary.total_rows,
331 total_elapsed_ns: duration_ns(summary.total_elapsed),
332 mean_rows: summary.mean_rows(),
333 mean_elapsed_ns: duration_ns(summary.mean_elapsed()),
334 })
335 .collect(),
336 }
337}
338
339async fn write_result_file(path: &PathBuf, result: &ScanbenchResult) -> error::Result<()> {
340 let content = serde_json::to_vec_pretty(result).context(error::SerdeJsonSnafu)?;
341 fs::write(path, content).await.context(error::FileIoSnafu)
342}
343
344fn validate_scan_config_suite(
345 mut configs: Vec<ScanConfig>,
346 iterations: usize,
347) -> error::Result<Vec<ScanConfig>> {
348 if iterations != 1 {
349 return Err(error::IllegalConfigSnafu {
350 msg: format!(
351 "--iterations cannot be used with --scan-configs (got {iterations}); the array length defines the run count"
352 ),
353 }
354 .build());
355 }
356 if configs.is_empty() {
357 return Err(error::IllegalConfigSnafu {
358 msg: "--scan-configs file must contain at least one scan config".to_string(),
359 }
360 .build());
361 }
362
363 let mut names = HashSet::with_capacity(configs.len());
364 for (index, config) in configs.iter_mut().enumerate() {
365 let name = match config.name.take() {
366 Some(name) => {
367 let name = name.trim();
368 if name.is_empty() {
369 return Err(error::IllegalConfigSnafu {
370 msg: format!("scan config at index {index} has an empty name"),
371 }
372 .build());
373 }
374 name.to_string()
375 }
376 None => format!("query-{:03}", index + 1),
377 };
378 if !names.insert(name.clone()) {
379 return Err(error::IllegalConfigSnafu {
380 msg: format!("duplicate scan config name '{name}' at index {index}"),
381 }
382 .build());
383 }
384 config.name = Some(name);
385 }
386
387 Ok(configs)
388}
389
390fn resolve_series_row_selector(
391 scan_config: &ScanConfig,
392) -> error::Result<Option<TimeSeriesRowSelector>> {
393 match scan_config.series_row_selector.as_deref() {
394 Some("last_row") => Ok(Some(TimeSeriesRowSelector::LastRow { after_merge: false })),
395 Some(other) => Err(error::IllegalConfigSnafu {
396 msg: format!("Unknown series_row_selector '{other}'"),
397 }
398 .build()),
399 None => Ok(None),
400 }
401}
402
403fn resolve_scan_configs(
404 config_set: &ScanConfigSet,
405 metadata: &RegionMetadata,
406) -> error::Result<Vec<ResolvedScanConfig>> {
407 config_set
408 .configs
409 .iter()
410 .enumerate()
411 .map(|(index, config)| {
412 let name = config
413 .name
414 .clone()
415 .unwrap_or_else(|| format!("query-{:03}", index + 1));
416 let projection = resolve_projection(config, Some(metadata)).map_err(|err| {
417 error::IllegalConfigSnafu {
418 msg: format!("invalid scan config at index {index} ('{name}'): {err}"),
419 }
420 .build()
421 })?;
422 let filters = resolve_filters(config, metadata).map_err(|err| {
423 error::IllegalConfigSnafu {
424 msg: format!("invalid scan config at index {index} ('{name}'): {err}"),
425 }
426 .build()
427 })?;
428 let series_row_selector = resolve_series_row_selector(config).map_err(|err| {
429 error::IllegalConfigSnafu {
430 msg: format!("invalid scan config at index {index} ('{name}'): {err}"),
431 }
432 .build()
433 })?;
434 Ok(ResolvedScanConfig {
435 name,
436 projection,
437 filters,
438 series_row_selector,
439 })
440 })
441 .collect()
442}
443
444fn resolve_projection(
445 scan_config: &ScanConfig,
446 metadata: Option<&RegionMetadata>,
447) -> error::Result<Option<Vec<usize>>> {
448 if scan_config.projection.is_some() && scan_config.projection_names.is_some() {
449 return Err(error::IllegalConfigSnafu {
450 msg: "scan config cannot contain both 'projection' and 'projection_names'".to_string(),
451 }
452 .build());
453 }
454
455 if let Some(projection) = &scan_config.projection {
456 return Ok(Some(projection.clone()));
457 }
458
459 if let Some(projection_names) = &scan_config.projection_names {
460 let metadata = metadata.context(error::IllegalConfigSnafu {
461 msg: "Missing region metadata while resolving 'projection_names'".to_string(),
462 })?;
463 let available_columns = metadata
464 .column_metadatas
465 .iter()
466 .map(|column| column.column_schema.name.as_str())
467 .collect::<Vec<_>>()
468 .join(", ");
469 let projection = projection_names
470 .iter()
471 .map(|name| {
472 metadata
473 .column_index_by_name(name)
474 .with_context(|| error::IllegalConfigSnafu {
475 msg: format!(
476 "Unknown column '{}' in projection_names, available columns: [{}]",
477 name, available_columns
478 ),
479 })
480 })
481 .collect::<error::Result<Vec<_>>>()?;
482 return Ok(Some(projection));
483 }
484
485 Ok(None)
486}
487
488struct LiteralTypeCaster {
490 schema: DFSchemaRef,
491}
492
493impl TreeNodeRewriter for LiteralTypeCaster {
494 type Node = DfExpr;
495
496 fn f_up(&mut self, expr: DfExpr) -> datafusion_common::Result<Transformed<DfExpr>> {
497 let DfExpr::BinaryExpr(BinaryExpr { left, op, right }) = &expr else {
498 return Ok(Transformed::no(expr));
499 };
500
501 if !matches!(
502 op,
503 Operator::Eq
504 | Operator::NotEq
505 | Operator::Lt
506 | Operator::LtEq
507 | Operator::Gt
508 | Operator::GtEq
509 ) {
510 return Ok(Transformed::no(expr));
511 }
512
513 let (col_expr, lit_expr, col_left) = match (left.as_ref(), right.as_ref()) {
514 (col @ DfExpr::Column(_), lit @ DfExpr::Literal(_, _)) => (col, lit, true),
515 (lit @ DfExpr::Literal(_, _), col @ DfExpr::Column(_)) => (col, lit, false),
516 _ => return Ok(Transformed::no(expr)),
517 };
518
519 let col_type = col_expr.get_type(self.schema.as_ref())?;
520 let DfExpr::Literal(scalar, _) = lit_expr else {
521 unreachable!()
522 };
523
524 if scalar.data_type() == col_type {
525 return Ok(Transformed::no(expr));
526 }
527
528 let lit_array = scalar.to_array()?;
529 let casted = compute::cast(lit_array.as_ref(), &col_type).map_err(|e| {
530 datafusion_common::DataFusionError::Internal(format!(
531 "Failed to cast literal {:?} to {:?}: {}",
532 scalar, col_type, e
533 ))
534 })?;
535 let casted_scalar = ScalarValue::try_from_array(&casted, 0)?;
536
537 let new_lit = DfExpr::Literal(casted_scalar, None);
538 let (new_left, new_right) = if col_left {
539 (left.clone(), Box::new(new_lit))
540 } else {
541 (Box::new(new_lit), right.clone())
542 };
543
544 Ok(Transformed::yes(DfExpr::BinaryExpr(BinaryExpr {
545 left: new_left,
546 op: *op,
547 right: new_right,
548 })))
549 }
550}
551
552fn convert_literal_types(
553 exprs: Vec<DfExpr>,
554 schema: &DFSchemaRef,
555) -> datafusion_common::Result<Vec<DfExpr>> {
556 use datafusion_common::tree_node::TreeNode;
557
558 let mut caster = LiteralTypeCaster {
559 schema: schema.clone(),
560 };
561 exprs
562 .into_iter()
563 .map(|e| e.rewrite(&mut caster).map(|x| x.data))
564 .collect()
565}
566
567fn resolve_filters(
568 scan_config: &ScanConfig,
569 metadata: &RegionMetadata,
570) -> error::Result<Vec<DfExpr>> {
571 let Some(filters) = &scan_config.filters else {
572 return Ok(Vec::new());
573 };
574
575 let df_schema = metadata
576 .schema
577 .arrow_schema()
578 .clone()
579 .to_dfschema()
580 .map_err(|e| {
581 error::IllegalConfigSnafu {
582 msg: format!("Failed to convert region schema to DataFusion schema: {e}"),
583 }
584 .build()
585 })?;
586
587 let state = SessionStateBuilder::new()
588 .with_config(Default::default())
589 .with_runtime_env(Default::default())
590 .with_default_features()
591 .build();
592
593 let exprs: Vec<DfExpr> = filters
594 .iter()
595 .enumerate()
596 .map(|(idx, filter)| {
597 let mut parser = SqlParser::new(&GenericDialect {})
598 .try_with_sql(filter)
599 .map_err(|e| {
600 error::IllegalConfigSnafu {
601 msg: format!("Invalid filter at index {idx} ('{filter}'): {e}"),
602 }
603 .build()
604 })?;
605
606 let sql_expr = parser.parse_expr().map_err(|e| {
607 error::IllegalConfigSnafu {
608 msg: format!("Invalid filter at index {idx} ('{filter}'): {e}"),
609 }
610 .build()
611 })?;
612
613 state
614 .create_logical_expr_from_sql_expr(
615 SqlExprWithAlias {
616 expr: sql_expr,
617 alias: None,
618 },
619 &df_schema,
620 )
621 .map_err(|e| {
622 error::IllegalConfigSnafu {
623 msg: format!(
624 "Failed to convert filter at index {idx} ('{filter}') to logical expr: {e}"
625 ),
626 }
627 .build()
628 })
629 })
630 .collect::<error::Result<Vec<_>>>()?;
631
632 let df_schema_ref = Arc::new(df_schema);
633 convert_literal_types(exprs, &df_schema_ref).map_err(|e| {
634 error::IllegalConfigSnafu {
635 msg: format!("Failed to convert filter expression types: {e}"),
636 }
637 .build()
638 })
639}
640
641fn noop_partition_expr_fetcher() -> mito2::region::opener::PartitionExprFetcherRef {
642 struct NoopPartitionExprFetcher;
643
644 #[async_trait::async_trait]
645 impl mito2::region::opener::PartitionExprFetcher for NoopPartitionExprFetcher {
646 async fn fetch_expr(&self, _region_id: RegionId) -> Option<String> {
647 None
648 }
649 }
650
651 Arc::new(NoopPartitionExprFetcher)
652}
653
654struct EngineComponents {
655 data_home: String,
656 mito_config: MitoConfig,
657 object_store_manager: Arc<ObjectStoreManager>,
658 schema_metadata_manager: Arc<SchemaMetadataManager>,
659 file_ref_manager: Arc<FileReferenceManager>,
660 partition_expr_fetcher: mito2::region::opener::PartitionExprFetcherRef,
661}
662
663impl EngineComponents {
664 async fn build<S: store_api::logstore::LogStore>(
665 self,
666 log_store: Arc<S>,
667 ) -> error::Result<MitoEngine> {
668 MitoEngine::new(
669 &self.data_home,
670 self.mito_config,
671 log_store,
672 self.object_store_manager,
673 self.schema_metadata_manager,
674 self.file_ref_manager,
675 self.partition_expr_fetcher,
676 Plugins::default(),
677 )
678 .await
679 .map_err(BoxedError::new)
680 .context(error::BuildCliSnafu)
681 }
682}
683
684fn mock_schema_metadata_manager() -> Arc<SchemaMetadataManager> {
685 let kv_backend = Arc::new(MemoryKvBackend::new());
686 let table_schema_cache = Arc::new(new_table_schema_cache(
687 "table_schema_name_cache".to_string(),
688 CacheBuilder::default().build(),
689 kv_backend.clone(),
690 ));
691 let schema_cache = Arc::new(new_schema_cache(
692 "schema_cache".to_string(),
693 CacheBuilder::default().build(),
694 kv_backend.clone(),
695 ));
696 Arc::new(SchemaMetadataManager::new(table_schema_cache, schema_cache))
697}
698
699#[derive(Debug, PartialEq)]
701enum ScanbenchLogStore<'a> {
702 RaftEngine(&'a RaftEngineConfig),
703 Kafka(&'a DatanodeKafkaConfig),
704 Noop,
705}
706
707fn select_log_store(
710 wal_config: &DatanodeWalConfig,
711 enable_wal: bool,
712) -> error::Result<ScanbenchLogStore<'_>> {
713 match wal_config {
714 DatanodeWalConfig::RaftEngine(config) if enable_wal => {
715 Ok(ScanbenchLogStore::RaftEngine(config))
716 }
717 DatanodeWalConfig::Kafka(config) if enable_wal => Ok(ScanbenchLogStore::Kafka(config)),
718 DatanodeWalConfig::ObjectStore(_) if enable_wal => error::IllegalConfigSnafu {
719 msg: "object store WAL is not wired into scanbench yet; disable WAL or use another provider",
720 }
721 .fail(),
722 _ => Ok(ScanbenchLogStore::Noop),
723 }
724}
725
726impl ScanbenchCommand {
727 async fn load_scan_config_set(&self) -> error::Result<ScanConfigSet> {
728 match (&self.scan_config, &self.scan_configs) {
729 (Some(_), Some(_)) => Err(error::IllegalConfigSnafu {
730 msg: "--scan-config and --scan-configs are mutually exclusive".to_string(),
731 }
732 .build()),
733 (Some(path), None) => {
734 let content = tokio::fs::read_to_string(path)
735 .await
736 .context(error::FileIoSnafu)?;
737 let config =
738 serde_json::from_str::<ScanConfig>(&content).context(error::SerdeJsonSnafu)?;
739 Ok(ScanConfigSet {
740 configs: vec![config],
741 is_suite: false,
742 })
743 }
744 (None, Some(path)) => {
745 if self.iterations != 1 {
746 return Err(error::IllegalConfigSnafu {
747 msg: format!(
748 "--iterations cannot be used with --scan-configs (got {}); the array length defines the run count",
749 self.iterations
750 ),
751 }
752 .build());
753 }
754 let content = tokio::fs::read_to_string(path)
755 .await
756 .context(error::FileIoSnafu)?;
757 let configs = serde_json::from_str::<Vec<ScanConfig>>(&content)
758 .context(error::SerdeJsonSnafu)?;
759 Ok(ScanConfigSet {
760 configs: validate_scan_config_suite(configs, self.iterations)?,
761 is_suite: true,
762 })
763 }
764 (None, None) => Ok(ScanConfigSet {
765 configs: vec![ScanConfig::default()],
766 is_suite: false,
767 }),
768 }
769 }
770
771 pub async fn run(&self) -> error::Result<()> {
772 if self.verbose {
773 common_telemetry::init_default_ut_logging();
774 }
775
776 println!("{}", "Starting scanbench...".cyan().bold());
777 let benchmark_started_at_unix_ms = started_at_unix_ms();
778
779 let scan_config_set = self.load_scan_config_set().await?;
780
781 let region_id = parse_region_id(&self.region_id)?;
782 let path_type = parse_path_type(&self.path_type)?;
783 println!(
784 "{} Region ID: {} (u64: {})",
785 "✓".green(),
786 self.region_id,
787 region_id.as_u64()
788 );
789
790 let (store_cfg, mito_config, wal_config) = parse_config(&self.config)?;
792 println!("{} Config parsed", "✓".green());
793
794 let object_store = build_object_store(&store_cfg).await?;
795 println!("{} Object store initialized", "✓".green());
796
797 let object_store_manager =
798 Arc::new(ObjectStoreManager::new("default", object_store.clone()));
799
800 let schema_metadata_manager = mock_schema_metadata_manager();
802 let file_ref_manager = Arc::new(FileReferenceManager::new(None));
803 let partition_expr_fetcher = noop_partition_expr_fetcher();
804
805 let components = EngineComponents {
807 data_home: store_cfg.data_home.clone(),
808 mito_config,
809 object_store_manager,
810 schema_metadata_manager,
811 file_ref_manager,
812 partition_expr_fetcher,
813 };
814
815 let engine = match select_log_store(&wal_config, self.enable_wal)? {
816 ScanbenchLogStore::RaftEngine(raft_engine_config) => {
817 let data_home = normalize_dir(&store_cfg.data_home);
818 let wal_dir = match &raft_engine_config.dir {
819 Some(dir) => dir.clone(),
820 None => format!("{}{WAL_DIR}", data_home),
821 };
822 fs::create_dir_all(&wal_dir).await.map_err(|e| {
823 error::IllegalConfigSnafu {
824 msg: format!("failed to create WAL directory {}: {e}", wal_dir),
825 }
826 .build()
827 })?;
828 let log_store = Arc::new(
829 RaftEngineLogStore::try_new(wal_dir, raft_engine_config)
830 .await
831 .map_err(BoxedError::new)
832 .context(error::BuildCliSnafu)?,
833 );
834 println!("{} Using RaftEngine WAL", "✓".green());
835 components.build(log_store).await?
836 }
837 ScanbenchLogStore::Kafka(kafka_config) => {
838 let log_store = Arc::new(
839 KafkaLogStore::try_new(kafka_config)
840 .await
841 .map_err(BoxedError::new)
842 .context(error::BuildCliSnafu)?,
843 );
844 println!("{} Using Kafka WAL", "✓".green());
845 components.build(log_store).await?
846 }
847 ScanbenchLogStore::Noop => {
848 let log_store = Arc::new(NoopLogStore);
849 println!(
850 "{} Using NoopLogStore (enable_wal={})",
851 "✓".green(),
852 self.enable_wal
853 );
854 components.build(log_store).await?
855 }
856 };
857
858 let open_request = RegionOpenRequest {
860 engine: "mito".to_string(),
861 table_dir: self.table_dir.clone(),
862 path_type,
863 options: HashMap::default(),
864 skip_wal_replay: !self.enable_wal,
865 checkpoint: None,
866 requirements: Default::default(),
867 };
868
869 engine
870 .handle_request(region_id, RegionRequest::Open(open_request))
871 .await
872 .map_err(BoxedError::new)
873 .context(error::BuildCliSnafu)?;
874 println!("{} Region opened", "✓".green());
875
876 let metadata = engine
877 .get_metadata(region_id)
878 .await
879 .map_err(BoxedError::new)
880 .context(error::BuildCliSnafu)?;
881 let scan_configs = resolve_scan_configs(&scan_config_set, &metadata)?;
882
883 let distribution = match self.scanner.as_str() {
885 "seq" => None,
886 "unordered" => Some(TimeSeriesDistribution::TimeWindowed),
887 "series" => Some(TimeSeriesDistribution::PerSeries),
888 other => {
889 return Err(error::IllegalConfigSnafu {
890 msg: format!(
891 "Unknown scanner type '{}', expected: seq, unordered, series",
892 other
893 ),
894 }
895 .build());
896 }
897 };
898
899 let run_count = scan_config_set.run_count(self.iterations);
900 if scan_config_set.is_suite {
901 println!(
902 "{} Scanner: {}, Parallelism: {}, Queries: {}",
903 "ℹ".blue(),
904 self.scanner,
905 self.parallelism,
906 run_count,
907 );
908 } else {
909 println!(
910 "{} Scanner: {}, Parallelism: {}, Iterations: {}",
911 "ℹ".blue(),
912 self.scanner,
913 self.parallelism,
914 run_count,
915 );
916 }
917
918 #[cfg(unix)]
920 let mut profiler_guard = if self.pprof_file.is_some() && !self.pprof_after_warmup {
921 println!("{} Starting profiling...", "⚡".yellow());
922 Some(
923 pprof::ProfilerGuardBuilder::default()
924 .frequency(99)
925 .blocklist(&["libc", "libgcc", "pthread", "vdso"])
926 .build()
927 .map_err(|e| {
928 BoxedError::new(PlainError::new(
929 format!("Failed to start profiler: {e}"),
930 StatusCode::Unexpected,
931 ))
932 })
933 .context(error::BuildCliSnafu)?,
934 )
935 } else {
936 None
937 };
938
939 #[cfg(not(unix))]
940 if self.pprof_file.is_some() {
941 eprintln!(
942 "{}: Profiling is not supported on this platform",
943 "Warning".yellow()
944 );
945 }
946
947 let mut total_rows_all = 0u64;
948 let mut total_elapsed_all = std::time::Duration::ZERO;
949 let mut run_results = Vec::with_capacity(if self.result_file.is_some() {
950 run_count
951 } else {
952 0
953 });
954 let mut query_summaries = scan_configs
955 .iter()
956 .map(|config| QueryRunSummary::new(config.name.clone()))
957 .collect::<Vec<_>>();
958 let collect_scanner_explain = self.verbose || self.result_file.is_some();
959
960 for iteration in 0..run_count {
961 let query_index = scan_config_set.query_index(iteration);
962 let scan_config = &scan_configs[query_index];
963 let request = ScanRequest {
964 projection: scan_config.projection.clone(),
965 filters: scan_config.filters.clone(),
966 series_row_selector: scan_config.series_row_selector,
967 distribution,
968 ..Default::default()
969 };
970
971 let start = Instant::now();
972
973 let mut scanner = engine
975 .handle_query(region_id, request)
976 .await
977 .map_err(BoxedError::new)
978 .context(error::BuildCliSnafu)?;
979
980 let original_partitions = scanner.properties().partitions.clone();
982 let total_ranges: usize = original_partitions.iter().map(|p| p.len()).sum();
983
984 if self.verbose {
985 println!(
986 " {} Original partitions: {}, total ranges: {}",
987 "ℹ".blue(),
988 original_partitions.len(),
989 total_ranges
990 );
991 }
992
993 if self.parallelism > 1 {
994 let all_ranges: Vec<_> = original_partitions.into_iter().flatten().collect();
996
997 let mut partitions =
999 ParallelizeScan::assign_partition_range(all_ranges, self.parallelism);
1000
1001 for partition in &mut partitions {
1003 partition.sort_by_key(|a| a.start);
1004 }
1005
1006 scanner
1007 .prepare(
1008 PrepareRequest::default()
1009 .with_ranges(partitions)
1010 .with_target_partitions(self.parallelism),
1011 )
1012 .map_err(BoxedError::new)
1013 .context(error::BuildCliSnafu)?;
1014 }
1015
1016 let num_partitions = scanner.properties().partitions.len();
1018 let ctx = QueryScanContext {
1019 explain_verbose: collect_scanner_explain,
1020 };
1021 let metrics_set = ExecutionPlanMetricsSet::new();
1022
1023 let mut scan_futures = FuturesUnordered::new();
1024 let setup_elapsed = start.elapsed();
1025 let scan_start = Instant::now();
1026
1027 for partition_idx in 0..num_partitions {
1028 let mut stream = scanner
1029 .scan_partition(&ctx, &metrics_set, partition_idx)
1030 .map_err(BoxedError::new)
1031 .context(error::BuildCliSnafu)?;
1032
1033 scan_futures.push(tokio::spawn(async move {
1034 let partition_start = Instant::now();
1035 let mut rows = 0u64;
1036 let mut batches = 0u64;
1037 let mut array_mem_size = 0u64;
1038 let mut estimated_size = 0u64;
1039 let mut first_batch_elapsed = None;
1040 while let Some(batch_result) = stream.next().await {
1041 match batch_result {
1042 Ok(batch) => {
1043 if first_batch_elapsed.is_none() {
1044 first_batch_elapsed = Some(partition_start.elapsed());
1045 }
1046 batches += 1;
1047 rows += batch.num_rows() as u64;
1048 let df_batch = batch.df_record_batch();
1049 array_mem_size += df_batch.get_array_memory_size() as u64;
1050 estimated_size +=
1051 mito2::memtable::record_batch_estimated_size(df_batch) as u64;
1052 }
1053 Err(e) => {
1054 return Err(BoxedError::new(e));
1055 }
1056 }
1057 }
1058 Ok::<PartitionScanStats, BoxedError>(PartitionScanStats {
1059 partition: partition_idx,
1060 rows,
1061 batches,
1062 array_mem_size,
1063 estimated_size,
1064 first_batch_elapsed,
1065 elapsed: partition_start.elapsed(),
1066 })
1067 }));
1068 }
1069
1070 let mut total_rows = 0u64;
1071 let mut total_batches = 0u64;
1072 let mut total_array_mem_size = 0u64;
1073 let mut total_estimated_size = 0u64;
1074 let mut partition_stats = Vec::with_capacity(num_partitions);
1075 while let Some(task) = scan_futures.next().await {
1076 let result = task
1077 .map_err(|e| {
1078 BoxedError::new(PlainError::new(
1079 format!("scan task failed: {e}"),
1080 StatusCode::Unexpected,
1081 ))
1082 })
1083 .context(error::BuildCliSnafu)?;
1084 let stats = result.context(error::BuildCliSnafu)?;
1085 total_rows += stats.rows;
1086 total_batches += stats.batches;
1087 total_array_mem_size += stats.array_mem_size;
1088 total_estimated_size += stats.estimated_size;
1089 partition_stats.push(stats);
1090 }
1091 let scan_elapsed = scan_start.elapsed();
1092
1093 let elapsed = start.elapsed();
1094 total_rows_all += total_rows;
1095 total_elapsed_all += elapsed;
1096 query_summaries[query_index].record(total_rows, elapsed);
1097
1098 let query_display = if scan_config_set.is_suite {
1099 format!(
1100 " [query {}/{}: {}]",
1101 query_index + 1,
1102 scan_configs.len(),
1103 scan_config.name
1104 )
1105 } else {
1106 String::new()
1107 };
1108
1109 println!(
1110 " [iter {}]{} {} rows in {:?} ({} partitions), array_mem_size: {}, estimated_size: {}",
1111 iteration + 1,
1112 query_display,
1113 total_rows.to_string().cyan(),
1114 elapsed,
1115 num_partitions,
1116 format_bytes(total_array_mem_size),
1117 format_bytes(total_estimated_size),
1118 );
1119
1120 if collect_scanner_explain {
1121 partition_stats.sort_unstable_by_key(|stats| stats.partition);
1122 }
1123
1124 if self.verbose {
1125 for stats in &partition_stats {
1126 let first_batch = stats
1127 .first_batch_elapsed
1128 .map(|elapsed| format!("{elapsed:?}"))
1129 .unwrap_or_else(|| "n/a".to_string());
1130 println!(
1131 " partition {}: rows={}, batches={}, first_batch={}, elapsed={:?}, array_mem_size={}, estimated_size={}",
1132 stats.partition,
1133 stats.rows,
1134 stats.batches,
1135 first_batch,
1136 stats.elapsed,
1137 format_bytes(stats.array_mem_size),
1138 format_bytes(stats.estimated_size),
1139 );
1140 }
1141 if !partition_stats.is_empty() {
1142 let total_partition_elapsed = partition_stats
1143 .iter()
1144 .map(|stats| stats.elapsed)
1145 .sum::<Duration>();
1146 let mean_partition_elapsed =
1147 total_partition_elapsed / partition_stats.len() as u32;
1148 let max_partition_elapsed = partition_stats
1149 .iter()
1150 .map(|stats| stats.elapsed)
1151 .max()
1152 .unwrap_or_default();
1153 let skew = max_partition_elapsed.as_secs_f64()
1154 / mean_partition_elapsed.as_secs_f64().max(f64::EPSILON);
1155 println!(
1156 " {} Timing: setup={:?}, scan={:?}, mean_partition={:?}, max_partition={:?}, partition_skew={:.2}x",
1157 "ℹ".blue(),
1158 setup_elapsed,
1159 scan_elapsed,
1160 mean_partition_elapsed,
1161 max_partition_elapsed,
1162 skew,
1163 );
1164 }
1165 }
1166
1167 let scanner_explain = if collect_scanner_explain {
1168 format!("{}", VerboseScannerDisplay(scanner.as_ref()))
1169 } else {
1170 String::new()
1171 };
1172 if self.verbose {
1173 println!(" {} Scanner explain: {}", "ℹ".blue(), scanner_explain);
1174 }
1175
1176 if self.result_file.is_some() {
1177 let source_config = &scan_config_set.configs[query_index];
1178 run_results.push(ScanRunResult {
1179 iteration: iteration + 1,
1180 query_index: query_index + 1,
1181 name: scan_config.name.clone(),
1182 config: NormalizedScanConfig {
1183 name: scan_config.name.clone(),
1184 projection: scan_config.projection.clone(),
1185 filters: source_config.filters.clone().unwrap_or_default(),
1186 series_row_selector: source_config.series_row_selector.clone(),
1187 },
1188 rows: total_rows,
1189 batches: total_batches,
1190 setup_elapsed_ns: duration_ns(setup_elapsed),
1191 scan_elapsed_ns: duration_ns(scan_elapsed),
1192 elapsed_ns: duration_ns(elapsed),
1193 array_mem_size_bytes: total_array_mem_size,
1194 estimated_size_bytes: total_estimated_size,
1195 partitions: partition_stats
1196 .iter()
1197 .map(|stats| PartitionResult {
1198 partition: stats.partition,
1199 rows: stats.rows,
1200 batches: stats.batches,
1201 array_mem_size_bytes: stats.array_mem_size,
1202 estimated_size_bytes: stats.estimated_size,
1203 first_batch_elapsed_ns: stats.first_batch_elapsed.map(duration_ns),
1204 elapsed_ns: duration_ns(stats.elapsed),
1205 })
1206 .collect(),
1207 scanner_explain,
1208 });
1209 }
1210
1211 #[cfg(unix)]
1213 if iteration == 0
1214 && self.pprof_after_warmup
1215 && self.pprof_file.is_some()
1216 && profiler_guard.is_none()
1217 {
1218 println!(
1219 "{} Starting profiling after warmup iteration...",
1220 "⚡".yellow()
1221 );
1222 profiler_guard = Some(
1223 pprof::ProfilerGuardBuilder::default()
1224 .frequency(99)
1225 .blocklist(&["libc", "libgcc", "pthread", "vdso"])
1226 .build()
1227 .map_err(|e| {
1228 BoxedError::new(PlainError::new(
1229 format!("Failed to start profiler: {e}"),
1230 StatusCode::Unexpected,
1231 ))
1232 })
1233 .context(error::BuildCliSnafu)?,
1234 );
1235 }
1236 }
1237
1238 #[cfg(unix)]
1240 if let (Some(guard), Some(pprof_file)) = (profiler_guard, &self.pprof_file) {
1241 println!("{} Generating flamegraph...", "🔥".yellow());
1242 match guard.report().build() {
1243 Ok(report) => {
1244 let mut flamegraph_data = Vec::new();
1245 if let Err(e) = report.flamegraph(&mut flamegraph_data) {
1246 println!("{}: Failed to generate flamegraph: {}", "Error".red(), e);
1247 } else if let Err(e) = std::fs::write(pprof_file, flamegraph_data) {
1248 println!(
1249 "{}: Failed to write flamegraph to {}: {}",
1250 "Error".red(),
1251 pprof_file.display(),
1252 e
1253 );
1254 } else {
1255 println!(
1256 "{} Flamegraph saved to {}",
1257 "✓".green(),
1258 pprof_file.display().to_string().cyan()
1259 );
1260 }
1261 }
1262 Err(e) => {
1263 println!("{}: Failed to generate pprof report: {}", "Error".red(), e);
1264 }
1265 }
1266 }
1267
1268 if scan_config_set.is_suite {
1270 let avg_elapsed = total_elapsed_all / run_count as u32;
1271 let avg_rows = total_rows_all / run_count as u64;
1272 println!(
1273 "\n{} Overall average: {} rows in {:?} over {} queries",
1274 "Summary".green().bold(),
1275 avg_rows.to_string().cyan(),
1276 avg_elapsed,
1277 run_count,
1278 );
1279 println!("{} Per-query:", "Summary".green().bold());
1280 for (index, summary) in query_summaries.iter().enumerate() {
1281 println!(
1282 " [{}] {}: runs={}, mean_rows={}, mean_elapsed={:?}",
1283 index + 1,
1284 summary.name,
1285 summary.runs,
1286 summary.mean_rows(),
1287 summary.mean_elapsed(),
1288 );
1289 }
1290 } else if run_count > 1 {
1291 let avg_elapsed = total_elapsed_all / run_count as u32;
1292 let avg_rows = total_rows_all / run_count as u64;
1293 println!(
1294 "\n{} Average: {} rows in {:?} over {} iterations",
1295 "Summary".green().bold(),
1296 avg_rows.to_string().cyan(),
1297 avg_elapsed,
1298 run_count,
1299 );
1300 }
1301
1302 if let Some(result_file) = &self.result_file {
1303 let result = ScanbenchResult {
1304 format_version: RESULT_FORMAT_VERSION,
1305 started_at_unix_ms: benchmark_started_at_unix_ms,
1306 benchmark: BenchmarkMetadata {
1307 scanner: self.scanner.clone(),
1308 region_id: self.region_id.clone(),
1309 region_id_u64: region_id.as_u64(),
1310 table_dir: self.table_dir.clone(),
1311 path_type: self.path_type.clone(),
1312 parallelism: self.parallelism,
1313 enable_wal: self.enable_wal,
1314 config_mode: if scan_config_set.is_suite {
1315 "suite".to_string()
1316 } else {
1317 "single".to_string()
1318 },
1319 run_count,
1320 },
1321 runs: run_results,
1322 summary: result_summary(total_rows_all, total_elapsed_all, &query_summaries),
1323 };
1324 write_result_file(result_file, &result).await?;
1325 println!(
1326 "{} Results saved to {}",
1327 "✓".green(),
1328 result_file.display().to_string().cyan()
1329 );
1330 }
1331
1332 println!("\n{}", "Benchmark completed!".green().bold());
1333 Ok(())
1334 }
1335}
1336
1337#[cfg(test)]
1338mod tests {
1339 use common_wal::config::DatanodeWalConfig;
1340 use common_wal::config::object_store::ObjectStoreWalConfig;
1341 use datatypes::prelude::ConcreteDataType;
1342 use datatypes::schema::ColumnSchema;
1343 use sqlparser::ast::{BinaryOperator, Expr};
1344 use sqlparser::dialect::GenericDialect;
1345 use sqlparser::parser::Parser;
1346 use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder};
1347 use store_api::storage::RegionId;
1348
1349 use super::{
1350 BenchmarkMetadata, BenchmarkResultSummary, NormalizedScanConfig, PartitionResult,
1351 ScanConfig, ScanRunResult, ScanbenchLogStore, ScanbenchResult, resolve_filters,
1352 resolve_projection, select_log_store, validate_scan_config_suite, write_result_file,
1353 };
1354 use crate::error;
1355
1356 #[test]
1357 fn test_select_log_store() {
1358 let object_store = DatanodeWalConfig::ObjectStore(ObjectStoreWalConfig::default());
1359 let err = select_log_store(&object_store, true).unwrap_err();
1360 assert!(matches!(err, error::Error::IllegalConfig { .. }));
1361 assert_eq!(
1362 ScanbenchLogStore::Noop,
1363 select_log_store(&object_store, false).unwrap()
1364 );
1365
1366 assert_eq!(
1367 ScanbenchLogStore::Noop,
1368 select_log_store(&DatanodeWalConfig::Noop, true).unwrap()
1369 );
1370 let raft_engine = DatanodeWalConfig::default();
1371 assert!(matches!(
1372 select_log_store(&raft_engine, true).unwrap(),
1373 ScanbenchLogStore::RaftEngine(_)
1374 ));
1375 assert_eq!(
1376 ScanbenchLogStore::Noop,
1377 select_log_store(&raft_engine, false).unwrap()
1378 );
1379 }
1380
1381 #[test]
1382 fn test_parse_scan_config_projection_names() {
1383 let json = r#"{"projection_names":["host","ts"]}"#;
1384 let config: ScanConfig = serde_json::from_str(json).unwrap();
1385
1386 assert_eq!(
1387 config.projection_names,
1388 Some(vec!["host".to_string(), "ts".to_string()])
1389 );
1390 assert_eq!(config.projection, None);
1391 }
1392
1393 #[test]
1394 fn test_resolve_projection_by_indexes() -> error::Result<()> {
1395 let config = ScanConfig {
1396 name: None,
1397 projection: Some(vec![0, 2]),
1398 projection_names: None,
1399 filters: None,
1400 series_row_selector: None,
1401 };
1402
1403 let projection = resolve_projection(&config, None)?;
1404 assert_eq!(projection, Some(vec![0, 2]));
1405 Ok(())
1406 }
1407
1408 #[test]
1409 fn test_resolve_projection_by_names_without_metadata() {
1410 let config = ScanConfig {
1411 name: None,
1412 projection: None,
1413 projection_names: Some(vec!["cpu".to_string(), "host".to_string()]),
1414 filters: None,
1415 series_row_selector: None,
1416 };
1417
1418 let err = resolve_projection(&config, None).unwrap_err();
1419 assert!(
1420 err.to_string()
1421 .contains("Missing region metadata while resolving 'projection_names'")
1422 );
1423 }
1424
1425 #[test]
1426 fn test_resolve_projection_conflict_fields() {
1427 let config = ScanConfig {
1428 name: None,
1429 projection: Some(vec![0]),
1430 projection_names: Some(vec!["host".to_string()]),
1431 filters: None,
1432 series_row_selector: None,
1433 };
1434
1435 let err = resolve_projection(&config, None).unwrap_err();
1436 let msg = err.to_string();
1437 assert!(msg.contains("projection"));
1438 assert!(msg.contains("projection_names"));
1439 }
1440
1441 #[test]
1442 fn test_sqlparser_parse_expr_string() {
1443 let dialect = GenericDialect {};
1444 let mut parser = Parser::new(&dialect)
1445 .try_with_sql("host = 'web-1' AND cpu > 80")
1446 .unwrap();
1447
1448 let expr = parser.parse_expr().unwrap();
1449
1450 match expr {
1451 Expr::BinaryOp { op, .. } => assert_eq!(op, BinaryOperator::And),
1452 other => panic!("expected BinaryOp, got: {other:?}"),
1453 }
1454 }
1455
1456 #[test]
1457 fn test_resolve_filters_uint32_type_conversion() {
1458 use api::v1::SemanticType;
1459
1460 let mut builder = RegionMetadataBuilder::new(RegionId::new(1, 0));
1461 builder
1462 .push_column_metadata(ColumnMetadata {
1463 column_schema: ColumnSchema::new(
1464 "table_id",
1465 ConcreteDataType::uint32_datatype(),
1466 false,
1467 ),
1468 semantic_type: SemanticType::Tag,
1469 column_id: 1,
1470 })
1471 .push_column_metadata(ColumnMetadata {
1472 column_schema: ColumnSchema::new(
1473 "ts",
1474 ConcreteDataType::timestamp_millisecond_datatype(),
1475 false,
1476 ),
1477 semantic_type: SemanticType::Timestamp,
1478 column_id: 2,
1479 })
1480 .primary_key(vec![1]);
1481 let metadata = builder.build().unwrap();
1482
1483 let config = ScanConfig {
1484 name: None,
1485 projection: None,
1486 projection_names: None,
1487 filters: Some(vec!["table_id = 1117".to_string()]),
1488 series_row_selector: None,
1489 };
1490
1491 let exprs = resolve_filters(&config, &metadata).unwrap();
1492 assert_eq!(exprs.len(), 1);
1493 let expr_str = format!("{}", exprs[0]);
1495 assert!(
1496 expr_str.contains("UInt32(1117)"),
1497 "Expected UInt32(1117) in expression, got: {expr_str}"
1498 );
1499 }
1500
1501 #[test]
1502 fn test_parse_scan_config_filters() {
1503 let json = r#"{"filters":["host = 'web-1'","cpu > 80"]}"#;
1504 let config: ScanConfig = serde_json::from_str(json).unwrap();
1505
1506 assert_eq!(
1507 config.filters,
1508 Some(vec!["host = 'web-1'".to_string(), "cpu > 80".to_string()])
1509 );
1510 }
1511
1512 #[test]
1513 fn test_parse_and_validate_scan_config_suite() {
1514 let json = r#"[
1515 {"name":" cold ","filters":["host = 'web-1'"]},
1516 {"projection_names":["host","cpu"]}
1517 ]"#;
1518 let configs: Vec<ScanConfig> = serde_json::from_str(json).unwrap();
1519 let configs = validate_scan_config_suite(configs, 1).unwrap();
1520
1521 assert_eq!(2, configs.len());
1522 assert_eq!(Some("cold"), configs[0].name.as_deref());
1523 assert_eq!(Some("query-002"), configs[1].name.as_deref());
1524 assert_eq!(
1525 Some(&vec!["host = 'web-1'".to_string()]),
1526 configs[0].filters.as_ref()
1527 );
1528 }
1529
1530 #[test]
1531 fn test_validate_scan_config_suite_rejects_invalid_input() {
1532 let empty = validate_scan_config_suite(Vec::new(), 1)
1533 .unwrap_err()
1534 .to_string();
1535 assert!(empty.contains("at least one"));
1536
1537 let iterations = validate_scan_config_suite(vec![ScanConfig::default()], 2)
1538 .unwrap_err()
1539 .to_string();
1540 assert!(iterations.contains("--iterations"));
1541
1542 let configs: Vec<ScanConfig> =
1543 serde_json::from_str(r#"[{"name":"query"},{"name":" query "}]"#).unwrap();
1544 let duplicate = validate_scan_config_suite(configs, 1)
1545 .unwrap_err()
1546 .to_string();
1547 assert!(duplicate.contains("duplicate"));
1548
1549 let configs: Vec<ScanConfig> = serde_json::from_str(r#"[{"name":" "}]"#).unwrap();
1550 let blank = validate_scan_config_suite(configs, 1)
1551 .unwrap_err()
1552 .to_string();
1553 assert!(blank.contains("empty name"));
1554 }
1555
1556 #[tokio::test]
1557 async fn test_write_result_file_overwrites_complete_result() {
1558 let dir = tempfile::tempdir().unwrap();
1559 let path = dir.path().join("scanbench-result.json");
1560 tokio::fs::write(&path, b"old result").await.unwrap();
1561
1562 let result = ScanbenchResult {
1563 format_version: 1,
1564 started_at_unix_ms: 42,
1565 benchmark: BenchmarkMetadata {
1566 scanner: "seq".to_string(),
1567 region_id: "1024:0".to_string(),
1568 region_id_u64: RegionId::new(1024, 0).as_u64(),
1569 table_dir: "greptime/public/1024".to_string(),
1570 path_type: "bare".to_string(),
1571 parallelism: 8,
1572 enable_wal: false,
1573 config_mode: "suite".to_string(),
1574 run_count: 1,
1575 },
1576 runs: vec![ScanRunResult {
1577 iteration: 1,
1578 query_index: 1,
1579 name: "cpu-host-1".to_string(),
1580 config: NormalizedScanConfig {
1581 name: "cpu-host-1".to_string(),
1582 projection: Some(vec![1, 2]),
1583 filters: vec!["hostname = 'host_1'".to_string()],
1584 series_row_selector: None,
1585 },
1586 rows: 100,
1587 batches: 2,
1588 setup_elapsed_ns: 10,
1589 scan_elapsed_ns: 80,
1590 elapsed_ns: 90,
1591 array_mem_size_bytes: 1024,
1592 estimated_size_bytes: 512,
1593 partitions: vec![PartitionResult {
1594 partition: 0,
1595 rows: 100,
1596 batches: 2,
1597 array_mem_size_bytes: 1024,
1598 estimated_size_bytes: 512,
1599 first_batch_elapsed_ns: Some(20),
1600 elapsed_ns: 80,
1601 }],
1602 scanner_explain: "SeqScan: region=1024(0)".to_string(),
1603 }],
1604 summary: BenchmarkResultSummary {
1605 runs: 1,
1606 total_rows: 100,
1607 total_elapsed_ns: 90,
1608 mean_rows: 100,
1609 mean_elapsed_ns: 90,
1610 queries: vec![],
1611 },
1612 };
1613
1614 write_result_file(&path, &result).await.unwrap();
1615 let content = tokio::fs::read(&path).await.unwrap();
1616 let actual: serde_json::Value = serde_json::from_slice(&content).unwrap();
1617
1618 assert_eq!(1, actual["format_version"]);
1619 assert_eq!("suite", actual["benchmark"]["config_mode"]);
1620 assert_eq!(2, actual["runs"][0]["config"]["projection"][1]);
1621 assert_eq!(100, actual["runs"][0]["partitions"][0]["rows"]);
1622 assert_eq!(
1623 "SeqScan: region=1024(0)",
1624 actual["runs"][0]["scanner_explain"]
1625 );
1626 assert_eq!(90, actual["summary"]["mean_elapsed_ns"]);
1627 }
1628}