mirror of
https://github.com/GreptimeTeam/greptimedb.git
synced 2026-09-05 21:18:57 +00:00
feat(cmd): add parquet development tools (#8939)
* feat(cmd): add parquet metadata development tool Signed-off-by: evenyag <realevenyag@gmail.com> * feat(cmd): add parquet rewrite development tool Signed-off-by: evenyag <realevenyag@gmail.com> * feat(cmd): add SST replacement development tool Signed-off-by: evenyag <realevenyag@gmail.com> * feat(cmd): support local parquet files in parquetbench Signed-off-by: evenyag <realevenyag@gmail.com> * fix(cmd): clean up parquet development tools Signed-off-by: evenyag <realevenyag@gmail.com> * refactor(cmd): share datanode tool utilities Signed-off-by: evenyag <realevenyag@gmail.com> * fix(cmd): satisfy parquet tool lints Signed-off-by: evenyag <realevenyag@gmail.com> * docs: document parquet development tools Signed-off-by: evenyag <realevenyag@gmail.com> * refactor(cmd): rename SST replacement module Signed-off-by: evenyag <realevenyag@gmail.com> * fix(cmd): validate parquet rewrite options Signed-off-by: evenyag <realevenyag@gmail.com> --------- Signed-off-by: evenyag <realevenyag@gmail.com>
This commit is contained in:
Yingwen
@@ -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 <COMMAND> --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 <ROWS>` 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
|
||||
`<table-id>:<region-number>`. `--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 <SVG>` 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.
|
||||
+1
-1
@@ -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"]
|
||||
|
||||
@@ -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)?)
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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() {
|
||||
|
||||
@@ -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();
|
||||
}
|
||||
}
|
||||
@@ -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();
|
||||
}
|
||||
}
|
||||
@@ -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, ®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<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, ®ion_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]
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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, ®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<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());
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user