From c34bec197caccdd89593f1e89c033556da8639fd Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Thu, 27 Aug 2026 12:01:27 +0800 Subject: [PATCH] fix(query): share record batch forwarding for DML Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/query/src/datafusion.rs | 142 ++++++++++++++++++++++++------------ 1 file changed, 95 insertions(+), 47 deletions(-) diff --git a/src/query/src/datafusion.rs b/src/query/src/datafusion.rs index 1eb91a6d0f..e6e879a584 100644 --- a/src/query/src/datafusion.rs +++ b/src/query/src/datafusion.rs @@ -30,7 +30,7 @@ use common_function::function::FunctionContext; use common_function::function_factory::ScalarFunctionFactory; use common_query::{Output, OutputData, OutputMeta}; use common_recordbatch::adapter::{RecordBatchStreamAdapter, RegionQueryStatCounters}; -use common_recordbatch::{EmptyRecordBatchStream, SendableRecordBatchStream}; +use common_recordbatch::{EmptyRecordBatchStream, RecordBatch, SendableRecordBatchStream}; use common_telemetry::tracing; use datafusion::catalog::TableFunction; use datafusion::dataframe::DataFrame; @@ -81,6 +81,29 @@ pub const QUERY_PARALLELISM_HINT: &str = "query_parallelism"; /// Whether to fallback to the original plan when failed to push down. pub const QUERY_FALLBACK_HINT: &str = "query_fallback"; +// An unbounded queue keeps draining source RPCs while mutation RPCs on a shared +// HTTP/2 connection are pending, trading bounded memory for request liveness. +async fn forward_record_batches( + mut stream: SendableRecordBatchStream, + batch_tx: tokio::sync::mpsc::UnboundedSender>, +) -> Result<()> { + while let Some(batch) = stream.next().await { + match batch.context(CreateRecordBatchSnafu) { + Ok(batch) => { + if batch_tx.send(Ok(batch)).is_err() { + break; + } + tokio::task::yield_now().await; + } + Err(error) => { + let _ = batch_tx.send(Err(error)); + break; + } + } + } + Ok(()) +} + fn query_load_region_id(plan: &Arc) -> Option { let mut region_id = None; let mut stack = vec![plan.clone()]; @@ -206,7 +229,7 @@ impl DatafusionQueryEngine { let Output { data, meta } = self .exec_query_plan((*dml.input).clone(), query_ctx.clone()) .await?; - let mut stream = match data { + let stream = match data { OutputData::RecordBatches(batches) => batches.as_stream(), OutputData::Stream(stream) => stream, _ => unreachable!(), @@ -217,53 +240,31 @@ impl DatafusionQueryEngine { match dml.op { WriteOp::Insert(_) => { - // An unbounded queue keeps draining source RPCs while sink RPCs on a shared - // HTTP/2 connection are pending, trading bounded memory for request liveness. - let (batch_tx, mut batch_rx) = tokio::sync::mpsc::unbounded_channel(); - let producer = async move { - while let Some(batch) = stream.next().await { - let batch = batch.context(CreateRecordBatchSnafu); - let is_err = batch.is_err(); - if batch_tx.send(batch).is_err() { - break; - } - if is_err { - break; - } - tokio::task::yield_now().await; - } - Ok::<_, crate::error::Error>(()) - }; - let consumer = async { - while let Some(batch) = batch_rx.recv().await { - let batch = batch?; - let column_vectors = batch - .column_vectors(&table_name.to_string(), table.schema()) - .map_err(BoxedError::new) - .context(QueryExecutionSnafu)?; - // We ignore the insert op. - let output = self - .insert(&table_name, column_vectors, query_ctx.clone()) - .await?; - let (rows, cost) = output.extract_rows_and_cost(); - affected_rows += rows; - insert_cost += cost; - } - Ok::<_, crate::error::Error>(()) - }; - try_join(producer, consumer).await?; + let (batch_tx, batch_rx) = tokio::sync::mpsc::unbounded_channel(); + let producer = forward_record_batches(stream, batch_tx); + let consumer = self.consume_insert_record_batches( + batch_rx, + &table_name, + table.schema(), + query_ctx.clone(), + ); + let ((), (rows, cost)) = try_join(producer, consumer).await?; + affected_rows += rows; + insert_cost += cost; } WriteOp::Delete => { - while let Some(batch) = stream.next().await { - let batch = batch.context(CreateRecordBatchSnafu)?; - let column_vectors = batch - .column_vectors(&table_name.to_string(), table.schema()) - .map_err(BoxedError::new) - .context(QueryExecutionSnafu)?; - affected_rows += self - .delete(&table_name, &table, column_vectors, query_ctx.clone()) - .await?; - } + // Keep DELETE on the same producer/consumer schedule as INSERT so the source + // stream can continue draining while mutation RPCs are pending. + let (batch_tx, batch_rx) = tokio::sync::mpsc::unbounded_channel(); + let producer = forward_record_batches(stream, batch_tx); + let consumer = self.consume_delete_record_batches( + batch_rx, + &table_name, + &table, + query_ctx.clone(), + ); + let (rows, ()) = try_join(consumer, producer).await?; + affected_rows += rows; } _ => unreachable!("guarded by the 'ensure!' at the beginning"), } @@ -273,6 +274,53 @@ impl DatafusionQueryEngine { )) } + async fn consume_insert_record_batches( + &self, + mut batch_rx: tokio::sync::mpsc::UnboundedReceiver>, + table_name: &ResolvedTableReference, + table_schema: Arc, + query_ctx: QueryContextRef, + ) -> Result<(usize, usize)> { + let mut affected_rows = 0; + let mut insert_cost = 0; + while let Some(batch) = batch_rx.recv().await { + let batch = batch?; + let column_vectors = batch + .column_vectors(&table_name.to_string(), table_schema.clone()) + .map_err(BoxedError::new) + .context(QueryExecutionSnafu)?; + // We ignore the insert op. + let output = self + .insert(table_name, column_vectors, query_ctx.clone()) + .await?; + let (rows, cost) = output.extract_rows_and_cost(); + affected_rows += rows; + insert_cost += cost; + } + Ok((affected_rows, insert_cost)) + } + + async fn consume_delete_record_batches( + &self, + mut batch_rx: tokio::sync::mpsc::UnboundedReceiver>, + table_name: &ResolvedTableReference, + table: &TableRef, + query_ctx: QueryContextRef, + ) -> Result { + let mut affected_rows = 0; + while let Some(batch) = batch_rx.recv().await { + let batch = batch?; + let column_vectors = batch + .column_vectors(&table_name.to_string(), table.schema()) + .map_err(BoxedError::new) + .context(QueryExecutionSnafu)?; + affected_rows += self + .delete(table_name, table, column_vectors, query_ctx.clone()) + .await?; + } + Ok(affected_rows) + } + #[tracing::instrument(skip_all)] async fn delete( &self,