From 76efeb5e2aecea2ec17842fd6e10544d20c668d1 Mon Sep 17 00:00:00 2001 From: Hengfei Yang Date: Tue, 4 Jun 2024 21:37:52 +0800 Subject: [PATCH] fix: reduce CPU const for merging (#3673) --- src/config/src/utils/record_batch_ext.rs | 31 ++++---- src/infra/src/db/postgres.rs | 5 ++ src/ingester/src/writer.rs | 2 +- src/job/files/parquet.rs | 4 +- src/service/compact/merge.rs | 96 ++++++++++++++++++++++-- src/service/search/mod.rs | 11 +-- 6 files changed, 120 insertions(+), 29 deletions(-) diff --git a/src/config/src/utils/record_batch_ext.rs b/src/config/src/utils/record_batch_ext.rs index 789bffdf46..aeebb985e5 100644 --- a/src/config/src/utils/record_batch_ext.rs +++ b/src/config/src/utils/record_batch_ext.rs @@ -484,6 +484,7 @@ pub fn format_recordbatch_by_schema(schema: Arc, batch: RecordBatch) -> pub fn concat_batches( schema: Arc, mut batches: Vec, + less_memory_mode: bool, ) -> Result { // When schema is empty, sum the number of the rows of all batches if schema.fields().is_empty() { @@ -492,25 +493,29 @@ pub fn concat_batches( options.row_count = Some(num_rows); return RecordBatch::try_new_with_options(schema.clone(), vec![], &options); } - if batches.is_empty() { return Ok(RecordBatch::new_empty(schema.clone())); } + let field_num = schema.fields().len(); let mut arrays = Vec::with_capacity(field_num); for i in 0..field_num { - // let array = arrow::compute::concat( - // &batches.iter() - // .map(|batch| batch.column(i).as_ref()) - // .collect::>(), - // )?; - let mut array = Vec::with_capacity(batches.len()); - for batch in &mut batches { - let i = i - arrays.len(); - let column = batch.remove_column(i); - array.push(column); - } - let array = arrow::compute::concat(&array.iter().map(|c| c.as_ref()).collect::>())?; + let array = if less_memory_mode { + let mut new_batches = Vec::with_capacity(batches.len()); + for batch in &mut batches { + let i = i - arrays.len(); + let column = batch.remove_column(i); + new_batches.push(column); + } + arrow::compute::concat(&new_batches.iter().map(|c| c.as_ref()).collect::>())? + } else { + arrow::compute::concat( + &batches + .iter() + .map(|batch| batch.column(i).as_ref()) + .collect::>(), + )? + }; arrays.push(array); } RecordBatch::try_new(schema.clone(), arrays) diff --git a/src/infra/src/db/postgres.rs b/src/infra/src/db/postgres.rs index e922766813..191e3df749 100644 --- a/src/infra/src/db/postgres.rs +++ b/src/infra/src/db/postgres.rs @@ -164,6 +164,11 @@ impl super::Db for PostgresDb { let mut tx = pool.begin().await?; let lock_key = format!("get_for_update_{}", key); let lock_id = config::utils::hash::gxhash::new().sum64(&lock_key); + let lock_id = if lock_id > i64::MAX as u64 { + (lock_id >> 1) as i64 + } else { + lock_id as i64 + }; let lock_sql = format!("SELECT pg_advisory_xact_lock({lock_id})"); if let Err(e) = sqlx::query(&lock_sql).execute(&mut *tx).await { if let Err(e) = tx.rollback().await { diff --git a/src/ingester/src/writer.rs b/src/ingester/src/writer.rs index 7c5a42221a..6abf2ec688 100644 --- a/src/ingester/src/writer.rs +++ b/src/ingester/src/writer.rs @@ -243,7 +243,7 @@ impl Writer { wal.write(&entry_bytes, false).context(WalSnafu)?; drop(wal); // write into memtable - let mem = self.memtable.read().await; + let mem = self.memtable.write().await; mem.write(schema, entry).await?; drop(mem); } diff --git a/src/job/files/parquet.rs b/src/job/files/parquet.rs index 3e5e2a187c..9386e26087 100644 --- a/src/job/files/parquet.rs +++ b/src/job/files/parquet.rs @@ -53,7 +53,7 @@ use crate::{ common::{infra::wal, meta::stream::SchemaRecords}, job::files::idx::write_to_disk, service::{ - compact::merge::{generate_inverted_idx_recordbatch, merge_parquet_files}, + compact::merge::{generate_inverted_idx_recordbatch, merge_parquet_files_v1}, db, schema::generate_schema_for_defined_schema_fields, search::datafusion::{ @@ -593,7 +593,7 @@ async fn merge_files( ) .await } else if stream_type == StreamType::Logs { - merge_parquet_files(thread_id, tmp_dir.name(), schema.clone()).await + merge_parquet_files_v1(thread_id, tmp_dir.name(), schema.clone()).await } else { merge_parquet_files_by_datafusion(tmp_dir.name(), stream_type, &stream_name, schema.clone()) .await diff --git a/src/service/compact/merge.rs b/src/service/compact/merge.rs index 5f0aed85bc..2a88884fec 100644 --- a/src/service/compact/merge.rs +++ b/src/service/compact/merge.rs @@ -685,7 +685,7 @@ pub async fn merge_files( let start = std::time::Instant::now(); let merge_result = if stream_type == StreamType::Logs { - merge_parquet_files(thread_id, tmp_dir.name(), schema_latest.clone()).await + merge_parquet_files_v2(thread_id, tmp_dir.name(), schema_latest.clone()).await } else { datafusion::exec::merge_parquet_files( tmp_dir.name(), @@ -1028,7 +1028,92 @@ pub fn generate_inverted_idx_recordbatch( inverted_idx_batches } -pub async fn merge_parquet_files( +// v1 use more memory but faster +pub async fn merge_parquet_files_v1( + thread_id: usize, + trace_id: &str, + schema: Arc, +) -> ::datafusion::error::Result<(Arc, Vec)> { + let start = std::time::Instant::now(); + + // 1. get record batches from tmpfs + let temp_files = infra::cache::tmpfs::list(trace_id, "parquet").map_err(|e| { + log::error!( + "[INGESTER:JOB:{thread_id}] merge small files failed at getting temp files. Error {}", + e + ); + DataFusionError::Execution(e.to_string()) + })?; + + let mut record_batches = Vec::new(); + for file in temp_files { + let bytes = infra::cache::tmpfs::get(&file.location).map_err(|e| { + log::error!( + "[INGESTER:JOB:{thread_id}] merge small files failed at reading temp files to bytes. Error {}", + e + ); + DataFusionError::Execution(e.to_string()) + })?; + + let (_, batches) = read_recordbatch_from_bytes(&bytes).await.map_err(|e| { + log::error!("[INGESTER:JOB:{thread_id}] read_recordbatch_from_bytes error",); + log::error!( + "[INGESTER:JOB:{thread_id}] read_recordbatch_from_bytes error for file: {}, err: {}", + file.location, + e + ); + DataFusionError::Execution(e.to_string()) + })?; + record_batches.extend(batches); + } + + // format recordbatch + let record_batches = record_batches + .into_iter() + .map(|b| format_recordbatch_by_schema(schema.clone(), b)) + .collect::>(); + + // 2. concatenate all record batches into one single RecordBatch + let concatenated_record_batch = arrow::compute::concat_batches(&schema, &record_batches)?; + + // 3. sort concatenated record batch by timestamp col in desc order + let sort_indices = arrow::compute::sort_to_indices( + concatenated_record_batch + .column_by_name(&get_config().common.column_timestamp) + .ok_or_else(|| { + log::error!( + "[INGESTER:JOB:{thread_id}] merge small files failed to find _timestamp column from merged record batch.", + ); + DataFusionError::Execution( + "No _timestamp column found in merged record batch".to_string(), + ) + })?, + Some(arrow_schema::SortOptions { + descending: false, + nulls_first: true, + }), + None, + )?; + + let sorted_columns = concatenated_record_batch + .columns() + .iter() + .map(|c| arrow::compute::take(c, &sort_indices, None)) + .collect::, _>>()?; + + let final_record_batch = RecordBatch::try_new(schema.clone(), sorted_columns)?; + let schema = final_record_batch.schema(); + + log::info!( + "merge_parquet_files took {} ms", + start.elapsed().as_millis() + ); + + Ok((schema, vec![final_record_batch])) +} + +// v2 use less memory but slower +pub async fn merge_parquet_files_v2( thread_id: usize, trace_id: &str, mut schema: Arc, @@ -1073,7 +1158,7 @@ pub async fn merge_parquet_files( .collect::>(); // 2. concatenate all record batches into one single RecordBatch - let mut concated_record_batch = concat_batches(schema.clone(), record_batches)?; + let mut concated_record_batch = concat_batches(schema.clone(), record_batches, true)?; // 3. delete all the null columns let num_rows = concated_record_batch.num_rows(); @@ -1109,11 +1194,6 @@ pub async fn merge_parquet_files( None, )?; - // let sorted_columns = concated_record_batch - // .columns() - // .into_iter() - // .map(|c| arrow::compute::take(c, &sort_indices, None)) - // .collect::, _>>()?; let batch_columns_len = concated_record_batch.columns().len(); let mut sorted_columns = Vec::with_capacity(batch_columns_len); for i in 0..batch_columns_len { diff --git a/src/service/search/mod.rs b/src/service/search/mod.rs index 420ba3d477..8b5037311e 100644 --- a/src/service/search/mod.rs +++ b/src/service/search/mod.rs @@ -20,7 +20,8 @@ use config::{ get_config, ider, meta::{ search, - stream::{FileKey, StreamType}, usage::{RequestStats, UsageType}, + stream::{FileKey, StreamType}, + usage::{RequestStats, UsageType}, }, utils::str::find, }; @@ -33,8 +34,6 @@ use opentelemetry::trace::TraceContextExt; use proto::cluster_rpc; use regex::Regex; use tracing_opentelemetry::OpenTelemetrySpanExt; -use super::usage::report_request_usage_stats; - #[cfg(feature = "enterprise")] use { hashbrown::HashSet, @@ -45,6 +44,7 @@ use { #[cfg(not(feature = "enterprise"))] use {std::sync::Arc, tokio::sync::Mutex}; +use super::usage::report_request_usage_stats; use crate::{ common::{infra::cluster as infra_cluster, meta::stream::StreamParams}, handler::grpc::request::search::Searcher, @@ -158,7 +158,7 @@ pub async fn search( search::SearchEventType::Other => (false, None), }, None => (false, None), - }; + }; if report_usage { let stream_name = match config::meta::sql::Sql::new(&req_query.sql) { @@ -190,7 +190,8 @@ pub async fn search( ) .await; } - Ok(res)}, + Ok(res) + } Err(e) => Err(e), } }