starred/openobserve
1
0
Fork 0
mirror of https://github.com/openobserve/openobserve.git synced 2026-11-05 08:09:22 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: reduce CPU const for merging (#3673)

Browse Source
This commit is contained in:
Hengfei Yang
2024-06-04 21:37:52 +08:00
committed by GitHub
parent 99eceb4db0
commit 76efeb5e2a
6 changed files with 120 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/config/src/utils/record_batch_ext.rs
+18 -13
View File
@@ -484,6 +484,7 @@ pub fn format_recordbatch_by_schema(schema: Arc<Schema>, batch: RecordBatch) ->
pub fn concat_batches(
schema: Arc<Schema>,
mut batches: Vec<RecordBatch>,
less_memory_mode: bool,
) -> Result<RecordBatch, ArrowError> {
// 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::<Vec<_>>(),
// )?;
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::<Vec<_>>())?;
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::<Vec<_>>())?
} else {
arrow::compute::concat(
&batches
.iter()
.map(|batch| batch.column(i).as_ref())
.collect::<Vec<_>>(),
)?
};
arrays.push(array);
}
RecordBatch::try_new(schema.clone(), arrays)
src/infra/src/db/postgres.rs
+5
View File
@@ -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 {
src/ingester/src/writer.rs
+1 -1
View File
@@ -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);
}
src/job/files/parquet.rs
+2 -2
View File
@@ -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
src/service/compact/merge.rs
+88 -8
View File
@@ -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<Schema>,
) -> ::datafusion::error::Result<(Arc<Schema>, Vec<RecordBatch>)> {
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::<Vec<_>>();
// 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::<std::result::Result<Vec<_>, _>>()?;
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<Schema>,
@@ -1073,7 +1158,7 @@ pub async fn merge_parquet_files(
.collect::<Vec<_>>();
// 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::<std::result::Result<Vec<_>, _>>()?;
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 {
src/service/search/mod.rs
+6 -5
View File
@@ -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),
}
}