diff --git a/src/mito2/src/compaction/compactor.rs b/src/mito2/src/compaction/compactor.rs index 8ad1cdf0b22..fb2172b4875 100644 --- a/src/mito2/src/compaction/compactor.rs +++ b/src/mito2/src/compaction/compactor.rs @@ -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(); diff --git a/src/mito2/src/compaction/window.rs b/src/mito2/src/compaction/window.rs index a9f3f021b84..b9053466444 100644 --- a/src/mito2/src/compaction/window.rs +++ b/src/mito2/src/compaction/window.rs @@ -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, } diff --git a/src/mito2/src/flush.rs b/src/mito2/src/flush.rs index b7c8513afc7..c0dee32bd50 100644 --- a/src/mito2/src/flush.rs +++ b/src/mito2/src/flush.rs @@ -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 { diff --git a/src/mito2/src/memtable.rs b/src/mito2/src/memtable.rs index 8c107d15f79..8d98ec64340 100644 --- a/src/mito2/src/memtable.rs +++ b/src/mito2/src/memtable.rs @@ -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 { diff --git a/src/mito2/src/memtable/bulk.rs b/src/mito2/src/memtable/bulk.rs index 1b67d487a5f..57a7fcb07e6 100644 --- a/src/mito2/src/memtable/bulk.rs +++ b/src/mito2/src/memtable/bulk.rs @@ -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, + compact_dispatcher: Option>, + 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::>>>()?; @@ -1257,6 +1294,7 @@ impl MemtableCompactor { encode_row_threshold: usize, encode_bytes_threshold: usize, row_group_size: usize, + float_field_encoding: FloatFieldEncoding, ) -> Result> { 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) -> 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, )) } diff --git a/src/mito2/src/memtable/bulk/part.rs b/src/mito2/src/memtable/bulk/part.rs index 2a0c42b7035..bbf80a1571f 100644 --- a/src/mito2/src/memtable/bulk/part.rs +++ b/src/mito2/src/memtable/bulk/part.rs @@ -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 { + 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 { // 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, diff --git a/src/mito2/src/region/options.rs b/src/mito2/src/region/options.rs index 48480f34754..90c49a913b3 100644 --- a/src/mito2/src/region/options.rs +++ b/src/mito2/src/region/options.rs @@ -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, #[serde_as(as = "DisplayFromStr")] preserve_row_sequence: bool, + #[serde(rename = "experimental_sst_float_field_encoding")] + #[serde_as(as = "NoneAsEmptyString")] + float_field_encoding: Option, } 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); } diff --git a/src/mito2/src/sst/parquet.rs b/src/mito2/src/sst/parquet.rs index 3d90a360925..1f809e9bc97 100644 --- a/src/mito2/src/sst/parquet.rs +++ b/src/mito2/src/sst/parquet.rs @@ -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, + /// 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::().unwrap(); + let actual_f64 = actual.column(1).as_any().downcast_ref::().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, diff --git a/src/mito2/src/sst/parquet/writer.rs b/src/mito2/src/sst/parquet/writer.rs index b9d5016173e..9124fd4bf65 100644 --- a/src/mito2/src/sst/parquet/writer.rs +++ b/src/mito2/src/sst/parquet/writer.rs @@ -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( diff --git a/src/store-api/src/mito_engine_options.rs b/src/store-api/src/mito_engine_options.rs index 6890f06cc32..cab6eacc8a5 100644 --- a/src/store-api/src/mito_engine_options.rs +++ b/src/store-api/src/mito_engine_options.rs @@ -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::(), + Ok(FloatFieldEncoding::Default) + ); + assert_eq!( + "byte_stream_split".parse::(), + Ok(FloatFieldEncoding::ByteStreamSplit) + ); + assert_eq!( + serde_json::from_str::("\"byte_stream_split\"").unwrap(), + FloatFieldEncoding::ByteStreamSplit + ); + assert!(serde_json::from_str::("\"unknown\"").is_err()); + } } diff --git a/src/table/src/requests.rs b/src/table/src/requests.rs index aea45130ed9..76726b1f041 100644 --- a/src/table/src/requests.rs +++ b/src/table/src/requests.rs @@ -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::().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 {