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 <killme2008@gmail.com>

* chore(mito2): drop the answered bulk encoder compression TODO

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>

---------

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>
This commit is contained in:
dennis zhuang
2026-10-20 07:46:56 +00:00
committed by GitHub
parent 41ec5e0763
commit 977e76d0a4
4 changed files with 64 additions and 17 deletions
+41 -1
View File
@@ -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;
+2 -2
View File
@@ -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());
+19
View File
@@ -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<usize> =
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,
+2 -14
View File
@@ -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,