From 35593bdff73c368e58a19d9d2ed579af6c0ebd4e Mon Sep 17 00:00:00 2001 From: discord9 Date: Thu, 9 Jul 2026 19:42:18 +0800 Subject: [PATCH] test: normalize remote write perf fixtures Signed-off-by: discord9 --- .github/scripts/query-regression-run.py | 11 +- .github/workflows/query-regression.yml | 7 +- src/cmd/Cargo.toml | 9 +- src/cmd/src/bin/parquet_footer_inspector.rs | 281 -------- src/cmd/src/bin/prom_remote_write_fixture.rs | 426 ------------ src/cmd/src/bin/query_perf_fixture.rs | 652 +----------------- src/cmd/src/bin/query_perf_fixture/case.rs | 470 +++++++++++++ .../src/bin/query_perf_fixture/direct_sst.rs | 528 ++++++++++++++ .../bin/query_perf_fixture/inspect_footer.rs | 181 +++++ src/cmd/src/bin/query_perf_fixture/mod.rs | 133 ++++ .../query_perf_fixture/prom_remote_write.rs | 208 ++++++ src/cmd/src/bin/query_perf_fixture/util.rs | 23 + tests/perf/AGENTS.md | 8 +- tests/perf/README.md | 60 +- .../prom_remote_write_smoke/case.toml | 28 +- tests/perf/query_regression_runner.py | 325 ++++++--- 16 files changed, 1835 insertions(+), 1515 deletions(-) delete mode 100644 src/cmd/src/bin/parquet_footer_inspector.rs delete mode 100644 src/cmd/src/bin/prom_remote_write_fixture.rs create mode 100644 src/cmd/src/bin/query_perf_fixture/case.rs create mode 100644 src/cmd/src/bin/query_perf_fixture/direct_sst.rs create mode 100644 src/cmd/src/bin/query_perf_fixture/inspect_footer.rs create mode 100644 src/cmd/src/bin/query_perf_fixture/mod.rs create mode 100644 src/cmd/src/bin/query_perf_fixture/prom_remote_write.rs create mode 100644 src/cmd/src/bin/query_perf_fixture/util.rs diff --git a/.github/scripts/query-regression-run.py b/.github/scripts/query-regression-run.py index b89857be81f..07593d7b2f5 100644 --- a/.github/scripts/query-regression-run.py +++ b/.github/scripts/query-regression-run.py @@ -26,6 +26,7 @@ from pathlib import Path DEFAULT_CASES = [ "tests/perf/query_cases/smoke_direct_sst/case.toml", + "tests/perf/query_cases/prom_remote_write_smoke/case.toml", "tests/perf/query_cases/promql_pushdown_7913/case.toml", "tests/perf/query_cases/sql_topk_order_by/case.toml", "tests/perf/query_cases/sql_aggregate_order_by/case.toml", @@ -106,8 +107,6 @@ def run_case(args: argparse.Namespace, case_path: Path, work_dir: Path) -> int: str(candidate_bin), "--fixture-generator", str(fixture_generator), - "--remote-write-generator", - str(args.remote_write_generator), "--work-dir", str(work_dir), "--http-timeout", @@ -160,11 +159,6 @@ def main() -> int: parser.add_argument("--work-dir", default=Path("query-regression-work"), type=Path) parser.add_argument("--http-timeout", default=os.environ.get("HTTP_TIMEOUT", "300")) parser.add_argument("--allow-large-fixture", default=os.environ.get("ALLOW_LARGE_FIXTURE", "false")) - parser.add_argument( - "--remote-write-generator", - type=Path, - default=configured_path(os.environ.get("REMOTE_WRITE_GENERATOR")), - ) parser.add_argument( "--summary-script", type=Path, @@ -177,9 +171,6 @@ def main() -> int: parser.add_argument("--candidate-ref", default=os.environ.get("CANDIDATE_REF", "")) parser.add_argument("--github-output", default=os.environ.get("GITHUB_OUTPUT")) args = parser.parse_args() - if args.remote_write_generator is None: - args.remote_write_generator = args.candidate_src / "target" / profile_dir(args.cargo_profile) / "prom_remote_write_fixture" - try: cases = split_cases(args.cases or [os.environ.get("CASE_PATHS", "all")]) except ValueError as err: diff --git a/.github/workflows/query-regression.yml b/.github/workflows/query-regression.yml index a898be6ad29..b87d81fa42d 100644 --- a/.github/workflows/query-regression.yml +++ b/.github/workflows/query-regression.yml @@ -103,7 +103,7 @@ jobs: - name: Build base greptime working-directory: src run: | - cargo build --profile "${CARGO_PROFILE}" -p cmd --bin greptime + cargo build --profile "${CARGO_PROFILE}" -p cmd --bin greptime --features dev-tools target_dir="${CARGO_PROFILE}" if [[ "${CARGO_PROFILE}" == "dev" ]]; then target_dir="debug" @@ -122,7 +122,7 @@ jobs: - name: Build candidate greptime and fixture generators working-directory: src run: | - cargo build --profile "${CARGO_PROFILE}" -p cmd --bin greptime --bin query_perf_fixture --bin prom_remote_write_fixture + cargo build --profile "${CARGO_PROFILE}" -p cmd --bin greptime --bin query_perf_fixture --features dev-tools target_dir="${CARGO_PROFILE}" if [[ "${CARGO_PROFILE}" == "dev" ]]; then target_dir="debug" @@ -132,8 +132,6 @@ jobs: "${GITHUB_WORKSPACE}/query-regression-bins/candidate/greptime" cp "${CARGO_TARGET_DIR}/${target_dir}/query_perf_fixture" \ "${GITHUB_WORKSPACE}/query-regression-bins/candidate/query_perf_fixture" - cp "${CARGO_TARGET_DIR}/${target_dir}/prom_remote_write_fixture" \ - "${GITHUB_WORKSPACE}/query-regression-bins/candidate/prom_remote_write_fixture" - name: Run query regression id: run @@ -146,7 +144,6 @@ jobs: BASE_BIN: ${{ github.workspace }}/query-regression-bins/base/greptime CANDIDATE_BIN: ${{ github.workspace }}/query-regression-bins/candidate/greptime FIXTURE_GENERATOR: ${{ github.workspace }}/query-regression-bins/candidate/query_perf_fixture - REMOTE_WRITE_GENERATOR: ${{ github.workspace }}/query-regression-bins/candidate/prom_remote_write_fixture run: >- uv run --no-project python src/.github/scripts/query-regression-run.py --base-src src diff --git a/src/cmd/Cargo.toml b/src/cmd/Cargo.toml index 2e859585711..7b77d221765 100644 --- a/src/cmd/Cargo.toml +++ b/src/cmd/Cargo.toml @@ -12,14 +12,7 @@ path = "src/bin/greptime.rs" [[bin]] name = "query_perf_fixture" path = "src/bin/query_perf_fixture.rs" - -[[bin]] -name = "prom_remote_write_fixture" -path = "src/bin/prom_remote_write_fixture.rs" - -[[bin]] -name = "parquet_footer_inspector" -path = "src/bin/parquet_footer_inspector.rs" +required-features = ["dev-tools"] [features] default = [ diff --git a/src/cmd/src/bin/parquet_footer_inspector.rs b/src/cmd/src/bin/parquet_footer_inspector.rs deleted file mode 100644 index 4c8fed2fd4d..00000000000 --- a/src/cmd/src/bin/parquet_footer_inspector.rs +++ /dev/null @@ -1,281 +0,0 @@ -// 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. - -#![allow(clippy::print_stdout)] - -use std::collections::BTreeSet; -use std::fs::{self, File}; -use std::path::{Path, PathBuf}; - -use clap::Parser; -use parquet::file::reader::{FileReader, SerializedFileReader}; -use serde::Serialize; - -#[derive(Debug, Parser)] -#[command(about = "Inspect local Parquet/SST footer metadata as JSON")] -struct Args { - #[arg(long)] - root: PathBuf, - #[arg(long, default_value = "greptime_value")] - column: String, - #[arg(long)] - include_metadata_files: bool, -} - -#[derive(Debug, Serialize)] -struct Report { - root: String, - summary: Summary, - files: Vec, -} - -#[derive(Debug, Default, Serialize)] -struct Summary { - column: String, - file_count: usize, - files_with_column: usize, - total_file_size: u64, - total_rows: i64, - column_compressed_size: i64, - column_uncompressed_size: i64, - unique_encodings: Vec, -} - -#[derive(Debug, Serialize)] -struct FileReport { - path: String, - relative_path: String, - file_size: u64, - num_rows: i64, - num_row_groups: usize, - columns: Vec, -} - -#[derive(Debug, Serialize)] -struct ColumnChunkReport { - row_group_index: usize, - column_index: usize, - column_path: String, - encodings: Vec, - compression: String, - compressed_size: i64, - uncompressed_size: i64, - num_values: i64, -} - -fn main() -> Result<(), Box> { - let args = Args::parse(); - let files = collect_parquet_files(&args.root, args.include_metadata_files)?; - let reports = files - .iter() - .map(|path| inspect_file(&args.root, path, &args.column)) - .collect::, _>>()?; - let summary = summarize(&args.column, &reports); - println!( - "{}", - serde_json::to_string_pretty(&Report { - root: args.root.display().to_string(), - summary, - files: reports, - })? - ); - Ok(()) -} - -fn collect_parquet_files( - root: &Path, - include_metadata_files: bool, -) -> Result, Box> { - let mut files = Vec::new(); - collect_parquet_files_inner(root, root, include_metadata_files, &mut files)?; - files.sort(); - Ok(files) -} - -fn collect_parquet_files_inner( - root: &Path, - path: &Path, - include_metadata_files: bool, - files: &mut Vec, -) -> Result<(), Box> { - if path.is_dir() { - for entry in fs::read_dir(path)? { - let entry = entry?; - collect_parquet_files_inner(root, &entry.path(), include_metadata_files, files)?; - } - } else if is_data_parquet(root, path, include_metadata_files) { - files.push(path.to_path_buf()); - } - Ok(()) -} - -fn is_data_parquet(root: &Path, path: &Path, include_metadata_files: bool) -> bool { - let is_parquet = path.extension().is_some_and(|ext| ext == "parquet"); - let relative_path = path.strip_prefix(root).unwrap_or(path); - is_parquet && (include_metadata_files || !has_metadata_segment(relative_path)) -} - -fn has_metadata_segment(path: &Path) -> bool { - path.components() - .any(|component| component.as_os_str() == "metadata") -} - -fn inspect_file( - root: &Path, - path: &Path, - column: &str, -) -> Result> { - let file_size = fs::metadata(path)?.len(); - let reader = SerializedFileReader::new(File::open(path)?)?; - let metadata = reader.metadata().file_metadata(); - let row_groups = reader.metadata().row_groups(); - let mut columns = Vec::new(); - for (row_group_index, row_group) in row_groups.iter().enumerate() { - for (column_index, chunk) in row_group.columns().iter().enumerate() { - let column_path = chunk.column_path().string(); - if column_path == column { - columns.push(ColumnChunkReport { - row_group_index, - column_index, - column_path, - encodings: chunk - .encodings() - .map(|encoding| format!("{encoding:?}")) - .collect(), - compression: format!("{:?}", chunk.compression()), - compressed_size: chunk.compressed_size(), - uncompressed_size: chunk.uncompressed_size(), - num_values: chunk.num_values(), - }); - } - } - } - Ok(FileReport { - path: path.display().to_string(), - relative_path: path - .strip_prefix(root) - .unwrap_or(path) - .display() - .to_string(), - file_size, - num_rows: metadata.num_rows(), - num_row_groups: row_groups.len(), - columns, - }) -} - -fn summarize(column: &str, files: &[FileReport]) -> Summary { - let mut encodings = BTreeSet::new(); - let mut summary = Summary { - column: column.to_string(), - file_count: files.len(), - ..Default::default() - }; - for file in files { - summary.total_file_size += file.file_size; - summary.total_rows += file.num_rows; - if !file.columns.is_empty() { - summary.files_with_column += 1; - } - for column in &file.columns { - summary.column_compressed_size += column.compressed_size; - summary.column_uncompressed_size += column.uncompressed_size; - encodings.extend(column.encodings.iter().cloned()); - } - } - summary.unique_encodings = encodings.into_iter().collect(); - summary -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn filters_metadata_segment_by_default() { - let root = Path::new("/tmp/metadata/root"); - assert!(is_data_parquet( - root, - Path::new("/tmp/metadata/root/data/a.parquet"), - false - )); - assert!(!is_data_parquet( - root, - Path::new("/tmp/metadata/root/region/metadata/a.parquet"), - false - )); - assert!(is_data_parquet( - root, - Path::new("/tmp/metadata/root/region/metadata/a.parquet"), - true - )); - assert!(!is_data_parquet( - root, - Path::new("/tmp/metadata/root/region/data/a.sst"), - false - )); - } - - #[test] - fn summarizes_column_chunks() { - let files = vec![ - FileReport { - path: "a.parquet".to_string(), - relative_path: "a.parquet".to_string(), - file_size: 100, - num_rows: 10, - num_row_groups: 1, - columns: vec![ColumnChunkReport { - row_group_index: 0, - column_index: 1, - column_path: "greptime_value".to_string(), - encodings: vec!["PLAIN".to_string(), "RLE".to_string()], - compression: "ZSTD".to_string(), - compressed_size: 30, - uncompressed_size: 80, - num_values: 10, - }], - }, - FileReport { - path: "b.parquet".to_string(), - relative_path: "b.parquet".to_string(), - file_size: 50, - num_rows: 5, - num_row_groups: 1, - columns: vec![ColumnChunkReport { - row_group_index: 0, - column_index: 1, - column_path: "greptime_value".to_string(), - encodings: vec!["BYTE_STREAM_SPLIT".to_string(), "RLE".to_string()], - compression: "ZSTD".to_string(), - compressed_size: 10, - uncompressed_size: 20, - num_values: 5, - }], - }, - ]; - - let summary = summarize("greptime_value", &files); - assert_eq!(summary.file_count, 2); - assert_eq!(summary.files_with_column, 2); - assert_eq!(summary.total_file_size, 150); - assert_eq!(summary.total_rows, 15); - assert_eq!(summary.column_compressed_size, 40); - assert_eq!(summary.column_uncompressed_size, 100); - assert_eq!( - summary.unique_encodings, - vec!["BYTE_STREAM_SPLIT", "PLAIN", "RLE"] - ); - } -} diff --git a/src/cmd/src/bin/prom_remote_write_fixture.rs b/src/cmd/src/bin/prom_remote_write_fixture.rs deleted file mode 100644 index 89fd20622bf..00000000000 --- a/src/cmd/src/bin/prom_remote_write_fixture.rs +++ /dev/null @@ -1,426 +0,0 @@ -// 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. - -#![allow(clippy::print_stdout)] - -use std::time::{Duration, Instant}; - -use api::prom_store::remote::{Label, Sample, TimeSeries, WriteRequest}; -use clap::{Parser, ValueEnum}; -use prost::Message; -use serde_json::json; -use servers::prom_store::snappy_compress; - -#[derive(Debug, Parser)] -#[command(about = "Generate a deterministic Prometheus remote-write fixture")] -struct Args { - #[arg(long, default_value = "http://127.0.0.1:4000/v1/prometheus/write")] - endpoint: String, - #[arg(long, default_value = "public")] - database: String, - #[arg(long)] - metric: String, - #[arg(long, default_value = "greptime_physical_table")] - physical_table: String, - #[arg(long, default_value_t = 8)] - series_count: u64, - #[arg(long, default_value_t = 30)] - samples_per_series: u64, - #[arg(long, default_value_t = 1_704_067_200_000)] - start_unix_millis: i64, - #[arg(long, default_value_t = 15_000)] - step_millis: i64, - #[arg(long, alias = "batch-size", default_value_t = 8)] - chunk_series_count: u64, - #[arg(long, default_value_t = 60)] - timeout_seconds: u64, - #[arg(long, default_value_t = ValuePattern::Linear)] - value_pattern: ValuePattern, - #[arg(long, default_value_t = 0.0)] - value_base: f64, - #[arg(long, default_value_t = 0.125)] - value_step: f64, - #[arg(long, default_value_t = 97)] - value_cardinality: u64, - #[arg(long, default_value_t = 0)] - value_seed: u64, - #[arg(long, default_value_t = 8)] - value_run_length: u64, - #[arg(long, default_value_t = 100)] - value_stall_every: u64, - #[arg(long, default_value_t = 16)] - value_stall_length: u64, - #[arg(long, default_value_t = 5)] - value_mixed_every: u64, - #[arg(long, default_value_t = 0)] - value_sample_offset: u64, - #[arg(long)] - value_total_samples_per_series: Option, -} - -#[derive(Clone, Copy, Debug, Eq, PartialEq, ValueEnum)] -#[value(rename_all = "snake_case")] -enum ValuePattern { - Linear, - Constant, - Modulo, - Unique, - SeededRandom, - RunLength, - QuantizedSignal, - SignalWithSporadicStalls, - MixedSignalRepeated, -} - -impl std::fmt::Display for ValuePattern { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - Self::Linear => write!(f, "linear"), - Self::Constant => write!(f, "constant"), - Self::Modulo => write!(f, "modulo"), - Self::Unique => write!(f, "unique"), - Self::SeededRandom => write!(f, "seeded_random"), - Self::RunLength => write!(f, "run_length"), - Self::QuantizedSignal => write!(f, "quantized_signal"), - Self::SignalWithSporadicStalls => write!(f, "signal_with_sporadic_stalls"), - Self::MixedSignalRepeated => write!(f, "mixed_signal_repeated"), - } - } -} - -#[tokio::main] -async fn main() -> Result<(), Box> { - let args = Args::parse(); - let started = Instant::now(); - let client = reqwest::Client::builder() - .timeout(Duration::from_secs(args.timeout_seconds)) - .build()?; - let mut batches = 0_u64; - let mut rows = 0_u64; - let mut http_statuses = Vec::new(); - - let chunk = args.chunk_series_count.max(1); - for first in (0..args.series_count).step_by(chunk as usize) { - let last = (first + chunk).min(args.series_count); - let timeseries = (first..last) - .map(|series_idx| make_series(&args, series_idx)) - .collect::>(); - rows += timeseries - .iter() - .map(|ts| ts.samples.len() as u64) - .sum::(); - let body = snappy_compress( - &WriteRequest { - timeseries, - ..Default::default() - } - .encode_to_vec(), - )?; - let response = client - .post(&args.endpoint) - .header("Content-Encoding", "snappy") - .body(body) - .send() - .await?; - let status = response.status(); - http_statuses.push(status.as_u16()); - if !status.is_success() { - let body_text = response.text().await.unwrap_or_default(); - return Err(format!("remote-write failed with status {status}: {body_text}").into()); - } - batches += 1; - } - - println!( - "{}", - json!({ - "status": "ok", - "endpoint": args.endpoint, - "database": args.database, - "metric": args.metric, - "physical_table": args.physical_table, - "series_count": args.series_count, - "samples_per_series": args.samples_per_series, - "rows": rows, - "samples_written": rows, - "batches": batches, - "elapsed_seconds": started.elapsed().as_secs_f64(), - "http_statuses": http_statuses, - "value": { - "pattern": args.value_pattern.to_string(), - "base": args.value_base, - "step": args.value_step, - "cardinality": args.value_cardinality, - "seed": args.value_seed, - "run_length": args.value_run_length, - "stall_every": args.value_stall_every, - "stall_length": args.value_stall_length, - "mixed_every": args.value_mixed_every, - "sample_offset": args.value_sample_offset, - "total_samples_per_series": args.value_total_samples_per_series.unwrap_or(args.samples_per_series), - }, - }) - ); - Ok(()) -} - -fn make_series(args: &Args, series_idx: u64) -> TimeSeries { - let mut labels = vec![ - label("__name__", &args.metric), - label("x_greptime_database", &args.database), - label("x_greptime_physical_table", &args.physical_table), - label("host", &format!("host{:04}", series_idx % 1024)), - label("instance", &format!("instance{:06}", series_idx)), - ]; - labels.sort_unstable_by(|left, right| left.name.cmp(&right.name)); - let samples = (0..args.samples_per_series) - .map(|sample_idx| Sample { - value: deterministic_value(args, series_idx, sample_idx), - timestamp: args.start_unix_millis + (sample_idx as i64 * args.step_millis), - }) - .collect(); - TimeSeries { - labels, - samples, - ..Default::default() - } -} - -fn label(name: &str, value: &str) -> Label { - Label { - name: name.to_string(), - value: value.to_string(), - } -} - -fn deterministic_value(args: &Args, series_idx: u64, sample_idx: u64) -> f64 { - let ordinal = value_ordinal(args, series_idx, sample_idx); - let series_sample_ordinal = value_series_sample_ordinal(args, sample_idx); - match args.value_pattern { - ValuePattern::Linear => linear_value(args, series_idx, series_sample_ordinal), - ValuePattern::Constant => args.value_base, - ValuePattern::Modulo => { - args.value_base + (ordinal % args.value_cardinality.max(1)) as f64 * args.value_step - } - ValuePattern::Unique => args.value_base + ordinal as f64 * args.value_step, - ValuePattern::SeededRandom => { - let bucket = splitmix64(ordinal ^ args.value_seed) % args.value_cardinality.max(1); - args.value_base + bucket as f64 * args.value_step - } - ValuePattern::RunLength => { - let bucket = (ordinal / args.value_run_length.max(1)) % args.value_cardinality.max(1); - args.value_base + bucket as f64 * args.value_step - } - ValuePattern::QuantizedSignal => { - let bucket = ((series_idx % args.value_cardinality.max(1)) - + (series_sample_ordinal / args.value_run_length.max(1))) - % args.value_cardinality.max(1); - args.value_base + bucket as f64 * args.value_step - } - ValuePattern::SignalWithSporadicStalls => { - let stall_every = args.value_stall_every.max(1); - let stall_length = args.value_stall_length.min(stall_every); - let phase = series_sample_ordinal % stall_every; - if phase < stall_length { - linear_value(args, series_idx, series_sample_ordinal - phase) - } else { - linear_value(args, series_idx, series_sample_ordinal) - } - } - ValuePattern::MixedSignalRepeated => { - if series_sample_ordinal.is_multiple_of(args.value_mixed_every.max(1)) { - args.value_base - } else { - linear_value(args, series_idx, series_sample_ordinal) - } - } - } -} - -fn linear_value(args: &Args, series_idx: u64, series_sample_ordinal: u64) -> f64 { - args.value_base + (series_idx % 97) as f64 + series_sample_ordinal as f64 * args.value_step -} - -fn value_ordinal(args: &Args, series_idx: u64, sample_idx: u64) -> u64 { - let total_samples_per_series = args - .value_total_samples_per_series - .unwrap_or(args.samples_per_series); - series_idx * total_samples_per_series + args.value_sample_offset + sample_idx -} - -fn value_series_sample_ordinal(args: &Args, sample_idx: u64) -> u64 { - args.value_sample_offset + sample_idx -} - -fn splitmix64(mut value: u64) -> u64 { - value = value.wrapping_add(0x9e3779b97f4a7c15); - let mut mixed = value; - mixed = (mixed ^ (mixed >> 30)).wrapping_mul(0xbf58476d1ce4e5b9); - mixed = (mixed ^ (mixed >> 27)).wrapping_mul(0x94d049bb133111eb); - mixed ^ (mixed >> 31) -} - -#[cfg(test)] -mod tests { - use super::*; - - fn args(value_pattern: ValuePattern) -> Args { - Args { - endpoint: "http://127.0.0.1:4000/v1/prometheus/write".to_string(), - database: "public".to_string(), - metric: "metric".to_string(), - physical_table: "greptime_physical_table".to_string(), - series_count: 8, - samples_per_series: 30, - start_unix_millis: 1_704_067_200_000, - step_millis: 15_000, - chunk_series_count: 8, - timeout_seconds: 60, - value_pattern, - value_base: 0.0, - value_step: 0.125, - value_cardinality: 97, - value_seed: 0, - value_run_length: 8, - value_stall_every: 100, - value_stall_length: 16, - value_mixed_every: 5, - value_sample_offset: 0, - value_total_samples_per_series: None, - } - } - - #[test] - fn default_linear_matches_old_formula() { - let args = args(ValuePattern::Linear); - for series_idx in [0, 1, 96, 97, 194] { - for sample_idx in [0, 1, 10, 29] { - assert_eq!( - deterministic_value(&args, series_idx, sample_idx), - ((series_idx % 97) as f64) + (sample_idx as f64 * 0.125) - ); - } - } - } - - #[test] - fn constant_uses_base() { - let mut args = args(ValuePattern::Constant); - args.value_base = 42.5; - assert_eq!(deterministic_value(&args, 7, 11), 42.5); - } - - #[test] - fn modulo_uses_global_ordinal_and_cardinality() { - let mut args = args(ValuePattern::Modulo); - args.samples_per_series = 3; - args.value_base = 10.0; - args.value_step = 2.0; - args.value_cardinality = 4; - assert_eq!(deterministic_value(&args, 1, 2), 12.0); - assert_eq!(deterministic_value(&args, 2, 0), 14.0); - } - - #[test] - fn unique_uses_global_ordinal() { - let mut args = args(ValuePattern::Unique); - args.samples_per_series = 10; - args.value_base = 1.0; - args.value_step = 0.5; - assert_eq!(deterministic_value(&args, 2, 3), 12.5); - } - - #[test] - fn unique_uses_chunk_offset_and_total_samples() { - let mut args = args(ValuePattern::Unique); - args.samples_per_series = 3; - args.value_total_samples_per_series = Some(10); - args.value_sample_offset = 3; - args.value_step = 1.0; - assert_eq!(deterministic_value(&args, 2, 0), 23.0); - } - - #[test] - fn seeded_random_is_deterministic() { - let mut args = args(ValuePattern::SeededRandom); - args.value_cardinality = 5; - args.value_seed = 123; - let first = deterministic_value(&args, 4, 5); - assert_eq!(first, deterministic_value(&args, 4, 5)); - args.value_seed = 124; - assert_ne!(first, deterministic_value(&args, 4, 5)); - } - - #[test] - fn run_length_repeats_each_bucket_for_configured_window() { - let mut args = args(ValuePattern::RunLength); - args.samples_per_series = 10; - args.value_step = 1.0; - args.value_cardinality = 4; - args.value_run_length = 3; - assert_eq!(deterministic_value(&args, 0, 0), 0.0); - assert_eq!(deterministic_value(&args, 0, 2), 0.0); - assert_eq!(deterministic_value(&args, 0, 3), 1.0); - assert_eq!(deterministic_value(&args, 1, 0), 3.0); - } - - #[test] - fn quantized_signal_uses_series_local_quantized_steps() { - let mut args = args(ValuePattern::QuantizedSignal); - args.value_base = 10.0; - args.value_step = 2.0; - args.value_cardinality = 5; - args.value_run_length = 4; - assert_eq!(deterministic_value(&args, 2, 0), 14.0); - assert_eq!(deterministic_value(&args, 2, 3), 14.0); - assert_eq!(deterministic_value(&args, 2, 4), 16.0); - } - - #[test] - fn quantized_signal_respects_chunk_offset() { - let mut args = args(ValuePattern::QuantizedSignal); - args.value_step = 1.0; - args.value_cardinality = 8; - args.value_run_length = 4; - args.value_sample_offset = 4; - assert_eq!(deterministic_value(&args, 0, 0), 1.0); - assert_eq!(deterministic_value(&args, 0, 3), 1.0); - assert_eq!(deterministic_value(&args, 0, 4), 2.0); - } - - #[test] - fn sporadic_stalls_hold_value_within_stall_window() { - let mut args = args(ValuePattern::SignalWithSporadicStalls); - args.value_step = 1.0; - args.value_stall_every = 5; - args.value_stall_length = 2; - assert_eq!(deterministic_value(&args, 0, 0), 0.0); - assert_eq!(deterministic_value(&args, 0, 1), 0.0); - assert_eq!(deterministic_value(&args, 0, 2), 2.0); - assert_eq!(deterministic_value(&args, 0, 5), 5.0); - assert_eq!(deterministic_value(&args, 0, 6), 5.0); - } - - #[test] - fn mixed_signal_repeated_inserts_periodic_base_value() { - let mut args = args(ValuePattern::MixedSignalRepeated); - args.value_base = 100.0; - args.value_step = 1.0; - args.value_mixed_every = 3; - assert_eq!(deterministic_value(&args, 0, 0), 100.0); - assert_eq!(deterministic_value(&args, 0, 1), 101.0); - assert_eq!(deterministic_value(&args, 0, 2), 102.0); - assert_eq!(deterministic_value(&args, 0, 3), 100.0); - } -} diff --git a/src/cmd/src/bin/query_perf_fixture.rs b/src/cmd/src/bin/query_perf_fixture.rs index dfc3a664a03..d0e185038fa 100644 --- a/src/cmd/src/bin/query_perf_fixture.rs +++ b/src/cmd/src/bin/query_perf_fixture.rs @@ -14,656 +14,10 @@ #![allow(clippy::print_stderr, clippy::print_stdout)] -//! Generic direct-readable-SST fixture generator for query performance labs. -//! -//! This is a bounded MVP: it reads one table/one region case TOML, builds -//! synthetic region metadata, writes real Mito readable SST parquet files, and -//! emits a manifest checkpoint plus validation metadata. The first case uses -//! `sst_format = "flat"`; primary-key output is also accepted for early lab use -//! because the writer can convert the same flat input batch to primary-key SSTs. -//! -//! TODO(query-perf): add multi-region generation and seed-manifest/catalog DB -//! integration so the generated fixture can be mounted directly by a full -//! GreptimeDB open/query path without manual metadata wiring. - -use std::collections::HashMap; -use std::fs; -use std::io::{BufWriter, Write}; -use std::num::{NonZeroU64, NonZeroUsize}; -use std::path::{Path, PathBuf}; -use std::sync::Arc; - -use api::v1::{OpType, SemanticType}; -use async_trait::async_trait; -use clap::Parser; -use datatypes::arrow::array::{ - ArrayRef, BinaryDictionaryBuilder, Float64Array, RecordBatch, StringDictionaryBuilder, - TimestampMillisecondArray, TimestampNanosecondArray, UInt8Array, UInt64Array, -}; -use datatypes::arrow::datatypes::UInt32Type; -use datatypes::prelude::ConcreteDataType; -use datatypes::schema::{ColumnSchema, SkippingIndexOptions}; -use mito_codec::row_converter::{DensePrimaryKeyCodec, PrimaryKeyCodecExt, SortField}; -use mito2::access_layer::{FilePathProvider, Metrics, WriteType}; -use mito2::config::IndexConfig; -use mito2::manifest::action::{RegionCheckpoint, RegionManifest, RemovedFilesRecord}; -use mito2::read::FlatSource; -use mito2::sst::file::{FileMeta, RegionFileId}; -use mito2::sst::index::{Indexer, IndexerBuilder}; -use mito2::sst::parquet::writer::ParquetWriter; -use mito2::sst::parquet::{SstInfo, WriteOptions}; -use mito2::sst::{ - DEFAULT_WRITE_BUFFER_SIZE, FlatSchemaOptions, FormatType, to_flat_sst_arrow_schema, -}; -use object_store::ObjectStore; -use object_store::services::Fs as FsBuilder; -use serde::Deserialize; -use store_api::codec::PrimaryKeyEncoding; -use store_api::metadata::{ColumnMetadata, RegionMetadata, RegionMetadataRef}; -use store_api::path_utils::region_name; -use store_api::region_request::PathType; -use store_api::storage::{FileId, RegionId}; - -#[derive(Parser, Debug)] -#[command(name = "query_perf_fixture")] -#[command(about = "Generate direct-readable SST fixtures for query performance cases")] -struct Args { - /// TOML case file, for example tests/perf/query_cases/.../case.toml. - #[arg(long, value_name = "PATH")] - case: PathBuf, - - /// Output directory (creates object-store/, manifest/, files.jsonl, summary.json). - #[arg(long, value_name = "DIR")] - out_dir: PathBuf, - - /// Region ID to synthesize. Defaults to table id 1024, region 0. - #[arg(long, default_value = "4398046511104")] - region_id: u64, - - /// Table directory relative to object-store root. Defaults to data/{database}/{table}/. - #[arg(long, value_name = "DIR")] - table_dir: Option, - - /// Table name to generate when the case contains multiple [[tables]]. - #[arg(long, value_name = "NAME")] - table: Option, - - /// Manifest/checkpoint version. - #[arg(long, default_value = "1000000")] - checkpoint_version: u64, - - /// Safety flag when sst_count exceeds 1000. - #[arg(long)] - allow_large: bool, - - /// Print plan only. - #[arg(long)] - dry_run: bool, -} - -#[derive(Debug, Deserialize)] -struct CaseFile { - scenario: Scenario, -} - -#[derive(Debug, Deserialize)] -#[serde(tag = "kind")] -enum Scenario { - #[serde(rename = "direct_readable_sst")] - DirectReadableSst(DirectReadableSstScenario), -} - -#[derive(Debug, Deserialize)] -struct DirectReadableSstScenario { - #[serde(default)] - seed: Option, - tables: Vec, - layout: LayoutConfig, -} - -impl Scenario { - fn kind(&self) -> &'static str { - match self { - Scenario::DirectReadableSst(_) => "direct_readable_sst", - } - } - - fn direct_readable_sst(&self) -> &DirectReadableSstScenario { - match self { - Scenario::DirectReadableSst(scenario) => scenario, - } - } -} - -#[derive(Debug, Deserialize)] -struct TableConfig { - database: String, - name: String, - engine: String, - #[serde(default)] - append_mode: Option, - #[serde(default)] - sst_format: Option, - primary_key: Vec, - time_index: String, - columns: Vec, -} - -#[derive(Debug, Deserialize)] -struct ColumnConfig { - name: String, - #[serde(rename = "type")] - ty: String, - semantic: String, - distribution: Option, -} - -#[derive(Debug, Deserialize)] -#[serde(tag = "kind")] -enum Distribution { - #[serde(rename = "cardinality")] - Cardinality { - values: NonZeroUsize, - prefix: String, - }, - #[serde(rename = "deterministic_wave")] - DeterministicWave { min: f64, max: f64 }, -} - -#[derive(Debug, Deserialize)] -struct LayoutConfig { - regions: usize, - sst_count: usize, - rows_per_sst: usize, - row_group_size: usize, - series_count: NonZeroUsize, - start_unix_nanos: i64, - step_nanos: i64, - time_range_layout: String, - series_layout: String, -} - -struct NoopIndexBuilder; - -#[async_trait] -impl IndexerBuilder for NoopIndexBuilder { - async fn build(&self, _file_id: FileId, _index_version: u64) -> Indexer { - Indexer::default() - } -} - -#[derive(Clone)] -struct FixedPathProvider { - table_dir: String, -} - -impl FilePathProvider for FixedPathProvider { - fn build_index_file_path(&self, file_id: RegionFileId) -> String { - mito2::sst::location::index_file_path_legacy(&self.table_dir, file_id, PathType::Bare) - } - fn build_index_file_path_with_version( - &self, - index_id: mito2::sst::file::RegionIndexId, - ) -> String { - mito2::sst::location::index_file_path(&self.table_dir, index_id, PathType::Bare) - } - fn build_sst_file_path(&self, file_id: RegionFileId) -> String { - mito2::sst::location::sst_file_path(&self.table_dir, file_id, PathType::Bare) - } -} - -fn semantic_type(s: &str) -> SemanticType { - match s.to_ascii_lowercase().as_str() { - "tag" => SemanticType::Tag, - "field" => SemanticType::Field, - "timestamp" => SemanticType::Timestamp, - other => panic!("unsupported semantic {other}"), - } -} - -fn concrete_type(ty: &str) -> ConcreteDataType { - match ty.to_ascii_uppercase().as_str() { - "STRING" => ConcreteDataType::string_datatype(), - "DOUBLE" => ConcreteDataType::float64_datatype(), - "UINT64" => ConcreteDataType::uint64_datatype(), - "TIMESTAMP(9)" => ConcreteDataType::timestamp_nanosecond_datatype(), - "TIMESTAMP(3)" => ConcreteDataType::timestamp_millisecond_datatype(), - other => panic!("unsupported column type {other}"), - } -} - -fn build_region_metadata(table: &TableConfig, region_id: RegionId) -> RegionMetadata { - let mut builder = store_api::metadata::RegionMetadataBuilder::new(region_id); - for (idx, col) in table.columns.iter().enumerate() { - let mut schema = ColumnSchema::new( - &col.name, - concrete_type(&col.ty), - col.semantic != "timestamp", - ); - if col.semantic == "tag" { - schema = schema.with_inverted_index(true); - } - if col.semantic == "tag" || col.semantic == "field" { - schema = schema - .with_skipping_options(SkippingIndexOptions { - granularity: 1, - ..Default::default() - }) - .expect("valid skipping index options for query perf fixture"); - } - if col.name == table.time_index { - schema = schema.with_time_index(true); - } - builder.push_column_metadata(ColumnMetadata { - column_schema: schema, - semantic_type: semantic_type(&col.semantic), - column_id: idx as u32, - }); - } - let pk_ids = table - .primary_key - .iter() - .map(|name| { - table - .columns - .iter() - .position(|c| c.name == *name) - .unwrap_or_else(|| panic!("primary key column {name} not found")) as u32 - }) - .collect(); - builder.primary_key(pk_ids); - builder.primary_key_encoding(PrimaryKeyEncoding::Dense); - builder - .build() - .expect("region metadata should be valid for fixture table config") -} - -fn encode_dense_primary_key(table: &TableConfig, values: &HashMap) -> Vec { - let fields = table - .primary_key - .iter() - .enumerate() - .map(|(idx, _)| { - ( - idx as u32, - SortField::new(ConcreteDataType::string_datatype()), - ) - }) - .collect(); - let converter = DensePrimaryKeyCodec::with_fields(fields); - converter - .encode( - table - .primary_key - .iter() - .map(|name| datatypes::value::ValueRef::String(values[name].as_str())), - ) - .expect("dense primary key encoding should match fixture primary key fields") -} - -fn tag_value(col: &ColumnConfig, series: usize) -> String { - match col - .distribution - .as_ref() - .expect("tag distribution is required") - { - Distribution::Cardinality { values, prefix } => { - format!("{}{}", prefix, series % values.get()) - } - _ => panic!("tag column {} requires cardinality distribution", col.name), - } -} - -fn series_for_row(layout: &LayoutConfig, sst_idx: usize, base_row: usize, row: usize) -> usize { - match layout.series_layout.as_str() { - "round_robin" | "timestamp_major" => (base_row + row) % layout.series_count.get(), - "per_sst" => sst_idx % layout.series_count.get(), - other => panic!("unsupported series_layout {other}"), - } -} - -fn timestamp_for_row(layout: &LayoutConfig, base_row: usize, row: usize) -> i64 { - let logical_row = base_row + row; - let timestamp_step = if layout.series_layout == "timestamp_major" { - logical_row / layout.series_count.get() - } else { - logical_row - }; - - layout.start_unix_nanos + (timestamp_step as i64 * layout.step_nanos) -} - -fn wave_value(min: f64, max: f64, row: usize) -> f64 { - let span = max - min; - let phase = (row % 1024) as f64 / 1023.0; - min + span * (0.5 - 0.5 * (std::f64::consts::TAU * phase).cos()) -} - -fn generate_record_batch( - table: &TableConfig, - metadata: &RegionMetadataRef, - layout: &LayoutConfig, - sst_idx: usize, - sequence: u64, -) -> RecordBatch { - let flat_schema = to_flat_sst_arrow_schema(metadata, &FlatSchemaOptions::default()); - let rows = layout.rows_per_sst; - let mut columns = Vec::new(); - let base_row = sst_idx * rows; - - for col in &table.columns { - match (col.semantic.as_str(), col.ty.to_ascii_uppercase().as_str()) { - ("tag", "STRING") => { - let mut b = StringDictionaryBuilder::::new(); - for row in 0..rows { - let series = series_for_row(layout, sst_idx, base_row, row); - let value = tag_value(col, series); - b.append_value(value); - } - columns.push(Arc::new(b.finish()) as ArrayRef); - } - ("field", "DOUBLE") => { - let (min, max) = match col.distribution.as_ref() { - Some(Distribution::DeterministicWave { min, max }) => (*min, *max), - _ => (0.0, 1.0), - }; - columns.push(Arc::new(Float64Array::from( - (0..rows) - .map(|r| wave_value(min, max, base_row + r)) - .collect::>(), - )) as ArrayRef); - } - ("field", "UINT64") => columns.push(Arc::new(UInt64Array::from( - (0..rows).map(|r| (base_row + r) as u64).collect::>(), - )) as ArrayRef), - ("timestamp", "TIMESTAMP(9)") => columns.push(Arc::new(TimestampNanosecondArray::from( - (0..rows) - .map(|r| timestamp_for_row(layout, base_row, r)) - .collect::>(), - )) as ArrayRef), - ("timestamp", "TIMESTAMP(3)") => { - columns.push(Arc::new(TimestampMillisecondArray::from( - (0..rows) - .map(|r| timestamp_for_row(layout, base_row, r) / 1_000_000) - .collect::>(), - )) as ArrayRef) - } - other => panic!("unsupported column combination {other:?}"), - } - } - - let mut pk_builder = BinaryDictionaryBuilder::::new(); - for row in 0..rows { - let series = series_for_row(layout, sst_idx, base_row, row); - let values = table - .columns - .iter() - .filter(|c| c.semantic == "tag") - .map(|c| (c.name.clone(), tag_value(c, series))) - .collect(); - pk_builder - .append(encode_dense_primary_key(table, &values)) - .expect("append generated dense primary key to Arrow dictionary builder"); - } - columns.push(Arc::new(pk_builder.finish())); - columns.push(Arc::new(UInt64Array::from_value(sequence, rows))); - columns.push(Arc::new(UInt8Array::from_value(OpType::Put as u8, rows))); - RecordBatch::try_new(flat_schema, columns) - .expect("generated fixture columns should match flat SST Arrow schema") -} - -fn file_meta_from_sst_info( - info: &SstInfo, - region_id: RegionId, - file_id: FileId, - sequence: u64, -) -> FileMeta { - FileMeta { - region_id, - file_id, - time_range: info.time_range, - level: 0, - file_size: info.file_size, - max_row_group_uncompressed_size: info.max_row_group_uncompressed_size, - available_indexes: Default::default(), - indexes: Default::default(), - index_file_size: 0, - index_version: 0, - num_rows: info.num_rows as u64, - num_row_groups: info.num_row_groups, - sequence: NonZeroU64::new(sequence), - partition_expr: None, - num_series: info.num_series, - primary_key_min: None, - primary_key_max: None, - } -} - -fn deterministic_file_id(seed: u64, index: usize) -> FileId { - let seed_hi = ((seed >> 16) & 0xffff) as u16; - let seed_lo = (seed & 0xffff) as u16; - FileId::parse_str(&format!( - "00000000-0000-{seed_hi:04x}-{seed_lo:04x}-{index:012x}" - )) - .expect("deterministic file id must be a valid UUID") -} - -fn remap_sst_file( - object_store_dir: &Path, - table_dir: &str, - region_id: RegionId, - old_file_id: FileId, - new_file_id: FileId, -) { - if old_file_id == new_file_id { - return; - } - let old_path = object_store_dir.join(mito2::sst::location::sst_file_path( - table_dir, - RegionFileId::new(region_id, old_file_id), - PathType::Bare, - )); - let new_path = object_store_dir.join(mito2::sst::location::sst_file_path( - table_dir, - RegionFileId::new(region_id, new_file_id), - PathType::Bare, - )); - if let Some(parent) = new_path.parent() { - fs::create_dir_all(parent).expect("failed to create remapped SST parent directory"); - } - fs::rename(&old_path, &new_path).unwrap_or_else(|err| { - panic!( - "failed to remap SST file {} -> {}: {err}", - old_path.display(), - new_path.display() - ) - }); -} - -fn case_name_from_path(path: &Path) -> String { - path.parent() - .and_then(Path::file_name) - .and_then(|name| name.to_str()) - .unwrap_or("query_perf_case") - .to_string() -} +#[path = "query_perf_fixture/mod.rs"] +mod query_perf_fixture; #[tokio::main] async fn main() { - let args = Args::parse(); - let case_text = fs::read_to_string(&args.case).expect("failed to read query perf case TOML"); - let case: CaseFile = toml::from_str(&case_text).expect("failed to parse query perf case TOML"); - let case_name = case_name_from_path(&args.case); - let scenario = case.scenario.direct_readable_sst(); - if scenario.tables.is_empty() || scenario.layout.regions != 1 { - panic!("fixture generator supports one or more tables and exactly one region per table"); - } - if scenario.layout.time_range_layout != "non_overlapping_per_sst" { - panic!("MVP supports time_range_layout=non_overlapping_per_sst"); - } - if scenario.layout.series_layout == "timestamp_major" - && !scenario - .layout - .rows_per_sst - .is_multiple_of(scenario.layout.series_count.get()) - { - panic!( - "series_layout=timestamp_major requires rows_per_sst to be divisible by series_count" - ); - } - if scenario.layout.sst_count > 1000 && !args.allow_large { - panic!("sst_count exceeds 1000; pass --allow-large"); - } - let seed = scenario.seed.unwrap_or(0); - let table_index = match args.table.as_deref() { - Some(name) => scenario - .tables - .iter() - .position(|table| table.name == name) - .unwrap_or_else(|| panic!("--table {name} was not found in case {case_name}")), - None if scenario.tables.len() == 1 => 0, - None => panic!( - "case {} contains {} tables; pass --table to choose one", - case_name, - scenario.tables.len() - ), - }; - let table = &scenario.tables[table_index]; - if table.engine != "mito" { - panic!("MVP supports engine=mito"); - } - let region_id = RegionId::from(args.region_id); - let table_dir = args - .table_dir - .unwrap_or_else(|| format!("data/{}/{}/", table.database, table.name)); - if Path::new(&table_dir).is_absolute() { - panic!("--table-dir must be relative"); - } - if table_dir.trim_end_matches('/').rsplit('/').next() - == Some(region_name(region_id.table_id(), region_id.region_sequence()).as_str()) - { - panic!("--table-dir must be table-level directory"); - } - let region_dir = - mito2::sst::location::region_dir_from_table_dir(&table_dir, region_id, PathType::Bare); - let format = match table.sst_format.as_deref().unwrap_or("flat") { - "flat" => FormatType::Flat, - "primary_key" | "primary-key" => FormatType::PrimaryKey, - other => panic!("unsupported sst_format {other}"), - }; - println!( - "query_perf_fixture case={} scenario={} table={}.{} ssts={} rows_per_sst={} format={format:?}", - case_name, - case.scenario.kind(), - table.database, - table.name, - scenario.layout.sst_count, - scenario.layout.rows_per_sst - ); - println!( - "out_dir={} table_dir={} region_dir={}", - args.out_dir.display(), - table_dir, - region_dir - ); - if args.dry_run { - return; - } - - let obj_store_dir = args.out_dir.join("object-store"); - let manifest_dir = args.out_dir.join("manifest"); - fs::create_dir_all(&obj_store_dir).expect("failed to create fixture object-store directory"); - fs::create_dir_all(&manifest_dir).expect("failed to create fixture manifest directory"); - let ostorage = ObjectStore::new(FsBuilder::default().root(&obj_store_dir.to_string_lossy())) - .expect("failed to create filesystem object store for fixture output") - .finish(); - let metadata: RegionMetadataRef = Arc::new(build_region_metadata(table, region_id)); - let mut files = HashMap::with_capacity(scenario.layout.sst_count); - let mut next_file_index = 1; - for i in 0..scenario.layout.sst_count { - let sequence = 1000 + i as u64; - let batch = generate_record_batch(table, &metadata, &scenario.layout, i, sequence); - let source = - FlatSource::new_iter(batch.schema(), Box::new(vec![batch].into_iter().map(Ok))); - let mut metrics = Metrics::new(WriteType::Flush); - let mut writer = ParquetWriter::new_with_object_store( - ostorage.clone(), - metadata.clone(), - IndexConfig::default(), - NoopIndexBuilder, - FixedPathProvider { - table_dir: table_dir.clone(), - }, - &mut metrics, - ) - .await; - let opts = WriteOptions { - write_buffer_size: DEFAULT_WRITE_BUFFER_SIZE, - row_group_size: scenario.layout.row_group_size, - max_file_size: None, - }; - let infos = match format { - FormatType::Flat => writer.write_all_flat(source, Some(sequence), &opts).await, - FormatType::PrimaryKey => { - writer - .write_all_flat_as_primary_key(source, Some(sequence), &opts) - .await - } - } - .expect("failed to write fixture SST through Mito parquet writer"); - for info in infos { - let file_id = deterministic_file_id( - seed.wrapping_add((table_index as u64) << 16), - next_file_index, - ); - next_file_index += 1; - remap_sst_file(&obj_store_dir, &table_dir, region_id, info.file_id, file_id); - files.insert( - file_id, - file_meta_from_sst_info(&info, region_id, file_id, sequence), - ); - } - } - let overall_seq = files - .values() - .filter_map(|m| m.sequence) - .map(|s| s.get()) - .max() - .unwrap_or(0); - let manifest = RegionManifest { - metadata, - files, - removed_files: RemovedFilesRecord::default(), - flushed_entry_id: overall_seq, - flushed_sequence: overall_seq, - committed_sequence: Some(overall_seq), - manifest_version: args.checkpoint_version, - truncated_entry_id: None, - compaction_time_window: None, - sst_format: format, - append_mode: table.append_mode, - }; - let checkpoint = RegionCheckpoint { - last_version: args.checkpoint_version, - compacted_actions: manifest.files.len(), - checkpoint: Some(manifest.clone()), - }; - let checkpoint_bytes = checkpoint - .encode() - .expect("failed to encode fixture region checkpoint"); - let checkpoint_path = manifest_dir.join(format!("{:020}.checkpoint", args.checkpoint_version)); - fs::write(&checkpoint_path, &checkpoint_bytes).expect("failed to write checkpoint file"); - fs::write(manifest_dir.join("_last_checkpoint"), serde_json::to_vec_pretty(&serde_json::json!({ "size": checkpoint_bytes.len(), "version": args.checkpoint_version, "checksum": null, "extend_metadata": {} })).expect("failed to serialize _last_checkpoint metadata")).expect("failed to write _last_checkpoint file"); - let files_jsonl_path = args.out_dir.join("files.jsonl"); - let mut jsonl = BufWriter::new( - fs::File::create(&files_jsonl_path).expect("failed to create fixture files.jsonl"), - ); - let mut entries: Vec<_> = manifest.files.iter().collect(); - entries.sort_by_key(|(id, _)| id.to_string()); - for (file_id, meta) in entries { - writeln!(jsonl, "{}", serde_json::to_string(&serde_json::json!({ "file_id": file_id.to_string(), "region_id": meta.region_id.as_u64(), "object_path": mito2::sst::location::sst_file_path(&table_dir, RegionFileId::new(meta.region_id, *file_id), PathType::Bare), "time_range_start": meta.time_range.0.value(), "time_range_end": meta.time_range.1.value(), "num_rows": meta.num_rows, "num_row_groups": meta.num_row_groups, "file_size": meta.file_size, "num_series": meta.num_series, "sequence": meta.sequence.map(|s| s.get()) })).expect("failed to serialize fixture SST metadata JSONL entry")).expect("failed to write fixture SST metadata JSONL entry"); - } - jsonl.flush().expect("failed to flush fixture files.jsonl"); - fs::write(args.out_dir.join("summary.json"), serde_json::to_vec_pretty(&serde_json::json!({ "case": case_name, "seed": seed, "table_index": table_index, "table": table.name, "database": table.database, "region_id": region_id.as_u64(), "table_dir": table_dir, "region_dir": region_dir, "sst_format": format!("{format:?}"), "sst_count": scenario.layout.sst_count, "rows_per_sst": scenario.layout.rows_per_sst, "row_group_size": scenario.layout.row_group_size, "total_rows": scenario.layout.sst_count * scenario.layout.rows_per_sst, "checkpoint_path": checkpoint_path, "files_jsonl_path": files_jsonl_path, "readback_validated": false, "metadata_source": "synthetic" })).expect("failed to serialize fixture summary")).expect("failed to write fixture summary.json"); - println!("Done. wrote {} SST file entries", manifest.files.len()); + query_perf_fixture::run().await; } diff --git a/src/cmd/src/bin/query_perf_fixture/case.rs b/src/cmd/src/bin/query_perf_fixture/case.rs new file mode 100644 index 00000000000..7dce735d7b9 --- /dev/null +++ b/src/cmd/src/bin/query_perf_fixture/case.rs @@ -0,0 +1,470 @@ +// 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::num::NonZeroUsize; + +use clap::ValueEnum; +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Deserialize, Serialize)] +pub(super) struct CaseFile { + pub(super) scenario: Scenario, +} + +#[derive(Debug, Deserialize, Serialize)] +#[serde(tag = "kind")] +pub(super) enum Scenario { + #[serde(rename = "direct_readable_sst")] + DirectReadableSst(DirectReadableSstScenario), + #[serde(rename = "prom_remote_write_then_query")] + PromRemoteWriteThenQuery(PromRemoteWriteThenQueryScenario), +} + +#[derive(Debug, Deserialize, Serialize)] +pub(super) struct DirectReadableSstScenario { + #[serde(default)] + pub(super) seed: Option, + pub(super) tables: Vec, + pub(super) layout: LayoutConfig, +} + +#[derive(Debug, Deserialize, Serialize)] +pub(super) struct PromRemoteWriteThenQueryScenario { + #[serde(default)] + pub(super) queries: Vec, + pub(super) remote_write: PromRemoteWritePlan, +} + +#[derive(Debug, Deserialize, Serialize)] +pub(super) struct PromRemoteWritePlan { + #[serde(default = "default_database")] + pub(super) database: String, + #[serde(alias = "metric_name")] + pub(super) metric: String, + #[serde(default = "default_physical_table")] + pub(super) physical_table: String, + #[serde(default = "default_series_count")] + pub(super) series_count: u64, + #[serde(default = "default_samples_per_series")] + pub(super) samples_per_series: u64, + #[serde(default = "default_start_unix_millis")] + pub(super) start_unix_millis: i64, + #[serde(default = "default_step_millis")] + pub(super) step_millis: i64, + #[serde(default = "default_chunk_series_count", alias = "batch_size")] + pub(super) chunk_series_count: u64, + #[serde(default = "default_timeout_seconds")] + pub(super) timeout_seconds: u64, + #[serde(default)] + pub(super) sample_chunk_size: Option, + #[serde(default = "default_flush_every_sample_chunks")] + pub(super) flush_every_sample_chunks: u64, + #[serde(default = "default_visibility_timeout_seconds")] + pub(super) visibility_timeout_seconds: u64, + #[serde(default)] + pub(super) prom_store: PromStoreConfig, + #[serde(default)] + pub(super) value: ValueConfig, + #[serde(default)] + pub(super) storage: Option, + #[serde(default)] + pub(super) read_bench: Option, +} + +pub(super) fn default_database() -> String { + "public".to_string() +} +pub(super) fn default_physical_table() -> String { + "greptime_physical_table".to_string() +} +pub(super) fn default_series_count() -> u64 { + 8 +} +pub(super) fn default_samples_per_series() -> u64 { + 30 +} +pub(super) fn default_start_unix_millis() -> i64 { + 1_704_067_200_000 +} +pub(super) fn default_step_millis() -> i64 { + 15_000 +} +pub(super) fn default_chunk_series_count() -> u64 { + 8 +} +pub(super) fn default_timeout_seconds() -> u64 { + 60 +} +pub(super) fn default_flush_every_sample_chunks() -> u64 { + 1 +} +pub(super) fn default_visibility_timeout_seconds() -> u64 { + 30 +} + +#[derive(Debug, Deserialize, Serialize)] +pub(super) struct PromStoreConfig { + #[serde(default = "default_pending_rows_flush_interval")] + pub(super) pending_rows_flush_interval: String, + #[serde(default = "default_max_batch_rows")] + pub(super) max_batch_rows: u64, + #[serde(default = "default_max_concurrent_flushes")] + pub(super) max_concurrent_flushes: u64, + #[serde(default = "default_worker_channel_capacity")] + pub(super) worker_channel_capacity: u64, + #[serde(default = "default_max_inflight_requests")] + pub(super) max_inflight_requests: u64, +} + +impl Default for PromStoreConfig { + fn default() -> Self { + Self { + pending_rows_flush_interval: default_pending_rows_flush_interval(), + max_batch_rows: default_max_batch_rows(), + max_concurrent_flushes: default_max_concurrent_flushes(), + worker_channel_capacity: default_worker_channel_capacity(), + max_inflight_requests: default_max_inflight_requests(), + } + } +} + +pub(super) fn default_pending_rows_flush_interval() -> String { + "1s".to_string() +} +pub(super) fn default_max_batch_rows() -> u64 { + 100000 +} +pub(super) fn default_max_concurrent_flushes() -> u64 { + 256 +} +pub(super) fn default_worker_channel_capacity() -> u64 { + 65526 +} +pub(super) fn default_max_inflight_requests() -> u64 { + 3000 +} + +#[derive(Debug, Clone, Copy, Deserialize, Serialize, ValueEnum, Eq, PartialEq)] +#[serde(rename_all = "snake_case")] +#[value(rename_all = "snake_case")] +pub(super) enum ValuePattern { + Linear, + Constant, + Modulo, + Unique, + SeededRandom, + RunLength, + QuantizedSignal, + SignalWithSporadicStalls, + MixedSignalRepeated, +} + +impl std::fmt::Display for ValuePattern { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "{}", + serde_json::to_value(self).unwrap().as_str().unwrap() + ) + } +} + +#[derive(Debug, Deserialize, Serialize)] +pub(super) struct ValueConfig { + #[serde(default = "default_value_pattern")] + pub(super) pattern: ValuePattern, + #[serde(default)] + pub(super) base: f64, + #[serde(default = "default_value_step")] + pub(super) step: f64, + #[serde(default = "default_value_cardinality")] + pub(super) cardinality: u64, + #[serde(default)] + pub(super) seed: u64, + #[serde(default = "default_value_run_length")] + pub(super) run_length: u64, + #[serde(default = "default_value_stall_every")] + pub(super) stall_every: u64, + #[serde(default = "default_value_stall_length")] + pub(super) stall_length: u64, + #[serde(default = "default_value_mixed_every")] + pub(super) mixed_every: u64, +} +impl Default for ValueConfig { + fn default() -> Self { + Self { + pattern: default_value_pattern(), + base: 0.0, + step: default_value_step(), + cardinality: default_value_cardinality(), + seed: 0, + run_length: default_value_run_length(), + stall_every: default_value_stall_every(), + stall_length: default_value_stall_length(), + mixed_every: default_value_mixed_every(), + } + } +} +pub(super) fn default_value_pattern() -> ValuePattern { + ValuePattern::Linear +} +pub(super) fn default_value_step() -> f64 { + 0.125 +} +pub(super) fn default_value_cardinality() -> u64 { + 97 +} +pub(super) fn default_value_run_length() -> u64 { + 8 +} +pub(super) fn default_value_stall_every() -> u64 { + 100 +} +pub(super) fn default_value_stall_length() -> u64 { + 16 +} +pub(super) fn default_value_mixed_every() -> u64 { + 5 +} + +#[derive(Debug, Deserialize, Serialize)] +pub(super) struct StorageConfig { + #[serde(default = "default_true")] + pub(super) inspect: bool, + #[serde(default = "default_storage_column")] + pub(super) column: String, + #[serde(default)] + pub(super) root_suffix: Option, + #[serde(default)] + pub(super) include_metadata_files: bool, + #[serde(default = "default_min_files")] + pub(super) min_files: u64, + #[serde(default = "default_min_files")] + pub(super) min_files_with_column: u64, + #[serde(default)] + pub(super) require_encodings: Vec, + #[serde(default)] + pub(super) forbid_encodings: Vec, + #[serde(default)] + pub(super) max_total_file_size_bytes: Option, + #[serde(default)] + pub(super) max_column_compressed_size_bytes: Option, + #[serde(default)] + pub(super) max_column_uncompressed_size_bytes: Option, + #[serde(default)] + pub(super) max_candidate_total_file_size_regression_pct: Option, + #[serde(default)] + pub(super) max_candidate_column_compressed_size_regression_pct: Option, + #[serde(default)] + pub(super) max_candidate_column_uncompressed_size_regression_pct: Option, + #[serde(skip_deserializing, default)] + pub(super) planned_thresholds: Vec, +} + +#[derive(Debug, Serialize)] +pub(super) struct StorageThresholdPlan { + pub(super) threshold: String, + pub(super) status: String, + pub(super) value: serde_json::Value, +} + +impl StorageConfig { + pub(super) fn populate_planned_thresholds(&mut self) { + let mut planned = Vec::new(); + planned.push(StorageThresholdPlan::new("min_files", self.min_files)); + planned.push(StorageThresholdPlan::new( + "min_files_with_column", + self.min_files_with_column, + )); + if !self.require_encodings.is_empty() { + planned.push(StorageThresholdPlan::new( + "require_encodings", + &self.require_encodings, + )); + } + if !self.forbid_encodings.is_empty() { + planned.push(StorageThresholdPlan::new( + "forbid_encodings", + &self.forbid_encodings, + )); + } + macro_rules! push_optional { + ($name:literal, $value:expr) => { + if let Some(value) = $value { + planned.push(StorageThresholdPlan::new($name, value)); + } + }; + } + push_optional!("max_total_file_size_bytes", self.max_total_file_size_bytes); + push_optional!( + "max_column_compressed_size_bytes", + self.max_column_compressed_size_bytes + ); + push_optional!( + "max_column_uncompressed_size_bytes", + self.max_column_uncompressed_size_bytes + ); + push_optional!( + "max_candidate_total_file_size_regression_pct", + self.max_candidate_total_file_size_regression_pct + ); + push_optional!( + "max_candidate_column_compressed_size_regression_pct", + self.max_candidate_column_compressed_size_regression_pct + ); + push_optional!( + "max_candidate_column_uncompressed_size_regression_pct", + self.max_candidate_column_uncompressed_size_regression_pct + ); + self.planned_thresholds = planned; + } +} + +impl StorageThresholdPlan { + fn new(threshold: &str, value: T) -> Self { + Self { + threshold: threshold.to_string(), + status: "planned".to_string(), + value: serde_json::to_value(value).expect("storage threshold value must serialize"), + } + } +} +pub(super) fn default_true() -> bool { + true +} +pub(super) fn default_storage_column() -> String { + "greptime_value".to_string() +} +pub(super) fn default_min_files() -> u64 { + 1 +} + +#[derive(Debug, Deserialize, Serialize, Clone)] +pub(super) struct ReadBenchConfig { + #[serde(default = "default_true")] + pub(super) enabled: bool, + #[serde(default = "default_true")] + pub(super) parquetbench: bool, + #[serde(default = "default_true")] + pub(super) scanbench: bool, + #[serde(default = "default_iterations")] + pub(super) iterations: u64, + #[serde(default)] + pub(super) projection: Vec, + #[serde(default = "default_parquet_reader")] + pub(super) parquet_reader: String, + #[serde(default = "default_scan_scanner")] + pub(super) scan_scanner: String, + #[serde(default = "default_parallelism")] + pub(super) parallelism: u64, + #[serde(default)] + pub(super) max_files: Option, + #[serde(flatten)] + pub(super) thresholds: HashMap, +} + +impl Default for ReadBenchConfig { + fn default() -> Self { + Self { + enabled: true, + parquetbench: true, + scanbench: true, + iterations: default_iterations(), + projection: vec![], + parquet_reader: default_parquet_reader(), + scan_scanner: default_scan_scanner(), + parallelism: default_parallelism(), + max_files: None, + thresholds: HashMap::new(), + } + } +} + +pub(super) fn default_iterations() -> u64 { + 7 +} +pub(super) fn default_parquet_reader() -> String { + "direct".to_string() +} +pub(super) fn default_scan_scanner() -> String { + "seq".to_string() +} +pub(super) fn default_parallelism() -> u64 { + 1 +} + +impl Scenario { + pub(super) fn kind(&self) -> &'static str { + match self { + Scenario::DirectReadableSst(_) => "direct_readable_sst", + Scenario::PromRemoteWriteThenQuery(_) => "prom_remote_write_then_query", + } + } + + pub(super) fn direct_readable_sst(&self) -> &DirectReadableSstScenario { + match self { + Scenario::DirectReadableSst(scenario) => scenario, + _ => panic!("scenario is not direct_readable_sst"), + } + } +} + +#[derive(Debug, Deserialize, Serialize)] +pub(super) struct TableConfig { + pub(super) database: String, + pub(super) name: String, + pub(super) engine: String, + #[serde(default)] + pub(super) append_mode: Option, + #[serde(default)] + pub(super) sst_format: Option, + pub(super) primary_key: Vec, + pub(super) time_index: String, + pub(super) columns: Vec, +} + +#[derive(Debug, Deserialize, Serialize)] +pub(super) struct ColumnConfig { + pub(super) name: String, + #[serde(rename = "type")] + pub(super) ty: String, + pub(super) semantic: String, + pub(super) distribution: Option, +} + +#[derive(Debug, Deserialize, Serialize)] +#[serde(tag = "kind")] +pub(super) enum Distribution { + #[serde(rename = "cardinality")] + Cardinality { + values: NonZeroUsize, + prefix: String, + }, + #[serde(rename = "deterministic_wave")] + DeterministicWave { min: f64, max: f64 }, +} + +#[derive(Debug, Deserialize, Serialize)] +pub(super) struct LayoutConfig { + pub(super) regions: usize, + pub(super) sst_count: usize, + pub(super) rows_per_sst: usize, + pub(super) row_group_size: usize, + pub(super) series_count: NonZeroUsize, + pub(super) start_unix_nanos: i64, + pub(super) step_nanos: i64, + pub(super) time_range_layout: String, + pub(super) series_layout: String, +} diff --git a/src/cmd/src/bin/query_perf_fixture/direct_sst.rs b/src/cmd/src/bin/query_perf_fixture/direct_sst.rs new file mode 100644 index 00000000000..61dedd4cd4a --- /dev/null +++ b/src/cmd/src/bin/query_perf_fixture/direct_sst.rs @@ -0,0 +1,528 @@ +// 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::fs::{self}; +use std::io::{BufWriter, Write}; +use std::num::NonZeroU64; +use std::path::Path; +use std::sync::Arc; + +use api::v1::{OpType, SemanticType}; +use async_trait::async_trait; +use datatypes::arrow::array::{ + ArrayRef, BinaryDictionaryBuilder, Float64Array, RecordBatch, StringDictionaryBuilder, + TimestampMillisecondArray, TimestampNanosecondArray, UInt8Array, UInt64Array, +}; +use datatypes::arrow::datatypes::UInt32Type; +use datatypes::prelude::ConcreteDataType; +use datatypes::schema::{ColumnSchema, SkippingIndexOptions}; +use mito_codec::row_converter::{DensePrimaryKeyCodec, PrimaryKeyCodecExt, SortField}; +use mito2::access_layer::{FilePathProvider, Metrics, WriteType}; +use mito2::config::IndexConfig; +use mito2::manifest::action::{RegionCheckpoint, RegionManifest, RemovedFilesRecord}; +use mito2::read::FlatSource; +use mito2::sst::file::{FileMeta, RegionFileId}; +use mito2::sst::index::{Indexer, IndexerBuilder}; +use mito2::sst::parquet::writer::ParquetWriter; +use mito2::sst::parquet::{SstInfo, WriteOptions}; +use mito2::sst::{ + DEFAULT_WRITE_BUFFER_SIZE, FlatSchemaOptions, FormatType, to_flat_sst_arrow_schema, +}; +use object_store::ObjectStore; +use object_store::services::Fs as FsBuilder; +use store_api::codec::PrimaryKeyEncoding; +use store_api::metadata::{ColumnMetadata, RegionMetadata, RegionMetadataRef}; +use store_api::path_utils::region_name; +use store_api::region_request::PathType; +use store_api::storage::{FileId, RegionId}; + +use super::util::case_name_from_path; +use super::*; + +struct NoopIndexBuilder; + +#[async_trait] +impl IndexerBuilder for NoopIndexBuilder { + async fn build(&self, _file_id: FileId, _index_version: u64) -> Indexer { + Indexer::default() + } +} + +#[derive(Clone)] +struct FixedPathProvider { + table_dir: String, +} + +impl FilePathProvider for FixedPathProvider { + fn build_index_file_path(&self, file_id: RegionFileId) -> String { + mito2::sst::location::index_file_path_legacy(&self.table_dir, file_id, PathType::Bare) + } + fn build_index_file_path_with_version( + &self, + index_id: mito2::sst::file::RegionIndexId, + ) -> String { + mito2::sst::location::index_file_path(&self.table_dir, index_id, PathType::Bare) + } + fn build_sst_file_path(&self, file_id: RegionFileId) -> String { + mito2::sst::location::sst_file_path(&self.table_dir, file_id, PathType::Bare) + } +} + +fn semantic_type(s: &str) -> SemanticType { + match s.to_ascii_lowercase().as_str() { + "tag" => SemanticType::Tag, + "field" => SemanticType::Field, + "timestamp" => SemanticType::Timestamp, + other => panic!("unsupported semantic {other}"), + } +} + +fn concrete_type(ty: &str) -> ConcreteDataType { + match ty.to_ascii_uppercase().as_str() { + "STRING" => ConcreteDataType::string_datatype(), + "DOUBLE" => ConcreteDataType::float64_datatype(), + "UINT64" => ConcreteDataType::uint64_datatype(), + "TIMESTAMP(9)" => ConcreteDataType::timestamp_nanosecond_datatype(), + "TIMESTAMP(3)" => ConcreteDataType::timestamp_millisecond_datatype(), + other => panic!("unsupported column type {other}"), + } +} + +fn build_region_metadata(table: &TableConfig, region_id: RegionId) -> RegionMetadata { + let mut builder = store_api::metadata::RegionMetadataBuilder::new(region_id); + for (idx, col) in table.columns.iter().enumerate() { + let mut schema = ColumnSchema::new( + &col.name, + concrete_type(&col.ty), + col.semantic != "timestamp", + ); + if col.semantic == "tag" { + schema = schema.with_inverted_index(true); + } + if col.semantic == "tag" || col.semantic == "field" { + schema = schema + .with_skipping_options(SkippingIndexOptions { + granularity: 1, + ..Default::default() + }) + .expect("valid skipping index options for query perf fixture"); + } + if col.name == table.time_index { + schema = schema.with_time_index(true); + } + builder.push_column_metadata(ColumnMetadata { + column_schema: schema, + semantic_type: semantic_type(&col.semantic), + column_id: idx as u32, + }); + } + let pk_ids = table + .primary_key + .iter() + .map(|name| { + table + .columns + .iter() + .position(|c| c.name == *name) + .unwrap_or_else(|| panic!("primary key column {name} not found")) as u32 + }) + .collect(); + builder.primary_key(pk_ids); + builder.primary_key_encoding(PrimaryKeyEncoding::Dense); + builder + .build() + .expect("region metadata should be valid for fixture table config") +} + +fn encode_dense_primary_key(table: &TableConfig, values: &HashMap) -> Vec { + let fields = table + .primary_key + .iter() + .enumerate() + .map(|(idx, _)| { + ( + idx as u32, + SortField::new(ConcreteDataType::string_datatype()), + ) + }) + .collect(); + let converter = DensePrimaryKeyCodec::with_fields(fields); + converter + .encode( + table + .primary_key + .iter() + .map(|name| datatypes::value::ValueRef::String(values[name].as_str())), + ) + .expect("dense primary key encoding should match fixture primary key fields") +} + +fn tag_value(col: &ColumnConfig, series: usize) -> String { + match col + .distribution + .as_ref() + .expect("tag distribution is required") + { + Distribution::Cardinality { values, prefix } => { + format!("{}{}", prefix, series % values.get()) + } + _ => panic!("tag column {} requires cardinality distribution", col.name), + } +} + +fn series_for_row(layout: &LayoutConfig, sst_idx: usize, base_row: usize, row: usize) -> usize { + match layout.series_layout.as_str() { + "round_robin" | "timestamp_major" => (base_row + row) % layout.series_count.get(), + "per_sst" => sst_idx % layout.series_count.get(), + other => panic!("unsupported series_layout {other}"), + } +} + +fn timestamp_for_row(layout: &LayoutConfig, base_row: usize, row: usize) -> i64 { + let logical_row = base_row + row; + let timestamp_step = if layout.series_layout == "timestamp_major" { + logical_row / layout.series_count.get() + } else { + logical_row + }; + + layout.start_unix_nanos + (timestamp_step as i64 * layout.step_nanos) +} + +fn wave_value(min: f64, max: f64, row: usize) -> f64 { + let span = max - min; + let phase = (row % 1024) as f64 / 1023.0; + min + span * (0.5 - 0.5 * (std::f64::consts::TAU * phase).cos()) +} + +fn generate_record_batch( + table: &TableConfig, + metadata: &RegionMetadataRef, + layout: &LayoutConfig, + sst_idx: usize, + sequence: u64, +) -> RecordBatch { + let flat_schema = to_flat_sst_arrow_schema(metadata, &FlatSchemaOptions::default()); + let rows = layout.rows_per_sst; + let mut columns = Vec::new(); + let base_row = sst_idx * rows; + + for col in &table.columns { + match (col.semantic.as_str(), col.ty.to_ascii_uppercase().as_str()) { + ("tag", "STRING") => { + let mut b = StringDictionaryBuilder::::new(); + for row in 0..rows { + let series = series_for_row(layout, sst_idx, base_row, row); + let value = tag_value(col, series); + b.append_value(value); + } + columns.push(Arc::new(b.finish()) as ArrayRef); + } + ("field", "DOUBLE") => { + let (min, max) = match col.distribution.as_ref() { + Some(Distribution::DeterministicWave { min, max }) => (*min, *max), + _ => (0.0, 1.0), + }; + columns.push(Arc::new(Float64Array::from( + (0..rows) + .map(|r| wave_value(min, max, base_row + r)) + .collect::>(), + )) as ArrayRef); + } + ("field", "UINT64") => columns.push(Arc::new(UInt64Array::from( + (0..rows).map(|r| (base_row + r) as u64).collect::>(), + )) as ArrayRef), + ("timestamp", "TIMESTAMP(9)") => columns.push(Arc::new(TimestampNanosecondArray::from( + (0..rows) + .map(|r| timestamp_for_row(layout, base_row, r)) + .collect::>(), + )) as ArrayRef), + ("timestamp", "TIMESTAMP(3)") => { + columns.push(Arc::new(TimestampMillisecondArray::from( + (0..rows) + .map(|r| timestamp_for_row(layout, base_row, r) / 1_000_000) + .collect::>(), + )) as ArrayRef) + } + other => panic!("unsupported column combination {other:?}"), + } + } + + let mut pk_builder = BinaryDictionaryBuilder::::new(); + for row in 0..rows { + let series = series_for_row(layout, sst_idx, base_row, row); + let values = table + .columns + .iter() + .filter(|c| c.semantic == "tag") + .map(|c| (c.name.clone(), tag_value(c, series))) + .collect(); + pk_builder + .append(encode_dense_primary_key(table, &values)) + .expect("append generated dense primary key to Arrow dictionary builder"); + } + columns.push(Arc::new(pk_builder.finish())); + columns.push(Arc::new(UInt64Array::from_value(sequence, rows))); + columns.push(Arc::new(UInt8Array::from_value(OpType::Put as u8, rows))); + RecordBatch::try_new(flat_schema, columns) + .expect("generated fixture columns should match flat SST Arrow schema") +} + +fn file_meta_from_sst_info( + info: &SstInfo, + region_id: RegionId, + file_id: FileId, + sequence: u64, +) -> FileMeta { + FileMeta { + region_id, + file_id, + time_range: info.time_range, + level: 0, + file_size: info.file_size, + max_row_group_uncompressed_size: info.max_row_group_uncompressed_size, + available_indexes: Default::default(), + indexes: Default::default(), + index_file_size: 0, + index_version: 0, + num_rows: info.num_rows as u64, + num_row_groups: info.num_row_groups, + sequence: NonZeroU64::new(sequence), + partition_expr: None, + num_series: info.num_series, + primary_key_min: None, + primary_key_max: None, + } +} + +fn deterministic_file_id(seed: u64, index: usize) -> FileId { + let seed_hi = ((seed >> 16) & 0xffff) as u16; + let seed_lo = (seed & 0xffff) as u16; + FileId::parse_str(&format!( + "00000000-0000-{seed_hi:04x}-{seed_lo:04x}-{index:012x}" + )) + .expect("deterministic file id must be a valid UUID") +} + +fn remap_sst_file( + object_store_dir: &Path, + table_dir: &str, + region_id: RegionId, + old_file_id: FileId, + new_file_id: FileId, +) { + if old_file_id == new_file_id { + return; + } + let old_path = object_store_dir.join(mito2::sst::location::sst_file_path( + table_dir, + RegionFileId::new(region_id, old_file_id), + PathType::Bare, + )); + let new_path = object_store_dir.join(mito2::sst::location::sst_file_path( + table_dir, + RegionFileId::new(region_id, new_file_id), + PathType::Bare, + )); + if let Some(parent) = new_path.parent() { + fs::create_dir_all(parent).expect("failed to create remapped SST parent directory"); + } + fs::rename(&old_path, &new_path).unwrap_or_else(|err| { + panic!( + "failed to remap SST file {} -> {}: {err}", + old_path.display(), + new_path.display() + ) + }); +} +pub(super) async fn run_direct_sst(args: DirectArgs) { + let case_path = args.case.as_ref().expect("--case is required"); + let out_dir = args.out_dir.as_ref().expect("--out-dir is required"); + let case_text = fs::read_to_string(case_path).expect("failed to read query perf case TOML"); + let case: CaseFile = toml::from_str(&case_text).expect("failed to parse query perf case TOML"); + let case_name = case_name_from_path(case_path); + let scenario = case.scenario.direct_readable_sst(); + if scenario.tables.is_empty() || scenario.layout.regions != 1 { + panic!("fixture generator supports one or more tables and exactly one region per table"); + } + if scenario.layout.time_range_layout != "non_overlapping_per_sst" { + panic!("MVP supports time_range_layout=non_overlapping_per_sst"); + } + if scenario.layout.series_layout == "timestamp_major" + && !scenario + .layout + .rows_per_sst + .is_multiple_of(scenario.layout.series_count.get()) + { + panic!( + "series_layout=timestamp_major requires rows_per_sst to be divisible by series_count" + ); + } + if scenario.layout.sst_count > 1000 && !args.allow_large { + panic!("sst_count exceeds 1000; pass --allow-large"); + } + let seed = scenario.seed.unwrap_or(0); + let table_index = match args.table.as_deref() { + Some(name) => scenario + .tables + .iter() + .position(|table| table.name == name) + .unwrap_or_else(|| panic!("--table {name} was not found in case {case_name}")), + None if scenario.tables.len() == 1 => 0, + None => panic!( + "case {} contains {} tables; pass --table to choose one", + case_name, + scenario.tables.len() + ), + }; + let table = &scenario.tables[table_index]; + if table.engine != "mito" { + panic!("MVP supports engine=mito"); + } + let region_id = RegionId::from(args.region_id); + let table_dir = args + .table_dir + .unwrap_or_else(|| format!("data/{}/{}/", table.database, table.name)); + if Path::new(&table_dir).is_absolute() { + panic!("--table-dir must be relative"); + } + if table_dir.trim_end_matches('/').rsplit('/').next() + == Some(region_name(region_id.table_id(), region_id.region_sequence()).as_str()) + { + panic!("--table-dir must be table-level directory"); + } + let region_dir = + mito2::sst::location::region_dir_from_table_dir(&table_dir, region_id, PathType::Bare); + let format = match table.sst_format.as_deref().unwrap_or("flat") { + "flat" => FormatType::Flat, + "primary_key" | "primary-key" => FormatType::PrimaryKey, + other => panic!("unsupported sst_format {other}"), + }; + println!( + "query_perf_fixture case={} scenario={} table={}.{} ssts={} rows_per_sst={} format={format:?}", + case_name, + case.scenario.kind(), + table.database, + table.name, + scenario.layout.sst_count, + scenario.layout.rows_per_sst + ); + println!( + "out_dir={} table_dir={} region_dir={}", + out_dir.display(), + table_dir, + region_dir + ); + if args.dry_run { + return; + } + + let obj_store_dir = out_dir.join("object-store"); + let manifest_dir = out_dir.join("manifest"); + fs::create_dir_all(&obj_store_dir).expect("failed to create fixture object-store directory"); + fs::create_dir_all(&manifest_dir).expect("failed to create fixture manifest directory"); + let ostorage = ObjectStore::new(FsBuilder::default().root(&obj_store_dir.to_string_lossy())) + .expect("failed to create filesystem object store for fixture output") + .finish(); + let metadata: RegionMetadataRef = Arc::new(build_region_metadata(table, region_id)); + let mut files = HashMap::with_capacity(scenario.layout.sst_count); + let mut next_file_index = 1; + for i in 0..scenario.layout.sst_count { + let sequence = 1000 + i as u64; + let batch = generate_record_batch(table, &metadata, &scenario.layout, i, sequence); + let source = + FlatSource::new_iter(batch.schema(), Box::new(vec![batch].into_iter().map(Ok))); + let mut metrics = Metrics::new(WriteType::Flush); + let mut writer = ParquetWriter::new_with_object_store( + ostorage.clone(), + metadata.clone(), + IndexConfig::default(), + NoopIndexBuilder, + FixedPathProvider { + table_dir: table_dir.clone(), + }, + &mut metrics, + ) + .await; + let opts = WriteOptions { + write_buffer_size: DEFAULT_WRITE_BUFFER_SIZE, + row_group_size: scenario.layout.row_group_size, + max_file_size: None, + }; + let infos = match format { + FormatType::Flat => writer.write_all_flat(source, Some(sequence), &opts).await, + FormatType::PrimaryKey => { + writer + .write_all_flat_as_primary_key(source, Some(sequence), &opts) + .await + } + } + .expect("failed to write fixture SST through Mito parquet writer"); + for info in infos { + let file_id = deterministic_file_id( + seed.wrapping_add((table_index as u64) << 16), + next_file_index, + ); + next_file_index += 1; + remap_sst_file(&obj_store_dir, &table_dir, region_id, info.file_id, file_id); + files.insert( + file_id, + file_meta_from_sst_info(&info, region_id, file_id, sequence), + ); + } + } + let overall_seq = files + .values() + .filter_map(|m| m.sequence) + .map(|s| s.get()) + .max() + .unwrap_or(0); + let manifest = RegionManifest { + metadata, + files, + removed_files: RemovedFilesRecord::default(), + flushed_entry_id: overall_seq, + flushed_sequence: overall_seq, + committed_sequence: Some(overall_seq), + manifest_version: args.checkpoint_version, + truncated_entry_id: None, + compaction_time_window: None, + sst_format: format, + append_mode: table.append_mode, + }; + let checkpoint = RegionCheckpoint { + last_version: args.checkpoint_version, + compacted_actions: manifest.files.len(), + checkpoint: Some(manifest.clone()), + }; + let checkpoint_bytes = checkpoint + .encode() + .expect("failed to encode fixture region checkpoint"); + let checkpoint_path = manifest_dir.join(format!("{:020}.checkpoint", args.checkpoint_version)); + fs::write(&checkpoint_path, &checkpoint_bytes).expect("failed to write checkpoint file"); + fs::write(manifest_dir.join("_last_checkpoint"), serde_json::to_vec_pretty(&serde_json::json!({ "size": checkpoint_bytes.len(), "version": args.checkpoint_version, "checksum": null, "extend_metadata": {} })).expect("failed to serialize _last_checkpoint metadata")).expect("failed to write _last_checkpoint file"); + let files_jsonl_path = out_dir.join("files.jsonl"); + let mut jsonl = BufWriter::new( + fs::File::create(&files_jsonl_path).expect("failed to create fixture files.jsonl"), + ); + let mut entries: Vec<_> = manifest.files.iter().collect(); + entries.sort_by_key(|(id, _)| id.to_string()); + for (file_id, meta) in entries { + writeln!(jsonl, "{}", serde_json::to_string(&serde_json::json!({ "file_id": file_id.to_string(), "region_id": meta.region_id.as_u64(), "object_path": mito2::sst::location::sst_file_path(&table_dir, RegionFileId::new(meta.region_id, *file_id), PathType::Bare), "time_range_start": meta.time_range.0.value(), "time_range_end": meta.time_range.1.value(), "num_rows": meta.num_rows, "num_row_groups": meta.num_row_groups, "file_size": meta.file_size, "num_series": meta.num_series, "sequence": meta.sequence.map(|s| s.get()) })).expect("failed to serialize fixture SST metadata JSONL entry")).expect("failed to write fixture SST metadata JSONL entry"); + } + jsonl.flush().expect("failed to flush fixture files.jsonl"); + fs::write(out_dir.join("summary.json"), serde_json::to_vec_pretty(&serde_json::json!({ "case": case_name, "seed": seed, "table_index": table_index, "table": table.name, "database": table.database, "region_id": region_id.as_u64(), "table_dir": table_dir, "region_dir": region_dir, "sst_format": format!("{format:?}"), "sst_count": scenario.layout.sst_count, "rows_per_sst": scenario.layout.rows_per_sst, "row_group_size": scenario.layout.row_group_size, "total_rows": scenario.layout.sst_count * scenario.layout.rows_per_sst, "checkpoint_path": checkpoint_path, "files_jsonl_path": files_jsonl_path, "readback_validated": false, "metadata_source": "synthetic" })).expect("failed to serialize fixture summary")).expect("failed to write fixture summary.json"); + println!("Done. wrote {} SST file entries", manifest.files.len()); +} diff --git a/src/cmd/src/bin/query_perf_fixture/inspect_footer.rs b/src/cmd/src/bin/query_perf_fixture/inspect_footer.rs new file mode 100644 index 00000000000..a52005c4e30 --- /dev/null +++ b/src/cmd/src/bin/query_perf_fixture/inspect_footer.rs @@ -0,0 +1,181 @@ +// 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::BTreeSet; +use std::fs::{self, File}; +use std::path::{Path, PathBuf}; + +use clap::Args as ClapArgs; +use parquet::file::reader::{FileReader, SerializedFileReader}; +use serde::Serialize; + +#[derive(Debug, ClapArgs)] +pub(super) struct InspectFooterArgs { + #[arg(long)] + root: PathBuf, + #[arg(long, default_value = "greptime_value")] + column: String, + #[arg(long)] + include_metadata_files: bool, +} + +#[derive(Debug, Serialize)] +struct FooterReport { + root: String, + summary: FooterSummary, + files: Vec, +} +#[derive(Debug, Default, Serialize)] +struct FooterSummary { + column: String, + file_count: usize, + files_with_column: usize, + total_file_size: u64, + total_rows: i64, + column_compressed_size: i64, + column_uncompressed_size: i64, + unique_encodings: Vec, +} +#[derive(Debug, Serialize)] +struct FooterFileReport { + path: String, + relative_path: String, + file_size: u64, + num_rows: i64, + num_row_groups: usize, + columns: Vec, +} +#[derive(Debug, Serialize)] +struct FooterColumnChunkReport { + row_group_index: usize, + column_index: usize, + column_path: String, + encodings: Vec, + compression: String, + compressed_size: i64, + uncompressed_size: i64, + num_values: i64, +} + +pub(super) fn run_inspect_footer( + args: InspectFooterArgs, +) -> Result<(), Box> { + let files = collect_parquet_files(&args.root, args.include_metadata_files)?; + let reports = files + .iter() + .map(|p| inspect_file(&args.root, p, &args.column)) + .collect::, _>>()?; + println!( + "{}", + serde_json::to_string_pretty(&FooterReport { + root: args.root.display().to_string(), + summary: summarize_footer(&args.column, &reports), + files: reports + })? + ); + Ok(()) +} +fn collect_parquet_files( + root: &Path, + include_metadata_files: bool, +) -> Result, Box> { + let mut files = Vec::new(); + collect_parquet_files_inner(root, root, include_metadata_files, &mut files)?; + files.sort(); + Ok(files) +} +fn collect_parquet_files_inner( + root: &Path, + path: &Path, + include_metadata_files: bool, + files: &mut Vec, +) -> Result<(), Box> { + if path.is_dir() { + for entry in fs::read_dir(path)? { + collect_parquet_files_inner(root, &entry?.path(), include_metadata_files, files)?; + } + } else if path.extension().is_some_and(|ext| ext == "parquet") + && (include_metadata_files + || !path + .strip_prefix(root) + .unwrap_or(path) + .components() + .any(|c| c.as_os_str() == "metadata")) + { + files.push(path.to_path_buf()); + } + Ok(()) +} +fn inspect_file( + root: &Path, + path: &Path, + column: &str, +) -> Result> { + let file_size = fs::metadata(path)?.len(); + let reader = SerializedFileReader::new(File::open(path)?)?; + let metadata = reader.metadata().file_metadata(); + let row_groups = reader.metadata().row_groups(); + let mut columns = Vec::new(); + for (row_group_index, rg) in row_groups.iter().enumerate() { + for (column_index, chunk) in rg.columns().iter().enumerate() { + let column_path = chunk.column_path().string(); + if column_path == column { + columns.push(FooterColumnChunkReport { + row_group_index, + column_index, + column_path, + encodings: chunk.encodings().map(|e| format!("{e:?}")).collect(), + compression: format!("{:?}", chunk.compression()), + compressed_size: chunk.compressed_size(), + uncompressed_size: chunk.uncompressed_size(), + num_values: chunk.num_values(), + }); + } + } + } + Ok(FooterFileReport { + path: path.display().to_string(), + relative_path: path + .strip_prefix(root) + .unwrap_or(path) + .display() + .to_string(), + file_size, + num_rows: metadata.num_rows(), + num_row_groups: row_groups.len(), + columns, + }) +} +fn summarize_footer(column: &str, files: &[FooterFileReport]) -> FooterSummary { + let mut encodings = BTreeSet::new(); + let mut s = FooterSummary { + column: column.to_string(), + file_count: files.len(), + ..Default::default() + }; + for file in files { + s.total_file_size += file.file_size; + s.total_rows += file.num_rows; + if !file.columns.is_empty() { + s.files_with_column += 1; + } + for c in &file.columns { + s.column_compressed_size += c.compressed_size; + s.column_uncompressed_size += c.uncompressed_size; + encodings.extend(c.encodings.iter().cloned()); + } + } + s.unique_encodings = encodings.into_iter().collect(); + s +} diff --git a/src/cmd/src/bin/query_perf_fixture/mod.rs b/src/cmd/src/bin/query_perf_fixture/mod.rs new file mode 100644 index 00000000000..00cef327c78 --- /dev/null +++ b/src/cmd/src/bin/query_perf_fixture/mod.rs @@ -0,0 +1,133 @@ +// 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. + +mod case; +mod direct_sst; +mod inspect_footer; +mod prom_remote_write; +mod util; + +use std::fs; +use std::path::PathBuf; + +use case::*; +use clap::{Args as ClapArgs, Parser, Subcommand}; +use direct_sst::run_direct_sst; +use inspect_footer::{InspectFooterArgs, run_inspect_footer}; +use prom_remote_write::{PromRemoteWriteArgs, run_prom_remote_write}; +use serde_json::json; + +#[derive(Parser, Debug)] +#[command(name = "query_perf_fixture")] +#[command(about = "Query performance fixture helper")] +struct Args { + #[command(subcommand)] + command: Option, + + #[command(flatten)] + legacy: DirectArgs, +} + +#[derive(Subcommand, Debug)] +enum Command { + DirectSst(DirectArgs), + PromRemoteWrite(PromRemoteWriteArgs), + InspectFooter(InspectFooterArgs), + Plan(PlanArgs), +} + +#[derive(ClapArgs, Debug, Clone)] +struct DirectArgs { + /// TOML case file, for example tests/perf/query_cases/.../case.toml. + #[arg(long, value_name = "PATH")] + case: Option, + + /// Output directory (creates object-store/, manifest/, files.jsonl, summary.json). + #[arg(long, value_name = "DIR")] + out_dir: Option, + + /// Region ID to synthesize. Defaults to table id 1024, region 0. + #[arg(long, default_value = "4398046511104")] + region_id: u64, + + /// Table directory relative to object-store root. Defaults to data/{database}/{table}/. + #[arg(long, value_name = "DIR")] + table_dir: Option, + + /// Table name to generate when the case contains multiple [[tables]]. + #[arg(long, value_name = "NAME")] + table: Option, + + /// Manifest/checkpoint version. + #[arg(long, default_value = "1000000")] + checkpoint_version: u64, + + /// Safety flag when sst_count exceeds 1000. + #[arg(long)] + allow_large: bool, + + /// Print plan only. + #[arg(long)] + dry_run: bool, +} + +#[derive(ClapArgs, Debug)] +struct PlanArgs { + #[arg(long, value_name = "PATH")] + case: PathBuf, +} + +pub async fn run() { + let args = Args::parse(); + match args.command { + Some(Command::DirectSst(direct)) => run_direct_sst(direct).await, + Some(Command::PromRemoteWrite(rw)) => run_prom_remote_write(rw) + .await + .expect("prom remote write failed"), + Some(Command::InspectFooter(inspect)) => { + run_inspect_footer(inspect).expect("inspect footer failed") + } + Some(Command::Plan(plan)) => run_plan(plan).expect("plan failed"), + None => run_direct_sst(args.legacy).await, + } +} + +fn run_plan(args: PlanArgs) -> Result<(), Box> { + let case_text = fs::read_to_string(&args.case)?; + let mut case: CaseFile = toml::from_str(&case_text)?; + if let Scenario::PromRemoteWriteThenQuery(s) = &mut case.scenario + && let Some(storage) = &mut s.remote_write.storage + { + storage.populate_planned_thresholds(); + if storage.inspect { + let mut rb = s.remote_write.read_bench.clone().unwrap_or_default(); + if rb.enabled && rb.projection.is_empty() { + rb.projection = vec![storage.column.clone()]; + } + s.remote_write.read_bench = Some(rb); + } else if s + .remote_write + .read_bench + .as_ref() + .is_some_and(|rb| rb.enabled) + { + return Err("scenario.remote_write.read_bench requires scenario.remote_write.storage.inspect = true".into()); + } + } + println!( + "{}", + serde_json::to_string_pretty(&json!({"schema_version": 1, "scenario": case.scenario}))? + ); + Ok(()) +} diff --git a/src/cmd/src/bin/query_perf_fixture/prom_remote_write.rs b/src/cmd/src/bin/query_perf_fixture/prom_remote_write.rs new file mode 100644 index 00000000000..25a32572901 --- /dev/null +++ b/src/cmd/src/bin/query_perf_fixture/prom_remote_write.rs @@ -0,0 +1,208 @@ +// 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 clap::Args as ClapArgs; +use prost::Message; +use serde_json::json; +use servers::prom_store::snappy_compress; + +use super::ValuePattern; + +#[derive(Debug, ClapArgs)] +pub(super) struct PromRemoteWriteArgs { + #[arg(long, default_value = "http://127.0.0.1:4000/v1/prometheus/write")] + endpoint: String, + #[arg(long, default_value = "public")] + database: String, + #[arg(long)] + metric: String, + #[arg(long, default_value = "greptime_physical_table")] + physical_table: String, + #[arg(long, default_value_t = 8)] + series_count: u64, + #[arg(long, default_value_t = 30)] + samples_per_series: u64, + #[arg(long, default_value_t = 1_704_067_200_000)] + start_unix_millis: i64, + #[arg(long, default_value_t = 15_000)] + step_millis: i64, + #[arg(long, alias = "batch-size", default_value_t = 8)] + chunk_series_count: u64, + #[arg(long, default_value_t = 60)] + timeout_seconds: u64, + #[arg(long, default_value_t = ValuePattern::Linear)] + value_pattern: ValuePattern, + #[arg(long, default_value_t = 0.0)] + value_base: f64, + #[arg(long, default_value_t = 0.125)] + value_step: f64, + #[arg(long, default_value_t = 97)] + value_cardinality: u64, + #[arg(long, default_value_t = 0)] + value_seed: u64, + #[arg(long, default_value_t = 8)] + value_run_length: u64, + #[arg(long, default_value_t = 100)] + value_stall_every: u64, + #[arg(long, default_value_t = 16)] + value_stall_length: u64, + #[arg(long, default_value_t = 5)] + value_mixed_every: u64, + #[arg(long, default_value_t = 0)] + value_sample_offset: u64, + #[arg(long)] + value_total_samples_per_series: Option, +} + +pub(super) async fn run_prom_remote_write( + args: PromRemoteWriteArgs, +) -> Result<(), Box> { + use api::prom_store::remote::WriteRequest; + let started = std::time::Instant::now(); + let client = reqwest::Client::builder() + .timeout(std::time::Duration::from_secs(args.timeout_seconds)) + .build()?; + let mut batches = 0_u64; + let mut rows = 0_u64; + let mut http_statuses = Vec::new(); + for first in (0..args.series_count).step_by(args.chunk_series_count.max(1) as usize) { + let last = (first + args.chunk_series_count.max(1)).min(args.series_count); + let timeseries = (first..last) + .map(|series_idx| make_prom_series(&args, series_idx)) + .collect::>(); + rows += timeseries + .iter() + .map(|ts| ts.samples.len() as u64) + .sum::(); + let body = snappy_compress( + &WriteRequest { + timeseries, + ..Default::default() + } + .encode_to_vec(), + )?; + let response = client + .post(&args.endpoint) + .header("Content-Encoding", "snappy") + .body(body) + .send() + .await?; + let status = response.status(); + http_statuses.push(status.as_u16()); + if !status.is_success() { + return Err(format!( + "remote-write failed with status {status}: {}", + response.text().await.unwrap_or_default() + ) + .into()); + } + batches += 1; + } + println!( + "{}", + json!({"status":"ok","endpoint":args.endpoint,"database":args.database,"metric":args.metric,"physical_table":args.physical_table,"series_count":args.series_count,"samples_per_series":args.samples_per_series,"rows":rows,"samples_written":rows,"batches":batches,"elapsed_seconds":started.elapsed().as_secs_f64(),"http_statuses":http_statuses,"value":{"pattern":args.value_pattern.to_string(),"base":args.value_base,"step":args.value_step,"cardinality":args.value_cardinality,"seed":args.value_seed,"run_length":args.value_run_length,"stall_every":args.value_stall_every,"stall_length":args.value_stall_length,"mixed_every":args.value_mixed_every,"sample_offset":args.value_sample_offset,"total_samples_per_series":args.value_total_samples_per_series.unwrap_or(args.samples_per_series)}}) + ); + Ok(()) +} + +fn make_prom_series( + args: &PromRemoteWriteArgs, + series_idx: u64, +) -> api::prom_store::remote::TimeSeries { + let mut labels = vec![ + prom_label("__name__", &args.metric), + prom_label("x_greptime_database", &args.database), + prom_label("x_greptime_physical_table", &args.physical_table), + prom_label("host", &format!("host{:04}", series_idx % 1024)), + prom_label("instance", &format!("instance{:06}", series_idx)), + ]; + labels.sort_unstable_by(|l, r| l.name.cmp(&r.name)); + let samples = (0..args.samples_per_series) + .map(|sample_idx| api::prom_store::remote::Sample { + value: deterministic_prom_value(args, series_idx, sample_idx), + timestamp: args.start_unix_millis + (sample_idx as i64 * args.step_millis), + }) + .collect(); + api::prom_store::remote::TimeSeries { + labels, + samples, + ..Default::default() + } +} +fn prom_label(name: &str, value: &str) -> api::prom_store::remote::Label { + api::prom_store::remote::Label { + name: name.to_string(), + value: value.to_string(), + } +} +fn deterministic_prom_value(args: &PromRemoteWriteArgs, series_idx: u64, sample_idx: u64) -> f64 { + let ordinal = series_idx + * args + .value_total_samples_per_series + .unwrap_or(args.samples_per_series) + + args.value_sample_offset + + sample_idx; + let local = args.value_sample_offset + sample_idx; + match args.value_pattern { + ValuePattern::Linear => { + args.value_base + (series_idx % 97) as f64 + local as f64 * args.value_step + } + ValuePattern::Constant => args.value_base, + ValuePattern::Modulo => { + args.value_base + (ordinal % args.value_cardinality.max(1)) as f64 * args.value_step + } + ValuePattern::Unique => args.value_base + ordinal as f64 * args.value_step, + ValuePattern::SeededRandom => { + args.value_base + + (splitmix64(ordinal ^ args.value_seed) % args.value_cardinality.max(1)) as f64 + * args.value_step + } + ValuePattern::RunLength => { + args.value_base + + ((ordinal / args.value_run_length.max(1)) % args.value_cardinality.max(1)) as f64 + * args.value_step + } + ValuePattern::QuantizedSignal => { + args.value_base + + (((series_idx % args.value_cardinality.max(1)) + + (local / args.value_run_length.max(1))) + % args.value_cardinality.max(1)) as f64 + * args.value_step + } + ValuePattern::SignalWithSporadicStalls => { + let every = args.value_stall_every.max(1); + let phase = local % every; + let effective = if phase < args.value_stall_length.min(every) { + local - phase + } else { + local + }; + args.value_base + (series_idx % 97) as f64 + effective as f64 * args.value_step + } + ValuePattern::MixedSignalRepeated => { + if local.is_multiple_of(args.value_mixed_every.max(1)) { + args.value_base + } else { + args.value_base + (series_idx % 97) as f64 + local as f64 * args.value_step + } + } + } +} +fn splitmix64(mut value: u64) -> u64 { + value = value.wrapping_add(0x9e3779b97f4a7c15); + let mut mixed = value; + mixed = (mixed ^ (mixed >> 30)).wrapping_mul(0xbf58476d1ce4e5b9); + mixed = (mixed ^ (mixed >> 27)).wrapping_mul(0x94d049bb133111eb); + mixed ^ (mixed >> 31) +} diff --git a/src/cmd/src/bin/query_perf_fixture/util.rs b/src/cmd/src/bin/query_perf_fixture/util.rs new file mode 100644 index 00000000000..2c86baa6839 --- /dev/null +++ b/src/cmd/src/bin/query_perf_fixture/util.rs @@ -0,0 +1,23 @@ +// 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::path::Path; + +pub(super) fn case_name_from_path(path: &Path) -> String { + path.parent() + .and_then(Path::file_name) + .and_then(|name| name.to_str()) + .unwrap_or("query_perf_case") + .to_string() +} diff --git a/tests/perf/AGENTS.md b/tests/perf/AGENTS.md index eb95544d22d..7f64b283bdc 100644 --- a/tests/perf/AGENTS.md +++ b/tests/perf/AGENTS.md @@ -11,14 +11,14 @@ docs together. - `[case]` is report metadata only. `[scenario]` is the executable regression configuration and must include `kind`, data layout, tables, queries, and - thresholds. Keep scenario parsing explicit in both - `tests/perf/query_regression_runner.py` and - `src/cmd/src/bin/query_perf_fixture.rs`. + thresholds. Rust owns case schema, defaults, validation, and normalized plan + output through `query_perf_fixture plan`; `tests/perf/query_regression_runner.py` + should only orchestrate from that normalized JSON. - Keep the direct-SST generator generic. Issue-specific behavior belongs in case files and thresholds, not in Rust generator logic. - Before pushing perf harness changes, run at least: - `uv run --no-project python -m py_compile .github/scripts/query-regression-run.py .github/scripts/query-regression-summary.py .github/scripts/query-regression-pr-metadata.py tests/perf/query_regression_runner.py` - `cargo fmt --all -- --check` - - `cargo build -p cmd --bin query_perf_fixture` + - `cargo build -p cmd --bin query_perf_fixture --features dev-tools` - dry-run the Python runner and Rust fixture generator against all built-in cases when the DSL or workflow case selection changes. diff --git a/tests/perf/README.md b/tests/perf/README.md index f7a7abfe8eb..62b33b09626 100644 --- a/tests/perf/README.md +++ b/tests/perf/README.md @@ -55,20 +55,22 @@ kind = "prom_remote_write_then_query" database = "public" metric = "prom_remote_write_smoke" physical_table = "greptime_physical_table" -series_count = 8 -samples_per_series = 30 +series_count = 512 +samples_per_series = 1440 +sample_chunk_size = 480 +flush_every_sample_chunks = 1 start_unix_millis = 1_704_067_200_000 -step_millis = 15_000 -chunk_series_count = 4 -visibility_timeout_seconds = 30 +step_millis = 60_000 +chunk_series_count = 128 +timeout_seconds = 180 +visibility_timeout_seconds = 120 # Optional: split by time so each helper invocation writes only this many samples # per series, then periodically flush the physical metric table to produce # multiple time-interval SSTs. -# sample_chunk_size = 300 -# flush_every_sample_chunks = 1 [scenario.remote_write.prom_store] pending_rows_flush_interval = "1s" +max_batch_rows = 100000 ``` Cases can optionally add `[scenario.remote_write.value]` to control the generated @@ -102,11 +104,19 @@ ingestion, the runner passes the sample offset and total sample count to the helper so non-linear value patterns use a stable global/per-series ordinal across chunks. +Case schema, value distribution defaults, storage defaults, and read-bench +defaults are owned by Rust. The Python runner calls +`query_perf_fixture plan --case ` and orchestrates the normalized JSON. +The same helper exposes `direct-sst`, `prom-remote-write`, and `inspect-footer` +subcommands; the old direct invocation (`query_perf_fixture --case ... --out-dir +...`) remains compatible. + Remote-write cases that need to validate storage output can add `[scenario.remote_write.storage]`. When present, the scenario body becomes: remote-write → flush → visibility check → Parquet footer/size/encoding inspection -→ query. The inspector is a small Rust helper (`parquet_footer_inspector`) that -reads local Parquet footers directly; it does not require Python `pyarrow`. +→ optional read-bench → query. Footer inspection is handled by +`query_perf_fixture inspect-footer`, which reads local Parquet footers directly; +it does not require Python `pyarrow`. ```toml [scenario.remote_write.storage] @@ -140,17 +150,24 @@ By default the inspector root is the target datanode data home, which is suitabl for a fresh single-case data directory. For reused or more complex data homes, set `root_suffix` to a path relative to the datanode data home to narrow inspection. +When storage inspection is enabled, `[scenario.remote_write.read_bench]` defaults +to enabled with both datanode `parquetbench` and `scanbench`. Disable it with +`enabled = false`. `parquetbench` measures per-SST reader cost, `scanbench` +measures region scan cost, and query measurements still exercise the SQL/TQL +frontend path. Treat all performance conclusions as release-only; debug builds +are suitable only for command wiring and correctness checks. + The runner creates the configured database if needed, writes a per-target frontend config enabling `[prom_store]` with metric engine storage and a non-zero `pending_rows_flush_interval`, and validates that the logical metric table reaches `series_count * samples_per_series` rows before trusting the query measurements. -Use `--remote-write-generator /path/to/prom_remote_write_fixture` to provide the -Rust payload helper. Use `--storage-inspector /path/to/parquet_footer_inspector` -when `[scenario.remote_write.storage]` is enabled. `--fixture-only` is rejected -for remote-write cases; use `--dry-run` for planning. +Use `--fixture-generator /path/to/query_perf_fixture` to provide the Rust helper. +Deprecated `--remote-write-generator` and `--storage-inspector` options are kept +only for CLI compatibility. `--fixture-only` is rejected for remote-write cases; +use `--dry-run` for planning. Large manual remote-write cases can set `sample_chunk_size` to split ingestion by -time. For each chunk, the runner invokes `prom_remote_write_fixture` with the +time. For each chunk, the runner invokes `query_perf_fixture prom-remote-write` with the same series cardinality but a shorter `--samples-per-series` and an advanced `--start-unix-millis`. `flush_every_sample_chunks` controls periodic `ADMIN FLUSH_TABLE('')` calls; with `flush_every_sample_chunks = 1`, @@ -160,6 +177,12 @@ need multi-window file distribution should span multiple compaction windows. If `sample_chunk_size` is omitted, the runner keeps the older single-helper-invocation behavior and flushes once at the end. +`tests/perf/query_cases/prom_remote_write_smoke/case.toml` is the default +remote-write coverage case. It writes 512 series × 1440 samples, splits the day +into three flushed chunks, and requires multiple landed SSTs so storage +inspection and the read-bench supplement have enough data to reduce incidental +query noise while staying much smaller than the manual 7913 case. + `tests/perf/query_cases/prom_remote_write_7913/case.toml` is a larger manual case for issue #7913. It writes 8192 series × 20160 samples through remote-write in 1440-sample daily time chunks, flushing after each chunk before running 1d/7d/14d @@ -317,7 +340,8 @@ as a fallback. Manifest files are reflinked or copied, not hardlinked. Fixture generator smoke test: ```bash -cargo run -p cmd --bin query_perf_fixture -- \ +cargo run -p cmd --bin query_perf_fixture --features dev-tools -- \ + direct-sst \ --case tests/perf/query_cases/smoke_direct_sst/case.toml \ --out-dir /tmp/query-perf-smoke ``` @@ -345,7 +369,7 @@ uv run --no-project python tests/perf/query_regression_runner.py \ --case tests/perf/query_cases/prom_remote_write_smoke/case.toml \ --base-bin /path/to/base/greptime \ --candidate-bin /path/to/candidate/greptime \ - --remote-write-generator /path/to/prom_remote_write_fixture \ + --fixture-generator /path/to/query_perf_fixture \ --work-dir /tmp/query-perf-remote-write \ --dry-run ``` @@ -360,6 +384,10 @@ query regression runs. It builds its own binaries for now: candidate checkout - runner and summary formatter from the candidate checkout +The workflow builds `cmd` binaries with the `dev-tools` feature so the helper +binary and the datanode `parquetbench`/`scanbench` subcommands stay out of normal +release builds but remain available to this opt-in perf job. + The workflow runs automatically only for non-draft PRs labeled `query-regression` (on label/ready-for-review/reopen events, not every push). PR runs build base/candidate once and then run the default case set with diff --git a/tests/perf/query_cases/prom_remote_write_smoke/case.toml b/tests/perf/query_cases/prom_remote_write_smoke/case.toml index b946fbc8f51..e92dae98dd9 100644 --- a/tests/perf/query_cases/prom_remote_write_smoke/case.toml +++ b/tests/perf/query_cases/prom_remote_write_smoke/case.toml @@ -1,6 +1,6 @@ [case] name = "prom_remote_write_smoke" -description = "Tiny Prometheus remote-write ingestion then TQL query smoke case" +description = "Moderate Prometheus remote-write storage/read-bench case with multiple flushed SSTs" [scenario] kind = "prom_remote_write_then_query" @@ -9,22 +9,30 @@ kind = "prom_remote_write_then_query" database = "public" metric = "prom_remote_write_smoke" physical_table = "greptime_physical_table" -series_count = 8 -samples_per_series = 30 +series_count = 512 +samples_per_series = 1440 +sample_chunk_size = 480 +flush_every_sample_chunks = 1 start_unix_millis = 1_704_067_200_000 -step_millis = 15_000 -chunk_series_count = 4 -timeout_seconds = 30 +step_millis = 60_000 +chunk_series_count = 128 +timeout_seconds = 180 +visibility_timeout_seconds = 120 [scenario.remote_write.prom_store] pending_rows_flush_interval = "1s" -max_batch_rows = 10000 +max_batch_rows = 100000 + +[scenario.remote_write.storage] +column = "greptime_value" +min_files = 2 +min_files_with_column = 2 [[scenario.queries]] -name = "selector_5m_smoke" +name = "selector_1d_storage_smoke" kind = "tql" -query = "TQL ANALYZE VERBOSE (1704067200, 1704067650, '15s') prom_remote_write_smoke{host=~'host.*'}" -warmup = 0 +query = "TQL ANALYZE VERBOSE (1704067200, 1704153600, '5m') prom_remote_write_smoke{host=~'host.*'}" +warmup = 1 iterations = 2 [scenario.queries.thresholds] diff --git a/tests/perf/query_regression_runner.py b/tests/perf/query_regression_runner.py index 97f9f3914ee..b469a146596 100644 --- a/tests/perf/query_regression_runner.py +++ b/tests/perf/query_regression_runner.py @@ -66,8 +66,8 @@ def parse_args() -> argparse.Namespace: p.add_argument("--base-bin", required=True, type=Path) p.add_argument("--candidate-bin", required=True, type=Path) p.add_argument("--fixture-generator", type=Path) - p.add_argument("--remote-write-generator", type=Path, help="prom_remote_write_fixture helper binary") - p.add_argument("--storage-inspector", type=Path, help="parquet_footer_inspector helper binary") + p.add_argument("--remote-write-generator", type=Path, help="deprecated: use --fixture-generator query_perf_fixture") + p.add_argument("--storage-inspector", type=Path, help="deprecated: use --fixture-generator query_perf_fixture") p.add_argument("--work-dir", required=True, type=Path) p.add_argument("--fixture-cache-dir", type=Path, help="persistent directory for generated fixtures, keyed by case content") p.add_argument("--reuse-fixture", action="store_true") @@ -84,6 +84,16 @@ def load_case(path: Path) -> dict[str, Any]: return tomllib.load(f) +def load_normalized_case(case_path: Path, fixture_generator: Path | None) -> dict[str, Any]: + if fixture_generator is None: + raise ValueError("--fixture-generator query_perf_fixture is required for Rust-owned case planning") + result = run_command([str(fixture_generator), "plan", "--case", str(case_path)]) + if result["returncode"] != 0: + raise RuntimeError(f"query_perf_fixture plan failed: {result['stderr'][:2000]}") + plan = json.loads(result["stdout"]) + return {"case": load_case(case_path).get("case", {}), "scenario": plan["scenario"], "schema_version": plan.get("schema_version")} + + def require_binary(path: Path, name: str, *, dry_run: bool) -> None: if dry_run: return @@ -158,11 +168,9 @@ def scenario(case: dict[str, Any]) -> dict[str, Any]: def case_tables(case: dict[str, Any]) -> list[dict[str, Any]]: value = scenario(case) if value.get("kind") == "prom_remote_write_then_query": - remote = value.get("remote_write") or {} - metric = remote.get("metric") or remote.get("metric_name") - database = remote.get("database", "public") - if not metric: - raise ValueError("remote-write scenario requires scenario.remote_write.metric") + remote = value["remote_write"] + metric = remote["metric"] + database = remote["database"] return [{"database": database, "name": metric, "engine": "metric", "validate_show_create_engine": False}] tables = value.get("tables") or [] if not tables or value.get("layout", {}).get("regions") != 1: @@ -549,7 +557,7 @@ def generate_fixture(generator: Path | None, case_path: Path, fixture_dir: Path, if fixture_dir.exists() and not reuse_fixture: shutil.rmtree(fixture_dir) fixture_dir.mkdir(parents=True, exist_ok=True) - cmd = [str(generator), "--case", str(case_path), "--out-dir", str(fixture_dir)] + cmd = [str(generator), "direct-sst", "--case", str(case_path), "--out-dir", str(fixture_dir)] if table_name is not None: cmd += ["--table", table_name] if region_id is not None: @@ -779,28 +787,29 @@ def enforce_thresholds(case: dict[str, Any], base: dict[str, Any], candidate: di def storage_config(remote: dict[str, Any]) -> dict[str, Any] | None: - value = remote.get("storage") + value = remote["storage"] if not isinstance(value, dict): return None - if value.get("inspect", True) is False: + if value["inspect"] is False: return None return value -def run_storage_inspection(inspector: Path | None, target: RunTarget, storage: dict[str, Any], *, dry_run: bool) -> dict[str, Any]: - if inspector is None: +def run_storage_inspection(helper: Path | None, target: RunTarget, storage: dict[str, Any], *, dry_run: bool) -> dict[str, Any]: + if helper is None: if not dry_run: - raise ValueError("--storage-inspector is required when scenario.remote_write.storage.inspect is enabled") - inspector = Path("parquet_footer_inspector") + raise ValueError("--fixture-generator query_perf_fixture is required when storage inspection is enabled") + helper = Path("query_perf_fixture") root = target.datanode_data_dir if storage.get("root_suffix"): root = root / str(storage["root_suffix"]) cmd = [ - str(inspector), + str(helper), + "inspect-footer", "--root", str(root), - "--column", str(storage.get("column", "greptime_value")), + "--column", str(storage["column"]), ] - if storage.get("include_metadata_files", False): + if storage["include_metadata_files"]: cmd.append("--include-metadata-files") if dry_run: return {"status": "dry-run", "cmd": cmd, "root": str(root)} @@ -843,7 +852,7 @@ def enforce_storage_thresholds(storage: dict[str, Any] | None, base: dict[str, A results.append({"target": target_name, "threshold": name, "status": "passed" if ok else "failed", "actual": actual, "limit": limit}) def add_min(name: str, field: str) -> None: - limit = storage.get(name, 1 if name in ("min_files", "min_files_with_column") else None) + limit = storage[name] if limit is None: return for target_name, summary in target_summaries: @@ -870,9 +879,9 @@ def enforce_storage_thresholds(storage: dict[str, Any] | None, base: dict[str, A add_abs("max_column_uncompressed_size_bytes", "column_uncompressed_size") for target_name, summary in target_summaries: encodings = set(summary.get("unique_encodings") or []) - for required in storage.get("require_encodings", []): + for required in storage["require_encodings"]: results.append({"target": target_name, "threshold": "require_encodings", "encoding": required, "status": "passed" if required in encodings else "failed"}) - for forbidden in storage.get("forbid_encodings", []): + for forbidden in storage["forbid_encodings"]: results.append({"target": target_name, "threshold": "forbid_encodings", "encoding": forbidden, "status": "passed" if forbidden not in encodings else "failed"}) add_cmp("max_candidate_total_file_size_regression_pct", "total_file_size") add_cmp("max_candidate_column_compressed_size_regression_pct", "column_compressed_size") @@ -883,23 +892,135 @@ def enforce_storage_thresholds(storage: dict[str, Any] | None, base: dict[str, A def planned_storage_thresholds(storage: dict[str, Any] | None) -> list[dict[str, Any]]: if not storage: return [] - keys = [ - "min_files", - "min_files_with_column", - "require_encodings", - "forbid_encodings", - "max_total_file_size_bytes", - "max_column_compressed_size_bytes", - "max_column_uncompressed_size_bytes", - "max_candidate_total_file_size_regression_pct", - "max_candidate_column_compressed_size_regression_pct", - "max_candidate_column_uncompressed_size_regression_pct", - ] - planned = [] - for key in keys: - if key in storage or key in ("min_files", "min_files_with_column"): - planned.append({"threshold": key, "status": "planned", "value": storage.get(key, 1)}) - return planned + return list(storage["planned_thresholds"]) + + +REGION_DIR_RE = re.compile(r"^(\d+)_(\d{10})$") + + +def inspection_report(inspection: dict[str, Any] | None) -> dict[str, Any]: + if not inspection: + return {} + report = inspection.get("summary") + return report if isinstance(report, dict) else {} + + +def inspected_relative_path(target: RunTarget, inspection: dict[str, Any], relative_path: str) -> Path: + report = inspection_report(inspection) + root_text = report.get("root") or inspection.get("root") + if not root_text: + return Path(relative_path) + root = Path(root_text).resolve(strict=False) + data_home = target.datanode_data_dir.resolve(strict=False) + try: + root_suffix = root.relative_to(data_home) + except ValueError as e: + raise RuntimeError(f"storage inspection root {root} is not under datanode data home {data_home}") from e + return root_suffix / relative_path + + +def parse_sst_bench_target(target: RunTarget, inspection: dict[str, Any], file: dict[str, Any], *, dry_run: bool) -> dict[str, str] | None: + relative_path = str(file.get("relative_path") or "") + if dry_run and (not relative_path or "<" in relative_path): + return { + "relative_path": relative_path or "/_/data/.parquet", + "table_dir": "/", + "region_id": ":", + "path_type": "data", + "file_id": "", + } + if not relative_path: + return None + full_relative = inspected_relative_path(target, inspection, relative_path) + parts = full_relative.parts + if not parts or not parts[-1].endswith(".parquet"): + return None + region_index = next((idx for idx, part in enumerate(parts) if REGION_DIR_RE.match(part)), None) + if region_index is None or region_index == 0: + if dry_run: + return { + "relative_path": relative_path, + "table_dir": "/", + "region_id": ":", + "path_type": "data", + "file_id": Path(parts[-1]).stem, + } + return None + region_match = REGION_DIR_RE.match(parts[region_index]) + assert region_match is not None + file_index = len(parts) - 1 + path_type = "bare" + if region_index + 1 < file_index and parts[region_index + 1] in ("data", "metadata"): + path_type = parts[region_index + 1] + if path_type == "metadata": + return None + table_dir = "/".join(parts[:region_index]).rstrip("/") + "/" + return { + "relative_path": str(full_relative), + "table_dir": table_dir, + "region_id": f"{int(region_match.group(1))}:{int(region_match.group(2))}", + "path_type": path_type, + "file_id": Path(parts[-1]).stem, + } + + +def run_read_bench(target: RunTarget, read_bench: dict[str, Any] | None, storage_inspection: dict[str, Any] | None, *, dry_run: bool) -> dict[str, Any]: + if not read_bench or read_bench["enabled"] is False: + return {"status": "skipped", "reason": "read_bench disabled"} + run_parquetbench = read_bench["parquetbench"] + run_scanbench = read_bench["scanbench"] + if not run_parquetbench and not run_scanbench: + return {"status": "skipped", "reason": "parquetbench and scanbench disabled"} + if storage_inspection is None: + raise ValueError("read_bench requires storage inspection") + bench_dir = target.work_dir / "read_bench" + config_toml = bench_dir / "bench.toml" + scan_json = bench_dir / "scan.json" + files = inspection_report(storage_inspection).get("files") or [] + ssts = [f for f in files if f.get("relative_path", "").endswith(".parquet") and f.get("columns")] + if dry_run and not ssts: + ssts = [{"relative_path": ".parquet"}] + if read_bench["max_files"] is not None: + ssts = ssts[: int(read_bench["max_files"])] + bench_targets = [parsed for f in ssts if (parsed := parse_sst_bench_target(target, storage_inspection, f, dry_run=dry_run)) is not None] + if not bench_targets: + return {"status": "failed", "reason": "no inspected data SST files available for read_bench"} + config_text = f'[storage]\ndata_home = "{target.datanode_data_dir}"\ntype = "File"\n\n[[region_engine]]\n[region_engine.mito]\n' + commands = [] + parquet_runs = [] + scan_runs = [] + iterations = str(int(read_bench["iterations"])) + if run_parquetbench: + for bench_target in bench_targets: + cmd = [str(target.binary), "datanode", "parquetbench", "--config", str(config_toml), "--region-id", bench_target["region_id"], "--table-dir", bench_target["table_dir"], "--file-id", bench_target["file_id"], "--scan-config", str(scan_json), "--path-type", bench_target["path_type"], "--iterations", iterations, "--reader", str(read_bench["parquet_reader"])] + commands.append(cmd) + parquet_runs.append({**bench_target, "cmd": cmd, "status": "dry-run" if dry_run else "planned"}) + regions: dict[tuple[str, str, str], list[str]] = {} + if run_scanbench: + for bench_target in bench_targets: + regions.setdefault((bench_target["table_dir"], bench_target["region_id"], bench_target["path_type"]), []).append(bench_target["relative_path"]) + for (table_dir, region_id, path_type), region_files in regions.items(): + cmd = [str(target.binary), "datanode", "scanbench", "--config", str(config_toml), "--region-id", region_id, "--table-dir", table_dir, "--scan-config", str(scan_json), "--path-type", path_type, "--scanner", str(read_bench["scan_scanner"]), "--parallelism", str(int(read_bench["parallelism"])), "--iterations", iterations] + commands.append(cmd) + scan_runs.append({"table_dir": table_dir, "region_id": region_id, "path_type": path_type, "files": region_files, "cmd": cmd, "status": "dry-run" if dry_run else "planned"}) + if dry_run: + return {"status": "dry-run", "config_path": str(config_toml), "scan_config_path": str(scan_json), "commands": commands, "parquetbench": parquet_runs, "scanbench": scan_runs} + bench_dir.mkdir(parents=True, exist_ok=True) + config_toml.write_text(config_text) + scan_json.write_text(json.dumps({"projection_names": read_bench["projection"]}, indent=2) + "\n") + avg_re = re.compile(r"Average duration[^0-9]*([0-9.]+)\s*ms", re.I) + for runs in (parquet_runs, scan_runs): + for run in runs: + result = run_command(run["cmd"]) + run.update(result) + run["status"] = "ok" if result["returncode"] == 0 else "failed" + m = avg_re.search(result.get("stdout", "")) + if m: + run["average_ms"] = float(m.group(1)) + def median(runs: list[dict[str, Any]]) -> float | None: + vals = [r["average_ms"] for r in runs if "average_ms" in r] + return statistics.median(vals) if vals else None + return {"status": "ok" if all(r.get("status") == "ok" for r in parquet_runs + scan_runs) else "failed", "config_path": str(config_toml), "scan_config_path": str(scan_json), "parquetbench": parquet_runs, "scanbench": scan_runs, "aggregate": {"parquetbench_median_average_ms": median(parquet_runs), "scanbench_median_average_ms": median(scan_runs)}} def write_json(path: Path, payload: dict[str, Any]) -> None: @@ -916,18 +1037,18 @@ def require_fresh_work_dirs(targets: list[RunTarget], *, reuse_work_dir: bool, d def write_frontend_prom_config(target: RunTarget, remote: dict[str, Any]) -> Path: - prom = remote.get("prom_store") or {} + prom = remote["prom_store"] path = target.work_dir / "frontend-prom-store.toml" path.parent.mkdir(parents=True, exist_ok=True) path.write_text( "[prom_store]\n" "enable = true\n" "with_metric_engine = true\n" - f"pending_rows_flush_interval = {json.dumps(str(prom.get('pending_rows_flush_interval', '1s')))}\n" - f"max_batch_rows = {int(prom.get('max_batch_rows', 100000))}\n" - f"max_concurrent_flushes = {int(prom.get('max_concurrent_flushes', 256))}\n" - f"worker_channel_capacity = {int(prom.get('worker_channel_capacity', 65526))}\n" - f"max_inflight_requests = {int(prom.get('max_inflight_requests', 3000))}\n" + f"pending_rows_flush_interval = {json.dumps(str(prom['pending_rows_flush_interval']))}\n" + f"max_batch_rows = {int(prom['max_batch_rows'])}\n" + f"max_concurrent_flushes = {int(prom['max_concurrent_flushes'])}\n" + f"worker_channel_capacity = {int(prom['worker_channel_capacity'])}\n" + f"max_inflight_requests = {int(prom['max_inflight_requests'])}\n" ) return path @@ -935,47 +1056,38 @@ def write_frontend_prom_config(target: RunTarget, remote: dict[str, Any]) -> Pat def run_remote_write(generator: Path | None, target: RunTarget, remote: dict[str, Any], *, dry_run: bool) -> dict[str, Any]: if generator is None: if not dry_run: - raise ValueError("--remote-write-generator is required for prom_remote_write_then_query unless --dry-run is set") - generator = Path("prom_remote_write_fixture") - metric = remote.get("metric") or remote.get("metric_name") - physical_table = remote.get("physical_table", "greptime_physical_table") + raise ValueError("--fixture-generator query_perf_fixture is required for prom_remote_write_then_query") + generator = Path("query_perf_fixture") + metric = remote["metric"] + physical_table = remote["physical_table"] cmd = [ str(generator), + "prom-remote-write", "--endpoint", f"http://127.0.0.1:{target.http_port}/v1/prometheus/write", - "--database", remote.get("database", "public"), + "--database", remote["database"], "--metric", metric, "--physical-table", physical_table, - "--series-count", str(int(remote.get("series_count", 8))), - "--samples-per-series", str(int(remote.get("samples_per_series", 30))), - "--start-unix-millis", str(int(remote.get("start_unix_millis", 1_704_067_200_000))), - "--step-millis", str(int(remote.get("step_millis", 15_000))), - "--chunk-series-count", str(int(remote.get("chunk_series_count", remote.get("batch_size", 8)))), - "--timeout-seconds", str(int(remote.get("timeout_seconds", 60))), + "--series-count", str(int(remote["series_count"])), + "--samples-per-series", str(int(remote["samples_per_series"])), + "--start-unix-millis", str(int(remote["start_unix_millis"])), + "--step-millis", str(int(remote["step_millis"])), + "--chunk-series-count", str(int(remote["chunk_series_count"])), + "--timeout-seconds", str(int(remote["timeout_seconds"])), ] - value = remote.get("value") - if isinstance(value, dict): - if "pattern" in value: - cmd.extend(["--value-pattern", str(value["pattern"])]) - if "base" in value: - cmd.extend(["--value-base", str(value["base"])]) - if "step" in value: - cmd.extend(["--value-step", str(value["step"])]) - if "cardinality" in value: - cmd.extend(["--value-cardinality", str(int(value["cardinality"]))]) - if "seed" in value: - cmd.extend(["--value-seed", str(int(value["seed"]))]) - if "run_length" in value: - cmd.extend(["--value-run-length", str(int(value["run_length"]))]) - if "stall_every" in value: - cmd.extend(["--value-stall-every", str(int(value["stall_every"]))]) - if "stall_length" in value: - cmd.extend(["--value-stall-length", str(int(value["stall_length"]))]) - if "mixed_every" in value: - cmd.extend(["--value-mixed-every", str(int(value["mixed_every"]))]) - if "sample_offset" in remote: - cmd.extend(["--value-sample-offset", str(int(remote["sample_offset"]))]) - if "total_samples_per_series" in remote: - cmd.extend(["--value-total-samples-per-series", str(int(remote["total_samples_per_series"]))]) + value = remote["value"] + cmd.extend(["--value-pattern", str(value["pattern"])]) + cmd.extend(["--value-base", str(value["base"])]) + cmd.extend(["--value-step", str(value["step"])]) + cmd.extend(["--value-cardinality", str(int(value["cardinality"]))]) + cmd.extend(["--value-seed", str(int(value["seed"]))]) + cmd.extend(["--value-run-length", str(int(value["run_length"]))]) + cmd.extend(["--value-stall-every", str(int(value["stall_every"]))]) + cmd.extend(["--value-stall-length", str(int(value["stall_length"]))]) + cmd.extend(["--value-mixed-every", str(int(value["mixed_every"]))]) + if "sample_offset" in remote: + cmd.extend(["--value-sample-offset", str(int(remote["sample_offset"]))]) + if "total_samples_per_series" in remote: + cmd.extend(["--value-total-samples-per-series", str(int(remote["total_samples_per_series"]))]) if dry_run: return {"status": "dry-run", "cmd": cmd} result = run_command(cmd) @@ -1011,23 +1123,23 @@ def flush_remote_physical_table(target: RunTarget, db: str, physical_table: str, def run_remote_write_ingestion(generator: Path | None, target: RunTarget, remote: dict[str, Any], args: argparse.Namespace, *, dry_run: bool) -> tuple[dict[str, Any], list[dict[str, Any]]]: - sample_chunk_size = remote.get("sample_chunk_size") - db = remote.get("database", "public") - physical_table = remote.get("physical_table", "greptime_physical_table") + sample_chunk_size = remote["sample_chunk_size"] + db = remote["database"] + physical_table = remote["physical_table"] if sample_chunk_size is None: rw = run_remote_write(generator, target, remote, dry_run=dry_run) flush = flush_remote_physical_table(target, db, physical_table, args, dry_run=dry_run, reason="final") return rw, [flush] - total_samples = int(remote.get("samples_per_series", 30)) + total_samples = int(remote["samples_per_series"]) chunk_samples = int(sample_chunk_size) if chunk_samples <= 0: raise ValueError("scenario.remote_write.sample_chunk_size must be positive") - flush_every = int(remote.get("flush_every_sample_chunks", 1)) + flush_every = int(remote["flush_every_sample_chunks"]) if flush_every <= 0: raise ValueError("scenario.remote_write.flush_every_sample_chunks must be positive") - start = int(remote.get("start_unix_millis", 1_704_067_200_000)) - step = int(remote.get("step_millis", 15_000)) + start = int(remote["start_unix_millis"]) + step = int(remote["step_millis"]) chunks: list[dict[str, Any]] = [] flushes: list[dict[str, Any]] = [] chunk_index = 0 @@ -1052,8 +1164,8 @@ def run_remote_write_ingestion(generator: Path | None, target: RunTarget, remote flushes.append(flush_remote_physical_table(target, db, physical_table, args, dry_run=dry_run, reason="final", chunk_index=chunk_index)) aggregate = summarize_remote_write_chunks(chunks) if dry_run: - series_count = int(remote.get("series_count", 8)) - chunk_series_count = int(remote.get("chunk_series_count", remote.get("batch_size", 8))) + series_count = int(remote["series_count"]) + chunk_series_count = int(remote["chunk_series_count"]) aggregate["rows"] = series_count * total_samples aggregate["samples_written"] = aggregate["rows"] aggregate["batches"] = chunk_index * ((series_count + chunk_series_count - 1) // chunk_series_count) @@ -1068,7 +1180,7 @@ def run_remote_write_ingestion(generator: Path | None, target: RunTarget, remote def expected_remote_write_rows(remote: dict[str, Any]) -> int: - return int(remote.get("series_count", 8)) * int(remote.get("samples_per_series", 30)) + return int(remote["series_count"]) * int(remote["samples_per_series"]) def extract_count_value(result: dict[str, Any]) -> int | None: @@ -1117,19 +1229,18 @@ def poll_expected_count(target: RunTarget, table_name: str, db: str, expected_ro def run_remote_write_scenario(args: argparse.Namespace, case: dict[str, Any], case_path: Path, targets: list[RunTarget], report: dict[str, Any]) -> None: - remote = scenario(case).get("remote_write") or {} + remote = scenario(case)["remote_write"] tables = case_tables(case) if args.fixture_only: raise ValueError("--fixture-only is not supported for prom_remote_write_then_query; use --dry-run for planning") - if args.remote_write_generator is not None: - require_binary(args.remote_write_generator, "remote-write generator", dry_run=args.dry_run) + helper = args.fixture_generator or args.remote_write_generator + if helper is not None: + require_binary(helper, "query_perf_fixture", dry_run=args.dry_run) elif not args.dry_run: - raise ValueError("--remote-write-generator is required for prom_remote_write_then_query") + raise ValueError("--fixture-generator is required for prom_remote_write_then_query") storage = storage_config(remote) - if storage and args.storage_inspector is not None: - require_binary(args.storage_inspector, "storage inspector", dry_run=args.dry_run) - elif storage and not args.dry_run: - raise ValueError("--storage-inspector is required when scenario.remote_write.storage.inspect is enabled") + if storage and helper is None and not args.dry_run: + raise ValueError("--fixture-generator is required when scenario.remote_write.storage.inspect is enabled") clusters: list[DistributedCluster] = [] target_results = [] storage_results = [] @@ -1141,29 +1252,32 @@ def run_remote_write_scenario(args: argparse.Namespace, case: dict[str, Any], ca clusters.append(cluster) if not args.dry_run: cluster.start_all(config_path) - db = remote.get("database", "public") + db = remote["database"] create_database = {"status": "dry-run", "database": db} if not args.dry_run: create_database = http_post_sql(target.http_port, f"CREATE DATABASE IF NOT EXISTS {sql_ident(db)}", "public", args.http_timeout) - rw, flushes = run_remote_write_ingestion(args.remote_write_generator, target, remote, args, dry_run=args.dry_run) + rw, flushes = run_remote_write_ingestion(helper, target, remote, args, dry_run=args.dry_run) visibility = {"status": "dry-run"} query_result = {"validation": [], "validation_errors": [], "measurements": [], "status": "planned"} if not args.dry_run: - visibility = poll_expected_count(target, tables[0]["name"], db, expected_remote_write_rows(remote), float(remote.get("visibility_timeout_seconds", 30)), args.http_timeout) + visibility = poll_expected_count(target, tables[0]["name"], db, expected_remote_write_rows(remote), float(remote["visibility_timeout_seconds"]), args.http_timeout) storage_inspection = None if storage: - storage_inspection = run_storage_inspection(args.storage_inspector, target, storage, dry_run=args.dry_run) + storage_inspection = run_storage_inspection(helper or args.storage_inspector, target, storage, dry_run=args.dry_run) storage_results.append(storage_inspection) + read_bench_result = run_read_bench(target, remote["read_bench"], storage_inspection, dry_run=args.dry_run) if not args.dry_run: query_result = run_queries(target, case, tables, args.http_timeout) tr = {"name": target.name, "binary": str(target.binary), "work_dir": str(target.work_dir), "components": cluster.component_report(), "frontend_config": str(config_path), "create_database": create_database, "remote_write": rw, "flushes": flushes, "flush": flushes[-1] if flushes else None, "visibility": visibility, **query_result} if storage_inspection is not None: tr["storage_inspection"] = storage_inspection + tr["read_bench"] = read_bench_result flushes_ok = all(flush.get("ok") for flush in flushes) storage_ok = storage_inspection is None or args.dry_run or storage_inspection.get("status") == "ok" - remote_checks_ok = args.dry_run or (create_database.get("ok") and flushes_ok and visibility.get("ok") and visibility.get("row_count_ok") and storage_ok) + read_bench_ok = args.dry_run or read_bench_result.get("status") in ("ok", "skipped") + remote_checks_ok = args.dry_run or (create_database.get("ok") and flushes_ok and visibility.get("ok") and visibility.get("row_count_ok") and storage_ok and read_bench_ok) if not remote_checks_ok: - tr.setdefault("validation_errors", []).append({"phase": "remote_write_visibility_storage", "create_database_ok": create_database.get("ok"), "flushes_ok": flushes_ok, "visibility_ok": visibility.get("ok"), "row_count_ok": visibility.get("row_count_ok"), "storage_ok": storage_ok, "create_database": create_database, "flushes": flushes, "visibility": visibility}) + tr.setdefault("validation_errors", []).append({"phase": "remote_write_visibility_storage", "create_database_ok": create_database.get("ok"), "flushes_ok": flushes_ok, "visibility_ok": visibility.get("ok"), "row_count_ok": visibility.get("row_count_ok"), "storage_ok": storage_ok, "read_bench_ok": read_bench_ok, "create_database": create_database, "flushes": flushes, "visibility": visibility, "read_bench": read_bench_result}) tr["status"] = "planned" if args.dry_run else ("measured" if remote_checks_ok and query_result["status"] == "ok" else "failed") write_json(target.report_path, tr) report["targets"].append(tr) @@ -1181,7 +1295,9 @@ def run_remote_write_scenario(args: argparse.Namespace, case: dict[str, Any], ca def main() -> int: args = parse_args() case_path = args.case.resolve() - case = load_case(case_path) + if args.fixture_generator is not None: + require_binary(args.fixture_generator, "fixture generator", dry_run=False) + case = load_normalized_case(case_path, args.fixture_generator) scenario_config = scenario(case) scenario_kind = scenario_config.get("kind") tables = case_tables(case) @@ -1191,9 +1307,6 @@ def main() -> int: require_binary(args.candidate_bin, "candidate", dry_run=args.dry_run or args.fixture_only) if scenario_kind == "direct_readable_sst" and not args.dry_run and args.fixture_generator is None: raise ValueError("--fixture-generator is required unless --dry-run is set") - if args.fixture_generator is not None: - require_binary(args.fixture_generator, "fixture generator", dry_run=args.dry_run) - ports = allocate_ports(16) targets = [make_target("base", args.base_bin.resolve(), work_root, ports[:8]), make_target("candidate", args.candidate_bin.resolve(), work_root, ports[8:])] require_fresh_work_dirs(targets, reuse_work_dir=args.reuse_work_dir, dry_run=args.dry_run, fixture_only=args.fixture_only)