diff --git a/src/metric-engine/src/engine/bulk_insert.rs b/src/metric-engine/src/engine/bulk_insert.rs index d700d858acf..a2b39d10c9b 100644 --- a/src/metric-engine/src/engine/bulk_insert.rs +++ b/src/metric-engine/src/engine/bulk_insert.rs @@ -241,14 +241,20 @@ mod tests { use common_error::ext::ErrorExt; use common_query::prelude::{greptime_timestamp, greptime_value}; use common_recordbatch::RecordBatches; + use datafusion::parquet::basic::Encoding; use datatypes::arrow::array::{Float64Array, StringArray, TimestampMillisecondArray}; use datatypes::arrow::datatypes::{DataType, Field, Schema as ArrowSchema, TimeUnit}; use datatypes::arrow::record_batch::RecordBatch; use mito2::config::MitoConfig; + use mito2::sst::parquet::metadata::MetadataLoader; + use mito2::sst::parquet::reader::MetadataCacheMetrics; use store_api::metric_engine_consts::PRIMARY_KEY_ENCODING; + use store_api::mito_engine_options::EXPERIMENTAL_SST_FLOAT_FIELD_ENCODING; use store_api::path_utils::table_dir; use store_api::region_engine::RegionEngine; - use store_api::region_request::{RegionBulkInsertsRequest, RegionPutRequest, RegionRequest}; + use store_api::region_request::{ + RegionBulkInsertsRequest, RegionFlushRequest, RegionPutRequest, RegionRequest, + }; use store_api::storage::{RegionId, ScanRequest}; use super::record_batch_to_ipc; @@ -550,16 +556,68 @@ mod tests { #[tokio::test] async fn test_bulk_insert_sparse_encoding() { let env = TestEnv::new().await; - env.init_metric_region().await; + let physical_region_id = env.default_physical_region_id(); + env.create_physical_region( + physical_region_id, + &TestEnv::default_table_dir(), + vec![( + EXPERIMENTAL_SST_FLOAT_FIELD_ENCODING.to_string(), + "byte_stream_split".to_string(), + )], + ) + .await; let logical_region_id = env.default_logical_region_id(); + env.create_logical_region(physical_region_id, logical_region_id) + .await; - let request = build_bulk_request(logical_region_id, build_logical_batch(0, 4)); + let rows = 4; + let request = build_bulk_request(logical_region_id, build_logical_batch(0, rows)); let response = env .metric() .handle_request(logical_region_id, request) .await .unwrap(); - assert_eq!(response.affected_rows, 4); + assert_eq!(response.affected_rows, rows); + + let data_region_id = crate::utils::to_data_region_id(physical_region_id); + env.mito() + .handle_request( + data_region_id, + RegionRequest::Flush(RegionFlushRequest::default()), + ) + .await + .unwrap(); + + let region = env.mito().find_region(data_region_id).unwrap(); + let entry = region + .manifest_sst_entries() + .await + .into_iter() + .find(|entry| entry.visible && entry.file_path.ends_with(".parquet")) + .unwrap(); + let mut cache_metrics = MetadataCacheMetrics::default(); + let footer = MetadataLoader::new( + region.access_layer().object_store().clone(), + &entry.file_path, + entry.file_size, + ) + .load(&mut cache_metrics) + .await + .unwrap(); + let field_column_index = footer + .file_metadata() + .schema_descr() + .columns() + .iter() + .position(|column| column.name() == greptime_value()) + .unwrap(); + assert!(!footer.row_groups().is_empty()); + assert!(footer.row_groups().iter().all(|row_group| { + row_group + .column(field_column_index) + .encodings() + .any(|encoding| encoding == Encoding::BYTE_STREAM_SPLIT) + })); let stream = env .metric() @@ -567,7 +625,7 @@ mod tests { .await .unwrap(); let batches = RecordBatches::try_collect(stream).await.unwrap(); - assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 4); + assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), rows); } #[tokio::test]