From 977e76d0a4100ce9b6f2ccd9e171769ba362428a Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Tue, 20 Oct 2026 07:46:56 +0000 Subject: [PATCH] fix(mito2): apply SST internal column encodings to bulk-encoded parts (#9735) * fix(mito2): apply SST internal column encodings to bulk-encoded parts Bulk memtable parts are encoded to parquet and flushed as SSTs as-is, but the bulk encoder did not apply the time index / __sequence DELTA_BINARY_PACKED encoding (with dictionary disabled) and the uncompressed op_type setting that the regular SST writer uses. Share these settings through one helper. Signed-off-by: Dennis Zhuang * chore(mito2): drop the answered bulk encoder compression TODO Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/mito2/src/memtable/bulk.rs | 42 ++++++++++++++++++++++++++++- src/mito2/src/memtable/bulk/part.rs | 4 +-- src/mito2/src/sst/parquet.rs | 19 +++++++++++++ src/mito2/src/sst/parquet/writer.rs | 16 ++--------- 4 files changed, 64 insertions(+), 17 deletions(-) diff --git a/src/mito2/src/memtable/bulk.rs b/src/mito2/src/memtable/bulk.rs index 7402eeb0582..847a12db5ae 100644 --- a/src/mito2/src/memtable/bulk.rs +++ b/src/mito2/src/memtable/bulk.rs @@ -1692,9 +1692,10 @@ mod tests { use datatypes::schema::ColumnSchema; use datatypes::types::json_type::{JsonNativeType, JsonObjectType}; use mito_codec::row_converter::build_primary_key_codec; - use parquet::basic::{Encoding, Type as PhysicalType}; + use parquet::basic::{Compression, Encoding, Type as PhysicalType}; use serde_json::json; use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder, RegionMetadataRef}; + use store_api::storage::consts::{OP_TYPE_COLUMN_NAME, SEQUENCE_COLUMN_NAME}; use super::*; use crate::memtable::bulk::part::BulkPartConverter; @@ -3060,6 +3061,45 @@ mod tests { Ok(()) } + #[test] + fn test_bulk_part_encoder_matches_sst_writer_internal_encodings() -> WhateverResult<()> { + let metadata = metadata_for_test(); + let bulk_part = create_bulk_part_with_converter( + "key", + 0, + vec![1000, 2000, 3000], + vec![Some(1.0), Some(2.0), Some(3.0)], + 100, + )?; + let encoder = BulkPartEncoder::new(metadata, DEFAULT_ROW_GROUP_SIZE)?; + let encoded_part = encoder.encode_part(&bulk_part)?.unwrap(); + let columns = encoded_part + .metadata() + .parquet_metadata + .row_group(0) + .columns(); + + for name in ["ts", SEQUENCE_COLUMN_NAME] { + let column = columns + .iter() + .find(|column| column.column_path().string() == name) + .unwrap(); + assert!( + column + .encodings() + .any(|encoding| encoding == Encoding::DELTA_BINARY_PACKED), + "{name}" + ); + assert!(column.dictionary_page_offset().is_none(), "{name}"); + } + let op_type = columns + .iter() + .find(|column| column.column_path().string() == OP_TYPE_COLUMN_NAME) + .unwrap(); + assert_eq!(Compression::UNCOMPRESSED, op_type.compression()); + Ok(()) + } + #[test] fn test_encoded_part_batch_size_uses_largest_uncompressed_row_group() { const NUM_ROWS: usize = 14; diff --git a/src/mito2/src/memtable/bulk/part.rs b/src/mito2/src/memtable/bulk/part.rs index 27fc7983513..9bcc4238b45 100644 --- a/src/mito2/src/memtable/bulk/part.rs +++ b/src/mito2/src/memtable/bulk/part.rs @@ -74,7 +74,7 @@ use crate::sst::parquet::flat_format::primary_key_column_index; use crate::sst::parquet::format::{PrimaryKeyArray, PrimaryKeyArrayBuilder}; use crate::sst::parquet::{ COLUMN_INDEX_TRUNCATE_LENGTH, PARQUET_METADATA_KEY, SstInfo, apply_column_compression, - apply_float_field_encoding, + apply_float_field_encoding, apply_internal_column_encoding, }; const INIT_DICT_VALUE_CAPACITY: usize = 8; @@ -1339,7 +1339,6 @@ impl BulkPartEncoder { let key_value_meta = parquet::file::metadata::KeyValue::new(PARQUET_METADATA_KEY.to_string(), json); - // TODO(yingwen): Do we need compression? let mut props = WriterProperties::builder() .set_key_value_metadata(Some(vec![key_value_meta])) .set_write_batch_size(row_group_size) @@ -1347,6 +1346,7 @@ impl BulkPartEncoder { .set_compression(Compression::ZSTD(ZstdLevel::default())) .set_column_index_truncate_length(COLUMN_INDEX_TRUNCATE_LENGTH) .set_statistics_truncate_length(None); + props = apply_internal_column_encoding(props, &metadata); props = apply_float_field_encoding(props, &metadata, float_field_encoding); props = apply_column_compression(props, &metadata)?; let writer_props = Some(props.build()); diff --git a/src/mito2/src/sst/parquet.rs b/src/mito2/src/sst/parquet.rs index 79ad957b75a..b1b459a595a 100644 --- a/src/mito2/src/sst/parquet.rs +++ b/src/mito2/src/sst/parquet.rs @@ -28,6 +28,7 @@ use parquet::schema::types::ColumnPath; use snafu::ResultExt; use store_api::metadata::RegionMetadataRef; use store_api::mito_engine_options::FloatFieldEncoding; +use store_api::storage::consts::{OP_TYPE_COLUMN_NAME, SEQUENCE_COLUMN_NAME}; use store_api::storage::{ColumnId, FileId}; use crate::error::{CompressionOptionsSnafu, Result, WriteParquetSnafu}; @@ -93,6 +94,24 @@ pub const DEFAULT_ROW_GROUP_SIZE: usize = 100 * 1024; pub(crate) const COLUMN_INDEX_TRUNCATE_LENGTH: Option = parquet::file::properties::DEFAULT_COLUMN_INDEX_TRUNCATE_LENGTH; +/// Applies the encodings of the time index and internal columns shared by every SST writer. +pub(crate) fn apply_internal_column_encoding( + builder: WriterPropertiesBuilder, + metadata: &RegionMetadataRef, +) -> WriterPropertiesBuilder { + let ts_col = ColumnPath::new(vec![ + 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) +} + /// Applies the per-column options declared with `COMPRESSION WITH (...)`. pub(crate) fn apply_column_compression( mut builder: WriterPropertiesBuilder, diff --git a/src/mito2/src/sst/parquet/writer.rs b/src/mito2/src/sst/parquet/writer.rs index 9d2defcceb3..69a89e5802b 100644 --- a/src/mito2/src/sst/parquet/writer.rs +++ b/src/mito2/src/sst/parquet/writer.rs @@ -39,11 +39,9 @@ use parquet::arrow::AsyncArrowWriter; use parquet::basic::{Compression, Encoding, ZstdLevel}; use parquet::file::metadata::KeyValue; use parquet::file::properties::WriterProperties; -use parquet::schema::types::ColumnPath; use smallvec::smallvec; use snafu::{OptionExt, ResultExt}; use store_api::metadata::RegionMetadataRef; -use store_api::storage::consts::{OP_TYPE_COLUMN_NAME, SEQUENCE_COLUMN_NAME}; use store_api::storage::{FileId, SequenceNumber}; use tokio::io::AsyncWrite; use tokio_util::compat::{Compat, FuturesAsyncWriteCompatExt}; @@ -63,7 +61,7 @@ use crate::sst::parquet::flat_format::{ use crate::sst::parquet::format::{PrimaryKeyArray, PrimaryKeyWriteFormat}; use crate::sst::parquet::{ COLUMN_INDEX_TRUNCATE_LENGTH, PARQUET_METADATA_KEY, SstInfo, WriteOptions, - apply_column_compression, apply_float_field_encoding, + apply_column_compression, apply_float_field_encoding, apply_internal_column_encoding, }; use crate::sst::{ DEFAULT_WRITE_CONCURRENCY, FlatSchemaOptions, SeriesEstimator, maybe_wrap_schema, @@ -547,17 +545,7 @@ where .set_max_row_group_row_count(Some(opts.row_group_size)) .set_column_index_truncate_length(COLUMN_INDEX_TRUNCATE_LENGTH) .set_statistics_truncate_length(None); - 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_internal_column_encoding(props_builder, &self.metadata); let props_builder = apply_float_field_encoding( props_builder, &self.metadata,