From 74440ca050691ed00492d9dd37a3f362df1a8a83 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Fri, 4 Sep 2026 16:52:51 +0800 Subject: [PATCH] fix(mito2): correct float encoding checks Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/mito2/src/memtable/bulk/part.rs | 4 +- src/mito2/src/region/options.rs | 11 ++-- src/mito2/src/sst/parquet.rs | 85 +++++++++++++++++++++++------ src/mito2/src/sst/parquet/writer.rs | 7 ++- src/table/src/requests.rs | 9 ++- 5 files changed, 83 insertions(+), 33 deletions(-) diff --git a/src/mito2/src/memtable/bulk/part.rs b/src/mito2/src/memtable/bulk/part.rs index bbf80a1571f..1d015bdeaf1 100644 --- a/src/mito2/src/memtable/bulk/part.rs +++ b/src/mito2/src/memtable/bulk/part.rs @@ -71,9 +71,7 @@ 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, apply_float_field_encoding, -}; +use crate::sst::parquet::{PARQUET_METADATA_KEY, SstInfo, apply_float_field_encoding}; const INIT_DICT_VALUE_CAPACITY: usize = 8; diff --git a/src/mito2/src/region/options.rs b/src/mito2/src/region/options.rs index 90c49a913b3..f0d32e1daf9 100644 --- a/src/mito2/src/region/options.rs +++ b/src/mito2/src/region/options.rs @@ -33,8 +33,7 @@ use store_api::metric_engine_consts::{ MEMTABLE_PARTITION_TREE_PRIMARY_KEY_ENCODING, PRIMARY_KEY_ENCODING, }; use store_api::mito_engine_options::{ - COMPACTION_OVERRIDE, EXPERIMENTAL_SST_FLOAT_FIELD_ENCODING, FloatFieldEncoding, - MAX_ROW_GROUP_ROW_COUNT_LIMIT, + COMPACTION_OVERRIDE, FloatFieldEncoding, MAX_ROW_GROUP_ROW_COUNT_LIMIT, }; use store_api::storage::{ColumnId, RegionId}; use strum::EnumString; @@ -279,9 +278,7 @@ impl RegionOptions { sst_format = Some(FormatType::Flat); } - let float_field_encoding = options - .float_field_encoding - .unwrap_or_default(); + let float_field_encoding = options.float_field_encoding.unwrap_or_default(); let compaction_override_flag = options_map .get(COMPACTION_OVERRIDE) @@ -598,7 +595,9 @@ mod tests { use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; use common_wal::options::KafkaWalOptions; - use store_api::mito_engine_options::{SKIP_WAL_KEY, WRITE_BUFFER_SIZE_KEY}; + use store_api::mito_engine_options::{ + EXPERIMENTAL_SST_FLOAT_FIELD_ENCODING, SKIP_WAL_KEY, WRITE_BUFFER_SIZE_KEY, + }; use super::*; diff --git a/src/mito2/src/sst/parquet.rs b/src/mito2/src/sst/parquet.rs index 1f809e9bc97..d5c6adfc8b3 100644 --- a/src/mito2/src/sst/parquet.rs +++ b/src/mito2/src/sst/parquet.rs @@ -17,12 +17,13 @@ use std::collections::BTreeMap; use std::sync::Arc; +use api::v1::SemanticType; 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::metadata::RegionMetadataRef; use store_api::mito_engine_options::FloatFieldEncoding; use store_api::storage::{ColumnId, FileId}; @@ -180,8 +181,8 @@ mod tests { 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::mito_engine_options::FloatFieldEncoding; use store_api::region_request::PathType; use store_api::storage::{ColumnSchema, RegionId}; use table::predicate::Predicate; @@ -278,32 +279,65 @@ mod tests { 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), + 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(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"))); + 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(); + 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"))); @@ -313,15 +347,32 @@ mod tests { 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 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(); + .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(); + 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 { diff --git a/src/mito2/src/sst/parquet/writer.rs b/src/mito2/src/sst/parquet/writer.rs index 9124fd4bf65..9888b4ead53 100644 --- a/src/mito2/src/sst/parquet/writer.rs +++ b/src/mito2/src/sst/parquet/writer.rs @@ -554,8 +554,11 @@ where .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 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/table/src/requests.rs b/src/table/src/requests.rs index 76726b1f041..0e05153c377 100644 --- a/src/table/src/requests.rs +++ b/src/table/src/requests.rs @@ -775,11 +775,10 @@ mod tests { .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()); + assert!( + TableOptions::try_from_iter([(EXPERIMENTAL_SST_FLOAT_FIELD_ENCODING, "invalid",)]) + .is_err() + ); } #[test]