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)
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                ..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            // Start profiling after the first iteration (warmup) if pprof_after_warmup is set
1213            #[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        // Stop profiling and generate flamegraph if enabled
1240        #[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        // Summary
1270        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        // The expression should contain a UInt32 literal after type conversion.
1495        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}