Files
greptimedb/src/cmd/src/datanode/scanbench.rs
T
Yingwen 0cc83c4570 feat(cmd): add parquet development tools (#8939)
* feat(cmd): add parquet metadata development tool

Signed-off-by: evenyag <realevenyag@gmail.com>

* feat(cmd): add parquet rewrite development tool

Signed-off-by: evenyag <realevenyag@gmail.com>

* feat(cmd): add SST replacement development tool

Signed-off-by: evenyag <realevenyag@gmail.com>

* feat(cmd): support local parquet files in parquetbench

Signed-off-by: evenyag <realevenyag@gmail.com>

* fix(cmd): clean up parquet development tools

Signed-off-by: evenyag <realevenyag@gmail.com>

* refactor(cmd): share datanode tool utilities

Signed-off-by: evenyag <realevenyag@gmail.com>

* fix(cmd): satisfy parquet tool lints

Signed-off-by: evenyag <realevenyag@gmail.com>

* docs: document parquet development tools

Signed-off-by: evenyag <realevenyag@gmail.com>

* refactor(cmd): rename SST replacement module

Signed-off-by: evenyag <realevenyag@gmail.com>

* fix(cmd): validate parquet rewrite options

Signed-off-by: evenyag <realevenyag@gmail.com>

---------

Signed-off-by: evenyag <realevenyag@gmail.com>
2026-08-24 13:34:30 +00:00

906 lines
31 KiB
Rust
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
// Copyright 2023 Greptime Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Instant;
use clap::Parser;
use colored::Colorize;
use common_base::Plugins;
use common_error::ext::{BoxedError, PlainError};
use common_error::status_code::StatusCode;
use common_meta::cache::{new_schema_cache, new_table_schema_cache};
use common_meta::key::SchemaMetadataManager;
use common_meta::kv_backend::memory::MemoryKvBackend;
use common_wal::config::DatanodeWalConfig;
use datafusion::execution::SessionStateBuilder;
use datafusion::logical_expr::{BinaryExpr, Expr as DfExpr, ExprSchemable, Operator};
use datafusion_common::tree_node::{Transformed, TreeNodeRewriter};
use datafusion_common::{DFSchemaRef, ScalarValue, ToDFSchema};
use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet;
use datatypes::arrow::compute;
use futures::StreamExt;
use futures::stream::FuturesUnordered;
use log_store::kafka::log_store::KafkaLogStore;
use log_store::noop::log_store::NoopLogStore;
use log_store::raft_engine::log_store::RaftEngineLogStore;
use mito2::config::MitoConfig;
use mito2::engine::MitoEngine;
use mito2::sst::file_ref::FileReferenceManager;
use moka::future::CacheBuilder;
use object_store::manager::ObjectStoreManager;
use object_store::util::normalize_dir;
use query::optimizer::parallelize_scan::ParallelizeScan;
use serde::Deserialize;
use snafu::{OptionExt, ResultExt};
use sqlparser::ast::ExprWithAlias as SqlExprWithAlias;
use sqlparser::dialect::GenericDialect;
use sqlparser::parser::Parser as SqlParser;
use store_api::metadata::RegionMetadata;
use store_api::path_utils::WAL_DIR;
use store_api::region_engine::{PrepareRequest, QueryScanContext, RegionEngine};
use store_api::region_request::{RegionOpenRequest, RegionRequest};
use store_api::storage::{RegionId, ScanRequest, TimeSeriesDistribution, TimeSeriesRowSelector};
use tokio::fs;
use crate::datanode::tool_util::{
build_object_store, format_bytes, parse_config, parse_path_type, parse_region_id,
};
use crate::error;
/// Scan benchmark command - benchmarks scanning a region directly from storage.
#[derive(Debug, Parser)]
pub struct ScanbenchCommand {
/// Path to config TOML file (same format as standalone/datanode config)
#[clap(long, value_name = "FILE")]
config: PathBuf,
/// Region ID: either numeric u64 (e.g. "4398046511104") or "table_id:region_num" (e.g. "1024:0")
#[clap(long)]
region_id: String,
/// Table directory relative to data home (e.g. "data/greptime/public/1024/")
#[clap(long)]
table_dir: String,
/// Scanner type: seq, unordered, series
#[clap(long, default_value = "seq")]
scanner: String,
/// Path to scan request JSON config file (optional)
#[clap(long, value_name = "FILE")]
scan_config: Option<PathBuf>,
/// Number of partitions for parallel scan (simulates parallelism)
#[clap(long, default_value = "1")]
parallelism: usize,
/// Number of iterations for benchmarking
#[clap(long, default_value = "1")]
iterations: usize,
/// Path type for the region: bare, data, metadata
#[clap(long, default_value = "bare")]
path_type: String,
/// Verbose output
#[clap(short, long, default_value_t = false)]
verbose: bool,
/// Output pprof flamegraph
#[clap(long, value_name = "FILE")]
pprof_file: Option<PathBuf>,
/// Enable WAL replay when opening the region.
#[clap(long, default_value_t = false)]
enable_wal: bool,
/// Start pprof after the first iteration (use first iteration as warmup).
#[clap(long, default_value_t = false)]
pprof_after_warmup: bool,
}
/// JSON config for scan request parameters.
#[derive(Debug, Deserialize, Default)]
struct ScanConfig {
projection: Option<Vec<usize>>,
projection_names: Option<Vec<String>>,
filters: Option<Vec<String>>,
series_row_selector: Option<String>,
}
fn resolve_projection(
scan_config: &ScanConfig,
metadata: Option<&RegionMetadata>,
) -> error::Result<Option<Vec<usize>>> {
if scan_config.projection.is_some() && scan_config.projection_names.is_some() {
return Err(error::IllegalConfigSnafu {
msg: "scan config cannot contain both 'projection' and 'projection_names'".to_string(),
}
.build());
}
if let Some(projection) = &scan_config.projection {
return Ok(Some(projection.clone()));
}
if let Some(projection_names) = &scan_config.projection_names {
let metadata = metadata.context(error::IllegalConfigSnafu {
msg: "Missing region metadata while resolving 'projection_names'".to_string(),
})?;
let available_columns = metadata
.column_metadatas
.iter()
.map(|column| column.column_schema.name.as_str())
.collect::<Vec<_>>()
.join(", ");
let projection = projection_names
.iter()
.map(|name| {
metadata
.column_index_by_name(name)
.with_context(|| error::IllegalConfigSnafu {
msg: format!(
"Unknown column '{}' in projection_names, available columns: [{}]",
name, available_columns
),
})
})
.collect::<error::Result<Vec<_>>>()?;
return Ok(Some(projection));
}
Ok(None)
}
/// Rewrites literal values in comparison expressions to match the column's arrow type.
struct LiteralTypeCaster {
schema: DFSchemaRef,
}
impl TreeNodeRewriter for LiteralTypeCaster {
type Node = DfExpr;
fn f_up(&mut self, expr: DfExpr) -> datafusion_common::Result<Transformed<DfExpr>> {
let DfExpr::BinaryExpr(BinaryExpr { left, op, right }) = &expr else {
return Ok(Transformed::no(expr));
};
if !matches!(
op,
Operator::Eq
| Operator::NotEq
| Operator::Lt
| Operator::LtEq
| Operator::Gt
| Operator::GtEq
) {
return Ok(Transformed::no(expr));
}
let (col_expr, lit_expr, col_left) = match (left.as_ref(), right.as_ref()) {
(col @ DfExpr::Column(_), lit @ DfExpr::Literal(_, _)) => (col, lit, true),
(lit @ DfExpr::Literal(_, _), col @ DfExpr::Column(_)) => (col, lit, false),
_ => return Ok(Transformed::no(expr)),
};
let col_type = col_expr.get_type(self.schema.as_ref())?;
let DfExpr::Literal(scalar, _) = lit_expr else {
unreachable!()
};
if scalar.data_type() == col_type {
return Ok(Transformed::no(expr));
}
let lit_array = scalar.to_array()?;
let casted = compute::cast(lit_array.as_ref(), &col_type).map_err(|e| {
datafusion_common::DataFusionError::Internal(format!(
"Failed to cast literal {:?} to {:?}: {}",
scalar, col_type, e
))
})?;
let casted_scalar = ScalarValue::try_from_array(&casted, 0)?;
let new_lit = DfExpr::Literal(casted_scalar, None);
let (new_left, new_right) = if col_left {
(left.clone(), Box::new(new_lit))
} else {
(Box::new(new_lit), right.clone())
};
Ok(Transformed::yes(DfExpr::BinaryExpr(BinaryExpr {
left: new_left,
op: *op,
right: new_right,
})))
}
}
fn convert_literal_types(
exprs: Vec<DfExpr>,
schema: &DFSchemaRef,
) -> datafusion_common::Result<Vec<DfExpr>> {
use datafusion_common::tree_node::TreeNode;
let mut caster = LiteralTypeCaster {
schema: schema.clone(),
};
exprs
.into_iter()
.map(|e| e.rewrite(&mut caster).map(|x| x.data))
.collect()
}
fn resolve_filters(
scan_config: &ScanConfig,
metadata: &RegionMetadata,
) -> error::Result<Vec<DfExpr>> {
let Some(filters) = &scan_config.filters else {
return Ok(Vec::new());
};
let df_schema = metadata
.schema
.arrow_schema()
.clone()
.to_dfschema()
.map_err(|e| {
error::IllegalConfigSnafu {
msg: format!("Failed to convert region schema to DataFusion schema: {e}"),
}
.build()
})?;
let state = SessionStateBuilder::new()
.with_config(Default::default())
.with_runtime_env(Default::default())
.with_default_features()
.build();
let exprs: Vec<DfExpr> = filters
.iter()
.enumerate()
.map(|(idx, filter)| {
let mut parser = SqlParser::new(&GenericDialect {})
.try_with_sql(filter)
.map_err(|e| {
error::IllegalConfigSnafu {
msg: format!("Invalid filter at index {idx} ('{filter}'): {e}"),
}
.build()
})?;
let sql_expr = parser.parse_expr().map_err(|e| {
error::IllegalConfigSnafu {
msg: format!("Invalid filter at index {idx} ('{filter}'): {e}"),
}
.build()
})?;
state
.create_logical_expr_from_sql_expr(
SqlExprWithAlias {
expr: sql_expr,
alias: None,
},
&df_schema,
)
.map_err(|e| {
error::IllegalConfigSnafu {
msg: format!(
"Failed to convert filter at index {idx} ('{filter}') to logical expr: {e}"
),
}
.build()
})
})
.collect::<error::Result<Vec<_>>>()?;
let df_schema_ref = Arc::new(df_schema);
convert_literal_types(exprs, &df_schema_ref).map_err(|e| {
error::IllegalConfigSnafu {
msg: format!("Failed to convert filter expression types: {e}"),
}
.build()
})
}
fn noop_partition_expr_fetcher() -> mito2::region::opener::PartitionExprFetcherRef {
struct NoopPartitionExprFetcher;
#[async_trait::async_trait]
impl mito2::region::opener::PartitionExprFetcher for NoopPartitionExprFetcher {
async fn fetch_expr(&self, _region_id: RegionId) -> Option<String> {
None
}
}
Arc::new(NoopPartitionExprFetcher)
}
struct EngineComponents {
data_home: String,
mito_config: MitoConfig,
object_store_manager: Arc<ObjectStoreManager>,
schema_metadata_manager: Arc<SchemaMetadataManager>,
file_ref_manager: Arc<FileReferenceManager>,
partition_expr_fetcher: mito2::region::opener::PartitionExprFetcherRef,
}
impl EngineComponents {
async fn build<S: store_api::logstore::LogStore>(
self,
log_store: Arc<S>,
) -> error::Result<MitoEngine> {
MitoEngine::new(
&self.data_home,
self.mito_config,
log_store,
self.object_store_manager,
self.schema_metadata_manager,
self.file_ref_manager,
self.partition_expr_fetcher,
Plugins::default(),
)
.await
.map_err(BoxedError::new)
.context(error::BuildCliSnafu)
}
}
fn mock_schema_metadata_manager() -> Arc<SchemaMetadataManager> {
let kv_backend = Arc::new(MemoryKvBackend::new());
let table_schema_cache = Arc::new(new_table_schema_cache(
"table_schema_name_cache".to_string(),
CacheBuilder::default().build(),
kv_backend.clone(),
));
let schema_cache = Arc::new(new_schema_cache(
"schema_cache".to_string(),
CacheBuilder::default().build(),
kv_backend.clone(),
));
Arc::new(SchemaMetadataManager::new(table_schema_cache, schema_cache))
}
impl ScanbenchCommand {
pub async fn run(&self) -> error::Result<()> {
if self.verbose {
common_telemetry::init_default_ut_logging();
}
println!("{}", "Starting scanbench...".cyan().bold());
let region_id = parse_region_id(&self.region_id)?;
let path_type = parse_path_type(&self.path_type)?;
println!(
"{} Region ID: {} (u64: {})",
"✓".green(),
self.region_id,
region_id.as_u64()
);
// Parse config and build object store
let (store_cfg, mito_config, wal_config) = parse_config(&self.config)?;
println!("{} Config parsed", "✓".green());
let object_store = build_object_store(&store_cfg).await?;
println!("{} Object store initialized", "✓".green());
let object_store_manager =
Arc::new(ObjectStoreManager::new("default", object_store.clone()));
// Create mock dependencies
let schema_metadata_manager = mock_schema_metadata_manager();
let file_ref_manager = Arc::new(FileReferenceManager::new(None));
let partition_expr_fetcher = noop_partition_expr_fetcher();
// Create MitoEngine with appropriate log store
let components = EngineComponents {
data_home: store_cfg.data_home.clone(),
mito_config,
object_store_manager,
schema_metadata_manager,
file_ref_manager,
partition_expr_fetcher,
};
let engine = match &wal_config {
DatanodeWalConfig::RaftEngine(raft_engine_config) if self.enable_wal => {
let data_home = normalize_dir(&store_cfg.data_home);
let wal_dir = match &raft_engine_config.dir {
Some(dir) => dir.clone(),
None => format!("{}{WAL_DIR}", data_home),
};
fs::create_dir_all(&wal_dir).await.map_err(|e| {
error::IllegalConfigSnafu {
msg: format!("failed to create WAL directory {}: {e}", wal_dir),
}
.build()
})?;
let log_store = Arc::new(
RaftEngineLogStore::try_new(wal_dir, raft_engine_config)
.await
.map_err(BoxedError::new)
.context(error::BuildCliSnafu)?,
);
println!("{} Using RaftEngine WAL", "✓".green());
components.build(log_store).await?
}
DatanodeWalConfig::Kafka(kafka_config) if self.enable_wal => {
let log_store = Arc::new(
KafkaLogStore::try_new(kafka_config, None)
.await
.map_err(BoxedError::new)
.context(error::BuildCliSnafu)?,
);
println!("{} Using Kafka WAL", "✓".green());
components.build(log_store).await?
}
_ => {
let log_store = Arc::new(NoopLogStore);
println!(
"{} Using NoopLogStore (enable_wal={})",
"✓".green(),
self.enable_wal
);
components.build(log_store).await?
}
};
// Open region
let open_request = RegionOpenRequest {
engine: "mito".to_string(),
table_dir: self.table_dir.clone(),
path_type,
options: HashMap::default(),
skip_wal_replay: !self.enable_wal,
checkpoint: None,
requirements: Default::default(),
};
engine
.handle_request(region_id, RegionRequest::Open(open_request))
.await
.map_err(BoxedError::new)
.context(error::BuildCliSnafu)?;
println!("{} Region opened", "✓".green());
// Load scan config
let scan_config = if let Some(path) = &self.scan_config {
let content = tokio::fs::read_to_string(path)
.await
.context(error::FileIoSnafu)?;
serde_json::from_str::<ScanConfig>(&content).context(error::SerdeJsonSnafu)?
} else {
ScanConfig::default()
};
let metadata = engine
.get_metadata(region_id)
.await
.map_err(BoxedError::new)
.context(error::BuildCliSnafu)?;
let projection = resolve_projection(&scan_config, Some(&metadata))?;
let filters = resolve_filters(&scan_config, &metadata)?;
// Build scan request
let distribution = match self.scanner.as_str() {
"seq" => None,
"unordered" => Some(TimeSeriesDistribution::TimeWindowed),
"series" => Some(TimeSeriesDistribution::PerSeries),
other => {
return Err(error::IllegalConfigSnafu {
msg: format!(
"Unknown scanner type '{}', expected: seq, unordered, series",
other
),
}
.build());
}
};
let series_row_selector = match scan_config.series_row_selector.as_deref() {
Some("last_row") => Some(TimeSeriesRowSelector::LastRow),
Some(other) => {
return Err(error::IllegalConfigSnafu {
msg: format!("Unknown series_row_selector '{}'", other),
}
.build());
}
None => None,
};
println!(
"{} Scanner: {}, Parallelism: {}, Iterations: {}",
"".blue(),
self.scanner,
self.parallelism,
self.iterations,
);
// Start profiling if pprof_file is specified (unless pprof_after_warmup is set)
#[cfg(unix)]
let mut profiler_guard = if self.pprof_file.is_some() && !self.pprof_after_warmup {
println!("{} Starting profiling...", "⚡".yellow());
Some(
pprof::ProfilerGuardBuilder::default()
.frequency(99)
.blocklist(&["libc", "libgcc", "pthread", "vdso"])
.build()
.map_err(|e| {
BoxedError::new(PlainError::new(
format!("Failed to start profiler: {e}"),
StatusCode::Unexpected,
))
})
.context(error::BuildCliSnafu)?,
)
} else {
None
};
#[cfg(not(unix))]
if self.pprof_file.is_some() {
eprintln!(
"{}: Profiling is not supported on this platform",
"Warning".yellow()
);
}
let mut total_rows_all = 0u64;
let mut total_elapsed_all = std::time::Duration::ZERO;
for iteration in 0..self.iterations {
let request = ScanRequest {
projection: projection.clone(),
filters: filters.clone(),
series_row_selector,
distribution,
..Default::default()
};
let start = Instant::now();
// Get scanner
let mut scanner = engine
.handle_query(region_id, request)
.await
.map_err(BoxedError::new)
.context(error::BuildCliSnafu)?;
// Get partition ranges and apply parallelism
let original_partitions = scanner.properties().partitions.clone();
let total_ranges: usize = original_partitions.iter().map(|p| p.len()).sum();
if self.verbose {
println!(
" {} Original partitions: {}, total ranges: {}",
"".blue(),
original_partitions.len(),
total_ranges
);
}
if self.parallelism > 1 {
// Flatten all ranges
let all_ranges: Vec<_> = original_partitions.into_iter().flatten().collect();
// Distribute ranges across partitions
let mut partitions =
ParallelizeScan::assign_partition_range(all_ranges, self.parallelism);
// Sort ranges within each partition by start time ascending
for partition in &mut partitions {
partition.sort_by_key(|a| a.start);
}
scanner
.prepare(
PrepareRequest::default()
.with_ranges(partitions)
.with_target_partitions(self.parallelism),
)
.map_err(BoxedError::new)
.context(error::BuildCliSnafu)?;
}
// Scan all partitions
let num_partitions = scanner.properties().partitions.len();
let ctx = QueryScanContext {
explain_verbose: self.verbose,
};
let metrics_set = ExecutionPlanMetricsSet::new();
let mut scan_futures = FuturesUnordered::new();
for partition_idx in 0..num_partitions {
let mut stream = scanner
.scan_partition(&ctx, &metrics_set, partition_idx)
.map_err(BoxedError::new)
.context(error::BuildCliSnafu)?;
scan_futures.push(tokio::spawn(async move {
let mut rows = 0u64;
let mut array_mem_size = 0u64;
let mut estimated_size = 0u64;
while let Some(batch_result) = stream.next().await {
match batch_result {
Ok(batch) => {
rows += batch.num_rows() as u64;
let df_batch = batch.df_record_batch();
array_mem_size += df_batch.get_array_memory_size() as u64;
estimated_size +=
mito2::memtable::record_batch_estimated_size(df_batch) as u64;
}
Err(e) => {
return Err(BoxedError::new(e));
}
}
}
Ok::<(u64, u64, u64), BoxedError>((rows, array_mem_size, estimated_size))
}));
}
let mut total_rows = 0u64;
let mut total_array_mem_size = 0u64;
let mut total_estimated_size = 0u64;
while let Some(task) = scan_futures.next().await {
let result = task
.map_err(|e| {
BoxedError::new(PlainError::new(
format!("scan task failed: {e}"),
StatusCode::Unexpected,
))
})
.context(error::BuildCliSnafu)?;
let (rows, array_mem_size, estimated_size) =
result.context(error::BuildCliSnafu)?;
total_rows += rows;
total_array_mem_size += array_mem_size;
total_estimated_size += estimated_size;
}
let elapsed = start.elapsed();
total_rows_all += total_rows;
total_elapsed_all += elapsed;
println!(
" [iter {}] {} rows in {:?} ({} partitions), array_mem_size: {}, estimated_size: {}",
iteration + 1,
total_rows.to_string().cyan(),
elapsed,
num_partitions,
format_bytes(total_array_mem_size),
format_bytes(total_estimated_size),
);
// Start profiling after the first iteration (warmup) if pprof_after_warmup is set
#[cfg(unix)]
if iteration == 0
&& self.pprof_after_warmup
&& self.pprof_file.is_some()
&& profiler_guard.is_none()
{
println!(
"{} Starting profiling after warmup iteration...",
"⚡".yellow()
);
profiler_guard = Some(
pprof::ProfilerGuardBuilder::default()
.frequency(99)
.blocklist(&["libc", "libgcc", "pthread", "vdso"])
.build()
.map_err(|e| {
BoxedError::new(PlainError::new(
format!("Failed to start profiler: {e}"),
StatusCode::Unexpected,
))
})
.context(error::BuildCliSnafu)?,
);
}
}
// Stop profiling and generate flamegraph if enabled
#[cfg(unix)]
if let (Some(guard), Some(pprof_file)) = (profiler_guard, &self.pprof_file) {
println!("{} Generating flamegraph...", "🔥".yellow());
match guard.report().build() {
Ok(report) => {
let mut flamegraph_data = Vec::new();
if let Err(e) = report.flamegraph(&mut flamegraph_data) {
println!("{}: Failed to generate flamegraph: {}", "Error".red(), e);
} else if let Err(e) = std::fs::write(pprof_file, flamegraph_data) {
println!(
"{}: Failed to write flamegraph to {}: {}",
"Error".red(),
pprof_file.display(),
e
);
} else {
println!(
"{} Flamegraph saved to {}",
"✓".green(),
pprof_file.display().to_string().cyan()
);
}
}
Err(e) => {
println!("{}: Failed to generate pprof report: {}", "Error".red(), e);
}
}
}
// Summary
if self.iterations > 1 {
let avg_elapsed = total_elapsed_all / self.iterations as u32;
let avg_rows = total_rows_all / self.iterations as u64;
println!(
"\n{} Average: {} rows in {:?} over {} iterations",
"Summary".green().bold(),
avg_rows.to_string().cyan(),
avg_elapsed,
self.iterations,
);
}
println!("\n{}", "Benchmark completed!".green().bold());
Ok(())
}
}
#[cfg(test)]
mod tests {
use datatypes::prelude::ConcreteDataType;
use datatypes::schema::ColumnSchema;
use sqlparser::ast::{BinaryOperator, Expr};
use sqlparser::dialect::GenericDialect;
use sqlparser::parser::Parser;
use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder};
use store_api::storage::RegionId;
use super::{ScanConfig, resolve_filters, resolve_projection};
use crate::error;
#[test]
fn test_parse_scan_config_projection_names() {
let json = r#"{"projection_names":["host","ts"]}"#;
let config: ScanConfig = serde_json::from_str(json).unwrap();
assert_eq!(
config.projection_names,
Some(vec!["host".to_string(), "ts".to_string()])
);
assert_eq!(config.projection, None);
}
#[test]
fn test_resolve_projection_by_indexes() -> error::Result<()> {
let config = ScanConfig {
projection: Some(vec![0, 2]),
projection_names: None,
filters: None,
series_row_selector: None,
};
let projection = resolve_projection(&config, None)?;
assert_eq!(projection, Some(vec![0, 2]));
Ok(())
}
#[test]
fn test_resolve_projection_by_names_without_metadata() {
let config = ScanConfig {
projection: None,
projection_names: Some(vec!["cpu".to_string(), "host".to_string()]),
filters: None,
series_row_selector: None,
};
let err = resolve_projection(&config, None).unwrap_err();
assert!(
err.to_string()
.contains("Missing region metadata while resolving 'projection_names'")
);
}
#[test]
fn test_resolve_projection_conflict_fields() {
let config = ScanConfig {
projection: Some(vec![0]),
projection_names: Some(vec!["host".to_string()]),
filters: None,
series_row_selector: None,
};
let err = resolve_projection(&config, None).unwrap_err();
let msg = err.to_string();
assert!(msg.contains("projection"));
assert!(msg.contains("projection_names"));
}
#[test]
fn test_sqlparser_parse_expr_string() {
let dialect = GenericDialect {};
let mut parser = Parser::new(&dialect)
.try_with_sql("host = 'web-1' AND cpu > 80")
.unwrap();
let expr = parser.parse_expr().unwrap();
match expr {
Expr::BinaryOp { op, .. } => assert_eq!(op, BinaryOperator::And),
other => panic!("expected BinaryOp, got: {other:?}"),
}
}
#[test]
fn test_resolve_filters_uint32_type_conversion() {
use api::v1::SemanticType;
let mut builder = RegionMetadataBuilder::new(RegionId::new(1, 0));
builder
.push_column_metadata(ColumnMetadata {
column_schema: ColumnSchema::new(
"table_id",
ConcreteDataType::uint32_datatype(),
false,
),
semantic_type: SemanticType::Tag,
column_id: 1,
})
.push_column_metadata(ColumnMetadata {
column_schema: ColumnSchema::new(
"ts",
ConcreteDataType::timestamp_millisecond_datatype(),
false,
),
semantic_type: SemanticType::Timestamp,
column_id: 2,
})
.primary_key(vec![1]);
let metadata = builder.build().unwrap();
let config = ScanConfig {
projection: None,
projection_names: None,
filters: Some(vec!["table_id = 1117".to_string()]),
series_row_selector: None,
};
let exprs = resolve_filters(&config, &metadata).unwrap();
assert_eq!(exprs.len(), 1);
// The expression should contain a UInt32 literal after type conversion.
let expr_str = format!("{}", exprs[0]);
assert!(
expr_str.contains("UInt32(1117)"),
"Expected UInt32(1117) in expression, got: {expr_str}"
);
}
#[test]
fn test_parse_scan_config_filters() {
let json = r#"{"filters":["host = 'web-1'","cpu > 80"]}"#;
let config: ScanConfig = serde_json::from_str(json).unwrap();
assert_eq!(
config.filters,
Some(vec!["host = 'web-1'".to_string(), "cpu > 80".to_string()])
);
}
}