rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-10-12 23:12:28 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(mito2): add opt-in byte stream split encoding

Browse Source 
Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>
This commit is contained in:
discord9
2026-09-04 19:23:03 +08:00
parent 27b0ca6676
commit 994f465a6d
11 changed files with 371 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/mito2/src/compaction/compactor.rs
+1
View File
@@ -655,6 +655,7 @@ where
write_buffer_size: compaction_region.engine_config.sst_write_buffer_size,
max_file_size: picker_output.max_file_size,
row_group_size: compaction_region.region_options.row_group_size(),
float_field_encoding: compaction_region.region_options.float_field_encoding,
};
let merger = self.merger.clone();
let compaction_region = compaction_region.clone();
src/mito2/src/compaction/window.rs
+1
View File
@@ -337,6 +337,7 @@ mod tests {
primary_key_encoding: None,
write_buffer_size: None,
preserve_row_sequence: false,
float_field_encoding: Default::default(),
},
compaction_time_window: None,
}
src/mito2/src/flush.rs
+1
View File
@@ -451,6 +451,7 @@ impl RegionFlushTask {
let mut write_opts = WriteOptions {
write_buffer_size: self.engine_config.sst_write_buffer_size,
float_field_encoding: version.options.float_field_encoding,
..Default::default()
};
if let Some(row_group_size) = self.row_group_size {
src/mito2/src/memtable.rs
+2
View File
@@ -473,6 +473,7 @@ impl MemtableBuilderProvider {
BulkMemtableBuilder::new(self.write_buffer_manager.clone(), !dedup, merge_mode)
.with_config(config.clone())
.with_row_group_size(options.row_group_size())
.with_float_field_encoding(options.float_field_encoding)
.with_compact_dispatcher(self.compact_dispatcher.clone()),
),
Some(MemtableOptions::TimeSeries) => Arc::new(TimeSeriesMemtableBuilder::new(
@@ -496,6 +497,7 @@ impl MemtableBuilderProvider {
merge_mode,
)
.with_row_group_size(options.row_group_size())
.with_float_field_encoding(options.float_field_encoding)
.with_compact_dispatcher(self.compact_dispatcher.clone());
if let Some(MemtableOptions::Bulk(config)) = &options.memtable {
src/mito2/src/memtable/bulk.rs
+52 -2
View File
@@ -40,6 +40,7 @@ use rayon::prelude::*;
use serde::{Deserialize, Serialize};
use serde_with::{DisplayFromStr, serde_as};
use store_api::metadata::RegionMetadataRef;
use store_api::mito_engine_options::FloatFieldEncoding;
use store_api::storage::{ColumnId, FileId, RegionId, SequenceRange};
use tokio::sync::Semaphore;
@@ -410,6 +411,7 @@ pub struct BulkMemtable {
merge_mode: MergeMode,
/// Max number of rows in a parquet row group for encoded parts.
row_group_size: usize,
float_field_encoding: FloatFieldEncoding,
}
impl std::fmt::Debug for BulkMemtable {
@@ -655,11 +657,13 @@ impl Memtable for BulkMemtable {
id,
self.config.clone(),
self.row_group_size,
self.float_field_encoding,
))),
compact_dispatcher: self.compact_dispatcher.clone(),
append_mode: self.append_mode,
merge_mode: self.merge_mode,
row_group_size: self.row_group_size,
float_field_encoding: self.float_field_encoding,
})
}
@@ -722,6 +726,32 @@ impl BulkMemtable {
append_mode: bool,
merge_mode: MergeMode,
row_group_size: usize,
) -> Self {
Self::new_with_row_group_size_and_encoding(
id,
config,
metadata,
write_buffer_manager,
compact_dispatcher,
append_mode,
merge_mode,
row_group_size,
FloatFieldEncoding::default(),
)
}
/// Creates a new BulkMemtable with the given row group size and float encoding.
#[allow(clippy::too_many_arguments)]
pub fn new_with_row_group_size_and_encoding(
id: MemtableId,
config: BulkMemtableConfig,
metadata: RegionMetadataRef,
write_buffer_manager: Option<WriteBufferManagerRef>,
compact_dispatcher: Option<Arc<CompactDispatcher>>,
append_mode: bool,
merge_mode: MergeMode,
row_group_size: usize,
float_field_encoding: FloatFieldEncoding,
) -> Self {
let config = config.sanitize();
let region_id = metadata.region_id;
@@ -740,11 +770,13 @@ impl BulkMemtable {
id,
config,
row_group_size,
float_field_encoding,
))),
compact_dispatcher,
append_mode,
merge_mode,
row_group_size,
float_field_encoding,
}
}
@@ -1154,6 +1186,7 @@ struct MemtableCompactor {
config: BulkMemtableConfig,
/// Max number of rows in a parquet row group for encoded parts.
row_group_size: usize,
float_field_encoding: FloatFieldEncoding,
}
impl MemtableCompactor {
@@ -1163,12 +1196,14 @@ impl MemtableCompactor {
memtable_id: MemtableId,
config: BulkMemtableConfig,
row_group_size: usize,
float_field_encoding: FloatFieldEncoding,
) -> Self {
Self {
region_id,
memtable_id,
config,
row_group_size,
float_field_encoding,
}
}
@@ -1208,6 +1243,7 @@ impl MemtableCompactor {
let encode_row_threshold = self.config.encode_row_threshold;
let encode_bytes_threshold = self.config.encode_bytes_threshold;
let row_group_size = self.row_group_size;
let float_field_encoding = self.float_field_encoding;
// Merge all groups in parallel
let merged_parts = collected
@@ -1222,6 +1258,7 @@ impl MemtableCompactor {
encode_row_threshold,
encode_bytes_threshold,
row_group_size,
float_field_encoding,
)
})
.collect::<Result<Vec<Option<MergedPart>>>>()?;
@@ -1257,6 +1294,7 @@ impl MemtableCompactor {
encode_row_threshold: usize,
encode_bytes_threshold: usize,
row_group_size: usize,
float_field_encoding: FloatFieldEncoding,
) -> Result<Option<MergedPart>> {
if parts_to_merge.is_empty() {
return Ok(None);
@@ -1357,7 +1395,11 @@ impl MemtableCompactor {
if estimated_total_rows > encode_row_threshold
|| estimated_total_bytes > encode_bytes_threshold
{
let encoder = BulkPartEncoder::new(metadata.clone(), row_group_size)?;
let encoder = BulkPartEncoder::new_with_float_field_encoding(
metadata.clone(),
row_group_size,
float_field_encoding,
)?;
let mut metrics = BulkPartEncodeMetrics::default();
let encoded_part = encoder.encode_record_batch_iter(
boxed_iter,
@@ -1483,6 +1525,7 @@ pub struct BulkMemtableBuilder {
merge_mode: MergeMode,
/// Max number of rows in a parquet row group for encoded parts.
row_group_size: usize,
float_field_encoding: FloatFieldEncoding,
}
impl Default for BulkMemtableBuilder {
@@ -1494,6 +1537,7 @@ impl Default for BulkMemtableBuilder {
append_mode: false,
merge_mode: MergeMode::default(),
row_group_size: DEFAULT_ROW_GROUP_SIZE,
float_field_encoding: FloatFieldEncoding::default(),
}
}
}
@@ -1525,6 +1569,11 @@ impl BulkMemtableBuilder {
self
}
pub fn with_float_field_encoding(mut self, encoding: FloatFieldEncoding) -> Self {
self.float_field_encoding = encoding;
self
}
/// Sets the compact dispatcher.
pub fn with_compact_dispatcher(mut self, compact_dispatcher: Arc<CompactDispatcher>) -> Self {
self.compact_dispatcher = Some(compact_dispatcher);
@@ -1539,7 +1588,7 @@ impl BulkMemtableBuilder {
impl MemtableBuilder for BulkMemtableBuilder {
fn build(&self, id: MemtableId, metadata: &RegionMetadataRef) -> MemtableRef {
Arc::new(BulkMemtable::new_with_row_group_size(
Arc::new(BulkMemtable::new_with_row_group_size_and_encoding(
id,
self.config.clone(),
metadata.clone(),
@@ -1548,6 +1597,7 @@ impl MemtableBuilder for BulkMemtableBuilder {
self.append_mode,
self.merge_mode,
self.row_group_size,
self.float_field_encoding,
))
}
src/mito2/src/memtable/bulk/part.rs
+21 -11
View File
@@ -53,6 +53,7 @@ use smallvec::SmallVec;
use snafu::{OptionExt, ResultExt};
use store_api::codec::PrimaryKeyEncoding;
use store_api::metadata::{RegionMetadata, RegionMetadataRef};
use store_api::mito_engine_options::FloatFieldEncoding;
use store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME;
use store_api::storage::{ColumnId, FileId, SequenceNumber, SequenceRange};
@@ -70,7 +71,9 @@ use crate::sst::SeriesEstimator;
use crate::sst::index::IndexOutput;
use crate::sst::parquet::flat_format::primary_key_column_index;
use crate::sst::parquet::format::{PrimaryKeyArray, PrimaryKeyArrayBuilder};
use crate::sst::parquet::{PARQUET_METADATA_KEY, SstInfo};
use crate::sst::parquet::{
PARQUET_METADATA_KEY, SstInfo, apply_float_field_encoding,
};
const INIT_DICT_VALUE_CAPACITY: usize = 8;
@@ -1306,22 +1309,29 @@ pub struct BulkPartEncoder {
impl BulkPartEncoder {
pub fn new(metadata: RegionMetadataRef, row_group_size: usize) -> Result<BulkPartEncoder> {
Self::new_with_float_field_encoding(metadata, row_group_size, FloatFieldEncoding::default())
}
pub fn new_with_float_field_encoding(
metadata: RegionMetadataRef,
row_group_size: usize,
float_field_encoding: FloatFieldEncoding,
) -> Result<BulkPartEncoder> {
// TODO(yingwen): Skip arrow schema if needed.
let json = metadata.to_json().context(InvalidMetadataSnafu)?;
let key_value_meta =
parquet::file::metadata::KeyValue::new(PARQUET_METADATA_KEY.to_string(), json);
// TODO(yingwen): Do we need compression?
let writer_props = Some(
WriterProperties::builder()
.set_key_value_metadata(Some(vec![key_value_meta]))
.set_write_batch_size(row_group_size)
.set_max_row_group_row_count(Some(row_group_size))
.set_compression(Compression::ZSTD(ZstdLevel::default()))
.set_column_index_truncate_length(None)
.set_statistics_truncate_length(None)
.build(),
);
let mut props = WriterProperties::builder()
.set_key_value_metadata(Some(vec![key_value_meta]))
.set_write_batch_size(row_group_size)
.set_max_row_group_row_count(Some(row_group_size))
.set_compression(Compression::ZSTD(ZstdLevel::default()))
.set_column_index_truncate_length(None)
.set_statistics_truncate_length(None);
props = apply_float_field_encoding(props, &metadata, float_field_encoding);
let writer_props = Some(props.build());
Ok(Self {
metadata,
src/mito2/src/region/options.rs
+40 -1
View File
@@ -32,7 +32,10 @@ use store_api::codec::PrimaryKeyEncoding;
use store_api::metric_engine_consts::{
MEMTABLE_PARTITION_TREE_PRIMARY_KEY_ENCODING, PRIMARY_KEY_ENCODING,
};
use store_api::mito_engine_options::{COMPACTION_OVERRIDE, MAX_ROW_GROUP_ROW_COUNT_LIMIT};
use store_api::mito_engine_options::{
COMPACTION_OVERRIDE, EXPERIMENTAL_SST_FLOAT_FIELD_ENCODING, FloatFieldEncoding,
MAX_ROW_GROUP_ROW_COUNT_LIMIT,
};
use store_api::storage::{ColumnId, RegionId};
use strum::EnumString;
@@ -127,6 +130,9 @@ pub struct RegionOptions {
/// sequence metadata.
#[serde(default, skip_serializing_if = "is_false")]
pub preserve_row_sequence: bool,
/// Encoding for direct floating-point field columns in Parquet SSTs.
#[serde(default)]
pub float_field_encoding: FloatFieldEncoding,
}
fn is_false(value: &bool) -> bool {
@@ -273,6 +279,10 @@ impl RegionOptions {
sst_format = Some(FormatType::Flat);
}
let float_field_encoding = options
.float_field_encoding
.unwrap_or_default();
let compaction_override_flag = options_map
.get(COMPACTION_OVERRIDE)
.map(|v| matches!(v.to_lowercase().as_str(), "true" | "1"))
@@ -308,6 +318,7 @@ impl RegionOptions {
primary_key_encoding,
write_buffer_size: options.write_buffer_size,
preserve_row_sequence: options.preserve_row_sequence,
float_field_encoding,
};
opts.validate()?;
@@ -424,6 +435,9 @@ struct RegionOptionsWithoutEnum {
max_row_group_row_count: Option<usize>,
#[serde_as(as = "DisplayFromStr")]
preserve_row_sequence: bool,
#[serde(rename = "experimental_sst_float_field_encoding")]
#[serde_as(as = "NoneAsEmptyString")]
float_field_encoding: Option<FloatFieldEncoding>,
}
impl Default for RegionOptionsWithoutEnum {
@@ -440,6 +454,7 @@ impl Default for RegionOptionsWithoutEnum {
sst_format: options.sst_format,
max_row_group_row_count: options.max_row_group_row_count,
preserve_row_sequence: options.preserve_row_sequence,
float_field_encoding: Some(options.float_field_encoding),
}
}
}
@@ -601,6 +616,26 @@ mod tests {
assert_eq!(RegionOptions::default(), options);
}
#[test]
fn test_float_field_encoding_defaults_and_parses() {
let options = RegionOptions::try_from_options(RegionId::new(0, 0), &make_map(&[])).unwrap();
assert_eq!(FloatFieldEncoding::Default, options.float_field_encoding);
let map = make_map(&[(EXPERIMENTAL_SST_FLOAT_FIELD_ENCODING, "default")]);
let options = RegionOptions::try_from_options(RegionId::new(0, 0), &map).unwrap();
assert_eq!(FloatFieldEncoding::Default, options.float_field_encoding);
let map = make_map(&[(EXPERIMENTAL_SST_FLOAT_FIELD_ENCODING, "byte_stream_split")]);
let options = RegionOptions::try_from_options(RegionId::new(0, 0), &map).unwrap();
assert_eq!(
FloatFieldEncoding::ByteStreamSplit,
options.float_field_encoding
);
let map = make_map(&[(EXPERIMENTAL_SST_FLOAT_FIELD_ENCODING, "invalid")]);
assert!(RegionOptions::try_from_options(RegionId::new(0, 0), &map).is_err());
}
#[test]
fn test_with_ttl() {
let map = make_map(&[("ttl", "7d")]);
@@ -995,6 +1030,7 @@ mod tests {
primary_key_encoding: None,
write_buffer_size: None,
preserve_row_sequence: false,
float_field_encoding: FloatFieldEncoding::default(),
};
assert_eq!(expect, options);
}
@@ -1050,6 +1086,7 @@ mod tests {
primary_key_encoding: None,
write_buffer_size: Some(ReadableSize::mb(128)),
preserve_row_sequence: true,
float_field_encoding: FloatFieldEncoding::default(),
};
let region_options_json_str = serde_json::to_string(&options).unwrap();
assert!(region_options_json_str.contains("preserve_row_sequence"));
@@ -1061,6 +1098,7 @@ mod tests {
let got: RegionOptions = serde_json::from_str(old_region_options_json_str).unwrap();
assert_eq!(None, got.write_buffer_size);
assert!(!got.preserve_row_sequence);
assert_eq!(FloatFieldEncoding::Default, got.float_field_encoding);
let default_json = serde_json::to_value(RegionOptions::default()).unwrap();
assert!(default_json.get(WRITE_BUFFER_SIZE_KEY).is_none());
@@ -1121,6 +1159,7 @@ mod tests {
primary_key_encoding: None,
write_buffer_size: None,
preserve_row_sequence: false,
float_field_encoding: FloatFieldEncoding::default(),
};
assert_eq!(options, got);
}
src/mito2/src/sst/parquet.rs
+154 -3
View File
@@ -20,6 +20,10 @@ use std::sync::Arc;
use common_base::readable_size::ReadableSize;
use datatypes::json::JsonSettings;
use parquet::file::metadata::ParquetMetaData;
use parquet::file::properties::WriterPropertiesBuilder;
use parquet::schema::types::ColumnPath;
use store_api::metadata::{RegionMetadataRef, SemanticType};
use store_api::mito_engine_options::FloatFieldEncoding;
use store_api::storage::{ColumnId, FileId};
use crate::sst::DEFAULT_WRITE_BUFFER_SIZE;
@@ -72,6 +76,27 @@ pub(crate) struct Json2TargetLayout {
/// batching without changing the row group layout of newly written SSTs.
pub const DEFAULT_ROW_GROUP_SIZE: usize = 100 * 1024;
/// Applies the configured encoding to direct floating-point field columns.
pub(crate) fn apply_float_field_encoding(
mut builder: WriterPropertiesBuilder,
metadata: &RegionMetadataRef,
encoding: FloatFieldEncoding,
) -> WriterPropertiesBuilder {
if encoding == FloatFieldEncoding::ByteStreamSplit {
for column in &metadata.column_metadatas {
if column.semantic_type == SemanticType::Field
&& column.column_schema.data_type.is_float()
{
let path = ColumnPath::new(vec![column.column_schema.name.clone()]);
builder = builder
.set_column_encoding(path.clone(), parquet::basic::Encoding::BYTE_STREAM_SPLIT)
.set_column_dictionary_enabled(path, false);
}
}
}
builder
}
/// Parquet write options.
#[derive(Debug, Clone)]
pub struct WriteOptions {
@@ -83,6 +108,8 @@ pub struct WriteOptions {
/// Note: This is not a hard limit as we can only observe the file size when
/// ArrowWrite writes to underlying writers.
pub max_file_size: Option<usize>,
/// Encoding policy for direct floating-point field columns.
pub float_field_encoding: FloatFieldEncoding,
}
impl Default for WriteOptions {
@@ -91,6 +118,7 @@ impl Default for WriteOptions {
write_buffer_size: DEFAULT_WRITE_BUFFER_SIZE,
row_group_size: DEFAULT_ROW_GROUP_SIZE,
max_file_size: None,
float_field_encoding: FloatFieldEncoding::default(),
}
}
}
@@ -125,6 +153,7 @@ mod tests {
use std::sync::Arc;
use api::v1::{OpType, SemanticType};
use bytes::Bytes;
use common_function::function::FunctionRef;
use common_function::function_factory::ScalarFunctionFactory;
use common_function::scalars::matches::MatchesFunction;
@@ -135,19 +164,23 @@ mod tests {
use datafusion_expr::{BinaryExpr, Expr, Literal, Operator, col, lit};
use datatypes::arrow;
use datatypes::arrow::array::{
ArrayRef, AsArray, BinaryDictionaryBuilder, RecordBatch, StringArray,
StringDictionaryBuilder, TimestampMillisecondArray, UInt8Array, UInt64Array,
Array, ArrayRef, AsArray, BinaryDictionaryBuilder, Float32Array, Float64Array, Int32Array,
RecordBatch, StringArray, StringDictionaryBuilder, TimestampMillisecondArray, UInt8Array,
UInt64Array,
};
use datatypes::arrow::datatypes::{DataType, Field, Schema, UInt32Type};
use datatypes::arrow::util::pretty::pretty_format_batches;
use datatypes::prelude::ConcreteDataType;
use datatypes::schema::{FulltextAnalyzer, FulltextBackend, FulltextOptions};
use object_store::ObjectStore;
use parquet::arrow::AsyncArrowWriter;
use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
use parquet::arrow::{ArrowWriter, AsyncArrowWriter};
use parquet::basic::{Compression, Encoding, ZstdLevel};
use parquet::file::metadata::{KeyValue, PageIndexPolicy};
use parquet::file::properties::WriterProperties;
use parquet::schema::types::ColumnPath;
use store_api::codec::PrimaryKeyEncoding;
use store_api::mito_engine_options::FloatFieldEncoding;
use store_api::metadata::{ColumnMetadata, RegionMetadata, RegionMetadataBuilder};
use store_api::region_request::PathType;
use store_api::storage::{ColumnSchema, RegionId};
@@ -187,6 +220,124 @@ mod tests {
const FILE_DIR: &str = "/";
const REGION_ID: RegionId = RegionId::new(0, 0);
#[test]
fn test_float_field_encoding_properties_and_roundtrip() {
let mut metadata_builder = RegionMetadataBuilder::new(REGION_ID);
metadata_builder
.push_column_metadata(ColumnMetadata {
column_schema: ColumnSchema::new("f32", ConcreteDataType::float32_datatype(), true),
semantic_type: SemanticType::Field,
column_id: 0,
})
.push_column_metadata(ColumnMetadata {
column_schema: ColumnSchema::new("f64", ConcreteDataType::float64_datatype(), true),
semantic_type: SemanticType::Field,
column_id: 1,
})
.push_column_metadata(ColumnMetadata {
column_schema: ColumnSchema::new("tag", ConcreteDataType::float32_datatype(), true),
semantic_type: SemanticType::Tag,
column_id: 2,
})
.push_column_metadata(ColumnMetadata {
column_schema: ColumnSchema::new("i32", ConcreteDataType::int32_datatype(), true),
semantic_type: SemanticType::Field,
column_id: 3,
})
.push_column_metadata(ColumnMetadata {
column_schema: ColumnSchema::new(
"ts",
ConcreteDataType::timestamp_millisecond_datatype(),
false,
),
semantic_type: SemanticType::Timestamp,
column_id: 4,
});
let metadata = Arc::new(metadata_builder.build().unwrap());
let f32_values = [
Some(1.5_f32),
Some(2.0),
Some(0.0),
Some(-0.0),
Some(f32::INFINITY),
Some(f32::from_bits(0x7fc0_1234)),
None,
];
let f64_values = [
Some(1.5_f64),
Some(2.0),
Some(0.0),
Some(-0.0),
Some(f64::NEG_INFINITY),
Some(f64::from_bits(0x7ff8_0000_0000_1234)),
None,
];
let schema = Arc::new(Schema::new(vec![
Field::new("f32", DataType::Float32, true),
Field::new("f64", DataType::Float64, true),
Field::new("tag", DataType::Float32, true),
Field::new("i32", DataType::Int32, true),
Field::new("ts", DataType::Timestamp(TimeUnit::Millisecond, None), false),
]));
let batch = RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(Float32Array::from(f32_values.to_vec())) as ArrayRef,
Arc::new(Float64Array::from(f64_values.to_vec())) as ArrayRef,
Arc::new(Float32Array::from(vec![Some(1.0), None, Some(-0.0), Some(2.0), Some(3.0), Some(4.0), None])),
Arc::new(Int32Array::from(vec![Some(1), None, Some(3), Some(4), Some(5), Some(6), None])),
Arc::new(TimestampMillisecondArray::from_iter_values(0..7)),
],
)
.unwrap();
let path = |name: &str| ColumnPath::new(vec![name.to_string()]);
let bss = apply_float_field_encoding(WriterProperties::builder(), &metadata,
FloatFieldEncoding::ByteStreamSplit);
assert_eq!(Some(Encoding::BYTE_STREAM_SPLIT), bss.clone().build().encoding(&path("f32")));
assert_eq!(Some(Encoding::BYTE_STREAM_SPLIT), bss.clone().build().encoding(&path("f64")));
assert!(!bss.clone().build().dictionary_enabled(&path("f32")));
assert!(!bss.clone().build().dictionary_enabled(&path("f64")));
assert_eq!(None, bss.clone().build().encoding(&path("i32")));
assert!(bss.clone().build().dictionary_enabled(&path("tag")));
let default = apply_float_field_encoding(WriterProperties::builder().set_encoding(Encoding::PLAIN),
&metadata, FloatFieldEncoding::Default).build();
assert_eq!(Some(Encoding::PLAIN), default.encoding(&path("f32")));
assert!(default.dictionary_enabled(&path("f32")));
assert!(default.dictionary_enabled(&path("f64")));
let mut bytes = Vec::new();
let mut writer = ArrowWriter::try_new(&mut bytes, schema, Some(bss.build())).unwrap();
writer.write(&batch).unwrap();
let footer = writer.finish().unwrap();
for name in ["f32", "f64"] {
let column = footer.row_groups()[0].columns().iter()
.find(|column| column.column_path().string() == name).unwrap();
assert!(column.encodings().any(|encoding| encoding == Encoding::BYTE_STREAM_SPLIT));
}
let mut reader = ParquetRecordBatchReaderBuilder::try_new(Bytes::from(bytes))
.unwrap().build().unwrap();
let actual = reader.next().unwrap().unwrap();
let actual_f32 = actual.column(0).as_any().downcast_ref::<Float32Array>().unwrap();
let actual_f64 = actual.column(1).as_any().downcast_ref::<Float64Array>().unwrap();
for (index, value) in f32_values.into_iter().enumerate() {
assert_eq!(value.is_none(), actual_f32.is_null(index));
if let Some(value) = value {
assert_eq!(value.to_bits(), actual_f32.value(index).to_bits());
}
}
for (index, value) in f64_values.into_iter().enumerate() {
assert_eq!(value.is_none(), actual_f64.is_null(index));
if let Some(value) = value {
assert_eq!(value.to_bits(), actual_f64.value(index).to_bits());
}
}
assert!(actual.column(2).is_null(1));
assert!(actual.column(3).is_null(1));
}
#[derive(Clone)]
struct FixedPathProvider {
region_file_id: RegionFileId,
src/mito2/src/sst/parquet/writer.rs
+17 -27
View File
@@ -38,7 +38,7 @@ use object_store::{FuturesAsyncWriter, ObjectStore};
use parquet::arrow::AsyncArrowWriter;
use parquet::basic::{Compression, Encoding, ZstdLevel};
use parquet::file::metadata::KeyValue;
use parquet::file::properties::{WriterProperties, WriterPropertiesBuilder};
use parquet::file::properties::WriterProperties;
use parquet::schema::types::ColumnPath;
use smallvec::smallvec;
use snafu::{OptionExt, ResultExt};
@@ -61,7 +61,9 @@ use crate::sst::parquet::flat_format::{
FlatWriteFormat, primary_key_column_index, time_index_column_index,
};
use crate::sst::parquet::format::{PrimaryKeyArray, PrimaryKeyWriteFormat};
use crate::sst::parquet::{PARQUET_METADATA_KEY, SstInfo, WriteOptions};
use crate::sst::parquet::{
PARQUET_METADATA_KEY, SstInfo, WriteOptions, apply_float_field_encoding,
};
use crate::sst::{
DEFAULT_WRITE_BUFFER_SIZE, DEFAULT_WRITE_CONCURRENCY, FlatSchemaOptions, SeriesEstimator,
maybe_wrap_schema,
@@ -473,29 +475,6 @@ where
Ok(results)
}
/// Customizes per-column config according to schema and maybe column cardinality.
fn customize_column_config(
builder: WriterPropertiesBuilder,
region_metadata: &RegionMetadataRef,
) -> WriterPropertiesBuilder {
let ts_col = ColumnPath::new(vec![
region_metadata
.time_index_column()
.column_schema
.name
.clone(),
]);
let seq_col = ColumnPath::new(vec![SEQUENCE_COLUMN_NAME.to_string()]);
let op_type_col = ColumnPath::new(vec![OP_TYPE_COLUMN_NAME.to_string()]);
builder
.set_column_encoding(seq_col.clone(), Encoding::DELTA_BINARY_PACKED)
.set_column_dictionary_enabled(seq_col, false)
.set_column_encoding(ts_col.clone(), Encoding::DELTA_BINARY_PACKED)
.set_column_dictionary_enabled(ts_col, false)
.set_column_compression(op_type_col, Compression::UNCOMPRESSED)
}
async fn append_flat_batch(
&mut self,
batch: &RecordBatch,
@@ -564,8 +543,19 @@ where
.set_max_row_group_row_count(Some(opts.row_group_size))
.set_column_index_truncate_length(None)
.set_statistics_truncate_length(None);
let props_builder = Self::customize_column_config(props_builder, &self.metadata);
let ts_col = ColumnPath::new(vec![
self.metadata.time_index_column().column_schema.name.clone(),
]);
let seq_col = ColumnPath::new(vec![SEQUENCE_COLUMN_NAME.to_string()]);
let op_type_col = ColumnPath::new(vec![OP_TYPE_COLUMN_NAME.to_string()]);
let props_builder = props_builder
.set_column_encoding(seq_col.clone(), Encoding::DELTA_BINARY_PACKED)
.set_column_dictionary_enabled(seq_col, false)
.set_column_encoding(ts_col.clone(), Encoding::DELTA_BINARY_PACKED)
.set_column_dictionary_enabled(ts_col, false)
.set_column_compression(op_type_col, Compression::UNCOMPRESSED);
let props_builder =
apply_float_field_encoding(props_builder, &self.metadata, opts.float_field_encoding);
let writer_props = props_builder.build();
let sst_file_path = self.path_provider.build_sst_file_path(RegionFileId::new(
src/store-api/src/mito_engine_options.rs
+46
View File
@@ -74,6 +74,31 @@ pub const MAX_ROW_GROUP_ROW_COUNT: &str = "max_row_group_row_count";
pub const MAX_ROW_GROUP_ROW_COUNT_LIMIT: usize = 10 * 1024 * 1024;
/// Option key for preserving per-row sequence numbers through flush and compaction.
pub const PRESERVE_ROW_SEQUENCE: &str = "preserve_row_sequence";
/// Option key for experimental Parquet float field encoding.
pub const EXPERIMENTAL_SST_FLOAT_FIELD_ENCODING: &str = "experimental_sst_float_field_encoding";
/// Encoding policy for direct floating-point field columns in Parquet SSTs.
#[derive(
Debug,
Default,
Clone,
Copy,
PartialEq,
Eq,
Hash,
serde::Serialize,
serde::Deserialize,
strum::EnumString,
)]
#[serde(rename_all = "snake_case")]
#[strum(serialize_all = "snake_case")]
pub enum FloatFieldEncoding {
/// The existing Parquet writer behavior.
#[default]
Default,
/// Parquet byte-stream-split encoding.
ByteStreamSplit,
}
// Note: Adding new options here should also check if this option should be removed in [metric_engine::engine::create::region_options_for_metadata_region].
/// Returns true if the `key` is a valid option key for the mito engine.
@@ -107,6 +132,7 @@ pub fn is_mito_engine_option_key(key: &str) -> bool {
SST_FORMAT_KEY,
MAX_ROW_GROUP_ROW_COUNT,
PRESERVE_ROW_SEQUENCE,
EXPERIMENTAL_SST_FLOAT_FIELD_ENCODING,
]
.contains(&key)
}
@@ -155,6 +181,26 @@ mod tests {
assert!(is_mito_engine_option_key("append_mode"));
assert!(is_mito_engine_option_key("max_row_group_row_count"));
assert!(is_mito_engine_option_key("preserve_row_sequence"));
assert!(is_mito_engine_option_key(
EXPERIMENTAL_SST_FLOAT_FIELD_ENCODING
));
assert!(!is_mito_engine_option_key("foo"));
}
#[test]
fn test_float_field_encoding_values() {
assert_eq!(
"default".parse::<FloatFieldEncoding>(),
Ok(FloatFieldEncoding::Default)
);
assert_eq!(
"byte_stream_split".parse::<FloatFieldEncoding>(),
Ok(FloatFieldEncoding::ByteStreamSplit)
);
assert_eq!(
serde_json::from_str::<FloatFieldEncoding>("\"byte_stream_split\"").unwrap(),
FloatFieldEncoding::ByteStreamSplit
);
assert!(serde_json::from_str::<FloatFieldEncoding>("\"unknown\"").is_err());
}
}
src/table/src/requests.rs
+36 -5
View File
@@ -36,11 +36,11 @@ use store_api::metric_engine_consts::{
LOGICAL_TABLE_METADATA_KEY, PHYSICAL_TABLE_METADATA_KEY, is_metric_engine_option_key,
};
use store_api::mito_engine_options::{
APPEND_MODE_KEY, COMPACTION_TYPE, MEMTABLE_BULK_ENCODE_BYTES_THRESHOLD,
MEMTABLE_BULK_ENCODE_ROW_THRESHOLD, MEMTABLE_BULK_MAX_MERGE_GROUPS,
MEMTABLE_BULK_MERGE_THRESHOLD, MEMTABLE_TYPE, MERGE_MODE_KEY, SST_FORMAT_KEY,
TWCS_FALLBACK_TO_LOCAL, TWCS_MAX_OUTPUT_FILE_SIZE, TWCS_TIME_WINDOW, TWCS_TRIGGER_FILE_NUM,
is_mito_engine_option_key,
APPEND_MODE_KEY, COMPACTION_TYPE, EXPERIMENTAL_SST_FLOAT_FIELD_ENCODING, FloatFieldEncoding,
MEMTABLE_BULK_ENCODE_BYTES_THRESHOLD, MEMTABLE_BULK_ENCODE_ROW_THRESHOLD,
MEMTABLE_BULK_MAX_MERGE_GROUPS, MEMTABLE_BULK_MERGE_THRESHOLD, MEMTABLE_TYPE, MERGE_MODE_KEY,
SST_FORMAT_KEY, TWCS_FALLBACK_TO_LOCAL, TWCS_MAX_OUTPUT_FILE_SIZE, TWCS_TIME_WINDOW,
TWCS_TRIGGER_FILE_NUM, is_mito_engine_option_key,
};
use store_api::region_request::{SetRegionOption, UnsetRegionOption};
@@ -210,6 +210,16 @@ impl TableOptions {
options.ttl = Some(ttl_value);
}
if let Some(encoding) = kvs.get(EXPERIMENTAL_SST_FLOAT_FIELD_ENCODING) {
encoding.parse::<FloatFieldEncoding>().map_err(|_| {
ParseTableOptionSnafu {
key: EXPERIMENTAL_SST_FLOAT_FIELD_ENCODING,
value: encoding,
}
.build()
})?;
}
if let Some(skip_wal) = kvs.get(SKIP_WAL_KEY) {
options.skip_wal = skip_wal.parse().map_err(|_| {
ParseTableOptionSnafu {
@@ -727,6 +737,7 @@ mod tests {
assert!(validate_table_option(WRITE_BUFFER_SIZE_KEY));
assert!(validate_table_option(STORAGE_KEY));
assert!(validate_table_option(MEMTABLE_BULK_MERGE_THRESHOLD));
assert!(validate_table_option(EXPERIMENTAL_SST_FLOAT_FIELD_ENCODING));
assert!(validate_table_option(REPARTITION_COLUMN_HINT_KEY));
assert!(!validate_table_option("foo"));
@@ -751,6 +762,26 @@ mod tests {
assert!(!validate_database_option("foo"));
}
#[test]
fn test_parse_float_field_encoding_option() {
let options = TableOptions::try_from_iter([(
EXPERIMENTAL_SST_FLOAT_FIELD_ENCODING,
"byte_stream_split",
)])
.unwrap();
assert_eq!(
options
.extra_options
.get(EXPERIMENTAL_SST_FLOAT_FIELD_ENCODING),
Some(&"byte_stream_split".to_string())
);
assert!(TableOptions::try_from_iter([(
EXPERIMENTAL_SST_FLOAT_FIELD_ENCODING,
"invalid",
)])
.is_err());
}
#[test]
fn test_serialize_table_options() {
let options = TableOptions {