From 742e47a534cb2d4a8b00e9323a43ba8bb7721105 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Wed, 30 Sep 2026 08:40:27 +0000 Subject: [PATCH] fix(mito2): truncate parquet column index min/max for SST writes (#9420) * fix(mito2): truncate parquet column index min/max for SST writes The SST writer and the bulk part encoder disabled column index truncation together with statistics truncation in #6977. Mito never reads the column index, but every page of a large string or binary column still stored its full min and max values, uncompressed. Truncate the column index to the parquet default of 64 bytes and keep chunk statistics untruncated, since primary key pruning decodes them. Signed-off-by: Dennis Zhuang * docs(mito2): note that the column index must not be used to decode primary keys Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang (cherry picked from commit a314ac428144b32decdd85866011e3d29f0025b0) Signed-off-by: WenyXu --- src/mito2/src/memtable/bulk/part.rs | 6 +- src/mito2/src/sst/parquet.rs | 114 ++++++++++++++++++++++++++++ src/mito2/src/sst/parquet/writer.rs | 5 +- 3 files changed, 121 insertions(+), 4 deletions(-) diff --git a/src/mito2/src/memtable/bulk/part.rs b/src/mito2/src/memtable/bulk/part.rs index 7e50fbbd4cb..444cc59e317 100644 --- a/src/mito2/src/memtable/bulk/part.rs +++ b/src/mito2/src/memtable/bulk/part.rs @@ -71,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, apply_float_field_encoding}; +use crate::sst::parquet::{ + COLUMN_INDEX_TRUNCATE_LENGTH, PARQUET_METADATA_KEY, SstInfo, apply_float_field_encoding, +}; const INIT_DICT_VALUE_CAPACITY: usize = 8; @@ -1334,7 +1336,7 @@ impl BulkPartEncoder { .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_column_index_truncate_length(COLUMN_INDEX_TRUNCATE_LENGTH) .set_statistics_truncate_length(None); props = apply_float_field_encoding(props, &metadata, float_field_encoding); let writer_props = Some(props.build()); diff --git a/src/mito2/src/sst/parquet.rs b/src/mito2/src/sst/parquet.rs index 301842bde72..9404d4484fc 100644 --- a/src/mito2/src/sst/parquet.rs +++ b/src/mito2/src/sst/parquet.rs @@ -77,6 +77,18 @@ pub(crate) struct Json2TargetLayout { /// batching without changing the row group layout of newly written SSTs. pub const DEFAULT_ROW_GROUP_SIZE: usize = 100 * 1024; +/// Truncation length of min/max values in the parquet column index. +/// +/// Mito never reads the column index (it is skipped or stripped on load), so truncating +/// it is safe. Chunk statistics stay untruncated because pruning decodes primary keys +/// from them. Without truncation every page of a large string column stores its whole +/// min and max values, uncompressed. +/// +/// Truncated values are only bounds: code that starts reading the column index must not +/// decode primary keys from it. +pub(crate) const COLUMN_INDEX_TRUNCATE_LENGTH: Option = + parquet::file::properties::DEFAULT_COLUMN_INDEX_TRUNCATE_LENGTH; + /// Applies the configured encoding to direct floating-point field columns. pub(crate) fn apply_float_field_encoding( mut builder: WriterPropertiesBuilder, @@ -901,6 +913,82 @@ mod tests { .await; } + #[tokio::test] + async fn test_column_index_truncates_large_field_values() { + let mut env = TestEnv::new().await; + let object_store = env.init_object_store_manager(); + let handle = sst_file_handle(0, 1000); + let file_path = FixedPathProvider { + region_file_id: handle.file_id(), + }; + let metadata = build_test_binary_test_region_metadata(); + // A tag long enough that the encoded primary key exceeds the 64-byte truncation length. + let tag = "t".repeat(100); + let values: Vec> = (0..64) + .map(|i| format!("{i:08}").into_bytes().repeat(4 * 1024)) + .collect(); + let batch = new_record_batch_with_binary_values(&tag, &values); + let mut metrics = Metrics::new(WriteType::Flush); + let mut writer = ParquetWriter::new_with_object_store( + object_store.clone(), + metadata.clone(), + IndexConfig::default(), + NoopIndexBuilder, + file_path, + &mut metrics, + ) + .await; + writer + .write_all_flat_as_primary_key( + new_flat_source_from_record_batches(vec![batch]), + None, + &WriteOptions::default(), + ) + .await + .unwrap(); + + let path = handle.file_path(FILE_DIR, PathType::Bare); + let bytes = object_store.read(&path).await.unwrap().to_bytes(); + let options = parquet::arrow::arrow_reader::ArrowReaderOptions::new() + .with_page_index_policy(PageIndexPolicy::Required); + let builder = + ParquetRecordBatchReaderBuilder::try_new_with_options(bytes, options).unwrap(); + let parquet_meta = builder.metadata().clone(); + let schema = parquet_meta.file_metadata().schema_descr(); + let column = |name: &str| { + (0..schema.num_columns()) + .find(|i| schema.column(*i).name() == name) + .unwrap() + }; + + let field_index = &parquet_meta.column_index().unwrap()[0][column("field_0")]; + let parquet::file::page_index::column_index::ColumnIndexMetaData::BYTE_ARRAY(field_index) = + field_index + else { + panic!("unexpected column index: {field_index:?}"); + }; + let mut pages = 0; + for (min, max) in field_index + .min_values_iter() + .zip(field_index.max_values_iter()) + { + assert!(min.unwrap().len() <= 64); + assert!(max.unwrap().len() <= 64); + pages += 1; + } + assert!(pages > 1, "values should span several pages"); + + // Chunk statistics keep exact primary keys; pruning decodes them. + let pk_stats = parquet_meta + .row_group(0) + .column(column(store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME)) + .statistics() + .unwrap(); + let pk = new_primary_key(&[&tag]); + assert!(pk.len() > 64); + assert_eq!(pk_stats.max_bytes_opt().unwrap(), pk.as_slice()); + } + #[rstest::rstest] #[tokio::test] async fn test_write_multiple_files(#[values(1024, 4096)] write_buffer_size: usize) { @@ -1475,6 +1563,32 @@ mod tests { assert!(cached.contains_row_group(3)); } + fn new_record_batch_with_binary_values(tag: &str, values: &[Vec]) -> RecordBatch { + let metadata = build_test_binary_test_region_metadata(); + let flat_schema = to_flat_sst_arrow_schema(&metadata, &FlatSchemaOptions::default()); + let num_rows = values.len(); + let mut tag_0_builder = StringDictionaryBuilder::::new(); + let mut pk_builder = BinaryDictionaryBuilder::::new(); + let pk = new_primary_key(&[tag]); + for _ in 0..num_rows { + tag_0_builder.append_value(tag); + pk_builder.append(&pk).unwrap(); + } + let columns: Vec = vec![ + Arc::new(tag_0_builder.finish()), + Arc::new(datatypes::arrow::array::BinaryArray::from_iter_values( + values.iter().map(|v| v.as_slice()), + )), + Arc::new(TimestampMillisecondArray::from_iter_values( + 0..num_rows as i64, + )), + Arc::new(pk_builder.finish()), + Arc::new(UInt64Array::from_value(1000, num_rows)), + Arc::new(UInt8Array::from_value(OpType::Put as u8, num_rows)), + ]; + RecordBatch::try_new(flat_schema, columns).unwrap() + } + fn new_record_batch_with_binary(tags: &[&str], start: usize, end: usize) -> RecordBatch { assert!(end >= start); let metadata = build_test_binary_test_region_metadata(); diff --git a/src/mito2/src/sst/parquet/writer.rs b/src/mito2/src/sst/parquet/writer.rs index b9ef15add34..932654dfdd8 100644 --- a/src/mito2/src/sst/parquet/writer.rs +++ b/src/mito2/src/sst/parquet/writer.rs @@ -62,7 +62,8 @@ use crate::sst::parquet::flat_format::{ }; use crate::sst::parquet::format::{PrimaryKeyArray, PrimaryKeyWriteFormat}; use crate::sst::parquet::{ - PARQUET_METADATA_KEY, SstInfo, WriteOptions, apply_float_field_encoding, + COLUMN_INDEX_TRUNCATE_LENGTH, PARQUET_METADATA_KEY, SstInfo, WriteOptions, + apply_float_field_encoding, }; use crate::sst::{ DEFAULT_WRITE_CONCURRENCY, FlatSchemaOptions, SeriesEstimator, maybe_wrap_schema, @@ -544,7 +545,7 @@ where .set_compression(Compression::ZSTD(ZstdLevel::default())) .set_encoding(Encoding::PLAIN) .set_max_row_group_row_count(Some(opts.row_group_size)) - .set_column_index_truncate_length(None) + .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(),