Skip to main content

cmd/datanode/
scanbench.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use 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
68/// Displays a scanner using DataFusion's verbose explain format.
69struct 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/// Scan benchmark command - benchmarks scanning a region directly from storage.
88#[derive(Debug, Parser)]
89pub struct ScanbenchCommand {
90    /// Path to config TOML file (same format as standalone/datanode config)
91    #[clap(long, value_name = "FILE")]
92    config: PathBuf,
93
94    /// Region ID: either numeric u64 (e.g. "4398046511104") or "table_id:region_num" (e.g. "1024:0")
95    #[clap(long)]
96    region_id: String,
97
98    /// Table directory relative to data home (e.g. "data/greptime/public/1024/")
99    #[clap(long)]
100    table_dir: String,
101
102    /// Scanner type: seq, unordered, series
103    #[clap(long, default_value = "seq")]
104    scanner: String,
105
106    /// Path to scan request JSON config file (optional)
107    #[clap(long, value_name = "FILE", conflicts_with = "scan_configs")]
108    scan_config: Option<PathBuf>,
109
110    /// Path to a JSON array of scan request configs, each executed once
111    #[clap(long, value_name = "FILE", conflicts_with = "scan_config")]
112    scan_configs: Option<PathBuf>,
113
114    /// Number of partitions for parallel scan (simulates parallelism)
115    #[clap(long, default_value = "1")]
116    parallelism: usize,
117
118    /// Number of iterations for benchmarking
119    #[clap(long, default_value = "1")]
120    iterations: usize,
121
122    /// Path type for the region: bare, data, metadata
123    #[clap(long, default_value = "bare")]
124    path_type: String,
125
126    /// Verbose output
127    #[clap(short, long, default_value_t = false)]
128    verbose: bool,
129
130    /// Output pprof flamegraph
131    #[clap(long, value_name = "FILE")]
132    pprof_file: Option<PathBuf>,
133
134    /// Output structured benchmark and analyze results as JSON
135    #[clap(long, value_name = "FILE")]
136    result_file: Option<PathBuf>,
137
138    /// Enable WAL replay when opening the region.
139    #[clap(long, default_value_t = false)]
140    enable_wal: bool,
141
142    /// Start pprof after the first iteration (use first iteration as warmup).
143    #[clap(long, default_value_t = false)]
144    pprof_after_warmup: bool,
145}
146
147/// JSON config for scan request parameters.
148#[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
488/// Rewrites literal values in comparison expressions to match the column's arrow type.
489struct 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/// The log store scanbench opens for a WAL config.
700#[derive(Debug, PartialEq)]
701enum ScanbenchLogStore<'a> {
702    RaftEngine(&'a RaftEngineConfig),
703    Kafka(&'a DatanodeKafkaConfig),
704    Noop,
705}
706
707/// Selects the log store for the WAL config. WAL providers that scanbench cannot
708/// build are rejected instead of falling back to the Noop log store.
709fn 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        // Parse config and build object store
791        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        // Create mock dependencies
801        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        // Create MitoEngine with appropriate log store
806        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, None)
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        // Open region
859        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        // Build scan request
884        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        // Start profiling if pprof_file is specified (unless pprof_after_warmup is set)
919        #[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            // Get scanner
974            let mut scanner = engine
975                .handle_query(region_id, request)
976                .await
977                .map_err(BoxedError::new)
978                .context(error::BuildCliSnafu)?;
979
980            // Get partition ranges and apply parallelism
981            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                // Flatten all ranges
995                let all_ranges: Vec<_> = original_partitions.into_iter().flatten().collect();
996
997                // Distribute ranges across partitions
998                let mut partitions =
999                    ParallelizeScan::assign_partition_range(all_ranges, self.parallelism);
1000
1001                // Sort ranges within each partition by start time ascending
1002                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            // Scan all partitions
1017            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            // Start profiling after the first iteration (warmup) if pprof_after_warmup is set
1212            #[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        // Stop profiling and generate flamegraph if enabled
1239        #[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        // Summary
1269        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        // The expression should contain a UInt32 literal after type conversion.
1494        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}