diff --git a/src/service/search/grpc/wal.rs b/src/service/search/grpc/wal.rs index 21fe6c7021..7efe5acfad 100644 --- a/src/service/search/grpc/wal.rs +++ b/src/service/search/grpc/wal.rs @@ -449,16 +449,19 @@ pub async fn search_memtable( let mut tables = Vec::new(); let start = std::time::Instant::now(); - for (schema, mut record_batches) in batch_groups { + for (schema, record_batches) in batch_groups { if record_batches.is_empty() { continue; } let diff_fields = generate_search_schema_diff(&schema, &latest_schema_map); - - for batch in record_batches.iter_mut() { - *batch = adapt_batch(latest_schema.clone(), batch); + let mut adapt_batches = Vec::with_capacity(record_batches.len()); + for batch in record_batches { + adapt_batches.push(adapt_batch(latest_schema.clone(), batch)); } + let record_batches = adapt_batches; + + tokio::task::coop::consume_budget().await; // merge small batches into big batches let mut merge_groupes = Vec::new(); @@ -482,6 +485,8 @@ pub async fn search_memtable( .map(|group| concat_batches(group[0].schema().clone(), group).unwrap()) .collect::>(); + tokio::task::coop::consume_budget().await; + let table = match NewMemTable::try_new( record_batches[0].schema().clone(), vec![record_batches], @@ -638,7 +643,7 @@ async fn get_file_list( .await } -pub fn adapt_batch(latest_schema: Arc, batch: &RecordBatch) -> RecordBatch { +pub fn adapt_batch(latest_schema: Arc, batch: RecordBatch) -> RecordBatch { let batch_schema = &*batch.schema(); let batch_cols = batch.columns().to_vec(); diff --git a/src/service/search/sql.rs b/src/service/search/sql.rs index 96c5538c5f..48594e1277 100644 --- a/src/service/search/sql.rs +++ b/src/service/search/sql.rs @@ -215,13 +215,7 @@ impl Sql { ); } else { for (stream, schema) in total_schemas.iter() { - let columns = match columns.get(stream) { - Some(columns) => columns.clone(), - None => { - used_schemas.insert(stream.clone(), schema.clone()); - continue; - } - }; + let columns = columns.get(stream).cloned().unwrap_or(Default::default()); let fields = generate_schema_fields(columns, schema, match_visitor.match_items.is_some()); let schema = Schema::new(fields).with_metadata(schema.schema().metadata().clone());