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 ..Default::default()
1021 };
1022 let metrics_set = ExecutionPlanMetricsSet::new();
1023
1024 let mut scan_futures = FuturesUnordered::new();
1025 let setup_elapsed = start.elapsed();
1026 let scan_start = Instant::now();
1027
1028 for partition_idx in 0..num_partitions {
1029 let mut stream = scanner
1030 .scan_partition(&ctx, &metrics_set, partition_idx)
1031 .map_err(BoxedError::new)
1032 .context(error::BuildCliSnafu)?;
1033
1034 scan_futures.push(tokio::spawn(async move {
1035 let partition_start = Instant::now();
1036 let mut rows = 0u64;
1037 let mut batches = 0u64;
1038 let mut array_mem_size = 0u64;
1039 let mut estimated_size = 0u64;
1040 let mut first_batch_elapsed = None;
1041 while let Some(batch_result) = stream.next().await {
1042 match batch_result {
1043 Ok(batch) => {
1044 if first_batch_elapsed.is_none() {
1045 first_batch_elapsed = Some(partition_start.elapsed());
1046 }
1047 batches += 1;
1048 rows += batch.num_rows() as u64;
1049 let df_batch = batch.df_record_batch();
1050 array_mem_size += df_batch.get_array_memory_size() as u64;
1051 estimated_size +=
1052 mito2::memtable::record_batch_estimated_size(df_batch) as u64;
1053 }
1054 Err(e) => {
1055 return Err(BoxedError::new(e));
1056 }
1057 }
1058 }
1059 Ok::<PartitionScanStats, BoxedError>(PartitionScanStats {
1060 partition: partition_idx,
1061 rows,
1062 batches,
1063 array_mem_size,
1064 estimated_size,
1065 first_batch_elapsed,
1066 elapsed: partition_start.elapsed(),
1067 })
1068 }));
1069 }
1070
1071 let mut total_rows = 0u64;
1072 let mut total_batches = 0u64;
1073 let mut total_array_mem_size = 0u64;
1074 let mut total_estimated_size = 0u64;
1075 let mut partition_stats = Vec::with_capacity(num_partitions);
1076 while let Some(task) = scan_futures.next().await {
1077 let result = task
1078 .map_err(|e| {
1079 BoxedError::new(PlainError::new(
1080 format!("scan task failed: {e}"),
1081 StatusCode::Unexpected,
1082 ))
1083 })
1084 .context(error::BuildCliSnafu)?;
1085 let stats = result.context(error::BuildCliSnafu)?;
1086 total_rows += stats.rows;
1087 total_batches += stats.batches;
1088 total_array_mem_size += stats.array_mem_size;
1089 total_estimated_size += stats.estimated_size;
1090 partition_stats.push(stats);
1091 }
1092 let scan_elapsed = scan_start.elapsed();
1093
1094 let elapsed = start.elapsed();
1095 total_rows_all += total_rows;
1096 total_elapsed_all += elapsed;
1097 query_summaries[query_index].record(total_rows, elapsed);
1098
1099 let query_display = if scan_config_set.is_suite {
1100 format!(
1101 " [query {}/{}: {}]",
1102 query_index + 1,
1103 scan_configs.len(),
1104 scan_config.name
1105 )
1106 } else {
1107 String::new()
1108 };
1109
1110 println!(
1111 " [iter {}]{} {} rows in {:?} ({} partitions), array_mem_size: {}, estimated_size: {}",
1112 iteration + 1,
1113 query_display,
1114 total_rows.to_string().cyan(),
1115 elapsed,
1116 num_partitions,
1117 format_bytes(total_array_mem_size),
1118 format_bytes(total_estimated_size),
1119 );
1120
1121 if collect_scanner_explain {
1122 partition_stats.sort_unstable_by_key(|stats| stats.partition);
1123 }
1124
1125 if self.verbose {
1126 for stats in &partition_stats {
1127 let first_batch = stats
1128 .first_batch_elapsed
1129 .map(|elapsed| format!("{elapsed:?}"))
1130 .unwrap_or_else(|| "n/a".to_string());
1131 println!(
1132 " partition {}: rows={}, batches={}, first_batch={}, elapsed={:?}, array_mem_size={}, estimated_size={}",
1133 stats.partition,
1134 stats.rows,
1135 stats.batches,
1136 first_batch,
1137 stats.elapsed,
1138 format_bytes(stats.array_mem_size),
1139 format_bytes(stats.estimated_size),
1140 );
1141 }
1142 if !partition_stats.is_empty() {
1143 let total_partition_elapsed = partition_stats
1144 .iter()
1145 .map(|stats| stats.elapsed)
1146 .sum::<Duration>();
1147 let mean_partition_elapsed =
1148 total_partition_elapsed / partition_stats.len() as u32;
1149 let max_partition_elapsed = partition_stats
1150 .iter()
1151 .map(|stats| stats.elapsed)
1152 .max()
1153 .unwrap_or_default();
1154 let skew = max_partition_elapsed.as_secs_f64()
1155 / mean_partition_elapsed.as_secs_f64().max(f64::EPSILON);
1156 println!(
1157 " {} Timing: setup={:?}, scan={:?}, mean_partition={:?}, max_partition={:?}, partition_skew={:.2}x",
1158 "ℹ".blue(),
1159 setup_elapsed,
1160 scan_elapsed,
1161 mean_partition_elapsed,
1162 max_partition_elapsed,
1163 skew,
1164 );
1165 }
1166 }
1167
1168 let scanner_explain = if collect_scanner_explain {
1169 format!("{}", VerboseScannerDisplay(scanner.as_ref()))
1170 } else {
1171 String::new()
1172 };
1173 if self.verbose {
1174 println!(" {} Scanner explain: {}", "ℹ".blue(), scanner_explain);
1175 }
1176
1177 if self.result_file.is_some() {
1178 let source_config = &scan_config_set.configs[query_index];
1179 run_results.push(ScanRunResult {
1180 iteration: iteration + 1,
1181 query_index: query_index + 1,
1182 name: scan_config.name.clone(),
1183 config: NormalizedScanConfig {
1184 name: scan_config.name.clone(),
1185 projection: scan_config.projection.clone(),
1186 filters: source_config.filters.clone().unwrap_or_default(),
1187 series_row_selector: source_config.series_row_selector.clone(),
1188 },
1189 rows: total_rows,
1190 batches: total_batches,
1191 setup_elapsed_ns: duration_ns(setup_elapsed),
1192 scan_elapsed_ns: duration_ns(scan_elapsed),
1193 elapsed_ns: duration_ns(elapsed),
1194 array_mem_size_bytes: total_array_mem_size,
1195 estimated_size_bytes: total_estimated_size,
1196 partitions: partition_stats
1197 .iter()
1198 .map(|stats| PartitionResult {
1199 partition: stats.partition,
1200 rows: stats.rows,
1201 batches: stats.batches,
1202 array_mem_size_bytes: stats.array_mem_size,
1203 estimated_size_bytes: stats.estimated_size,
1204 first_batch_elapsed_ns: stats.first_batch_elapsed.map(duration_ns),
1205 elapsed_ns: duration_ns(stats.elapsed),
1206 })
1207 .collect(),
1208 scanner_explain,
1209 });
1210 }
1211
1212 #[cfg(unix)]
1214 if iteration == 0
1215 && self.pprof_after_warmup
1216 && self.pprof_file.is_some()
1217 && profiler_guard.is_none()
1218 {
1219 println!(
1220 "{} Starting profiling after warmup iteration...",
1221 "⚡".yellow()
1222 );
1223 profiler_guard = Some(
1224 pprof::ProfilerGuardBuilder::default()
1225 .frequency(99)
1226 .blocklist(&["libc", "libgcc", "pthread", "vdso"])
1227 .build()
1228 .map_err(|e| {
1229 BoxedError::new(PlainError::new(
1230 format!("Failed to start profiler: {e}"),
1231 StatusCode::Unexpected,
1232 ))
1233 })
1234 .context(error::BuildCliSnafu)?,
1235 );
1236 }
1237 }
1238
1239 #[cfg(unix)]
1241 if let (Some(guard), Some(pprof_file)) = (profiler_guard, &self.pprof_file) {
1242 println!("{} Generating flamegraph...", "🔥".yellow());
1243 match guard.report().build() {
1244 Ok(report) => {
1245 let mut flamegraph_data = Vec::new();
1246 if let Err(e) = report.flamegraph(&mut flamegraph_data) {
1247 println!("{}: Failed to generate flamegraph: {}", "Error".red(), e);
1248 } else if let Err(e) = std::fs::write(pprof_file, flamegraph_data) {
1249 println!(
1250 "{}: Failed to write flamegraph to {}: {}",
1251 "Error".red(),
1252 pprof_file.display(),
1253 e
1254 );
1255 } else {
1256 println!(
1257 "{} Flamegraph saved to {}",
1258 "✓".green(),
1259 pprof_file.display().to_string().cyan()
1260 );
1261 }
1262 }
1263 Err(e) => {
1264 println!("{}: Failed to generate pprof report: {}", "Error".red(), e);
1265 }
1266 }
1267 }
1268
1269 if scan_config_set.is_suite {
1271 let avg_elapsed = total_elapsed_all / run_count as u32;
1272 let avg_rows = total_rows_all / run_count as u64;
1273 println!(
1274 "\n{} Overall average: {} rows in {:?} over {} queries",
1275 "Summary".green().bold(),
1276 avg_rows.to_string().cyan(),
1277 avg_elapsed,
1278 run_count,
1279 );
1280 println!("{} Per-query:", "Summary".green().bold());
1281 for (index, summary) in query_summaries.iter().enumerate() {
1282 println!(
1283 " [{}] {}: runs={}, mean_rows={}, mean_elapsed={:?}",
1284 index + 1,
1285 summary.name,
1286 summary.runs,
1287 summary.mean_rows(),
1288 summary.mean_elapsed(),
1289 );
1290 }
1291 } else if run_count > 1 {
1292 let avg_elapsed = total_elapsed_all / run_count as u32;
1293 let avg_rows = total_rows_all / run_count as u64;
1294 println!(
1295 "\n{} Average: {} rows in {:?} over {} iterations",
1296 "Summary".green().bold(),
1297 avg_rows.to_string().cyan(),
1298 avg_elapsed,
1299 run_count,
1300 );
1301 }
1302
1303 if let Some(result_file) = &self.result_file {
1304 let result = ScanbenchResult {
1305 format_version: RESULT_FORMAT_VERSION,
1306 started_at_unix_ms: benchmark_started_at_unix_ms,
1307 benchmark: BenchmarkMetadata {
1308 scanner: self.scanner.clone(),
1309 region_id: self.region_id.clone(),
1310 region_id_u64: region_id.as_u64(),
1311 table_dir: self.table_dir.clone(),
1312 path_type: self.path_type.clone(),
1313 parallelism: self.parallelism,
1314 enable_wal: self.enable_wal,
1315 config_mode: if scan_config_set.is_suite {
1316 "suite".to_string()
1317 } else {
1318 "single".to_string()
1319 },
1320 run_count,
1321 },
1322 runs: run_results,
1323 summary: result_summary(total_rows_all, total_elapsed_all, &query_summaries),
1324 };
1325 write_result_file(result_file, &result).await?;
1326 println!(
1327 "{} Results saved to {}",
1328 "✓".green(),
1329 result_file.display().to_string().cyan()
1330 );
1331 }
1332
1333 println!("\n{}", "Benchmark completed!".green().bold());
1334 Ok(())
1335 }
1336}
1337
1338#[cfg(test)]
1339mod tests {
1340 use common_wal::config::DatanodeWalConfig;
1341 use common_wal::config::object_store::ObjectStoreWalConfig;
1342 use datatypes::prelude::ConcreteDataType;
1343 use datatypes::schema::ColumnSchema;
1344 use sqlparser::ast::{BinaryOperator, Expr};
1345 use sqlparser::dialect::GenericDialect;
1346 use sqlparser::parser::Parser;
1347 use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder};
1348 use store_api::storage::RegionId;
1349
1350 use super::{
1351 BenchmarkMetadata, BenchmarkResultSummary, NormalizedScanConfig, PartitionResult,
1352 ScanConfig, ScanRunResult, ScanbenchLogStore, ScanbenchResult, resolve_filters,
1353 resolve_projection, select_log_store, validate_scan_config_suite, write_result_file,
1354 };
1355 use crate::error;
1356
1357 #[test]
1358 fn test_select_log_store() {
1359 let object_store = DatanodeWalConfig::ObjectStore(ObjectStoreWalConfig::default());
1360 let err = select_log_store(&object_store, true).unwrap_err();
1361 assert!(matches!(err, error::Error::IllegalConfig { .. }));
1362 assert_eq!(
1363 ScanbenchLogStore::Noop,
1364 select_log_store(&object_store, false).unwrap()
1365 );
1366
1367 assert_eq!(
1368 ScanbenchLogStore::Noop,
1369 select_log_store(&DatanodeWalConfig::Noop, true).unwrap()
1370 );
1371 let raft_engine = DatanodeWalConfig::default();
1372 assert!(matches!(
1373 select_log_store(&raft_engine, true).unwrap(),
1374 ScanbenchLogStore::RaftEngine(_)
1375 ));
1376 assert_eq!(
1377 ScanbenchLogStore::Noop,
1378 select_log_store(&raft_engine, false).unwrap()
1379 );
1380 }
1381
1382 #[test]
1383 fn test_parse_scan_config_projection_names() {
1384 let json = r#"{"projection_names":["host","ts"]}"#;
1385 let config: ScanConfig = serde_json::from_str(json).unwrap();
1386
1387 assert_eq!(
1388 config.projection_names,
1389 Some(vec!["host".to_string(), "ts".to_string()])
1390 );
1391 assert_eq!(config.projection, None);
1392 }
1393
1394 #[test]
1395 fn test_resolve_projection_by_indexes() -> error::Result<()> {
1396 let config = ScanConfig {
1397 name: None,
1398 projection: Some(vec![0, 2]),
1399 projection_names: None,
1400 filters: None,
1401 series_row_selector: None,
1402 };
1403
1404 let projection = resolve_projection(&config, None)?;
1405 assert_eq!(projection, Some(vec![0, 2]));
1406 Ok(())
1407 }
1408
1409 #[test]
1410 fn test_resolve_projection_by_names_without_metadata() {
1411 let config = ScanConfig {
1412 name: None,
1413 projection: None,
1414 projection_names: Some(vec!["cpu".to_string(), "host".to_string()]),
1415 filters: None,
1416 series_row_selector: None,
1417 };
1418
1419 let err = resolve_projection(&config, None).unwrap_err();
1420 assert!(
1421 err.to_string()
1422 .contains("Missing region metadata while resolving 'projection_names'")
1423 );
1424 }
1425
1426 #[test]
1427 fn test_resolve_projection_conflict_fields() {
1428 let config = ScanConfig {
1429 name: None,
1430 projection: Some(vec![0]),
1431 projection_names: Some(vec!["host".to_string()]),
1432 filters: None,
1433 series_row_selector: None,
1434 };
1435
1436 let err = resolve_projection(&config, None).unwrap_err();
1437 let msg = err.to_string();
1438 assert!(msg.contains("projection"));
1439 assert!(msg.contains("projection_names"));
1440 }
1441
1442 #[test]
1443 fn test_sqlparser_parse_expr_string() {
1444 let dialect = GenericDialect {};
1445 let mut parser = Parser::new(&dialect)
1446 .try_with_sql("host = 'web-1' AND cpu > 80")
1447 .unwrap();
1448
1449 let expr = parser.parse_expr().unwrap();
1450
1451 match expr {
1452 Expr::BinaryOp { op, .. } => assert_eq!(op, BinaryOperator::And),
1453 other => panic!("expected BinaryOp, got: {other:?}"),
1454 }
1455 }
1456
1457 #[test]
1458 fn test_resolve_filters_uint32_type_conversion() {
1459 use api::v1::SemanticType;
1460
1461 let mut builder = RegionMetadataBuilder::new(RegionId::new(1, 0));
1462 builder
1463 .push_column_metadata(ColumnMetadata {
1464 column_schema: ColumnSchema::new(
1465 "table_id",
1466 ConcreteDataType::uint32_datatype(),
1467 false,
1468 ),
1469 semantic_type: SemanticType::Tag,
1470 column_id: 1,
1471 })
1472 .push_column_metadata(ColumnMetadata {
1473 column_schema: ColumnSchema::new(
1474 "ts",
1475 ConcreteDataType::timestamp_millisecond_datatype(),
1476 false,
1477 ),
1478 semantic_type: SemanticType::Timestamp,
1479 column_id: 2,
1480 })
1481 .primary_key(vec![1]);
1482 let metadata = builder.build().unwrap();
1483
1484 let config = ScanConfig {
1485 name: None,
1486 projection: None,
1487 projection_names: None,
1488 filters: Some(vec!["table_id = 1117".to_string()]),
1489 series_row_selector: None,
1490 };
1491
1492 let exprs = resolve_filters(&config, &metadata).unwrap();
1493 assert_eq!(exprs.len(), 1);
1494 let expr_str = format!("{}", exprs[0]);
1496 assert!(
1497 expr_str.contains("UInt32(1117)"),
1498 "Expected UInt32(1117) in expression, got: {expr_str}"
1499 );
1500 }
1501
1502 #[test]
1503 fn test_parse_scan_config_filters() {
1504 let json = r#"{"filters":["host = 'web-1'","cpu > 80"]}"#;
1505 let config: ScanConfig = serde_json::from_str(json).unwrap();
1506
1507 assert_eq!(
1508 config.filters,
1509 Some(vec!["host = 'web-1'".to_string(), "cpu > 80".to_string()])
1510 );
1511 }
1512
1513 #[test]
1514 fn test_parse_and_validate_scan_config_suite() {
1515 let json = r#"[
1516 {"name":" cold ","filters":["host = 'web-1'"]},
1517 {"projection_names":["host","cpu"]}
1518 ]"#;
1519 let configs: Vec<ScanConfig> = serde_json::from_str(json).unwrap();
1520 let configs = validate_scan_config_suite(configs, 1).unwrap();
1521
1522 assert_eq!(2, configs.len());
1523 assert_eq!(Some("cold"), configs[0].name.as_deref());
1524 assert_eq!(Some("query-002"), configs[1].name.as_deref());
1525 assert_eq!(
1526 Some(&vec!["host = 'web-1'".to_string()]),
1527 configs[0].filters.as_ref()
1528 );
1529 }
1530
1531 #[test]
1532 fn test_validate_scan_config_suite_rejects_invalid_input() {
1533 let empty = validate_scan_config_suite(Vec::new(), 1)
1534 .unwrap_err()
1535 .to_string();
1536 assert!(empty.contains("at least one"));
1537
1538 let iterations = validate_scan_config_suite(vec![ScanConfig::default()], 2)
1539 .unwrap_err()
1540 .to_string();
1541 assert!(iterations.contains("--iterations"));
1542
1543 let configs: Vec<ScanConfig> =
1544 serde_json::from_str(r#"[{"name":"query"},{"name":" query "}]"#).unwrap();
1545 let duplicate = validate_scan_config_suite(configs, 1)
1546 .unwrap_err()
1547 .to_string();
1548 assert!(duplicate.contains("duplicate"));
1549
1550 let configs: Vec<ScanConfig> = serde_json::from_str(r#"[{"name":" "}]"#).unwrap();
1551 let blank = validate_scan_config_suite(configs, 1)
1552 .unwrap_err()
1553 .to_string();
1554 assert!(blank.contains("empty name"));
1555 }
1556
1557 #[tokio::test]
1558 async fn test_write_result_file_overwrites_complete_result() {
1559 let dir = tempfile::tempdir().unwrap();
1560 let path = dir.path().join("scanbench-result.json");
1561 tokio::fs::write(&path, b"old result").await.unwrap();
1562
1563 let result = ScanbenchResult {
1564 format_version: 1,
1565 started_at_unix_ms: 42,
1566 benchmark: BenchmarkMetadata {
1567 scanner: "seq".to_string(),
1568 region_id: "1024:0".to_string(),
1569 region_id_u64: RegionId::new(1024, 0).as_u64(),
1570 table_dir: "greptime/public/1024".to_string(),
1571 path_type: "bare".to_string(),
1572 parallelism: 8,
1573 enable_wal: false,
1574 config_mode: "suite".to_string(),
1575 run_count: 1,
1576 },
1577 runs: vec![ScanRunResult {
1578 iteration: 1,
1579 query_index: 1,
1580 name: "cpu-host-1".to_string(),
1581 config: NormalizedScanConfig {
1582 name: "cpu-host-1".to_string(),
1583 projection: Some(vec![1, 2]),
1584 filters: vec!["hostname = 'host_1'".to_string()],
1585 series_row_selector: None,
1586 },
1587 rows: 100,
1588 batches: 2,
1589 setup_elapsed_ns: 10,
1590 scan_elapsed_ns: 80,
1591 elapsed_ns: 90,
1592 array_mem_size_bytes: 1024,
1593 estimated_size_bytes: 512,
1594 partitions: vec![PartitionResult {
1595 partition: 0,
1596 rows: 100,
1597 batches: 2,
1598 array_mem_size_bytes: 1024,
1599 estimated_size_bytes: 512,
1600 first_batch_elapsed_ns: Some(20),
1601 elapsed_ns: 80,
1602 }],
1603 scanner_explain: "SeqScan: region=1024(0)".to_string(),
1604 }],
1605 summary: BenchmarkResultSummary {
1606 runs: 1,
1607 total_rows: 100,
1608 total_elapsed_ns: 90,
1609 mean_rows: 100,
1610 mean_elapsed_ns: 90,
1611 queries: vec![],
1612 },
1613 };
1614
1615 write_result_file(&path, &result).await.unwrap();
1616 let content = tokio::fs::read(&path).await.unwrap();
1617 let actual: serde_json::Value = serde_json::from_slice(&content).unwrap();
1618
1619 assert_eq!(1, actual["format_version"]);
1620 assert_eq!("suite", actual["benchmark"]["config_mode"]);
1621 assert_eq!(2, actual["runs"][0]["config"]["projection"][1]);
1622 assert_eq!(100, actual["runs"][0]["partitions"][0]["rows"]);
1623 assert_eq!(
1624 "SeqScan: region=1024(0)",
1625 actual["runs"][0]["scanner_explain"]
1626 );
1627 assert_eq!(90, actual["summary"]["mean_elapsed_ns"]);
1628 }
1629}