rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-09-08 06:29:04 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(cmd): add parquet development tools (#8939)

Browse Source 
* feat(cmd): add parquet metadata development tool

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

* feat(cmd): add parquet rewrite development tool

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

* feat(cmd): add SST replacement development tool

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

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

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

* fix(cmd): clean up parquet development tools

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

* refactor(cmd): share datanode tool utilities

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

* fix(cmd): satisfy parquet tool lints

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

* docs: document parquet development tools

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

* refactor(cmd): rename SST replacement module

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

* fix(cmd): validate parquet rewrite options

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

---------

Signed-off-by: evenyag <realevenyag@gmail.com>
This commit is contained in:
Yingwen
2026-08-24 13:34:30 +00:00
committed by GitHub
parent 1c5eabcbbf
commit 0cc83c4570
11 changed files with 2126 additions and 263 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/cmd/Cargo.toml
+1 -1
View File
@@ -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"]
src/cmd/src/bin/greptime.rs
+6
View File
@@ -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)?)
src/cmd/src/datanode.rs
+46
View File
@@ -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);
}
}
}
}
src/cmd/src/datanode/objbench.rs
+8 -99
View File
@@ -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<PathBuf>,
}
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::<u64>()
})
.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<FileDirComponents> {
})
}
pub(crate) fn extract_region_metadata(
file_path: &str,
meta: &parquet::file::metadata::ParquetMetaData,
) -> error::Result<RegionMetadataRef> {
let kvs: Option<&Vec<KeyValue>> = 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<ObjectStore> {
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() {
src/cmd/src/datanode/parquet_meta.rs
+340
View File
@@ -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<String, Option<String>>,
row_groups: Vec<RowGroupMetaView>,
}
#[derive(Debug, Serialize)]
struct RowGroupMetaView {
index: usize,
num_rows: i64,
uncompressed_size: i64,
compressed_size: i64,
compression_ratio: Option<f64>,
data_pages: Option<usize>,
dictionary_page_bytes: i64,
columns: Vec<ColumnChunkMetaView>,
}
#[derive(Debug, Serialize)]
struct ColumnChunkMetaView {
index: usize,
path: String,
physical_type: String,
encodings: Vec<String>,
compression: String,
num_values: i64,
uncompressed_size: i64,
compressed_size: i64,
compression_ratio: Option<f64>,
data_page_offset: i64,
dictionary_page_offset: Option<i64>,
dictionary_page_bytes: Option<i64>,
data_pages: Option<usize>,
has_statistics: bool,
column_index_offset: Option<i64>,
column_index_length: Option<i32>,
offset_index_offset: Option<i64>,
offset_index_length: Option<i32>,
bloom_filter_offset: Option<i64>,
bloom_filter_length: Option<i32>,
}
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<i64>,
data_page_offset: i64,
) -> Option<i64> {
dictionary_page_offset.map(|offset| data_page_offset.saturating_sub(offset))
}
fn compression_ratio(uncompressed: i64, compressed: i64) -> Option<f64> {
(uncompressed > 0).then(|| compressed as f64 / uncompressed as f64)
}
fn format_ratio(ratio: Option<f64>) -> String {
ratio
.map(|value| format!("{value:.4}"))
.unwrap_or_else(|| "unknown".to_string())
}
fn format_optional_usize(value: Option<usize>) -> String {
value
.map(|value| value.to_string())
.unwrap_or_else(|| "unknown".to_string())
}
fn format_optional_i64(value: Option<i64>) -> String {
value
.map(|value| value.to_string())
.unwrap_or_else(|| "none".to_string())
}
fn format_optional_i32(value: Option<i32>) -> 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();
}
}
src/cmd/src/datanode/parquet_rewrite.rs
+579
View File
@@ -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<PathBuf>,
/// Path to writer properties TOML in rewrite mode.
#[clap(long, value_name = "FILE")]
properties: Option<PathBuf>,
/// Dump writer properties TOML inferred from the input parquet file.
#[clap(long, value_name = "FILE")]
dump_properties: Option<PathBuf>,
/// Number of rows per record batch.
#[clap(long)]
batch_size: Option<usize>,
/// 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<ColumnConfig>,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
#[serde(default, deny_unknown_fields)]
struct WriterConfig {
dictionary_enabled: Option<bool>,
compression: Option<CompressionConfig>,
compression_level: Option<u32>,
encoding: Option<EncodingConfig>,
max_row_group_row_count: Option<usize>,
data_page_size_limit: Option<usize>,
data_page_row_count_limit: Option<usize>,
dictionary_page_size_limit: Option<usize>,
}
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
#[serde(default, deny_unknown_fields)]
struct ColumnConfig {
path: Vec<String>,
dictionary_enabled: Option<bool>,
compression: Option<CompressionConfig>,
compression_level: Option<u32>,
encoding: Option<EncodingConfig>,
}
#[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<RewriteProperties> {
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<Vec<parquet::file::metadata::KeyValue>>,
) -> error::Result<WriterProperties> {
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> {
File::open(path).context(error::FileIoSnafu)
}
fn create_file(path: &Path) -> error::Result<File> {
File::create(path).context(error::FileIoSnafu)
}
fn illegal_config<T>(msg: String) -> error::Result<T> {
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<u32>,
) -> error::Result<Compression> {
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<EncodingConfig> {
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();
}
}
src/cmd/src/datanode/parquetbench.rs
+292 -107
View File
@@ -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<PathBuf>,
/// 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<String>,
/// Table directory relative to data home (e.g. "data/greptime/public/1024/")
#[clap(long)]
table_dir: String,
table_dir: Option<String>,
/// SST file id to benchmark.
#[clap(long)]
file_id: String,
file_id: Option<String>,
/// Local parquet SST file to benchmark with the direct reader.
#[clap(long, value_name = "FILE")]
file_path: Option<PathBuf>,
/// 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<PathType>,
table_dir: Option<String>,
}
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, &region_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<ParquetbenchInput> {
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<ParquetScanConfig> {
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<ParquetbenchSource> {
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, &region_id, table_dir, &file_id, path_type).await,
}
}
fn build_local_file_source(file_path: &Path) -> error::Result<ParquetbenchSource> {
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<ParquetbenchSource> {
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<usize, String> {
Ok(batch_size)
}
fn parse_region_id(s: &str) -> error::Result<RegionId> {
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<PathType> {
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]
src/cmd/src/datanode/scanbench.rs
+4 -56
View File
@@ -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<RegionId> {
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<PathType> {
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,
src/cmd/src/datanode/sst_replace.rs
+361
View File
@@ -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<PathBuf>,
/// Object-store parquet path used as the replacement.
#[clap(long, value_name = "PATH", conflicts_with = "replacement_file")]
replacement_object: Option<String>,
/// 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, &region_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<Replacement<'_>> {
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<Vec<u8>> {
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<parquet::file::metadata::ParquetMetaData, Box<dyn std::error::Error + Send + Sync>> {
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<Option<PathType>> {
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());
}
}
src/cmd/src/datanode/tool_util.rs
+236
View File
@@ -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<ObjectStore> {
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<RegionMetadataRef> {
let key_values: Option<&Vec<KeyValue>> = 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<RegionId> {
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> {
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<PathType> {
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::<u64>()
})
.max()
.unwrap_or(0)
}
#[cfg(feature = "dev-tools")]
pub(crate) fn load_local_parquet_metadata(path: &Path) -> error::Result<ParquetMetaData> {
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);
}
}