From 0cc83c45708e442d929c22084fc0e6dda9c76f4a Mon Sep 17 00:00:00 2001 From: Yingwen Date: Mon, 24 Aug 2026 13:34:30 +0000 Subject: [PATCH] feat(cmd): add parquet development tools (#8939) * feat(cmd): add parquet metadata development tool Signed-off-by: evenyag * feat(cmd): add parquet rewrite development tool Signed-off-by: evenyag * feat(cmd): add SST replacement development tool Signed-off-by: evenyag * feat(cmd): support local parquet files in parquetbench Signed-off-by: evenyag * fix(cmd): clean up parquet development tools Signed-off-by: evenyag * refactor(cmd): share datanode tool utilities Signed-off-by: evenyag * fix(cmd): satisfy parquet tool lints Signed-off-by: evenyag * docs: document parquet development tools Signed-off-by: evenyag * refactor(cmd): rename SST replacement module Signed-off-by: evenyag * fix(cmd): validate parquet rewrite options Signed-off-by: evenyag --------- Signed-off-by: evenyag --- docs/parquet-tools.md | 253 +++++++++++ src/cmd/Cargo.toml | 2 +- src/cmd/src/bin/greptime.rs | 6 + src/cmd/src/datanode.rs | 46 ++ src/cmd/src/datanode/objbench.rs | 107 +---- src/cmd/src/datanode/parquet_meta.rs | 340 ++++++++++++++ src/cmd/src/datanode/parquet_rewrite.rs | 579 ++++++++++++++++++++++++ src/cmd/src/datanode/parquetbench.rs | 399 +++++++++++----- src/cmd/src/datanode/scanbench.rs | 60 +-- src/cmd/src/datanode/sst_replace.rs | 361 +++++++++++++++ src/cmd/src/datanode/tool_util.rs | 236 ++++++++++ 11 files changed, 2126 insertions(+), 263 deletions(-) create mode 100644 docs/parquet-tools.md create mode 100644 src/cmd/src/datanode/parquet_meta.rs create mode 100644 src/cmd/src/datanode/parquet_rewrite.rs create mode 100644 src/cmd/src/datanode/sst_replace.rs create mode 100644 src/cmd/src/datanode/tool_util.rs diff --git a/docs/parquet-tools.md b/docs/parquet-tools.md new file mode 100644 index 0000000000..8528f96943 --- /dev/null +++ b/docs/parquet-tools.md @@ -0,0 +1,253 @@ +# Parquet Development Tools + +GreptimeDB provides the following datanode CLI commands for inspecting, +rewriting, benchmarking, and replacing Parquet SST files: + +- `parquet-meta`: inspect a local Parquet file. +- `parquet-rewrite`: rewrite a local Parquet file with different writer + properties. +- `parquetbench`: benchmark reads from a local GreptimeDB SST or an SST in a + configured object store. +- `sst-replace`: replace an existing Mito SST object and update its manifest + metadata. + +These are development and recovery tools, not stable user-facing interfaces. +Always test a workflow on disposable data before using it on an important +region. + +## Build + +`parquet-meta`, `parquet-rewrite`, and `sst-replace` require the `dev-tools` +feature. The same build also includes `parquetbench`: + +```bash +cargo build -p cmd --bin greptime --features dev-tools +``` + +The examples below use: + +```bash +GREPTIME=./target/debug/greptime +``` + +Run `$GREPTIME datanode --help` for the authoritative argument list. + +## Inspect a Parquet file + +`parquet-meta` reads a local Parquet footer and optional page indexes. It reports +file, row-group, and column metadata, including compression, encodings, sizes, +page offsets, statistics presence, and bloom-filter/index offsets. + +Text output: + +```bash +$GREPTIME datanode parquet-meta \ + --input /tmp/source.parquet +``` + +Machine-readable JSON output: + +```bash +$GREPTIME datanode parquet-meta \ + --input /tmp/source.parquet \ + --format json > /tmp/source-meta.json +``` + +`--format` accepts `text` (the default) or `json`. This command can inspect a +general Parquet file; it does not require GreptimeDB SST key-value metadata. + +## Rewrite a Parquet file + +`parquet-rewrite` has two modes. First, dump a TOML properties file inferred +from an existing file: + +```bash +$GREPTIME datanode parquet-rewrite \ + --input /tmp/source.parquet \ + --dump-properties /tmp/writer-properties.toml +``` + +Review and edit the generated file. Its shape is: + +```toml +[writer] +dictionary_enabled = true +compression = "zstd" +max_row_group_row_count = 8192 + +[[columns]] +path = ["host"] +dictionary_enabled = true +compression = "zstd" +encoding = "plain" +``` + +Supported compression names are `uncompressed`, `snappy`, `gzip`, `lzo`, +`brotli`, `lz4`, `zstd`, and `lz4-raw`. Supported encodings are `plain`, +`delta-binary-packed`, `delta-length-byte-array`, `delta-byte-array`, and +`byte-stream-split`. + +The `[writer]` table also accepts `compression_level`, +`data_page_size_limit`, `data_page_row_count_limit`, and +`dictionary_page_size_limit`. A `[[columns]]` entry overrides dictionary, +compression, compression level, or encoding for its column path. Unknown TOML +fields are rejected. Compression levels are supported for `gzip`, `brotli`, and +`zstd`. A column-level `compression_level` without `compression` inherits the +writer compression codec; it is rejected when there is no writer codec to +inherit. + +Rewrite the data using the edited properties: + +```bash +$GREPTIME datanode parquet-rewrite \ + --input /tmp/source.parquet \ + --properties /tmp/writer-properties.toml \ + --output /tmp/rewritten.parquet +``` + +Use `--batch-size ` to control reader batch size. Output and dumped +properties files are not replaced unless `--overwrite` is passed. The output or +dump path must not be exactly the same path as the input, even with +`--overwrite`. This check does not resolve symbolic links, hard links, or other +spellings of the same path. + +The rewrite decodes and writes the Arrow record batches and copies the Parquet +key-value metadata, but it creates a new physical Parquet layout. Inspect and +validate the result before using it as an SST: + +```bash +$GREPTIME datanode parquet-meta \ + --input /tmp/rewritten.parquet +``` + +The dumped properties are inferred primarily from the first row group. Review +them when the source uses different properties across row groups. + +## Benchmark an SST + +`parquetbench` expects GreptimeDB region metadata embedded in the SST. It can +read a local file with the direct reader: + +```bash +$GREPTIME datanode parquetbench \ + --file-path /tmp/rewritten.parquet \ + --reader direct \ + --iterations 5 \ + --batch-size 8192 +``` + +Local-file mode cannot be combined with `--config`, `--region-id`, +`--table-dir`, or `--file-id`, and it does not support the `flat-prune` reader. + +To benchmark an SST in the object store configured for a datanode or standalone +deployment: + +```bash +$GREPTIME datanode parquetbench \ + --config /path/to/datanode.toml \ + --region-id 1024:0 \ + --table-dir data/greptime/public/1024 \ + --file-id 00020380-009c-426d-953e-b4e34c15af34 \ + --path-type bare \ + --reader flat-prune \ + --iterations 5 +``` + +Region mode requires all four of `--config`, `--region-id`, `--table-dir`, and +`--file-id`. `--region-id` accepts either the packed unsigned integer or +`:`. `--path-type` accepts `bare`, `data`, or +`metadata` and defaults to `bare`. + +An optional scan configuration selects columns and row groups: + +```json +{ + "projection_names": ["host", "value", "ts"], + "row_groups": [0, 2] +} +``` + +```bash +$GREPTIME datanode parquetbench \ + --file-path /tmp/rewritten.parquet \ + --scan-config /tmp/parquet-scan.json \ + --iterations 5 +``` + +Use `--pk-as-binary` to expose `__primary_key` as binary with the direct reader. +On Unix, `--pprof-file ` writes a flamegraph. Add +`--pprof-after-warmup` and use at least two iterations to exclude the first +iteration from profiling. + +## Replace an existing region SST + +> **Warning:** `sst-replace` mutates both an SST object and its region manifest. +> Stop the datanode that owns the region and back up the target SST and manifest +> before using `--confirm`. The SST write happens before the manifest update, so +> an interrupted or failed operation may require restoring the backup. + +`sst-replace` replaces the contents of an existing SST file ID. It does not add +a new file ID to a region. The command requires the replacement to have the same +row count and row-group count recorded in the manifest when those manifest +values are nonzero. It does not validate schema or row contents. The replacement +is loaded into memory in full, so ensure the machine has enough memory for the +SST. + +Start with the default dry run using a local replacement file: + +```bash +$GREPTIME datanode sst-replace \ + --config /path/to/datanode.toml \ + --region-id 1024:0 \ + --table-dir data/greptime/public/1024 \ + --file-id 00020380-009c-426d-953e-b4e34c15af34 \ + --replacement-file /tmp/rewritten.parquet +``` + +The dry run locates the manifest and target SST, reads and validates the +replacement footer, and prints the old and new sizes without writing anything. +`--path-type` defaults to `auto`, which probes `bare`, `data`, and `metadata`. +Specify the path type if the file ID is visible in more than one manifest. + +The replacement can instead be read from the configured object store: + +```bash +$GREPTIME datanode sst-replace \ + --config /path/to/datanode.toml \ + --region-id 1024:0 \ + --table-dir data/greptime/public/1024 \ + --file-id 00020380-009c-426d-953e-b4e34c15af34 \ + --replacement-object staging/rewritten.parquet \ + --path-type bare +``` + +After reviewing the dry-run output and confirming that the datanode is stopped, +repeat the exact command with `--confirm`: + +```bash +$GREPTIME datanode sst-replace \ + --config /path/to/datanode.toml \ + --region-id 1024:0 \ + --table-dir data/greptime/public/1024 \ + --file-id 00020380-009c-426d-953e-b4e34c15af34 \ + --replacement-file /tmp/rewritten.parquet \ + --path-type bare \ + --confirm +``` + +The confirmed operation overwrites the existing SST object, recalculates file +size and row-group statistics, and appends a manifest edit for the existing file +ID. Restart the datanode and validate queries against the region before removing +the backup. + +## Recommended rewrite and replacement workflow + +1. Back up the source SST and its region manifest. +2. Inspect the source with `parquet-meta`. +3. Dump and edit writer properties with `parquet-rewrite`. +4. Rewrite to a new local file; never rewrite directly over the source SST. +5. Inspect the output and benchmark it with `parquetbench`. +6. Stop the owning datanode. +7. Run `sst-replace` without `--confirm` and review its resolved target and + statistics. +8. Repeat with `--confirm`, restart the datanode, and validate the region. diff --git a/src/cmd/Cargo.toml b/src/cmd/Cargo.toml index 10865b48a1..579e5b1454 100644 --- a/src/cmd/Cargo.toml +++ b/src/cmd/Cargo.toml @@ -27,7 +27,7 @@ default = [ "meta-srv/mysql_kvbackend", ] enterprise = ["common-meta/enterprise", "frontend/enterprise", "meta-srv/enterprise"] -# Developer-only helper binary gate for `query_perf_fixture`. +# Developer-only helper binaries and diagnostic datanode commands. # Kept out of `default` so normal/release builds don't compile them. dev-tools = [] mysql-object-store = ["object-store/mysql-object-store"] diff --git a/src/cmd/src/bin/greptime.rs b/src/cmd/src/bin/greptime.rs index 957e227355..f5f4a59048 100644 --- a/src/cmd/src/bin/greptime.rs +++ b/src/cmd/src/bin/greptime.rs @@ -118,6 +118,12 @@ async fn start(cli: Command) -> Result<()> { datanode::SubCommand::Objbench(ref bench) => bench.run().await, datanode::SubCommand::Scanbench(ref bench) => bench.run().await, datanode::SubCommand::Parquetbench(ref bench) => bench.run().await, + #[cfg(feature = "dev-tools")] + datanode::SubCommand::ParquetMeta(ref cmd) => cmd.run().await, + #[cfg(feature = "dev-tools")] + datanode::SubCommand::ParquetRewrite(ref cmd) => cmd.run().await, + #[cfg(feature = "dev-tools")] + datanode::SubCommand::SstReplace(ref cmd) => cmd.run().await, }, SubCommand::Flownode(cmd) => { cmd.build(cmd.load_options(&cli.global_options)?) diff --git a/src/cmd/src/datanode.rs b/src/cmd/src/datanode.rs index c33be6c6f3..9cb8a8c66c 100644 --- a/src/cmd/src/datanode.rs +++ b/src/cmd/src/datanode.rs @@ -15,10 +15,20 @@ pub mod builder; #[allow(clippy::print_stdout)] pub(crate) mod objbench; +#[cfg(feature = "dev-tools")] +#[allow(clippy::print_stdout)] +mod parquet_meta; +#[cfg(feature = "dev-tools")] +#[allow(clippy::print_stdout)] +mod parquet_rewrite; #[allow(clippy::print_stdout)] pub mod parquetbench; #[allow(clippy::print_stdout)] pub mod scanbench; +#[cfg(feature = "dev-tools")] +#[allow(clippy::print_stdout)] +mod sst_replace; +mod tool_util; use std::path::Path; use std::time::Duration; @@ -39,8 +49,14 @@ use tracing_appender::non_blocking::WorkerGuard; use crate::App; use crate::datanode::builder::InstanceBuilder; use crate::datanode::objbench::ObjbenchCommand; +#[cfg(feature = "dev-tools")] +use crate::datanode::parquet_meta::ParquetMetaCommand; +#[cfg(feature = "dev-tools")] +use crate::datanode::parquet_rewrite::ParquetRewriteCommand; use crate::datanode::parquetbench::ParquetbenchCommand; use crate::datanode::scanbench::ScanbenchCommand; +#[cfg(feature = "dev-tools")] +use crate::datanode::sst_replace::SstReplaceCommand; use crate::error::{ LoadLayeredConfigSnafu, MissingConfigSnafu, Result, ShutdownDatanodeSnafu, StartDatanodeSnafu, }; @@ -114,6 +130,12 @@ impl Command { // Bench commands are standalone utilities and don't need to load DatanodeOptions. SubCommand::Objbench(_) | SubCommand::Scanbench(_) => Self::default_bench_options(), SubCommand::Parquetbench(_) => Self::default_bench_options(), + #[cfg(feature = "dev-tools")] + SubCommand::ParquetMeta(_) => Self::default_bench_options(), + #[cfg(feature = "dev-tools")] + SubCommand::ParquetRewrite(_) => Self::default_bench_options(), + #[cfg(feature = "dev-tools")] + SubCommand::SstReplace(_) => Self::default_bench_options(), } } @@ -139,6 +161,15 @@ pub enum SubCommand { Scanbench(ScanbenchCommand), /// Benchmark scanning a single parquet SST. Parquetbench(ParquetbenchCommand), + /// Display metadata of a parquet file. + #[cfg(feature = "dev-tools")] + ParquetMeta(ParquetMetaCommand), + /// Rewrite a parquet file with different writer properties. + #[cfg(feature = "dev-tools")] + ParquetRewrite(ParquetRewriteCommand), + /// Replace a Mito region SST and update its manifest metadata. + #[cfg(feature = "dev-tools")] + SstReplace(SstReplaceCommand), } impl SubCommand { @@ -160,6 +191,21 @@ impl SubCommand { cmd.run().await?; std::process::exit(0); } + #[cfg(feature = "dev-tools")] + SubCommand::ParquetMeta(cmd) => { + cmd.run().await?; + std::process::exit(0); + } + #[cfg(feature = "dev-tools")] + SubCommand::ParquetRewrite(cmd) => { + cmd.run().await?; + std::process::exit(0); + } + #[cfg(feature = "dev-tools")] + SubCommand::SstReplace(cmd) => { + cmd.run().await?; + std::process::exit(0); + } } } } diff --git a/src/cmd/src/datanode/objbench.rs b/src/cmd/src/datanode/objbench.rs index f33734687d..a01cad8468 100644 --- a/src/cmd/src/datanode/objbench.rs +++ b/src/cmd/src/datanode/objbench.rs @@ -18,8 +18,6 @@ use std::time::Instant; use clap::Parser; use colored::Colorize; -use datanode::config::RegionEngineConfig; -use datanode::store; use futures::stream; use mito2::access_layer::{ AccessLayer, AccessLayerRef, Metrics, OperationType, SstWriteRequest, WriteType, @@ -32,19 +30,20 @@ use mito2::sst::file::{FileHandle, FileMeta}; use mito2::sst::file_purger::{FilePurger, FilePurgerRef}; use mito2::sst::index::intermediate::IntermediateManager; use mito2::sst::index::puffin_manager::PuffinManagerFactory; +use mito2::sst::parquet::WriteOptions; use mito2::sst::parquet::reader::ParquetReaderBuilder; -use mito2::sst::parquet::{PARQUET_METADATA_KEY, WriteOptions}; use mito2::worker::write_cache_from_config; use object_store::ObjectStore; -use parquet::file::metadata::{FooterTail, KeyValue}; +use parquet::file::metadata::FooterTail; use regex::Regex; use snafu::OptionExt; -use store_api::metadata::{RegionMetadata, RegionMetadataRef}; use store_api::path_utils::region_name; use store_api::region_request::PathType; use store_api::storage::FileId; -use crate::datanode::{StorageConfig, StorageConfigWrapper}; +use crate::datanode::tool_util::{ + build_object_store, extract_region_metadata, max_row_group_uncompressed_size, parse_config, +}; use crate::error; /// Object storage benchmark command @@ -67,46 +66,6 @@ pub struct ObjbenchCommand { pub pprof_file: Option, } -pub(crate) fn parse_config( - config_path: &PathBuf, -) -> error::Result<( - StorageConfig, - MitoConfig, - common_wal::config::DatanodeWalConfig, -)> { - let cfg_str = std::fs::read_to_string(config_path).map_err(|e| { - error::IllegalConfigSnafu { - msg: format!("failed to read config {}: {e}", config_path.display()), - } - .build() - })?; - - let store_cfg: StorageConfigWrapper = toml::from_str(&cfg_str).map_err(|e| { - error::IllegalConfigSnafu { - msg: format!("failed to parse config {}: {e}", config_path.display()), - } - .build() - })?; - - let wal_config = store_cfg.wal; - let storage_config = store_cfg.storage; - let mito_engine_config = store_cfg - .region_engine - .into_iter() - .filter_map(|c| { - if let RegionEngineConfig::Mito(mito) = c { - Some(mito) - } else { - None - } - }) - .next() - .with_context(|| error::IllegalConfigSnafu { - msg: format!("Engine config not found in {:?}", config_path), - })?; - Ok((storage_config, mito_engine_config, wal_config)) -} - impl ObjbenchCommand { pub async fn run(&self) -> error::Result<()> { if self.verbose { @@ -154,17 +113,7 @@ impl ObjbenchCommand { let region_meta = extract_region_metadata(&self.source, &parquet_meta)?; let num_rows = parquet_meta.file_metadata().num_rows() as u64; let num_row_groups = parquet_meta.num_row_groups() as u64; - let max_row_group_uncompressed_size: u64 = parquet_meta - .row_groups() - .iter() - .map(|rg| { - rg.columns() - .iter() - .map(|c| c.uncompressed_size() as u64) - .sum::() - }) - .max() - .unwrap_or(0); + let max_row_group_uncompressed_size = max_row_group_uncompressed_size(&parquet_meta); println!( "{} Metadata loaded - rows: {}, size: {} bytes", @@ -485,47 +434,6 @@ fn parse_file_dir_components(path: &str) -> error::Result { }) } -pub(crate) fn extract_region_metadata( - file_path: &str, - meta: &parquet::file::metadata::ParquetMetaData, -) -> error::Result { - let kvs: Option<&Vec> = meta.file_metadata().key_value_metadata(); - let Some(kvs) = kvs else { - return Err(error::IllegalConfigSnafu { - msg: format!("{file_path}: missing parquet key_value metadata"), - } - .build()); - }; - let json = kvs - .iter() - .find(|kv| kv.key == PARQUET_METADATA_KEY) - .and_then(|kv| kv.value.as_ref()) - .ok_or_else(|| { - error::IllegalConfigSnafu { - msg: format!("{file_path}: key {PARQUET_METADATA_KEY} not found or empty"), - } - .build() - })?; - let region: RegionMetadata = RegionMetadata::from_json(json).map_err(|e| { - error::IllegalConfigSnafu { - msg: format!("invalid region metadata json: {e}"), - } - .build() - })?; - Ok(Arc::new(region)) -} - -pub(crate) async fn build_object_store(sc: &StorageConfig) -> error::Result { - store::new_object_store(sc.store.clone(), &sc.data_home) - .await - .map_err(|e| { - error::IllegalConfigSnafu { - msg: format!("Failed to build object store: {e:?}"), - } - .build() - }) -} - async fn build_access_layer_simple( components: &FileDirComponents, object_store: ObjectStore, @@ -666,7 +574,8 @@ mod tests { use common_base::readable_size::ReadableSize; use store_api::region_request::PathType; - use crate::datanode::objbench::{parse_config, parse_file_dir_components}; + use crate::datanode::objbench::parse_file_dir_components; + use crate::datanode::tool_util::parse_config; #[test] fn test_parse_dir() { diff --git a/src/cmd/src/datanode/parquet_meta.rs b/src/cmd/src/datanode/parquet_meta.rs new file mode 100644 index 0000000000..425da45e3c --- /dev/null +++ b/src/cmd/src/datanode/parquet_meta.rs @@ -0,0 +1,340 @@ +// 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::BTreeMap; +use std::path::{Path, PathBuf}; + +use clap::{Parser, ValueEnum}; +use parquet::file::metadata::ParquetMetaData; +use serde::Serialize; +use snafu::ResultExt; + +use crate::datanode::tool_util::{compression_name, load_local_parquet_metadata}; +use crate::error; + +/// Display parquet file metadata. +#[derive(Debug, Parser)] +pub struct ParquetMetaCommand { + /// Path to input parquet file. + #[clap(long, value_name = "FILE")] + input: PathBuf, + + /// Output format. + #[clap(long, value_enum, default_value = "text")] + format: MetaOutputFormat, +} + +#[derive(Debug, Clone, Copy, ValueEnum)] +enum MetaOutputFormat { + Text, + Json, +} + +#[derive(Debug, Serialize)] +struct FileMetaView { + input: String, + num_rows: i64, + num_row_groups: usize, + num_columns: usize, + key_value_metadata: BTreeMap>, + row_groups: Vec, +} + +#[derive(Debug, Serialize)] +struct RowGroupMetaView { + index: usize, + num_rows: i64, + uncompressed_size: i64, + compressed_size: i64, + compression_ratio: Option, + data_pages: Option, + dictionary_page_bytes: i64, + columns: Vec, +} + +#[derive(Debug, Serialize)] +struct ColumnChunkMetaView { + index: usize, + path: String, + physical_type: String, + encodings: Vec, + compression: String, + num_values: i64, + uncompressed_size: i64, + compressed_size: i64, + compression_ratio: Option, + data_page_offset: i64, + dictionary_page_offset: Option, + dictionary_page_bytes: Option, + data_pages: Option, + has_statistics: bool, + column_index_offset: Option, + column_index_length: Option, + offset_index_offset: Option, + offset_index_length: Option, + bloom_filter_offset: Option, + bloom_filter_length: Option, +} + +impl ParquetMetaCommand { + pub async fn run(&self) -> error::Result<()> { + let metadata = load_local_parquet_metadata(&self.input)?; + let view = build_file_meta_view(&self.input, &metadata); + + match self.format { + MetaOutputFormat::Text => print_meta_text(&view), + MetaOutputFormat::Json => { + let json = serde_json::to_string_pretty(&view).context(error::SerdeJsonSnafu)?; + println!("{json}"); + } + } + + Ok(()) + } +} + +fn build_file_meta_view(path: &Path, metadata: &ParquetMetaData) -> FileMetaView { + let key_value_metadata = metadata + .file_metadata() + .key_value_metadata() + .into_iter() + .flatten() + .map(|kv| (kv.key.clone(), kv.value.clone())) + .collect(); + let offset_index = metadata.offset_index(); + let row_groups = metadata + .row_groups() + .iter() + .enumerate() + .map(|(row_group_idx, row_group)| { + let columns: Vec<_> = row_group + .columns() + .iter() + .enumerate() + .map(|(column_idx, column)| { + let data_pages = offset_index + .and_then(|index| index.get(row_group_idx)) + .and_then(|columns| columns.get(column_idx)) + .map(|index| index.page_locations().len()); + let dictionary_page_bytes = dictionary_page_bytes( + column.dictionary_page_offset(), + column.data_page_offset(), + ); + + ColumnChunkMetaView { + index: column_idx, + path: column.column_path().string(), + physical_type: format!("{:?}", column.column_type()), + encodings: column.encodings().map(|enc| format!("{enc:?}")).collect(), + compression: compression_name(column.compression()).to_string(), + num_values: column.num_values(), + uncompressed_size: column.uncompressed_size(), + compressed_size: column.compressed_size(), + compression_ratio: compression_ratio( + column.uncompressed_size(), + column.compressed_size(), + ), + data_page_offset: column.data_page_offset(), + dictionary_page_offset: column.dictionary_page_offset(), + dictionary_page_bytes, + data_pages, + has_statistics: column.statistics().is_some(), + column_index_offset: column.column_index_offset(), + column_index_length: column.column_index_length(), + offset_index_offset: column.offset_index_offset(), + offset_index_length: column.offset_index_length(), + bloom_filter_offset: column.bloom_filter_offset(), + bloom_filter_length: column.bloom_filter_length(), + } + }) + .collect(); + let data_pages = if columns.iter().all(|column| column.data_pages.is_some()) { + Some( + columns + .iter() + .map(|column| column.data_pages.unwrap_or_default()) + .sum(), + ) + } else { + None + }; + let dictionary_page_bytes = columns + .iter() + .map(|column| column.dictionary_page_bytes.unwrap_or_default()) + .sum(); + + RowGroupMetaView { + index: row_group_idx, + num_rows: row_group.num_rows(), + uncompressed_size: row_group.total_byte_size(), + compressed_size: row_group.compressed_size(), + compression_ratio: compression_ratio( + row_group.total_byte_size(), + row_group.compressed_size(), + ), + data_pages, + dictionary_page_bytes, + columns, + } + }) + .collect(); + + FileMetaView { + input: path.display().to_string(), + num_rows: metadata.file_metadata().num_rows(), + num_row_groups: metadata.num_row_groups(), + num_columns: metadata.file_metadata().schema_descr().num_columns(), + key_value_metadata, + row_groups, + } +} + +fn print_meta_text(view: &FileMetaView) { + println!("file: {}", view.input); + println!("rows: {}", view.num_rows); + println!("row_groups: {}", view.num_row_groups); + println!("columns: {}", view.num_columns); + println!("key_value_metadata: {}", view.key_value_metadata.len()); + for row_group in &view.row_groups { + println!( + "row_group[{}]: rows={}, uncompressed_size={}, compressed_size={}, compression_ratio={}, data_pages={}, dictionary_page_bytes={}", + row_group.index, + row_group.num_rows, + row_group.uncompressed_size, + row_group.compressed_size, + format_ratio(row_group.compression_ratio), + format_optional_usize(row_group.data_pages), + row_group.dictionary_page_bytes, + ); + for column in &row_group.columns { + println!( + " column[{}] path={}: type={}, compression={}, encodings=[{}], values={}, uncompressed_size={}, compressed_size={}, compression_ratio={}, data_pages={}, dictionary_page_offset={}, dictionary_page_bytes={}, data_page_offset={}, statistics={}, column_index={}/{}, offset_index={}/{}, bloom_filter={}/{}", + column.index, + column.path, + column.physical_type, + column.compression, + column.encodings.join(","), + column.num_values, + column.uncompressed_size, + column.compressed_size, + format_ratio(column.compression_ratio), + format_optional_usize(column.data_pages), + format_optional_i64(column.dictionary_page_offset), + column + .dictionary_page_bytes + .map(|value| value.to_string()) + .unwrap_or_else(|| "none".to_string()), + column.data_page_offset, + column.has_statistics, + format_optional_i64(column.column_index_offset), + format_optional_i32(column.column_index_length), + format_optional_i64(column.offset_index_offset), + format_optional_i32(column.offset_index_length), + format_optional_i64(column.bloom_filter_offset), + format_optional_i32(column.bloom_filter_length), + ); + } + } +} + +fn dictionary_page_bytes( + dictionary_page_offset: Option, + data_page_offset: i64, +) -> Option { + dictionary_page_offset.map(|offset| data_page_offset.saturating_sub(offset)) +} + +fn compression_ratio(uncompressed: i64, compressed: i64) -> Option { + (uncompressed > 0).then(|| compressed as f64 / uncompressed as f64) +} + +fn format_ratio(ratio: Option) -> String { + ratio + .map(|value| format!("{value:.4}")) + .unwrap_or_else(|| "unknown".to_string()) +} + +fn format_optional_usize(value: Option) -> String { + value + .map(|value| value.to_string()) + .unwrap_or_else(|| "unknown".to_string()) +} + +fn format_optional_i64(value: Option) -> String { + value + .map(|value| value.to_string()) + .unwrap_or_else(|| "none".to_string()) +} + +fn format_optional_i32(value: Option) -> String { + value + .map(|value| value.to_string()) + .unwrap_or_else(|| "none".to_string()) +} + +#[cfg(test)] +mod tests { + use std::fs::File; + use std::sync::Arc; + + use datatypes::arrow::array::{Int32Array, RecordBatch, StringArray}; + use datatypes::arrow::datatypes::{DataType, Field, Schema}; + use parquet::arrow::ArrowWriter; + use parquet::file::metadata::KeyValue; + use parquet::file::properties::WriterProperties; + use tempfile::tempdir; + + use super::*; + + #[test] + fn test_meta_view_unknown_page_count() { + let dir = tempdir().unwrap(); + let path = dir.path().join("input.parquet"); + write_test_parquet(&path); + + let metadata = load_local_parquet_metadata(&path).unwrap(); + let view = build_file_meta_view(&path, &metadata); + + assert_eq!(view.num_row_groups, 1); + assert!(view.row_groups[0].uncompressed_size > 0); + assert!(view.row_groups[0].compressed_size > 0); + assert!(view.row_groups[0].compression_ratio.is_some()); + assert_eq!(view.row_groups[0].columns.len(), 2); + } + + fn write_test_parquet(path: &Path) { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("host", DataType::Utf8, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3, 4])), + Arc::new(StringArray::from(vec!["a", "a", "b", "b"])), + ], + ) + .unwrap(); + let props = WriterProperties::builder() + .set_key_value_metadata(Some(vec![KeyValue::new( + "greptime:test".to_string(), + "value".to_string(), + )])) + .build(); + let mut writer = + ArrowWriter::try_new(File::create(path).unwrap(), schema, Some(props)).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + } +} diff --git a/src/cmd/src/datanode/parquet_rewrite.rs b/src/cmd/src/datanode/parquet_rewrite.rs new file mode 100644 index 0000000000..47346aac88 --- /dev/null +++ b/src/cmd/src/datanode/parquet_rewrite.rs @@ -0,0 +1,579 @@ +// 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::fs::File; +use std::path::{Path, PathBuf}; + +use clap::Parser; +use datatypes::arrow::record_batch::RecordBatchReader; +use parquet::arrow::ArrowWriter; +use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; +use parquet::basic::{BrotliLevel, Compression, Encoding, GzipLevel, ZstdLevel}; +use parquet::file::metadata::ParquetMetaData; +use parquet::file::properties::WriterProperties; +use parquet::file::reader::FileReader; +use parquet::file::serialized_reader::SerializedFileReader; +use parquet::schema::types::ColumnPath; +use serde::{Deserialize, Serialize}; +use snafu::ResultExt; + +use crate::datanode::tool_util::load_local_parquet_metadata; +use crate::error; + +/// Read and rewrite a parquet file with different writer properties. +#[derive(Debug, Parser)] +pub struct ParquetRewriteCommand { + /// Path to input parquet file. + #[clap(long, value_name = "FILE")] + input: PathBuf, + + /// Path to output parquet file in rewrite mode. + #[clap(long, value_name = "FILE")] + output: Option, + + /// Path to writer properties TOML in rewrite mode. + #[clap(long, value_name = "FILE")] + properties: Option, + + /// Dump writer properties TOML inferred from the input parquet file. + #[clap(long, value_name = "FILE")] + dump_properties: Option, + + /// Number of rows per record batch. + #[clap(long)] + batch_size: Option, + + /// Overwrite output files. + #[clap(long, default_value_t = false)] + overwrite: bool, +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)] +#[serde(default, deny_unknown_fields)] +struct RewriteProperties { + writer: WriterConfig, + columns: Vec, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)] +#[serde(default, deny_unknown_fields)] +struct WriterConfig { + dictionary_enabled: Option, + compression: Option, + compression_level: Option, + encoding: Option, + max_row_group_row_count: Option, + data_page_size_limit: Option, + data_page_row_count_limit: Option, + dictionary_page_size_limit: Option, +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)] +#[serde(default, deny_unknown_fields)] +struct ColumnConfig { + path: Vec, + dictionary_enabled: Option, + compression: Option, + compression_level: Option, + encoding: Option, +} + +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "kebab-case")] +enum CompressionConfig { + Uncompressed, + Snappy, + Gzip, + Lzo, + Brotli, + Lz4, + Zstd, + Lz4Raw, +} + +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "kebab-case")] +enum EncodingConfig { + Plain, + DeltaBinaryPacked, + DeltaLengthByteArray, + DeltaByteArray, + ByteStreamSplit, +} + +impl ParquetRewriteCommand { + pub async fn run(&self) -> error::Result<()> { + match (&self.dump_properties, &self.output, &self.properties) { + (Some(path), None, None) => self.dump_properties(path), + (None, Some(output), Some(properties)) => self.rewrite(output, properties), + (Some(_), Some(_), _) | (Some(_), _, Some(_)) => illegal_config( + "use either --dump-properties or rewrite mode, not both".to_string(), + ), + (None, _, _) => illegal_config( + "rewrite mode requires --output and --properties; config dump mode requires --dump-properties".to_string(), + ), + } + } + + fn dump_properties(&self, path: &Path) -> error::Result<()> { + ensure_distinct_paths(&self.input, path)?; + ensure_can_write(path, self.overwrite)?; + let metadata = load_local_parquet_metadata(&self.input)?; + let properties = infer_rewrite_properties(&metadata); + let content = toml::to_string_pretty(&properties).map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("failed to serialize writer properties: {e}"), + } + .build() + })?; + std::fs::write(path, content).context(error::FileIoSnafu)?; + println!("Wrote writer properties to {}", path.display()); + Ok(()) + } + + fn rewrite(&self, output: &Path, properties: &Path) -> error::Result<()> { + ensure_distinct_paths(&self.input, output)?; + ensure_can_write(output, self.overwrite)?; + + let props = load_rewrite_properties(properties)?; + let input_reader = SerializedFileReader::new(open_file(&self.input)?) + .map_err(|e| parquet_error("read source parquet metadata", &self.input, e))?; + let key_value_metadata = input_reader + .metadata() + .file_metadata() + .key_value_metadata() + .cloned(); + + let mut reader_builder = ParquetRecordBatchReaderBuilder::try_new(open_file(&self.input)?) + .map_err(|e| parquet_error("open source parquet", &self.input, e))?; + if let Some(batch_size) = self.batch_size { + if batch_size == 0 { + return illegal_config("--batch-size must be greater than 0".to_string()); + } + reader_builder = reader_builder.with_batch_size(batch_size); + } + let reader = reader_builder + .build() + .map_err(|e| parquet_error("build parquet reader", &self.input, e))?; + let schema = reader.schema(); + + let writer_props = build_writer_properties(props, key_value_metadata)?; + let mut writer = ArrowWriter::try_new(create_file(output)?, schema, Some(writer_props)) + .map_err(|e| parquet_error("create parquet writer", output, e))?; + + for batch in reader { + let batch = + batch.map_err(|e| parquet_error("read parquet batch", &self.input, e.into()))?; + writer + .write(&batch) + .map_err(|e| parquet_error("write parquet batch", output, e))?; + } + writer + .close() + .map_err(|e| parquet_error("close parquet writer", output, e))?; + + println!("Wrote parquet file to {}", output.display()); + Ok(()) + } +} + +fn infer_rewrite_properties(metadata: &ParquetMetaData) -> RewriteProperties { + let first_column = metadata + .row_groups() + .first() + .and_then(|row_group| row_group.columns().first()); + let compression = first_column.map(|column| compression_to_config(column.compression())); + let max_row_group_row_count = metadata + .row_groups() + .first() + .and_then(|row_group| usize::try_from(row_group.num_rows()).ok()); + let dictionary_enabled = first_column + .map(|column| column.dictionary_page_offset().is_some()) + .or(Some(true)); + + let columns = metadata + .file_metadata() + .schema_descr() + .columns() + .iter() + .enumerate() + .map(|(idx, column)| { + let first_chunk = metadata + .row_groups() + .first() + .and_then(|row_group| row_group.columns().get(idx)); + ColumnConfig { + path: column.path().parts().to_vec(), + dictionary_enabled: first_chunk + .map(|chunk| chunk.dictionary_page_offset().is_some()), + compression: first_chunk.map(|chunk| compression_to_config(chunk.compression())), + compression_level: None, + encoding: first_chunk.and_then(infer_data_encoding), + } + }) + .collect(); + + RewriteProperties { + writer: WriterConfig { + dictionary_enabled, + compression, + compression_level: None, + encoding: first_column.and_then(infer_data_encoding), + max_row_group_row_count, + data_page_size_limit: None, + data_page_row_count_limit: None, + dictionary_page_size_limit: None, + }, + columns, + } +} + +fn load_rewrite_properties(path: &Path) -> error::Result { + let content = std::fs::read_to_string(path).context(error::FileIoSnafu)?; + toml::from_str(&content).map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("failed to parse writer properties {}: {e}", path.display()), + } + .build() + }) +} + +fn build_writer_properties( + config: RewriteProperties, + key_value_metadata: Option>, +) -> error::Result { + let mut builder = WriterProperties::builder().set_key_value_metadata(key_value_metadata); + let writer_compression = config.writer.compression; + if let Some(dictionary_enabled) = config.writer.dictionary_enabled { + builder = builder.set_dictionary_enabled(dictionary_enabled); + } + match (writer_compression, config.writer.compression_level) { + (Some(compression), level) => { + builder = builder.set_compression(to_parquet_compression(compression, level)?); + } + (None, Some(_)) => { + return illegal_config( + "writer compression_level requires writer compression".to_string(), + ); + } + (None, None) => {} + } + if let Some(encoding) = config.writer.encoding { + builder = builder.set_encoding(to_parquet_encoding(encoding)); + } + if let Some(max_row_group_row_count) = config.writer.max_row_group_row_count { + if max_row_group_row_count == 0 { + return illegal_config("max_row_group_row_count must be greater than 0".to_string()); + } + builder = builder.set_max_row_group_row_count(Some(max_row_group_row_count)); + } + if let Some(data_page_size_limit) = config.writer.data_page_size_limit { + builder = builder.set_data_page_size_limit(data_page_size_limit); + } + if let Some(data_page_row_count_limit) = config.writer.data_page_row_count_limit { + builder = builder.set_data_page_row_count_limit(data_page_row_count_limit); + } + if let Some(dictionary_page_size_limit) = config.writer.dictionary_page_size_limit { + builder = builder.set_dictionary_page_size_limit(dictionary_page_size_limit); + } + + for column in config.columns { + if column.path.is_empty() { + return illegal_config("column path must not be empty".to_string()); + } + let path = ColumnPath::new(column.path); + if let Some(dictionary_enabled) = column.dictionary_enabled { + builder = builder.set_column_dictionary_enabled(path.clone(), dictionary_enabled); + } + if column.compression.is_some() || column.compression_level.is_some() { + let Some(compression) = column.compression.or(writer_compression) else { + return illegal_config(format!( + "compression_level for column {} requires column or writer compression", + path.string() + )); + }; + builder = builder.set_column_compression( + path.clone(), + to_parquet_compression(compression, column.compression_level)?, + ); + } + if let Some(encoding) = column.encoding { + builder = builder.set_column_encoding(path, to_parquet_encoding(encoding)); + } + } + + Ok(builder.build()) +} + +fn ensure_distinct_paths(input: &Path, output: &Path) -> error::Result<()> { + if input == output { + return illegal_config(format!( + "input and output paths must be different: {}", + input.display() + )); + } + Ok(()) +} + +fn ensure_can_write(path: &Path, overwrite: bool) -> error::Result<()> { + if !overwrite && path.exists() { + return illegal_config(format!( + "{} already exists; pass --overwrite to replace it", + path.display() + )); + } + Ok(()) +} + +fn open_file(path: &Path) -> error::Result { + File::open(path).context(error::FileIoSnafu) +} + +fn create_file(path: &Path) -> error::Result { + File::create(path).context(error::FileIoSnafu) +} + +fn illegal_config(msg: String) -> error::Result { + error::IllegalConfigSnafu { msg }.fail() +} + +fn parquet_error( + action: &'static str, + path: &Path, + error: parquet::errors::ParquetError, +) -> error::Error { + error::IllegalConfigSnafu { + msg: format!("{action} failed for {}: {error}", path.display()), + } + .build() +} + +fn to_parquet_compression( + compression: CompressionConfig, + level: Option, +) -> error::Result { + if level.is_some() + && !matches!( + compression, + CompressionConfig::Gzip | CompressionConfig::Brotli | CompressionConfig::Zstd + ) + { + return illegal_config(format!( + "compression level is not supported for {}", + compression.name() + )); + } + + Ok(match compression { + CompressionConfig::Uncompressed => Compression::UNCOMPRESSED, + CompressionConfig::Snappy => Compression::SNAPPY, + CompressionConfig::Gzip => Compression::GZIP(match level { + Some(level) => GzipLevel::try_new(level).map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("invalid gzip compression level {level}: {e}"), + } + .build() + })?, + None => GzipLevel::default(), + }), + CompressionConfig::Lzo => Compression::LZO, + CompressionConfig::Brotli => Compression::BROTLI(match level { + Some(level) => BrotliLevel::try_new(level).map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("invalid brotli compression level {level}: {e}"), + } + .build() + })?, + None => BrotliLevel::default(), + }), + CompressionConfig::Lz4 => Compression::LZ4, + CompressionConfig::Zstd => Compression::ZSTD(match level { + Some(level) => { + let converted_level = i32::try_from(level).map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("invalid zstd compression level {level}: {e}"), + } + .build() + })?; + ZstdLevel::try_new(converted_level).map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("invalid zstd compression level {level}: {e}"), + } + .build() + })? + } + None => ZstdLevel::default(), + }), + CompressionConfig::Lz4Raw => Compression::LZ4_RAW, + }) +} + +impl CompressionConfig { + fn name(self) -> &'static str { + match self { + CompressionConfig::Uncompressed => "uncompressed", + CompressionConfig::Snappy => "snappy", + CompressionConfig::Gzip => "gzip", + CompressionConfig::Lzo => "lzo", + CompressionConfig::Brotli => "brotli", + CompressionConfig::Lz4 => "lz4", + CompressionConfig::Zstd => "zstd", + CompressionConfig::Lz4Raw => "lz4-raw", + } + } +} + +fn to_parquet_encoding(encoding: EncodingConfig) -> Encoding { + match encoding { + EncodingConfig::Plain => Encoding::PLAIN, + EncodingConfig::DeltaBinaryPacked => Encoding::DELTA_BINARY_PACKED, + EncodingConfig::DeltaLengthByteArray => Encoding::DELTA_LENGTH_BYTE_ARRAY, + EncodingConfig::DeltaByteArray => Encoding::DELTA_BYTE_ARRAY, + EncodingConfig::ByteStreamSplit => Encoding::BYTE_STREAM_SPLIT, + } +} + +fn infer_data_encoding( + column: &parquet::file::metadata::ColumnChunkMetaData, +) -> Option { + let encodings: Vec<_> = column.encodings().collect(); + if encodings.contains(&Encoding::DELTA_BINARY_PACKED) { + Some(EncodingConfig::DeltaBinaryPacked) + } else if encodings.contains(&Encoding::DELTA_LENGTH_BYTE_ARRAY) { + Some(EncodingConfig::DeltaLengthByteArray) + } else if encodings.contains(&Encoding::DELTA_BYTE_ARRAY) { + Some(EncodingConfig::DeltaByteArray) + } else if encodings.contains(&Encoding::BYTE_STREAM_SPLIT) { + Some(EncodingConfig::ByteStreamSplit) + } else if encodings.contains(&Encoding::PLAIN) { + Some(EncodingConfig::Plain) + } else { + None + } +} + +fn compression_to_config(compression: Compression) -> CompressionConfig { + match compression { + Compression::UNCOMPRESSED => CompressionConfig::Uncompressed, + Compression::SNAPPY => CompressionConfig::Snappy, + Compression::GZIP(_) => CompressionConfig::Gzip, + Compression::LZO => CompressionConfig::Lzo, + Compression::BROTLI(_) => CompressionConfig::Brotli, + Compression::LZ4 => CompressionConfig::Lz4, + Compression::ZSTD(_) => CompressionConfig::Zstd, + Compression::LZ4_RAW => CompressionConfig::Lz4Raw, + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use datatypes::arrow::array::{Int32Array, RecordBatch, StringArray}; + use datatypes::arrow::datatypes::{DataType, Field, Schema}; + use parquet::file::metadata::KeyValue; + use tempfile::tempdir; + + use super::*; + + #[test] + fn test_dump_and_rewrite_preserves_key_value_metadata_and_disables_dictionary() { + let dir = tempdir().unwrap(); + let input = dir.path().join("input.parquet"); + let output = dir.path().join("output.parquet"); + write_test_parquet(&input, true); + + let metadata = load_local_parquet_metadata(&input).unwrap(); + let mut properties = infer_rewrite_properties(&metadata); + properties.columns[1].dictionary_enabled = Some(false); + let props_path = dir.path().join("props.toml"); + std::fs::write(&props_path, toml::to_string(&properties).unwrap()).unwrap(); + + let command = ParquetRewriteCommand { + input: input.clone(), + output: Some(output.clone()), + properties: Some(props_path), + dump_properties: None, + batch_size: Some(2), + overwrite: false, + }; + command + .rewrite(&output, command.properties.as_ref().unwrap()) + .unwrap(); + + let rewritten = load_local_parquet_metadata(&output).unwrap(); + assert_eq!(rewritten.file_metadata().num_rows(), 4); + let key_values = rewritten + .file_metadata() + .key_value_metadata() + .cloned() + .unwrap_or_default(); + assert!(key_values.iter().any(|kv| kv.key == "greptime:test")); + assert!( + rewritten + .row_groups() + .iter() + .all(|row_group| row_group.column(1).dictionary_page_offset().is_none()) + ); + } + + #[test] + fn test_column_compression_level_inherits_writer_compression() { + let path = ColumnPath::new(vec!["host".to_string()]); + let config = RewriteProperties { + writer: WriterConfig { + compression: Some(CompressionConfig::Zstd), + compression_level: Some(1), + ..Default::default() + }, + columns: vec![ColumnConfig { + path: vec!["host".to_string()], + compression_level: Some(3), + ..Default::default() + }], + }; + + let properties = build_writer_properties(config, None).unwrap(); + assert_eq!( + properties.compression(&path), + Compression::ZSTD(ZstdLevel::try_new(3).unwrap()) + ); + } + + fn write_test_parquet(path: &Path, dictionary_enabled: bool) { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("host", DataType::Utf8, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3, 4])), + Arc::new(StringArray::from(vec!["a", "a", "b", "b"])), + ], + ) + .unwrap(); + let props = WriterProperties::builder() + .set_dictionary_enabled(dictionary_enabled) + .set_key_value_metadata(Some(vec![KeyValue::new( + "greptime:test".to_string(), + "value".to_string(), + )])) + .build(); + let mut writer = + ArrowWriter::try_new(File::create(path).unwrap(), schema, Some(props)).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + } +} diff --git a/src/cmd/src/datanode/parquetbench.rs b/src/cmd/src/datanode/parquetbench.rs index 2b46a775f7..09ec50f48e 100644 --- a/src/cmd/src/datanode/parquetbench.rs +++ b/src/cmd/src/datanode/parquetbench.rs @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::path::PathBuf; +use std::path::{Path, PathBuf}; use std::sync::Arc; use std::time::{Duration, Instant}; @@ -42,7 +42,10 @@ use store_api::region_request::PathType; use store_api::storage::consts::{PRIMARY_KEY_COLUMN_NAME, is_internal_column}; use store_api::storage::{ColumnId, FileId, RegionId}; -use crate::datanode::objbench::{build_object_store, extract_region_metadata, parse_config}; +use crate::datanode::tool_util::{ + build_object_store, extract_region_metadata, format_bytes, parse_config, parse_file_id, + parse_path_type, parse_region_id, +}; use crate::error; const DEFAULT_READ_BATCH_SIZE: usize = 8 * 1024; @@ -52,19 +55,23 @@ const DEFAULT_READ_BATCH_SIZE: usize = 8 * 1024; pub struct ParquetbenchCommand { /// Path to config TOML file (same format as standalone/datanode config) #[clap(long, value_name = "FILE")] - config: PathBuf, + config: Option, /// Region ID: either numeric u64 (e.g. "4398046511104") or "table_id:region_num" (e.g. "1024:0") #[clap(long)] - region_id: String, + region_id: Option, /// Table directory relative to data home (e.g. "data/greptime/public/1024/") #[clap(long)] - table_dir: String, + table_dir: Option, /// SST file id to benchmark. #[clap(long)] - file_id: String, + file_id: Option, + + /// Local parquet SST file to benchmark with the direct reader. + #[clap(long, value_name = "FILE")] + file_path: Option, /// Path to scan request JSON config file (supports projection_names only) #[clap(long, value_name = "FILE")] @@ -129,6 +136,32 @@ struct IterationStats { elapsed: Duration, } +#[derive(Debug)] +enum ParquetbenchInput { + LocalFile { + file_path: PathBuf, + }, + Region { + config: PathBuf, + region_id: String, + table_dir: String, + file_id: String, + path_type: PathType, + }, +} + +struct ParquetbenchSource { + object_store: object_store::ObjectStore, + file_path: String, + display_path: String, + region_id_label: String, + region_id: RegionId, + file_id: FileId, + region_file_id: RegionFileId, + path_type: Option, + table_dir: Option, +} + impl ParquetbenchCommand { pub async fn run(&self) -> error::Result<()> { if self.verbose { @@ -145,41 +178,40 @@ impl ParquetbenchCommand { .fail(); } - let region_id = parse_region_id(&self.region_id)?; - let path_type = parse_path_type(&self.path_type)?; - let file_id = FileId::parse_str(&self.file_id).map_err(|e| { - error::IllegalConfigSnafu { - msg: format!("invalid file_id '{}': {}", self.file_id, e), - } - .build() - })?; - let region_file_id = RegionFileId::new(region_id, file_id); - let file_path = sst_file_path(&self.table_dir, region_file_id, path_type); + let input = self.resolve_input()?; + let mut source = build_source(input).await?; - let (store_cfg, _mito_config, _wal_config) = parse_config(&self.config)?; - let object_store = build_object_store(&store_cfg).await?; - - let file_size = object_store - .stat(&file_path) + let file_size = source + .object_store + .stat(&source.file_path) .await .map_err(|e| { error::IllegalConfigSnafu { - msg: format!("stat failed for {}: {}", file_path, e), + msg: format!("stat failed for {}: {}", source.display_path, e), } .build() })? .content_length(); let mut metadata_metrics = MetadataCacheMetrics::default(); - let parquet_meta = MetadataLoader::new(object_store.clone(), &file_path, file_size) - .load(&mut metadata_metrics) - .await - .map_err(|e| { - error::IllegalConfigSnafu { - msg: format!("read parquet metadata failed for {}: {:?}", file_path, e), - } - .build() - })?; - let region_meta = extract_region_metadata(&file_path, &parquet_meta)?; + let parquet_meta = + MetadataLoader::new(source.object_store.clone(), &source.file_path, file_size) + .load(&mut metadata_metrics) + .await + .map_err(|e| { + error::IllegalConfigSnafu { + msg: format!( + "read parquet metadata failed for {}: {:?}", + source.display_path, e + ), + } + .build() + })?; + let region_meta = extract_region_metadata(&source.display_path, &parquet_meta)?; + if source.table_dir.is_none() { + source.region_id = region_meta.region_id; + source.region_id_label = source.region_id.as_u64().to_string(); + source.region_file_id = RegionFileId::new(source.region_id, source.file_id); + } let scan_config = self.load_scan_config().await?; let projection = if self.reader == ReaderMode::Direct { resolve_projection_names(&scan_config, ®ion_meta)? @@ -216,10 +248,10 @@ impl ParquetbenchCommand { println!( "{} Region ID: {} (u64: {})", "✓".green(), - self.region_id, - region_id.as_u64() + source.region_id_label, + source.region_id.as_u64() ); - println!("{} File path: {}", "✓".green(), file_path.cyan()); + println!("{} File path: {}", "✓".green(), source.display_path.cyan()); println!( "{} Columns: {}", "✓".green(), @@ -307,8 +339,8 @@ impl ParquetbenchCommand { let mut schema_printed = false; let file_handle = FileHandle::new( FileMeta { - region_id, - file_id, + region_id: source.region_id, + file_id: source.file_id, time_range: Default::default(), level: 0, file_size, @@ -332,9 +364,9 @@ impl ParquetbenchCommand { let stats = match self.reader { ReaderMode::Direct => { run_direct_iteration( - object_store.clone(), - file_path.clone(), - region_file_id, + source.object_store.clone(), + source.file_path.clone(), + source.region_file_id, parquet_meta.clone(), projection.clone(), row_groups.clone(), @@ -345,9 +377,19 @@ impl ParquetbenchCommand { } ReaderMode::FlatPrune => { run_flat_prune_iteration( - object_store.clone(), - self.table_dir.clone(), - path_type, + source.object_store.clone(), + source.table_dir.clone().ok_or_else(|| { + error::IllegalConfigSnafu { + msg: "flat-prune reader requires --table-dir".to_string(), + } + .build() + })?, + source.path_type.ok_or_else(|| { + error::IllegalConfigSnafu { + msg: "flat-prune reader requires --path-type".to_string(), + } + .build() + })?, file_handle.clone(), region_meta.clone(), projection_column_ids.clone(), @@ -450,6 +492,65 @@ impl ParquetbenchCommand { Ok(()) } + fn resolve_input(&self) -> error::Result { + let has_region_args = self.config.is_some() + || self.region_id.is_some() + || self.table_dir.is_some() + || self.file_id.is_some(); + + if let Some(file_path) = &self.file_path { + if self.reader == ReaderMode::FlatPrune { + return Err(error::IllegalConfigSnafu { + msg: "--file-path currently supports only --reader direct".to_string(), + } + .build()); + } + if has_region_args { + return Err(error::IllegalConfigSnafu { + msg: "--file-path cannot be used with --config, --region-id, --table-dir, or --file-id".to_string(), + } + .build()); + } + return Ok(ParquetbenchInput::LocalFile { + file_path: file_path.clone(), + }); + } + + let config = self.config.clone().ok_or_else(|| { + error::IllegalConfigSnafu { + msg: "missing --config unless --file-path is specified".to_string(), + } + .build() + })?; + let region_id = self.region_id.clone().ok_or_else(|| { + error::IllegalConfigSnafu { + msg: "missing --region-id unless --file-path is specified".to_string(), + } + .build() + })?; + let table_dir = self.table_dir.clone().ok_or_else(|| { + error::IllegalConfigSnafu { + msg: "missing --table-dir unless --file-path is specified".to_string(), + } + .build() + })?; + let file_id = self.file_id.clone().ok_or_else(|| { + error::IllegalConfigSnafu { + msg: "missing --file-id unless --file-path is specified".to_string(), + } + .build() + })?; + let path_type = parse_path_type(&self.path_type)?; + + Ok(ParquetbenchInput::Region { + config, + region_id, + table_dir, + file_id, + path_type, + }) + } + async fn load_scan_config(&self) -> error::Result { if let Some(path) = &self.scan_config { let content = tokio::fs::read_to_string(path) @@ -462,6 +563,117 @@ impl ParquetbenchCommand { } } +async fn build_source(input: ParquetbenchInput) -> error::Result { + match input { + ParquetbenchInput::LocalFile { file_path } => build_local_file_source(&file_path), + ParquetbenchInput::Region { + config, + region_id, + table_dir, + file_id, + path_type, + } => build_region_source(&config, ®ion_id, table_dir, &file_id, path_type).await, + } +} + +fn build_local_file_source(file_path: &Path) -> error::Result { + let file_path = std::fs::canonicalize(file_path).map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("invalid --file-path {}: {e}", file_path.display()), + } + .build() + })?; + if !file_path.is_file() { + return Err(error::IllegalConfigSnafu { + msg: format!("--file-path {} is not a file", file_path.display()), + } + .build()); + } + let parent = file_path.parent().ok_or_else(|| { + error::IllegalConfigSnafu { + msg: format!( + "--file-path {} has no parent directory", + file_path.display() + ), + } + .build() + })?; + let file_name = file_path + .file_name() + .and_then(|name| name.to_str()) + .ok_or_else(|| { + error::IllegalConfigSnafu { + msg: format!("invalid UTF-8 file name in {}", file_path.display()), + } + .build() + })? + .to_string(); + let object_store = object_store::ObjectStore::new(object_store::services::Fs::default().root( + parent.to_str().ok_or_else(|| { + error::IllegalConfigSnafu { + msg: format!("invalid UTF-8 parent directory in {}", file_path.display()), + } + .build() + })?, + )) + .map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("failed to build local file object store: {e:?}"), + } + .build() + })? + .finish(); + + let file_id = file_path + .file_stem() + .and_then(|stem| stem.to_str()) + .and_then(|stem| FileId::parse_str(stem).ok()) + .unwrap_or_else(FileId::random); + let region_id = RegionId::new(0, 0); + let region_file_id = RegionFileId::new(region_id, file_id); + let display_path = file_path.display().to_string(); + + Ok(ParquetbenchSource { + object_store, + file_path: file_name, + display_path, + region_id_label: region_id.as_u64().to_string(), + region_id, + file_id, + region_file_id, + path_type: None, + table_dir: None, + }) +} + +async fn build_region_source( + config: &Path, + region_id: &str, + table_dir: String, + file_id: &str, + path_type: PathType, +) -> error::Result { + let region = parse_region_id(region_id)?; + let file_id = parse_file_id(file_id)?; + let region_file_id = RegionFileId::new(region, file_id); + let file_path = sst_file_path(&table_dir, region_file_id, path_type); + + let (store_cfg, _mito_config, _wal_config) = parse_config(config)?; + let object_store = build_object_store(&store_cfg).await?; + + Ok(ParquetbenchSource { + object_store, + display_path: file_path.clone(), + file_path, + region_id_label: region_id.to_string(), + region_id: region, + file_id, + region_file_id, + path_type: Some(path_type), + table_dir: Some(table_dir), + }) +} + #[allow(clippy::too_many_arguments)] async fn run_direct_iteration( object_store: object_store::ObjectStore, @@ -773,60 +985,6 @@ fn parse_batch_size(s: &str) -> Result { Ok(batch_size) } -fn parse_region_id(s: &str) -> error::Result { - if s.contains(':') { - let parts: Vec<&str> = s.splitn(2, ':').collect(); - let table_id: u32 = parts[0].parse().map_err(|e| { - error::IllegalConfigSnafu { - msg: format!("invalid table_id in region_id '{}': {}", s, e), - } - .build() - })?; - let region_num: u32 = parts[1].parse().map_err(|e| { - error::IllegalConfigSnafu { - msg: format!("invalid region_num in region_id '{}': {}", s, e), - } - .build() - })?; - Ok(RegionId::new(table_id, region_num)) - } else { - let id: u64 = s.parse().map_err(|e| { - error::IllegalConfigSnafu { - msg: format!("invalid region_id '{}': {}", s, e), - } - .build() - })?; - Ok(RegionId::from_u64(id)) - } -} - -fn parse_path_type(s: &str) -> error::Result { - match s.to_lowercase().as_str() { - "bare" => Ok(PathType::Bare), - "data" => Ok(PathType::Data), - "metadata" => Ok(PathType::Metadata), - _ => Err(error::IllegalConfigSnafu { - msg: format!("invalid path_type '{}', expected: bare, data, metadata", s), - } - .build()), - } -} - -fn format_bytes(bytes: u64) -> String { - const KIB: u64 = 1024; - const MIB: u64 = 1024 * KIB; - const GIB: u64 = 1024 * MIB; - if bytes >= GIB { - format!("{:.2} GiB", bytes as f64 / GIB as f64) - } else if bytes >= MIB { - format!("{:.2} MiB", bytes as f64 / MIB as f64) - } else if bytes >= KIB { - format!("{:.2} KiB", bytes as f64 / KIB as f64) - } else { - format!("{} B", bytes) - } -} - fn format_rate(rate: f64) -> String { if !rate.is_finite() { return "inf rows".to_string(); @@ -851,6 +1009,25 @@ mod tests { use super::*; + fn test_command() -> ParquetbenchCommand { + ParquetbenchCommand { + config: None, + region_id: None, + table_dir: None, + file_id: None, + file_path: None, + scan_config: None, + iterations: 1, + batch_size: DEFAULT_READ_BATCH_SIZE, + path_type: "bare".to_string(), + verbose: false, + pprof_file: None, + pprof_after_warmup: false, + pk_as_binary: false, + reader: ReaderMode::Direct, + } + } + fn new_test_metadata() -> RegionMetadata { let mut builder = RegionMetadataBuilder::new(RegionId::new(1, 0)); builder @@ -882,19 +1059,27 @@ mod tests { } #[test] - fn test_parse_region_id() { - assert_eq!(parse_region_id("1024:7").unwrap(), RegionId::new(1024, 7)); - assert_eq!( - parse_region_id(&RegionId::new(1, 2).as_u64().to_string()).unwrap(), - RegionId::new(1, 2) - ); + fn test_resolve_input_accepts_direct_file_for_direct_reader() { + let mut command = test_command(); + command.file_path = Some(PathBuf::from("/tmp/source.parquet")); + + let input = command.resolve_input().unwrap(); + match input { + ParquetbenchInput::LocalFile { file_path } => { + assert_eq!(file_path, PathBuf::from("/tmp/source.parquet")); + } + ParquetbenchInput::Region { .. } => panic!("expected local file input"), + } } #[test] - fn test_parse_path_type() { - assert_eq!(parse_path_type("bare").unwrap(), PathType::Bare); - assert_eq!(parse_path_type("data").unwrap(), PathType::Data); - assert_eq!(parse_path_type("metadata").unwrap(), PathType::Metadata); + fn test_resolve_input_rejects_mixed_input_modes() { + let mut command = test_command(); + command.file_path = Some(PathBuf::from("/tmp/source.parquet")); + command.config = Some(PathBuf::from("config.toml")); + + let err = command.resolve_input().unwrap_err(); + assert!(err.to_string().contains("--file-path cannot be used with")); } #[test] diff --git a/src/cmd/src/datanode/scanbench.rs b/src/cmd/src/datanode/scanbench.rs index 4909467f48..1dbee025a7 100644 --- a/src/cmd/src/datanode/scanbench.rs +++ b/src/cmd/src/datanode/scanbench.rs @@ -52,11 +52,13 @@ use sqlparser::parser::Parser as SqlParser; use store_api::metadata::RegionMetadata; use store_api::path_utils::WAL_DIR; use store_api::region_engine::{PrepareRequest, QueryScanContext, RegionEngine}; -use store_api::region_request::{PathType, RegionOpenRequest, RegionRequest}; +use store_api::region_request::{RegionOpenRequest, RegionRequest}; use store_api::storage::{RegionId, ScanRequest, TimeSeriesDistribution, TimeSeriesRowSelector}; use tokio::fs; -use crate::datanode::objbench::{build_object_store, parse_config}; +use crate::datanode::tool_util::{ + build_object_store, format_bytes, parse_config, parse_path_type, parse_region_id, +}; use crate::error; /// Scan benchmark command - benchmarks scanning a region directly from storage. @@ -164,60 +166,6 @@ fn resolve_projection( Ok(None) } -fn format_bytes(bytes: u64) -> String { - const KIB: u64 = 1024; - const MIB: u64 = 1024 * KIB; - const GIB: u64 = 1024 * MIB; - if bytes >= GIB { - format!("{:.2} GiB", bytes as f64 / GIB as f64) - } else if bytes >= MIB { - format!("{:.2} MiB", bytes as f64 / MIB as f64) - } else if bytes >= KIB { - format!("{:.2} KiB", bytes as f64 / KIB as f64) - } else { - format!("{} B", bytes) - } -} - -fn parse_region_id(s: &str) -> error::Result { - if s.contains(':') { - let parts: Vec<&str> = s.splitn(2, ':').collect(); - let table_id: u32 = parts[0].parse().map_err(|e| { - error::IllegalConfigSnafu { - msg: format!("invalid table_id in region_id '{}': {}", s, e), - } - .build() - })?; - let region_num: u32 = parts[1].parse().map_err(|e| { - error::IllegalConfigSnafu { - msg: format!("invalid region_num in region_id '{}': {}", s, e), - } - .build() - })?; - Ok(RegionId::new(table_id, region_num)) - } else { - let id: u64 = s.parse().map_err(|e| { - error::IllegalConfigSnafu { - msg: format!("invalid region_id '{}': {}", s, e), - } - .build() - })?; - Ok(RegionId::from_u64(id)) - } -} - -fn parse_path_type(s: &str) -> error::Result { - match s.to_lowercase().as_str() { - "bare" => Ok(PathType::Bare), - "data" => Ok(PathType::Data), - "metadata" => Ok(PathType::Metadata), - _ => Err(error::IllegalConfigSnafu { - msg: format!("invalid path_type '{}', expected: bare, data, metadata", s), - } - .build()), - } -} - /// Rewrites literal values in comparison expressions to match the column's arrow type. struct LiteralTypeCaster { schema: DFSchemaRef, diff --git a/src/cmd/src/datanode/sst_replace.rs b/src/cmd/src/datanode/sst_replace.rs new file mode 100644 index 0000000000..df13fcd6f2 --- /dev/null +++ b/src/cmd/src/datanode/sst_replace.rs @@ -0,0 +1,361 @@ +// 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::PathBuf; + +use clap::Parser; +use colored::Colorize; +use mito2::manifest::action::{RegionEdit, RegionMetaAction, RegionMetaActionList}; +use mito2::manifest::manager::{RegionManifestManager, RegionManifestOptions}; +use mito2::region::ManifestStats; +use mito2::sst::file::RegionFileId; +use mito2::sst::location::{region_dir_from_table_dir, sst_file_path}; +use object_store::ObjectStore; +use parquet::file::FOOTER_SIZE; +use parquet::file::metadata::{FooterTail, ParquetMetaDataReader}; +use store_api::region_request::PathType; + +use crate::datanode::tool_util::{ + build_object_store, max_row_group_uncompressed_size, parse_config, parse_file_id, + parse_path_type, parse_region_id, +}; +use crate::error; + +/// Replace a mito region SST and update the corresponding manifest metadata. +#[derive(Debug, Parser)] +pub struct SstReplaceCommand { + /// Path to config TOML file (same format as standalone/datanode config). + #[clap(long, value_name = "FILE")] + config: PathBuf, + + /// Region ID: either numeric u64 (e.g. "4398046511104") or "table_id:region_num" (e.g. "1024:0"). + #[clap(long)] + region_id: String, + + /// Table directory relative to data home (e.g. "data/greptime/public/1024/"). + #[clap(long)] + table_dir: String, + + /// SST file id to replace. + #[clap(long)] + file_id: String, + + /// Local parquet file used as the replacement. + #[clap(long, value_name = "FILE", conflicts_with = "replacement_object")] + replacement_file: Option, + + /// Object-store parquet path used as the replacement. + #[clap(long, value_name = "PATH", conflicts_with = "replacement_file")] + replacement_object: Option, + + /// Path type for the region: auto, bare, data, metadata. + #[clap(long, default_value = "auto")] + path_type: String, + + /// Actually overwrite the SST object and append a manifest delta. + #[clap(long, default_value_t = false)] + confirm: bool, + + /// Verbose output. + #[clap(short, long, default_value_t = false)] + verbose: bool, +} + +impl SstReplaceCommand { + pub async fn run(&self) -> error::Result<()> { + if self.verbose { + common_telemetry::init_default_ut_logging(); + } + + println!("{}", "Starting sst-replace...".cyan().bold()); + + let region_id = parse_region_id(&self.region_id)?; + let file_id = parse_file_id(&self.file_id)?; + let path_type = parse_optional_path_type(&self.path_type)?; + let replacement = self.replacement()?; + + let (store_cfg, mito_config, _wal_config) = parse_config(&self.config)?; + let object_store = build_object_store(&store_cfg).await?; + println!("{} Object store initialized", "[ok]".green()); + + let candidates = path_type + .map(|path_type| vec![path_type]) + .unwrap_or_else(|| vec![PathType::Bare, PathType::Data, PathType::Metadata]); + let mut found = None; + let mut existing_manifests = Vec::new(); + for candidate in candidates { + let stats = ManifestStats::default(); + let region_dir = region_dir_from_table_dir(&self.table_dir, region_id, candidate); + let manifest_opts = + RegionManifestOptions::new(&mito_config, ®ion_dir, &object_store); + let Some(manifest_manager) = RegionManifestManager::open(manifest_opts, &stats) + .await + .map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("open {} manifest failed: {e:?}", path_type_name(candidate)), + } + .build() + })? + else { + continue; + }; + existing_manifests.push(path_type_name(candidate)); + let manifest = manifest_manager.manifest(); + if let Some(file) = manifest.files.get(&file_id) { + if found.is_some() { + return error::IllegalConfigSnafu { + msg: format!( + "file {} exists in multiple manifests under region {}; specify --path-type", + file_id, region_id + ), + } + .fail(); + } + found = Some((candidate, manifest_manager, file.clone())); + } + } + let Some((path_type, mut manifest_manager, old_file)) = found else { + let suffix = if existing_manifests.is_empty() { + "no manifest found".to_string() + } else { + format!( + "manifests found for [{}], but none contains file {}", + existing_manifests.join(", "), + file_id + ) + }; + return error::IllegalConfigSnafu { + msg: format!("region manifest not found for {}: {}", region_id, suffix), + } + .fail(); + }; + + let region_file_id = RegionFileId::new(old_file.region_id, old_file.file_id); + let target_path = sst_file_path(&self.table_dir, region_file_id, path_type); + let target_stat = object_store.stat(&target_path).await.map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("stat target SST {} failed: {e}", target_path), + } + .build() + })?; + let replacement_bytes = load_replacement_bytes(&object_store, replacement).await?; + let new_file_size = replacement_bytes.len() as u64; + let parquet_meta = decode_parquet_metadata(&replacement_bytes).map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("read replacement parquet metadata failed: {e}"), + } + .build() + })?; + let new_num_rows = parquet_meta.file_metadata().num_rows() as u64; + let new_num_row_groups = parquet_meta.num_row_groups() as u64; + validate_replacement(&old_file, new_num_rows, new_num_row_groups)?; + + let mut new_file = old_file.clone(); + new_file.file_size = new_file_size; + new_file.num_rows = new_num_rows; + new_file.num_row_groups = new_num_row_groups; + new_file.max_row_group_uncompressed_size = max_row_group_uncompressed_size(&parquet_meta); + + println!("{} Region: {}", "[ok]".green(), region_id); + println!("{} Target SST: {}", "[ok]".green(), target_path.cyan()); + println!( + "{} Size: {} -> {} bytes", + "[ok]".green(), + target_stat.content_length(), + new_file_size + ); + println!( + "{} Rows: {}, row groups: {}", + "[ok]".green(), + new_num_rows, + new_num_row_groups + ); + + if !self.confirm { + println!( + "{} Dry run only. Re-run with --confirm to overwrite the SST and update the manifest.", + "[dry-run]".yellow() + ); + return Ok(()); + } + + object_store + .write(&target_path, replacement_bytes) + .await + .map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("write replacement SST {} failed: {e}", target_path), + } + .build() + })?; + + let edit = RegionEdit { + files_to_add: vec![new_file], + files_to_remove: vec![], + timestamp_ms: None, + compaction_time_window: None, + flushed_entry_id: None, + flushed_sequence: None, + committed_sequence: None, + }; + let version = manifest_manager + .update( + RegionMetaActionList::with_action(RegionMetaAction::Edit(edit)), + false, + ) + .await + .map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("update manifest failed: {e:?}"), + } + .build() + })?; + + println!( + "{} Replacement committed, manifest version {}", + "[ok]".green(), + version + ); + Ok(()) + } + + fn replacement(&self) -> error::Result> { + match (&self.replacement_file, &self.replacement_object) { + (Some(path), None) => Ok(Replacement::Local(path)), + (None, Some(path)) => Ok(Replacement::Object(path)), + _ => error::IllegalConfigSnafu { + msg: "specify exactly one of --replacement-file or --replacement-object" + .to_string(), + } + .fail(), + } + } +} + +enum Replacement<'a> { + Local(&'a PathBuf), + Object(&'a str), +} + +async fn load_replacement_bytes( + object_store: &ObjectStore, + replacement: Replacement<'_>, +) -> error::Result> { + match replacement { + Replacement::Local(path) => tokio::fs::read(path).await.map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("read replacement file {} failed: {e}", path.display()), + } + .build() + }), + Replacement::Object(path) => { + object_store + .read(path) + .await + .map(|b| b.to_vec()) + .map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("read replacement object {} failed: {e}", path), + } + .build() + }) + } + } +} + +fn validate_replacement( + old_file: &mito2::sst::file::FileMeta, + new_num_rows: u64, + new_num_row_groups: u64, +) -> error::Result<()> { + if old_file.num_rows != 0 && old_file.num_rows != new_num_rows { + return error::IllegalConfigSnafu { + msg: format!( + "replacement row count mismatch: manifest has {}, replacement has {}", + old_file.num_rows, new_num_rows + ), + } + .fail(); + } + if old_file.num_row_groups != 0 && old_file.num_row_groups != new_num_row_groups { + return error::IllegalConfigSnafu { + msg: format!( + "replacement row group count mismatch: manifest has {}, replacement has {}", + old_file.num_row_groups, new_num_row_groups + ), + } + .fail(); + } + Ok(()) +} + +fn decode_parquet_metadata( + data: &[u8], +) -> Result> { + if data.len() < FOOTER_SIZE { + return Err("file too small".into()); + } + let footer_start = data.len() - FOOTER_SIZE; + let mut footer = [0; FOOTER_SIZE]; + footer.copy_from_slice(&data[footer_start..]); + let footer = FooterTail::try_new(&footer)?; + let metadata_len = footer.metadata_length(); + if footer_start < metadata_len { + return Err("invalid footer/metadata length".into()); + } + let metadata_start = footer_start - metadata_len; + Ok(ParquetMetaDataReader::decode_metadata( + &data[metadata_start..footer_start], + )?) +} + +fn parse_optional_path_type(value: &str) -> error::Result> { + match value.to_lowercase().as_str() { + "auto" => Ok(None), + _ => parse_path_type(value).map(Some).map_err(|_| { + error::IllegalConfigSnafu { + msg: format!("invalid path_type '{value}', expected: auto, bare, data, metadata"), + } + .build() + }), + } +} + +fn path_type_name(path_type: PathType) -> &'static str { + match path_type { + PathType::Bare => "bare", + PathType::Data => "data", + PathType::Metadata => "metadata", + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_validate_replacement_counts() { + let old_file = mito2::sst::file::FileMeta { + num_rows: 4, + num_row_groups: 1, + ..Default::default() + }; + assert!(validate_replacement(&old_file, 4, 1).is_ok()); + assert!(validate_replacement(&old_file, 3, 1).is_err()); + assert!(validate_replacement(&old_file, 4, 2).is_err()); + + let legacy_file = mito2::sst::file::FileMeta::default(); + assert!(validate_replacement(&legacy_file, 4, 1).is_ok()); + } +} diff --git a/src/cmd/src/datanode/tool_util.rs b/src/cmd/src/datanode/tool_util.rs new file mode 100644 index 0000000000..3b16415e89 --- /dev/null +++ b/src/cmd/src/datanode/tool_util.rs @@ -0,0 +1,236 @@ +// 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. + +#[cfg(feature = "dev-tools")] +use std::fs::File; +use std::path::Path; +use std::sync::Arc; + +use common_wal::config::DatanodeWalConfig; +use datanode::config::RegionEngineConfig; +use datanode::store; +use mito2::config::MitoConfig; +use object_store::ObjectStore; +#[cfg(feature = "dev-tools")] +use parquet::basic::Compression; +use parquet::file::metadata::{KeyValue, ParquetMetaData}; +#[cfg(feature = "dev-tools")] +use parquet::file::metadata::{PageIndexPolicy, ParquetMetaDataReader}; +use snafu::OptionExt; +#[cfg(feature = "dev-tools")] +use snafu::ResultExt; +use store_api::metadata::{RegionMetadata, RegionMetadataRef}; +use store_api::region_request::PathType; +use store_api::storage::{FileId, RegionId}; + +use crate::datanode::{StorageConfig, StorageConfigWrapper}; +use crate::error; + +pub(crate) fn parse_config( + config_path: &Path, +) -> error::Result<(StorageConfig, MitoConfig, DatanodeWalConfig)> { + let cfg_str = std::fs::read_to_string(config_path).map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("failed to read config {}: {e}", config_path.display()), + } + .build() + })?; + + let store_cfg: StorageConfigWrapper = toml::from_str(&cfg_str).map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("failed to parse config {}: {e}", config_path.display()), + } + .build() + })?; + + let wal_config = store_cfg.wal; + let storage_config = store_cfg.storage; + let mito_engine_config = store_cfg + .region_engine + .into_iter() + .find_map(|config| match config { + RegionEngineConfig::Mito(mito) => Some(mito), + _ => None, + }) + .with_context(|| error::IllegalConfigSnafu { + msg: format!("Engine config not found in {:?}", config_path), + })?; + + Ok((storage_config, mito_engine_config, wal_config)) +} + +pub(crate) async fn build_object_store(config: &StorageConfig) -> error::Result { + store::new_object_store(config.store.clone(), &config.data_home) + .await + .map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("Failed to build object store: {e:?}"), + } + .build() + }) +} + +pub(crate) fn extract_region_metadata( + file_path: &str, + metadata: &ParquetMetaData, +) -> error::Result { + let key_values: Option<&Vec> = metadata.file_metadata().key_value_metadata(); + let Some(key_values) = key_values else { + return Err(error::IllegalConfigSnafu { + msg: format!("{file_path}: missing parquet key_value metadata"), + } + .build()); + }; + let json = key_values + .iter() + .find(|key_value| key_value.key == mito2::sst::parquet::PARQUET_METADATA_KEY) + .and_then(|key_value| key_value.value.as_ref()) + .ok_or_else(|| { + error::IllegalConfigSnafu { + msg: format!( + "{file_path}: key {} not found or empty", + mito2::sst::parquet::PARQUET_METADATA_KEY + ), + } + .build() + })?; + let region = RegionMetadata::from_json(json).map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("invalid region metadata json: {e}"), + } + .build() + })?; + Ok(Arc::new(region)) +} + +pub(crate) fn parse_region_id(value: &str) -> error::Result { + if let Some((table_id, region_number)) = value.split_once(':') { + let table_id = table_id.parse().map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("invalid table_id in region_id '{value}': {e}"), + } + .build() + })?; + let region_number = region_number.parse().map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("invalid region_num in region_id '{value}': {e}"), + } + .build() + })?; + Ok(RegionId::new(table_id, region_number)) + } else { + value.parse().map(RegionId::from_u64).map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("invalid region_id '{value}': {e}"), + } + .build() + }) + } +} + +pub(crate) fn parse_file_id(value: &str) -> error::Result { + FileId::parse_str(value).map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("invalid file_id '{value}': {e}"), + } + .build() + }) +} + +pub(crate) fn parse_path_type(value: &str) -> error::Result { + match value.to_lowercase().as_str() { + "bare" => Ok(PathType::Bare), + "data" => Ok(PathType::Data), + "metadata" => Ok(PathType::Metadata), + _ => Err(error::IllegalConfigSnafu { + msg: format!("invalid path_type '{value}', expected: bare, data, metadata"), + } + .build()), + } +} + +pub(crate) fn format_bytes(bytes: u64) -> String { + const KIB: u64 = 1024; + const MIB: u64 = 1024 * KIB; + const GIB: u64 = 1024 * MIB; + if bytes >= GIB { + format!("{:.2} GiB", bytes as f64 / GIB as f64) + } else if bytes >= MIB { + format!("{:.2} MiB", bytes as f64 / MIB as f64) + } else if bytes >= KIB { + format!("{:.2} KiB", bytes as f64 / KIB as f64) + } else { + format!("{bytes} B") + } +} + +pub(crate) fn max_row_group_uncompressed_size(metadata: &ParquetMetaData) -> u64 { + metadata + .row_groups() + .iter() + .map(|row_group| { + row_group + .columns() + .iter() + .map(|column| column.uncompressed_size() as u64) + .sum::() + }) + .max() + .unwrap_or(0) +} + +#[cfg(feature = "dev-tools")] +pub(crate) fn load_local_parquet_metadata(path: &Path) -> error::Result { + let file = File::open(path).context(error::FileIoSnafu)?; + ParquetMetaDataReader::new() + .with_page_index_policy(PageIndexPolicy::Optional) + .parse_and_finish(&file) + .map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("read parquet metadata failed for {}: {e}", path.display()), + } + .build() + }) +} + +#[cfg(feature = "dev-tools")] +pub(crate) fn compression_name(compression: Compression) -> &'static str { + match compression { + Compression::UNCOMPRESSED => "uncompressed", + Compression::SNAPPY => "snappy", + Compression::GZIP(_) => "gzip", + Compression::LZO => "lzo", + Compression::BROTLI(_) => "brotli", + Compression::LZ4 => "lz4", + Compression::ZSTD(_) => "zstd", + Compression::LZ4_RAW => "lz4-raw", + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_parse_region_and_path_type() { + assert_eq!(parse_region_id("1024:7").unwrap(), RegionId::new(1024, 7)); + assert_eq!( + parse_region_id(&RegionId::new(1, 2).as_u64().to_string()).unwrap(), + RegionId::new(1, 2) + ); + assert_eq!(parse_path_type("bare").unwrap(), PathType::Bare); + assert_eq!(parse_path_type("data").unwrap(), PathType::Data); + assert_eq!(parse_path_type("metadata").unwrap(), PathType::Metadata); + } +}